1 // Copyright 2019 The SwiftShader Authors. All Rights Reserved.
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 #ifndef sw_ComputeProgram_hpp
16 #define sw_ComputeProgram_hpp
17 
18 #include "SpirvShader.hpp"
19 
20 #include "Reactor/Coroutine.hpp"
21 #include "Vulkan/VkDescriptorSet.hpp"
22 #include "Vulkan/VkPipeline.hpp"
23 
24 #include <functional>
25 
26 namespace vk {
27 class Device;
28 class PipelineLayout;
29 }  // namespace vk
30 
31 namespace sw {
32 
33 using namespace rr;
34 
35 class DescriptorSetsLayout;
36 struct Constants;
37 
38 // ComputeProgram builds a SPIR-V compute shader.
39 class ComputeProgram : public Coroutine<SpirvShader::YieldResult(
40                            void *data,
41                            int32_t workgroupX,
42                            int32_t workgroupY,
43                            int32_t workgroupZ,
44                            void *workgroupMemory,
45                            int32_t firstSubgroup,
46                            int32_t subgroupCount)>
47 {
48 public:
49 	ComputeProgram(vk::Device *device, SpirvShader const *spirvShader, vk::PipelineLayout const *pipelineLayout, const vk::DescriptorSet::Bindings &descriptorSets);
50 
51 	virtual ~ComputeProgram();
52 
53 	// generate builds the shader program.
54 	void generate();
55 
56 	// run executes the compute shader routine for all workgroups.
57 	void run(
58 	    vk::DescriptorSet::Array const &descriptorSetObjects,
59 	    vk::DescriptorSet::Bindings const &descriptorSetBindings,
60 	    vk::DescriptorSet::DynamicOffsets const &descriptorDynamicOffsets,
61 	    vk::Pipeline::PushConstantStorage const &pushConstants,
62 	    uint32_t baseGroupX, uint32_t baseGroupY, uint32_t baseGroupZ,
63 	    uint32_t groupCountX, uint32_t groupCountY, uint32_t groupCountZ);
64 
65 protected:
66 	void emit(SpirvRoutine *routine);
67 	void setWorkgroupBuiltins(Pointer<Byte> data, SpirvRoutine *routine, Int workgroupID[3]);
68 	void setSubgroupBuiltins(Pointer<Byte> data, SpirvRoutine *routine, Int workgroupID[3], SIMD::Int localInvocationIndex, Int subgroupIndex);
69 
70 	struct Data
71 	{
72 		vk::DescriptorSet::Bindings descriptorSets;
73 		vk::DescriptorSet::DynamicOffsets descriptorDynamicOffsets;
74 		uint4 numWorkgroups;               // [x, y, z, 0]
75 		uint4 workgroupSize;               // [x, y, z, 0]
76 		uint32_t invocationsPerSubgroup;   // SPIR-V: "SubgroupSize"
77 		uint32_t subgroupsPerWorkgroup;    // SPIR-V: "NumSubgroups"
78 		uint32_t invocationsPerWorkgroup;  // Total number of invocations per workgroup.
79 		vk::Pipeline::PushConstantStorage pushConstants;
80 		const Constants *constants;
81 	};
82 
83 	vk::Device *const device;
84 	SpirvShader const *const shader;
85 	vk::PipelineLayout const *const pipelineLayout;
86 	const vk::DescriptorSet::Bindings &descriptorSets;
87 };
88 
89 }  // namespace sw
90 
91 #endif  // sw_ComputeProgram_hpp
92