1 | //===- MLIRContext.cpp - MLIR Type Classes --------------------------------===// |
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 | #include "mlir/IR/MLIRContext.h" |
10 | #include "AffineExprDetail.h" |
11 | #include "AffineMapDetail.h" |
12 | #include "AttributeDetail.h" |
13 | #include "IntegerSetDetail.h" |
14 | #include "TypeDetail.h" |
15 | #include "mlir/IR/Action.h" |
16 | #include "mlir/IR/AffineExpr.h" |
17 | #include "mlir/IR/AffineMap.h" |
18 | #include "mlir/IR/Attributes.h" |
19 | #include "mlir/IR/BuiltinAttributes.h" |
20 | #include "mlir/IR/BuiltinDialect.h" |
21 | #include "mlir/IR/Diagnostics.h" |
22 | #include "mlir/IR/Dialect.h" |
23 | #include "mlir/IR/ExtensibleDialect.h" |
24 | #include "mlir/IR/IntegerSet.h" |
25 | #include "mlir/IR/Location.h" |
26 | #include "mlir/IR/OpImplementation.h" |
27 | #include "mlir/IR/OperationSupport.h" |
28 | #include "mlir/IR/Types.h" |
29 | #include "llvm/ADT/DenseMap.h" |
30 | #include "llvm/ADT/DenseSet.h" |
31 | #include "llvm/ADT/SmallString.h" |
32 | #include "llvm/ADT/StringSet.h" |
33 | #include "llvm/ADT/Twine.h" |
34 | #include "llvm/Support/Allocator.h" |
35 | #include "llvm/Support/CommandLine.h" |
36 | #include "llvm/Support/Compiler.h" |
37 | #include "llvm/Support/Debug.h" |
38 | #include "llvm/Support/Mutex.h" |
39 | #include "llvm/Support/RWMutex.h" |
40 | #include "llvm/Support/ThreadPool.h" |
41 | #include "llvm/Support/raw_ostream.h" |
42 | #include <memory> |
43 | #include <optional> |
44 | |
45 | #define DEBUG_TYPE "mlircontext" |
46 | |
47 | using namespace mlir; |
48 | using namespace mlir::detail; |
49 | |
50 | //===----------------------------------------------------------------------===// |
51 | // MLIRContext CommandLine Options |
52 | //===----------------------------------------------------------------------===// |
53 | |
54 | namespace { |
55 | /// This struct contains command line options that can be used to initialize |
56 | /// various bits of an MLIRContext. This uses a struct wrapper to avoid the need |
57 | /// for global command line options. |
58 | struct MLIRContextOptions { |
59 | llvm::cl::opt<bool> disableThreading{ |
60 | "mlir-disable-threading" , |
61 | llvm::cl::desc("Disable multi-threading within MLIR, overrides any " |
62 | "further call to MLIRContext::enableMultiThreading()" )}; |
63 | |
64 | llvm::cl::opt<bool> printOpOnDiagnostic{ |
65 | "mlir-print-op-on-diagnostic" , |
66 | llvm::cl::desc("When a diagnostic is emitted on an operation, also print " |
67 | "the operation as an attached note" ), |
68 | llvm::cl::init(Val: true)}; |
69 | |
70 | llvm::cl::opt<bool> printStackTraceOnDiagnostic{ |
71 | "mlir-print-stacktrace-on-diagnostic" , |
72 | llvm::cl::desc("When a diagnostic is emitted, also print the stack trace " |
73 | "as an attached note" )}; |
74 | }; |
75 | } // namespace |
76 | |
77 | static llvm::ManagedStatic<MLIRContextOptions> clOptions; |
78 | |
79 | static bool isThreadingGloballyDisabled() { |
80 | #if LLVM_ENABLE_THREADS != 0 |
81 | return clOptions.isConstructed() && clOptions->disableThreading; |
82 | #else |
83 | return true; |
84 | #endif |
85 | } |
86 | |
87 | /// Register a set of useful command-line options that can be used to configure |
88 | /// various flags within the MLIRContext. These flags are used when constructing |
89 | /// an MLIR context for initialization. |
90 | void mlir::registerMLIRContextCLOptions() { |
91 | // Make sure that the options struct has been initialized. |
92 | *clOptions; |
93 | } |
94 | |
95 | //===----------------------------------------------------------------------===// |
96 | // Locking Utilities |
97 | //===----------------------------------------------------------------------===// |
98 | |
99 | namespace { |
100 | /// Utility writer lock that takes a runtime flag that specifies if we really |
101 | /// need to lock. |
102 | struct ScopedWriterLock { |
103 | ScopedWriterLock(llvm::sys::SmartRWMutex<true> &mutexParam, bool shouldLock) |
104 | : mutex(shouldLock ? &mutexParam : nullptr) { |
105 | if (mutex) |
106 | mutex->lock(); |
107 | } |
108 | ~ScopedWriterLock() { |
109 | if (mutex) |
110 | mutex->unlock(); |
111 | } |
112 | llvm::sys::SmartRWMutex<true> *mutex; |
113 | }; |
114 | } // namespace |
115 | |
116 | //===----------------------------------------------------------------------===// |
117 | // MLIRContextImpl |
118 | //===----------------------------------------------------------------------===// |
119 | |
120 | namespace mlir { |
121 | /// This is the implementation of the MLIRContext class, using the pImpl idiom. |
122 | /// This class is completely private to this file, so everything is public. |
123 | class MLIRContextImpl { |
124 | public: |
125 | //===--------------------------------------------------------------------===// |
126 | // Debugging |
127 | //===--------------------------------------------------------------------===// |
128 | |
129 | /// An action handler for handling actions that are dispatched through this |
130 | /// context. |
131 | std::function<void(function_ref<void()>, const tracing::Action &)> |
132 | actionHandler; |
133 | |
134 | //===--------------------------------------------------------------------===// |
135 | // Diagnostics |
136 | //===--------------------------------------------------------------------===// |
137 | DiagnosticEngine diagEngine; |
138 | |
139 | //===--------------------------------------------------------------------===// |
140 | // Options |
141 | //===--------------------------------------------------------------------===// |
142 | |
143 | /// In most cases, creating operation in unregistered dialect is not desired |
144 | /// and indicate a misconfiguration of the compiler. This option enables to |
145 | /// detect such use cases |
146 | bool allowUnregisteredDialects = false; |
147 | |
148 | /// Enable support for multi-threading within MLIR. |
149 | bool threadingIsEnabled = true; |
150 | |
151 | /// Track if we are currently executing in a threaded execution environment |
152 | /// (like the pass-manager): this is only a debugging feature to help reducing |
153 | /// the chances of data races one some context APIs. |
154 | #ifndef NDEBUG |
155 | std::atomic<int> multiThreadedExecutionContext{0}; |
156 | #endif |
157 | |
158 | /// If the operation should be attached to diagnostics printed via the |
159 | /// Operation::emit methods. |
160 | bool printOpOnDiagnostic = true; |
161 | |
162 | /// If the current stack trace should be attached when emitting diagnostics. |
163 | bool printStackTraceOnDiagnostic = false; |
164 | |
165 | //===--------------------------------------------------------------------===// |
166 | // Other |
167 | //===--------------------------------------------------------------------===// |
168 | |
169 | /// This points to the ThreadPool used when processing MLIR tasks in parallel. |
170 | /// It can't be nullptr when multi-threading is enabled. Otherwise if |
171 | /// multi-threading is disabled, and the threadpool wasn't externally provided |
172 | /// using `setThreadPool`, this will be nullptr. |
173 | llvm::ThreadPoolInterface *threadPool = nullptr; |
174 | |
175 | /// In case where the thread pool is owned by the context, this ensures |
176 | /// destruction with the context. |
177 | std::unique_ptr<llvm::ThreadPoolInterface> ownedThreadPool; |
178 | |
179 | /// An allocator used for AbstractAttribute and AbstractType objects. |
180 | llvm::BumpPtrAllocator abstractDialectSymbolAllocator; |
181 | |
182 | /// This is a mapping from operation name to the operation info describing it. |
183 | llvm::StringMap<std::unique_ptr<OperationName::Impl>> operations; |
184 | |
185 | /// A vector of operation info specifically for registered operations. |
186 | llvm::DenseMap<TypeID, RegisteredOperationName> registeredOperations; |
187 | llvm::StringMap<RegisteredOperationName> registeredOperationsByName; |
188 | |
189 | /// This is a sorted container of registered operations for a deterministic |
190 | /// and efficient `getRegisteredOperations` implementation. |
191 | SmallVector<RegisteredOperationName, 0> sortedRegisteredOperations; |
192 | |
193 | /// This is a list of dialects that are created referring to this context. |
194 | /// The MLIRContext owns the objects. These need to be declared after the |
195 | /// registered operations to ensure correct destruction order. |
196 | DenseMap<StringRef, std::unique_ptr<Dialect>> loadedDialects; |
197 | DialectRegistry dialectsRegistry; |
198 | |
199 | /// A mutex used when accessing operation information. |
200 | llvm::sys::SmartRWMutex<true> operationInfoMutex; |
201 | |
202 | //===--------------------------------------------------------------------===// |
203 | // Affine uniquing |
204 | //===--------------------------------------------------------------------===// |
205 | |
206 | // Affine expression, map and integer set uniquing. |
207 | StorageUniquer affineUniquer; |
208 | |
209 | //===--------------------------------------------------------------------===// |
210 | // Type uniquing |
211 | //===--------------------------------------------------------------------===// |
212 | |
213 | DenseMap<TypeID, AbstractType *> registeredTypes; |
214 | StorageUniquer typeUniquer; |
215 | |
216 | /// This is a mapping from type name to the abstract type describing it. |
217 | /// It is used by `AbstractType::lookup` to get an `AbstractType` from a name. |
218 | /// As this map needs to be populated before `StringAttr` is loaded, we |
219 | /// cannot use `StringAttr` as the key. The context does not take ownership |
220 | /// of the key, so the `StringRef` must outlive the context. |
221 | llvm::DenseMap<StringRef, AbstractType *> nameToType; |
222 | |
223 | /// Cached Type Instances. |
224 | Float8E5M2Type f8E5M2Ty; |
225 | Float8E4M3FNType f8E4M3FNTy; |
226 | Float8E5M2FNUZType f8E5M2FNUZTy; |
227 | Float8E4M3FNUZType f8E4M3FNUZTy; |
228 | Float8E4M3B11FNUZType f8E4M3B11FNUZTy; |
229 | BFloat16Type bf16Ty; |
230 | Float16Type f16Ty; |
231 | FloatTF32Type tf32Ty; |
232 | Float32Type f32Ty; |
233 | Float64Type f64Ty; |
234 | Float80Type f80Ty; |
235 | Float128Type f128Ty; |
236 | IndexType indexTy; |
237 | IntegerType int1Ty, int8Ty, int16Ty, int32Ty, int64Ty, int128Ty; |
238 | NoneType noneType; |
239 | |
240 | //===--------------------------------------------------------------------===// |
241 | // Attribute uniquing |
242 | //===--------------------------------------------------------------------===// |
243 | |
244 | DenseMap<TypeID, AbstractAttribute *> registeredAttributes; |
245 | StorageUniquer attributeUniquer; |
246 | |
247 | /// This is a mapping from attribute name to the abstract attribute describing |
248 | /// it. It is used by `AbstractType::lookup` to get an `AbstractType` from a |
249 | /// name. |
250 | /// As this map needs to be populated before `StringAttr` is loaded, we |
251 | /// cannot use `StringAttr` as the key. The context does not take ownership |
252 | /// of the key, so the `StringRef` must outlive the context. |
253 | llvm::DenseMap<StringRef, AbstractAttribute *> nameToAttribute; |
254 | |
255 | /// Cached Attribute Instances. |
256 | BoolAttr falseAttr, trueAttr; |
257 | UnitAttr unitAttr; |
258 | UnknownLoc unknownLocAttr; |
259 | DictionaryAttr emptyDictionaryAttr; |
260 | StringAttr emptyStringAttr; |
261 | |
262 | /// Map of string attributes that may reference a dialect, that are awaiting |
263 | /// that dialect to be loaded. |
264 | llvm::sys::SmartMutex<true> dialectRefStrAttrMutex; |
265 | DenseMap<StringRef, SmallVector<StringAttrStorage *>> |
266 | dialectReferencingStrAttrs; |
267 | |
268 | /// A distinct attribute allocator that allocates every time since the |
269 | /// address of the distinct attribute storage serves as unique identifier. The |
270 | /// allocator is thread safe and frees the allocated storage after its |
271 | /// destruction. |
272 | DistinctAttributeAllocator distinctAttributeAllocator; |
273 | |
274 | public: |
275 | MLIRContextImpl(bool threadingIsEnabled) |
276 | : threadingIsEnabled(threadingIsEnabled) { |
277 | if (threadingIsEnabled) { |
278 | ownedThreadPool = std::make_unique<llvm::DefaultThreadPool>(); |
279 | threadPool = ownedThreadPool.get(); |
280 | } |
281 | } |
282 | ~MLIRContextImpl() { |
283 | for (auto typeMapping : registeredTypes) |
284 | typeMapping.second->~AbstractType(); |
285 | for (auto attrMapping : registeredAttributes) |
286 | attrMapping.second->~AbstractAttribute(); |
287 | } |
288 | }; |
289 | } // namespace mlir |
290 | |
291 | MLIRContext::MLIRContext(Threading setting) |
292 | : MLIRContext(DialectRegistry(), setting) {} |
293 | |
294 | MLIRContext::MLIRContext(const DialectRegistry ®istry, Threading setting) |
295 | : impl(new MLIRContextImpl(setting == Threading::ENABLED && |
296 | !isThreadingGloballyDisabled())) { |
297 | // Initialize values based on the command line flags if they were provided. |
298 | if (clOptions.isConstructed()) { |
299 | printOpOnDiagnostic(enable: clOptions->printOpOnDiagnostic); |
300 | printStackTraceOnDiagnostic(enable: clOptions->printStackTraceOnDiagnostic); |
301 | } |
302 | |
303 | // Pre-populate the registry. |
304 | registry.appendTo(destination&: impl->dialectsRegistry); |
305 | |
306 | // Ensure the builtin dialect is always pre-loaded. |
307 | getOrLoadDialect<BuiltinDialect>(); |
308 | |
309 | // Initialize several common attributes and types to avoid the need to lock |
310 | // the context when accessing them. |
311 | |
312 | //// Types. |
313 | /// Floating-point Types. |
314 | impl->f8E5M2Ty = TypeUniquer::get<Float8E5M2Type>(this); |
315 | impl->f8E4M3FNTy = TypeUniquer::get<Float8E4M3FNType>(this); |
316 | impl->f8E5M2FNUZTy = TypeUniquer::get<Float8E5M2FNUZType>(this); |
317 | impl->f8E4M3FNUZTy = TypeUniquer::get<Float8E4M3FNUZType>(this); |
318 | impl->f8E4M3B11FNUZTy = TypeUniquer::get<Float8E4M3B11FNUZType>(this); |
319 | impl->bf16Ty = TypeUniquer::get<BFloat16Type>(this); |
320 | impl->f16Ty = TypeUniquer::get<Float16Type>(this); |
321 | impl->tf32Ty = TypeUniquer::get<FloatTF32Type>(this); |
322 | impl->f32Ty = TypeUniquer::get<Float32Type>(this); |
323 | impl->f64Ty = TypeUniquer::get<Float64Type>(this); |
324 | impl->f80Ty = TypeUniquer::get<Float80Type>(this); |
325 | impl->f128Ty = TypeUniquer::get<Float128Type>(this); |
326 | /// Index Type. |
327 | impl->indexTy = TypeUniquer::get<IndexType>(this); |
328 | /// Integer Types. |
329 | impl->int1Ty = TypeUniquer::get<IntegerType>(this, 1, IntegerType::Signless); |
330 | impl->int8Ty = TypeUniquer::get<IntegerType>(this, 8, IntegerType::Signless); |
331 | impl->int16Ty = |
332 | TypeUniquer::get<IntegerType>(this, 16, IntegerType::Signless); |
333 | impl->int32Ty = |
334 | TypeUniquer::get<IntegerType>(this, 32, IntegerType::Signless); |
335 | impl->int64Ty = |
336 | TypeUniquer::get<IntegerType>(this, 64, IntegerType::Signless); |
337 | impl->int128Ty = |
338 | TypeUniquer::get<IntegerType>(this, 128, IntegerType::Signless); |
339 | /// None Type. |
340 | impl->noneType = TypeUniquer::get<NoneType>(this); |
341 | |
342 | //// Attributes. |
343 | //// Note: These must be registered after the types as they may generate one |
344 | //// of the above types internally. |
345 | /// Unknown Location Attribute. |
346 | impl->unknownLocAttr = AttributeUniquer::get<UnknownLoc>(this); |
347 | /// Bool Attributes. |
348 | impl->falseAttr = IntegerAttr::getBoolAttrUnchecked(impl->int1Ty, false); |
349 | impl->trueAttr = IntegerAttr::getBoolAttrUnchecked(impl->int1Ty, true); |
350 | /// Unit Attribute. |
351 | impl->unitAttr = AttributeUniquer::get<UnitAttr>(this); |
352 | /// The empty dictionary attribute. |
353 | impl->emptyDictionaryAttr = DictionaryAttr::getEmptyUnchecked(this); |
354 | /// The empty string attribute. |
355 | impl->emptyStringAttr = StringAttr::getEmptyStringAttrUnchecked(this); |
356 | |
357 | // Register the affine storage objects with the uniquer. |
358 | impl->affineUniquer |
359 | .registerParametricStorageType<AffineBinaryOpExprStorage>(); |
360 | impl->affineUniquer |
361 | .registerParametricStorageType<AffineConstantExprStorage>(); |
362 | impl->affineUniquer.registerParametricStorageType<AffineDimExprStorage>(); |
363 | impl->affineUniquer.registerParametricStorageType<AffineMapStorage>(); |
364 | impl->affineUniquer.registerParametricStorageType<IntegerSetStorage>(); |
365 | } |
366 | |
367 | MLIRContext::~MLIRContext() = default; |
368 | |
369 | /// Copy the specified array of elements into memory managed by the provided |
370 | /// bump pointer allocator. This assumes the elements are all PODs. |
371 | template <typename T> |
372 | static ArrayRef<T> copyArrayRefInto(llvm::BumpPtrAllocator &allocator, |
373 | ArrayRef<T> elements) { |
374 | auto result = allocator.Allocate<T>(elements.size()); |
375 | std::uninitialized_copy(elements.begin(), elements.end(), result); |
376 | return ArrayRef<T>(result, elements.size()); |
377 | } |
378 | |
379 | //===----------------------------------------------------------------------===// |
380 | // Action Handling |
381 | //===----------------------------------------------------------------------===// |
382 | |
383 | void MLIRContext::registerActionHandler(HandlerTy handler) { |
384 | getImpl().actionHandler = std::move(handler); |
385 | } |
386 | |
387 | /// Dispatch the provided action to the handler if any, or just execute it. |
388 | void MLIRContext::executeActionInternal(function_ref<void()> actionFn, |
389 | const tracing::Action &action) { |
390 | assert(getImpl().actionHandler); |
391 | getImpl().actionHandler(actionFn, action); |
392 | } |
393 | |
394 | bool MLIRContext::hasActionHandler() { return (bool)getImpl().actionHandler; } |
395 | |
396 | //===----------------------------------------------------------------------===// |
397 | // Diagnostic Handlers |
398 | //===----------------------------------------------------------------------===// |
399 | |
400 | /// Returns the diagnostic engine for this context. |
401 | DiagnosticEngine &MLIRContext::getDiagEngine() { return getImpl().diagEngine; } |
402 | |
403 | //===----------------------------------------------------------------------===// |
404 | // Dialect and Operation Registration |
405 | //===----------------------------------------------------------------------===// |
406 | |
407 | void MLIRContext::appendDialectRegistry(const DialectRegistry ®istry) { |
408 | if (registry.isSubsetOf(rhs: impl->dialectsRegistry)) |
409 | return; |
410 | |
411 | assert(impl->multiThreadedExecutionContext == 0 && |
412 | "appending to the MLIRContext dialect registry while in a " |
413 | "multi-threaded execution context" ); |
414 | registry.appendTo(destination&: impl->dialectsRegistry); |
415 | |
416 | // For the already loaded dialects, apply any possible extensions immediately. |
417 | registry.applyExtensions(ctx: this); |
418 | } |
419 | |
420 | const DialectRegistry &MLIRContext::getDialectRegistry() { |
421 | return impl->dialectsRegistry; |
422 | } |
423 | |
424 | /// Return information about all registered IR dialects. |
425 | std::vector<Dialect *> MLIRContext::getLoadedDialects() { |
426 | std::vector<Dialect *> result; |
427 | result.reserve(n: impl->loadedDialects.size()); |
428 | for (auto &dialect : impl->loadedDialects) |
429 | result.push_back(x: dialect.second.get()); |
430 | llvm::array_pod_sort(Start: result.begin(), End: result.end(), |
431 | Compare: [](Dialect *const *lhs, Dialect *const *rhs) -> int { |
432 | return (*lhs)->getNamespace() < (*rhs)->getNamespace(); |
433 | }); |
434 | return result; |
435 | } |
436 | std::vector<StringRef> MLIRContext::getAvailableDialects() { |
437 | std::vector<StringRef> result; |
438 | for (auto dialect : impl->dialectsRegistry.getDialectNames()) |
439 | result.push_back(x: dialect); |
440 | return result; |
441 | } |
442 | |
443 | /// Get a registered IR dialect with the given namespace. If none is found, |
444 | /// then return nullptr. |
445 | Dialect *MLIRContext::getLoadedDialect(StringRef name) { |
446 | // Dialects are sorted by name, so we can use binary search for lookup. |
447 | auto it = impl->loadedDialects.find(Val: name); |
448 | return (it != impl->loadedDialects.end()) ? it->second.get() : nullptr; |
449 | } |
450 | |
451 | Dialect *MLIRContext::getOrLoadDialect(StringRef name) { |
452 | Dialect *dialect = getLoadedDialect(name); |
453 | if (dialect) |
454 | return dialect; |
455 | DialectAllocatorFunctionRef allocator = |
456 | impl->dialectsRegistry.getDialectAllocator(name); |
457 | return allocator ? allocator(this) : nullptr; |
458 | } |
459 | |
460 | /// Get a dialect for the provided namespace and TypeID: abort the program if a |
461 | /// dialect exist for this namespace with different TypeID. Returns a pointer to |
462 | /// the dialect owned by the context. |
463 | Dialect * |
464 | MLIRContext::getOrLoadDialect(StringRef dialectNamespace, TypeID dialectID, |
465 | function_ref<std::unique_ptr<Dialect>()> ctor) { |
466 | auto &impl = getImpl(); |
467 | // Get the correct insertion position sorted by namespace. |
468 | auto dialectIt = impl.loadedDialects.try_emplace(Key: dialectNamespace, Args: nullptr); |
469 | |
470 | if (dialectIt.second) { |
471 | LLVM_DEBUG(llvm::dbgs() |
472 | << "Load new dialect in Context " << dialectNamespace << "\n" ); |
473 | #ifndef NDEBUG |
474 | if (impl.multiThreadedExecutionContext != 0) |
475 | llvm::report_fatal_error( |
476 | reason: "Loading a dialect (" + dialectNamespace + |
477 | ") while in a multi-threaded execution context (maybe " |
478 | "the PassManager): this can indicate a " |
479 | "missing `dependentDialects` in a pass for example." ); |
480 | #endif // NDEBUG |
481 | // loadedDialects entry is initialized to nullptr, indicating that the |
482 | // dialect is currently being loaded. Re-lookup the address in |
483 | // loadedDialects because the table might have been rehashed by recursive |
484 | // dialect loading in ctor(). |
485 | std::unique_ptr<Dialect> &dialectOwned = |
486 | impl.loadedDialects[dialectNamespace] = ctor(); |
487 | Dialect *dialect = dialectOwned.get(); |
488 | assert(dialect && "dialect ctor failed" ); |
489 | |
490 | // Refresh all the identifiers dialect field, this catches cases where a |
491 | // dialect may be loaded after identifier prefixed with this dialect name |
492 | // were already created. |
493 | auto stringAttrsIt = impl.dialectReferencingStrAttrs.find(Val: dialectNamespace); |
494 | if (stringAttrsIt != impl.dialectReferencingStrAttrs.end()) { |
495 | for (StringAttrStorage *storage : stringAttrsIt->second) |
496 | storage->referencedDialect = dialect; |
497 | impl.dialectReferencingStrAttrs.erase(I: stringAttrsIt); |
498 | } |
499 | |
500 | // Apply any extensions to this newly loaded dialect. |
501 | impl.dialectsRegistry.applyExtensions(dialect); |
502 | return dialect; |
503 | } |
504 | |
505 | #ifndef NDEBUG |
506 | if (dialectIt.first->second == nullptr) |
507 | llvm::report_fatal_error( |
508 | reason: "Loading (and getting) a dialect (" + dialectNamespace + |
509 | ") while the same dialect is still loading: use loadDialect instead " |
510 | "of getOrLoadDialect." ); |
511 | #endif // NDEBUG |
512 | |
513 | // Abort if dialect with namespace has already been registered. |
514 | std::unique_ptr<Dialect> &dialect = dialectIt.first->second; |
515 | if (dialect->getTypeID() != dialectID) |
516 | llvm::report_fatal_error(reason: "a dialect with namespace '" + dialectNamespace + |
517 | "' has already been registered" ); |
518 | |
519 | return dialect.get(); |
520 | } |
521 | |
522 | bool MLIRContext::isDialectLoading(StringRef dialectNamespace) { |
523 | auto it = getImpl().loadedDialects.find(Val: dialectNamespace); |
524 | // nullptr indicates that the dialect is currently being loaded. |
525 | return it != getImpl().loadedDialects.end() && it->second == nullptr; |
526 | } |
527 | |
528 | DynamicDialect *MLIRContext::getOrLoadDynamicDialect( |
529 | StringRef dialectNamespace, function_ref<void(DynamicDialect *)> ctor) { |
530 | auto &impl = getImpl(); |
531 | // Get the correct insertion position sorted by namespace. |
532 | auto dialectIt = impl.loadedDialects.find(Val: dialectNamespace); |
533 | |
534 | if (dialectIt != impl.loadedDialects.end()) { |
535 | if (auto *dynDialect = dyn_cast<DynamicDialect>(Val: dialectIt->second.get())) |
536 | return dynDialect; |
537 | llvm::report_fatal_error(reason: "a dialect with namespace '" + dialectNamespace + |
538 | "' has already been registered" ); |
539 | } |
540 | |
541 | LLVM_DEBUG(llvm::dbgs() << "Load new dynamic dialect in Context " |
542 | << dialectNamespace << "\n" ); |
543 | #ifndef NDEBUG |
544 | if (impl.multiThreadedExecutionContext != 0) |
545 | llvm::report_fatal_error( |
546 | reason: "Loading a dynamic dialect (" + dialectNamespace + |
547 | ") while in a multi-threaded execution context (maybe " |
548 | "the PassManager): this can indicate a " |
549 | "missing `dependentDialects` in a pass for example." ); |
550 | #endif |
551 | |
552 | auto name = StringAttr::get(this, dialectNamespace); |
553 | auto *dialect = new DynamicDialect(name, this); |
554 | (void)getOrLoadDialect(name, dialect->getTypeID(), [dialect, ctor]() { |
555 | ctor(dialect); |
556 | return std::unique_ptr<DynamicDialect>(dialect); |
557 | }); |
558 | // This is the same result as `getOrLoadDialect` (if it didn't failed), |
559 | // since it has the same TypeID, and TypeIDs are unique. |
560 | return dialect; |
561 | } |
562 | |
563 | void MLIRContext::loadAllAvailableDialects() { |
564 | for (StringRef name : getAvailableDialects()) |
565 | getOrLoadDialect(name); |
566 | } |
567 | |
568 | llvm::hash_code MLIRContext::getRegistryHash() { |
569 | llvm::hash_code hash(0); |
570 | // Factor in number of loaded dialects, attributes, operations, types. |
571 | hash = llvm::hash_combine(args: hash, args: impl->loadedDialects.size()); |
572 | hash = llvm::hash_combine(args: hash, args: impl->registeredAttributes.size()); |
573 | hash = llvm::hash_combine(args: hash, args: impl->registeredOperations.size()); |
574 | hash = llvm::hash_combine(args: hash, args: impl->registeredTypes.size()); |
575 | return hash; |
576 | } |
577 | |
578 | bool MLIRContext::allowsUnregisteredDialects() { |
579 | return impl->allowUnregisteredDialects; |
580 | } |
581 | |
582 | void MLIRContext::allowUnregisteredDialects(bool allowing) { |
583 | assert(impl->multiThreadedExecutionContext == 0 && |
584 | "changing MLIRContext `allow-unregistered-dialects` configuration " |
585 | "while in a multi-threaded execution context" ); |
586 | impl->allowUnregisteredDialects = allowing; |
587 | } |
588 | |
589 | /// Return true if multi-threading is enabled by the context. |
590 | bool MLIRContext::isMultithreadingEnabled() { |
591 | return impl->threadingIsEnabled && llvm::llvm_is_multithreaded(); |
592 | } |
593 | |
594 | /// Set the flag specifying if multi-threading is disabled by the context. |
595 | void MLIRContext::disableMultithreading(bool disable) { |
596 | // This API can be overridden by the global debugging flag |
597 | // --mlir-disable-threading |
598 | if (isThreadingGloballyDisabled()) |
599 | return; |
600 | assert(impl->multiThreadedExecutionContext == 0 && |
601 | "changing MLIRContext `disable-threading` configuration while " |
602 | "in a multi-threaded execution context" ); |
603 | |
604 | impl->threadingIsEnabled = !disable; |
605 | |
606 | // Update the threading mode for each of the uniquers. |
607 | impl->affineUniquer.disableMultithreading(disable); |
608 | impl->attributeUniquer.disableMultithreading(disable); |
609 | impl->typeUniquer.disableMultithreading(disable); |
610 | |
611 | // Destroy thread pool (stop all threads) if it is no longer needed, or create |
612 | // a new one if multithreading was re-enabled. |
613 | if (disable) { |
614 | // If the thread pool is owned, explicitly set it to nullptr to avoid |
615 | // keeping a dangling pointer around. If the thread pool is externally |
616 | // owned, we don't do anything. |
617 | if (impl->ownedThreadPool) { |
618 | assert(impl->threadPool); |
619 | impl->threadPool = nullptr; |
620 | impl->ownedThreadPool.reset(); |
621 | } |
622 | } else if (!impl->threadPool) { |
623 | // The thread pool isn't externally provided. |
624 | assert(!impl->ownedThreadPool); |
625 | impl->ownedThreadPool = std::make_unique<llvm::DefaultThreadPool>(); |
626 | impl->threadPool = impl->ownedThreadPool.get(); |
627 | } |
628 | } |
629 | |
630 | void MLIRContext::setThreadPool(llvm::ThreadPoolInterface &pool) { |
631 | assert(!isMultithreadingEnabled() && |
632 | "expected multi-threading to be disabled when setting a ThreadPool" ); |
633 | impl->threadPool = &pool; |
634 | impl->ownedThreadPool.reset(); |
635 | enableMultithreading(); |
636 | } |
637 | |
638 | unsigned MLIRContext::getNumThreads() { |
639 | if (isMultithreadingEnabled()) { |
640 | assert(impl->threadPool && |
641 | "multi-threading is enabled but threadpool not set" ); |
642 | return impl->threadPool->getMaxConcurrency(); |
643 | } |
644 | // No multithreading or active thread pool. Return 1 thread. |
645 | return 1; |
646 | } |
647 | |
648 | llvm::ThreadPoolInterface &MLIRContext::getThreadPool() { |
649 | assert(isMultithreadingEnabled() && |
650 | "expected multi-threading to be enabled within the context" ); |
651 | assert(impl->threadPool && |
652 | "multi-threading is enabled but threadpool not set" ); |
653 | return *impl->threadPool; |
654 | } |
655 | |
656 | void MLIRContext::enterMultiThreadedExecution() { |
657 | #ifndef NDEBUG |
658 | ++impl->multiThreadedExecutionContext; |
659 | #endif |
660 | } |
661 | void MLIRContext::exitMultiThreadedExecution() { |
662 | #ifndef NDEBUG |
663 | --impl->multiThreadedExecutionContext; |
664 | #endif |
665 | } |
666 | |
667 | /// Return true if we should attach the operation to diagnostics emitted via |
668 | /// Operation::emit. |
669 | bool MLIRContext::shouldPrintOpOnDiagnostic() { |
670 | return impl->printOpOnDiagnostic; |
671 | } |
672 | |
673 | /// Set the flag specifying if we should attach the operation to diagnostics |
674 | /// emitted via Operation::emit. |
675 | void MLIRContext::printOpOnDiagnostic(bool enable) { |
676 | assert(impl->multiThreadedExecutionContext == 0 && |
677 | "changing MLIRContext `print-op-on-diagnostic` configuration while in " |
678 | "a multi-threaded execution context" ); |
679 | impl->printOpOnDiagnostic = enable; |
680 | } |
681 | |
682 | /// Return true if we should attach the current stacktrace to diagnostics when |
683 | /// emitted. |
684 | bool MLIRContext::shouldPrintStackTraceOnDiagnostic() { |
685 | return impl->printStackTraceOnDiagnostic; |
686 | } |
687 | |
688 | /// Set the flag specifying if we should attach the current stacktrace when |
689 | /// emitting diagnostics. |
690 | void MLIRContext::printStackTraceOnDiagnostic(bool enable) { |
691 | assert(impl->multiThreadedExecutionContext == 0 && |
692 | "changing MLIRContext `print-stacktrace-on-diagnostic` configuration " |
693 | "while in a multi-threaded execution context" ); |
694 | impl->printStackTraceOnDiagnostic = enable; |
695 | } |
696 | |
697 | /// Return information about all registered operations. |
698 | ArrayRef<RegisteredOperationName> MLIRContext::getRegisteredOperations() { |
699 | return impl->sortedRegisteredOperations; |
700 | } |
701 | |
702 | bool MLIRContext::isOperationRegistered(StringRef name) { |
703 | return RegisteredOperationName::lookup(name, ctx: this).has_value(); |
704 | } |
705 | |
706 | void Dialect::addType(TypeID typeID, AbstractType &&typeInfo) { |
707 | auto &impl = context->getImpl(); |
708 | assert(impl.multiThreadedExecutionContext == 0 && |
709 | "Registering a new type kind while in a multi-threaded execution " |
710 | "context" ); |
711 | auto *newInfo = |
712 | new (impl.abstractDialectSymbolAllocator.Allocate<AbstractType>()) |
713 | AbstractType(std::move(typeInfo)); |
714 | if (!impl.registeredTypes.insert(KV: {typeID, newInfo}).second) |
715 | llvm::report_fatal_error(reason: "Dialect Type already registered." ); |
716 | if (!impl.nameToType.insert(KV: {newInfo->getName(), newInfo}).second) |
717 | llvm::report_fatal_error(reason: "Dialect Type with name " + newInfo->getName() + |
718 | " is already registered." ); |
719 | } |
720 | |
721 | void Dialect::addAttribute(TypeID typeID, AbstractAttribute &&attrInfo) { |
722 | auto &impl = context->getImpl(); |
723 | assert(impl.multiThreadedExecutionContext == 0 && |
724 | "Registering a new attribute kind while in a multi-threaded execution " |
725 | "context" ); |
726 | auto *newInfo = |
727 | new (impl.abstractDialectSymbolAllocator.Allocate<AbstractAttribute>()) |
728 | AbstractAttribute(std::move(attrInfo)); |
729 | if (!impl.registeredAttributes.insert(KV: {typeID, newInfo}).second) |
730 | llvm::report_fatal_error(reason: "Dialect Attribute already registered." ); |
731 | if (!impl.nameToAttribute.insert(KV: {newInfo->getName(), newInfo}).second) |
732 | llvm::report_fatal_error(reason: "Dialect Attribute with name " + |
733 | newInfo->getName() + " is already registered." ); |
734 | } |
735 | |
736 | //===----------------------------------------------------------------------===// |
737 | // AbstractAttribute |
738 | //===----------------------------------------------------------------------===// |
739 | |
740 | /// Get the dialect that registered the attribute with the provided typeid. |
741 | const AbstractAttribute &AbstractAttribute::lookup(TypeID typeID, |
742 | MLIRContext *context) { |
743 | const AbstractAttribute *abstract = lookupMutable(typeID, context); |
744 | if (!abstract) |
745 | llvm::report_fatal_error(reason: "Trying to create an Attribute that was not " |
746 | "registered in this MLIRContext." ); |
747 | return *abstract; |
748 | } |
749 | |
750 | AbstractAttribute *AbstractAttribute::lookupMutable(TypeID typeID, |
751 | MLIRContext *context) { |
752 | auto &impl = context->getImpl(); |
753 | return impl.registeredAttributes.lookup(Val: typeID); |
754 | } |
755 | |
756 | std::optional<std::reference_wrapper<const AbstractAttribute>> |
757 | AbstractAttribute::lookup(StringRef name, MLIRContext *context) { |
758 | MLIRContextImpl &impl = context->getImpl(); |
759 | const AbstractAttribute *type = impl.nameToAttribute.lookup(Val: name); |
760 | |
761 | if (!type) |
762 | return std::nullopt; |
763 | return {*type}; |
764 | } |
765 | |
766 | //===----------------------------------------------------------------------===// |
767 | // OperationName |
768 | //===----------------------------------------------------------------------===// |
769 | |
770 | OperationName::Impl::Impl(StringRef name, Dialect *dialect, TypeID typeID, |
771 | detail::InterfaceMap interfaceMap) |
772 | : Impl(StringAttr::get(dialect->getContext(), name), dialect, typeID, |
773 | std::move(interfaceMap)) {} |
774 | |
775 | OperationName::OperationName(StringRef name, MLIRContext *context) { |
776 | MLIRContextImpl &ctxImpl = context->getImpl(); |
777 | |
778 | // Check for an existing name in read-only mode. |
779 | bool isMultithreadingEnabled = context->isMultithreadingEnabled(); |
780 | if (isMultithreadingEnabled) { |
781 | // Check the registered info map first. In the overwhelmingly common case, |
782 | // the entry will be in here and it also removes the need to acquire any |
783 | // locks. |
784 | auto registeredIt = ctxImpl.registeredOperationsByName.find(Key: name); |
785 | if (LLVM_LIKELY(registeredIt != ctxImpl.registeredOperationsByName.end())) { |
786 | impl = registeredIt->second.impl; |
787 | return; |
788 | } |
789 | |
790 | llvm::sys::SmartScopedReader<true> contextLock(ctxImpl.operationInfoMutex); |
791 | auto it = ctxImpl.operations.find(Key: name); |
792 | if (it != ctxImpl.operations.end()) { |
793 | impl = it->second.get(); |
794 | return; |
795 | } |
796 | } |
797 | |
798 | // Acquire a writer-lock so that we can safely create the new instance. |
799 | ScopedWriterLock lock(ctxImpl.operationInfoMutex, isMultithreadingEnabled); |
800 | |
801 | auto it = ctxImpl.operations.insert(KV: {name, nullptr}); |
802 | if (it.second) { |
803 | auto nameAttr = StringAttr::get(context, name); |
804 | it.first->second = std::make_unique<UnregisteredOpModel>( |
805 | nameAttr, nameAttr.getReferencedDialect(), TypeID::get<void>(), |
806 | detail::InterfaceMap()); |
807 | } |
808 | impl = it.first->second.get(); |
809 | } |
810 | |
811 | StringRef OperationName::getDialectNamespace() const { |
812 | if (Dialect *dialect = getDialect()) |
813 | return dialect->getNamespace(); |
814 | return getStringRef().split(Separator: '.').first; |
815 | } |
816 | |
817 | LogicalResult |
818 | OperationName::UnregisteredOpModel::foldHook(Operation *, ArrayRef<Attribute>, |
819 | SmallVectorImpl<OpFoldResult> &) { |
820 | return failure(); |
821 | } |
822 | void OperationName::UnregisteredOpModel::getCanonicalizationPatterns( |
823 | RewritePatternSet &, MLIRContext *) {} |
824 | bool OperationName::UnregisteredOpModel::hasTrait(TypeID) { return false; } |
825 | |
826 | OperationName::ParseAssemblyFn |
827 | OperationName::UnregisteredOpModel::getParseAssemblyFn() { |
828 | llvm::report_fatal_error(reason: "getParseAssemblyFn hook called on unregistered op" ); |
829 | } |
830 | void OperationName::UnregisteredOpModel::populateDefaultAttrs( |
831 | const OperationName &, NamedAttrList &) {} |
832 | void OperationName::UnregisteredOpModel::printAssembly( |
833 | Operation *op, OpAsmPrinter &p, StringRef defaultDialect) { |
834 | p.printGenericOp(op); |
835 | } |
836 | LogicalResult |
837 | OperationName::UnregisteredOpModel::verifyInvariants(Operation *) { |
838 | return success(); |
839 | } |
840 | LogicalResult |
841 | OperationName::UnregisteredOpModel::verifyRegionInvariants(Operation *) { |
842 | return success(); |
843 | } |
844 | |
845 | std::optional<Attribute> |
846 | OperationName::UnregisteredOpModel::getInherentAttr(Operation *op, |
847 | StringRef name) { |
848 | auto dict = dyn_cast_or_null<DictionaryAttr>(getPropertiesAsAttr(op)); |
849 | if (!dict) |
850 | return std::nullopt; |
851 | if (Attribute attr = dict.get(name)) |
852 | return attr; |
853 | return std::nullopt; |
854 | } |
855 | void OperationName::UnregisteredOpModel::setInherentAttr(Operation *op, |
856 | StringAttr name, |
857 | Attribute value) { |
858 | auto dict = dyn_cast_or_null<DictionaryAttr>(getPropertiesAsAttr(op)); |
859 | assert(dict); |
860 | NamedAttrList attrs(dict); |
861 | attrs.set(name, value); |
862 | *op->getPropertiesStorage().as<Attribute *>() = |
863 | attrs.getDictionary(op->getContext()); |
864 | } |
865 | void OperationName::UnregisteredOpModel::populateInherentAttrs( |
866 | Operation *op, NamedAttrList &attrs) {} |
867 | LogicalResult OperationName::UnregisteredOpModel::verifyInherentAttrs( |
868 | OperationName opName, NamedAttrList &attributes, |
869 | function_ref<InFlightDiagnostic()> emitError) { |
870 | return success(); |
871 | } |
872 | int OperationName::UnregisteredOpModel::getOpPropertyByteSize() { |
873 | return sizeof(Attribute); |
874 | } |
875 | void OperationName::UnregisteredOpModel::initProperties( |
876 | OperationName opName, OpaqueProperties storage, OpaqueProperties init) { |
877 | new (storage.as<Attribute *>()) Attribute(); |
878 | } |
879 | void OperationName::UnregisteredOpModel::deleteProperties( |
880 | OpaqueProperties prop) { |
881 | prop.as<Attribute *>()->~Attribute(); |
882 | } |
883 | void OperationName::UnregisteredOpModel::populateDefaultProperties( |
884 | OperationName opName, OpaqueProperties properties) {} |
885 | LogicalResult OperationName::UnregisteredOpModel::setPropertiesFromAttr( |
886 | OperationName opName, OpaqueProperties properties, Attribute attr, |
887 | function_ref<InFlightDiagnostic()> emitError) { |
888 | *properties.as<Attribute *>() = attr; |
889 | return success(); |
890 | } |
891 | Attribute |
892 | OperationName::UnregisteredOpModel::getPropertiesAsAttr(Operation *op) { |
893 | return *op->getPropertiesStorage().as<Attribute *>(); |
894 | } |
895 | void OperationName::UnregisteredOpModel::copyProperties(OpaqueProperties lhs, |
896 | OpaqueProperties rhs) { |
897 | *lhs.as<Attribute *>() = *rhs.as<Attribute *>(); |
898 | } |
899 | bool OperationName::UnregisteredOpModel::compareProperties( |
900 | OpaqueProperties lhs, OpaqueProperties rhs) { |
901 | return *lhs.as<Attribute *>() == *rhs.as<Attribute *>(); |
902 | } |
903 | llvm::hash_code |
904 | OperationName::UnregisteredOpModel::hashProperties(OpaqueProperties prop) { |
905 | return llvm::hash_combine(args: *prop.as<Attribute *>()); |
906 | } |
907 | |
908 | //===----------------------------------------------------------------------===// |
909 | // RegisteredOperationName |
910 | //===----------------------------------------------------------------------===// |
911 | |
912 | std::optional<RegisteredOperationName> |
913 | RegisteredOperationName::lookup(TypeID typeID, MLIRContext *ctx) { |
914 | auto &impl = ctx->getImpl(); |
915 | auto it = impl.registeredOperations.find(Val: typeID); |
916 | if (it != impl.registeredOperations.end()) |
917 | return it->second; |
918 | return std::nullopt; |
919 | } |
920 | |
921 | std::optional<RegisteredOperationName> |
922 | RegisteredOperationName::lookup(StringRef name, MLIRContext *ctx) { |
923 | auto &impl = ctx->getImpl(); |
924 | auto it = impl.registeredOperationsByName.find(Key: name); |
925 | if (it != impl.registeredOperationsByName.end()) |
926 | return it->getValue(); |
927 | return std::nullopt; |
928 | } |
929 | |
930 | void RegisteredOperationName::insert( |
931 | std::unique_ptr<RegisteredOperationName::Impl> ownedImpl, |
932 | ArrayRef<StringRef> attrNames) { |
933 | RegisteredOperationName::Impl *impl = ownedImpl.get(); |
934 | MLIRContext *ctx = impl->getDialect()->getContext(); |
935 | auto &ctxImpl = ctx->getImpl(); |
936 | assert(ctxImpl.multiThreadedExecutionContext == 0 && |
937 | "registering a new operation kind while in a multi-threaded execution " |
938 | "context" ); |
939 | |
940 | // Register the attribute names of this operation. |
941 | MutableArrayRef<StringAttr> cachedAttrNames; |
942 | if (!attrNames.empty()) { |
943 | cachedAttrNames = MutableArrayRef<StringAttr>( |
944 | ctxImpl.abstractDialectSymbolAllocator.Allocate<StringAttr>( |
945 | Num: attrNames.size()), |
946 | attrNames.size()); |
947 | for (unsigned i : llvm::seq<unsigned>(0, attrNames.size())) |
948 | new (&cachedAttrNames[i]) StringAttr(StringAttr::get(ctx, attrNames[i])); |
949 | impl->attributeNames = cachedAttrNames; |
950 | } |
951 | StringRef name = impl->getName().strref(); |
952 | // Insert the operation info if it doesn't exist yet. |
953 | auto it = ctxImpl.operations.insert(KV: {name, nullptr}); |
954 | it.first->second = std::move(ownedImpl); |
955 | |
956 | // Update the registered info for this operation. |
957 | auto emplaced = ctxImpl.registeredOperations.try_emplace( |
958 | Key: impl->getTypeID(), Args: RegisteredOperationName(impl)); |
959 | assert(emplaced.second && "operation name registration must be successful" ); |
960 | auto emplacedByName = ctxImpl.registeredOperationsByName.try_emplace( |
961 | Key: name, Args: RegisteredOperationName(impl)); |
962 | (void)emplacedByName; |
963 | assert(emplacedByName.second && |
964 | "operation name registration must be successful" ); |
965 | |
966 | // Add emplaced operation name to the sorted operations container. |
967 | RegisteredOperationName &value = emplaced.first->second; |
968 | ctxImpl.sortedRegisteredOperations.insert( |
969 | I: llvm::upper_bound(Range&: ctxImpl.sortedRegisteredOperations, Value&: value, |
970 | C: [](auto &lhs, auto &rhs) { |
971 | return lhs.getIdentifier().compare( |
972 | rhs.getIdentifier()); |
973 | }), |
974 | Elt: value); |
975 | } |
976 | |
977 | //===----------------------------------------------------------------------===// |
978 | // AbstractType |
979 | //===----------------------------------------------------------------------===// |
980 | |
981 | const AbstractType &AbstractType::lookup(TypeID typeID, MLIRContext *context) { |
982 | const AbstractType *type = lookupMutable(typeID, context); |
983 | if (!type) |
984 | llvm::report_fatal_error( |
985 | reason: "Trying to create a Type that was not registered in this MLIRContext." ); |
986 | return *type; |
987 | } |
988 | |
989 | AbstractType *AbstractType::lookupMutable(TypeID typeID, MLIRContext *context) { |
990 | auto &impl = context->getImpl(); |
991 | return impl.registeredTypes.lookup(Val: typeID); |
992 | } |
993 | |
994 | std::optional<std::reference_wrapper<const AbstractType>> |
995 | AbstractType::lookup(StringRef name, MLIRContext *context) { |
996 | MLIRContextImpl &impl = context->getImpl(); |
997 | const AbstractType *type = impl.nameToType.lookup(Val: name); |
998 | |
999 | if (!type) |
1000 | return std::nullopt; |
1001 | return {*type}; |
1002 | } |
1003 | |
1004 | //===----------------------------------------------------------------------===// |
1005 | // Type uniquing |
1006 | //===----------------------------------------------------------------------===// |
1007 | |
1008 | /// Returns the storage uniquer used for constructing type storage instances. |
1009 | /// This should not be used directly. |
1010 | StorageUniquer &MLIRContext::getTypeUniquer() { return getImpl().typeUniquer; } |
1011 | |
1012 | Float8E5M2Type Float8E5M2Type::get(MLIRContext *context) { |
1013 | return context->getImpl().f8E5M2Ty; |
1014 | } |
1015 | Float8E4M3FNType Float8E4M3FNType::get(MLIRContext *context) { |
1016 | return context->getImpl().f8E4M3FNTy; |
1017 | } |
1018 | Float8E5M2FNUZType Float8E5M2FNUZType::get(MLIRContext *context) { |
1019 | return context->getImpl().f8E5M2FNUZTy; |
1020 | } |
1021 | Float8E4M3FNUZType Float8E4M3FNUZType::get(MLIRContext *context) { |
1022 | return context->getImpl().f8E4M3FNUZTy; |
1023 | } |
1024 | Float8E4M3B11FNUZType Float8E4M3B11FNUZType::get(MLIRContext *context) { |
1025 | return context->getImpl().f8E4M3B11FNUZTy; |
1026 | } |
1027 | BFloat16Type BFloat16Type::get(MLIRContext *context) { |
1028 | return context->getImpl().bf16Ty; |
1029 | } |
1030 | Float16Type Float16Type::get(MLIRContext *context) { |
1031 | return context->getImpl().f16Ty; |
1032 | } |
1033 | FloatTF32Type FloatTF32Type::get(MLIRContext *context) { |
1034 | return context->getImpl().tf32Ty; |
1035 | } |
1036 | Float32Type Float32Type::get(MLIRContext *context) { |
1037 | return context->getImpl().f32Ty; |
1038 | } |
1039 | Float64Type Float64Type::get(MLIRContext *context) { |
1040 | return context->getImpl().f64Ty; |
1041 | } |
1042 | Float80Type Float80Type::get(MLIRContext *context) { |
1043 | return context->getImpl().f80Ty; |
1044 | } |
1045 | Float128Type Float128Type::get(MLIRContext *context) { |
1046 | return context->getImpl().f128Ty; |
1047 | } |
1048 | |
1049 | /// Get an instance of the IndexType. |
1050 | IndexType IndexType::get(MLIRContext *context) { |
1051 | return context->getImpl().indexTy; |
1052 | } |
1053 | |
1054 | /// Return an existing integer type instance if one is cached within the |
1055 | /// context. |
1056 | static IntegerType |
1057 | getCachedIntegerType(unsigned width, |
1058 | IntegerType::SignednessSemantics signedness, |
1059 | MLIRContext *context) { |
1060 | if (signedness != IntegerType::Signless) |
1061 | return IntegerType(); |
1062 | |
1063 | switch (width) { |
1064 | case 1: |
1065 | return context->getImpl().int1Ty; |
1066 | case 8: |
1067 | return context->getImpl().int8Ty; |
1068 | case 16: |
1069 | return context->getImpl().int16Ty; |
1070 | case 32: |
1071 | return context->getImpl().int32Ty; |
1072 | case 64: |
1073 | return context->getImpl().int64Ty; |
1074 | case 128: |
1075 | return context->getImpl().int128Ty; |
1076 | default: |
1077 | return IntegerType(); |
1078 | } |
1079 | } |
1080 | |
1081 | IntegerType IntegerType::get(MLIRContext *context, unsigned width, |
1082 | IntegerType::SignednessSemantics signedness) { |
1083 | if (auto cached = getCachedIntegerType(width, signedness, context)) |
1084 | return cached; |
1085 | return Base::get(context, width, signedness); |
1086 | } |
1087 | |
1088 | IntegerType |
1089 | IntegerType::getChecked(function_ref<InFlightDiagnostic()> emitError, |
1090 | MLIRContext *context, unsigned width, |
1091 | SignednessSemantics signedness) { |
1092 | if (auto cached = getCachedIntegerType(width, signedness, context)) |
1093 | return cached; |
1094 | return Base::getChecked(emitError, context, width, signedness); |
1095 | } |
1096 | |
1097 | /// Get an instance of the NoneType. |
1098 | NoneType NoneType::get(MLIRContext *context) { |
1099 | if (NoneType cachedInst = context->getImpl().noneType) |
1100 | return cachedInst; |
1101 | // Note: May happen when initializing the singleton attributes of the builtin |
1102 | // dialect. |
1103 | return Base::get(context); |
1104 | } |
1105 | |
1106 | //===----------------------------------------------------------------------===// |
1107 | // Attribute uniquing |
1108 | //===----------------------------------------------------------------------===// |
1109 | |
1110 | /// Returns the storage uniquer used for constructing attribute storage |
1111 | /// instances. This should not be used directly. |
1112 | StorageUniquer &MLIRContext::getAttributeUniquer() { |
1113 | return getImpl().attributeUniquer; |
1114 | } |
1115 | |
1116 | /// Initialize the given attribute storage instance. |
1117 | void AttributeUniquer::initializeAttributeStorage(AttributeStorage *storage, |
1118 | MLIRContext *ctx, |
1119 | TypeID attrID) { |
1120 | storage->initializeAbstractAttribute(abstractAttr: AbstractAttribute::lookup(typeID: attrID, context: ctx)); |
1121 | } |
1122 | |
1123 | BoolAttr BoolAttr::get(MLIRContext *context, bool value) { |
1124 | return value ? context->getImpl().trueAttr : context->getImpl().falseAttr; |
1125 | } |
1126 | |
1127 | UnitAttr UnitAttr::get(MLIRContext *context) { |
1128 | return context->getImpl().unitAttr; |
1129 | } |
1130 | |
1131 | UnknownLoc UnknownLoc::get(MLIRContext *context) { |
1132 | return context->getImpl().unknownLocAttr; |
1133 | } |
1134 | |
1135 | DistinctAttrStorage * |
1136 | detail::DistinctAttributeUniquer::allocateStorage(MLIRContext *context, |
1137 | Attribute referencedAttr) { |
1138 | return context->getImpl().distinctAttributeAllocator.allocate(referencedAttr); |
1139 | } |
1140 | |
1141 | /// Return empty dictionary. |
1142 | DictionaryAttr DictionaryAttr::getEmpty(MLIRContext *context) { |
1143 | return context->getImpl().emptyDictionaryAttr; |
1144 | } |
1145 | |
1146 | void StringAttrStorage::initialize(MLIRContext *context) { |
1147 | // Check for a dialect namespace prefix, if there isn't one we don't need to |
1148 | // do any additional initialization. |
1149 | auto dialectNamePair = value.split(Separator: '.'); |
1150 | if (dialectNamePair.first.empty() || dialectNamePair.second.empty()) |
1151 | return; |
1152 | |
1153 | // If one exists, we check to see if this dialect is loaded. If it is, we set |
1154 | // the dialect now, if it isn't we record this storage for initialization |
1155 | // later if the dialect ever gets loaded. |
1156 | if ((referencedDialect = context->getLoadedDialect(name: dialectNamePair.first))) |
1157 | return; |
1158 | |
1159 | MLIRContextImpl &impl = context->getImpl(); |
1160 | llvm::sys::SmartScopedLock<true> lock(impl.dialectRefStrAttrMutex); |
1161 | impl.dialectReferencingStrAttrs[dialectNamePair.first].push_back(Elt: this); |
1162 | } |
1163 | |
1164 | /// Return an empty string. |
1165 | StringAttr StringAttr::get(MLIRContext *context) { |
1166 | return context->getImpl().emptyStringAttr; |
1167 | } |
1168 | |
1169 | //===----------------------------------------------------------------------===// |
1170 | // AffineMap uniquing |
1171 | //===----------------------------------------------------------------------===// |
1172 | |
1173 | StorageUniquer &MLIRContext::getAffineUniquer() { |
1174 | return getImpl().affineUniquer; |
1175 | } |
1176 | |
1177 | AffineMap AffineMap::getImpl(unsigned dimCount, unsigned symbolCount, |
1178 | ArrayRef<AffineExpr> results, |
1179 | MLIRContext *context) { |
1180 | auto &impl = context->getImpl(); |
1181 | auto *storage = impl.affineUniquer.get<AffineMapStorage>( |
1182 | initFn: [&](AffineMapStorage *storage) { storage->context = context; }, args&: dimCount, |
1183 | args&: symbolCount, args&: results); |
1184 | return AffineMap(storage); |
1185 | } |
1186 | |
1187 | /// Check whether the arguments passed to the AffineMap::get() are consistent. |
1188 | /// This method checks whether the highest index of dimensional identifier |
1189 | /// present in result expressions is less than `dimCount` and the highest index |
1190 | /// of symbolic identifier present in result expressions is less than |
1191 | /// `symbolCount`. |
1192 | LLVM_ATTRIBUTE_UNUSED static bool |
1193 | willBeValidAffineMap(unsigned dimCount, unsigned symbolCount, |
1194 | ArrayRef<AffineExpr> results) { |
1195 | int64_t maxDimPosition = -1; |
1196 | int64_t maxSymbolPosition = -1; |
1197 | getMaxDimAndSymbol(exprsList: ArrayRef<ArrayRef<AffineExpr>>(results), maxDim&: maxDimPosition, |
1198 | maxSym&: maxSymbolPosition); |
1199 | if ((maxDimPosition >= dimCount) || (maxSymbolPosition >= symbolCount)) { |
1200 | LLVM_DEBUG( |
1201 | llvm::dbgs() |
1202 | << "maximum dimensional identifier position in result expression must " |
1203 | "be less than `dimCount` and maximum symbolic identifier position " |
1204 | "in result expression must be less than `symbolCount`\n" ); |
1205 | return false; |
1206 | } |
1207 | return true; |
1208 | } |
1209 | |
1210 | AffineMap AffineMap::get(MLIRContext *context) { |
1211 | return getImpl(/*dimCount=*/0, /*symbolCount=*/0, /*results=*/{}, context); |
1212 | } |
1213 | |
1214 | AffineMap AffineMap::get(unsigned dimCount, unsigned symbolCount, |
1215 | MLIRContext *context) { |
1216 | return getImpl(dimCount, symbolCount, /*results=*/{}, context); |
1217 | } |
1218 | |
1219 | AffineMap AffineMap::get(unsigned dimCount, unsigned symbolCount, |
1220 | AffineExpr result) { |
1221 | assert(willBeValidAffineMap(dimCount, symbolCount, {result})); |
1222 | return getImpl(dimCount, symbolCount, results: {result}, context: result.getContext()); |
1223 | } |
1224 | |
1225 | AffineMap AffineMap::get(unsigned dimCount, unsigned symbolCount, |
1226 | ArrayRef<AffineExpr> results, MLIRContext *context) { |
1227 | assert(willBeValidAffineMap(dimCount, symbolCount, results)); |
1228 | return getImpl(dimCount, symbolCount, results, context); |
1229 | } |
1230 | |
1231 | //===----------------------------------------------------------------------===// |
1232 | // Integer Sets: these are allocated into the bump pointer, and are immutable. |
1233 | // Unlike AffineMap's, these are uniqued only if they are small. |
1234 | //===----------------------------------------------------------------------===// |
1235 | |
1236 | IntegerSet IntegerSet::get(unsigned dimCount, unsigned symbolCount, |
1237 | ArrayRef<AffineExpr> constraints, |
1238 | ArrayRef<bool> eqFlags) { |
1239 | // The number of constraints can't be zero. |
1240 | assert(!constraints.empty()); |
1241 | assert(constraints.size() == eqFlags.size()); |
1242 | |
1243 | auto &impl = constraints[0].getContext()->getImpl(); |
1244 | auto *storage = impl.affineUniquer.get<IntegerSetStorage>( |
1245 | initFn: [](IntegerSetStorage *) {}, args&: dimCount, args&: symbolCount, args&: constraints, args&: eqFlags); |
1246 | return IntegerSet(storage); |
1247 | } |
1248 | |
1249 | //===----------------------------------------------------------------------===// |
1250 | // StorageUniquerSupport |
1251 | //===----------------------------------------------------------------------===// |
1252 | |
1253 | /// Utility method to generate a callback that can be used to generate a |
1254 | /// diagnostic when checking the construction invariants of a storage object. |
1255 | /// This is defined out-of-line to avoid the need to include Location.h. |
1256 | llvm::unique_function<InFlightDiagnostic()> |
1257 | mlir::detail::getDefaultDiagnosticEmitFn(MLIRContext *ctx) { |
1258 | return [ctx] { return emitError(UnknownLoc::get(ctx)); }; |
1259 | } |
1260 | llvm::unique_function<InFlightDiagnostic()> |
1261 | mlir::detail::getDefaultDiagnosticEmitFn(const Location &loc) { |
1262 | return [=] { return emitError(loc); }; |
1263 | } |
1264 | |