[mlir][affine] Fold affine.min with constant zero expressions

[antiagainst]

Summary

Such affine.min ops can be folded into constant zero op, given
that affine ops works on the index type, which by definition
cannot be negative.

Depends On D103243

[ftynse]

I'm not sure that values of index type cannot be negative. Index are signless, so the interpretation of the most significant bit depends on the operation.

[ftynse]

https://github.com/llvm/llvm-project/blob/86627be23312bd227e5afa88c206771a9aaf6589/mlir/test/Conversion/AffineToStandard/lower-affine.mlir#L158 relies on having index with negative value, for example, and so does https://github.com/llvm/llvm-project/blob/main/mlir/test/Conversion/StandardToSPIRV/std-ops-to-spirv.mlir#L525.

The only restriction on non-negative values I know is in the RHS of floordiv/ceildiv/mod of affine expressions, but that is not related to the type.

[antiagainst]

Ah, I think you are right. I just checked the definition for index type; indeed it's just spec'ed as "signless integer whose size is equal to the natural machine word of the target". I was under the wrong impression due to its close relationship with indices/sizes. Thanks for pointing out!! :)

What I'm trying to do is clean up some of the affine.mins that are used as indices/sizes (as shown in the modified test cases), that shouldn't be negative (otherwise I think there is a problem). Do you know what's the proper way to do that?

[bondhugula]

index type can be negative as well - very much!

[ftynse]

I'm not sure we restrict indices to be positive either... If there is an op that clearly requires values to be non-negative and has undefined behavior otherwise, we could follow a C-compiler approach and assume the value passed to this op must not be negative in valid programs. This isn't something we usually do in MLIR though, so I would think carefully before introducing such transformations.

A more robust, but possibly less aggressive approach, is to reason about the sets of values that each result of the affine map used in affine.min can take. I see two cases:

• symbolic-only: there is an affine result expression that being subtracted from other expressions always results in a non-negative constant expression, i.e. is symbolically known to be less than others; for example, in (d0 + 3, d0, d0 - 1, d0), we can try (d0) - (d0 - 1) = 1, (d0 + 3) - (d0 - 1) = 4 and prove that d0 - 1 is the smallest expression always, and replace affine.min with affine.apply; I think we have a similar simplification to remove affine loops that never iterate.
• integer set based: given (partial) knowledge about the sets of values each dimension or symbol can take, we can compute the set of values each affine map result can take (or, combining with the previous case, each difference between expressions can take); if all values in one of the sets are less than or equal to any value in all the other sets, then that expression from which the set is constructed is the minimum and we can simplify again (for the difference case, we want the set of possible differences to be non-negative). For example, if we know that d0 is a loop induction variable going from 1 to N, we can prove that d0 is the minimum of (d0, 2*d0 - 1, 3*d0). This subsumes the symbolic approach, but relies on more complex things internally.

[antiagainst]

Thanks a lot for the explanation @ftynse, really appreciated that. And sorry for the super late reply. I think what you said in the above is actually implemented in mlir. So I'll just close this.

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

static OpFoldResult foldMinMaxOp(T op, ArrayRef<Attribute> operands) {
   }
 
   // Otherwise, completely fold the op into a constant.
   auto resultIt = std::is_same<T, AffineMinOp>::value
                       ? std::min_element(results.begin(), results.end())
                       : std::max_element(results.begin(), results.end());
   if (resultIt == results.end())
     return {};
