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/BuiltinAttributes.h"
18 #include "mlir/IR/BuiltinTypes.h"
19 #include "mlir/IR/IntegerSet.h"
20 #include "mlir/IR/MLIRContext.h"
21 #include "mlir/Support/StorageUniquer.h"
22 #include "llvm/ADT/APFloat.h"
23 #include "llvm/ADT/PointerIntPair.h"
24 #include "llvm/Support/TrailingObjects.h"
25
26 namespace mlir {
27 namespace detail {
28
29 //===----------------------------------------------------------------------===//
30 // Elements Attributes
31 //===----------------------------------------------------------------------===//
32
33 /// Return the bit width which DenseElementsAttr should use for this type.
getDenseElementBitWidth(Type eltType)34 inline size_t getDenseElementBitWidth(Type eltType) {
35 // Align the width for complex to 8 to make storage and interpretation easier.
36 if (ComplexType comp = eltType.dyn_cast<ComplexType>())
37 return llvm::alignTo<8>(getDenseElementBitWidth(comp.getElementType())) * 2;
38 if (eltType.isIndex())
39 return IndexType::kInternalStorageBitWidth;
40 return eltType.getIntOrFloatBitWidth();
41 }
42
43 /// An attribute representing a reference to a dense vector or tensor object.
44 struct DenseElementsAttributeStorage : public AttributeStorage {
45 public:
DenseElementsAttributeStorageDenseElementsAttributeStorage46 DenseElementsAttributeStorage(ShapedType ty, bool isSplat)
47 : AttributeStorage(ty), isSplat(isSplat) {}
48
49 bool isSplat;
50 };
51
52 /// An attribute representing a reference to a dense vector or tensor object.
53 struct DenseIntOrFPElementsAttrStorage : public DenseElementsAttributeStorage {
54 DenseIntOrFPElementsAttrStorage(ShapedType ty, ArrayRef<char> data,
55 bool isSplat = false)
DenseElementsAttributeStorageDenseIntOrFPElementsAttrStorage56 : DenseElementsAttributeStorage(ty, isSplat), data(data) {}
57
58 struct KeyTy {
59 KeyTy(ShapedType type, ArrayRef<char> data, llvm::hash_code hashCode,
60 bool isSplat = false)
typeDenseIntOrFPElementsAttrStorage::KeyTy61 : type(type), data(data), hashCode(hashCode), isSplat(isSplat) {}
62
63 /// The type of the dense elements.
64 ShapedType type;
65
66 /// The raw buffer for the data storage.
67 ArrayRef<char> data;
68
69 /// The computed hash code for the storage data.
70 llvm::hash_code hashCode;
71
72 /// A boolean that indicates if this data is a splat or not.
73 bool isSplat;
74 };
75
76 /// Compare this storage instance with the provided key.
77 bool operator==(const KeyTy &key) const {
78 if (key.type != getType())
79 return false;
80
81 // For boolean splats we need to explicitly check that the first bit is the
82 // same. Boolean values are packed at the bit level, and even though a splat
83 // is detected the rest of the bits in the first byte may differ from the
84 // splat value.
85 if (key.type.getElementType().isInteger(1)) {
86 if (key.isSplat != isSplat)
87 return false;
88 if (isSplat)
89 return (key.data.front() & 1) == data.front();
90 }
91
92 // Otherwise, we can default to just checking the data.
93 return key.data == data;
94 }
95
96 /// Construct a key from a shaped type, raw data buffer, and a flag that
97 /// signals if the data is already known to be a splat. Callers to this
98 /// function are expected to tag preknown splat values when possible, e.g. one
99 /// element shapes.
getKeyDenseIntOrFPElementsAttrStorage100 static KeyTy getKey(ShapedType ty, ArrayRef<char> data, bool isKnownSplat) {
101 // Handle an empty storage instance.
102 if (data.empty())
103 return KeyTy(ty, data, 0);
104
105 // If the data is already known to be a splat, the key hash value is
106 // directly the data buffer.
107 if (isKnownSplat)
108 return KeyTy(ty, data, llvm::hash_value(data), isKnownSplat);
109
110 // Otherwise, we need to check if the data corresponds to a splat or not.
111
112 // Handle the simple case of only one element.
113 size_t numElements = ty.getNumElements();
114 assert(numElements != 1 && "splat of 1 element should already be detected");
115
116 // Handle boolean values directly as they are packed to 1-bit.
117 if (ty.getElementType().isInteger(1) == 1)
118 return getKeyForBoolData(ty, data, numElements);
119
120 size_t elementWidth = getDenseElementBitWidth(ty.getElementType());
121 // Non 1-bit dense elements are padded to 8-bits.
122 size_t storageSize = llvm::divideCeil(elementWidth, CHAR_BIT);
123 assert(((data.size() / storageSize) == numElements) &&
124 "data does not hold expected number of elements");
125
126 // Create the initial hash value with just the first element.
127 auto firstElt = data.take_front(storageSize);
128 auto hashVal = llvm::hash_value(firstElt);
129
130 // Check to see if this storage represents a splat. If it doesn't then
131 // combine the hash for the data starting with the first non splat element.
132 for (size_t i = storageSize, e = data.size(); i != e; i += storageSize)
133 if (memcmp(data.data(), &data[i], storageSize))
134 return KeyTy(ty, data, llvm::hash_combine(hashVal, data.drop_front(i)));
135
136 // Otherwise, this is a splat so just return the hash of the first element.
137 return KeyTy(ty, firstElt, hashVal, /*isSplat=*/true);
138 }
139
140 /// Construct a key with a set of boolean data.
getKeyForBoolDataDenseIntOrFPElementsAttrStorage141 static KeyTy getKeyForBoolData(ShapedType ty, ArrayRef<char> data,
142 size_t numElements) {
143 ArrayRef<char> splatData = data;
144 bool splatValue = splatData.front() & 1;
145
146 // Helper functor to generate a KeyTy for a boolean splat value.
147 auto generateSplatKey = [=] {
148 return KeyTy(ty, data.take_front(1),
149 llvm::hash_value(ArrayRef<char>(splatValue ? 1 : 0)),
150 /*isSplat=*/true);
151 };
152
153 // Handle the case where the potential splat value is 1 and the number of
154 // elements is non 8-bit aligned.
155 size_t numOddElements = numElements % CHAR_BIT;
156 if (splatValue && numOddElements != 0) {
157 // Check that all bits are set in the last value.
158 char lastElt = splatData.back();
159 if (lastElt != llvm::maskTrailingOnes<unsigned char>(numOddElements))
160 return KeyTy(ty, data, llvm::hash_value(data));
161
162 // If this is the only element, the data is known to be a splat.
163 if (splatData.size() == 1)
164 return generateSplatKey();
165 splatData = splatData.drop_back();
166 }
167
168 // Check that the data buffer corresponds to a splat of the proper mask.
169 char mask = splatValue ? ~0 : 0;
170 return llvm::all_of(splatData, [mask](char c) { return c == mask; })
171 ? generateSplatKey()
172 : KeyTy(ty, data, llvm::hash_value(data));
173 }
174
175 /// Hash the key for the storage.
hashKeyDenseIntOrFPElementsAttrStorage176 static llvm::hash_code hashKey(const KeyTy &key) {
177 return llvm::hash_combine(key.type, key.hashCode);
178 }
179
180 /// Construct a new storage instance.
181 static DenseIntOrFPElementsAttrStorage *
constructDenseIntOrFPElementsAttrStorage182 construct(AttributeStorageAllocator &allocator, KeyTy key) {
183 // If the data buffer is non-empty, we copy it into the allocator with a
184 // 64-bit alignment.
185 ArrayRef<char> copy, data = key.data;
186 if (!data.empty()) {
187 char *rawData = reinterpret_cast<char *>(
188 allocator.allocate(data.size(), alignof(uint64_t)));
189 std::memcpy(rawData, data.data(), data.size());
190
191 // If this is a boolean splat, make sure only the first bit is used.
192 if (key.isSplat && key.type.getElementType().isInteger(1))
193 rawData[0] &= 1;
194 copy = ArrayRef<char>(rawData, data.size());
195 }
196
197 return new (allocator.allocate<DenseIntOrFPElementsAttrStorage>())
198 DenseIntOrFPElementsAttrStorage(key.type, copy, key.isSplat);
199 }
200
201 ArrayRef<char> data;
202 };
203
204 /// An attribute representing a reference to a dense vector or tensor object
205 /// containing strings.
206 struct DenseStringElementsAttrStorage : public DenseElementsAttributeStorage {
207 DenseStringElementsAttrStorage(ShapedType ty, ArrayRef<StringRef> data,
208 bool isSplat = false)
DenseElementsAttributeStorageDenseStringElementsAttrStorage209 : DenseElementsAttributeStorage(ty, isSplat), data(data) {}
210
211 struct KeyTy {
212 KeyTy(ShapedType type, ArrayRef<StringRef> data, llvm::hash_code hashCode,
213 bool isSplat = false)
typeDenseStringElementsAttrStorage::KeyTy214 : type(type), data(data), hashCode(hashCode), isSplat(isSplat) {}
215
216 /// The type of the dense elements.
217 ShapedType type;
218
219 /// The raw buffer for the data storage.
220 ArrayRef<StringRef> data;
221
222 /// The computed hash code for the storage data.
223 llvm::hash_code hashCode;
224
225 /// A boolean that indicates if this data is a splat or not.
226 bool isSplat;
227 };
228
229 /// Compare this storage instance with the provided key.
230 bool operator==(const KeyTy &key) const {
231 if (key.type != getType())
232 return false;
233
234 // Otherwise, we can default to just checking the data. StringRefs compare
235 // by contents.
236 return key.data == data;
237 }
238
239 /// Construct a key from a shaped type, StringRef data buffer, and a flag that
240 /// signals if the data is already known to be a splat. Callers to this
241 /// function are expected to tag preknown splat values when possible, e.g. one
242 /// element shapes.
getKeyDenseStringElementsAttrStorage243 static KeyTy getKey(ShapedType ty, ArrayRef<StringRef> data,
244 bool isKnownSplat) {
245 // Handle an empty storage instance.
246 if (data.empty())
247 return KeyTy(ty, data, 0);
248
249 // If the data is already known to be a splat, the key hash value is
250 // directly the data buffer.
251 if (isKnownSplat)
252 return KeyTy(ty, data, llvm::hash_value(data.front()), isKnownSplat);
253
254 // Handle the simple case of only one element.
255 assert(ty.getNumElements() != 1 &&
256 "splat of 1 element should already be detected");
257
258 // Create the initial hash value with just the first element.
259 const auto &firstElt = data.front();
260 auto hashVal = llvm::hash_value(firstElt);
261
262 // Check to see if this storage represents a splat. If it doesn't then
263 // combine the hash for the data starting with the first non splat element.
264 for (size_t i = 1, e = data.size(); i != e; i++)
265 if (!firstElt.equals(data[i]))
266 return KeyTy(ty, data, llvm::hash_combine(hashVal, data.drop_front(i)));
267
268 // Otherwise, this is a splat so just return the hash of the first element.
269 return KeyTy(ty, data.take_front(), hashVal, /*isSplat=*/true);
270 }
271
272 /// Hash the key for the storage.
hashKeyDenseStringElementsAttrStorage273 static llvm::hash_code hashKey(const KeyTy &key) {
274 return llvm::hash_combine(key.type, key.hashCode);
275 }
276
277 /// Construct a new storage instance.
278 static DenseStringElementsAttrStorage *
constructDenseStringElementsAttrStorage279 construct(AttributeStorageAllocator &allocator, KeyTy key) {
280 // If the data buffer is non-empty, we copy it into the allocator with a
281 // 64-bit alignment.
282 ArrayRef<StringRef> copy, data = key.data;
283 if (data.empty()) {
284 return new (allocator.allocate<DenseStringElementsAttrStorage>())
285 DenseStringElementsAttrStorage(key.type, copy, key.isSplat);
286 }
287
288 int numEntries = key.isSplat ? 1 : data.size();
289
290 // Compute the amount data needed to store the ArrayRef and StringRef
291 // contents.
292 size_t dataSize = sizeof(StringRef) * numEntries;
293 for (int i = 0; i < numEntries; i++)
294 dataSize += data[i].size();
295
296 char *rawData = reinterpret_cast<char *>(
297 allocator.allocate(dataSize, alignof(uint64_t)));
298
299 // Setup a mutable array ref of our string refs so that we can update their
300 // contents.
301 auto mutableCopy = MutableArrayRef<StringRef>(
302 reinterpret_cast<StringRef *>(rawData), numEntries);
303 auto *stringData = rawData + numEntries * sizeof(StringRef);
304
305 for (int i = 0; i < numEntries; i++) {
306 memcpy(stringData, data[i].data(), data[i].size());
307 mutableCopy[i] = StringRef(stringData, data[i].size());
308 stringData += data[i].size();
309 }
310
311 copy =
312 ArrayRef<StringRef>(reinterpret_cast<StringRef *>(rawData), numEntries);
313
314 return new (allocator.allocate<DenseStringElementsAttrStorage>())
315 DenseStringElementsAttrStorage(key.type, copy, key.isSplat);
316 }
317
318 ArrayRef<StringRef> data;
319 };
320
321 //===----------------------------------------------------------------------===//
322 // StringAttr
323 //===----------------------------------------------------------------------===//
324
325 struct StringAttrStorage : public AttributeStorage {
StringAttrStorageStringAttrStorage326 StringAttrStorage(StringRef value, Type type)
327 : AttributeStorage(type), value(value), referencedDialect(nullptr) {}
328
329 /// The hash key is a tuple of the parameter types.
330 using KeyTy = std::pair<StringRef, Type>;
331 bool operator==(const KeyTy &key) const {
332 return value == key.first && getType() == key.second;
333 }
hashKeyStringAttrStorage334 static ::llvm::hash_code hashKey(const KeyTy &key) {
335 return DenseMapInfo<KeyTy>::getHashValue(key);
336 }
337
338 /// Define a construction method for creating a new instance of this
339 /// storage.
constructStringAttrStorage340 static StringAttrStorage *construct(AttributeStorageAllocator &allocator,
341 const KeyTy &key) {
342 return new (allocator.allocate<StringAttrStorage>())
343 StringAttrStorage(allocator.copyInto(key.first), key.second);
344 }
345
346 /// Initialize the storage given an MLIRContext.
347 void initialize(MLIRContext *context);
348
349 /// The raw string value.
350 StringRef value;
351 /// If the string value contains a dialect namespace prefix (e.g.
352 /// dialect.blah), this is the dialect referenced.
353 Dialect *referencedDialect;
354 };
355
356 } // namespace detail
357 } // namespace mlir
358
359 #endif // ATTRIBUTEDETAIL_H_
360