[mlir][scf] Relax type requirement on for

[jpienaar]

Summary

scf.for loop was restricted to only operate on Index type since
splitting out from affine.for. Relax requirement to allow for signless
integer types additionally. This allows specifying explicitly different
bitwidths for different loops as well as specialize from index to iN
while still using scf.for.

[ftynse]

Do we need to also update the lowering to LLVM? It may assume type is always index and end up using a different bitwidth.

[Mogball]

+1 to what Alex said, but otherwise LGT

[jpienaar]

In D145288#4169046, @ftynse wrote:

Do we need to also update the lowering to LLVM? It may assume type is always index and end up using a different bitwidth.

I checked the final assembly output and saw different outputs as I changed width here. But let me double check the code.

[jpienaar]

In D145288#4169715, @jpienaar wrote:

In D145288#4169046, @ftynse wrote:

Do we need to also update the lowering to LLVM? It may assume type is always index and end up using a different bitwidth.

I checked the final assembly output and saw different outputs as I changed width here. But let me double check the code.

SCFToControlFlow.cpp is handling this by way of using type inference method (and the constraints added on ForOp here with matching element type)
SCFToGPU.cpp seems to be only operating on AffineForOp
SCFToSPIRV.cpp is handling this by explicit element type queries
VectorToGPU.cpp looks also not to depend on requiring index type

I see no other lowering from scf::ForOp conversion side.

[jpienaar]

Fix formatting

