WIP: Refactor memory hierarchy.

maleadt · maleadt · commit 6edf8ac3dd2a · 2025-06-19T15:35:28.000+02:00
diff --git a/LocalPreferences.toml b/LocalPreferences.toml
@@ -1,7 +1,7 @@
 [OpenCL]
 # Which memory back-end to use for unspecified CLArray allocations. This can be:
 # - "usm": Unified Shared Memory (`cl_intel_unified_shared_memory`)
-# - "bda": Buffer Device Address (`cl_mem` + `cl_ext_buffer_device_address`)
+# - "bda": plain buffers (`cl_mem` + `cl_ext_buffer_device_address`)
 # - "svm": Shared Virtual Memory (coarse-grained)
 # If unspecified, the default will be used based on the platform and device capabilities.
 #default_memory_backend="..."
diff --git a/lib/cl/CL.jl b/lib/cl/CL.jl
@@ -20,8 +20,7 @@ include("device.jl")
 include("context.jl")
 include("cmdqueue.jl")
 include("event.jl")
-include("buffer.jl")
-include("memory/memory.jl")
+include("memory.jl")
 include("program.jl")
 include("kernel.jl")
 
diff --git a/lib/cl/kernel.jl b/lib/cl/kernel.jl
@@ -79,7 +79,7 @@ function set_arg!(k::Kernel, idx::Integer, arg::AbstractMemory)
         clSetKernelArgSVMPointer(k, idx - 1, pointer(arg))
     elseif arg isa UnifiedMemory
         clSetKernelArgMemPointerINTEL(k, idx - 1, pointer(arg))
-    elseif arg isa BufferDeviceMemory
+    elseif arg isa Buffer
         clSetKernelArgDevicePointerEXT(k, idx - 1, pointer(arg))
     else
         error("Unknown memory type")
@@ -203,7 +203,7 @@ function call(
 
             if memory isa SharedVirtualMemory
                 push!(svm_pointers, ptr)
-            elseif memory isa BufferDeviceMemory
+            elseif memory isa Buffer
                 push!(bda_pointers, ptr)
             elseif memory isa UnifiedDeviceMemory
                 device_access = true
diff --git a/lib/cl/memory/bda.jl b/lib/cl/memory/bda.jl
diff --git a/lib/cl/memory/buffer.jl b/lib/cl/memory/buffer.jl
@@ -1,92 +1,17 @@
-# 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
-
-release(mem::AbstractMemoryObject) = clReleaseMemObject(mem)
-
-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[])
-    elseif s == :device_address
-        result = Ref{cl_mem_device_address_ext}()
-        clGetMemObjectInfo(mem, CL_MEM_DEVICE_ADDRESS_EXT, sizeof(cl_mem_device_address_ext), result, C_NULL)
-        return CLPtr{Cvoid}(result[])
-    else
-        return getfield(mem, s)
-    end
-end
-
-# convenience functions
-context(mem::AbstractMemoryObject) = mem.context
-Base.pointer(mem::AbstractMemoryObject) = mem.pointer
-
-#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
 
 struct Buffer <: AbstractMemoryObject
     id::cl_mem
+    ptr::Union{Nothing,CLPtr{Cvoid}}
     bytesize::Int
+    context::Context
 end
 
+Buffer() = Buffer(C_NULL, nothing, 0, context())
+
+Base.pointer(buf::Buffer) = @something buf.ptr error("Buffer does not have a device private address")
 Base.sizeof(buf::Buffer) = buf.bytesize
+context(buf::Buffer) = buf.context
 
 
 ## constructors
