1//===- AttributeDetail.h - MLIR Affine Map details Class --------*- C++ -*-===//
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// This holds implementation details of Attribute.
10//
11//===----------------------------------------------------------------------===//
12
13#ifndef ATTRIBUTEDETAIL_H_
14#define ATTRIBUTEDETAIL_H_
15
16#include "mlir/IR/AffineMap.h"
17#include "mlir/IR/AttributeSupport.h"
18#include "mlir/IR/BuiltinAttributes.h"
19#include "mlir/IR/BuiltinTypes.h"
20#include "mlir/IR/IntegerSet.h"
21#include "mlir/IR/MLIRContext.h"
22#include "mlir/Support/StorageUniquer.h"
23#include "mlir/Support/ThreadLocalCache.h"
24#include "llvm/ADT/APFloat.h"
25#include "llvm/ADT/PointerIntPair.h"
26#include "llvm/Support/TrailingObjects.h"
27
28namespace mlir {
29namespace detail {
30
31//===----------------------------------------------------------------------===//
32// Elements Attributes
33//===----------------------------------------------------------------------===//
34
35/// Return the bit width which DenseElementsAttr should use for this type.
36inline size_t getDenseElementBitWidth(Type eltType) {
37 // Align the width for complex to 8 to make storage and interpretation easier.
38 if (ComplexType comp = llvm::dyn_cast<ComplexType>(eltType))
39 return llvm::alignTo<8>(getDenseElementBitWidth(comp.getElementType())) * 2;
40 if (eltType.isIndex())
41 return IndexType::kInternalStorageBitWidth;
42 return eltType.getIntOrFloatBitWidth();
43}
44
45/// An attribute representing a reference to a dense vector or tensor object.
46struct DenseElementsAttributeStorage : public AttributeStorage {
47public:
48 DenseElementsAttributeStorage(ShapedType type, bool isSplat)
49 : type(type), isSplat(isSplat) {}
50
51 ShapedType type;
52 bool isSplat;
53};
54
55/// An attribute representing a reference to a dense vector or tensor object.
56struct DenseIntOrFPElementsAttrStorage : public DenseElementsAttributeStorage {
57 DenseIntOrFPElementsAttrStorage(ShapedType ty, ArrayRef<char> data,
58 bool isSplat = false)
59 : DenseElementsAttributeStorage(ty, isSplat), data(data) {}
60
61 struct KeyTy {
62 KeyTy(ShapedType type, ArrayRef<char> data, llvm::hash_code hashCode,
63 bool isSplat = false)
64 : type(type), data(data), hashCode(hashCode), isSplat(isSplat) {}
65
66 /// The type of the dense elements.
67 ShapedType type;
68
69 /// The raw buffer for the data storage.
70 ArrayRef<char> data;
71
72 /// The computed hash code for the storage data.
73 llvm::hash_code hashCode;
74
75 /// A boolean that indicates if this data is a splat or not.
76 bool isSplat;
77 };
78
79 /// Compare this storage instance with the provided key.
80 bool operator==(const KeyTy &key) const {
81 return key.type == type && key.data == data;
82 }
83
84 /// Construct a key from a shaped type, raw data buffer, and a flag that
85 /// signals if the data is already known to be a splat. Callers to this
86 /// function are expected to tag preknown splat values when possible, e.g. one
87 /// element shapes.
88 static KeyTy getKey(ShapedType ty, ArrayRef<char> data, bool isKnownSplat) {
89 // Handle an empty storage instance.
90 if (data.empty())
91 return KeyTy(ty, data, 0);
92
93 // If the data is already known to be a splat, the key hash value is
94 // directly the data buffer.
95 bool isBoolData = ty.getElementType().isInteger(1);
96 if (isKnownSplat) {
97 if (isBoolData)
98 return getKeyForSplatBoolData(ty, data[0] != 0);
99 return KeyTy(ty, data, llvm::hash_value(S: data), isKnownSplat);
100 }
101
102 // Otherwise, we need to check if the data corresponds to a splat or not.
103
104 // Handle the simple case of only one element.
105 size_t numElements = ty.getNumElements();
106 assert(numElements != 1 && "splat of 1 element should already be detected");
107
108 // Handle boolean values directly as they are packed to 1-bit.
109 if (isBoolData)
110 return getKeyForBoolData(ty, data, numElements);
111
112 size_t elementWidth = getDenseElementBitWidth(ty.getElementType());
113 // Non 1-bit dense elements are padded to 8-bits.
114 size_t storageSize = llvm::divideCeil(Numerator: elementWidth, CHAR_BIT);
115 assert(((data.size() / storageSize) == numElements) &&
116 "data does not hold expected number of elements");
117
118 // Create the initial hash value with just the first element.
119 auto firstElt = data.take_front(N: storageSize);
120 auto hashVal = llvm::hash_value(firstElt);
121
122 // Check to see if this storage represents a splat. If it doesn't then
123 // combine the hash for the data starting with the first non splat element.
124 for (size_t i = storageSize, e = data.size(); i != e; i += storageSize)
125 if (memcmp(s1: data.data(), s2: &data[i], n: storageSize))
126 return KeyTy(ty, data, llvm::hash_combine(hashVal, data.drop_front(N: i)));
127
128 // Otherwise, this is a splat so just return the hash of the first element.
129 return KeyTy(ty, firstElt, hashVal, /*isSplat=*/true);
130 }
131
132 /// Construct a key with a set of boolean data.
133 static KeyTy getKeyForBoolData(ShapedType ty, ArrayRef<char> data,
134 size_t numElements) {
135 ArrayRef<char> splatData = data;
136 bool splatValue = splatData.front() & 1;
137
138 // Check the simple case where the data matches the known splat value.
139 if (splatData == ArrayRef<char>(splatValue ? kSplatTrue : kSplatFalse))
140 return getKeyForSplatBoolData(ty, splatValue);
141
142 // Handle the case where the potential splat value is 1 and the number of
143 // elements is non 8-bit aligned.
144 size_t numOddElements = numElements % CHAR_BIT;
145 if (splatValue && numOddElements != 0) {
146 // Check that all bits are set in the last value.
147 char lastElt = splatData.back();
148 if (lastElt != llvm::maskTrailingOnes<unsigned char>(N: numOddElements))
149 return KeyTy(ty, data, llvm::hash_value(S: data));
150
151 // If this is the only element, the data is known to be a splat.
152 if (splatData.size() == 1)
153 return getKeyForSplatBoolData(ty, splatValue);
154 splatData = splatData.drop_back();
155 }
156
157 // Check that the data buffer corresponds to a splat of the proper mask.
158 char mask = splatValue ? ~0 : 0;
159 return llvm::all_of(Range&: splatData, P: [mask](char c) { return c == mask; })
160 ? getKeyForSplatBoolData(ty, splatValue)
161 : KeyTy(ty, data, llvm::hash_value(S: data));
162 }
163
164 /// Return a key to use for a boolean splat of the given value.
165 static KeyTy getKeyForSplatBoolData(ShapedType type, bool splatValue) {
166 const char &splatData = splatValue ? kSplatTrue : kSplatFalse;
167 return KeyTy(type, splatData, llvm::hash_value(value: splatData),
168 /*isSplat=*/true);
169 }
170
171 /// Hash the key for the storage.
172 static llvm::hash_code hashKey(const KeyTy &key) {
173 return llvm::hash_combine(key.type, key.hashCode);
174 }
175
176 /// Construct a new storage instance.
177 static DenseIntOrFPElementsAttrStorage *
178 construct(AttributeStorageAllocator &allocator, KeyTy key) {
179 // If the data buffer is non-empty, we copy it into the allocator with a
180 // 64-bit alignment.
181 ArrayRef<char> copy, data = key.data;
182 if (!data.empty()) {
183 char *rawData = reinterpret_cast<char *>(
184 allocator.allocate(size: data.size(), alignment: alignof(uint64_t)));
185 std::memcpy(dest: rawData, src: data.data(), n: data.size());
186 copy = ArrayRef<char>(rawData, data.size());
187 }
188
189 return new (allocator.allocate<DenseIntOrFPElementsAttrStorage>())
190 DenseIntOrFPElementsAttrStorage(key.type, copy, key.isSplat);
191 }
192
193 ArrayRef<char> data;
194
195 /// The values used to denote a boolean splat value.
196 // This is not using constexpr declaration due to compilation failure
197 // encountered with MSVC where it would inline these values, which makes it
198 // unsafe to refer by reference in KeyTy.
199 static const char kSplatTrue;
200 static const char kSplatFalse;
201};
202
203/// An attribute representing a reference to a dense vector or tensor object
204/// containing strings.
205struct DenseStringElementsAttrStorage : public DenseElementsAttributeStorage {
206 DenseStringElementsAttrStorage(ShapedType ty, ArrayRef<StringRef> data,
207 bool isSplat = false)
208 : DenseElementsAttributeStorage(ty, isSplat), data(data) {}
209
210 struct KeyTy {
211 KeyTy(ShapedType type, ArrayRef<StringRef> data, llvm::hash_code hashCode,
212 bool isSplat = false)
213 : type(type), data(data), hashCode(hashCode), isSplat(isSplat) {}
214
215 /// The type of the dense elements.
216 ShapedType type;
217
218 /// The raw buffer for the data storage.
219 ArrayRef<StringRef> data;
220
221 /// The computed hash code for the storage data.
222 llvm::hash_code hashCode;
223
224 /// A boolean that indicates if this data is a splat or not.
225 bool isSplat;
226 };
227
228 /// Compare this storage instance with the provided key.
229 bool operator==(const KeyTy &key) const {
230 if (key.type != type)
231 return false;
232
233 // Otherwise, we can default to just checking the data. StringRefs compare
234 // by contents.
235 return key.data == data;
236 }
237
238 /// Construct a key from a shaped type, StringRef data buffer, and a flag that
239 /// signals if the data is already known to be a splat. Callers to this
240 /// function are expected to tag preknown splat values when possible, e.g. one
241 /// element shapes.
242 static KeyTy getKey(ShapedType ty, ArrayRef<StringRef> data,
243 bool isKnownSplat) {
244 // Handle an empty storage instance.
245 if (data.empty())
246 return KeyTy(ty, data, 0);
247
248 // If the data is already known to be a splat, the key hash value is
249 // directly the data buffer.
250 if (isKnownSplat)
251 return KeyTy(ty, data, llvm::hash_value(S: data.front()), isKnownSplat);
252
253 // Handle the simple case of only one element.
254 assert(ty.getNumElements() != 1 &&
255 "splat of 1 element should already be detected");
256
257 // Create the initial hash value with just the first element.
258 const auto &firstElt = data.front();
259 auto hashVal = llvm::hash_value(S: firstElt);
260
261 // Check to see if this storage represents a splat. If it doesn't then
262 // combine the hash for the data starting with the first non splat element.
263 for (size_t i = 1, e = data.size(); i != e; i++)
264 if (!firstElt.equals(RHS: data[i]))
265 return KeyTy(ty, data, llvm::hash_combine(args: hashVal, args: data.drop_front(N: i)));
266
267 // Otherwise, this is a splat so just return the hash of the first element.
268 return KeyTy(ty, data.take_front(), hashVal, /*isSplat=*/true);
269 }
270
271 /// Hash the key for the storage.
272 static llvm::hash_code hashKey(const KeyTy &key) {
273 return llvm::hash_combine(key.type, key.hashCode);
274 }
275
276 /// Construct a new storage instance.
277 static DenseStringElementsAttrStorage *
278 construct(AttributeStorageAllocator &allocator, KeyTy key) {
279 // If the data buffer is non-empty, we copy it into the allocator with a
280 // 64-bit alignment.
281 ArrayRef<StringRef> copy, data = key.data;
282 if (data.empty()) {
283 return new (allocator.allocate<DenseStringElementsAttrStorage>())
284 DenseStringElementsAttrStorage(key.type, copy, key.isSplat);
285 }
286
287 int numEntries = key.isSplat ? 1 : data.size();
288
289 // Compute the amount data needed to store the ArrayRef and StringRef
290 // contents.
291 size_t dataSize = sizeof(StringRef) * numEntries;
292 for (int i = 0; i < numEntries; i++)
293 dataSize += data[i].size();
294
295 char *rawData = reinterpret_cast<char *>(
296 allocator.allocate(size: dataSize, alignment: alignof(uint64_t)));
297
298 // Setup a mutable array ref of our string refs so that we can update their
299 // contents.
300 auto mutableCopy = MutableArrayRef<StringRef>(
301 reinterpret_cast<StringRef *>(rawData), numEntries);
302 auto *stringData = rawData + numEntries * sizeof(StringRef);
303
304 for (int i = 0; i < numEntries; i++) {
305 memcpy(dest: stringData, src: data[i].data(), n: data[i].size());
306 mutableCopy[i] = StringRef(stringData, data[i].size());
307 stringData += data[i].size();
308 }
309
310 copy =
311 ArrayRef<StringRef>(reinterpret_cast<StringRef *>(rawData), numEntries);
312
313 return new (allocator.allocate<DenseStringElementsAttrStorage>())
314 DenseStringElementsAttrStorage(key.type, copy, key.isSplat);
315 }
316
317 ArrayRef<StringRef> data;
318};
319
320//===----------------------------------------------------------------------===//
321// StringAttr
322//===----------------------------------------------------------------------===//
323
324struct StringAttrStorage : public AttributeStorage {
325 StringAttrStorage(StringRef value, Type type)
326 : type(type), value(value), referencedDialect(nullptr) {}
327
328 /// The hash key is a tuple of the parameter types.
329 using KeyTy = std::pair<StringRef, Type>;
330 bool operator==(const KeyTy &key) const {
331 return value == key.first && type == key.second;
332 }
333 static ::llvm::hash_code hashKey(const KeyTy &key) {
334 return DenseMapInfo<KeyTy>::getHashValue(PairVal: key);
335 }
336
337 /// Define a construction method for creating a new instance of this
338 /// storage.
339 static StringAttrStorage *construct(AttributeStorageAllocator &allocator,
340 const KeyTy &key) {
341 return new (allocator.allocate<StringAttrStorage>())
342 StringAttrStorage(allocator.copyInto(str: key.first), key.second);
343 }
344
345 /// Initialize the storage given an MLIRContext.
346 void initialize(MLIRContext *context);
347
348 /// The type of the string.
349 Type type;
350 /// The raw string value.
351 StringRef value;
352 /// If the string value contains a dialect namespace prefix (e.g.
353 /// dialect.blah), this is the dialect referenced.
354 Dialect *referencedDialect;
355};
356
357//===----------------------------------------------------------------------===//
358// DistinctAttr
359//===----------------------------------------------------------------------===//
360
361/// An attribute to store a distinct reference to another attribute.
362struct DistinctAttrStorage : public AttributeStorage {
363 using KeyTy = Attribute;
364
365 DistinctAttrStorage(Attribute referencedAttr)
366 : referencedAttr(referencedAttr) {}
367
368 /// Returns the referenced attribute as key.
369 KeyTy getAsKey() const { return KeyTy(referencedAttr); }
370
371 /// The referenced attribute.
372 Attribute referencedAttr;
373};
374
375/// A specialized attribute uniquer for distinct attributes that always
376/// allocates since the distinct attribute instances use the address of their
377/// storage as unique identifier.
378class DistinctAttributeUniquer {
379public:
380 /// Creates a distinct attribute storage. Allocates every time since the
381 /// address of the storage serves as unique identifier.
382 template <typename T, typename... Args>
383 static T get(MLIRContext *context, Args &&...args) {
384 static_assert(std::is_same_v<typename T::ImplType, DistinctAttrStorage>,
385 "expects a distinct attribute storage");
386 DistinctAttrStorage *storage = DistinctAttributeUniquer::allocateStorage(
387 context, referencedAttr: std::forward<Args>(args)...);
388 storage->initializeAbstractAttribute(
389 AbstractAttribute::lookup(DistinctAttr::getTypeID(), context));
390 return storage;
391 }
392
393private:
394 /// Allocates a distinct attribute storage.
395 static DistinctAttrStorage *allocateStorage(MLIRContext *context,
396 Attribute referencedAttr);
397};
398
399/// An allocator for distinct attribute storage instances. It uses thread local
400/// bump pointer allocators stored in a thread local cache to ensure the storage
401/// is freed after the destruction of the distinct attribute allocator.
402class DistinctAttributeAllocator {
403public:
404 DistinctAttributeAllocator() = default;
405
406 DistinctAttributeAllocator(DistinctAttributeAllocator &&) = delete;
407 DistinctAttributeAllocator(const DistinctAttributeAllocator &) = delete;
408 DistinctAttributeAllocator &
409 operator=(const DistinctAttributeAllocator &) = delete;
410
411 /// Allocates a distinct attribute storage using a thread local bump pointer
412 /// allocator to enable synchronization free parallel allocations.
413 DistinctAttrStorage *allocate(Attribute referencedAttr) {
414 return new (allocatorCache.get().Allocate<DistinctAttrStorage>())
415 DistinctAttrStorage(referencedAttr);
416 }
417
418private:
419 ThreadLocalCache<llvm::BumpPtrAllocator> allocatorCache;
420};
421} // namespace detail
422} // namespace mlir
423
424#endif // ATTRIBUTEDETAIL_H_
425

source code of mlir/lib/IR/AttributeDetail.h