Add support for constant memory.

Author: S-D-R

lib/cudadrv/module/global.jl

@@ -4,8 +4,7 @@
 # - should be more dict-like: get and setindex(::name), haskey(::name)
 # - globals(::Type)?
 
-export
-    CuGlobal, get, set
+export CuGlobal, get, set, CuGlobalArray
 
 
 """
@@ -62,3 +61,35 @@ function Base.setindex!(var::CuGlobal{T}, val::T) where T
     val_ref = Ref{T}(val)
     cuMemcpyHtoD_v2(var, val_ref, var.buf.bytesize)
 end
+
+"""
+    CuGlobalArray{T}(mod::CuModule, name::String, len::Integer)
+
+Acquires a global array variable handle from a named global in a module.
+"""
+struct CuGlobalArray{T} # TODO: the functionality provided by this struct can most likely be merged into CuGlobal{T}
+    buf::Mem.DeviceBuffer
+
+    function CuGlobalArray{T}(mod::CuModule, name::String, len::Integer) where T
+        ptr_ref = Ref{CuPtr{Cvoid}}()
+        nbytes_ref = Ref{Csize_t}()
+        cuModuleGetGlobal_v2(ptr_ref, nbytes_ref, mod, name)
+        if nbytes_ref[] != (sizeof(T) * len)
+            throw(ArgumentError("size of global array '$name' ($(nbytes_ref[])) does not match given size (sizeof($T) * $length)"))
+        end
+        buf = Mem.DeviceBuffer(ptr_ref[], nbytes_ref[])
+
+        return new{T}(buf)
+    end
+end
+
+Base.eltype(::Type{CuGlobalArray{T}}) where {T} = T
+
+Base.sizeof(global_array::CuGlobalArray{T}) where T = sizeof(global_array.buf)
+
+function Base.copyto!(global_array::CuGlobalArray{T}, src::Array{T}) where T
+    if sizeof(src) != sizeof(global_array)
+        throw(DimensionMismatch("size of `src` ($(sizeof(src))) does not match global array ($(sizeof(global_array)))"))
+    end
+    cuMemcpyHtoD_v2(global_array.buf, src, sizeof(src))
+end

src/CUDA.jl

@@ -60,6 +60,7 @@ include("array.jl")
 include("gpuarrays.jl")
 include("utilities.jl")
 include("texture.jl")
+include("memory_constant.jl")
 
 # array libraries
 include("../lib/complex.jl")

src/compiler/gpucompiler.jl

@@ -35,6 +35,7 @@ function GPUCompiler.finish_module!(job::CUDACompilerJob, mod::LLVM.Module)
            Tuple{CompilerJob{PTXCompilerTarget}, typeof(mod)},
            job, mod)
     emit_exception_flag!(mod)
+    emit_constant_memory_initializer!(mod)
 end
 
 function GPUCompiler.link_libraries!(job::CUDACompilerJob, mod::LLVM.Module,

src/device/intrinsics.jl

@@ -12,6 +12,7 @@ include("intrinsics/memory_dynamic.jl")
 include("intrinsics/atomics.jl")
 include("intrinsics/misc.jl")
 include("intrinsics/wmma.jl")
+include("intrinsics/memory_constant.jl")
 
 # functionality from libdevice
 #

src/device/intrinsics/memory_constant.jl

@@ -0,0 +1,72 @@
+# Constant Memory
+
+export CuDeviceConstantMemory
+
+"""
+    CuDeviceConstantMemory{T,N,Name,Shape}
+
+The device-side counterpart of [`CuConstantMemory{T,N}`](@ref). This type should not be used
+directly except in the case of `CuConstantMemory` global variables, where it denotes the
+type of the relevant kernel argument.
+
+Note that the `Name` and `Shape` type variables are implementation details and it
+discouraged to use them directly. Instead use [name(::CuConstantMemory)](@ref) and
+[Base.size(::CuConstantMemory)](@ref) respectively.
+"""
+struct CuDeviceConstantMemory{T,N,Name,Shape} <: AbstractArray{T,N} end
+
+"""
+Get the name of underlying global variable of this `CuDeviceConstantMemory`.
+"""
+name(::CuDeviceConstantMemory{T,N,Name,Shape}) where {T,N,Name,Shape} = Name
+
+Base.:(==)(A::CuDeviceConstantMemory, B::CuDeviceConstantMemory) = name(A) == name(B)
+Base.hash(A::CuDeviceConstantMemory, h::UInt) = hash(name(A), h)
+
+Base.size(::CuDeviceConstantMemory{T,N,Name,Shape}) where {T,N,Name,Shape} = Shape
+
+Base.@propagate_inbounds Base.getindex(A::CuDeviceConstantMemory, i::Integer) = constmemref(A, i)
+
+Base.IndexStyle(::Type{<:CuDeviceConstantMemory}) = Base.IndexLinear()
+
+@inline function constmemref(A::CuDeviceConstantMemory{T,N,Name,Shape}, index::Integer) where {T,N,Name,Shape}
+    @boundscheck checkbounds(A, index)
+    len = length(A)
+    return read_constant_mem(Val(Name), index, T, Val(len))
+end
+
+@generated function read_constant_mem(::Val{global_name}, index::Integer, ::Type{T}, ::Val{len}) where {global_name,T,len}
+    JuliaContext() do ctx
+        # define LLVM types
+        T_int = convert(LLVMType, Int, ctx)
+        T_result = convert(LLVMType, T, ctx)
+
+        # define function and get LLVM module
+        param_types = [T_int]
+        llvm_f, _ = create_function(T_result, param_types)
+        mod = LLVM.parent(llvm_f)
+
+        # create a constant memory global variable
+        T_global = LLVM.ArrayType(T_result, len)
+        global_var = GlobalVariable(mod, T_global, string(global_name), AS.Constant)
+        linkage!(global_var, LLVM.API.LLVMExternalLinkage) # NOTE: external linkage is the default
+        extinit!(global_var, true)
+        # TODO: global_var alignment?
+
+        # generate IR
+        Builder(ctx) do builder
+            entry = BasicBlock(llvm_f, "entry", ctx)
+            position!(builder, entry)
+
+            typed_ptr = inbounds_gep!(builder, global_var, [ConstantInt(0, ctx), parameters(llvm_f)[1]])
+            ld = load!(builder, typed_ptr)
+
+            metadata(ld)[LLVM.MD_tbaa] = tbaa_addrspace(AS.Constant, ctx)
+
+            ret!(builder, ld)
+        end
+
+        # call the function
+        call_function(llvm_f, T, Tuple{Int}, :((Int(index - one(index))),))
+    end
+end

src/memory_constant.jl

@@ -0,0 +1,137 @@
+export CuConstantMemory
+
+# Map a constant memory name to its array value
+const constant_memory_initializer = Dict{Symbol,WeakRef}()
+
+"""
+    CuConstantMemory{T,N}(value::Array{T,N})
+    CuConstantMemory{T}(::UndefInitializer, dims::Integer...)
+    CuConstantMemory{T}(::UndefInitializer, dims::Dims{N})
+
+Construct an `N`-dimensional constant memory array of type `T`, where `isbits(T)`.
+
+Note that `deepcopy` will be called on the `value` constructor argument, meaning that
+mutations to the original `value` or its elements after construction will not be
+reflected in the value of `CuConstantMemory`.
+
+The `UndefInitializer` constructors behave exactly like the regular `Array` version,
+i.e. the value of `CuConstantMemory` will be completely random when using them.
+
+Unlike in CUDA C, structs cannot be put directly into constant memory. This feature can
+be emulated however by wrapping the struct inside of a 1-element array.
+
+When using `CuConstantMemory` as a global variable it is required to pass it as an argument
+to a kernel, where the argument is of type [`CuDeviceConstantMemory{T,N}`](@ref).
+When using `CuConstantMemory` as a local variable that is captured by a kernel closure
+this is not required, and it can be used directly like any other captured variable
+without passing it as an argument.
+
+In cases where the same kernel object gets called mutiple times, and it is desired to mutate
+the value of a `CuConstantMemory` variable in this kernel between calls, please refer
+to [`Base.copyto!(const_mem::CuConstantMemory{T}, value::Array{T}, kernel::HostKernel)`](@ref)
+"""
+struct CuConstantMemory{T,N} <: AbstractArray{T,N}
+    name::Symbol
+    value::Array{T,N}
+
+    function CuConstantMemory(value::Array{T,N}) where {T,N}
+        # TODO: add finalizer that removes the relevant entry from constant_memory_initializer?
+        Base.isbitstype(T) || throw(ArgumentError("CuConstantMemory only supports bits types"))
+        name = gensym("constant_memory")
+        name = GPUCompiler.safe_name(string(name))
+        name = Symbol(name)
+        val = deepcopy(value)
+        constant_memory_initializer[name] = WeakRef(val)
+        return new{T,N}(name, val)
+    end
+end
+
+CuConstantMemory{T}(::UndefInitializer, dims::Integer...) where {T} =
+    CuConstantMemory(Array{T}(undef, dims))
+CuConstantMemory{T}(::UndefInitializer, dims::Dims{N}) where {T,N} =
+    CuConstantMemory(Array{T,N}(undef, dims))
+
+Base.size(A::CuConstantMemory) = size(A.value)
+
+Base.getindex(A::CuConstantMemory, i::Integer) = Base.getindex(A.value, i)
+Base.setindex!(A::CuConstantMemory, v, i::Integer) = Base.setindex!(A.value, v, i)
+Base.IndexStyle(::Type{<:CuConstantMemory}) = Base.IndexLinear()
+
+Adapt.adapt_storage(::Adaptor, A::CuConstantMemory{T,N}) where {T,N} = 
+    CuDeviceConstantMemory{T,N,A.name,size(A.value)}()
+
+
+"""
+Given a `kernel` returned by `@cuda`, copy `value` into `const_mem` for subsequent calls to this `kernel`.
+If `const_mem` is not used within `kernel`, an error will be thrown.
+"""
+function Base.copyto!(const_mem::CuConstantMemory{T}, value::Array{T}, kernel::HostKernel) where T
+    # TODO: add bool argument to also change the value field of const_mem?
+    if size(const_mem) != size(value)
+        throw(DimensionMismatch("size of `value` does not match size of constant memory"))
+    end
+
+    global_array = CuGlobalArray{T}(kernel.mod, string(const_mem.name), length(const_mem))
+    copyto!(global_array, value)
+end
+
+
+function emit_constant_memory_initializer!(mod::LLVM.Module)
+    for global_var in globals(mod)
+        T_global = llvmtype(global_var)
+        if addrspace(T_global) == AS.Constant
+            constant_memory_name = Symbol(LLVM.name(global_var))
+            if !haskey(constant_memory_initializer, constant_memory_name)
+                continue # non user defined constant memory, most likely from the CUDA runtime
+            end
+
+            arr = constant_memory_initializer[constant_memory_name].value
+            @assert !isnothing(arr) "calling kernel containing garbage collected constant memory"
+
+            flattened_arr = reduce(vcat, arr)
+            ctx = LLVM.context(mod)
+            typ = eltype(eltype(T_global))
+
+            # TODO: have a look at how julia converts structs to llvm:
+            #       https://github.com/JuliaLang/julia/blob/80ace52b03d9476f3d3e6ff6da42f04a8df1cf7b/src/cgutils.cpp#L572
+            #       this only seems to emit a type though
+            if isa(typ, LLVM.IntegerType) || isa(typ, LLVM.FloatingPointType)
+                init = ConstantArray(flattened_arr, ctx)
+            elseif isa(typ, LLVM.ArrayType) # a struct with every field of the same type gets optimized to an array
+                constant_arrays = LLVM.Constant[]
+                for x in flattened_arr
+                    fields = collect(map(name->getfield(x, name), fieldnames(typeof(x))))
+                    constant_array = ConstantArray(fields, ctx)
+                    push!(constant_arrays, constant_array)
+                end
+                init = ConstantArray(typ, constant_arrays)
+            elseif isa(typ, LLVM.StructType)
+                constant_structs = LLVM.Constant[]
+                for x in flattened_arr
+                    constants = LLVM.Constant[]
+                    for fieldname in fieldnames(typeof(x))
+                        field = getfield(x, fieldname)
+                        if isa(field, Bool)
+                            # NOTE: Bools get compiled to i8 instead of the more "correct" type i1
+                            push!(constants, ConstantInt(LLVM.Int8Type(ctx), field))
+                        elseif isa(field, Integer)
+                            push!(constants, ConstantInt(field, ctx))
+                        elseif isa(field, AbstractFloat)
+                            push!(constants, ConstantFP(field, ctx))
+                        else
+                            throw(error("constant memory does not currently support structs with non-primitive fields ($(typeof(x)).$fieldname::$(typeof(field)))"))
+                        end
+                    end
+                    const_struct = ConstantStruct(typ, constants)
+                    push!(constant_structs, const_struct)
+                end
+                init = ConstantArray(typ, constant_structs)
+            else
+                # unreachable, but let's be safe and throw a nice error message just in case
+                throw(error("could not emit initializer for constant memory of type $typ"))
+            end
+            
+            initializer!(global_var, init)
+        end
+    end
+end