[Mlir-commits] [mlir] [MLIR] Add fusability query to TilingInterface (PR #166502)

[Quinn Dawkins]

================
@@ -359,6 +359,52 @@ struct LinalgOpTilingInterface
     /// Inline the op payload and store the result.
     return inlinePayload(builder, linalgOp, ivs, indexedValues);
   }
+
+  bool isOpFusableWithConsumerSlice(Operation *op, unsigned resultNumber,
+                                    ArrayRef<OpFoldResult> offsets,
+                                    ArrayRef<OpFoldResult> sizes) const {
+    return !cast<LinalgOp>(op).getShapesToLoopsMap();
+  }
+
+  bool isOpFusableWithProducerSlices(
+      Operation *op, ArrayRef<unsigned> operandNumbers,
+      ArrayRef<SmallVector<OpFoldResult>> allOffsets,
+      ArrayRef<SmallVector<OpFoldResult>> allSizes) const {
+
+    auto linalgOp = cast<LinalgOp>(op);
----------------
qedawkins wrote:

The ShapesToLoopsMap check is for the same reason as above.

> In theory it is feasible to fuse with any producer, irrespective of what slice of the operand is used

This is certainly not true for all operations. Not all tensor operands and dimensions are partitionable. For example, the inner most dimension of a `tensor.gather` must have size equal to the number of gathered dims, meaning we can never fuse into a producer that partitions that dimension. Additionally when fusing along multiple operands, if their slices don't line up, the feasibility of fusion depends on the loop type producing it. In the case of scf.for and scf.forall, it is impossible to fuse without fissioning the loop since we would need to relayout one or more of the mismatched operands.

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