426 lines · c
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-exception6//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 26namespace mlir {27namespace detail {28 29//===----------------------------------------------------------------------===//30// Elements Attributes31//===----------------------------------------------------------------------===//32 33/// Return the bit width which DenseElementsAttr should use for this type.34inline 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>(eltType))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.44struct DenseElementsAttributeStorage : public AttributeStorage {45public: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.54struct 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 that83 /// signals if the data is already known to be a splat. Callers to this84 /// function are expected to tag preknown splat values when possible, e.g. one85 /// 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 is92 // directly the data buffer.93 bool isBoolData = ty.getElementType().isInteger(1);94 if (isKnownSplat) {95 if (isBoolData)96 return getKeyForSplatBoolData(ty, data[0] != 0);97 return KeyTy(ty, data, llvm::hash_value(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(ty.getElementType());111 // Non 1-bit dense elements are padded to 8-bits.112 size_t storageSize = llvm::divideCeil(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(storageSize);118 auto hashVal = llvm::hash_value(firstElt);119 120 // Check to see if this storage represents a splat. If it doesn't then121 // 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(data.data(), &data[i], storageSize))124 return KeyTy(ty, data, llvm::hash_combine(hashVal, data.drop_front(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(ty, splatValue);139 140 // Handle the case where the potential splat value is 1 and the number of141 // 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>(numOddElements))147 return KeyTy(ty, data, llvm::hash_value(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(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(splatData, [mask](char c) { return c == mask; })158 ? getKeyForSplatBoolData(ty, splatValue)159 : KeyTy(ty, data, llvm::hash_value(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(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(key.type, 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 a178 // 64-bit alignment.179 ArrayRef<char> copy, data = key.data;180 if (!data.empty()) {181 char *rawData = reinterpret_cast<char *>(182 allocator.allocate(data.size(), alignof(uint64_t)));183 std::memcpy(rawData, data.data(), 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 failure195 // encountered with MSVC where it would inline these values, which makes it196 // 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 object202/// containing strings.203struct 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 compare232 // by contents.233 return key.data == data;234 }235 236 /// Construct a key from a shaped type, StringRef data buffer, and a flag that237 /// signals if the data is already known to be a splat. Callers to this238 /// function are expected to tag preknown splat values when possible, e.g. one239 /// 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 is247 // directly the data buffer.248 if (isKnownSplat)249 return KeyTy(ty, data, llvm::hash_value(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(firstElt);258 259 // Check to see if this storage represents a splat. If it doesn't then260 // 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(hashVal, data.drop_front(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(key.type, 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 a278 // 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 StringRef288 // 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(dataSize, alignof(uint64_t)));295 296 // Setup a mutable array ref of our string refs so that we can update their297 // 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(stringData, data[i].data(), 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// StringAttr320//===----------------------------------------------------------------------===//321 322struct 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(key);333 }334 335 /// Define a construction method for creating a new instance of this336 /// storage.337 static StringAttrStorage *construct(AttributeStorageAllocator &allocator,338 const KeyTy &key) {339 return new (allocator.allocate<StringAttrStorage>())340 StringAttrStorage(allocator.copyInto(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// DistinctAttr357//===----------------------------------------------------------------------===//358 359/// An attribute to store a distinct reference to another attribute.360struct 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 always374/// allocates since the distinct attribute instances use the address of their375/// storage as unique identifier.376class DistinctAttributeUniquer {377public:378 /// Creates a distinct attribute storage. Allocates every time since the379 /// 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, std::forward<Args>(args)...);386 storage->initializeAbstractAttribute(387 AbstractAttribute::lookup(DistinctAttr::getTypeID(), context));388 return storage;389 }390 391private: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 synchronized398/// BumpPtrAllocator to ensure thread-safety. The allocated storage is deleted399/// when the DistinctAttributeAllocator is destroyed.400class DistinctAttributeAllocator final {401public: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 414private:415 /// Used to allocate distict attribute storages. The managed memory is freed416 /// 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 detail423} // namespace mlir424 425#endif // ATTRIBUTEDETAIL_H_426