[mlir][Linalg] Properly use distribution options.

[nicolasvasilache]

Summary

Let tiling to scf.for actually use the distribution method.
For now only Cyclic is supported.

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

 
 /// 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);
 
 struct LinalgBaseTilingPattern : public RewritePattern {
+  // Entry point to match any LinalgOp OpInterface.
+  LinalgBaseTilingPattern(LinalgTilingOptions options,
+                          LinalgMarker marker = LinalgMarker(),
+                          PatternBenefit benefit = 1);
+  // Entry point to match a specific Linalg op.
   LinalgBaseTilingPattern(StringRef opName, MLIRContext *context,
                           LinalgTilingOptions options,
                           LinalgMarker marker = LinalgMarker(),
                           PatternBenefit benefit = 1);
   LogicalResult
   matchAndRewriteBase(Operation *op, PatternRewriter &rewriter,
                       SmallVectorImpl<Value> &tensorResults) const;
 

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

 
 /// Linalg base tiling pattern.
 mlir::linalg::LinalgBaseTilingPattern::LinalgBaseTilingPattern(
     StringRef opName, MLIRContext *context, LinalgTilingOptions options,
     LinalgMarker marker, PatternBenefit benefit)
     : RewritePattern(opName, {}, benefit, context), marker(marker),
       options(options) {}
 
+mlir::linalg::LinalgBaseTilingPattern::LinalgBaseTilingPattern(
+    LinalgTilingOptions options, LinalgMarker marker, PatternBenefit benefit)
+    : RewritePattern(benefit, MatchAnyOpTypeTag()), marker(marker),
+      options(options) {}
+
 LogicalResult mlir::linalg::LinalgBaseTilingPattern::matchAndRewriteBase(
     Operation *op, PatternRewriter &rewriter,
     SmallVectorImpl<Value> &tensorResults) const {
   LinalgOp linalgOp = dyn_cast<LinalgOp>(op);
   if (!linalgOp)
     return failure();
   if (failed(marker.checkAndNotify(rewriter, linalgOp)))
     return failure();

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

 
 /// 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> distributionOptions) {
-  // Create procInfo so it dominate loops, if appropriate.
+  // Create procInfo so it dominates 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>{});
+    procInfo = distributionOptions->procInfo(builder, loc, loopRanges);
 
   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.
+  // Only supports cyclic distribution for now.
   for (auto it : llvm::zip(loopNest.loops, procInfo,
                            distributionOptions->distributionMethod))
+    if (std::get<2>(it) == DistributionMethod::Cyclic)
       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/test/lib/Transforms/TestLinalgTransforms.cpp

patterns.insert<LinalgTilingPattern<MatmulOp>>(
             .setLoopType(LinalgTilingLoopType::ParallelLoops)
             .setDistributionOptions(cyclicNprocsMixed3),
         LinalgMarker(Identifier::get("distribute6", context),
                      Identifier::get("after_distribute6", context)));
   }
 
   {
     LinalgLoopDistributionOptions cyclicNprocsEqNiters;
-    cyclicNprocsEqNiters.distributionMethod.resize(
-        2, DistributionMethod::CyclicNumProcsEqNumIters);
+    cyclicNprocsEqNiters.distributionMethod.resize(2,
+                                                   DistributionMethod::Cyclic);
     cyclicNprocsEqNiters.procInfo =
         getGpuProcIds<gpu::BlockIdOp, gpu::GridDimOp>;
     patterns.insert<LinalgTilingPattern<MatmulOp>>(
         context,
         LinalgTilingOptions()
             .setTileSizes({8, 8, 4})
             .setLoopType(LinalgTilingLoopType::Loops)
             .setDistributionOptions(cyclicNprocsEqNiters),