Made BatchedCuArray, SharedCuArray, SharedScalar and BatchedStruct Adapt.jl compatible

Author: sihanyu03

GeneralisedFilters/Project.toml

@@ -6,6 +6,7 @@ authors = ["THargreaves <tim.hargreaves@icloud.com>", "Charles Knipp <charleskni
 [deps]
 AbstractMCMC = "80f14c24-f653-4e6a-9b94-39d6b0f70001"
 AcceleratedKernels = "6a4ca0a5-0e36-4168-a932-d9be78d558f1"
+Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 DataStructures = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
@@ -25,6 +26,7 @@ StructArrays = "09ab397b-f2b6-538f-b94a-2f83cf4a842a"
 [compat]
 AbstractMCMC = "5"
 AcceleratedKernels = "0.3, 0.4"
+Adapt = "4.5.0"
 Aqua = "0.8"
 CUDA = "5"
 DataStructures = "0.18.20, 0.19"

GeneralisedFilters/src/batching/operations.jl

@@ -6,17 +6,17 @@ import LinearAlgebra: norm
 # =============================================================================
 
 # Type aliases for BatchedStruct-wrapped matrices
-const BatchedAdjoint{T,M} = BatchedStruct{
-    Adjoint{T,CuArray{T,2,M}},@NamedTuple{parent::BatchedCuMatrix{T,M}}
+const BatchedAdjoint{T,A<:AbstractArray{T,3}} = BatchedStruct{
+    <:Adjoint{T,<:AbstractArray{T,2}},@NamedTuple{parent::BatchedCuMatrix{T,A}}
 }
-const BatchedTranspose{T,M} = BatchedStruct{
-    Transpose{T,CuArray{T,2,M}},@NamedTuple{parent::BatchedCuMatrix{T,M}}
+const BatchedTranspose{T,A<:AbstractArray{T,3}} = BatchedStruct{
+    <:Transpose{T,<:AbstractArray{T,2}},@NamedTuple{parent::BatchedCuMatrix{T,A}}
 }
-const SharedAdjoint{T,M} = BatchedStruct{
-    Adjoint{T,CuArray{T,2,M}},@NamedTuple{parent::SharedCuMatrix{T,M}}
+const SharedAdjoint{T,A<:AbstractArray{T,2}} = BatchedStruct{
+    <:Adjoint{T,<:AbstractArray{T,2}},@NamedTuple{parent::SharedCuMatrix{T,A}}
 }
-const SharedTranspose{T,M} = BatchedStruct{
-    Transpose{T,CuArray{T,2,M}},@NamedTuple{parent::SharedCuMatrix{T,M}}
+const SharedTranspose{T,A<:AbstractArray{T,2}} = BatchedStruct{
+    <:Transpose{T,<:AbstractArray{T,2}},@NamedTuple{parent::SharedCuMatrix{T,A}}
 }
 
 # Union of all GEMM-compatible matrix types

GeneralisedFilters/src/batching/types.jl

@@ -1,3 +1,4 @@
+using Adapt
 using CUDA
 using LinearAlgebra:
     Adjoint, Transpose, LowerTriangular, UpperTriangular, UniformScaling, Cholesky
@@ -13,45 +14,50 @@ export BatchedStruct
 # =============================================================================
 
 """
-    BatchedCuArray{T, NE, NB, NT, M} <: AbstractArray{CuArray{T,NE,M}, NB}
+    BatchedCuArray{T,NE,NB,NT,A<:AbstractArray{T,NT}} <: AbstractArray{Any,NB}
 
-An `NB`-dimensional batch of `NE`-dimensional CuArrays, stored as a single contiguous
-`CuArray{T, NT, M}` where `NT = NE + NB`.
+An `NB`-dimensional batch of `NE`-dimensional arrays, stored as a single contiguous
+`NT`-dimensional array `data` where `NT = NE + NB`.
 
 - `NE`: number of element dimensions (the "inner" array shape)
 - `NB`: number of batch dimensions
 - `NT`: total number of dimensions (`NE + NB`); required explicitly because Julia's type
   system cannot express arithmetic on type parameters
+- `A`: storage array type
 
 The first `NE` dimensions index within each element; the last `NB` dimensions index across
 the batch.
 
+This type is generic over the storage array type so that it can participate in `Adapt.jl`
+transformations. In the user-facing intended usage, `data` is a `CuArray{T, NT, M}`.
+
 # Common aliases
-- `BatchedCuMatrix{T,M}` = `BatchedCuArray{T,2,1,3,M}` — a vector of matrices
-- `BatchedCuVector{T,M}` = `BatchedCuArray{T,1,1,2,M}` — a vector of vectors
+- `BatchedCuMatrix{T,A}` = `BatchedCuArray{T,2,1,3,A}` — a vector of matrices
+- `BatchedCuVector{T,A}` = `BatchedCuArray{T,1,1,2,A}` — a vector of vectors
 """
-struct BatchedCuArray{T,NE,NB,NT,M} <: AbstractArray{CuArray{T,NE,M},NB}
-    data::CuArray{T,NT,M}
+struct BatchedCuArray{T,NE,NB,NT,A<:AbstractArray{T,NT}} <: AbstractArray{Any,NB}
+    data::A
 
-    function BatchedCuArray{T,NE,NB,NT,M}(data::CuArray{T,NT,M}) where {T,NE,NB,NT,M}
+    function BatchedCuArray{T,NE,NB,NT,A}(data::A) where {T,NE,NB,NT,A<:AbstractArray{T,NT}}
         NE + NB == NT || error("NE ($NE) + NB ($NB) must equal ndims(data) ($NT)")
-        return new{T,NE,NB,NT,M}(data)
+        return new{T,NE,NB,NT,A}(data)
     end
 end
 
 # Convenience constructor: infer T and M, require explicit NE and NB
-function BatchedCuArray{T,NE,NB}(data::CuArray{T,NT,M}) where {T,NE,NB,NT,M}
+function BatchedCuArray{T,NE,NB}(data::A) where {T,NE,NB,A<:AbstractArray{T}}
+    NT = ndims(data)
     NE + NB == NT || error("NE ($NE) + NB ($NB) must equal ndims(data) ($NT)")
-    return BatchedCuArray{T,NE,NB,NT,M}(data)
+    return BatchedCuArray{T,NE,NB,NT,A}(data)
 end
 
 # Common case aliases
-const BatchedCuMatrix{T,M} = BatchedCuArray{T,2,1,3,M}
-const BatchedCuVector{T,M} = BatchedCuArray{T,1,1,2,M}
+const BatchedCuMatrix{T,A<:AbstractArray{T,3}} = BatchedCuArray{T,2,1,3,A}
+const BatchedCuVector{T,A<:AbstractArray{T,2}} = BatchedCuArray{T,1,1,2,A}
 
 # Constructors for aliased cases
-BatchedCuMatrix(data::CuArray{T,3,M}) where {T,M} = BatchedCuArray{T,2,1,3,M}(data)
-BatchedCuVector(data::CuArray{T,2,M}) where {T,M} = BatchedCuArray{T,1,1,2,M}(data)
+BatchedCuMatrix(data::A) where {T,A<:AbstractArray{T,3}} = BatchedCuArray{T,2,1,3,A}(data)
+BatchedCuVector(data::A) where {T,A<:AbstractArray{T,2}} = BatchedCuArray{T,1,1,2,A}(data)
 
 const BatchedArray = BatchedCuArray
 
@@ -71,47 +77,69 @@ end
 
 batch_size(x::BatchedCuArray) = length(x)
 
+# Adapting BatchedCuArray to bitstype
+function Adapt.adapt_structure(
+    to,
+    x::BatchedCuArray{T,NE,NB,NT,A},
+) where {T,NE,NB,NT,A}
+    data_adapted = Adapt.adapt(to, x.data)
+    return BatchedCuArray{T,NE,NB,NT,typeof(data_adapted)}(data_adapted)
+end
+
 # =============================================================================
 # Shared Types (same data reused across all batch elements)
 # =============================================================================
 
 """
-    SharedCuArray{T, InnerN, BatchN, M} <: AbstractArray{CuArray{T,InnerN,M}, BatchN}
+    SharedCuArray{T,InnerN,BatchN,A<:AbstractArray{T,InnerN}} <: AbstractArray{Any,BatchN}
 
-A batch of CuArrays where every element is the same underlying `CuArray{T,InnerN,M}`.
+A batch of arrays where every element is the same underlying array.
 Unlike `Ref(array)`, this type carries an explicit batch size and satisfies the
 `AbstractArray` contract honestly.
 
 Use `Ref(array)` when the batch size is unknown or irrelevant (e.g. during broadcast
 setup). Use `SharedCuArray` when you need a proper `AbstractArray` with a known size.
 
+This type is generic over the storage array type so that it can participate in `Adapt.jl`
+transformations. In the user-facing intended usage, `data` is a `CuArray{T,InnerN,M}`.
+
 # Common aliases
-- `SharedCuMatrix{T,M}` = `SharedCuArray{T,2,1,M}`
-- `SharedCuVector{T,M}` = `SharedCuArray{T,1,1,M}`
+- `SharedCuMatrix{T,A}` = `SharedCuArray{T,2,1,A}`
+- `SharedCuVector{T,A}` = `SharedCuArray{T,1,1,A}`
 """
-struct SharedCuArray{T,InnerN,BatchN,M} <: AbstractArray{CuArray{T,InnerN,M},BatchN}
-    data::CuArray{T,InnerN,M}
+struct SharedCuArray{T,InnerN,BatchN,A<:AbstractArray{T,InnerN}} <: AbstractArray{Any,BatchN}
+    data::A
     batchsize::NTuple{BatchN,Int}
 end
 
 # Outer constructor: accept a plain Int for the common 1D-batch case
-function SharedCuArray{T,InnerN,1,M}(data::CuArray{T,InnerN,M}, N::Int) where {T,InnerN,M}
-    return SharedCuArray{T,InnerN,1,M}(data, (N,))
+function SharedCuArray{T,InnerN,1,A}(data::A, N::Int) where {T,InnerN,A<:AbstractArray{T,InnerN}}
+    return SharedCuArray{T,InnerN,1,A}(data, (N,))
 end
 
-const SharedCuMatrix{T,M} = SharedCuArray{T,2,1,M}
-const SharedCuVector{T,M} = SharedCuArray{T,1,1,M}
+const SharedCuMatrix{T,A<:AbstractArray{T,2}} = SharedCuArray{T,2,1,A}
+const SharedCuVector{T,A<:AbstractArray{T,1}} = SharedCuArray{T,1,1,A}
+
+# Constructors for aliased cases
+SharedCuMatrix(data::A, N::Int) where {T,A<:AbstractArray{T,2}} = SharedCuArray{T,2,1,A}(data, N)
+SharedCuVector(data::A, N::Int) where {T,A<:AbstractArray{T,1}} = SharedCuArray{T,1,1,A}(data, N)
 
 const SharedArray = SharedCuArray
 
+Base.eltype(::Type{<:BatchedCuArray{T,NE}}) where {T,NE} = AbstractArray{T,NE}
+Base.eltype(::Type{<:SharedCuArray{T,InnerN}}) where {T,InnerN} = AbstractArray{T,InnerN}
+
 """
-    Shared(data::CuArray, N::Int) -> SharedCuArray
+    Shared(data::AbstractArray, N::Int) -> SharedCuArray
 
-Convenience constructor: create a `SharedCuArray` from a CuArray with an explicit
+Convenience constructor: create a `SharedCuArray` from an arrat with an explicit
 1D batch size `N`.
+
+The underlying storage is generic to support `Adapt.jl` transformations, but in
+the user-facing intended interface `A` is type `CuArray`
 """
-Shared(x::CuArray{T,2,M}, N::Int) where {T,M} = SharedCuArray{T,2,1,M}(x, (N,))
-Shared(x::CuArray{T,1,M}, N::Int) where {T,M} = SharedCuArray{T,1,1,M}(x, (N,))
+Shared(x::A, N::Int) where {T,A<:AbstractArray{T,2}} = SharedCuArray{T,2,1,A}(x, (N,))
+Shared(x::A, N::Int) where {T,A<:AbstractArray{T,1}} = SharedCuArray{T,1,1,A}(x, (N,))
 
 Base.IndexStyle(::Type{<:SharedCuArray}) = Base.IndexCartesian()
 
@@ -129,6 +157,15 @@ end
 
 batch_size(x::SharedCuArray) = length(x)
 
+# Adapting SharedCuArray to bitstype
+function Adapt.adapt_structure(
+    to,
+    x::SharedCuArray{T,InnerN,BatchN,A},
+) where {T,InnerN,BatchN,A<:AbstractArray{T,InnerN}}
+    data_adapted = Adapt.adapt(to, x.data)
+    return SharedCuArray{T,InnerN,BatchN,typeof(data_adapted)}(data_adapted, x.batchsize)
+end
+
 # =============================================================================
 # SharedScalar: a scalar value shared across all batch elements
 # =============================================================================
@@ -149,6 +186,9 @@ Base.:(==)(x::SharedScalar, y) = x.value == y
 Base.:(==)(x, y::SharedScalar) = x == y.value
 Base.:(==)(x::SharedScalar, y::SharedScalar) = x.value == y.value
 
+# Adapting SharedScalar to bitstype
+Adapt.@adapt_structure SharedScalar
+
 # =============================================================================
 # BatchedStruct - Custom wrapper for batched composite types
 # =============================================================================
@@ -271,6 +311,15 @@ function Base.show(io::IO, ::MIME"text/plain", x::BatchedStruct{T}) where {T}
     end
 end
 
+# Adapting BatchedStruct to bitstype
+function Adapt.adapt_structure(
+    to,
+    x::BatchedStruct{T,C},
+) where {T,C<:NamedTuple}
+    comps_adapted = Adapt.adapt(to, x.components)
+    return BatchedStruct{T}(comps_adapted)
+end
+
 # =============================================================================
 # Union Types for Dispatch
 # =============================================================================