distributed matmul using Tullio.jl

tullio_matmul.jl

using Tullio
using Dagger
using Dagger: DArray, chunks, Thunk, ArrayDomain

struct Zero end
Base.zero(::Type{Thunk}) = delayed(()->Zero())()
Base.zero(::Type{Union{Zero, ArrayDomain}}) = Zero()
Base.:(+)(z::Zero, x) = x
Base.:(+)(x, z::Zero) = x
Base.:(+)(z::Zero, x::Zero) = z

function matmul(A::DArray, B::DArray)
    # Make the output domains
    dA = A.subdomains
    dB = B.subdomains

    dC = Union{Zero, ArrayDomain}[Zero() for i=1:size(dA, 1), j=1:size(dB, 2)]

    combine(da, db) = ArrayDomain(da.indexes[1], db.indexes[2])
    @tullio dC[i,j] = combine(dA[i,k], dB[k,j]) threads=false

    # Make the output chunks
    cA = chunks(A)
    cB = chunks(B)
    on_each = delayed((a,b) -> @tullio(c[i,j] := a[i, k] * b[k, j]))
    plus = (x,y) -> delayed(+)(x,y)
    cC = [zero(Thunk) for i=1:size(cA, 1), j=1:size(cB, 2)]
    @tullio (plus) (cC[i,j] = on_each(cA[i,k], cB[k,j])) threads=false

    # Make the output DArray
    DArray(promote_type(eltype(A), eltype(B)),
           combine(A.domain, B.domain), map(identity, dC), cC, cat)
end

A = compute(rand(Blocks(1000,100), 10000,1000));
D = compute(rand(Blocks(100,100), 1000,100));

matmul(A, D)

using Test
@test collect(A * D) ≈ collect(matmul(A, D))