[mlir][affine] Generalize canonicalization for one-element delinearize (#114871)

There was an existing canonicalization pattern for delinearize_index
that would remove `affine.delinearize_index %linear into (%basis)`
under some conditions. However, the delinearize_index means that this
rewrite is always permissisible.

Author: krzysz00

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

@@ -4605,74 +4605,22 @@ struct DropUnitExtentBasis
   }
 };
 
-/// Drop delinearization pattern related to loops in the following way
+/// Drop delinearization with a single basis element
 ///
-/// ```
-/// <loop>(%iv) = (%c0) to (%ub) step (%c1) {
-///   %0 = affine.delinearize_index %iv into (%ub) : index
-///   <some_use>(%0)
-/// }
-/// ```
-///
-/// can be canonicalized to
-///
-/// ```
-/// <loop>(%iv) = (%c0) to (%ub) step (%c1) {
-///   <some_use>(%iv)
-/// }
-/// ```
-struct DropDelinearizeOfSingleLoop
+/// By definition, `delinearize_index %linear into (%basis)` is
+/// `%linear floorDiv 1` (since `1` is the product of the basis elememts,
+/// ignoring the 0th one, and since there is no previous division we need
+/// to use the remainder of). Therefore, a single-element `delinearize`
+/// can be replaced by the underlying linear index.
+struct DropDelinearizeOneBasisElement
     : public OpRewritePattern<affine::AffineDelinearizeIndexOp> {
   using OpRewritePattern::OpRewritePattern;
 
   LogicalResult matchAndRewrite(affine::AffineDelinearizeIndexOp delinearizeOp,
                                 PatternRewriter &rewriter) const override {
     if (delinearizeOp.getStaticBasis().size() != 1)
       return failure();
-    auto basis = delinearizeOp.getMixedBasis();
-
-    // Check that the `linear_index` is an induction variable.
-    auto inductionVar = dyn_cast<BlockArgument>(delinearizeOp.getLinearIndex());
-    if (!inductionVar)
-      return failure();
-
-    // Check that the parent is a `LoopLikeOpInterface`.
-    auto loopLikeOp = dyn_cast<LoopLikeOpInterface>(
-        inductionVar.getParentRegion()->getParentOp());
-    if (!loopLikeOp)
-      return failure();
-
-    // Check that loop is unit-rank and that the `linear_index` is the induction
-    // variable.
-    auto inductionVars = loopLikeOp.getLoopInductionVars();
-    if (!inductionVars || inductionVars->size() != 1 ||
-        inductionVars->front() != inductionVar) {
-      return rewriter.notifyMatchFailure(
-          delinearizeOp, "`linear_index` is not loop induction variable");
-    }
-
-    // Check that the upper-bound is the basis.
-    auto upperBounds = loopLikeOp.getLoopUpperBounds();
-    if (!upperBounds || upperBounds->size() != 1 ||
-        upperBounds->front() != basis.front()) {
-      return rewriter.notifyMatchFailure(delinearizeOp,
-                                         "`basis` is not upper bound");
-    }
-
-    // Check that the lower bound is zero.
-    auto lowerBounds = loopLikeOp.getLoopLowerBounds();
-    if (!lowerBounds || lowerBounds->size() != 1 ||
-        !isZeroIndex(lowerBounds->front())) {
-      return rewriter.notifyMatchFailure(delinearizeOp,
-                                         "loop lower bound is not zero");
-    }
-
-    // Check that the step is one.
-    auto steps = loopLikeOp.getLoopSteps();
-    if (!steps || steps->size() != 1 || !isConstantIntValue(steps->front(), 1))
-      return rewriter.notifyMatchFailure(delinearizeOp, "loop step is not one");
-
-    rewriter.replaceOp(delinearizeOp, inductionVar);
+    rewriter.replaceOp(delinearizeOp, delinearizeOp.getLinearIndex());
     return success();
   }
 };
@@ -4681,7 +4629,7 @@ struct DropDelinearizeOfSingleLoop
 
 void affine::AffineDelinearizeIndexOp::getCanonicalizationPatterns(
     RewritePatternSet &patterns, MLIRContext *context) {
-  patterns.insert<DropDelinearizeOfSingleLoop, DropUnitExtentBasis>(context);
+  patterns.insert<DropDelinearizeOneBasisElement, DropUnitExtentBasis>(context);
 }
 
 //===----------------------------------------------------------------------===//

mlir/test/Dialect/Affine/canonicalize.mlir

@@ -1517,6 +1517,7 @@ func.func @drop_single_loop_delinearize(%arg0 : index, %arg1 : index) -> index {
 // -----
 
 // CHECK-LABEL: func @delinearize_non_induction_variable
+// CHECK-NOT: affine.delinearize
 func.func @delinearize_non_induction_variable(%arg0: memref<?xi32>, %i : index, %t0 : index, %t1 : index, %t2 : index) -> index {
   %1 = affine.apply affine_map<(d0)[s0, s1, s2] -> (d0 + s0 + s1 * 64 + s2 * 128)>(%i)[%t0, %t1, %t2]
   %2 = affine.delinearize_index %1 into (1024) : index
@@ -1526,6 +1527,7 @@ func.func @delinearize_non_induction_variable(%arg0: memref<?xi32>, %i : index,
 // -----
 
 // CHECK-LABEL: func @delinearize_non_loop_like
+// CHECK-NOT: affine.delinearize
 func.func @delinearize_non_loop_like(%arg0: memref<?xi32>, %i : index) -> index {
   %2 = affine.delinearize_index %i into (1024) : index
   return %2 : index