[mlir][Linalg] Refactor Linalg creation of loops to allow passing iterArgs - NFC

[nicolasvasilache]

Summary

This revision changes the signatures of helper function that Linalg uses to create loops so that they can also take iterArgs.
iterArgs are asserted empty to ensure no functional change.
This is a mechanical change in preparation of tiling on linalg on tensors to avoid polluting the implementation with an NFC change.

[nicolasvasilache]

Fix compilation.

mlir/include/mlir/Dialect/Linalg/Utils/Utils.h

 /// is used to generate the body of the innermost loop. It is passed a range
 /// of loop induction variables.
 template <typename LoopTy>
 struct GenerateLoopNest {
   using IndexedValueTy =
       typename std::conditional<std::is_same<LoopTy, AffineForOp>::value,
                                 AffineIndexedValue, StdIndexedValue>::type;
 
-  static void doit(ArrayRef<SubViewOp::Range> loopRanges,
-                   ArrayRef<Attribute> iteratorTypes,
-                   function_ref<void(ValueRange)> bodyBuilderFn,
-                   Optional<LinalgLoopDistributionOptions> = None);
+  static void
+  doit(ArrayRef<SubViewOp::Range> loopRanges, ValueRange iterArgInitValues,
+       function_ref<scf::ValueVector(ValueRange, ValueRange)> bodyBuilderFn,
 };
 
 } // namespace linalg
 } // namespace mlir
 
 #endif // MLIR_DIALECT_LINALG_UTILS_H_

mlir/include/mlir/Dialect/SCF/EDSC/Builders.h

 scf::ValueVector loopNestBuilder(ValueRange lbs, ValueRange ubs,
                                  ValueRange steps,
                                  function_ref<void(ValueRange)> fun = nullptr);
 scf::ValueVector loopNestBuilder(Value lb, Value ub, Value step,
                                  function_ref<void(Value)> fun = nullptr);
 scf::ValueVector loopNestBuilder(
     Value lb, Value ub, Value step, ValueRange iterArgInitValues,
     function_ref<scf::ValueVector(Value, ValueRange)> fun = nullptr);
