Broadcasting issue in MMM operation

[mtsokol]

Hi @willow-ahrens,

Here's a reproduction code for broadcasting issue that I found while implementing SDDMM:

using Finch

LEN = 10;
a_raw = rand(LEN, LEN - 5) * 10;
b_raw = rand(LEN, LEN - 5) * 10;
c_raw = rand(LEN, LEN) * 10;

a = lazy(swizzle(Tensor(a_raw), 1, 2));
b = lazy(swizzle(Tensor(b_raw), 1, 2));
c = lazy(swizzle(Tensor(c_raw), 1, 2));

# doesn't equal
plan = broadcast(*, broadcast(*, c[:, :, nothing], a[:, nothing, :]), b[nothing, :, :]);
# works correctly
# plan = c[:, :, nothing] .* a[:, nothing, :] .* b[nothing, :, :];

result = compute(plan);

actual = reshape(c_raw, 10, 10, 1) .* reshape(a_raw, 10, 1, 5) .* reshape(b_raw, 1, 10, 5);
# other notation
# actual = broadcast(*, broadcast(*, reshape(c_raw, 10, 10, 1), reshape(a_raw, 10, 1, 5)), reshape(b_raw, 1, 10, 5));

isequal(result, actual)

[mtsokol]

@willow-ahrens Here's the output of the debug mode for .* broadcasting that works correctly:

plan = c[:, :, nothing] .* a[:, nothing, :] .* b[nothing, :, :];
compute(plan, verbose=true);

``` Executing: :(function var"##compute#378"(prgm) begin V = (((((((((((((prgm.children[1]).children[2]).children[2]).children[1]).children[2]).children[2]).children[1]).children[1]).children[1]).children[1]).children[1]).children[1]).children[2]).tns.val::Tensor{DenseLevel{Int64, DenseLevel{Int64, ElementLevel{0.0, Float64, Int64, Vector{Float64}}}}} V_2 = (((((((((((((prgm.children[1]).children[2]).children[2]).children[1]).children[2]).children[3]).children[1]).children[1]).children[1]).children[1]).children[1]).children[1]).children[2]).tns.val::Tensor{DenseLevel{Int64, DenseLevel{Int64, ElementLevel{0.0, Float64, Int64, Vector{Float64}}}}} V_3 = ((((((((((((prgm.children[1]).children[2]).children[2]).children[1]).children[3]).children[1]).children[1]).children[1]).children[1]).children[1]).children[1]).children[2]).tns.val::Tensor{DenseLevel{Int64, DenseLevel{Int64, ElementLevel{0.0, Float64, Int64, Vector{Float64}}}}} A0 = V::Tensor{DenseLevel{Int64, DenseLevel{Int64, ElementLevel{0.0, Float64, Int64, Vector{Float64}}}}} A2 = V_2::Tensor{DenseLevel{Int64, DenseLevel{Int64, ElementLevel{0.0, Float64, Int64, Vector{Float64}}}}} A4 = V_3::Tensor{DenseLevel{Int64, DenseLevel{Int64, ElementLevel{0.0, Float64, Int64, Vector{Float64}}}}} A6 = Tensor(Dense(Dense(Dense(Element{0.0, Float64}()))))::Tensor{DenseLevel{Int64, DenseLevel{Int64, DenseLevel{Int64, ElementLevel{0.0, Float64, Int64, Vector{Float64}}}}}} @finch mode = :fast begin A6 .= 0.0 for i20 = _ for i12 = _ for i11 = _ A6[i11, i12, i20] << Finch.FinchNotation.InitWriter{0.0}() >>= (*)((*)(A0[i11, i12], A2[i11, i20]), A4[i12, i20]) end end end return A6 end return (A6,) end end) ```

And here's the output of the debug mode for broadcast broadcasting (used internally by finch-tensor) that we're debugging:

plan = broadcast(*, broadcast(*, c[:, :, nothing], a[:, nothing, :]), b[nothing, :, :]);
compute(plan, verbose=true);

``` Executing: :(function var"##compute#421"(prgm) begin V = (((((((((((((((((prgm.children[1]).children[2]).children[2]).children[1]).children[2]).children[1]).children[1]).children[2]).children[1]).children[2]).children[1]).children[1]).children[1]).children[1]).children[1]).children[1]).children[2]).tns.val::Tensor{DenseLevel{Int64, DenseLevel{Int64, ElementLevel{0.0, Float64, Int64, Vector{Float64}}}}} V_2 = (((((((((((((((((prgm.children[1]).children[2]).children[2]).children[1]).children[2]).children[1]).children[1]).children[2]).children[1]).children[3]).children[1]).children[1]).children[1]).children[1]).children[1]).children[1]).children[2]).tns.val::Tensor{DenseLevel{Int64, DenseLevel{Int64, ElementLevel{0.0, Float64, Int64, Vector{Float64}}}}} V_3 = ((((((((((((prgm.children[1]).children[2]).children[2]).children[1]).children[3]).children[1]).children[1]).children[1]).children[1]).children[1]).children[1]).children[2]).tns.val::Tensor{DenseLevel{Int64, DenseLevel{Int64, ElementLevel{0.0, Float64, Int64, Vector{Float64}}}}} A0 = V::Tensor{DenseLevel{Int64, DenseLevel{Int64, ElementLevel{0.0, Float64, Int64, Vector{Float64}}}}} A2 = V_2::Tensor{DenseLevel{Int64, DenseLevel{Int64, ElementLevel{0.0, Float64, Int64, Vector{Float64}}}}} A5 = V_3::Tensor{DenseLevel{Int64, DenseLevel{Int64, ElementLevel{0.0, Float64, Int64, Vector{Float64}}}}} A7 = Tensor(Dense(Dense(Dense(Element{0.0, Float64}()))))::Tensor{DenseLevel{Int64, DenseLevel{Int64, DenseLevel{Int64, ElementLevel{0.0, Float64, Int64, Vector{Float64}}}}}} @finch mode = :fast begin A7 .= 0.0 for i31 = _ for i30 = _ for i21 = _ A7[i21, i30, i31] << Finch.FinchNotation.InitWriter{0.0}() >>= (*)((*)(A0[i21, 1], A2[i21, 1]), A5[i30, i31]) end end end return A7 end return (A7,) end end) ```

---

I think the key difference (that I found with https://www.diffchecker.com) is in:

A6[i11, i12, i20] << Finch.FinchNotation.InitWriter{0.0}() >>= (*)((*)(A0[i11, i12], A2[i11, i20]), A4[i12, i20])
vs
A7[i21, i30, i31] << Finch.FinchNotation.InitWriter{0.0}() >>= (*)((*)(A0[i21, 1], A2[i21, 1]), A5[i30, i31])

For some reason for broadcast(...) a 1 was placed there instead of the index. WDYT?

[hameerabbasi]

I think it's during broadcasting that an index would be replaced with 1, right @willow-ahrens? That would mean that the "broadcast indices" are somehow at an incorrect location.