[mlir][Linalg] Fix and properly test CodegenStrategy API

[nicolasvasilache]

Summary

Fix a bug that was introduced where calling the codegen strategy with actual concrete C++ Op types did not trigger the expected behavior.
Also introduce a test for the behavior that was missing.

[bkramer]

thanks

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

 /// SFINAE: Enqueue helper for ConcreteOpType that have a `getOperationName`.
 template <template <typename> class PatternType, typename ConcreteOpType,
           typename OptionsType,
           typename = std::enable_if_t<std::is_member_function_pointer<
               decltype(&ConcreteOpType::getOperationName)>::value>>
 void sfinae_enqueue(OwningRewritePatternList &patterList, OptionsType options,
                     MLIRContext *context, StringRef opName,
                     linalg::LinalgTransformationFilter m) {
-  assert(opName.empty() ||
-         opName == ConcreteOpType::getOperationName() &&
-             "explicit name must match ConcreteOpType::getOperationName");
+  assert(opName == ConcreteOpType::getOperationName() &&
   patterList.insert<PatternType<ConcreteOpType>>(context, options, m);
 }
 
 /// SFINAE: Enqueue helper for OpType that do not have a `getOperationName`
 /// (e.g. LinalgOp, other interfaces, Operation*).
 template <template <typename> class PatternType, typename OpType,
           typename OptionsType>
 void sfinae_enqueue(OwningRewritePatternList &patterList, OptionsType options,
                     MLIRContext *context, StringRef opName,
                     linalg::LinalgTransformationFilter m) {
   assert(!opName.empty() && "opName must not be empty");
   patterList.insert<PatternType<OpType>>(opName, context, options, m);
 }
 
 /// Promotion transformation enqueues a particular stage-1 pattern for
 /// `Tile<LinalgOpType>`with the appropriate `options`.
 template <typename LinalgOpType>
 struct Tile : public Transformation {
   explicit Tile(linalg::LinalgTilingOptions options,
                 linalg::LinalgTransformationFilter::FilterFunction f = nullptr)
-      : Transformation(f), opName(""), options(options) {}
+      : Transformation(f), opName(LinalgOpType::getOperationName()),
+        options(options) {}
 
   Tile(StringRef name, linalg::LinalgTilingOptions options,
        linalg::LinalgTransformationFilter::FilterFunction f = nullptr)
       : Transformation(f), opName(name), options(options) {}
 
   OwningRewritePatternList
   buildRewritePatterns(MLIRContext *context,
                        linalg::LinalgTransformationFilter m) override {

 
 /// Promotion transformation enqueues a particular stage-1 pattern for
 /// `Promote<LinalgOpType>`with the appropriate `options`.
 template <typename LinalgOpType>
 struct Promote : public Transformation {
   explicit Promote(
       linalg::LinalgPromotionOptions options,
       linalg::LinalgTransformationFilter::FilterFunction f = nullptr)
-      : Transformation(f), opName(""), options(options) {}
+      : Transformation(f), opName(LinalgOpType::getOperationName()),
+        options(options) {}
 
   Promote(StringRef name, linalg::LinalgPromotionOptions options,
           linalg::LinalgTransformationFilter::FilterFunction f = nullptr)
       : Transformation(f), opName(name), options(options) {}
 
   OwningRewritePatternList
   buildRewritePatterns(MLIRContext *context,
                        linalg::LinalgTransformationFilter m) override {

 /// Vectorization transformation enqueues a particular stage-1 pattern for
 /// `LinalgVectorizationPattern<LinalgOpType>` as well as copy to vector
 /// transfer rewrite forwarding patterns.
 template <typename LinalgOpType>
 struct Vectorize : public Transformation {
   explicit Vectorize(
       linalg::LinalgVectorizationOptions options,
       linalg::LinalgTransformationFilter::FilterFunction f = nullptr)
-      : Transformation(f), opName(""), options(options) {}
+      : Transformation(f), opName(LinalgOpType::getOperationName()),
+        options(options) {}
 
   Vectorize(StringRef name, linalg::LinalgVectorizationOptions options,
             linalg::LinalgTransformationFilter::FilterFunction f = nullptr)
       : Transformation(f), opName(name), options(options) {}
 
   OwningRewritePatternList
   buildRewritePatterns(MLIRContext *context,
                        linalg::LinalgTransformationFilter m) override {

mlir/test/Dialect/Linalg/codegen-strategy.mlir

+// Test that both anchor-op name and MatmulOp-based codegen strategy produce the same result.
+// RUN: mlir-opt %s -test-linalg-codegen-strategy="tile-sizes=2,4,8 vectorize vectorize-contraction-to=matrixintrinsics unroll-vector-transfers=true" | FileCheck %s
+// RUN: mlir-opt %s -test-linalg-codegen-strategy="tile-sizes=16,32,64 promote promote-full-tile-pad register-tile-sizes=2,4,8 vectorize vectorize-contraction-to=outerproduct split-transfers=true unroll-vector-transfers=false" | FileCheck %s --check-prefix=OUTER
 // RUN: mlir-opt %s -test-linalg-codegen-strategy="anchor-op=linalg.matmul tile-sizes=2,4,8 vectorize vectorize-contraction-to=matrixintrinsics unroll-vector-transfers=true" | FileCheck %s
 // RUN: mlir-opt %s -test-linalg-codegen-strategy="anchor-op=linalg.matmul tile-sizes=16,32,64 promote promote-full-tile-pad register-tile-sizes=2,4,8 vectorize vectorize-contraction-to=outerproduct split-transfers=true unroll-vector-transfers=false" | FileCheck %s --check-prefix=OUTER
 
 // CHECK-LABEL: func @matmul(
 // OUTER-LABEL: func @matmul(
 func @matmul(%A: memref<1584x1584xf32>, %B: memref<1584x1584xf32>, %C: memref<1584x1584xf32>) {
   linalg.matmul
    ins(%A, %B: memref<1584x1584xf32>, memref<1584x1584xf32>)

mlir/test/lib/Transforms/TestLinalgCodegenStrategy.cpp

void getDependentDialects(DialectRegistry &registry) const override {
     // clang-format on
   }
 
   template <typename LinalgNamedOp>
   void applyStrategyToNamedLinalgOp();
 
   void runOnFunction() override;
 
+  template <typename OpType>
+  void runStrategy(LinalgTilingOptions tilingOptions,
+                   LinalgTilingOptions registerTilingOptions,
+                   vector::VectorContractLowering vectorContractLowering,
+                   vector::VectorTransferSplit vectorTransferSplit);
+
   ListOption<int64_t> tileSizes{*this, "tile-sizes",
                                 llvm::cl::MiscFlags::CommaSeparated,
                                 llvm::cl::desc("Specifies the tile sizes.")};
   Option<bool> promote{
       *this, "promote",
       llvm::cl::desc("Promote the tile into a small aligned memory buffer."),
       llvm::cl::init(false)};
   Option<bool> promoteFullTile{
Option<std::string> anchorOpName{
           "\tlinalg.matmul: anchor on linalg.matmul\n"
           "\tlinalg.matmul_column_major: anchor on linalg.matmul_column_major\n"
           "\tlinalg.copy: anchor on linalg.copy\n"
           "\tlinalg.fill: anchor on linalg.fill\n"),
       llvm::cl::init("")};
 };
 } // end anonymous namespace
 
+template <>
+void TestLinalgCodegenStrategy::runStrategy<LinalgOp>(
+    LinalgTilingOptions tilingOptions,
+    LinalgTilingOptions registerTilingOptions,
+    vector::VectorContractLowering vectorContractLowering,
+    vector::VectorTransferSplit vectorTransferSplit) {
+  assert(!anchorOpName.empty());
+  CodegenStrategy strategy;
+  strategy.tileIf<LinalgOp>(!tileSizes.empty(), anchorOpName, tilingOptions)
+      .promoteIf<LinalgOp>(promote, anchorOpName,
+                           LinalgPromotionOptions()
+                               .setAlignment(16)
+                               .setUseFullTileBuffersByDefault(promoteFullTile))
+      .tileIf<LinalgOp>(!registerTileSizes.empty(), anchorOpName,
+                        registerTilingOptions)
+      .promoteIf<LinalgOp>(
+          registerPromote, anchorOpName,
+          LinalgPromotionOptions()
+              .setAlignment(16)
+              .setUseFullTileBuffersByDefault(registerPromoteFullTile))
+      .vectorizeIf<LinalgOp>(vectorize, anchorOpName)
+      .setVectorTransformsOptions(
+          vector::VectorTransformsOptions()
+              .setVectorTransformsOptions(vectorContractLowering)
+              .setVectorTransferSplit(vectorTransferSplit))
+      .setVectorTransferToSCFOptions(
+          VectorTransferToSCFOptions().setUnroll(unrollVectorTransfers));
+  strategy.transform(getFunction());
+}
+
+template <typename OpType>
+void TestLinalgCodegenStrategy::runStrategy(
+    LinalgTilingOptions tilingOptions,
+    LinalgTilingOptions registerTilingOptions,
+    vector::VectorContractLowering vectorContractLowering,
+    vector::VectorTransferSplit vectorTransferSplit) {
+  CodegenStrategy strategy;
+  strategy.tileIf<OpType>(!tileSizes.empty(), tilingOptions)
+      .template promoteIf<OpType>(
+          promote, LinalgPromotionOptions()
+                       .setAlignment(16)
+                       .setUseFullTileBuffersByDefault(promoteFullTile))
+      .template tileIf<OpType>(!registerTileSizes.empty(),
+                               registerTilingOptions)
+      .template promoteIf<OpType>(
+          registerPromote,
+          LinalgPromotionOptions()
+              .setAlignment(16)
+              .setUseFullTileBuffersByDefault(registerPromoteFullTile))
+      .template vectorizeIf<OpType>(vectorize)
+      .setVectorTransformsOptions(
+          vector::VectorTransformsOptions()
+              .setVectorTransformsOptions(vectorContractLowering)
+              .setVectorTransferSplit(vectorTransferSplit))
+      .setVectorTransferToSCFOptions(
+          VectorTransferToSCFOptions().setUnroll(unrollVectorTransfers));
+  strategy.transform(getFunction());
+}
+
 /// Apply transformations specified as patterns.
 void TestLinalgCodegenStrategy::runOnFunction() {
-  linalg::LinalgTransformationFilter::FilterFunction filterOpName =
-      [&](Operation *op) -> LogicalResult {
-    return success(op->getName().getStringRef() == anchorOpName);
-  };
   LinalgTilingOptions tilingOptions;
   if (!tileSizes.empty())
     tilingOptions = tilingOptions.setTileSizes(tileSizes);
 
   LinalgTilingOptions registerTilingOptions;
   if (!registerTileSizes.empty())
     registerTilingOptions =
         registerTilingOptions.setTileSizes(registerTileSizes);
 
   vector::VectorContractLowering vectorContractLowering =
       llvm::StringSwitch<vector::VectorContractLowering>(
           vectorizeContractionTo.getValue())
           .Case("matrixintrinsics", vector::VectorContractLowering::Matmul)
           .Case("dot", vector::VectorContractLowering::Dot)
           .Case("outerproduct", vector::VectorContractLowering::OuterProduct)
           .Default(vector::VectorContractLowering::OuterProduct);
   vector::VectorTransferSplit vectorTransferSplit =
       llvm::StringSwitch<vector::VectorTransferSplit>(
           splitVectorTransfersTo.getValue())
           .Case("none", vector::VectorTransferSplit::None)
           .Case("linalg-copy", vector::VectorTransferSplit::LinalgCopy)
           .Case("vector-transfers", vector::VectorTransferSplit::VectorTransfer)
           .Default(vector::VectorTransferSplit::None);
 
-  CodegenStrategy strategy;
-  strategy.tileIf<LinalgOp>(!tileSizes.empty(), anchorOpName, tilingOptions)
-      .promoteIf<LinalgOp>(promote, anchorOpName,
-                           LinalgPromotionOptions()
-                               .setAlignment(16)
-                               .setUseFullTileBuffersByDefault(promoteFullTile),
-                           filterOpName)
-      .tileIf<LinalgOp>(!registerTileSizes.empty(), anchorOpName,
-                        registerTilingOptions)
-      .promoteIf<LinalgOp>(
-          registerPromote, anchorOpName,
-          LinalgPromotionOptions()
-              .setAlignment(16)
-              .setUseFullTileBuffersByDefault(registerPromoteFullTile))
-      .vectorizeIf<LinalgOp>(vectorize, anchorOpName)
-      .setVectorTransformsOptions(
-          vector::VectorTransformsOptions()
-              .setVectorTransformsOptions(vectorContractLowering)
-              .setVectorTransferSplit(vectorTransferSplit))
-      .setVectorTransferToSCFOptions(
-          VectorTransferToSCFOptions().setUnroll(unrollVectorTransfers));
-  strategy.transform(getFunction());
+  // If no anchorOpNameis specified, just test that strategy applies properly to
+  // linalg::MatmulOp.
+  if (anchorOpName.empty())
+    runStrategy<linalg::MatmulOp>(tilingOptions, registerTilingOptions,
+                                  vectorContractLowering, vectorTransferSplit);
+  else
+    runStrategy<LinalgOp>(tilingOptions, registerTilingOptions,
+                          vectorContractLowering, vectorTransferSplit);
 }
 
 namespace mlir {
 namespace test {
 void registerTestLinalgCodegenStrategy() {
   PassRegistration<TestLinalgCodegenStrategy> testLinalgCodegenStrategyPass(
       "test-linalg-codegen-strategy", "Test Linalg Codegen Strategy.");
 }
 } // namespace test
 } // namespace mlir