Add TBLIS implementation for tensortrace

Author: lkdvos

src/TensorOperationsTBLIS.jl

@@ -64,13 +64,27 @@ function TensorOperations.tensorcontract!(C::StridedArray{T}, pC::Index2Tuple,
     return C
 end
 
-# partial traces do not exist in tblis afaik -> use default implementation
-# TODO: implement full trace
 function TensorOperations.tensortrace!(C::StridedArray{T}, pC::Index2Tuple,
                                        A::StridedArray{T}, pA::Index2Tuple, conjA::Symbol,
                                        α::Number, β::Number,
                                        ::tblisBackend) where {T<:BlasFloat}
-    return tensortrace!(C, pC, A, pA, conjA, α, β)
+    TensorOperations.argcheck_tensortrace(C, pC, A, pA)
+    TensorOperations.dimcheck_tensortrace(C, pC, A, pA)
+
+    rmul!(C, β) # TODO: is it possible to use tblis scaling here?
+    szC = ndims(C) == 0 ? Int[] : collect(size(C))
+    strC = ndims(C) == 0 ? Int[] : collect(strides(C))
+    C_tblis = tblis_tensor(C, szC, strC)
+
+    szA = collect(size(A))
+    strA = collect(strides(A))
+    A_tblis = tblis_tensor(conjA == :C ? conj(A) : A, szA, strA, α)
+
+    einA, einC = TensorOperations.trace_labels(pC, pA...)
+
+    tblis_tensor_add(A_tblis, string(einA...), C_tblis, string(einC...))
+
+    return C
 end
 
 end # module TensorOperationsTBLIS

test/runtests.jl

@@ -27,26 +27,26 @@ using LinearAlgebra: norm
         @test collect(E2) ≈ E1
     end
 
-    # @testset "tensortrace" begin
-    #     A = randn(Float32, (5, 10, 10))
-    #     @tensor B1[a] := A[a, b′, b′]
-    #     @tensor B2[a] := CuArray(A)[a, b′, b′]
-    #     @test collect(B2) ≈ B1
-
-    #     C = randn(ComplexF32, (3, 20, 5, 3, 20, 4, 5))
-    #     @tensor D1[e, a, d] := C[a, b, c, d, b, e, c]
-    #     @tensor D2[e, a, d] := CuArray(C)[a, b, c, d, b, e, c]
-    #     @test collect(D2) ≈ D1
-
-    #     @tensor D3[a, e, d] := conj(C[a, b, c, d, b, e, c])
-    #     @tensor D4[a, e, d] := conj(CuArray(C)[a, b, c, d, b, e, c])
-    #     @test collect(D4) ≈ D3
-
-    #     α = randn(ComplexF32)
-    #     @tensor D5[d, e, a] := α * C[a, b, c, d, b, e, c]
-    #     @tensor D6[d, e, a] := α * CuArray(C)[a, b, c, d, b, e, c]
-    #     @test collect(D6) ≈ D5
-    # end
+    @testset "tensortrace" begin
+        A = randn(Float32, (5, 10, 10))
+        @tensor B1[a] := A[a, b′, b′]
+        @tensor backend = tblis B2[a] := A[a, b′, b′]
+        @test B2 ≈ B1
+
+        C = randn(ComplexF32, (3, 20, 5, 3, 20, 4, 5))
+        @tensor D1[e, a, d] := C[a, b, c, d, b, e, c]
+        @tensor backend = tblis D2[e, a, d] := C[a, b, c, d, b, e, c]
+        @test D2 ≈ D1
+
+        @tensor D3[a, e, d] := conj(C[a, b, c, d, b, e, c])
+        @tensor backend = tblis D4[a, e, d] := conj(C[a, b, c, d, b, e, c])
+        @test D4 ≈ D3
+
+        α = randn(ComplexF32)
+        @tensor D5[d, e, a] := α * C[a, b, c, d, b, e, c]
+        @tensor backend = tblis D6[d, e, a] := α * C[a, b, c, d, b, e, c]
+        @test D6 ≈ D5
+    end
 
     @testset "tensorcontract" begin
         A = randn(Float32, (3, 20, 5, 3, 4))