the three removed stablehlo passes are necessary.

We patch them back.

The passes are
mhlo-legalize-control-flow
mhlo-legalize-to-std
hlo-legalize-sort

tensorflow/mlir-hlo@4a640be#diff-ef0d7e30da19a396ba036405a9ef636f8b1be194618b0a90f4602671fc2ef34d
tensorflow/mlir-hlo@2a5e267#diff-f8c7cb07b43593403e00e0dbf9983f0186b4eb70368cc99af3b924061f1ea46f

Author: paul0403

frontend/catalyst/pipelines.py

@@ -184,10 +184,10 @@ def get_hlo_lowering_stage(_options: CompileOptions) -> List[str]:
         "canonicalize",
         "func.func(chlo-legalize-to-hlo)",
         "stablehlo-legalize-to-hlo",
-        # "func.func(mhlo-legalize-control-flow)",
+        "func.func(mhlo-legalize-control-flow)",
         "func.func(hlo-legalize-to-linalg)",
-        # "func.func(mhlo-legalize-to-std)",
-        # "func.func(hlo-legalize-sort)",
+        "func.func(mhlo-legalize-to-std)",
+        "func.func(hlo-legalize-sort)",
         "convert-to-signless",
         "canonicalize",
         "scatter-lowering",

mlir/Makefile

@@ -104,6 +104,17 @@ mhlo:
 	@if patch -p0 --dry-run -N < $(MK_DIR)/patches/mhlo-remove-shardy.patch > /dev/null 2>&1; then \
 		patch -p0 < $(MK_DIR)/patches/mhlo-remove-shardy.patch; \
 	fi
+	# Patch MHLO passes removed from upstream
+	mkdir $(MK_DIR)/mlir-hlo/mhlo/transforms/legalize_control_flow
+	mkdir $(MK_DIR)/mlir-hlo/mhlo/transforms/legalize_sort
+	mkdir $(MK_DIR)/mlir-hlo/mhlo/transforms/legalize_to_standard
+	touch $(MK_DIR)/mlir-hlo/mhlo/transforms/legalize_control_flow/legalize_control_flow.cc
+	touch $(MK_DIR)/mlir-hlo/mhlo/transforms/legalize_sort/legalize_sort.cc
+	touch $(MK_DIR)/mlir-hlo/mhlo/transforms/legalize_to_standard/legalize_to_standard.cc
+	touch $(MK_DIR)/mlir-hlo/mhlo/transforms/legalize_to_standard/legalize_to_standard_patterns.td
+	@if patch -p0 --dry-run -N --verbose < $(MK_DIR)/patches/mhlo-add-back-necessary-passes.patch > /dev/null 2>&1; then \
+		patch -p0 < $(MK_DIR)/patches/mhlo-add-back-necessary-passes.patch; \
+	fi
 	cmake -G Ninja -S mlir-hlo -B $(MHLO_BUILD_DIR) \
 		-DCMAKE_BUILD_TYPE=$(BUILD_TYPE) \
 		-DLLVM_ENABLE_ASSERTIONS=ON \

mlir/lib/Driver/Pipelines.cpp

@@ -42,10 +42,10 @@ void createHloLoweringPipeline(OpPassManager &pm)
     pm.addPass(mlir::createCanonicalizerPass());
     pm.addNestedPass<mlir::func::FuncOp>(mhlo::createChloLegalizeToHloPass());
     pm.addPass(mlir::mhlo::createStablehloLegalizeToHloPass());
-    // pm.addNestedPass<mlir::func::FuncOp>(mhlo::createLegalizeControlFlowPass());
+    pm.addNestedPass<mlir::func::FuncOp>(mhlo::createLegalizeControlFlowPass());
     pm.addNestedPass<mlir::func::FuncOp>(mhlo::createLegalizeHloToLinalgPass());
-    // pm.addNestedPass<mlir::func::FuncOp>(mhlo::createLegalizeToStdPass());
-    // pm.addNestedPass<mlir::func::FuncOp>(mhlo::createLegalizeSortPass());
+    pm.addNestedPass<mlir::func::FuncOp>(mhlo::createLegalizeToStdPass());
+    pm.addNestedPass<mlir::func::FuncOp>(mhlo::createLegalizeSortPass());
     pm.addPass(mlir::mhlo::createConvertToSignlessPass());
     pm.addPass(mlir::createCanonicalizerPass());
     pm.addPass(catalyst::createScatterLoweringPass());