[mlir][Linalg] Allow distribution `scf.parallel` loops generated in Linalg to processors.

[mravishankar]

Summary

This changes adds infrastructure to distribute the loops generated in
Linalg to processors at the time of generation. This addresses use
case where the instantiation of loop is done just to distribute
them. The option to distribute is added to TilingOptions for now and
will allow specifying the distribution as a transformation option,
just like tiling and promotion are specified as options.

[ThomasRaoux]

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

187–189

Does this need to be an affine expression? Nicolas has some patch in flight to allow folding the affine.min created by tiling, I'm not sure if it needs it. Hopefully Nicolas can comment.

[nicolasvasilache]

Look great Mahesh, thanks for working on this !

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

187–189

+1

268

nit: I would reformat this under a // clang-format off to make it "look more like loops".

290

Nit: while we're still under EDSC here, it should look quite nicer to use: Value slt(Value, Value) etc.
Here and below.

296

same nit re "look more like loops"

309

same nit re "look more like loops"

[mravishankar]

Using affine maps for bounds/step update.

[mravishankar]

Removing errant tabs

[mravishankar]

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

268

Not sure I follow. Looks like a loop to me already, i.e.,

OperationBuilder(...) {

}

296

This too looks like a loop to me...

conditionBuilder() {
  ...
}

309

There is no loop here :)

mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h

LinalgTilingOptions &setInterchange(ArrayRef<unsigned> interchange) {
     return *this;
   }
   /// The type of tile loops to generate.
   LinalgTilingLoopType loopType{LinalgTilingLoopType::Loops};
   LinalgTilingOptions &setLoopType(LinalgTilingLoopType lt) {
     loopType = lt;
     return *this;
   }
+
+  /// When specified, specifies distribution of generated tile loops to
+  /// processors.
+  Optional<LinalgLoopDistributionOptions> distribution = None;
+  LinalgTilingOptions &
+  setDistributionOptions(LinalgLoopDistributionOptions &distributionOptions) {
+    distribution = distributionOptions;
+    return *this;
+  }
 };
 
 /// Canonicalization patterns relevant to apply after tiling patterns. These are
 /// applied automatically by the tiling pass but need to be applied manually
 /// when tiling is called programmatically.
 OwningRewritePatternList
 getLinalgTilingCanonicalizationPatterns(MLIRContext *ctx);
 

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

 void applyPermutationToVector(SmallVector<T, N> &inVec,
                               ArrayRef<unsigned> permutation) {
   SmallVector<T, N> auxVec(inVec.size());
   for (unsigned i = 0; i < permutation.size(); ++i)
     auxVec[i] = inVec[permutation[i]];
   inVec = auxVec;
 }
 
+/// Scheme used to distribute loops to processors.
+enum class DistributionMethod {
+  /// Cyclic distribution where no assumption is made about the dynamic
+  /// relationship between number of processors and number of iterations of the
+  /// distributed loop. Distributes the following loop
+  ///
+  /// scf.parallel (%iv) = (%lb) to (%ub) step (%step)
+  ///
+  /// to
+  ///
+  /// scf.parallel(%iv)= (%lb + %procId * %step) to (%ub) step (%step * %nprocs)
+  Cyclic = 0,
+
+  /// Cyclic distribtuion where the number of processors can be assumed to be
+  /// more than or equal to the number of iterations of the distributed loop. In
+  /// such cases, a simple in-bounds check is enough (instead of materializing a
+  /// loop). Distributes the following loop
+  ///
+  /// scf.parallel (%iv) = (%lb) to (%ub) step (%step)
+  ///
+  /// to
+  ///
+  /// %iv = %lb + %procId * %step
+  /// %cond = cmpi "slt", %iv, %ub
+  /// scf.if %cond {
+  ///   ...
+  /// }
+  CyclicNumProcsGeNumIters = 1,
+
+  /// Cyclic distribtuion where the number of processors can be assumed to be
+  ///  equal to the number of iterations of the distributed loop. In such cases,
+  ///  no bounds check is needed. Distributes the following loop
+  ///
+  /// scf.parallel (%iv) = (%lb) to (%ub) step (%step)
+  ///
+  /// to
+  ///
+  /// %iv = %lb + %procId * %step
+  CyclicNumProcsEqNumIters = 2
+};
+
+/// Callback function type used to get processor ID, and number of processors
+/// used for distribution.
+struct ProcInfo {
+  Value procId;
+  Value nprocs;
+};
+using ProcInfoCallBackFn =
+    std::function<ProcInfo(OpBuilder &b, Location loc, unsigned loopNum)>;
+
+/// Options that allow distribution of loops generated in Linalg transforms to
+/// processors while generating the loops.
+struct LinalgLoopDistributionOptions {
+  /// Callback function that returns the Value for processor ID, and number of
+  /// processors used to execute a given loop.
+  ProcInfoCallBackFn procInfo;
+  /// Specification of how to distribute the `scf.parallel` loops that are
+  /// generated. As the `scf.parallel` loop is generated, the elements of this
+  /// vector is used (from left to right) and the specified distribution is
+  /// applied. If the vector is less than the number of `scf.parallel` loops
+  /// generated, then no distribution is applied.
+  SmallVector<DistributionMethod, 0> distributionMethod = {};
+};
+
 /// Utility class used to generate nested loops with ranges described by
 /// `loopRanges` and loop type described by the `iteratorTypes`. `bodyBuilderFn`
 /// 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);
