Refactoring GPU extensions

.github/workflows/test_itensors_base_ubuntu.yml

@@ -33,17 +33,13 @@ jobs:
         with:
           version: ${{ matrix.version }}
           arch: ${{ matrix.arch }}
-      - name: Install Julia dependencies
-        shell: julia --project=monorepo {0}
+      - name: Install Julia dependencies and run tests
+        shell: julia {0}
         run: |
           using Pkg;
-          Pkg.develop(path=".");
+          Pkg.activate(temp=true)
           Pkg.develop(path="./NDTensors");
-      - name: Run the tests
-        shell: julia --project=monorepo {0}
-        run: |
-          using Pkg;
-          # https://github.com/JuliaLang/Pkg.jl/pull/1226
+          Pkg.develop(path=".");
           Pkg.test("ITensors"; coverage=true, test_args=["base"])
       - uses: julia-actions/julia-uploadcodecov@latest
         env:

.github/workflows/test_ndtensors.yml

@@ -30,12 +30,11 @@ jobs:
         with:
           version: ${{ matrix.version }}
           arch: ${{ matrix.arch }}
-      - name: Install Julia dependencies
-        shell: julia --project=monorepo {0}
+      - name: Install Julia dependencies and run tests
+        shell: julia --depwarn=yes {0}
         run: |
           using Pkg;
-          Pkg.develop(path=".");
+          Pkg.activate(temp=true);
           Pkg.develop(path="./NDTensors");
-      - name: Run the tests
-        run: |
-          julia --project=monorepo --depwarn=yes -e 'using Pkg; Pkg.test("NDTensors")'
+          Pkg.develop(path=".");
+          Pkg.test("NDTensors");

NDTensors/ext/NDTensorsAMDGPUExt/NDTensorsAMDGPUExt.jl

@@ -1,5 +1,6 @@
 module NDTensorsAMDGPUExt
 
+include("append.jl")
 include("copyto.jl")
 include("set_types.jl")
 include("adapt.jl")

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/NDTensorsAMDGPUExt/append.jl

@@ -0,0 +1,8 @@
+using GPUArraysCore: @allowscalar
+using AMDGPU: ROCArray
+using NDTensors.Expose: Exposed, unexpose
+
+## Warning this append function uses scalar indexing and is therefore extremely slow
+function Base.append!(Ecollection::Exposed{<:ROCArray}, collections...)
+  return @allowscalar append!(unexpose(Ecollection), collections...)
+end

NDTensors/ext/NDTensorsCUDAExt/NDTensorsCUDAExt.jl

