[mlir][Affine][NFC] Define AffineForOp operands in ODS

[matthias-springer]

Modernize affine dialect ops: Define LB, UB, step and inits as operands in TableGen.

[llvmbot]

@llvm/pr-subscribers-mlir-gpu
@llvm/pr-subscribers-mlir-affine

@llvm/pr-subscribers-mlir

Changes

Modernize affine dialect ops: Define LB, UB, step and inits as operands in TableGen.

Note: There are still some operands that have not been updated to TableGen yet.

---

Patch is 26.86 KiB, truncated to 20.00 KiB below, full version: https://github.com/llvm/llvm-project/pull/67694.diff

11 Files Affected:

• (modified) mlir/include/mlir/Dialect/Affine/IR/AffineOps.td (+11-20)
• (modified) mlir/lib/Conversion/AffineToStandard/AffineToStandard.cpp (+2-1)
• (modified) mlir/lib/Conversion/SCFToGPU/SCFToGPU.cpp (+1-1)
• (modified) mlir/lib/Dialect/Affine/Analysis/AffineStructures.cpp (+1-1)
• (modified) mlir/lib/Dialect/Affine/Analysis/LoopAnalysis.cpp (+1-1)
• (modified) mlir/lib/Dialect/Affine/IR/AffineOps.cpp (+37-42)
• (modified) mlir/lib/Dialect/Affine/Transforms/PipelineDataTransfer.cpp (+1-1)
• (modified) mlir/lib/Dialect/Affine/Transforms/SuperVectorize.cpp (+3-3)
• (modified) mlir/lib/Dialect/Affine/Utils/LoopUtils.cpp (+22-21)
• (modified) mlir/lib/Dialect/Affine/Utils/Utils.cpp (+2-2)
• (modified) mlir/test/Dialect/Affine/ops.mlir (+2-2)

diff --git a/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td b/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td
index d8ef0506d0822d7..43ad0a82ff58f07 100644
--- a/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td
+++ b/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td
@@ -117,7 +117,8 @@ def AffineApplyOp : Affine_Op<"apply", [Pure]> {
 }
 
 def AffineForOp : Affine_Op<"for",
-    [AutomaticAllocationScope, ImplicitAffineTerminator, ConditionallySpeculatable,
+    [AttrSizedOperandSegments, AutomaticAllocationScope,
+     ImplicitAffineTerminator, ConditionallySpeculatable,
      RecursiveMemoryEffects, DeclareOpInterfaceMethods<LoopLikeOpInterface,
      ["getSingleInductionVar", "getSingleLowerBound", "getSingleStep",
       "getSingleUpperBound", "replaceWithAdditionalYields"]>,
@@ -147,7 +148,7 @@ def AffineForOp : Affine_Op<"for",
     from a lower bound to an upper bound by a stride. The stride, represented by
     `step`, is a positive constant integer which defaults to "1" if not present.
     The lower and upper bounds specify a half-open range: the range includes the
-    lower bound but does not include the upper bound.
+    lower bound but does not include the upper bound.
 
     The lower and upper bounds of a `affine.for` operation are represented as an
     application of an affine mapping to a list of SSA values passed to the map.
@@ -229,7 +230,10 @@ def AffineForOp : Affine_Op<"for",
     explicitly present. The number and types of the "affine.for" results must
     match the initial values in the `iter_args` binding and the yield operands.
   }];
-  let arguments = (ins Variadic<AnyType>);
+  let arguments = (ins Variadic<Index>:$lowerBoundOperands,
+                       Variadic<Index>:$upperBoundOperands,
+                       Variadic<AnyType>:$inits,
+                       IndexAttr:$step);
   let results = (outs Variadic<AnyType>:$results);
   let regions = (region SizedRegion<1>:$region);
 
@@ -259,19 +263,10 @@ def AffineForOp : Affine_Op<"for",
     Block::BlockArgListType getRegionIterArgs() {
       return getBody()->getArguments().drop_front();
     }
-    Operation::operand_range getInits() {
-      return getOperands().drop_front(getNumControlOperands());
-    }
 
     // TODO: provide iterators for the lower and upper bound operands
     // if the current access via getLowerBound(), getUpperBound() is too slow.
 
