1 //
2 // Copyright 2020 Pixar
3 //
4 // Licensed under the Apache License, Version 2.0 (the "Apache License")
5 // with the following modification; you may not use this file except in
6 // compliance with the Apache License and the following modification to it:
7 // Section 6. Trademarks. is deleted and replaced with:
8 //
9 // 6. Trademarks. This License does not grant permission to use the trade
10 // names, trademarks, service marks, or product names of the Licensor
11 // and its affiliates, except as required to comply with Section 4(c) of
12 // the License and to reproduce the content of the NOTICE file.
13 //
14 // You may obtain a copy of the Apache License at
15 //
16 // http://www.apache.org/licenses/LICENSE-2.0
17 //
18 // Unless required by applicable law or agreed to in writing, software
19 // distributed under the Apache License with the above modification is
20 // distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
21 // KIND, either express or implied. See the Apache License for the specific
22 // language governing permissions and limitations under the Apache License.
23 //
24 #include "pxr/base/tf/diagnostic.h"
25 #include "pxr/imaging/hgiVulkan/api.h"
26 #include "pxr/imaging/hgiVulkan/device.h"
27 #include "pxr/imaging/hgiVulkan/diagnostic.h"
28 #include "pxr/imaging/hgiVulkan/shaderCompiler.h"
29 #include "pxr/imaging/hgiVulkan/spirv_reflect.h"
30
31 #include <shaderc/shaderc.hpp>
32
33 #include <unordered_map>
34
35
36 PXR_NAMESPACE_OPEN_SCOPE
37
38
39 static shaderc_shader_kind
_GetShaderStage(HgiShaderStage stage)40 _GetShaderStage(HgiShaderStage stage)
41 {
42 switch(stage) {
43 case HgiShaderStageVertex:
44 return shaderc_glsl_vertex_shader;
45 case HgiShaderStageTessellationControl:
46 return shaderc_glsl_tess_control_shader;
47 case HgiShaderStageTessellationEval:
48 return shaderc_glsl_tess_evaluation_shader;
49 case HgiShaderStageGeometry:
50 return shaderc_glsl_geometry_shader;
51 case HgiShaderStageFragment:
52 return shaderc_glsl_fragment_shader;
53 case HgiShaderStageCompute:
54 return shaderc_glsl_compute_shader;
55 }
56
57 TF_CODING_ERROR("Unknown stage");
58 return shaderc_glsl_infer_from_source;
59 }
60
61 bool
HgiVulkanCompileGLSL(const char * name,const char * shaderCodes[],uint8_t numShaderCodes,HgiShaderStage stage,std::vector<unsigned int> * spirvOUT,std::string * errors)62 HgiVulkanCompileGLSL(
63 const char* name,
64 const char* shaderCodes[],
65 uint8_t numShaderCodes,
66 HgiShaderStage stage,
67 std::vector<unsigned int>* spirvOUT,
68 std::string* errors)
69 {
70 if (numShaderCodes==0 || !spirvOUT) {
71 if (errors) {
72 errors->append("No shader to compile %s", name);
73 }
74 return false;
75 }
76
77 std::string source;
78 for (uint8_t i=0; i<numShaderCodes; ++i) {
79 source += shaderCodes[i];
80 }
81
82 shaderc::CompileOptions options;
83 options.SetTargetEnvironment(shaderc_target_env_vulkan,
84 shaderc_env_version_vulkan_1_0);
85 options.SetTargetSpirv(shaderc_spirv_version_1_0);
86
87 shaderc_shader_kind const kind = _GetShaderStage(stage);
88
89 shaderc::Compiler compiler;
90 shaderc::SpvCompilationResult result =
91 compiler.CompileGlslToSpv(source, kind, name, options);
92
93 if (result.GetCompilationStatus() != shaderc_compilation_status_success) {
94 *errors = result.GetErrorMessage();
95 return false;
96 }
97
98 spirvOUT->assign(result.cbegin(), result.cend());
99
100 return true;
101 }
102
103 static bool
_VerifyResults(SpvReflectShaderModule * module,SpvReflectResult const & result)104 _VerifyResults(SpvReflectShaderModule* module, SpvReflectResult const& result)
105 {
106 if (!TF_VERIFY(result == SPV_REFLECT_RESULT_SUCCESS)) {
107 spvReflectDestroyShaderModule(module);
108 return false;
109 }
110
111 return true;
112 }
113
114 static VkDescriptorSetLayout
_CreateDescriptorSetLayout(HgiVulkanDevice * device,VkDescriptorSetLayoutCreateInfo const & createInfo,std::string const & debugName)115 _CreateDescriptorSetLayout(
116 HgiVulkanDevice* device,
117 VkDescriptorSetLayoutCreateInfo const& createInfo,
118 std::string const& debugName)
119 {
120 VkDescriptorSetLayout layout = nullptr;
121 TF_VERIFY(
122 vkCreateDescriptorSetLayout(
123 device->GetVulkanDevice(),
124 &createInfo,
125 HgiVulkanAllocator(),
126 &layout) == VK_SUCCESS
127 );
128
129 // Debug label
130 if (!debugName.empty()) {
131 std::string debugLabel = "DescriptorSetLayout " + debugName;
132 HgiVulkanSetDebugName(
133 device,
134 (uint64_t)layout,
135 VK_OBJECT_TYPE_DESCRIPTOR_SET_LAYOUT,
136 debugLabel.c_str());
137 }
138
139 return layout;
140 }
141
142 HgiVulkanDescriptorSetInfoVector
HgiVulkanGatherDescriptorSetInfo(std::vector<unsigned int> const & spirv)143 HgiVulkanGatherDescriptorSetInfo(
144 std::vector<unsigned int> const& spirv)
145 {
146 // This code is based on main_descriptors.cpp in the SPIRV-Reflect repo.
147
148 SpvReflectShaderModule module = {};
149 SpvReflectResult result = spvReflectCreateShaderModule(
150 spirv.size()*sizeof(uint32_t), spirv.data(), &module);
151 if (!_VerifyResults(&module, result)) {
152 return HgiVulkanDescriptorSetInfoVector();
153 }
154
155 uint32_t count = 0;
156 result = spvReflectEnumerateDescriptorSets(&module, &count, NULL);
157 if (!_VerifyResults(&module, result)) {
158 return HgiVulkanDescriptorSetInfoVector();
159 }
160
161 std::vector<SpvReflectDescriptorSet*> sets(count);
162 result = spvReflectEnumerateDescriptorSets(&module, &count, sets.data());
163 if (!_VerifyResults(&module, result)) {
164 return HgiVulkanDescriptorSetInfoVector();
165 }
166
167 // Generate all necessary data structures to create a VkDescriptorSetLayout
168 // for each descriptor set in this shader.
169 std::vector<HgiVulkanDescriptorSetInfo> infos(sets.size());
170
171 for (size_t s = 0; s < sets.size(); s++) {
172 SpvReflectDescriptorSet const& reflSet = *(sets[s]);
173 HgiVulkanDescriptorSetInfo& info = infos[s];
174 info.bindings.resize(reflSet.binding_count);
175
176 for (uint32_t b = 0; b < reflSet.binding_count; b++) {
177 SpvReflectDescriptorBinding const& reflBinding =
178 *(reflSet.bindings[b]);
179
180 VkDescriptorSetLayoutBinding& layoutBinding = info.bindings[b];
181 layoutBinding.binding = reflBinding.binding;
182 layoutBinding.descriptorType =
183 static_cast<VkDescriptorType>(reflBinding.descriptor_type);
184 layoutBinding.descriptorCount = 1;
185
186 for (uint32_t d = 0; d < reflBinding.array.dims_count; d++) {
187 layoutBinding.descriptorCount *= reflBinding.array.dims[d];
188 }
189 layoutBinding.stageFlags =
190 static_cast<VkShaderStageFlagBits>(module.shader_stage);
191 }
192
193 info.setNumber = reflSet.set;
194 }
195
196 spvReflectDestroyShaderModule(&module);
197
198 return infos;
199 }
200
201 VkDescriptorSetLayoutVector
HgiVulkanMakeDescriptorSetLayouts(HgiVulkanDevice * device,std::vector<HgiVulkanDescriptorSetInfoVector> const & infos,std::string const & debugName)202 HgiVulkanMakeDescriptorSetLayouts(
203 HgiVulkanDevice* device,
204 std::vector<HgiVulkanDescriptorSetInfoVector> const& infos,
205 std::string const& debugName)
206 {
207 std::unordered_map<uint32_t, HgiVulkanDescriptorSetInfo> mergedInfos;
208
209 // Merge the binding info of each of the infos such that the resource
210 // bindings information for each of the shader stage modules is merged
211 // together. For example a vertex shader may have different buffers and
212 // textures bound than a fragment shader. We merge them all together to
213 // create the descriptor set layout for that shader program.
214
215 for (HgiVulkanDescriptorSetInfoVector const& infoVec : infos) {
216 for (HgiVulkanDescriptorSetInfo const& info : infoVec) {
217
218 // Get the set (or create one)
219 HgiVulkanDescriptorSetInfo& trg = mergedInfos[info.setNumber];
220
221 for (VkDescriptorSetLayoutBinding const& bi : info.bindings) {
222
223 // If two shader modules have the same binding information for
224 // a specific resource, we only want to insert it once.
225 // For example both the vertex shaders and fragment shader may
226 // have a texture bound at the same binding index.
227
228 VkDescriptorSetLayoutBinding* dst = nullptr;
229 for (VkDescriptorSetLayoutBinding& bind : trg.bindings) {
230 if (bind.binding == bi.binding) {
231 dst = &bind;
232 break;
233 }
234 }
235
236 // It is a new binding we haven't seen before. Add it
237 if (!dst) {
238 trg.setNumber = info.setNumber;
239 trg.bindings.push_back(bi);
240 dst = &trg.bindings.back();
241 }
242
243 // If a vertex module and a fragment module access the same
244 // resource, we need to merge the stageFlags.
245 dst->stageFlags |= bi.stageFlags;
246 }
247 }
248 }
249
250 // Generate the VkDescriptorSetLayoutCreateInfos for the bindings we merged.
251 for (auto& pair : mergedInfos) {
252 HgiVulkanDescriptorSetInfo& info = pair.second;
253 info.createInfo.sType =
254 VK_STRUCTURE_TYPE_DESCRIPTOR_SET_LAYOUT_CREATE_INFO;
255 info.createInfo.bindingCount = info.bindings.size();
256 info.createInfo.pBindings = info.bindings.data();
257 }
258
259 // Create VkDescriptorSetLayouts out of the merge infos above.
260 VkDescriptorSetLayoutVector layouts;
261
262 for (auto const& pair : mergedInfos) {
263 HgiVulkanDescriptorSetInfo const& info = pair.second;
264 VkDescriptorSetLayout layout = _CreateDescriptorSetLayout(
265 device, info.createInfo, debugName);
266 layouts.push_back(layout);
267 }
268
269 return layouts;
270 }
271
272 PXR_NAMESPACE_CLOSE_SCOPE
273