[mlir][Linalg] Add support for tileAndDistribute on tensors.

[nicolasvasilache]

Summary

scf.parallel is currently not a good fit for tiling on tensors.
Instead provide a path to parallelism directly through scf.for.
For now, this transformation ignores the distribution scheme and always does a block-cyclic mapping (where block is the tile size).

[mravishankar]

Change is fairly straight-forward, but not sure what the issue with scf.parallel is. Is it the semantics of the op or the implementation of the distribution logic. If it is the latter, then maybe I can take a look there. Change looks fine as is though.

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

 #include "mlir/IR/Builders.h"
 #include "mlir/IR/Types.h"
 
 namespace mlir {
 namespace edsc {
 
 /// 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,
+scf::LoopNest loopNestBuilder(ValueRange lbs, ValueRange ubs,
                                  ValueRange steps,
                                  function_ref<void(ValueRange)> fun = nullptr);
-scf::ValueVector loopNestBuilder(Value lb, Value ub, Value step,
+scf::LoopNest loopNestBuilder(Value lb, Value ub, Value step,
                                  function_ref<void(Value)> fun = nullptr);
-scf::ValueVector loopNestBuilder(
+scf::LoopNest loopNestBuilder(
     Value lb, Value ub, Value step, ValueRange iterArgInitValues,
     function_ref<scf::ValueVector(Value, ValueRange)> fun = nullptr);
-scf::ValueVector loopNestBuilder(
+scf::LoopNest 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

mlir/include/mlir/Dialect/SCF/SCF.h

 ForOp getForInductionVarOwner(Value val);
 
 /// Returns the parallel loop parent of an induction variable. If the provided
 /// value is not an induction variable, then return nullptr.
 ParallelOp getParallelForInductionVarOwner(Value val);
 
 /// An owning vector of values, handy to return from functions.
 using ValueVector = std::vector<Value>;
+using LoopVector = std::vector<scf::ForOp>;
+struct LoopNest {
+  ResultRange getResults() { return loops.front().getResults(); }
+  LoopVector loops;
+};
 
 /// Creates a perfect nest of "for" loops, i.e. all loops but the innermost
 /// contain only another loop and a terminator. The lower, upper bounds and
 /// steps are provided as `lbs`, `ubs` and `steps`, which are expected to be of
 /// the same size. `iterArgs` points to the initial values of the loop iteration
 /// arguments, which will be forwarded through the nest to the innermost loop.
 /// The body of the loop is populated using `bodyBuilder`, which accepts an
 /// ordered list of induction variables of all loops, followed by a list of
 /// iteration arguments of the innermost loop, in the same order as provided to
 /// `iterArgs`. This function is expected to return as many values as
 /// `iterArgs`, of the same type and in the same order, that will be treated as
 /// yielded from the loop body and forwarded back through the loop nest. If the
 /// function is not provided, the loop nest is not expected to have iteration
 /// arguments, the body of the innermost loop will be left empty, containing
-/// only the zero-operand terminator. Returns the values yielded by the
-/// outermost loop. If bound arrays are empty, the body builder will be called
+/// only the zero-operand terminator. Returns the LoopNest containing the list
+/// of perfectly nest scf::ForOp build during the call.
+/// If bound arrays are empty, the body builder will be called
 /// once to construct the IR outside of the loop with an empty list of induction
 /// variables.
-ValueVector buildLoopNest(
+LoopNest buildLoopNest(
     OpBuilder &builder, Location loc, ValueRange lbs, ValueRange ubs,
     ValueRange steps, ValueRange iterArgs,
     function_ref<ValueVector(OpBuilder &, Location, ValueRange, ValueRange)>
         bodyBuilder = nullptr);
 
 /// A convenience version for building loop nests without iteration arguments
 /// (like for reductions). Does not take the initial value of reductions or
 /// expect the body building functions to return their current value.
-ValueVector buildLoopNest(OpBuilder &builder, Location loc, ValueRange lbs,
+/// The built nested scf::For are captured in `capturedLoops` when non-null.
+LoopNest buildLoopNest(OpBuilder &builder, Location loc, ValueRange lbs,
                           ValueRange ubs, ValueRange steps,
                           function_ref<void(OpBuilder &, Location, ValueRange)>
                               bodyBuilder = nullptr);
 
 } // end namespace scf
 } // end namespace mlir
 #endif // MLIR_DIALECT_SCF_H_

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

 #include "mlir/Dialect/SCF/EDSC/Builders.h"
 #include "mlir/Dialect/SCF/SCF.h"
 #include "mlir/Dialect/StandardOps/IR/Ops.h"
 #include "mlir/IR/AffineExpr.h"
 #include "mlir/IR/AffineMap.h"
 #include "mlir/IR/Matchers.h"
 #include "mlir/IR/OpImplementation.h"
 #include "mlir/Pass/Pass.h"
+#include "mlir/Transforms/LoopUtils.h"
 
 using namespace mlir;
 using namespace mlir::linalg;
 using namespace mlir::scf;
 
 Optional<RegionMatcher::BinaryOpKind>
 RegionMatcher::matchAsScalarBinaryOp(GenericOp op) {
   auto &region = op.region();

 }
 
 /// Specialization to build an scf "for" nest.
 template <>
 void GenerateLoopNest<scf::ForOp>::doit(
     ArrayRef<Range> loopRanges, ValueRange iterArgInitValues,
     ArrayRef<Attribute> iteratorTypes,
     function_ref<scf::ValueVector(ValueRange, ValueRange)> bodyBuilderFn,
-    Optional<LinalgLoopDistributionOptions>) {
+    Optional<LinalgLoopDistributionOptions> distributionOptions) {
+  // Create procInfo so it dominate loops, if appropriate.
+  OpBuilder &builder = edsc::ScopedContext::getBuilderRef();
+  Location loc = edsc::ScopedContext::getLocation();
+  SmallVector<ProcInfo, 2> procInfo;
+  if (distributionOptions.hasValue())
+    procInfo = distributionOptions->procInfo(builder, loc, ArrayRef<Range>{});
+
   SmallVector<Value, 4> lbs, ubs, steps;
   unpackRanges(loopRanges, lbs, ubs, steps);
+  LoopNest loopNest =
       edsc::loopNestBuilder(lbs, ubs, steps, iterArgInitValues, bodyBuilderFn);
+
+  if (!distributionOptions.hasValue() || loopNest.loops.empty())
+    return;
+
+  // TODO: support distributionMethod, which is currently ignored.
+  for (auto it : llvm::zip(loopNest.loops, procInfo,
+                           distributionOptions->distributionMethod))
+    mapLoopToProcessorIds(std::get<0>(it), std::get<1>(it).procId,
+                          std::get<1>(it).nprocs);
 }
 
 /// Specialization to build affine "for" nest.
 template <>
 void GenerateLoopNest<AffineForOp>::doit(
     ArrayRef<Range> loopRanges, ValueRange iterArgInitValues,
     ArrayRef<Attribute> iteratorTypes,
     function_ref<scf::ValueVector(ValueRange, ValueRange)> bodyBuilderFn,

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

 //===- Builders.cpp - MLIR Declarative Builder Classes --------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/SCF/EDSC/Builders.h"
 #include "mlir/Dialect/SCF/SCF.h"
 #include "mlir/IR/AffineExpr.h"
 #include "mlir/IR/AffineMap.h"
 
 using namespace mlir;
 using namespace mlir::edsc;
 
-mlir::scf::ValueVector
+mlir::scf::LoopNest
 mlir::edsc::loopNestBuilder(ValueRange lbs, ValueRange ubs, ValueRange steps,
                             function_ref<void(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, [&](OpBuilder &builder, Location loc, ValueRange ivs) {
         ScopedContext context(builder, loc);
         if (fun)
           fun(ivs);
       });
 }
 
-mlir::scf::ValueVector
+mlir::scf::LoopNest
 mlir::edsc::loopNestBuilder(Value lb, Value ub, Value step,
                             function_ref<void(Value)> fun) {
   // Delegates to the ValueRange-based version by wrapping the lambda.
   auto wrapper = [&](ValueRange ivs) {
     assert(ivs.size() == 1);
     if (fun)
       fun(ivs[0]);
   };
   return loopNestBuilder(ValueRange(lb), ValueRange(ub), ValueRange(step),
                          wrapper);
 }
 
-mlir::scf::ValueVector mlir::edsc::loopNestBuilder(
+mlir::scf::LoopNest mlir::edsc::loopNestBuilder(
     Value lb, Value ub, Value step, ValueRange iterArgInitValues,
     function_ref<scf::ValueVector(Value, 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(), lb, ub,
       step, iterArgInitValues,
       [&](OpBuilder &builder, Location loc, ValueRange ivs, ValueRange args) {
         assert(ivs.size() == 1 && "expected one induction variable");
         ScopedContext context(builder, loc);
         if (fun)
           return fun(ivs[0], args);
         return scf::ValueVector(iterArgInitValues.begin(),
                                 iterArgInitValues.end());
       });
 }
 
-mlir::scf::ValueVector mlir::edsc::loopNestBuilder(
+mlir::scf::LoopNest 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,

mlir/lib/Dialect/SCF/SCF.cpp

if (!lb || !ub || !step || step.getValue().getSExtValue() == 0) {
     return;
   }
 
   countPerRegion[0] =
       ceilDiv(ub.getValue().getSExtValue() - lb.getValue().getSExtValue(),
               step.getValue().getSExtValue());
 }
 
-ValueVector mlir::scf::buildLoopNest(
+LoopNest mlir::scf::buildLoopNest(
     OpBuilder &builder, Location loc, ValueRange lbs, ValueRange ubs,
     ValueRange steps, ValueRange iterArgs,
     function_ref<ValueVector(OpBuilder &, Location, ValueRange, ValueRange)>
         bodyBuilder) {
   assert(lbs.size() == ubs.size() &&
          "expected the same number of lower and upper bounds");
   assert(lbs.size() == steps.size() &&
          "expected the same number of lower bounds and steps");
 
   // If there are no bounds, call the body-building function and return early.
   if (lbs.empty()) {
     ValueVector results =
         bodyBuilder ? bodyBuilder(builder, loc, ValueRange(), iterArgs)
                     : ValueVector();
     assert(results.size() == iterArgs.size() &&
            "loop nest body must return as many values as loop has iteration "
            "arguments");
-    return results;
+    return LoopNest();
   }
 
   // First, create the loop structure iteratively using the body-builder
   // callback of `ForOp::build`. Do not create `YieldOp`s yet.
   OpBuilder::InsertionGuard guard(builder);
   SmallVector<scf::ForOp, 4> loops;
   SmallVector<Value, 4> ivs;
   loops.reserve(lbs.size());
ValueVector results = bodyBuilder
                                           loops.back().getRegionIterArgs())
                             : ValueVector();
   assert(results.size() == iterArgs.size() &&
          "loop nest body must return as many values as loop has iteration "
          "arguments");
   builder.setInsertionPointToEnd(loops.back().getBody());
   builder.create<scf::YieldOp>(loc, results);
 
-  // Return the results of the outermost loop.
-  return ValueVector(loops.front().result_begin(), loops.front().result_end());
+  // Return the loops.
+  LoopNest res;
+  res.loops.assign(loops.begin(), loops.end());
+  return res;
 }
 
-ValueVector mlir::scf::buildLoopNest(
+LoopNest mlir::scf::buildLoopNest(
     OpBuilder &builder, Location loc, ValueRange lbs, ValueRange ubs,
     ValueRange steps,
     function_ref<void(OpBuilder &, Location, ValueRange)> bodyBuilder) {
   // Delegate to the main function by wrapping the body builder.
   return buildLoopNest(builder, loc, lbs, ubs, steps, llvm::None,
                        [&bodyBuilder](OpBuilder &nestedBuilder,
                                       Location nestedLoc, ValueRange ivs,
                                       ValueRange) -> ValueVector {

mlir/test/Dialect/Linalg/tile-and-distribute.mlir

 //      CHECK: scf.parallel (%[[ARG3.*]]) = (%[[LBY]]) to (%{{.*}}) step (%[[STEPY]])
 //      CHECK:   scf.for %[[ARG4:.*]] =
 //      CHECK:     %[[SV1:.*]] = subview %[[ARG0]][%[[ARG3]], %[[ARG4]]]
 //      CHECK:     %[[OFFSETX:.*]] = affine.apply #[[MAP0]]()[%[[BIDX]]]
 //      CHECK:     %[[SV2:.*]] = subview %[[ARG1]][%[[ARG4]], %[[OFFSETX]]]
 //      CHECK:     %[[OFFSETX_2:.*]] = affine.apply #[[MAP0]]()[%[[BIDX]]]
 //      CHECK:     %[[SV3:.*]] = subview %[[ARG2]][%[[ARG3]], %[[OFFSETX_2]]]
 //      CHECK:     linalg.matmul ins(%[[SV1]], %[[SV2]]{{.*}} outs(%[[SV3]]
+
+// -----
+
+// CHECK-LABEL: func @matmul_tensors(
+// CHECK-SAME:    %[[TA:[0-9a-z]+]]: tensor<?x?xf32>
+// CHECK-SAME:    %[[TB:[0-9a-z]+]]: tensor<?x?xf32>
+// CHECK-SAME:    %[[TC:[0-9a-z]+]]: tensor<?x?xf32>) -> tensor<?x?xf32> {
+func @matmul_tensors(
+  %arg0: tensor<?x?xf32>, %arg1: tensor<?x?xf32>, %arg2: tensor<?x?xf32>)
+    -> tensor<?x?xf32> {
+//      CHECK: %[[C8:.*]] = constant 8 : index
+//      CHECK: %[[BIDY:.*]] = "gpu.block_id"() {dimension = "y"}
+//      CHECK: %[[NBLOCKSY:.*]] = "gpu.grid_dim"() {dimension = "y"}
+//      CHECK: %[[BIDX:.*]] = "gpu.block_id"() {dimension = "x"}
+//      CHECK: %[[NBLOCKSX:.*]] = "gpu.grid_dim"() {dimension = "x"}
+//      CHECK: %[[LBY:.*]] = muli %[[BIDY]], %[[C8]] : index
+//      CHECK: %[[STEPY:.*]] = muli %[[NBLOCKSY]], %[[C8]] : index
+//      CHECK: %[[TD0:.*]] = scf.for {{.*}} to {{.*}} step {{.*}} iter_args(%[[TC0:.*]] = %[[TC]]) -> (tensor<?x?xf32>) {
+//      CHECK: %[[LBX:.*]] = muli %[[BIDX]], %[[C8]] : index
+//      CHECK: %[[STEPX:.*]] = muli %[[NBLOCKSX]], %[[C8]] : index
+//      CHECK:   %[[TD1:.*]] = scf.for {{.*}} to {{.*}} step {{.*}} iter_args(%[[TC1:.*]] = %[[TC0]]) -> (tensor<?x?xf32>) {
+//      CHECK:     %[[TD2:.*]] = scf.for {{.*}} to {{.*}} step {{.*}} iter_args(%[[TC2:.*]] = %[[TC1]]) -> (tensor<?x?xf32>) {
+//      CHECK:       %[[sTA:.*]] = subtensor %[[TA]][{{.*}}] : tensor<?x?xf32> to tensor<?x?xf32>
+//      CHECK:       %[[sTB:.*]] = subtensor %[[TB]][{{.*}}] : tensor<?x?xf32> to tensor<?x?xf32>
+//      CHECK:       %[[sTC:.*]] = subtensor %[[TC2]][{{.*}}] : tensor<?x?xf32> to tensor<?x?xf32>
+//      CHECK:       %[[sTD:.*]] = linalg.matmul ins(%[[sTA]], %[[sTB]] : tensor<?x?xf32>, tensor<?x?xf32>)
+// CHECK-SAME:                                  init(%[[sTC]] : tensor<?x?xf32>)  -> tensor<?x?xf32>
+//      CHECK:       %[[TD:.*]] = subtensor_insert %[[sTD]] into %[[TC2]][{{.*}}]  : tensor<?x?xf32> into tensor<?x?xf32>
+//      CHECK:       scf.yield %[[TD]] : tensor<?x?xf32>
+//      CHECK:     scf.yield %[[TD2]] : tensor<?x?xf32>
+//      CHECK:   scf.yield %[[TD1]] : tensor<?x?xf32>
+  %0 = linalg.matmul {__internal_linalg_transform__ = "tensors_distribute1"}
+       ins(%arg0, %arg1: tensor<?x?xf32>, tensor<?x?xf32>)
+      init(%arg2: tensor<?x?xf32>)
+    -> tensor<?x?xf32>
+
+//      CHECK: return %[[TD0]] : tensor<?x?xf32>
+  return %0 : tensor<?x?xf32>
+}
+

mlir/test/EDSC/builder-api-test.cpp

TEST_FUNC(builder_loop_for_yield) {
   Value init1 = std_constant_float(llvm::APFloat(2.0f), f32Type);
   Value a(f.getArgument(0)), b(f.getArgument(1)), c(f.getArgument(2)),
       d(f.getArgument(3));
   using namespace edsc::op;
   auto results = loopNestBuilder(a - b, c + d, a, {init0, init1},
                                  [&](Value iv, ValueRange args) {
                                    Value sum = args[0] + args[1];
                                    return scf::ValueVector{args[1], sum};
-                                 });
+                                 }).getResults();
   results[0] + results[1];
 
   // clang-format off
   // CHECK-LABEL: func @builder_loop_for_yield(%{{.*}}: index, %{{.*}}: index, %{{.*}}: index, %{{.*}}: index) {
   // CHECK:     [[init0:%.*]] = constant
   // CHECK:     [[init1:%.*]] = constant
   // CHECK-DAG:    [[r0:%[0-9]+]] = affine.apply affine_map<()[s0, s1] -> (s0 - s1)>()[%{{.*}}, %{{.*}}]
   // CHECK-DAG:    [[r1:%[0-9]+]] = affine.apply affine_map<()[s0, s1] -> (s0 + s1)>()[%{{.*}}, %{{.*}}]

mlir/test/lib/Transforms/TestLinalgTransforms.cpp

patterns.insert<LinalgTilingPattern<MatmulOp>>(
         context,
         LinalgTilingOptions()
             .setTileSizes({8, 8, 4})
             .setLoopType(LinalgTilingLoopType::ParallelLoops)
             .setDistributionOptions(cyclicNprocsMixed3),
         LinalgMarker(Identifier::get("distribute6", context),
                      Identifier::get("after_distribute6", context)));
   }
+
+  {
+    LinalgLoopDistributionOptions cyclicNprocsEqNiters;
+    cyclicNprocsEqNiters.distributionMethod.resize(
+        2, DistributionMethod::CyclicNumProcsEqNumIters);
+    cyclicNprocsEqNiters.procInfo =
+        getGpuProcIds<gpu::BlockIdOp, gpu::GridDimOp>;
+    patterns.insert<LinalgTilingPattern<MatmulOp>>(
+        context,
+        LinalgTilingOptions()
+            .setTileSizes({8, 8, 4})
+            .setLoopType(LinalgTilingLoopType::Loops)
+            .setDistributionOptions(cyclicNprocsEqNiters),
+        LinalgMarker(Identifier::get("tensors_distribute1", context),
+                     Identifier::get("tensors_after_distribute1", context)));
+  }
 }
 
 static void
 applyMatmulToVectorPatterns(FuncOp funcOp,
                             bool testMatmulToVectorPatterns1dTiling,
                             bool testMatmulToVectorPatterns2dTiling) {
   MLIRContext *ctx = funcOp.getContext();
   SmallVector<OwningRewritePatternList, 4> stage1Patterns;