+                   function_ref<void(ValueRange)> bodyBuilderFn,
+                   Optional<LinalgLoopDistributionOptions> = None);
 };
 
 } // namespace linalg
 } // namespace mlir
 
 #endif // MLIR_DIALECT_LINALG_UTILS_H_

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, iteratorTypes,
+      [&](ValueRange localIvs) {
         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);
-      });
+      },
+      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;
   loops.reserve(ivs.size());
   for (auto iv : ivs) {
+    if (iv.isa<BlockArgument>()) {
       loops.push_back(iv.cast<BlockArgument>().getOwner()->getParentOp());
       assert(loops.back() && "no owner found for induction variable!");
+    } else {
+      // TODO: Instead of doing this, try to recover the ops used instead of the
+      // loop.
+      loops.push_back(nullptr);
+    }
   }
   return TiledLinalgOp{res, loops};
 }
 
 Optional<TiledLinalgOp>
 mlir::linalg::tileLinalgOp(OpBuilder &b, LinalgOp op,
                            const LinalgTilingOptions &options) {
   if (options.loopType == LinalgTilingLoopType::Loops)

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

 //===- Utils.cpp - Utilities to support the Linalg dialect ----------------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 //
 // This file implements utilities for the Linalg dialect.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Linalg/Utils/Utils.h"
+
+#include "mlir/Dialect/Affine/EDSC/Intrinsics.h"
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Linalg/IR/LinalgOps.h"
 #include "mlir/Dialect/Linalg/IR/LinalgTypes.h"
 #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"

 
 namespace mlir {
 namespace linalg {
 
 /// 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) {
+    function_ref<void(ValueRange)> bodyBuilderFn,
+    Optional<LinalgLoopDistributionOptions>) {
   SmallVector<Value, 4> lbs, ubs, steps;
   unpackRanges(loopRanges, lbs, ubs, steps);
   edsc::loopNestBuilder(lbs, ubs, steps, bodyBuilderFn);
 }
 
 /// Specialization to build affine "for" nest.
 template <>
 void GenerateLoopNest<AffineForOp>::doit(
     ArrayRef<SubViewOp::Range> loopRanges, ArrayRef<Attribute> iteratorTypes,
-    function_ref<void(ValueRange)> bodyBuilderFn) {
+    function_ref<void(ValueRange)> bodyBuilderFn,
+    Optional<LinalgLoopDistributionOptions>) {
   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);
 }
 
-/// Generates a loop nest consisting of scf.parallel and scf.for, depending on
-/// the `iteratorTypes.` Consecutive parallel loops create a single scf.parallel
-/// operation; each sequential loop creates a new scf.for operation. The body
-/// of the innermost loop is populated by `bodyBuilderFn` that accepts a range
-/// of induction variables for all loops. `ivStorage` is used to store the
-/// partial list of induction variables.
+/// 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+;
+  using edsc::op::operator*;
+  lb = lb + (procId * step);
+  step = nprocs * step;
+}
+
+/// Generates a loop nest consisting of scf.parallel and scf.for, depending
+/// on the `iteratorTypes.` Consecutive parallel loops create a single
+/// scf.parallel operation; each sequential loop creates a new scf.for
+/// operation. The body of the innermost loop is populated by
+/// `bodyBuilderFn` that accepts a range of induction variables for all
+/// loops. `ivStorage` is used to store the partial list of induction
+/// variables.
 // TODO: this function can be made iterative instead. However, it
 // will have at most as many recursive calls as nested loops, which rarely
 // exceeds 10.
 static void
 generateParallelLoopNest(ValueRange lbs, ValueRange ubs, ValueRange steps,
                          ArrayRef<Attribute> iteratorTypes,
                          function_ref<void(ValueRange)> bodyBuilderFn,
-                         SmallVectorImpl<Value> &ivStorage) {
+                         SmallVectorImpl<Value> &ivStorage,
+                         ArrayRef<DistributionMethod> distributionMethod = {}) {
   assert(lbs.size() == ubs.size());
   assert(lbs.size() == steps.size());
   assert(lbs.size() == iteratorTypes.size());
 
   // If there are no (more) loops to be generated, generate the body and be
   // done with it.
   if (iteratorTypes.empty())
     return bodyBuilderFn(ivStorage);
 
   // Find the outermost parallel loops and drop their types from the list.
   unsigned nLoops = iteratorTypes.size();
-  iteratorTypes = iteratorTypes.drop_while(isParallelIteratorType);
-  unsigned nOuterPar = nLoops - iteratorTypes.size();
+  unsigned nOuterPar =
+      nLoops - iteratorTypes.drop_while(isParallelIteratorType).size();
 
   // If there are no outer parallel loops, generate one sequential loop and
   // recurse. Note that we wouldn't have dropped anything from `iteratorTypes`
   // in this case.
   if (nOuterPar == 0) {
     edsc::loopNestBuilder(lbs[0], ubs[0], steps[0], [&](Value iv) {
       ivStorage.push_back(iv);
       generateParallelLoopNest(lbs.drop_front(), ubs.drop_front(),
                                steps.drop_front(), iteratorTypes.drop_front(),
-                               bodyBuilderFn, ivStorage);
+                               bodyBuilderFn, ivStorage, distributionMethod);
     });
     return;
   }
