Skip to content

Commit 93d07b9

Browse files
committed
Address comments
1 parent e5a56f5 commit 93d07b9

File tree

1 file changed

+11
-3
lines changed

1 file changed

+11
-3
lines changed

mlir/lib/Dialect/Linalg/Transforms/TilingInterfaceImpl.cpp

Lines changed: 11 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -324,6 +324,13 @@ struct LinalgOpTilingInterface
324324
// External Model for implementing `PartialReductionInterface` for `LinalgOp`s.
325325
//===----------------------------------------------------------------------===//
326326

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.
327334
static AffineMap getPartialResultAffineMap(LinalgOp linalgOp,
328335
ArrayRef<int> reductionDims,
329336
unsigned resultNumber) {
@@ -491,9 +498,10 @@ struct LinalgOpPartialReductionInterface
491498
SmallVector<Operation *> mergeOperations;
492499
SmallVector<Value> replacements;
493500
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.
497505
AffineMap partialMap =
498506
getPartialResultAffineMap(linalgOp, reductionDims, idx);
499507
SmallVector<int64_t> partialReductionDims;

0 commit comments

Comments
 (0)