1 //===- TransformDialect.h - Transform Dialect Definition --------*- C++ -*-===//
2 //
3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4 // See https://llvm.org/LICENSE.txt for license information.
5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6 //
7 //===----------------------------------------------------------------------===//
8 
9 #ifndef MLIR_DIALECT_TRANSFORM_IR_TRANSFORMDIALECT_H
10 #define MLIR_DIALECT_TRANSFORM_IR_TRANSFORMDIALECT_H
11 
12 #include "mlir/Dialect/Transform/IR/TransformInterfaces.h"
13 #include "mlir/IR/Dialect.h"
14 #include "mlir/IR/PatternMatch.h"
15 #include "mlir/Support/LLVM.h"
16 #include "llvm/ADT/StringMap.h"
17 
18 namespace mlir {
19 namespace transform {
20 #ifndef NDEBUG
21 namespace detail {
22 /// Asserts that the operations provided as template arguments implement the
23 /// TransformOpInterface and MemoryEffectsOpInterface. This must be a dynamic
24 /// assertion since interface implementations may be registered at runtime.
25 template <typename OpTy>
checkImplementsTransformInterface(MLIRContext * context)26 static inline void checkImplementsTransformInterface(MLIRContext *context) {
27   // Since the operation is being inserted into the Transform dialect and the
28   // dialect does not implement the interface fallback, only check for the op
29   // itself having the interface implementation.
30   RegisteredOperationName opName =
31       *RegisteredOperationName::lookup(OpTy::getOperationName(), context);
32   assert((opName.hasInterface<TransformOpInterface>() ||
33           opName.hasTrait<OpTrait::IsTerminator>()) &&
34          "non-terminator ops injected into the transform dialect must "
35          "implement TransformOpInterface");
36   assert(opName.hasInterface<MemoryEffectOpInterface>() &&
37          "ops injected into the transform dialect must implement "
38          "MemoryEffectsOpInterface");
39 }
40 } // namespace detail
41 #endif // NDEBUG
42 } // namespace transform
43 } // namespace mlir
44 
45 #include "mlir/Dialect/Transform/IR/TransformDialect.h.inc"
46 
47 namespace mlir {
48 namespace transform {
49 
50 /// Base class for extensions of the Transform dialect that supports injecting
51 /// operations into the Transform dialect at load time. Concrete extensions are
52 /// expected to derive this class and register operations in the constructor.
53 /// They can be registered with the DialectRegistry and automatically applied
54 /// to the Transform dialect when it is loaded.
55 ///
56 /// Derived classes are expected to define a `void init()` function in which
57 /// they can call various protected methods of the base class to register
58 /// extension operations and declare their dependencies.
59 ///
60 /// By default, the extension is configured both for construction of the
61 /// Transform IR and for its application to some payload. If only the
62 /// construction is desired, the extension can be switched to "build-only" mode
63 /// that avoids loading the dialects that are only necessary for transforming
64 /// the payload. To perform the switch, the extension must be wrapped into the
65 /// `BuildOnly` class template (see below) when it is registered, as in:
66 ///
67 ///    dialectRegistry.addExtension<BuildOnly<MyTransformDialectExt>>();
68 ///
69 /// instead of:
70 ///
71 ///    dialectRegistry.addExtension<MyTransformDialectExt>();
72 ///
73 /// Derived classes must reexport the constructor of this class or otherwise
74 /// forward its boolean argument to support this behavior.
75 template <typename DerivedTy, typename... ExtraDialects>
76 class TransformDialectExtension
77     : public DialectExtension<DerivedTy, TransformDialect, ExtraDialects...> {
78   using Initializer = std::function<void(TransformDialect *)>;
79   using DialectLoader = std::function<void(MLIRContext *)>;
80 
81 public:
82   /// Extension application hook. Actually loads the dependent dialects and
83   /// registers the additional operations. Not expected to be called directly.
apply(MLIRContext * context,TransformDialect * transformDialect,ExtraDialects * ...)84   void apply(MLIRContext *context, TransformDialect *transformDialect,
85              ExtraDialects *...) const final {
86     for (const DialectLoader &loader : dialectLoaders)
87       loader(context);
88 
89     // Only load generated dialects if the user intends to apply
90     // transformations specified by the extension.
91     if (!buildOnly)
92       for (const DialectLoader &loader : generatedDialectLoaders)
93         loader(context);
94 
95     for (const Initializer &init : opInitializers)
96       init(transformDialect);
97     transformDialect->mergeInPDLMatchHooks(std::move(pdlMatchConstraintFns));
98   }
99 
100 protected:
101   using Base = TransformDialectExtension<DerivedTy, ExtraDialects...>;
102 
103   /// Extension constructor. The argument indicates whether to skip generated
104   /// dialects when applying the extension.
105   explicit TransformDialectExtension(bool buildOnly = false)
buildOnly(buildOnly)106       : buildOnly(buildOnly) {
107     static_cast<DerivedTy *>(this)->init();
108   }
109 
110   /// Hook for derived classes to inject constructor behavior.
init()111   void init() {}
112 
113   /// Injects the operations into the Transform dialect. The operations must
114   /// implement the TransformOpInterface and MemoryEffectsOpInterface, and the
115   /// implementations must be already available when the operation is injected.
116   template <typename... OpTys>
registerTransformOps()117   void registerTransformOps() {
118     opInitializers.push_back([](TransformDialect *transformDialect) {
119       transformDialect->addOperationsChecked<OpTys...>();
120     });
121   }
122 
123   /// Declares that this Transform dialect extension depends on the dialect
124   /// provided as template parameter. When the Transform dialect is loaded,
125   /// dependent dialects will be loaded as well. This is intended for dialects
126   /// that contain attributes and types used in creation and canonicalization of
127   /// the injected operations, similarly to how the dialect definition may list
128   /// dependent dialects. This is *not* intended for dialects entities from
129   /// which may be produced when applying the transformations specified by ops
130   /// registered by this extension.
131   template <typename DialectTy>
declareDependentDialect()132   void declareDependentDialect() {
133     dialectLoaders.push_back(
134         [](MLIRContext *context) { context->loadDialect<DialectTy>(); });
135   }
136 
137   /// Declares that the transformations associated with the operations
138   /// registered by this dialect extension may produce operations from the
139   /// dialect provided as template parameter while processing payload IR that
140   /// does not contain the operations from said dialect. This is similar to
141   /// dependent dialects of a pass. These dialects will be loaded along with the
142   /// transform dialect unless the extension is in the build-only mode.
143   template <typename DialectTy>
declareGeneratedDialect()144   void declareGeneratedDialect() {
145     generatedDialectLoaders.push_back(
146         [](MLIRContext *context) { context->loadDialect<DialectTy>(); });
147   }
148 
149   /// Injects the named constraint to make it available for use with the
150   /// PDLMatchOp in the transform dialect.
registerPDLMatchConstraintFn(StringRef name,PDLConstraintFunction && fn)151   void registerPDLMatchConstraintFn(StringRef name,
152                                     PDLConstraintFunction &&fn) {
153     pdlMatchConstraintFns.try_emplace(name,
154                                       std::forward<PDLConstraintFunction>(fn));
155   }
156   template <typename ConstraintFnTy>
registerPDLMatchConstraintFn(StringRef name,ConstraintFnTy && fn)157   void registerPDLMatchConstraintFn(StringRef name, ConstraintFnTy &&fn) {
158     pdlMatchConstraintFns.try_emplace(
159         name, ::mlir::detail::pdl_function_builder::buildConstraintFn(
160                   std::forward<ConstraintFnTy>(fn)));
161   }
162 
163 private:
164   SmallVector<Initializer> opInitializers;
165 
166   /// Callbacks loading the dependent dialects, i.e. the dialect needed for the
167   /// extension ops.
168   SmallVector<DialectLoader> dialectLoaders;
169 
170   /// Callbacks loading the generated dialects, i.e. the dialects produced when
171   /// applying the transformations.
172   SmallVector<DialectLoader> generatedDialectLoaders;
173 
174   /// A list of constraints that should be made available to PDL patterns
175   /// processed by PDLMatchOp in the Transform dialect.
176   ///
177   /// Declared as mutable so its contents can be moved in the `apply` const
178   /// method, which is only called once.
179   mutable llvm::StringMap<PDLConstraintFunction> pdlMatchConstraintFns;
180 
181   /// Indicates that the extension is in build-only mode.
182   bool buildOnly;
183 };
184 
185 /// A wrapper for transform dialect extensions that forces them to be
186 /// constructed in the build-only mode.
187 template <typename DerivedTy>
188 class BuildOnly : public DerivedTy {
189 public:
BuildOnly()190   BuildOnly() : DerivedTy(/*buildOnly=*/true) {}
191 };
192 
193 } // namespace transform
194 } // namespace mlir
195 
196 #endif // MLIR_DIALECT_TRANSFORM_IR_TRANSFORMDIALECT_H
197