[IFRT] Consolidate dialect registration to a single util function.

xla/python/ifrt/ir/tests/BUILD

@@ -46,9 +46,8 @@ xla_cc_binary(
     srcs = ["ifrt-opt.cc"],
     deps = [
         "//xla/mlir_hlo:hlo_dialect_registration",
-        "//xla/python/ifrt/ir",
-        "//xla/python/ifrt/ir/transforms:built_in_spmd_expansions",
         "//xla/python/ifrt/ir/transforms:passes",
+        "//xla/python/ifrt/support:module_parsing",
         "@llvm-project//mlir:AllPassesAndDialects",
         "@llvm-project//mlir:IR",
         "@llvm-project//mlir:MlirOptLib",
@@ -64,19 +63,16 @@ cc_library(
     deps = [
         "//xla:status_macros",
         "//xla/mlir/utils:error_util",
-        "//xla/mlir_hlo:hlo_dialect_registration",
         "//xla/python/ifrt",
         "//xla/python/ifrt:test_util",
-        "//xla/python/ifrt/ir",
         "//xla/python/ifrt/ir:sharding_param",
-        "//xla/python/ifrt/ir/transforms:built_in_spmd_expansions",
+        "//xla/python/ifrt/support:module_parsing",
         "//xla/tsl/concurrency:ref_count",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/strings",
         "@com_google_absl//absl/strings:str_format",
         "@com_google_absl//absl/types:span",
-        "@llvm-project//mlir:AllPassesAndDialects",
         "@llvm-project//mlir:IR",
         "@llvm-project//mlir:Parser",
         "@tsl//tsl/platform:statusor",

xla/python/ifrt/ir/tests/executable_impl_test_base.cc

@@ -25,23 +25,19 @@ limitations under the License.
 #include "absl/strings/string_view.h"
 #include "absl/types/span.h"
 #include "mlir/IR/BuiltinOps.h"
-#include "mlir/IR/DialectRegistry.h"
 #include "mlir/IR/MLIRContext.h"
 #include "mlir/IR/OwningOpRef.h"
-#include "mlir/InitAllDialects.h"
 #include "mlir/Parser/Parser.h"
 #include "xla/mlir/utils/error_util.h"
-#include "xla/mlir_hlo/mhlo/IR/register.h"
 #include "xla/python/ifrt/array.h"
 #include "xla/python/ifrt/device.h"
 #include "xla/python/ifrt/device_list.h"
 #include "xla/python/ifrt/dtype.h"
-#include "xla/python/ifrt/ir/ifrt_dialect.h"
 #include "xla/python/ifrt/ir/sharding_param.h"
-#include "xla/python/ifrt/ir/transforms/built_in_spmd_expansions.h"
 #include "xla/python/ifrt/memory.h"
 #include "xla/python/ifrt/shape.h"
 #include "xla/python/ifrt/sharding.h"
+#include "xla/python/ifrt/support/module_parsing.h"
 #include "xla/python/ifrt/test_util.h"
 #include "xla/status_macros.h"
 #include "xla/tsl/concurrency/ref_count.h"
@@ -53,13 +49,7 @@ namespace test_util {
 
 IfrtIrExecutableImplTestBase::IfrtIrExecutableImplTestBase() {
   mlir::registerMLIRContextCLOptions();
-
-  mlir::DialectRegistry registry;
-  mlir::registerAllDialects(registry);
-  mlir::mhlo::registerAllMhloDialects(registry);
-  registry.insert<xla::ifrt::IfrtDialect>();
-  xla::ifrt::AttachBuiltInSpmdExpansions(registry);
-  mlir_context_.appendDialectRegistry(registry);
+  xla::ifrt::support::RegisterMlirDialects(mlir_context_);
 }
 
 void IfrtIrExecutableImplTestBase::SetUp() {

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

@@ -14,25 +14,18 @@ limitations under the License.
 ==============================================================================*/
 
 #include "mlir/IR/DialectRegistry.h"
-#include "mlir/InitAllDialects.h"
 #include "mlir/InitAllPasses.h"
 #include "mlir/Tools/mlir-opt/MlirOptMain.h"
 #include "xla/mlir_hlo/mhlo/IR/register.h"
-#include "xla/python/ifrt/ir/ifrt_dialect.h"
-#include "xla/python/ifrt/ir/transforms/built_in_spmd_expansions.h"
 #include "xla/python/ifrt/ir/transforms/passes.h"
+#include "xla/python/ifrt/support/module_parsing.h"
 
 int main(int argc, char** argv) {
-  mlir::DialectRegistry registry;
-  mlir::registerAllDialects(registry);
-  mlir::mhlo::registerAllMhloDialects(registry);
-  registry.insert<xla::ifrt::IfrtDialect>();
   mlir::registerAllPasses();
-  xla::ifrt::registerIfrtIrPasses();
   xla::ifrt::RegisterIfrtPassesAndPipelines();
-
-  xla::ifrt::AttachBuiltInSpmdExpansions(registry);
+  mlir::DialectRegistry registry;
+  xla::ifrt::support::InitializeMlirDialectRegistry(registry);
 
   return mlir::asMainReturnCode(
-      mlir::MlirOptMain(argc, argv, "IFRT dialect driver\n", registry));
+      mlir::MlirOptMain(argc, argv, "IFRT IR dialect driver\n", registry));
 }

xla/python/ifrt/support/BUILD

@@ -21,13 +21,10 @@ cc_library(
         "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/strings:str_format",
         "@com_google_absl//absl/strings:string_view",
-        "@llvm-project//mlir:ArithDialect",
-        "@llvm-project//mlir:FuncDialect",
+        "@llvm-project//mlir:AllPassesAndDialects",
         "@llvm-project//mlir:FuncExtensions",
         "@llvm-project//mlir:IR",
-        "@llvm-project//mlir:MLProgramDialect",
         "@llvm-project//mlir:Parser",
-        "@llvm-project//mlir:ShapeDialect",
         "@shardy//shardy/dialect/sdy/ir:register",
         "@stablehlo//:register",
     ],

xla/python/ifrt/support/module_parsing.cc

@@ -19,14 +19,11 @@ limitations under the License.
 #include "absl/status/statusor.h"
 #include "absl/strings/str_format.h"
 #include "absl/strings/string_view.h"
-#include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Func/Extensions/AllExtensions.h"
-#include "mlir/Dialect/Func/IR/FuncOps.h"
-#include "mlir/Dialect/MLProgram/IR/MLProgram.h"
-#include "mlir/Dialect/Shape/IR/Shape.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/MLIRContext.h"
 #include "mlir/IR/OwningOpRef.h"
+#include "mlir/InitAllDialects.h"
 #include "mlir/Parser/Parser.h"
 #include "shardy/dialect/sdy/ir/register.h"
 #include "stablehlo/dialect/Register.h"
@@ -39,18 +36,19 @@ namespace xla {
 namespace ifrt {
 namespace support {
 
-void RegisterMlirDialects(mlir::MLIRContext& context) {
-  mlir::DialectRegistry registry;
-  registry.insert<mlir::arith::ArithDialect>();
-  registry.insert<mlir::func::FuncDialect>();
-  registry.insert<mlir::ml_program::MLProgramDialect>();
-  registry.insert<mlir::shape::ShapeDialect>();
+void InitializeMlirDialectRegistry(mlir::DialectRegistry& registry) {
   registry.insert<xla::ifrt::IfrtDialect>();
+  mlir::registerAllDialects(registry);
   mlir::func::registerAllExtensions(registry);
   mlir::mhlo::registerAllMhloDialects(registry);
   mlir::sdy::registerAllDialects(registry);
   mlir::stablehlo::registerAllDialects(registry);
   xla::ifrt::AttachBuiltInSpmdExpansions(registry);
+}
+
+void RegisterMlirDialects(mlir::MLIRContext& context) {
+  mlir::DialectRegistry registry;
+  InitializeMlirDialectRegistry(registry);
   context.appendDialectRegistry(registry);
 }
 

xla/python/ifrt/support/module_parsing.h

@@ -26,6 +26,10 @@ namespace xla {
 namespace ifrt {
 namespace support {
 
+// Initializes the given MLIR dialect registry with dialects that are required
+// by IFRT IR passes.
+void InitializeMlirDialectRegistry(mlir::DialectRegistry& registry);
+
 // Registers all dialects required by IFRT IR modules.
 void RegisterMlirDialects(mlir::MLIRContext& context);