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