[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 to
postpone the loading of the generated dialects until the pass pipeline
involving the transform dialect is actually run, instead of loading them along
with the extension.

Depends On D128528

[ftynse]

Replaced by https://reviews.llvm.org/D130289 with a different implementation.

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

 //===- TransformDialect.h - Transform Dialect Definition --------*- C++ -*-===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_DIALECT_TRANSFORM_IR_TRANSFORMDIALECT_H
 #define MLIR_DIALECT_TRANSFORM_IR_TRANSFORMDIALECT_H
 
 #include "mlir/Dialect/Transform/IR/TransformInterfaces.h"
-#include "mlir/IR/Dialect.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/Support/LLVM.h"
 #include "llvm/ADT/StringMap.h"
 
 namespace mlir {
 namespace transform {
 #ifndef NDEBUG
 namespace detail {
public:
   /// 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);
     for (const Initializer &init : opInitializers)
       init(transformDialect);
     transformDialect->mergeInPDLMatchHooks(std::move(pdlMatchConstraintFns));
+    llvm::append_range(transformDialect->generatedDialects, dialectRegistrars);
   }
 
 protected:
   /// 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. This is **not** intended for dialects that
+  /// contain IR entities produced in the Payload IR by transformations included
+  /// in this extension, use `declareGeneratedDialect` instead.
   template <typename DialectTy>
   void declareDependentDialect() {
     dialectLoaders.push_back(
         [](MLIRContext *context) { context->loadDialect<DialectTy>(); });
   }
 
+  /// Declares that the transformations provided by this Transform dialect
+  /// extension may generate IR entities (attributes, operations or types) from
+  /// the dialect provided as template parameter. Passes that apply Transform
+  /// dialect operations to payload IR can use these to list the dialects they
+  /// "depend on" in the sense of creating IR entities, as expected by the pass
+  /// manager. This is **not** intended for dialects entities from which are
+  /// used in the Transform IR operations, use `declareDependentDialect`
+  /// instead.
+  template <typename DialectTy>
+  void declareGeneratedDialect() {
+    dialectRegistrars.push_back(
+        [](DialectRegistry &registry) { registry.insert<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;
   SmallVector<DialectLoader> dialectLoaders;
+  SmallVector<TransformDialect::DialectRegistrar> dialectRegistrars;
 
   /// A list of constraints that should be made availble 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;
 };
 
 } // namespace transform
 } // namespace mlir
 
 #endif // MLIR_DIALECT_TRANSFORM_IR_TRANSFORMDIALECT_H

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

def Transform_Dialect : Dialect {
   ];
 
   let extraClassDeclaration = [{
       /// Returns the named PDL constraint functions available in the dialect
       /// as a map from their name to the function.
       const ::llvm::StringMap<::mlir::PDLConstraintFunction> &
       getPDLConstraintHooks() const;
 
+      /// Appends dialects entities from which can be generated in the Payload
+      /// IR when Transform ops from any of the registered extensions are
+      /// applied.
+      void registerGeneratedDialects(DialectRegistry &registry) const;
+
     private:
       template <typename OpTy>
       void addOperationIfNotRegistered() {
         Optional<RegisteredOperationName> opName =
             RegisteredOperationName::lookup(OpTy::getOperationName(),
                                             getContext());
         if (!opName)
           return addOperations<OpTy>();
private:
       template <typename, typename...>
       friend class TransformDialectExtension;
 
       /// Takes ownership of the named PDL constraint function from the given
       /// map and makes them available for use by the operations in the dialect.
       void mergeInPDLMatchHooks(
           ::llvm::StringMap<::mlir::PDLConstraintFunction> &&constraintFns);
 
+      /// Type of callbacks adding dialects into the dialect registry.
+      using DialectRegistrar = std::function<void(DialectRegistry &)>;
+
       /// A container for PDL constraint function that can be used by
       /// operations in this dialect.
       PDLPatternModule pdlMatchHooks;
+
+      /// Callbacks registering the dialects that can be produced when applying
+      /// the transformation operations registered by this dialect and its
+      /// extensions.
+      SmallVector<DialectRegistrar> generatedDialects;
   }];
 }
 
 // Base class for ops that belong to the tranfsorm dialect. Ops defined in
 // extensions of this dialect may also use this.
 class TransformDialectOp<string mnemonic, list<Trait> traits = []>
     : Op<Transform_Dialect, mnemonic, traits>;
 

mlir/lib/Dialect/Linalg/TransformOps/CMakeLists.txt

 add_mlir_dialect_library(MLIRLinalgTransformOps
   LinalgTransformOps.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/Linalg/TransformOps
 
   DEPENDS
   MLIRLinalgTransformOpsIncGen
 
   LINK_LIBS PUBLIC
+  MLIRArithmeticDialect
   MLIRIR
   MLIRLinalgDialect
   MLIRLinalgTransforms
   MLIRParser
   MLIRPDLDialect
   MLIRSideEffectInterfaces
   MLIRTransformDialect
   MLIRVectorDialect
   )

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

 //===- LinalgTransformOps.cpp - Implementation of Linalg transform 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/Linalg/TransformOps/LinalgTransformOps.h"
 
+#include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
 #include "mlir/Dialect/Linalg/IR/Linalg.h"
 #include "mlir/Dialect/Linalg/Transforms/Transforms.h"
 #include "mlir/Dialect/PDL/IR/PDL.h"
 #include "mlir/Dialect/PDL/IR/PDLTypes.h"
 #include "mlir/Dialect/Transform/IR/TransformDialect.h"
 #include "mlir/Parser/Parser.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 

 /// 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<pdl::PDLDialect>();
-    declareDependentDialect<scf::SCFDialect>();
-    declareDependentDialect<vector::VectorDialect>();
+
+    declareGeneratedDialect<arith::ArithmeticDialect>();
+    declareGeneratedDialect<scf::SCFDialect>();
+    declareGeneratedDialect<vector::VectorDialect>();
     registerTransformOps<
 #define GET_OP_LIST
 #include "mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp.inc"
         >();
   }
 };
+
 } // namespace
 
 #define GET_OP_CLASSES
 #include "mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp.inc"
 
 void mlir::linalg::registerTransformDialectExtension(
     DialectRegistry &registry) {
   registry.addExtensions<LinalgTransformDialectExtension>();
 }

