Merge pull request #524 from JuliaGPU/tb/ptx_trap

PTX: Improve handling of trap

Author: maleadt

Project.toml

@@ -1,7 +1,7 @@
 name = "GPUCompiler"
 uuid = "61eb1bfa-7361-4325-ad38-22787b887f55"
 authors = ["Tim Besard <[email protected]>"]
-version = "0.24.5"
+version = "0.25.0"
 
 [deps]
 ExprTools = "e2ba6199-217a-4e67-a87a-7c52f15ade04"

src/precompile.jl

@@ -52,7 +52,6 @@ end
 function _precompile_()
     ccall(:jl_generating_output, Cint, ()) == 1 || return nothing
     @assert precompile(Tuple{typeof(GPUCompiler.assign_args!),Expr,Vector{Any}})
-    @assert precompile(Tuple{typeof(GPUCompiler.lower_trap!),LLVM.Module})
     @assert precompile(Tuple{typeof(GPUCompiler.lower_unreachable!),LLVM.Function})
     @assert precompile(Tuple{typeof(GPUCompiler.lower_gc_frame!),LLVM.Function})
     @assert precompile(Tuple{typeof(GPUCompiler.lower_throw!),LLVM.Module})

src/ptx.jl

@@ -187,7 +187,6 @@ end
 
 function finish_ir!(@nospecialize(job::CompilerJob{PTXCompilerTarget}),
                     mod::LLVM.Module, entry::LLVM.Function)
-    lower_trap!(mod)
     for f in functions(mod)
         lower_unreachable!(f)
     end
@@ -246,36 +245,6 @@ end
 
 ## LLVM passes
 
-# replace calls to `trap` with inline assembly calling `exit`, which isn't fatal
-function lower_trap!(mod::LLVM.Module)
-    job = current_job::CompilerJob
-    changed = false
-    @timeit_debug to "lower trap" begin
-
-    if haskey(functions(mod), "llvm.trap")
-        trap = functions(mod)["llvm.trap"]
-
-        # inline assembly to exit a thread
-        exit_ft = LLVM.FunctionType(LLVM.VoidType())
-        exit = InlineAsm(exit_ft, "exit;", "", true)
-
-        for use in uses(trap)
-            val = user(use)
-            if isa(val, LLVM.CallInst)
-                @dispose builder=IRBuilder() begin
-                    position!(builder, val)
-                    call!(builder, exit_ft, exit)
-                end
-                unsafe_delete!(LLVM.parent(val), val)
-                changed = true
-            end
-        end
-    end
-
-    end
-    return changed
-end
-
 # lower `unreachable` to `exit` so that the emitted PTX has correct control flow
 #
 # During back-end compilation, `ptxas` inserts instructions to manage the harware's
@@ -328,10 +297,14 @@ end
 # `bar.sync` cannot be executed divergently on Pascal hardware or earlier.
 #
 # To avoid these fall-through successors that change the control flow,
-# we replace `unreachable` instructions with a call to `exit`. This informs
-# `ptxas` that the thread exits, and allows it to correctly construct a CFG,
-# and consequently correctly determine the divergence regions as intended.
+# we replace `unreachable` instructions with a call to `trap` and `exit`. This
+# informs `ptxas` that the thread exits, and allows it to correctly construct a
+# CFG, and consequently correctly determine the divergence regions as intended.
+# Note that we first emit a call to `trap`, so that the behaviour is the same
+# as before.
 function lower_unreachable!(f::LLVM.Function)
+    mod = LLVM.parent(f)
+
     # TODO:
     # - if unreachable blocks have been merged, we still may be jumping from different
     #   divergent regions, potentially causing the same problem as above:
@@ -375,6 +348,12 @@ function lower_unreachable!(f::LLVM.Function)
     # inline assembly to exit a thread
     exit_ft = LLVM.FunctionType(LLVM.VoidType())
     exit = InlineAsm(exit_ft, "exit;", "", true)
+    trap_ft = LLVM.FunctionType(LLVM.VoidType())
+    trap = if haskey(functions(mod), "llvm.trap")
+        functions(mod)["llvm.trap"]
+    else
+        LLVM.Function(mod, "llvm.trap", trap_ft)
+    end
 
     # rewrite the unreachable terminators
     @dispose builder=IRBuilder() begin
@@ -384,6 +363,7 @@ function lower_unreachable!(f::LLVM.Function)
             @assert inst isa LLVM.UnreachableInst
 
             position!(builder, inst)
+            call!(builder, trap_ft, trap)
             call!(builder, exit_ft, exit)
         end
     end