-  return IntegerAttr::get(IndexType::get(op.getContext()), *resultIt);
+  return Builder(op).getIndexAttr(*resultIt);
 }
 
 /// Remove duplicated expressions in affine min/max ops.
 template <typename T>
 struct DeduplicateAffineMinMaxExpressions : public OpRewritePattern<T> {
   using OpRewritePattern<T>::OpRewritePattern;
 
   LogicalResult matchAndRewrite(T affineOp,

 
 //===----------------------------------------------------------------------===//
 // AffineMinOp
 //===----------------------------------------------------------------------===//
 //
 //   %0 = affine.min (d0) -> (1000, d0 + 512) (%i0)
 //
 
+static OpFoldResult foldMinZeroOp(AffineMinOp minOp) {
+  auto isCstZeroExpr = [](AffineExpr expr) {
+    if (auto cst = expr.dyn_cast<AffineConstantExpr>())
+      return cst.getValue() == 0;
+    return false;
+  };
+
+  if (llvm::any_of(minOp.getAffineMap().getResults(), isCstZeroExpr))
+    return Builder(minOp).getIndexAttr(0);
+  return {};
+}
+
 OpFoldResult AffineMinOp::fold(ArrayRef<Attribute> operands) {
+  if (auto result = foldMinZeroOp(*this))
+    return result;
   return foldMinMaxOp(*this, operands);
 }
 
 void AffineMinOp::getCanonicalizationPatterns(RewritePatternSet &patterns,
                                               MLIRContext *context) {
   patterns.add<DeduplicateAffineMinMaxExpressions<AffineMinOp>,
                MergeAffineMinMaxOp<AffineMinOp>, SimplifyAffineOp<AffineMinOp>>(
       context);

mlir/test/Dialect/Affine/canonicalize.mlir

affine.for %i = 0 to 1024 {
 
     // Map remains the same, but operand changes on composition.
     %copy = affine.apply affine_map<(d0) -> (d0)> (%i)
     %1 = affine.vector_load %A[%copy] : memref<1024xf32>, vector<8xf32>
     "prevent.dce"(%1) : (vector<8xf32>) -> ()
   }
   return
 }
+
+// -----
+
+// CHECK-LABEL: func @fold_affine_min_zero
+func @fold_affine_min_zero(%i0: index, %i1: index, %i2: index) -> index {
+  // CHECK: %[[ZERO:.+]] = constant 0 : index
+  %0 = affine.min affine_map<(d0, d1)[s0] -> (d0, d1 + s0, 0)> (%i0, %i1)[%i2]
+  // CHECK: return %[[ZERO]]
+  return %0: index
+}

mlir/test/Dialect/Linalg/tile-and-fuse-tensors.mlir

%for1 = scf.for %iv1 = %c0 to %d1 step %c8 iter_args(%arg1 = %arg0) -> tensor<?x?x?x?xf32> {
       %insert = subtensor_insert %add into %arg1[%iv0, %iv1, 0, 0] [%d0_size, %d1_size, %d2, %d3] [1, 1, 1, 1]  : tensor<?x?x?x?xf32> into tensor<?x?x?x?xf32>
       scf.yield %insert : tensor<?x?x?x?xf32>
     }
     scf.yield %for1 : tensor<?x?x?x?xf32>
   }
   return %for0 : tensor<?x?x?x?xf32>
 }
 
-// CHECK-DAG: #[[DIM2_UNUSED_LOW_MAP:.+]] = affine_map<()[s0] -> (s0, 0)>
-// CHECK-DAG: #[[DIM2_USED_LOW_MAP:.+]] = affine_map<()[s0, s1, s2] -> (s0, s1 - s2)>
+// CHECK-DAG: #[[DIM2_USED_LOW_MAP:.+]] = affine_map<()[s0, s1] -> (s0, s1)>
 // CHECK-DAG: #[[DIM2_UNUSED_HIGH_MAP:.+]] = affine_map<()[s0, s1, s2, s3] -> (s0, s0 - s1 + s2 + s3)>
