make tensor products faster

[vpuri3]

with current @views based implementation:

SciMLOperators.jl/src/sciml.jl

Lines 584 to 618 in 746c0d5

	function LinearAlgebra.mul!(v::AbstractVecOrMat, L::TensorProductOperator, u::AbstractVecOrMat)
	@assert L.isset "cache needs to be set up for operator of type $(typeof(L)).
	set up cache by calling cache_operator(L::AbstractSciMLOperator, u::AbstractArray)"
	mi, ni = size(L.inner)
	mo, no = size(L.outer)
	k = size(u, 2)
	C = L.cache
	U = _reshape(u, (ni, no*k))
	"""
	v .= kron(B, A) * u
	V .= A * U * B'
	"""
	# C .= A * U
	mul!(C, L.inner, U)
	# V .= U * B' <===> V' .= B * C'
	if k>1
	V = _reshape(v, (mi, mo, k))
	C = _reshape(C, (mi, no, k))
	@views for i=1:k
	mul!(transpose(V[:,:,i]), L.outer, transpose(C[:,:,i]))
	end
	else
	V = _reshape(v, (mi, mo))
	C = _reshape(C, (mi, no))
	mul!(transpose(V), L.outer, transpose(C))
	end
	v
	end

using SciMLOperators, LinearAlgebra
using BenchmarkTools

A = TensorProductOperator(rand(12,12), rand(12,12), rand(12,12))

u = rand(12^3, 100)
v = rand(12^3, 100)

A = cache_operator(A, u)

mul!(v, A, u) # dunny
@btime mul!($v, $A, $u); # 16.901 ms (348801 allocations: 45.51 MiB)

B = convert(AbstractMatrix, A)
mul!(v, B, u) # dummy
@btime mul!($v, $B, $u); # 10.337 ms (0 allocations: 0 bytes)

julia> versioninfo()
Julia Version 1.8.0-rc1
Commit 6368fdc6565 (2022-05-27 18:33 UTC)
Platform Info:
  OS: macOS (x86_64-apple-darwin21.4.0)
  CPU: 4 × Intel(R) Core(TM) i5-5257U CPU @ 2.70GHz
  WORD_SIZE: 64
  LIBM: libopenlibm
  LLVM: libLLVM-13.0.1 (ORCJIT, broadwell)
  Threads: 4 on 4 virtual cores
Environment:
  JULIA_NUM_PRECOMPILE_TASKS = 4
  JULIA_DEPOT_PATH = /Users/vp/.julia
  JULIA_NUM_THREADS = 4

[vpuri3]

using SciMLOperators, LinearAlgebra
using BenchmarkTools

A = TensorProductOperator(rand(12,12), rand(12,12), rand(12,12))

u = [rand(12^3) for i=1:100]
v = [rand(12^3) for i=1:100]

A = cache_operator(A, u[1])

mul!(v[1], A, u[1]) # dummy
@btime for i=1:length(u); mul!($v[i], $A, $u[i]); end; # 7.333 ms (3301 allocations: 27.05 MiB)

B = convert(AbstractMatrix, A)
mul!(v[1], B, u[1]) # dummy
@btime for i=1:length(u); mul!($v[i], $B, $u[i]); end; # 129.053 ms (101 allocations: 3.16 KiB)

so applying the operator to individual vectors is indeed faster. likely because we are avoiding the views call

[vpuri3]

@chriselrod @ChrisRackauckas can you take a look please?

[vpuri3]

i think the 'slow' part is that we are applying the operator to column vectors in u one at a time

SciMLOperators.jl/src/sciml.jl

Lines 608 to 610 in 746c0d5

	@views for i=1:k
	mul!(transpose(V[:,:,i]), L.outer, transpose(C[:,:,i]))
	end

[vpuri3]

there's definitely something im doing wrong because we are slower than LinearMaps, and LinearMaps has allocations in their mul!.

julia> using LinearMaps, LinearAlgebra, BenchmarkTools; A = ⊗(rand(12,12), rand(12,12), rand(12,12)); u=rand(12^3,100); v=copy(u); @btime mul!($v, $A, $u);
  5.142 ms (8500 allocations: 27.11 MiB)

https://github.com/JuliaLinearAlgebra/LinearMaps.jl/blob/master/src/kronecker.jl#L117-L131

[vpuri3]

@danielwe

[chriselrod]

i think the 'slow' part is that we are applying the operator to column vectors in u one at a time

SciMLOperators.jl/src/sciml.jl

Lines 608 to 610 in 746c0d5

	@views for i=1:k
	mul!(transpose(V[:,:,i]), L.outer, transpose(C[:,:,i]))
	end

You could try replacing that with an @turbo or @tturbo loop.

[vpuri3]

@turbo didn't recognize mul! when i checked. and L.outer is not necessarily an AbstractMatrix (for eg it is an AbstractSciMLOperator we nest ops TensorProductOperator(A, B, C) === TensorProductOperator(A, TPPOp(B,C))) so it can't necessarily be indexed. So we can't get rid of the mul! really unless we can come up with a general implementation for TensorProductOperator with an arbitrary number of arguments

[chriselrod]

I meant delete mul!, replacing it with loops

[ChrisRackauckas]

You can't assume operators are indexable, it's not part of the interface (for a good reason, it may be a matrix-free operator). If you need to do that, then we can have a trait for the fast indexing (extending the ArrayInterface one), but I'm not sure that's useful since in most cases it should be matrix-free

[chriselrod]

No indexing, but views are okay?

We could have a special overload for certain types.

[vpuri3]

@chriselrod im not sure how view without indexing would work.

SciMLOperators.jl/src/sciml.jl

Lines 608 to 610 in 746c0d5

	@views for i=1:k
	mul!(transpose(V[:,:,i]), L.outer, transpose(C[:,:,i]))
	end

in this case, the slices C[:,:,i], V[:,:,i] are passed into mul! as subarrays. But L.outer isn't guaranteed to be indexable.

[vpuri3]

in #59 i've written an implementation with using permutedims! that is faster than LinearMaps.jl and doesn't allocate. could you take a look and suggest improvements?