[NDTensorsMetalExt] Update for Unwrap module

NDTensors/ext/NDTensorsMetalExt/NDTensorsMetalExt.jl

@@ -2,7 +2,7 @@ module NDTensorsMetalExt
 
 using Adapt
 using Functors
-using LinearAlgebra: LinearAlgebra, Transpose, mul!, qr, eigen, svd
+using LinearAlgebra: LinearAlgebra, Adjoint, Transpose, mul!, qr, eigen, svd
 using NDTensors
 using NDTensors.SetParameters
 using NDTensors.Unwrap: qr_positive, ql_positive, ql
@@ -22,4 +22,5 @@ include("copyto.jl")
 include("append.jl")
 include("permutedims.jl")
 include("mul.jl")
+
 end

NDTensors/ext/NDTensorsMetalExt/adapt.jl

@@ -1,3 +1,5 @@
+NDTensors.cpu(e::Exposed{<:MtlArray}) = adapt(Array, e)
+
 function mtl(xs; storage=DefaultStorageMode)
  return adapt(set_storagemode(MtlArray, storage), xs)
 end

NDTensors/ext/NDTensorsMetalExt/copyto.jl

@@ -1,13 +1,23 @@
-# Catches a bug in `copyto!` in Metal backend.
-function NDTensors.copyto!(
- ::Type{<:MtlArray}, dest::AbstractArray, ::Type{<:MtlArray}, src::SubArray
+function Base.copy(src::Exposed{<:MtlArray,<:Base.ReshapedArray})
+ return reshape(copy(parent(src)), size(unexpose(src)))
+end
+
+function Base.copy(
+ src::Exposed{
+ <:MtlArray,<:SubArray{<:Any,<:Any,<:Base.ReshapedArray{<:Any,<:Any,<:Adjoint}}
+ },
 )
- return Base.copyto!(dest, copy(src))
+ return copy(@view copy(expose(parent(src)))[parentindices(unexpose(src))...])
+end
+
+# Catches a bug in `copyto!` in Metal backend.
+function Base.copyto!(dest::Exposed{<:MtlArray}, src::Exposed{<:MtlArray,<:SubArray})
+ return copyto!(dest, expose(copy(src)))
 end
 
 # Catches a bug in `copyto!` in Metal backend.
-function NDTensors.copyto!(
- ::Type{<:MtlArray}, dest::AbstractArray, ::Type{<:MtlArray}, src::Base.ReshapedArray
+function Base.copyto!(
+ dest::Exposed{<:MtlArray}, src::Exposed{<:MtlArray,<:Base.ReshapedArray}
 )
- return NDTensors.copyto!(dest, parent(src))
+ return copyto!(dest, expose(parent(src)))
 end

NDTensors/ext/NDTensorsMetalExt/indexing.jl

@@ -6,3 +6,8 @@ function Base.setindex!(E::Exposed{<:MtlArray}, x::Number)
  Metal.@allowscalar unexpose(E)[] = x
  return unexpose(E)
 end
+
+# Shared with `CuArray`. Move to `NDTensorsGPUArraysCoreExt`?
+function Base.getindex(E::Exposed{<:MtlArray,<:Adjoint}, i, j)
+ return (expose(parent(E))[j, i])'
+end

NDTensors/ext/NDTensorsMetalExt/linearalgebra.jl

@@ -23,8 +23,9 @@ function LinearAlgebra.eigen(A::Exposed{<:MtlMatrix})
 end
 
 function LinearAlgebra.svd(A::Exposed{<:MtlMatrix}; kwargs...)
- U, S, V = svd(expose(NDTensors.cpu(A)); kwargs...)
- return adapt(unwrap_type(A), U),
- adapt(set_ndims(unwrap_type(A), ndims(S)), S),
- adapt(unwrap_type(A), V)
+ Ucpu, Scpu, Vcpu = svd(expose(NDTensors.cpu(A)); kwargs...)
+ U = adapt(unwrap_type(A), Ucpu)
+ S = adapt(set_ndims(unwrap_type(A), ndims(Scpu)), Scpu)
+ V = adapt(unwrap_type(A), Vcpu)
+ return U, S, V
 end

NDTensors/ext/NDTensorsMetalExt/permutedims.jl

@@ -1,4 +1,4 @@
-function permutedims!(
+function Base.permutedims!(
  Edest::Exposed{<:MtlArray,<:Base.ReshapedArray}, Esrc::Exposed{<:MtlArray}, perm
 )
  Aperm = permutedims(Esrc, perm)

NDTensors/src/Unwrap/src/Unwrap.jl

@@ -3,6 +3,7 @@ using SimpleTraits
 using LinearAlgebra
 using Base: ReshapedArray
 using StridedViews
+using Adapt: Adapt, adapt, adapt_structure
 
 include("expose.jl")
 include("iswrappedarray.jl")
@@ -16,6 +17,7 @@ include("functions/copyto.jl")
 include("functions/linearalgebra.jl")
 include("functions/mul.jl")
 include("functions/permutedims.jl")
+include("functions/adapt.jl")
 
 export IsWrappedArray,
  is_wrapped_array, parenttype, unwrap_type, expose, Exposed, unexpose, cpu

NDTensors/src/Unwrap/src/functions/adapt.jl

@@ -0,0 +1,8 @@
+Adapt.adapt(to, x::Exposed) = adapt_structure(to, x)
+Adapt.adapt_structure(to, x::Exposed) = adapt_structure(to, unexpose(x))
+
+# https://github.com/JuliaGPU/Adapt.jl/pull/51
+# TODO: Remove once https://github.com/JuliaGPU/Adapt.jl/issues/71 is addressed.
+function Adapt.adapt_structure(to, A::Exposed{<:Any,<:Hermitian})
+ return Hermitian(adapt(to, parent(unexpose(A))), Symbol(unexpose(A).uplo))
+end

NDTensors/src/Unwrap/test/runtests.jl

@@ -5,7 +5,9 @@ using LinearAlgebra
 
 include("../../../test/device_list.jl")
 @testset "Testing Unwrap" for dev in devices_list(ARGS)
- v = dev(Vector{Float64}(undef, 10))
+ elt = Float32
+
+ v = dev(Vector{elt}(undef, 10))
  vt = transpose(v)
  va = v'
 
@@ -37,7 +39,7 @@ include("../../../test/device_list.jl")
  @test typeof(Et) == Exposed{m_type,LinearAlgebra.Transpose{e_type,m_type}}
  @test typeof(Ea) == Exposed{m_type,LinearAlgebra.Adjoint{e_type,m_type}}
 
- o = dev(Vector{Float32})(undef, 1)
+ o = dev(Vector{elt})(undef, 1)
  expose(o)[] = 2
  @test expose(o)[] == 2
 
@@ -56,8 +58,8 @@ include("../../../test/device_list.jl")
  q, r = Unwrap.qr_positive(expose(mp))
  @test q * r ≈ mp
 
- square = dev(rand(Float64, (10, 10)))
- square = (square + transpose(square)) ./ 2.0
+ square = dev(rand(elt, (10, 10)))
+ square = (square + transpose(square)) / 2
  ## CUDA only supports Hermitian or Symmetric eigen decompositions
  ## So I symmetrize square and call symetric here
  l, U = eigen(expose(Symmetric(square)))
@@ -66,25 +68,59 @@ include("../../../test/device_list.jl")
  U, S, V, = svd(expose(mp))
  @test U * Diagonal(S) * V' ≈ mp
 
- cm = dev(fill!(Matrix{Float64}(undef, (2, 2)), 0.0))
+ cm = dev(fill!(Matrix{elt}(undef, (2, 2)), 0.0))
  mul!(expose(cm), expose(mp), expose(mp'), 1.0, 0.0)
  @test cm ≈ mp * mp'
 
  @test permutedims(expose(mp), (2, 1)) == transpose(mp)
- fill!(mt, 3.0)
+ fill!(mt, 3)
  permutedims!(expose(m), expose(mt), (2, 1))
- @test norm(m) == sqrt(3^2 * 10)
+ @test norm(m) ≈ sqrt(3^2 * 10)
  @test size(m) == (5, 2)
  permutedims!(expose(m), expose(mt), (2, 1), +)
  @test size(m) == (5, 2)
- @test norm(m) == sqrt(6^2 * 10)
+ @test norm(m) ≈ sqrt(6^2 * 10)
 
  m = reshape(m, (5, 2, 1))
  mt = fill!(similar(m), 3.0)
  m = permutedims(expose(m), (2, 1, 3))
  @test size(m) == (2, 5, 1)
  permutedims!(expose(m), expose(mt), (2, 1, 3))
- @test norm(m) == sqrt(3^2 * 10)
+ @test norm(m) ≈ sqrt(3^2 * 10)
  permutedims!(expose(m), expose(mt), (2, 1, 3), -)
  @test norm(m) == 0
+
+ x = dev(rand(elt, 4, 4))
+ y = dev(rand(elt, 4, 4))
+ copyto!(expose(y), expose(x))
+ @test y == x
+
+ y = dev(rand(elt, 4, 4))
+ x = Base.ReshapedArray(dev(rand(elt, 16)), (4, 4), ())
+ copyto!(expose(y), expose(x))
+ @test NDTensors.cpu(y) == NDTensors.cpu(x)
+ @test NDTensors.cpu(copy(expose(x))) == NDTensors.cpu(x)
+
+ y = dev(rand(elt, 4, 4))
+ x = @view dev(rand(elt, 8, 8))[1:4, 1:4]
+ copyto!(expose(y), expose(x))
+ @test y == x
+ @test copy(x) == x
+
+ y = dev(randn(elt, 16))
+ x = reshape(dev(randn(elt, 4, 4))', 16)
+ copyto!(expose(y), expose(x))
+ @test y == x
+ @test copy(x) == x
+
+ y = dev(randn(elt, 8))
+ x = @view reshape(dev(randn(elt, 8, 8))', 64)[1:8]
+ copyto!(expose(y), expose(x))
+ @test y == x
+ @test copy(x) == x
+
+ y = Base.ReshapedArray(dev(randn(elt, 16)), (4, 4), ())
+ x = dev(randn(elt, 4, 4))
+ permutedims!(expose(y), expose(x), (2, 1))
+ @test NDTensors.cpu(y) == transpose(NDTensors.cpu(x))
 end

NDTensors/src/abstractarray/append.jl

@@ -1,6 +1,7 @@
 # NDTensors.append!
 # Used to circumvent issues with some GPU backends like Metal
 # not supporting `resize!`.
+# TODO: Change this over to use `expose`.
 function append!!(collection, collections...)
  return append!!(unwrap_type(collection), collection, collections...)
 end

NDTensors/src/blocksparse/linearalgebra.jl

@@ -181,7 +181,7 @@ function svd(
  if sU == -1
  Ub *= -1
  end
- copyto!(blockview(U, blockU), Ub)
+ copyto!(expose(blockview(U, blockU)), expose(Ub))
 
  blockviewS = blockview(S, blockS)
  # TODO: Replace `data` with `diagview`.
@@ -200,7 +200,7 @@ function svd(
  if (sV * sVP) == -1
  Vb *= -1
  end
- copyto!(blockview(V, blockV), Vb)
+ copyto!(expose(blockview(V, blockV)), expose(Vb))
  end
  return U, S, V, Spectrum(d, truncerr)
 end