+scf::ValueVector loopNestBuilder(
+    ValueRange lbs, ValueRange ubs, ValueRange steps,
+    ValueRange iterArgInitValues,
+    function_ref<scf::ValueVector(ValueRange, ValueRange)> fun = nullptr);
 
 /// Adapters for building if conditions using the builder and the location
 /// stored in ScopedContext. 'thenBody' is mandatory, 'elseBody' can be omitted
 /// if the condition should not have an 'else' part.
 /// When `ifOp` is specified, the scf::IfOp is captured. This is particularly
 /// convenient for 0-result conditions.
 ValueRange conditionBuilder(TypeRange results, Value condition,
                             function_ref<scf::ValueVector()> thenBody,

mlir/lib/Dialect/Linalg/Transforms/Loops.cpp

Optional<LinalgLoops> linalgOpToLoopsImpl(Operation *op, OpBuilder &builder) {
   auto maps = llvm::to_vector<8>(
       llvm::map_range(mapsRange, [](AffineMapAttr a) { return a.getValue(); }));
   SmallVector<Value, 8> sizes = getViewSizes(builder, linalgOp);
   AffineMap map = concatAffineMaps(maps);
   auto loopRanges = emitLoopRanges(scope.getBuilderRef(), scope.getLocation(),
                                    map, getViewSizes(builder, linalgOp));
   SmallVector<Value, 4> allIvs;
   GenerateLoopNest<LoopTy>::doit(
-      loopRanges, linalgOp.iterator_types().getValue(), [&](ValueRange ivs) {
+      loopRanges, /*iterInitArgs*/ {}, linalgOp.iterator_types().getValue(),
+      [&](ValueRange ivs, ValueRange iterArgs) -> scf::ValueVector {
+        assert(iterArgs.empty() && "unexpected iterArgs");
         allIvs.append(ivs.begin(), ivs.end());
         emitScalarImplementation<IndexedValueTy>(allIvs, linalgOp);
+        return scf::ValueVector{};
       });
   // Number of loop ops might be different from the number of ivs since some
   // loops like affine.parallel and scf.parallel have multiple ivs.
   llvm::SetVector<Operation *> loopSet;
   for (Value iv : allIvs) {
     if (!iv)
       return {};
     // The induction variable is a block argument of the entry block of the

mlir/lib/Dialect/Linalg/Transforms/Tiling.cpp

Optional<TiledLinalgOp> static tileLinalgOpImpl(
   // 3. Create the tiled loops.
   LinalgOp res = op;
   SmallVector<Value, 4> ivs;
   SmallVector<Attribute, 4> iteratorTypes =
       llvm::to_vector<4>(op.iterator_types().cast<ArrayAttr>().getValue());
   if (!options.interchangeVector.empty())
     applyPermutationToVector(iteratorTypes, options.interchangeVector);
   GenerateLoopNest<LoopTy>::doit(
-      loopRanges, iteratorTypes,
-      [&](ValueRange localIvs) {
+      loopRanges, /*iterArgInitValues*/ {}, iteratorTypes,
+      [&](ValueRange localIvs, ValueRange iterArgs) -> scf::ValueVector {
         auto &b = ScopedContext::getBuilderRef();
         auto loc = ScopedContext::getLocation();
         ivs.assign(localIvs.begin(), localIvs.end());
         SmallVector<Value, 4> ivValues(ivs.begin(), ivs.end());
 
         // If we have to apply a permutation to the tiled loop nest, we have to
         // reorder the induction variables This permutation is the right one
         // assuming that loopRanges have previously been permuted by
         // (i,j,k)->(k,i,j) So this permutation should be the inversePermutation
         // of that one: (d0,d1,d2)->(d2,d0,d1)
         if (!options.interchangeVector.empty())
           ivValues = applyMapToValues(b, loc, invPermutationMap, ivValues);
 
         auto views = makeTiledViews(b, loc, op, viewSizesToLoopsMap, ivValues,
                                     tileSizes, allViewSizes);
         auto operands = getAssumedNonViewOperands(op);
         views.append(operands.begin(), operands.end());
         res = op.clone(b, loc, views);
+        return scf::ValueVector{};
       },
       options.distribution);
 
   // 4. Transforms index arguments of `linalg.generic` w.r.t. to the tiling.
   transformIndexedGenericOpIndices(b, res, ivs, loopIndexToRangeIndex);
 
   // 5. Gather the newly created loops and return them with the new op.
   SmallVector<Operation *, 8> loops;

mlir/lib/Dialect/Linalg/Utils/Utils.cpp

if (!invertedMap)
     return {};
   return applyMapToValues(builder, linalgOp.getLoc(), invertedMap, viewSizes,
                           folder);
 }
 
 /// Specialization to build an scf "for" nest.
 template <>
 void GenerateLoopNest<scf::ForOp>::doit(
-    ArrayRef<SubViewOp::Range> loopRanges, ArrayRef<Attribute> iteratorTypes,
-    function_ref<void(ValueRange)> bodyBuilderFn,
+    ArrayRef<SubViewOp::Range> loopRanges, ValueRange iterArgInitValues,
+    ArrayRef<Attribute> iteratorTypes,
+    function_ref<scf::ValueVector(ValueRange, ValueRange)> bodyBuilderFn,
     Optional<LinalgLoopDistributionOptions>) {
   SmallVector<Value, 4> lbs, ubs, steps;
   unpackRanges(loopRanges, lbs, ubs, steps);
-  edsc::loopNestBuilder(lbs, ubs, steps, bodyBuilderFn);
+  edsc::loopNestBuilder(lbs, ubs, steps, iterArgInitValues, bodyBuilderFn);
 }
 
 /// Specialization to build affine "for" nest.
 template <>
 void GenerateLoopNest<AffineForOp>::doit(
-    ArrayRef<SubViewOp::Range> loopRanges, ArrayRef<Attribute> iteratorTypes,
-    function_ref<void(ValueRange)> bodyBuilderFn,
+    ArrayRef<SubViewOp::Range> loopRanges, ValueRange iterArgInitValues,
+    ArrayRef<Attribute> iteratorTypes,
+    function_ref<scf::ValueVector(ValueRange, ValueRange)> bodyBuilderFn,
     Optional<LinalgLoopDistributionOptions>) {
+  assert(iterArgInitValues.empty() && "unexpected AffineForOp init values");
   SmallVector<Value, 4> lbs, ubs, steps;
   unpackRanges(loopRanges, lbs, ubs, steps);
 
   // Affine loops require constant steps.
   SmallVector<int64_t, 4> constantSteps;
   constantSteps.reserve(steps.size());
   for (Value v : steps) {
     auto op = v.getDefiningOp<ConstantIndexOp>();
     assert(op && "Affine loops require constant steps");
     constantSteps.push_back(op.getValue());
   }
 
-  edsc::affineLoopNestBuilder(lbs, ubs, constantSteps, bodyBuilderFn);
+  auto bodyBuilderWithoutIterArgsFn = [&](ValueRange ivs) {
+    bodyBuilderFn(ivs, {});
+  };
+  edsc::affineLoopNestBuilder(lbs, ubs, constantSteps,
+                              bodyBuilderWithoutIterArgsFn);
 }
 
 /// Update the `lb`, `ub` and `step` to get per processor `lb`, `ub` and `step`.
 static void updateBoundsForCyclicDistribution(OpBuilder &builder, Location loc,
                                               Value procId, Value nprocs,
                                               Value &lb, Value &ub,
                                               Value &step) {
   using edsc::op::operator+;
generateParallelLoopNest(
         bodyBuilderFn, ivStorage, distributionMethod.drop_front(numProcessed));
     return;
   }
 }
 
 /// Specialization for generating a mix of parallel and sequential scf loops.
 template <>
 void GenerateLoopNest<scf::ParallelOp>::doit(
-    ArrayRef<SubViewOp::Range> loopRanges, ArrayRef<Attribute> iteratorTypes,
-    function_ref<void(ValueRange)> bodyBuilderFn,
+    ArrayRef<SubViewOp::Range> loopRanges, ValueRange iterArgInitValues,
+    ArrayRef<Attribute> iteratorTypes,
+    function_ref<scf::ValueVector(ValueRange, ValueRange)> bodyBuilderFn,
     Optional<LinalgLoopDistributionOptions> distributionOptions) {
+  assert(iterArgInitValues.empty() && "unexpected ParallelOp init values");
   // This function may be passed more iterator types than ranges.
   assert(iteratorTypes.size() >= loopRanges.size() &&
          "expected iterator type for all ranges");
   iteratorTypes = iteratorTypes.take_front(loopRanges.size());
   SmallVector<Value, 8> lbsStorage, ubsStorage, stepsStorage, ivs;
   unsigned numLoops = iteratorTypes.size();
   ivs.reserve(numLoops);
   lbsStorage.reserve(numLoops);
for (auto iteratorType : enumerate(iteratorTypes)) {
         updateBoundsForCyclicDistribution(builder, loc, procInfo[index].procId,
                                           procInfo[index].nprocs, lbsStorage[i],
                                           ubsStorage[i], stepsStorage[i]);
         index++;
       }
     }
   }
   ValueRange lbs(lbsStorage), ubs(ubsStorage), steps(stepsStorage);
-  generateParallelLoopNest(lbs, ubs, steps, iteratorTypes, bodyBuilderFn, ivs,
+  auto bodyBuilderWithoutIterArgsFn = [&](ValueRange ivs) {
+    bodyBuilderFn(ivs, {});
+  };
+  generateParallelLoopNest(lbs, ubs, steps, iteratorTypes,
+                           bodyBuilderWithoutIterArgsFn, ivs,
                            distributionMethod);
 
   assert(ivs.size() == iteratorTypes.size() && "did not generate enough loops");
 }
 
 } // namespace linalg
 } // namespace mlir

mlir/lib/Dialect/SCF/EDSC/Builders.cpp

return mlir::scf::buildLoopNest(
         ScopedContext context(builder, loc);
         if (fun)
           return fun(ivs[0], args);
         return scf::ValueVector(iterArgInitValues.begin(),
                                 iterArgInitValues.end());
       });
 }
 
