resolve conjugation flags in strided implementations

Project.toml

@@ -1,7 +1,7 @@
 name = "TensorOperations"
 uuid = "6aa20fa7-93e2-5fca-9bc0-fbd0db3c71a2"
 authors = ["Lukas Devos <[email protected]>", "Maarten Van Damme <[email protected]>", "Jutho Haegeman <[email protected]>"]
-version = "5.0.1"
+version = "5.0.2"
 
 [deps]
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"

src/TensorOperations.jl

@@ -58,8 +58,8 @@ include("indexnotation/tensormacros.jl")
 include("implementation/functions.jl")
 include("implementation/ncon.jl")
 include("implementation/abstractarray.jl")
-include("implementation/diagonal.jl")
 include("implementation/strided.jl")
+include("implementation/diagonal.jl")
 include("implementation/base.jl")
 include("implementation/indices.jl")
 include("implementation/allocator.jl")

src/implementation/abstractarray.jl

@@ -160,8 +160,3 @@ function dimcheck_tensorcontract(C::AbstractArray,
         throw(DimensionMismatch("non-matching sizes in uncontracted dimensions"))
     return nothing
 end
-
-#-------------------------------------------------------------------------------------------
-# Utility functions
-#-------------------------------------------------------------------------------------------
-flag2op(flag::Bool) = flag ? conj : identity

src/implementation/diagonal.jl

@@ -10,10 +10,15 @@ function tensorcontract!(C::AbstractArray,
     argcheck_tensorcontract(C, A, pA, B, pB, pAB)
     dimcheck_tensorcontract(C, A, pA, B, pB, pAB)
 
-    _diagtensorcontract!(StridedView(C),
-                         StridedView(A), pA, conjA,
-                         StridedView(B.diag), pB, conjB,
-                         pAB, α, β)
+    if conjA && conjB
+        _diagtensorcontract!(SV(C), conj(SV(A)), pA, conj(SV(B.diag)), pB, pAB, α, β)
+    elseif conjA
+        _diagtensorcontract!(SV(C), conj(SV(A)), pA, SV(B.diag), pB, pAB, α, β)
+    elseif conjB
+        _diagtensorcontract!(SV(C), SV(A), pA, conj(SV(B.diag)), pB, pAB, α, β)
+    else
+        _diagtensorcontract!(SV(C), SV(A), pA, SV(B.diag), pB, pAB, α, β)
+    end
     return C
 end
 
@@ -35,10 +40,15 @@ function tensorcontract!(C::AbstractArray,
     rpAB = (TupleTools.getindices(indCinoBA, tpAB[1]),
             TupleTools.getindices(indCinoBA, tpAB[2]))
 
-    _diagtensorcontract!(StridedView(C),
-                         StridedView(B), rpB, conjB,
-                         StridedView(A.diag), rpA, conjA,
-                         rpAB, α, β)
+    if conjA && conjB
+        _diagtensorcontract!(SV(C), conj(SV(B)), rpB, conj(SV(A.diag)), rpA, rpAB, α, β)
+    elseif conjA
+        _diagtensorcontract!(SV(C), SV(B), rpB, conj(SV(A.diag)), rpA, rpAB, α, β)
+    elseif conjB
+        _diagtensorcontract!(SV(C), conj(SV(B)), rpB, SV(A.diag), rpA, rpAB, α, β)
+    else
+        _diagtensorcontract!(SV(C), SV(B), rpB, SV(A.diag), rpA, rpAB, α, β)
+    end
     return C
 end
 
@@ -50,40 +60,16 @@ function tensorcontract!(C::AbstractArray,
                          ::StridedNative, allocator=DefaultAllocator())
     argcheck_tensorcontract(C, A, pA, B, pB, pAB)
     dimcheck_tensorcontract(C, A, pA, B, pB, pAB)
-    if numin(pA) == 1 # matrix multiplication
-        scale!(C, β)
-        β = one(β)
-
-        A2 = sreshape(flag2op(conjA)(StridedView(A.diag)), (length(A.diag), 1))
-        B2 = sreshape(flag2op(conjB)(StridedView(B.diag)), (length(B.diag), 1))
-        # take a view of the diagonal elements of C, having strides 1 + length(diag)
-        totsize = (length(A.diag),)
-        C2 = StridedView(C, totsize, (sum(strides(C)),))
-
-    elseif numin(pA) == 2 # trace
-        A2 = flag2op(conjA)(StridedView(A.diag, (length(A.diag),)))
-        B2 = flag2op(conjB)(StridedView(B.diag, (length(B.diag),)))
-        totsize = (length(A.diag),)
-        C2 = sreshape(StridedView(C), (1,))
-
-    else # outer product
-        scale!(C, β)
-        β = one(β)
 
-        A2 = sreshape(StridedView(A.diag), (length(A.diag), 1))
-        B2 = sreshape(StridedView(B.diag), (1, length(A.diag)))
-
-        C3 = permutedims(StridedView(C), invperm(linearize(pAB)))
-        strC = strides(C3)
-        newstrides = (strC[1] + strC[2], strC[3] + strC[4])
-        totsize = (length(A2), length(B2))
-        C2 = StridedView(C3.parent, totsize, newstrides, C3.offset, C3.op)
+    if conjA && conjB
+        _diagdiagcontract!(SV(C), conj(SV(A.diag)), pA, conj(SV(B.diag)), pB, pAB, α, β)
+    elseif conjA
+        _diagdiagcontract!(SV(C), conj(SV(A.diag)), pA, SV(B.diag), pB, pAB, α, β)
+    elseif conjB
+        _diagdiagcontract!(SV(C), SV(A.diag), pA, conj(SV(B.diag)), pB, pAB, α, β)
+    else
+        _diagdiagcontract!(SV(C), SV(A.diag), pA, SV(B.diag), pB, pAB, α, β)
     end
-
-    op1 = Base.Fix2(scale, α) ∘ *
-    op2 = Base.Fix2(scale, β)
-    Strided._mapreducedim!(op1, +, op2, totsize, (C2, A2, B2))
-
     return C
 end
 
@@ -96,41 +82,49 @@ function tensorcontract!(C::Diagonal,
     argcheck_tensorcontract(C, A, pA, B, pB, pAB)
     dimcheck_tensorcontract(C, A, pA, B, pB, pAB)
 
-    A2 = flag2op(conjA)(StridedView(A.diag))
-    B2 = flag2op(conjB)(StridedView(B.diag))
+    A2 = StridedView(A.diag)
+    B2 = StridedView(B.diag)
     C2 = StridedView(C.diag)
 
-    C2 .= C2 .* β .+ A2 .* B2 .* α
+    if conjA && conjB
+        C2 .= C2 .* β .+ conj.(A2 .* B2) .* α
+    elseif conjA
+        C2 .= C2 .* β .+ conj.(A2) .* B2 .* α
+    elseif conjB
+        C2 .= C2 .* β .+ A2 .* conj.(B2) .* α
+    else
+        C2 .= C2 .* β .+ A2 .* B2 .* α
+    end
     return C
 end
 
 function _diagtensorcontract!(C::StridedView,
-                              A::StridedView, pA::Index2Tuple, conjA::Bool,
-                              Bdiag::StridedView, pB::Index2Tuple, conjB::Bool,
+                              A::StridedView, pA::Index2Tuple,
+                              Bdiag::StridedView, pB::Index2Tuple,
                               pAB::Index2Tuple, α::Number, β::Number)
     sizeA = i -> size(A, i)
     csizeA = sizeA.(pA[2])
     osizeA = sizeA.(pA[1])
 
     if numin(pB) == 1 # => numin(A) == numout(B) == 1
         totsize = (osizeA..., csizeA...)
-        A2 = flag2op(conjA)(permutedims(A, linearize(pA)))
-        B2 = flag2op(conjB)(sreshape(Bdiag, ((one.(osizeA))..., csizeA...)))
+        A2 = permutedims(A, linearize(pA))
+        B2 = sreshape(Bdiag, ((one.(osizeA))..., csizeA...))
         C2 = permutedims(C, invperm(linearize(pAB)))
 
     elseif numin(pB) == 0
         strideA = i -> stride(A, i)
         newstrides = (strideA.(pA[1])..., strideA(pA[2][1]) + strideA(pA[2][2]))
         totsize = (osizeA..., csizeA[1])
-        A2 = flag2op(conjA)(StridedView(A.parent, totsize, newstrides, A.offset, A.op))
-        B2 = flag2op(conjB)(sreshape(Bdiag, ((one.(osizeA))..., csizeA[1])))
+        A2 = StridedView(A.parent, totsize, newstrides, A.offset, A.op)
+        B2 = sreshape(Bdiag, ((one.(osizeA))..., csizeA[1]))
         C2 = permutedims(C, invperm(linearize(pAB)))
 
     else # numout(pB) == 2 # direct product
         scale!(C, β)
         β = one(β)
-        A2 = flag2op(conjA)(sreshape(permutedims(A, linearize(pA)), (osizeA..., 1)))
-        B2 = flag2op(conjB)(sreshape(Bdiag, ((one.(osizeA))..., length(Bdiag))))
+        A2 = sreshape(permutedims(A, linearize(pA)), (osizeA..., 1))
+        B2 = sreshape(Bdiag, ((one.(osizeA))..., length(Bdiag)))
 
         C3 = permutedims(C, invperm(linearize(pAB)))
         sC = strides(C3)
@@ -145,3 +139,44 @@ function _diagtensorcontract!(C::StridedView,
 
     return C
 end
+
+function _diagdiagcontract!(C::StridedView,
+                            Adiag::StridedView, pA::Index2Tuple,
+                            Bdiag::StridedView, pB::Index2Tuple,
+                            pAB::Index2Tuple, α::Number, β::Number)
+    if numin(pA) == 1 # matrix multiplication
+        scale!(C, β)
+        β = one(β)
+
+        A2 = sreshape(Adiag, (length(Adiag), 1))
+        B2 = sreshape(Bdiag, (length(Bdiag), 1))
+        # take a view of the diagonal elements of C, having strides 1 + length(diag)
+        totsize = (length(Adiag),)
+        C2 = StridedView(C.parent, totsize, (sum(strides(C)),))
+
+    elseif numin(pA) == 2 # trace
+        A2 = Adiag
+        B2 = Bdiag
+        totsize = (length(Adiag),)
+        C2 = sreshape(C, (1,))
+
+    else # outer product
+        scale!(C, β)
+        β = one(β)
+
+        A2 = sreshape(Adiag, (length(Adiag), 1))
+        B2 = sreshape(Bdiag, (1, length(Adiag)))
+
+        C3 = permutedims(C, invperm(linearize(pAB)))
+        strC = strides(C3)
+        newstrides = (strC[1] + strC[2], strC[3] + strC[4])
+        totsize = (length(A2), length(B2))
+        C2 = StridedView(C3.parent, totsize, newstrides, C3.offset, C3.op)
+    end
+
+    op1 = Base.Fix2(scale, α) ∘ *
+    op2 = Base.Fix2(scale, β)
+    Strided._mapreducedim!(op1, +, op2, totsize, (C2, A2, B2))
+
+    return C
+end