ITensor / NDTensors.jl

A Julia package for n-dimensional sparse tensors.
Apache License 2.0
27 stars 7 forks source link

Small contraction optimizations #32

Closed mtfishman closed 4 years ago

mtfishman commented 4 years ago

Oddly enough, for very small tensor contractions, using reshape to reshape the Tensor storage data into a Matrix that gets passed to Blas.gemm! takes up a non-negligible amount of time. This appears to be a known issue: https://github.com/JuliaLang/julia/issues/36313

The most common place this is used is the array(::Tensor) constructor, here: https://github.com/ITensor/NDTensors.jl/blob/master/src/dense.jl#L355, but it is also used directly here: https://github.com/ITensor/NDTensors.jl/blob/master/src/dense.jl#L652

An alternative might be to use a Base.ReshapedArray instead. It looks like it can still be passed to BLAS, for example:

julia> V = randn(6)
6-element Array{Float64,1}:
 -0.26695949899994226
 -0.13359169415563152
  0.37079765916385093
  0.6749298904891178
 -0.8679197316877514
  0.2343964793982253

julia> @btime reshape($V, (3, 2));
  31.608 ns (1 allocation: 64 bytes)

julia> @btime Base.ReshapedArray($V, (3, 2), ());
  4.781 ns (1 allocation: 32 bytes)

julia> A = Base.ReshapedArray(V, (3, 2), ())
3×2 reshape(::Array{Float64,1}, 3, 2) with eltype Float64:
 -0.266959   0.67493
 -0.133592  -0.86792
  0.370798   0.234396

julia> A isa StridedArray
true

julia> BLAS.gemm('N', 'T', A, A)
3×3 Array{Float64,2}:
  0.526798   -0.550121   0.0592132
 -0.550121    0.771131  -0.252973
  0.0592132  -0.252973   0.192433
mtfishman commented 4 years ago

Closed by #34.