Allow controlling compilation target versions. (#430)

Author: maleadt

src/compiler/compilation.jl

@@ -75,15 +75,22 @@ function compiler_config(dev; kwargs...)
     end
     return config
 end
-@noinline function _compiler_config(dev; kernel=true, name=nothing, always_inline=false, kwargs...)
-    # TODO: configure the compiler target based on the device
-
-    macos = macos_version()
-    metal = metal_support()
-    # we support down to macOS 13, which supports AIR 2.5
-    # so always target that version for now
-    air = v"2.5"
-    @assert air <= air_support()
+@noinline function _compiler_config(dev; kernel=true, name=nothing, always_inline=false,
+                                         macos=nothing, air=nothing, metal=nothing,
+                                         kwargs...)
+    # determine the versions of things to target
+    if macos === nothing
+        macos = macos_version()
+    end
+    if metal === nothing
+        metal = metal_support()
+    end
+    if air === nothing
+        # we support down to macOS 13, which supports AIR 2.5
+        # so always target that version for now
+        air = v"2.5"
+        @assert air <= air_support()
+    end
 
     # create GPUCompiler objects
     target = MetalCompilerTarget(; macos, air, metal, kwargs...)

src/compiler/execution.jl

@@ -4,7 +4,7 @@ export @metal
 ## high-level @metal interface
 
 const MACRO_KWARGS = [:launch]
-const COMPILER_KWARGS = [:kernel, :name, :always_inline]
+const COMPILER_KWARGS = [:kernel, :name, :always_inline, :macos, :air, :metal]
 const LAUNCH_KWARGS = [:groups, :threads, :queue]
 
 """
@@ -171,6 +171,10 @@ const mtlfunction_lock = ReentrantLock()
 Low-level interface to compile a function invocation for the currently-active GPU, returning
 a callable kernel object. For a higher-level interface, use [`@metal`](@ref).
 
+The following keyword arguments are supported:
+- `macos`, `metal` and `air`: to override the macOS OS, Metal language and AIR bitcode
+   versions used during compilation. Value should be a valid version number.
+
 The output of this function is automatically cached, i.e. you can simply call `mtlfunction`
 in a hot path without degrading performance. New code will be generated automatically when
 the function changes, or when different types or keyword arguments are provided.

src/compiler/reflection.jl

@@ -27,7 +27,7 @@ signature to `io` which defaults to `stdout`.
 
 See also: [`@device_code_agx`](@ref)
 """
-function code_agx(io::IO, @nospecialize(func), @nospecialize(types),
+function code_agx(io::IO, @nospecialize(func::Base.Callable), @nospecialize(types),
                   kernel::Bool=true; kwargs...)
     compiler_kwargs, kwargs = split_kwargs_runtime(kwargs, COMPILER_KWARGS)
     source = methodinstance(typeof(func), Base.to_tuple_type(types))
@@ -170,7 +170,7 @@ function disassemble(path)
     return String(take!(io))
 end
 
-code_agx(@nospecialize(func), @nospecialize(types); kwargs...) =
+code_agx(@nospecialize(func::Base.Callable), @nospecialize(types); kwargs...) =
     code_agx(stdout, func, types; kwargs...)
 
 const code_native = code_agx

test/execution.jl

@@ -21,6 +21,10 @@ dummy() = return
     @metal groups=(1,1) dummy()
     @metal groups=(1,1,1) dummy()
 
+    @metal macos=Metal.macos_version() dummy()
+    @metal metal=Metal.metal_support() dummy()
+    @metal air=Metal.air_support() dummy()
+
     @test_throws InexactError @metal threads=(-2) dummy()
     @test_throws InexactError @metal groups=(-2) dummy()
     @test_throws ArgumentError @metal threads=(1025) dummy()
@@ -95,6 +99,21 @@ end
         @metal name="mykernel" dummy()
     end)))
 
+    # set macOS, AIR, and Metal versions
+    let
+        @test occursin("""!{!"Metal", i32 3, i32 2, i32 1}""",
+                       sprint(io->Metal.code_llvm(io, dummy, Tuple{}; metal=v"3.2.1",
+                                                      dump_module=true, kernel=true)))
+
+        @test occursin("!{i32 3, i32 2, i32 1}",
+                       sprint(io->Metal.code_llvm(io, dummy, Tuple{}; air=v"3.2.1",
+                                                      dump_module=true, kernel=true)))
+
+        @test occursin("""!"SDK Version", [3 x i32] [i32 3, i32 2, i32 1]}""",
+                       sprint(io->Metal.code_llvm(io, dummy, Tuple{}; macos=v"3.2.1",
+                                                      dump_module=true, kernel=true)))
+    end
+
     @test Metal.return_type(identity, Tuple{Int}) === Int
     @test Metal.return_type(sin, Tuple{Float32}) === Float32
     @test Metal.return_type(getindex, Tuple{MtlDeviceArray{Float32,1,1},Int32}) === Float32