-    /// Returns operands for the lower bound map.
-    operand_range getLowerBoundOperands();
-
-    /// Returns operands for the upper bound map.
-    operand_range getUpperBoundOperands();
-
     /// Returns operands for the lower and upper bound maps with the operands
     /// for the lower bound map in front of those for the upper bound map.
     operand_range getControlOperands();
@@ -283,19 +278,17 @@ def AffineForOp : Affine_Op<"for",
     AffineBound getUpperBound();
 
     /// Returns loop step.
-    int64_t getStep() {
-      return ::llvm::cast<IntegerAttr>(*(*this)->getInherentAttr(getStepAttrStrName())).getInt();
-    }
+    int64_t getStepAsInt() { return getStep().getSExtValue(); }
 
     /// Returns affine map for the lower bound.
     AffineMap getLowerBoundMap() { return getLowerBoundMapAttr().getValue(); }
     AffineMapAttr getLowerBoundMapAttr() {
-      return ::llvm::cast<AffineMapAttr>(*(*this)->getInherentAttr(getLowerBoundAttrStrName()));
+      return ::llvm::cast<AffineMapAttr>(getOperation()->getAttr(getLowerBoundAttrStrName()));
     }
     /// Returns affine map for the upper bound. The upper bound is exclusive.
     AffineMap getUpperBoundMap() { return getUpperBoundMapAttr().getValue(); }
     AffineMapAttr getUpperBoundMapAttr() {
-      return ::llvm::cast<AffineMapAttr>(*(*this)->getInherentAttr(getUpperBoundAttrStrName()));
+      return ::llvm::cast<AffineMapAttr>(getOperation()->getAttr(getUpperBoundAttrStrName()));
     }
 
     /// Set lower bound. The new bound must have the same number of operands as
@@ -314,9 +307,7 @@ def AffineForOp : Affine_Op<"for",
     /// Set loop step.
     void setStep(int64_t step) {
       assert(step > 0 && "step has to be a positive integer constant");
-      auto *context = getLowerBoundMap().getContext();
-      (*this)->setAttr(StringAttr::get(context, getStepAttrStrName()),
-                       IntegerAttr::get(IndexType::get(context), step));
+      setStep(APInt(/*numBits=*/64, step, /*isSigned=*/true));
     }
 
     /// Returns number of region arguments for loop-carried values.
