Add device compilation static assertions

Author: maleadt

src/GPUCompiler.jl

@@ -77,6 +77,7 @@ include("driver.jl")
 
 # other reusable functionality
 include("execution.jl")
+include("static_assert.jl")
 include("reflection.jl")
 
 include("precompile.jl")

src/static_assert.jl

@@ -0,0 +1,65 @@
+export @static_assert
+
+"""
+    @static_assert condition message
+
+Require `condition` to be proven true while compiling device code. Unlike
+`Base.@static`, the condition is evaluated after target-specific values have
+been propagated through LLVM. A condition that remains dynamic is rejected;
+this is not a runtime assertion mechanism.
+
+`message` must be a string literal. It is reported with the device-code
+backtrace if the assertion cannot be proven.
+"""
+macro static_assert(condition, message)
+    message isa String || throw(ArgumentError("@static_assert message must be a string literal"))
+    marker = static_assert_marker(message)
+    return quote
+        if !$(esc(condition))
+            $marker
+        end
+        nothing
+    end
+end
+
+const STATIC_ASSERT_MARKER = "gpu_static_assert"
+const STATIC_ASSERTION = "static assertion failed"
+
+function static_assert_marker(message::String)
+    LLVM.Context() do _
+        entry, entry_type = LLVM.Interop.create_function()
+        mod = LLVM.parent(entry)
+        @dispose builder=IRBuilder() begin
+            block = BasicBlock(entry, "entry")
+            position!(builder, block)
+
+            # LLVM.jl owns the pointer representation and creates an anonymous private
+            # string global, just like LLVM's annotation helpers.
+            string = globalstring_ptr!(builder, message)
+            marker_type = LLVM.FunctionType(LLVM.VoidType(), [value_type(string)])
+            marker = LLVM.Function(mod, STATIC_ASSERT_MARKER, marker_type)
+            call!(builder, marker_type, marker, [string])
+            ret!(builder)
+        end
+        return LLVM.Interop.call_function(entry, Nothing, Tuple{})
+    end
+end
+
+function static_assert_message(inst::LLVM.CallInst)
+    try
+        value = only(arguments(inst))
+        while value isa ConstantExpr
+            value = first(operands(value))
+        end
+        value isa GlobalVariable || return nothing
+        initializer = LLVM.initializer(value)
+        initializer isa ConstantDataSequential || initializer isa ConstantArray || return nothing
+        values = initializer isa ConstantDataSequential ? collect(initializer) : operands(initializer)
+        bytes = UInt8[convert(UInt8, byte) for byte in values]
+        !isempty(bytes) && bytes[end] == 0x00 && pop!(bytes)
+        return String(bytes)
+    catch err
+        err isa ArgumentError || err isa BoundsError || rethrow()
+        return nothing
+    end
+end

src/validation.jl

@@ -135,12 +135,15 @@ const DYNAMIC_CALL     = "dynamic function invocation"
 function Base.showerror(io::IO, err::InvalidIRError)
     print(io, "InvalidIRError: compiling ", err.job.source, " resulted in invalid LLVM IR")
     for (kind, bt, meta) in err.errors
-        printstyled(io, "\nReason: unsupported $kind"; color=:red)
+        prefix = kind == STATIC_ASSERTION ? "Reason: $kind" : "Reason: unsupported $kind"
+        printstyled(io, "\n$prefix"; color=:red)
         if meta !== nothing
             if kind == RUNTIME_FUNCTION || kind == UNKNOWN_FUNCTION || kind == POINTER_FUNCTION || kind == DYNAMIC_CALL || kind == CCALL_FUNCTION || kind == LAZY_FUNCTION
                 printstyled(io, " (call to ", meta, ")"; color=:red)
             elseif kind == DELAYED_BINDING
                 printstyled(io, " (use of '", meta, "')"; color=:red)
+            elseif kind == STATIC_ASSERTION
+                printstyled(io, " (", meta, ")"; color=:red)
             end
         end
         Base.show_backtrace(io, bt)
@@ -225,7 +228,9 @@ function check_ir!(job, errors::Vector{IRError}, inst::LLVM.CallInst)
         fn = LLVM.name(dest)
 
         # some special handling for runtime functions that we don't implement
-        if fn == "jl_get_binding_or_error" || fn == "ijl_get_binding_or_error"
+        if fn == STATIC_ASSERT_MARKER
+            push!(errors, (STATIC_ASSERTION, bt, static_assert_message(inst)))
+        elseif fn == "jl_get_binding_or_error" || fn == "ijl_get_binding_or_error"
             try
                 m, sym = arguments(inst)
                 sym = first(operands(sym::ConstantExpr))::ConstantInt

test/native.jl

@@ -848,6 +848,79 @@ end
     end
 end
 
+@testset "static assertions" begin
+    mod = @eval module $(gensym())
+        using ..GPUCompiler
+        kernel() = (@static_assert true "this should disappear"; return)
+    end
+
+    llvm = sprint(io -> Native.code_llvm(io, mod.kernel, Tuple{}; dump_module=true))
+    @test !occursin(GPUCompiler.STATIC_ASSERT_MARKER, llvm)
+    @test Native.code_execution(mod.kernel, Tuple{}) !== nothing
+
+    mod = @eval module $(gensym())
+        using ..GPUCompiler
+        kernel() = (@static_assert false "the target is too old"; return)
+    end
+    @test_throws_message(InvalidIRError,
+                         Native.code_execution(mod.kernel, Tuple{})) do msg
+        occursin(GPUCompiler.STATIC_ASSERTION, msg) &&
+        occursin("the target is too old", msg) &&
+        occursin("kernel", msg)
+    end
+
+    mod = @eval module $(gensym())
+        using ..GPUCompiler
+        function kernel(condition)
+            @static_assert condition "condition was not proven"
+            return
+        end
+    end
+    @test_throws_message(InvalidIRError,
+                         Native.code_execution(mod.kernel, Tuple{Bool}; opt_level=0)) do msg
+        occursin(GPUCompiler.STATIC_ASSERTION, msg) &&
+        occursin("condition was not proven", msg)
+    end
+
+    mod = @eval module $(gensym())
+        using ..GPUCompiler
+        function kernel()
+            if false
+                @static_assert false "dead assertion"
+            end
+            return
+        end
+    end
+    @test Native.code_execution(mod.kernel, Tuple{}; opt_level=0) !== nothing
+
+    mod = @eval module $(gensym())
+        using ..GPUCompiler
+        function kernel(condition)
+            @static_assert condition "first failure"
+            @static_assert condition "second failure"
+            return
+        end
+    end
+    @test_throws_message(InvalidIRError,
+                         Native.code_execution(mod.kernel, Tuple{Bool})) do msg
+        occursin("first failure", msg) && occursin("second failure", msg) &&
+        !occursin("unknown function", msg)
+    end
+
+    mod = @eval module $(gensym())
+        using ..GPUCompiler
+        @inline assertion() = @static_assert false "inlined failure"
+        kernel() = (assertion(); return)
+    end
+    @test_throws_message(InvalidIRError,
+                         Native.code_execution(mod.kernel, Tuple{})) do msg
+        occursin("inlined failure", msg) &&
+        occursin("assertion", msg) && occursin("kernel", msg)
+    end
+
+    @test_throws ArgumentError macroexpand(mod, :(@static_assert true string("message")))
+end
+
 @testset "invalid LLVM IR (ccall)" begin
     mod = @eval module $(gensym())
         function foobar(p)