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