+mlir::scf::ValueVector mlir::edsc::loopNestBuilder(
+    ValueRange lbs, ValueRange ubs, ValueRange steps,
+    ValueRange iterArgInitValues,
+    function_ref<scf::ValueVector(ValueRange, ValueRange)> fun) {
+  // Delegates actual construction to scf::buildLoopNest by wrapping `fun` into
+  // the expected function interface.
+  assert(ScopedContext::getContext() && "EDSC ScopedContext not set up");
+  return mlir::scf::buildLoopNest(
+      ScopedContext::getBuilderRef(), ScopedContext::getLocation(), lbs, ubs,
+      steps, iterArgInitValues,
+      [&](OpBuilder &builder, Location loc, ValueRange ivs, ValueRange args) {
+        ScopedContext context(builder, loc);
+        if (fun)
+          return fun(ivs, args);
+        return scf::ValueVector(iterArgInitValues.begin(),
+                                iterArgInitValues.end());
+      });
+}
+
 static std::function<void(OpBuilder &, Location)>
 wrapIfBody(function_ref<scf::ValueVector()> body, TypeRange expectedTypes) {
   (void)expectedTypes;
   return [=](OpBuilder &builder, Location loc) {
     ScopedContext context(builder, loc);
     scf::ValueVector returned = body();
     assert(ValueRange(returned).getTypes() == expectedTypes &&
            "'if' body builder returned values of unexpected type");