@@ -5083,6 +5083,31 @@ SmallVector<OpFoldResult> AffineLinearizeIndexOp::getPaddedBasis() {
50835083 return ret;
50845084}
50855085
5086+ namespace mlir {
5087+ namespace affine {
5088+ OpFoldResult computeProduct (Location loc, OpBuilder &builder,
5089+ ArrayRef<OpFoldResult> terms) {
5090+ int64_t nDynamic = 0 ;
5091+ SmallVector<Value> dynamicPart;
5092+ AffineExpr result = builder.getAffineConstantExpr (1 );
5093+ for (OpFoldResult term : terms) {
5094+ if (!term)
5095+ return term;
5096+ std::optional<int64_t > maybeConst = getConstantIntValue (term);
5097+ if (maybeConst) {
5098+ result = result * builder.getAffineConstantExpr (*maybeConst);
5099+ } else {
5100+ dynamicPart.push_back (cast<Value>(term));
5101+ result = result * builder.getAffineSymbolExpr (nDynamic++);
5102+ }
5103+ }
5104+ if (auto constant = dyn_cast<AffineConstantExpr>(result))
5105+ return getAsIndexOpFoldResult (builder.getContext (), constant.getValue ());
5106+ return builder.create <AffineApplyOp>(loc, result, dynamicPart).getResult ();
5107+ }
5108+ } // namespace affine
5109+ } // namespace mlir
5110+
50865111namespace {
50875112// / Rewrite `affine.linearize_index disjoint [%...a, %x, %...b] by (%...c, 1,
50885113// / %...d)` to `affine.linearize_index disjoint [%...a, %...b] by (%...c,
@@ -5142,27 +5167,6 @@ struct DropLinearizeUnitComponentsIfDisjointOrZero final
51425167 }
51435168};
51445169
5145- OpFoldResult computeProduct (Location loc, OpBuilder &builder,
5146- ArrayRef<OpFoldResult> terms) {
5147- int64_t nDynamic = 0 ;
5148- SmallVector<Value> dynamicPart;
5149- AffineExpr result = builder.getAffineConstantExpr (1 );
5150- for (OpFoldResult term : terms) {
5151- if (!term)
5152- return term;
5153- std::optional<int64_t > maybeConst = getConstantIntValue (term);
5154- if (maybeConst) {
5155- result = result * builder.getAffineConstantExpr (*maybeConst);
5156- } else {
5157- dynamicPart.push_back (cast<Value>(term));
5158- result = result * builder.getAffineSymbolExpr (nDynamic++);
5159- }
5160- }
5161- if (auto constant = dyn_cast<AffineConstantExpr>(result))
5162- return getAsIndexOpFoldResult (builder.getContext (), constant.getValue ());
5163- return builder.create <AffineApplyOp>(loc, result, dynamicPart).getResult ();
5164- }
5165-
51665170// / If conseceutive outputs of a delinearize_index are linearized with the same
51675171// / bounds, canonicalize away the redundant arithmetic.
51685172// /
@@ -5309,7 +5313,7 @@ struct CancelLinearizeOfDelinearizePortion final
53095313 // We use the slice from the linearize's basis above because of the
53105314 // "bounds inferred from `disjoint`" case above.
53115315 OpFoldResult newSize =
5312- computeProduct (linearizeOp.getLoc (), rewriter, basisToMerge);
5316+ affine:: computeProduct (linearizeOp.getLoc (), rewriter, basisToMerge);
53135317
53145318 // Trivial case where we can just skip past the delinearize all together
53155319 if (m.length == m.delinearize .getNumResults ()) {
0 commit comments