| //===- TypeDetail.h - MLIR Type storage details -----------------*- C++ -*-===// |
| // |
| // Part of the MLIR Project, under the Apache License v2.0 with LLVM Exceptions. |
| // See https://llvm.org/LICENSE.txt for license information. |
| // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception |
| // |
| //===----------------------------------------------------------------------===// |
| // |
| // This holds implementation details of Type. |
| // |
| //===----------------------------------------------------------------------===// |
| #ifndef TYPEDETAIL_H_ |
| #define TYPEDETAIL_H_ |
| |
| #include "mlir/IR/AffineMap.h" |
| #include "mlir/IR/Identifier.h" |
| #include "mlir/IR/MLIRContext.h" |
| #include "mlir/IR/TypeSupport.h" |
| #include "mlir/IR/Types.h" |
| #include "llvm/Support/TrailingObjects.h" |
| |
| namespace mlir { |
| |
| class MLIRContext; |
| |
| namespace detail { |
| |
| /// Opaque Type Storage and Uniquing. |
| struct OpaqueTypeStorage : public TypeStorage { |
| OpaqueTypeStorage(Identifier dialectNamespace, StringRef typeData) |
| : dialectNamespace(dialectNamespace), typeData(typeData) {} |
| |
| /// The hash key used for uniquing. |
| using KeyTy = std::pair<Identifier, StringRef>; |
| bool operator==(const KeyTy &key) const { |
| return key == KeyTy(dialectNamespace, typeData); |
| } |
| |
| static OpaqueTypeStorage *construct(TypeStorageAllocator &allocator, |
| const KeyTy &key) { |
| StringRef tyData = allocator.copyInto(key.second); |
| return new (allocator.allocate<OpaqueTypeStorage>()) |
| OpaqueTypeStorage(key.first, tyData); |
| } |
| |
| // The dialect namespace. |
| Identifier dialectNamespace; |
| |
| // The parser type data for this opaque type. |
| StringRef typeData; |
| }; |
| |
| /// Integer Type Storage and Uniquing. |
| struct IntegerTypeStorage : public TypeStorage { |
| IntegerTypeStorage(unsigned width) : width(width) {} |
| |
| /// The hash key used for uniquing. |
| using KeyTy = unsigned; |
| bool operator==(const KeyTy &key) const { return key == width; } |
| |
| static IntegerTypeStorage *construct(TypeStorageAllocator &allocator, |
| KeyTy bitwidth) { |
| return new (allocator.allocate<IntegerTypeStorage>()) |
| IntegerTypeStorage(bitwidth); |
| } |
| |
| unsigned width; |
| }; |
| |
| /// Function Type Storage and Uniquing. |
| struct FunctionTypeStorage : public TypeStorage { |
| FunctionTypeStorage(unsigned numInputs, unsigned numResults, |
| Type const *inputsAndResults) |
| : TypeStorage(numInputs), numResults(numResults), |
| inputsAndResults(inputsAndResults) {} |
| |
| /// The hash key used for uniquing. |
| using KeyTy = std::pair<ArrayRef<Type>, ArrayRef<Type>>; |
| bool operator==(const KeyTy &key) const { |
| return key == KeyTy(getInputs(), getResults()); |
| } |
| |
| /// Construction. |
| static FunctionTypeStorage *construct(TypeStorageAllocator &allocator, |
| const KeyTy &key) { |
| ArrayRef<Type> inputs = key.first, results = key.second; |
| |
| // Copy the inputs and results into the bump pointer. |
| SmallVector<Type, 16> types; |
| types.reserve(inputs.size() + results.size()); |
| types.append(inputs.begin(), inputs.end()); |
| types.append(results.begin(), results.end()); |
| auto typesList = allocator.copyInto(ArrayRef<Type>(types)); |
| |
| // Initialize the memory using placement new. |
| return new (allocator.allocate<FunctionTypeStorage>()) |
| FunctionTypeStorage(inputs.size(), results.size(), typesList.data()); |
| } |
| |
| ArrayRef<Type> getInputs() const { |
| return ArrayRef<Type>(inputsAndResults, getSubclassData()); |
| } |
| ArrayRef<Type> getResults() const { |
| return ArrayRef<Type>(inputsAndResults + getSubclassData(), numResults); |
| } |
| |
| unsigned numResults; |
| Type const *inputsAndResults; |
| }; |
| |
| /// Shaped Type Storage. |
| struct ShapedTypeStorage : public TypeStorage { |
| ShapedTypeStorage(Type elementTy, unsigned subclassData = 0) |
| : TypeStorage(subclassData), elementType(elementTy) {} |
| |
| /// The hash key used for uniquing. |
| using KeyTy = Type; |
| bool operator==(const KeyTy &key) const { return key == elementType; } |
| |
| Type elementType; |
| }; |
| |
| /// Vector Type Storage and Uniquing. |
| struct VectorTypeStorage : public ShapedTypeStorage { |
| VectorTypeStorage(unsigned shapeSize, Type elementTy, |
| const int64_t *shapeElements) |
| : ShapedTypeStorage(elementTy, shapeSize), shapeElements(shapeElements) {} |
| |
| /// The hash key used for uniquing. |
| using KeyTy = std::pair<ArrayRef<int64_t>, Type>; |
| bool operator==(const KeyTy &key) const { |
| return key == KeyTy(getShape(), elementType); |
| } |
| |
| /// Construction. |
| static VectorTypeStorage *construct(TypeStorageAllocator &allocator, |
| const KeyTy &key) { |
| // Copy the shape into the bump pointer. |
| ArrayRef<int64_t> shape = allocator.copyInto(key.first); |
| |
| // Initialize the memory using placement new. |
| return new (allocator.allocate<VectorTypeStorage>()) |
| VectorTypeStorage(shape.size(), key.second, shape.data()); |
| } |
| |
| ArrayRef<int64_t> getShape() const { |
| return ArrayRef<int64_t>(shapeElements, getSubclassData()); |
| } |
| |
| const int64_t *shapeElements; |
| }; |
| |
| struct RankedTensorTypeStorage : public ShapedTypeStorage { |
| RankedTensorTypeStorage(unsigned shapeSize, Type elementTy, |
| const int64_t *shapeElements) |
| : ShapedTypeStorage(elementTy, shapeSize), shapeElements(shapeElements) {} |
| |
| /// The hash key used for uniquing. |
| using KeyTy = std::pair<ArrayRef<int64_t>, Type>; |
| bool operator==(const KeyTy &key) const { |
| return key == KeyTy(getShape(), elementType); |
| } |
| |
| /// Construction. |
| static RankedTensorTypeStorage *construct(TypeStorageAllocator &allocator, |
| const KeyTy &key) { |
| // Copy the shape into the bump pointer. |
| ArrayRef<int64_t> shape = allocator.copyInto(key.first); |
| |
| // Initialize the memory using placement new. |
| return new (allocator.allocate<RankedTensorTypeStorage>()) |
| RankedTensorTypeStorage(shape.size(), key.second, shape.data()); |
| } |
| |
| ArrayRef<int64_t> getShape() const { |
| return ArrayRef<int64_t>(shapeElements, getSubclassData()); |
| } |
| |
| const int64_t *shapeElements; |
| }; |
| |
| struct UnrankedTensorTypeStorage : public ShapedTypeStorage { |
| using ShapedTypeStorage::KeyTy; |
| using ShapedTypeStorage::ShapedTypeStorage; |
| |
| /// Construction. |
| static UnrankedTensorTypeStorage *construct(TypeStorageAllocator &allocator, |
| Type elementTy) { |
| return new (allocator.allocate<UnrankedTensorTypeStorage>()) |
| UnrankedTensorTypeStorage(elementTy); |
| } |
| }; |
| |
| struct MemRefTypeStorage : public ShapedTypeStorage { |
| MemRefTypeStorage(unsigned shapeSize, Type elementType, |
| const int64_t *shapeElements, const unsigned numAffineMaps, |
| AffineMap const *affineMapList, const unsigned memorySpace) |
| : ShapedTypeStorage(elementType, shapeSize), shapeElements(shapeElements), |
| numAffineMaps(numAffineMaps), affineMapList(affineMapList), |
| memorySpace(memorySpace) {} |
| |
| /// The hash key used for uniquing. |
| // MemRefs are uniqued based on their shape, element type, affine map |
| // composition, and memory space. |
| using KeyTy = |
| std::tuple<ArrayRef<int64_t>, Type, ArrayRef<AffineMap>, unsigned>; |
| bool operator==(const KeyTy &key) const { |
| return key == KeyTy(getShape(), elementType, getAffineMaps(), memorySpace); |
| } |
| |
| /// Construction. |
| static MemRefTypeStorage *construct(TypeStorageAllocator &allocator, |
| const KeyTy &key) { |
| // Copy the shape into the bump pointer. |
| ArrayRef<int64_t> shape = allocator.copyInto(std::get<0>(key)); |
| |
| // Copy the affine map composition into the bump pointer. |
| ArrayRef<AffineMap> affineMapComposition = |
| allocator.copyInto(std::get<2>(key)); |
| |
| // Initialize the memory using placement new. |
| return new (allocator.allocate<MemRefTypeStorage>()) |
| MemRefTypeStorage(shape.size(), std::get<1>(key), shape.data(), |
| affineMapComposition.size(), |
| affineMapComposition.data(), std::get<3>(key)); |
| } |
| |
| ArrayRef<int64_t> getShape() const { |
| return ArrayRef<int64_t>(shapeElements, getSubclassData()); |
| } |
| |
| ArrayRef<AffineMap> getAffineMaps() const { |
| return ArrayRef<AffineMap>(affineMapList, numAffineMaps); |
| } |
| |
| /// An array of integers which stores the shape dimension sizes. |
| const int64_t *shapeElements; |
| /// The number of affine maps in the 'affineMapList' array. |
| const unsigned numAffineMaps; |
| /// List of affine maps in the memref's layout/index map composition. |
| AffineMap const *affineMapList; |
| /// Memory space in which data referenced by memref resides. |
| const unsigned memorySpace; |
| }; |
| |
| /// Unranked MemRef is a MemRef with unknown rank. |
| /// Only element type and memory space are known |
| struct UnrankedMemRefTypeStorage : public ShapedTypeStorage { |
| |
| UnrankedMemRefTypeStorage(Type elementTy, const unsigned memorySpace) |
| : ShapedTypeStorage(elementTy), memorySpace(memorySpace) {} |
| |
| /// The hash key used for uniquing. |
| using KeyTy = std::tuple<Type, unsigned>; |
| bool operator==(const KeyTy &key) const { |
| return key == KeyTy(elementType, memorySpace); |
| } |
| |
| /// Construction. |
| static UnrankedMemRefTypeStorage *construct(TypeStorageAllocator &allocator, |
| const KeyTy &key) { |
| |
| // Initialize the memory using placement new. |
| return new (allocator.allocate<UnrankedMemRefTypeStorage>()) |
| UnrankedMemRefTypeStorage(std::get<0>(key), std::get<1>(key)); |
| } |
| /// Memory space in which data referenced by memref resides. |
| const unsigned memorySpace; |
| }; |
| |
| /// Complex Type Storage. |
| struct ComplexTypeStorage : public TypeStorage { |
| ComplexTypeStorage(Type elementType) : elementType(elementType) {} |
| |
| /// The hash key used for uniquing. |
| using KeyTy = Type; |
| bool operator==(const KeyTy &key) const { return key == elementType; } |
| |
| /// Construction. |
| static ComplexTypeStorage *construct(TypeStorageAllocator &allocator, |
| Type elementType) { |
| return new (allocator.allocate<ComplexTypeStorage>()) |
| ComplexTypeStorage(elementType); |
| } |
| |
| Type elementType; |
| }; |
| |
| /// A type representing a collection of other types. |
| struct TupleTypeStorage final |
| : public TypeStorage, |
| public llvm::TrailingObjects<TupleTypeStorage, Type> { |
| using KeyTy = ArrayRef<Type>; |
| |
| TupleTypeStorage(unsigned numTypes) : TypeStorage(numTypes) {} |
| |
| /// Construction. |
| static TupleTypeStorage *construct(TypeStorageAllocator &allocator, |
| ArrayRef<Type> key) { |
| // Allocate a new storage instance. |
| auto byteSize = TupleTypeStorage::totalSizeToAlloc<Type>(key.size()); |
| auto rawMem = allocator.allocate(byteSize, alignof(TupleTypeStorage)); |
| auto result = ::new (rawMem) TupleTypeStorage(key.size()); |
| |
| // Copy in the element types into the trailing storage. |
| std::uninitialized_copy(key.begin(), key.end(), |
| result->getTrailingObjects<Type>()); |
| return result; |
| } |
| |
| bool operator==(const KeyTy &key) const { return key == getTypes(); } |
| |
| /// Return the number of held types. |
| unsigned size() const { return getSubclassData(); } |
| |
| /// Return the held types. |
| ArrayRef<Type> getTypes() const { |
| return {getTrailingObjects<Type>(), size()}; |
| } |
| }; |
| |
| } // namespace detail |
| } // namespace mlir |
| #endif // TYPEDETAIL_H_ |