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

source code of mlir/lib/IR/MLIRContext.cpp