[mlir][tensor][transform][WIP] Run test cases via registered patterns

[springerm]

Summary

Proof of concept. The registered patterns infrastructure can be used to reduce the amount of boilerplate code in TestTensorTransforms.cpp.

mlir/include/mlir/Dialect/Transform/IR/TransformOps.h

 #include "mlir/IR/PatternMatch.h"
 #include "mlir/IR/SymbolTable.h"
 #include "mlir/Interfaces/CallInterfaces.h"
 #include "mlir/Interfaces/CastInterfaces.h"
 #include "mlir/Interfaces/ControlFlowInterfaces.h"
 
 namespace mlir {
 namespace transform {
-class ApplyPatternsOp;
 
 enum class FailurePropagationMode : uint32_t;
 class FailurePropagationModeAttr;
 
 /// A builder function that populates the body of a SequenceOp.
 using SequenceBodyBuilderFn = ::llvm::function_ref<void(
     ::mlir::OpBuilder &, ::mlir::Location, ::mlir::BlockArgument)>;
 using SequenceBodyBuilderArgsFn =
public:
 
   /// A function that populates a `RewritePatternSet`.
   using PopulatePatternsFn = std::function<void(RewritePatternSet &)>;
 
   /// Registers patterns with the specified identifier. The identifier should
   /// be prefixed with the dialect to which the patterns belong.
   void registerPatterns(StringRef identifier, PopulatePatternsFn &&fn);
 
-protected:
-  friend class ApplyPatternsOp;
-
   /// Returns "true" if patterns are registered with the specified identifier.
   bool hasPatterns(StringAttr identifier) const;
 
   /// Populates the given pattern set with the specified patterns.
   void populatePatterns(StringAttr identifier,
                         RewritePatternSet &patternSet) const;
 
 private:

mlir/lib/Dialect/Tensor/TransformOps/TensorTransformOps.cpp

addDialectDataInitializer<transform::PatternRegistry>(
           registry.registerPatterns(
               "tensor.reassociative_reshape_folding",
               tensor::populateReassociativeReshapeFoldingPatterns);
           registry.registerPatterns("tensor.fold_tensor_empty",
                                     tensor::populateFoldTensorEmptyPatterns);
           registry.registerPatterns(
               "tensor.fold_into_pack_and_unpack",
               tensor::populateFoldIntoPackAndUnpackPatterns);
+          registry.registerPatterns("tensor.simplify_tensor_pack",
+                                    tensor::populateSimplifyTensorPack);
         });
   }
 };
 } // namespace
 
 #define GET_OP_CLASSES
 #include "mlir/Dialect/Tensor/TransformOps/TensorTransformOps.cpp.inc"
 
 void mlir::tensor::registerTransformDialectExtension(
     DialectRegistry &registry) {
   registry.addExtensions<TensorTransformDialectExtension>();
 }

mlir/test/Dialect/Tensor/drop-redundant-insert-slice-rank-expansion.mlir

-// RUN: mlir-opt -split-input-file -test-tensor-transform-patterns=test-drop-redundant-insert-slice-rank-expansion %s | FileCheck %s
+// RUN: mlir-opt -split-input-file -test-registered-patterns="patterns=tensor.drop_redundant_insert_slice_rank_expansion" %s | FileCheck %s
 
 // CHECK-LABEL: func @test_drop_rank_expansion(
 //  CHECK-SAME:     %[[src:.*]]: tensor<128x480xf32>,
 //       CHECK:   %[[extract:.*]] = tensor.extract_slice %[[src]][0, 0] [123, 456] [1, 1] : tensor<128x480xf32> to tensor<123x456xf32>
 //       CHECK:   return %[[extract]]
 func.func @test_drop_rank_expansion(%src: tensor<128x480xf32>, %dest: tensor<1x1x128x480xf32>) -> tensor<123x456xf32> {
   %inserted_slice = tensor.insert_slice %src into %dest[0, 0, 0, 0] [1, 1, 128, 480] [1, 1, 1, 1] : tensor<128x480xf32> into tensor<1x1x128x480xf32>
   %extracted_slice = tensor.extract_slice %inserted_slice[0, 0, 0, 0] [1, 1, 123, 456] [1, 1, 1, 1] : tensor<1x1x128x480xf32> to tensor<123x456xf32>
   return %extracted_slice : tensor<123x456xf32>
 }

mlir/test/Dialect/Tensor/fold-consecutive-insert-extract-slice.mlir

