Clean up the TFR Decompose context implementation

This cl did the code refactoring and resolved the tsan error for the tests.

PiperOrigin-RevId: 337314570
Change-Id: I72454b187ed0e6d4a4cb5a8381ace73d55bbad69

Author: liufengdb

tensorflow/compiler/mlir/tfr/BUILD

@@ -232,9 +232,9 @@ tf_py_test(
     data = ["//tensorflow/compiler/mlir/tfr/resources:decomposition_lib"],
     python_version = "PY3",
     tags = [
+        "no_pip",
         "no_windows",  # TODO(b/170752141)
         "nomac",  # TODO(b/170752141)
-        "notsan",  # TODO(b/170752141)
     ],
     deps = [
         "//tensorflow/compiler/mlir/tfr/resources:composite_ops",
@@ -263,9 +263,9 @@ tf_py_test(
     data = ["//tensorflow/compiler/mlir/tfr/resources:decomposition_lib"],
     python_version = "PY3",
     tags = [
+        "no_pip",
         "no_windows",  # TODO(b/170752141)
         "nomac",  # TODO(b/170752141)
-        "notsan",  # TODO(b/170862433)
     ],
     deps = [
         "//tensorflow/compiler/mlir/tfr/resources:composite_ops",

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

@@ -14,10 +14,12 @@ limitations under the License.
 ==============================================================================*/
 #include "tensorflow/compiler/mlir/tfr/integration/graph_decompose_pass.h"
 
+#include "mlir/IR/MLIRContext.h"  // from @llvm-project
 #include "tensorflow/compiler/mlir/tfr/integration/tfr_decompose_ctx.h"
 #include "tensorflow/stream_executor/lib/statusor.h"
 
 namespace tensorflow {
+namespace tfr {
 
 bool GraphDecomposePass::IsEnabled(const ConfigProto& config_proto) const {
   const char* tfr_lib_env_val = getenv(std::string(kTFRLibEnv).c_str());
@@ -27,14 +29,18 @@ bool GraphDecomposePass::IsEnabled(const ConfigProto& config_proto) const {
 Status GraphDecomposePass::Run(const ConfigProto& config_proto,
                                mlir::ModuleOp module) {
   if (!IsEnabled(config_proto)) {
-    VLOG(1) << "Skipping Graph Decomposition Pass, decompositin library was "
-               "not found";
+    LOG_FIRST_N(INFO, 1) << "Skipping Graph Decomposition Pass, decompositin "
+                            "library was not found";
     return Status::OK();
   }
-  VLOG(1) << "Run Graph Decomposition Passes";
-  TF_ASSIGN_OR_RETURN(ctx_, TFRDecomposeContext::Get(module.getContext()));
-  TF_RETURN_IF_ERROR(ctx_->Decompose(module));
-  return ctx_->Destroy();
+
+  LOG_FIRST_N(INFO, 1) << "Run Graph Decomposition Passes";
+
+  TF_RETURN_IF_ERROR(DecomposeGraph(module));
+
+  LOG_FIRST_N(INFO, 1) << "Finish Graph Decomposition Passes";
+
+  return Status::OK();
 }
 
 namespace {
@@ -45,4 +51,5 @@ static mlir_pass_registration::MlirOptimizationPassRegistration
                                        std::make_unique<GraphDecomposePass>());
 }  // namespace
 
+}  // namespace tfr
 }  // namespace tensorflow

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

@@ -21,6 +21,7 @@ limitations under the License.
 #include "tensorflow/stream_executor/lib/statusor.h"
 
 namespace tensorflow {
+namespace tfr {
 
 // An optimization pass that decompose the composite ops in a module according
 // to the decomposition library. Currently the decomposition library is loaded
@@ -37,11 +38,9 @@ class GraphDecomposePass : public MlirOptimizationPass {
   // This should be used as a thin mapper around mlir::ModulePass::runOnModule
   // API integrated with the Tensorflow runtime.
   Status Run(const ConfigProto& config_proto, mlir::ModuleOp module) override;
-
- private:
-  std::unique_ptr<TFRDecomposeContext> ctx_;
 };
 
+}  // namespace tfr
 }  // namespace tensorflow
 
 #endif  // TENSORFLOW_COMPILER_MLIR_TFR_INTEGRATION_GRAPH_DECOMPOSE_PASS_H_

tensorflow/compiler/mlir/tfr/integration/graph_decompose_test.py

@@ -35,14 +35,6 @@
 
 class GraphDecomposeTest(test.TestCase):
 
-  def setUp(self):
-    os.environ['TF_MLIR_TFR_LIB_DIR'] = 'tensorflow/compiler/mlir/tfr/resources'
-    super(GraphDecomposeTest, self).setUp()
-
-  def tearDown(self):
-    del os.environ['TF_MLIR_TFR_LIB_DIR']
-    super(GraphDecomposeTest, self).tearDown()
-
   def testAddN(self):
     add = def_function.function(gen_composite_ops.my_add_n)
     t1 = constant_op.constant([[1.0, 2.0], [3.0, 4.0]])
@@ -86,5 +78,6 @@ def biasd_dense_elu(x, y, z):
 
 
 if __name__ == '__main__':
+  os.environ['TF_MLIR_TFR_LIB_DIR'] = 'tensorflow/compiler/mlir/tfr/resources'
   ops.enable_eager_execution()
   test.main()

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

@@ -21,13 +21,14 @@ limitations under the License.
 #include "tensorflow/stream_executor/lib/statusor.h"
 
 namespace tensorflow {
+namespace tfr {
 
 Status CompositeOpExpansion::Run(EagerOperation* orig_op,
                                  std::unique_ptr<EagerOperation>* out_op) {
   if (!IsEnabled()) return Status::OK();
   if (orig_op->Device() != kVariantDeviceNull) return Status::OK();
 
-  VLOG(1) << "Run Node Expansion Passes";
+  LOG_FIRST_N(INFO, 1) << "Run Node Expansion Passes";
 
   // Get the FunctionDef and insert that into the context
   const NodeDef& ndef = orig_op->MutableAttrs()->BuildNodeDef();
@@ -40,7 +41,7 @@ Status CompositeOpExpansion::Run(EagerOperation* orig_op,
   std::string fname =
       absl::StrCat("_expanded_", ndef.name(), "_", std::to_string(x));
   if (!ctx.FindFunctionByName(fname)) {
-    TF_ASSIGN_OR_RETURN(auto func, TFRDecomposeContext::Expand(ndef, fname));
+    TF_ASSIGN_OR_RETURN(auto func, ExpandNode(ndef, fname));
     TF_RETURN_IF_ERROR(ctx.AddFunctionDef(func));
   }
 
@@ -55,11 +56,14 @@ Status CompositeOpExpansion::Run(EagerOperation* orig_op,
   new_op->MutableAttrs()->CopyAttributes(orig_op->Attrs());
   out_op->reset(new_op);
 
-  VLOG(1) << "Rewrite the op to call function: " << fname;
+  LOG_FIRST_N(INFO, 1)
+      << "Finish Node Expansion Passes. Rewrite the op to call function: "
+      << fname;
 
   return Status::OK();
 }
 
 REGISTER_REWRITE(EagerOpRewriteRegistry::POST_PLACEMENT, CompositeOpExpansion);
 
+}  // namespace tfr
 }  // namespace tensorflow

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

@@ -20,6 +20,7 @@ limitations under the License.
 #include "tensorflow/stream_executor/lib/statusor.h"
 
 namespace tensorflow {
+namespace tfr {
 
 // An optimization pass that decompose the composite ops in a module according
 // to the decomposition library. Currently the decomposition library is loaded
@@ -42,6 +43,7 @@ class CompositeOpExpansion : public EagerOpRewrite {
   }
 };
 
+}  // namespace tfr
 }  // namespace tensorflow
 
 #endif  // TENSORFLOW_COMPILER_MLIR_TFR_INTEGRATION_NODE_EXPANSION_PASS_H_

tensorflow/compiler/mlir/tfr/integration/node_expansion_test.py

@@ -34,14 +34,6 @@
 
 class NodeExpansionTest(test.TestCase):
 
-  def setUp(self):
-    os.environ['TF_MLIR_TFR_LIB_DIR'] = 'tensorflow/compiler/mlir/tfr/resources'
-    super(NodeExpansionTest, self).setUp()
-
-  def tearDown(self):
-    del os.environ['TF_MLIR_TFR_LIB_DIR']
-    super(NodeExpansionTest, self).tearDown()
-
   def testAddN(self):
     t1 = constant_op.constant([[1.0, 2.0], [3.0, 4.0]])
     t2 = constant_op.constant([[1.0, 2.0], [3.0, 4.0]])
@@ -81,5 +73,6 @@ def biasd_dense_elu(x, y, z):
 
 
 if __name__ == '__main__':
+  os.environ['TF_MLIR_TFR_LIB_DIR'] = 'tensorflow/compiler/mlir/tfr/resources'
   ops.enable_eager_execution()
   test.main()

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

@@ -54,6 +54,7 @@ limitations under the License.
 #include "tensorflow/stream_executor/lib/statusor.h"
 
 namespace tensorflow {
+namespace tfr {
 
 const char* const kTFRLibEnv = "TF_MLIR_TFR_LIB_DIR";
 
@@ -66,6 +67,10 @@ StatusOr<std::unique_ptr<TFRDecomposeContext>> TFRDecomposeContext::Get(
   string composite_mlir_dir = io::JoinPath(env->GetRunfilesDir(), tfr_lib_dir);
   std::vector<string> files;
   TF_RETURN_IF_ERROR(env->GetChildren(composite_mlir_dir, &files));
+  if (files.empty()) {
+    return errors::Internal(absl::StrCat(
+        "Failed to find the decomposition lib from path ", composite_mlir_dir));
+  }
   std::string tfr_raw_text;
   for (const auto& file : files) {
     string fullpath = io::JoinPath(composite_mlir_dir, file);
@@ -76,15 +81,15 @@ StatusOr<std::unique_ptr<TFRDecomposeContext>> TFRDecomposeContext::Get(
     }
   }
 
-  auto ctx = TFRDecomposeContext::Get(tfr_raw_text, mlir_ctx);
+  auto ctx = TFRDecomposeContext::GetFromText(tfr_raw_text, mlir_ctx);
   if (!ctx) {
     return errors::Internal(absl::StrCat(
         "Failed to load the imported decomposition lib: ", tfr_raw_text));
   }
   return ctx;
 }
 
-std::unique_ptr<TFRDecomposeContext> TFRDecomposeContext::Get(
+std::unique_ptr<TFRDecomposeContext> TFRDecomposeContext::GetFromText(
     StringPiece tfr_raw_text, mlir::MLIRContext* mlir_ctx) {
   mlir_ctx->allowUnregisteredDialects(/*allow=*/true);
   // Load dialects involved in the conversion
@@ -105,20 +110,22 @@ std::unique_ptr<TFRDecomposeContext> TFRDecomposeContext::Get(
   llvm::SourceMgr source_mgr;
   source_mgr.AddNewSourceBuffer(std::move(memory_buffer), llvm::SMLoc());
   mlir::OwningModuleRef module = mlir::parseSourceFile(source_mgr, mlir_ctx);
+  // The MLIRContext owns the module
+  auto module_op = module.release();
 
   // Create the context
-  return absl::make_unique<TFRDecomposeContext>(std::move(module));
+  return absl::make_unique<TFRDecomposeContext>(module_op);
 }
 
-StatusOr<FunctionDef> TFRDecomposeContext::Decompose(const NodeDef& node_def,
-                                                     StringPiece func_name) {
+StatusOr<FunctionDef> TFRDecomposeContext::ExpandNode(const NodeDef& node_def,
+                                                      StringPiece 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();
+  mlir::MLIRContext* context = tfr_module_.getContext();
   llvm::SmallVector<mlir::Type, 4> input_tys, output_tys;
   mlir::Builder builder(context);
   for (auto ty : input_dtys) {
@@ -159,15 +166,8 @@ StatusOr<FunctionDef> TFRDecomposeContext::Decompose(const NodeDef& node_def,
   mlir::Operation* tf_op = op_builder.createOperation(op_state);
   op_builder.create<mlir::ReturnOp>(loc, tf_op->getResults());
 
-  if (failed(mlir::verify(module))) {
-    return errors::Internal(absl::StrCat(
-        "Failed to verify the imported NodeDef: ", node_def.DebugString()));
-  }
-
-  // Call the decompose passes by using the external symbol table.
-  if (failed(pm_.run(module))) {
-    return errors::Internal("Failed to run the decompose passes.");
-  }
+  // Run the decompose passes on the module
+  TF_RETURN_IF_ERROR(DecomposeGraph(module));
 
   // Export the result as a FunctionDef.
   FunctionDef func_def;
@@ -177,43 +177,46 @@ StatusOr<FunctionDef> TFRDecomposeContext::Decompose(const NodeDef& node_def,
   return func_def;
 }
 
-Status TFRDecomposeContext::Decompose(mlir::ModuleOp user_module) {
+Status TFRDecomposeContext::DecomposeGraph(mlir::ModuleOp user_module) {
   // Call the decompose passes by using the external symbol table.
   if (failed(pm_.run(user_module))) {
     return errors::Internal("Failed to run the decompose passes.");
   }
   return Status::OK();
 }
 
-StatusOr<FunctionDef> TFRDecomposeContext::Expand(const NodeDef& node_def,
-                                                  StringPiece func_name) {
-  mlir::MLIRContext mlir_ctx;
-  mlir_ctx.allowUnregisteredDialects(/*allow=*/true);
-  TF_ASSIGN_OR_RETURN(auto ctx, Get(&mlir_ctx));
-  return ctx->Decompose(node_def, func_name);
-}
-
-Status TFRDecomposeContext::Destroy() {
-  tfr_module_.release().erase();
-  return Status::OK();
-}
-
 // Constructor of the decompose context.
-TFRDecomposeContext::TFRDecomposeContext(mlir::OwningModuleRef tfr_module)
-    : tfr_module_(std::move(tfr_module)), pm_(tfr_module_->getContext()) {
+TFRDecomposeContext::TFRDecomposeContext(mlir::ModuleOp tfr_module)
+    : tfr_module_(tfr_module), pm_(tfr_module_.getContext()) {
   mlir::OpPassManager& func_pm = pm_.nest<mlir::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_.get()));
+  func_pm.addPass(mlir::TFR::CreateDecomposeTFOpsPass(tfr_module_));
   func_pm.addPass(mlir::TFR::CreateRaiseToTFOpsPass(
-      tfr_module_.get(), /*materialize_derived_attrs=*/true));
+      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(); }
+
+StatusOr<FunctionDef> ExpandNode(const NodeDef& node_def,
+                                 StringPiece func_name) {
+  mlir::MLIRContext mlir_ctx;
+  TF_ASSIGN_OR_RETURN(auto ctx, TFRDecomposeContext::Get(&mlir_ctx));
+  return ctx->ExpandNode(node_def, func_name);
+}
+
+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