[NDTensorsCUDAExt] Remove and test for scalar indexing

NDTensors/ext/NDTensorsCUDAExt/NDTensorsCUDAExt.jl

@@ -5,7 +5,7 @@ using NDTensors.SetParameters
 using NDTensors.Unwrap
 using Adapt
 using Functors
-using LinearAlgebra
+using LinearAlgebra: LinearAlgebra, Adjoint, Transpose
 
 if isdefined(Base, :get_extension)
   using CUDA
@@ -24,4 +24,5 @@ include("iscu.jl")
 include("adapt.jl")
 include("indexing.jl")
 include("linearalgebra.jl")
+include("mul.jl")
 end

NDTensors/ext/NDTensorsCUDAExt/imports.jl

@@ -4,3 +4,4 @@ import NDTensors:
 import NDTensors.SetParameters: nparameters, get_parameter, set_parameter, default_parameter
 
 import .CUDA: CuArrayAdaptor
+import LinearAlgebra: mul!

NDTensors/ext/NDTensorsCUDAExt/mul.jl

@@ -0,0 +1,27 @@
+# This was calling generic matrix multiplication.
+# TODO: Raise an issue with `CUDA.jl`.
+function mul!(
+  CM::Exposed{<:CuArray,<:LinearAlgebra.Transpose},
+  AM::Exposed{<:CuArray},
+  BM::Exposed{<:CuArray},
+  α,
+  β,
+)
+  return mul!(parent(CM), transpose(BM), transpose(AM), α, β)
+
+  return unexpose(CM)
+end
+
+# This was calling generic matrix multiplication.
+# TODO: Raise an issue with `CUDA.jl`.
+function mul!(
+  CM::Exposed{<:CuArray,<:LinearAlgebra.Adjoint},
+  AM::Exposed{<:CuArray},
+  BM::Exposed{<:CuArray},
+  α,
+  β,
+)
+  return mul!(parent(CM), BM', AM', α, β)
+
+  return unexpose(CM)
+end

NDTensors/ext/NDTensorsMetalExt/mul.jl

@@ -10,3 +10,13 @@ function LinearAlgebra.mul!(
   mul!(transpose(CM), transpose(BM), transpose(AM), α, β)
   return unexpose(CM)
 end
+
+# This was calling generic matrix multiplication.
+# TODO: Raise an issue with `Metal.jl`.
+function LinearAlgebra.mul!(
+  CM::Exposed{<:MtlArray,<:Adjoint}, AM::Exposed{<:MtlArray}, BM::Exposed{<:MtlArray}, α, β
+)
+  return mul!(parent(CM), BM', AM', α, β)
+
+  return unexpose(CM)
+end

NDTensors/src/Unwrap/src/functions/abstractarray.jl

@@ -2,6 +2,8 @@ parent(E::Exposed) = parent(unexpose(E))
 
 transpose(E::Exposed) = transpose(unexpose(E))
 
+adjoint(E::Exposed) = adjoint(unexpose(E))
+
 cpu(E::Exposed) = cpu(unexpose(E))
 
 getindex(E::Exposed) = unexpose(E)[]

NDTensors/src/Unwrap/src/import.jl

@@ -1,4 +1,5 @@
 import Base:
+  adjoint,
   permutedims,
   permutedims!,
   copy,

NDTensors/src/Unwrap/test/runtests.jl

@@ -138,4 +138,33 @@ include("../../../test/device_list.jl")
   x = dev(randn(elt, 4, 4))
   permutedims!(expose(y), expose(x), (2, 1))
   @test NDTensors.cpu(y) == transpose(NDTensors.cpu(x))
+
+  ##########################################
+  ### Testing an issue with CUDA&Metal transpose/adjoint mul
+  A = dev(randn(Float64, (3, 2)))
+  B = dev(randn(Float64, (3, 4)))
+  C = dev(randn(Float64, (4, 2)))
+  Cp = copy(C)
+  if (dev == NDTensors.cu)
+    CUDA.allowscalar(false)
+  end
+  ## This fails with scalar indexing 
+  #mul!(transpose(C), transpose(A), B, 1.0, 0.0)
+  mul!(C, transpose(B), A, 1.0, 0.0)
+  mul!(expose(transpose(Cp)), expose(transpose(A)), expose(B), 1.0, 0.0)
+  @test C ≈ Cp
+  Cp = fill!(similar(C), 0.0)
+  ## Try calling mul!! with transposes to verify that code works
+  NDTensors.mul!!(transpose(Cp), transpose(A), B, 1.0, 0.0)
+  @test C ≈ Cp
+
+  Cp = fill!(similar(C), 0.0)
+  ## This fails with scalar indexing 
+  #mul!(C', A', B, 1.0, 0)
+  mul!(C, B', A, 1.0, 0.0)
+  mul!(expose(Cp'), expose(A'), expose(B), 1.0, 0)
+  @test C ≈ Cp
+  Cp = fill!(similar(C), 0.0)
+  NDTensors.mul!!(Cp', A', B, 1.0, 0.0)
+  @test Cp ≈ C
 end