Refactoring GPU extensions

NDTensors/ext/NDTensorsAMDGPUExt/adapt.jl

@@ -22,7 +22,7 @@ function Adapt.adapt_storage(adaptor::ROCArrayAdaptor, xs::AbstractArray)
 end
 
 function NDTensors.adapt_storagetype(
-  adaptor::ROCArrayAdaptor, xs::Type{EmptyStorage{ElT,StoreT}}
+  adaptor::ROCArrayAdaptor, ::Type{EmptyStorage{ElT,StoreT}}
 ) where {ElT,StoreT}
   roctype = set_type_parameters(
     ROCVector, (eltype, storagemode), (ElT, storagemode(adaptor))

NDTensors/ext/NDTensorsCUDAExt/NDTensorsCUDAExt.jl

@@ -1,15 +1,4 @@
 module NDTensorsCUDAExt
-
-using NDTensors
-using NDTensors.Expose
-using Adapt
-using Functors
-using LinearAlgebra: LinearAlgebra, Adjoint, Transpose, mul!, svd
-using CUDA
-using CUDA.CUBLAS
-using CUDA.CUSOLVER
-
-include("imports.jl")
 include("default_kwargs.jl")
 include("copyto.jl")
 include("set_types.jl")

NDTensors/ext/NDTensorsCUDAExt/adapt.jl

@@ -1,24 +1,26 @@
-using NDTensors.TypeParameterAccessors: TypeParameterAccessors
-using NDTensors.GPUArraysCoreExtensions: storagemode
+using Adapt: Adapt
+using CUDA: CUDA, CuArray, CuVector
+using Functors: fmap
+using NDTensors: NDTensors, EmptyStorage, adapt_storagetype, emptytype
 using NDTensors.CUDAExtensions: CUDAExtensions, CuArrayAdaptor
+using NDTensors.GPUArraysCoreExtensions: storagemode
+using NDTensors.TypeParameterAccessors:
+  TypeParameterAccessors, default_type_parameter, set_type_parameters, type_parameters
 
-## TODO make this work for unified. This works but overwrites CUDA's adapt_storage. This fails for emptystorage...
-function CUDAExtensions.cu(xs; unified::Bool=false)
-  return fmap(
-    x -> adapt(CuArrayAdaptor{unified ? Mem.UnifiedBuffer : Mem.DeviceBuffer}(), x), xs
-  )
+function CUDAExtensions.cu(xs; storagemode=default_type_parameter(CuArray, storagemode))
+  return fmap(x -> adapt(CuArrayAdaptor{storagemode}(), x), xs)
 end
 
+## Could do this generically
 function Adapt.adapt_storage(adaptor::CuArrayAdaptor, xs::AbstractArray)
-  ElT = eltype(xs)
-  BufT = storagemode(adaptor)
-  N = ndims(xs)
-  return isbits(xs) ? xs : adapt(CuArray{ElT,N,BufT}, xs)
+  params = (type_parameters(xs, (eltype, ndims))..., storagemode(adaptor))
+  cutype = set_type_parameters(CuArray, (eltype, ndims, storagemode), params)
+  return isbits(xs) ? xs : adapt(cutype, xs)
 end
 
 function NDTensors.adapt_storagetype(
-  adaptor::CuArrayAdaptor, xs::Type{EmptyStorage{ElT,StoreT}}
+  adaptor::CuArrayAdaptor, ::Type{EmptyStorage{ElT,StoreT}}
 ) where {ElT,StoreT}
-  BufT = storagemode(adaptor)
-  return NDTensors.emptytype(NDTensors.adapt_storagetype(CuVector{ElT,BufT}, StoreT))
+  cutype = set_type_parameters(CuVector, (eltype, storagemode), (ElT, storagemode(adaptor)))
+  return emptytype(adapt_storagetype(cutype, StoreT))
 end

NDTensors/ext/NDTensorsCUDAExt/copyto.jl

@@ -1,3 +1,7 @@
+using CUDA: CuArray
+using NDTensors.Expose: Exposed, expose, parent, unexpose
+using LinearAlgebra: Adjoint
+
 # Same definition as `MtlArray`.
 function Base.copy(src::Exposed{<:CuArray,<:Base.ReshapedArray})
   return reshape(copy(parent(src)), size(unexpose(src)))

NDTensors/ext/NDTensorsCUDAExt/default_kwargs.jl

@@ -1 +1,4 @@
+using CUDA: CuArray
+using NDTensors: NDTensors
+
 NDTensors.default_svd_alg(::Type{<:CuArray}, a) = "qr_algorithm"

NDTensors/ext/NDTensorsCUDAExt/indexing.jl

@@ -1,9 +1,15 @@
+using CUDA: CuArray
+using GPUArraysCore: @allowscalar
+using NDTensors: NDTensors
+using NDTensors.Expose: Exposed, expose, parent, unexpose
+
+using GPUArraysCore: @allowscalar
 function Base.getindex(E::Exposed{<:CuArray})
-  return CUDA.@allowscalar unexpose(E)[]
+  return @allowscalar unexpose(E)[]
 end
 
 function Base.setindex!(E::Exposed{<:CuArray}, x::Number)
-  CUDA.@allowscalar unexpose(E)[] = x
+  @allowscalar unexpose(E)[] = x
   return unexpose(E)
 end
 

NDTensors/ext/NDTensorsCUDAExt/iscu.jl

@@ -1 +1,4 @@
-iscu(::Type{<:CuArray}) = true
+using CUDA: CuArray
+using NDTensors: NDTensors
+
+NDTensors.iscu(::Type{<:CuArray}) = true

NDTensors/ext/NDTensorsCUDAExt/linearalgebra.jl

@@ -1,3 +1,10 @@
+using Adapt: adapt
+using CUDA: CUDA, CuMatrix
+using LinearAlgebra: Adjoint, svd
+using NDTensors: NDTensors
+using NDTensors.Expose: Expose, expose, ql, ql_positive
+using NDTensors.GPUArraysCoreExtensions: cpu
+using NDTensors.TypeParameterAccessors: unwrap_array_type
 function NDTensors.svd_catch_error(A::CuMatrix; alg::String="jacobi_algorithm")
   if alg == "jacobi_algorithm"
     alg = CUDA.CUSOLVER.JacobiAlgorithm()
@@ -42,9 +49,6 @@ function NDTensors.svd_catch_error(A::CuMatrix, ::CUDA.CUSOLVER.QRAlgorithm)
   return USV
 end
 
-using NDTensors.GPUArraysCoreExtensions: cpu
-using NDTensors.Expose: Expose, expose, ql, ql_positive
-using NDTensors.TypeParameterAccessors: unwrap_array_type
 ## TODO currently AMDGPU doesn't have ql so make a ql function
 function Expose.ql(A::Exposed{<:CuMatrix})
   Q, L = ql(expose(cpu(A)))

NDTensors/ext/NDTensorsCUDAExt/mul.jl

@@ -1,3 +1,7 @@
+using CUDA: CuArray
+using LinearAlgebra: LinearAlgebra, mul!, transpose
+using NDTensors.Expose: Exposed, expose, parent, unexpose
+
 # This was calling generic matrix multiplication.
 # TODO: Raise an issue with `CUDA.jl`.
 function LinearAlgebra.mul!(

NDTensors/ext/NDTensorsCUDAExt/permutedims.jl

@@ -1,3 +1,6 @@
+using CUDA: CuArray
+using NDTensors.Expose: Exposed, expose, parent, unexpose
+
 function Base.permutedims!(
   Edest::Exposed{<:CuArray,<:Base.ReshapedArray}, Esrc::Exposed{<:CuArray}, perm
 )

NDTensors/ext/NDTensorsCUDAExt/set_types.jl

@@ -1,4 +1,5 @@
 # TypeParameterAccessors definitions
+using CUDA: CUDA, CuArray
 using NDTensors.TypeParameterAccessors: TypeParameterAccessors, Position
 using NDTensors.GPUArraysCoreExtensions: storagemode
 

NDTensors/src/blocksparse/blocksparsetensor.jl

@@ -256,6 +256,7 @@ end
 # Returns the offset of the new block added.
 # XXX rename to insertblock!, no need to return offset
 using .TypeParameterAccessors: unwrap_array_type
+using .Expose: expose
 function insertblock_offset!(T::BlockSparseTensor{ElT,N}, newblock::Block{N}) where {ElT,N}
   newdim = blockdim(T, newblock)
   newoffset = nnz(T)
@@ -264,7 +265,7 @@ function insertblock_offset!(T::BlockSparseTensor{ElT,N}, newblock::Block{N}) wh
   new_data = generic_zeros(unwrap_array_type(T), newdim)
   # TODO: `append!` is broken on `Metal` since `resize!`
   # isn't implemented.
-  append!(data(T), new_data)
+  append!(expose(data(T)), new_data)
   return newoffset
 end
 
@@ -725,6 +726,7 @@ end
 # <fermions>
 permfactor(perm, block, inds) = 1
 
+using .TypeParameterAccessors: set_type_parameters, parenttype
 function permutedims!(
   R::BlockSparseTensor{<:Number,N},
   T::BlockSparseTensor{<:Number,N},
@@ -751,17 +753,20 @@ function permutedims!(
       # Rblock doesn't exist
       block_size = permute(size(Tblock), perm)
       # TODO: Make GPU friendly.
-      Rblock = tensor(Dense(zeros(eltype(R), block_size)), block_size)
+      DenseT = set_type_parameters(Dense, (eltype, parenttype), (eltype(R), datatype(R)))
+      Rblock = tensor(generic_zeros(DenseT, block_size), block_size)
     elseif !Tblock_exists
       # Tblock doesn't exist
       block_size = permute(size(Rblock), invperm(perm))
       # TODO: Make GPU friendly.
-      Tblock = tensor(Dense(zeros(eltype(T), block_size)), block_size)
+      DenseT = set_type_parameters(Dense, (eltype, parenttype), (eltype(T), datatype(T)))
+      Tblock = tensor(generic_zeros(DenseT, block_size), block_size)
     end
     permutedims!(Rblock, Tblock, perm, f_fac)
     if !Rblock_exists
       # Set missing nonzero block
-      if !iszero(Rblock)
+      ## To make sure no allowscalar issue grab the data
+      if !iszero(data(Rblock))
         R[block] = Rblock
       end
     end

NDTensors/src/blocksparse/linearalgebra.jl

@@ -1,4 +1,5 @@
 using .TypeParameterAccessors: unwrap_array_type
+using .Expose: expose
 const BlockSparseMatrix{ElT,StoreT,IndsT} = BlockSparseTensor{ElT,2,StoreT,IndsT}
 const DiagBlockSparseMatrix{ElT,StoreT,IndsT} = DiagBlockSparseTensor{ElT,2,StoreT,IndsT}
 const DiagMatrix{ElT,StoreT,IndsT} = DiagTensor{ElT,2,StoreT,IndsT}
@@ -68,7 +69,7 @@ function svd(
     # TODO: call this a function `diagonal`, i.e.:
     # https://github.com/JuliaLang/julia/issues/30250
     # or make `diag(::Tensor)` return a view by default.
-    append!(d, data(Sb))
+    append!(expose(d), data(Sb))
   end
 
   # Square the singular values to get
@@ -234,14 +235,14 @@ function LinearAlgebra.eigen(
   Db, Vb = eigen(expose(blockT))
   Ds = [Db]
   Vs = [Vb]
-  append!(d, abs.(data(Db)))
+  append!(expose(d), abs.(data(Db)))
   for (n, b) in enumerate(eachnzblock(T))
     n == 1 && continue
     blockT = blockview(T, b)
     Db, Vb = eigen(expose(blockT))
     push!(Ds, Db)
     push!(Vs, Vb)
-    append!(d, abs.(data(Db)))
+    append!(expose(d), abs.(data(Db)))
   end
 
   dropblocks = Int[]

NDTensors/src/dense/fill.jl

@@ -22,18 +22,27 @@ function generic_randn(StoreT::Type{<:Dense}, dim::Integer=0)
 end
 
 function generic_zeros(
-  StoreT::Type{<:Dense{ElT,DataT}}, dim::Integer=0
+  StoreT::Type{<:Dense{ElT,DataT}}, dim::Integer
 ) where {DataT<:AbstractArray,ElT}
   @assert ElT == eltype(DataT)
   data = generic_zeros(DataT, dim)
   StoreT = set_datatype(StoreT, typeof(data))
   return StoreT(data)
 end
 
-function generic_zeros(StoreT::Type{<:Dense{ElT}}, dim::Integer=0) where {ElT}
+function generic_zeros(
+  StoreT::Type{<:Dense{ElT,DataT}}, dims
+) where {DataT<:AbstractArray,ElT}
+  @assert ElT == eltype(DataT)
+  data = generic_zeros(DataT, dim(dims))
+  StoreT = set_datatype(StoreT, typeof(data))
+  return StoreT(data)
+end
+
+function generic_zeros(StoreT::Type{<:Dense{ElT}}, dim) where {ElT}
   return generic_zeros(default_storagetype(ElT), dim)
 end
 
-function generic_zeros(StoreT::Type{<:Dense}, dim::Integer=0)
+function generic_zeros(StoreT::Type{<:Dense}, dim)
   return generic_zeros(default_storagetype(), dim)
 end

NDTensors/src/imports.jl

@@ -28,11 +28,11 @@ for lib in [
   :BaseExtensions,
   :UnspecifiedTypes,
   :TypeParameterAccessors,
+  :Expose,
   :GPUArraysCoreExtensions,
   :AMDGPUExtensions,
   :CUDAExtensions,
   :MetalExtensions,
-  :Expose,
   :BroadcastMapConversion,
   :RankFactorization,
   :Sectors,

NDTensors/src/lib/Expose/src/Expose.jl

@@ -12,6 +12,7 @@ include("import.jl")
 ## using that type
 ## Exposed based functions
 include("functions/abstractarray.jl")
+include("functions/append.jl")
 include("functions/copyto.jl")
 include("functions/linearalgebra.jl")
 include("functions/mul.jl")

NDTensors/src/lib/Expose/src/functions/abstractarray.jl

@@ -1,13 +1,8 @@
-using NDTensors.GPUArraysCoreExtensions: GPUArraysCoreExtensions, cpu
-
 parent(E::Exposed) = parent(unexpose(E))
 
 transpose(E::Exposed) = transpose(unexpose(E))
 
 adjoint(E::Exposed) = adjoint(unexpose(E))
-
-GPUArraysCoreExtensions.cpu(E::Exposed) = cpu(unexpose(E))
-
 getindex(E::Exposed) = unexpose(E)[]
 
 function setindex!(E::Exposed, x::Number)

NDTensors/src/lib/Expose/src/functions/append.jl

@@ -0,0 +1,3 @@
+function Base.append!(Ecollection::Exposed, collections...)
+  return append!(unexpose(Ecollection), collections...)
+end

NDTensors/src/lib/Expose/test/runtests.jl

@@ -17,6 +17,7 @@ using LinearAlgebra:
 using GPUArraysCore: @allowscalar
 include(joinpath(pkgdir(NDTensors), "test", "NDTensorsTestUtils", "NDTensorsTestUtils.jl"))
 using .NDTensorsTestUtils: devices_list
+using NDTensors.GPUArraysCoreExtensions: cpu
 
 @testset "Testing Expose $dev, $elt" for dev in devices_list(ARGS),
   elt in (Float32, ComplexF32)
@@ -235,5 +236,26 @@ using .NDTensorsTestUtils: devices_list
   zero(C)
   mul!!(C, B, A, true, false)
   @test cpu(C) ≈ Cp
+
+  ##################################
+  ### Add test for append! to address scalar indexing in GPUs
+
+  A = dev(randn(elt, 10))
+  Ap = copy(A)
+  B = randn(elt, 3)
+  C = append!(expose(A), B)
+
+  @test length(C) == 13
+  @test sum(C) ≈ sum(Ap) + sum(B)
+
+  A = Ap
+  B = dev(randn(elt, 29))
+  Bp = copy(B)
+  C = append!(expose(B), A)
+  @test length(C) == 39
+  @test sum(C) ≈ sum(Bp) + sum(Ap)
+  @allowscalar for i in 1:length(B)
+    C[i] == B[i]
+  end
 end
 end

NDTensors/src/lib/GPUArraysCoreExtensions/src/gpuarrayscore.jl

@@ -1,3 +1,5 @@
+using GPUArraysCore: AbstractGPUArray, @allowscalar
+using NDTensors.Expose: Exposed, unexpose
 using NDTensors.TypeParameterAccessors:
   TypeParameterAccessors, type_parameter, set_type_parameter
 
@@ -13,3 +15,9 @@ function set_storagemode(type::Type, param)
 end
 
 function cpu end
+
+cpu(E::Exposed) = cpu(unexpose(E))
+
+function Base.append!(Ecollection::Exposed{<:AbstractGPUArray}, collections...)
+  return @allowscalar append!(unexpose(Ecollection), collections...)
+end

NDTensors/test/ITensors/TestITensorDMRG/TestITensorDMRG.jl

@@ -17,9 +17,9 @@ is_broken(dev, elt::Type, conserve_qns::Val) = false
 
 ## Disable blocksparse GPU testing on CUDA, Metal and ROC backends.
 ## Currently Metal fails because we are waiting for `resize!` to be added. Should be in the next metal release
-## CUDA fails because there is no defined `append!`.
 ## ROC fails because TODO determine again why roc fails.
-is_broken(dev::typeof(cu), elt::Type, conserve_qns::Val{true}) = true
+
+#is_broken(dev::typeof(cu), elt::Type, conserve_qns::Val{true}) = true
 is_broken(dev::typeof(mtl), elt::Type, conserve_qns::Val{true}) = true
 is_broken(dev::typeof(roc), elt::Type, conserve_qns::Val{true}) = true
 

src/imports.jl

@@ -110,7 +110,7 @@ import LinearAlgebra:
   tr,
   transpose
 
-using ITensors.NDTensors.Expose: cpu
+using ITensors.NDTensors.GPUArraysCoreExtensions: cpu
 
 using ITensors.NDTensors:
   Algorithm,