[mlir] Transform dialect: separate dependent and generated dialects

[ftynse]

Summary

In the Transform dialect extensions, provide the separate mechanism to
declare dependent dialects (the dialects the transform IR depends on)
and the generated dialects (the dialects the payload IR may be
transformed into). This allows the Transform dialect clients that are
only constructing the transform IR to avoid loading the dialects
relevant for the payload IR along with the Transform dialect itself,
thus decreasing the build/link time.

[ftynse]

This replaces https://reviews.llvm.org/D128530.

[ftynse]

Add Bazel

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

 namespace mlir {
 namespace transform {
 
 /// Base class for extensions of the Transform dialect that supports injecting
 /// operations into the Transform dialect at load time. Concrete extensions are
 /// expected to derive this class and register operations in the constructor.
 /// They can be registered with the DialectRegistry and automatically applied
 /// to the Transform dialect when it is loaded.
+///
+/// Derived classes are expected to define a `void init()` function in which
+/// they can call various protected methods of the base class to register
+/// extension operations and declare their dependencies.
+///
+/// By default, the extension is configured both for construction of the
+/// Transform IR and for its application to some payload. If only the
+/// construction is desired, the extension can be switched to "build-only" mode
+/// that avoids loading the dialects that are only necessary for transforming
+/// the payload. To perform the switch, the extension must be wrapped into the
+/// `BuildOnly` class template (see below) when it is registered, as in:
+///
+///    dialectRegistry.addExtension<BuildOnly<MyTransformDialectExt>>();
+///
+/// instead of:
+///
+///    dialectRegistry.addExtension<MyTransformDialectExt>();
+///
+/// Derived classes must reexport the constructor of this class or otherwise
+/// forward its boolean argument to support this behavior.
 template <typename DerivedTy, typename... ExtraDialects>
 class TransformDialectExtension
     : public DialectExtension<DerivedTy, TransformDialect, ExtraDialects...> {
   using Initializer = std::function<void(TransformDialect *)>;
   using DialectLoader = std::function<void(MLIRContext *)>;
 
 public:
   /// Extension application hook. Actually loads the dependent dialects and
   /// registers the additional operations. Not expected to be called directly.
   void apply(MLIRContext *context, TransformDialect *transformDialect,
              ExtraDialects *...) const final {
     for (const DialectLoader &loader : dialectLoaders)
       loader(context);
+
+    // Only load generated dialects if the user intends to apply
+    // transformations specified by the extension.
+    if (!buildOnly)
+      for (const DialectLoader &loader : generatedDialectLoaders)
+        loader(context);
+
     for (const Initializer &init : opInitializers)
       init(transformDialect);
     transformDialect->mergeInPDLMatchHooks(std::move(pdlMatchConstraintFns));
   }
 
 protected:
+  using Base = TransformDialectExtension<DerivedTy, ExtraDialects...>;
+
+  /// Extension constructor. The argument indicates whether to skip generated
+  /// dialects when applying the extension.
+  explicit TransformDialectExtension(bool buildOnly = false)
+      : buildOnly(buildOnly) {
+    static_cast<DerivedTy *>(this)->init();
+  }
+
+  /// Hook for derived classes to inject constructor behavior.
+  void init() {}
+
   /// Injects the operations into the Transform dialect. The operations must
   /// implement the TransformOpInterface and MemoryEffectsOpInterface, and the
   /// implementations must be already available when the operation is injected.
   template <typename... OpTys>
   void registerTransformOps() {
     opInitializers.push_back([](TransformDialect *transformDialect) {
       transformDialect->addOperationsChecked<OpTys...>();
     });
   }
 
   /// Declares that this Transform dialect extension depends on the dialect
   /// provided as template parameter. When the Transform dialect is loaded,
   /// dependent dialects will be loaded as well. This is intended for dialects
   /// that contain attributes and types used in creation and canonicalization of
-  /// the injected operations.
+  /// the injected operations, similarly to how the dialect definition may list
+  /// dependent dialects. This is *not* intended for dialects entities from
+  /// which may be produced when applying the transformations specified by ops
+  /// registered by this extension.
   template <typename DialectTy>
   void declareDependentDialect() {
     dialectLoaders.push_back(
         [](MLIRContext *context) { context->loadDialect<DialectTy>(); });
   }
 
+  /// Declares that the transformations associated with the operations
+  /// registered by this dialect extension may produce operations from the
+  /// dialect provided as template parameter while processing payload IR that
+  /// does not contain the operations from said dialect. This is similar to
+  /// dependent dialects of a pass. These dialects will be loaded along with the
+  /// transform dialect unless the extension is in the build-only mode.
+  template <typename DialectTy>
+  void declareGeneratedDialect() {
+    generatedDialectLoaders.push_back(
+        [](MLIRContext *context) { context->loadDialect<DialectTy>(); });
+  }
+
   /// Injects the named constraint to make it available for use with the
   /// PDLMatchOp in the transform dialect.
   void registerPDLMatchConstraintFn(StringRef name,
                                     PDLConstraintFunction &&fn) {
     pdlMatchConstraintFns.try_emplace(name,
                                       std::forward<PDLConstraintFunction>(fn));
   }
   template <typename ConstraintFnTy>
   void registerPDLMatchConstraintFn(StringRef name, ConstraintFnTy &&fn) {
     pdlMatchConstraintFns.try_emplace(
         name, ::mlir::detail::pdl_function_builder::buildConstraintFn(
                   std::forward<ConstraintFnTy>(fn)));
   }
 
 private:
   SmallVector<Initializer> opInitializers;
+
+  /// Callbacks loading the dependent dialects, i.e. the dialect needed for the
+  /// extension ops.
   SmallVector<DialectLoader> dialectLoaders;
 
-  /// A list of constraints that should be made availble to PDL patterns
+  /// Callbacks loading the generated dialects, i.e. the dialects produced when
+  /// applying the transformations.
+  SmallVector<DialectLoader> generatedDialectLoaders;
+
+  /// A list of constraints that should be made available to PDL patterns
   /// processed by PDLMatchOp in the Transform dialect.
   ///
   /// Declared as mutable so its contents can be moved in the `apply` const
   /// method, which is only called once.
   mutable llvm::StringMap<PDLConstraintFunction> pdlMatchConstraintFns;
+
+  /// Indicates that the extension is in build-only mode.
+  bool buildOnly;
+};
+
+/// A wrapper for transform dialect extensions that forces them to be
+/// constructed in the build-only mode.
+template <typename DerivedTy>
+class BuildOnly : public DerivedTy {
+public:
+  BuildOnly() : DerivedTy(/*buildOnly=*/true) {}
 };
 
 } // namespace transform
 } // namespace mlir
 
 #endif // MLIR_DIALECT_TRANSFORM_IR_TRANSFORMDIALECT_H

