1 // Copyright 2019 The Amber Authors.
2 //
3 // Licensed under the Apache License, Version 2.0 (the "License");
4 // you may not use this file except in compliance with the License.
5 // You may obtain a copy of the License at
6 //
7 //     http://www.apache.org/licenses/LICENSE-2.0
8 //
9 // Unless required by applicable law or agreed to in writing, software
10 // distributed under the License is distributed on an "AS IS" BASIS,
11 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 // See the License for the specific language governing permissions and
13 // limitations under the License.
14 
15 #include "src/clspv_helper.h"
16 
17 #include <unordered_map>
18 #include <utility>
19 
20 #include "clspv/ArgKind.h"
21 #include "clspv/Compiler.h"
22 #include "clspv/Sampler.h"
23 #include "spirv-tools/libspirv.hpp"
24 #include "spirv-tools/optimizer.hpp"
25 #include "spirv/unified1/NonSemanticClspvReflection.h"
26 #include "spirv/unified1/spirv.hpp"
27 
28 using amber::Pipeline;
29 
30 namespace {
31 
32 struct ReflectionHelper {
33   Pipeline::ShaderInfo* shader_info = nullptr;
34   Pipeline* pipeline = nullptr;
35   uint32_t uint_id = 0;
36   std::unordered_map<uint32_t, std::string> strings;
37   std::unordered_map<uint32_t, uint32_t> constants;
38   std::string error_message;
39 };
40 
GetArgKindFromExtInst(uint32_t value)41 Pipeline::ShaderInfo::DescriptorMapEntry::Kind GetArgKindFromExtInst(
42     uint32_t value) {
43   switch (static_cast<NonSemanticClspvReflectionInstructions>(value)) {
44     case NonSemanticClspvReflectionArgumentStorageBuffer:
45     case NonSemanticClspvReflectionConstantDataStorageBuffer:
46       return Pipeline::ShaderInfo::DescriptorMapEntry::Kind::SSBO;
47     case NonSemanticClspvReflectionArgumentUniform:
48     case NonSemanticClspvReflectionConstantDataUniform:
49       return Pipeline::ShaderInfo::DescriptorMapEntry::Kind::UBO;
50     case NonSemanticClspvReflectionArgumentPodStorageBuffer:
51       return Pipeline::ShaderInfo::DescriptorMapEntry::Kind::POD;
52     case NonSemanticClspvReflectionArgumentPodUniform:
53       return Pipeline::ShaderInfo::DescriptorMapEntry::Kind::POD_UBO;
54     case NonSemanticClspvReflectionArgumentPodPushConstant:
55       return Pipeline::ShaderInfo::DescriptorMapEntry::Kind::POD_PUSHCONSTANT;
56     case NonSemanticClspvReflectionArgumentSampledImage:
57       return Pipeline::ShaderInfo::DescriptorMapEntry::Kind::RO_IMAGE;
58     case NonSemanticClspvReflectionArgumentStorageImage:
59       return Pipeline::ShaderInfo::DescriptorMapEntry::Kind::WO_IMAGE;
60     case NonSemanticClspvReflectionArgumentSampler:
61       return Pipeline::ShaderInfo::DescriptorMapEntry::Kind::SAMPLER;
62     case NonSemanticClspvReflectionArgumentWorkgroup:
63     default:
64       return Pipeline::ShaderInfo::DescriptorMapEntry::Kind::SSBO;
65   }
66 }
67 
ParseExtendedInst(ReflectionHelper * helper,const spv_parsed_instruction_t * inst)68 spv_result_t ParseExtendedInst(ReflectionHelper* helper,
69                                const spv_parsed_instruction_t* inst) {
70   auto ext_inst = inst->words[inst->operands[3].offset];
71   switch (ext_inst) {
72     case NonSemanticClspvReflectionKernel: {
73       // Remap the name string to the declaration's result id.
74       const auto& name = helper->strings[inst->words[inst->operands[5].offset]];
75       helper->strings[inst->result_id] = name;
76       break;
77     }
78     case NonSemanticClspvReflectionArgumentInfo: {
79       // Remap the name string to the info's result id.
80       const auto& name = helper->strings[inst->words[inst->operands[4].offset]];
81       helper->strings[inst->result_id] = name;
82       break;
83     }
84     case NonSemanticClspvReflectionArgumentStorageBuffer:
85     case NonSemanticClspvReflectionArgumentUniform:
86     case NonSemanticClspvReflectionArgumentSampledImage:
87     case NonSemanticClspvReflectionArgumentStorageImage:
88     case NonSemanticClspvReflectionArgumentSampler: {
89       // These arguments have descriptor set and binding.
90       auto kernel_id = inst->words[inst->operands[4].offset];
91       auto ordinal_id = inst->words[inst->operands[5].offset];
92       auto ds_id = inst->words[inst->operands[6].offset];
93       auto binding_id = inst->words[inst->operands[7].offset];
94       std::string arg_name;
95       if (inst->num_operands == 9) {
96         arg_name = helper->strings[inst->words[inst->operands[8].offset]];
97       }
98       auto kind = GetArgKindFromExtInst(ext_inst);
99       Pipeline::ShaderInfo::DescriptorMapEntry entry{
100           arg_name,
101           kind,
102           helper->constants[ds_id],
103           helper->constants[binding_id],
104           helper->constants[ordinal_id],
105           /* offset */ 0,
106           /* size */ 0};
107       helper->shader_info->AddDescriptorEntry(helper->strings[kernel_id],
108                                               std::move(entry));
109       break;
110     }
111     case NonSemanticClspvReflectionArgumentPodStorageBuffer:
112     case NonSemanticClspvReflectionArgumentPodUniform: {
113       // These arguments have descriptor set, binding, offset and size.
114       auto kernel_id = inst->words[inst->operands[4].offset];
115       auto ordinal_id = inst->words[inst->operands[5].offset];
116       auto ds_id = inst->words[inst->operands[6].offset];
117       auto binding_id = inst->words[inst->operands[7].offset];
118       auto offset_id = inst->words[inst->operands[8].offset];
119       auto size_id = inst->words[inst->operands[9].offset];
120       std::string arg_name;
121       if (inst->num_operands == 11) {
122         arg_name = helper->strings[inst->words[inst->operands[10].offset]];
123       }
124       auto kind = GetArgKindFromExtInst(ext_inst);
125       Pipeline::ShaderInfo::DescriptorMapEntry entry{
126           arg_name,
127           kind,
128           helper->constants[ds_id],
129           helper->constants[binding_id],
130           helper->constants[ordinal_id],
131           helper->constants[offset_id],
132           helper->constants[size_id]};
133       helper->shader_info->AddDescriptorEntry(helper->strings[kernel_id],
134                                               std::move(entry));
135       break;
136     }
137     case NonSemanticClspvReflectionArgumentPodPushConstant: {
138       // These arguments have offset and size.
139       auto kernel_id = inst->words[inst->operands[4].offset];
140       auto ordinal_id = inst->words[inst->operands[5].offset];
141       auto offset_id = inst->words[inst->operands[6].offset];
142       auto size_id = inst->words[inst->operands[7].offset];
143       std::string arg_name;
144       if (inst->num_operands == 9) {
145         arg_name = helper->strings[inst->words[inst->operands[8].offset]];
146       }
147       auto kind = GetArgKindFromExtInst(ext_inst);
148       Pipeline::ShaderInfo::DescriptorMapEntry entry{
149           arg_name,
150           kind,
151           /* descriptor set */ 0,
152           /* binding */ 0,
153           helper->constants[ordinal_id],
154           helper->constants[offset_id],
155           helper->constants[size_id]};
156       helper->shader_info->AddDescriptorEntry(helper->strings[kernel_id],
157                                               std::move(entry));
158       break;
159     }
160     case NonSemanticClspvReflectionArgumentWorkgroup:
161       helper->error_message = "Workgroup arguments are not currently supported";
162       return SPV_ERROR_INVALID_DATA;
163     case NonSemanticClspvReflectionConstantDataStorageBuffer:
164     case NonSemanticClspvReflectionConstantDataUniform:
165       helper->error_message =
166           "Constant descriptor entries are not currently supported";
167       return SPV_ERROR_INVALID_DATA;
168     case NonSemanticClspvReflectionSpecConstantWorkgroupSize:
169     case NonSemanticClspvReflectionSpecConstantGlobalOffset:
170     case NonSemanticClspvReflectionSpecConstantWorkDim:
171       // Nothing to do. Amber currently requires script authors to know the
172       // spec ids and use them directly.
173       break;
174     case NonSemanticClspvReflectionPushConstantGlobalOffset: {
175       auto offset_id = inst->words[inst->operands[4].offset];
176       auto size_id = inst->words[inst->operands[5].offset];
177       Pipeline::ShaderInfo::PushConstant push_constant{
178           Pipeline::ShaderInfo::PushConstant::PushConstantType::kGlobalOffset,
179           helper->constants[offset_id], helper->constants[size_id]};
180       helper->shader_info->AddPushConstant(std::move(push_constant));
181       break;
182     }
183     case NonSemanticClspvReflectionPushConstantRegionOffset: {
184       auto offset_id = inst->words[inst->operands[4].offset];
185       auto size_id = inst->words[inst->operands[5].offset];
186       Pipeline::ShaderInfo::PushConstant push_constant{
187           Pipeline::ShaderInfo::PushConstant::PushConstantType::kRegionOffset,
188           helper->constants[offset_id], helper->constants[size_id]};
189       helper->shader_info->AddPushConstant(std::move(push_constant));
190       break;
191     }
192     case NonSemanticClspvReflectionPushConstantEnqueuedLocalSize:
193     case NonSemanticClspvReflectionPushConstantGlobalSize:
194     case NonSemanticClspvReflectionPushConstantNumWorkgroups:
195     case NonSemanticClspvReflectionPushConstantRegionGroupOffset:
196       helper->error_message = "Unsupported push constant";
197       return SPV_ERROR_INVALID_DATA;
198     case NonSemanticClspvReflectionLiteralSampler: {
199       auto ds_id = inst->words[inst->operands[4].offset];
200       auto binding_id = inst->words[inst->operands[5].offset];
201       auto mask_id = inst->words[inst->operands[6].offset];
202       helper->pipeline->AddSampler(helper->constants[mask_id],
203                                    helper->constants[ds_id],
204                                    helper->constants[binding_id]);
205       break;
206     } break;
207   }
208 
209   return SPV_SUCCESS;
210 }
211 
ParseReflection(void * user_data,const spv_parsed_instruction_t * inst)212 spv_result_t ParseReflection(void* user_data,
213                              const spv_parsed_instruction_t* inst) {
214   auto* helper = reinterpret_cast<ReflectionHelper*>(user_data);
215   switch (inst->opcode) {
216     case spv::OpTypeInt:
217       if (inst->words[inst->operands[1].offset] == 32 &&
218           inst->words[inst->operands[2].offset] == 0) {
219         // Track the result id of OpTypeInt 32 0.
220         helper->uint_id = inst->result_id;
221       }
222       break;
223     case spv::OpConstant:
224       if (inst->words[inst->operands[0].offset] == helper->uint_id) {
225         // Record the values for all uint32_t constants.
226         uint32_t value = inst->words[inst->operands[2].offset];
227         helper->constants[inst->result_id] = value;
228       }
229       break;
230     case spv::OpString: {
231       // Track all strings.
232       std::string value =
233           reinterpret_cast<const char*>(inst->words + inst->operands[1].offset);
234       helper->strings[inst->result_id] = value;
235       break;
236     }
237     case spv::OpExtInst:
238       if (inst->ext_inst_type ==
239           SPV_EXT_INST_TYPE_NONSEMANTIC_CLSPVREFLECTION) {
240         return ParseExtendedInst(helper, inst);
241       }
242       break;
243   }
244   return SPV_SUCCESS;
245 }
246 
247 }  // anonymous namespace
248 
249 namespace amber {
250 namespace clspvhelper {
251 
Compile(Pipeline::ShaderInfo * shader_info,Pipeline * pipeline,spv_target_env env,std::vector<uint32_t> * generated_binary)252 Result Compile(Pipeline::ShaderInfo* shader_info,
253                Pipeline* pipeline,
254                spv_target_env env,
255                std::vector<uint32_t>* generated_binary) {
256   const auto& src_str = shader_info->GetShader()->GetData();
257   std::string options;
258   for (const auto& option : shader_info->GetCompileOptions()) {
259     options += option + " ";
260   }
261   if (clspv::CompileFromSourceString(src_str, /* sampler map */ "", options,
262                                      generated_binary)) {
263     return Result("Clspv compile failed");
264   }
265 
266   // Parse the reflection instructions.
267   ReflectionHelper helper;
268   helper.shader_info = shader_info;
269   helper.pipeline = pipeline;
270   spv_context context(spvContextCreate(env));
271   if (spvBinaryParse(context, &helper, generated_binary->data(),
272                      generated_binary->size(), nullptr, ParseReflection,
273                      nullptr)) {
274     return Result(helper.error_message);
275   }
276 
277   // Strip the reflection instructions to avoid requiring the implementation to
278   // support VK_KHR_shader_non_semantic_info.
279   spvtools::Optimizer opt(env);
280   opt.RegisterPass(spvtools::CreateStripReflectInfoPass());
281   std::vector<uint32_t> stripped;
282   if (!opt.Run(generated_binary->data(), generated_binary->size(), &stripped)) {
283     return Result("failed to strip reflection instructions");
284   }
285   generated_binary->swap(stripped);
286 
287   return Result();
288 }
289 
290 }  // namespace clspvhelper
291 }  // namespace amber
292