TensorOperations.jl icon indicating copy to clipboard operation
TensorOperations.jl copied to clipboard

StaticArrays

Open cortner opened this issue 8 years ago • 15 comments

This must have been asked before but I couldn’t find a reference: the package doesn’t seem to support StaticArrays, is there a plan / possibility to do so?

cortner avatar Nov 27 '17 18:11 cortner

This package is mostly intended for big arrays with strided memory layout. In principle, the @tensor parsing and lowering part is reusable and can be made to work with any array or tensor class, provided a few basic methods are implemented. But within this package, those methods are only implemented for StridedArray.

Jutho avatar Nov 28 '17 22:11 Jutho

It is an interesting question - the syntax provided by TensorOperations is very convenient, even if the array is small (and there's also block matrix structures).

andyferris avatar Nov 29 '17 06:11 andyferris

I am using Einsum.jl for now, which works in-place (MArrays) but not with SArrays, which is ok for now. My concern still is that Einsum.jl kind of seems semi-abandoned with TensorOperations becoming the dominant package for this?

cortner avatar Nov 29 '17 14:11 cortner

What's the error for StaticArrays here? Is it permutedims?

andyferris avatar Nov 29 '17 21:11 andyferris

What should the ideal implementation for StaticArrays look like? Just plain loops, since sizes are typically small? I haven't used StaticArrays yet.

Jutho avatar Nov 29 '17 21:11 Jutho

All the code in StaticArrays.jl is manually unrolled by generated functions, which tend to also generate linear (not Cartesian) indices statically (all of this lets LLVM use SIMD).

I generally assume or the O(length) operations can be completely unrolled (that would be the equivalent of permutedims and reshape here, but I haven't implemented permutedims in StaticArrays yet but am happy to provide). The matrix multiplication stuff is a bit complex since complete unrolling isn't optimal for larger (more than 50-100) elements, so it has code to unroll 2 of the 3 loops for multiplication of larger matrices.

Is there a generic AbstractArray path in this package? I could provide fast permutedims, reshape and *. One catch is that fast permutedims and reshape need special argument singleton types for permutations and sizes for the static compilation (perm::Val and size::Size).

andyferris avatar Nov 29 '17 21:11 andyferris

No it's really only StridedArray. Permutations are not stored in compile time information (at least not in the latest master version), since they are dealt with using stride calculations at run time. I always assumed that approach to be faster than having to compile a new function for every new permutation.

Note that this package does not use permutedims at all, the reason being that I find permutedims from Julia Base way to slow, even for plain arrays. I have my own permute function, which is actually more general and is just called add!.

Over at https://github.com/Jutho/Strided.jl I am actually experimenting with having a multithreaded implementation of my own permute function, which is even much more general (but still assumes a strided memory layout).

Jutho avatar Nov 29 '17 21:11 Jutho

Multithreaded - cool!

Did you ever consider contributing your faster code to Base, and putting it in permutedims!(::StridedArray, perm) in base/multidimensional.jl?

Obviously StaticArrays is the complete opposite end of the spectrum here, and has annoying interface problems as well...

andyferris avatar Nov 29 '17 22:11 andyferris

What's the error for StaticArrays here? Is it permutedims?

using TensorOperations, StaticArrays
a = @SVector rand(3)
b = @SMatrix rand(3, 3)
c = @MVector zeros(3)
@tensor c[i] = b[i, j] * a[j]

gives the output:

ERROR: MethodError: no method matching contract!(::Int64, ::StaticArrays.SArray{Tuple{3,3},Float64,2,9}, ::Type{Val{:N}}, ::SVector{3,Float64}, ::Type{Val{:N}}, ::Int64, ::MVector{3,Float64}, ::Tuple{Int64}, ::Tuple{Int64}, ::Tuple{}, ::Tuple{Int64}, ::Tuple{Int64})
Closest candidates are:
  contract!(::Any, ::Union{Base.ReshapedArray{T,N,A,MI} where MI<:Tuple{Vararg{Base.MultiplicativeInverses.SignedMultiplicativeInverse{Int64},N} where N} where A<:Union{DenseArray, SubArray{T,N,P,I,true} where I<:Tuple{Union{Base.Slice, UnitRange},Vararg{Any,N} where N} where P where N where T}, DenseArray{T,N}, SubArray{T,N,A,I,L} where L} where I<:Tuple{Vararg{Union{Base.AbstractCartesianIndex, Int64, Range{Int64}},N} where N} where A<:Union{Base.ReshapedArray{T,N,A,MI} where MI<:Tuple{Vararg{Base.MultiplicativeInverses.SignedMultiplicativeInverse{Int64},N} where N} where A<:Union{DenseArray, SubArray{T,N,P,I,true} where I<:Tuple{Union{Base.Slice, UnitRange},Vararg{Any,N} where N} where P where N where T} where N where T, DenseArray} where N where T, ::Type{Val{CA}}, ::Union{Base.ReshapedArray{T,N,A,MI} where MI<:Tuple{Vararg{Base.MultiplicativeInverses.SignedMultiplicativeInverse{Int64},N} where N} where A<:Union{DenseArray, SubArray{T,N,P,I,true} where I<:Tuple{Union{Base.Slice, UnitRange},Vararg{Any,N} where N} where P where N where T}, DenseArray{T,N}, SubArray{T,N,A,I,L} where L} where I<:Tuple{Vararg{Union{Base.AbstractCartesianIndex, Int64, Range{Int64}},N} where N} where A<:Union{Base.ReshapedArray{T,N,A,MI} where MI<:Tuple{Vararg{Base.MultiplicativeInverses.SignedMultiplicativeInverse{Int64},N} where N} where A<:Union{DenseArray, SubArray{T,N,P,I,true} where I<:Tuple{Union{Base.Slice, UnitRange},Vararg{Any,N} where N} where P where N where T} where N where T, DenseArray} where N where T, ::Type{Val{CB}}, ::Any, ::Union{Base.ReshapedArray{TC<:Union{Complex{Float32}, Complex{Float64}, Float32, Float64},N,A,MI} where MI<:Tuple{Vararg{Base.MultiplicativeInverses.SignedMultiplicativeInverse{Int64},N} where N} where A<:Union{DenseArray, SubArray{T,N,P,I,true} where I<:Tuple{Union{Base.Slice, UnitRange},Vararg{Any,N} where N} where P where N where T}, DenseArray{TC<:Union{Complex{Float32}, Complex{Float64}, Float32, Float64},N}, SubArray{TC<:Union{Complex{Float32}, Complex{Float64}, Float32, Float64},N,A,I,L} where L} where I<:Tuple{Vararg{Union{Base.AbstractCartesianIndex, Int64, Range{Int64}},N} where N} where A<:Union{Base.ReshapedArray{T,N,A,MI} where MI<:Tuple{Vararg{Base.MultiplicativeInverses.SignedMultiplicativeInverse{Int64},N} where N} where A<:Union{DenseArray, SubArray{T,N,P,I,true} where I<:Tuple{Union{Base.Slice, UnitRange},Vararg{Any,N} where N} where P where N where T} where N where T, DenseArray} where N, ::Any, ::Any, ::Any, ::Any, ::Any) where {CA, CB, TC<:Union{Complex{Float32}, Complex{Float64}, Float32, Float64}} at /Users/ortner/.julia/v0.6/TensorOperations/src/implementation/stridedarray.jl:101
  contract!(::Any, ::Union{Base.ReshapedArray{T,N,A,MI} where MI<:Tuple{Vararg{Base.MultiplicativeInverses.SignedMultiplicativeInverse{Int64},N} where N} where A<:Union{DenseArray, SubArray{T,N,P,I,true} where I<:Tuple{Union{Base.Slice, UnitRange},Vararg{Any,N} where N} where P where N where T}, DenseArray{T,N}, SubArray{T,N,A,I,L} where L} where I<:Tuple{Vararg{Union{Base.AbstractCartesianIndex, Int64, Range{Int64}},N} where N} where A<:Union{Base.ReshapedArray{T,N,A,MI} where MI<:Tuple{Vararg{Base.MultiplicativeInverses.SignedMultiplicativeInverse{Int64},N} where N} where A<:Union{DenseArray, SubArray{T,N,P,I,true} where I<:Tuple{Union{Base.Slice, UnitRange},Vararg{Any,N} where N} where P where N where T} where N where T, DenseArray} where N where T, ::Type{Val{CA}}, ::Union{Base.ReshapedArray{T,N,A,MI} where MI<:Tuple{Vararg{Base.MultiplicativeInverses.SignedMultiplicativeInverse{Int64},N} where N} where A<:Union{DenseArray, SubArray{T,N,P,I,true} where I<:Tuple{Union{Base.Slice, UnitRange},Vararg{Any,N} where N} where P where N where T}, DenseArray{T,N}, SubArray{T,N,A,I,L} where L} where I<:Tuple{Vararg{Union{Base.AbstractCartesianIndex, Int64, Range{Int64}},N} where N} where A<:Union{Base.ReshapedArray{T,N,A,MI} where MI<:Tuple{Vararg{Base.MultiplicativeInverses.SignedMultiplicativeInverse{Int64},N} where N} where A<:Union{DenseArray, SubArray{T,N,P,I,true} where I<:Tuple{Union{Base.Slice, UnitRange},Vararg{Any,N} where N} where P where N where T} where N where T, DenseArray} where N where T, ::Type{Val{CB}}, ::Any, ::Union{Base.ReshapedArray{TC<:Union{Complex{Float32}, Complex{Float64}, Float32, Float64},N,A,MI} where MI<:Tuple{Vararg{Base.MultiplicativeInverses.SignedMultiplicativeInverse{Int64},N} where N} where A<:Union{DenseArray, SubArray{T,N,P,I,true} where I<:Tuple{Union{Base.Slice, UnitRange},Vararg{Any,N} where N} where P where N where T}, DenseArray{TC<:Union{Complex{Float32}, Complex{Float64}, Float32, Float64},N}, SubArray{TC<:Union{Complex{Float32}, Complex{Float64}, Float32, Float64},N,A,I,L} where L} where I<:Tuple{Vararg{Union{Base.AbstractCartesianIndex, Int64, Range{Int64}},N} where N} where A<:Union{Base.ReshapedArray{T,N,A,MI} where MI<:Tuple{Vararg{Base.MultiplicativeInverses.SignedMultiplicativeInverse{Int64},N} where N} where A<:Union{DenseArray, SubArray{T,N,P,I,true} where I<:Tuple{Union{Base.Slice, UnitRange},Vararg{Any,N} where N} where P where N where T} where N where T, DenseArray} where N, ::Any, ::Any, ::Any, ::Any, ::Any, ::Type{Val{:BLAS}}) where {CA, CB, TC<:Union{Complex{Float32}, Complex{Float64}, Float32, Float64}} at /Users/ortner/.julia/v0.6/TensorOperations/src/implementation/stridedarray.jl:101
  contract!(::Any, ::Union{Base.ReshapedArray{T,N,A,MI} where MI<:Tuple{Vararg{Base.MultiplicativeInverses.SignedMultiplicativeInverse{Int64},N} where N} where A<:Union{DenseArray, SubArray{T,N,P,I,true} where I<:Tuple{Union{Base.Slice, UnitRange},Vararg{Any,N} where N} where P where N where T}, DenseArray{T,N}, SubArray{T,N,A,I,L} where L} where I<:Tuple{Vararg{Union{Base.AbstractCartesianIndex, Int64, Range{Int64}},N} where N} where A<:Union{Base.ReshapedArray{T,N,A,MI} where MI<:Tuple{Vararg{Base.MultiplicativeInverses.SignedMultiplicativeInverse{Int64},N} where N} where A<:Union{DenseArray, SubArray{T,N,P,I,true} where I<:Tuple{Union{Base.Slice, UnitRange},Vararg{Any,N} where N} where P where N where T} where N where T, DenseArray} where N where T, ::Type{Val{CA}}, ::Union{Base.ReshapedArray{T,N,A,MI} where MI<:Tuple{Vararg{Base.MultiplicativeInverses.SignedMultiplicativeInverse{Int64},N} where N} where A<:Union{DenseArray, SubArray{T,N,P,I,true} where I<:Tuple{Union{Base.Slice, UnitRange},Vararg{Any,N} where N} where P where N where T}, DenseArray{T,N}, SubArray{T,N,A,I,L} where L} where I<:Tuple{Vararg{Union{Base.AbstractCartesianIndex, Int64, Range{Int64}},N} where N} where A<:Union{Base.ReshapedArray{T,N,A,MI} where MI<:Tuple{Vararg{Base.MultiplicativeInverses.SignedMultiplicativeInverse{Int64},N} where N} where A<:Union{DenseArray, SubArray{T,N,P,I,true} where I<:Tuple{Union{Base.Slice, UnitRange},Vararg{Any,N} where N} where P where N where T} where N where T, DenseArray} where N where T, ::Type{Val{CB}}, ::Any, ::Union{Base.ReshapedArray{T,N,A,MI} where MI<:Tuple{Vararg{Base.MultiplicativeInverses.SignedMultiplicativeInverse{Int64},N} where N} where A<:Union{DenseArray, SubArray{T,N,P,I,true} where I<:Tuple{Union{Base.Slice, UnitRange},Vararg{Any,N} where N} where P where N where T}, DenseArray{T,N}, SubArray{T,N,A,I,L} where L} where I<:Tuple{Vararg{Union{Base.AbstractCartesianIndex, Int64, Range{Int64}},N} where N} where A<:Union{Base.ReshapedArray{T,N,A,MI} where MI<:Tuple{Vararg{Base.MultiplicativeInverses.SignedMultiplicativeInverse{Int64},N} where N} where A<:Union{DenseArray, SubArray{T,N,P,I,true} where I<:Tuple{Union{Base.Slice, UnitRange},Vararg{Any,N} where N} where P where N where T} where N where T, DenseArray} where N where T, ::Any, ::Any, ::Any, ::Any, ::Any) where {CA, CB} at /Users/ortner/.julia/v0.6/TensorOperations/src/implementation/stridedarray.jl:198
  ...
Stacktrace:
 [1] macro expansion at /Users/ortner/.julia/v0.6/TensorOperations/src/indexnotation/product.jl:69 [inlined]
 [2] deindexify!(::MVector{3,Float64}, ::TensorOperations.ProductOfIndexedObjects{(:i, :j),(:j,),:N,:N,StaticArrays.SArray{Tuple{3,3},Float64,2,9},SVector{3,Float64},Int64,Int64}, ::TensorOperations.Indices{(:i,)}, ::Int64) at /Users/ortner/.julia/v0.6/TensorOperations/src/indexnotation/product.jl:63

cortner avatar Nov 29 '17 22:11 cortner

and the allocating version:

@tensor d[i] := b[i, j] * a[j]

leads to

ERROR: MethodError: no method matching similar_from_indices(::Type{Float64}, ::Tuple{Int64}, ::StaticArrays.SArray{Tuple{3,3},Float64,2,9}, ::SVector{3,Float64}, ::Type{Val{:N}}, ::Type{Val{:N}})
Closest candidates are:
  similar_from_indices(::Type{T}, ::Tuple{Vararg{Int64,N}}, ::Union{Base.ReshapedArray{T,N,A,MI} where MI<:Tuple{Vararg{Base.MultiplicativeInverses.SignedMultiplicativeInverse{Int64},N} where N} where A<:Union{DenseArray, SubArray{T,N,P,I,true} where I<:Tuple{Union{Base.Slice, UnitRange},Vararg{Any,N} where N} where P where N where T}, DenseArray{T,N}, SubArray{T,N,A,I,L} where L} where I<:Tuple{Vararg{Union{Base.AbstractCartesianIndex, Int64, Range{Int64}},N} where N} where A<:Union{Base.ReshapedArray{T,N,A,MI} where MI<:Tuple{Vararg{Base.MultiplicativeInverses.SignedMultiplicativeInverse{Int64},N} where N} where A<:Union{DenseArray, SubArray{T,N,P,I,true} where I<:Tuple{Union{Base.Slice, UnitRange},Vararg{Any,N} where N} where P where N where T} where N where T, DenseArray} where N where T, ::Union{Base.ReshapedArray{T,N,A,MI} where MI<:Tuple{Vararg{Base.MultiplicativeInverses.SignedMultiplicativeInverse{Int64},N} where N} where A<:Union{DenseArray, SubArray{T,N,P,I,true} where I<:Tuple{Union{Base.Slice, UnitRange},Vararg{Any,N} where N} where P where N where T}, DenseArray{T,N}, SubArray{T,N,A,I,L} where L} where I<:Tuple{Vararg{Union{Base.AbstractCartesianIndex, Int64, Range{Int64}},N} where N} where A<:Union{Base.ReshapedArray{T,N,A,MI} where MI<:Tuple{Vararg{Base.MultiplicativeInverses.SignedMultiplicativeInverse{Int64},N} where N} where A<:Union{DenseArray, SubArray{T,N,P,I,true} where I<:Tuple{Union{Base.Slice, UnitRange},Vararg{Any,N} where N} where P where N where T} where N where T, DenseArray} where N where T, ::Type{Val{CA}}, ::Type{Val{CB}}) where {T, N, CA, CB} at /Users/ortner/.julia/v0.6/TensorOperations/src/auxiliary/stridedarray.jl:41
  similar_from_indices(::Type{T}, ::Tuple{Vararg{Int64,N}}, ::Union{Base.ReshapedArray{T,N,A,MI} where MI<:Tuple{Vararg{Base.MultiplicativeInverses.SignedMultiplicativeInverse{Int64},N} where N} where A<:Union{DenseArray, SubArray{T,N,P,I,true} where I<:Tuple{Union{Base.Slice, UnitRange},Vararg{Any,N} where N} where P where N where T}, DenseArray{T,N}, SubArray{T,N,A,I,L} where L} where I<:Tuple{Vararg{Union{Base.AbstractCartesianIndex, Int64, Range{Int64}},N} where N} where A<:Union{Base.ReshapedArray{T,N,A,MI} where MI<:Tuple{Vararg{Base.MultiplicativeInverses.SignedMultiplicativeInverse{Int64},N} where N} where A<:Union{DenseArray, SubArray{T,N,P,I,true} where I<:Tuple{Union{Base.Slice, UnitRange},Vararg{Any,N} where N} where P where N where T} where N where T, DenseArray} where N where T, ::Union{Base.ReshapedArray{T,N,A,MI} where MI<:Tuple{Vararg{Base.MultiplicativeInverses.SignedMultiplicativeInverse{Int64},N} where N} where A<:Union{DenseArray, SubArray{T,N,P,I,true} where I<:Tuple{Union{Base.Slice, UnitRange},Vararg{Any,N} where N} where P where N where T}, DenseArray{T,N}, SubArray{T,N,A,I,L} where L} where I<:Tuple{Vararg{Union{Base.AbstractCartesianIndex, Int64, Range{Int64}},N} where N} where A<:Union{Base.ReshapedArray{T,N,A,MI} where MI<:Tuple{Vararg{Base.MultiplicativeInverses.SignedMultiplicativeInverse{Int64},N} where N} where A<:Union{DenseArray, SubArray{T,N,P,I,true} where I<:Tuple{Union{Base.Slice, UnitRange},Vararg{Any,N} where N} where P where N where T} where N where T, DenseArray} where N where T, ::Type{Val{CA}}) where {T, N, CA} at /Users/ortner/.julia/v0.6/TensorOperations/src/auxiliary/stridedarray.jl:41
  similar_from_indices(::Type{T}, ::Tuple{Vararg{Int64,N}}, ::Union{Base.ReshapedArray{T,N,A,MI} where MI<:Tuple{Vararg{Base.MultiplicativeInverses.SignedMultiplicativeInverse{Int64},N} where N} where A<:Union{DenseArray, SubArray{T,N,P,I,true} where I<:Tuple{Union{Base.Slice, UnitRange},Vararg{Any,N} where N} where P where N where T}, DenseArray{T,N}, SubArray{T,N,A,I,L} where L} where I<:Tuple{Vararg{Union{Base.AbstractCartesianIndex, Int64, Range{Int64}},N} where N} where A<:Union{Base.ReshapedArray{T,N,A,MI} where MI<:Tuple{Vararg{Base.MultiplicativeInverses.SignedMultiplicativeInverse{Int64},N} where N} where A<:Union{DenseArray, SubArray{T,N,P,I,true} where I<:Tuple{Union{Base.Slice, UnitRange},Vararg{Any,N} where N} where P where N where T} where N where T, DenseArray} where N where T) where {T, N} at /Users/ortner/.julia/v0.6/TensorOperations/src/auxiliary/stridedarray.jl:27
  ...
Stacktrace:
 [1] macro expansion at /Users/ortner/.julia/v0.6/TensorOperations/src/indexnotation/product.jl:58 [inlined]
 [2] deindexify at /Users/ortner/.julia/v0.6/TensorOperations/src/indexnotation/product.jl:51 [inlined] (repeats 2 times)

cortner avatar Nov 29 '17 22:11 cortner

with Einsum, the allocating version creates an Array instead of a SArray (ideally) or MArray

cortner avatar Nov 29 '17 22:11 cortner

Did you ever consider contributing your faster code to Base, and putting it in permutedims!(::StridedArray, perm) in base/multidimensional.jl?

The problem is that it's actually quite a bit of code, but once you have it it can immediately do more general things. It's basically a recursive blocking strategy (like the transpose! in Base, which I contributed long time ago), but then for arbitrary ranks and strides and permutations.

I certainly think that the new code in Base (since the introduction of PermutedDimsArray) has actually had a negative impact on the speed. My favorite (but worst case) example is (timing both allocating and non-allocating versions)

julia> A=randn(ntuple(n->10,8));
julia> @time B=permutedims(A,(8,7,6,5,4,3,2,1));
  3.814514 seconds (80.71 k allocations: 766.801 MiB, 1.30% gc time)
julia> @time permutedims!(B,A,(8,7,6,5,4,3,2,1));
  4.298002 seconds (8 allocations: 416 bytes)

(this used to be around 2 seconds) whereas

julia> @time @tensor B[:] := A[-8,-7,-6,-5,-4,-3,-2,-1];
  0.715068 seconds (25 allocations: 762.941 MiB, 0.30% gc time)
julia> @time @tensor B[:] = A[-8,-7,-6,-5,-4,-3,-2,-1];
  0.312276 seconds (19 allocations: 1.516 KiB)

Jutho avatar Nov 29 '17 22:11 Jutho

@cortner, similar_from_indices and contract! are two of the functions one needs to implement to make TensorOperations work on custom array types, together with add! and maybe trace!.

Jutho avatar Nov 29 '17 22:11 Jutho

ok - so what I ask is feasible in principle and possibly not even too difficult? It is not urgent for me, but if you are willing leave this issue open then at some point I might break down and have a go at it.

cortner avatar Nov 29 '17 22:11 cortner

Sure, no problem to leave this open.

I think it's certainly feasible, I am not sure whether you will be able to get optimal code for StaticArrays (i.e. completely/partially unrolled as described by @andyferris ), as the necessary information about the permutations won't be available at compile time.

Not sure where the implementation needs to go. I haven't worked with conditional dependencies yet, and I would prefer not to require StaticArrays for TensorOperations. But I guess this is already possible in Julia?

Jutho avatar Nov 29 '17 22:11 Jutho