1 //===-- LayoutUtils.cpp - Decorate composite type with layout information -===//
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 file implements Utilities used to get alignment and layout information
10 // for types in SPIR-V dialect.
11 //
12 //===----------------------------------------------------------------------===//
13 #include "mlir/Dialect/SPIRV/LayoutUtils.h"
14 #include "mlir/Dialect/SPIRV/SPIRVDialect.h"
15 #include "mlir/Dialect/SPIRV/SPIRVTypes.h"
16
17 using namespace mlir;
18
19 spirv::StructType
decorateType(spirv::StructType structType)20 VulkanLayoutUtils::decorateType(spirv::StructType structType) {
21 Size size = 0;
22 Size alignment = 1;
23 return decorateType(structType, size, alignment);
24 }
25
26 spirv::StructType
decorateType(spirv::StructType structType,VulkanLayoutUtils::Size & size,VulkanLayoutUtils::Size & alignment)27 VulkanLayoutUtils::decorateType(spirv::StructType structType,
28 VulkanLayoutUtils::Size &size,
29 VulkanLayoutUtils::Size &alignment) {
30 if (structType.getNumElements() == 0) {
31 return structType;
32 }
33
34 SmallVector<Type, 4> memberTypes;
35 SmallVector<spirv::StructType::OffsetInfo, 4> offsetInfo;
36 SmallVector<spirv::StructType::MemberDecorationInfo, 4> memberDecorations;
37
38 Size structMemberOffset = 0;
39 Size maxMemberAlignment = 1;
40
41 for (uint32_t i = 0, e = structType.getNumElements(); i < e; ++i) {
42 Size memberSize = 0;
43 Size memberAlignment = 1;
44
45 auto memberType =
46 decorateType(structType.getElementType(i), memberSize, memberAlignment);
47 structMemberOffset = llvm::alignTo(structMemberOffset, memberAlignment);
48 memberTypes.push_back(memberType);
49 offsetInfo.push_back(
50 static_cast<spirv::StructType::OffsetInfo>(structMemberOffset));
51 // If the member's size is the max value, it must be the last member and it
52 // must be a runtime array.
53 assert(memberSize != std::numeric_limits<Size>().max() ||
54 (i + 1 == e &&
55 structType.getElementType(i).isa<spirv::RuntimeArrayType>()));
56 // According to the Vulkan spec:
57 // "A structure has a base alignment equal to the largest base alignment of
58 // any of its members."
59 structMemberOffset += memberSize;
60 maxMemberAlignment = std::max(maxMemberAlignment, memberAlignment);
61 }
62
63 // According to the Vulkan spec:
64 // "The Offset decoration of a member must not place it between the end of a
65 // structure or an array and the next multiple of the alignment of that
66 // structure or array."
67 size = llvm::alignTo(structMemberOffset, maxMemberAlignment);
68 alignment = maxMemberAlignment;
69 structType.getMemberDecorations(memberDecorations);
70
71 if (!structType.isIdentified())
72 return spirv::StructType::get(memberTypes, offsetInfo, memberDecorations);
73
74 // Identified structs are uniqued by identifier so it is not possible
75 // to create 2 structs with the same name but different decorations.
76 return nullptr;
77 }
78
decorateType(Type type,VulkanLayoutUtils::Size & size,VulkanLayoutUtils::Size & alignment)79 Type VulkanLayoutUtils::decorateType(Type type, VulkanLayoutUtils::Size &size,
80 VulkanLayoutUtils::Size &alignment) {
81 if (type.isa<spirv::ScalarType>()) {
82 alignment = getScalarTypeAlignment(type);
83 // Vulkan spec does not specify any padding for a scalar type.
84 size = alignment;
85 return type;
86 }
87 if (auto structType = type.dyn_cast<spirv::StructType>())
88 return decorateType(structType, size, alignment);
89 if (auto arrayType = type.dyn_cast<spirv::ArrayType>())
90 return decorateType(arrayType, size, alignment);
91 if (auto vectorType = type.dyn_cast<VectorType>())
92 return decorateType(vectorType, size, alignment);
93 if (auto arrayType = type.dyn_cast<spirv::RuntimeArrayType>()) {
94 size = std::numeric_limits<Size>().max();
95 return decorateType(arrayType, alignment);
96 }
97 llvm_unreachable("unhandled SPIR-V type");
98 }
99
decorateType(VectorType vectorType,VulkanLayoutUtils::Size & size,VulkanLayoutUtils::Size & alignment)100 Type VulkanLayoutUtils::decorateType(VectorType vectorType,
101 VulkanLayoutUtils::Size &size,
102 VulkanLayoutUtils::Size &alignment) {
103 const auto numElements = vectorType.getNumElements();
104 auto elementType = vectorType.getElementType();
105 Size elementSize = 0;
106 Size elementAlignment = 1;
107
108 auto memberType = decorateType(elementType, elementSize, elementAlignment);
109 // According to the Vulkan spec:
110 // 1. "A two-component vector has a base alignment equal to twice its scalar
111 // alignment."
112 // 2. "A three- or four-component vector has a base alignment equal to four
113 // times its scalar alignment."
114 size = elementSize * numElements;
115 alignment = numElements == 2 ? elementAlignment * 2 : elementAlignment * 4;
116 return VectorType::get(numElements, memberType);
117 }
118
decorateType(spirv::ArrayType arrayType,VulkanLayoutUtils::Size & size,VulkanLayoutUtils::Size & alignment)119 Type VulkanLayoutUtils::decorateType(spirv::ArrayType arrayType,
120 VulkanLayoutUtils::Size &size,
121 VulkanLayoutUtils::Size &alignment) {
122 const auto numElements = arrayType.getNumElements();
123 auto elementType = arrayType.getElementType();
124 Size elementSize = 0;
125 Size elementAlignment = 1;
126
127 auto memberType = decorateType(elementType, elementSize, elementAlignment);
128 // According to the Vulkan spec:
129 // "An array has a base alignment equal to the base alignment of its element
130 // type."
131 size = elementSize * numElements;
132 alignment = elementAlignment;
133 return spirv::ArrayType::get(memberType, numElements, elementSize);
134 }
135
decorateType(spirv::RuntimeArrayType arrayType,VulkanLayoutUtils::Size & alignment)136 Type VulkanLayoutUtils::decorateType(spirv::RuntimeArrayType arrayType,
137 VulkanLayoutUtils::Size &alignment) {
138 auto elementType = arrayType.getElementType();
139 Size elementSize = 0;
140
141 auto memberType = decorateType(elementType, elementSize, alignment);
142 return spirv::RuntimeArrayType::get(memberType, elementSize);
143 }
144
145 VulkanLayoutUtils::Size
getScalarTypeAlignment(Type scalarType)146 VulkanLayoutUtils::getScalarTypeAlignment(Type scalarType) {
147 // According to the Vulkan spec:
148 // 1. "A scalar of size N has a scalar alignment of N."
149 // 2. "A scalar has a base alignment equal to its scalar alignment."
150 // 3. "A scalar, vector or matrix type has an extended alignment equal to its
151 // base alignment."
152 auto bitWidth = scalarType.getIntOrFloatBitWidth();
153 if (bitWidth == 1)
154 return 1;
155 return bitWidth / 8;
156 }
157
isLegalType(Type type)158 bool VulkanLayoutUtils::isLegalType(Type type) {
159 auto ptrType = type.dyn_cast<spirv::PointerType>();
160 if (!ptrType) {
161 return true;
162 }
163
164 auto storageClass = ptrType.getStorageClass();
165 auto structType = ptrType.getPointeeType().dyn_cast<spirv::StructType>();
166 if (!structType) {
167 return true;
168 }
169
170 switch (storageClass) {
171 case spirv::StorageClass::Uniform:
172 case spirv::StorageClass::StorageBuffer:
173 case spirv::StorageClass::PushConstant:
174 case spirv::StorageClass::PhysicalStorageBuffer:
175 return structType.hasOffset() || !structType.getNumElements();
176 default:
177 return true;
178 }
179 }
180