-// CHECK-DAG: #[[DIM2_OFFSET_MAP:.+]] = affine_map<()[s0] -> (-s0)>
-// CHECK-DAG: #[[DIM3_SIZE_MAP:.+]] = affine_map<()[s0, s1, s2] -> (s0 - s1 - s2)>
-// CHECK-DAG: #[[DIM3_UNUSED_HIGH_MAP:.+]] = affine_map<()[s0, s1] -> (0, -s0 + s1)>
-// CHECK-DAG: #[[DIM3_USED_HIGH_MAP:.+]] = affine_map<()[s0, s1] -> (s0, -s1)>
-// CHECK-DAG: #[[DIM0_UNUSED_LOW_MAP:.+]] = affine_map<(d0)[s0] -> (4, -d0 + s0)>
-// CHECK-DAG: #[[DIM0_USED_LOW_MAP:.+]] = affine_map<(d0)[s0] -> (0, 4, -d0 + s0)>
-// CHECK-DAG: #[[DIM0_UNUSED_HIGH_MAP:.+]] = affine_map<(d0, d1)[s0] -> (0, -d0 - d1 + s0)>
-// CHECK-DAG: #[[DIM0_USED_HIGH_MAP:.+]] = affine_map<(d0, d1)[s0] -> (-d1, 4, -d0 + s0)>
-// CHECK-DAG: #[[DIM0_SIZE_MAP:.+]] = affine_map<(d0, d1, d2) -> (d0 - d1 - d2)>
+// CHECK-DAG: #[[DIM2_USED_HIGH_MAP:.+]] = affine_map<()[s0, s1, s2] -> (s0, s1 - s2)>
+// CHECK-DAG: #[[DIM2_SIZE_MAP:.+]] = affine_map<()[s0, s1, s2] -> (s0 - s1 - s2)>
+// CHECK-DAG: #[[DIM0_SIZE_MAP:.+]] = affine_map<(d0)[s0] -> (4, -d0 + s0)>
 // CHECK-DAG: #[[DIM1_TILE_SIZE_MAP:.+]] = affine_map<(d0)[s0] -> (8, -d0 + s0)>
 // CHECK-DAG: #[[DIM1_UNUSED_LOW_MAP:.+]] = affine_map<(d0)[s0] -> (s0, d0)>
 // CHECK-DAG: #[[DIM1_USED_MAP:.+]] = affine_map<(d0, d1)[s0, s1] -> (-d1 + s1, 8, -d0 + s0)>
 // CHECK-DAG: #[[DIM1_UNUSED_HIGH_MAP:.+]] = affine_map<(d0, d1)[s0, s1, s2] -> (s0, -d0 - d1 + s0 + s1 + s2)>
 // CHECK-DAG: #[[DIM1_OFFSET_MAP:.+]] = affine_map<(d0, d1) -> (d0 - d1)>
+// CHECK-DAG: #[[DIM1_SIZE_MAP:.+]] = affine_map<(d0, d1, d2) -> (d0 - d1 - d2)>
 
 //      CHECK: func @pad_generic_dynamic
 // CHECK-SAME:   %[[SMALL_INPUT:.+]]: tensor<?x?x?x?xf32>, %[[LARGE_INPUT:.+]]: tensor<?x?x?x?xf32>
 // CHECK-SAME:   %[[INPUT_LOW_PAD:.+]]: index, %[[INPUT_HIGH_PAD:.+]]: index
 
 //  CHECK-DAG:   %[[ZERO:.+]] = constant 0.000000e+00 : f32
 //  CHECK-DAG:   %[[C0:.+]] = constant 0 : index
 //  CHECK-DAG:   %[[C1:.+]] = constant 1 : index
 //  CHECK-DAG:   %[[C2:.+]] = constant 2 : index
 //  CHECK-DAG:   %[[C3:.+]] = constant 3 : index
 //      CHECK:   %[[LARGE_DIM0:.+]] = memref.dim %[[LARGE_INPUT]], %[[C0]]
 //      CHECK:   %[[LARGE_DIM1:.+]] = memref.dim %[[LARGE_INPUT]], %[[C1]]
 //      CHECK:   %[[LARGE_DIM2:.+]] = memref.dim %[[LARGE_INPUT]], %[[C2]]
 //      CHECK:   %[[LARGE_DIM3:.+]] = memref.dim %[[LARGE_INPUT]], %[[C3]]
-//      CHECK:   %[[SMALL_DIM0:.+]] = memref.dim %[[SMALL_INPUT]], %[[C0]]
 //      CHECK:   %[[SMALL_DIM1:.+]] = memref.dim %[[SMALL_INPUT]], %[[C1]]
 //      CHECK:   %[[SMALL_DIM2:.+]] = memref.dim %[[SMALL_INPUT]], %[[C2]]
-//      CHECK:   %[[DIM2_UNUSED_LOW_PAD:.+]] = affine.min #[[DIM2_UNUSED_LOW_MAP]]()[%[[INPUT_LOW_PAD]]]
-//      CHECK:   %[[DIM2_USED_LOW_PAD:.+]] = affine.min #[[DIM2_USED_LOW_MAP]]()[%[[LARGE_DIM2]], %[[INPUT_LOW_PAD]], %[[DIM2_UNUSED_LOW_PAD]]]
+//      CHECK:   %[[DIM2_USED_LOW_PAD:.+]] = affine.min #[[DIM2_USED_LOW_MAP]]()[%[[LARGE_DIM2]], %[[INPUT_LOW_PAD]]]
 //      CHECK:   %[[DIM2_UNUSED_HIGH_PAD:.+]] = affine.min #[[DIM2_UNUSED_HIGH_MAP]]()[%[[INPUT_HIGH_PAD]], %[[LARGE_DIM2]], %[[SMALL_DIM2]], %[[INPUT_LOW_PAD]]]
