| 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 | |
| 28 | namespace mlir { |
| 29 | namespace detail { |
| 30 | |
| 31 | //===----------------------------------------------------------------------===// |
| 32 | // Elements Attributes |
| 33 | //===----------------------------------------------------------------------===// |
| 34 | |
| 35 | /// Return the bit width which DenseElementsAttr should use for this type. |
| 36 | inline 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. |
| 46 | struct DenseElementsAttributeStorage : public AttributeStorage { |
| 47 | public: |
| 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. |
| 56 | struct 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. |
| 205 | struct 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 != 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 | |
| 324 | struct 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. |
| 362 | struct 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. |
| 378 | class DistinctAttributeUniquer { |
| 379 | public: |
| 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 | |
| 393 | private: |
| 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. |
| 402 | class DistinctAttributeAllocator { |
| 403 | public: |
| 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 | |
| 418 | private: |
| 419 | ThreadLocalCache<llvm::BumpPtrAllocator> allocatorCache; |
| 420 | }; |
| 421 | } // namespace detail |
| 422 | } // namespace mlir |
| 423 | |
| 424 | #endif // ATTRIBUTEDETAIL_H_ |
| 425 | |