Add support for the LLVM SPIR-V back-end.

src/spirv.jl

@@ -4,21 +4,45 @@
 # https://github.com/KhronosGroup/LLVM-SPIRV-Backend/blob/master/llvm/docs/SPIR-V-Backend.rst
 # https://github.com/KhronosGroup/SPIRV-LLVM-Translator/blob/master/docs/SPIRVRepresentationInLLVM.rst
 
-const SPIRV_LLVM_Translator_unified_jll = LazyModule("SPIRV_LLVM_Translator_unified_jll", UUID("85f0d8ed-5b39-5caa-b1ae-7472de402361"))
-const SPIRV_Tools_jll = LazyModule("SPIRV_Tools_jll", UUID("6ac6d60f-d740-5983-97d7-a4482c0689f4"))
+const SPIRV_LLVM_Backend_jll =
+    LazyModule("SPIRV_LLVM_Backend_jll",
+               UUID("4376b9bf-cff8-51b6-bb48-39421dff0d0c"))
+const SPIRV_LLVM_Translator_unified_jll =
+    LazyModule("SPIRV_LLVM_Translator_unified_jll",
+               UUID("85f0d8ed-5b39-5caa-b1ae-7472de402361"))
+const SPIRV_Tools_jll =
+    LazyModule("SPIRV_Tools_jll",
+               UUID("6ac6d60f-d740-5983-97d7-a4482c0689f4"))
 
 
 ## target
 
 export SPIRVCompilerTarget
 
 Base.@kwdef struct SPIRVCompilerTarget <: AbstractCompilerTarget
+    version::Union{Nothing,VersionNumber} = nothing
     extensions::Vector{String} = []
     supports_fp16::Bool = true
     supports_fp64::Bool = true
+
+    backend::Symbol = isavailable(SPIRV_LLVM_Backend_jll) ? :llvm : :khronos
+    # XXX: these don't really belong in the _target_ struct
+    validate::Bool = false
+    optimize::Bool = false
 end
 
-llvm_triple(::SPIRVCompilerTarget) = Int===Int64 ? "spir64-unknown-unknown" : "spirv-unknown-unknown"
+function llvm_triple(target::SPIRVCompilerTarget)
+    if target.backend == :llvm
+        architecture = Int===Int64 ? "spirv64" : "spirv32"  # could also be "spirv" for logical addressing
+        subarchitecture = target.version === nothing ? "" : "v$(target.version.major).$(target.version.minor)"
+        vendor = "unknown"  # could also be AMD
+        os = "unknown"
+        environment = "unknown"
+        return "$architecture$subarchitecture-$vendor-$os-$environment"
+    elseif target.backend == :khronos
+        return Int===Int64 ? "spir64-unknown-unknown" : "spirv-unknown-unknown"
+    end
+end
 
 # SPIRV is not supported by our LLVM builds, so we can't get a target machine
 llvm_machine(::SPIRVCompilerTarget) = nothing
@@ -32,7 +56,8 @@ llvm_datalayout(::SPIRVCompilerTarget) = Int===Int64 ?
 
 # TODO: encode debug build or not in the compiler job
 #       https://github.com/JuliaGPU/CUDAnative.jl/issues/368
-runtime_slug(job::CompilerJob{SPIRVCompilerTarget}) = "spirv"
+runtime_slug(job::CompilerJob{SPIRVCompilerTarget}) =
+    "spirv-" * String(job.config.target.backend)
 
 function finish_module!(job::CompilerJob{SPIRVCompilerTarget}, mod::LLVM.Module, entry::LLVM.Function)
     # update calling convention
@@ -90,47 +115,57 @@ end
     # (SPIRV-LLVM-Translator#1140)
     rm_freeze!(mod)
 
-
     # translate to SPIR-V
     input = tempname(cleanup=false) * ".bc"
     translated = tempname(cleanup=false) * ".spv"
-    options = `--spirv-debug-info-version=ocl-100`
-    if !isempty(job.config.target.extensions)
-        str = join(map(ext->"+$ext", job.config.target.extensions), ",")
-        options = `$options --spirv-ext=$str`
-    end
     write(input, mod)
-    let cmd = `$(SPIRV_LLVM_Translator_unified_jll.llvm_spirv()) $options -o $translated $input`
-        proc = run(ignorestatus(cmd))
-        if !success(proc)
-            error("""Failed to translate LLVM code to SPIR-V.
-                     If you think this is a bug, please file an issue and attach $(input).""")
+    if job.config.target.backend === :llvm
+        cmd = `$(SPIRV_LLVM_Backend_jll.llc()) $input -filetype=obj -o $translated`
+
+        if !isempty(job.config.target.extensions)
+            str = join(map(ext->"+$ext", job.config.target.extensions), ",")
+            cmd = `$(cmd) -spirv-ext=$str`
+        end
+    elseif job.config.target.backend === :khronos
+        cmd = `$(SPIRV_LLVM_Translator_unified_jll.llvm_spirv()) -o $translated $input --spirv-debug-info-version=ocl-100`
+
+        if !isempty(job.config.target.extensions)
+            str = join(map(ext->"+$ext", job.config.target.extensions), ",")
+            cmd = `$(cmd) --spirv-ext=$str`
+        end
+
+        if job.config.target.version !== nothing
+            cmd = `$(cmd) --spirv-max-version=$(job.config.target.version.major).$(job.config.target.version.minor)`
         end
     end
