[MLIR] Add support for empty IVs to affine.parallel

[flaub]

Summary

Allow support for specifying empty IVs in an affine.parallel.

For example:

affine.parallel () = () to () {
  affine.yield
}

[bondhugula]

LGTM. Just a request to update documentation.

mlir/test/Dialect/Affine/ops.mlir

187–189

I think the documentation needs to be updated to say this is allowed.

[flaub]

Update docs

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

 }
 
 def AffineParallelOp : Affine_Op<"parallel",
     [ImplicitAffineTerminator, RecursiveSideEffects,
      DeclareOpInterfaceMethods<LoopLikeOpInterface>, MemRefsNormalizable]> {
   let summary = "multi-index parallel band operation";
   let description = [{
     The "affine.parallel" operation represents a hyper-rectangular affine
-    parallel band, defining multiple SSA values for its induction variables. It
-    has one region capturing the parallel band body. The induction variables are
-    represented as arguments of this region. These SSA values always have type
-    index, which is the size of the machine word. The strides, represented by
-    steps, are positive constant integers 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. The body region must
-    contain exactly one block that terminates with "affine.yield".
+    parallel band, defining zero or more SSA values for its induction variables.
+    It has one region capturing the parallel band body. The induction variables
+    are represented as arguments of this region. These SSA values always have
+    type index, which is the size of the machine word. The strides, represented
+    by steps, are positive constant integers 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. The body
+    region must contain exactly one block that terminates with "affine.yield".
 
     The lower and upper bounds of a parallel operation are represented as an
     application of an affine mapping to a list of SSA values passed to the map.
     The same restrictions hold for these SSA values as for all bindings of SSA
     values to dimensions and symbols. The list of expressions in each map is
     interpreted according to the respective bounds group attribute. If a single
     expression belongs to the group, then the result of this expression is taken
     as a lower(upper) bound of the corresponding loop induction variable. If

mlir/lib/Dialect/Affine/IR/AffineOps.cpp

 }
 
 void AffineParallelOp::build(OpBuilder &builder, OperationState &result,
                              TypeRange resultTypes,
                              ArrayRef<AtomicRMWKind> reductions,
                              ArrayRef<AffineMap> lbMaps, ValueRange lbArgs,
                              ArrayRef<AffineMap> ubMaps, ValueRange ubArgs,
                              ArrayRef<int64_t> steps) {
-  assert(!lbMaps.empty() && "expected the lower bound map to be non-empty");
-  assert(!ubMaps.empty() && "expected the upper bound map to be non-empty");
   assert(llvm::all_of(lbMaps,
                       [lbMaps](AffineMap m) {
                         return m.getNumDims() == lbMaps[0].getNumDims() &&
                                m.getNumSymbols() == lbMaps[0].getNumSymbols();
                       }) &&
          "expected all lower bounds maps to have the same number of dimensions "
          "and symbols");
   assert(llvm::all_of(ubMaps,
                       [ubMaps](AffineMap m) {
                         return m.getNumDims() == ubMaps[0].getNumDims() &&
                                m.getNumSymbols() == ubMaps[0].getNumSymbols();
                       }) &&
          "expected all upper bounds maps to have the same number of dimensions "
          "and symbols");
-  assert(lbMaps[0].getNumInputs() == lbArgs.size() &&
+  assert(lbMaps.empty() ||
+         lbMaps[0].getNumInputs() == lbArgs.size() &&
              "expected lower bound maps to have as many inputs as lower bound "
              "operands");
-  assert(ubMaps[0].getNumInputs() == ubArgs.size() &&
+  assert(ubMaps.empty() ||
+         ubMaps[0].getNumInputs() == ubArgs.size() &&
              "expected upper bound maps to have as many inputs as upper bound "
              "operands");
 
   result.addTypes(resultTypes);
 
   // Convert the reductions to integer attributes.
   SmallVector<Attribute, 4> reductionAttrs;
   for (AtomicRMWKind reduction : reductions)
     reductionAttrs.push_back(
         builder.getI64IntegerAttr(static_cast<int64_t>(reduction)));
   result.addAttribute(getReductionsAttrName(),
                       builder.getArrayAttr(reductionAttrs));
 
   // Concatenates maps defined in the same input space (same dimensions and
   // symbols), assumes there is at least one map.
-  auto concatMapsSameInput = [](ArrayRef<AffineMap> maps,
+  auto concatMapsSameInput = [&builder](ArrayRef<AffineMap> maps,
                                         SmallVectorImpl<int32_t> &groups) {
+    if (maps.empty())
+      return AffineMap::get(builder.getContext());
     SmallVector<AffineExpr> exprs;
     groups.reserve(groups.size() + maps.size());
     exprs.reserve(maps.size());
     for (AffineMap m : maps) {
       llvm::append_range(exprs, m.getResults());
       groups.push_back(m.getNumResults());
     }
-    assert(!maps.empty() && "expected a non-empty list of maps");
     return AffineMap::get(maps[0].getNumDims(), maps[0].getNumSymbols(), exprs,
                           maps[0].getContext());
   };
 
   // Set up the bounds.
   SmallVector<int32_t> lbGroups, ubGroups;
   AffineMap lbMap = concatMapsSameInput(lbMaps, lbGroups);
   AffineMap ubMap = concatMapsSameInput(ubMaps, ubGroups);
   result.addAttribute(getLowerBoundsMapAttrName(), AffineMapAttr::get(lbMap));
   result.addAttribute(getLowerBoundsGroupsAttrName(),
-                      builder.getI32VectorAttr(lbGroups));
+                      builder.getI32TensorAttr(lbGroups));
   result.addAttribute(getUpperBoundsMapAttrName(), AffineMapAttr::get(ubMap));
   result.addAttribute(getUpperBoundsGroupsAttrName(),
-                      builder.getI32VectorAttr(ubGroups));
+                      builder.getI32TensorAttr(ubGroups));
   result.addAttribute(getStepsAttrName(), builder.getI64ArrayAttr(steps));
   result.addOperands(lbArgs);
   result.addOperands(ubArgs);
 
   // Create a region and a block for the body.
   auto *bodyRegion = result.addRegion();
   auto *body = new Block();
   // Add all the block arguments.

 
 OpBuilder AffineParallelOp::getBodyBuilder() {
   return OpBuilder(getBody(), std::prev(getBody()->end()));
 }
 
 void AffineParallelOp::setLowerBounds(ValueRange lbOperands, AffineMap map) {
   assert(lbOperands.size() == map.getNumInputs() &&
          "operands to map must match number of inputs");
-  assert(map.getNumResults() >= 1 && "bounds map has at least one result");
 
   auto ubOperands = getUpperBoundsOperands();
 
   SmallVector<Value, 4> newOperands(lbOperands);
   newOperands.append(ubOperands.begin(), ubOperands.end());
   (*this)->setOperands(newOperands);
 
   lowerBoundsMapAttr(AffineMapAttr::get(map));
 }
 
 void AffineParallelOp::setUpperBounds(ValueRange ubOperands, AffineMap map) {
   assert(ubOperands.size() == map.getNumInputs() &&
          "operands to map must match number of inputs");
-  assert(map.getNumResults() >= 1 && "bounds map has at least one result");
 
   SmallVector<Value, 4> newOperands(getLowerBoundsOperands());
   newOperands.append(ubOperands.begin(), ubOperands.end());
   (*this)->setOperands(newOperands);
 
   upperBoundsMapAttr(AffineMapAttr::get(map));
 }
 
StringRef groupsName = kind == MinMaxKind::Min
                              : AffineParallelOp::getLowerBoundsGroupsAttrName();
 
   if (failed(parser.parseLParen()))
     return failure();
 
   if (succeeded(parser.parseOptionalRParen())) {
     result.addAttribute(
         mapName, AffineMapAttr::get(parser.getBuilder().getEmptyAffineMap()));
-    result.addAttribute(groupsName, parser.getBuilder().getI32VectorAttr({}));
+    result.addAttribute(groupsName, parser.getBuilder().getI32TensorAttr({}));
     return success();
   }
 
   SmallVector<AffineExpr> flatExprs;
   SmallVector<SmallVector<OpAsmParser::OperandType>> flatDimOperands;
   SmallVector<SmallVector<OpAsmParser::OperandType>> flatSymOperands;
   SmallVector<int32_t> numMapsPerGroup;
   SmallVector<OpAsmParser::OperandType> mapOperands;
static ParseResult parseAffineMapWithMinMax(OpAsmParser &parser,
 
   Builder &builder = parser.getBuilder();
   auto flatMap = AffineMap::get(totalNumDims, totalNumSyms, flatExprs,
                                 parser.getBuilder().getContext());
   flatMap = flatMap.replaceDimsAndSymbols(
       dimRplacements, symRepacements, dimOperands.size(), symOperands.size());
 
   result.addAttribute(mapName, AffineMapAttr::get(flatMap));
-  result.addAttribute(groupsName, builder.getI32VectorAttr(numMapsPerGroup));
+  result.addAttribute(groupsName, builder.getI32TensorAttr(numMapsPerGroup));
   return success();
 }
 
 //
 // operation ::= `affine.parallel` `(` ssa-ids `)` `=` parallel-bound
 //               `to` parallel-bound steps? region attr-dict?
 // steps     ::= `steps` `(` integer-literals `)`
 //

mlir/test/Dialect/Affine/ops.mlir

affine.parallel (%i, %j, %k) = (max(%a, %b), %b, max(%a, %c))
                               to (%c, min(%c, %d), %b) {
     affine.yield
   }
   return
 }
 
 // -----
 
+// CHECK-LABEL: @parallel_no_ivs
+func @parallel_no_ivs() {
+  // CHECK: affine.parallel () = () to ()
+  affine.parallel () = () to () {
+    affine.yield
+  }
+  return
+}
+
+// -----
+
 // CHECK-LABEL: func @affine_if
 func @affine_if() -> f32 {
   // CHECK: %[[ZERO:.*]] = constant {{.*}} : f32
   %zero = constant 0.0 : f32
   // CHECK: %[[OUT:.*]] = affine.if {{.*}}() -> f32 {
   %0 = affine.if affine_set<() : ()> () -> f32 {
     // CHECK: affine.yield %[[ZERO]] : f32
     affine.yield %zero : f32