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

[Lewuathe]

Summary

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.

[Lewuathe]

Add newline.

[ftynse]

mlir/test/Dialect/Affine/canonicalize.mlir

1234–1236

Nit: s/%min/%max.

mlir/include/mlir/IR/AffineMap.h

public:
   bool operator==(AffineMap other) const { return other.map == map; }
   bool operator!=(AffineMap other) const { return !(other.map == map); }
 
   /// Returns true if this affine map is an identity affine map.
   /// An identity affine map corresponds to an identity affine function on the
   /// 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;
 
   /// Returns true if this affine map is a minor identity up to broadcasted
   /// dimensions which are indicated by value 0 in the result. If
   /// `broadcastedDims` is not null, it will be populated with the indices of
   /// the broadcasted dimensions in the result array.

mlir/lib/Dialect/Affine/IR/AffineOps.cpp

static_assert(llvm::is_one_of<T, AffineMinOp, AffineMaxOp>::value,
                 "expected affine min or max op");
 
   // Fold the affine map.
   // TODO: Fold more cases:
   // min(some_affine, some_affine + constant, ...), etc.
   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.
     if (foldedMap == op.getMap())
       return {};
     op->setAttr("map", AffineMapAttr::get(foldedMap));
     return op.getResult();
   }

mlir/lib/IR/AffineMap.cpp

bool AffineMap::isIdentity() const {
   for (unsigned i = 0, numDims = getNumDims(); i < numDims; ++i) {
     auto expr = results[i].dyn_cast<AffineDimExpr>();
     if (!expr || expr.getPosition() != i)
       return false;
   }
   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;
 }
 
 bool AffineMap::isSingleConstant() const {
   return getNumResults() == 1 && getResult(0).isa<AffineConstantExpr>();
 }
 

mlir/test/Dialect/Affine/canonicalize.mlir

func.func @regression_do_not_perform_invalid_replacements(%arg0: index) {
   // CHECK-BOTTOM-UP: %[[apply:.*]] = affine.apply #[[$map]]()[%{{.*}}]
   %0 = affine.apply affine_map<(d0) -> (-d0 + 40961)>(%arg0)
   %1 = affine.apply affine_map<(d0)[s0] -> (d0 * (s0 ceildiv 512))>(%arg0)[%0]
   //           CHECK: "test.foo"(%[[apply]])
   // CHECK-BOTTOM-UP: "test.foo"(%[[apply]])
   "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
+}