This is an attempt to use Tullio.jl to run tensor expressions on Dagger-based distributed arrays.
- You need this branch of Tullio.jl
- DArray struct is how we represent a distributed array
@dtulliomacro is in this file
julia> using DaggerArrays
julia> A = distribute(reshape(1:24, (4,6)), Blocks(2,3))
4×6 DArray{Int64, 2} with 2×2 chunks
julia> B = distribute(reshape(1:24, (6,4)), Blocks(3,2))
6×4 DArray{Int64, 2} with 2×2 chunks
julia> @dtullio C[i,k] := A[i,j] * B[j,k]
4×4 DArray{Any, 2} with 2×2 chunksjulia> @generated matmul(A, B) = :(@dtullio C[i,k] := A[i,j] * B[j,k])
matmul (generic function with 1 method)
julia> matmul(A,B)
ERROR: The function body AST defined by this @generated function is not pure. This likely means it contains a closure or comprehension.
Stacktrace:
[1] top-level scope at REPL[24]:1
[2] run_repl(::REPL.AbstractREPL, ::Any) at /build/julia/src/julia-1.5.2/usr/share/julia/stdlib/v1.5/REPL/src/REPL.jl:288This is required to implement dimension-independent operations
julia> A = distribute(reshape(1:24, (4,6)), Blocks(2,3))
4×6 DArray{Int64, 2} with 2×2 chunks
julia> B = distribute(reshape(1:24, (6,4)), Blocks(6,2))
6×4 DArray{Int64, 2} with 1×2 chunks
julia> @dtullio C[i,k] := A[i,j] * B[j,k]
ERROR: range of index j must agree
Stacktrace:
[1] top-level scope at /home/shashi/.julia/dev/Tullio/src/macro.jl:777
[2] top-level scope at /home/shashi/.julia/dev/DaggerArrays/src/tullio.jl:103
[3] run_repl(::REPL.AbstractREPL, ::Any) at /build/julia/src/julia-1.5.2/usr/share/julia/stdlib/v1.5/REPL/src/REPL.jl:288(Should auto-slice B to match the layout of A)