Iteration Order for Masked Operations, Iteration Space Sharing

[sohnryang]

Thanks for the question.

• Suppose two operations in a subgraph produce output matrices of sizes 256x256 and 128x256, respectively. Can I use an execution granularity of 128x128 and simply mask/silently ignore the execution when the spatial grid exceeds the boundaries of the smaller matrix? -> Yes, this is valid. The execution granularity [w, h, k] sets the tile size, not the grid dimensions. Different operations in the same subgraph may have different tile counts based on their output sizes; each operation's output is fully tiled and computed. This is standard output-stationary scheduling, as shown in several PROBLEM.md examples.

• Additionally, is it allowed for a single subgraph to have multiple output operations of different types—specifically, one being a MatMul and the other a Pointwise operation? -> yes, as long as your scheduling produced all results at least once.

Originally posted by @yarongmu-google in #58

According to the clarification, it is possible to mask some of tile executions to respect the tensor boundaries. In this case, how do we specify the iteration order? Do we simply specify the iteration order for the tensor with the most iteration?

Also, the clarification in #20 introduced a notion of "sharing a same iteration space" and it seems like this is related to the issue presented in #58. Can you clarify what that notion exactly means?

[yarongmu-google]

Thanks for the question.

Given that multiple outputs of ops have caused a lot of confusion, I just removed all such cases from the benchmarks (both released and unreleased) - hope that makes things easier.

[sohnryang]

The question still stands even if all operations have single output. Consider a computation graph that creates tensors with different dimensions:

graph LR
    Tensor0(("Tensor[0]<br>256x128"))
    Tensor1(("Tensor[1]<br>128x256"))
    Tensor2(("Tensor[2]<br>256x256"))
    Tensor3(("Tensor[3]<br>128x256"))
    Op0["Op[0]<br>MatMul"]
    Op1["Op[1]<br>Pointwise"]

    Tensor0 --> Op0
    Tensor1 --> Op0
    Op0 --> Tensor2
    Tensor1 --> Op1
    Op1 --> Tensor3

Loading

Suppose that we fuse the whole graph into one subgraph, and choose [128, 128, 128] as execution granularity, which is valid according to the clarification in #58. How do we set iteration order for this subgraph? For example, if I set [0, 2, 3, 1] as iteration order then iteration 2 and 3 are simply "masked off" for the computation of Tensor[3]? In this example I could somehow guess the no-op iteration because at least one axis of the output tensors matched. How about this graph?

graph LR
    Tensor0(("Tensor[0]<br>256x128"))
    Tensor1(("Tensor[1]<br>128x256"))
    Tensor2(("Tensor[2]<br>256x256"))
    Tensor3(("Tensor[3]<br>256x512"))
    Tensor4(("Tensor[4]<br>128x256"))
    Tensor5(("Tensor[5]<br>256x512"))
    Op0["Op[0]<br>MatMul"]
    Op1["Op[1]<br>Pointwise"]
    Op2["Op[2]<br>MatMul"]

    Tensor0 --> Op0
    Tensor1 --> Op0
    Op0 --> Tensor2
    Tensor1 --> Op1
    Op1 --> Tensor4
    Tensor2 --> Op2
    Tensor3 --> Op2
    Op2 --> Tensor5

Loading

If we set execution granularity as [128, 128, 256] and iteration order as [0, 4, 1, 5, 2, 6, 3, 7], which iterations compute parts of Tensor[4] and which iterations don't?

[papp-pal-andras]

This is a great question, I'd also be interested in it. Allowing to mask parts of an operation seems to raise several questions, and the iteration space is then not shared anymore.