Add a utility pass for writing atom programs and main IFRT func to fi…

…les. PiperOrigin-RevId: 714569854
openxla · Jan 12, 2025 · c1c63f1 · c1c63f1
1 parent 0cb8b51
commit c1c63f1
Show file tree

Hide file tree

Showing 36 changed files with 167 additions and 30 deletions.
diff --git a/xla/python/ifrt/ir/tests/BUILD b/xla/python/ifrt/ir/tests/BUILD
@@ -51,6 +51,7 @@ xla_cc_binary(
         "@llvm-project//mlir:AllPassesAndDialects",
         "@llvm-project//mlir:IR",
         "@llvm-project//mlir:MlirOptLib",
+        "@tsl//tsl/platform:platform_port",
     ],
 )
 

diff --git a/xla/python/ifrt/ir/tests/ifrt-opt.cc b/xla/python/ifrt/ir/tests/ifrt-opt.cc
@@ -39,6 +39,7 @@ limitations under the License.
 #include "xla/python/ifrt/mock.h"
 #include "xla/python/ifrt/shape.h"
 #include "xla/python/ifrt/support/module_parsing.h"
+#include "tsl/platform/init_main.h"
 
 namespace xla {
 namespace ifrt {
@@ -117,6 +118,8 @@ class TestChildExecutableCompiler : public AtomProgramCompiler {
 }  // namespace xla
 
 int main(int argc, char** argv) {
+  tsl::port::InitMain(argv[0], &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<

diff --git a/xla/python/ifrt/ir/tests/ifrt_compile_and_propagate_shardings.mlir b/xla/python/ifrt/ir/tests/ifrt_compile_and_propagate_shardings.mlir
@@ -1,4 +1,4 @@
-// RUN: ifrt-opt %s -ifrt-compile-and-propagate-shardings -split-input-file | FileCheck %s
+// RUN: ifrt-opt %s -- -ifrt-compile-and-propagate-shardings -split-input-file | FileCheck %s
 
 !array = !ifrt.array<tensor<2x2xi32>,
                      #ifrt.sharding_param<2x1 to [0] on 2>, [0, 1]>

diff --git a/xla/python/ifrt/ir/tests/ifrt_compile_atom_program.mlir b/xla/python/ifrt/ir/tests/ifrt_compile_atom_program.mlir
@@ -1,4 +1,4 @@
-// RUN: ifrt-opt %s -ifrt-compile-atom-program -split-input-file | FileCheck %s
+// RUN: ifrt-opt %s -- -ifrt-compile-atom-program -split-input-file | FileCheck %s
 
 // CHECK-LABEL: @call_hlo
 !array = !ifrt.array<tensor<2x2xi32>,

diff --git a/xla/python/ifrt/ir/tests/ifrt_duplicated_callee_elimination.mlir b/xla/python/ifrt/ir/tests/ifrt_duplicated_callee_elimination.mlir
@@ -1,4 +1,4 @@
-// RUN: ifrt-opt %s -ifrt-duplicated-callee-elimination | FileCheck %s
+// RUN: ifrt-opt %s -- -ifrt-duplicated-callee-elimination | FileCheck %s
 
 // CHECK-LABEL: @main
 func.func @main(%arg0: !ifrt.array<tensor<2x2xi32>,

diff --git a/xla/python/ifrt/ir/tests/ifrt_lower_atom_program_metadata_to_xla.mlir b/xla/python/ifrt/ir/tests/ifrt_lower_atom_program_metadata_to_xla.mlir
@@ -1,4 +1,4 @@
-// RUN: ifrt-opt %s -ifrt-lower-atom-program-metadata-to-xla -split-input-file -verify-diagnostics | FileCheck %s
+// RUN: ifrt-opt %s -- -ifrt-lower-atom-program-metadata-to-xla -split-input-file -verify-diagnostics | FileCheck %s
 
 // CHECK-LABEL: @arg_metadata
 module @arg_metadata attributes {ifrt.num_devices = 2} {

diff --git a/xla/python/ifrt/ir/tests/ifrt_lower_mpmd_reshard_to_call.mlir b/xla/python/ifrt/ir/tests/ifrt_lower_mpmd_reshard_to_call.mlir
@@ -1,4 +1,4 @@
-// RUN: ifrt-opt %s -ifrt-lower-mpmd-reshard-to-call -split-input-file -verify-diagnostics | FileCheck %s
+// RUN: ifrt-opt %s -- -ifrt-lower-mpmd-reshard-to-call -split-input-file -verify-diagnostics | FileCheck %s
 
 !array0 = !ifrt.array<tensor<2x2xi32>,
                       #ifrt.sharding_param<2x1 to [0] on 2>, [0, 1]>

diff --git a/xla/python/ifrt/ir/tests/ifrt_merge_reshards.mlir b/xla/python/ifrt/ir/tests/ifrt_merge_reshards.mlir
@@ -1,4 +1,4 @@
-// RUN: ifrt-opt %s -ifrt-merge-reshards | FileCheck %s
+// RUN: ifrt-opt %s -- -ifrt-merge-reshards | FileCheck %s
 
 #sharding = #ifrt.sharding_param<2 to [0] on 2>
 !array0 = !ifrt.array<tensor<2xi32>, #sharding, [0,1]>

diff --git a/xla/python/ifrt/ir/tests/ifrt_outline_atom_program_to_module.mlir b/xla/python/ifrt/ir/tests/ifrt_outline_atom_program_to_module.mlir
@@ -1,4 +1,4 @@
-// RUN: ifrt-opt %s -ifrt-outline-atom-program-to-module -split-input-file -verify-diagnostics | FileCheck %s
+// RUN: ifrt-opt %s -- -ifrt-outline-atom-program-to-module -split-input-file -verify-diagnostics | FileCheck %s
 
 !array = !ifrt.array<tensor<2x2xi32>,
                      #ifrt.sharding_param<2x1 to [0] on 2>, [0,1]>

diff --git a/xla/python/ifrt/ir/tests/ifrt_populate_atom_program_metadata.mlir b/xla/python/ifrt/ir/tests/ifrt_populate_atom_program_metadata.mlir
@@ -1,4 +1,4 @@
-// RUN: ifrt-opt %s -ifrt-populate-atom-program-metadata -ifrt-duplicated-callee-elimination -symbol-dce -split-input-file | FileCheck %s
+// RUN: ifrt-opt %s -- -ifrt-populate-atom-program-metadata -ifrt-duplicated-callee-elimination -symbol-dce -split-input-file | FileCheck %s
 
 !array = !ifrt.array<tensor<2x2xi32>,
                      #ifrt.sharding_param<2x1 to [0] on 2>, [0,1],

diff --git a/xla/python/ifrt/ir/tests/ifrt_precompile_atom_program_preprocessing.mlir b/xla/python/ifrt/ir/tests/ifrt_precompile_atom_program_preprocessing.mlir
@@ -1,4 +1,4 @@
-// RUN: ifrt-opt %s -ifrt-precompile-atom-program-preprocessing='platform_names=tpu,tpu' -split-input-file -verify-diagnostics | FileCheck %s
+// RUN: ifrt-opt %s -- -ifrt-precompile-atom-program-preprocessing='platform_names=tpu,tpu' -split-input-file -verify-diagnostics | FileCheck %s
 
 #sharding = #ifrt.sharding_param<2x1 to [0] on 2>
 !array = !ifrt.array<tensor<2x2xi32>, #sharding, [0, 1]>

diff --git a/xla/python/ifrt/ir/tests/ifrt_remove_attrs_from_other_dialects.mlir b/xla/python/ifrt/ir/tests/ifrt_remove_attrs_from_other_dialects.mlir
@@ -1,4 +1,4 @@
-// RUN: ifrt-opt %s -ifrt-remove-attrs-from-other-dialects -split-input-file | FileCheck %s
+// RUN: ifrt-opt %s -- -ifrt-remove-attrs-from-other-dialects -split-input-file | FileCheck %s
 
 !array = !ifrt.array<tensor<2x2xi32>, #ifrt.sharding_param<1x1 to [0] on 2>,
                      [0,1]>

diff --git a/xla/python/ifrt/ir/tests/ifrt_remove_ifrt_attrs.mlir b/xla/python/ifrt/ir/tests/ifrt_remove_ifrt_attrs.mlir
@@ -1,4 +1,4 @@
-// RUN: ifrt-opt %s -ifrt-remove-ifrt-attrs | FileCheck %s
+// RUN: ifrt-opt %s -- -ifrt-remove-ifrt-attrs | FileCheck %s
 
 // CHECK-LABEL: @ifrt_attributes_are_removed
 // CHECK-NOT: ifrt

diff --git a/xla/python/ifrt/ir/tests/ifrt_reshard_to_copy_arrays.mlir b/xla/python/ifrt/ir/tests/ifrt_reshard_to_copy_arrays.mlir
@@ -1,4 +1,4 @@
-// RUN: ifrt-opt %s -ifrt-reshard-to-copy-arrays -verify-diagnostics -split-input-file | FileCheck %s
+// RUN: ifrt-opt %s -- -ifrt-reshard-to-copy-arrays -verify-diagnostics -split-input-file | FileCheck %s
 
 !array0 = !ifrt.array<tensor<2x2xi32>, #ifrt.sharding_param<2x1 to [0] on 2>,
                      [0,1]>

diff --git a/xla/python/ifrt/ir/tests/ifrt_verify_device_type_consistency.mlir b/xla/python/ifrt/ir/tests/ifrt_verify_device_type_consistency.mlir
@@ -1,4 +1,4 @@
-// RUN: ifrt-opt %s -ifrt-verify-device-type-consistency='platform_names=tpu,tpu,cpu,tpu,cpu,cuda,cuda' -split-input-file -verify-diagnostics | FileCheck %s
+// RUN: ifrt-opt %s -- -ifrt-verify-device-type-consistency='platform_names=tpu,tpu,cpu,tpu,cpu,cuda,cuda' -split-input-file -verify-diagnostics | FileCheck %s
 
 // CHECK-LABEL: @good_call_multiple
 #sharding = #ifrt.sharding_param<2 to [0] on 2>

diff --git a/xla/python/ifrt/ir/tests/ifrt_verify_donation.mlir b/xla/python/ifrt/ir/tests/ifrt_verify_donation.mlir
@@ -1,4 +1,4 @@
-// RUN: ifrt-opt %s -ifrt-verify-donation -split-input-file -verify-diagnostics | FileCheck %s
+// RUN: ifrt-opt %s -- -ifrt-verify-donation -split-input-file -verify-diagnostics | FileCheck %s
 
 !array0 = !ifrt.array<tensor<2xi32>,
                       #ifrt.sharding_param<2 to [0] on 2>, [0, 1]>

diff --git a/xla/python/ifrt/ir/tests/ifrt_verify_sharding_specified.mlir b/xla/python/ifrt/ir/tests/ifrt_verify_sharding_specified.mlir
@@ -1,4 +1,4 @@
-// RUN: ifrt-opt %s -ifrt-verify-sharding-specified -split-input-file -verify-diagnostics | FileCheck %s
+// RUN: ifrt-opt %s -- -ifrt-verify-sharding-specified -split-input-file -verify-diagnostics | FileCheck %s
 
 // CHECK-LABEL: @good_arrays
 #sharding = #ifrt.sharding_param<2 to [0] on 2>

diff --git a/xla/python/ifrt/ir/tests/spmd_expansion.mlir b/xla/python/ifrt/ir/tests/spmd_expansion.mlir
@@ -1,4 +1,4 @@
-// RUN: ifrt-opt %s -spmd-expansion -split-input-file -verify-diagnostics | FileCheck %s
+// RUN: ifrt-opt %s -- -spmd-expansion -split-input-file -verify-diagnostics | FileCheck %s
 
 #device = #ifrt<devices[0,1]>
 #sharding = #ifrt.sharding_param<2x1 to [0] on 2>

diff --git a/xla/python/ifrt/ir/tests/spmd_interface_verification.mlir b/xla/python/ifrt/ir/tests/spmd_interface_verification.mlir
@@ -1,4 +1,4 @@
-// RUN: ifrt-opt %s -split-input-file -spmd-expandable-interface-verification='excluded-dialects=arith' -verify-diagnostics
+// RUN: ifrt-opt %s -- -split-input-file -spmd-expandable-interface-verification='excluded-dialects=arith' -verify-diagnostics
 
 module @good_return_only {
   func.func @main(

diff --git a/xla/python/ifrt/ir/tests/verify_array.mlir b/xla/python/ifrt/ir/tests/verify_array.mlir
@@ -1,4 +1,4 @@
-// RUN: ifrt-opt %s -split-input-file -verify-diagnostics
+// RUN: ifrt-opt %s -- -split-input-file -verify-diagnostics
 
 func.func @good_array() {
   /// Dim 0 of the tensor is sharded into 4 slices.

diff --git a/xla/python/ifrt/ir/tests/verify_assemble.mlir b/xla/python/ifrt/ir/tests/verify_assemble.mlir
@@ -1,4 +1,4 @@
-// RUN: ifrt-opt %s -split-input-file -verify-diagnostics
+// RUN: ifrt-opt %s -- -split-input-file -verify-diagnostics
 
 func.func @good_assemble(
     %arg0: !ifrt.array<tensor<2x2xi32>,

diff --git a/xla/python/ifrt/ir/tests/verify_attrs.mlir b/xla/python/ifrt/ir/tests/verify_attrs.mlir
@@ -1,4 +1,4 @@
-// RUN: ifrt-opt %s -split-input-file -verify-diagnostics
+// RUN: ifrt-opt %s -- -split-input-file -verify-diagnostics
 
 func.func @good_function_attr() attributes {ifrt.function} {
   return

diff --git a/xla/python/ifrt/ir/tests/verify_call.mlir b/xla/python/ifrt/ir/tests/verify_call.mlir
@@ -1,4 +1,4 @@
-// RUN: ifrt-opt %s -split-input-file -verify-diagnostics
+// RUN: ifrt-opt %s -- -split-input-file -verify-diagnostics
 
 func.func @good_call(
     %arg0: !ifrt.array<tensor<2x2xi32>, #ifrt.sharding_param<1x1 to [0] on 2>,

diff --git a/xla/python/ifrt/ir/tests/verify_call_loaded_executable.mlir b/xla/python/ifrt/ir/tests/verify_call_loaded_executable.mlir
@@ -1,4 +1,4 @@
-// RUN: ifrt-opt %s -split-input-file -verify-diagnostics
+// RUN: ifrt-opt %s -- -split-input-file -verify-diagnostics
 
 func.func @good(
     %arg0: !ifrt.array<tensor<2x2xi32>, #ifrt.sharding_param<1x1 to [0] on 2>,

diff --git a/xla/python/ifrt/ir/tests/verify_copy_arrays.mlir b/xla/python/ifrt/ir/tests/verify_copy_arrays.mlir
@@ -1,4 +1,4 @@
-// RUN: ifrt-opt %s -split-input-file -verify-diagnostics
+// RUN: ifrt-opt %s -- -split-input-file -verify-diagnostics
 
 !array0 = !ifrt.array<tensor<2x4xi32>,
                       #ifrt.sharding_param<1x2 to [0] on 2>, [0,1]>

diff --git a/xla/python/ifrt/ir/tests/verify_disassemble.mlir b/xla/python/ifrt/ir/tests/verify_disassemble.mlir
@@ -1,4 +1,4 @@
-// RUN: ifrt-opt %s -split-input-file -verify-diagnostics
+// RUN: ifrt-opt %s -- -split-input-file -verify-diagnostics
 
 func.func @good_disassemble(
     %arg0: !ifrt.array<tensor<2x4xi32>,

diff --git a/xla/python/ifrt/ir/tests/verify_loaded_executable.mlir b/xla/python/ifrt/ir/tests/verify_loaded_executable.mlir
@@ -1,4 +1,4 @@
-// RUN: ifrt-opt %s -split-input-file -verify-diagnostics
+// RUN: ifrt-opt %s -- -split-input-file -verify-diagnostics
 
 ifrt.LoadedExecutable @good on devices [0,1]
     : (!ifrt.array<tensor<2x2xi32>, #ifrt.sharding_param<1x1 to [0] on 2>,

diff --git a/xla/python/ifrt/ir/tests/verify_remap_arrays.mlir b/xla/python/ifrt/ir/tests/verify_remap_arrays.mlir
@@ -1,4 +1,4 @@
-// RUN: ifrt-opt %s -split-input-file -verify-diagnostics
+// RUN: ifrt-opt %s -- -split-input-file -verify-diagnostics
 
 !array0 = !ifrt.array<tensor<2x2xi32>,
                       #ifrt.sharding_param<1x1 to [0] on 1>, [0]>

diff --git a/xla/python/ifrt/ir/tests/verify_reshard.mlir b/xla/python/ifrt/ir/tests/verify_reshard.mlir
@@ -1,4 +1,4 @@
-// RUN: ifrt-opt %s -split-input-file -verify-diagnostics
+// RUN: ifrt-opt %s -- -split-input-file -verify-diagnostics
 
 func.func @good_reshard(
     %arg0: !ifrt.array<tensor<2x2xi32>, #ifrt.sharding_param<1x1 to [0] on 2>,

diff --git a/xla/python/ifrt/ir/tests/vifrt/ifrt_legalize_to_vifrt.0_1_0.mlir b/xla/python/ifrt/ir/tests/vifrt/ifrt_legalize_to_vifrt.0_1_0.mlir
@@ -1,4 +1,4 @@
-// RUN: ifrt-opt %s --ifrt-legalize-to-vifrt --symbol-dce --mlir-print-op-generic -split-input-file | FileCheck %s
+// RUN: ifrt-opt %s -- --ifrt-legalize-to-vifrt --symbol-dce --mlir-print-op-generic -split-input-file | FileCheck %s
 // RUN: ifrt-translate --serialize --ifrt_version=0.1.0 --atom_program_version=1.8.0 --strip_debuginfo %s | ifrt-translate --deserialize --strip_debuginfo | ifrt-opt > %t.0
 // RUN: ifrt-opt %s > %t.1
 // RUN: diff %t.0 %t.1

diff --git a/xla/python/ifrt/ir/tests/vifrt/ifrt_legalize_to_vifrt.mlir b/xla/python/ifrt/ir/tests/vifrt/ifrt_legalize_to_vifrt.mlir
@@ -1,4 +1,4 @@
-// RUN: ifrt-opt %s --ifrt-legalize-to-vifrt --symbol-dce --mlir-print-op-generic -split-input-file | FileCheck %s
+// RUN: ifrt-opt %s -- --ifrt-legalize-to-vifrt --symbol-dce --mlir-print-op-generic -split-input-file | FileCheck %s
 // RUN: ifrt-translate --serialize --ifrt_version=current --atom_program_version=current %s | ifrt-translate --deserialize | ifrt-opt > %t.0
 // RUN: ifrt-opt %s > %t.1
 // RUN: diff %t.0 %t.1

diff --git a/xla/python/ifrt/ir/tests/vifrt/ifrt_legalize_to_vifrt_invalid.mlir b/xla/python/ifrt/ir/tests/vifrt/ifrt_legalize_to_vifrt_invalid.mlir
@@ -1,4 +1,4 @@
-// RUN: ifrt-opt --ifrt-legalize-to-vifrt --symbol-dce --split-input-file -verify-diagnostics %s
+// RUN: ifrt-opt -- --ifrt-legalize-to-vifrt --symbol-dce --split-input-file -verify-diagnostics %s
 
 !array_t0 = !ifrt.array<tensor<2x4xi32>,
                         #ifrt.sharding_param<1x1 to [0] on 2>, [0,1]>

diff --git a/xla/python/ifrt/ir/transforms/BUILD b/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",
     ],

diff --git a/xla/python/ifrt/ir/transforms/ifrt_dump_atom_programs_pass.cc b/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
diff --git a/xla/python/ifrt/ir/transforms/passes.h b/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,

diff --git a/xla/python/ifrt/ir/transforms/passes.td b/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"];