Emit trap before exit when lowering unreachable instructions.

We only insert exit to correct the CFG, but shouldn't execute it.
Make sure to insert trap first so that the behavior of the unreachable
remains the same.

Author: maleadt

src/ptx.jl

@@ -297,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:
@@ -344,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
@@ -353,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