[Compiler] Make raise a keyword argument (#797)

* [Compiler] Make `raise` a boolean keyword argument

* [Compiler] Allow `raise` to be a string too

* Use `task_local_storage` to set whether we're raising or not

* Add super basic tests for `raise` keyword argument

* [Compiler] Make a stack of raising contexts

---------

Co-authored-by: William Moses <wmoses@google.com>

Author: giordano

ext/ReactantCUDAExt.jl

@@ -3,6 +3,7 @@ module ReactantCUDAExt
 using CUDA
 using Reactant:
     Reactant, TracedRArray, AnyTracedRArray, AnyConcretePJRTArray, MLIR, TracedRNumber
+using Reactant.Compiler: raising
 using ReactantCore: @trace
 using GPUCompiler: GPUCompiler
 using KernelAbstractions: KernelAbstractions
@@ -294,7 +295,7 @@ function ka_with_reactant(ndrange, workgroupsize, obj, args...)
 
     # figure out the optimal workgroupsize automatically
     if KA.workgroupsize(obj) <: KA.DynamicSize && workgroupsize === nothing
-        if !Reactant.Compiler.PartitionKA[] || Reactant.Compiler.Raise[]
+        if !Reactant.Compiler.PartitionKA[] || raising()
             threads = prod(ndrange)
         else
             config = CUDA.launch_configuration(kernel.fun; max_threads=prod(ndrange))

src/Compiler.jl

@@ -578,7 +578,37 @@ end
 const DEBUG_KERNEL = Ref{Bool}(false)
 const DUMP_LLVMIR = Ref{Bool}(false)
 
-const Raise = Ref{Bool}(false)
+function activate_raising!(is_raising::Bool)
+    stack = get!(task_local_storage(), :reactant_is_raising) do
+        Bool[]
+    end
+    push!(stack, is_raising)
+    return nothing
+end
+
+function deactivate_raising!(is_raising::Bool)
+    key = :reactant_is_raising
+    is_raising === last(task_local_storage(key)) ||
+        error("Deactivating wrong Reactant raising context")
+    return pop!(task_local_storage(key))
+end
+
+function raising(; throw_error::Bool=true)
+    key = :reactant_is_raising
+    if !(haskey(task_local_storage(), key) && !Base.isempty(task_local_storage(key)))
+        throw_error && error("No Reactant raising context")
+    end
+    return last(task_local_storage(key)::Vector{Bool})
+end
+
+function raising!(f, is_raising::Bool)
+    activate_raising!(is_raising)
+    try
+        return f()
+    finally
+        deactivate_raising!(is_raising)
+    end
+end
 
 function compile_mlir!(
     mod,
@@ -605,6 +635,7 @@ function compile_mlir!(
     no_nan::Bool=false,
     backend="gpu",
     fn_kwargs=(),
+    raise::Union{Bool,String}=false,
 )
     # Explicitly don't use block! to avoid creating a closure, which creates
     # both compile-time and relocatability issues
@@ -614,9 +645,16 @@ function compile_mlir!(
     activate_callcache!(callcache)
     activate_sdycache!(sdycache)
 
+    # Save in the TLS whether we are raising.  We identify that condition by
+    # checking whether the user set an explicit list of passes, or chose
+    # `raise=true` to use the default passes.
+    is_raising = raise isa String || raise
+    activate_raising!(is_raising)
+
     mlir_fn_res = try
         Reactant.TracedUtils.make_mlir_fn(f, args, fn_kwargs, "main", true)
     finally
+        deactivate_raising!(is_raising)
         deactivate_sdycache!(sdycache)
         deactivate_callcache!(callcache)
         MLIR.IR.deactivate!(MLIR.IR.body(mod))
@@ -648,14 +686,14 @@ function compile_mlir!(
         )
         @assert curesulthandler !== nothing
         curesulthandler = Base.reinterpret(UInt, curesulthandler)
-        kern = if Raise[]
+        kern = if is_raising
             "lower-kernel{backend=cpu},symbol-dce,canonicalize"
         else
             "lower-kernel,canonicalize"
         end
         jit = "lower-jit{debug=true cuResultHandlerPtr=$curesulthandler cuOptLevel=$(cuOptLevel[]) cubinFormat=$(cubinFormat[]) indexBitWidth=$(cuindexBitWidth[])  cubinChip=$(cubinChip[]) cubinFeatures=$(cubinFeatures()) run_init=true toolkitPath=$toolkit},symbol-dce"
     else
-        kern = if Raise[]
+        kern = if is_raising
             "lower-kernel{backend=cpu},symbol-dce,canonicalize"
         else
             "lower-kernel,canonicalize"
@@ -666,8 +704,13 @@ function compile_mlir!(
     opt_passes = optimization_passes(; no_nan, sroa=true)
     opt_passes2 = optimization_passes(; no_nan, sroa=false)
 
-    raise = if Raise[]
-	"canonicalize,llvm-to-memref-access,canonicalize,convert-llvm-to-cf,canonicalize,enzyme-lift-cf-to-scf,canonicalize,func.func(canonicalize-loops),canonicalize-scf-for,canonicalize,affine-cfg,canonicalize,func.func(canonicalize-loops),canonicalize,llvm-to-affine-access,canonicalize,delinearize-indexing,canonicalize,raise-affine-to-stablehlo,arith-raise{stablehlo=true}," * opt_passes2   
+    raise_passes = if raise isa String
+        # Raising passes were specified
+        raise
+    elseif raise
+        # Raise enabled but use default passes
+        "canonicalize,llvm-to-memref-access,canonicalize,convert-llvm-to-cf,canonicalize,enzyme-lift-cf-to-scf,canonicalize,func.func(canonicalize-loops),canonicalize-scf-for,canonicalize,affine-cfg,canonicalize,func.func(canonicalize-loops),canonicalize,llvm-to-affine-access,canonicalize,delinearize-indexing,canonicalize,raise-affine-to-stablehlo,arith-raise{stablehlo=true}," *
+        opt_passes2
     else
         "canonicalize"
     end
@@ -686,7 +729,7 @@ function compile_mlir!(
                     "enzyme-simplify-math",
                     opt_passes2,
                     kern,
-                    raise,
+                    raise_passes,
                     jit,
                 ],
                 ',',
@@ -723,7 +766,7 @@ function compile_mlir!(
                     "enzyme-simplify-math",
                     opt_passes2,
                     kern,
-                    raise,
+                    raise_passes,
                 ],
                 ',',
             ),
@@ -787,7 +830,7 @@ function compile_mlir!(
                     "enzyme-simplify-math",
                     opt_passes2,
                     kern,
-                    raise,
+                    raise_passes,
                     jit,
                 ],
                 ',',
@@ -804,7 +847,7 @@ function compile_mlir!(
                 [
                     "canonicalize,remove-unnecessary-enzyme-ops,enzyme-simplify-math",
                     kern,
-                    raise,
+                    raise_passes,
                     jit,
                 ],
                 ',',
@@ -891,7 +934,7 @@ See also [`@code_xla`](@ref), [`@code_mhlo`](@ref).
 """
 macro code_hlo(args...)
     default_options = Dict{Symbol,Any}(
-        :optimize => true, :no_nan => false, :client => nothing
+        :optimize => true, :no_nan => false, :client => nothing, :raise => false
     )
     compile_expr, (; compiled) = compile_call_expr(
         __module__, compile_mlir, default_options, args...
@@ -915,7 +958,7 @@ See also [`@code_xla`](@ref), [`@code_hlo`](@ref).
 """
 macro code_mhlo(args...)
     default_options = Dict{Symbol,Any}(
-        :optimize => true, :no_nan => false, :client => nothing
+        :optimize => true, :no_nan => false, :client => nothing, :raise => false
     )
     compile_expr, (; compiled) = compile_call_expr(
         __module__, compile_xla, default_options, args...
@@ -939,7 +982,7 @@ See also [`@code_mhlo`](@ref), [`@code_hlo`](@ref).
 """
 macro code_xla(args...)
     default_options = Dict{Symbol,Any}(
-        :optimize => true, :no_nan => false, :client => nothing
+        :optimize => true, :no_nan => false, :client => nothing, :raise => false
     )
     compile_expr, (; compiled) = compile_call_expr(
         __module__, compile_xla, default_options, args...
@@ -961,7 +1004,11 @@ end
 """
 macro compile(args...)
     default_options = Dict{Symbol,Any}(
-        :optimize => true, :sync => false, :no_nan => false, :client => nothing
+        :optimize => true,
+        :sync => false,
+        :no_nan => false,
+        :client => nothing,
+        :raise => false,
     )
     return esc(first(compile_call_expr(__module__, compile, default_options, args...)))
 end
@@ -973,7 +1020,11 @@ Run @compile f(args..) then immediately execute it
 """
 macro jit(args...)
     default_options = Dict{Symbol,Any}(
-        :optimize => true, :sync => false, :no_nan => false, :client => nothing
+        :optimize => true,
+        :sync => false,
+        :no_nan => false,
+        :client => nothing,
+        :raise => false,
     )
     compile_expr, (; compiled, args) = compile_call_expr(
         __module__, compile, default_options, args...
@@ -988,14 +1039,14 @@ macro jit(args...)
     #! format: on
 end
 
-function compile_call_expr(mod, compiler, options, args...)
+function compile_call_expr(mod, compiler, options::Dict, args...)
     while length(args) > 1
         option, args = args[1], args[2:end]
         if !Meta.isexpr(option, :(=))
             error("Invalid option $(option)")
         else
             option_name = option.args[1]
-            @assert haskey(options, option_name) "Invalid option $(option_name)"
+            @assert haskey(options, option_name) "Invalid option name '$(option_name)'. Valid options are $(join(keys(options), ", "))"
             options[option_name] = option.args[2]
         end
     end

test/basic.jl

@@ -920,3 +920,14 @@ end
     @test contains(hlo, "HloModule")
     @test contains(hlo, "sine")
 end
+
+@testset "Raise keyword" begin
+    v = randn(Float32, 16)
+    rv = Reactant.to_rarray(v)
+    @test sin.(v) ≈ @jit raise = true sin.(rv)
+    @test cos.(v) ≈ @jit raise = false cos.(rv)
+    @test exp.(v) ≈ @jit raise = "canonicalize" exp.(rv)
+    @test_throws Reactant.MLIR.IR.AddPipelineException @jit raise = "this_pass-does_not_ExisT" exp.(
+        rv
+    )
+end