diff --git a/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td b/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td index f5ca24389065..e2eab1fb2178 100644 --- a/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td +++ b/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td @@ -1083,6 +1083,9 @@ def AffineDelinearizeIndexOp : Affine_Op<"delinearize_index", [Pure]> { %indices_2 = affine.apply #map2()[%linear_index] ``` + In other words, `%0:3 = affine.delinearize_index %x into (B, C)` produces + `%0 = {%x / (B * C), (%x mod (B * C)) / C, %x mod C}`. + The basis may either contain `N` or `N-1` elements, where `N` is the number of results. If there are N basis elements, the first one will not be used during computations, but may be used during analysis and canonicalization to eliminate terms from @@ -1098,7 +1101,12 @@ def AffineDelinearizeIndexOp : Affine_Op<"delinearize_index", [Pure]> { %0:3 = affine.delinearize_index %linear_index into (244, 244) : index, index ``` - Note that, due to the constraints of affine maps, all the basis elements must + Note that, for symmetry with `getPaddedBasis()`, if `hasOuterBound` is `true` + when one of the `OpFoldResult` builders is called but the first element of the + basis is `nullptr`, that first element is ignored and the builder proceeds as if + there was no outer bound. + + Due to the constraints of affine maps, all the basis elements must be strictly positive. A dynamic basis element being 0 or negative causes undefined behavior. }]; @@ -1136,6 +1144,11 @@ def AffineDelinearizeIndexOp : Affine_Op<"delinearize_index", [Pure]> { /// Return a vector that contains the basis of the operation, removing /// the outer bound if one is present. SmallVector getEffectiveBasis(); + + /// Return the vector with one basis element per result of the operation. If + /// there is no outer bound specified, the leading entry of this result will be + /// nullptr. + SmallVector getPaddedBasis(); }]; let hasVerifier = 1; @@ -1160,6 +1173,9 @@ def AffineLinearizeIndexOp : Affine_Op<"linearize_index", sum(i = 0 to N-1) %idx_i * product(j = i + 1 to N-1) B_j ``` + In other words, `%0 = affine.linearize_index [%z, %y, %x] by (Z, Y, X)` + gives `%0 = %x + %y * X + %z * X * Y`, or `%0 = %x + X * (%y + Y * (%z))`. + The basis may either have `N` or `N-1` elements, where `N` is the number of inputs to linearize_index. If `N` inputs are provided, the first one is not used in computation, but may be used during analysis or canonicalization as a bound @@ -1168,6 +1184,10 @@ def AffineLinearizeIndexOp : Affine_Op<"linearize_index", If all `N` basis elements are provided, the linearize_index operation is said to "have an outer bound". + As a convenience, and for symmetry with `getPaddedBasis()`, ifg the first + element of a set of `OpFoldResult`s passed to the builders of this operation is + `nullptr`, that element is ignored. + If the `disjoint` property is present, this is an optimization hint that, for all `i`, `0 <= %idx_i < B_i` - that is, no index affects any other index, except that `%idx_0` may be negative to make the index as a whole negative. @@ -1224,6 +1244,11 @@ def AffineLinearizeIndexOp : Affine_Op<"linearize_index", /// Return a vector that contains the basis of the operation, removing /// the outer bound if one is present. SmallVector getEffectiveBasis(); + + /// Return the vector with one basis element per index operand of the operation. + /// If there is no outer bound specified, the leading entry of this basis will be + /// nullptr. + SmallVector getPaddedBasis(); }]; let hasVerifier = 1; diff --git a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp index dceebbfec586..b45829bcf6d2 100644 --- a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp +++ b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp @@ -4520,6 +4520,10 @@ void AffineDelinearizeIndexOp::build(OpBuilder &odsBuilder, OperationState &odsState, Value linearIndex, ValueRange basis, bool hasOuterBound) { + if (hasOuterBound && !basis.empty() && basis.front() == nullptr) { + hasOuterBound = false; + basis = basis.drop_front(); + } SmallVector dynamicBasis; SmallVector staticBasis; dispatchIndexOpFoldResults(getAsOpFoldResult(basis), dynamicBasis, @@ -4533,6 +4537,10 @@ void AffineDelinearizeIndexOp::build(OpBuilder &odsBuilder, Value linearIndex, ArrayRef basis, bool hasOuterBound) { + if (hasOuterBound && !basis.empty() && basis.front() == OpFoldResult()) { + hasOuterBound = false; + basis = basis.drop_front(); + } SmallVector dynamicBasis; SmallVector staticBasis; dispatchIndexOpFoldResults(basis, dynamicBasis, staticBasis); @@ -4654,6 +4662,13 @@ SmallVector AffineDelinearizeIndexOp::getEffectiveBasis() { return getMixedValues(getStaticBasis(), getDynamicBasis(), builder); } +SmallVector AffineDelinearizeIndexOp::getPaddedBasis() { + SmallVector ret = getMixedBasis(); + if (!hasOuterBound()) + ret.insert(ret.begin(), OpFoldResult()); + return ret; +} + namespace { // Drops delinearization indices that correspond to unit-extent basis @@ -4672,25 +4687,27 @@ struct DropUnitExtentBasis return zero.value(); }; - bool hasOuterBound = delinearizeOp.hasOuterBound(); // Replace all indices corresponding to unit-extent basis with 0. // Remaining basis can be used to get a new `affine.delinearize_index` op. SmallVector newBasis; - for (auto [index, basis] : llvm::enumerate(delinearizeOp.getMixedBasis())) { - std::optional basisVal = getConstantIntValue(basis); + for (auto [index, basis] : + llvm::enumerate(delinearizeOp.getPaddedBasis())) { + std::optional basisVal = + basis ? getConstantIntValue(basis) : std::nullopt; if (basisVal && *basisVal == 1) - replacements[index + (hasOuterBound ? 0 : 1)] = getZero(); + replacements[index] = getZero(); else newBasis.push_back(basis); } - if (newBasis.size() == delinearizeOp.getStaticBasis().size()) + if (newBasis.size() == delinearizeOp.getNumResults()) return rewriter.notifyMatchFailure(delinearizeOp, "no unit basis elements"); - if (!newBasis.empty() || !hasOuterBound) { + if (!newBasis.empty()) { + // Will drop the leading nullptr from `basis` if there was no outer bound. auto newDelinearizeOp = rewriter.create( - loc, delinearizeOp.getLinearIndex(), newBasis, hasOuterBound); + loc, delinearizeOp.getLinearIndex(), newBasis); int newIndex = 0; // Map back the new delinearized indices to the values they replace. for (auto &replacement : replacements) { @@ -4871,6 +4888,8 @@ void AffineLinearizeIndexOp::build(OpBuilder &odsBuilder, OperationState &odsState, ValueRange multiIndex, ValueRange basis, bool disjoint) { + if (!basis.empty() && basis.front() == Value()) + basis = basis.drop_front(); SmallVector dynamicBasis; SmallVector staticBasis; dispatchIndexOpFoldResults(getAsOpFoldResult(basis), dynamicBasis, @@ -4883,6 +4902,8 @@ void AffineLinearizeIndexOp::build(OpBuilder &odsBuilder, ValueRange multiIndex, ArrayRef basis, bool disjoint) { + if (!basis.empty() && basis.front() == OpFoldResult()) + basis = basis.drop_front(); SmallVector dynamicBasis; SmallVector staticBasis; dispatchIndexOpFoldResults(basis, dynamicBasis, staticBasis); @@ -4965,7 +4986,14 @@ SmallVector AffineLinearizeIndexOp::getEffectiveBasis() { builder); } - return ::mlir::getMixedValues(getStaticBasis(), getDynamicBasis(), builder); + return getMixedValues(getStaticBasis(), getDynamicBasis(), builder); +} + +SmallVector AffineLinearizeIndexOp::getPaddedBasis() { + SmallVector ret = getMixedBasis(); + if (!hasOuterBound()) + ret.insert(ret.begin(), OpFoldResult()); + return ret; } namespace { @@ -5027,38 +5055,228 @@ struct DropLinearizeUnitComponentsIfDisjointOrZero final } }; -/// Cancel out linearize_index(delinearize_index(x, B), B). +/// Return the product of `terms`, creating an `affine.apply` if any of them are +/// non-constant values. If any of `terms` is `nullptr`, return `nullptr`. +static OpFoldResult computeProduct(Location loc, OpBuilder &builder, + ArrayRef terms) { + int64_t nDynamic = 0; + SmallVector dynamicPart; + AffineExpr result = builder.getAffineConstantExpr(1); + for (OpFoldResult term : terms) { + if (!term) + return term; + std::optional maybeConst = getConstantIntValue(term); + if (maybeConst) { + result = result * builder.getAffineConstantExpr(*maybeConst); + } else { + dynamicPart.push_back(term.get()); + result = result * builder.getAffineSymbolExpr(nDynamic++); + } + } + if (auto constant = dyn_cast(result)) + return getAsIndexOpFoldResult(builder.getContext(), constant.getValue()); + return builder.create(loc, result, dynamicPart).getResult(); +} + +/// If conseceutive outputs of a delinearize_index are linearized with the same +/// bounds, canonicalize away the redundant arithmetic. /// -/// That is, rewrite +/// That is, if we have /// ``` -/// %0:N = affine.delinearize_index %x by (%b1, %b2, ... %bN) -/// %y = affine.linearize_index [%0#0, %0#1, ... %0#(N-1)] by (%b1, %b2, ... -/// %bN) +/// %s:N = affine.delinearize_index %x into (...a, B1, B2, ... BK, ...b) +/// %t = affine.linearize_index [...c, %s#I, %s#(I + 1), ... %s#(I+K-1), ...d] +/// by (...e, B1, B2, ..., BK, ...f) /// ``` -/// to replacing `%y` with `%x`. -struct CancelLinearizeOfDelinearizeExact final +/// +/// We can rewrite this to +/// ``` +/// B = B1 * B2 ... BK +/// %sMerged:(N-K+1) affine.delinearize_index %x into (...a, B, ...b) +/// %t = affine.linearize_index [...c, %s#I, ...d] by (...e, B, ...f) +/// ``` +/// where we replace all results of %s unaffected by the change with results +/// from %sMerged. +/// +/// As a special case, if all results of the delinearize are merged in this way +/// we can replace those usages with %x, thus cancelling the delinearization +/// entirely, as in +/// ``` +/// %s:3 = affine.delinearize_index %x into (2, 4, 8) +/// %t = affine.linearize_index [%s#0, %s#1, %s#2, %c0] by (2, 4, 8, 16) +/// ``` +/// becoming `%t = affine.linearize_index [%x, %c0] by (64, 16)` +struct CancelLinearizeOfDelinearizePortion final : OpRewritePattern { using OpRewritePattern::OpRewritePattern; +private: + // Struct representing a case where the cancellation pattern + // applies. A `Match` means that `length` inputs to the linearize operation + // starting at `linStart` can be cancelled with `length` outputs of + // `delinearize`, starting from `delinStart`. + struct Match { + AffineDelinearizeIndexOp delinearize; + unsigned linStart = 0; + unsigned delinStart = 0; + unsigned length = 0; + }; + +public: LogicalResult matchAndRewrite(affine::AffineLinearizeIndexOp linearizeOp, PatternRewriter &rewriter) const override { - auto delinearizeOp = linearizeOp.getMultiIndex() - .front() - .getDefiningOp(); - if (!delinearizeOp) - return rewriter.notifyMatchFailure( - linearizeOp, "last entry doesn't come from a delinearize"); + SmallVector matches; - if (linearizeOp.getEffectiveBasis() != delinearizeOp.getEffectiveBasis()) - return rewriter.notifyMatchFailure( - linearizeOp, "basis of linearize and delinearize don't match exactly " - "(excluding outer bounds)"); + const SmallVector linBasis = linearizeOp.getPaddedBasis(); + ArrayRef linBasisRef = linBasis; - if (delinearizeOp.getResults() != linearizeOp.getMultiIndex()) - return rewriter.notifyMatchFailure( - linearizeOp, "not all indices come from delinearize"); + ValueRange multiIndex = linearizeOp.getMultiIndex(); + unsigned numLinArgs = multiIndex.size(); + unsigned linArgIdx = 0; + // We only want to replace one run from the same delinearize op per + // pattern invocation lest we run into invalidation issues. + llvm::SmallPtrSet alreadyMatchedDelinearize; + while (linArgIdx < numLinArgs) { + auto asResult = dyn_cast(multiIndex[linArgIdx]); + if (!asResult) { + linArgIdx++; + continue; + } + + auto delinearizeOp = + dyn_cast(asResult.getOwner()); + if (!delinearizeOp) { + linArgIdx++; + continue; + } + + /// Result 0 of the delinearize and argument 0 of the linearize can + /// leave their maximum value unspecified. However, even if this happens + /// we can still sometimes start the match process. Specifically, if + /// - The argument we're matching is result 0 and argument 0 (so the + /// bounds don't matter). For example, + /// + /// %0:2 = affine.delinearize_index %x into (8) : index, index + /// %1 = affine.linearize_index [%s#0, %s#1, ...] (8, ...) + /// allows cancellation + /// - The delinearization doesn't specify a bound, but the linearization + /// is `disjoint`, which asserts that the bound on the linearization is + /// correct. + unsigned delinArgIdx = asResult.getResultNumber(); + SmallVector delinBasis = delinearizeOp.getPaddedBasis(); + OpFoldResult firstDelinBound = delinBasis[delinArgIdx]; + OpFoldResult firstLinBound = linBasis[linArgIdx]; + bool boundsMatch = firstDelinBound == firstLinBound; + bool bothAtFront = linArgIdx == 0 && delinArgIdx == 0; + bool knownByDisjoint = + linearizeOp.getDisjoint() && delinArgIdx == 0 && !firstDelinBound; + if (!boundsMatch && !bothAtFront && !knownByDisjoint) { + linArgIdx++; + continue; + } + + unsigned j = 1; + unsigned numDelinOuts = delinearizeOp.getNumResults(); + for (; j + linArgIdx < numLinArgs && j + delinArgIdx < numDelinOuts; + ++j) { + if (multiIndex[linArgIdx + j] != + delinearizeOp.getResult(delinArgIdx + j)) + break; + if (linBasis[linArgIdx + j] != delinBasis[delinArgIdx + j]) + break; + } + // If there're multiple matches against the same delinearize_index, + // only rewrite the first one we find to prevent invalidations. The next + // ones will be taken care of by subsequent pattern invocations. + if (j <= 1 || !alreadyMatchedDelinearize.insert(delinearizeOp).second) { + linArgIdx++; + continue; + } + matches.push_back(Match{delinearizeOp, linArgIdx, delinArgIdx, j}); + linArgIdx += j; + } + + if (matches.empty()) + return rewriter.notifyMatchFailure( + linearizeOp, "no run of delinearize outputs to deal with"); + + // Record all the delinearize replacements so we can do them after creating + // the new linearization operation, since the new operation might use + // outputs of something we're replacing. + SmallVector> delinearizeReplacements; + + SmallVector newIndex; + newIndex.reserve(numLinArgs); + SmallVector newBasis; + newBasis.reserve(numLinArgs); + unsigned prevMatchEnd = 0; + for (Match m : matches) { + unsigned gap = m.linStart - prevMatchEnd; + llvm::append_range(newIndex, multiIndex.slice(prevMatchEnd, gap)); + llvm::append_range(newBasis, linBasisRef.slice(prevMatchEnd, gap)); + // Update here so we don't forget this during early continues + prevMatchEnd = m.linStart + m.length; + + PatternRewriter::InsertionGuard g(rewriter); + rewriter.setInsertionPoint(m.delinearize); + + ArrayRef basisToMerge = + linBasisRef.slice(m.linStart, m.length); + // We use the slice from the linearize's basis above because of the + // "bounds inferred from `disjoint`" case above. + OpFoldResult newSize = + computeProduct(linearizeOp.getLoc(), rewriter, basisToMerge); + + // Trivial case where we can just skip past the delinearize all together + if (m.length == m.delinearize.getNumResults()) { + newIndex.push_back(m.delinearize.getLinearIndex()); + newBasis.push_back(newSize); + // Pad out set of replacements so we don't do anything with this one. + delinearizeReplacements.push_back(SmallVector()); + continue; + } + + SmallVector newDelinResults; + SmallVector newDelinBasis = m.delinearize.getPaddedBasis(); + newDelinBasis.erase(newDelinBasis.begin() + m.delinStart, + newDelinBasis.begin() + m.delinStart + m.length); + newDelinBasis.insert(newDelinBasis.begin() + m.delinStart, newSize); + auto newDelinearize = rewriter.create( + m.delinearize.getLoc(), m.delinearize.getLinearIndex(), + newDelinBasis); + + // Since there may be other uses of the indices we just merged together, + // create a residual affine.delinearize_index that delinearizes the + // merged output into its component parts. + Value combinedElem = newDelinearize.getResult(m.delinStart); + auto residualDelinearize = rewriter.create( + m.delinearize.getLoc(), combinedElem, basisToMerge); + + // Swap all the uses of the unaffected delinearize outputs to the new + // delinearization so that the old code can be removed if this + // linearize_index is the only user of the merged results. + llvm::append_range(newDelinResults, + newDelinearize.getResults().take_front(m.delinStart)); + llvm::append_range(newDelinResults, residualDelinearize.getResults()); + llvm::append_range( + newDelinResults, + newDelinearize.getResults().drop_front(m.delinStart + 1)); + + delinearizeReplacements.push_back(newDelinResults); + newIndex.push_back(combinedElem); + newBasis.push_back(newSize); + } + llvm::append_range(newIndex, multiIndex.drop_front(prevMatchEnd)); + llvm::append_range(newBasis, linBasisRef.drop_front(prevMatchEnd)); + rewriter.replaceOpWithNewOp( + linearizeOp, newIndex, newBasis, linearizeOp.getDisjoint()); + + for (auto [m, newResults] : + llvm::zip_equal(matches, delinearizeReplacements)) { + if (newResults.empty()) + continue; + rewriter.replaceOp(m.delinearize, newResults); + } - rewriter.replaceOp(linearizeOp, delinearizeOp.getLinearIndex()); return success(); } }; @@ -5096,7 +5314,7 @@ struct DropLinearizeLeadingZero final void affine::AffineLinearizeIndexOp::getCanonicalizationPatterns( RewritePatternSet &patterns, MLIRContext *context) { - patterns.add(context); } diff --git a/mlir/test/Dialect/Affine/canonicalize.mlir b/mlir/test/Dialect/Affine/canonicalize.mlir index 717004eb50c0..a9ac13ad7162 100644 --- a/mlir/test/Dialect/Affine/canonicalize.mlir +++ b/mlir/test/Dialect/Affine/canonicalize.mlir @@ -1917,12 +1917,12 @@ func.func @linearize_one_element_basis(%arg0: index, %arg1: index) -> index { // ----- -// CHECK-LABEL: func @cancel_linearize_denearize_exact( +// CHECK-LABEL: func @cancel_linearize_delinearize_exact( // CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: index, // CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]: index, // CHECK-SAME: %[[ARG2:[a-zA-Z0-9]+]]: index) // CHECK: return %[[ARG0]] -func.func @cancel_linearize_denearize_exact(%arg0: index, %arg1: index, %arg2: index) -> index { +func.func @cancel_linearize_delinearize_exact(%arg0: index, %arg1: index, %arg2: index) -> index { %0:3 = affine.delinearize_index %arg0 into (%arg1, 4, %arg2) : index, index, index %1 = affine.linearize_index [%0#0, %0#1, %0#2] by (%arg1, 4, %arg2) : index return %1 : index @@ -1930,12 +1930,12 @@ func.func @cancel_linearize_denearize_exact(%arg0: index, %arg1: index, %arg2: i // ----- -// CHECK-LABEL: func @cancel_linearize_denearize_linearize_extra_bound( +// CHECK-LABEL: func @cancel_linearize_delinearize_linearize_extra_bound( // CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: index, // CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]: index, // CHECK-SAME: %[[ARG2:[a-zA-Z0-9]+]]: index) // CHECK: return %[[ARG0]] -func.func @cancel_linearize_denearize_linearize_extra_bound(%arg0: index, %arg1: index, %arg2: index) -> index { +func.func @cancel_linearize_delinearize_linearize_extra_bound(%arg0: index, %arg1: index, %arg2: index) -> index { %0:3 = affine.delinearize_index %arg0 into (4, %arg2) : index, index, index %1 = affine.linearize_index [%0#0, %0#1, %0#2] by (%arg1, 4, %arg2) : index return %1 : index @@ -1943,12 +1943,12 @@ func.func @cancel_linearize_denearize_linearize_extra_bound(%arg0: index, %arg1: // ----- -// CHECK-LABEL: func @cancel_linearize_denearize_delinearize_extra_bound( +// CHECK-LABEL: func @cancel_linearize_delinearize_delinearize_extra_bound( // CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: index, // CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]: index, // CHECK-SAME: %[[ARG2:[a-zA-Z0-9]+]]: index) // CHECK: return %[[ARG0]] -func.func @cancel_linearize_denearize_delinearize_extra_bound(%arg0: index, %arg1: index, %arg2: index) -> index { +func.func @cancel_linearize_delinearize_delinearize_extra_bound(%arg0: index, %arg1: index, %arg2: index) -> index { %0:3 = affine.delinearize_index %arg0 into (%arg1, 4, %arg2) : index, index, index %1 = affine.linearize_index [%0#0, %0#1, %0#2] by (4, %arg2) : index return %1 : index @@ -1956,31 +1956,252 @@ func.func @cancel_linearize_denearize_delinearize_extra_bound(%arg0: index, %arg // ----- -// Don't cancel because the values from the delinearize aren't used in order -// CHECK-LABEL: func @no_cancel_linearize_denearize_permuted( +// CHECK-LABEL: func @cancel_linearize_delinearize_head( // CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: index, -// CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]: index, -// CHECK-SAME: %[[ARG2:[a-zA-Z0-9]+]]: index) -// CHECK: %[[DELIN:.+]]:3 = affine.delinearize_index %[[ARG0]] into (%[[ARG1]], 4, %[[ARG2]]) -// CHECK: %[[LIN:.+]] = affine.linearize_index [%[[DELIN]]#0, %[[DELIN]]#2, %[[DELIN]]#1] by (%[[ARG1]], 4, %[[ARG2]]) +// CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]: index) +// CHECK: %[[DELIN:.+]]:2 = affine.delinearize_index %[[ARG0]] into (12, 8) +// CHECK: %[[LIN:.+]] = affine.linearize_index [%[[DELIN]]#0, %[[ARG1]]] by (12, 16) // CHECK: return %[[LIN]] -func.func @no_cancel_linearize_denearize_permuted(%arg0: index, %arg1: index, %arg2: index) -> index { - %0:3 = affine.delinearize_index %arg0 into (%arg1, 4, %arg2) : index, index, index - %1 = affine.linearize_index [%0#0, %0#2, %0#1] by (%arg1, 4, %arg2) : index +func.func @cancel_linearize_delinearize_head(%arg0: index, %arg1: index) -> index { + %0:3 = affine.delinearize_index %arg0 into (3, 4, 8) : index, index, index + %1 = affine.linearize_index [%0#0, %0#1, %arg1] by (3, 4, 16) : index return %1 : index } // ----- +// CHECK-LABEL: func @cancel_linearize_delinearize_head_delinearize_unbounded( +// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: index, +// CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]: index) +// CHECK: %[[DELIN:.+]]:2 = affine.delinearize_index %[[ARG0]] into (12, 8) +// CHECK: %[[LIN:.+]] = affine.linearize_index [%[[DELIN]]#0, %[[ARG1]]] by (12, 16) +// CHECK: return %[[LIN]] +func.func @cancel_linearize_delinearize_head_delinearize_unbounded(%arg0: index, %arg1: index) -> index { + %0:3 = affine.delinearize_index %arg0 into (4, 8) : index, index, index + %1 = affine.linearize_index [%0#0, %0#1, %arg1] by (3, 4, 16) : index + return %1 : index +} + +// ----- + +// CHECK-LABEL: func @cancel_linearize_delinearize_head_linearize_unbounded( +// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: index, +// CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]: index) +// CHECK: %[[DELIN:.+]]:2 = affine.delinearize_index %[[ARG0]] into (8) +// CHECK: %[[LIN:.+]] = affine.linearize_index [%[[DELIN]]#0, %[[ARG1]]] by (16) +// CHECK: return %[[LIN]] +func.func @cancel_linearize_delinearize_head_linearize_unbounded(%arg0: index, %arg1: index) -> index { + %0:3 = affine.delinearize_index %arg0 into (3, 4, 8) : index, index, index + %1 = affine.linearize_index [%0#0, %0#1, %arg1] by (4, 16) : index + return %1 : index +} + +// ----- + +// CHECK-LABEL: func @cancel_linearize_delinearize_head_both_unbounded( +// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: index, +// CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]: index) +// CHECK: %[[DELIN:.+]]:2 = affine.delinearize_index %[[ARG0]] into (8) +// CHECK: %[[LIN:.+]] = affine.linearize_index [%[[DELIN]]#0, %[[ARG1]]] by (16) +// CHECK: return %[[LIN]] +func.func @cancel_linearize_delinearize_head_both_unbounded(%arg0: index, %arg1: index) -> index { + %0:3 = affine.delinearize_index %arg0 into (4, 8) : index, index, index + %1 = affine.linearize_index [%0#0, %0#1, %arg1] by (4, 16) : index + return %1 : index +} + +// ----- + +// CHECK-LABEL: func @cancel_linearize_delinearize_tail( +// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: index, +// CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]: index) +// CHECK: %[[DELIN:.+]]:2 = affine.delinearize_index %[[ARG0]] into (3, 32) +// CHECK: %[[LIN:.+]] = affine.linearize_index [%[[ARG1]], %[[DELIN]]#1] by (5, 32) +// CHECK: return %[[LIN]] +func.func @cancel_linearize_delinearize_tail(%arg0: index, %arg1: index) -> index { + %0:3 = affine.delinearize_index %arg0 into (3, 4, 8) : index, index, index + %1 = affine.linearize_index [%arg1, %0#1, %0#2] by (5, 4, 8) : index + return %1 : index +} + +// ----- + +// CHECK-LABEL: func @cancel_linearize_delinearize_middle_exact( +// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: index, +// CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]: index, +// CHECK-SAME: %[[ARG2:[a-zA-z0-9]+]]: index) +// CHECK: %[[LIN:.+]] = affine.linearize_index [%[[ARG1]], %[[ARG0]], %[[ARG2]]] by (9, 30, 7) +// CHECK: return %[[LIN]] +func.func @cancel_linearize_delinearize_middle_exact(%arg0: index, %arg1: index, %arg2: index) -> index { + %0:3 = affine.delinearize_index %arg0 into (2, 3, 5) : index, index, index + %1 = affine.linearize_index [%arg1, %0#0, %0#1, %0#2, %arg2] by (9, 2, 3, 5, 7) : index + return %1 : index +} + +// ----- + +// CHECK: #[[$MAP:.+]] = affine_map<()[s0, s1] -> ((s0 * s1) * 16)> + +// CHECK-LABEL: func @cancel_linearize_delinearize_middle_exact_dynamic_basis( +// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: index, +// CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]: index, +// CHECK-SAME: %[[ARG2:[a-zA-z0-9]+]]: index) +// CHECK: %[[C1:.+]] = arith.constant 1 : index +// CHECK: %[[SIZEPROD:.+]] = affine.apply #[[$MAP]]()[%[[ARG1]], %[[ARG2]]] +// CHECK: %[[LIN:.+]] = affine.linearize_index [%[[C1]], %[[ARG0]], %[[C1]]] by (3, %[[SIZEPROD]], 4) +// CHECK: return %[[LIN]] +func.func @cancel_linearize_delinearize_middle_exact_dynamic_basis(%arg0: index, %arg1: index, %arg2: index) -> index { + %c1 = arith.constant 1 : index + %0:4 = affine.delinearize_index %arg0 into (2, %arg1, %arg2, 8) : index, index, index, index + %1 = affine.linearize_index [%c1, %0#0, %0#1, %0#2, %0#3, %c1] by (3, 2, %arg1, %arg2, 8, 4) : index + return %1 : index +} + +// ----- + +// CHECK-LABEL: func @cancel_linearize_delinearize_middle_exact_delinearize_unbounded_disjoint( +// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: index, +// CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]: index, +// CHECK-SAME: %[[ARG2:[a-zA-z0-9]+]]: index) +// CHECK: %[[LIN:.+]] = affine.linearize_index disjoint [%[[ARG1]], %[[ARG0]], %[[ARG2]]] by (9, 30, 7) +// CHECK: return %[[LIN]] +func.func @cancel_linearize_delinearize_middle_exact_delinearize_unbounded_disjoint(%arg0: index, %arg1: index, %arg2: index) -> index { + %0:3 = affine.delinearize_index %arg0 into (3, 5) : index, index, index + %1 = affine.linearize_index disjoint [%arg1, %0#0, %0#1, %0#2, %arg2] by (9, 2, 3, 5, 7) : index + return %1 : index +} + +// ----- + +// Unlike in the test above, the linerize indices aren't asserted to be disjoint, so +// we can't know if the `2` from the basis is a correct bound. +// CHECK-LABEL: func @dont_cancel_linearize_delinearize_middle_exact_delinearize_unbounded( +// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: index, +// CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]: index, +// CHECK-SAME: %[[ARG2:[a-zA-z0-9]+]]: index) +// CHECK: %[[DELIN:.+]]:2 = affine.delinearize_index %[[ARG0]] into (3) +// CHECK: %[[LIN:.+]] = affine.linearize_index [%[[ARG1]], %[[DELIN]]#0, %[[DELIN]]#1, %[[ARG2]]] by (9, 2, 3, 7) +// CHECK: return %[[LIN]] + +func.func @dont_cancel_linearize_delinearize_middle_exact_delinearize_unbounded(%arg0: index, %arg1: index, %arg2: index) -> index { + %0:2 = affine.delinearize_index %arg0 into (3) : index, index + %1 = affine.linearize_index [%arg1, %0#0, %0#1, %arg2] by (9, 2, 3, 7) : index + return %1 : index +} + +// ----- + +// The presence of a `disjoint` here tells us that the "unbounded" term on the +// delinearization can't have been above 2. +// CHECK-LABEL: func @cancel_linearize_delinearize_middle_delinearize_unbounded_disjoint_implied_bound( +// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: index, +// CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]: index, +// CHECK-SAME: %[[ARG2:[a-zA-z0-9]+]]: index) +// CHECK: %[[DELIN:.+]]:2 = affine.delinearize_index %[[ARG0]] into (6, 5) +// CHECK: %[[LIN:.+]] = affine.linearize_index disjoint [%[[ARG1]], %[[DELIN]]#0, %[[ARG2]]] by (9, 6, 7) +// CHECK: return %[[LIN]] +func.func @cancel_linearize_delinearize_middle_delinearize_unbounded_disjoint_implied_bound(%arg0: index, %arg1: index, %arg2: index) -> index { + %0:3 = affine.delinearize_index %arg0 into (3, 5) : index, index, index + %1 = affine.linearize_index disjoint [%arg1, %0#0, %0#1, %arg2] by (9, 2, 3, 7) : index + return %1 : index +} + +// ----- + +// CHECK-LABEL: func @cancel_linearize_delinearize_multiple_matches( +// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: index, +// CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]: index) +// CHECK: %[[C0:.+]] = arith.constant 0 +// CHECK: %[[DELIN:.+]]:4 = affine.delinearize_index %[[ARG0]] into (4, 16, 4, 64) +// CHECK: %[[LIN:.+]] = affine.linearize_index [%[[ARG1]], %[[DELIN]]#1, %[[C0]], %[[DELIN]]#3] by (4, 16, 4, 64) +// CHECK: return %[[LIN]] +func.func @cancel_linearize_delinearize_multiple_matches(%arg0: index, %arg1: index) -> index { + %c0 = arith.constant 0 : index + %0:7 = affine.delinearize_index %arg0 into (4, 4, 4, 4, 4, 4, 4) : index, index, index, index, index, index, index + %1 = affine.linearize_index [%arg1, %0#1, %0#2, %c0, %0#4, %0#5, %0#6] by (4, 4, 4, 4, 4, 4, 4) : index + return %1 : index +} + +// ----- + +// CHECK-LABEL: func @cancel_linearize_delinearize_multiple_delinearizes( +// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: index, +// CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]: index) +// CHECK: %[[LIN:.+]] = affine.linearize_index [%[[ARG0]], %[[ARG1]]] by (32, 32) +// CHECK: return %[[LIN]] +func.func @cancel_linearize_delinearize_multiple_delinearizes(%arg0: index, %arg1: index) -> index { + %0:2 = affine.delinearize_index %arg0 into (4, 8) : index, index + %1:2 = affine.delinearize_index %arg1 into (2, 16) : index, index + %2 = affine.linearize_index [%0#0, %0#1, %1#0, %1#1] by (4, 8, 2, 16) : index + return %2 : index +} + +// ----- + +// Don't cancel because the values from the delinearize aren't used in order +// CHECK-LABEL: func @no_cancel_linearize_delinearize_permuted( +// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: index, +// CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]: index, +// CHECK-SAME: %[[ARG2:[a-zA-Z0-9]+]]: index) +// CHECK: %[[DELIN:.+]]:3 = affine.delinearize_index %[[ARG0]] into (%[[ARG1]], 4, %[[ARG2]]) +// CHECK: %[[LIN:.+]] = affine.linearize_index [%[[DELIN]]#0, %[[DELIN]]#2, %[[DELIN]]#1] by (%[[ARG1]], %[[ARG2]], 4) +// CHECK: return %[[LIN]] +func.func @no_cancel_linearize_delinearize_permuted(%arg0: index, %arg1: index, %arg2: index) -> index { + %0:3 = affine.delinearize_index %arg0 into (%arg1, 4, %arg2) : index, index, index + %1 = affine.linearize_index [%0#0, %0#2, %0#1] by (%arg1, %arg2, 4) : index + return %1 : index +} + +// ----- + +// CHECK: #[[$MAP:.+]] = affine_map<()[s0] -> (s0 * 3)> +// But these cancel because they're a contiguous segment +// CHECK-LABEL: func @partial_cancel_linearize_delinearize_not_fully_permuted( +// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: index, +// CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]: index, +// CHECK-SAME: %[[ARG2:[a-zA-Z0-9]+]]: index) +// CHECK: %[[SIZEPROD:.+]] = affine.apply #[[$MAP]]()[%[[ARG2]]] +// CHECK: %[[DELIN:.+]]:3 = affine.delinearize_index %[[ARG0]] into (%[[ARG1]], 4, %[[SIZEPROD]]) +// CHECK: %[[LIN:.+]] = affine.linearize_index [%[[DELIN]]#0, %[[DELIN]]#2, %[[DELIN]]#1] by (%[[ARG1]], %[[SIZEPROD]], 4) +// CHECK: return %[[LIN]] +func.func @partial_cancel_linearize_delinearize_not_fully_permuted(%arg0: index, %arg1: index, %arg2: index) -> index { + %0:4 = affine.delinearize_index %arg0 into (%arg1, 4, %arg2, 3) : index, index, index, index + %1 = affine.linearize_index [%0#0, %0#2, %0#3, %0#1] by (%arg1, %arg2, 3, 4) : index + return %1 : index +} + +// ----- + +// Ensure we don't get SSA errors when creating new `affine.delinearize` operations. +// CHECK-LABEL: func @cancel_linearize_delinearize_placement +// CHECK-SAME: (%[[ARG0:.+]]: index) +// CHECK: %[[C0:.+]] = arith.constant 0 : index +// CHECK: %[[NEW_DELIN:.+]]:2 = affine.delinearize_index %[[ARG0]] into (8, 32) : index, index +// CHECK-NEXT: %[[DELIN_PART:.+]]:2 = affine.delinearize_index %[[NEW_DELIN]]#1 into (8, 4) : index, index +// CHECK-NEXT: %[[L1:.+]] = affine.linearize_index disjoint [%[[DELIN_PART]]#1, %[[NEW_DELIN]]#0, %[[C0]], %[[C0]]] by (4, 8, 4, 8) +// CHECK-NEXT: %[[L2:.+]] = affine.linearize_index disjoint [%[[NEW_DELIN]]#1, %[[C0]], %[[C0]]] by (32, 8, 4) +// CHECK-NEXT: %[[L3:.+]] = affine.linearize_index disjoint [%[[DELIN_PART]]#0, %[[NEW_DELIN]]#0, %[[C0]], %[[C0]]] by (8, 8, 4, 4) +// CHECK-NEXT: return %[[L1]], %[[L2]], %[[L3]] +func.func @cancel_linearize_delinearize_placement(%arg0: index) -> (index, index, index) { + %c0 = arith.constant 0 : index + %0:3 = affine.delinearize_index %arg0 into (8, 8, 4) : index, index, index + %1 = affine.linearize_index disjoint [%0#2, %0#0, %c0, %c0] by (4, 8, 4, 8) : index + %2 = affine.linearize_index disjoint [%0#1, %0#2, %c0, %c0] by (8, 4, 8, 4) : index + %3 = affine.linearize_index disjoint [%0#1, %0#0, %c0, %c0] by (8, 8, 4, 4) : index + return %1, %2, %3 : index, index, index +} + +// ----- + // Won't cancel because the linearize and delinearize are using a different basis -// CHECK-LABEL: func @no_cancel_linearize_denearize_different_basis( +// CHECK-LABEL: func @no_cancel_linearize_delinearize_different_basis( // CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: index, // CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]: index, // CHECK-SAME: %[[ARG2:[a-zA-Z0-9]+]]: index) // CHECK: %[[DELIN:.+]]:3 = affine.delinearize_index %[[ARG0]] into (%[[ARG1]], 4, %[[ARG2]]) // CHECK: %[[LIN:.+]] = affine.linearize_index [%[[DELIN]]#0, %[[DELIN]]#1, %[[DELIN]]#2] by (%[[ARG1]], 8, %[[ARG2]]) // CHECK: return %[[LIN]] -func.func @no_cancel_linearize_denearize_different_basis(%arg0: index, %arg1: index, %arg2: index) -> index { +func.func @no_cancel_linearize_delinearize_different_basis(%arg0: index, %arg1: index, %arg2: index) -> index { %0:3 = affine.delinearize_index %arg0 into (%arg1, 4, %arg2) : index, index, index %1 = affine.linearize_index [%0#0, %0#1, %0#2] by (%arg1, 8, %arg2) : index return %1 : index