[Mlir-commits] [mlir] [mlir][PartialReductionTilingInterface] Generalize implementation of `tileUsingSCF` for `ReductionTilingStrategy::PartialOuterReduction`. (PR #143467)

[llvmlistbot at llvm.org]

================
@@ -327,23 +328,48 @@ struct LinalgOpTilingInterface
 // External Model for implementing `PartialReductionInterface` for `LinalgOp`s.
 //===----------------------------------------------------------------------===//
 
-/// Return an AffineMap for a partial result for the given result number,
-/// assuming the partial tiling strategy is outer-reduction loop +
-/// inner-parallel tile. The returned AffineMap can be used as the replacement
-/// AffineMap for the inner-parallel tile linalg op for the given result number.
-///
-/// The new AffineMap is the old AffineMap with reduction dimensions appended
-/// at end.
-static AffineMap getPartialResultAffineMap(LinalgOp linalgOp,
-                                           ArrayRef<int> reductionDims,
-                                           unsigned resultNumber) {
-  AffineMap map =
-      linalgOp.getMatchingIndexingMap(linalgOp.getDpsInitOperand(resultNumber));
-  for (int redPos : reductionDims) {
-    map = map.insertResult(getAffineDimExpr(redPos, linalgOp.getContext()),
-                           map.getNumResults());
+/// Return an AffineMaps to use for the `outs` operands of the linalg op
+/// generated for partial results. The new AffineMap is the AffineMap of the
+/// untiled op with reduction dimensions appended at end in order in which they
+/// were specified during tiling.
----------------
MaheshRavishankar wrote:

We need to have a fixed way of doing this even if the order does not matter. So I used this as a way to allow caller to control it (if it matters, but I agree it shouldnt).

W.R.T to where to insert the dimensions, it might be a batch dimension, but I am not totally clear that it should always go to the outermost. If needed doing that "swap" could be a post tiling transformations (using something like  [`interchangeGenericOp`](https://github.com/llvm/llvm-project/blob/users/MaheshRavishankar/splitKGeneralizeOuterReduction/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h#L788) or something. 

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