-//      CHECK:   %[[DIM2_USED_HIGH_PAD:.+]] = affine.min #[[DIM2_USED_LOW_MAP]]()[%[[LARGE_DIM2]], %[[INPUT_HIGH_PAD]], %[[DIM2_UNUSED_HIGH_PAD]]]
-//      CHECK:   %[[DIM2_SRC_OFFSET:.+]] = affine.apply #[[DIM2_OFFSET_MAP]]()[%[[DIM2_UNUSED_LOW_PAD]]]
-//      CHECK:   %[[DIM2_SRC_SIZE:.+]] = affine.apply #[[DIM3_SIZE_MAP]]()[%[[LARGE_DIM2]], %[[DIM2_USED_LOW_PAD]], %[[DIM2_USED_HIGH_PAD]]]
-//      CHECK:   %[[SMALL_DIM3:.+]] = memref.dim %[[SMALL_INPUT]], %[[C3]]
-//      CHECK:   %[[DIM3_USED_LOW_PAD:.+]] = affine.min #[[DIM2_UNUSED_LOW_MAP]]()[%[[LARGE_DIM3]]]
-//      CHECK:   %[[DIM3_UNUSED_HIGH_PAD:.+]] = affine.min #[[DIM3_UNUSED_HIGH_MAP]]()[%[[LARGE_DIM3]], %[[SMALL_DIM3]]]
-//      CHECK:   %[[DIM3_USED_HIGH_PAD:.+]] = affine.min #[[DIM3_USED_HIGH_MAP]]()[%[[LARGE_DIM3]], %[[DIM3_UNUSED_HIGH_PAD]]]
-//      CHECK:   %[[DIM3_SRC_SIZE:.+]] = affine.apply #[[DIM3_SIZE_MAP]]()[%[[LARGE_DIM3]], %[[DIM3_USED_LOW_PAD]], %[[DIM3_USED_HIGH_PAD]]]
+//      CHECK:   %[[DIM2_USED_HIGH_PAD:.+]] = affine.min #[[DIM2_USED_HIGH_MAP]]()[%[[LARGE_DIM2]], %[[INPUT_HIGH_PAD]], %[[DIM2_UNUSED_HIGH_PAD]]]
+//      CHECK:   %[[DIM2_SRC_SIZE:.+]] = affine.apply #[[DIM2_SIZE_MAP]]()[%[[LARGE_DIM2]], %[[DIM2_USED_LOW_PAD]], %[[DIM2_USED_HIGH_PAD]]]
 //      CHECK:   scf.for %[[IV0:[a-z0-9]+]]