diff --git a/mlir/lib/Conversion/AffineToStandard/AffineToStandard.cpp b/mlir/lib/Conversion/AffineToStandard/AffineToStandard.cpp
index 456fd487352fb5d..7dbbf015182f39f 100644
--- a/mlir/lib/Conversion/AffineToStandard/AffineToStandard.cpp
+++ b/mlir/lib/Conversion/AffineToStandard/AffineToStandard.cpp
@@ -156,7 +156,8 @@ class AffineForLowering : public OpRewritePattern<AffineForOp> {
     Location loc = op.getLoc();
     Value lowerBound = lowerAffineLowerBound(op, rewriter);
     Value upperBound = lowerAffineUpperBound(op, rewriter);
-    Value step = rewriter.create<arith::ConstantIndexOp>(loc, op.getStep());
+    Value step =
+        rewriter.create<arith::ConstantIndexOp>(loc, op.getStepAsInt());
     auto scfForOp = rewriter.create<scf::ForOp>(loc, lowerBound, upperBound,
                                                 step, op.getInits());
     rewriter.eraseBlock(scfForOp.getBody());
diff --git a/mlir/lib/Conversion/SCFToGPU/SCFToGPU.cpp b/mlir/lib/Conversion/SCFToGPU/SCFToGPU.cpp
index aff620ad0e70608..94d875d678df293 100644
--- a/mlir/lib/Conversion/SCFToGPU/SCFToGPU.cpp
+++ b/mlir/lib/Conversion/SCFToGPU/SCFToGPU.cpp
@@ -86,7 +86,7 @@ static Operation::operand_range getUpperBoundOperands(AffineForOp forOp) {
 // materialize a corresponding constant using builder.
 static Value getOrCreateStep(AffineForOp forOp, OpBuilder &builder) {
   return builder.create<arith::ConstantIndexOp>(forOp.getLoc(),
-                                                forOp.getStep());
+                                                forOp.getStepAsInt());
 }
 
 // Get a Value for the loop lower bound.  If the value requires computation,
diff --git a/mlir/lib/Dialect/Affine/Analysis/AffineStructures.cpp b/mlir/lib/Dialect/Affine/Analysis/AffineStructures.cpp
index 5f32505690263fc..6ed3ba14fe15229 100644
--- a/mlir/lib/Dialect/Affine/Analysis/AffineStructures.cpp
+++ b/mlir/lib/Dialect/Affine/Analysis/AffineStructures.cpp
@@ -74,7 +74,7 @@ FlatAffineValueConstraints::addAffineForOpDomain(AffineForOp forOp) {
     return failure();
   }
 
-  int64_t step = forOp.getStep();
+  int64_t step = forOp.getStepAsInt();
   if (step != 1) {
     if (!forOp.hasConstantLowerBound())
       LLVM_DEBUG(forOp.emitWarning("domain conservatively approximated"));
diff --git a/mlir/lib/Dialect/Affine/Analysis/LoopAnalysis.cpp b/mlir/lib/Dialect/Affine/Analysis/LoopAnalysis.cpp
index c97e99c0a0c19da..d56db64eac08261 100644
--- a/mlir/lib/Dialect/Affine/Analysis/LoopAnalysis.cpp
+++ b/mlir/lib/Dialect/Affine/Analysis/LoopAnalysis.cpp
@@ -39,7 +39,7 @@ void mlir::affine::getTripCountMapAndOperands(
     AffineForOp forOp, AffineMap *tripCountMap,
     SmallVectorImpl<Value> *tripCountOperands) {
   MLIRContext *context = forOp.getContext();
-  int64_t step = forOp.getStep();
+  int64_t step = forOp.getStepAsInt();
   int64_t loopSpan;
   if (forOp.hasConstantBounds()) {
     int64_t lb = forOp.getConstantLowerBound();
diff --git a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
index 6c060c90e24af82..c2dd84e6e67ca7d 100644
--- a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
+++ b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
@@ -613,11 +613,11 @@ static int64_t getLargestKnownDivisor(AffineExpr e, ArrayRef<Value> operands) {
   // LoopLikeOpInterface.
   if (AffineForOp forOp = getForInductionVarOwner(operand)) {
     if (forOp.hasConstantLowerBound() && forOp.getConstantLowerBound() == 0) {
-      operandDivisor = forOp.getStep();
+      operandDivisor = forOp.getStepAsInt();
     } else {
       uint64_t lbLargestKnownDivisor =
           forOp.getLowerBoundMap().getLargestKnownDivisorOfMapExprs();
-      operandDivisor = std::gcd(lbLargestKnownDivisor, forOp.getStep());
+      operandDivisor = std::gcd(lbLargestKnownDivisor, forOp.getStepAsInt());
     }
   }
   return operandDivisor;
@@ -695,7 +695,7 @@ static std::optional<int64_t> getUpperBound(Value iv) {
   if (forOp.hasConstantLowerBound()) {
     return forOp.getConstantUpperBound() - 1 -
            (forOp.getConstantUpperBound() - forOp.getConstantLowerBound() - 1) %
-               forOp.getStep();
+               forOp.getStepAsInt();
   }
   return forOp.getConstantUpperBound() - 1;
 }
@@ -1897,6 +1897,13 @@ void AffineForOp::build(OpBuilder &builder, OperationState &result,
          "upper bound operand count does not match the affine map");
   assert(step > 0 && "step has to be a positive integer constant");
 
+  // Set variadic segment sizes.
+  result.addAttribute(
+      getOperandSegmentSizeAttr(),
+      builder.getDenseI32ArrayAttr({static_cast<int32_t>(lbOperands.size()),
+                                    static_cast<int32_t>(ubOperands.size()),
+                                    static_cast<int32_t>(iterArgs.size())}));
+
   for (Value val : iterArgs)
     result.addTypes(val.getType());
 
@@ -2083,10 +2090,16 @@ ParseResult AffineForOp::parse(OpAsmParser &parser, OperationState &result) {
     return failure();
 
   // Parse loop bounds.
-  if (parseBound(/*isLower=*/true, result, parser) ||
-      parser.parseKeyword("to", " between bounds") ||
-      parseBound(/*isLower=*/false, result, parser))
+  int64_t numOperands = result.operands.size();
+  if (parseBound(/*isLower=*/true, result, parser))
+    return failure();
+  int64_t numLbOperands = result.operands.size() - numOperands;
+  if (parser.parseKeyword("to", " between bounds"))
     return failure();
+  numOperands = result.operands.size();
+  if (parseBound(/*isLower=*/false, result, parser))
+    return failure();
+  int64_t numUbOperands = result.operands.size() - numOperands;
 
   // Parse the optional loop step, we default to 1 if one is not present.
   if (parser.parseOptionalKeyword("step")) {
@@ -2130,6 +2143,12 @@ ParseResult AffineForOp::parse(OpAsmParser &parser, OperationState &result) {
     }
   }
 
+  result.addAttribute(
+      getOperandSegmentSizeAttr(),
+      builder.getDenseI32ArrayAttr({static_cast<int32_t>(numLbOperands),
+                                    static_cast<int32_t>(numUbOperands),
+                                    static_cast<int32_t>(operands.size())}));
+
   // Parse the body region.
   Region *body = result.addRegion();
   if (regionArgs.size() != result.types.size() + 1)
@@ -2202,8 +2221,8 @@ void AffineForOp::print(OpAsmPrinter &p) {
   p << " to ";
   printBound(getUpperBoundMapAttr(), getUpperBoundOperands(), "min", p);
 
-  if (getStep() != 1)
-    p << " step " << getStep();
+  if (getStepAsInt() != 1)
+    p << " step " << getStepAsInt();
 
   bool printBlockTerminators = false;
   if (getNumIterOperands() > 0) {
@@ -2221,10 +2240,10 @@ void AffineForOp::print(OpAsmPrinter &p) {
   p << ' ';
   p.printRegion(getRegion(), /*printEntryBlockArgs=*/false,
                 printBlockTerminators);
-  p.printOptionalAttrDict((*this)->getAttrs(),
-                          /*elidedAttrs=*/{getLowerBoundAttrStrName(),
-                                           getUpperBoundAttrStrName(),
-                                           getStepAttrStrName()});
+  p.printOptionalAttrDict(
+      (*this)->getAttrs(),
+      /*elidedAttrs=*/{getLowerBoundAttrStrName(), getUpperBoundAttrStrName(),
+                       getStepAttrStrName(), getOperandSegmentSizeAttr()});
 }
 
 /// Fold the constant bounds of a loop.
@@ -2307,7 +2326,7 @@ static LogicalResult canonicalizeLoopBounds(AffineForOp forOp) {
 namespace {
 /// Returns constant trip count in trivial cases.
 static std::optional<uint64_t> getTrivialConstantTripCount(AffineForOp forOp) {
-  int64_t step = forOp.getStep();
+  int64_t step = forOp.getStepAsInt();
   if (!forOp.hasConstantBounds() || step <= 0)
     return std::nullopt;
   int64_t lb = forOp.getConstantLowerBound();
@@ -2452,28 +2471,14 @@ AffineBound AffineForOp::getUpperBound() {
 void AffineForOp::setLowerBound(ValueRange lbOperands, AffineMap map) {
   assert(lbOperands.size() == map.getNumInputs());
   assert(map.getNumResults() >= 1 && "bound map has at least one result");
-
-  SmallVector<Value, 4> newOperands(lbOperands.begin(), lbOperands.end());
-
-  auto ubOperands = getUpperBoundOperands();
-  newOperands.append(ubOperands.begin(), ubOperands.end());
-  auto iterOperands = getInits();
-  newOperands.append(iterOperands.begin(), iterOperands.end());
-  (*this)->setOperands(newOperands);
-
+  getLowerBoundOperandsMutable().assign(lbOperands);
   (*this)->setAttr(getLowerBoundAttrStrName(), AffineMapAttr::get(map));
 }
 
 void AffineForOp::setUpperBound(ValueRange ubOperands, AffineMap map) {
   assert(ubOperands.size() == map.getNumInputs());
   assert(map.getNumResults() >= 1 && "bound map has at least one result");
-
-  SmallVector<Value, 4> newOperands(getLowerBoundOperands());
-  newOperands.append(ubOperands.begin(), ubOperands.end());
-  auto iterOperands = getInits();
-  newOperands.append(iterOperands.begin(), iterOperands.end());
-  (*this)->setOperands(newOperands);
-
+  getUpperBoundOperandsMutable().assign(ubOperands);
   (*this)->setAttr(getUpperBoundAttrStrName(), AffineMapAttr::get(map));
 }
 
@@ -2519,16 +2524,6 @@ void AffineForOp::setConstantUpperBound(int64_t value) {
   setUpperBound({}, AffineMap::getConstantMap(value, getContext()));
 }
 
-AffineForOp::operand_range AffineForOp::getLowerBoundOperands() {
-  return {operand_begin(), operand_begin() + getLowerBoundMap().getNumInputs()};
-}
-
-AffineForOp::operand_range AffineForOp::getUpperBoundOperands() {
-  return {operand_begin() + getLowerBoundMap().getNumInputs(),
-          operand_begin() + getLowerBoundMap().getNumInputs() +
-              getUpperBoundMap().getNumInputs()};
-}
-
 AffineForOp::operand_range AffineForOp::getControlOperands() {
   return {operand_begin(), operand_begin() + getLowerBoundMap().getNumInputs() +
                                getUpperBoundMap().getNumInputs()};
@@ -2565,7 +2560,7 @@ std::optional<OpFoldResult> AffineForOp::getSingleLowerBound() {
 
 std::optional<OpFoldResult> AffineForOp::getSingleStep() {
   OpBuilder b(getContext());
-  return OpFoldResult(b.getI64IntegerAttr(getStep()));
+  return OpFoldResult(b.getI64IntegerAttr(getStepAsInt()));
 }
 
 std::optional<OpFoldResult> AffineForOp::getSingleUpperBound() {
@@ -2633,8 +2628,8 @@ Speculation::Speculatability AffineForOp::getSpeculatability() {
   //
   // For Step != 1, the loop may not terminate.  We can add more smarts here if
   // needed.
-  return getStep() == 1 ? Speculation::RecursivelySpeculatable
-                        : Speculation::NotSpeculatable;
+  return getStepAsInt() == 1 ? Speculation::RecursivelySpeculatable
+                             : Speculation::NotSpeculatable;
 }
 
 /// Returns true if the provided value is the induction variable of a
diff --git a/mlir/lib/Dialect/Affine/Transforms/PipelineDataTransfer.cpp b/mlir/lib/Dialect/Affine/Transforms/PipelineDataTransfer.cpp
index 7029251a3720cb7..d169aca2c971f24 100644
--- a/mlir/lib/Dialect/Affine/Transforms/PipelineDataTransfer.cpp
+++ b/mlir/lib/Dialect/Affine/Transforms/PipelineDataTransfer.cpp
@@ -107,7 +107,7 @@ static bool doubleBuffer(Value oldMemRef, AffineForOp forOp) {
 
   // Create 'iv mod 2' value to index the leading dimension.
   auto d0 = bInner.getAffineDimExpr(0);
-  int64_t step = forOp.getStep();
+  int64_t step = forOp.getStepAsInt();
   auto modTwoMap =
       AffineMap::get(/*dimCount=*/1, /*symbolCount=*/0, d0.floorDiv(step) % 2);
   auto ivModTwoOp = bInner.create<AffineApplyOp>(forOp.getLoc(), modTwoMap,
diff --git a/mlir/lib/Dialect/Affine/Transforms/SuperVectorize.cpp b/mlir/lib/Dialect/Affine/Transforms/SuperVectorize.cpp
index 85c2602aa266d67..f07ed626a916384 100644
--- a/mlir/lib/Dialect/Affine/Transforms/SuperVectorize.cpp
+++ b/mlir/lib/Dialect/Affine/Transforms/SuperVectorize.cpp
@@ -1021,7 +1021,7 @@ static Value createMask(AffineForOp vecForOp, VectorizationState &state) {
   if (vecForOp.hasConstantBounds()) {
     int64_t originalTripCount =
         vecForOp.getConstantUpperBound() - vecForOp.getConstantLowerBound();
-    if (originalTripCount % vecForOp.getStep() == 0)
+    if (originalTripCount % vecForOp.getStepAsInt() == 0)
       return nullptr;
   }
 
@@ -1296,9 +1296,9 @@ static Operation *vectorizeAffineForOp(AffineForOp forOp,
     unsigned vectorDim = loopToVecDimIt->second;
     assert(vectorDim < strategy.vectorSizes.size() && "vector dim overflow");
     int64_t forOpVecFactor = strategy.vectorSizes[vectorDim];
-    newStep = forOp.getStep() * forOpVecFactor;
+    newStep = forOp.getStepAsInt() * forOpVecFactor;
   } else {
-    newStep = forOp.getStep();
+    newStep = forOp.getStepAsInt();
   }
 
   // Get information about reduction kinds.
diff --git a/mlir/lib/Dialect/Affine/Utils/LoopUtils.cpp b/mlir/lib/Dialect/Affine/Utils/LoopUtils.cpp
index 9d8ed9b4ac93387..fb8a0a7c330cf22 100644
--- a/mlir/lib/Dialect/Affine/Utils/LoopUtils.cpp
+++ b/mlir/lib/Dialect/Affine/Utils/LoopUtils.cpp
@@ -78,7 +78,7 @@ getCleanupLoopLowerBound(AffineForOp forOp, unsigned unrollFactor,
   // these affine.apply's make up the cleanup loop lower bound.
   SmallVector<AffineExpr, 4> bumpExprs(tripCountMap.getNumResults());
   SmallVector<Value, 4> bumpValues(tripCountMap.getNumResults());
-  int64_t step = forOp.getStep();
+  int64_t step = forOp.getStepAsInt();
   for (unsigned i = 0, e = tripCountMap.getNumResults(); i < e; i++) {
     auto tripCountExpr = tripCountMap.getResult(i);
     bumpExprs[i] = (tripCountExpr - tripCountExpr % unrollFactor) * step;
@@ -189,8 +189,9 @@ static AffineForOp generateShiftedLoop(
   assert(lbMap.getNumInputs() == lbOperands.size());
   assert(ubMap.getNumInputs() == ubOperands.size());
 
-  auto loopChunk = b.create<AffineForOp>(srcForOp.getLoc(), lbOperands, lbMap,
-                                         ubOperands, ubMap, srcForOp.getStep());
+  auto loopChunk =
+      b.create<AffineForOp>(srcForOp.getLoc(), lbOperands, lbMap, ubOperands,
+                            ubMap, srcForOp.getStepAsInt());
   auto loopChunkIV = loopChunk.getInductionVar();
   auto srcIV = srcForOp.getInductionVar();
 
@@ -208,7 +209,7 @@ static AffineForOp generateShiftedLoop(
       auto ivRemap = bodyBuilder.create<AffineApplyOp>(
           srcForOp.getLoc(),
           bodyBuilder.getSingleDimShiftAffineMap(
-              -static_cast<int64_t>(srcForOp.getStep() * shift)),
+              -static_cast<int64_t>(srcForOp.getStepAsInt() * shift)),
           loopChunkIV);
       operandMap.map(srcIV, ivRemap);
     } else {
@@ -255,7 +256,7 @@ LogicalResult mlir::affine::affineForOpBodySkew(AffineForOp forOp,
   assert(isOpwiseShiftValid(forOp, shifts) &&
          "shifts will lead to an invalid transformation\n");
 
-  int64_t step = forOp.getStep();
+  int64_t step = forOp.getStepAsInt();
 
   unsigned numChildOps = shifts.size();
 
@@ -535,7 +536,7 @@ static void setIntraTileBoundsParametric(OpBuilder &b, AffineForOp origLoop,
       origLowerBoundExpr);
 
   // Get the origLoopStep as an affine expression.
-  AffineExpr origLoopStep = b.getAffineConstantExpr(origLoop.getStep());
+  AffineExpr origLoopStep = b.getAffineConstantExpr(origLoop.getStepAsInt());
 
   // Insert ub as inter-tile ((loop IV - origlb) * tilingParameter) +
   // (tilingParameter * origLoopStep) + origlb.
@@ -557,7 +558,7 @@ static void setIntraTileBoundsParametric(OpBuilder &b, AffineForOp origLoop,
   newIntraTileLoop.setUpperBound(ubOperands, ubMap);
 
   // Original loop step must be preserved.
-  newIntraTileLoop.setStep(origLoop.getStep());
+  newIntraTileLoop.setStep(origLoop.getStepAsInt());
 }
 
 /// Set lower and upper bounds of inter-tile loops for parametric tiling.
@@ -654,7 +655,7 @@ static void setInterTileBoundsParametric(OpBuilder &b, AffineForOp origLoop,
   newLoop.setUpperBound(ubOperands, ubMap);
 
   // Original loop step must be preserved.
-  newLoop.setStep(origLoop.getStep());
+  newLoop.setStep(origLoop.getStepAsInt());
 }
 
 /// Constructs and sets new loop bounds after tiling for the...
[truncated]