-// RUN: mlir-opt -split-input-file -test-tensor-transform-patterns=test-fold-consecutive-insert-extract-slice -canonicalize -mlir-print-local-scope %s | FileCheck %s
+// RUN: mlir-opt -split-input-file -test-registered-patterns="patterns=tensor.merge_consecutive_insert_extract_slice" -canonicalize -mlir-print-local-scope %s | FileCheck %s
 
 func.func @extract_slice_same_rank(
     %src: tensor<?x?x?x?xf32>, %offset0: index, %offset1: index, %size0: index, %size1: index) -> tensor<8x16x32x?xf32> {
   %0 = tensor.extract_slice %src[0, 1, 2, %offset0] [128, 128, 128, %size0] [1, 1, 1, 1] : tensor<?x?x?x?xf32> to tensor<128x128x128x?xf32>
   %1 = tensor.extract_slice %0[7, 8, 9, %offset1] [8, 16, 32, %size1] [1, 1, 1, 1] : tensor<128x128x128x?xf32> to tensor<8x16x32x?xf32>
   return %1: tensor<8x16x32x?xf32>
 }
 

mlir/test/Dialect/Tensor/fold-empty-op.mlir

-// RUN: mlir-opt -split-input-file -test-tensor-transform-patterns=test-empty-op-folding  %s | FileCheck %s
+// RUN: mlir-opt -split-input-file -test-registered-patterns="patterns=tensor.fold_tensor_empty"  %s | FileCheck %s
 
 func.func @empty_reshape_expansion(%arg0 : index) -> tensor<2x3x5x4x?x7xf32> {
   %0 = tensor.empty(%arg0) : tensor<6x5x?xf32>
   %1 = tensor.expand_shape %0 [[0, 1], [2], [3, 4, 5]]
       : tensor<6x5x?xf32> into tensor<2x3x5x4x?x7xf32>
   return %1 : tensor<2x3x5x4x?x7xf32>
 }
 //      CHECK: #[[MAP:.+]] = affine_map<()[s0] -> (s0 floordiv 28)>

mlir/test/Dialect/Tensor/fold-into-pack-and-unpack.mlir

-// RUN: mlir-opt -split-input-file -test-tensor-transform-patterns=test-fold-into-pack-and-unpack  %s | FileCheck %s
+// RUN: mlir-opt -split-input-file -test-registered-patterns="patterns=tensor.fold_into_pack_and_unpack"  %s | FileCheck %s
 
 func.func @fold_unpack_slice(%arg0 : tensor<?x?x8x4xf32>, %arg1 : tensor<?x?xf32>,
     %arg2 : index, %arg3 : index) -> tensor<?x?xf32> {
   %0 = tensor.unpack %arg0 inner_dims_pos = [0, 1] inner_tiles = [8, 4] into %arg1
       : tensor<?x?x8x4xf32> -> tensor<?x?xf32>
   %1 = tensor.extract_slice %0[0, 0] [%arg2, %arg3] [1, 1] : tensor<?x?xf32> to tensor<?x?xf32>
   return %1 : tensor<?x?xf32>
 }

mlir/test/Dialect/Tensor/fold-reassociative-reshapes.mlir

-// RUN: mlir-opt -split-input-file -test-tensor-transform-patterns=test-reassociative-reshape-folding %s | FileCheck %s
+// RUN: mlir-opt -split-input-file -test-registered-patterns="patterns=tensor.reassociative_reshape_folding" %s | FileCheck %s
 
 // CHECK-LABEL: func @expand_shape_of_rank_reducing_extract(
 //  CHECK-SAME:     %[[t:.*]]: tensor<?x?x?x?xf32>
 //   CHECK-DAG:   %[[extract1:.*]] = tensor.extract_slice %{{.*}}[0, 0, 0, 0] [%{{.*}}, 1, 1, 5] [1, 1, 1, 1] : tensor<?x?x?x?xf32> to tensor<?x1x1x5xf32>
 //   CHECK-DAG:   %[[extract2:.*]] = tensor.extract_slice %{{.*}}[0, 0, 0, 0] [%{{.*}}, 1, 1, 5] [1, 1, 1, 1] : tensor<?x?x?x?xf32> to tensor<?x1x1x5xf32>
 //       CHECK:   return %[[extract1]], %[[extract2]]
 func.func @expand_shape_of_rank_reducing_extract(
     %t: tensor<?x?x?x?xf32>, %idx: index)

mlir/test/Dialect/Tensor/simplify-tensor-pack.mlir

