[XLA:GPU] Refactor Triton pass dumping and improve error logging.

Extracts the logic for setting up MLIR pass dumping to a helper class and uses it in both `CompileTritonToLLVM` and `LowerXTileToTriton`.

The public API is minimized to a single function EnableIRPrintingIfRequested.

PiperOrigin-RevId: 853236668

Author: loislo

xla/backends/gpu/codegen/triton/BUILD

@@ -320,13 +320,13 @@ cc_library(
         ":lowering_util",
         ":support",
         "//xla:autotuning_proto_cc",
-        "//xla:status_macros",
         "//xla:util",
         "//xla:xla_data_proto_cc",
         "//xla:xla_proto_cc",
         "//xla/backends/gpu/codegen/emitters/ir:xla_gpu",
         "//xla/backends/gpu/codegen/triton/ir:triton_xla",
         "//xla/backends/gpu/codegen/triton/transforms:passes",
+        "//xla/codegen:ir_printing",
         "//xla/codegen/emitters/ir:xla",
         "//xla/codegen/emitters/transforms:passes",
         "//xla/codegen/tiling:symbolic_tile_analysis",
@@ -383,7 +383,6 @@ cc_library(
         "@llvm-project//mlir:NVVMToLLVMIRTranslation",
         "@llvm-project//mlir:Pass",
         "@llvm-project//mlir:ROCDLToLLVMIRTranslation",
-        "@llvm-project//mlir:SCFToControlFlow",
         "@llvm-project//mlir:Support",
         "@llvm-project//mlir:TensorDialect",
         "@llvm-project//mlir:ToLLVMIRTranslation",

xla/backends/gpu/codegen/triton/triton_gemm_fusion_test.cc

@@ -388,16 +388,25 @@ ENTRY e {
   }
   DebugOptions debug_options = verified_module->config().debug_options();
   debug_options.set_xla_dump_to(output_directory);
-  debug_options.set_xla_dump_emitter_re("triton-fusion");
+  debug_options.set_xla_dump_emitter_re("triton");
   verified_module->mutable_config().set_debug_options(debug_options);
 
   EXPECT_TRUE(RunAndCompare(std::move(verified_module),
                             ErrorSpec{/*aabs=*/1e-3, /*arel=*/1e-3}));
 
   std::vector<std::string> paths;
-  TF_EXPECT_OK(tsl::Env::Default()->GetMatchingPaths(
-      tsl::io::JoinPath(output_directory, "*.triton-passes.log"), &paths));
+  EXPECT_OK(tsl::Env::Default()->GetMatchingPaths(
+      tsl::io::JoinPath(output_directory, "*.xtile-to-triton.txt"), &paths));
   EXPECT_EQ(paths.size(), 1);
+  size_t file_size = 0;
+  EXPECT_OK(tsl::Env::Default()->GetFileSize(paths[0], &file_size));
+  EXPECT_GT(file_size, 10);
+  EXPECT_OK(tsl::Env::Default()->GetMatchingPaths(
+      tsl::io::JoinPath(output_directory, "*.triton-to-llvm.txt"), &paths));
+  EXPECT_EQ(paths.size(), 1);
+  file_size = 0;
+  EXPECT_OK(tsl::Env::Default()->GetFileSize(paths[0], &file_size));
+  EXPECT_GT(file_size, 10);
 }
 
 TEST_F(TritonGemmTest, DotWithPredFromCompareProducesCorrectResult) {

xla/backends/gpu/codegen/triton/xtile_compiler.cc

@@ -17,9 +17,7 @@ limitations under the License.
 
 #include <cstdint>
 #include <memory>
-#include <optional>
 #include <string>
-#include <system_error>  // NOLINT
 #include <utility>
 #include <vector>
 
@@ -38,15 +36,13 @@ limitations under the License.
 #include "llvm/IR/Module.h"
 #include "llvm/Linker/Linker.h"
 #include "llvm/Support/Debug.h"
-#include "llvm/Support/FileSystem.h"
 #include "llvm/Support/LogicalResult.h"
 #include "llvm/Support/raw_ostream.h"
 #include "llvm/TargetParser/Triple.h"
 #include "mlir/Conversion/AffineToStandard/AffineToStandard.h"
 #include "mlir/Conversion/ArithToLLVM/ArithToLLVM.h"
 #include "mlir/Conversion/ControlFlowToLLVM/ControlFlowToLLVM.h"
 #include "mlir/Conversion/IndexToLLVM/IndexToLLVM.h"
-#include "mlir/Conversion/SCFToControlFlow/SCFToControlFlow.h"
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Func/Extensions/InlinerExtension.h"
@@ -69,7 +65,6 @@ limitations under the License.
 #include "mlir/IR/OwningOpRef.h"
 #include "mlir/IR/Value.h"
 #include "mlir/IR/Verifier.h"
-#include "mlir/Pass/Pass.h"
 #include "mlir/Pass/PassManager.h"
 #include "mlir/Support/LLVM.h"
 #include "mlir/Support/LogicalResult.h"
@@ -90,6 +85,7 @@ limitations under the License.
 #include "xla/backends/gpu/codegen/triton/transforms/passes.h"
 #include "xla/codegen/emitters/ir/xla_dialect.h"
 #include "xla/codegen/emitters/transforms/passes.h"
+#include "xla/codegen/ir_printing.h"
 #include "xla/codegen/xtile/ir/transforms/passes.h"
 #include "xla/codegen/xtile/ir/xtile_dialect.h"
 #include "xla/hlo/builder/xla_builder.h"
@@ -107,7 +103,6 @@ limitations under the License.
 #include "xla/service/gpu/model/triton_emitter_constraints.h"
 #include "xla/service/hlo_module_config.h"
 #include "xla/service/llvm_ir/llvm_util.h"
-#include "xla/status_macros.h"
 #include "xla/stream_executor/cuda/cuda_compute_capability.h"
 #include "xla/stream_executor/device_description.h"
 #include "xla/stream_executor/gpu/tma_metadata.h"
@@ -119,7 +114,6 @@ limitations under the License.
 #include "xla/util.h"
 #include "xla/xla.pb.h"
 #include "xla/xla_data.pb.h"
-#include "tsl/platform/path.h"
 #include "triton/Dialect/Triton/IR/Dialect.h"
 #include "triton/Dialect/TritonGPU/IR/Dialect.h"
 
@@ -374,52 +368,10 @@ absl::StatusOr<TritonWrapperResult> CompileTritonToLLVM(
   should_verify = true;
 #endif
 
-  bool should_dump_mlir_passes =
-      hlo_config.debug_options().xla_enable_dumping() &&
-      DumpingEnabledForHloModule(hlo_module) &&
-      DumpingEnabledForEmitter("triton-fusion", hlo_config.debug_options());
-
   mlir::PassManager pm(&mlir_context);
+  EnableIRPrintingIfRequested(pm, &mlir_context, hlo_module, kernel_name,
+                              "triton-to-llvm");
   pm.enableVerifier(should_verify);
-
-  std::optional<llvm::raw_fd_ostream> log_stream;
-  if (should_dump_mlir_passes) {
-    std::string outputs_dir = hlo_config.debug_options().xla_dump_to();
-    if (outputs_dir == "sponge") {
-      if (!tsl::io::GetTestUndeclaredOutputsDir(&outputs_dir)) {
-        LOG(ERROR) << "Failed to get test undeclared outputs dir. Lets skip "
-                      "dumping triton passes.";
-        outputs_dir = "";
-      }
-    }
-    if (!outputs_dir.empty()) {
-      const std::string basename =
-          absl::StrCat(absl::string_view(tsl::io::Basename(hlo_module.name())),
-                       ".", kernel_name, ".triton-passes.log");
-      std::string path = tsl::io::JoinPath(outputs_dir, basename);
-      std::error_code err;
-      log_stream.emplace(path, err, llvm::sys::fs::OF_None);
-      if (err) {
-        log_stream.reset();
-        LOG(ERROR) << "Failed to dump triton passes to " << path << ": "
-                   << err.message();
-      } else {
-        pm.getContext()->disableMultithreading();
-        auto print_always = [](mlir::Pass*, mlir::Operation*) { return true; };
-        pm.enableIRPrinting(/*shouldPrintBeforePass=*/print_always,
-                            /*shouldPrintAfterPass=*/print_always,
-                            /*printModuleScope=*/true,
-                            /*printAfterOnlyOnChange=*/false,
-                            /*printAfterOnlyOnFailure=*/true, *log_stream);
-      }
-    } else {
-      LOG(ERROR)
-          << "--xla_dump_emitter_re=triton-fusion is set, but neither "
-          << "the environment variable TEST_UNDECLARED_OUTPUTS_DIR nor the "
-          << "flag --xla_dump_to is set, so the llvm dumps are disabled.";
-    }
-  }
-
   CreateTritonXlaPipeline(&pm, gpu_cc, /*rewrite_int4=*/is_xla_fusion,
                           block_level_parameters.is_tma_allowed,
                           block_level_parameters.num_stages);
@@ -537,8 +489,11 @@ absl::Status LowerXTileToTriton(mlir::ModuleOp xtile_dialect_module,
                                 const HloFusionInstruction& fusion,
                                 const se::DeviceDescription& device_info) {
   {
+    const HloModule& hlo_module = *fusion.GetModule();
     // Convert xTile ops to Triton ops.
     mlir::PassManager pm(&mlir_context);
+    EnableIRPrintingIfRequested(pm, &mlir_context, hlo_module, fusion.name(),
+                                "xtile-to-triton");
     // Disable verifier because the Triton code may be invalid due to the
     // unsupported types.
     pm.enableVerifier(/*enabled=*/false);
@@ -566,33 +521,40 @@ absl::Status LowerXTileToTriton(mlir::ModuleOp xtile_dialect_module,
     }
   }
 
-  if (fusion.GetModule()
-          ->config()
-          .debug_options()
-          .xla_gpu_experimental_scaled_dot_with_triton()) {
-    // Convert unsupported types before verification.
+  {
+    if (fusion.GetModule()
+            ->config()
+            .debug_options()
+            .xla_gpu_experimental_scaled_dot_with_triton()) {
+      // Convert unsupported types before verification.
+      mlir::PassManager pm(&mlir_context);
+
+      EnableIRPrintingIfRequested(pm, &mlir_context, *fusion.GetModule(),
+                                  fusion.name(),
+                                  "convert-scaled-dot-unsupported-types");
+      pm.addPass(
+          mlir::triton::xla::CreateTritonXLAConvertUnsupportedTypesPass());
+      if (mlir::failed(pm.run(xtile_dialect_module))) {
+        return CreateInternalError(
+            "Failed to fix unsupported types in Triton module for fusion:",
+            &fusion, xtile_dialect_module);
+      }
+    }
+
+    if (mlir::failed(mlir::verify(xtile_dialect_module))) {
+      return CreateInternalError("Failed to verify Triton module for fusion:",
+                                 &fusion, xtile_dialect_module);
+    }
     mlir::PassManager pm(&mlir_context);
-    pm.addPass(mlir::triton::xla::CreateTritonXLAConvertUnsupportedTypesPass());
+    EnableIRPrintingIfRequested(pm, &mlir_context, *fusion.GetModule(),
+                                fusion.name(), "canonicalize-cse");
+    pm.addPass(mlir::createCanonicalizerPass());
+    pm.addPass(mlir::createCSEPass());
     if (mlir::failed(pm.run(xtile_dialect_module))) {
-      return CreateInternalError(
-          "Failed to fix unsupported types in Triton module for fusion:",
-          &fusion, xtile_dialect_module);
+      return CreateInternalError("Failed to create Triton module for fusion:",
+                                 &fusion, xtile_dialect_module);
     }
   }
-
-  if (mlir::failed(mlir::verify(xtile_dialect_module))) {
-    return CreateInternalError("Failed to verify Triton module for fusion:",
-                               &fusion, xtile_dialect_module);
-  }
-  mlir::PassManager pm(&mlir_context);
-
-  pm.addPass(mlir::createCanonicalizerPass());
-  pm.addPass(mlir::createCSEPass());
-  if (mlir::failed(pm.run(xtile_dialect_module))) {
-    return CreateInternalError("Failed to create Triton module for fusion:",
-                               &fusion, xtile_dialect_module);
-  }
-
   return absl::OkStatus();
 }
 

xla/codegen/BUILD

@@ -35,6 +35,24 @@ cc_library(
     ],
 )
 
+cc_library(
+    name = "ir_printing",
+    srcs = ["ir_printing.cc"],
+    hdrs = ["ir_printing.h"],
+    deps = [
+        "//xla/hlo/ir:hlo",
+        "//xla/service:dump",
+        "//xla/service:hlo_module_config",
+        "@com_google_absl//absl/log",
+        "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/strings:string_view",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Pass",
+        "@tsl//tsl/platform:path",
+    ],
+)
+
 cc_library(
     name = "kernel_spec",
     srcs = ["kernel_spec.cc"],