Add OpenCL runtime support

Abstract runtime functionality by HSA (TODO: OCL)
Use Requires to load OpenCL bindings
Allow choosing runtime via environment variables

Author: jpsamaroo

Project.toml

@@ -10,6 +10,7 @@ HSARuntime = "2c364e2c-59fb-59c3-96f3-194112e690e0"
 InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240"
 LLVM = "929cbde3-209d-540e-8aea-75f648917ca0"
 Libdl = "8f399da3-3557-5675-b5ff-fb832c97cbdb"
+Requires = "ae029012-a4dd-5104-9daa-d747884805df"
 TimerOutputs = "a759f4b9-e2f1-59dc-863e-4aeb61b1ea8f"
 
 [compat]

src/AMDGPUnative.jl

@@ -7,6 +7,14 @@ using Adapt
 using TimerOutputs
 using DataStructures
 using Libdl
+using Requires
+
+@enum DeviceRuntime HSA OCL
+const RUNTIME = Ref{DeviceRuntime}(HSA)
+if get(ENV, "AMDGPUNATIVE_OPENCL", "") != ""
+    RUNTIME[] = OCL
+end
+include("runtime.jl")
 
 const configured = HSARuntime.configured
 
@@ -29,6 +37,7 @@ include("reflection.jl")
 
 function __init__()
     check_deps()
+    @require OpenCL="08131aa3-fb12-5dee-8b74-c09406e224a2" include("opencl.jl")
     __init_compiler__()
 end
 

src/compiler/common.jl

@@ -4,7 +4,7 @@ struct CompilerJob
     # core invocation
     f::Base.Callable
     tt::DataType
-    agent::HSAAgent
+    device::RuntimeDevice
     kernel::Bool
 
     # optional properties
@@ -14,10 +14,10 @@ struct CompilerJob
     maxregs::Union{Nothing,Integer}
     name::Union{Nothing,String}
 