@@ -130,7 +55,16 @@ function Buffer(sz::Int, flags::Integer, hostbuf=nothing;
     if err_code[] != CL_SUCCESS
         throw(CLError(err_code[]))
     end
-    return Buffer(mem_id, sz)
+
+    ptr = if device_private_address
+        ptr_ref = Ref{cl_mem_device_address_ext}()
+        clGetMemObjectInfo(mem_id, CL_MEM_DEVICE_ADDRESS_EXT, sizeof(cl_mem_device_address_ext), ptr_ref, C_NULL)
+        CLPtr{Cvoid}(ptr_ref[])
+    else
+        nothing
+    end
+
+    return Buffer(mem_id, ptr, sz, context())
 end
 
 # allocated buffer
diff --git a/lib/cl/memory/memory.jl b/lib/cl/memory/memory.jl
diff --git a/lib/cl/memory/svm.jl b/lib/cl/memory/svm.jl
@@ -1,4 +1,4 @@
-struct SharedVirtualMemory <: AbstractMemory
+struct SharedVirtualMemory <: AbstractPointerMemory
     ptr::CLPtr{Cvoid}
     bytesize::Int
     context::Context
diff --git a/lib/cl/memory/usm.jl b/lib/cl/memory/usm.jl
@@ -1,4 +1,4 @@
-abstract type UnifiedMemory <: AbstractMemory end
+abstract type UnifiedMemory <: AbstractPointerMemory end
 
 function usm_free(mem::UnifiedMemory; blocking::Bool = false)
     if blocking
diff --git a/lib/cl/state.jl b/lib/cl/state.jl
@@ -199,10 +199,10 @@ function default_memory_backend(dev::Device)
 
     backend = if backend_str == "usm"
         USMBackend()
-    elseif backend_str == "bda"
-        BDABackend()
     elseif backend_str == "svm"
         SVMBackend()
+    elseif backend_str == "bda"
+        BDABackend()
     else
         error("Unknown memory backend '$backend_str' requested")
     end
diff --git a/src/array.jl b/src/array.jl
@@ -99,7 +99,7 @@ function memory_type()
     elseif cl.memory_backend() == cl.SVMBackend()
         return cl.SharedVirtualMemory
     elseif cl.memory_backend() == cl.BDABackend()
-        return cl.BufferDeviceMemory
+        return cl.Buffer
     end
 end
 CLArray{T, N}(::UndefInitializer, dims::Dims{N}) where {T, N} =
@@ -175,11 +175,14 @@ context(A::CLArray) = cl.context(A.data[].mem)
 memtype(x::CLArray) = memtype(typeof(x))
 memtype(::Type{<:CLArray{<:Any, <:Any, M}}) where {M} = @isdefined(M) ? M : Any
 
-is_device(a::CLArray) = memtype(a) == cl.UnifiedDeviceMemory
-is_shared(a::CLArray) = memtype(a) == cl.UnifiedSharedMemory
-is_host(a::CLArray) = memtype(a) == cl.UnifiedHostMemory
-is_svm(a::CLArray) = memtype(a) == cl.SharedVirtualMemory
-is_bda(a::CLArray) = memtype(a) == cl.BufferDeviceMemory
+# can we read this array from the device (i.e. derive a CLPtr)?
+is_device(a::CLArray) =
+    memtype(a) in (cl.UnifiedDeviceMemory, cl.UnifiedSharedMemory, cl.SharedVirtualMemory, cl.Buffer)
+is_shared(a::CLArray) =
+    memtype(a) in (cl.UnifiedSharedMemory, cl.SharedVirtualMemory)
+is_host(a::CLArray) =
+    memtype(a) in (cl.UnifiedHostMemory, cl.UnifiedSharedMemory, cl.SharedVirtualMemory)
+
 
 ## derived types
 
@@ -283,13 +286,16 @@ end
 ## interop with libraries
 
 function Base.unsafe_convert(::Type{Ptr{T}}, x::CLArray{T}) where {T}
-    if is_device(x)
+    if !is_host(x)
         throw(ArgumentError("cannot take the CPU address of a $(typeof(x))"))
     end
     return convert(Ptr{T}, x.data[]) + x.offset * Base.elsize(x)
 end
 
 function Base.unsafe_convert(::Type{CLPtr{T}}, x::CLArray{T}) where {T}
+    if !is_device(x)
+        throw(ArgumentError("cannot take the device address of a $(typeof(x))"))
+    end
     return convert(CLPtr{T}, x.data[]) + x.offset * Base.elsize(x)
 end
 
diff --git a/src/memory.jl b/src/memory.jl
@@ -113,7 +113,7 @@ end
 
 
 ## public interface
-function managed_alloc(t::Type{T}, bytes::Int; kwargs...) where T <: cl.AbstractMemory
+function managed_alloc(t::Type{T}, bytes::Int; kwargs...) where T
     if bytes == 0
         return Managed(T())
     else
@@ -142,12 +142,13 @@ function alloc(::Type{cl.SharedVirtualMemory}, bytes::Int; alignment::Int = 0)
     return Managed(mem)
 end
 
-function alloc(::Type{cl.BufferDeviceMemory}, bytes::Int; alignment::Int = 0)
-    mem = cl.bda_alloc(bytes; alignment)
-    return Managed(mem)
+function alloc(::Type{cl.Buffer}, bytes::Int; alignment::Int = 0)
+    # TODO: use alignment
+    buf = cl.Buffer(bytes; device_private_address = true)
+    return Managed(buf)
 end
 
-function free(managed::Managed{<:cl.AbstractMemory})
+function free(managed::Managed)
     sizeof(managed) == 0 && return
     mem = managed.mem
     cl.context!(cl.context(mem)) do
@@ -162,10 +163,10 @@ function free(managed::Managed{<:cl.AbstractMemory})
 
         if mem isa cl.SharedVirtualMemory
             cl.svm_free(mem)
-        elseif mem isa cl.BufferDeviceMemory
-            cl.bda_free(mem)
-        else
+        elseif mem isa cl.UnifiedMemory
             cl.usm_free(mem)
+        else
+            cl.release(mem)
         end
     end
 
diff --git a/src/util.jl b/src/util.jl
@@ -85,10 +85,10 @@ function versioninfo(io::IO=stdout)
                 suffix = backend == cl.default_memory_backend(device) ? "*" : ""
                 if backend isa cl.SVMBackend
                     push!(tags, "svm"*suffix)
-                elseif backend isa cl.BDABackend
-                    push!(tags, "bda"*suffix)
                 elseif backend isa cl.USMBackend
                     push!(tags, "usm"*suffix)
+                elseif backend isa cl.BDABackend
+                    push!(tags, "bda"*suffix)
                 end
             end
             ## relevant extensions
diff --git a/test/memory.jl b/test/memory.jl
@@ -17,8 +17,8 @@
         buf = create_test_buffer()
 
         expectations = [
-            (:mem_type, cl.CL_MEM_OBJECT_BUFFER),
-            (:mem_flags, (:rw, :copy)),
+            (:type, cl.CL_MEM_OBJECT_BUFFER),
+            (:flags, (:rw, :copy)),
             (:size, sizeof(buf)),
             (:reference_count, 1),
             (:map_count, 0)

Original file line number	Diff line number	Diff line change
`@@ -1,4 +1,4 @@`
`1`		`-struct SharedVirtualMemory <: AbstractMemory`
	`1`	`+struct SharedVirtualMemory <: AbstractPointerMemory`
`2`	`2`	`ptr::CLPtr{Cvoid}`
`3`	`3`	`bytesize::Int`
`4`	`4`	`context::Context`
Original file line number	Diff line number	Diff line change
`@@ -1,4 +1,4 @@`
`1`		`-abstract type UnifiedMemory <: AbstractMemory end`
	`1`	`+abstract type UnifiedMemory <: AbstractPointerMemory end`
`2`	`2`
`3`	`3`	`function usm_free(mem::UnifiedMemory; blocking::Bool = false)`
`4`	`4`	`if blocking`