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 "LocationDetail.h"
15 #include "TypeDetail.h"
16 #include "mlir/IR/AffineExpr.h"
17 #include "mlir/IR/AffineMap.h"
18 #include "mlir/IR/Attributes.h"
19 #include "mlir/IR/BuiltinDialect.h"
20 #include "mlir/IR/BuiltinOps.h"
21 #include "mlir/IR/Diagnostics.h"
22 #include "mlir/IR/Dialect.h"
23 #include "mlir/IR/Identifier.h"
24 #include "mlir/IR/IntegerSet.h"
25 #include "mlir/IR/Location.h"
26 #include "mlir/IR/OpImplementation.h"
27 #include "mlir/IR/Types.h"
28 #include "mlir/Support/ThreadLocalCache.h"
29 #include "llvm/ADT/DenseMap.h"
30 #include "llvm/ADT/DenseSet.h"
31 #include "llvm/ADT/SetVector.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/Debug.h"
37 #include "llvm/Support/RWMutex.h"
38 #include "llvm/Support/raw_ostream.h"
39 #include <memory>
40 
41 #define DEBUG_TYPE "mlircontext"
42 
43 using namespace mlir;
44 using namespace mlir::detail;
45 
46 using llvm::hash_combine;
47 using llvm::hash_combine_range;
48 
49 //===----------------------------------------------------------------------===//
50 // MLIRContext CommandLine Options
51 //===----------------------------------------------------------------------===//
52 
53 namespace {
54 /// This struct contains command line options that can be used to initialize
55 /// various bits of an MLIRContext. This uses a struct wrapper to avoid the need
56 /// for global command line options.
57 struct MLIRContextOptions {
58   llvm::cl::opt<bool> disableThreading{
59       "mlir-disable-threading",
60       llvm::cl::desc("Disabling multi-threading within MLIR")};
61 
62   llvm::cl::opt<bool> printOpOnDiagnostic{
63       "mlir-print-op-on-diagnostic",
64       llvm::cl::desc("When a diagnostic is emitted on an operation, also print "
65                      "the operation as an attached note"),
66       llvm::cl::init(true)};
67 
68   llvm::cl::opt<bool> printStackTraceOnDiagnostic{
69       "mlir-print-stacktrace-on-diagnostic",
70       llvm::cl::desc("When a diagnostic is emitted, also print the stack trace "
71                      "as an attached note")};
72 };
73 } // end anonymous namespace
74 
75 static llvm::ManagedStatic<MLIRContextOptions> clOptions;
76 
77 /// Register a set of useful command-line options that can be used to configure
78 /// various flags within the MLIRContext. These flags are used when constructing
79 /// an MLIR context for initialization.
80 void mlir::registerMLIRContextCLOptions() {
81   // Make sure that the options struct has been initialized.
82   *clOptions;
83 }
84 
85 //===----------------------------------------------------------------------===//
86 // Locking Utilities
87 //===----------------------------------------------------------------------===//
88 
89 namespace {
90 /// Utility reader lock that takes a runtime flag that specifies if we really
91 /// need to lock.
92 struct ScopedReaderLock {
93   ScopedReaderLock(llvm::sys::SmartRWMutex<true> &mutexParam, bool shouldLock)
94       : mutex(shouldLock ? &mutexParam : nullptr) {
95     if (mutex)
96       mutex->lock_shared();
97   }
98   ~ScopedReaderLock() {
99     if (mutex)
100       mutex->unlock_shared();
101   }
102   llvm::sys::SmartRWMutex<true> *mutex;
103 };
104 /// Utility writer lock that takes a runtime flag that specifies if we really
105 /// need to lock.
106 struct ScopedWriterLock {
107   ScopedWriterLock(llvm::sys::SmartRWMutex<true> &mutexParam, bool shouldLock)
108       : mutex(shouldLock ? &mutexParam : nullptr) {
109     if (mutex)
110       mutex->lock();
111   }
112   ~ScopedWriterLock() {
113     if (mutex)
114       mutex->unlock();
115   }
116   llvm::sys::SmartRWMutex<true> *mutex;
117 };
118 } // end anonymous namespace.
119 
120 //===----------------------------------------------------------------------===//
121 // AffineMap and IntegerSet hashing
122 //===----------------------------------------------------------------------===//
123 
124 /// A utility function to safely get or create a uniqued instance within the
125 /// given set container.
126 template <typename ValueT, typename DenseInfoT, typename KeyT,
127           typename ConstructorFn>
128 static ValueT safeGetOrCreate(DenseSet<ValueT, DenseInfoT> &container,
129                               KeyT &&key, llvm::sys::SmartRWMutex<true> &mutex,
130                               bool threadingIsEnabled,
131                               ConstructorFn &&constructorFn) {
132   // Check for an existing instance in read-only mode.
133   if (threadingIsEnabled) {
134     llvm::sys::SmartScopedReader<true> instanceLock(mutex);
135     auto it = container.find_as(key);
136     if (it != container.end())
137       return *it;
138   }
139 
140   // Acquire a writer-lock so that we can safely create the new instance.
141   ScopedWriterLock instanceLock(mutex, threadingIsEnabled);
142 
143   // Check for an existing instance again here, because another writer thread
144   // may have already created one. Otherwise, construct a new instance.
145   auto existing = container.insert_as(ValueT(), key);
146   if (existing.second)
147     return *existing.first = constructorFn();
148   return *existing.first;
149 }
150 
151 namespace {
152 struct AffineMapKeyInfo : DenseMapInfo<AffineMap> {
153   // Affine maps are uniqued based on their dim/symbol counts and affine
154   // expressions.
155   using KeyTy = std::tuple<unsigned, unsigned, ArrayRef<AffineExpr>>;
156   using DenseMapInfo<AffineMap>::isEqual;
157 
158   static unsigned getHashValue(const AffineMap &key) {
159     return getHashValue(
160         KeyTy(key.getNumDims(), key.getNumSymbols(), key.getResults()));
161   }
162 
163   static unsigned getHashValue(KeyTy key) {
164     return hash_combine(
165         std::get<0>(key), std::get<1>(key),
166         hash_combine_range(std::get<2>(key).begin(), std::get<2>(key).end()));
167   }
168 
169   static bool isEqual(const KeyTy &lhs, AffineMap rhs) {
170     if (rhs == getEmptyKey() || rhs == getTombstoneKey())
171       return false;
172     return lhs == std::make_tuple(rhs.getNumDims(), rhs.getNumSymbols(),
173                                   rhs.getResults());
174   }
175 };
176 
177 struct IntegerSetKeyInfo : DenseMapInfo<IntegerSet> {
178   // Integer sets are uniqued based on their dim/symbol counts, affine
179   // expressions appearing in the LHS of constraints, and eqFlags.
180   using KeyTy =
181       std::tuple<unsigned, unsigned, ArrayRef<AffineExpr>, ArrayRef<bool>>;
182   using DenseMapInfo<IntegerSet>::isEqual;
183 
184   static unsigned getHashValue(const IntegerSet &key) {
185     return getHashValue(KeyTy(key.getNumDims(), key.getNumSymbols(),
186                               key.getConstraints(), key.getEqFlags()));
187   }
188 
189   static unsigned getHashValue(KeyTy key) {
190     return hash_combine(
191         std::get<0>(key), std::get<1>(key),
192         hash_combine_range(std::get<2>(key).begin(), std::get<2>(key).end()),
193         hash_combine_range(std::get<3>(key).begin(), std::get<3>(key).end()));
194   }
195 
196   static bool isEqual(const KeyTy &lhs, IntegerSet rhs) {
197     if (rhs == getEmptyKey() || rhs == getTombstoneKey())
198       return false;
199     return lhs == std::make_tuple(rhs.getNumDims(), rhs.getNumSymbols(),
200                                   rhs.getConstraints(), rhs.getEqFlags());
201   }
202 };
203 } // end anonymous namespace.
204 
205 //===----------------------------------------------------------------------===//
206 // MLIRContextImpl
207 //===----------------------------------------------------------------------===//
208 
209 namespace mlir {
210 /// This is the implementation of the MLIRContext class, using the pImpl idiom.
211 /// This class is completely private to this file, so everything is public.
212 class MLIRContextImpl {
213 public:
214   //===--------------------------------------------------------------------===//
215   // Identifier uniquing
216   //===--------------------------------------------------------------------===//
217 
218   // Identifier allocator and mutex for thread safety.
219   llvm::BumpPtrAllocator identifierAllocator;
220   llvm::sys::SmartRWMutex<true> identifierMutex;
221 
222   //===--------------------------------------------------------------------===//
223   // Diagnostics
224   //===--------------------------------------------------------------------===//
225   DiagnosticEngine diagEngine;
226 
227   //===--------------------------------------------------------------------===//
228   // Options
229   //===--------------------------------------------------------------------===//
230 
231   /// In most cases, creating operation in unregistered dialect is not desired
232   /// and indicate a misconfiguration of the compiler. This option enables to
233   /// detect such use cases
234   bool allowUnregisteredDialects = false;
235 
236   /// Enable support for multi-threading within MLIR.
237   bool threadingIsEnabled = true;
238 
239   /// Track if we are currently executing in a threaded execution environment
240   /// (like the pass-manager): this is only a debugging feature to help reducing
241   /// the chances of data races one some context APIs.
242 #ifndef NDEBUG
243   std::atomic<int> multiThreadedExecutionContext{0};
244 #endif
245 
246   /// If the operation should be attached to diagnostics printed via the
247   /// Operation::emit methods.
248   bool printOpOnDiagnostic = true;
249 
250   /// If the current stack trace should be attached when emitting diagnostics.
251   bool printStackTraceOnDiagnostic = false;
252 
253   //===--------------------------------------------------------------------===//
254   // Other
255   //===--------------------------------------------------------------------===//
256 
257   /// This is a list of dialects that are created referring to this context.
258   /// The MLIRContext owns the objects.
259   DenseMap<StringRef, std::unique_ptr<Dialect>> loadedDialects;
260   DialectRegistry dialectsRegistry;
261 
262   /// This is a mapping from operation name to AbstractOperation for registered
263   /// operations.
264   llvm::StringMap<AbstractOperation> registeredOperations;
265 
266   /// Identifiers are uniqued by string value and use the internal string set
267   /// for storage.
268   llvm::StringSet<llvm::BumpPtrAllocator &> identifiers;
269   /// A thread local cache of identifiers to reduce lock contention.
270   ThreadLocalCache<llvm::StringMap<llvm::StringMapEntry<llvm::NoneType> *>>
271       localIdentifierCache;
272 
273   /// An allocator used for AbstractAttribute and AbstractType objects.
274   llvm::BumpPtrAllocator abstractDialectSymbolAllocator;
275 
276   //===--------------------------------------------------------------------===//
277   // Affine uniquing
278   //===--------------------------------------------------------------------===//
279 
280   // Affine allocator and mutex for thread safety.
281   llvm::BumpPtrAllocator affineAllocator;
282   llvm::sys::SmartRWMutex<true> affineMutex;
283 
284   // Affine map uniquing.
285   using AffineMapSet = DenseSet<AffineMap, AffineMapKeyInfo>;
286   AffineMapSet affineMaps;
287 
288   // Integer set uniquing.
289   using IntegerSets = DenseSet<IntegerSet, IntegerSetKeyInfo>;
290   IntegerSets integerSets;
291 
292   // Affine expression uniquing.
293   StorageUniquer affineUniquer;
294 
295   //===--------------------------------------------------------------------===//
296   // Type uniquing
297   //===--------------------------------------------------------------------===//
298 
299   DenseMap<TypeID, const AbstractType *> registeredTypes;
300   StorageUniquer typeUniquer;
301 
302   /// Cached Type Instances.
303   BFloat16Type bf16Ty;
304   Float16Type f16Ty;
305   Float32Type f32Ty;
306   Float64Type f64Ty;
307   IndexType indexTy;
308   IntegerType int1Ty, int8Ty, int16Ty, int32Ty, int64Ty, int128Ty;
309   NoneType noneType;
310 
311   //===--------------------------------------------------------------------===//
312   // Attribute uniquing
313   //===--------------------------------------------------------------------===//
314 
315   DenseMap<TypeID, const AbstractAttribute *> registeredAttributes;
316   StorageUniquer attributeUniquer;
317 
318   /// Cached Attribute Instances.
319   BoolAttr falseAttr, trueAttr;
320   UnitAttr unitAttr;
321   UnknownLoc unknownLocAttr;
322   DictionaryAttr emptyDictionaryAttr;
323 
324 public:
325   MLIRContextImpl() : identifiers(identifierAllocator) {}
326   ~MLIRContextImpl() {
327     for (auto typeMapping : registeredTypes)
328       typeMapping.second->~AbstractType();
329     for (auto attrMapping : registeredAttributes)
330       attrMapping.second->~AbstractAttribute();
331   }
332 };
333 } // end namespace mlir
334 
335 MLIRContext::MLIRContext() : impl(new MLIRContextImpl()) {
336   // Initialize values based on the command line flags if they were provided.
337   if (clOptions.isConstructed()) {
338     disableMultithreading(clOptions->disableThreading);
339     printOpOnDiagnostic(clOptions->printOpOnDiagnostic);
340     printStackTraceOnDiagnostic(clOptions->printStackTraceOnDiagnostic);
341   }
342 
343   // Ensure the builtin dialect is always pre-loaded.
344   getOrLoadDialect<BuiltinDialect>();
345 
346   // Initialize several common attributes and types to avoid the need to lock
347   // the context when accessing them.
348 
349   //// Types.
350   /// Floating-point Types.
351   impl->bf16Ty = TypeUniquer::get<BFloat16Type>(this);
352   impl->f16Ty = TypeUniquer::get<Float16Type>(this);
353   impl->f32Ty = TypeUniquer::get<Float32Type>(this);
354   impl->f64Ty = TypeUniquer::get<Float64Type>(this);
355   /// Index Type.
356   impl->indexTy = TypeUniquer::get<IndexType>(this);
357   /// Integer Types.
358   impl->int1Ty = TypeUniquer::get<IntegerType>(this, 1, IntegerType::Signless);
359   impl->int8Ty = TypeUniquer::get<IntegerType>(this, 8, IntegerType::Signless);
360   impl->int16Ty =
361       TypeUniquer::get<IntegerType>(this, 16, IntegerType::Signless);
362   impl->int32Ty =
363       TypeUniquer::get<IntegerType>(this, 32, IntegerType::Signless);
364   impl->int64Ty =
365       TypeUniquer::get<IntegerType>(this, 64, IntegerType::Signless);
366   impl->int128Ty =
367       TypeUniquer::get<IntegerType>(this, 128, IntegerType::Signless);
368   /// None Type.
369   impl->noneType = TypeUniquer::get<NoneType>(this);
370 
371   //// Attributes.
372   //// Note: These must be registered after the types as they may generate one
373   //// of the above types internally.
374   /// Bool Attributes.
375   impl->falseAttr = AttributeUniquer::get<IntegerAttr>(
376                         this, impl->int1Ty, APInt(/*numBits=*/1, false))
377                         .cast<BoolAttr>();
378   impl->trueAttr = AttributeUniquer::get<IntegerAttr>(
379                        this, impl->int1Ty, APInt(/*numBits=*/1, true))
380                        .cast<BoolAttr>();
381   /// Unit Attribute.
382   impl->unitAttr = AttributeUniquer::get<UnitAttr>(this);
383   /// Unknown Location Attribute.
384   impl->unknownLocAttr = AttributeUniquer::get<UnknownLoc>(this);
385   /// The empty dictionary attribute.
386   impl->emptyDictionaryAttr =
387       AttributeUniquer::get<DictionaryAttr>(this, ArrayRef<NamedAttribute>());
388 
389   // Register the affine storage objects with the uniquer.
390   impl->affineUniquer
391       .registerParametricStorageType<AffineBinaryOpExprStorage>();
392   impl->affineUniquer
393       .registerParametricStorageType<AffineConstantExprStorage>();
394   impl->affineUniquer.registerParametricStorageType<AffineDimExprStorage>();
395 }
396 
397 MLIRContext::~MLIRContext() {}
398 
399 /// Copy the specified array of elements into memory managed by the provided
400 /// bump pointer allocator.  This assumes the elements are all PODs.
401 template <typename T>
402 static ArrayRef<T> copyArrayRefInto(llvm::BumpPtrAllocator &allocator,
403                                     ArrayRef<T> elements) {
404   auto result = allocator.Allocate<T>(elements.size());
405   std::uninitialized_copy(elements.begin(), elements.end(), result);
406   return ArrayRef<T>(result, elements.size());
407 }
408 
409 //===----------------------------------------------------------------------===//
410 // Diagnostic Handlers
411 //===----------------------------------------------------------------------===//
412 
413 /// Returns the diagnostic engine for this context.
414 DiagnosticEngine &MLIRContext::getDiagEngine() { return getImpl().diagEngine; }
415 
416 //===----------------------------------------------------------------------===//
417 // Dialect and Operation Registration
418 //===----------------------------------------------------------------------===//
419 
420 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(impl->loadedDialects.size());
428   for (auto &dialect : impl->loadedDialects)
429     result.push_back(dialect.second.get());
430   llvm::array_pod_sort(result.begin(), result.end(),
431                        [](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)
439     result.push_back(dialect.first);
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(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   return impl->dialectsRegistry.loadByName(name, this);
456 }
457 
458 /// Get a dialect for the provided namespace and TypeID: abort the program if a
459 /// dialect exist for this namespace with different TypeID. Returns a pointer to
460 /// the dialect owned by the context.
461 Dialect *
462 MLIRContext::getOrLoadDialect(StringRef dialectNamespace, TypeID dialectID,
463                               function_ref<std::unique_ptr<Dialect>()> ctor) {
464   auto &impl = getImpl();
465   // Get the correct insertion position sorted by namespace.
466   std::unique_ptr<Dialect> &dialect = impl.loadedDialects[dialectNamespace];
467 
468   if (!dialect) {
469     LLVM_DEBUG(llvm::dbgs()
470                << "Load new dialect in Context " << dialectNamespace << "\n");
471 #ifndef NDEBUG
472     if (impl.multiThreadedExecutionContext != 0)
473       llvm::report_fatal_error(
474           "Loading a dialect (" + dialectNamespace +
475           ") while in a multi-threaded execution context (maybe "
476           "the PassManager): this can indicate a "
477           "missing `dependentDialects` in a pass for example.");
478 #endif
479     dialect = ctor();
480     assert(dialect && "dialect ctor failed");
481     return dialect.get();
482   }
483 
484   // Abort if dialect with namespace has already been registered.
485   if (dialect->getTypeID() != dialectID)
486     llvm::report_fatal_error("a dialect with namespace '" + dialectNamespace +
487                              "' has already been registered");
488 
489   return dialect.get();
490 }
491 
492 bool MLIRContext::allowsUnregisteredDialects() {
493   return impl->allowUnregisteredDialects;
494 }
495 
496 void MLIRContext::allowUnregisteredDialects(bool allowing) {
497   impl->allowUnregisteredDialects = allowing;
498 }
499 
500 /// Return true if multi-threading is disabled by the context.
501 bool MLIRContext::isMultithreadingEnabled() {
502   return impl->threadingIsEnabled && llvm::llvm_is_multithreaded();
503 }
504 
505 /// Set the flag specifying if multi-threading is disabled by the context.
506 void MLIRContext::disableMultithreading(bool disable) {
507   impl->threadingIsEnabled = !disable;
508 
509   // Update the threading mode for each of the uniquers.
510   impl->affineUniquer.disableMultithreading(disable);
511   impl->attributeUniquer.disableMultithreading(disable);
512   impl->typeUniquer.disableMultithreading(disable);
513 }
514 
515 void MLIRContext::enterMultiThreadedExecution() {
516 #ifndef NDEBUG
517   ++impl->multiThreadedExecutionContext;
518 #endif
519 }
520 void MLIRContext::exitMultiThreadedExecution() {
521 #ifndef NDEBUG
522   --impl->multiThreadedExecutionContext;
523 #endif
524 }
525 
526 /// Return true if we should attach the operation to diagnostics emitted via
527 /// Operation::emit.
528 bool MLIRContext::shouldPrintOpOnDiagnostic() {
529   return impl->printOpOnDiagnostic;
530 }
531 
532 /// Set the flag specifying if we should attach the operation to diagnostics
533 /// emitted via Operation::emit.
534 void MLIRContext::printOpOnDiagnostic(bool enable) {
535   impl->printOpOnDiagnostic = enable;
536 }
537 
538 /// Return true if we should attach the current stacktrace to diagnostics when
539 /// emitted.
540 bool MLIRContext::shouldPrintStackTraceOnDiagnostic() {
541   return impl->printStackTraceOnDiagnostic;
542 }
543 
544 /// Set the flag specifying if we should attach the current stacktrace when
545 /// emitting diagnostics.
546 void MLIRContext::printStackTraceOnDiagnostic(bool enable) {
547   impl->printStackTraceOnDiagnostic = enable;
548 }
549 
550 /// Return information about all registered operations.  This isn't very
551 /// efficient, typically you should ask the operations about their properties
552 /// directly.
553 std::vector<AbstractOperation *> MLIRContext::getRegisteredOperations() {
554   // We just have the operations in a non-deterministic hash table order. Dump
555   // into a temporary array, then sort it by operation name to get a stable
556   // ordering.
557   llvm::StringMap<AbstractOperation> &registeredOps =
558       impl->registeredOperations;
559 
560   std::vector<AbstractOperation *> result;
561   result.reserve(registeredOps.size());
562   for (auto &elt : registeredOps)
563     result.push_back(&elt.second);
564   llvm::array_pod_sort(
565       result.begin(), result.end(),
566       [](AbstractOperation *const *lhs, AbstractOperation *const *rhs) {
567         return (*lhs)->name.compare((*rhs)->name);
568       });
569 
570   return result;
571 }
572 
573 bool MLIRContext::isOperationRegistered(StringRef name) {
574   return impl->registeredOperations.count(name);
575 }
576 
577 void Dialect::addType(TypeID typeID, AbstractType &&typeInfo) {
578   auto &impl = context->getImpl();
579   assert(impl.multiThreadedExecutionContext == 0 &&
580          "Registering a new type kind while in a multi-threaded execution "
581          "context");
582   auto *newInfo =
583       new (impl.abstractDialectSymbolAllocator.Allocate<AbstractType>())
584           AbstractType(std::move(typeInfo));
585   if (!impl.registeredTypes.insert({typeID, newInfo}).second)
586     llvm::report_fatal_error("Dialect Type already registered.");
587 }
588 
589 void Dialect::addAttribute(TypeID typeID, AbstractAttribute &&attrInfo) {
590   auto &impl = context->getImpl();
591   assert(impl.multiThreadedExecutionContext == 0 &&
592          "Registering a new attribute kind while in a multi-threaded execution "
593          "context");
594   auto *newInfo =
595       new (impl.abstractDialectSymbolAllocator.Allocate<AbstractAttribute>())
596           AbstractAttribute(std::move(attrInfo));
597   if (!impl.registeredAttributes.insert({typeID, newInfo}).second)
598     llvm::report_fatal_error("Dialect Attribute already registered.");
599 }
600 
601 //===----------------------------------------------------------------------===//
602 // AbstractAttribute
603 //===----------------------------------------------------------------------===//
604 
605 /// Get the dialect that registered the attribute with the provided typeid.
606 const AbstractAttribute &AbstractAttribute::lookup(TypeID typeID,
607                                                    MLIRContext *context) {
608   auto &impl = context->getImpl();
609   auto it = impl.registeredAttributes.find(typeID);
610   if (it == impl.registeredAttributes.end())
611     llvm::report_fatal_error("Trying to create an Attribute that was not "
612                              "registered in this MLIRContext.");
613   return *it->second;
614 }
615 
616 //===----------------------------------------------------------------------===//
617 // AbstractOperation
618 //===----------------------------------------------------------------------===//
619 
620 ParseResult AbstractOperation::parseAssembly(OpAsmParser &parser,
621                                              OperationState &result) const {
622   return parseAssemblyFn(parser, result);
623 }
624 
625 /// Look up the specified operation in the operation set and return a pointer
626 /// to it if present. Otherwise, return a null pointer.
627 const AbstractOperation *AbstractOperation::lookup(StringRef opName,
628                                                    MLIRContext *context) {
629   auto &impl = context->getImpl();
630   auto it = impl.registeredOperations.find(opName);
631   if (it != impl.registeredOperations.end())
632     return &it->second;
633   return nullptr;
634 }
635 
636 void AbstractOperation::insert(
637     StringRef name, Dialect &dialect, OperationProperties opProperties,
638     TypeID typeID, ParseAssemblyFn parseAssembly, PrintAssemblyFn printAssembly,
639     VerifyInvariantsFn verifyInvariants, FoldHookFn foldHook,
640     GetCanonicalizationPatternsFn getCanonicalizationPatterns,
641     detail::InterfaceMap &&interfaceMap, HasTraitFn hasTrait) {
642   AbstractOperation opInfo(name, dialect, opProperties, typeID, parseAssembly,
643                            printAssembly, verifyInvariants, foldHook,
644                            getCanonicalizationPatterns, std::move(interfaceMap),
645                            hasTrait);
646 
647   auto &impl = dialect.getContext()->getImpl();
648   assert(impl.multiThreadedExecutionContext == 0 &&
649          "Registering a new operation kind while in a multi-threaded execution "
650          "context");
651   if (!impl.registeredOperations.insert({name, std::move(opInfo)}).second) {
652     llvm::errs() << "error: operation named '" << name
653                  << "' is already registered.\n";
654     abort();
655   }
656 }
657 
658 AbstractOperation::AbstractOperation(
659     StringRef name, Dialect &dialect, OperationProperties opProperties,
660     TypeID typeID, ParseAssemblyFn parseAssembly, PrintAssemblyFn printAssembly,
661     VerifyInvariantsFn verifyInvariants, FoldHookFn foldHook,
662     GetCanonicalizationPatternsFn getCanonicalizationPatterns,
663     detail::InterfaceMap &&interfaceMap, HasTraitFn hasTrait)
664     : name(Identifier::get(name, dialect.getContext())), dialect(dialect),
665       typeID(typeID), opProperties(opProperties),
666       interfaceMap(std::move(interfaceMap)), foldHookFn(foldHook),
667       getCanonicalizationPatternsFn(getCanonicalizationPatterns),
668       hasTraitFn(hasTrait), parseAssemblyFn(parseAssembly),
669       printAssemblyFn(printAssembly), verifyInvariantsFn(verifyInvariants) {}
670 
671 //===----------------------------------------------------------------------===//
672 // AbstractType
673 //===----------------------------------------------------------------------===//
674 
675 const AbstractType &AbstractType::lookup(TypeID typeID, MLIRContext *context) {
676   auto &impl = context->getImpl();
677   auto it = impl.registeredTypes.find(typeID);
678   if (it == impl.registeredTypes.end())
679     llvm::report_fatal_error(
680         "Trying to create a Type that was not registered in this MLIRContext.");
681   return *it->second;
682 }
683 
684 //===----------------------------------------------------------------------===//
685 // Identifier uniquing
686 //===----------------------------------------------------------------------===//
687 
688 /// Return an identifier for the specified string.
689 Identifier Identifier::get(StringRef str, MLIRContext *context) {
690   // Check invariants after seeing if we already have something in the
691   // identifier table - if we already had it in the table, then it already
692   // passed invariant checks.
693   assert(!str.empty() && "Cannot create an empty identifier");
694   assert(str.find('\0') == StringRef::npos &&
695          "Cannot create an identifier with a nul character");
696 
697   auto &impl = context->getImpl();
698   if (!context->isMultithreadingEnabled())
699     return Identifier(&*impl.identifiers.insert(str).first);
700 
701   // Check for an existing instance in the local cache.
702   auto *&localEntry = (*impl.localIdentifierCache)[str];
703   if (localEntry)
704     return Identifier(localEntry);
705 
706   // Check for an existing identifier in read-only mode.
707   {
708     llvm::sys::SmartScopedReader<true> contextLock(impl.identifierMutex);
709     auto it = impl.identifiers.find(str);
710     if (it != impl.identifiers.end()) {
711       localEntry = &*it;
712       return Identifier(localEntry);
713     }
714   }
715 
716   // Acquire a writer-lock so that we can safely create the new instance.
717   llvm::sys::SmartScopedWriter<true> contextLock(impl.identifierMutex);
718   auto it = impl.identifiers.insert(str).first;
719   localEntry = &*it;
720   return Identifier(localEntry);
721 }
722 
723 //===----------------------------------------------------------------------===//
724 // Type uniquing
725 //===----------------------------------------------------------------------===//
726 
727 /// Returns the storage uniquer used for constructing type storage instances.
728 /// This should not be used directly.
729 StorageUniquer &MLIRContext::getTypeUniquer() { return getImpl().typeUniquer; }
730 
731 BFloat16Type BFloat16Type::get(MLIRContext *context) {
732   return context->getImpl().bf16Ty;
733 }
734 Float16Type Float16Type::get(MLIRContext *context) {
735   return context->getImpl().f16Ty;
736 }
737 Float32Type Float32Type::get(MLIRContext *context) {
738   return context->getImpl().f32Ty;
739 }
740 Float64Type Float64Type::get(MLIRContext *context) {
741   return context->getImpl().f64Ty;
742 }
743 
744 /// Get an instance of the IndexType.
745 IndexType IndexType::get(MLIRContext *context) {
746   return context->getImpl().indexTy;
747 }
748 
749 /// Return an existing integer type instance if one is cached within the
750 /// context.
751 static IntegerType
752 getCachedIntegerType(unsigned width,
753                      IntegerType::SignednessSemantics signedness,
754                      MLIRContext *context) {
755   if (signedness != IntegerType::Signless)
756     return IntegerType();
757 
758   switch (width) {
759   case 1:
760     return context->getImpl().int1Ty;
761   case 8:
762     return context->getImpl().int8Ty;
763   case 16:
764     return context->getImpl().int16Ty;
765   case 32:
766     return context->getImpl().int32Ty;
767   case 64:
768     return context->getImpl().int64Ty;
769   case 128:
770     return context->getImpl().int128Ty;
771   default:
772     return IntegerType();
773   }
774 }
775 
776 IntegerType IntegerType::get(unsigned width, MLIRContext *context) {
777   return get(width, IntegerType::Signless, context);
778 }
779 
780 IntegerType IntegerType::get(unsigned width,
781                              IntegerType::SignednessSemantics signedness,
782                              MLIRContext *context) {
783   if (auto cached = getCachedIntegerType(width, signedness, context))
784     return cached;
785   return Base::get(context, width, signedness);
786 }
787 
788 IntegerType IntegerType::getChecked(unsigned width, Location location) {
789   return getChecked(width, IntegerType::Signless, location);
790 }
791 
792 IntegerType IntegerType::getChecked(unsigned width,
793                                     SignednessSemantics signedness,
794                                     Location location) {
795   if (auto cached =
796           getCachedIntegerType(width, signedness, location->getContext()))
797     return cached;
798   return Base::getChecked(location, width, signedness);
799 }
800 
801 /// Get an instance of the NoneType.
802 NoneType NoneType::get(MLIRContext *context) {
803   if (NoneType cachedInst = context->getImpl().noneType)
804     return cachedInst;
805   // Note: May happen when initializing the singleton attributes of the builtin
806   // dialect.
807   return Base::get(context);
808 }
809 
810 //===----------------------------------------------------------------------===//
811 // Attribute uniquing
812 //===----------------------------------------------------------------------===//
813 
814 /// Returns the storage uniquer used for constructing attribute storage
815 /// instances. This should not be used directly.
816 StorageUniquer &MLIRContext::getAttributeUniquer() {
817   return getImpl().attributeUniquer;
818 }
819 
820 /// Initialize the given attribute storage instance.
821 void AttributeUniquer::initializeAttributeStorage(AttributeStorage *storage,
822                                                   MLIRContext *ctx,
823                                                   TypeID attrID) {
824   storage->initialize(AbstractAttribute::lookup(attrID, ctx));
825 
826   // If the attribute did not provide a type, then default to NoneType.
827   if (!storage->getType())
828     storage->setType(NoneType::get(ctx));
829 }
830 
831 BoolAttr BoolAttr::get(bool value, MLIRContext *context) {
832   return value ? context->getImpl().trueAttr : context->getImpl().falseAttr;
833 }
834 
835 UnitAttr UnitAttr::get(MLIRContext *context) {
836   return context->getImpl().unitAttr;
837 }
838 
839 Location UnknownLoc::get(MLIRContext *context) {
840   return context->getImpl().unknownLocAttr;
841 }
842 
843 /// Return empty dictionary.
844 DictionaryAttr DictionaryAttr::getEmpty(MLIRContext *context) {
845   return context->getImpl().emptyDictionaryAttr;
846 }
847 
848 //===----------------------------------------------------------------------===//
849 // AffineMap uniquing
850 //===----------------------------------------------------------------------===//
851 
852 StorageUniquer &MLIRContext::getAffineUniquer() {
853   return getImpl().affineUniquer;
854 }
855 
856 AffineMap AffineMap::getImpl(unsigned dimCount, unsigned symbolCount,
857                              ArrayRef<AffineExpr> results,
858                              MLIRContext *context) {
859   auto &impl = context->getImpl();
860   auto key = std::make_tuple(dimCount, symbolCount, results);
861 
862   // Safely get or create an AffineMap instance.
863   return safeGetOrCreate(
864       impl.affineMaps, key, impl.affineMutex, impl.threadingIsEnabled, [&] {
865         auto *res = impl.affineAllocator.Allocate<detail::AffineMapStorage>();
866 
867         // Copy the results into the bump pointer.
868         results = copyArrayRefInto(impl.affineAllocator, results);
869 
870         // Initialize the memory using placement new.
871         new (res)
872             detail::AffineMapStorage{dimCount, symbolCount, results, context};
873         return AffineMap(res);
874       });
875 }
876 
877 AffineMap AffineMap::get(MLIRContext *context) {
878   return getImpl(/*dimCount=*/0, /*symbolCount=*/0, /*results=*/{}, context);
879 }
880 
881 AffineMap AffineMap::get(unsigned dimCount, unsigned symbolCount,
882                          MLIRContext *context) {
883   return getImpl(dimCount, symbolCount, /*results=*/{}, context);
884 }
885 
886 AffineMap AffineMap::get(unsigned dimCount, unsigned symbolCount,
887                          AffineExpr result) {
888   return getImpl(dimCount, symbolCount, {result}, result.getContext());
889 }
890 
891 AffineMap AffineMap::get(unsigned dimCount, unsigned symbolCount,
892                          ArrayRef<AffineExpr> results, MLIRContext *context) {
893   return getImpl(dimCount, symbolCount, results, context);
894 }
895 
896 //===----------------------------------------------------------------------===//
897 // Integer Sets: these are allocated into the bump pointer, and are immutable.
898 // Unlike AffineMap's, these are uniqued only if they are small.
899 //===----------------------------------------------------------------------===//
900 
901 IntegerSet IntegerSet::get(unsigned dimCount, unsigned symbolCount,
902                            ArrayRef<AffineExpr> constraints,
903                            ArrayRef<bool> eqFlags) {
904   // The number of constraints can't be zero.
905   assert(!constraints.empty());
906   assert(constraints.size() == eqFlags.size());
907 
908   auto &impl = constraints[0].getContext()->getImpl();
909 
910   // A utility function to construct a new IntegerSetStorage instance.
911   auto constructorFn = [&] {
912     auto *res = impl.affineAllocator.Allocate<detail::IntegerSetStorage>();
913 
914     // Copy the results and equality flags into the bump pointer.
915     constraints = copyArrayRefInto(impl.affineAllocator, constraints);
916     eqFlags = copyArrayRefInto(impl.affineAllocator, eqFlags);
917 
918     // Initialize the memory using placement new.
919     new (res)
920         detail::IntegerSetStorage{dimCount, symbolCount, constraints, eqFlags};
921     return IntegerSet(res);
922   };
923 
924   // If this instance is uniqued, then we handle it separately so that multiple
925   // threads may simultaneously access existing instances.
926   if (constraints.size() < IntegerSet::kUniquingThreshold) {
927     auto key = std::make_tuple(dimCount, symbolCount, constraints, eqFlags);
928     return safeGetOrCreate(impl.integerSets, key, impl.affineMutex,
929                            impl.threadingIsEnabled, constructorFn);
930   }
931 
932   // Otherwise, acquire a writer-lock so that we can safely create the new
933   // instance.
934   ScopedWriterLock affineLock(impl.affineMutex, impl.threadingIsEnabled);
935   return constructorFn();
936 }
937 
938 //===----------------------------------------------------------------------===//
939 // StorageUniquerSupport
940 //===----------------------------------------------------------------------===//
941 
942 /// Utility method to generate a default location for use when checking the
943 /// construction invariants of a storage object. This is defined out-of-line to
944 /// avoid the need to include Location.h.
945 const AttributeStorage *
946 mlir::detail::generateUnknownStorageLocation(MLIRContext *ctx) {
947   return reinterpret_cast<const AttributeStorage *>(
948       ctx->getImpl().unknownLocAttr.getAsOpaquePointer());
949 }
950