From 5aff4c624cdacba62f819c3132785e011bea4c05 Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Mon, 3 Nov 2025 18:30:29 -0500 Subject: [PATCH 1/2] fix: add checks for running shardy without raising --- src/Compiler.jl | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/Compiler.jl b/src/Compiler.jl index 0639c142d5..bbc7e7c650 100644 --- a/src/Compiler.jl +++ b/src/Compiler.jl @@ -1738,6 +1738,8 @@ function compile_mlir!( () end + legal_to_run_shardy_passes = compile_options.optimization_passes !== :before_raise + if compile_options.optimization_passes === :all run_pass_pipeline!( mod, @@ -2173,7 +2175,7 @@ function compile_mlir!( # shardy passes use_shardy_partitioner = false result_shardings = missing - if is_sharded + if is_sharded && legal_to_run_shardy_passes module_op = copy(MLIR.IR.Operation(mod)) mod_copied = MLIR.IR.Module(module_op) From 2ef33b875789e290174c856e66a374c7ef7984d2 Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Mon, 3 Nov 2025 18:34:04 -0500 Subject: [PATCH 2/2] fix: only run in the all case --- src/Compiler.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/Compiler.jl b/src/Compiler.jl index bbc7e7c650..bdbbf230f2 100644 --- a/src/Compiler.jl +++ b/src/Compiler.jl @@ -1738,7 +1738,7 @@ function compile_mlir!( () end - legal_to_run_shardy_passes = compile_options.optimization_passes !== :before_raise + legal_to_run_shardy_passes = compile_options.optimization_passes === :all if compile_options.optimization_passes === :all run_pass_pipeline!(