[mlir][linalg] adapt named op generalization to work with captures.

[gysit]

Summary

Instead of always running the region builder check if the generalized op has a region attached. If yes inline the existing region instead of calling the region builder. This change circumvents a problem with named operations that have a region builder taking captures and the generalization pass not knowing about this captures.

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

 #include "mlir/IR/PatternMatch.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/Support/Debug.h"
 
 #define DEBUG_TYPE "linalg-generalization"
 
 using namespace mlir;
+using namespace mlir::linalg;
 
 // Creates a linalg.generic op from the given `namedOp`. Returns a null op if
 // the given `namedOp` does not have a region builder.
-static linalg::GenericOp createGenericOpFromNamedOp(linalg::LinalgOp namedOp,
-                                                    OpBuilder &builder) {
+static GenericOp createGenericOpFromNamedOp(LinalgOp namedOp,
+                                            PatternRewriter &rewriter) {
+  SmallVector<AffineMap> indexingMaps = namedOp.getIndexingMaps();
+  SmallVector<StringRef> iterators = llvm::to_vector<4>(
+      namedOp.iterator_types().getAsValueRange<StringAttr>());
+  SmallVector<RankedTensorType> resultTypes = namedOp.getOutputTensorTypes();
+  SmallVector<Type> types(resultTypes.begin(), resultTypes.end());
+
+  // Inline the existing region if the named operation has a region attached.
+  if (namedOp->getNumRegions() == 1) {
+    GenericOp genericOp = rewriter.create<GenericOp>(
+        namedOp.getLoc(), types, namedOp.getInputs(), namedOp.getOutputs(),
+        indexingMaps, iterators);
+    rewriter.inlineRegionBefore(namedOp->getRegion(0), genericOp.region(),
+                                genericOp.region().begin());
+    return genericOp;
+  }
+
+  // Otherwise use the region builder to generate a new region.
+  // TODO: Remove this path once all linag operations have a region attached.
   auto regionBuilder = namedOp.getRegionBuilder();
   if (!regionBuilder) {
     LLVM_DEBUG(llvm::dbgs() << "no region builder for op: " << namedOp << "\n");
     return nullptr;
   }
-
-  SmallVector<AffineMap, 4> indexingMaps = namedOp.getIndexingMaps();
-  auto iterators = llvm::to_vector<4>(
-      namedOp.iterator_types().getAsValueRange<StringAttr>());
-  auto resultTypes = namedOp.getOutputTensorTypes();
-  SmallVector<Type, 4> types(resultTypes.begin(), resultTypes.end());
-
-  return builder.create<linalg::GenericOp>(
+  return rewriter.create<GenericOp>(
       namedOp.getLoc(), types, namedOp.getInputs(), namedOp.getOutputs(),
       indexingMaps, iterators,
       [&regionBuilder](OpBuilder &bodyBuilder, Location loc, ValueRange) {
         edsc::ScopedContext scope(bodyBuilder, loc);
         regionBuilder(*bodyBuilder.getBlock(), /*captures=*/{});
       });
 }
 
 namespace {
 
 /// Base class for all linalg generalization patterns. A subclass must provide
 /// the following method:
-///   linalg::GenericOp createGenericOp(RootOp, PatternRewriter &)
+///   GenericOp createGenericOp(RootOp, PatternRewriter &)
 /// for creating the generic op.
 // TODO: remove this pattern after migrating all manually-written named ops
 // into auto-generated ones.
 template <typename ConcretePattern, typename RootOp>
 struct LinalgGeneralizationPattern : OpRewritePattern<RootOp> {
   LinalgGeneralizationPattern(MLIRContext *context,
-                              linalg::LinalgTransformationFilter marker,
+                              LinalgTransformationFilter marker,
                               PatternBenefit benefit = 1)
       : OpRewritePattern<RootOp>(context, benefit), marker(std::move(marker)) {}
 
   LogicalResult matchAndRewrite(RootOp rootOp,
                                 PatternRewriter &rewriter) const override {
-    auto linalgOp = dyn_cast<linalg::LinalgOp>(rootOp.getOperation());
+    auto linalgOp = dyn_cast<LinalgOp>(rootOp.getOperation());
     if (!linalgOp)
       return failure();
     if (failed(marker.checkAndNotify(rewriter, linalgOp)))
       return failure();
 
     auto *pattern = static_cast<const ConcretePattern *>(this);
-    linalg::GenericOp genericOp = pattern->createGenericOp(rootOp, rewriter);
+    GenericOp genericOp = pattern->createGenericOp(rootOp, rewriter);
     if (!genericOp)
       return failure();
 
     rewriter.replaceOp(rootOp, genericOp.getResults());
     marker.replaceLinalgTransformationFilter(rewriter,
                                              genericOp.getOperation());
     return success();
   }
 
 private:
-  linalg::LinalgTransformationFilter marker;
+  LinalgTransformationFilter marker;
 };
 
 struct GeneralizeConvOp
-    : public LinalgGeneralizationPattern<GeneralizeConvOp, linalg::ConvOp> {
+    : public LinalgGeneralizationPattern<GeneralizeConvOp, ConvOp> {
   using LinalgGeneralizationPattern::LinalgGeneralizationPattern;
 
-  linalg::GenericOp createGenericOp(linalg::ConvOp, OpBuilder &rewriter) const;
+  GenericOp createGenericOp(ConvOp convOp, OpBuilder &builder) const;
 };
 
 /// Catch-all pattern for converting all named ops with a region builder into
 /// linalg.generic.
 struct LinalgNamedOpGeneralizationPattern : RewritePattern {
   LinalgNamedOpGeneralizationPattern(MLIRContext *context,
-                                     linalg::LinalgTransformationFilter marker,
+                                     LinalgTransformationFilter marker,
                                      PatternBenefit benefit = 1)
       : RewritePattern(MatchAnyOpTypeTag(), benefit, context),
         marker(std::move(marker)) {}
 
   LogicalResult matchAndRewrite(Operation *rootOp,
                                 PatternRewriter &rewriter) const override {
-    auto linalgOp = dyn_cast<linalg::LinalgOp>(rootOp);
+    auto linalgOp = dyn_cast<LinalgOp>(rootOp);
     if (!linalgOp)
       return failure();
     if (failed(marker.checkAndNotify(rewriter, linalgOp)))
       return failure();
 
     // No nothing to do for linalg.generic and linalg.indexed_generic.
-    if (isa<linalg::GenericOp, linalg::IndexedGenericOp>(rootOp))
+    if (isa<GenericOp, IndexedGenericOp>(rootOp))
       return failure();
 
-    linalg::GenericOp genericOp =
-        createGenericOpFromNamedOp(linalgOp, rewriter);
+    GenericOp genericOp = createGenericOpFromNamedOp(linalgOp, rewriter);
     if (!genericOp)
       return failure();
 
     rewriter.replaceOp(rootOp, genericOp.getResults());
     marker.replaceLinalgTransformationFilter(rewriter,
                                              genericOp.getOperation());
     return success();
   }
 
 private:
-  linalg::LinalgTransformationFilter marker;
+  LinalgTransformationFilter marker;
 };
 
 struct LinalgGeneralizationPass
     : public LinalgGeneralizationBase<LinalgGeneralizationPass> {
   void runOnFunction() override;
 };
 
 } // namespace
 
 void LinalgGeneralizationPass::runOnFunction() {
   FuncOp func = getFunction();
   RewritePatternSet patterns(&getContext());
-  linalg::populateLinalgConvGeneralizationPatterns(patterns);
-  linalg::populateLinalgNamedOpsGeneralizationPatterns(patterns);
+  populateLinalgConvGeneralizationPatterns(patterns);
+  populateLinalgNamedOpsGeneralizationPatterns(patterns);
   (void)applyPatternsAndFoldGreedily(func.getBody(), std::move(patterns));
 }
 
-linalg::GenericOp GeneralizeConvOp::createGenericOp(linalg::ConvOp convOp,
-                                                    OpBuilder &builder) const {
+GenericOp GeneralizeConvOp::createGenericOp(ConvOp convOp,
   SmallVector<AffineMap, 4> indexingMaps = convOp.getIndexingMaps();
   auto iterators =
       llvm::to_vector<4>(convOp.iterator_types().getAsValueRange<StringAttr>());
-  return builder.create<linalg::GenericOp>(
+  return builder.create<GenericOp>(
       convOp.getLoc(), /*resultTensorTypes=*/ArrayRef<Type>(),
       convOp.getInputBuffers(), convOp.getOutputBuffers(), indexingMaps,
       iterators,
       [](OpBuilder &bodyBuilder, Location bodyLoc, ValueRange bodyArgs) {
         Value mul =
             bodyBuilder.create<MulFOp>(bodyLoc, bodyArgs[0], bodyArgs[1]);
         Value add = bodyBuilder.create<AddFOp>(bodyLoc, mul, bodyArgs[2]);
-        bodyBuilder.create<linalg::YieldOp>(bodyLoc, add);
+        bodyBuilder.create<YieldOp>(bodyLoc, add);
       });
 }
 
 void mlir::linalg::populateLinalgConvGeneralizationPatterns(
-    RewritePatternSet &patterns, linalg::LinalgTransformationFilter marker) {
+    RewritePatternSet &patterns, LinalgTransformationFilter marker) {
   patterns.add<GeneralizeConvOp>(patterns.getContext(), marker);
 }
 
 void mlir::linalg::populateLinalgNamedOpsGeneralizationPatterns(
-    RewritePatternSet &patterns, linalg::LinalgTransformationFilter marker) {
+    RewritePatternSet &patterns, LinalgTransformationFilter marker) {
   patterns.add<LinalgNamedOpGeneralizationPattern>(patterns.getContext(),
                                                    marker);
 }
 
 std::unique_ptr<OperationPass<FuncOp>> mlir::createLinalgGeneralizationPass() {
   return std::make_unique<LinalgGeneralizationPass>();
 }

mlir/test/Dialect/Linalg/generalize-named-ops.mlir

 // CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel", "parallel", "reduction", "reduction"]}
 // CHECK-SAME: ins(%{{.+}}, %{{.+}} : memref<?x?x?x?xf32>, memref<2x3xf32>)
 // CHECK-SAME: outs(%{{.+}} : memref<?x?x?x?xf32>)
 
 // CHECK:         ^{{.+}}(%[[BBARG0:.+]]: f32, %[[BBARG1:.+]]: f32, %[[BBARG2:.+]]: f32)
 // CHECK-NEXT:      %[[CMP:.+]] = cmpf olt, %[[BBARG0]], %[[BBARG2]] : f32
 // CHECK-NEXT:      %[[RES:.+]] = select %[[CMP]], %[[BBARG0]], %[[BBARG2]] : f32
 // CHECK-NEXT:      linalg.yield %[[RES]] : f32
+
+// -----
+
+func @generalize_fill(%output: memref<?x?xf32>, %value : f32) {
+  linalg.fill(%output, %value) : memref<?x?xf32>, f32
+  return
+}
+
+// CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1) -> (d0, d1)>
+
+// CHECK: func @generalize_fill
+// CHECK-SAME: (%[[ARG0:.+]]: memref<?x?xf32>, %[[VAL:.+]]: f32)
+
+// CHECK: linalg.generic
+// CHECK-SAME: indexing_maps = [#[[MAP0]]]
+// CHECK-SAME: iterator_types = ["parallel", "parallel"]}
+// CHECK-SAME: outs(%{{.+}} : memref<?x?xf32>)
+
+// CHECK:         ^{{.+}}(%[[BBARG0:.+]]: f32)
+// CHECK-NEXT:      linalg.yield %[[VAL]] : f32