mlir/include/mlir/Dialect/SCF/IR/SCFOps.td

 //===----------------------------------------------------------------------===//
 // ForOp
 //===----------------------------------------------------------------------===//
 
 def ForOp : SCF_Op<"for",
       [AutomaticAllocationScope, DeclareOpInterfaceMethods<LoopLikeOpInterface,
        ["getSingleInductionVar", "getSingleLowerBound", "getSingleStep",
         "getSingleUpperBound"]>,
+       AllTypesMatch<["lowerBound", "upperBound", "step"]>,
        ConditionallySpeculatable,
        DeclareOpInterfaceMethods<RegionBranchOpInterface>,
        SingleBlockImplicitTerminator<"scf::YieldOp">,
        RecursiveMemoryEffects]> {
   let summary = "for operation";
   let description = [{
-    The "scf.for" operation represents a loop taking 3 SSA value as operands
+    The `scf.for` operation represents a loop taking 3 SSA value as operands
     that represent the lower bound, upper bound and step respectively. The
     operation defines an SSA value for its induction variable. It has one
     region capturing the loop body. The induction variable is represented as an
-    argument of this region. This SSA value always has type index, which is the
-    size of the machine word. The step is a value of type index, required to be
-    positive.
-    The lower and upper bounds specify a half-open range: the range includes
-    the lower bound but does not include the upper bound.
+    argument of this region. This SSA value is a signless integer or index.
+    The step is a value of same type but required to be positive. The lower and
+    upper bounds specify a half-open range: the range includes the lower bound
+    but does not include the upper bound.
 
     The body region must contain exactly one block that terminates with
-    "scf.yield". Calling ForOp::build will create such a region and insert
+    `scf.yield`. Calling ForOp::build will create such a region and insert
     the terminator implicitly if none is defined, so will the parsing even in
     cases when it is absent from the custom format. For example:
 
     ```mlir
+    // Index case.
     scf.for %iv = %lb to %ub step %step {
       ... // body
     }
+    ...
+    // Integer case.
+    scf.for %iv_32 = %lb_32 to %ub_32 step %step_32 : i32 {
+      ... // body
+    }
     ```
 
     `scf.for` can also operate on loop-carried variables and returns the final
     values after loop termination. The initial values of the variables are
-    passed as additional SSA operands to the "scf.for" following the 3 loop
+    passed as additional SSA operands to the `scf.for` following the 3 loop
     control SSA values mentioned above (lower bound, upper bound and step). The
     operation region has an argument for the induction variable, followed by
     one argument for each loop-carried variable, representing the value of the
     variable at the current iteration.
 
-    The region must terminate with a "scf.yield" that passes the current
+    The region must terminate with a `scf.yield` that passes the current
     values of all loop-carried variables to the next iteration, or to the
-    "scf.for" result, if at the last iteration. The static type of a
+    `scf.for` result, if at the last iteration. The static type of a
     loop-carried variable may not change with iterations; its runtime type is
     allowed to change. Note, that when the loop-carried variables are present,
     calling ForOp::build will not insert the terminator implicitly. The caller
-    must insert "scf.yield" in that case.
+    must insert `scf.yield` in that case.
 
-    "scf.for" results hold the final values after the last iteration.
+    `scf.for` results hold the final values after the last iteration.
     For example, to sum-reduce a memref:
 
     ```mlir
     func.func @reduce(%buffer: memref<1024xf32>, %lb: index,
                       %ub: index, %step: index) -> (f32) {
       // Initial sum set to 0.
       %sum_0 = arith.constant 0.0 : f32
       // iter_args binds initial values to the loop's region arguments.
       %sum = scf.for %iv = %lb to %ub step %step
           iter_args(%sum_iter = %sum_0) -> (f32) {
         %t = load %buffer[%iv] : memref<1024xf32>
         %sum_next = arith.addf %sum_iter, %t : f32
         // Yield current iteration sum to next iteration %sum_iter or to %sum
         // if final iteration.
         scf.yield %sum_next : f32
       }
       return %sum : f32
     }
     ```
 
-    If the "scf.for" defines any values, a yield must be explicitly present.
-    The number and types of the "scf.for" results must match the initial
-    values in the "iter_args" binding and the yield operands.
+    If the `scf.for` defines any values, a yield must be explicitly present.
+    The number and types of the `scf.for` results must match the initial
+    values in the `iter_args` binding and the yield operands.
 
-    Another example with a nested "scf.if" (see "scf.if" for details) to
+    Another example with a nested `scf.if` (see `scf.if` for details) to
     perform conditional reduction:
 
     ```mlir
     func.func @conditional_reduce(%buffer: memref<1024xf32>, %lb: index,
                                   %ub: index, %step: index) -> (f32) {
       %sum_0 = arith.constant 0.0 : f32
       %c0 = arith.constant 0.0 : f32
       %sum = scf.for %iv = %lb to %ub step %step
           iter_args(%sum_iter = %sum_0) -> (f32) {
         %t = load %buffer[%iv] : memref<1024xf32>
         %cond = arith.cmpf "ugt", %t, %c0 : f32
         %sum_next = scf.if %cond -> (f32) {
           %new_sum = arith.addf %sum_iter, %t : f32
           scf.yield %new_sum : f32
         } else {
           scf.yield %sum_iter : f32
         }
         scf.yield %sum_next : f32
       }
       return %sum : f32
     }
     ```
   }];
-  let arguments = (ins Index:$lowerBound,
-                       Index:$upperBound,
-                       Index:$step,
+  let arguments = (ins AnySignlessIntegerOrIndex:$lowerBound,
+                       AnySignlessIntegerOrIndex:$upperBound,
+                       AnySignlessIntegerOrIndex:$step,
                        Variadic<AnyType>:$initArgs);
   let results = (outs Variadic<AnyType>:$results);
   let regions = (region SizedRegion<1>:$region);
 
   let skipDefaultBuilders = 1;
   let builders = [
     OpBuilder<(ins "Value":$lowerBound, "Value":$upperBound, "Value":$step,
       CArg<"ValueRange", "std::nullopt">:$iterArgs,

mlir/lib/Dialect/SCF/IR/SCF.cpp

 
 void ForOp::build(OpBuilder &builder, OperationState &result, Value lb,
                   Value ub, Value step, ValueRange iterArgs,
                   BodyBuilderFn bodyBuilder) {
   result.addOperands({lb, ub, step});
   result.addOperands(iterArgs);
   for (Value v : iterArgs)
     result.addTypes(v.getType());
+  Type t = lb.getType();
   Region *bodyRegion = result.addRegion();
   bodyRegion->push_back(new Block);
   Block &bodyBlock = bodyRegion->front();
-  bodyBlock.addArgument(builder.getIndexType(), result.location);
+  bodyBlock.addArgument(t, result.location);
   for (Value v : iterArgs)
     bodyBlock.addArgument(v.getType(), v.getLoc());
 
   // Create the default terminator if the builder is not provided and if the
   // iteration arguments are not provided. Otherwise, leave this to the caller
   // because we don't know which values to return from the loop.
   if (iterArgs.empty() && !bodyBuilder) {
     ForOp::ensureTerminator(*bodyRegion, builder, result.location);
if (getNumIterOperands() != opNumResults)
     return emitOpError(
         "mismatch in number of loop-carried values and defined values");
   return success();
 }
 
 LogicalResult ForOp::verifyRegions() {
   // Check that the body defines as single block argument for the induction
   // variable.
-  auto *body = getBody();
-  if (!body->getArgument(0).getType().isIndex())
+  if (getInductionVar().getType() != getLowerBound().getType())
     return emitOpError(
-        "expected body first argument to be an index argument for "
-        "the induction variable");
+        "expected induction variable to be same type as bounds and step");
 
   auto opNumResults = getNumResults();
   if (opNumResults == 0)
     return success();
 
   if (getNumRegionIterArgs() != opNumResults)
     return emitOpError(
         "mismatch in number of basic block args and defined values");
 
   auto iterOperands = getIterOperands();
   auto iterArgs = getRegionIterArgs();
   auto opResults = getResults();
   unsigned i = 0;
   for (auto e : llvm::zip(iterOperands, iterArgs, opResults)) {
     if (std::get<0>(e).getType() != std::get<2>(e).getType())
       return emitOpError() << "types mismatch between " << i
                            << "th iter operand and defined value";
     if (std::get<1>(e).getType() != std::get<2>(e).getType())
       return emitOpError() << "types mismatch between " << i
                            << "th iter region arg and defined value";
 
-    i++;
+    ++i;
   }
   return success();
 }
 
 std::optional<Value> ForOp::getSingleInductionVar() {
   return getInductionVar();
 }
 
void ForOp::print(OpAsmPrinter &p) {
   p << " " << getInductionVar() << " = " << getLowerBound() << " to "
     << getUpperBound() << " step " << getStep();
 
   printInitializationList(p, getRegionIterArgs(), getIterOperands(),
                           " iter_args");
   if (!getIterOperands().empty())
     p << " -> (" << getIterOperands().getTypes() << ')';
   p << ' ';
+  if (Type t = getInductionVar().getType(); !t.isIndex())
+    p << " : " << t << ' ';
   p.printRegion(getRegion(),
                 /*printEntryBlockArgs=*/false,
                 /*printBlockTerminators=*/hasIterOperands());
   p.printOptionalAttrDict((*this)->getAttrs());
 }
 
 ParseResult ForOp::parse(OpAsmParser &parser, OperationState &result) {
   auto &builder = parser.getBuilder();
-  Type indexType = builder.getIndexType();
+  Type type;
 
   OpAsmParser::Argument inductionVariable;
-  inductionVariable.type = indexType;
   OpAsmParser::UnresolvedOperand lb, ub, step;
 
   // Parse the induction variable followed by '='.
-  if (parser.parseArgument(inductionVariable) || parser.parseEqual() ||
+  if (parser.parseOperand(inductionVariable.ssaName) || parser.parseEqual() ||
       // Parse loop bounds.
-      parser.parseOperand(lb) ||
-      parser.resolveOperand(lb, indexType, result.operands) ||
-      parser.parseKeyword("to") || parser.parseOperand(ub) ||
-      parser.resolveOperand(ub, indexType, result.operands) ||
-      parser.parseKeyword("step") || parser.parseOperand(step) ||
-      parser.resolveOperand(step, indexType, result.operands))
+      parser.parseOperand(lb) || parser.parseKeyword("to") ||
+      parser.parseOperand(ub) || parser.parseKeyword("step") ||
+      parser.parseOperand(step))
+    return failure();
+  // Parse optional type, else assume Index.
+  if (parser.parseOptionalColon())
+    type = builder.getIndexType();
+  else if (parser.parseType(type))
+    return failure();
+  inductionVariable.type = type;
+  if (parser.resolveOperand(lb, type, result.operands) ||
+      parser.resolveOperand(ub, type, result.operands) ||
+      parser.resolveOperand(step, type, result.operands))
     return failure();
 
   // Parse the optional initial iteration arguments.
   SmallVector<OpAsmParser::Argument, 4> regionArgs;
   SmallVector<OpAsmParser::UnresolvedOperand, 4> operands;
   regionArgs.push_back(inductionVariable);
 
   if (succeeded(parser.parseOptionalKeyword("iter_args"))) {

mlir/test/Dialect/SCF/invalid.mlir

 // RUN: mlir-opt -allow-unregistered-dialect %s -split-input-file -verify-diagnostics
 
 func.func @loop_for_lb(%arg0: f32, %arg1: index) {
-  // expected-error@+1 {{operand #0 must be index}}
+  // expected-error@+1 {{operand #0 must be signless integer or index}}
   "scf.for"(%arg0, %arg1, %arg1) ({}) : (f32, index, index) -> ()
   return
 }
 
 // -----
 
 func.func @loop_for_ub(%arg0: f32, %arg1: index) {
-  // expected-error@+1 {{operand #1 must be index}}
+  // expected-error@+1 {{operand #1 must be signless integer or index}}
   "scf.for"(%arg1, %arg0, %arg1) ({}) : (index, f32, index) -> ()
   return
 }
 
 // -----
 
 func.func @loop_for_step(%arg0: f32, %arg1: index) {
-  // expected-error@+1 {{operand #2 must be index}}
+  // expected-error@+1 {{operand #2 must be signless integer or index}}
   "scf.for"(%arg1, %arg1, %arg0) ({}) : (index, index, f32) -> ()
   return
 }
 
 // -----
 
+func.func @loop_for_mismatch(%arg0: i32, %arg1: index) {
+  // expected-error@+1 {{all of {lowerBound, upperBound, step} have same type}}
+  "scf.for"(%arg1, %arg0, %arg1) ({}) : (index, i32, index) -> ()
+  return
+}
+
+// -----
+
 func.func @loop_for_step_positive(%arg0: index) {
   // expected-error@+2 {{constant step operand must be positive}}
   %c0 = arith.constant 0 : index
   "scf.for"(%arg0, %arg0, %c0) ({
     ^bb0(%arg1: index):
       scf.yield
   }) : (index, index, index) -> ()
   return
"scf.for"(%arg0, %arg0, %arg0) (
     }
   ) : (index, index, index) -> ()
   return
 }
 
 // -----
 
 func.func @loop_for_single_index_argument(%arg0: index) {
-  // expected-error@+1 {{op expected body first argument to be an index argument for the induction variable}}
+  // expected-error@+1 {{expected induction variable to be same type as bounds}}
   "scf.for"(%arg0, %arg0, %arg0) (
     {
     ^bb0(%i0 : f32):
       scf.yield
     }
   ) : (index, index, index) -> ()
   return
 }

mlir/test/Dialect/SCF/ops.mlir

 //  CHECK-NEXT:   scf.for %{{.*}} = %{{.*}} to %{{.*}} step %{{.*}} {
 //  CHECK-NEXT:     scf.for %{{.*}} = %{{.*}} to %{{.*}} step %{{.*}} {
 //  CHECK-NEXT:       %{{.*}} = arith.cmpi slt, %{{.*}}, %{{.*}} : index
 //  CHECK-NEXT:       %{{.*}} = arith.select %{{.*}}, %{{.*}}, %{{.*}} : index
 //  CHECK-NEXT:       %{{.*}} = arith.cmpi sge, %{{.*}}, %{{.*}} : index
 //  CHECK-NEXT:       %{{.*}} = arith.select %{{.*}}, %{{.*}}, %{{.*}} : index
 //  CHECK-NEXT:       scf.for %{{.*}} = %{{.*}} to %{{.*}} step %{{.*}} {
 
+func.func @std_for_i32(%arg0 : i32, %arg1 : i32, %arg2 : i32) {
+  scf.for %i0 = %arg0 to %arg1 step %arg2 : i32 {
+    scf.for %i1 = %arg0 to %arg1 step %arg2 : i32 {
+    }
+  }
+  return
+}
+// CHECK-LABEL: func @std_for_i32(
+//  CHECK-NEXT:   scf.for %{{.*}} = %{{.*}} to %{{.*}} step %{{.*}} : i32 {
+//  CHECK-NEXT:     scf.for %{{.*}} = %{{.*}} to %{{.*}} step %{{.*}} : i32 {
+
 func.func @std_if(%arg0: i1, %arg1: f32) {
   scf.if %arg0 {
     %0 = arith.addf %arg1, %arg1 : f32
   }
   return
 }
 // CHECK-LABEL: func @std_if(
 //  CHECK-NEXT:   scf.if %{{.*}} {