[Mlir-commits] [mlir] [mlir][linalg] Implement TilingInterface for winograd operators (PR #96184)

[llvmlistbot at llvm.org]

================
@@ -2922,6 +3021,129 @@ LogicalResult WinogradInputTransformOp::verify() {
   return success();
 }
 
+SmallVector<Range>
+WinogradInputTransformOp::getIterationDomain(OpBuilder &builder) {
+  Location loc = getLoc();
+  IntegerAttr zeroAttr = builder.getIndexAttr(0);
+  IntegerAttr oneAttr = builder.getIndexAttr(1);
+  Value output = getOutput();
+  int64_t outputRank = getOutputOperandRank();
+  SmallVector<Range> loopBounds(outputRank - 2);
+  for (unsigned dim = 2; dim < outputRank; ++dim) {
+    loopBounds[dim - 2].offset = zeroAttr;
+    loopBounds[dim - 2].size = getDimValue(builder, loc, output, dim);
+    loopBounds[dim - 2].stride = oneAttr;
+  }
+  return loopBounds;
+}
+
----------------
Max191 wrote:

I would suggest still including the outermost 2 alpha dims in the iteration domain, even though they shouldn't be tiled. The reason for this is so that they can be tiled and fused with other ops that implement tililng interface. For example, say there is an element-wise op that performs an `arith.truncf` after the winograd op. Ideally, we would want to fuse that truncf in with the innermost loop of the winograd op, so that the truncf can happen in register, rather than being stored to memory and loaded by a later loop to do the truncf.

Having all dimensions of the tensor be part of the tiling interface makes it easier to match iteration domains for successive ops when doing tiling + fusion.

You can add the dimensions to the iteration domain, and add an assert on the `sizes` and `offsets` in the tiled implementation to ensure that the tile sizes and offsets are always taking the full dimension at those dims. That way we ensure tiling does not touch those dimensions, but the iteration domain will include it.

https://github.com/llvm/llvm-project/pull/96184