Add inliner pass to CompileSerializedMlirToXlaHlo.

Graphs generated via model parallelism has sharding annotations but does not propagate them into functions as it is assumed functions are inlined.

PiperOrigin-RevId: 346637211
Change-Id: Ifc9a95b6036fd11770e12893b374940a248589c2

Author: andyly

tensorflow/compiler/mlir/tensorflow/utils/compile_mlir_util.cc

@@ -271,7 +271,8 @@ static void RegisterDialects(mlir::DialectRegistry& registry) {
 void CreateConvertMlirToXlaHloPipeline(
     mlir::OpPassManager& pm, llvm::StringRef device_type,
     llvm::MutableArrayRef<std::unique_ptr<mlir::Pass>>
-        custom_legalization_passes) {
+        custom_legalization_passes,
+    bool inline_after_legalization) {
   pm.addPass(mlir::TF::CreateTFFunctionalControlFlowToRegions());
   pm.addNestedPass<mlir::FuncOp>(mlir::createCanonicalizerPass());
   // Run shape inference pass before tensorlist decomposition to get buffer
@@ -319,6 +320,9 @@ void CreateConvertMlirToXlaHloPipeline(
   pm.addNestedPass<mlir::FuncOp>(mlir::mhlo::createLegalizeTFPass(
       /*allow_partial_conversion=*/false, /*legalize_chlo=*/true,
       /*tf2xla_fallback_device_type=*/device_type));
+
+  if (inline_after_legalization) pm.addPass(mlir::createInlinerPass());
+
   // In order to export to XLA, we must sink constants to control flow regions,
   // since XLA uses functional control flow.
   pm.addNestedPass<mlir::FuncOp>(
@@ -331,11 +335,13 @@ Status ConvertMLIRToXlaComputation(
     bool return_tuple,
     const XlaHelpers::ShapeRepresentationFn shape_representation_fn,
     llvm::MutableArrayRef<std::unique_ptr<mlir::Pass>>
-        custom_legalization_passes) {
+        custom_legalization_passes,
+    bool inline_after_legalization) {
   mlir::PassManager tf2xla(module_op.getContext());
   applyTensorflowAndCLOptions(tf2xla);
   CreateConvertMlirToXlaHloPipeline(tf2xla, device_type,
-                                    custom_legalization_passes);
+                                    custom_legalization_passes,
+                                    inline_after_legalization);
 
   if (VLOG_IS_ON(1)) {
     // Print the whole module after each pass which requires disabling
@@ -373,7 +379,8 @@ Status CompileMlirToXlaHlo(
     XlaHelpers::ShapeRepresentationFn shape_representation_fn,
     XlaCompilationResult* compilation_result,
     llvm::MutableArrayRef<std::unique_ptr<mlir::Pass>>
-        custom_legalization_passes) {
+        custom_legalization_passes,
+    bool inline_after_legalization) {
   if (VLOG_IS_ON(1))
     tensorflow::DumpMlirOpToFile("mlir_compile_before", module_op);
 
@@ -391,7 +398,7 @@ Status CompileMlirToXlaHlo(
   TF_RETURN_IF_ERROR(ConvertMLIRToXlaComputation(
       module_op, device_type, compilation_result->computation.get(),
       use_tuple_args, use_return_tuple, shape_representation_fn,
-      custom_legalization_passes));
+      custom_legalization_passes, inline_after_legalization));
 
   // Construct mapping from XlaComputation's arg to input edges of execute
   // node.
@@ -434,7 +441,8 @@ Status CompileSerializedMlirToXlaHlo(
   return CompileMlirToXlaHlo(
       mlir_module.get(), tensor_or_resource_shapes, device_type, use_tuple_args,
       /*use_return_tuple=*/true, /*use_resource_updates_for_aliases=*/false,
-      shape_representation_fn, compilation_result, custom_legalization_passes);
+      shape_representation_fn, compilation_result, custom_legalization_passes,
+      /*inline_after_legalization=*/true);
 }
 
 // Rewrites the given module with specified args. For each of the constant args,
@@ -535,7 +543,8 @@ Status CompileGraphToXlaHlo(
   auto status = CompileMlirToXlaHlo(
       module_op, arg_shapes, device_type, use_tuple_args, use_return_tuple,
       /*use_resource_updates_for_aliases=*/true, shape_representation_fn,
-      compilation_result, custom_legalization_passes);
+      compilation_result, custom_legalization_passes,
+      /*inline_after_legalization=*/false);
   compilation_result->input_mapping = remaining_params;
   return status;
 }

tensorflow/compiler/mlir/tensorflow/utils/compile_mlir_util.h

@@ -39,7 +39,8 @@ namespace tensorflow {
 void CreateConvertMlirToXlaHloPipeline(
     mlir::OpPassManager& pm, llvm::StringRef device_type,
     llvm::MutableArrayRef<std::unique_ptr<mlir::Pass>>
-        custom_legalization_passes);
+        custom_legalization_passes,
+    bool inline_after_legalization);
 
 // Lowers MLIR module to XLA HLO inside an XlaComputation. The input module
 // should only contain operations in tf dialect. If the input module contains
@@ -73,7 +74,8 @@ Status ConvertMLIRToXlaComputation(
     bool return_tuple,
     const XlaHelpers::ShapeRepresentationFn shape_representation_fn = nullptr,
     llvm::MutableArrayRef<std::unique_ptr<mlir::Pass>>
-        custom_legalization_passes = {});
+        custom_legalization_passes = {},
+    bool inline_after_legalization = false);
 
 // Helper struct representing argument tensor or resource handle shapes.
 struct TensorOrResourceShape {
@@ -83,14 +85,16 @@ struct TensorOrResourceShape {
 
 // Compiles a MLIR module into XLA HLO, generates all accompanying metadata and
 // stores them in CompilationResult.
+// TODO(hinsu): Migrate options to separate struct.
 Status CompileMlirToXlaHlo(
     mlir::ModuleOp module_op, llvm::ArrayRef<TensorOrResourceShape> arg_shapes,
     llvm::StringRef device_type, bool use_tuple_args, bool use_return_tuple,
     bool use_resource_updates_for_aliases,
     XlaHelpers::ShapeRepresentationFn shape_representation_fn,
     XlaCompilationResult* compilation_result,
     llvm::MutableArrayRef<std::unique_ptr<mlir::Pass>>
-        custom_legalization_passes);
+        custom_legalization_passes,
+    bool inline_after_legalization);
 
 // Compiles a serialized MLIR module into XLA HLO, generates all accompanying
 // metadata and stores them in CompilationResult.

tensorflow/compiler/mlir/tensorflow/utils/compile_mlir_util_pass.cc

@@ -21,7 +21,7 @@ namespace {
 void CreateConvertMlirToXlaHloPipelineWithDefaults(mlir::OpPassManager& pm) {
   tensorflow::CreateConvertMlirToXlaHloPipeline(
       pm, /*device_type=*/"XLA_CPU_JIT",
-      /*custom_legalization_passes=*/{});
+      /*custom_legalization_passes=*/{}, /*inline_after_legalization=*/false);
 }
 
 mlir::PassPipelineRegistration<> pipeline(

tensorflow/compiler/mlir/tensorflow/utils/tf_xla_mlir_translate.cc

@@ -250,7 +250,7 @@ static mlir::LogicalResult MlirTfToHloTextTranslateFunction(
       module_op, arg_shapes, /*device_type=*/"XLA_CPU_JIT", emit_use_tuple_arg,
       emit_return_tuple, /*use_resource_updates_for_aliases=*/true,
       IdentityShapeRepresentationFn(), &compilation_result,
-      /*custom_legalization_passes=*/{});
+      /*custom_legalization_passes=*/{}, /*inline_after_legalization=*/false);
   if (!compilation_status.ok()) {
     LOG(ERROR) << "TF/XLA compilation failed: "
                << compilation_status.ToString();