/home/arjun/llvm-project/mlir/lib/IR/MLIRContext.cpp
Line | Count | Source (jump to first uncovered line) |
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. |
75 | 0 | void mlir::registerMLIRContextCLOptions() { |
76 | 0 | // Make sure that the options struct has been initialized. |
77 | 0 | *clOptions; |
78 | 0 | } |
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 { |
88 | 0 | BuiltinDialect(MLIRContext *context) : Dialect(/*name=*/"", context) { |
89 | 0 | addAttributes<AffineMapAttr, ArrayAttr, BoolAttr, DenseIntOrFPElementsAttr, |
90 | 0 | DenseStringElementsAttr, DictionaryAttr, FloatAttr, |
91 | 0 | SymbolRefAttr, IntegerAttr, IntegerSetAttr, OpaqueAttr, |
92 | 0 | OpaqueElementsAttr, SparseElementsAttr, StringAttr, TypeAttr, |
93 | 0 | UnitAttr>(); |
94 | 0 | addAttributes<CallSiteLoc, FileLineColLoc, FusedLoc, NameLoc, OpaqueLoc, |
95 | 0 | UnknownLoc>(); |
96 | 0 |
|
97 | 0 | addTypes<ComplexType, FloatType, FunctionType, IndexType, IntegerType, |
98 | 0 | MemRefType, UnrankedMemRefType, NoneType, OpaqueType, |
99 | 0 | RankedTensorType, TupleType, UnrankedTensorType, VectorType>(); |
100 | 0 |
|
101 | 0 | // TODO: These operations should be moved to a different dialect when they |
102 | 0 | // have been fully decoupled from the core. |
103 | 0 | addOperations<FuncOp, ModuleOp, ModuleTerminatorOp>(); |
104 | 0 | } |
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 { |
116 | | ScopedReaderLock(llvm::sys::SmartRWMutex<true> &mutexParam, bool shouldLock) |
117 | 0 | : mutex(shouldLock ? &mutexParam : nullptr) { |
118 | 0 | if (mutex) |
119 | 0 | mutex->lock_shared(); |
120 | 0 | } |
121 | 0 | ~ScopedReaderLock() { |
122 | 0 | if (mutex) |
123 | 0 | mutex->unlock_shared(); |
124 | 0 | } |
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 { |
130 | | ScopedWriterLock(llvm::sys::SmartRWMutex<true> &mutexParam, bool shouldLock) |
131 | 0 | : mutex(shouldLock ? &mutexParam : nullptr) { |
132 | 0 | if (mutex) |
133 | 0 | mutex->lock(); |
134 | 0 | } |
135 | 0 | ~ScopedWriterLock() { |
136 | 0 | if (mutex) |
137 | 0 | mutex->unlock(); |
138 | 0 | } |
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> |
151 | | static ValueT safeGetOrCreate(DenseSet<ValueT, DenseInfoT> &container, |
152 | | KeyT &&key, llvm::sys::SmartRWMutex<true> &mutex, |
153 | | bool threadingIsEnabled, |
154 | 0 | ConstructorFn &&constructorFn) { |
155 | 0 | // Check for an existing instance in read-only mode. |
156 | 0 | if (threadingIsEnabled) { |
157 | 0 | llvm::sys::SmartScopedReader<true> instanceLock(mutex); |
158 | 0 | auto it = container.find_as(key); |
159 | 0 | if (it != container.end()) |
160 | 0 | return *it; |
161 | 0 | } |
162 | 0 | |
163 | 0 | // Acquire a writer-lock so that we can safely create the new instance. |
164 | 0 | ScopedWriterLock instanceLock(mutex, threadingIsEnabled); |
165 | 0 |
|
166 | 0 | // Check for an existing instance again here, because another writer thread |
167 | 0 | // may have already created one. Otherwise, construct a new instance. |
168 | 0 | auto existing = container.insert_as(ValueT(), key); |
169 | 0 | if (existing.second) |
170 | 0 | return *existing.first = constructorFn(); |
171 | 0 | return *existing.first; |
172 | 0 | } Unexecuted instantiation: MLIRContext.cpp:_ZL15safeGetOrCreateIN4mlir9AffineMapEN12_GLOBAL__N_116AffineMapKeyInfoERSt5tupleIJjjN4llvm8ArrayRefINS0_10AffineExprEEEEEZNS1_7getImplEjjS8_PNS0_11MLIRContextEE3$_2ET_RNS5_8DenseSetISE_T0_EEOT1_RNS5_3sys12SmartRWMutexILb1EEEbOT2_ Unexecuted instantiation: MLIRContext.cpp:_ZL15safeGetOrCreateIN4mlir10IntegerSetEN12_GLOBAL__N_117IntegerSetKeyInfoERSt5tupleIJjjN4llvm8ArrayRefINS0_10AffineExprEEENS6_IbEEEERZNS1_3getEjjS8_S9_E3$_3ET_RNS5_8DenseSetISE_T0_EEOT1_RNS5_3sys12SmartRWMutexILb1EEEbOT2_ |
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 | | |
181 | 0 | static unsigned getHashValue(const AffineMap &key) { |
182 | 0 | return getHashValue( |
183 | 0 | KeyTy(key.getNumDims(), key.getNumSymbols(), key.getResults())); |
184 | 0 | } |
185 | | |
186 | 0 | static unsigned getHashValue(KeyTy key) { |
187 | 0 | return hash_combine( |
188 | 0 | std::get<0>(key), std::get<1>(key), |
189 | 0 | hash_combine_range(std::get<2>(key).begin(), std::get<2>(key).end())); |
190 | 0 | } |
191 | | |
192 | 0 | static bool isEqual(const KeyTy &lhs, AffineMap rhs) { |
193 | 0 | if (rhs == getEmptyKey() || rhs == getTombstoneKey()) |
194 | 0 | return false; |
195 | 0 | return lhs == std::make_tuple(rhs.getNumDims(), rhs.getNumSymbols(), |
196 | 0 | rhs.getResults()); |
197 | 0 | } |
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 | | |
207 | 0 | static unsigned getHashValue(const IntegerSet &key) { |
208 | 0 | return getHashValue(KeyTy(key.getNumDims(), key.getNumSymbols(), |
209 | 0 | key.getConstraints(), key.getEqFlags())); |
210 | 0 | } |
211 | | |
212 | 0 | static unsigned getHashValue(KeyTy key) { |
213 | 0 | return hash_combine( |
214 | 0 | std::get<0>(key), std::get<1>(key), |
215 | 0 | hash_combine_range(std::get<2>(key).begin(), std::get<2>(key).end()), |
216 | 0 | hash_combine_range(std::get<3>(key).begin(), std::get<3>(key).end())); |
217 | 0 | } |
218 | | |
219 | 0 | static bool isEqual(const KeyTy &lhs, IntegerSet rhs) { |
220 | 0 | if (rhs == getEmptyKey() || rhs == getTombstoneKey()) |
221 | 0 | return false; |
222 | 0 | return lhs == std::make_tuple(rhs.getNumDims(), rhs.getNumSymbols(), |
223 | 0 | rhs.getConstraints(), rhs.getEqFlags()); |
224 | 0 | } |
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 | | /// A general purpose mutex to lock access to parts of the context that do not |
274 | | /// have a more specific mutex, e.g. registry operations. |
275 | | llvm::sys::SmartRWMutex<true> contextMutex; |
276 | | |
277 | | /// This is a list of dialects that are created referring to this context. |
278 | | /// The MLIRContext owns the objects. |
279 | | std::vector<std::unique_ptr<Dialect>> dialects; |
280 | | |
281 | | /// This is a mapping from operation name to AbstractOperation for registered |
282 | | /// operations. |
283 | | llvm::StringMap<AbstractOperation> registeredOperations; |
284 | | |
285 | | /// This is a mapping from type id to Dialect for registered attributes and |
286 | | /// types. |
287 | | DenseMap<TypeID, Dialect *> registeredDialectSymbols; |
288 | | |
289 | | /// These are identifiers uniqued into this MLIRContext. |
290 | | llvm::StringSet<llvm::BumpPtrAllocator &> identifiers; |
291 | | |
292 | | //===--------------------------------------------------------------------===// |
293 | | // Affine uniquing |
294 | | //===--------------------------------------------------------------------===// |
295 | | |
296 | | // Affine allocator and mutex for thread safety. |
297 | | llvm::BumpPtrAllocator affineAllocator; |
298 | | llvm::sys::SmartRWMutex<true> affineMutex; |
299 | | |
300 | | // Affine map uniquing. |
301 | | using AffineMapSet = DenseSet<AffineMap, AffineMapKeyInfo>; |
302 | | AffineMapSet affineMaps; |
303 | | |
304 | | // Integer set uniquing. |
305 | | using IntegerSets = DenseSet<IntegerSet, IntegerSetKeyInfo>; |
306 | | IntegerSets integerSets; |
307 | | |
308 | | // Affine expression uniquing. |
309 | | StorageUniquer affineUniquer; |
310 | | |
311 | | //===--------------------------------------------------------------------===// |
312 | | // Type uniquing |
313 | | //===--------------------------------------------------------------------===// |
314 | | StorageUniquer typeUniquer; |
315 | | |
316 | | /// Cached Type Instances. |
317 | | FloatType bf16Ty, f16Ty, f32Ty, f64Ty; |
318 | | IndexType indexTy; |
319 | | IntegerType int1Ty, int8Ty, int16Ty, int32Ty, int64Ty, int128Ty; |
320 | | NoneType noneType; |
321 | | |
322 | | //===--------------------------------------------------------------------===// |
323 | | // Attribute uniquing |
324 | | //===--------------------------------------------------------------------===// |
325 | | StorageUniquer attributeUniquer; |
326 | | |
327 | | /// Cached Attribute Instances. |
328 | | BoolAttr falseAttr, trueAttr; |
329 | | UnitAttr unitAttr; |
330 | | UnknownLoc unknownLocAttr; |
331 | | DictionaryAttr emptyDictionaryAttr; |
332 | | |
333 | | public: |
334 | 0 | MLIRContextImpl() : identifiers(identifierAllocator) {} |
335 | | }; |
336 | | } // end namespace mlir |
337 | | |
338 | 0 | MLIRContext::MLIRContext() : impl(new MLIRContextImpl()) { |
339 | 0 | // Initialize values based on the command line flags if they were provided. |
340 | 0 | if (clOptions.isConstructed()) { |
341 | 0 | disableMultithreading(clOptions->disableThreading); |
342 | 0 | printOpOnDiagnostic(clOptions->printOpOnDiagnostic); |
343 | 0 | printStackTraceOnDiagnostic(clOptions->printStackTraceOnDiagnostic); |
344 | 0 | } |
345 | 0 |
|
346 | 0 | // Register dialects with this context. |
347 | 0 | new BuiltinDialect(this); |
348 | 0 | registerAllDialects(this); |
349 | 0 |
|
350 | 0 | // Initialize several common attributes and types to avoid the need to lock |
351 | 0 | // the context when accessing them. |
352 | 0 |
|
353 | 0 | //// Types. |
354 | 0 | /// Floating-point Types. |
355 | 0 | impl->bf16Ty = TypeUniquer::get<FloatType>(this, StandardTypes::BF16); |
356 | 0 | impl->f16Ty = TypeUniquer::get<FloatType>(this, StandardTypes::F16); |
357 | 0 | impl->f32Ty = TypeUniquer::get<FloatType>(this, StandardTypes::F32); |
358 | 0 | impl->f64Ty = TypeUniquer::get<FloatType>(this, StandardTypes::F64); |
359 | 0 | /// Index Type. |
360 | 0 | impl->indexTy = TypeUniquer::get<IndexType>(this, StandardTypes::Index); |
361 | 0 | /// Integer Types. |
362 | 0 | impl->int1Ty = TypeUniquer::get<IntegerType>(this, StandardTypes::Integer, 1, |
363 | 0 | IntegerType::Signless); |
364 | 0 | impl->int8Ty = TypeUniquer::get<IntegerType>(this, StandardTypes::Integer, 8, |
365 | 0 | IntegerType::Signless); |
366 | 0 | impl->int16Ty = TypeUniquer::get<IntegerType>(this, StandardTypes::Integer, |
367 | 0 | 16, IntegerType::Signless); |
368 | 0 | impl->int32Ty = TypeUniquer::get<IntegerType>(this, StandardTypes::Integer, |
369 | 0 | 32, IntegerType::Signless); |
370 | 0 | impl->int64Ty = TypeUniquer::get<IntegerType>(this, StandardTypes::Integer, |
371 | 0 | 64, IntegerType::Signless); |
372 | 0 | impl->int128Ty = TypeUniquer::get<IntegerType>(this, StandardTypes::Integer, |
373 | 0 | 128, IntegerType::Signless); |
374 | 0 | /// None Type. |
375 | 0 | impl->noneType = TypeUniquer::get<NoneType>(this, StandardTypes::None); |
376 | 0 |
|
377 | 0 | //// Attributes. |
378 | 0 | //// Note: These must be registered after the types as they may generate one |
379 | 0 | //// of the above types internally. |
380 | 0 | /// Bool Attributes. |
381 | 0 | // Note: The context is also used within the BoolAttrStorage. |
382 | 0 | impl->falseAttr = AttributeUniquer::get<BoolAttr>( |
383 | 0 | this, StandardAttributes::Bool, this, false); |
384 | 0 | impl->trueAttr = AttributeUniquer::get<BoolAttr>( |
385 | 0 | this, StandardAttributes::Bool, this, true); |
386 | 0 | /// Unit Attribute. |
387 | 0 | impl->unitAttr = |
388 | 0 | AttributeUniquer::get<UnitAttr>(this, StandardAttributes::Unit); |
389 | 0 | /// Unknown Location Attribute. |
390 | 0 | impl->unknownLocAttr = AttributeUniquer::get<UnknownLoc>( |
391 | 0 | this, StandardAttributes::UnknownLocation); |
392 | 0 | /// The empty dictionary attribute. |
393 | 0 | impl->emptyDictionaryAttr = AttributeUniquer::get<DictionaryAttr>( |
394 | 0 | this, StandardAttributes::Dictionary, ArrayRef<NamedAttribute>()); |
395 | 0 | } |
396 | | |
397 | 0 | 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 | 0 | ArrayRef<T> elements) { |
404 | 0 | auto result = allocator.Allocate<T>(elements.size()); |
405 | 0 | std::uninitialized_copy(elements.begin(), elements.end(), result); |
406 | 0 | return ArrayRef<T>(result, elements.size()); |
407 | 0 | } Unexecuted instantiation: MLIRContext.cpp:_ZL16copyArrayRefIntoIN4mlir10AffineExprEEN4llvm8ArrayRefIT_EERNS2_20BumpPtrAllocatorImplINS2_15MallocAllocatorELm4096ELm4096ELm128EEES5_ Unexecuted instantiation: MLIRContext.cpp:_ZL16copyArrayRefIntoIbEN4llvm8ArrayRefIT_EERNS0_20BumpPtrAllocatorImplINS0_15MallocAllocatorELm4096ELm4096ELm128EEES3_ |
408 | | |
409 | | //===----------------------------------------------------------------------===// |
410 | | // Diagnostic Handlers |
411 | | //===----------------------------------------------------------------------===// |
412 | | |
413 | | /// Returns the diagnostic engine for this context. |
414 | 0 | DiagnosticEngine &MLIRContext::getDiagEngine() { return getImpl().diagEngine; } |
415 | | |
416 | | //===----------------------------------------------------------------------===// |
417 | | // Dialect and Operation Registration |
418 | | //===----------------------------------------------------------------------===// |
419 | | |
420 | | /// Return information about all registered IR dialects. |
421 | 0 | std::vector<Dialect *> MLIRContext::getRegisteredDialects() { |
422 | 0 | // Lock access to the context registry. |
423 | 0 | ScopedReaderLock registryLock(impl->contextMutex, impl->threadingIsEnabled); |
424 | 0 | std::vector<Dialect *> result; |
425 | 0 | result.reserve(impl->dialects.size()); |
426 | 0 | for (auto &dialect : impl->dialects) |
427 | 0 | result.push_back(dialect.get()); |
428 | 0 | return result; |
429 | 0 | } |
430 | | |
431 | | /// Get a registered IR dialect with the given namespace. If none is found, |
432 | | /// then return nullptr. |
433 | 0 | Dialect *MLIRContext::getRegisteredDialect(StringRef name) { |
434 | 0 | // Lock access to the context registry. |
435 | 0 | ScopedReaderLock registryLock(impl->contextMutex, impl->threadingIsEnabled); |
436 | 0 |
|
437 | 0 | // Dialects are sorted by name, so we can use binary search for lookup. |
438 | 0 | auto it = llvm::lower_bound( |
439 | 0 | impl->dialects, name, |
440 | 0 | [](const auto &lhs, StringRef rhs) { return lhs->getNamespace() < rhs; }); |
441 | 0 | return (it != impl->dialects.end() && (*it)->getNamespace() == name) |
442 | 0 | ? (*it).get() |
443 | 0 | : nullptr; |
444 | 0 | } |
445 | | |
446 | | /// Register this dialect object with the specified context. The context |
447 | | /// takes ownership of the heap allocated dialect. |
448 | 0 | void Dialect::registerDialect(MLIRContext *context) { |
449 | 0 | auto &impl = context->getImpl(); |
450 | 0 | std::unique_ptr<Dialect> dialect(this); |
451 | 0 |
|
452 | 0 | // Lock access to the context registry. |
453 | 0 | ScopedWriterLock registryLock(impl.contextMutex, impl.threadingIsEnabled); |
454 | 0 |
|
455 | 0 | // Get the correct insertion position sorted by namespace. |
456 | 0 | auto insertPt = llvm::lower_bound( |
457 | 0 | impl.dialects, dialect, [](const auto &lhs, const auto &rhs) { |
458 | 0 | return lhs->getNamespace() < rhs->getNamespace(); |
459 | 0 | }); |
460 | 0 |
|
461 | 0 | // Abort if dialect with namespace has already been registered. |
462 | 0 | if (insertPt != impl.dialects.end() && |
463 | 0 | (*insertPt)->getNamespace() == getNamespace()) { |
464 | 0 | llvm::report_fatal_error("a dialect with namespace '" + getNamespace() + |
465 | 0 | "' has already been registered"); |
466 | 0 | } |
467 | 0 | impl.dialects.insert(insertPt, std::move(dialect)); |
468 | 0 | } |
469 | | |
470 | 0 | bool MLIRContext::allowsUnregisteredDialects() { |
471 | 0 | return impl->allowUnregisteredDialects; |
472 | 0 | } |
473 | | |
474 | 0 | void MLIRContext::allowUnregisteredDialects(bool allowing) { |
475 | 0 | impl->allowUnregisteredDialects = allowing; |
476 | 0 | } |
477 | | |
478 | | /// Return true if multi-threading is disabled by the context. |
479 | 0 | bool MLIRContext::isMultithreadingEnabled() { |
480 | 0 | return impl->threadingIsEnabled && llvm::llvm_is_multithreaded(); |
481 | 0 | } |
482 | | |
483 | | /// Set the flag specifying if multi-threading is disabled by the context. |
484 | 0 | void MLIRContext::disableMultithreading(bool disable) { |
485 | 0 | impl->threadingIsEnabled = !disable; |
486 | 0 |
|
487 | 0 | // Update the threading mode for each of the uniquers. |
488 | 0 | impl->affineUniquer.disableMultithreading(disable); |
489 | 0 | impl->attributeUniquer.disableMultithreading(disable); |
490 | 0 | impl->typeUniquer.disableMultithreading(disable); |
491 | 0 | } |
492 | | |
493 | | /// Return true if we should attach the operation to diagnostics emitted via |
494 | | /// Operation::emit. |
495 | 0 | bool MLIRContext::shouldPrintOpOnDiagnostic() { |
496 | 0 | return impl->printOpOnDiagnostic; |
497 | 0 | } |
498 | | |
499 | | /// Set the flag specifying if we should attach the operation to diagnostics |
500 | | /// emitted via Operation::emit. |
501 | 0 | void MLIRContext::printOpOnDiagnostic(bool enable) { |
502 | 0 | impl->printOpOnDiagnostic = enable; |
503 | 0 | } |
504 | | |
505 | | /// Return true if we should attach the current stacktrace to diagnostics when |
506 | | /// emitted. |
507 | 0 | bool MLIRContext::shouldPrintStackTraceOnDiagnostic() { |
508 | 0 | return impl->printStackTraceOnDiagnostic; |
509 | 0 | } |
510 | | |
511 | | /// Set the flag specifying if we should attach the current stacktrace when |
512 | | /// emitting diagnostics. |
513 | 0 | void MLIRContext::printStackTraceOnDiagnostic(bool enable) { |
514 | 0 | impl->printStackTraceOnDiagnostic = enable; |
515 | 0 | } |
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. |
520 | 0 | std::vector<AbstractOperation *> MLIRContext::getRegisteredOperations() { |
521 | 0 | std::vector<std::pair<StringRef, AbstractOperation *>> opsToSort; |
522 | 0 |
|
523 | 0 | { // Lock access to the context registry. |
524 | 0 | ScopedReaderLock registryLock(impl->contextMutex, impl->threadingIsEnabled); |
525 | 0 |
|
526 | 0 | // We just have the operations in a non-deterministic hash table order. Dump |
527 | 0 | // into a temporary array, then sort it by operation name to get a stable |
528 | 0 | // ordering. |
529 | 0 | llvm::StringMap<AbstractOperation> ®isteredOps = |
530 | 0 | impl->registeredOperations; |
531 | 0 |
|
532 | 0 | opsToSort.reserve(registeredOps.size()); |
533 | 0 | for (auto &elt : registeredOps) |
534 | 0 | opsToSort.push_back({elt.first(), &elt.second}); |
535 | 0 | } |
536 | 0 |
|
537 | 0 | llvm::array_pod_sort(opsToSort.begin(), opsToSort.end()); |
538 | 0 |
|
539 | 0 | std::vector<AbstractOperation *> result; |
540 | 0 | result.reserve(opsToSort.size()); |
541 | 0 | for (auto &elt : opsToSort) |
542 | 0 | result.push_back(elt.second); |
543 | 0 | return result; |
544 | 0 | } |
545 | | |
546 | 0 | bool MLIRContext::isOperationRegistered(StringRef name) { |
547 | 0 | // Lock access to the context registry. |
548 | 0 | ScopedReaderLock registryLock(impl->contextMutex, impl->threadingIsEnabled); |
549 | 0 |
|
550 | 0 | return impl->registeredOperations.count(name); |
551 | 0 | } |
552 | | |
553 | | void Dialect::addOperation(AbstractOperation opInfo) { |
554 | | assert((getNamespace().empty() || |
555 | | opInfo.name.split('.').first == getNamespace()) && |
556 | | "op name doesn't start with dialect namespace"); |
557 | | assert(&opInfo.dialect == this && "Dialect object mismatch"); |
558 | | auto &impl = context->getImpl(); |
559 | | |
560 | | // Lock access to the context registry. |
561 | | ScopedWriterLock registryLock(impl.contextMutex, impl.threadingIsEnabled); |
562 | | if (!impl.registeredOperations.insert({opInfo.name, opInfo}).second) { |
563 | | llvm::errs() << "error: operation named '" << opInfo.name |
564 | | << "' is already registered.\n"; |
565 | | abort(); |
566 | | } |
567 | | } |
568 | | |
569 | | /// Register a dialect-specific symbol(e.g. type) with the current context. |
570 | 0 | void Dialect::addSymbol(TypeID typeID) { |
571 | 0 | auto &impl = context->getImpl(); |
572 | 0 |
|
573 | 0 | // Lock access to the context registry. |
574 | 0 | ScopedWriterLock registryLock(impl.contextMutex, impl.threadingIsEnabled); |
575 | 0 | if (!impl.registeredDialectSymbols.insert({typeID, this}).second) { |
576 | 0 | llvm::errs() << "error: dialect symbol already registered.\n"; |
577 | 0 | abort(); |
578 | 0 | } |
579 | 0 | } |
580 | | |
581 | | /// Look up the specified operation in the operation set and return a pointer |
582 | | /// to it if present. Otherwise, return a null pointer. |
583 | | const AbstractOperation *AbstractOperation::lookup(StringRef opName, |
584 | 0 | MLIRContext *context) { |
585 | 0 | auto &impl = context->getImpl(); |
586 | 0 |
|
587 | 0 | // Lock access to the context registry. |
588 | 0 | ScopedReaderLock registryLock(impl.contextMutex, impl.threadingIsEnabled); |
589 | 0 | auto it = impl.registeredOperations.find(opName); |
590 | 0 | if (it != impl.registeredOperations.end()) |
591 | 0 | return &it->second; |
592 | 0 | return nullptr; |
593 | 0 | } |
594 | | |
595 | | //===----------------------------------------------------------------------===// |
596 | | // Identifier uniquing |
597 | | //===----------------------------------------------------------------------===// |
598 | | |
599 | | /// Return an identifier for the specified string. |
600 | 0 | Identifier Identifier::get(StringRef str, MLIRContext *context) { |
601 | 0 | auto &impl = context->getImpl(); |
602 | 0 |
|
603 | 0 | // Check for an existing identifier in read-only mode. |
604 | 0 | if (context->isMultithreadingEnabled()) { |
605 | 0 | llvm::sys::SmartScopedReader<true> contextLock(impl.identifierMutex); |
606 | 0 | auto it = impl.identifiers.find(str); |
607 | 0 | if (it != impl.identifiers.end()) |
608 | 0 | return Identifier(&*it); |
609 | 0 | } |
610 | 0 | |
611 | 0 | // Check invariants after seeing if we already have something in the |
612 | 0 | // identifier table - if we already had it in the table, then it already |
613 | 0 | // passed invariant checks. |
614 | 0 | assert(!str.empty() && "Cannot create an empty identifier"); |
615 | 0 | assert(str.find('\0') == StringRef::npos && |
616 | 0 | "Cannot create an identifier with a nul character"); |
617 | 0 |
|
618 | 0 | // Acquire a writer-lock so that we can safely create the new instance. |
619 | 0 | ScopedWriterLock contextLock(impl.identifierMutex, impl.threadingIsEnabled); |
620 | 0 | auto it = impl.identifiers.insert(str).first; |
621 | 0 | return Identifier(&*it); |
622 | 0 | } |
623 | | |
624 | | //===----------------------------------------------------------------------===// |
625 | | // Type uniquing |
626 | | //===----------------------------------------------------------------------===// |
627 | | |
628 | 0 | static Dialect &lookupDialectForSymbol(MLIRContext *ctx, TypeID typeID) { |
629 | 0 | auto &impl = ctx->getImpl(); |
630 | 0 | auto it = impl.registeredDialectSymbols.find(typeID); |
631 | 0 | if (it == impl.registeredDialectSymbols.end()) |
632 | 0 | llvm::report_fatal_error( |
633 | 0 | "Trying to create a type that was not registered in this MLIRContext."); |
634 | 0 | return *it->second; |
635 | 0 | } |
636 | | |
637 | | /// Returns the storage uniquer used for constructing type storage instances. |
638 | | /// This should not be used directly. |
639 | 0 | StorageUniquer &MLIRContext::getTypeUniquer() { return getImpl().typeUniquer; } |
640 | | |
641 | | /// Get the dialect that registered the type with the provided typeid. |
642 | 0 | Dialect &TypeUniquer::lookupDialectForType(MLIRContext *ctx, TypeID typeID) { |
643 | 0 | return lookupDialectForSymbol(ctx, typeID); |
644 | 0 | } |
645 | | |
646 | | FloatType FloatType::get(StandardTypes::Kind kind, MLIRContext *context) { |
647 | | assert(kindof(kind) && "Not a FP kind."); |
648 | | switch (kind) { |
649 | | case StandardTypes::BF16: |
650 | | return context->getImpl().bf16Ty; |
651 | | case StandardTypes::F16: |
652 | | return context->getImpl().f16Ty; |
653 | | case StandardTypes::F32: |
654 | | return context->getImpl().f32Ty; |
655 | | case StandardTypes::F64: |
656 | | return context->getImpl().f64Ty; |
657 | | default: |
658 | | llvm_unreachable("unexpected floating-point kind"); |
659 | | } |
660 | | } |
661 | | |
662 | | /// Get an instance of the IndexType. |
663 | 0 | IndexType IndexType::get(MLIRContext *context) { |
664 | 0 | return context->getImpl().indexTy; |
665 | 0 | } |
666 | | |
667 | | /// Return an existing integer type instance if one is cached within the |
668 | | /// context. |
669 | | static IntegerType |
670 | | getCachedIntegerType(unsigned width, |
671 | | IntegerType::SignednessSemantics signedness, |
672 | 0 | MLIRContext *context) { |
673 | 0 | if (signedness != IntegerType::Signless) |
674 | 0 | return IntegerType(); |
675 | 0 | |
676 | 0 | switch (width) { |
677 | 0 | case 1: |
678 | 0 | return context->getImpl().int1Ty; |
679 | 0 | case 8: |
680 | 0 | return context->getImpl().int8Ty; |
681 | 0 | case 16: |
682 | 0 | return context->getImpl().int16Ty; |
683 | 0 | case 32: |
684 | 0 | return context->getImpl().int32Ty; |
685 | 0 | case 64: |
686 | 0 | return context->getImpl().int64Ty; |
687 | 0 | case 128: |
688 | 0 | return context->getImpl().int128Ty; |
689 | 0 | default: |
690 | 0 | return IntegerType(); |
691 | 0 | } |
692 | 0 | } |
693 | | |
694 | 0 | IntegerType IntegerType::get(unsigned width, MLIRContext *context) { |
695 | 0 | return get(width, IntegerType::Signless, context); |
696 | 0 | } |
697 | | |
698 | | IntegerType IntegerType::get(unsigned width, |
699 | | IntegerType::SignednessSemantics signedness, |
700 | 0 | MLIRContext *context) { |
701 | 0 | if (auto cached = getCachedIntegerType(width, signedness, context)) |
702 | 0 | return cached; |
703 | 0 | return Base::get(context, StandardTypes::Integer, width, signedness); |
704 | 0 | } |
705 | | |
706 | 0 | IntegerType IntegerType::getChecked(unsigned width, Location location) { |
707 | 0 | return getChecked(width, IntegerType::Signless, location); |
708 | 0 | } |
709 | | |
710 | | IntegerType IntegerType::getChecked(unsigned width, |
711 | | SignednessSemantics signedness, |
712 | 0 | Location location) { |
713 | 0 | if (auto cached = |
714 | 0 | getCachedIntegerType(width, signedness, location->getContext())) |
715 | 0 | return cached; |
716 | 0 | return Base::getChecked(location, StandardTypes::Integer, width, signedness); |
717 | 0 | } |
718 | | |
719 | | /// Get an instance of the NoneType. |
720 | 0 | NoneType NoneType::get(MLIRContext *context) { |
721 | 0 | return context->getImpl().noneType; |
722 | 0 | } |
723 | | |
724 | | //===----------------------------------------------------------------------===// |
725 | | // Attribute uniquing |
726 | | //===----------------------------------------------------------------------===// |
727 | | |
728 | | /// Returns the storage uniquer used for constructing attribute storage |
729 | | /// instances. This should not be used directly. |
730 | 0 | StorageUniquer &MLIRContext::getAttributeUniquer() { |
731 | 0 | return getImpl().attributeUniquer; |
732 | 0 | } |
733 | | |
734 | | /// Initialize the given attribute storage instance. |
735 | | void AttributeUniquer::initializeAttributeStorage(AttributeStorage *storage, |
736 | | MLIRContext *ctx, |
737 | 0 | TypeID attrID) { |
738 | 0 | storage->initializeDialect(lookupDialectForSymbol(ctx, attrID)); |
739 | 0 |
|
740 | 0 | // If the attribute did not provide a type, then default to NoneType. |
741 | 0 | if (!storage->getType()) |
742 | 0 | storage->setType(NoneType::get(ctx)); |
743 | 0 | } |
744 | | |
745 | 0 | BoolAttr BoolAttr::get(bool value, MLIRContext *context) { |
746 | 0 | return value ? context->getImpl().trueAttr : context->getImpl().falseAttr; |
747 | 0 | } |
748 | | |
749 | 0 | UnitAttr UnitAttr::get(MLIRContext *context) { |
750 | 0 | return context->getImpl().unitAttr; |
751 | 0 | } |
752 | | |
753 | 0 | Location UnknownLoc::get(MLIRContext *context) { |
754 | 0 | return context->getImpl().unknownLocAttr; |
755 | 0 | } |
756 | | |
757 | | /// Return empty dictionary. |
758 | 0 | DictionaryAttr DictionaryAttr::getEmpty(MLIRContext *context) { |
759 | 0 | return context->getImpl().emptyDictionaryAttr; |
760 | 0 | } |
761 | | |
762 | | //===----------------------------------------------------------------------===// |
763 | | // AffineMap uniquing |
764 | | //===----------------------------------------------------------------------===// |
765 | | |
766 | 0 | StorageUniquer &MLIRContext::getAffineUniquer() { |
767 | 0 | return getImpl().affineUniquer; |
768 | 0 | } |
769 | | |
770 | | AffineMap AffineMap::getImpl(unsigned dimCount, unsigned symbolCount, |
771 | | ArrayRef<AffineExpr> results, |
772 | 0 | MLIRContext *context) { |
773 | 0 | auto &impl = context->getImpl(); |
774 | 0 | auto key = std::make_tuple(dimCount, symbolCount, results); |
775 | 0 |
|
776 | 0 | // Safely get or create an AffineMap instance. |
777 | 0 | return safeGetOrCreate( |
778 | 0 | impl.affineMaps, key, impl.affineMutex, impl.threadingIsEnabled, [&] { |
779 | 0 | auto *res = impl.affineAllocator.Allocate<detail::AffineMapStorage>(); |
780 | 0 |
|
781 | 0 | // Copy the results into the bump pointer. |
782 | 0 | results = copyArrayRefInto(impl.affineAllocator, results); |
783 | 0 |
|
784 | 0 | // Initialize the memory using placement new. |
785 | 0 | new (res) |
786 | 0 | detail::AffineMapStorage{dimCount, symbolCount, results, context}; |
787 | 0 | return AffineMap(res); |
788 | 0 | }); |
789 | 0 | } |
790 | | |
791 | 0 | AffineMap AffineMap::get(MLIRContext *context) { |
792 | 0 | return getImpl(/*dimCount=*/0, /*symbolCount=*/0, /*results=*/{}, context); |
793 | 0 | } |
794 | | |
795 | | AffineMap AffineMap::get(unsigned dimCount, unsigned symbolCount, |
796 | 0 | MLIRContext *context) { |
797 | 0 | return getImpl(dimCount, symbolCount, /*results=*/{}, context); |
798 | 0 | } |
799 | | |
800 | | AffineMap AffineMap::get(unsigned dimCount, unsigned symbolCount, |
801 | 0 | AffineExpr result) { |
802 | 0 | return getImpl(dimCount, symbolCount, {result}, result.getContext()); |
803 | 0 | } |
804 | | |
805 | | AffineMap AffineMap::get(unsigned dimCount, unsigned symbolCount, |
806 | 0 | ArrayRef<AffineExpr> results, MLIRContext *context) { |
807 | 0 | return getImpl(dimCount, symbolCount, results, context); |
808 | 0 | } |
809 | | |
810 | | //===----------------------------------------------------------------------===// |
811 | | // Integer Sets: these are allocated into the bump pointer, and are immutable. |
812 | | // Unlike AffineMap's, these are uniqued only if they are small. |
813 | | //===----------------------------------------------------------------------===// |
814 | | |
815 | | IntegerSet IntegerSet::get(unsigned dimCount, unsigned symbolCount, |
816 | | ArrayRef<AffineExpr> constraints, |
817 | 0 | ArrayRef<bool> eqFlags) { |
818 | 0 | // The number of constraints can't be zero. |
819 | 0 | assert(!constraints.empty()); |
820 | 0 | assert(constraints.size() == eqFlags.size()); |
821 | 0 |
|
822 | 0 | auto &impl = constraints[0].getContext()->getImpl(); |
823 | 0 |
|
824 | 0 | // A utility function to construct a new IntegerSetStorage instance. |
825 | 0 | auto constructorFn = [&] { |
826 | 0 | auto *res = impl.affineAllocator.Allocate<detail::IntegerSetStorage>(); |
827 | 0 |
|
828 | 0 | // Copy the results and equality flags into the bump pointer. |
829 | 0 | constraints = copyArrayRefInto(impl.affineAllocator, constraints); |
830 | 0 | eqFlags = copyArrayRefInto(impl.affineAllocator, eqFlags); |
831 | 0 |
|
832 | 0 | // Initialize the memory using placement new. |
833 | 0 | new (res) |
834 | 0 | detail::IntegerSetStorage{dimCount, symbolCount, constraints, eqFlags}; |
835 | 0 | return IntegerSet(res); |
836 | 0 | }; |
837 | 0 |
|
838 | 0 | // If this instance is uniqued, then we handle it separately so that multiple |
839 | 0 | // threads may simultaneously access existing instances. |
840 | 0 | if (constraints.size() < IntegerSet::kUniquingThreshold) { |
841 | 0 | auto key = std::make_tuple(dimCount, symbolCount, constraints, eqFlags); |
842 | 0 | return safeGetOrCreate(impl.integerSets, key, impl.affineMutex, |
843 | 0 | impl.threadingIsEnabled, constructorFn); |
844 | 0 | } |
845 | 0 | |
846 | 0 | // Otherwise, acquire a writer-lock so that we can safely create the new |
847 | 0 | // instance. |
848 | 0 | ScopedWriterLock affineLock(impl.affineMutex, impl.threadingIsEnabled); |
849 | 0 | return constructorFn(); |
850 | 0 | } |
851 | | |
852 | | //===----------------------------------------------------------------------===// |
853 | | // StorageUniquerSupport |
854 | | //===----------------------------------------------------------------------===// |
855 | | |
856 | | /// Utility method to generate a default location for use when checking the |
857 | | /// construction invariants of a storage object. This is defined out-of-line to |
858 | | /// avoid the need to include Location.h. |
859 | | const AttributeStorage * |
860 | 0 | mlir::detail::generateUnknownStorageLocation(MLIRContext *ctx) { |
861 | 0 | return reinterpret_cast<const AttributeStorage *>( |
862 | 0 | ctx->getImpl().unknownLocAttr.getAsOpaquePointer()); |
863 | 0 | } |