Fixes loading of extensions (#32)

* Add missing using

* Fix rng usage and remove machine argument

* Add runtime generated functions init to extensions

* Fix github docs link

* Add more docs to the kernel generation and add AMDGPU kernel

* WIP fix GPU ext

* Make the extensions load correctly

* Update Numa impl too

* Add warning that GPUs are experimental

* Include new file in the docs

docs/make.jl

@@ -32,7 +32,7 @@ makedocs(;
     sitename="ComputableDAGs.jl",
     format=Documenter.HTML(;
         prettyurls=get(ENV, "CI", "false") == "true",
-        canonical="https://ComputableDAGs.gitlab.io/ComputableDAGs.jl",
+        canonical="https://ComputableDAGs.github.io/ComputableDAGs.jl",
         assets=String[],
     ),
     pages=pages,

docs/src/lib/internals/devices.md

@@ -37,11 +37,10 @@ Pages = ["devices/numa/impl.jl"]
 Order = [:type, :function]
 ```
 
-### CUDA
-For CUDA functionality to be available, the `CUDA.jl` package must be installed separately, as it is only a weak dependency.
+### GPUs
 
-### ROCm
-For ROCm functionality to be available, the `AMDGPU.jl` package must be installed separately, as it is only a weak dependency.
-
-### oneAPI
-For oneAPI functionality to be available, the `oneAPI.jl` package must be installed separately, as it is only a weak dependency.
+```@autodocs
+Modules = [ComputableDAGs]
+Pages = ["devices/ext.jl"]
+Order = [:type]
+```

ext/AMDGPUExt.jl

@@ -1,8 +1,20 @@
 module AMDGPUExt
 
-using ComputableDAGs, AMDGPU
+using ComputableDAGs
+using UUIDs
+using AMDGPU
+
+function __init__()
+    @debug "Loading AMDGPUExt"
+
+    push!(ComputableDAGs.DEVICE_TYPES, ROCmGPU)
+    ComputableDAGs.CACHE_STRATEGIES[ROCmGPU] = [LocalVariables()]
+
+    return nothing
+end
 
 # include specialized AMDGPU functions here
 include("devices/rocm/impl.jl")
+include("devices/rocm/function.jl")
 
 end

ext/CUDAExt.jl

@@ -1,8 +1,17 @@
 module CUDAExt
 
 using ComputableDAGs
+using UUIDs
 using CUDA
-using RuntimeGeneratedFunctions
+
+function __init__()
+    @debug "Loading CUDAExt"
+
+    push!(ComputableDAGs.DEVICE_TYPES, CUDAGPU)
+    ComputableDAGs.CACHE_STRATEGIES[CUDAGPU] = [LocalVariables()]
+
+    return nothing
+end
 
 # include specialized CUDA functions here
 include("devices/cuda/impl.jl")

ext/devices/cuda/function.jl

@@ -1,14 +1,12 @@
-
-function ComputableDAGs.cuda_kernel(
-    graph::DAG, instance, machine::Machine, context_module::Module
-)
+function ComputableDAGs.kernel(::Type{CUDAGPU}, graph::DAG, instance)
+    machine = cpu_st()
     tape = ComputableDAGs.gen_tape(graph, instance, machine, context_module)
 
     init_caches = Expr(:block, tape.initCachesCode...)
     assign_inputs = Expr(:block, ComputableDAGs.expr_from_fc.(tape.inputAssignCode)...)
     code = Expr(:block, ComputableDAGs.expr_from_fc.(tape.computeCode)...)
 
-    function_id = ComputableDAGs.to_var_name(UUIDs.uuid1(rng[1]))
+    function_id = ComputableDAGs.to_var_name(UUIDs.uuid1(ComputableDAGs.rng[1]))
     res_sym = eval(
         ComputableDAGs.gen_access_expr(
             ComputableDAGs.entry_device(tape.machine), tape.outputSymbol
@@ -29,5 +27,5 @@ function ComputableDAGs.cuda_kernel(
         end"
     )
 
-    return RuntimeGeneratedFunction(@__MODULE__, context_module, expr)
+    return expr
 end

ext/devices/cuda/impl.jl

@@ -1,24 +1,7 @@
-"""
-    CUDAGPU <: AbstractGPU
-
-Representation of a specific CUDA GPU that code can run on. Implements the [`AbstractDevice`](@ref) interface.
-"""
-mutable struct CUDAGPU <: ComputableDAGs.AbstractGPU
-    device::Any # TODO: what's the cuda device type?
-    cacheStrategy::CacheStrategy
-    FLOPS::Float64
-end
-
-push!(ComputableDAGs.DEVICE_TYPES, CUDAGPU)
-
-ComputableDAGs.CACHE_STRATEGIES[CUDAGPU] = [LocalVariables()]
-
 ComputableDAGs.default_strategy(::Type{CUDAGPU}) = LocalVariables()
 
 function ComputableDAGs.measure_device!(device::CUDAGPU; verbose::Bool)
-    if verbose
-        println("Measuring CUDA GPU $(device.device)")
-    end
+    verbose && @info "Measuring CUDA GPU $(device.device)"
 
     # TODO implement
     return nothing
@@ -29,20 +12,16 @@ end
 
 Return a Vector of [`CUDAGPU`](@ref)s available on the current machine. If `verbose` is true, print some additional information.
 """
-function get_devices(::Type{CUDAGPU}; verbose::Bool=false)
+function ComputableDAGs.get_devices(::Type{CUDAGPU}; verbose::Bool=false)
     devices = Vector{ComputableDAGs.AbstractDevice}()
 
     if !CUDA.functional()
-        if verbose
-            println("CUDA.jl is non-functional")
-        end
+        @warn "The CUDA extension is loaded but CUDA.jl is non-functional"
         return devices
     end
 
     CUDADevices = CUDA.devices()
-    if verbose
-        println("Found $(length(CUDADevices)) CUDA devices")
-    end
+    verbose && @info "Found $(length(CUDADevices)) CUDA devices"
     for device in CUDADevices
         push!(devices, CUDAGPU(device, default_strategy(CUDAGPU), -1))
     end

ext/devices/oneapi/impl.jl

@@ -1,24 +1,7 @@
-"""
-    oneAPIGPU <: AbstractGPU
-
-Representation of a specific Intel GPU that code can run on. Implements the [`AbstractDevice`](@ref) interface.
-"""
-mutable struct oneAPIGPU <: ComputableDAGs.AbstractGPU
-    device::Any
-    cacheStrategy::CacheStrategy
-    FLOPS::Float64
-end
-
-push!(ComputableDAGs.DEVICE_TYPES, oneAPIGPU)
-
-ComputableDAGs.CACHE_STRATEGIES[oneAPIGPU] = [LocalVariables()]
-
 ComputableDAGs.default_strategy(::Type{oneAPIGPU}) = LocalVariables()
 
 function ComputableDAGs.measure_device!(device::oneAPIGPU; verbose::Bool)
-    if verbose
-        println("Measuring oneAPI GPU $(device.device)")
-    end
+    verbose && @info "Measuring oneAPI GPU $(device.device)"
 
     # TODO implement
     return nothing
@@ -29,20 +12,16 @@ end
 
 Return a Vector of [`oneAPIGPU`](@ref)s available on the current machine. If `verbose` is true, print some additional information.
 """
-function get_devices(::Type{oneAPIGPU}; verbose::Bool=false)
-    devices = Vector{AbstractDevice}()
+function ComputableDAGs.get_devices(::Type{oneAPIGPU}; verbose::Bool=false)
+    devices = Vector{ComputableDAGs.AbstractDevice}()
 
     if !oneAPI.functional()
-        if verbose
-            println("oneAPI is non-functional")
-        end
+        @warn "the oneAPI extension is loaded but oneAPI.jl is non-functional"
         return devices
     end
 
     oneAPIDevices = oneAPI.devices()
-    if verbose
-        println("Found $(length(oneAPIDevices)) oneAPI devices")
-    end
+    verbose && @info "Found $(length(oneAPIDevices)) oneAPI devices"
     for device in oneAPIDevices
         push!(devices, oneAPIGPU(device, default_strategy(oneAPIGPU), -1))
     end

ext/devices/rocm/function.jl

@@ -0,0 +1,31 @@
+function ComputableDAGs.kernel(::Type{ROCmGPU}, graph::DAG, instance)
+    machine = cpu_st()
+    tape = ComputableDAGs.gen_tape(graph, instance, machine, context_module)
+
+    init_caches = Expr(:block, tape.initCachesCode...)
+    assign_inputs = Expr(:block, ComputableDAGs.expr_from_fc.(tape.inputAssignCode)...)
+    code = Expr(:block, ComputableDAGs.expr_from_fc.(tape.computeCode)...)
+
+    function_id = ComputableDAGs.to_var_name(UUIDs.uuid1(ComputableDAGs.rng[1]))
+    res_sym = eval(
+        ComputableDAGs.gen_access_expr(
+            ComputableDAGs.entry_device(tape.machine), tape.outputSymbol
+        ),
+    )
+    expr = Meta.parse(
+        "function compute_$(function_id)(input_vector, output_vector, n::Int64)
+            id = (workgroupIdx().x - 1) * workgroupDim().x + workgroupIdx().x
+            if (id > n)  
+                return
+            end
+            @inline data_input = input_vector[id]
+            $(init_caches)
+            $(assign_inputs)
+            $code
+            @inline output_vector[id] = $res_sym
+            return nothing
+        end"
+    )
+
+    return expr
+end

ext/devices/rocm/impl.jl

@@ -1,26 +1,7 @@
-using AMDGPU
-
-"""
-    ROCmGPU <: AbstractGPU
-
-Representation of a specific AMD GPU that code can run on. Implements the [`AbstractDevice`](@ref) interface.
-"""
-mutable struct ROCmGPU <: ComputableDAGs.AbstractGPU
-    device::Any
-    cacheStrategy::CacheStrategy
-    FLOPS::Float64
-end
-
-push!(ComputableDAGs.DEVICE_TYPES, ROCmGPU)
-
-ComputableDAGs.CACHE_STRATEGIES[ROCmGPU] = [LocalVariables()]
-
 ComputableDAGs.default_strategy(::Type{ROCmGPU}) = LocalVariables()
 
 function ComputableDAGs.measure_device!(device::ROCmGPU; verbose::Bool)
-    if verbose
-        println("Measuring ROCm GPU $(device.device)")
-    end
+    verbose && @info "Measuring ROCm GPU $(device.device)"
 
     # TODO implement
     return nothing
@@ -31,20 +12,16 @@ end
 
 Return a Vector of [`ROCmGPU`](@ref)s available on the current machine. If `verbose` is true, print some additional information.
 """
-function get_devices(::Type{ROCmGPU}; verbose::Bool=false)
-    devices = Vector{AbstractDevice}()
+function ComputableDAGs.get_devices(::Type{ROCmGPU}; verbose::Bool=false)
+    devices = Vector{ComputableDAGs.AbstractDevice}()
 
     if !AMDGPU.functional()
-        if verbose
-            println("AMDGPU is non-functional")
-        end
+        @warn "The AMDGPU extension is loaded but AMDGPU.jl is non-functional"
         return devices
     end
 
     AMDDevices = AMDGPU.devices()
-    if verbose
-        println("Found $(length(AMDDevices)) AMD devices")
-    end
+    verbose && @info "Found $(length(AMDDevices)) AMD devices"
     for device in AMDDevices
         push!(devices, ROCmGPU(device, default_strategy(ROCmGPU), -1))
     end

ext/oneAPIExt.jl

@@ -1,6 +1,17 @@
 module oneAPIExt
 
-using ComputableDAGs, oneAPI
+using ComputableDAGs
+using UUIDs
+using oneAPI
+
+function __init__()
+    @debug "Loading oneAPIExt"
+
+    push!(ComputableDAGs.DEVICE_TYPES, oneAPIGPU)
+    ComputableDAGs.CACHE_STRATEGIES[oneAPIGPU] = [LocalVariables()]
+
+    return nothing
+end
 
 # include specialized oneAPI functions here
 include("devices/oneapi/impl.jl")

src/ComputableDAGs.jl

@@ -60,8 +60,8 @@ export get_machine_info, cpu_st
 export CacheStrategy, default_strategy
 export LocalVariables, Dictionary
 
-# CUDAExt
-export cuda_kernel
+# GPU Extensions
+export kernel, CUDAGPU, ROCmGPU, oneAPIGPU
 
 include("devices/interface.jl")
 include("task/type.jl")
@@ -124,6 +124,7 @@ include("devices/detect.jl")
 include("devices/impl.jl")
 
 include("devices/numa/impl.jl")
+include("devices/ext.jl")
 
 include("scheduler/interface.jl")
 include("scheduler/greedy.jl")

src/devices/ext.jl

@@ -0,0 +1,44 @@
+# file for struct definitions used by the extensions
+# since extensions can't export names themselves
+
+"""
+    CUDAGPU <: AbstractGPU
+
+Representation of a specific CUDA GPU that code can run on. Implements the [`AbstractDevice`](@ref) interface.
+
+!!! note
+    This requires CUDA to be loaded to be useful.
+"""
+mutable struct CUDAGPU <: AbstractGPU
+    device::Any # CuDevice
+    cacheStrategy::CacheStrategy
+    FLOPS::Float64
+end
+
+"""
+    oneAPIGPU <: AbstractGPU
+
+Representation of a specific Intel GPU that code can run on. Implements the [`AbstractDevice`](@ref) interface.
+
+!!! note
+    This requires oneAPI to be loaded to be useful.
+"""
+mutable struct oneAPIGPU <: AbstractGPU
+    device::Any # oneAPI.oneL0.ZeDevice
+    cacheStrategy::CacheStrategy
+    FLOPS::Float64
+end
+
+"""
+    ROCmGPU <: AbstractGPU
+
+Representation of a specific AMD GPU that code can run on. Implements the [`AbstractDevice`](@ref) interface.
+
+!!! note
+    This requires AMDGPU to be loaded to be useful.
+"""
+mutable struct ROCmGPU <: AbstractGPU
+    device::Any # HIPDevice
+    cacheStrategy::CacheStrategy
+    FLOPS::Float64
+end

src/devices/impl.jl

@@ -59,15 +59,6 @@ It is the simplest machine definition possible and produces a simple function wh
 """
 function cpu_st()
     return Machine(
-        [
-            ComputableDAGs.NumaNode(
-                0,
-                1,
-                ComputableDAGs.default_strategy(ComputableDAGs.NumaNode),
-                -1.0,
-                UUIDs.uuid1(),
-            ),
-        ],
-        [-1.0;;],
+        [NumaNode(0, 1, default_strategy(NumaNode), -1.0, UUIDs.uuid1())], [-1.0;;]
     )
 end