Metal: Drive the downgrader here, lower intrinsics during mcgen. (#829)

Author: maleadt

Project.toml

@@ -1,6 +1,6 @@
 name = "GPUCompiler"
 uuid = "61eb1bfa-7361-4325-ad38-22787b887f55"
-version = "1.16.1"
+version = "1.17.0"
 authors = ["Tim Besard <tim.besard@gmail.com>"]
 
 [workspace]
@@ -20,10 +20,14 @@ TOML = "fa267f1f-6049-4f14-aa54-33bafae1ed76"
 Tracy = "e689c965-62c8-4b79-b2c5-8359227902fd"
 UUIDs = "cf7118a7-6976-5b1a-9a39-7adc72f591a4"
 
+[weakdeps]
+LLVMDowngrader_jll = "f52de702-fb25-5922-94ba-81dd59b07444"
+
 [compat]
 ExprTools = "0.1"
 InteractiveUtils = "1"
 LLVM = "9.8.1"
+LLVMDowngrader_jll = "0.7"
 Libdl = "1"
 Logging = "1"
 PrecompileTools = "1"

src/metal.jl

@@ -1,5 +1,9 @@
 # implementation of the GPUCompiler interfaces for generating Metal code
 
+const LLVMDowngrader_jll =
+    LazyModule("LLVMDowngrader_jll",
+               UUID("f52de702-fb25-5922-94ba-81dd59b07444"))
+
 
 ## target info
 
@@ -64,7 +68,8 @@ function Base.hash(target::MetalCompilerTarget, h::UInt)
     h = hash(target.fastmath, h)
 end
 
-source_code(target::MetalCompilerTarget) = "text"
+# the canonical text representation is AIR assembly, i.e., LLVM 5 era textual IR
+source_code(target::MetalCompilerTarget) = "llvm"
 
 # Metal is not supported by our LLVM builds, so we can't get a target machine
 llvm_machine(::MetalCompilerTarget) = nothing
@@ -295,6 +300,17 @@ function finish_ir!(@nospecialize(job::CompilerJob{MetalCompilerTarget}), mod::L
         add_module_metadata!(job, mod)
     end
 
+    return functions(mod)[entry_fn]
+end
+
+# lowering of LLVM IR to AIR-compatible IR
+#
+# Metal does not have an LLVM back-end, so the lowering of target-independent LLVM IR into
+# target-specific constructs -- something that normally happens during instruction
+# selection -- is implemented here as IR-to-IR rewrites, run at the start of `mcgen`.
+# this keeps the `:llvm` output (e.g. `code_llvm`) close to what Julia generated, using
+# generic LLVM intrinsics, while the `:asm`/`:obj` outputs contain AIR intrinsics.
+function lower_air!(@nospecialize(job::CompilerJob{MetalCompilerTarget}), mod::LLVM.Module)
     # strip device-side `trap`s and rewrite `unreachable` into clean returns (#433, #370). this
     # runs post-`optimize!`, after the trap has finished serving as the optimizer guard; the pass
     # force-inlines throwing functions into the kernel first so the rewrite is sound, then scrubs
@@ -316,13 +332,15 @@ function finish_ir!(@nospecialize(job::CompilerJob{MetalCompilerTarget}), mod::L
         changed |= lower_llvm_intrinsics!(job, f)
     end
     if changed
-        # lowering may have introduced additional functions marked `alwaysinline`
+        # lowering may have introduced additional functions marked `alwaysinline`,
+        # and left dead declarations of the replaced LLVM intrinsics behind
         @dispose pb=NewPMPassBuilder() begin
             add!(pb, AlwaysInlinerPass())
             add!(pb, NewPMFunctionPassManager()) do fpm
                 add!(fpm, SimplifyCFGPass())
                 add!(fpm, instcombine_pass(job))
             end
+            add!(pb, StripDeadPrototypesPass())
             run!(pb, mod)
         end
     end
@@ -336,13 +354,28 @@ function finish_ir!(@nospecialize(job::CompilerJob{MetalCompilerTarget}), mod::L
         end
     end
 
-    return functions(mod)[entry_fn]
+    return
 end
 
 @unlocked function mcgen(job::CompilerJob{MetalCompilerTarget}, mod::LLVM.Module,
                          format=LLVM.API.LLVMObjectFile)
-    # our LLVM version does not support emitting Metal libraries
-    return nothing
+    # lower LLVM constructs that the AIR back-end does not support; this takes the place
+    # of instruction selection, as our LLVM does not have a Metal target machine.
+    lower_air!(job, mod)
+
+    if !isavailable(LLVMDowngrader_jll)
+        error("Metal machine-code generation requires the LLVMDowngrader_jll package, which should be installed and loaded first.")
+    end
+
+    # assemble to AIR, i.e., LLVM 5 era bitcode, as consumed by the metallib loader
+    air = run_tool(`$(LLVMDowngrader_jll.llvm_as()) --bitcode-version=5.0 -o -`, string(mod))
+
+    if format == LLVM.API.LLVMAssemblyFile
+        # disassemble the bitcode again to AIR assembly, i.e., LLVM 5 era textual IR
+        String(run_tool(`$(LLVMDowngrader_jll.llvm_dis_5()) -o -`, air))
+    else
+        air
+    end
 end
 
 

src/utils.jl

@@ -37,6 +37,37 @@ function Base.getproperty(lazy_mod::LazyModule, sym::Symbol)
 end
 
 
+## external tools
+
+# run an external tool (e.g. from a JLL package), feeding `input` to its standard input
+# and returning its standard output. throws on failure, including the tool's standard
+# error output in the exception. for tools that instead communicate through files, e.g.,
+# because the inputs should be preserved for error reporting, use `run` directly.
+function run_tool(cmd::Cmd, input)
+    stdin_pipe = Pipe()
+    stdout_pipe = Pipe()
+    stderr_pipe = Pipe()
+
+    proc = run(pipeline(cmd; stdin=stdin_pipe, stdout=stdout_pipe, stderr=stderr_pipe);
+               wait=false)
+    close(stdout_pipe.in)
+    close(stderr_pipe.in)
+
+    writer = @async begin
+        write(stdin_pipe, input)
+        close(stdin_pipe)
+    end
+    reader = @async read(stdout_pipe)
+    logger = @async read(stderr_pipe, String)
+
+    wait(proc)
+    if !success(proc)
+        error("Failed to run $(basename(cmd.exec[1])):\n" * fetch(logger))
+    end
+    fetch(reader)
+end
+
+
 ## safe logging
 
 using Logging

test/Project.toml

@@ -4,6 +4,7 @@ FileCheck = "4e644321-382b-4b05-b0b6-5d23c3d944fb"
 GPUCompiler = "61eb1bfa-7361-4325-ad38-22787b887f55"
 InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240"
 LLVM = "929cbde3-209d-540e-8aea-75f648917ca0"
+LLVMDowngrader_jll = "f52de702-fb25-5922-94ba-81dd59b07444"
 LLVM_jll = "86de99a1-58d6-5da7-8064-bd56ce2e322c"
 Logging = "56ddb016-857b-54e1-b83d-db4d58db5568"
 ParallelTestRunner = "d3525ed8-44d0-4b2c-a655-542cee43accc"

test/metal.jl

@@ -102,11 +102,18 @@ end
                           (NTuple{2, VecElement{Int64}}, NTuple{2, VecElement{Int64}}), x, y)
     end
 
+    # intrinsics survive the `:llvm` stage as-is; they are only lowered to AIR during mcgen
     @test @filecheck begin
         @check_label "define <2 x i64> @{{(julia|j)_foo_[0-9]+}}"
-        @check "air.max.s.v2i64"
+        @check "llvm.smax.v2i64"
         Metal.code_llvm(mod.foo, (NTuple{2, VecElement{Int64}}, NTuple{2, VecElement{Int64}}))
     end
+
+    @test @filecheck begin
+        @check_label "define <2 x i64> @{{(julia|j)_foo_[0-9]+}}"
+        @check "air.max.s.v2i64"
+        Metal.code_native(mod.foo, (NTuple{2, VecElement{Int64}}, NTuple{2, VecElement{Int64}}))
+    end
 end
 
 @testset "unsupported type detection" begin

test/runtests.jl

@@ -4,6 +4,7 @@ import GPUCompiler, LLVM
 const init_code = quote
     using GPUCompiler, LLVM
     using SPIRV_LLVM_Backend_jll, SPIRV_LLVM_Translator_jll, SPIRV_Tools_jll
+    using LLVMDowngrader_jll
 
     # include all helpers
     include(joinpath(@__DIR__, "helpers", "runtime.jl"))