mlir/lib/Dialect/Transform/IR/TransformDialect.cpp

void transform::TransformDialect::mergeInPDLMatchHooks(
   for (auto &it : constraintFns)
     pdlMatchHooks.registerConstraintFunction(it.getKey(), std::move(it.second));
 }
 
 const llvm::StringMap<PDLConstraintFunction> &
 transform::TransformDialect::getPDLConstraintHooks() const {
   return pdlMatchHooks.getConstraintFunctions();
 }
+
+void transform::TransformDialect::registerGeneratedDialects(
+    DialectRegistry &registry) const {
+  for (const DialectRegistrar &callback : generatedDialects)
+    callback(registry);
+}

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

 //===- TestTransformDialectInterpreter.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
 //
 //===----------------------------------------------------------------------===//
 //
 // This file defines a test pass that interprets Transform dialect operations in
 // the module.
 //
 //===----------------------------------------------------------------------===//
 
+#include "mlir/Dialect/Transform/IR/TransformDialect.h"
 #include "mlir/Dialect/Transform/IR/TransformInterfaces.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/Pass/Pass.h"
 
 using namespace mlir;
 
 namespace {
 /// Simple pass that applies transform dialect ops directly contained in a
transform::TransformState state(
             enableExpensiveChecks));
     for (auto op :
          module.getBody()->getOps<transform::TransformOpInterface>()) {
       if (failed(state.applyTransform(op).checkAndReport()))
         return signalPassFailure();
     }
   }
 
+  void getDependentDialects(DialectRegistry &registry,
+                            MLIRContext &context) const override {
+    auto *dialect = context.getOrLoadDialect<transform::TransformDialect>();
+    dialect->registerGeneratedDialects(registry);
+  }
+
   Option<bool> enableExpensiveChecks{
       *this, "enable-expensive-checks", llvm::cl::init(false),
       llvm::cl::desc("perform expensive checks to better report errors in the "
                      "transform IR")};
 };
 } // namespace
 
 namespace mlir {
 namespace test {
 /// Registers the test pass for applying transform dialect ops.
 void registerTestTransformDialectInterpreterPass() {
   PassRegistration<TestTransformDialectInterpreterPass> reg;
 }
 } // namespace test
 } // namespace mlir

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

cc_library(
     srcs = [
         "lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp",
     ],
     hdrs = [
         "include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.h",
     ],
     includes = ["include"],
     deps = [
+        ":ArithmeticDialect",
         ":IR",
         ":LinalgDialect",
         ":LinalgTransformOpsIncGen",
         ":LinalgTransforms",
         ":PDLDialect",
         ":Parser",
         ":SideEffectInterfaces",
         ":TransformDialect",