wehub-resource-sync
252 lines
10 KiB
C++
252 lines
10 KiB
C++
/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
|
|
|
|
Licensed under the Apache License, Version 2.0 (the "License");
|
|
you may not use this file except in compliance with the License.
|
|
You may obtain a copy of the License at
|
|
|
|
http://www.apache.org/licenses/LICENSE-2.0
|
|
|
|
Unless required by applicable law or agreed to in writing, software
|
|
distributed under the License is distributed on an "AS IS" BASIS,
|
|
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
See the License for the specific language governing permissions and
|
|
limitations under the License.
|
|
==============================================================================*/
|
|
|
|
#include <memory>
|
|
#include <string>
|
|
#include <vector>
|
|
|
|
#include "absl/strings/string_view.h"
|
|
#include "llvm/Support/CommandLine.h"
|
|
#include "llvm/Support/FormatVariadic.h"
|
|
#include "llvm/Support/MemoryBuffer.h"
|
|
#include "llvm/Support/SourceMgr.h"
|
|
#include "llvm/Support/ToolOutputFile.h"
|
|
#include "llvm/Support/raw_ostream.h"
|
|
#include "mlir/Dialect/Arith/IR/Arith.h" // from @llvm-project
|
|
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
|
|
#include "mlir/Dialect/Quant/IR/Quant.h" // from @llvm-project
|
|
#include "mlir/Dialect/Quant/IR/QuantTypes.h" // from @llvm-project
|
|
#include "mlir/IR/Attributes.h" // from @llvm-project
|
|
#include "mlir/IR/Builders.h" // from @llvm-project
|
|
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
|
|
#include "mlir/IR/BuiltinTypes.h" // from @llvm-project
|
|
#include "mlir/IR/Location.h" // from @llvm-project
|
|
#include "mlir/IR/MLIRContext.h" // from @llvm-project
|
|
#include "mlir/IR/Operation.h" // from @llvm-project
|
|
#include "mlir/IR/Types.h" // from @llvm-project
|
|
#include "mlir/IR/Value.h" // from @llvm-project
|
|
#include "mlir/Support/LogicalResult.h" // from @llvm-project
|
|
#include "mlir/Tools/mlir-translate/Translation.h" // from @llvm-project
|
|
#include "stablehlo/dialect/StablehloOps.h" // from @stablehlo
|
|
#include "stablehlo/dialect/VhloOps.h" // from @stablehlo
|
|
#include "tensorflow/compiler/mlir/lite/flatbuffer_export.h"
|
|
#include "tensorflow/compiler/mlir/lite/flatbuffer_import.h"
|
|
#include "tensorflow/compiler/mlir/lite/ir/tfl_ops.h"
|
|
#include "tensorflow/compiler/mlir/lite/quantization/ir/QuantOps.h"
|
|
#include "tensorflow/compiler/mlir/tensorflow/dialect_registration.h"
|
|
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
|
|
#include "tensorflow/compiler/mlir/tensorflow/translate/mlir_roundtrip_flags.h"
|
|
#include "tensorflow/compiler/mlir/tensorflow/translate/tools/parsers.h"
|
|
|
|
using llvm::cl::opt;
|
|
|
|
// Commandline flag to enable the control of flatbuffer import.
|
|
bool use_external_constant;
|
|
|
|
// Commandline flag to enable graph pruning.
|
|
bool experimental_prune_unreachable_nodes_unconditionally;
|
|
|
|
// NOLINTNEXTLINE
|
|
static opt<bool, true> use_external_constant_flag(
|
|
"use-external-constant",
|
|
llvm::cl::desc("Use external constant during flatbuffer import"),
|
|
llvm::cl::location(use_external_constant), llvm::cl::init(false));
|
|
|
|
// TODO(b/147111261): After the importer supports generic custom ops, we should
|
|
// change the flag to a more lightwise flag, e.g.
|
|
// "import_custom_ops_as_side_effect_free_ops", and let the MLIR DCE to prune
|
|
// the operations.
|
|
// NOLINTNEXTLINE
|
|
static opt<bool, true> experimental_prune_unreachable_nodes_unconditionally_flg(
|
|
"experimental-prune-unreachable-nodes-unconditionally",
|
|
llvm::cl::desc("Prune nodes that are not ancestors of the output nodes."),
|
|
llvm::cl::location(experimental_prune_unreachable_nodes_unconditionally),
|
|
llvm::cl::init(false));
|
|
|
|
// NOLINTNEXTLINE
|
|
static opt<std::string> input_arrays_flag(
|
|
"input-arrays",
|
|
llvm::cl::desc(
|
|
"List of input tensors, if different from the default inputs"),
|
|
llvm::cl::init(""));
|
|
|
|
// NOLINTNEXTLINE
|
|
static opt<std::string> output_arrays_flag(
|
|
"output-arrays",
|
|
llvm::cl::desc(
|
|
"List of output tensors, if different from the default outputs"),
|
|
llvm::cl::init(""));
|
|
|
|
using llvm::cl::opt;
|
|
|
|
// These command line flags enable control of the translation implementation.
|
|
bool emit_builtin_tflite_ops;
|
|
bool emit_custom_ops;
|
|
bool emit_select_tf_ops;
|
|
bool lower_tensor_list_ops;
|
|
bool strip_debug_info;
|
|
bool use_buffer_offset;
|
|
bool emit_stablehlo_ops;
|
|
bool disable_vhlo_to_stablehlo;
|
|
bool serialize_debug_metadata;
|
|
|
|
// NOLINTNEXTLINE
|
|
static opt<bool, true> emit_builtin_tflite_ops_flag(
|
|
"emit-builtin-tflite-ops",
|
|
llvm::cl::desc(
|
|
"Emit TFLite built in operations in the generated TFLite model"),
|
|
llvm::cl::location(emit_builtin_tflite_ops), llvm::cl::init(true));
|
|
|
|
// NOLINTNEXTLINE
|
|
static opt<bool, true> emit_select_tf_ops_flag(
|
|
"emit-select-tf-ops",
|
|
llvm::cl::desc(
|
|
"Emit Select TF operations (Flex ops) in the generated TFLite model"),
|
|
llvm::cl::location(emit_select_tf_ops), llvm::cl::init(false));
|
|
|
|
// NOLINTNEXTLINE
|
|
static opt<bool, true> emit_custom_ops_flag(
|
|
"emit-custom-ops",
|
|
llvm::cl::desc("Emit Custom operations in the generated TFLite model"),
|
|
llvm::cl::location(emit_custom_ops), llvm::cl::init(false));
|
|
|
|
// NOLINTNEXTLINE
|
|
static opt<bool, true> lower_tensor_list_ops_flag(
|
|
"lower-tensor-list-ops",
|
|
llvm::cl::desc("Lower the TensorList ops within the TFLite dialect"),
|
|
llvm::cl::location(lower_tensor_list_ops), llvm::cl::init(false));
|
|
|
|
// NOLINTNEXTLINE
|
|
static opt<bool, true> strip_debug_info_flag(
|
|
"strip-debug-info", llvm::cl::desc("Strip debug info during export"),
|
|
llvm::cl::location(strip_debug_info), llvm::cl::init(false));
|
|
|
|
// NOLINTNEXTLINE
|
|
static opt<bool, true> use_buffer_offset_flag(
|
|
"use-buffer-offset",
|
|
llvm::cl::desc("store constant buffers outside of Flatbuffers"),
|
|
llvm::cl::location(use_buffer_offset), llvm::cl::init(false));
|
|
|
|
// NOLINTNEXTLINE
|
|
static opt<bool, true> emit_stablehlo_ops_flag(
|
|
"emit-stablehlo-ops",
|
|
llvm::cl::desc("Whether serialize stablehlo ops or not"),
|
|
llvm::cl::location(emit_stablehlo_ops), llvm::cl::init(false));
|
|
|
|
// Flatbuffer import by default will also perform vhlo to stablehlo legalization
|
|
// to hide serialization detail from user, but for debug purpose we need to be
|
|
// able to dump raw vhlo ops as well
|
|
// NOLINTNEXTLINE
|
|
static opt<bool, true> disable_vhlo_to_stablehlo_flag(
|
|
"disable-vhlo-to-stablehlo",
|
|
llvm::cl::desc("Whether to deserialize to stablehlo ops or not"),
|
|
llvm::cl::location(disable_vhlo_to_stablehlo), llvm::cl::init(false));
|
|
|
|
// NOLINTNEXTLINE
|
|
static opt<bool, true> serialize_debug_metadata_flag(
|
|
"serialize-debug-metadata",
|
|
llvm::cl::desc("Whether to serialize debug metadata or not"),
|
|
llvm::cl::location(serialize_debug_metadata), llvm::cl::init(false));
|
|
|
|
// NOLINTNEXTLINE
|
|
static opt<bool> disable_buffer_deduping_flag(
|
|
"disable-buffer-deduping",
|
|
llvm::cl::desc("Whether to disable buffer deduping or not"),
|
|
llvm::cl::init(false));
|
|
|
|
namespace mlir {
|
|
namespace {
|
|
static OwningOpRef<mlir::ModuleOp> FlatBufferFileToMlirTrans(
|
|
llvm::SourceMgr* source_mgr, MLIRContext* context,
|
|
bool use_external_constant,
|
|
bool experimental_prune_unreachable_nodes_unconditionally) {
|
|
const llvm::MemoryBuffer* input =
|
|
source_mgr->getMemoryBuffer(source_mgr->getMainFileID());
|
|
std::string error;
|
|
auto loc =
|
|
mlir::FileLineColLoc::get(context, input->getBufferIdentifier(), 0, 0);
|
|
|
|
// Parses input/output names from command line options.
|
|
std::vector<std::string> inputs;
|
|
std::vector<std::string> outputs;
|
|
// Use output parser since we only have tensor names.
|
|
if (!tensorflow::ParseOutputArrayInfo(input_arrays_flag, &inputs).ok()) {
|
|
return emitError(loc, "parsing input array info failed ")
|
|
<< input_arrays_flag,
|
|
nullptr;
|
|
}
|
|
if (!tensorflow::ParseOutputArrayInfo(output_arrays_flag, &outputs).ok()) {
|
|
return emitError(loc, "parsing output array info failed ")
|
|
<< output_arrays_flag,
|
|
nullptr;
|
|
}
|
|
return tflite::FlatBufferToMlir(
|
|
absl::string_view(input->getBufferStart(), input->getBufferSize()),
|
|
context, loc, use_external_constant, inputs, outputs,
|
|
experimental_prune_unreachable_nodes_unconditionally,
|
|
disable_vhlo_to_stablehlo);
|
|
}
|
|
|
|
static LogicalResult MlirToFlatBufferFileTranslateFunction(
|
|
ModuleOp module, llvm::raw_ostream& output) {
|
|
std::string serialized_flatbuffer;
|
|
std::unique_ptr<tensorflow::OpOrArgNameMapper> op_or_arg_name_mapper;
|
|
if (strip_debug_info) {
|
|
op_or_arg_name_mapper =
|
|
std::make_unique<tensorflow::OpOrArgStripNameMapper>();
|
|
} else {
|
|
op_or_arg_name_mapper =
|
|
std::make_unique<tensorflow::OpOrArgLocNameMapper>();
|
|
}
|
|
tflite::FlatbufferExportOptions options;
|
|
options.converter_flags.set_force_select_tf_ops(!emit_builtin_tflite_ops);
|
|
options.converter_flags.set_enable_select_tf_ops(emit_select_tf_ops);
|
|
options.converter_flags.set_allow_custom_ops(emit_custom_ops);
|
|
options.converter_flags.set_use_buffer_offset(use_buffer_offset);
|
|
options.op_or_arg_name_mapper = op_or_arg_name_mapper.get();
|
|
options.converter_flags.set_serialize_debug_metadata(
|
|
serialize_debug_metadata);
|
|
options.disable_buffer_deduping = disable_buffer_deduping_flag.getValue();
|
|
if (!tflite::MlirToFlatBufferTranslateFunction(
|
|
module, options, &serialized_flatbuffer, emit_stablehlo_ops))
|
|
return mlir::failure();
|
|
|
|
output << serialized_flatbuffer;
|
|
return success();
|
|
}
|
|
} // namespace
|
|
|
|
static TranslateToMLIRRegistration FlatBufferFileToMlirTransReg(
|
|
"tflite-flatbuffer-to-mlir", "tflite-flatbuffer-to-mlir",
|
|
[](llvm::SourceMgr& source_mgr, MLIRContext* context) {
|
|
return FlatBufferFileToMlirTrans(
|
|
&source_mgr, context, use_external_constant,
|
|
experimental_prune_unreachable_nodes_unconditionally);
|
|
});
|
|
|
|
static TranslateFromMLIRRegistration MLIRToFlatBufferTranslate(
|
|
"mlir-to-tflite-flatbuffer", "mlir-to-tflite-flatbuffer",
|
|
MlirToFlatBufferFileTranslateFunction, [](DialectRegistry& registry) {
|
|
registry.insert<quant::QuantDialect,
|
|
quantfork::QuantizationForkDialect>();
|
|
mlir::RegisterAllTensorFlowDialects(registry);
|
|
registry.insert<TFL::TensorFlowLiteDialect>();
|
|
registry.insert<arith::ArithDialect>();
|
|
registry.insert<func::FuncDialect>();
|
|
registry.insert<mlir::vhlo::VhloDialect>();
|
|
registry.insert<mlir::stablehlo::StablehloDialect>();
|
|
});
|
|
} // namespace mlir
|