@@ -324,6 +324,13 @@ struct LinalgOpTilingInterface
324
324
// External Model for implementing `PartialReductionInterface` for `LinalgOp`s.
325
325
// ===----------------------------------------------------------------------===//
326
326
327
+ // / Return an AffineMap for a partial result for the given result number,
328
+ // / assuming the partial tiling strategy is outer-reduction loop +
329
+ // / inner-parallel tile. The returned AffineMap can be used as the replacement
330
+ // / AffineMap for the inner-parallel tile linalg op for the given result number.
331
+ // /
332
+ // / The new AffineMap is the old AffineMap with reduction dimensions appended
333
+ // / at end.
327
334
static AffineMap getPartialResultAffineMap (LinalgOp linalgOp,
328
335
ArrayRef<int > reductionDims,
329
336
unsigned resultNumber) {
@@ -491,9 +498,10 @@ struct LinalgOpPartialReductionInterface
491
498
SmallVector<Operation *> mergeOperations;
492
499
SmallVector<Value> replacements;
493
500
for (int idx : llvm::seq (numInits)) {
494
- // linalg.reduce's iteration space is the result's iteration space (and
495
- // not the operations iteration space). To account for this, permute the
496
- // reduction dimensions based on the partial result map.
501
+ // linalg.reduce's iteration space is the tiled result's iteration space
502
+ // (and not the tiled operation's iteration space). To account for this,
503
+ // permute the reduction dimensions based on the partial result map of the
504
+ // tiled result.
497
505
AffineMap partialMap =
498
506
getPartialResultAffineMap (linalgOp, reductionDims, idx);
499
507
SmallVector<int64_t > partialReductionDims;
0 commit comments