chore: import upstream snapshot with attribution

tensorflow/compiler/mlir/tfr/integration/tfr_decompose_ctx.cc

@@ -0,0 +1,247 @@
+/* Copyright 2020 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 "tensorflow/compiler/mlir/tfr/integration/tfr_decompose_ctx.h"
+
+#include <memory>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "absl/status/status.h"
+#include "absl/status/statusor.h"
+#include "absl/strings/str_cat.h"
+#include "absl/strings/string_view.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/ADT/StringRef.h"
+#include "llvm/Support/LogicalResult.h"
+#include "llvm/Support/MemoryBuffer.h"
+#include "llvm/Support/SMLoc.h"
+#include "llvm/Support/SourceMgr.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"  // from @llvm-project
+#include "mlir/Dialect/Func/Extensions/AllExtensions.h"  // from @llvm-project
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/Dialect/SCF/IR/SCF.h"  // from @llvm-project
+#include "mlir/Dialect/Shape/IR/Shape.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/OperationSupport.h"  // from @llvm-project
+#include "mlir/IR/OwningOpRef.h"  // from @llvm-project
+#include "mlir/IR/Types.h"  // from @llvm-project
+#include "mlir/IR/Verifier.h"  // from @llvm-project
+#include "mlir/Parser/Parser.h"  // from @llvm-project
+#include "mlir/Pass/PassManager.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_device.h"
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_dialect.h"
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_executor.h"
+#include "tensorflow/compiler/mlir/tensorflow/transforms/passes.h"
+#include "tensorflow/compiler/mlir/tensorflow/utils/convert_attr.h"
+#include "tensorflow/compiler/mlir/tensorflow/utils/convert_type.h"
+#include "tensorflow/compiler/mlir/tf2xla/api/v2/tf_executor_to_graph.h"
+#include "tensorflow/compiler/mlir/tfr/ir/tfr_ops.h"
+#include "tensorflow/compiler/mlir/tfr/passes/passes.h"
+#include "xla/tsl/platform/errors.h"
+#include "xla/tsl/platform/statusor.h"
+#include "tensorflow/core/framework/function.pb.h"
+#include "tensorflow/core/framework/node_def.pb.h"
+#include "tensorflow/core/framework/node_def_util.h"
+#include "tensorflow/core/framework/op.h"
+#include "tensorflow/core/framework/op_def.pb.h"
+#include "tensorflow/core/framework/types.h"
+#include "tensorflow/core/platform/env.h"
+#include "tensorflow/core/platform/errors.h"
+#include "tensorflow/core/platform/path.h"
+#include "tensorflow/core/platform/status.h"
+#include "tensorflow/core/platform/stringpiece.h"
+#include "tensorflow/core/platform/types.h"
+#include "tensorflow/core/util/env_var.h"
+
+namespace tensorflow {
+namespace tfr {
+
+const char* const kTFRLibEnv = "TF_MLIR_TFR_LIB_DIR";
+
+absl::StatusOr<std::unique_ptr<TFRDecomposeContext>> TFRDecomposeContext::Get(
+    mlir::MLIRContext* mlir_ctx) {
+  Env* env = Env::Default();
+  std::string tfr_lib_dir;
+  TF_RETURN_IF_ERROR(ReadStringFromEnvVar(
+      kTFRLibEnv, "tensorflow/compiler/mlir/tfr/resources", &tfr_lib_dir));
+  std::string composite_mlir_dir =
+      io::JoinPath(env->GetRunfilesDir(), tfr_lib_dir);
+  std::vector<std::string> files;
+  TF_RETURN_IF_ERROR(env->GetChildren(composite_mlir_dir, &files));
+  if (files.empty()) {
+    return absl::InternalError(absl::StrCat(
+        "Failed to find the decomposition lib from path ", composite_mlir_dir));
+  }
+  std::string tfr_raw_text;
+  for (const auto& file : files) {
+    std::string fullpath = io::JoinPath(composite_mlir_dir, file);
+    if (env->MatchPath(fullpath, io::JoinPath(composite_mlir_dir, "*.mlir"))) {
+      std::string text;
+      TF_RETURN_IF_ERROR(ReadFileToString(env, fullpath, &text));
+      tfr_raw_text.append(text);
+    }
+  }
+
+  auto ctx = TFRDecomposeContext::GetFromText(tfr_raw_text, mlir_ctx);
+  if (!ctx) {
+    return absl::InternalError(absl::StrCat(
+        "Failed to load the imported decomposition lib: ", tfr_raw_text));
+  }
+  return ctx;
+}
+
+std::unique_ptr<TFRDecomposeContext> TFRDecomposeContext::GetFromText(
+    absl::string_view tfr_raw_text, mlir::MLIRContext* mlir_ctx) {
+  mlir_ctx->allowUnregisteredDialects(/*allow=*/true);
+  // Load dialects involved in the conversion
+  mlir::DialectRegistry registry;
+  // clang-format off
+  registry.insert<mlir::arith::ArithDialect,
+                  mlir::func::FuncDialect,
+                  mlir::scf::SCFDialect,
+                  mlir::shape::ShapeDialect,
+                  mlir::TF::TensorFlowDialect,
+                  mlir::tf_device::TensorFlowDeviceDialect,
+                  mlir::tf_executor::TensorFlowExecutorDialect,
+                  mlir::TFR::TFRDialect>();
+  // clang-format on
+  mlir::func::registerAllExtensions(registry);
+  mlir_ctx->appendDialectRegistry(registry);
+  mlir_ctx->loadAllAvailableDialects();
+
+  // Load the TFR functions in a mlir::ModuleOp
+  auto memory_buffer = llvm::MemoryBuffer::getMemBuffer(
+      llvm::StringRef(tfr_raw_text.data(), tfr_raw_text.size()));
+  llvm::SourceMgr source_mgr;
+  source_mgr.AddNewSourceBuffer(std::move(memory_buffer), llvm::SMLoc());
+  mlir::OwningOpRef<mlir::ModuleOp> module =
+      mlir::parseSourceFile<mlir::ModuleOp>(source_mgr, mlir_ctx);
+  // The MLIRContext owns the module
+  auto module_op = module.release();
+
+  // Create the context
+  return std::make_unique<TFRDecomposeContext>(module_op);
+}
+
+absl::StatusOr<FunctionDef> TFRDecomposeContext::ExpandNode(
+    const NodeDef& node_def, absl::string_view func_name) {
+  const OpDef* op_def;
+  TF_RETURN_IF_ERROR(OpRegistry::Global()->LookUpOpDef(node_def.op(), &op_def));
+  DataTypeVector input_dtys, output_dtys;
+  TF_RETURN_IF_ERROR(InputTypesForNode(node_def, *op_def, &input_dtys));
+  TF_RETURN_IF_ERROR(OutputTypesForNode(node_def, *op_def, &output_dtys));
+
+  mlir::MLIRContext* context = tfr_module_.getContext();
+  llvm::SmallVector<mlir::Type, 4> input_tys, output_tys;
+  mlir::Builder builder(context);
+  for (auto ty : input_dtys) {
+    mlir::Type elt_ty;
+    TF_RETURN_IF_ERROR(ConvertDataType(ty, builder, &elt_ty));
+    mlir::TensorType mlir_ty = mlir::UnrankedTensorType::get(elt_ty);
+    input_tys.push_back(mlir_ty);
+  }
+  for (auto ty : output_dtys) {
+    mlir::Type elt_ty;
+    TF_RETURN_IF_ERROR(ConvertDataType(ty, builder, &elt_ty));
+    mlir::TensorType mlir_ty = mlir::UnrankedTensorType::get(elt_ty);
+    output_tys.push_back(mlir_ty);
+  }
+  llvm::SmallVector<mlir::NamedAttribute, 4> attrs;
+  for (const auto& attr : node_def.attr()) {
+    TF_ASSIGN_OR_RETURN(auto mlir_attr,
+                        ConvertAttributeValue(attr.second, &builder));
+    attrs.push_back({mlir::StringAttr::get(context, attr.first), mlir_attr});
+  }
+
+  mlir::Location loc = mlir::UnknownLoc::get(context);
+  mlir::ModuleOp module = mlir::ModuleOp::create(loc);
+  mlir::FunctionType func_type =
+      mlir::FunctionType::get(context, input_tys, output_tys);
+  llvm::StringRef func_name_str(func_name.data(), func_name.size());
+  auto func = mlir::func::FuncOp::create(loc, func_name_str, func_type, {});
+  module.push_back(func);
+  func.addEntryBlock();
+  mlir::OpBuilder op_builder(func.getBody());
+
+  // Create the TF op
+  const std::string tf_op_full_name = absl::StrCat("tf.", node_def.op());
+  mlir::OperationState op_state(loc, tf_op_full_name);
+  op_state.addOperands(func.getArguments());
+  op_state.addTypes(output_tys);
+  op_state.addAttributes(attrs);
+  mlir::Operation* tf_op = op_builder.create(op_state);
+  mlir::func::ReturnOp::create(op_builder, loc, tf_op->getResults());
+
+  // Run the decompose passes on the module
+  TF_RETURN_IF_ERROR(DecomposeGraph(module));
+
+  // Export the result as a FunctionDef.
+  FunctionDef func_def;
+  TF_RETURN_IF_ERROR(
+      tensorflow::tf2xla::v2::ConvertMlirFunctionToFunctionLibraryDef(
+          func, export_confs_, &func_def));
+  module.erase();
+  return func_def;
+}
+
+absl::Status TFRDecomposeContext::DecomposeGraph(mlir::ModuleOp user_module) {
+  // Call the decompose passes by using the external symbol table.
+  if (failed(pm_.run(user_module))) {
+    return absl::InternalError("Failed to run the decompose passes.");
+  }
+  return absl::OkStatus();
+}
+
+// Constructor of the decompose context.
+TFRDecomposeContext::TFRDecomposeContext(mlir::ModuleOp tfr_module)
+    : tfr_module_(tfr_module), pm_(tfr_module_.getContext()) {
+  mlir::OpPassManager& func_pm = pm_.nest<mlir::func::FuncOp>();
+
+  // Prepare the imported graph.
+  func_pm.addPass(mlir::CreateExecutorDialectToFunctionalConversionPass());
+
+  // Run TFR lowering, inlining and raising to tf.
+  func_pm.addPass(mlir::TFR::CreateDecomposeTFOpsPass(tfr_module_));
+  func_pm.addPass(mlir::TFR::CreateRaiseToTFOpsPass(
+      tfr_module_, /*materialize_derived_attrs=*/true));
+
+  // Prepare to be exported.
+  func_pm.addPass(mlir::CreateFunctionalToExecutorDialectConversionPass());
+  pm_.addPass(mlir::CreateBreakUpIslandsPass());
+}
+
+void TFRDecomposeContext::Destroy() { tfr_module_.erase(); }
+
+absl::StatusOr<FunctionDef> ExpandNode(const NodeDef& node_def,
+                                       absl::string_view func_name) {
+  mlir::MLIRContext mlir_ctx;
+  TF_ASSIGN_OR_RETURN(auto ctx, TFRDecomposeContext::Get(&mlir_ctx));
+  return ctx->ExpandNode(node_def, func_name);
+}
+
+absl::Status DecomposeGraph(mlir::ModuleOp user_module) {
+  mlir::MLIRContext* mlir_ctx = user_module.getContext();
+  TF_ASSIGN_OR_RETURN(auto ctx, TFRDecomposeContext::Get(mlir_ctx));
+  return ctx->DecomposeGraph(user_module);
+}
+
+}  // namespace tfr
+}  // namespace tensorflow

