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