Operator print tweak for schedule display

[mtsokol]

Hi @willow-ahrens,

I noticed that schedule printing gives an incorrect syntax for an operator located between operators.

using Finch

LEN = 150;
a_raw = rand(LEN, LEN - 5);
b_raw = rand(LEN, LEN - 5);
a = lazy(swizzle(Tensor(a_raw), 1, 2));
b = lazy(swizzle(Tensor(b_raw), 1, 2));

plan = sum(broadcast(*, a[:, nothing, :], b[nothing, :, :]), dims=(3,));
result = compute(plan, verbose=true);

gives:

:(function var"##compute#264"(prgm)
      begin
          V = (((((((((((((((prgm.children[1]).children[2]).children[2]).children[3]).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[3]).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}}}}}
          A0 = V::Tensor{DenseLevel{Int64, DenseLevel{Int64, ElementLevel{0.0, Float64, Int64, Vector{Float64}}}}}
          A1 = V_2::Tensor{DenseLevel{Int64, DenseLevel{Int64, ElementLevel{0.0, Float64, Int64, Vector{Float64}}}}}
          A2 = Tensor(Dense(Dense(Element{0.0, Float64}())))::Tensor{DenseLevel{Int64, DenseLevel{Int64, ElementLevel{0.0, Float64, Int64, Vector{Float64}}}}}
          @finch mode = :fast begin
                  A2 .= 0.0
                  for i1 = _
                      for i2 = _
                          for i0 = _
                              A2[i0, i2] << + >>= (*)(A0[i0, i1], A1[i2, i1])
                          end
                      end
                  end
                  return A2
              end
          return (A2,)
      end
  end)

When this output is pasted to the interpreter (as it's just julia code) it complains that the syntax is invalid:

ERROR: ParseError:
# Error @ REPL[17]:13:47
                          for i2 = _
                              A2[i2, i0] << + >>= (*)(A0[i0, i1], A1[i2, i1])
#                                             └─┘ ── invalid identifier
Stacktrace:
 [1] top-level scope
   @ none:1

With this fix the result is:

:(function var"##compute#264"(prgm)
      begin
          V = (((((((((((((((prgm.children[1]).children[2]).children[2]).children[3]).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[3]).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}}}}}
          A0 = V::Tensor{DenseLevel{Int64, DenseLevel{Int64, ElementLevel{0.0, Float64, Int64, Vector{Float64}}}}}
          A1 = V_2::Tensor{DenseLevel{Int64, DenseLevel{Int64, ElementLevel{0.0, Float64, Int64, Vector{Float64}}}}}
          A2 = Tensor(Dense(Dense(Element{0.0, Float64}())))::Tensor{DenseLevel{Int64, DenseLevel{Int64, ElementLevel{0.0, Float64, Int64, Vector{Float64}}}}}
          @finch mode = :fast begin
                  A2 .= 0.0
                  for i1 = _
                      for i2 = _
                          for i0 = _
                              A2[i0, i2] << (+) >>= (*)(A0[i0, i1], A1[i2, i1])
                          end
                      end
                  end
                  return A2
              end
          return (A2,)
      end
  end)

And it can be copied and reused directly. WDYT?

[mtsokol]

I think we could just always return Symbol(val) anyway - I see that it's no-op for functions:

julia> :(1 << $(Symbol(map)) >>= 1)
:(1 << map >>= 1)

[willow-ahrens]

I'm not sure this is always true. Different functions in different modules can have the same name.

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Files with missing lines	Coverage Δ
src/FinchLogic/nodes.jl	71.42% <ø> (ø)
src/scheduler/LogicCompiler.jl	80.90% <100.00%> (+0.53%)	⬆️

... and 1 file with indirect coverage changes

[willow-ahrens]

In the case of +, ideally we would return +=. However, if we must use the <<>> notation, this is helpful. i will merge this now, and feel free to file a PR to get it to use the incrementing assignments if you would like. There is a list of them in FinchNotation/syntax.jl

[willow-ahrens]

Thanks for the fix!