Btime for loading @tensor inside loop

Specific Domains Numerics
1

Hi, for the following code

using LinearAlgebra
using BenchmarkTools
using TensorOperations

# Initialization
A = rand(100, 100)
B = rand(100, 100)


B_3d = rand(100, 100, 100)


n=100
C_new = zeros(100, 100, 100)
@btime @tensor  C_new[a,c,d] := A[a,b] *B_3d[b,c,d]

@btime begin
for d = 1:n
    C_new[:,:,d] = LinearAlgebra.BLAS.gemm('N', 'N', A, B_3d[:,:,d])
end
end

@btime begin
for d = 1:n
   @tensor C_new[a,c,d] := A[a,b] *B_3d[b,c,d]
end
end

I got

  6.801 ms (3 allocations: 7.63 MiB)
  10.129 ms (701 allocations: 15.27 MiB)
  1.030 s (401 allocations: 762.95 MiB)

what is the reason for the third timing, that I tried to load @tensor as
a matrix multiplication, A[a,b] *B_3d[b,c,d] ā† for fixed d, is much slower than the first two?

2

Try

@btime begin
for d = 1:n
   Cview = @view C_new[:,:,d]
   Bview = @view B_3d[:,:,d]
   @tensor Cview[a,c] := A[a,b] *Bview[b,c]
end
end

The optimization isnā€™t the @view, I think you arenā€™t slicing at all, but doing

@btime begin
for _ = 1:n
   @tensor C_new[a,c,d] := A[a,b] *B_3d[b,c,d]
end
end

I have not tested your code to confirm it.

1 Like
3

Thanks. I got

ERROR: LoadError: syntax: unsupported assignment operator ":="

at the first line of

@btime begin
for d = 1:n
   Cview = @view C_new[a,c,d]
   Bview = @view B_3d[b,c,d]
   @tensor Cview := A[a,b] *Bview
end
end
4

Sorry, I fixed my previous comment.

5

Benchmarks:

julia> @btime @tensor $C_new[a,c,d] = $A[a,b] * $B_3d[b,c,d];
  936.132 Ī¼s (0 allocations: 0 bytes)

julia> f!(C_new, A, B_3d, n) = begin
       for d = 1:n
          LinearAlgebra.BLAS.gemm!('N', 'N', 1.0, A, @view(B_3d[:,:,d]), 0.0, @view(C_new[:,:,d]))
       end
       end; @btime f!($C_new, $A, $B_3d, $n)
  1.417 ms (0 allocations: 0 bytes)

julia> g!(C_new, A, B_3d, n) = begin
       for d = 1:n
          Cview = @view C_new[:,:,d]
          Bview = @view B_3d[:,:,d]
          @tensor Cview[a,c] = A[a,b] * Bview[b,c]
       end
       end; @btime g!($C_new, $A, $B_3d, $n)
  1.454 ms (100 allocations: 6.25 KiB)

julia> h!(C_new, A, B_3d, n) = begin
       for d = 1:n
          Cview = @view C_new[:,:,d]
          Bview = @view B_3d[:,:,d]
          mul!(Cview, A, Bview)
       end
       end; @btime h!($C_new, $A, $B_3d, $n)
  1.414 ms (0 allocations: 0 bytes)

Note that Iā€™m using = to be in place and avoid allocations, but g! still has them for some reason.

1 Like