TTM @einsum versus handwritten performance

[willow-ahrens]

julia> N = 700
700

julia> P = Tensor(Dense(SparseList(SparseList(Element(0.0)))), fsprand(N, N, N, 0.0005));

julia> E = Tensor(Dense(Dense(Element(0.0))), rand(N, N));

julia> @btime begin
           P = $P
           E = $E
           @einsum K1[n,j,k] = P[i,j,k]*E[i,n]
       end;
  17.489 s (678 allocations: 20.34 GiB)

julia> K1 = Tensor(Dense(SparseList(Dense(Element(0.0)))));

julia> tmp = Tensor(SparseByteMap(Element(0.0)));

julia> function mcc3_step1!(P, E, K1)
           @finch begin
               K1 .= 0
               for k=_, j=_
                   tmp .= 0
                   for n=_, i=_
                       tmp[n] += P[i,j,k]*E[i,n]
                   end
                   for n=_
                       K1[n,j,k] = tmp[n]
                   end
               end
           end
           return K1
       end
mcc3_step1! (generic function with 1 method)

julia> @btime mcc3_step1!($P, $E, $K1);
  747.221 ms (25 allocations: 3.94 KiB)

julia> function simple_mcc3_step1!(P, E, K1)
           @finch begin
               K1 .= 0
               for k=_, j=_
                   for n=_, i=_
                       K1[n,j,k] += P[i,j,k]*E[i,n]
                   end
               end
           end
           return K1
       end
simple_mcc3_step1! (generic function with 1 method)

julia> K1 = Tensor(Dense(SparseDict(Dense(Element(0.0)))));

julia> @btime simple_mcc3_step1!($P, $E, $K1);
  218.026 ms (18 allocations: 4.42 MiB)

[willow-ahrens]

the issue was @einsum K1[n,j,k] = P[i,j,k]*E[i,n] and not @einsum K1[n,j,k] += P[i,j,k]*E[i,n]