Add StaticDimsArray: GPU arrays with a mix of static and dynamic dimensions

Project.toml

@@ -18,13 +18,16 @@ Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
 ScopedValues = "7e506255-f358-4e82-b7e4-beb19740aa63"
 Serialization = "9e88b42a-f829-5b0c-bbe9-9e923198166b"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
+Static = "aedffcd0-7271-4cad-89d0-dc628f76c6d3"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 
 [weakdeps]
 JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819"
+StaticArrayInterface = "0d7ed370-da01-4f52-bd93-41d350b8b718"
 
 [extensions]
 JLD2Ext = "JLD2"
+StaticArrayInterfaceExt = "StaticArrayInterface"
 
 [compat]
 AcceleratedKernels = "0.5"
@@ -40,5 +43,7 @@ Reexport = "1"
 ScopedValues = "1"
 Serialization = "1"
 SparseArrays = "1"
+Static = "1"
+StaticArrayInterface = "1.5"
 Statistics = "1"
 julia = "1.10"

ext/StaticArrayInterfaceExt.jl

@@ -0,0 +1,15 @@
+module StaticArrayInterfaceExt
+
+# Static sizes of arrays with a mix of static and dynamic dimensions, for packages that build
+# on StaticArrayInterface.jl (e.g. StrideArraysCore.jl and LoopVectorization.jl).
+
+using GPUArrays: GPUArrays, StaticDimsArray, StaticDimsDeviceArray
+using StaticArrayInterface: StaticArrayInterface
+
+const AnyStaticDimsArray = Union{StaticDimsArray, StaticDimsDeviceArray}
+
+StaticArrayInterface.static_size(a::AnyStaticDimsArray) = GPUArrays.static_size(a)
+StaticArrayInterface.known_size(::Type{A}) where {A <: AnyStaticDimsArray} =
+    GPUArrays.known_size(A)
+
+end # module StaticArrayInterfaceExt

lib/JLArrays/src/mapreduce.jl

@@ -12,11 +12,14 @@ _host_source(A::Broadcast.Broadcasted) =   # (`copy` of a 0-dimensional one give
     Array(copyto!(similar(A, Broadcast.combine_eltypes(A.f, A.args)), A))
 _host_source(A) = adapt(ArrayNoCopy(), A)
 
