[mlir][Linalg] NFC - Simplify GenericNestLoop builder

[nicolasvasilache]

Summary

This revision trims unnecessary complexity.

[ftynse]

So red, much cleanup!

[ftynse]

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

39

This looks a bit fishy: why does it change from ParallelOp to ForOp?

[nicolasvasilache]

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

39

I blame phab's diff, this is the ForOp version, the ParallelFor is below.

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

 class SubViewOp;
 
 namespace scf {
 class ParallelOp;
 } // namespace scf
 
 namespace edsc {
 class AffineLoopNestBuilder;
+class LoopNestBuilder;
 class ParallelLoopNestBuilder;
 
-/// A LoopRangeBuilder is a generic NestedBuilder for scf.for operations.
-/// More specifically it is meant to be used as a temporary object for
-/// representing any nested MLIR construct that is "related to" an mlir::Value
-/// (for now an induction variable).
-class LoopRangeBuilder : public NestedBuilder {
-public:
-  /// Constructs a new scf.for and captures the associated induction
-  /// variable. A Value pointer is passed as the first argument and is the
-  /// *only* way to capture the loop induction variable.
-  LoopRangeBuilder(Value *iv, Value range);
-  LoopRangeBuilder(Value *iv, SubViewOp::Range range);
-
-  LoopRangeBuilder(const LoopRangeBuilder &) = delete;
-  LoopRangeBuilder(LoopRangeBuilder &&) = default;
-
-  LoopRangeBuilder &operator=(const LoopRangeBuilder &) = delete;
-  LoopRangeBuilder &operator=(LoopRangeBuilder &&) = default;
-
-  /// The only purpose of this operator is to serve as a sequence point so that
-  /// the evaluation of `fun` (which build IR snippets in a scoped fashion) is
-  /// scoped within a LoopRangeBuilder.
-  Value operator()(std::function<void(void)> fun = nullptr);
-};
-
-/// Helper class to sugar building scf.for loop nests from ranges.
-/// This is similar to edsc::AffineLoopNestBuilder except it works on ranges
-/// directly. In the current implementation it produces scf.for operations.
-class LoopNestRangeBuilder {
-public:
-  LoopNestRangeBuilder(MutableArrayRef<Value> ivs, ArrayRef<Value> ranges);
-  LoopNestRangeBuilder(MutableArrayRef<Value> ivs,
-                       ArrayRef<SubViewOp::Range> ranges);
-  Value operator()(std::function<void(void)> fun = nullptr);
-
-private:
-  SmallVector<LoopRangeBuilder, 4> loops;
-};
-
 /// Helper template class for building scf.for and affine.loop nests from
 /// ranges.
 template <typename LoopTy> class GenericLoopNestRangeBuilder {
 public:
   GenericLoopNestRangeBuilder(MutableArrayRef<Value> ivs,
-                              ArrayRef<Value> ranges);
-  GenericLoopNestRangeBuilder(MutableArrayRef<Value> ivs,
                               ArrayRef<SubViewOp::Range> ranges);
   void operator()(std::function<void(void)> fun = nullptr) { (*builder)(fun); }
 
 private:
   using LoopOrAffineLoopBuilder =
       typename std::conditional_t<std::is_same<LoopTy, AffineForOp>::value,
-                                  AffineLoopNestBuilder, LoopNestRangeBuilder>;
+                                  AffineLoopNestBuilder, LoopNestBuilder>;
   using BuilderType =
       typename std::conditional_t<std::is_same<LoopTy, scf::ParallelOp>::value,
                                   ParallelLoopNestBuilder,
                                   LoopOrAffineLoopBuilder>;
 
   std::unique_ptr<BuilderType> builder;
 };
 

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

 class LoopNestBuilder {
 public:
   LoopNestBuilder(Value *iv, Value lb, Value ub, Value step);
   LoopNestBuilder(Value *iv, Value lb, Value ub, Value step,
                   MutableArrayRef<Value> iterArgsHandles,
                   ValueRange iterArgsInitValues);
   LoopNestBuilder(MutableArrayRef<Value> ivs, ArrayRef<Value> lbs,
                   ArrayRef<Value> ubs, ArrayRef<Value> steps);
-  Operation::result_range operator()(std::function<void(void)> fun = nullptr);
+  ValueRange operator()(std::function<void(void)> fun = nullptr);
 
 private:
   SmallVector<LoopBuilder, 4> loops;
 };
 
 /// Adapters for building loop nests using the builder and the location stored
 /// in ScopedContext. Actual builders are in scf::buildLoopNest.
 scf::ValueVector loopNestBuilder(ValueRange lbs, ValueRange ubs,

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

 #include "mlir/IR/AffineExpr.h"
 
 using namespace mlir;
 using namespace mlir::edsc;
 using namespace mlir::edsc::intrinsics;
 using namespace mlir::linalg;
 using namespace mlir::scf;
 
-mlir::edsc::LoopRangeBuilder::LoopRangeBuilder(Value *iv, Value range) {
-  assert(range.getType() && "expected !linalg.range type");
-  assert(range.getDefiningOp() && "need operations to extract range parts");
-  auto rangeOp = cast<RangeOp>(range.getDefiningOp());
-  auto lb = rangeOp.min();
-  auto ub = rangeOp.max();
-  auto step = rangeOp.step();
-  ForOp forOp = OperationBuilder<ForOp>(lb, ub, step);
-  *iv = forOp.getInductionVar();
-  auto *body = forOp.getBody();
-  enter(body);
-}
-
-mlir::edsc::LoopRangeBuilder::LoopRangeBuilder(Value *iv,
-                                               SubViewOp::Range range) {
-  ForOp forOp = OperationBuilder<ForOp>(range.offset, range.size, range.stride);
-  *iv = forOp.getInductionVar();
-  auto *body = forOp.getBody();
-  enter(body);
-}
-
-Value mlir::edsc::LoopRangeBuilder::operator()(std::function<void(void)> fun) {
-  if (fun)
-    fun();
-  exit();
-  return Value();
-}
-
-mlir::edsc::LoopNestRangeBuilder::LoopNestRangeBuilder(
-    MutableArrayRef<Value> ivs, ArrayRef<SubViewOp::Range> ranges) {
-  loops.reserve(ranges.size());
-  for (unsigned i = 0, e = ranges.size(); i < e; ++i)
-    loops.emplace_back(&ivs[i], ranges[i]);
-  assert(loops.size() == ivs.size() && "Mismatch loops vs ivs size");
-}
-
-mlir::edsc::LoopNestRangeBuilder::LoopNestRangeBuilder(
-    MutableArrayRef<Value> ivs, ArrayRef<Value> ranges) {
-  loops.reserve(ranges.size());
-  for (unsigned i = 0, e = ranges.size(); i < e; ++i)
-    loops.emplace_back(&ivs[i], ranges[i]);
-  assert(loops.size() == ivs.size() && "Mismatch loops vs ivs size");
-}
-
-Value LoopNestRangeBuilder::LoopNestRangeBuilder::operator()(
-    std::function<void(void)> fun) {
-  if (fun)
-    fun();
-  for (auto &lit : reverse(loops)) {
-    lit({});
-  }
-  return Value();
-}
-
-namespace mlir {
-namespace edsc {
-
-static void unpackRanges(ArrayRef<Value> rangeOps, SmallVectorImpl<Value> &lbs,
-                         SmallVectorImpl<Value> &ubs,
-                         SmallVectorImpl<Value> &steps) {
-  for (Value range : rangeOps) {
-    assert(range.getType() && "expected linalg.range type");
-    assert(range.getDefiningOp() && "need operations to extract range parts");
-    RangeOp rangeOp = cast<RangeOp>(range.getDefiningOp());
-    lbs.emplace_back(rangeOp.min());
-    ubs.emplace_back(rangeOp.max());
-    steps.emplace_back(rangeOp.step());
-  }
-}
-
 static void unpackRanges(ArrayRef<SubViewOp::Range> ranges,
                          SmallVectorImpl<Value> &lbs,
                          SmallVectorImpl<Value> &ubs,
                          SmallVectorImpl<Value> &steps) {
   for (SubViewOp::Range range : ranges) {
     lbs.emplace_back(range.offset);
     ubs.emplace_back(range.size);
     steps.emplace_back(range.stride);
   }
 }
 
-template <>
-GenericLoopNestRangeBuilder<scf::ForOp>::GenericLoopNestRangeBuilder(
-    MutableArrayRef<Value> ivs, ArrayRef<SubViewOp::Range> ranges) {
-  builder = std::make_unique<LoopNestRangeBuilder>(ivs, ranges);
-}
-
-template <>
-GenericLoopNestRangeBuilder<AffineForOp>::GenericLoopNestRangeBuilder(
-    MutableArrayRef<Value> ivs, ArrayRef<SubViewOp::Range> ranges) {
-  SmallVector<Value, 4> lbs, ubs, steps;
-  unpackRanges(ranges, lbs, ubs, 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());
-  }
-  builder =
-      std::make_unique<AffineLoopNestBuilder>(ivs, lbs, ubs, constantSteps);
-}
+namespace mlir {
+namespace edsc {
 
 template <>
-GenericLoopNestRangeBuilder<scf::ParallelOp>::GenericLoopNestRangeBuilder(
+GenericLoopNestRangeBuilder<scf::ForOp>::GenericLoopNestRangeBuilder(
     MutableArrayRef<Value> ivs, ArrayRef<SubViewOp::Range> ranges) {
   SmallVector<Value, 4> lbs, ubs, steps;
   unpackRanges(ranges, lbs, ubs, steps);
-  builder = std::make_unique<ParallelLoopNestBuilder>(ivs, lbs, ubs, steps);
-}
-
-template <>
-GenericLoopNestRangeBuilder<scf::ForOp>::GenericLoopNestRangeBuilder(
-    MutableArrayRef<Value> ivs, ArrayRef<Value> ranges) {
-  builder = std::make_unique<LoopNestRangeBuilder>(ivs, ranges);
+  builder = std::make_unique<LoopNestBuilder>(ivs, lbs, ubs, steps);
 }
 
 template <>
 GenericLoopNestRangeBuilder<AffineForOp>::GenericLoopNestRangeBuilder(
-    MutableArrayRef<Value> ivs, ArrayRef<Value> ranges) {
+    MutableArrayRef<Value> ivs, ArrayRef<SubViewOp::Range> ranges) {
   SmallVector<Value, 4> lbs, ubs, steps;
   unpackRanges(ranges, lbs, ubs, 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());
   }
   builder =
       std::make_unique<AffineLoopNestBuilder>(ivs, lbs, ubs, constantSteps);
 }
 
 template <>
 GenericLoopNestRangeBuilder<scf::ParallelOp>::GenericLoopNestRangeBuilder(
-    MutableArrayRef<Value> ivs, ArrayRef<Value> ranges) {
+    MutableArrayRef<Value> ivs, ArrayRef<SubViewOp::Range> ranges) {
   SmallVector<Value, 4> lbs, ubs, steps;
   unpackRanges(ranges, lbs, ubs, steps);
   builder = std::make_unique<ParallelLoopNestBuilder>(ivs, lbs, ubs, steps);
 }
 
 } // namespace edsc
 } // namespace mlir
 

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

return permutation ? applyMapToValues(ScopedContext::getBuilderRef(),
                                         ScopedContext::getLocation(),
                                         permutation.getValue(), ivs)
                      : SmallVector<Value, 4>(ivs.begin(), ivs.end());
 }
 
 // Creates a number of ranges equal to the number of results in `map`.
 // The returned ranges correspond to the loop ranges, in the proper order, for
 // which new loops will be created.
-static SmallVector<Value, 4> emitLoopRanges(OpBuilder &b, Location loc,
-                                            AffineMap map,
-                                            ArrayRef<Value> allViewSizes);
-SmallVector<Value, 4> emitLoopRanges(OpBuilder &b, Location loc, AffineMap map,
-                                     ArrayRef<Value> allViewSizes) {
+static SmallVector<SubViewOp::Range, 4>
+emitLoopRanges(OpBuilder &b, Location loc, AffineMap map,
   // Apply `map` to get view sizes in loop order.
   auto sizes = applyMapToValues(b, loc, map, allViewSizes);
   // Create a new range with the applied tile sizes.
   ScopedContext scope(b, loc);
-  SmallVector<Value, 4> res;
+  SmallVector<SubViewOp::Range, 4> res;
   for (unsigned idx = 0, e = map.getNumResults(); idx < e; ++idx) {
-    res.push_back(
-        linalg_range(std_constant_index(0), sizes[idx], std_constant_index(1)));
+    res.push_back(SubViewOp::Range{std_constant_index(0), sizes[idx],
+                                   std_constant_index(1)});
   }
   return res;
 }
 
 template <typename IndexedValueType, typename OpType>
 static void inlineRegionAndEmitStore(OpType op, ArrayRef<Value> indexedValues,
                                      ArrayRef<SmallVector<Value, 8>> indexing,
                                      ArrayRef<Value> outputBuffers) {

 /// Helper struct to generate the loop nest for the op. This factored out here
 /// to be able to partially specialize this for different LoopTy.
 template <typename LoopTy, typename ConcreteOpTy>
 class GenerateLoopNest {
 public:
   using IndexedValueTy =
       typename std::conditional<std::is_same<LoopTy, AffineForOp>::value,
                                 AffineIndexedValue, StdIndexedValue>::type;
-  static void doit(ConcreteOpTy linalgOp, ArrayRef<Value> loopRanges,
+  static void doit(ConcreteOpTy linalgOp, ArrayRef<SubViewOp::Range> loopRanges,
                    MutableArrayRef<Value> allIvs) {
     GenericLoopNestRangeBuilder<LoopTy>(allIvs, loopRanges)([&] {
       SmallVector<Value, 4> allIvValues(allIvs.begin(), allIvs.end());
       LinalgScopedEmitter<IndexedValueTy,
                           ConcreteOpTy>::emitScalarImplementation(allIvValues,
                                                                   linalgOp);
     });
   }
 };
 
 /// Generates loop nest using scf.parallel. scf.parallel is only used for the
 /// outer parallel loops. All other loops are generated using scf.for
 /// operation.
 template <typename ConcreteOpTy>
 class GenerateLoopNest<scf::ParallelOp, ConcreteOpTy> {
 public:
   using IndexedValueTy = StdIndexedValue;
 
-  static void doit(ConcreteOpTy linalgOp, ArrayRef<Value> loopRanges,
+  static void doit(ConcreteOpTy linalgOp, ArrayRef<SubViewOp::Range> loopRanges,
                    MutableArrayRef<Value> allIvs) {
     // Only generate scf.parallel for outer consecutive "parallel"
     // iterator_types.
     // TODO(ravishankarm): Generate scf.parallel for all "parallel" iterator
     // types, not just the outer most ones. Also handle "reduction" iterator
     // types.
     auto nOuterPar = linalgOp.iterator_types()
                          .getValue()

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

 }
 
 mlir::edsc::LoopNestBuilder::LoopNestBuilder(Value *iv, Value lb, Value ub,
                                              Value step) {
   SmallVector<Value, 0> noArgs;
   loops.emplace_back(makeLoopBuilder(iv, lb, ub, step, noArgs, {}));
 }
 
-Operation::result_range
-mlir::edsc::LoopNestBuilder::LoopNestBuilder::operator()(
+ValueRange mlir::edsc::LoopNestBuilder::LoopNestBuilder::operator()(
     std::function<void(void)> fun) {
   if (fun)
     fun();
 
   for (auto &lit : reverse(loops))
     lit({});
 
+  if (!loops.empty())
     return loops[0].getOp()->getResults();
+  return {};
 }
 
 LoopBuilder mlir::edsc::makeParallelLoopBuilder(MutableArrayRef<Value> ivs,
                                                 ArrayRef<Value> lbs,
                                                 ArrayRef<Value> ubs,
                                                 ArrayRef<Value> steps) {
   scf::ParallelOp parallelOp = OperationBuilder<scf::ParallelOp>(
       SmallVector<Value, 4>(lbs.begin(), lbs.end()),