compiler/src/iree/compiler/Utils/FlatbufferUtils.cpp

// Copyright 2020 The IREE Authors
//
// Licensed under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception

#include "iree/compiler/Utils/FlatbufferUtils.h"

#include <assert.h>
#include <stdlib.h>

#include <cstdint>
#include <type_traits>

#include "mlir/IR/BuiltinTypes.h"

namespace mlir {
namespace iree_compiler {

// Combines all pages of the FlatBuffer builder into a single contiguous byte
// buffer and returns the result.
//
// NOTE: this is a alloc/copy. We need to have a single contiguous buffer to
// pass into the elements factory function and the data we have in the
// builder is paged. If we end up with a custom attribute type for this that
// does not support storage uniquing then we can directly allocate and copy
// the pages into the buffer without the extra copy.
static SmallVector<uint8_t, 32> cloneBufferIntoContiguousBytes(
    FlatbufferBuilder &fbb) {
  size_t packedSize = flatcc_builder_get_buffer_size(fbb);
  SmallVector<uint8_t, 32> packedData(packedSize);
  void *result =
      flatcc_builder_copy_buffer(fbb, packedData.data(), packedData.size());
  assert(result && "flatcc_emitter_t impl failed (non-default?)");
  (void)result;
  return packedData;
}

FlatbufferBuilder::FlatbufferBuilder() { flatcc_builder_init(&builder); }

FlatbufferBuilder::~FlatbufferBuilder() { flatcc_builder_clear(&builder); }

flatbuffers_uint8_vec_ref_t FlatbufferBuilder::streamUint8Vec(
    std::function<bool(raw_ostream &stream)> fn, size_t alignment) {
  flatcc_builder_start_vector(*this, 1, alignment, FLATBUFFERS_COUNT_MAX(1));
  raw_flatbuffer_uint8_vec_ostream stream(*this);
  if (!fn(stream)) {
    return 0;
  }
  stream.flush();
  return flatbuffers_uint8_vec_end(*this);
}

DenseIntElementsAttr FlatbufferBuilder::getBufferAttr(MLIRContext *context) {
  // We require direct access to the FlatBuffer bytes so we can pass them to
  // the attribute constructor (which needs to inspect them all for uniquing).
  auto bufferData = cloneBufferIntoContiguousBytes(*this);

  // NOTE: ew. OpaqueAttr may be better? It does equality checks but won't try
  // to unique and would let us get a mutable buffer out.
  return DenseIntElementsAttr::get(
      VectorType::get({static_cast<int64_t>(bufferData.size())},
                      IntegerType::get(context, 8)),
      std::move(bufferData));
}

LogicalResult FlatbufferBuilder::copyToStream(llvm::raw_ostream &output) {
  // NOTE: expected to be the default emitter.
  auto *E = reinterpret_cast<flatcc_emitter_t *>(
      flatcc_builder_get_emit_context(*this));

  if (!E->front) {
    return failure();
  }
  if (E->front == E->back) {
    output.write(reinterpret_cast<char *>(E->front_cursor), E->used);
    return success();
  }
  size_t len = FLATCC_EMITTER_PAGE_SIZE - E->front_left;
  output.write(reinterpret_cast<char *>(E->front_cursor), len);
  flatcc_emitter_page_t *p = E->front->next;
  while (p != E->back) {
    output.write(reinterpret_cast<char *>(p->page), FLATCC_EMITTER_PAGE_SIZE);
    p = p->next;
  }
  output.write(reinterpret_cast<char *>(p->page),
               FLATCC_EMITTER_PAGE_SIZE - E->back_left);
  return success();
}

LogicalResult FlatbufferBuilder::printJsonToStream(bool pretty,
                                                   bool includeDefaults,
                                                   print_json_fn_t printJsonFn,
                                                   llvm::raw_ostream &output) {
  // The printer requires direct access to the FlatBuffer bytes so clone here.
  auto bufferData = cloneBufferIntoContiguousBytes(*this);
  auto moduleData = ArrayRef<uint8_t>(bufferData.data(), bufferData.size())
                        .drop_front(sizeof(flatbuffers_uoffset_t));

  flatcc_json_printer_t printer;
  flatcc_json_printer_init_dynamic_buffer(&printer, /*buffer_size=*/0);
  flatcc_json_printer_set_indent(&printer, pretty ? 2 : 0);
  flatcc_json_printer_set_skip_default(&printer, !includeDefaults);
  flatcc_json_printer_set_force_default(&printer, includeDefaults);

  // Print into the dynamically-resizing buffer. May fail if OOM.
  int rv =
      printJsonFn(&printer, reinterpret_cast<const char *>(moduleData.data()),
                  moduleData.size());
  if (rv == -1) {
    flatcc_json_printer_clear(&printer);
    return failure();
  }

  // Take the buffer from the printer; note that it is 0 terminated and can be
  // used directly as a cstr if needed.
  size_t outputSize = 0;
  char *outputBytes = reinterpret_cast<char *>(
      flatcc_json_printer_finalize_dynamic_buffer(&printer, &outputSize));
  output.write(outputBytes, outputSize);
  free(outputBytes);

  return success();
}

}  // namespace iree_compiler
}  // namespace mlir