Revamp memory management. (#264)

Adds support for USM next to SVM, while also porting
several features from CUDA.jl (CLPtr, Managed, etc).

Co-authored-by: Tim Besard <tim.besard@gmail.com>

Author: VarLad

Project.toml

@@ -21,7 +21,7 @@ StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 Adapt = "4"
 GPUArrays = "11.2.1"
 GPUCompiler = "0.27, 1"
-KernelAbstractions = "0.9.1"
+KernelAbstractions = "0.9.2"
 LLVM = "9.1"
 LinearAlgebra = "1"
 OpenCL_jll = "=2024.5.8"

examples/vadd.jl

@@ -21,7 +21,7 @@ prog = cl.Program(; source) |> cl.build!
 kern = cl.Kernel(prog, "vadd")
 
 len = prod(dims)
-clcall(kern, Tuple{Ptr{Float32}, Ptr{Float32}, Ptr{Float32}},
+clcall(kern, Tuple{CLPtr{Float32}, CLPtr{Float32}, CLPtr{Float32}},
        d_a, d_b, d_c; global_size=(len,))
 c = Array(d_c)
 @test a+b ≈ c

lib/cl/CL.jl

@@ -1,5 +1,8 @@
 module cl
 
+using Printf
+
+include("pointer.jl")
 include("api.jl")
 
 # OpenCL wrapper objects are expected to have an `id` field containing a handle pointer
@@ -15,9 +18,8 @@ include("device.jl")
 include("context.jl")
 include("cmdqueue.jl")
 include("event.jl")
-include("memory.jl")
+include("memory/memory.jl")
 include("buffer.jl")
-include("svm.jl")
 include("program.jl")
 include("kernel.jl")
 

lib/cl/api.jl

@@ -69,6 +69,22 @@ function retry_reclaim(f, isfailed)
     ret
 end
 
+macro ext_ccall(ex)
+    # decode the expression
+    @assert Meta.isexpr(ex, :(::))
+    call, ret = ex.args
+    @assert Meta.isexpr(call, :call)
+    target, argexprs... = call.args
+    @assert Meta.isexpr(target, :(.))
+    _, fn = target.args
+
+    @gensym fptr
+    esc(quote
+        $fptr = $clGetExtensionFunctionAddressForPlatform(platform(), $fn)
+        @ccall $(Expr(:($), fptr))($(argexprs...))::$ret
+    end)
+end
+
 include("libopencl.jl")
 
 @static if Sys.iswindows()
@@ -176,22 +192,4 @@ function __init__()
     if !OpenCL_jll.is_available()
         @error "OpenCL_jll is not available for your platform, OpenCL.jl. will not work."
     end
-
-    # ensure that operations executed by the REPL back-end finish before returning,
-    # because displaying values happens on a different task
-    if isdefined(Base, :active_repl_backend) && !isnothing(Base.active_repl_backend)
-        push!(Base.active_repl_backend.ast_transforms, synchronize_opencl_tasks)
-    end
-end
-
-function synchronize_opencl_tasks(ex)
-    quote
-        try
-            $(ex)
-        finally
-            if haskey($task_local_storage(), :CLDevice)
-                $device_synchronize()
-            end
-        end
-    end
 end

lib/cl/buffer.jl

@@ -1,6 +1,79 @@
+# OpenCL Memory Object
+
+abstract type AbstractMemoryObject <: CLObject end
+
+#This should be implemented by all subtypes
+# type MemoryType <: AbstractMemoryObject
+#     id::cl_mem
+#     ...
+# end
+
+# for passing buffers to OpenCL APIs: use the underlying handle
+Base.unsafe_convert(::Type{cl_mem}, mem::AbstractMemoryObject) = mem.id
+
+# for passing buffers to kernels: keep the buffer, it's handled by `cl.set_arg!`
+Base.unsafe_convert(::Type{<:Ptr}, mem::AbstractMemoryObject) = mem
+
+Base.sizeof(mem::AbstractMemoryObject) = mem.size
+
+context(mem::AbstractMemoryObject) = mem.context
+
+function Base.getproperty(mem::AbstractMemoryObject, s::Symbol)
+    if s == :context
+        param = Ref{cl_context}()
+        clGetMemObjectInfo(mem, CL_MEM_CONTEXT, sizeof(cl_context), param, C_NULL)
+        return Context(param[], retain = true)
+    elseif s == :mem_type
+        result = Ref{cl_mem_object_type}()
+        clGetMemObjectInfo(mem, CL_MEM_TYPE, sizeof(cl_mem_object_type), result, C_NULL)
+        return result[]
+    elseif s == :mem_flags
+        result = Ref{cl_mem_flags}()
+        clGetMemObjectInfo(mem, CL_MEM_FLAGS, sizeof(cl_mem_flags), result, C_NULL)
+        mf = result[]
+        flags = Symbol[]
+        if (mf & CL_MEM_READ_WRITE) != 0
+            push!(flags, :rw)
+        end
+        if (mf & CL_MEM_WRITE_ONLY) != 0
+            push!(flags, :w)
+        end
+        if (mf & CL_MEM_READ_ONLY) != 0
+            push!(flags, :r)
+        end
+        if (mf & CL_MEM_USE_HOST_PTR) != 0
+            push!(flags, :use)
+        end
+        if (mf & CL_MEM_ALLOC_HOST_PTR) != 0
+            push!(flags, :alloc)
+        end
+        if (mf & CL_MEM_COPY_HOST_PTR) != 0
+            push!(flags, :copy)
+        end
+        return tuple(flags...)
+    elseif s == :size
+        result = Ref{Csize_t}()
+        clGetMemObjectInfo(mem, CL_MEM_SIZE, sizeof(Csize_t), result, C_NULL)
+        return result[]
+    elseif s == :reference_count
+        result = Ref{Cuint}()
+        clGetMemObjectInfo(mem, CL_MEM_REFERENCE_COUNT, sizeof(Cuint), result, C_NULL)
+        return Int(result[])
+    elseif s == :map_count
+        result = Ref{Cuint}()
+        clGetMemObjectInfo(mem, CL_MEM_MAP_COUNT, sizeof(Cuint), result, C_NULL)
+        return Int(result[])
+    else
+        return getfield(mem, s)
+    end
+end
+
+#TODO: enqueue_migrate_mem_objects(queue, mem_objects, flags=0, wait_for=None)
+#TODO: enqueue_migrate_mem_objects_ext(queue, mem_objects, flags=0, wait_for=None)
+
 # OpenCL.Buffer
 
-mutable struct Buffer{T} <: AbstractMemory
+mutable struct Buffer{T} <: AbstractMemoryObject
     const id::cl_mem
     const len::Int
 

lib/cl/cmdqueue.jl

@@ -2,15 +2,14 @@
 
 mutable struct CmdQueue <: CLObject
     const id::cl_command_queue
+    Base.@atomic valid::Bool
 
     function CmdQueue(q_id::cl_command_queue; retain::Bool=false)
-        q = new(q_id)
+        q = new(q_id, true)
         retain && clRetainCommandQueue(q)
         finalizer(q) do _
-          # this is to prevent `device_synchronize()` operating on freed queues.
-          # XXX: why does the WeakKeyDict contain freed objects?
-          delete!(cl.queues, q)
-          clReleaseCommandQueue(q)
+            Base.@atomic q.valid = false
+            clReleaseCommandQueue(q)
         end
         return q
     end

lib/cl/device.jl

@@ -190,6 +190,59 @@ function exec_capabilities(d::Device)
     )
 end
 
+function usm_supported(d::Device)
+    "cl_intel_unified_shared_memory" in d.extensions || return false
+    return true
+end
+
+function usm_capabilities(d::Device)
+    usm_supported(d) || throw(ArgumentError("Unified Shared Memory not supported on this device"))
+
+    function check_capability_bits(mask::cl_device_unified_shared_memory_capabilities_intel)
+        (;
+            access = mask & CL_UNIFIED_SHARED_MEMORY_ACCESS_INTEL != 0,
+            atomic_access = mask & CL_UNIFIED_SHARED_MEMORY_ATOMIC_ACCESS_INTEL != 0,
+            concurrent_access = mask & CL_UNIFIED_SHARED_MEMORY_CONCURRENT_ACCESS_INTEL != 0,
+            concurrent_atomic_access = mask & CL_UNIFIED_SHARED_MEMORY_CONCURRENT_ATOMIC_ACCESS_INTEL != 0,
+        )
+    end
+
+    host = Ref{cl_device_unified_shared_memory_capabilities_intel}()
+    device = Ref{cl_device_unified_shared_memory_capabilities_intel}()
+    single_device = Ref{cl_device_unified_shared_memory_capabilities_intel}()
+    shared = Ref{cl_device_unified_shared_memory_capabilities_intel}()
+    cross_device = Ref{cl_device_unified_shared_memory_capabilities_intel}()
+
+    clGetDeviceInfo(
+        d, CL_DEVICE_HOST_MEM_CAPABILITIES_INTEL,
+        sizeof(cl_device_unified_shared_memory_capabilities_intel), host, C_NULL
+    )
+    clGetDeviceInfo(
+        d, CL_DEVICE_DEVICE_MEM_CAPABILITIES_INTEL,
+        sizeof(cl_device_unified_shared_memory_capabilities_intel), device, C_NULL
+    )
+    clGetDeviceInfo(
+        d, CL_DEVICE_SINGLE_DEVICE_SHARED_MEM_CAPABILITIES_INTEL,
+        sizeof(cl_device_unified_shared_memory_capabilities_intel), single_device, C_NULL
+    )
+    clGetDeviceInfo(
+        d, CL_DEVICE_SHARED_SYSTEM_MEM_CAPABILITIES_INTEL,
+        sizeof(cl_device_unified_shared_memory_capabilities_intel), shared, C_NULL
+    )
+    clGetDeviceInfo(
+        d, CL_DEVICE_CROSS_DEVICE_SHARED_MEM_CAPABILITIES_INTEL,
+        sizeof(cl_device_unified_shared_memory_capabilities_intel), cross_device, C_NULL
+    )
+
+    return (;
+        host = check_capability_bits(host[]),
+        device = check_capability_bits(device[]),
+        single_device = check_capability_bits(single_device[]),
+        shared = check_capability_bits(shared[]),
+        cross_device = check_capability_bits(cross_device[]),
+    )
+end
+
 function svm_capabilities(d::Device)
     result = Ref{cl_device_svm_capabilities}()
     clGetDeviceInfo(d, CL_DEVICE_SVM_CAPABILITIES,