@@ -1,15 +1,5 @@
 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("append.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/append.jl

@@ -0,0 +1,8 @@
+using GPUArraysCore: @allowscalar
+using CUDA: CuArray
+using NDTensors.Expose: Exposed, unexpose
+
+## Warning this append function uses scalar indexing and is therefore extremely slow
+function Base.append!(Ecollection::Exposed{<:CuArray}, collections...)
+  return @allowscalar append!(unexpose(Ecollection), collections...)
+end

NDTensors/ext/NDTensorsCUDAExt/copyto.jl

@@ -1,3 +1,7 @@
+using CUDA: CuArray
+using NDTensors.Expose: Exposed, expose, 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,14 @@
+using CUDA: CuArray
+using GPUArraysCore: @allowscalar
+using NDTensors: NDTensors
+using NDTensors.Expose: Exposed, expose, unexpose
+
 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, 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, 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/ext/NDTensorsMetalExt/NDTensorsMetalExt.jl

@@ -1,14 +1,5 @@
 module NDTensorsMetalExt
 
-using Adapt
-using Functors
-using LinearAlgebra: LinearAlgebra, Adjoint, Transpose, mul!, qr, eigen, svd
-using NDTensors
-using NDTensors.Expose: qr_positive, ql_positive, ql
-
-using Metal
-
-include("imports.jl")
 include("adapt.jl")
 include("set_types.jl")
 include("indexing.jl")

NDTensors/ext/NDTensorsMetalExt/adapt.jl

@@ -1,17 +1,29 @@
-using NDTensors.MetalExtensions: MetalExtensions
-using NDTensors.GPUArraysCoreExtensions: GPUArraysCoreExtensions, set_storagemode
-using NDTensors.TypeParameterAccessors: specify_type_parameters, type_parameters
+using Adapt: Adapt, adapt
+using Functors: fmap
+using Metal: MtlArray, MtlVector, DefaultStorageMode
+using NDTensors: NDTensors, EmptyStorage, adapt_storagetype, emptytype
+using NDTensors.Expose: Exposed
+using NDTensors.MetalExtensions: MetalExtensions, MtlArrayAdaptor
+using NDTensors.GPUArraysCoreExtensions: GPUArraysCoreExtensions
+using NDTensors.TypeParameterAccessors: set_type_parameters, type_parameters
 
 GPUArraysCoreExtensions.cpu(e::Exposed{<:MtlArray}) = adapt(Array, e)
 
-function MetalExtensions.mtl(xs; storage=DefaultStorageMode)
-  return adapt(set_storagemode(MtlArray, storage), xs)
+function MetalExtensions.mtl(xs; storagemode=DefaultStorageMode)
+  return fmap(x -> adapt(MtlArrayAdaptor{storagemode}(), x), xs)
 end
 
-# More general than the version in Metal.jl
-## TODO Rewrite this using a custom `MtlArrayAdaptor` which will be written in  `MetalExtensions`.
-function Adapt.adapt_storage(arraytype::Type{<:MtlArray}, xs::AbstractArray)
-  params = type_parameters(xs)
-  arraytype_specified = specify_type_parameters(arraytype, params)
-  return isbitstype(typeof(xs)) ? xs : convert(arraytype_specified, xs)
+function Adapt.adapt_storage(adaptor::MtlArrayAdaptor, xs::AbstractArray)
+  new_parameters = (type_parameters(xs, (eltype, ndims))..., storagemode(adaptor))
+  mtltype = set_type_parameters(MtlArray, (eltype, ndims, storagemode), new_parameters)
+  return isbits(xs) ? xs : adapt(mtltype, xs)
+end
+
+function NDTensors.adapt_storagetype(
+  adaptor::MtlArrayAdaptor, ::Type{EmptyStorage{ElT,StoreT}}
+) where {ElT,StoreT}
+  mtltype = set_type_parameters(
+    MtlVector, (eltype, storagemode), (ElT, storagemode(adaptor))
+  )
+  return emptytype(adapt_storagetype(mtltype, StoreT))
 end

NDTensors/ext/NDTensorsMetalExt/append.jl

@@ -1,5 +1,10 @@
-# This circumvents an issues that `MtlArray` can't call `resize!`.
-# TODO: Raise an issue with Metal.jl.
-function NDTensors.append!!(::Type{<:MtlArray}, collection, collections...)
-  return vcat(collection, collections...)
+## Right now append! is broken on metal because of a missing resize! function
+## but make this available in the next release this will allow metal to work working
+using GPUArraysCore: @allowscalar
+using Metal: MtlArray
+using NDTensors.Expose: Exposed, unexpose
+
+## Warning this append function uses scalar indexing and is therefore extremely slow
+function Base.append!(Ecollection::Exposed{<:MtlArray}, collections...)
+  return @allowscalar append!(unexpose(Ecollection), collections...)
 end

NDTensors/ext/NDTensorsMetalExt/copyto.jl

@@ -1,3 +1,6 @@
+using Metal: MtlArray
+using NDTensors.Expose: Exposed, expose, unexpose
+
 function Base.copy(src::Exposed{<:MtlArray,<:Base.ReshapedArray})
   return reshape(copy(parent(src)), size(unexpose(src)))
 end

NDTensors/ext/NDTensorsMetalExt/indexing.jl

@@ -1,9 +1,14 @@
+using Metal: MtlArray
+using GPUArraysCore: @allowscalar
+using LinearAlgebra: Adjoint
+using NDTensors.Expose: Exposed, expose, unexpose
+
 function Base.getindex(E::Exposed{<:MtlArray})
-  return Metal.@allowscalar unexpose(E)[]
+  return @allowscalar unexpose(E)[]
 end
 
 function Base.setindex!(E::Exposed{<:MtlArray}, x::Number)
-  Metal.@allowscalar unexpose(E)[] = x
+  @allowscalar unexpose(E)[] = x
   return unexpose(E)
 end
 

NDTensors/ext/NDTensorsMetalExt/linearalgebra.jl

@@ -1,3 +1,6 @@
+using Metal: MtlMatrix
+using LinearAlgebra: LinearAlgebra, qr, eigen, svd
+using NDTensors.Expose: qr_positive, ql_positive, ql
 using NDTensors.TypeParameterAccessors:
   set_type_parameters, type_parameters, unwrap_array_type
 

NDTensors/ext/NDTensorsMetalExt/mul.jl

@@ -1,3 +1,5 @@
+using Metal: MtlArray
+using LinearAlgebra: LinearAlgebra, Adjoint, Transpose, mul!
 # This was calling generic matrix multiplication.
 # TODO: Raise an issue with `Metal.jl`.
 function LinearAlgebra.mul!(

NDTensors/ext/NDTensorsMetalExt/permutedims.jl

@@ -1,3 +1,4 @@
+using Metal: MtlArray
 ## Theres an issue in metal that `ReshapedArray' wrapped arrays cannot be permuted using
 ## permutedims (failing in that Metal uses scalar indexing)
 ## These functions are to address the problem in different instances of permutedims

NDTensors/ext/NDTensorsMetalExt/set_types.jl

@@ -1,3 +1,4 @@
+using Metal: Metal, MtlArray
 # `TypeParameterAccessors.jl` definitions.
 
 using NDTensors.TypeParameterAccessors: TypeParameterAccessors, Position, set_type_parameter