tensorflow/compiler/mlir/tfr/integration/tfr_decompose_ctx.h

@@ -0,0 +1,87 @@
+/* Copyright 2020 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.
+==============================================================================*/
+#ifndef TENSORFLOW_COMPILER_MLIR_TFR_INTEGRATION_TFR_DECOMPOSE_CTX_H_
+#define TENSORFLOW_COMPILER_MLIR_TFR_INTEGRATION_TFR_DECOMPOSE_CTX_H_
+
+#include "absl/status/status.h"
+#include "absl/status/statusor.h"
+#include "absl/strings/string_view.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/IR/MLIRContext.h"  // from @llvm-project
+#include "mlir/Pass/PassManager.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tensorflow/translate/mlir_roundtrip_flags.h"
+#include "xla/tsl/platform/statusor.h"
+#include "tensorflow/core/framework/function.pb.h"
+#include "tensorflow/core/platform/status.h"
+#include "tensorflow/core/platform/stringpiece.h"
+
+namespace tensorflow {
+namespace tfr {
+
+extern const char* const kTFRLibEnv;
+
+using tsl::StatusOr;
+
+// An wrapper for all the objects used to decompose a module (graph mode) and
+// node_def (eager mode). Note that this class owns the decomposition library.
+class TFRDecomposeContext {
+ public:
+  // The entry function to get a decompose context. All the required passes have
+  // been initialized.
+  static absl::StatusOr<std::unique_ptr<TFRDecomposeContext>> Get(
+      mlir::MLIRContext* mlir_ctx);
+
+  // Constructor of the decompose context. To share the decompose library, the
+  // whole decompose TFR function library is loaded.
+  explicit TFRDecomposeContext(mlir::ModuleOp tfr_module);
+
+  // Constructs the decompose context from the tfr text module and the mlir
+  // context. The tfr text module is added to the mlir context.
+  static std::unique_ptr<TFRDecomposeContext> GetFromText(
+      absl::string_view tfr_raw_text, mlir::MLIRContext* mlir_ctx);
+
+  // Decomposes the op in the NodeDef to a set of primitive ops according to the
+  // decompose library in the context. Wrap the decomposed result in a
+  // FunctionDef.
+  absl::StatusOr<FunctionDef> ExpandNode(const NodeDef& node_def,
+                                         absl::string_view func_name);
+
+  // Runs the decompose passes on the user_module.
+  absl::Status DecomposeGraph(mlir::ModuleOp user_module);
+
+  // Erases the tfr_module created.
+  void Destroy();
+
+ private:
+  mlir::ModuleOp tfr_module_;
+  mlir::PassManager pm_;
+
+  GraphExportConfig export_confs_;
+};
+
+// Decomposes the NodeDef to a set of primitive ops according to the decompose
+// library loaded. Wrap the decomposed result in a FunctionDef.
+absl::StatusOr<FunctionDef> ExpandNode(const NodeDef& node_def,
+                                       absl::string_view func_name);
+
+// Decomposes the ops in the ModuleOp to a set of primitive ops according to
+// decompose library in the context.
+absl::Status DecomposeGraph(mlir::ModuleOp user_module);
+
+}  // namespace tfr
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_COMPILER_MLIR_TFR_INTEGRATION_TFR_DECOMPOSE_CTX_H_

tensorflow/compiler/mlir/tfr/integration/tfr_decompose_ctx_test.cc

@@ -0,0 +1,159 @@
+/* Copyright 2020 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 "tensorflow/compiler/mlir/tfr/integration/tfr_decompose_ctx.h"
+
+#include <memory>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "absl/types/span.h"
+#include "mlir/IR/MLIRContext.h"  // from @llvm-project
+#include "xla/hlo/testlib/test.h"
+#include "tensorflow/core/framework/attr_value.pb.h"
+#include "tensorflow/core/framework/common_shape_fns.h"
+#include "tensorflow/core/framework/function.pb.h"
+#include "tensorflow/core/framework/graph.pb.h"
+#include "tensorflow/core/framework/node_def_builder.h"
+#include "tensorflow/core/framework/op.h"
+#include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/platform/test.h"
+
+using testing::ElementsAreArray;
+using testing::Test;
+using NodeAndType = std::pair<std::string, tensorflow::DataType>;
+
+namespace tensorflow {
+namespace {
+
+REGISTER_OP("MyAddN")
+    .Input("inputs: N * T")
+    .Output("sum: T")
+    .Attr("N: int >= 1")
+    .Attr("T: {numbertype, variant}")
+    .SetIsCommutative()
+    .SetIsAggregate()
+    .SetShapeFn(shape_inference::UnchangedShape);
+
+REGISTER_OP("RiscAddDummy")
+    .Input("x: T")
+    .Input("y: T")
+    .Output("z: T")
+    .Attr(
+        "T: {bfloat16, half, float, double, uint8, int8, int16, int32, int64, "
+        "complex64, complex128, string}")
+    .SetShapeFn(shape_inference::UnchangedShape);
+
+constexpr char tfr_raw_text[] = R"(
+
+tfr.func @tf__my_add_n(%values: !tfr.tensor_list,
+                       %n: i64 {tfr.name="N"}) -> !tfr.tensor {
+  %index = arith.constant 0 : index
+  %cst = arith.constant 1 : i64
+  %eq = arith.cmpi "eq", %n, %cst : i64
+  %v1 = tfr.get_element %values[%index] : (!tfr.tensor_list, index) -> !tfr.tensor
+  %res = scf.if %eq -> !tfr.tensor {
+    scf.yield %v1 : !tfr.tensor
+  } else {
+    %step = arith.index_cast %cst : i64 to index
+    %end = arith.index_cast %n : i64 to index
+    %reduce = scf.for %i = %step to %end step %step iter_args(%reduce_iter=%v1) -> !tfr.tensor {
+      %v = tfr.get_element %values[%i] : (!tfr.tensor_list, index) -> !tfr.tensor
+      %reduce_next =  tfr.call @tf__risc_add_dummy(%reduce_iter, %v) : (!tfr.tensor, !tfr.tensor) -> !tfr.tensor
+      scf.yield %reduce_next : !tfr.tensor
+    }
+    scf.yield %reduce : !tfr.tensor
+  }
+  tfr.return %res : !tfr.tensor
+}
+
+tfr.func @tf__my_add_n_(!tfr.tensor_list<N,T>, i64 {tfr.name="N"}) -> !tfr.tensor attributes{N,T}
+tfr.func @tf__risc_add_dummy_(!tfr.tensor<T>, !tfr.tensor<T>) -> !tfr.tensor<T> attributes{T}
+)";
+
+class TFRDecomposeContextTest : public Test {
+ protected:
+  void SetUp() override {
+    test_ctx_ = tfr::TFRDecomposeContext::GetFromText(tfr_raw_text, &ctx_);
+  }
+
+  void TearDown() override { test_ctx_->Destroy(); }
+
+  mlir::MLIRContext ctx_;
+  std::unique_ptr<tfr::TFRDecomposeContext> test_ctx_;
+};
+
+std::vector<NodeAndType> NodesSequenceOf(const FunctionDef& graph) {
+  std::vector<NodeAndType> nodes;
+  for (auto& node : graph.node_def()) {
+    nodes.push_back({node.op(), node.attr().at("T").type()});
+  }
+  return nodes;
+}
+
+TEST_F(TFRDecomposeContextTest, FLOAT_1_ins) {
+  std::vector<NodeDefBuilder::NodeOut> src_list;
+  src_list.emplace_back("input", 0, DT_FLOAT);
+  NodeDef test_node;
+  auto status = NodeDefBuilder("float_add", "MyAddN")
+                    .Input(src_list)
+                    .Finalize(&test_node);
+  EXPECT_TRUE(status.ok());
+  auto decomposed = test_ctx_->ExpandNode(test_node, "test");
+  EXPECT_TRUE(decomposed.ok());
+  std::vector<NodeAndType> expected_results{{"Identity", DT_FLOAT}};
+  EXPECT_THAT(NodesSequenceOf(decomposed.value()),
+              ElementsAreArray(expected_results));
+}
+
+TEST_F(TFRDecomposeContextTest, FLOAT_3_ins) {
+  std::vector<NodeDefBuilder::NodeOut> src_list;
+  src_list.emplace_back("in0", 0, DT_FLOAT);
+  src_list.emplace_back("in1", 0, DT_FLOAT);
+  src_list.emplace_back("in2", 0, DT_FLOAT);
+  NodeDef test_node;
+  auto status = NodeDefBuilder("float_add_3", "MyAddN")
+                    .Input(src_list)
+                    .Finalize(&test_node);
+  EXPECT_TRUE(status.ok());
+  auto decomposed = test_ctx_->ExpandNode(test_node, "test");
+  EXPECT_TRUE(decomposed.ok());
+
+  std::vector<NodeAndType> expected_results{{"RiscAddDummy", DT_FLOAT},
+                                            {"RiscAddDummy", DT_FLOAT}};
+  EXPECT_THAT(NodesSequenceOf(decomposed.value()),
+              ElementsAreArray(expected_results));
+}
+
+TEST_F(TFRDecomposeContextTest, INT32_3_ins) {
+  std::vector<NodeDefBuilder::NodeOut> src_list;
+  src_list.emplace_back("in0", 0, DT_INT32);
+  src_list.emplace_back("in1", 0, DT_INT32);
+  src_list.emplace_back("in2", 0, DT_INT32);
+  NodeDef test_node;
+  auto status =
+      NodeDefBuilder("int_add", "MyAddN").Input(src_list).Finalize(&test_node);
+  EXPECT_TRUE(status.ok());
+  auto decomposed = test_ctx_->ExpandNode(test_node, "test");
+  EXPECT_TRUE(decomposed.ok());
+
+  std::vector<NodeAndType> expected_results{{"RiscAddDummy", DT_INT32},
+                                            {"RiscAddDummy", DT_INT32}};
+  EXPECT_THAT(NodesSequenceOf(decomposed.value()),
+              ElementsAreArray(expected_results));
+}
+
+}  // namespace
+}  // namespace tensorflow