-//      CHECK:     %[[DIM0_UNUSED_LOW_PAD:.+]] = affine.min #[[DIM0_UNUSED_LOW_MAP]](%[[IV0]])[%[[LARGE_DIM0]]]
-//      CHECK:     %[[DIM0_USED_LOW_PAD:.+]] = affine.min #[[DIM0_USED_LOW_MAP]](%[[IV0]])[%[[LARGE_DIM0]]]
-//      CHECK:     %[[DIM0_UNUSED_HIGH_PAD:.+]] = affine.min #[[DIM0_UNUSED_HIGH_MAP]](%[[IV0]], %[[DIM0_UNUSED_LOW_PAD]])[%[[SMALL_DIM0]]]
-//      CHECK:     %[[DIM0_USED_HIGH_PAD:.+]] = affine.min #[[DIM0_USED_HIGH_MAP]](%[[IV0]], %[[DIM0_UNUSED_HIGH_PAD]])[%[[LARGE_DIM0]]]
-//      CHECK:     %[[DIM0_SRC_SIZE:.+]] = affine.apply #[[DIM0_SIZE_MAP]](%[[DIM0_UNUSED_LOW_PAD]], %[[DIM0_USED_LOW_PAD]], %[[DIM0_USED_HIGH_PAD]])
+//      CHECK:     %[[DIM0_SRC_SIZE:.+]] = affine.min #[[DIM0_SIZE_MAP]](%[[IV0]])[%[[LARGE_DIM0]]]
 //      CHECK:     scf.for %[[IV1:[a-z0-9]+]]
 //      CHECK:       %[[DIM1_TILE_SIZE:.+]] = affine.min #[[DIM1_TILE_SIZE_MAP]](%[[IV1]])[%[[LARGE_DIM1]]]
 //      CHECK:       %[[DIM1_UNUSED_LOW_PAD:.+]] = affine.min #[[DIM1_UNUSED_LOW_MAP]](%[[IV1]])[%[[INPUT_LOW_PAD]]]
 //      CHECK:       %[[DIM1_USED_LOW_PAD:.+]] = affine.min #[[DIM1_USED_MAP]](%[[IV1]], %[[DIM1_UNUSED_LOW_PAD]])[%[[LARGE_DIM1]], %[[INPUT_LOW_PAD]]]
 //      CHECK:       %[[DIM1_UNUSED_HIGH_PAD:.+]] = affine.min #[[DIM1_UNUSED_HIGH_MAP]](%[[IV1]], %[[DIM1_TILE_SIZE]])[%[[INPUT_HIGH_PAD]], %[[SMALL_DIM1]], %[[INPUT_LOW_PAD]]]
 //      CHECK:       %[[DIM1_USED_HIGH_PAD:.+]] = affine.min #[[DIM1_USED_MAP]](%[[IV1]], %[[DIM1_UNUSED_HIGH_PAD]])[%[[LARGE_DIM1]], %[[INPUT_HIGH_PAD]]]
 //      CHECK:       %[[DIM1_SRC_OFFSET:.+]] = affine.apply #[[DIM1_OFFSET_MAP]](%[[IV1]], %[[DIM1_UNUSED_LOW_PAD]])
-//      CHECK:       %[[DIM1_SRC_SIZE:.+]] = affine.apply #[[DIM0_SIZE_MAP]](%[[DIM1_TILE_SIZE]], %[[DIM1_USED_LOW_PAD]], %[[DIM1_USED_HIGH_PAD]])
+//      CHECK:       %[[DIM1_SRC_SIZE:.+]] = affine.apply #[[DIM1_SIZE_MAP]](%[[DIM1_TILE_SIZE]], %[[DIM1_USED_LOW_PAD]], %[[DIM1_USED_HIGH_PAD]])
 //      CHECK:       %[[SRC_SUBTENSOR:.+]] = subtensor %[[SMALL_INPUT]]
-// CHECK-SAME:         [%[[IV0]], %[[DIM1_SRC_OFFSET]], %[[DIM2_SRC_OFFSET]], 0]
-// CHECK-SAME:         [%[[DIM0_SRC_SIZE]], %[[DIM1_SRC_SIZE]], %[[DIM2_SRC_SIZE]], %[[DIM3_SRC_SIZE]]]
+// CHECK-SAME:         [%[[IV0]], %[[DIM1_SRC_OFFSET]], 0, 0]
+// CHECK-SAME:         [%[[DIM0_SRC_SIZE]], %[[DIM1_SRC_SIZE]], %[[DIM2_SRC_SIZE]], %[[LARGE_DIM3]]]
 //      CHECK:       %[[PAD:.+]] = linalg.pad_tensor %[[SRC_SUBTENSOR]]
-// CHECK-SAME:         low[%[[DIM0_USED_LOW_PAD]], %[[DIM1_USED_LOW_PAD]], %[[DIM2_USED_LOW_PAD]], %[[DIM3_USED_LOW_PAD]]]
-// CHECK-SAME:         high[%[[DIM0_USED_HIGH_PAD]], %[[DIM1_USED_HIGH_PAD]], %[[DIM2_USED_HIGH_PAD]], %[[DIM3_USED_HIGH_PAD]]]
+// CHECK-SAME:         low[%[[C0]], %[[DIM1_USED_LOW_PAD]], %[[DIM2_USED_LOW_PAD]], %[[C0]]]
+// CHECK-SAME:         high[%[[C0]], %[[DIM1_USED_HIGH_PAD]], %[[DIM2_USED_HIGH_PAD]], %[[C0]]]
 //      CHECK:         linalg.yield %[[ZERO]]
 //      CHECK:       linalg.generic
 // CHECK-SAME:         ins(%[[PAD]]