diff --git a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp index f384f454bc472..d8facddd63d38 100644 --- a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp +++ b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp @@ -4605,23 +4605,14 @@ struct DropUnitExtentBasis } }; -/// Drop delinearization pattern related to loops in the following way +/// Drop delinearization with a single basis element /// -/// ``` -/// (%iv) = (%c0) to (%ub) step (%c1) { -/// %0 = affine.delinearize_index %iv into (%ub) : index -/// (%0) -/// } -/// ``` -/// -/// can be canonicalized to -/// -/// ``` -/// (%iv) = (%c0) to (%ub) step (%c1) { -/// (%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 { using OpRewritePattern::OpRewritePattern; @@ -4629,50 +4620,7 @@ struct DropDelinearizeOfSingleLoop 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(delinearizeOp.getLinearIndex()); - if (!inductionVar) - return failure(); - - // Check that the parent is a `LoopLikeOpInterface`. - auto loopLikeOp = dyn_cast( - 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(context); + patterns.insert(context); } //===----------------------------------------------------------------------===// diff --git a/mlir/test/Dialect/Affine/canonicalize.mlir b/mlir/test/Dialect/Affine/canonicalize.mlir index d78c3b667589b..6c20b37267672 100644 --- a/mlir/test/Dialect/Affine/canonicalize.mlir +++ b/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, %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, %i : index, // ----- // CHECK-LABEL: func @delinearize_non_loop_like +// CHECK-NOT: affine.delinearize func.func @delinearize_non_loop_like(%arg0: memref, %i : index) -> index { %2 = affine.delinearize_index %i into (1024) : index return %2 : index