1 //===- TargetAndABI.cpp - SPIR-V target and ABI utilities -----------------===//
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 #include "mlir/Dialect/SPIRV/IR/TargetAndABI.h"
10 #include "mlir/Dialect/SPIRV/IR/SPIRVTypes.h"
11 #include "mlir/IR/Builders.h"
12 #include "mlir/IR/FunctionInterfaces.h"
13 #include "mlir/IR/Operation.h"
14 #include "mlir/IR/SymbolTable.h"
15 
16 using namespace mlir;
17 
18 //===----------------------------------------------------------------------===//
19 // TargetEnv
20 //===----------------------------------------------------------------------===//
21 
22 spirv::TargetEnv::TargetEnv(spirv::TargetEnvAttr targetAttr)
23     : targetAttr(targetAttr) {
24   for (spirv::Extension ext : targetAttr.getExtensions())
25     givenExtensions.insert(ext);
26 
27   // Add extensions implied by the current version.
28   for (spirv::Extension ext :
29        spirv::getImpliedExtensions(targetAttr.getVersion()))
30     givenExtensions.insert(ext);
31 
32   for (spirv::Capability cap : targetAttr.getCapabilities()) {
33     givenCapabilities.insert(cap);
34 
35     // Add capabilities implied by the current capability.
36     for (spirv::Capability c : spirv::getRecursiveImpliedCapabilities(cap))
37       givenCapabilities.insert(c);
38   }
39 }
40 
41 spirv::Version spirv::TargetEnv::getVersion() const {
42   return targetAttr.getVersion();
43 }
44 
45 bool spirv::TargetEnv::allows(spirv::Capability capability) const {
46   return givenCapabilities.count(capability);
47 }
48 
49 Optional<spirv::Capability>
50 spirv::TargetEnv::allows(ArrayRef<spirv::Capability> caps) const {
51   const auto *chosen = llvm::find_if(caps, [this](spirv::Capability cap) {
52     return givenCapabilities.count(cap);
53   });
54   if (chosen != caps.end())
55     return *chosen;
56   return llvm::None;
57 }
58 
59 bool spirv::TargetEnv::allows(spirv::Extension extension) const {
60   return givenExtensions.count(extension);
61 }
62 
63 Optional<spirv::Extension>
64 spirv::TargetEnv::allows(ArrayRef<spirv::Extension> exts) const {
65   const auto *chosen = llvm::find_if(exts, [this](spirv::Extension ext) {
66     return givenExtensions.count(ext);
67   });
68   if (chosen != exts.end())
69     return *chosen;
70   return llvm::None;
71 }
72 
73 spirv::Vendor spirv::TargetEnv::getVendorID() const {
74   return targetAttr.getVendorID();
75 }
76 
77 spirv::DeviceType spirv::TargetEnv::getDeviceType() const {
78   return targetAttr.getDeviceType();
79 }
80 
81 uint32_t spirv::TargetEnv::getDeviceID() const {
82   return targetAttr.getDeviceID();
83 }
84 
85 spirv::ResourceLimitsAttr spirv::TargetEnv::getResourceLimits() const {
86   return targetAttr.getResourceLimits();
87 }
88 
89 MLIRContext *spirv::TargetEnv::getContext() const {
90   return targetAttr.getContext();
91 }
92 
93 //===----------------------------------------------------------------------===//
94 // Utility functions
95 //===----------------------------------------------------------------------===//
96 
97 StringRef spirv::getInterfaceVarABIAttrName() {
98   return "spv.interface_var_abi";
99 }
100 
101 spirv::InterfaceVarABIAttr
102 spirv::getInterfaceVarABIAttr(unsigned descriptorSet, unsigned binding,
103                               Optional<spirv::StorageClass> storageClass,
104                               MLIRContext *context) {
105   return spirv::InterfaceVarABIAttr::get(descriptorSet, binding, storageClass,
106                                          context);
107 }
108 
109 bool spirv::needsInterfaceVarABIAttrs(spirv::TargetEnvAttr targetAttr) {
110   for (spirv::Capability cap : targetAttr.getCapabilities()) {
111     if (cap == spirv::Capability::Kernel)
112       return false;
113     if (cap == spirv::Capability::Shader)
114       return true;
115   }
116   return false;
117 }
118 
119 StringRef spirv::getEntryPointABIAttrName() { return "spv.entry_point_abi"; }
120 
121 spirv::EntryPointABIAttr
122 spirv::getEntryPointABIAttr(ArrayRef<int32_t> localSize, MLIRContext *context) {
123   if (localSize.empty())
124     return spirv::EntryPointABIAttr::get(nullptr, context);
125 
126   assert(localSize.size() == 3);
127   return spirv::EntryPointABIAttr::get(
128       DenseElementsAttr::get<int32_t>(
129           VectorType::get(3, IntegerType::get(context, 32)), localSize)
130           .cast<DenseIntElementsAttr>(),
131       context);
132 }
133 
134 spirv::EntryPointABIAttr spirv::lookupEntryPointABI(Operation *op) {
135   while (op && !isa<FunctionOpInterface>(op))
136     op = op->getParentOp();
137   if (!op)
138     return {};
139 
140   if (auto attr = op->getAttrOfType<spirv::EntryPointABIAttr>(
141           spirv::getEntryPointABIAttrName()))
142     return attr;
143 
144   return {};
145 }
146 
147 DenseIntElementsAttr spirv::lookupLocalWorkGroupSize(Operation *op) {
148   if (auto entryPoint = spirv::lookupEntryPointABI(op))
149     return entryPoint.local_size();
150 
151   return {};
152 }
153 
154 spirv::ResourceLimitsAttr
155 spirv::getDefaultResourceLimits(MLIRContext *context) {
156   // All the fields have default values. Here we just provide a nicer way to
157   // construct a default resource limit attribute.
158   return spirv::ResourceLimitsAttr ::get(
159       /*max_compute_shared_memory_size=*/nullptr,
160       /*max_compute_workgroup_invocations=*/nullptr,
161       /*max_compute_workgroup_size=*/nullptr,
162       /*subgroup_size=*/nullptr,
163       /*cooperative_matrix_properties_nv=*/nullptr, context);
164 }
165 
166 StringRef spirv::getTargetEnvAttrName() { return "spv.target_env"; }
167 
168 spirv::TargetEnvAttr spirv::getDefaultTargetEnv(MLIRContext *context) {
169   auto triple = spirv::VerCapExtAttr::get(spirv::Version::V_1_0,
170                                           {spirv::Capability::Shader},
171                                           ArrayRef<Extension>(), context);
172   return spirv::TargetEnvAttr::get(triple, spirv::Vendor::Unknown,
173                                    spirv::DeviceType::Unknown,
174                                    spirv::TargetEnvAttr::kUnknownDeviceID,
175                                    spirv::getDefaultResourceLimits(context));
176 }
177 
178 spirv::TargetEnvAttr spirv::lookupTargetEnv(Operation *op) {
179   while (op) {
180     op = SymbolTable::getNearestSymbolTable(op);
181     if (!op)
182       break;
183 
184     if (auto attr = op->getAttrOfType<spirv::TargetEnvAttr>(
185             spirv::getTargetEnvAttrName()))
186       return attr;
187 
188     op = op->getParentOp();
189   }
190 
191   return {};
192 }
193 
194 spirv::TargetEnvAttr spirv::lookupTargetEnvOrDefault(Operation *op) {
195   if (spirv::TargetEnvAttr attr = spirv::lookupTargetEnv(op))
196     return attr;
197 
198   return getDefaultTargetEnv(op->getContext());
199 }
200 
201 spirv::AddressingModel
202 spirv::getAddressingModel(spirv::TargetEnvAttr targetAttr) {
203   for (spirv::Capability cap : targetAttr.getCapabilities()) {
204     // TODO: Physical64 is hard-coded here, but some information should come
205     // from TargetEnvAttr to selected between Physical32 and Physical64.
206     if (cap == Capability::Kernel)
207       return spirv::AddressingModel::Physical64;
208   }
209   // Logical addressing doesn't need any capabilities so return it as default.
210   return spirv::AddressingModel::Logical;
211 }
212 
213 FailureOr<spirv::ExecutionModel>
214 spirv::getExecutionModel(spirv::TargetEnvAttr targetAttr) {
215   for (spirv::Capability cap : targetAttr.getCapabilities()) {
216     if (cap == spirv::Capability::Kernel)
217       return spirv::ExecutionModel::Kernel;
218     if (cap == spirv::Capability::Shader)
219       return spirv::ExecutionModel::GLCompute;
220   }
221   return failure();
222 }
223 
224 FailureOr<spirv::MemoryModel>
225 spirv::getMemoryModel(spirv::TargetEnvAttr targetAttr) {
226   for (spirv::Capability cap : targetAttr.getCapabilities()) {
227     if (cap == spirv::Capability::Addresses)
228       return spirv::MemoryModel::OpenCL;
229     if (cap == spirv::Capability::Shader)
230       return spirv::MemoryModel::GLSL450;
231   }
232   return failure();
233 }
234