-    CompilerJob(f, tt, agent, kernel; name=nothing,
+    CompilerJob(f, tt, device, kernel; name=nothing,
                 minthreads=nothing, maxthreads=nothing,
                 blocks_per_sm=nothing, maxregs=nothing) =
-        new(f, tt, agent, kernel, minthreads, maxthreads, blocks_per_sm,
+        new(f, tt, device, kernel, minthreads, maxthreads, blocks_per_sm,
             maxregs, name)
 end
 

src/compiler/driver.jl

@@ -4,11 +4,11 @@
 const compile_hook = Ref{Union{Nothing,Function}}(nothing)
 
 """
-    compile(target::Symbol, agent::HSAAgent, f, tt, kernel=true;
+    compile(target::Symbol, device::RuntimeDevice, f, tt, kernel=true;
             libraries=true, optimize=true, strip=false, strict=true, ...)
-Compile a function `f` invoked with types `tt` for agent `agent` to one of the
-following formats as specified by the `target` argument: `:julia` for Julia
-IR, `:llvm` for LLVM IR, `:gcn` for GCN assembly, and `:roc` for linked
+Compile a function `f` invoked with types `tt` for device `device` to one of
+the following formats as specified by the `target` argument: `:julia` for
+Julia IR, `:llvm` for LLVM IR, `:gcn` for GCN assembly, and `:roc` for linked
 objects. If the `kernel` flag is set, specialized code generation and
 optimization for kernel functions is enabled.
 The following keyword arguments are supported:
@@ -18,11 +18,11 @@ The following keyword arguments are supported:
 - `strict`: perform code validation either as early or as late as possible
 Other keyword arguments can be found in the documentation of [`rocfunction`](@ref).
 """
-compile(target::Symbol, agent::HSAAgent, @nospecialize(f::Core.Function),
+compile(target::Symbol, device::RuntimeDevice, @nospecialize(f::Core.Function),
                  @nospecialize(tt), kernel::Bool=true; libraries::Bool=true,
                  optimize::Bool=true, strip::Bool=false, strict::Bool=true, kwargs...) =
 
-                 compile(target, CompilerJob(f, tt, agent, kernel; kwargs...);
+                 compile(target, CompilerJob(f, tt, device, kernel; kwargs...);
                          libraries=libraries, optimize=optimize, strip=strip,
                          strict=strict)
 
@@ -88,7 +88,7 @@ function codegen(target::Symbol, job::CompilerJob; libraries::Bool=true,
     # always preload the runtime, and do so early; it cannot be part of any timing block
     # because it recurses into the compiler
     if libraries
-        runtime = load_runtime(job.agent)
+        runtime = load_runtime(job.device)
         runtime_fns = LLVM.name.(defs(runtime))
     end
 
@@ -105,7 +105,7 @@ function codegen(target::Symbol, job::CompilerJob; libraries::Bool=true,
             end
             =#
             # FIXME: Load this only when needed
-            device_libs = load_device_libs(job.agent)
+            device_libs = load_device_libs(job.device)
             for lib in device_libs
                 if need_library(ir, lib)
                     @timeit to[] "device library" link_device_lib!(job, ir, lib)

src/compiler/mcgen.jl

@@ -1,12 +1,12 @@
 # machine code generation
 
-function machine(agent::HSAAgent, triple::String)
+function machine(device::RuntimeDevice, triple::String)
     InitializeAMDGPUTarget()
     InitializeAMDGPUTargetInfo()
     t = Target(triple)
 
     InitializeAMDGPUTargetMC()
-    cpu = get_first_isa(agent) # TODO: Make this configurable
+    cpu = default_isa(device) # TODO: Make this configurable
     feat = ""
     tm = TargetMachine(t, triple, cpu, feat)
     asm_verbosity!(tm, true)
@@ -80,7 +80,7 @@ end
 
 function mcgen(job::CompilerJob, mod::LLVM.Module, f::LLVM.Function;
                output_format=LLVM.API.LLVMObjectFile)
-    tm = machine(job.agent, triple(mod))
+    tm = machine(job.device, triple(mod))
 
     InitializeAMDGPUAsmPrinter()
     return String(emit(tm, mod, output_format))

src/compiler/optim.jl

@@ -1,7 +1,7 @@
 # LLVM IR optimization
 
 function optimize!(job::CompilerJob, mod::LLVM.Module, entry::LLVM.Function)
-    tm = AMDGPUnative.machine(job.agent, triple(mod))
+    tm = AMDGPUnative.machine(job.device, triple(mod))
 
     if job.kernel
         entry = promote_kernel!(job, mod, entry)

src/compiler/rtlib.jl

@@ -30,10 +30,10 @@ const libcache = Dict{String, LLVM.Module}()
 
 # ROCm device library
 
-function load_device_libs(agent)
+function load_device_libs(device)
     device_libs_path === nothing && return
 
-    isa_short = replace(get_first_isa(agent), "gfx"=>"")
+    isa_short = replace(default_isa(device), "gfx"=>"")
     device_libs = LLVM.Module[]
     bitcode_files = (
         "hc.amdgcn.bc",
@@ -132,20 +132,20 @@ end
 
 ## functionality to build the runtime library
 
-function emit_function!(mod, agent, f, types, name)
+function emit_function!(mod, device, f, types, name)
     tt = Base.to_tuple_type(types)
-    new_mod, entry = codegen(:llvm, CompilerJob(f, tt, agent, #=kernel=# false);
+    new_mod, entry = codegen(:llvm, CompilerJob(f, tt, device, #=kernel=# false);
                              libraries=false, strict=false)
     LLVM.name!(entry, name)
     link!(mod, new_mod)
 end
 
-function build_runtime(agent)
+function build_runtime(device)
     mod = LLVM.Module("AMDGPUnative run-time library", JuliaContext())
 
     for method in values(Runtime.methods)
         try
-            emit_function!(mod, agent, method.def, method.types, method.llvm_name)
+            emit_function!(mod, device, method.def, method.types, method.llvm_name)
         catch err
             @warn method
         end
@@ -154,8 +154,8 @@ function build_runtime(agent)
     mod
 end
 
-function load_runtime(agent::HSAAgent)
-    isa = get_first_isa(agent)
+function load_runtime(device::RuntimeDevice)
+    isa = default_isa(device)
     name = "amdgpunative.$isa.bc"
     path = joinpath(@__DIR__, "..", "..", "deps", "runtime", name)
     mkpath(dirname(path))
@@ -167,7 +167,7 @@ function load_runtime(agent::HSAAgent)
             end
         else
             @info "Building the AMDGPUnative run-time library for your $isa device, this might take a while..."
-            lib = build_runtime(agent)
+            lib = build_runtime(device)
             open(path, "w") do io
                 write(io, lib)
             end

src/execution.jl

@@ -3,16 +3,17 @@
 export @roc, rocconvert, rocfunction
 
 struct Kernel{F,TT}
-    agent::HSAAgent
+    device::RuntimeDevice
     mod::ROCModule
     fun::ROCFunction
 end
 
 # `split_kwargs()` segregates keyword arguments passed to `@roc` into those
 # affecting the compiler, kernel execution, or both.
 function split_kwargs(kwargs)
-    compiler_kws = [:agent, :queue, :name]
-    call_kws     = [:groupsize, :gridsize, :agent, :queue]
+    # TODO: Alias groupsize and gridsize as threads and blocks, respectively
+    compiler_kws = [:device, :agent, :queue, :name]
+    call_kws     = [:groupsize, :gridsize, :device, :agent, :queue]
     compiler_kwargs = []
     call_kwargs = []
     for kwarg in kwargs
@@ -60,6 +61,23 @@ function assign_args!(code, args)
     return vars, var_exprs
 end
 
+function extract_device(;device=nothing, agent=nothing, kwargs...)
+    if device !== nothing
+        return device
+    elseif agent !== nothing
+        return agent
+    else
+        return default_device()
+    end
+end
+function extract_queue(device; queue=nothing, kwargs...)
+    if queue !== nothing
+        return queue
+    else
+        return default_queue(device)
+    end
+end
+
 # fast lookup of global world age
 world_age() = ccall(:jl_get_tls_world_age, UInt, ())
 
@@ -125,10 +143,10 @@ macro roc(ex...)
             GC.@preserve $(vars...) begin
                 local kernel_args = map(rocconvert, ($(var_exprs...),))
                 local kernel_tt = Tuple{Core.Typeof.(kernel_args)...}
-                local agent = get_default_agent()
-                local kernel = rocfunction(agent, $(esc(f)), kernel_tt;
+                local device = extract_device(; $(esc(call_kwargs)...))
+                local kernel = rocfunction(device, $(esc(f)), kernel_tt;
                                            $(map(esc, compiler_kwargs)...))
-                local queue = get_default_queue(agent)
+                local queue = extract_queue(device; $(esc(call_kwargs)...))
                 local signal = HSASignal()
                 kernel(queue, signal, kernel_args...; $(map(esc, call_kwargs)...))
                 wait(signal)
@@ -188,7 +206,7 @@ The output of this function is automatically cached, i.e. you can simply call
 generated automatically, when the function changes, or when different types or
 keyword arguments are provided.
 """
-@generated function rocfunction(agent::HSAAgent, f::Core.Function, tt::Type=Tuple{}; name=nothing, kwargs...)
+@generated function rocfunction(device::RuntimeDevice, f::Core.Function, tt::Type=Tuple{}; name=nothing, kwargs...)
     tt = Base.to_tuple_type(tt.parameters[1])
     sig = Base.signature_type(f, tt)
     t = Tuple(tt.parameters)
@@ -217,8 +235,8 @@ keyword arguments are provided.
 
         # compile the function
         if !haskey(compilecache, key)
-            fun, mod = compile(:roc, agent, f, tt; name=name, kwargs...)
-            kernel = Kernel{f,tt}(agent, mod, fun)
+            fun, mod = compile(:roc, device, f, tt; name=name, kwargs...)
+            kernel = Kernel{f,tt}(device, mod, fun)
             compilecache[key] = kernel
         end
 
@@ -227,9 +245,9 @@ keyword arguments are provided.
 end
 
 rocfunction(f::Core.Function, tt::Type=Tuple{}; kwargs...) =
-    rocfunction(get_default_agent(), f, tt; kwargs...)
+    rocfunction(default_device(), f, tt; kwargs...)
 
-@generated function call(kernel::Kernel{F,TT}, queue::HSAQueue,
+@generated function call(kernel::Kernel{F,TT}, queue::RuntimeQueue,
                          signal::HSASignal, args...; call_kwargs...) where {F,TT}
 
     sig = Base.signature_type(F, TT)

src/execution_utils.jl

@@ -101,23 +101,13 @@ end
         GC.@preserve $(arg_refs...) begin
             kernelParams = [$(arg_ptrs...)]
 
-            # link with ld.lld
-            ld_path = HSARuntime.ld_lld_path
-            @assert ld_path != "" "ld.lld was not found; cannot link kernel"
-            # TODO: Do this more idiomatically
-            io = open("/tmp/amdgpu-dump.o", "w")
-            write(io, f.mod.data)
-            close(io)
-            run(`$ld_path -shared -o /tmp/amdgpu.exe /tmp/amdgpu-dump.o`)
-            io = open("/tmp/amdgpu.exe", "r")
-            data = read(io)
-            close(io)
-
-            # generate executable and kernel instance
-            exe = HSAExecutable(queue.agent, data, f.entry)
-            kern = HSAKernelInstance(queue.agent, exe, f.entry, args)
-            HSARuntime.launch!(queue, kern, signal;
-                workgroup_size=groupsize, grid_size=gridsize)
+            # create executable and kernel instance
+            exe = create_executable(get_device(queue), data, f.entry)
+            kern = create_kernel(get_device(queue), exe, f.entry, args)
+
+            # launch kernel
+            launch_kernel(queue, kern, signal;
+                          workgroup_size=groupsize, grid_size=gridsize)
         end
     end).args)
 

src/opencl.jl

@@ -0,0 +1,3 @@
+# OpenCL runtime interface to AMDGPUnative
+
+include("opencl/args.jl")