-
-  // Generate a single parallel loop-nest operation for all outermost parallel
-  // loops and recurse.
+  if (distributionMethod.empty()) {
+    // Generate a single parallel loop-nest operation for all outermost
+    // parallel loops and recurse.
     edsc::OperationBuilder<scf::ParallelOp>(
         lbs.take_front(nOuterPar), ubs.take_front(nOuterPar),
         steps.take_front(nOuterPar),
         [&](OpBuilder &nestedBuilder, Location nestedLoc, ValueRange localIvs) {
           edsc::ScopedContext context(nestedBuilder, nestedLoc);
           ivStorage.append(localIvs.begin(), localIvs.end());
-        generateParallelLoopNest(lbs.drop_front(nOuterPar),
-                                 ubs.drop_front(nOuterPar),
-                                 steps.drop_front(nOuterPar), iteratorTypes,
-                                 bodyBuilderFn, ivStorage);
-      });
+          generateParallelLoopNest(
+              lbs.drop_front(nOuterPar), ubs.drop_front(nOuterPar),
+              steps.drop_front(nOuterPar), iteratorTypes.drop_front(nOuterPar),
+              bodyBuilderFn, ivStorage,
+              (distributionMethod.size() < nOuterPar)
+                  ? ArrayRef<DistributionMethod>()
+                  : distributionMethod.drop_front(nOuterPar));
+        });
+    return;
+  }
+
+  // Process all consecutive similarly distributed loops simultaneously.
+  DistributionMethod methodToUse = distributionMethod[0];
+  unsigned numProcessed = 1;
+  for (unsigned i = 1; i < nOuterPar && i < distributionMethod.size(); ++i) {
+    if (distributionMethod[i] != methodToUse)
+      break;
+    numProcessed++;
+  }
+
+  switch (methodToUse) {
+  case DistributionMethod::Cyclic: {
+    // Generate a single parallel loop-nest operation for all outermost
+    // parallel loops and recurse.
+    edsc::OperationBuilder<scf::ParallelOp>(
+        lbs.take_front(numProcessed), ubs.take_front(numProcessed),
+        steps.take_front(numProcessed),
+        [&](OpBuilder &nestedBuilder, Location nestedLoc, ValueRange localIvs) {
+          edsc::ScopedContext context(nestedBuilder, nestedLoc);
+          ivStorage.append(localIvs.begin(), localIvs.end());
+          generateParallelLoopNest(
+              lbs.drop_front(numProcessed), ubs.drop_front(numProcessed),
+              steps.drop_front(numProcessed),
+              iteratorTypes.drop_front(numProcessed), bodyBuilderFn, ivStorage,
+              (distributionMethod.size() < numProcessed)
+                  ? ArrayRef<DistributionMethod>()
+                  : distributionMethod.drop_front(numProcessed));
+        });
+    return;
+  }
+  case DistributionMethod::CyclicNumProcsGeNumIters: {
+    // Check (for the processed loops) that the iteration is in-bounds.
+    using edsc::op::slt;
+    using edsc::op::operator&&;
+    Value cond = slt(lbs[0], ubs[0]);
+    for (unsigned i = 1; i < numProcessed; ++i)
+      cond = cond && slt(lbs[i], ubs[i]);
+    ivStorage.append(lbs.begin(), std::next(lbs.begin(), numProcessed));
+    edsc::conditionBuilder(cond, [&]() {
+      generateParallelLoopNest(
+          lbs.drop_front(numProcessed), ubs.drop_front(numProcessed),
+          steps.drop_front(numProcessed),
+          iteratorTypes.drop_front(numProcessed), bodyBuilderFn, ivStorage,
+          distributionMethod.drop_front(numProcessed));
+    return;
+  }
+  case DistributionMethod::CyclicNumProcsEqNumIters:
+    // No check/loops needed here. Set the `%iv` to be the `%lb` and proceed
+    // with inner loop generation.
+    ivStorage.append(lbs.begin(), std::next(lbs.begin(), numProcessed));
+    generateParallelLoopNest(
+        lbs.drop_front(numProcessed), ubs.drop_front(numProcessed),
+        steps.drop_front(numProcessed), iteratorTypes.drop_front(numProcessed),
+        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) {
-  SmallVector<Value, 8> lbsStorage, ubsStorage, stepsStorage, ivs;
-  unpackRanges(loopRanges, lbsStorage, ubsStorage, stepsStorage);
-  ValueRange lbs(lbsStorage), ubs(ubsStorage), steps(stepsStorage);
-
+    function_ref<void(ValueRange)> bodyBuilderFn,
+    Optional<LinalgLoopDistributionOptions> distributionOptions) {
   // 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());
-  ivs.reserve(iteratorTypes.size());
-  generateParallelLoopNest(lbs, ubs, steps, iteratorTypes, bodyBuilderFn, ivs);
+  SmallVector<Value, 8> lbsStorage, ubsStorage, stepsStorage, ivs;
+  unsigned numLoops = iteratorTypes.size();
+  ivs.reserve(numLoops);
+  lbsStorage.reserve(numLoops);
+  ubsStorage.reserve(numLoops);
+  stepsStorage.reserve(numLoops);
+
+  // Get the loop lb, ub, and step.
+  unpackRanges(loopRanges, lbsStorage, ubsStorage, stepsStorage);
+
+  // Modify the lb, ub, and step based on the distribution options.
+  SmallVector<DistributionMethod, 0> distributionMethod;
+  if (distributionOptions) {
+    auto &options = distributionOptions.getValue();
+    unsigned index = 0;
+    OpBuilder &builder = edsc::ScopedContext::getBuilderRef();
+    Location loc = edsc::ScopedContext::getLocation();
+    distributionMethod.assign(distributionOptions->distributionMethod.begin(),
+                              distributionOptions->distributionMethod.end());
+    for (auto iteratorType : enumerate(iteratorTypes))
+      if (isParallelIteratorType(iteratorType.value()) &&
+          index < distributionMethod.size()) {
+        unsigned i = iteratorType.index();
+        ProcInfo procInfo = options.procInfo(builder, loc, index);
+        updateBoundsForCyclicDistribution(builder, loc, procInfo.procId,
+                                          procInfo.nprocs, lbsStorage[i],
+                                          ubsStorage[i], stepsStorage[i]);
+        index++;
+      }
+  }
+  ValueRange lbs(lbsStorage), ubs(ubsStorage), steps(stepsStorage);
+  generateParallelLoopNest(lbs, ubs, steps, iteratorTypes, bodyBuilderFn, ivs,
+                           distributionMethod);
+
   assert(ivs.size() == iteratorTypes.size() && "did not generate enough loops");
 }
 
 } // namespace linalg
 } // namespace mlir

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

+// RUN: mlir-opt %s -test-linalg-transform-patterns=test-tile-and-distribute-options -split-input-file | FileCheck %s
+
+func @gemm1(%a : memref<?x?xf32>, %b : memref<?x?xf32>, %c : memref<?x?xf32>)
+{
+  linalg.matmul %a, %b, %c {__internal_linalg_transform__ = "distribute1"}
+    : (memref<?x?xf32>, memref<?x?xf32>, memref<?x?xf32>)
+  return
+}
+//  CHECK-DAG: #[[MAP0:.*]] = affine_map<()[s0] -> (s0 * 8)>
+//      CHECK: func @gemm1(
+// CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]*]]: memref<?x?xf32>
+// CHECK-SAME: %[[ARG1:[a-zA-Z0-9_]*]]: memref<?x?xf32>
+// CHECK-SAME: %[[ARG2:[a-zA-Z0-9_]*]]: memref<?x?xf32>
+//      CHECK: %[[T1:.*]] = "gpu.block_id"() {dimension = "y"}
+//      CHECK: %[[T2:.*]] = "gpu.block_id"() {dimension = "x"}
+//      CHECK: scf.for %[[ARG3:.*]] =
+//      CHECK:   %[[T3:.*]] = affine.apply #[[MAP0]]()[%[[T1]]]
+//      CHECK:   %[[SV1:.*]] = subview %[[ARG0]][%[[T3]], %[[ARG3]]]
+//      CHECK:   %[[T11:.*]] = affine.apply #[[MAP0]]()[%[[T2]]]
+//      CHECK:   %[[SV2:.*]] = subview %[[ARG1]][%[[ARG3]], %[[T11]]]
+//      CHECK:   %[[T15:.*]] = affine.apply #[[MAP0]]()[%[[T1]]]
+//      CHECK:   %[[T18:.*]] = affine.apply #[[MAP0]]()[%[[T2]]]
+//      CHECK:   %[[SV3:.*]] = subview %[[ARG2]][%[[T15]], %[[T18]]]
+//      CHECK:   linalg.matmul %[[SV1]], %[[SV2]], %[[SV3]]
+
+// -----
+
+func @gemm2(%a : memref<?x?xf32>, %b : memref<?x?xf32>, %c : memref<?x?xf32>)
+{
+  linalg.matmul %a, %b, %c {__internal_linalg_transform__ = "distribute2"}
+    : (memref<?x?xf32>, memref<?x?xf32>, memref<?x?xf32>)
+  return
+}
+//  CHECK-DAG: #[[MAP0:.*]] = affine_map<()[s0] -> (s0 * 8)>
+//      CHECK: func @gemm2(
+// CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]*]]: memref<?x?xf32>
+// CHECK-SAME: %[[ARG1:[a-zA-Z0-9_]*]]: memref<?x?xf32>
+// CHECK-SAME: %[[ARG2:[a-zA-Z0-9_]*]]: memref<?x?xf32>
+//      CHECK: %[[T3:.*]] = "gpu.block_id"() {dimension = "y"}
+//      CHECK: %[[T4:.*]] = affine.apply #[[MAP0]]()[%[[T3]]]
+//      CHECK: %[[T5:.*]] = "gpu.block_id"() {dimension = "x"}
+//      CHECK: %[[T6:.*]] = affine.apply #[[MAP0]]()[%[[T5]]]
+//      CHECK: %[[T7:.*]] = cmpi "slt", %[[T4]], %{{.*}}
+//      CHECK: %[[T8:.*]] = cmpi "slt", %[[T6]], %{{.*}}
+//      CHECK: %[[T9:.*]] = and %[[T7]], %[[T8]]
+//      CHECK: scf.if %[[T9]]
+//      CHECK:   scf.for %[[ARG3:.*]] =
+//      CHECK:     %[[T10:.*]] = affine.apply #[[MAP0]]()[%[[T3]]]
+//      CHECK:     %[[SV1:.*]] = subview %[[ARG0]][%[[T10]], %[[ARG3]]]
+//      CHECK:     %[[T18:.*]] = affine.apply #[[MAP0]]()[%[[T5]]]
+//      CHECK:     %[[SV2:.*]] = subview %[[ARG1]][%[[ARG3]], %[[T18]]]
+//      CHECK:     %[[T22:.*]] = affine.apply #[[MAP0]]()[%[[T3]]]
+//      CHECK:     %[[T25:.*]] = affine.apply #[[MAP0]]()[%[[T5]]]
+//      CHECK:     %[[SV3:.*]] = subview %[[ARG2]][%[[T22]], %[[T25]]]
+//      CHECK:     linalg.matmul %[[SV1]], %[[SV2]], %[[SV3]]
+
+// -----
+
+func @gemm3(%a : memref<?x?xf32>, %b : memref<?x?xf32>, %c : memref<?x?xf32>)
+{
+  linalg.matmul %a, %b, %c {__internal_linalg_transform__ = "distribute3"}
+    : (memref<?x?xf32>, memref<?x?xf32>, memref<?x?xf32>)
+  return
+}
+//  CHECK-DAG: #[[MAP0:.*]] = affine_map<()[s0] -> (s0 * 8)>
+//      CHECK: func @gemm3(
+// CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]*]]: memref<?x?xf32>
+// CHECK-SAME: %[[ARG1:[a-zA-Z0-9_]*]]: memref<?x?xf32>
+// CHECK-SAME: %[[ARG2:[a-zA-Z0-9_]*]]: memref<?x?xf32>
+//      CHECK: %[[T3:.*]] = "gpu.block_id"() {dimension = "y"}
+//      CHECK: %[[T4:.*]] = "gpu.grid_dim"() {dimension = "y"}
+//      CHECK: %[[T5:.*]] = affine.apply #[[MAP0]]()[%[[T3]]]
+//      CHECK: %[[T6:.*]] = affine.apply #[[MAP0]]()[%[[T4]]]
+//      CHECK: %[[T7:.*]] = "gpu.block_id"() {dimension = "x"}
+//      CHECK: %[[T8:.*]] = "gpu.grid_dim"() {dimension = "x"}
+//      CHECK: %[[T9:.*]] = affine.apply #[[MAP0]]()[%[[T7]]]
+//      CHECK: %[[T10:.*]] = affine.apply #[[MAP0]]()[%[[T8]]]
+//      CHECK: scf.parallel (%[[ARG3:.*]], %[[ARG4:.*]]) = (%[[T5]], %[[T9]]) to (%{{.*}}, %{{.*}}) step (%[[T6]], %[[T10]])
+//      CHECK:   scf.for %[[ARG5:.*]] =
+//      CHECK:     %[[SV1:.*]] = subview %[[ARG0]][%[[ARG3]], %[[ARG5]]]
+//      CHECK:     %[[SV2:.*]] = subview %[[ARG1]][%[[ARG5]], %[[ARG4]]]
+//      CHECK:     %[[SV3:.*]] = subview %[[ARG2]][%[[ARG3]], %[[ARG4]]]
+//      CHECK:     linalg.matmul %[[SV1]], %[[SV2]], %[[SV3]]
+
+// -----
+
+func @gemm4(%a : memref<?x?xf32>, %b : memref<?x?xf32>, %c : memref<?x?xf32>)
+{
+  linalg.matmul %a, %b, %c {__internal_linalg_transform__ = "distribute4"}
+    : (memref<?x?xf32>, memref<?x?xf32>, memref<?x?xf32>)
+  return
+}
+//  CHECK-DAG: #[[MAP0:.*]] = affine_map<()[s0] -> (s0 * 8)>
+//      CHECK: func @gemm4(
+// CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]*]]: memref<?x?xf32>
+// CHECK-SAME: %[[ARG1:[a-zA-Z0-9_]*]]: memref<?x?xf32>
+// CHECK-SAME: %[[ARG2:[a-zA-Z0-9_]*]]: memref<?x?xf32>
+//      CHECK: %[[T2:.*]] = "gpu.block_id"() {dimension = "y"}
+//      CHECK: %[[T3:.*]] = "gpu.block_id"() {dimension = "x"}
+//      CHECK: %[[T4:.*]] = affine.apply #[[MAP0]]()[%[[T3]]]
+//      CHECK: %[[T5:.*]] = cmpi "slt", %[[T4]], %{{.*}}
+//      CHECK: scf.if %[[T5]]
+//      CHECK:   scf.for %[[ARG3:.*]] =
+//      CHECK:     %[[T6:.*]] = affine.apply #[[MAP0]]()[%[[T2]]]
+//      CHECK:     %[[SV1:.*]] = subview %[[ARG0]][%[[T6]], %[[ARG3]]]
+//      CHECK:     %[[T14:.*]] = affine.apply #[[MAP0]]()[%[[T3]]]
+//      CHECK:     %[[SV2:.*]] = subview %[[ARG1]][%[[ARG3]], %[[T14]]]
+//      CHECK:     %[[T18:.*]] = affine.apply #[[MAP0]]()[%[[T2]]]
+//      CHECK:     %[[T21:.*]] = affine.apply #[[MAP0]]()[%[[T3]]]
+//      CHECK:     %[[SV3:.*]] = subview %[[ARG2]][%[[T18]], %[[T21]]]
+//      CHECK:     linalg.matmul %[[SV1]], %[[SV2]], %[[SV3]]
+
+// -----
+
+func @gemm5(%a : memref<?x?xf32>, %b : memref<?x?xf32>, %c : memref<?x?xf32>)
+{
+  linalg.matmul %a, %b, %c {__internal_linalg_transform__ = "distribute5"}
+    : (memref<?x?xf32>, memref<?x?xf32>, memref<?x?xf32>)
+  return
+}
+//  CHECK-DAG: #[[MAP0:.*]] = affine_map<()[s0] -> (s0 * 8)>
+//      CHECK: func @gemm5(
+// CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]*]]: memref<?x?xf32>
+// CHECK-SAME: %[[ARG1:[a-zA-Z0-9_]*]]: memref<?x?xf32>
+// CHECK-SAME: %[[ARG2:[a-zA-Z0-9_]*]]: memref<?x?xf32>
+//      CHECK: %[[T3:.*]] = "gpu.block_id"() {dimension = "y"}
+//      CHECK: %[[T4:.*]] = affine.apply #[[MAP0]]()[%[[T3]]]
+//      CHECK: %[[T5:.*]] = "gpu.block_id"() {dimension = "x"}
+//      CHECK: %[[T6:.*]] = "gpu.grid_dim"() {dimension = "x"}
+//      CHECK: %[[T7:.*]] = affine.apply #[[MAP0]]()[%[[T5]]]
+//      CHECK: %[[T8:.*]] = affine.apply #[[MAP0]]()[%[[T6]]]
+//      CHECK: %[[T9:.*]] = cmpi "slt", %[[T4]], %{{.*}}
+//      CHECK: scf.if %[[T9]]
+//      CHECK:   scf.parallel (%[[ARG3.*]]) = (%[[T7]]) to (%{{.*}}) step (%[[T8]])
+//      CHECK:     scf.for %[[ARG4:.*]] =
+//      CHECK:      %[[T10:.*]] = affine.apply #[[MAP0]]()[%[[T3]]]
+//      CHECK:       %[[SV1:.*]] = subview %[[ARG0]][%[[T10]], %[[ARG4]]]
+//      CHECK:       %[[SV2:.*]] = subview %[[ARG1]][%[[ARG4]], %[[ARG3]]]
+//      CHECK:       %[[T21:.*]] = affine.apply #[[MAP0]]()[%[[T3]]]
+//      CHECK:       %[[SV3:.*]] = subview %[[ARG2]][%[[T21]], %[[ARG3]]]
+//      CHECK:       linalg.matmul %[[SV1]], %[[SV2]], %[[SV3]]
+
+// -----
+
+func @gemm6(%a : memref<?x?xf32>, %b : memref<?x?xf32>, %c : memref<?x?xf32>)
+{
+  linalg.matmul %a, %b, %c {__internal_linalg_transform__ = "distribute6"}
+    : (memref<?x?xf32>, memref<?x?xf32>, memref<?x?xf32>)
+  return
+}
+//  CHECK-DAG: #[[MAP0:.*]] = affine_map<()[s0] -> (s0 * 8)>
+//      CHECK: func @gemm6(
+// CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]*]]: memref<?x?xf32>
+// CHECK-SAME: %[[ARG1:[a-zA-Z0-9_]*]]: memref<?x?xf32>
+// CHECK-SAME: %[[ARG2:[a-zA-Z0-9_]*]]: memref<?x?xf32>
+//      CHECK: %[[T2:.*]] = "gpu.block_id"() {dimension = "y"}
+//      CHECK: %[[T3:.*]] = "gpu.grid_dim"() {dimension = "y"}
+//      CHECK: %[[T4:.*]] = affine.apply #[[MAP0]]()[%[[T2]]]
+//      CHECK: %[[T5:.*]] = affine.apply #[[MAP0]]()[%[[T3]]]
+//      CHECK: %[[T6:.*]] = "gpu.block_id"() {dimension = "x"}
+//      CHECK: scf.parallel (%[[ARG3.*]]) = (%[[T4]]) to (%{{.*}}) step (%[[T5]])
+//      CHECK:   scf.for %[[ARG4:.*]] =
+//      CHECK:     %[[SV1:.*]] = subview %[[ARG0]][%[[ARG3]], %[[ARG4]]]
+//      CHECK:     %[[T14:.*]] = affine.apply #[[MAP0]]()[%[[T6]]]
+//      CHECK:     %[[SV2:.*]] = subview %[[ARG1]][%[[ARG4]], %[[T14]]]
+//      CHECK:     %[[T20:.*]] = affine.apply #[[MAP0]]()[%[[T6]]]
+//      CHECK:     %[[SV3:.*]] = subview %[[ARG2]][%[[ARG3]], %[[T20]]]
+//      CHECK:     linalg.matmul %[[SV1]], %[[SV2]], %[[SV3]]

mlir/test/lib/Transforms/TestLinalgTransforms.cpp

 //===- TestLinalgTransforms.cpp - Test Linalg transformation patterns -----===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 //
 // This file implements logic for testing Linalg transformations.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
+#include "mlir/Dialect/GPU/GPUDialect.h"
 #include "mlir/Dialect/Linalg/IR/LinalgOps.h"
 #include "mlir/Dialect/Linalg/Transforms/Transforms.h"
 #include "mlir/Dialect/Linalg/Utils/Utils.h"
 #include "mlir/Dialect/Vector/VectorOps.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/Pass/Pass.h"
 
 #include "llvm/ADT/SetVector.h"
Option<bool> testMatmulToVectorPatterns2dTiling{
       *this, "test-matmul-to-vector-patterns-tile-2d",
       llvm::cl::desc(
           "Test a fused pass that applies patterns from matmul to vectors via "
           "2-d tiling"),
       llvm::cl::init(false)};
   Option<bool> testPromotionOptions{*this, "test-linalg-promotion-options",
                                     llvm::cl::desc("Test promotion options"),
                                     llvm::cl::init(false)};
+  Option<bool> testTileAndDistributionOptions{
+      *this, "test-tile-and-distribute-options",
+      llvm::cl::desc("Test tile and distribute options"),
+      llvm::cl::init(false)};
   Option<bool> testVectorTransferForwardingPatterns{
       *this, "test-vector-transfer-forwarding-patterns",
       llvm::cl::desc(
           "Test a fused pass that forwards linalg.copy to vector.transfer"),
       llvm::cl::init(false)};
   Option<bool> testGenericToVectorPattern{
       *this, "test-contraction-to-vector-patterns",
       llvm::cl::desc("Test a set of patterns that rewrite a linalg contraction "
static void applyPatterns(FuncOp funcOp) {
   // Linalg to loops patterns.
   //===--------------------------------------------------------------------===//
   patterns.insert<LinalgLoweringPattern<DotOp>>(
       ctx,
       /*loweringType=*/LinalgLoweringType::Loops,
       LinalgMarker(Identifier::get("REG", ctx)));
 
   //===--------------------------------------------------------------------===//
+  // Linalg distribution patterns.
+  //===--------------------------------------------------------------------===//
+  LinalgLoopDistributionOptions distributionOptions;
+
+  //===--------------------------------------------------------------------===//
   // Linalg to vector contraction patterns.
   //===--------------------------------------------------------------------===//
   patterns.insert<LinalgVectorizationPattern<MatmulOp>,
                   LinalgVectorizationPattern<FillOp>,
                   LinalgVectorizationPattern<CopyOp>,
                   LinalgVectorizationPattern<GenericOp>>(
       ctx, LinalgMarker(Identifier::get("VECTORIZE", ctx)));
 
patterns.insert<LinalgPromotionPattern<MatmulOp>>(
               },
               [](OpBuilder &b, Value src, Value dst) -> LogicalResult {
                 copyCallBackFn(b, src, dst, true);
                 return success();
               }),
       LinalgMarker(Identifier::get("PROMOTE", ctx)));
 }
 
+template <typename IdOp, typename NProcsOp>
+static ProcInfo getGpuProcIds(OpBuilder &b, Location loc, unsigned loopNum) {
+  Type indexType = b.getIndexType();
+  switch (loopNum) {
+  case 0:
+    return {b.create<IdOp>(loc, indexType, b.getStringAttr("y")),
+            b.create<NProcsOp>(loc, indexType, b.getStringAttr("y"))};
+  case 1:
+    return {b.create<IdOp>(loc, indexType, b.getStringAttr("x")),
+            b.create<NProcsOp>(loc, indexType, b.getStringAttr("x"))};
+  default:
+    llvm_unreachable("test patterns handles only upto 2-level nested loops");
+  }
+  return {nullptr, nullptr};
+}
+
+static void fillTileAndDistributePatterns(MLIRContext *context,
+                                          OwningRewritePatternList &patterns) {
+  {
+    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::ParallelLoops)
+            .setDistributionOptions(cyclicNprocsEqNiters),
+        LinalgMarker(Identifier::get("distribute1", context),
+                     Identifier::get("after_distribute1", context)));
+  }
+
+  {
+    LinalgLoopDistributionOptions cyclicNprocsGeNiters;
+    cyclicNprocsGeNiters.distributionMethod.resize(
+        2, DistributionMethod::CyclicNumProcsGeNumIters);
+    cyclicNprocsGeNiters.procInfo =
+        getGpuProcIds<gpu::BlockIdOp, gpu::GridDimOp>;
+    patterns.insert<LinalgTilingPattern<MatmulOp>>(
+        context,
+        LinalgTilingOptions()
+            .setTileSizes({8, 8, 4})
+            .setLoopType(LinalgTilingLoopType::ParallelLoops)
+            .setDistributionOptions(cyclicNprocsGeNiters),
+        LinalgMarker(Identifier::get("distribute2", context),
+                     Identifier::get("after_distribute2", context)));
+  }
+
+  {
+    LinalgLoopDistributionOptions cyclicNprocsDefault;
+    cyclicNprocsDefault.distributionMethod.resize(2,
+                                                  DistributionMethod::Cyclic);
+    cyclicNprocsDefault.procInfo =
+        getGpuProcIds<gpu::BlockIdOp, gpu::GridDimOp>;
+    patterns.insert<LinalgTilingPattern<MatmulOp>>(
+        context,
+        LinalgTilingOptions()
+            .setTileSizes({8, 8, 4})
+            .setLoopType(LinalgTilingLoopType::ParallelLoops)
+            .setDistributionOptions(cyclicNprocsDefault),
+        LinalgMarker(Identifier::get("distribute3", context),
+                     Identifier::get("after_distribute3", context)));
+  }
+
+  {
+    LinalgLoopDistributionOptions cyclicNprocsMixed1;
+    cyclicNprocsMixed1.distributionMethod = {
+        DistributionMethod::CyclicNumProcsEqNumIters,
+        DistributionMethod::CyclicNumProcsGeNumIters};
+    cyclicNprocsMixed1.procInfo = getGpuProcIds<gpu::BlockIdOp, gpu::GridDimOp>;
+    patterns.insert<LinalgTilingPattern<MatmulOp>>(
+        context,
+        LinalgTilingOptions()
+            .setTileSizes({8, 8, 4})
+            .setLoopType(LinalgTilingLoopType::ParallelLoops)
+            .setDistributionOptions(cyclicNprocsMixed1),
+        LinalgMarker(Identifier::get("distribute4", context),
+                     Identifier::get("after_distribute4", context)));
+  }
+
+  {
+    LinalgLoopDistributionOptions cyclicNprocsMixed2;
+    cyclicNprocsMixed2.distributionMethod = {
+        DistributionMethod::CyclicNumProcsGeNumIters,
+        DistributionMethod::Cyclic};
+    cyclicNprocsMixed2.procInfo = getGpuProcIds<gpu::BlockIdOp, gpu::GridDimOp>;
+    patterns.insert<LinalgTilingPattern<MatmulOp>>(
+        context,
+        LinalgTilingOptions()
+            .setTileSizes({8, 8, 4})
+            .setLoopType(LinalgTilingLoopType::ParallelLoops)
+            .setDistributionOptions(cyclicNprocsMixed2),
+        LinalgMarker(Identifier::get("distribute5", context),
+                     Identifier::get("after_distribute5", context)));
+  }
+
+  {
+    LinalgLoopDistributionOptions cyclicNprocsMixed3;
+    cyclicNprocsMixed3.distributionMethod = {
+        DistributionMethod::Cyclic,
+        DistributionMethod::CyclicNumProcsEqNumIters};
+    cyclicNprocsMixed3.procInfo = getGpuProcIds<gpu::BlockIdOp, gpu::GridDimOp>;
+
+    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)));
+  }
+}
+
 static void
 applyMatmulToVectorPatterns(FuncOp funcOp,
                             bool testMatmulToVectorPatterns1dTiling,
                             bool testMatmulToVectorPatterns2dTiling) {
   MLIRContext *ctx = funcOp.getContext();
   SmallVector<OwningRewritePatternList, 4> stage1Patterns;
   if (testMatmulToVectorPatterns1dTiling) {
     fillL1TilingAndMatmulToVectorPatterns(funcOp, Identifier::get("START", ctx),
void TestLinalgTransforms::runOnFunction() {
   std::unique_ptr<void, decltype(lambda)> cleanupGuard{(void *)1, lambda};
 
   if (testPromotionOptions) {
     OwningRewritePatternList patterns;
     fillPromotionCallBackPatterns(&getContext(), patterns);
     applyPatternsAndFoldGreedily(getFunction(), patterns);
     return;
   }
+  if (testTileAndDistributionOptions) {
+    OwningRewritePatternList patterns;
+    fillTileAndDistributePatterns(&getContext(), patterns);
+    applyPatternsAndFoldGreedily(getFunction(), patterns);
+    return;
+  }
   if (testPatterns)
     return applyPatterns(getFunction());
   if (testMatmulToVectorPatterns1dTiling || testMatmulToVectorPatterns2dTiling)
     return applyMatmulToVectorPatterns(getFunction(),
                                        testMatmulToVectorPatterns1dTiling,
                                        testMatmulToVectorPatterns2dTiling);
   if (testVectorTransferForwardingPatterns)
     return applyVectorTransferForwardingPatterns(getFunction());