Add an utility pass for writing atom programs and main IFRT func to f…

…iles.

PiperOrigin-RevId: 714569854

xla/python/ifrt/ir/tests/BUILD

@@ -32,6 +32,7 @@ xla_cc_binary(
     testonly = True,
     srcs = ["ifrt-opt.cc"],
     deps = [
+        "//tensorflow/compiler/mlir:init_mlir",
         "//xla/mlir_hlo:hlo_dialect_registration",
         "//xla/pjrt:pjrt_executable",
         "//xla/python/ifrt",

xla/python/ifrt/ir/tests/ifrt-opt.cc

@@ -28,6 +28,7 @@ limitations under the License.
 #include "mlir/IR/DialectRegistry.h"
 #include "mlir/InitAllPasses.h"
 #include "mlir/Tools/mlir-opt/MlirOptMain.h"
+#include "third_party/tensorflow/compiler/mlir/init_mlir.h"
 #include "xla/mlir_hlo/mhlo/IR/register.h"
 #include "xla/pjrt/pjrt_executable.h"
 #include "xla/python/ifrt/dtype.h"
@@ -117,6 +118,8 @@ class TestChildExecutableCompiler : public AtomProgramCompiler {
 }  // namespace xla
 
 int main(int argc, char** argv) {
+  tensorflow::InitMlir y(&argc, &argv);
+
   std::shared_ptr<xla::ifrt::AtomProgramCompiler> compiler =
       std::make_shared<xla::ifrt::TestChildExecutableCompiler>();
   auto compile_options = std::make_shared<absl::flat_hash_map<

xla/python/ifrt/ir/transforms/BUILD

@@ -33,6 +33,7 @@ cc_library(
         "ifrt_atom_programs_to_vhlo_pass.cc",
         "ifrt_compile_and_propagate_shardings_pass.cc",
         "ifrt_compile_atom_program_pass.cc",
+        "ifrt_dump_atom_programs_pass.cc",
         "ifrt_duplicated_callee_elimination_pass.cc",
         "ifrt_legalize_to_vifrt_pass.cc",
         "ifrt_lower_atom_program_metadata_to_xla_pass.cc",
@@ -87,6 +88,7 @@ cc_library(
         "//xla/service:hlo_proto_cc",
         "//xla/service/spmd/shardy:constants",
         "//xla/service/spmd/shardy:utils",
+        "//xla/tsl/platform:env",
         "@com_google_absl//absl/cleanup",
         "@com_google_absl//absl/container:btree",
         "@com_google_absl//absl/container:flat_hash_map",
@@ -114,6 +116,7 @@ cc_library(
         "@stablehlo//:stablehlo_serialization",
         "@tsl//tsl/platform:env",
         "@tsl//tsl/platform:fingerprint",
+        "@tsl//tsl/platform:path",
         "@tsl//tsl/platform:protobuf",
         "@tsl//tsl/platform:statusor",
     ],

xla/python/ifrt/ir/transforms/ifrt_dump_atom_programs_pass.cc

@@ -0,0 +1,117 @@
+/* Copyright 2025 The OpenXLA Authors.
+
+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 <memory>
+#include <string>
+
+#include "absl/container/flat_hash_set.h"
+#include "absl/log/check.h"
+#include "absl/strings/str_cat.h"
+#include "llvm/Support/Casting.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/IR/BuiltinOps.h"
+#include "mlir/IR/Operation.h"
+#include "mlir/IR/OperationSupport.h"
+#include "mlir/IR/SymbolTable.h"
+#include "mlir/IR/Visitors.h"
+#include "mlir/Pass/Pass.h"
+#include "xla/python/ifrt/ir/ifrt_ops.h"
+#include "xla/python/ifrt/ir/transforms/passes.h"
+#include "xla/python/ifrt/ir/transforms/utils.h"
+#include "xla/tsl/platform/env.h"
+#include "tsl/platform/path.h"
+
+namespace xla {
+namespace ifrt {
+
+namespace {
+
+#define GEN_PASS_DEF_IFRTDUMPATOMPROGRAMSPASS
+#include "xla/python/ifrt/ir/transforms/passes.h.inc"
+
+absl::Status DumpOperation(mlir::Operation* op, std::string dump_dir,
+                           std::string filename) {
+  std::string file_path =
+      tsl::io::JoinPath(dump_dir, absl::StrCat(filename, ".mlir"));
+  return tsl::WriteStringToFile(tsl::Env::Default(), file_path,
+                                OperationToString(op, mlir::OpPrintingFlags()));
+}
+
+class IfrtDumpAtomProgramsPass
+    : public impl::IfrtDumpAtomProgramsPassBase<IfrtDumpAtomProgramsPass> {
+ public:
+  using impl::IfrtDumpAtomProgramsPassBase<
+      IfrtDumpAtomProgramsPass>::IfrtDumpAtomProgramsPassBase;
+
+  void runOnOperation() override {
+    if (dump_dir.empty()) {
+      return signalPassFailure();
+    }
+
+    mlir::SymbolTableCollection symbol_table;
+    mlir::ModuleOp module_op = getOperation();
+    // Keeps track of the atom programs that have already been dumped.
+    absl::flat_hash_set<std::string> dumped_atom_program_names;
+
+    auto main_func = GetMainFunction(module_op);
+
+    // Clones the main function to ensure that the attribute aliases are
+    // preserved while printing. Otherwise, the op would be printed in its
+    // full form (i.e., every argument with the entire device list expanded)
+    // and would lead to large ifrt dump files.
+    auto cloned_main = main_func.clone();
+    if (auto status = DumpOperation(cloned_main, dump_dir, "ifrt_main_func");
+        !status.ok()) {
+      cloned_main.erase();
+      main_func->emitOpError()
+          << "failed to dump main func: " << status.ToString();
+      signalPassFailure();
+      return;
+    }
+    cloned_main.erase();
+
+    mlir::WalkResult result =
+        main_func.walk([&](CallOp call_op) -> mlir::WalkResult {
+          mlir::func::FuncOp callee = call_op.getCalleeOp(symbol_table);
+          CHECK(callee != nullptr);
+          auto atom_program_module =
+              llvm::cast<mlir::ModuleOp>(callee->getParentOp());
+          std::string atom_program_name =
+              atom_program_module.getSymNameAttr().str();
+          if (dumped_atom_program_names.insert(atom_program_name).second) {
+            if (auto status = DumpOperation(atom_program_module, dump_dir,
+                                            atom_program_name);
+                !status.ok()) {
+              return call_op->emitOpError()
+                     << "failed to dump atom program: " << status.ToString();
+            }
+          }
+          return mlir::WalkResult::advance();
+        });
+    if (result.wasInterrupted()) {
+      signalPassFailure();
+    }
+  }
+};
+
+}  // namespace
+
+std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
+CreateIfrtDumpAtomProgramsPass(IfrtDumpAtomProgramsPassOptions options) {
+  return std::make_unique<IfrtDumpAtomProgramsPass>(options);
+}
+
+}  // namespace ifrt
+}  // namespace xla

xla/python/ifrt/ir/transforms/passes.h

@@ -129,6 +129,9 @@ std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
 CreateIfrtVerifyDeviceTypeConsistencyPass(
     IfrtVerifyDeviceTypeConsistencyPassOptions options = {});
 
+std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
+CreateIfrtDumpAtomProgramsPass(IfrtDumpAtomProgramsPassOptions options = {});
+
 std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
 CreateIfrtAtomProgramsToVhloPass(
     tsl::protobuf::RepeatedPtrField<IfrtIrAtomProgramProto>* atom_programs,

xla/python/ifrt/ir/transforms/passes.td

@@ -474,6 +474,16 @@ This pass fails if
   ];
 }
 
+def IfrtDumpAtomProgramsPass
+    : Pass<"ifrt-dump-atom-programs", "mlir::ModuleOp"> {
+  let summary = "Extracts atom programs from module and dumps them to files.";
+  let constructor = "CreateIfrtDumpAtomProgramsPass()";
+  let options = [
+    Option<"dump_dir", "dump_dir", "std::string", "",
+           "The directory to dump the atom programs and the main function to.">,
+  ];
+}
+
 def IfrtLegalizeToVifrtPass : Pass<"ifrt-legalize-to-vifrt", "mlir::ModuleOp"> {
   let summary = "Legalize IFRT to VIFRT.";
   let dependentDialects = ["xla::ifrt::VifrtDialect"];