Use the SPIR-V translator as an in-process library

Project.toml

@@ -1,6 +1,6 @@
 name = "GPUCompiler"
 uuid = "61eb1bfa-7361-4325-ad38-22787b887f55"
-version = "2.7.0"
+version = "2.8.0"
 authors = ["Tim Besard <tim.besard@gmail.com>"]
 
 [workspace]
@@ -31,6 +31,7 @@ Highlights = "eafb193a-b7ab-5a9e-9068-77385905fa72"
 LLVMDowngrader_jll = "f52de702-fb25-5922-94ba-81dd59b07444"
 NVPTX_LLVM_Backend_jll = "ef6e0fe3-e6ef-59c0-bde6-4989574699e0"
 SPIRV_LLVM_Backend_jll = "4376b9bf-cff8-51b6-bb48-39421dff0d0c"
+SPIRV_LLVM_Translator_jll = "4a5d46fc-d8cf-5151-a261-86b458210efb"
 
 [extensions]
 HighlightsExt = "Highlights"
@@ -42,14 +43,15 @@ ExprTools = "0.1"
 Highlights = "0.6"
 InteractiveUtils = "1"
 LLVM = "9.13"
-LLVMDowngrader_jll = "0.10"
+LLVMDowngrader_jll = "0.11"
 Libdl = "1"
 Logging = "1"
 NVPTX_LLVM_Backend_jll = "23"
 PrecompileTools = "1.0.2"
 Preferences = "1"
 REPL = "1"
 SPIRV_LLVM_Backend_jll = "23"
+SPIRV_LLVM_Translator_jll = "23"
 ScopedValues = "1.5"
 TOML = "1"
 Tracy = "0.1.4"

src/spirv.jl

@@ -7,9 +7,6 @@
 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_LLVM_Translator_jll =
     LazyModule("SPIRV_LLVM_Translator_jll",
                UUID("4a5d46fc-d8cf-5151-a261-86b458210efb"))
@@ -156,6 +153,31 @@ struct SPIRVCompileOptions
     opt_level::Cint
 end
 
+# mirrors `LLVMSPIRVTranslateOptions` from libllvm_spirv.h
+struct LLVMSPIRVTranslateOptions
+    max_version_major::Cuint
+    max_version_minor::Cuint
+    extensions::Cstring
+    debug_info_version::Cint    # LLVMSPIRVDebugInfoVersion
+end
+const LLVMSPIRVDebugInfoOpenCL100 = Cint(1)
+
+# translate bitcode to SPIR-V through libllvm_spirv. unlike the back-ends' `Compile`, the
+# translator's entry point takes no diagnostic handler, so this doesn't use
+# `external_compile`.
+function translate(input::Vector{UInt8}, options::Ref{LLVMSPIRVTranslateOptions})
+    backend = ExternalBackend(SPIRV_LLVM_Translator_jll.libllvm_spirv, "LLVMSPIRV")
+    buffer = Ref{Ptr{Cvoid}}(C_NULL)
+    message = Ref{Cstring}(C_NULL)
+    status = @ccall $(api(backend, "Translate"))(input::Ptr{UInt8}, length(input)::Csize_t,
+                                                 options::Ptr{Cvoid},
+                                                 buffer::Ptr{Ptr{Cvoid}},
+                                                 message::Ptr{Cstring})::Cint
+    external_result(backend, status, "Failed to translate LLVM code to SPIR-V",
+                    message[], String[], input)
+    return take_buffer(backend, buffer[])
+end
+
 @unlocked function mcgen(job::CompilerJob{SPIRVCompilerTarget}, mod::LLVM.Module,
                          format=LLVM.API.LLVMAssemblyFile)
     target = job.config.target
@@ -187,34 +209,16 @@ end
                              "Failed to compile to SPIR-V with the SPIR-V back-end")
         end
     elseif target.backend === :khronos
-        translator = if isavailable(SPIRV_LLVM_Translator_jll)
-            SPIRV_LLVM_Translator_jll.llvm_spirv()
-        elseif isavailable(SPIRV_LLVM_Translator_unified_jll)
-            SPIRV_LLVM_Translator_unified_jll.llvm_spirv()
-        else
-            error("This functionality requires the SPIRV_LLVM_Translator_jll or SPIRV_LLVM_Translator_unified_jll package, which should be installed and loaded first.")
-        end
-        input_path = dump_input()
-        translated = tempname(cleanup=false) * ".spv"
-        cmd = `$translator -o $translated $input_path --spirv-debug-info-version=ocl-100`
-
-        if !isempty(target.extensions)
-            cmd = `$(cmd) --spirv-ext=$(target.extensions)`
-        end
-
-        if target.version !== nothing
-            cmd = `$(cmd) --spirv-max-version=$(target.version.major).$(target.version.minor)`
-        end
-        try
-            run(cmd)
-        catch e
-            error("""Failed to translate LLVM code to SPIR-V.
-                     If you think this is a bug, please file an issue and attach $(input_path).""")
+        # translate in-process through libllvm_spirv. like `llvm-spirv`, the translator
+        # takes the triple from the module and rejects unknown extensions.
+        version = something(target.version, v"0.0")
+        extensions = target.extensions
+        GC.@preserve extensions begin
+            options = Ref(LLVMSPIRVTranslateOptions(version.major, version.minor,
+                                                    Base.unsafe_convert(Cstring, extensions),
+                                                    LLVMSPIRVDebugInfoOpenCL100))
+            translate(input, options)
         end
-        rm(input_path)
-        code = read(translated)
-        rm(translated)
-        code
     else
         error("Unsupported SPIR-V back-end $(repr(target.backend)); expected :llvm or :khronos.")
     end

src/utils.jl

@@ -61,12 +61,13 @@ end
 
 ## external back-ends
 
-# The LLVM back-ends for PTX, GCN and SPIR-V and the bitcode downgrader ship as
-# symbol-hidden shared libraries with a small C API modelled after llvm-c: status as the
-# return value, an error message through an out-pointer on failure, results as opaque
-# memory buffers, and diagnostics that the tools used to print to stderr delivered through
-# a callback. The libraries share the shape of this API, differing only in the prefix of
-# their entry points, so this is implemented once against function pointers.
+# The LLVM back-ends for PTX, GCN and SPIR-V, the Khronos SPIR-V translator and the
+# bitcode downgrader ship as symbol-hidden shared libraries with a small C API modelled
+# after llvm-c: status as the return value, an error message through an out-pointer on
+# failure, results as opaque memory buffers, and diagnostics that the tools used to print
+# to stderr delivered through a callback. The libraries share the shape of this API,
+# differing only in the prefix of their entry points, so this is implemented once against
+# function pointers.
 
 struct ExternalBackend
     library::String     # path to the shared library, as exported by its JLL