Also dump MLIR module of failed compilation inside IFRT and PJRT (#1021)

Author: giordano

src/mlir/IR/Pass.jl

@@ -66,40 +66,58 @@ end
 
 const DUMP_MLIR_DIR = Ref{Union{Nothing,String}}(nothing)
 
-"""
-    run!(passManager, module)
-
-Run the provided `passManager` on the given `module`.
-"""
-function run!(pm::PassManager, mod::Module)
-    status = LogicalResult(@static if isdefined(API, :mlirPassManagerRunOnOp)
-        API.mlirPassManagerRunOnOp(pm, Operation(mod))
-    else
-        API.mlirPassManagerRun(pm, mod)
-    end)
-    if isfailure(status)
+# Utilities for dumping to a file the module of a failed compilation, useful for
+# debugging purposes.
+function compilation_failed_dump_mlir(mod::Module, pm::Union{Nothing,PassManager}=nothing)
+    try
         # If `DUMP_MLIR_DIR` is `nothing`, create a persistent new temp
         # directory, otherwise use the provided path.
         dir = if isnothing(DUMP_MLIR_DIR[])
+            mkpath(tempdir())
             mktempdir(; prefix="reactant_", cleanup=false)
         else
             DUMP_MLIR_DIR[]
         end
-        try
-            # Make sure the directory exists
-            mkpath(dir)
-            path = tempname(dir; cleanup=false) * ".mlir"
-            open(path, "w") do io
+        # Make sure the directory exists
+        mkpath(dir)
+        path = tempname(dir; cleanup=false) * ".mlir"
+        open(path, "w") do io
+            if !isnothing(pm)
                 println(io, "// Pass pipeline:")
                 print(io, "// ")
                 print_pass_pipeline(io, OpPassManager(pm))
                 println(io)
-                show(IOContext(io, :debug => true), mod)
             end
-            @error "Dumped module to " * path
-        catch err
-            @error "Couldn't save MLIR module" exception = err
+            show(IOContext(io, :debug => true), mod)
         end
+        @error "Compilation failed, MLIR module written to $(path)"
+    catch err
+        @error "Couldn't save MLIR module" exception = err
+    end
+end
+
+function try_compile_dump_mlir(f, mod::Module, pm=nothing)
+    try
+        f()
+    catch
+        compilation_failed_dump_mlir(mod, pm)
+        rethrow()
+    end
+end
+
+"""
+    run!(passManager, module)
+
+Run the provided `passManager` on the given `module`.
+"""
+function run!(pm::PassManager, mod::Module)
+    status = LogicalResult(@static if isdefined(API, :mlirPassManagerRunOnOp)
+        API.mlirPassManagerRunOnOp(pm, Operation(mod))
+    else
+        API.mlirPassManagerRun(pm, mod)
+    end)
+    if isfailure(status)
+        compilation_failed_dump_mlir(mod, pm)
         throw("failed to run pass manager on module")
     end
     return mod

src/xla/IFRT/LoadedExecutable.jl

@@ -85,16 +85,18 @@ function XLA.compile(
 )
     device_id = is_sharded ? Int64(-1) : Int64(XLA.device_ordinal(device))
     GC.@preserve client mod begin
-        exec = @ccall MLIR.API.mlir_c.ifrt_compile(
-            client.client::Ptr{Cvoid},
-            mod.module_::MLIR.API.MlirModule,
-            device_id::Clong,
-            is_sharded::Bool,
-            global_device_ids::Ptr{Clong},
-            length(global_device_ids)::Clong,
-            XLA.CUDA_DATA_DIR[]::Cstring,
-            use_shardy_partitioner::Bool,
-        )::Ptr{Cvoid}
+        exec = MLIR.IR.try_compile_dump_mlir(mod) do
+            @ccall MLIR.API.mlir_c.ifrt_compile(
+                client.client::Ptr{Cvoid},
+                mod.module_::MLIR.API.MlirModule,
+                device_id::Clong,
+                is_sharded::Bool,
+                global_device_ids::Ptr{Clong},
+                length(global_device_ids)::Clong,
+                XLA.CUDA_DATA_DIR[]::Cstring,
+                use_shardy_partitioner::Bool,
+            )::Ptr{Cvoid}
+        end
     end
     return LoadedExecutable(
         exec, num_outputs, num_parameters, is_sharded, num_replicas, num_partitions

src/xla/PJRT/LoadedExecutable.jl

@@ -79,16 +79,18 @@ function XLA.compile(
 )
     device_id = is_sharded ? Int64(-1) : Int64(XLA.device_ordinal(device))
     GC.@preserve client mod begin
-        exec = @ccall MLIR.API.mlir_c.ClientCompile(
-            client.client::Ptr{Cvoid},
-            mod.module_::MLIR.API.MlirModule,
-            device_id::Clong,
-            is_sharded::Bool,
-            global_device_ids::Ptr{Clong},
-            length(global_device_ids)::Clong,
-            XLA.CUDA_DATA_DIR[]::Cstring,
-            use_shardy_partitioner::Bool,
-        )::Ptr{Cvoid}
+        exec = MLIR.IR.try_compile_dump_mlir(mod) do
+            @ccall MLIR.API.mlir_c.ClientCompile(
+                client.client::Ptr{Cvoid},
+                mod.module_::MLIR.API.MlirModule,
+                device_id::Clong,
+                is_sharded::Bool,
+                global_device_ids::Ptr{Clong},
+                length(global_device_ids)::Clong,
+                XLA.CUDA_DATA_DIR[]::Cstring,
+                use_shardy_partitioner::Bool,
+            )::Ptr{Cvoid}
+        end
     end
     return LoadedExecutable(
         exec, num_outputs, num_parameters, is_sharded, num_replicas, num_partitions