-// RUN: mlir-opt -split-input-file -test-tensor-transform-patterns="test-simplify-pack-patterns" %s | FileCheck %s
+// RUN: mlir-opt -split-input-file -test-registered-patterns="patterns=tensor.simplify_tensor_pack" %s | FileCheck %s
 
 // CHECK: func.func @single_dim_packing(
 // CHECK-SAME: %[[ARG0:.+]]: tensor<256xf32>)
 // CHECK: %[[EXPANDED:.+]] = tensor.expand_shape %[[ARG0]] {{\[}}[0, 1]] : tensor<256xf32> into tensor<8x32xf32>
 // CHECK: return %[[EXPANDED]] : tensor<8x32xf32>
 func.func @single_dim_packing(%arg0: tensor<256xf32>) -> tensor<8x32xf32> {
   %empty = tensor.empty() : tensor<8x32xf32>
   %0 = tensor.pack %arg0 inner_dims_pos = [0] inner_tiles = [32] into %empty : tensor<256xf32> -> tensor<8x32xf32>

mlir/test/lib/Dialect/Tensor/TestTensorTransforms.cpp

struct TestTensorTransforms
 
   void runOnOperation() override;
 
   Option<bool> testFoldConstantExtractSlice{
       *this, "test-fold-constant-extract-slice",
       llvm::cl::desc("Test folding arith.constant and tensor.extract_slice"),
       llvm::cl::init(false)};
 
-  Option<bool> testFoldConsecutiveInsertExtractSlice{
-      *this, "test-fold-consecutive-insert-extract-slice",
-      llvm::cl::desc(
-          "Test folding consecutive tensor.insert_slice/tensor.extract_slice"),
-      llvm::cl::init(false)};
-
   Option<bool> testRewriteExtractSliceWithTiledCollapseShape{
       *this, "test-rewrite-extract-slice-from-collapse-shape",
       llvm::cl::desc("Test swapping tensor.extract_slice of a collapse_shape "
                      "with loop nest"),
       llvm::cl::init(false)};
 
-  Option<bool> testDropRedundantInsertSliceRankExpansion{
-      *this, "test-drop-redundant-insert-slice-rank-expansion",
-      llvm::cl::desc("Test dropping redundant insert_slice rank expansions"),
-      llvm::cl::init(false)};
-
-  Option<bool> testReassociativeReshapeFolding{
-      *this, "test-reassociative-reshape-folding",
-      llvm::cl::desc("Test folding of expand_shape/collapse_shape"),
-      llvm::cl::init(false)};
-
-  Option<bool> testEmptyOpFolding{
-      *this, "test-empty-op-folding",
-      llvm::cl::desc("Test folding of tensor.empty"), llvm::cl::init(false)};
-
-  Option<bool> testFoldIntoPackAndUnpack{
-      *this, "test-fold-into-pack-and-unpack",
-      llvm::cl::desc("Test folding ops into tensor.pack and tensor.unpack"),
-      llvm::cl::init(false)};
-
   Option<bool> useForeach{
       *this, "use-foreach",
       llvm::cl::desc(
           "Use the scf.forall operation when generating loop nests for "
           "the extract_slice of collapse_shape pattern"),
       llvm::cl::init(false)};
 
-  Option<bool> testSimplifyPackPatterns{
-      *this, "test-simplify-pack-patterns",
-      llvm::cl::desc("Test patterns to simplify tensor.pack"),
-      llvm::cl::init(false)};
-
   Option<bool> testTrackingListener{
       *this, "test-tracking-listener",
       llvm::cl::desc("Test tensor TrackingListener for the transform dialect"),
       llvm::cl::init(false)};
 };
 } // namespace
 
-static void applyReassociativeReshapeFoldingPatterns(Operation *rootOp) {
-  RewritePatternSet patterns(rootOp->getContext());
-  tensor::populateReassociativeReshapeFoldingPatterns(patterns);
-  (void)applyPatternsAndFoldGreedily(rootOp, std::move(patterns));
-}
-
-static void applyEmptyOpFoldingPatterns(Operation *rootOp) {
-  RewritePatternSet patterns(rootOp->getContext());
-  tensor::populateFoldTensorEmptyPatterns(patterns);
-  (void)applyPatternsAndFoldGreedily(rootOp, std::move(patterns));
-}
-
-static void applyFoldIntoPackAndUnpackPatterns(Operation *rootOp) {
-  RewritePatternSet patterns(rootOp->getContext());
-  tensor::populateFoldIntoPackAndUnpackPatterns(patterns);
-  (void)applyPatternsAndFoldGreedily(rootOp, std::move(patterns));
-}
-
 static void applyFoldConstantExtractSlicePatterns(Operation *rootOp) {
   RewritePatternSet patterns(rootOp->getContext());
   tensor::ControlConstantExtractSliceFusionFn controlFn =
       [](tensor::ExtractSliceOp op) {
         if (!op.getSource().hasOneUse())
           return false;
 
         auto resultType = cast<ShapedType>(op.getResult().getType());
         constexpr int64_t kConstantFoldingMaxNumElements = 1024;
         return resultType.getNumElements() <= kConstantFoldingMaxNumElements;
       };
 
   tensor::populateFoldConstantExtractSlicePatterns(patterns, controlFn);
   (void)applyPatternsAndFoldGreedily(rootOp, std::move(patterns));
 }
 
