[mlir][Linalg] Avoid using scf.parallel for non-parallel loops in Linalg ops.

[mravishankar]

Summary

Modifying the loop nest builder for generating scf.parallel loops to
not generate scf.parallel loops for non-parallel iterator types in
Linalg operations. The existing implementation incorrectly generated
scf.parallel for all tiled loops. It is rectified by refactoring logic
used while lowering to loops that accounted for this.

[nicolasvasilache]

I think there are a few issues at this time.
The GenericBuilder is introduced by Linalg because that is the place in spacetime that uses both affine and scf.
This CL ties the implementation of GenericBuilder to more Linalg internals which leaks more abstractions.
I have https://reviews.llvm.org/D80290 which retires a bunch of now unused stuff.
This should make this CL simpler after rebasing.

On the change itself, this does not seem to support interleaving.
I think I would rather try to go for a recursive helper function that takes ArrayRef slices and a leaf lambda but otherwise leaves GenericLoopNestBuilder untouched.
This should also evolve better when we have more loop types.

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

387

you also want to apply the interchange vector with applyPermutationToVector

mlir/test/Dialect/Linalg/tile_parallel_reduce.mlir

23

nit: newline

[mravishankar]

Addressing comments

[mravishankar]

In D80188#2046505, @nicolasvasilache wrote:

On the change itself, this does not seem to support interleaving.
I think I would rather try to go for a recursive helper function that takes ArrayRef slices and a leaf lambda but otherwise leaves GenericLoopNestBuilder untouched.
This should also evolve better when we have more loop types.

Modified this to not change the EDSC builders but added a Util struct that implements what is needed. PTAL

[mravishankar]

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

387

Can you give me more context here. I am not sure why I need to do that. I am just changing the class used. I dont follow why it is needed now, but wasnt there previously?

[nicolasvasilache]

Thanks for the cleanup, looks nice!
I think you also have an opportunity to fix it completely by also taking the iterator types into account with 3-5 more locs (and a test).
Approving conditioned on that.

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

387

I do not remember the chronology of when lowering to loops and the permutation argument was added to tiling but in the same way you are fixing the fact that not all loops are parallel after tiling, we also need to take the iterator types into account.

This also points at missing coverage in the test.

A few lines above we have:

if (!options.interchangeVector.empty())
  applyPermutationToVector(loopRanges, options.interchangeVector);

We should also have something along the lines of:

auto iteratorTypes = llvm::to_vector<4>(op.iterator_types().cast<ArrayAttr>().getValue());
if (!options.interchangeVector.empty())
  applyPermutationToVector(iteratorTypes, options.interchangeVector);

Could you please add a test that permutes and be sure that say your outer parallel loop becomes a sequential when appropriate?

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

129 ↗

(On Diff #265838)

A few interspersed comments would be welcome on the fact that this recursive impl terminates and where it supports interleaved bands of loop types.

136 ↗

(On Diff #265838)

seems like we could have an isParallelAttr in the Utils.h and this would become xxx.take_while(isParallelAttr).

156 ↗

(On Diff #265838)

nice!

[mravishankar]

Addressing comments, and adding a permutation test.

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

387

Thanks for pointing this out. Added a test to verify the behavior as well.

[mravishankar]

Addressing comments.

[jpienaar]

This appears to be causing build breakage https://buildkite.com/mlir/mlir-core/builds/5209#139e9028-2104-4f29-9a8e-29af92d23cff

[mravishankar]

In D80188#2057659, @jpienaar wrote:

This appears to be causing build breakage https://buildkite.com/mlir/mlir-core/builds/5209#139e9028-2104-4f29-9a8e-29af92d23cff

Sorry. This should fix it : https://reviews.llvm.org/D80657

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

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);
-  void operator()(std::function<void(void)> fun = nullptr) { (*builder)(fun); }
+                              ArrayRef<Value> ranges,
+                              ArrayRef<Attribute> iteratorTypes);
+  void operator()(std::function<void(void)> fun = nullptr);
 
 private:
   using LoopOrAffineLoopBuilder =
       typename std::conditional_t<std::is_same<LoopTy, AffineForOp>::value,
                                   AffineLoopNestBuilder, LoopNestRangeBuilder>;
   using BuilderType =
       typename std::conditional_t<std::is_same<LoopTy, scf::ParallelOp>::value,
                                   ParallelLoopNestBuilder,
                                   LoopOrAffineLoopBuilder>;
 
   std::unique_ptr<BuilderType> builder;
+
+  std::unique_ptr<LoopNestRangeBuilder> forLoopBuilder;
 };
 
 inline void defaultRegionBuilder(ArrayRef<BlockArgument> args) {}
 
 /// Build a `linalg.generic` op with the specified `inputs`, `outputs` and
 /// `region`.
 ///
 /// `otherValues` and `otherAttributes` may be passed and will be appended as

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

