MLIR  22.0.0git
MLIRContext.h
Go to the documentation of this file.
1 //===- MLIRContext.h - MLIR Global Context Class ----------------*- 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_IR_MLIRCONTEXT_H
10 #define MLIR_IR_MLIRCONTEXT_H
11 
12 #include "mlir/Support/LLVM.h"
13 #include "mlir/Support/TypeID.h"
14 #include "llvm/ADT/ArrayRef.h"
15 #include <functional>
16 #include <memory>
17 #include <vector>
18 
19 namespace llvm {
20 class ThreadPoolInterface;
21 } // namespace llvm
22 
23 namespace mlir {
24 namespace tracing {
25 class Action;
26 }
27 class DiagnosticEngine;
28 class Dialect;
29 class DialectRegistry;
30 class DynamicDialect;
31 class InFlightDiagnostic;
32 class Location;
33 class MLIRContextImpl;
34 class RegisteredOperationName;
35 class StorageUniquer;
36 class IRUnit;
37 namespace remark::detail {
38 class RemarkEngine;
39 } // namespace remark::detail
40 
41 /// MLIRContext is the top-level object for a collection of MLIR operations. It
42 /// holds immortal uniqued objects like types, and the tables used to unique
43 /// them.
44 ///
45 /// MLIRContext gets a redundant "MLIR" prefix because otherwise it ends up with
46 /// a very generic name ("Context") and because it is uncommon for clients to
47 /// interact with it.
48 ///
49 /// The context wrap some multi-threading facilities, and in particular by
50 /// default it will implicitly create a thread pool.
51 /// This can be undesirable if multiple context exists at the same time or if a
52 /// process will be long-lived and create and destroy contexts.
53 /// To control better thread spawning, an externally owned ThreadPool can be
54 /// injected in the context. For example:
55 ///
56 /// llvm::DefaultThreadPool myThreadPool;
57 /// while (auto *request = nextCompilationRequests()) {
58 /// MLIRContext ctx(registry, MLIRContext::Threading::DISABLED);
59 /// ctx.setThreadPool(myThreadPool);
60 /// processRequest(request, cxt);
61 /// }
62 ///
63 class MLIRContext {
64 public:
65  enum class Threading { DISABLED, ENABLED };
66  /// Create a new Context.
67  explicit MLIRContext(Threading multithreading = Threading::ENABLED);
68  explicit MLIRContext(const DialectRegistry &registry,
69  Threading multithreading = Threading::ENABLED);
71 
72  /// Return information about all IR dialects loaded in the context.
73  std::vector<Dialect *> getLoadedDialects();
74 
75  /// Return the dialect registry associated with this context.
77 
78  /// Append the contents of the given dialect registry to the registry
79  /// associated with this context.
80  void appendDialectRegistry(const DialectRegistry &registry);
81 
82  /// Return information about all available dialects in the registry in this
83  /// context.
84  std::vector<StringRef> getAvailableDialects();
85 
86  /// Get a registered IR dialect with the given namespace. If an exact match is
87  /// not found, then return nullptr.
88  Dialect *getLoadedDialect(StringRef name);
89 
90  /// Get a registered IR dialect for the given derived dialect type. The
91  /// derived type must provide a static 'getDialectNamespace' method.
92  template <typename T>
94  return static_cast<T *>(getLoadedDialect(T::getDialectNamespace()));
95  }
96 
97  /// Get (or create) a dialect for the given derived dialect type. The derived
98  /// type must provide a static 'getDialectNamespace' method.
99  template <typename T>
101  return static_cast<T *>(
102  getOrLoadDialect(T::getDialectNamespace(), TypeID::get<T>(), [this]() {
103  std::unique_ptr<T> dialect(new T(this));
104  return dialect;
105  }));
106  }
107 
108  /// Load a dialect in the context.
109  template <typename Dialect>
110  void loadDialect() {
111  // Do not load the dialect if it is currently loading. This can happen if a
112  // dialect initializer triggers loading the same dialect recursively.
113  if (!isDialectLoading(Dialect::getDialectNamespace()))
114  getOrLoadDialect<Dialect>();
115  }
116 
117  /// Load a list dialects in the context.
118  template <typename Dialect, typename OtherDialect, typename... MoreDialects>
119  void loadDialect() {
120  loadDialect<Dialect>();
121  loadDialect<OtherDialect, MoreDialects...>();
122  }
123 
124  /// Get (or create) a dynamic dialect for the given name.
126  getOrLoadDynamicDialect(StringRef dialectNamespace,
127  function_ref<void(DynamicDialect *)> ctor);
128 
129  /// Load all dialects available in the registry in this context.
131 
132  /// Get (or create) a dialect for the given derived dialect name.
133  /// The dialect will be loaded from the registry if no dialect is found.
134  /// If no dialect is loaded for this name and none is available in the
135  /// registry, returns nullptr.
136  Dialect *getOrLoadDialect(StringRef name);
137 
138  /// Return true if we allow to create operation for unregistered dialects.
139  [[nodiscard]] bool allowsUnregisteredDialects();
140 
141  /// Enables creating operations in unregistered dialects.
142  /// This option is **heavily discouraged**: it is convenient during testing
143  /// but it is not a good practice to use it in production code. Some system
144  /// invariants can be broken (like loading a dialect after creating
145  /// operations) without being caught by assertions or other means.
146  void allowUnregisteredDialects(bool allow = true);
147 
148  /// Return true if multi-threading is enabled by the context.
150 
151  /// Set the flag specifying if multi-threading is disabled by the context.
152  /// The command line debugging flag `--mlir-disable-threading` is overriding
153  /// this call and making it a no-op!
154  void disableMultithreading(bool disable = true);
155  void enableMultithreading(bool enable = true) {
156  disableMultithreading(!enable);
157  }
158 
159  /// Set a new thread pool to be used in this context. This method requires
160  /// that multithreading is disabled for this context prior to the call. This
161  /// allows to share a thread pool across multiple contexts, as well as
162  /// decoupling the lifetime of the threads from the contexts. The thread pool
163  /// must outlive the context. Multi-threading will be enabled as part of this
164  /// method.
165  /// The command line debugging flag `--mlir-disable-threading` will still
166  /// prevent threading from being enabled and threading won't be enabled after
167  /// this call in this case.
168  void setThreadPool(llvm::ThreadPoolInterface &pool);
169 
170  /// Return the number of threads used by the thread pool in this context. The
171  /// number of computed hardware threads can change over the lifetime of a
172  /// process based on affinity changes, so users should use the number of
173  /// threads actually in the thread pool for dispatching work. Returns 1 if
174  /// multithreading is disabled.
175  unsigned getNumThreads();
176 
177  /// Return the thread pool used by this context. This method requires that
178  /// multithreading be enabled within the context, and should generally not be
179  /// used directly. Users should instead prefer the threading utilities within
180  /// Threading.h.
181  llvm::ThreadPoolInterface &getThreadPool();
182 
183  /// Return true if we should attach the operation to diagnostics emitted via
184  /// Operation::emit.
186 
187  /// Set the flag specifying if we should attach the operation to diagnostics
188  /// emitted via Operation::emit.
189  void printOpOnDiagnostic(bool enable);
190 
191  /// Return true if we should attach the current stacktrace to diagnostics when
192  /// emitted.
194 
195  /// Set the flag specifying if we should attach the current stacktrace when
196  /// emitting diagnostics.
197  void printStackTraceOnDiagnostic(bool enable);
198 
199  /// Return a sorted array containing the information about all registered
200  /// operations.
202 
203  /// Return a sorted array containing the information for registered operations
204  /// filtered by dialect name.
206  getRegisteredOperationsByDialect(StringRef dialectName);
207 
208  /// Return true if this operation name is registered in this context.
209  bool isOperationRegistered(StringRef name);
210 
211  // This is effectively private given that only MLIRContext.cpp can see the
212  // MLIRContextImpl type.
213  MLIRContextImpl &getImpl() { return *impl; }
214 
215  /// Returns the diagnostic engine for this context.
217 
218  /// Returns the remark engine for this context, or nullptr if none has been
219  /// set.
221 
222  /// Set the remark engine for this context.
223  void setRemarkEngine(std::unique_ptr<remark::detail::RemarkEngine> engine);
224 
225  /// Returns the storage uniquer used for creating affine constructs.
227 
228  /// Returns the storage uniquer used for constructing type storage instances.
229  /// This should not be used directly.
231 
232  /// Returns the storage uniquer used for constructing attribute storage
233  /// instances. This should not be used directly.
235 
236  /// These APIs are tracking whether the context will be used in a
237  /// multithreading environment: this has no effect other than enabling
238  /// assertions on misuses of some APIs.
241 
242  /// Get a dialect for the provided namespace and TypeID: abort the program if
243  /// a dialect exist for this namespace with different TypeID. If a dialect has
244  /// not been loaded for this namespace/TypeID yet, use the provided ctor to
245  /// create one on the fly and load it. Returns a pointer to the dialect owned
246  /// by the context.
247  /// The use of this method is in general discouraged in favor of
248  /// 'getOrLoadDialect<DialectClass>()'.
249  Dialect *getOrLoadDialect(StringRef dialectNamespace, TypeID dialectID,
250  function_ref<std::unique_ptr<Dialect>()> ctor);
251 
252  /// Returns a hash of the registry of the context that may be used to give
253  /// a rough indicator of if the state of the context registry has changed. The
254  /// context registry correlates to loaded dialects and their entities
255  /// (attributes, operations, types, etc.).
256  llvm::hash_code getRegistryHash();
257 
258  //===--------------------------------------------------------------------===//
259  // Action API
260  //===--------------------------------------------------------------------===//
261 
262  /// Signatures for the action handler that can be registered with the context.
263  using HandlerTy =
264  std::function<void(function_ref<void()>, const tracing::Action &)>;
265 
266  /// Register a handler for handling actions that are dispatched through this
267  /// context. A nullptr handler can be set to disable a previously set handler.
268  void registerActionHandler(HandlerTy handler);
269 
270  /// Return true if a valid ActionHandler is set.
271  bool hasActionHandler();
272 
273  /// Dispatch the provided action to the handler if any, or just execute it.
274  void executeAction(function_ref<void()> actionFn,
275  const tracing::Action &action) {
276  if (LLVM_UNLIKELY(hasActionHandler()))
277  executeActionInternal(actionFn, action);
278  else
279  actionFn();
280  }
281 
282  /// Dispatch the provided action to the handler if any, or just execute it.
283  template <typename ActionTy, typename... Args>
284  void executeAction(function_ref<void()> actionFn, ArrayRef<IRUnit> irUnits,
285  Args &&...args) {
286  if (LLVM_UNLIKELY(hasActionHandler()))
287  executeActionInternal<ActionTy, Args...>(actionFn, irUnits,
288  std::forward<Args>(args)...);
289  else
290  actionFn();
291  }
292 
293 private:
294  /// Return true if the given dialect is currently loading.
295  bool isDialectLoading(StringRef dialectNamespace);
296 
297  /// Internal helper for the dispatch method.
298  void executeActionInternal(function_ref<void()> actionFn,
299  const tracing::Action &action);
300 
301  /// Internal helper for the dispatch method. We get here after checking that
302  /// there is a handler, for the purpose of keeping this code out-of-line. and
303  /// avoid calling the ctor for the Action unnecessarily.
304  template <typename ActionTy, typename... Args>
305  LLVM_ATTRIBUTE_NOINLINE void
306  executeActionInternal(function_ref<void()> actionFn, ArrayRef<IRUnit> irUnits,
307  Args &&...args) {
308  executeActionInternal(actionFn,
309  ActionTy(irUnits, std::forward<Args>(args)...));
310  }
311 
312  const std::unique_ptr<MLIRContextImpl> impl;
313 
314  MLIRContext(const MLIRContext &) = delete;
315  void operator=(const MLIRContext &) = delete;
316 };
317 
318 //===----------------------------------------------------------------------===//
319 // MLIRContext CommandLine Options
320 //===----------------------------------------------------------------------===//
321 
322 /// Register a set of useful command-line options that can be used to configure
323 /// various flags within the MLIRContext. These flags are used when constructing
324 /// an MLIR context for initialization.
326 
327 } // namespace mlir
328 
329 #endif // MLIR_IR_MLIRCONTEXT_H
This class is the main interface for diagnostics.
Definition: Diagnostics.h:414
The DialectRegistry maps a dialect namespace to a constructor for the matching dialect.
Dialects are groups of MLIR operations, types and attributes, as well as behavior associated with the...
Definition: Dialect.h:38
A dialect that can be defined at runtime.
This is the implementation of the MLIRContext class, using the pImpl idiom.
MLIRContext is the top-level object for a collection of MLIR operations.
Definition: MLIRContext.h:63
void appendDialectRegistry(const DialectRegistry &registry)
Append the contents of the given dialect registry to the registry associated with this context.
bool shouldPrintStackTraceOnDiagnostic()
Return true if we should attach the current stacktrace to diagnostics when emitted.
unsigned getNumThreads()
Return the number of threads used by the thread pool in this context.
bool isOperationRegistered(StringRef name)
Return true if this operation name is registered in this context.
MLIRContext(Threading multithreading=Threading::ENABLED)
Create a new Context.
void disableMultithreading(bool disable=true)
Set the flag specifying if multi-threading is disabled by the context.
void printStackTraceOnDiagnostic(bool enable)
Set the flag specifying if we should attach the current stacktrace when emitting diagnostics.
bool hasActionHandler()
Return true if a valid ActionHandler is set.
void setRemarkEngine(std::unique_ptr< remark::detail::RemarkEngine > engine)
Set the remark engine for this context.
void setThreadPool(llvm::ThreadPoolInterface &pool)
Set a new thread pool to be used in this context.
void executeAction(function_ref< void()> actionFn, const tracing::Action &action)
Dispatch the provided action to the handler if any, or just execute it.
Definition: MLIRContext.h:274
void enableMultithreading(bool enable=true)
Definition: MLIRContext.h:155
remark::detail::RemarkEngine * getRemarkEngine()
Returns the remark engine for this context, or nullptr if none has been set.
std::vector< Dialect * > getLoadedDialects()
Return information about all IR dialects loaded in the context.
ArrayRef< RegisteredOperationName > getRegisteredOperationsByDialect(StringRef dialectName)
Return a sorted array containing the information for registered operations filtered by dialect name.
void printOpOnDiagnostic(bool enable)
Set the flag specifying if we should attach the operation to diagnostics emitted via Operation::emit.
void executeAction(function_ref< void()> actionFn, ArrayRef< IRUnit > irUnits, Args &&...args)
Dispatch the provided action to the handler if any, or just execute it.
Definition: MLIRContext.h:284
void registerActionHandler(HandlerTy handler)
Register a handler for handling actions that are dispatched through this context.
ArrayRef< RegisteredOperationName > getRegisteredOperations()
Return a sorted array containing the information about all registered operations.
llvm::hash_code getRegistryHash()
Returns a hash of the registry of the context that may be used to give a rough indicator of if the st...
void enterMultiThreadedExecution()
These APIs are tracking whether the context will be used in a multithreading environment: this has no...
const DialectRegistry & getDialectRegistry()
Return the dialect registry associated with this context.
void loadDialect()
Load a dialect in the context.
Definition: MLIRContext.h:110
DynamicDialect * getOrLoadDynamicDialect(StringRef dialectNamespace, function_ref< void(DynamicDialect *)> ctor)
Get (or create) a dynamic dialect for the given name.
StorageUniquer & getAttributeUniquer()
Returns the storage uniquer used for constructing attribute storage instances.
StorageUniquer & getAffineUniquer()
Returns the storage uniquer used for creating affine constructs.
StorageUniquer & getTypeUniquer()
Returns the storage uniquer used for constructing type storage instances.
llvm::ThreadPoolInterface & getThreadPool()
Return the thread pool used by this context.
std::vector< StringRef > getAvailableDialects()
Return information about all available dialects in the registry in this context.
bool isMultithreadingEnabled()
Return true if multi-threading is enabled by the context.
void allowUnregisteredDialects(bool allow=true)
Enables creating operations in unregistered dialects.
bool allowsUnregisteredDialects()
Return true if we allow to create operation for unregistered dialects.
T * getLoadedDialect()
Get a registered IR dialect for the given derived dialect type.
Definition: MLIRContext.h:93
DiagnosticEngine & getDiagEngine()
Returns the diagnostic engine for this context.
std::function< void(function_ref< void()>, const tracing::Action &)> HandlerTy
Signatures for the action handler that can be registered with the context.
Definition: MLIRContext.h:264
void loadDialect()
Load a list dialects in the context.
Definition: MLIRContext.h:119
MLIRContextImpl & getImpl()
Definition: MLIRContext.h:213
T * getOrLoadDialect()
Get (or create) a dialect for the given derived dialect type.
Definition: MLIRContext.h:100
void exitMultiThreadedExecution()
bool shouldPrintOpOnDiagnostic()
Return true if we should attach the operation to diagnostics emitted via Operation::emit.
void loadAllAvailableDialects()
Load all dialects available in the registry in this context.
A utility class to get or create instances of "storage classes".
This class provides an efficient unique identifier for a specific C++ type.
Definition: TypeID.h:107
An action is a specific action that is to be taken by the compiler, that can be toggled and controlle...
Definition: Action.h:38
The OpAsmOpInterface, see OpAsmInterface.td for more details.
Definition: CallGraph.h:229
Action
The actions performed by @newSparseTensor.
Definition: Enums.h:146
Include the generated interface declarations.
void registerMLIRContextCLOptions()
Register a set of useful command-line options that can be used to configure various flags within the ...
Definition: MLIRContext.cpp:88