Eliminates a zero-initialization nest that is redundant before a GPU MMA matmul nest.
More...
Eliminates a zero-initialization nest that is redundant before a GPU MMA matmul nest.
Original: [init=zero] map i, j: Y[i, j] = 0 <MMA nest>: ... MmaFragmentLoad(c <- Y) ... EltwiseAdd(acc, c) ... MmaFragmentStore(Y)
After RedundantZeroInitElimination: <MMA nest>: ... GpuMmaFill(c) ... EltwiseAdd(acc, c) ... MmaFragmentStore(Y)
The C-fragment loads of Y in the MMA nest could only observe the zeros written by the init nest, so they are replaced by zero fills and the init nest is removed.
Prerequisites:
- The loop carries the zero-init annotation (ScheduleType_ZeroInit). The producer (naive MatMul/GEMM expansion) thereby certifies that the nest writes zeros to exactly the output region of the matmul nest directly following it.
- The nest is still a pure zero init, i.e. a perfect loop chain whose body is a single constant 0 -> assign -> Y block (it was not fused with other computation).
- The next sibling in the same sequence is a loop nest containing GPU MMA matmul nodes and fragment stores to Y.
- Within that nest, Y is only used as the memory operand of fragment stores and of C-fragment loads, and the C-fragment loads are only nested in maps and branches (no sequential loop could re-read stored partial results).
Note: the MMA nests produced by GpuMmaEinsumTransform never use the data already contained in Y beyond these C-fragment loads. If other MMA lowerings are matched in the future (e.g. an overwrite mode without C-fragment loads, or nests reading Y differently), the checks may need to be extended.
1.15.0