-GPUArrays._ak_mapreduce(f, op, A::Union{AnyJLArray, Broadcast.Broadcasted{<:JLArrayStyle}};
+# (as a `JLArray`, a `StaticDimsArray` over one is reduced on the host)
+const StaticDimsJLArray = GPUArrays.StaticDimsArray{<:Any, <:Any, <:Any, <:JLArray}
+
+GPUArrays._ak_mapreduce(f, op, A::Union{AnyJLArray, StaticDimsJLArray, Broadcast.Broadcasted{<:JLArrayStyle}};
                         backend=nothing, kwargs...) =
     GC.@preserve A @allowscalar GPUArrays.AK.mapreduce(f, op, _host_source(A); kwargs...)
 
-function GPUArrays._ak_mapreducedim!(f, op, R::AnyJLArray,
+function GPUArrays._ak_mapreducedim!(f, op, R::Union{AnyJLArray, StaticDimsJLArray},
                                      A::Union{AbstractArray, Broadcast.Broadcasted};
                                      backend=nothing, kwargs...)
     GC.@preserve R A @allowscalar GPUArrays.AK.mapreducedim!(

src/GPUArrays.jl

@@ -11,6 +11,7 @@ using Base.Cartesian
 
 using Adapt
 using LLVM.Interop: assume
+using Static: Static, StaticInt, dynamic, known
 
 using Reexport
 @reexport using GPUArraysCore
@@ -22,6 +23,7 @@ import AcceleratedKernels as AK
 # device functionality
 include("device/abstractarray.jl")
 include("device/sparse.jl")
+include("device/staticdims.jl")
 
 # host abstractions
 include("host/abstractarray.jl")
@@ -42,6 +44,7 @@ include("host/uniformscaling.jl")
 include("host/statistics.jl")
 include("host/sparse.jl")
 include("host/alloc_cache.jl")
+include("host/staticdims.jl")
 
 
 end # module

src/device/staticdims.jl

@@ -0,0 +1,67 @@
+# arrays with a mix of static and dynamic dimensions
+
+export StaticDimsDeviceArray
+
+# Each dimension is either a `StaticInt`, a compile-time constant that takes no space, or an
+# `Int`. E.g. `Tuple{StaticInt{4}, Int}` for an array of size `(4, n)`.
+const MixedDims{N} = NTuple{N, Union{Int, StaticInt}}
+
+# Column-major linear index of the (in-bounds) indices `I`, as a sum of strided terms:
+# `1 + (i1 - 1) + d1 * (i2 - 1) + (d1 * d2) * (i3 - 1)` for the dimensions `(d1, d2, d3)`.
+# The strides are products of `StaticInt`s as long as the dimensions are, i.e. constants.
+# (Unlike the Horner form `i1 + d1 * ((i2 - 1) + d2 * (i3 - 1))`, the terms of indices that
+# don't change in a loop can be hoisted out of it when the dimensions are dynamic.)
+# (Unrolled, as recursing over the dimensions with strides of growing types would not infer.)
+@inline @generated function static_linear_index(dims::Tuple, I::NTuple{N, Int}) where {N}
+    ex = quote
+        stride_1 = StaticInt(1)
+    end
+    for k in 2:N
+        push!(ex.args, :($(Symbol(:stride_, k)) = $(Symbol(:stride_, k - 1)) * dims[$(k - 1)]))
+    end
+    terms = [:(dynamic($(Symbol(:stride_, k))) * (I[$k] - 1)) for k in 1:N]
+    push!(ex.args, :(+(1, $(terms...))))
+    return ex
+end
+
+"""
+    StaticDimsDeviceArray{T,N,D,A} <: AbstractDeviceArray{T,N}
+
+The device-side counterpart of [`StaticDimsArray`](@ref): a vector `parent::A` (typically a
+device vector of the back-end) viewed as an `N`-dimensional array with the dimensions
+`dims::D`, a tuple of `Int`s and `StaticInt`s. If the parent is `isbits`, so is the array.
+"""
+struct StaticDimsDeviceArray{T, N, D <: MixedDims{N}, A <: AbstractVector{T}} <: AbstractDeviceArray{T, N}
+    parent::A
+    dims::D
+end
+
+Base.parent(a::StaticDimsDeviceArray) = a.parent
+
+# adapt the parent, e.g. so that KernelAbstractions' `@Const` reaches the back-end's read-only
+# loads (`ld.global.nc` with CUDA)
+Adapt.adapt_structure(to, a::StaticDimsDeviceArray) =
+    wrap_static_dims(adapt(to, parent(a)), a.dims)
+
+Base.size(a::StaticDimsDeviceArray) = map(dynamic, a.dims)
+Base.elsize(::Type{<:StaticDimsDeviceArray{T, N, D, A}}) where {T, N, D, A} = Base.elsize(A)
+
+@inline function Base.getindex(a::StaticDimsDeviceArray, i::Int)
+    @boundscheck checkbounds(a, i)
+    @inbounds a.parent[i]
+end
+@inline function Base.getindex(a::StaticDimsDeviceArray{T, N}, I::Vararg{Int, N}) where {T, N}
+    @boundscheck checkbounds(a, I...)
+    @inbounds a.parent[static_linear_index(a.dims, I)]
+end
+
+@inline function Base.setindex!(a::StaticDimsDeviceArray, x, i::Int)
+    @boundscheck checkbounds(a, i)
+    @inbounds a.parent[i] = x
+    return a
+end
+@inline function Base.setindex!(a::StaticDimsDeviceArray{T, N}, x, I::Vararg{Int, N}) where {T, N}
+    @boundscheck checkbounds(a, I...)
+    @inbounds a.parent[static_linear_index(a.dims, I)] = x
+    return a
+end

src/host/staticdims.jl

@@ -0,0 +1,115 @@
+# arrays with a mix of static and dynamic dimensions
+
+export StaticDimsArray
+
+"""
+    StaticDimsArray(parent::AbstractGPUArray, dims::Tuple)
+
+An `N`-dimensional GPU array that views the memory of `parent` with the dimensions `dims`, a
+tuple of `Int`s and `StaticInt`s from Static.jl, e.g. `(static(4), static(4), n)`. The static
+dimensions are part of the type, so in kernels the linear indices, strides and bounds checks
+involving them are computed from constants, while the other dimensions can change at runtime.
+
+Like other GPU arrays, a `StaticDimsArray` supports broadcasting, reductions and the other
+operations of GPUArrays; out-of-place operations return arrays of the back-end without the
+static dimensions. In kernels it becomes a [`StaticDimsDeviceArray`](@ref).
+
+The size information is available through `static_size` and `known_size` of
+StaticArrayInterface.jl, when it is loaded.
+
+```julia
+julia> a = StaticDimsArray(JLArray{Float32}(undef, 4 * 4 * 10), (static(4), static(4), 10));
+
+julia> size(a)
+(4, 4, 10)
+
+julia> typeof(a.dims)
+Tuple{StaticInt{4}, StaticInt{4}, Int64}
+```
+"""
+struct StaticDimsArray{T, N, D <: MixedDims{N}, A <: AbstractGPUArray{T, 1}} <: AbstractGPUArray{T, N}
+    parent::A
+    dims::D
+
+    function StaticDimsArray(parent::AbstractGPUArray{T, 1}, dims::MixedDims{N}) where {T, N}
+        if prod(map(dynamic, dims)) != length(parent)
+            throw(DimensionMismatch("parent of length $(length(parent)) cannot be viewed with dimensions $dims"))
+        end
+        new{T, N, typeof(dims), typeof(parent)}(parent, dims)
+    end
+end
+
+StaticDimsArray(parent::AbstractGPUArray, dims::Tuple{Vararg{Union{Integer, StaticInt}}}) =
+    StaticDimsArray(vec(parent), map(d -> d isa StaticInt ? d : Int(d), dims))
+
+# keep the dimensions of `a`, use different memory
+StaticDimsArray(a::StaticDimsArray, parent::AbstractGPUArray) = StaticDimsArray(parent, a.dims)
+
+"""
+    static_size(a)
+
+The dimensions of a [`StaticDimsArray`](@ref) or [`StaticDimsDeviceArray`](@ref), as a tuple
+of `StaticInt`s and `Int`s.
+"""
+static_size(a::Union{StaticDimsArray, StaticDimsDeviceArray}) = a.dims
+
+"""
+    known_size(::Type{<:StaticDimsArray})
+
+The dimensions that are known from the type of a [`StaticDimsArray`](@ref) or
+[`StaticDimsDeviceArray`](@ref), with `nothing` for the dynamic ones.
+"""
+known_size(::Type{<:Union{StaticDimsArray{T, N, D}, StaticDimsDeviceArray{T, N, D}}}) where {T, N, D} =
+    known(D)
+known_size(a::Union{StaticDimsArray, StaticDimsDeviceArray}) = known_size(typeof(a))
+
+Base.parent(a::StaticDimsArray) = a.parent
+Base.size(a::StaticDimsArray) = map(dynamic, a.dims)
+Base.elsize(::Type{<:StaticDimsArray{T, N, D, A}}) where {T, N, D, A} = Base.elsize(A)
+
+KernelAbstractions.get_backend(a::StaticDimsArray) = get_backend(parent(a))
+
+# the memory is the one of the parent, so is the alias detection
+storage(a::StaticDimsArray) = storage(parent(a))
+memory_location(a::StaticDimsArray) = memory_location(parent(a))
+
+# views, reshapes and reinterpretations derive from the parent, and lose the static dimensions
+derive(::Type{T}, a::StaticDimsArray, dims::Dims, offset::Int) where {T} =
+    derive(T, parent(a), dims, offset)
+
+# broadcasting: the style of the parent, for `N` dimensions
+function Broadcast.BroadcastStyle(::Type{<:StaticDimsArray{T, N, D, A}}) where {T, N, D, A}
+    style = Broadcast.BroadcastStyle(A)
+    typeof(style)(Val(N))
+end
+
+# `similar`: keep the static dimensions when the size does not change
+Base.similar(a::StaticDimsArray, ::Type{S}) where {S} =
+    StaticDimsArray(similar(parent(a), S), a.dims)
+function Base.similar(a::StaticDimsArray, ::Type{S}, dims::Dims) where {S}
+    if dims == size(a)
+        similar(a, S)
+    else
+        similar(parent(a), S, dims)
+    end
+end
+
+Base.copy(a::StaticDimsArray) = StaticDimsArray(copy(parent(a)), a.dims)
+
+# copies to and from the CPU go through the parent
+function Base.copyto!(dest::Array{T}, dstart::Integer, src::StaticDimsArray{T},
+                      sstart::Integer, n::Integer) where {T}
+    copyto!(dest, dstart, parent(src), sstart, n)
+end
+function Base.copyto!(dest::StaticDimsArray{T}, dstart::Integer, src::Array{T},
+                      sstart::Integer, n::Integer) where {T}
+    copyto!(parent(dest), dstart, src, sstart, n)
+    return dest
+end
+
+# kernel arguments (and other conversions of the memory)
+Adapt.adapt_structure(to, a::StaticDimsArray) = wrap_static_dims(adapt(to, parent(a)), a.dims)
+Adapt.adapt_structure(::ToArray, a::StaticDimsArray) = convert(Array, a)
+
+wrap_static_dims(parent::AbstractGPUArray, dims) = StaticDimsArray(parent, dims)
+wrap_static_dims(parent::AbstractVector, dims) = StaticDimsDeviceArray(parent, dims)

test/testsuite.jl

@@ -126,6 +126,7 @@ include("testsuite/construction.jl")
 include("testsuite/indexing.jl")
 include("testsuite/base.jl")
 include("testsuite/aliasing.jl")
+include("testsuite/staticdims.jl")
 include("testsuite/vector.jl")
 include("testsuite/reductions.jl")
 include("testsuite/sorting.jl")

test/testsuite/staticdims.jl

@@ -0,0 +1,78 @@
+@kernel function staticdims_index_kernel!(dest, src)
+    i, j, k = @index(Global, NTuple)
+    @inbounds dest[i, j, k] = src[i, j, k] + i + 10j + 100k
+end
+
+@testsuite "staticdims" (AT, eltypes)->begin
+    # a `StaticDimsArray` wraps GPU memory
+    AT <: AbstractGPUArray || return
+    static = GPUArrays.Static.static
+    StaticInt = GPUArrays.Static.StaticInt
+
+    T = Float32 in eltypes ? Float32 : first(eltypes)
+    dims = (static(4), static(3), 5)
+    host = rand(T, 4, 3, 5)
+    a = StaticDimsArray(AT(vec(host)), dims)
+
+    @testset "size" begin
+        @test a isa AbstractGPUArray{T, 3}
+        @test size(a) == (4, 3, 5)
+        @test GPUArrays.static_size(a) === dims
+        @test GPUArrays.known_size(typeof(a)) === (4, 3, nothing)
+        @test Array(a) == host
+        @test_throws DimensionMismatch StaticDimsArray(AT{T}(undef, 10), dims)
+        # a multidimensional parent is viewed as a vector
+        @test StaticDimsArray(AT(host), dims) == a
+    end
+
+    @testset "kernel arguments" begin
+        backend = get_backend(a)
+        b = similar(a)
+        @test b isa StaticDimsArray{T, 3, typeof(dims)}
+        staticdims_index_kernel!(backend)(b, a; ndrange = size(a))
+        expected = [host[i, j, k] + i + 10j + 100k for i in 1:4, j in 1:3, k in 1:5]
+        @test Array(b) ≈ expected
+    end
+
+    @testset "broadcasting" begin
+        b = similar(a)
+        b .= a .* 2
+        @test Array(b) ≈ host .* 2
+        @test b isa StaticDimsArray
+        # mixed with plain arrays of the back-end
+        b .= a .+ AT(host)
+        @test Array(b) ≈ host .+ host
+        # out of place: an array of the back-end
+        c = a .+ one(T)
+        @test Array(c) ≈ host .+ one(T)
+        fill!(b, zero(T))
+        @test all(iszero, Array(b))
+    end
+
+    @testset "reductions" begin
+        @test sum(a) ≈ sum(host)
+        @test Array(sum(a; dims = 3)) ≈ sum(host; dims = 3)
+        @test Array(maximum(a; dims = 1)) ≈ maximum(host; dims = 1)
+    end
+
+    @testset "linear indices" begin
+        # static and dynamic dimensions in any order
+        for dims in ((static(2), 3, static(4)), (2, static(3), 4), (2, 3, 4),
+                     (static(2), static(3), static(4)))
+            d = StaticDimsDeviceArray(collect(1:24), dims)
+            @test all(d[I] == LinearIndices((2, 3, 4))[I] for I in CartesianIndices((2, 3, 4)))
+        end
+    end
+
+    @testset "indexing" begin
+        @test (@allowscalar a[2, 3, 4]) == host[2, 3, 4]
+        @test Array(a[:, 2, :]) == host[:, 2, :]
+        @test Array(view(a, :, :, 2)) == host[:, :, 2]
+        @test Array(reshape(a, 12, 5)) == reshape(host, 12, 5)
+    end
+
+    @testset "aliasing" begin
+        @test Base.mightalias(a, parent(a))
+        @test !Base.mightalias(a, copy(a))
+    end
+end