From c723bd4bb7365ae828737befb9b22aff18506438 Mon Sep 17 00:00:00 2001 From: Matthias Gehre Date: Tue, 3 Dec 2024 13:31:31 +0100 Subject: [PATCH 1/3] computeSliceParameters: use full slice if affine exprs are non-monotonic --- mlir/include/mlir/IR/AffineExpr.h | 5 ++++ mlir/include/mlir/IR/AffineMap.h | 3 ++ mlir/lib/Dialect/Linalg/Utils/Utils.cpp | 7 +++-- mlir/lib/IR/AffineExpr.cpp | 33 ++++++++++++++++++++++ mlir/lib/IR/AffineMap.cpp | 4 +++ mlir/test/Dialect/Linalg/tile-tensors.mlir | 32 +++++++++++++++++++++ 6 files changed, 82 insertions(+), 2 deletions(-) diff --git a/mlir/include/mlir/IR/AffineExpr.h b/mlir/include/mlir/IR/AffineExpr.h index a93e74b449ceed..9abef1d70893c2 100644 --- a/mlir/include/mlir/IR/AffineExpr.h +++ b/mlir/include/mlir/IR/AffineExpr.h @@ -110,6 +110,11 @@ class AffineExpr { /// floordiv, ceildiv, and mod is only allowed w.r.t constants. bool isPureAffine() const; + /// Returns true if this expression is monotonic with respect to the + /// AffineDimExpr, i.e. increasing the value of any AffineDimExpr will never + /// decrease the value of the result. + bool isMonotonic() const; + /// Returns the greatest known integral divisor of this affine expression. The /// result is always positive. int64_t getLargestKnownDivisor() const; diff --git a/mlir/include/mlir/IR/AffineMap.h b/mlir/include/mlir/IR/AffineMap.h index 676da6d1764970..ca7d6b33dbd6a3 100644 --- a/mlir/include/mlir/IR/AffineMap.h +++ b/mlir/include/mlir/IR/AffineMap.h @@ -364,6 +364,9 @@ class AffineMap { /// Returns true if the AffineMap represents a symbol-less permutation map. bool isPermutation() const; + // Returns true if every result is monotonic. + bool isMonotonic() const; + /// Returns the map consisting of the `resultPos` subset. AffineMap getSubMap(ArrayRef resultPos) const; diff --git a/mlir/lib/Dialect/Linalg/Utils/Utils.cpp b/mlir/lib/Dialect/Linalg/Utils/Utils.cpp index 6a3f2fc5fbc496..8db6430d41ca10 100644 --- a/mlir/lib/Dialect/Linalg/Utils/Utils.cpp +++ b/mlir/lib/Dialect/Linalg/Utils/Utils.cpp @@ -616,7 +616,11 @@ computeSliceParameters(OpBuilder &builder, Location loc, Value valueToTile, sliceParams.strides.reserve(rank); for (unsigned r = 0; r < rank; ++r) { LLVM_DEBUG(llvm::dbgs() << "computeSliceParameters: for dim#" << r); - if (!isTiled(map.getSubMap({r}), tileSizes)) { + auto m = map.getSubMap({r}); + // The offset & size computation below only handles the case when + // the map is monotonic, i.e. the min and max values are attained at the + // lower and upper bounds of the iteration domain. + if (!isTiled(m, tileSizes) || !m.isMonotonic()) { sliceParams.offsets.push_back(builder.getIndexAttr(0)); OpFoldResult dim = createFoldedDimOp(builder, loc, valueToTile, r); sliceParams.sizes.push_back(dim); @@ -628,7 +632,6 @@ computeSliceParameters(OpBuilder &builder, Location loc, Value valueToTile, // Tiling creates a new slice at the proper index, the slice step is 1 // (i.e. the op does not subsample, stepping occurs in the loop). - auto m = map.getSubMap({r}); LLVM_DEBUG(llvm::dbgs() << "computeSliceParameters: submap: " << m << "\n"); IRRewriter rewriter(builder); OpFoldResult offset = makeComposedFoldedAffineApply(rewriter, loc, m, lbs); diff --git a/mlir/lib/IR/AffineExpr.cpp b/mlir/lib/IR/AffineExpr.cpp index fc7ede279643ed..3e01062c65274a 100644 --- a/mlir/lib/IR/AffineExpr.cpp +++ b/mlir/lib/IR/AffineExpr.cpp @@ -239,6 +239,39 @@ bool AffineExpr::isPureAffine() const { llvm_unreachable("Unknown AffineExpr"); } +static bool isNonNegativeConstant(AffineExpr expr) { + auto constant = dyn_cast(expr); + return constant && constant.getValue() >= 0; +} + +bool AffineExpr::isMonotonic() const { + switch (getKind()) { + case AffineExprKind::SymbolId: + case AffineExprKind::DimId: + case AffineExprKind::Constant: + return true; + case AffineExprKind::Add: { + auto op = llvm::cast(*this); + return op.getLHS().isMonotonic() && op.getRHS().isMonotonic(); + } + case AffineExprKind::Mul: { + // One operand must be a non-negative constant. + auto op = llvm::cast(*this); + return op.getLHS().isMonotonic() && op.getRHS().isMonotonic() && + (isNonNegativeConstant(op.getLHS()) || + isNonNegativeConstant(op.getRHS())); + } + case AffineExprKind::FloorDiv: + case AffineExprKind::CeilDiv: { + auto op = llvm::cast(*this); + return op.getLHS().isMonotonic() && isNonNegativeConstant(op.getRHS()); + } + case AffineExprKind::Mod: + return false; + } + llvm_unreachable("Unknown AffineExpr"); +} + // Returns the greatest known integral divisor of this affine expression. int64_t AffineExpr::getLargestKnownDivisor() const { AffineBinaryOpExpr binExpr(nullptr); diff --git a/mlir/lib/IR/AffineMap.cpp b/mlir/lib/IR/AffineMap.cpp index 5cbd0b090492b4..e6d8e9f970aba6 100644 --- a/mlir/lib/IR/AffineMap.cpp +++ b/mlir/lib/IR/AffineMap.cpp @@ -628,6 +628,10 @@ bool AffineMap::isPermutation() const { return isProjectedPermutation(); } +bool AffineMap::isMonotonic() const { + return all_of(getResults(), [](auto expr) { return expr.isMonotonic(); }); +} + AffineMap AffineMap::getSubMap(ArrayRef resultPos) const { SmallVector exprs; exprs.reserve(resultPos.size()); diff --git a/mlir/test/Dialect/Linalg/tile-tensors.mlir b/mlir/test/Dialect/Linalg/tile-tensors.mlir index 8f13c690704572..cacca8a637ab77 100644 --- a/mlir/test/Dialect/Linalg/tile-tensors.mlir +++ b/mlir/test/Dialect/Linalg/tile-tensors.mlir @@ -167,3 +167,35 @@ module attributes {transform.with_named_sequence} { transform.yield } } + +// ----- + +// CHECK-LABEL: func @non_monotonic_affine_expr +// CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]+]]: tensor +func.func @non_monotonic_affine_expr(%arg0 : tensor) -> tensor { + %c0 = arith.constant 0 : index + %0 = tensor.dim %arg0, %c0 : tensor + %empty = tensor.empty(%0) : tensor + + // CHECK: scf.for + // CHECK: %[[SIZE:[a-zA-Z0-9_]+]] = tensor.dim %[[ARG0]], + // CHECK: tensor.extract_slice %[[ARG0]][0] [%[[SIZE]]] [1] : tensor to tensor + %generic = linalg.generic + {indexing_maps = [affine_map<(d0) -> (d0 mod 3)>, + affine_map<(d0) -> (d0)>], + iterator_types = ["parallel"]} + ins(%arg0: tensor) + outs(%empty : tensor) { + ^bb0(%in : f32, %out: f32): + linalg.yield %in : f32 + } -> tensor + return %generic : tensor +} + +module attributes {transform.with_named_sequence} { + transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) { + %0 = transform.structured.match ops{["linalg.generic"]} in %arg1 : (!transform.any_op) -> !transform.any_op + %1, %loop = transform.structured.tile_using_for %0 tile_sizes [100] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) + transform.yield + } +} From fa056a80a4617a62ea43bc4c03afc6d15b1a0ad5 Mon Sep 17 00:00:00 2001 From: Matthias Gehre Date: Tue, 3 Dec 2024 14:40:03 +0100 Subject: [PATCH 2/3] Clarify naming --- mlir/include/mlir/IR/AffineExpr.h | 8 ++++---- mlir/include/mlir/IR/AffineMap.h | 5 +++-- mlir/lib/Dialect/Linalg/Utils/Utils.cpp | 2 +- mlir/lib/IR/AffineExpr.cpp | 11 +++++++---- mlir/lib/IR/AffineMap.cpp | 5 +++-- 5 files changed, 18 insertions(+), 13 deletions(-) diff --git a/mlir/include/mlir/IR/AffineExpr.h b/mlir/include/mlir/IR/AffineExpr.h index 9abef1d70893c2..28d00f1299f2f9 100644 --- a/mlir/include/mlir/IR/AffineExpr.h +++ b/mlir/include/mlir/IR/AffineExpr.h @@ -110,10 +110,10 @@ class AffineExpr { /// floordiv, ceildiv, and mod is only allowed w.r.t constants. bool isPureAffine() const; - /// Returns true if this expression is monotonic with respect to the - /// AffineDimExpr, i.e. increasing the value of any AffineDimExpr will never - /// decrease the value of the result. - bool isMonotonic() const; + /// Returns true if this expression is monotonicically increasing with respect + /// to the AffineDimExprs, i.e. increasing the value of any AffineDimExpr will + /// never decrease the value of the result. + bool isMonotonicallyIncreasing() const; /// Returns the greatest known integral divisor of this affine expression. The /// result is always positive. diff --git a/mlir/include/mlir/IR/AffineMap.h b/mlir/include/mlir/IR/AffineMap.h index ca7d6b33dbd6a3..e7d53b23539ccb 100644 --- a/mlir/include/mlir/IR/AffineMap.h +++ b/mlir/include/mlir/IR/AffineMap.h @@ -364,8 +364,9 @@ class AffineMap { /// Returns true if the AffineMap represents a symbol-less permutation map. bool isPermutation() const; - // Returns true if every result is monotonic. - bool isMonotonic() const; + // Returns true if every result is monotonically increasing. + // See AffineExpr::isMonotonicallyIncreasing(). + bool isComponentWiseMonotonicallyIncreasing() const; /// Returns the map consisting of the `resultPos` subset. AffineMap getSubMap(ArrayRef resultPos) const; diff --git a/mlir/lib/Dialect/Linalg/Utils/Utils.cpp b/mlir/lib/Dialect/Linalg/Utils/Utils.cpp index 8db6430d41ca10..e95fd65c95ed86 100644 --- a/mlir/lib/Dialect/Linalg/Utils/Utils.cpp +++ b/mlir/lib/Dialect/Linalg/Utils/Utils.cpp @@ -620,7 +620,7 @@ computeSliceParameters(OpBuilder &builder, Location loc, Value valueToTile, // The offset & size computation below only handles the case when // the map is monotonic, i.e. the min and max values are attained at the // lower and upper bounds of the iteration domain. - if (!isTiled(m, tileSizes) || !m.isMonotonic()) { + if (!isTiled(m, tileSizes) || !m.isComponentWiseMonotonicallyIncreasing()) { sliceParams.offsets.push_back(builder.getIndexAttr(0)); OpFoldResult dim = createFoldedDimOp(builder, loc, valueToTile, r); sliceParams.sizes.push_back(dim); diff --git a/mlir/lib/IR/AffineExpr.cpp b/mlir/lib/IR/AffineExpr.cpp index 3e01062c65274a..12ccadccf50ebc 100644 --- a/mlir/lib/IR/AffineExpr.cpp +++ b/mlir/lib/IR/AffineExpr.cpp @@ -244,7 +244,7 @@ static bool isNonNegativeConstant(AffineExpr expr) { return constant && constant.getValue() >= 0; } -bool AffineExpr::isMonotonic() const { +bool AffineExpr::isMonotonicallyIncreasing() const { switch (getKind()) { case AffineExprKind::SymbolId: case AffineExprKind::DimId: @@ -252,19 +252,22 @@ bool AffineExpr::isMonotonic() const { return true; case AffineExprKind::Add: { auto op = llvm::cast(*this); - return op.getLHS().isMonotonic() && op.getRHS().isMonotonic(); + return op.getLHS().isMonotonicallyIncreasing() && + op.getRHS().isMonotonicallyIncreasing(); } case AffineExprKind::Mul: { // One operand must be a non-negative constant. auto op = llvm::cast(*this); - return op.getLHS().isMonotonic() && op.getRHS().isMonotonic() && + return op.getLHS().isMonotonicallyIncreasing() && + op.getRHS().isMonotonicallyIncreasing() && (isNonNegativeConstant(op.getLHS()) || isNonNegativeConstant(op.getRHS())); } case AffineExprKind::FloorDiv: case AffineExprKind::CeilDiv: { auto op = llvm::cast(*this); - return op.getLHS().isMonotonic() && isNonNegativeConstant(op.getRHS()); + return op.getLHS().isMonotonicallyIncreasing() && + isNonNegativeConstant(op.getRHS()); } case AffineExprKind::Mod: return false; diff --git a/mlir/lib/IR/AffineMap.cpp b/mlir/lib/IR/AffineMap.cpp index e6d8e9f970aba6..f6061a7ea1e1e6 100644 --- a/mlir/lib/IR/AffineMap.cpp +++ b/mlir/lib/IR/AffineMap.cpp @@ -628,8 +628,9 @@ bool AffineMap::isPermutation() const { return isProjectedPermutation(); } -bool AffineMap::isMonotonic() const { - return all_of(getResults(), [](auto expr) { return expr.isMonotonic(); }); +bool AffineMap::isComponentWiseMonotonicallyIncreasing() const { + return all_of(getResults(), + [](auto expr) { return expr.isMonotonicallyIncreasing(); }); } AffineMap AffineMap::getSubMap(ArrayRef resultPos) const { From 79ef42aef43f70d7d37afc78c22cfddd14db7d2e Mon Sep 17 00:00:00 2001 From: Matthias Gehre Date: Tue, 3 Dec 2024 14:43:35 +0100 Subject: [PATCH 3/3] Fix comment --- mlir/lib/Dialect/Linalg/Utils/Utils.cpp | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/mlir/lib/Dialect/Linalg/Utils/Utils.cpp b/mlir/lib/Dialect/Linalg/Utils/Utils.cpp index e95fd65c95ed86..e8ace93fb07862 100644 --- a/mlir/lib/Dialect/Linalg/Utils/Utils.cpp +++ b/mlir/lib/Dialect/Linalg/Utils/Utils.cpp @@ -618,8 +618,8 @@ computeSliceParameters(OpBuilder &builder, Location loc, Value valueToTile, LLVM_DEBUG(llvm::dbgs() << "computeSliceParameters: for dim#" << r); auto m = map.getSubMap({r}); // The offset & size computation below only handles the case when - // the map is monotonic, i.e. the min and max values are attained at the - // lower and upper bounds of the iteration domain. + // the map is monotonically increasing, i.e. the min and max values are + // attained at the lower and upper bounds of the iteration domain. if (!isTiled(m, tileSizes) || !m.isComponentWiseMonotonicallyIncreasing()) { sliceParams.offsets.push_back(builder.getIndexAttr(0)); OpFoldResult dim = createFoldedDimOp(builder, loc, valueToTile, r);