[Mlir-commits] [mlir] 5b779a3 - [mlir][affine] Canonicalize single value affine.min/max

[Kai Sasaki]

Author: Kai Sasaki
Date: 2023-01-11T11:04:14+09:00
New Revision: 5b779a3df351a21b92f22e3716943c4be0bf5d50

URL: https://github.com/llvm/llvm-project/commit/5b779a3df351a21b92f22e3716943c4be0bf5d50
DIFF: https://github.com/llvm/llvm-project/commit/5b779a3df351a21b92f22e3716943c4be0bf5d50.diff

LOG: [mlir][affine] Canonicalize single value affine.min/max

Canonicalize identity affine.min/max to allow further optimization to follow the def-use chain of the given values. The reported issue is https://github.com/llvm/llvm-project/issues/59399.

Differential Revision: https://reviews.llvm.org/D141354

Added: 
    

Modified: 
    mlir/include/mlir/IR/AffineMap.h
    mlir/lib/Dialect/Affine/IR/AffineOps.cpp
    mlir/lib/IR/AffineMap.cpp
    mlir/test/Dialect/Affine/canonicalize.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/IR/AffineMap.h b/mlir/include/mlir/IR/AffineMap.h
index 3f484ba63e80..186c613b12fe 100644
--- a/mlir/include/mlir/IR/AffineMap.h
+++ b/mlir/include/mlir/IR/AffineMap.h
@@ -104,6 +104,10 @@ class AffineMap {
   /// dimensional identifiers.
   bool isIdentity() const;
 
+  /// Returns true if this affine map is an identity affine map on the symbol
+  /// identifiers.
+  bool isSymbolIdentity() const;
+
   /// Returns true if this affine map is a minor identity, i.e. an identity
   /// affine map (d0, ..., dn) -> (dp, ..., dn) on the most minor dimensions.
   bool isMinorIdentity() const;

diff  --git a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
index 1235934ce60a..be58118d3247 100644
--- a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
+++ b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
@@ -3042,6 +3042,9 @@ static OpFoldResult foldMinMaxOp(T op, ArrayRef<Attribute> operands) {
   SmallVector<int64_t, 2> results;
   auto foldedMap = op.getMap().partialConstantFold(operands, &results);
 
+  if (foldedMap.getNumSymbols() == 1 && foldedMap.isSymbolIdentity())
+    return op.getOperand(0);
+
   // If some of the map results are not constant, try changing the map in-place.
   if (results.empty()) {
     // If the map is the same, report that folding did not happen.

diff  --git a/mlir/lib/IR/AffineMap.cpp b/mlir/lib/IR/AffineMap.cpp
index 2b6d0bcf4193..c6873d32dadd 100644
--- a/mlir/lib/IR/AffineMap.cpp
+++ b/mlir/lib/IR/AffineMap.cpp
@@ -276,6 +276,18 @@ bool AffineMap::isIdentity() const {
   return true;
 }
 
+bool AffineMap::isSymbolIdentity() const {
+  if (getNumSymbols() != getNumResults())
+    return false;
+  ArrayRef<AffineExpr> results = getResults();
+  for (unsigned i = 0, numSymbols = getNumSymbols(); i < numSymbols; ++i) {
+    auto expr = results[i].dyn_cast<AffineDimExpr>();
+    if (!expr || expr.getPosition() != i)
+      return false;
+  }
+  return true;
+}
+
 bool AffineMap::isEmpty() const {
   return getNumDims() == 0 && getNumSymbols() == 0 && getNumResults() == 0;
 }

diff  --git a/mlir/test/Dialect/Affine/canonicalize.mlir b/mlir/test/Dialect/Affine/canonicalize.mlir
index 1dac40167db4..ddebab4e5f4b 100644
--- a/mlir/test/Dialect/Affine/canonicalize.mlir
+++ b/mlir/test/Dialect/Affine/canonicalize.mlir
@@ -1219,3 +1219,19 @@ func.func @regression_do_not_perform_invalid_replacements(%arg0: index) {
   "test.foo"(%1) : (index) -> ()
   return
 }
+
+// -----
+// CHECK-LABEL: func @min.oneval(%arg0: index)
+func.func @min.oneval(%arg0: index) -> index {
+  %min = affine.min affine_map<()[s0] -> (s0)> ()[%arg0]
+  // CHECK: return %arg0 : index
+  return %min: index
+}
+
+// -----
+// CHECK-LABEL: func @max.oneval(%arg0: index)
+func.func @max.oneval(%arg0: index) -> index {
+  %max = affine.max affine_map<()[s0] -> (s0)> ()[%arg0]
+  // CHECK: return %arg0 : index
+  return %max: index
+}