diff --git a/src/implementation/base.jl b/src/implementation/base.jl index ecd2aed9..2410436a 100644 --- a/src/implementation/base.jl +++ b/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 diff --git a/test/gpu.jl b/test/gpu.jl index 8b716906..81adb21f 100644 --- a/test/gpu.jl +++ b/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