Merge pull request #880 from JuliaGPU/tb/static_assert

Add a static assertion macro

Author: maleadt

Project.toml

@@ -1,6 +1,6 @@
 name = "GPUCompiler"
 uuid = "61eb1bfa-7361-4325-ad38-22787b887f55"
-version = "2.0.1"
+version = "2.1.0"
 authors = ["Tim Besard <tim.besard@gmail.com>"]
 
 [workspace]

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/ptx.jl

@@ -9,15 +9,6 @@ const NVPTX_LLVM_Backend_jll =
 
 export PTXCompilerTarget
 
-# Wire-format encoding of the feature set, stamped into the `sm_features` LLVM
-# global by `finish_module!` and read back by host-side runtime intrinsics (e.g.
-# CUDA.jl's `target_feature_set()`).
-@enum TargetFeatureSet::UInt32 begin
-    BaselineFeatures = 0
-    FamilyFeatures   = 1
-    ArchFeatures     = 2
-end
-
 Base.@kwdef struct PTXCompilerTarget <: AbstractCompilerTarget
     cap::VersionNumber
     ptx::VersionNumber = v"6.0" # for compatibility with older versions of CUDA.jl
@@ -151,26 +142,6 @@ function finish_module!(@nospecialize(job::CompilerJob{PTXCompilerTarget}),
             Metadata(ConstantInt(Int32(job.config.target.fastmath ? 1 : 0)))
     end
 
-    # emit the device capability and ptx isa version as constants in the module. this makes
-    # it possible to 'query' these in device code, relying on LLVM to optimize the checks
-    # away and generate static code. note that we only do so if there's actual uses of these
-    # variables; unconditionally creating a gvar would result in duplicate declarations.
-    sm_features = job.config.target.feature_set === :arch    ? ArchFeatures :
-                  job.config.target.feature_set === :family  ? FamilyFeatures :
-                                                               BaselineFeatures
-    for (name, value) in ["sm_major"    => job.config.target.cap.major,
-                          "sm_minor"    => job.config.target.cap.minor,
-                          "sm_features" => UInt32(sm_features),
-                          "ptx_major"   => job.config.target.ptx.major,
-                          "ptx_minor"   => job.config.target.ptx.minor]
-        if haskey(globals(mod), name)
-            gv = globals(mod)[name]
-            initializer!(gv, ConstantInt(LLVM.Int32Type(), value))
-            # change the linkage so that we can inline the value
-            linkage!(gv, LLVM.API.LLVMPrivateLinkage)
-        end
-    end
-
     # update calling convention
     if LLVM.version() >= v"8"
         for f in functions(mod)
@@ -244,14 +215,6 @@ function finish_module!(@nospecialize(job::CompilerJob{PTXCompilerTarget}),
         end
     end
 
-    # we emit properties (of the device and ptx isa) as private global constants,
-    # so run the optimizer so that they are inlined before the rest of the optimizer runs.
-    @dispose pb=NewPMPassBuilder() begin
-        add!(pb, RecomputeGlobalsAAPass())
-        add!(pb, GlobalOptPass())
-        run!(pb, mod, llvm_machine(job.config.target))
-    end
-
     return entry
 end
 

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)