Special case tensoradd! to bypass PermutedDimsArray for Diagonals

src/implementation/base.jl

@@ -28,6 +28,34 @@ function tensoradd!(
     end
     return C
 end
+
+# needed to avoid ambiguity error
+function tensoradd!(
+        C::AbstractArray,
+        A::Diagonal, pA::Index2Tuple, conjA::Bool,
+        α::Number, β::Number,
+        ::BaseView, allocator = DefaultAllocator()
+    )
+    argcheck_tensoradd(C, A, pA)
+    dimcheck_tensoradd(C, A, pA)
+
+    # can we assume that C is mutable?
+    # is there more functionality in base that we can use?
+    if conjA
+        if iszero(β)
+            C .= α .* conj.(A)
+        else
+            C .= β .* C .+ α .* conj.(A)
+        end
+    else
+        if iszero(β)
+            C .= α .* A
+        else
+            C .= β .* C .+ α .* A
+        end
+    end
+    return C
+end
 function tensoradd!(
         C::AbstractArray,
         A::AbstractArray, pA::Index2Tuple, conjA::Bool,
@@ -39,8 +67,13 @@ function tensoradd!(
 
     # can we assume that C is mutable?
     # is there more functionality in base that we can use?
-    Atemp = tensoralloc_add(eltype(A), A, pA, conjA, Val(true), allocator)
-    Ã = permutedims!(Atemp, A, linearize(pA))
+    Ã = if !isa(A, Diagonal)
+        Atemp = tensoralloc_add(eltype(A), A, pA, conjA, Val(true), allocator)
+        permutedims!(Atemp, A, linearize(pA))
+    else
+        Atemp = nothing
+        A
+    end
     if conjA
         if iszero(β)
             C .= α .* conj.(Ã)
@@ -54,7 +87,7 @@ function tensoradd!(
             C .= β .* C .+ α .* Ã
         end
     end
-    tensorfree!(Atemp, allocator)
+    !isa(A, Diagonal) && tensorfree!(Atemp, allocator)
     return C
 end
 

test/gpu.jl

@@ -54,6 +54,32 @@ bufferallocator(::Type{ROCArray}; kwargs...) = TensorOperations.AMDBufferAllocat
             tensoradd!(c, a, (p, ()), false, α, β, backend)
         end
 
+        T <: Real || @test compare(AT, C, A) do c, a
+            tensoradd!(c, a, (p, ()), true, α, β, backend)
+        end
+    end
+    # test Diagonal special case
+    sz = (8, 8)
+    p = (2, 1)
+    diag_backends = [BaseCopy(), BaseView()]
+    for backend in diag_backends, T in (Float32, ComplexF32)
+        A = Diagonal(randn(T, sz[1]))
+        C = randn(T, TupleTools.getindices(sz, p))
+
+        @test compare(AT, C, A) do c, a
+            tensoradd!(c, a, (p, ()), false, One(), Zero(), backend)
+        end
+
+        α = rand(T)
+        @test compare(AT, C, A) do c, a
+            tensoradd!(c, a, (p, ()), false, α, Zero(), backend)
+        end
+
+        β = rand(T)
+        @test compare(AT, C, A) do c, a
+            tensoradd!(c, a, (p, ()), false, α, β, backend)
+        end
+
         T <: Real || @test compare(AT, C, A) do c, a
             tensoradd!(c, a, (p, ()), true, α, β, backend)
         end