[mlir][Affine][NFC] Define AffineForOp operands in ODS (#67694)

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

Author: matthias-springer

mlir/include/mlir/Dialect/Affine/IR/AffineOps.h

@@ -499,27 +499,28 @@ class AffineBound {
   AffineForOp getAffineForOp() { return op; }
   AffineMap getMap() { return map; }
 
-  unsigned getNumOperands() { return opEnd - opStart; }
-  Value getOperand(unsigned idx) { return op.getOperand(opStart + idx); }
+  unsigned getNumOperands() { return operands.size(); }
+  Value getOperand(unsigned idx) {
+    return op.getOperand(operands.getBeginOperandIndex() + idx);
+  }
 
   using operand_iterator = AffineForOp::operand_iterator;
   using operand_range = AffineForOp::operand_range;
 
-  operand_iterator operandBegin() { return op.operand_begin() + opStart; }
-  operand_iterator operandEnd() { return op.operand_begin() + opEnd; }
+  operand_iterator operandBegin() { return operands.begin(); }
+  operand_iterator operandEnd() { return operands.end(); }
   operand_range getOperands() { return {operandBegin(), operandEnd()}; }
 
 private:
   // 'affine.for' operation that contains this bound.
   AffineForOp op;
-  // Start and end positions of this affine bound operands in the list of
-  // the containing 'affine.for' operation operands.
-  unsigned opStart, opEnd;
+  // Operands of the affine map.
+  OperandRange operands;
   // Affine map for this bound.
   AffineMap map;
 
-  AffineBound(AffineForOp op, unsigned opStart, unsigned opEnd, AffineMap map)
-      : op(op), opStart(opStart), opEnd(opEnd), map(map) {}
+  AffineBound(AffineForOp op, OperandRange operands, AffineMap map)
+      : op(op), operands(operands), map(map) {}
 
   friend class AffineForOp;
 };

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,12 @@ 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,
+                       AffineMapAttr:$lowerBoundMap,
+                       AffineMapAttr:$upperBoundMap,
+                       IndexAttr:$step);
   let results = (outs Variadic<AnyType>:$results);
   let regions = (region SizedRegion<1>:$region);
 
@@ -251,26 +257,10 @@ def AffineForOp : Affine_Op<"for",
     using BodyBuilderFn =
         function_ref<void(OpBuilder &, Location, Value, ValueRange)>;
 
-    static StringRef getStepAttrStrName() { return "step"; }
-    static StringRef getLowerBoundAttrStrName() { return "lower_bound"; }
-    static StringRef getUpperBoundAttrStrName() { return "upper_bound"; }
-
     BlockArgument getInductionVar() { return getBody()->getArgument(0); }
     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.
@@ -283,20 +273,7 @@ def AffineForOp : Affine_Op<"for",
     AffineBound getUpperBound();
 
     /// Returns loop step.
-    int64_t getStep() {
-      return ::llvm::cast<IntegerAttr>(*(*this)->getInherentAttr(getStepAttrStrName())).getInt();
-    }
-
-    /// Returns affine map for the lower bound.
-    AffineMap getLowerBoundMap() { return getLowerBoundMapAttr().getValue(); }
-    AffineMapAttr getLowerBoundMapAttr() {
-      return ::llvm::cast<AffineMapAttr>(*(*this)->getInherentAttr(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()));
-    }
+    int64_t getStepAsInt() { return getStep().getSExtValue(); }
 
     /// Set lower bound. The new bound must have the same number of operands as
     /// the current bound map. Otherwise, 'replaceForLowerBound' should be used.
@@ -305,18 +282,10 @@ def AffineForOp : Affine_Op<"for",
     /// current bound map. Otherwise, 'replaceForUpperBound' should be used.
     void setUpperBound(ValueRange operands, AffineMap map);
 
-    /// Set the lower bound map without changing operands.
-    void setLowerBoundMap(AffineMap map);
-
-    /// Set the upper bound map without changing operands.
-    void setUpperBoundMap(AffineMap map);
-
     /// 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.
@@ -325,7 +294,9 @@ def AffineForOp : Affine_Op<"for",
     }
 
     /// Number of operands controlling the loop: lb and ub.
-    unsigned getNumControlOperands() { return getOperation()->getNumOperands() - getNumIterOperands(); }
+    unsigned getNumControlOperands() {
+      return getOperation()->getNumOperands() - getNumIterOperands();
+    }
 
     /// Get the number of loop-carried values.
     unsigned getNumIterOperands();

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());

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,

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"));

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();