mlir/lib/Dialect/Bufferization/TransformOps/BufferizationTransformOps.cpp

 
 namespace {
 /// Registers new ops and declares PDL as dependent dialect since the additional
 /// ops are using PDL types for operands and results.
 class BufferizationTransformDialectExtension
     : public transform::TransformDialectExtension<
           BufferizationTransformDialectExtension> {
 public:
-  BufferizationTransformDialectExtension() {
-    declareDependentDialect<bufferization::BufferizationDialect>();
+  using Base::Base;
+
+  void init() {
     declareDependentDialect<pdl::PDLDialect>();
-    declareDependentDialect<memref::MemRefDialect>();
+
+    declareGeneratedDialect<bufferization::BufferizationDialect>();
+    declareGeneratedDialect<memref::MemRefDialect>();
+
     registerTransformOps<
 #define GET_OP_LIST
 #include "mlir/Dialect/Bufferization/TransformOps/BufferizationTransformOps.cpp.inc"
         >();
   }
 };
 } // namespace
 
 #define GET_OP_CLASSES
 #include "mlir/Dialect/Bufferization/TransformOps/BufferizationTransformOps.cpp.inc"
 
 void mlir::bufferization::registerTransformDialectExtension(
     DialectRegistry &registry) {
   registry.addExtensions<BufferizationTransformDialectExtension>();
 }

mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp

 
 namespace {
 /// Registers new ops and declares PDL as dependent dialect since the additional
 /// ops are using PDL types for operands and results.
 class LinalgTransformDialectExtension
     : public transform::TransformDialectExtension<
           LinalgTransformDialectExtension> {
 public:
-  LinalgTransformDialectExtension() {
-    declareDependentDialect<AffineDialect>();
-    declareDependentDialect<arith::ArithmeticDialect>();
+  using Base::Base;
+
+  void init() {
     declareDependentDialect<pdl::PDLDialect>();
-    declareDependentDialect<scf::SCFDialect>();
-    declareDependentDialect<vector::VectorDialect>();
+
+    declareGeneratedDialect<AffineDialect>();
+    declareGeneratedDialect<arith::ArithmeticDialect>();
+    declareGeneratedDialect<scf::SCFDialect>();
+    declareGeneratedDialect<vector::VectorDialect>();
+
     registerTransformOps<
 #define GET_OP_LIST
 #include "mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp.inc"
         >();
   }
 };
 } // namespace
 

mlir/lib/Dialect/SCF/TransformOps/SCFTransformOps.cpp

 //===- SCFTransformOps.cpp - Implementation of SCF transformation ops -----===//
 //
 // 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/SCF/TransformOps/SCFTransformOps.h"
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Dialect/PDL/IR/PDL.h"
 #include "mlir/Dialect/SCF/IR/SCF.h"
 #include "mlir/Dialect/SCF/Transforms/Patterns.h"
 #include "mlir/Dialect/SCF/Transforms/Transforms.h"
 #include "mlir/Dialect/SCF/Utils/Utils.h"
 #include "mlir/Dialect/Transform/IR/TransformDialect.h"
 #include "mlir/Dialect/Transform/IR/TransformInterfaces.h"
 #include "mlir/Dialect/Vector/IR/VectorOps.h"
 

 // Transform op registration
 //===----------------------------------------------------------------------===//
 
 namespace {
 class SCFTransformDialectExtension
     : public transform::TransformDialectExtension<
           SCFTransformDialectExtension> {
 public:
-  SCFTransformDialectExtension() {
-    declareDependentDialect<AffineDialect>();
-    declareDependentDialect<func::FuncDialect>();
+  using Base::Base;
+
+  void init() {
+    declareDependentDialect<pdl::PDLDialect>();
+
+    declareGeneratedDialect<AffineDialect>();
+    declareGeneratedDialect<func::FuncDialect>();
+
     registerTransformOps<
 #define GET_OP_LIST
 #include "mlir/Dialect/SCF/TransformOps/SCFTransformOps.cpp.inc"
         >();
   }
 };
 } // namespace
 
 #define GET_OP_CLASSES
 #include "mlir/Dialect/SCF/TransformOps/SCFTransformOps.cpp.inc"
 
 void mlir::scf::registerTransformDialectExtension(DialectRegistry &registry) {
   registry.addExtensions<SCFTransformDialectExtension>();
 }

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

 namespace {
 /// Test extension of the Transform dialect. Registers additional ops and
 /// declares PDL as dependent dialect since the additional ops are using PDL
 /// types for operands and results.
 class TestTransformDialectExtension
     : public transform::TransformDialectExtension<
           TestTransformDialectExtension> {
 public:
-  TestTransformDialectExtension() {
+  using Base::Base;
+
+  void init() {
     declareDependentDialect<pdl::PDLDialect>();
     registerTransformOps<TestTransformOp,
                          TestTransformUnrestrictedOpNoInterface,
 #define GET_OP_LIST
 #include "TestTransformDialectExtension.cpp.inc"
                          >();
   }
 };
 } // namespace
 
 #define GET_OP_CLASSES
 #include "TestTransformDialectExtension.cpp.inc"
 
 void ::test::registerTestTransformDialectExtension(DialectRegistry &registry) {
   registry.addExtensions<TestTransformDialectExtension>();
 }

mlir/unittests/Dialect/CMakeLists.txt

 add_mlir_unittest(MLIRDialectTests
   BroadcastShapeTest.cpp
 )
 target_link_libraries(MLIRDialectTests
   PRIVATE
   MLIRIR
   MLIRDialect)
 
 add_subdirectory(Affine)
 add_subdirectory(LLVMIR)
 add_subdirectory(MemRef)
 add_subdirectory(Quant)
 add_subdirectory(SparseTensor)
 add_subdirectory(SPIRV)
+add_subdirectory(Transform)
 add_subdirectory(Utils)

mlir/unittests/Dialect/Transform/BuildOnlyExtensionTest.cpp

+//===- BuildOnlyExtensionTest.cpp - unit test for transform extensions ----===//
+//
+// 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/Func/IR/FuncOps.h"
+#include "mlir/Dialect/Transform/IR/TransformDialect.h"
+#include "mlir/IR/DialectRegistry.h"
+#include "mlir/IR/MLIRContext.h"
+#include "gtest/gtest.h"
+
+using namespace mlir;
+using namespace mlir::transform;
+
+namespace {
+class Extension : public TransformDialectExtension<Extension> {
+public:
+  using Base::Base;
+  void init() { declareGeneratedDialect<func::FuncDialect>(); }
+};
+} // end namespace
+
+TEST(BuildOnlyExtensionTest, buildOnlyExtension) {
+  // Register the build-only version of the transform dialect extension. The
+  // func dialect is declared as generated so it should not be loaded along with
+  // the transform dialect.
+  DialectRegistry registry;
+  registry.addExtensions<BuildOnly<Extension>>();
+  MLIRContext ctx(registry);
+  ctx.getOrLoadDialect<TransformDialect>();
+  ASSERT_FALSE(ctx.getLoadedDialect<func::FuncDialect>());
+}
+
+TEST(BuildOnlyExtensionTest, buildAndApplyExtension) {
+  // Register the full version of the transform dialect extension. The func
+  // dialect should be loaded along with the transform dialect.
+  DialectRegistry registry;
+  registry.addExtensions<Extension>();
+  MLIRContext ctx(registry);
+  ctx.getOrLoadDialect<TransformDialect>();
+  ASSERT_TRUE(ctx.getLoadedDialect<func::FuncDialect>());
+}

mlir/unittests/Dialect/Transform/CMakeLists.txt

+add_mlir_unittest(MLIRTransformDialectTests
+  BuildOnlyExtensionTest.cpp
+)
+target_link_libraries(MLIRTransformDialectTests
+  PRIVATE
+  MLIRFuncDialect
+  MLIRTransformDialect
+)

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

deps = [
         "//llvm:gtest_main",
         "//mlir:SPIRVDeserialization",
         "//mlir:SPIRVDialect",
         "//mlir:SPIRVSerialization",
     ],
 )
 
 cc_test(
+    name = "transform_dialect_tests",
+    size = "small",
+    srcs = glob([
+        "Dialect/Transform/*.cpp",
+        "Dialect/Transform/*.h",
+    ]),
+    deps = [
+        "//llvm:TestingSupport",
+        "//llvm:gtest_main",
+        "//mlir:FuncDialect",
+        "//mlir:TransformDialect",
+    ],
+)
+
+cc_test(
     name = "dialect_utils_tests",
     size = "small",
     srcs = glob([
         "Dialect/Utils/*.cpp",
         "Dialect/Utils/*.h",
     ]),
     deps = [
         "//llvm:TestingSupport",