Value LoopNestRangeBuilder::LoopNestRangeBuilder::operator()(
   return Value();
 }
 
 namespace mlir {
 namespace edsc {
 
 template <>
 GenericLoopNestRangeBuilder<scf::ForOp>::GenericLoopNestRangeBuilder(
-    MutableArrayRef<Value> ivs, ArrayRef<Value> ranges) {
+    MutableArrayRef<Value> ivs, ArrayRef<Value> ranges,
+    ArrayRef<Attribute> iteratorTypes) {
   builder = std::make_unique<LoopNestRangeBuilder>(ivs, ranges);
 }
 
 template <>
+void GenericLoopNestRangeBuilder<scf::ForOp>::operator()(
+    std::function<void(void)> fun) {
+  (*builder)(fun);
+}
+
+template <>
 GenericLoopNestRangeBuilder<AffineForOp>::GenericLoopNestRangeBuilder(
-    MutableArrayRef<Value> ivs, ArrayRef<Value> ranges) {
+    MutableArrayRef<Value> ivs, ArrayRef<Value> ranges,
+    ArrayRef<Attribute> iteratorTypes) {
   SmallVector<Value, 4> lbs;
   SmallVector<Value, 4> ubs;
   SmallVector<int64_t, 4> steps;
   for (Value range : ranges) {
     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());
   }
   builder = std::make_unique<AffineLoopNestBuilder>(ivs, lbs, ubs, steps);
 }
 
 template <>
+void GenericLoopNestRangeBuilder<AffineForOp>::operator()(
+    std::function<void(void)> fun) {
+  (*builder)(fun);
+}
+
+template <>
 GenericLoopNestRangeBuilder<scf::ParallelOp>::GenericLoopNestRangeBuilder(
-    MutableArrayRef<Value> ivs, ArrayRef<Value> ranges) {
+    MutableArrayRef<Value> ivs, ArrayRef<Value> ranges,
+    ArrayRef<Attribute> iteratorTypes) {
   SmallVector<Value, 4> lbs, ubs, steps;
-  for (Value range : ranges) {
-    assert(range.getType() && "expected linalg.range type");
-    assert(range.getDefiningOp() && "need operations to extract range parts");
-    RangeOp rangeOp = cast<RangeOp>(range.getDefiningOp());
+  size_t nOuterPar = iteratorTypes
+                         .take_while([](Attribute attr) {
+                           return attr.cast<StringAttr>().getValue() ==
+                                  getParallelIteratorTypeName();
+                         })
+                         .size();
+  nOuterPar = std::min(nOuterPar, ivs.size());
+  for (size_t i = 0; i != nOuterPar; ++i) {
+    assert(ranges[i].getType() && "expected linalg.range type");
+    assert(ranges[i].getDefiningOp() &&
+           "need operations to extract range parts");
+    RangeOp rangeOp = cast<RangeOp>(ranges[i].getDefiningOp());
     lbs.emplace_back(rangeOp.min());
     ubs.emplace_back(rangeOp.max());
     steps.emplace_back(rangeOp.step());
   }
-  builder = std::make_unique<ParallelLoopNestBuilder>(ivs, lbs, ubs, steps);
+  if (nOuterPar)
+    builder = std::make_unique<ParallelLoopNestBuilder>(
+        ivs.take_front(nOuterPar), lbs, ubs, steps);
+  if (nOuterPar != ivs.size())
+    forLoopBuilder = std::make_unique<LoopNestRangeBuilder>(
+        ivs.drop_front(nOuterPar), ranges.drop_front(nOuterPar));
+}
+
+template <>
+void GenericLoopNestRangeBuilder<scf::ParallelOp>::operator()(
+    std::function<void(void)> fun) {
+  auto forLoopFn = forLoopBuilder ? [&]() { (*forLoopBuilder)(fun); } : fun;
+  if (builder)
+    (*builder)(forLoopFn);
+  else if (forLoopBuilder)
+    (*forLoopBuilder)(forLoopFn);
 }
 
 } // namespace edsc
 } // namespace mlir
 
 Operation *mlir::edsc::makeGenericLinalgOp(
     ArrayRef<IteratorType> iteratorTypes, ArrayRef<StructuredIndexed> inputs,
     ArrayRef<StructuredIndexed> outputs,

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

 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,
                    MutableArrayRef<Value> allIvs) {
-    GenericLoopNestRangeBuilder<LoopTy>(allIvs, loopRanges)([&] {
+    GenericLoopNestRangeBuilder<LoopTy>(
+        allIvs, loopRanges, linalgOp.iterator_types().getValue())([&] {
       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,
-                   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()
-                         .take_while([](Attribute attr) {
-                           return attr.cast<StringAttr>().getValue() ==
-                                  getParallelIteratorTypeName();
-                         })
-                         .size();
-    // If there are no outer parallel loops, then number of loop ops is same as
-    // the number of loops, and they are all scf.for ops.
-    if (nOuterPar) {
-      GenericLoopNestRangeBuilder<scf::ParallelOp>(
-          allIvs.take_front(nOuterPar), loopRanges.take_front(nOuterPar))([&] {
-        GenericLoopNestRangeBuilder<scf::ForOp>(
-            allIvs.drop_front(nOuterPar),
-            loopRanges.drop_front(nOuterPar))([&] {
-          SmallVector<Value, 4> allIvValues(allIvs.begin(), allIvs.end());
-          LinalgScopedEmitter<StdIndexedValue, ConcreteOpTy>::
-              emitScalarImplementation(allIvValues, linalgOp);
-        });
-      });
-    } else {
-      // If there are no parallel loops then fallback to generating all scf.for
-      // operations.
-      GenericLoopNestRangeBuilder<scf::ForOp>(allIvs, loopRanges)([&] {
-        SmallVector<Value, 4> allIvValues(allIvs.begin(), allIvs.end());
-        LinalgScopedEmitter<StdIndexedValue,
-                            ConcreteOpTy>::emitScalarImplementation(allIvValues,
-                                                                    linalgOp);
-      });
-    }
-  }
-};
 } // namespace
 
 template <typename LoopTy, typename ConcreteOpTy>
 Optional<LinalgLoops> linalgOpToLoopsImpl(Operation *op, OpBuilder &builder) {
   using Impl = GenerateLoopNest<LoopTy, ConcreteOpTy>;
   using IndexedValueTy =
       typename GenerateLoopNest<LoopTy, ConcreteOpTy>::IndexedValueTy;
 

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

Optional<TiledLinalgOp> static tileLinalgOpImpl(
   LinalgOp res = op;
   SmallVector<Value, 4> ivs(loopRanges.size());
   // Convert SubViewOp::Range to linalg_range.
   SmallVector<Value, 4> linalgRanges;
   for (auto &range : loopRanges) {
     linalgRanges.push_back(
         linalg_range(range.offset, range.size, range.stride));
   }
-  GenericLoopNestRangeBuilder<LoopTy>(ivs, linalgRanges)([&] {
+  GenericLoopNestRangeBuilder<LoopTy>(ivs, linalgRanges,
+                                      op.iterator_types().getValue())([&] {
     auto &b = ScopedContext::getBuilderRef();
     auto loc = ScopedContext::getLocation();
     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

mlir/test/Dialect/Linalg/tile_parallel_reduce.mlir

+// RUN: mlir-opt %s -linalg-tile-to-parallel-loops="linalg-tile-sizes=2,4,8" -split-input-file | FileCheck %s
+
+func @gemm(%arg0 : memref<?x?xf32>,
+           %arg1 : memref<?x?xf32>,
+     %arg2 : memref<?x?xf32>)
+{
+  linalg.matmul(%arg0, %arg1, %arg2)
+    : memref<?x?xf32>, memref<?x?xf32>, memref<?x?xf32>
+  return
+}
+// CHECK-LABEL: func @gemm
+//   CHECK-DAG:   %[[C2:.*]] = constant 2 : index
+//   CHECK-DAG:   %[[C4:.*]] = constant 4 : index
+//   CHECK-DAG:   %[[C8:.*]] = constant 8 : index
+//       CHECK:   scf.parallel (%[[ARG3:.*]], %[[ARG4:.*]]) =
+//  CHECK-SAME:     step (%[[C2]], %[[C4]])
+//       CHECK:     scf.for %[[ARG5:.*]] =
+//  CHECK-SAME:       step %[[C8]]
+//       CHECK:       %[[SV1:.*]] = subview %{{.*}}[%[[ARG3]], %[[ARG5]]]
+//       CHECK:       %[[SV2:.*]] = subview %{{.*}}[%[[ARG5]], %[[ARG4]]]
+//       CHECK:       %[[SV3:.*]] = subview %{{.*}}[%[[ARG3]], %[[ARG4]]]
+//       CHECK:       linalg.matmul(%[[SV1]], %[[SV2]], %[[SV3]])
 No newline at end of file

mlir/test/Dialect/Linalg/transform-patterns.mlir

linalg.matvec(%A, %x, %y) : memref<?x?xf32, offset: ?, strides: [?, 1]>,
                               memref<?xf32, offset: ?, strides: [1]>,
                               memref<?xf32, offset: ?, strides: [1]>
   return
 }
 // CHECK-LABEL: func @matvec
 // CHECK-DAG:     %[[c0:.*]] = constant 0 : index
 // CHECK-DAG:     %[[c5:.*]] = constant 5 : index
 // CHECK-DAG:     %[[c6:.*]] = constant 6 : index
-// CHECK:         scf.parallel {{.*}} step (%[[c5]], %[[c6]])
+// CHECK:         scf.parallel {{.*}} step (%[[c5]])
+// CHECK:           scf.for {{.*}} step %[[c6]]
 // CHECK:             linalg.matvec({{.*}}, {{.*}}, {{.*}}) : memref<?x?xf32, #[[STRIDED_2D]]>, memref<?xf32, #[[STRIDED_1D]]>, memref<?xf32, #[[STRIDED_1D]]>
 
 func @matmul(%A: memref<?x?xf32, offset: ?, strides: [?, 1]>,
              %B: memref<?x?xf32, offset: ?, strides: [?, 1]>,
              %C: memref<?x?xf32, offset: ?, strides: [?, 1]>) {
   linalg.matmul(%A, %B, %C) : memref<?x?xf32, offset: ?, strides: [?, 1]>,
                               memref<?x?xf32, offset: ?, strides: [?, 1]>,
                               memref<?x?xf32, offset: ?, strides: [?, 1]>