Put the MLIR emitter behind a flag.

PiperOrigin-RevId: 609346365
tensorflow · Feb 22, 2024 · 4d1da7b · 4d1da7b
1 parent ce97af5
commit 4d1da7b
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 2 deletions.
diff --git a/third_party/xla/xla/service/gpu/fusions/fusions.cc b/third_party/xla/xla/service/gpu/fusions/fusions.cc
@@ -189,7 +189,13 @@ absl::StatusOr<std::unique_ptr<FusionInterface>> GetFusionEmitter(
         return *std::move(copy_fusion);
       }
 
-      if (mlir_converter::IsHloConversionSupported(
+      if (analysis.fusion_roots()
+              .front()
+              ->GetModule()
+              ->config()
+              .debug_options()
+              .xla_gpu_enable_mlir_emitters() &&
+          mlir_converter::IsHloConversionSupported(
               analysis.fusion(),
               fusion_info.analysis().device_info().gpu_compute_capability())) {
         return std::make_unique<MlirLoopFusion>(analysis);

diff --git a/third_party/xla/xla/xla.proto b/third_party/xla/xla/xla.proto
@@ -695,7 +695,9 @@ message DebugOptions {
   // will use for p2p operations.
   int64 xla_gpu_nccl_p2p_max_nchannels = 274;
 
-  // Next id: 275
+  bool xla_gpu_enable_mlir_emitters = 275;
+
+  // Next id: 276
 
   // Extra options to pass to the compilation backend (e.g. LLVM); specific
   // interpretation of these values is left to the backend.