[NDTensorsMetalExt] Fix more scalar indexing issues on Metal

NDTensors/ext/NDTensorsCUDAExt/permutedims.jl

@@ -5,3 +5,16 @@ function Base.permutedims!(
   copyto!(expose(parent(Edest)), expose(Aperm))
   return unexpose(Edest)
 end
+
+## Found an issue in CUDA where if Edest is a reshaped{<:Adjoint}
+## .= can fail. So instead force Esrc into the shape of parent(Edest)
+function Base.permutedims!(
+  Edest::Exposed{<:CuArray,<:Base.ReshapedArray{<:Any,<:Any,<:Adjoint}},
+  Esrc::Exposed{<:CuArray},
+  perm,
+  f,
+)
+  Aperm = reshape(permutedims(Esrc, perm), size(parent(Edest)))
+  parent(Edest) .= f.(parent(Edest), Aperm)
+  return unexpose(Edest)
+end

NDTensors/ext/NDTensorsMetalExt/mul.jl

@@ -19,3 +19,22 @@ function LinearAlgebra.mul!(
   mul!(CM', BM', AM', α, β)
   return unexpose(CM)
 end
+
+## Fix issue in Metal.jl where it cannot distinguish Transpose{Reshape{Adjoint{MtlArray}}}
+## as a MtlArray and calls generic matmul
+function LinearAlgebra.mul!(
+  CM::Exposed{<:MtlArray},
+  AM::Exposed{<:MtlArray},
+  BM::Exposed{
+    <:MtlArray,
+    <:LinearAlgebra.Transpose{
+      <:Any,<:Base.ReshapedArray{<:Any,<:Any,<:LinearAlgebra.Adjoint}
+    },
+  },
+  α,
+  β,
+)
+  B = copy(expose(parent(BM)))
+  mul!(CM, AM, expose(transpose(B)), α, β)
+  return unexpose(CM)
+end

NDTensors/ext/NDTensorsMetalExt/permutedims.jl

@@ -1,7 +1,37 @@
+## 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
+function Base.permutedims(E::Exposed{<:MtlArray,<:Base.ReshapedArray}, perm)
+  A = copy(E)
+  return permutedims(A, perm)
+end
+
 function Base.permutedims!(
   Edest::Exposed{<:MtlArray,<:Base.ReshapedArray}, Esrc::Exposed{<:MtlArray}, perm
 )
   Aperm = permutedims(Esrc, perm)
   copyto!(expose(parent(Edest)), expose(Aperm))
   return unexpose(Edest)
 end
+
+function Base.permutedims!(
+  Edest::Exposed{<:MtlArray}, Esrc::Exposed{<:MtlArray,<:Base.ReshapedArray}, perm
+)
+  Aperm = permutedims(Esrc, perm)
+  copyto!(Edest, expose(Aperm))
+  return unexpose(Edest)
+end
+
+## To get around the Metal issue here we copy and permute Esrc,
+## then we reshape Esrc to the size of Edest's parent
+## and broadcast into the parent.
+function Base.permutedims!(
+  Edest::Exposed{<:MtlArray,<:Base.ReshapedArray},
+  Esrc::Exposed{<:MtlArray,<:Base.ReshapedArray},
+  perm,
+  f,
+)
+  Aperm = reshape(permutedims(Esrc, perm), size(parent(Edest)))
+  parent(Edest) .= f.(parent(Edest), Aperm)
+  return unexpose(Edest)
+end

NDTensors/ext/examples/NDTensorCUDA.jl

@@ -1,13 +1,14 @@
-using CUDA
 using NDTensors
-
-using ITensors
-using Test
-
-using Zygote
+using CUDA: CUDA, CuVector, cu, reshape
+using ITensors:
+  Index, ITensor, randomMPO, randomMPS, inner, orthogonalize, qr, siteinds, svd
+using Test: @test
+using Zygote: gradient
 
 function main()
   # using ITensorGPU
+  cpu = NDTensors.cpu
+  gpu = NDTensors.cu
   # Here is an example of how to utilize NDTensors based tensors with CUDA datatypes
   i = Index(2)
   j = Index(5)
@@ -18,10 +19,9 @@ function main()
   dim2 = (j, k)
 
   # Create  2 ITensors with CUDA backends (These will be made simpiler by randomITensor(CuVector) soon)
-  A = ITensor(NDTensors.generic_randn(CuVector, dim(dim1)), dim1)
-  B = ITensor(NDTensors.generic_randn(CuVector, dim(dim2)), dim2)
+  A = ITensor(randomTensor(CuVector, dim1))
+  B = ITensor(randomTensor(CuVector, dim2))
   # Contract the two tensors
-  cpu = NDTensors.cpu
   C = A * B
   A = cpu(A)
   B = cpu(B)
@@ -36,8 +36,8 @@ function main()
   fill!(B, randn())
 
   # Convert the ITensors to GPU
-  cA = NDTensors.cu(A)
-  cB = NDTensors.cu(B)
+  cA = gpu(A)
+  cB = gpu(B)
 
   #Check that backend of contraction is GPU
   @test A * A ≈ cpu(cA * cA)
@@ -47,11 +47,8 @@ function main()
 
   dim3 = (l, k)
   dim4 = (i,)
-  cC = ITensor(
-    NDTensors.generic_randn(CuVector{Float64,CUDA.Mem.DeviceBuffer}, dim(dim3)), dim3
-  )
-  cC = NDTensors.cu(ITensor(NDTensors.generic_randn(Vector{Float64}, dim(dim3)), dim3))
-  cD = ITensor(Tensor(CuVector, dim4))
+  cC = ITensor(randomTensor(CuVector{Float64,CUDA.Mem.DeviceBuffer}, dim3))
+  cD = ITensor(Tensor(CuVector{Float32}, dim4))
   fill!(cD, randn())
 
   # Create a function of 4 tensors on GPU
@@ -61,20 +58,17 @@ function main()
   #Currently this code fails with CUDA.allowscalar(false)
   # Because of outer calling the _gemm! function which calls a 
   # generic implementation
-  @allowscalar grad = gradient(f, cA, cB, cC, cD)
-  @allowscalar @test NDTensors.cpu(cB * cC * cD) ≈ NDTensors.cpu(grad[1])
-  @allowscalar @test (cB * cC * cD) ≈ grad[1]
+  grad = gradient(f, cA, cB, cC, cD)
+  @test cpu(cB * cC * cD) ≈ cpu(grad[1])
+  @test (cB * cC * cD) ≈ grad[1]
   # Create a tuple of indices
-  decomp = (
-    dim(NDTensors.ind(grad[1], 1)),
-    dim(NDTensors.ind(grad[1], 2)) * dim(NDTensors.ind(grad[1], 3)),
-  )
+  decomp = (dim(ind(grad[1], 1)), dim(ind(grad[1], 2)) * dim(ind(grad[1], 3)))
   # Reshape the CuVector of data into a matrix
-  cuTensor_data = CUDA.reshape(NDTensors.data(storage(grad[1])), decomp)
+  cuTensor_data = reshape(array(grad[1]), decomp)
   # Use cuBLAS to compute SVD of data
   U, S, V = svd(cuTensor_data)
-  decomp = (dim(NDTensors.ind(grad[2], 1)), dim(NDTensors.ind(grad[2], 2)))
-  cuTensor_data = CUDA.reshape(NDTensors.data(storage(grad[2])), decomp)
+  decomp = size(array(grad[2]))
+  cuTensor_data = reshape(array(grad[2]), decomp)
   U, S, V = svd(cuTensor_data)
 
   # These things can take up lots of memory, look at memory usage here
@@ -87,33 +81,33 @@ function main()
   CUDA.memory_status()
 
   # Its possible to compute QR of GPU tensor
-  cq = ITensors.qr(cA, (i,), (j, l))
-  q = ITensors.qr(A, (i,), (j, l))
+  cq = qr(cA, (i,), (j, l))
   A ≈ cpu(cq[1]) * cpu(cq[2])
 
   ## SVD does not yet work with CUDA backend, see above on
   ## Converting ITensors to vectors and calling CUDA svd function
   ## CuVectors...
   #ITensors.svd(A, (i,), (j, l))
 
-  s = ITensors.siteinds("S=1/2", 8)
+  s = siteinds("S=1/2", 8)
   m = randomMPS(s; linkdims=4)
-  cm = NDTensors.cu(m)
+  cm = gpu(m)
 
   @test inner(cm', cm) ≈ inner(m', m)
 
   H = randomMPO(s)
-  cH = NDTensors.cu(H)
+  cH = gpu(H)
   @test inner(cm', cH, cm) ≈ inner(m', H, m)
 
   m = orthogonalize(m, 1)
-  cm = NDTensors.cu(orthogonalize(cm, 1))
+  cm = gpu(orthogonalize(cm, 1))
   @test inner(m', m) ≈ inner(cm', cm)
 
   H = orthogonalize(H, 1)
-  cH = NDTensors.cu(cH)
+  cH = gpu(cH)
 
   @test inner(cm', cH, cm) ≈ inner(m', H, m)
 end
 
+## running the main function with Float64
 main()

NDTensors/ext/examples/NDTensorMetal.jl

@@ -1,11 +1,13 @@
-using Metal
+using Metal: MtlVector, mtl
 using NDTensors
 
-using ITensors
-using Test
-using Zygote
+using ITensors: ITensor, Index, randomITensor
+using Test: @test
+using Zygote: gradient
 
 function main()
+  cpu = NDTensors.cpu
+  gpu = NDTensors.mtl
   # Here is an example of how to utilize NDTensors based tensors with CUDA datatypes
   i = Index(20)
   j = Index(5)
@@ -15,27 +17,26 @@ function main()
   dim1 = (i, j, l)
   dim2 = (j, k)
 
-  cA = ITensor(NDTensors.generic_randn(MtlVector{Float32}, dim(dim1)), dim1)
-  cB = ITensor(NDTensors.generic_randn(MtlVector{Float32}, dim(dim2)), dim2)
+  ## MtlArrays only support Float32 arithmatic
+  cA = ITensor(randomTensor(MtlVector{Float32}, dim1))
+  cB = ITensor(randomTensor(MtlVector{Float32}, dim2))
   cC = cA * cB
 
-  cpu = NDTensors.cpu
   A = cpu(cA)
   B = cpu(cB)
 
   @test A * B ≈ cpu(cC)
 
-  #C = A * B
-
   dim3 = (l, k)
   dim4 = (i,)
 
-  cC = mtl(randomITensor(Float32, dim3))
-  cD = mtl(randomITensor(Float32, dim4))
+  cC = gpu(randomITensor(Float32, dim3))
+  cD = gpu(randomITensor(Float32, dim4))
 
   f(A, B, C, D) = (A * B * C * D)[]
 
-  return grad = gradient(f, cA, cB, cC, cD)
+  grad = gradient(f, cA, cB, cC, cD)
+  @test grad[2] ≈ cA * cC * cD
 end
 
 main()

NDTensors/src/blocksparse/blocksparsetensor.jl

@@ -49,7 +49,7 @@ Construct a block sparse tensor with uninitialized memory
 from indices and locations of non-zero blocks.
 """
 function BlockSparseTensor(::UndefInitializer, blockoffsets, inds)
-  return BlockSparseTensor(Float64, undef, blockoffsets, inds)
+  return BlockSparseTensor(default_eltype(), undef, blockoffsets, inds)
 end
 
 function BlockSparseTensor(
@@ -65,15 +65,15 @@ function BlockSparseTensor(eltype::Type{<:Number}, blockoffsets::BlockOffsets, i
 end
 
 function BlockSparseTensor(blockoffsets::BlockOffsets, inds)
-  return BlockSparseTensor(Float64, blockoffsets, inds)
+  return BlockSparseTensor(default_eltype(), blockoffsets, inds)
 end
 
 """
     BlockSparseTensor(inds)
 
 Construct a block sparse tensor with no blocks.
 """
-BlockSparseTensor(inds) = BlockSparseTensor(Float64, inds)
+BlockSparseTensor(inds) = BlockSparseTensor(default_eltype(), inds)
 
 function BlockSparseTensor(datatype::Type{<:AbstractArray}, inds)
   return BlockSparseTensor(datatype, BlockOffsets{length(inds)}(), inds)
@@ -99,7 +99,7 @@ Construct a block sparse tensor with the specified blocks.
 Defaults to setting structurally non-zero blocks to zero.
 """
 function BlockSparseTensor(blocks::Vector{BlockT}, inds) where {BlockT<:Union{Block,NTuple}}
-  return BlockSparseTensor(Float64, blocks, inds)
+  return BlockSparseTensor(default_eltype(), blocks, inds)
 end
 
 function BlockSparseTensor(
@@ -160,7 +160,7 @@ function randomBlockSparseTensor(blocks::Vector, inds)
 end
 
 function randomBlockSparseTensor(rng::AbstractRNG, blocks::Vector, inds)
-  return randomBlockSparseTensor(rng, Float64, blocks, inds)
+  return randomBlockSparseTensor(rng, default_eltype(), blocks, inds)
 end
 
 """
@@ -176,6 +176,12 @@ function BlockSparseTensor(
   return BlockSparseTensor(blocks, inds)
 end
 
+function BlockSparseTensor{ElT}(
+  blocks::Vector{BlockT}, inds::Vararg{BlockDim,N}
+) where {ElT<:Number,BlockT<:Union{Block{N},NTuple{N,<:Integer}}} where {N}
+  return BlockSparseTensor(ElT, blocks, inds)
+end
+
 function zeros(
   tensor::BlockSparseTensor{ElT,N}, blockoffsets::BlockOffsets{N}, inds
 ) where {ElT,N}

NDTensors/src/dense/tensoralgebra/outer.jl

@@ -21,12 +21,9 @@ function outer!(
   v1 = data(T1)
   v2 = data(T2)
   RM = reshape(R, length(v1), length(v2))
-  ## Potential fix is call reshape on array
-  #RM = reshape(array(R), length(v1), length(v2))
-  #RM .= v1 .* transpose(v2)
-  #mul!(RM, v1, transpose(v2))
-  _gemm!('N', 'T', one(ElR), v1, v2, zero(ElR), RM)
-  #mul!!(RM, v1, transpose(v2), one(ElR), zero(ElR))
+  ## There is no _gemm! defined for CUDA or Metal so it calls 
+  ## generic matmul. Replace with mul!! to call correct mul! (ger)
+  mul!!(array(RM), v1, transpose(v2), one(ElR), zero(ElR))
   return R
 end