+    proc = run(ignorestatus(cmd))
+    if !success(proc)
+        error("""Failed to translate LLVM code to SPIR-V.
+                 If you think this is a bug, please file an issue and attach $(input).""")
+    end
 
     # validate
-    # XXX: parameterize this on the `validate` driver argument
-    # XXX: our code currently doesn't pass the validator
-    #if Base.JLOptions().debug_level >= 2
-    #    cmd = `$(SPIRV_Tools_jll.spirv_val()) $translated`
-    #    proc = run(ignorestatus(cmd))
-    #    if !success(proc)
-    #        error("""Failed to validate generated SPIR-V.
-    #                 If you think this is a bug, please file an issue and attach $(input) and $(translated).""")
-    #    end
-    #end
+    if job.config.target.validate
+       cmd = `$(SPIRV_Tools_jll.spirv_val()) $translated`
+       proc = run(ignorestatus(cmd))
+       if !success(proc)
+           error("""Failed to validate generated SPIR-V.
+                    If you think this is a bug, please file an issue and attach $(input) and $(translated).""")
+       end
+    end
 
     # optimize
-    # XXX: parameterize this on the `optimize` driver argument
-    # XXX: the optimizer segfaults on some of our code
     optimized = tempname(cleanup=false) * ".spv"
-    #let cmd = `$(SPIRV_Tools_jll.spirv_opt()) -O --skip-validation $translated -o $optimized`
-    #    proc = run(ignorestatus(cmd))
-    #    if !success(proc)
-    #        error("""Failed to optimize generated SPIR-V.
-    #                 If you think this is a bug, please file an issue and attach $(input) and $(translated).""")
-    #    end
-    #end
+    if job.config.target.optimize
+        cmd = `$(SPIRV_Tools_jll.spirv_opt()) -O --skip-validation $translated -o $optimized`
+       proc = run(ignorestatus(cmd))
+       if !success(proc)
+           error("""Failed to optimize generated SPIR-V.
+                    If you think this is a bug, please file an issue and attach $(input) and $(translated).""")
+       end
+    else
+        cp(translated, optimized)
+    end
 
     output = if format == LLVM.API.LLVMObjectFile
         read(translated)
@@ -141,7 +176,7 @@ end
 
     rm(input)
     rm(translated)
-    #rm(optimized)
+    rm(optimized)
 
     return output
 end

src/utils.jl

@@ -30,6 +30,8 @@ struct LazyModule
     LazyModule(name, uuid) = new(Base.PkgId(uuid, name))
 end
 
+isavailable(lazy_mod::LazyModule) = haskey(Base.loaded_modules, getfield(lazy_mod, :pkg))
+
 function Base.getproperty(lazy_mod::LazyModule, sym::Symbol)
     pkg = getfield(lazy_mod, :pkg)
     mod = get(Base.loaded_modules, pkg, nothing)

test/Project.toml

@@ -9,6 +9,7 @@ PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a"
 Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
 REPL = "3fa0cd96-eef1-5676-8a61-b3b8758bbffb"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
+SPIRV_LLVM_Backend_jll = "4376b9bf-cff8-51b6-bb48-39421dff0d0c"
 SPIRV_LLVM_Translator_unified_jll = "85f0d8ed-5b39-5caa-b1ae-7472de402361"
 SPIRV_Tools_jll = "6ac6d60f-d740-5983-97d7-a4482c0689f4"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"

test/helpers/spirv.jl

@@ -8,9 +8,11 @@ GPUCompiler.runtime_module(::CompilerJob{<:Any,CompilerParams}) = TestRuntime
 
 function create_job(@nospecialize(func), @nospecialize(types);
                    kernel::Bool=false, always_inline=false,
-                   supports_fp16=true, supports_fp64=true, kwargs...)
+                   supports_fp16=true, supports_fp64=true,
+                   backend::Symbol, kwargs...)
     source = methodinstance(typeof(func), Base.to_tuple_type(types), Base.get_world_counter())
-    target = SPIRVCompilerTarget(; supports_fp16, supports_fp64)
+    target = SPIRVCompilerTarget(; backend, validate=true, optimize=true,
+                                   supports_fp16, supports_fp64)
     params = CompilerParams()
     config = CompilerConfig(target, params; kernel, always_inline)
     CompilerJob(source, config), kwargs

test/runtests.jl

@@ -116,10 +116,6 @@ end
         push!(skip_tests, "metal")
     end
 end
-if !(SPIRV_LLVM_Translator_unified_jll.is_available() && SPIRV_Tools_jll.is_available())
-    # SPIRV needs it's tools to be available
-    push!(skip_tests, "spirv")
-end
 if VERSION < v"1.11"
     append!(skip_tests, ["ptx/precompile", "native/precompile"])
 end