-static void applyFoldConsecutiveInsertExtractSlicePatterns(Operation *rootOp) {
-  RewritePatternSet patterns(rootOp->getContext());
-  tensor::populateMergeConsecutiveInsertExtractSlicePatterns(patterns);
-  (void)applyPatternsAndFoldGreedily(rootOp, std::move(patterns));
-}
-
-static void
-applyDropRedundantInsertSliceRankExpansionPatterns(Operation *rootOp) {
-  RewritePatternSet patterns(rootOp->getContext());
-  tensor::populateDropRedundantInsertSliceRankExpansionPatterns(patterns);
-  (void)applyPatternsAndFoldGreedily(rootOp, std::move(patterns));
-}
-
-static void applySimplifyPackPatterns(Operation *rootOp) {
-  RewritePatternSet patterns(rootOp->getContext());
-  tensor::populateSimplifyTensorPack(patterns);
-  (void)applyPatternsAndFoldGreedily(rootOp, std::move(patterns));
-}
-
 namespace {
 /// Base pattern to rewrite  a `tensor.collapse_shape -> tensor.extract_slice`.
 /// The `tensor.extract_slice` is replaced by a loop or gather operation that
 /// stitches together the desired tile from slices of the source of the collapse
 /// shape op.
 struct RewriteExtractSliceFromCollapseShapeBase
     : public OpRewritePattern<tensor::ExtractSliceOp> {
   RewriteExtractSliceFromCollapseShapeBase(MLIRContext *context)
static LogicalResult testTrackingListenerReplacements(Operation *rootOp) {
   }
 
   replacement->emitRemark("replacement found");
   return success();
 }
 
 void TestTensorTransforms::runOnOperation() {
   Operation *rootOp = getOperation();
-  if (testSimplifyPackPatterns)
-    applySimplifyPackPatterns(rootOp);
   if (testFoldConstantExtractSlice)
     applyFoldConstantExtractSlicePatterns(rootOp);
-  if (testFoldConsecutiveInsertExtractSlice)
-    applyFoldConsecutiveInsertExtractSlicePatterns(rootOp);
-  if (testDropRedundantInsertSliceRankExpansion)
-    applyDropRedundantInsertSliceRankExpansionPatterns(rootOp);
-  if (testReassociativeReshapeFolding)
-    applyReassociativeReshapeFoldingPatterns(rootOp);
-  if (testEmptyOpFolding)
-    applyEmptyOpFoldingPatterns(rootOp);
-  if (testFoldIntoPackAndUnpack)
-    applyFoldIntoPackAndUnpackPatterns(rootOp);
   if (testRewriteExtractSliceWithTiledCollapseShape) {
     if (failed(
             applyRewriteExtractFromCollapseShapePatterns(rootOp, useForeach)))
       return signalPassFailure();
   }
   if (testTrackingListener)
     if (failed(testTrackingListenerReplacements(rootOp)))
       return signalPassFailure();

mlir/test/lib/Dialect/Transform/CMakeLists.txt

 set(LLVM_TARGET_DEFINITIONS TestTransformDialectExtension.td)
 mlir_tablegen(TestTransformDialectExtension.h.inc -gen-op-decls)
 mlir_tablegen(TestTransformDialectExtension.cpp.inc -gen-op-defs)
 mlir_tablegen(TestTransformDialectExtensionTypes.h.inc -gen-typedef-decls -typedefs-dialect=transform)
 mlir_tablegen(TestTransformDialectExtensionTypes.cpp.inc -gen-typedef-defs -typedefs-dialect=transform)
 add_public_tablegen_target(MLIRTestTransformDialectExtensionIncGen)
 
 add_mlir_library(MLIRTestTransformDialect
   TestTransformDialectExtension.cpp
   TestTransformDialectInterpreter.cpp
+  TestRegisteredPatterns.cpp
   TestTransformStateExtension.cpp
 
   EXCLUDE_FROM_LIBMLIR
 
   DEPENDS
   MLIRTestTransformDialectExtensionIncGen
 
   LINK_LIBS PUBLIC
   MLIRIR
   MLIRPass
   MLIRPDLDialect
   MLIRTransformDialect
   MLIRTransformDialectTransforms
   MLIRTransformPDLExtension
+  MLIRTransforms
 )

mlir/test/lib/Dialect/Transform/TestRegisteredPatterns.cpp

+//===- TestRegisteredPatterns.cpp -----------------------------------------===//
+//
+// 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/Transform/IR/TransformOps.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+
+using namespace mlir;
+
+namespace {
+struct TestRegisteredPatternsPass
+    : public PassWrapper<TestRegisteredPatternsPass, OperationPass<>> {
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestRegisteredPatternsPass)
+
+  TestRegisteredPatternsPass() = default;
+  TestRegisteredPatternsPass(const TestRegisteredPatternsPass &pass)
+      : PassWrapper(pass) {}
+
+  StringRef getArgument() const final { return "test-registered-patterns"; }
+  StringRef getDescription() const final {
+    return "Test greedy application of registered patterns.";
+  }
+
+  void getDependentDialects(DialectRegistry &registry) const override {
+    registry.insert<transform::TransformDialect>();
+  }
+
+  void runOnOperation() override;
+
+  ListOption<std::string> patterns{
+      *this, "patterns", llvm::cl::desc("Pattens that should be applied")};
+};
+} // namespace
+
+void TestRegisteredPatternsPass::runOnOperation() {
+  MLIRContext *ctx = getOperation()->getContext();
+  Builder builder(ctx);
+  const auto &registry = ctx->getLoadedDialect<transform::TransformDialect>()
+                             ->getExtraData<transform::PatternRegistry>();
+  RewritePatternSet set(ctx);
+  for (const std::string &identifier : patterns) {
+    registry.populatePatterns(builder.getStringAttr(identifier), set);
+  }
+
+  LogicalResult result =
+      applyPatternsAndFoldGreedily(getOperation(), std::move(set));
+  if (failed(result)) {
+    getOperation()->emitOpError("greedy pattern application failed");
+    return signalPassFailure();
+  }
+}
+
+namespace mlir {
+namespace test {
+void registerTestRegisteredPatternsPass() {
+  PassRegistration<TestRegisteredPatternsPass> reg;
+}
+} // namespace test
+} // namespace mlir

mlir/tools/mlir-opt/mlir-opt.cpp

 void registerTestMemRefStrideCalculation();
 void registerTestOneToNTypeConversionPass();
 void registerTestOpaqueLoc();
 void registerTestPadFusion();
 void registerTestPDLByteCodePass();
 void registerTestPDLLPasses();
 void registerTestPreparationPassWithAllowedMemrefResults();
 void registerTestRecursiveTypesPass();
+void registerTestRegisteredPatternsPass();
 void registerTestSCFUtilsPass();
 void registerTestSCFWhileOpBuilderPass();
 void registerTestShapeMappingPass();
 void registerTestSliceAnalysisPass();
 void registerTestTensorCopyInsertionPass();
 void registerTestTensorTransforms();
 void registerTestTilingInterface();
 void registerTestTopologicalSortAnalysisPass();
#endif
   mlir::test::registerTestMemRefDependenceCheck();
   mlir::test::registerTestMemRefStrideCalculation();
   mlir::test::registerTestOneToNTypeConversionPass();
   mlir::test::registerTestOpaqueLoc();
   mlir::test::registerTestPadFusion();
   mlir::test::registerTestPDLByteCodePass();
   mlir::test::registerTestPDLLPasses();
   mlir::test::registerTestRecursiveTypesPass();
+  mlir::test::registerTestRegisteredPatternsPass();
   mlir::test::registerTestSCFUtilsPass();
   mlir::test::registerTestSCFWhileOpBuilderPass();
   mlir::test::registerTestShapeMappingPass();
   mlir::test::registerTestSliceAnalysisPass();
   mlir::test::registerTestTensorCopyInsertionPass();
   mlir::test::registerTestTensorTransforms();
   mlir::test::registerTestTilingInterface();
   mlir::test::registerTestTopologicalSortAnalysisPass();

utils/bazel/llvm-project-overlay/mlir/test/BUILD.bazel

deps = [
         "//llvm:Support",
         "//mlir:BytecodeOpInterface",
         "//mlir:IR",
         "//mlir:PDLDialect",
         "//mlir:Pass",
         "//mlir:TransformDialect",
         "//mlir:TransformDialectTransforms",
         "//mlir:TransformPDLExtension",
+        "//mlir:Transforms",
     ],
 )
 
 cc_library(
     name = "TestDialect",
     srcs = glob(["lib/Dialect/Test/*.cpp"]),
     hdrs = glob(["lib/Dialect/Test/*.h"]),
     includes = [