test/setup.jl

@@ -1,6 +1,6 @@
 using Distributed, Test, GPUCompiler, LLVM
 
-using SPIRV_LLVM_Translator_unified_jll, SPIRV_Tools_jll
+using SPIRV_LLVM_Backend_jll, SPIRV_LLVM_Translator_unified_jll, SPIRV_Tools_jll
 
 # include all helpers
 include(joinpath(@__DIR__, "helpers", "runtime.jl"))

test/spirv.jl

@@ -1,14 +1,16 @@
+for backend in (:khronos, :llvm)
+
 @testset "IR" begin
 
 @testset "kernel functions" begin
 @testset "calling convention" begin
     kernel() = return
 
-    ir = sprint(io->SPIRV.code_llvm(io, kernel, Tuple{}; dump_module=true))
+    ir = sprint(io->SPIRV.code_llvm(io, kernel, Tuple{}; backend, dump_module=true))
     @test !occursin("spir_kernel", ir)
 
     ir = sprint(io->SPIRV.code_llvm(io, kernel, Tuple{};
-                                    dump_module=true, kernel=true))
+                                    backend, dump_module=true, kernel=true))
     @test occursin("spir_kernel", ir)
 end
 
@@ -18,19 +20,20 @@ end
         kernel(x) = return
     end
 
-    ir = sprint(io->SPIRV.code_llvm(io, mod.kernel, Tuple{Tuple{Int}}))
+    ir = sprint(io->SPIRV.code_llvm(io, mod.kernel, Tuple{Tuple{Int}}; backend))
     @test occursin(r"@\w*kernel\w*\(({ i64 }|\[1 x i64\])\*", ir) ||
           occursin(r"@\w*kernel\w*\(ptr", ir)
 
-    ir = sprint(io->SPIRV.code_llvm(io, mod.kernel, Tuple{Tuple{Int}}; kernel=true))
+    ir = sprint(io->SPIRV.code_llvm(io, mod.kernel, Tuple{Tuple{Int}};
+                                    backend, kernel=true))
     @test occursin(r"@\w*kernel\w*\(.*{ ({ i64 }|\[1 x i64\]) }\*.+byval", ir) ||
           occursin(r"@\w*kernel\w*\(ptr byval", ir)
 end
 
 @testset "byval bug" begin
     # byval added alwaysinline, which could conflict with noinline and fail verification
     @noinline kernel() = return
-    SPIRV.code_llvm(devnull, kernel, Tuple{}; kernel=true)
+    SPIRV.code_llvm(devnull, kernel, Tuple{}; backend, kernel=true)
     @test "We did not crash!" != ""
 end
 end
@@ -44,26 +47,29 @@ end
         end
     end
 
-    ir = sprint(io->SPIRV.code_llvm(io, mod.kernel, Tuple{Ptr{Float16}, Float16}; validate=true))
+    ir = sprint(io->SPIRV.code_llvm(io, mod.kernel, Tuple{Ptr{Float16}, Float16};
+                                    backend, validate=true))
     @test occursin("store half", ir)
 
-    ir = sprint(io->SPIRV.code_llvm(io, mod.kernel, Tuple{Ptr{Float32}, Float32}; validate=true))
+    ir = sprint(io->SPIRV.code_llvm(io, mod.kernel, Tuple{Ptr{Float32}, Float32};
+                                    backend, validate=true))
     @test occursin("store float", ir)
 
-    ir = sprint(io->SPIRV.code_llvm(io, mod.kernel, Tuple{Ptr{Float64}, Float64}; validate=true))
+    ir = sprint(io->SPIRV.code_llvm(io, mod.kernel, Tuple{Ptr{Float64}, Float64};
+                                    backend, validate=true))
     @test occursin("store double", ir)
 
     @test_throws_message(InvalidIRError,
                          SPIRV.code_llvm(devnull, mod.kernel, Tuple{Ptr{Float16}, Float16};
-                                         supports_fp16=false, validate=true)) do msg
+                                         backend, supports_fp16=false, validate=true)) do msg
         occursin("unsupported use of half value", msg) &&
         occursin("[1] unsafe_store!", msg) &&
         occursin("[2] kernel", msg)
     end
 
     @test_throws_message(InvalidIRError,
                          SPIRV.code_llvm(devnull, mod.kernel, Tuple{Ptr{Float64}, Float64};
-                                         supports_fp64=false, validate=true)) do msg
+                                         backend, supports_fp64=false, validate=true)) do msg
         occursin("unsupported use of double value", msg) &&
         occursin("[1] unsafe_store!", msg) &&
         occursin("[2] kernel", msg)
@@ -82,8 +88,10 @@ end
         return
     end
 
-    asm = sprint(io->SPIRV.code_native(io, kernel, Tuple{Bool}; kernel=true))
+    asm = sprint(io->SPIRV.code_native(io, kernel, Tuple{Bool}; backend, kernel=true))
     @test occursin(r"OpFunctionCall %void %(julia|j)_error", asm)
 end
 
 end
+
+end