| // Copyright 2019 The SwiftShader Authors. All Rights Reserved. |
| // |
| // Licensed under the Apache License, Version 2.0 (the "License"); |
| // you may not use this file except in compliance with the License. |
| // You may obtain a copy of the License at |
| // |
| // http://www.apache.org/licenses/LICENSE-2.0 |
| // |
| // Unless required by applicable law or agreed to in writing, software |
| // distributed under the License is distributed on an "AS IS" BASIS, |
| // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| // See the License for the specific language governing permissions and |
| // limitations under the License. |
| |
| #include "ComputeProgram.hpp" |
| |
| #include "Constants.hpp" |
| #include "System/Debug.hpp" |
| #include "Vulkan/VkDevice.hpp" |
| #include "Vulkan/VkPipelineLayout.hpp" |
| |
| #include "marl/defer.h" |
| #include "marl/trace.h" |
| #include "marl/waitgroup.h" |
| |
| #include <queue> |
| |
| namespace sw { |
| |
| ComputeProgram::ComputeProgram(vk::Device *device, std::shared_ptr<SpirvShader> shader, const vk::PipelineLayout *pipelineLayout, const vk::DescriptorSet::Bindings &descriptorSets) |
| : device(device) |
| , shader(shader) |
| , pipelineLayout(pipelineLayout) |
| , descriptorSets(descriptorSets) |
| { |
| } |
| |
| ComputeProgram::~ComputeProgram() |
| { |
| } |
| |
| void ComputeProgram::generate() |
| { |
| MARL_SCOPED_EVENT("ComputeProgram::generate"); |
| |
| SpirvRoutine routine(pipelineLayout); |
| shader->emitProlog(&routine); |
| emit(&routine); |
| shader->emitEpilog(&routine); |
| } |
| |
| void ComputeProgram::setWorkgroupBuiltins(Pointer<Byte> data, SpirvRoutine *routine, Int workgroupID[3]) |
| { |
| // TODO(b/146486064): Consider only assigning these to the SpirvRoutine iff they are ever going to be read. |
| routine->numWorkgroups = *Pointer<Int4>(data + OFFSET(Data, numWorkgroups)); |
| routine->workgroupID = Insert(Insert(Insert(Int4(0), workgroupID[0], 0), workgroupID[1], 1), workgroupID[2], 2); |
| routine->workgroupSize = *Pointer<Int4>(data + OFFSET(Data, workgroupSize)); |
| routine->subgroupsPerWorkgroup = *Pointer<Int>(data + OFFSET(Data, subgroupsPerWorkgroup)); |
| routine->invocationsPerSubgroup = *Pointer<Int>(data + OFFSET(Data, invocationsPerSubgroup)); |
| |
| routine->setInputBuiltin(shader.get(), spv::BuiltInNumWorkgroups, [&](const SpirvShader::BuiltinMapping &builtin, Array<SIMD::Float> &value) { |
| value[builtin.FirstComponent + 0] = As<SIMD::Float>(SIMD::Int(routine->numWorkgroups.x)); |
| value[builtin.FirstComponent + 1] = As<SIMD::Float>(SIMD::Int(routine->numWorkgroups.y)); |
| value[builtin.FirstComponent + 2] = As<SIMD::Float>(SIMD::Int(routine->numWorkgroups.z)); |
| }); |
| |
| routine->setInputBuiltin(shader.get(), spv::BuiltInWorkgroupId, [&](const SpirvShader::BuiltinMapping &builtin, Array<SIMD::Float> &value) { |
| value[builtin.FirstComponent + 0] = As<SIMD::Float>(SIMD::Int(workgroupID[0])); |
| value[builtin.FirstComponent + 1] = As<SIMD::Float>(SIMD::Int(workgroupID[1])); |
| value[builtin.FirstComponent + 2] = As<SIMD::Float>(SIMD::Int(workgroupID[2])); |
| }); |
| |
| routine->setInputBuiltin(shader.get(), spv::BuiltInWorkgroupSize, [&](const SpirvShader::BuiltinMapping &builtin, Array<SIMD::Float> &value) { |
| value[builtin.FirstComponent + 0] = As<SIMD::Float>(SIMD::Int(routine->workgroupSize.x)); |
| value[builtin.FirstComponent + 1] = As<SIMD::Float>(SIMD::Int(routine->workgroupSize.y)); |
| value[builtin.FirstComponent + 2] = As<SIMD::Float>(SIMD::Int(routine->workgroupSize.z)); |
| }); |
| |
| routine->setInputBuiltin(shader.get(), spv::BuiltInNumSubgroups, [&](const SpirvShader::BuiltinMapping &builtin, Array<SIMD::Float> &value) { |
| value[builtin.FirstComponent] = As<SIMD::Float>(SIMD::Int(routine->subgroupsPerWorkgroup)); |
| }); |
| |
| routine->setInputBuiltin(shader.get(), spv::BuiltInSubgroupSize, [&](const SpirvShader::BuiltinMapping &builtin, Array<SIMD::Float> &value) { |
| value[builtin.FirstComponent] = As<SIMD::Float>(SIMD::Int(routine->invocationsPerSubgroup)); |
| }); |
| |
| routine->setImmutableInputBuiltins(shader.get()); |
| } |
| |
| void ComputeProgram::setSubgroupBuiltins(Pointer<Byte> data, SpirvRoutine *routine, Int workgroupID[3], SIMD::Int localInvocationIndex, Int subgroupIndex) |
| { |
| Int4 numWorkgroups = *Pointer<Int4>(data + OFFSET(Data, numWorkgroups)); |
| Int4 workgroupSize = *Pointer<Int4>(data + OFFSET(Data, workgroupSize)); |
| |
| Int workgroupSizeX = workgroupSize.x; |
| Int workgroupSizeY = workgroupSize.y; |
| |
| SIMD::Int localInvocationID[3]; |
| { |
| SIMD::Int idx = localInvocationIndex; |
| localInvocationID[2] = idx / SIMD::Int(workgroupSizeX * workgroupSizeY); |
| idx -= localInvocationID[2] * SIMD::Int(workgroupSizeX * workgroupSizeY); // modulo |
| localInvocationID[1] = idx / SIMD::Int(workgroupSizeX); |
| idx -= localInvocationID[1] * SIMD::Int(workgroupSizeX); // modulo |
| localInvocationID[0] = idx; |
| } |
| |
| Int4 wgID = Insert(Insert(Insert(Int4(0), workgroupID[0], 0), workgroupID[1], 1), workgroupID[2], 2); |
| auto localBase = workgroupSize * wgID; |
| SIMD::Int globalInvocationID[3]; |
| globalInvocationID[0] = SIMD::Int(Extract(localBase, 0)) + localInvocationID[0]; |
| globalInvocationID[1] = SIMD::Int(Extract(localBase, 1)) + localInvocationID[1]; |
| globalInvocationID[2] = SIMD::Int(Extract(localBase, 2)) + localInvocationID[2]; |
| |
| routine->localInvocationIndex = localInvocationIndex; |
| routine->subgroupIndex = subgroupIndex; |
| routine->localInvocationID[0] = localInvocationID[0]; |
| routine->localInvocationID[1] = localInvocationID[1]; |
| routine->localInvocationID[2] = localInvocationID[2]; |
| routine->globalInvocationID[0] = globalInvocationID[0]; |
| routine->globalInvocationID[1] = globalInvocationID[1]; |
| routine->globalInvocationID[2] = globalInvocationID[2]; |
| |
| routine->setInputBuiltin(shader.get(), spv::BuiltInLocalInvocationIndex, [&](const SpirvShader::BuiltinMapping &builtin, Array<SIMD::Float> &value) { |
| ASSERT(builtin.SizeInComponents == 1); |
| value[builtin.FirstComponent] = As<SIMD::Float>(localInvocationIndex); |
| }); |
| |
| routine->setInputBuiltin(shader.get(), spv::BuiltInSubgroupId, [&](const SpirvShader::BuiltinMapping &builtin, Array<SIMD::Float> &value) { |
| ASSERT(builtin.SizeInComponents == 1); |
| value[builtin.FirstComponent] = As<SIMD::Float>(SIMD::Int(subgroupIndex)); |
| }); |
| |
| routine->setInputBuiltin(shader.get(), spv::BuiltInLocalInvocationId, [&](const SpirvShader::BuiltinMapping &builtin, Array<SIMD::Float> &value) { |
| for(uint32_t component = 0; component < builtin.SizeInComponents; component++) |
| { |
| value[builtin.FirstComponent + component] = |
| As<SIMD::Float>(localInvocationID[component]); |
| } |
| }); |
| |
| routine->setInputBuiltin(shader.get(), spv::BuiltInGlobalInvocationId, [&](const SpirvShader::BuiltinMapping &builtin, Array<SIMD::Float> &value) { |
| for(uint32_t component = 0; component < builtin.SizeInComponents; component++) |
| { |
| value[builtin.FirstComponent + component] = |
| As<SIMD::Float>(globalInvocationID[component]); |
| } |
| }); |
| } |
| |
| void ComputeProgram::emit(SpirvRoutine *routine) |
| { |
| Pointer<Byte> device = Arg<0>(); |
| Pointer<Byte> data = Arg<1>(); |
| Int workgroupX = Arg<2>(); |
| Int workgroupY = Arg<3>(); |
| Int workgroupZ = Arg<4>(); |
| Pointer<Byte> workgroupMemory = Arg<5>(); |
| Int firstSubgroup = Arg<6>(); |
| Int subgroupCount = Arg<7>(); |
| |
| routine->device = device; |
| routine->descriptorSets = data + OFFSET(Data, descriptorSets); |
| routine->descriptorDynamicOffsets = data + OFFSET(Data, descriptorDynamicOffsets); |
| routine->pushConstants = data + OFFSET(Data, pushConstants); |
| routine->constants = device + OFFSET(vk::Device, constants); |
| routine->workgroupMemory = workgroupMemory; |
| |
| Int invocationsPerWorkgroup = *Pointer<Int>(data + OFFSET(Data, invocationsPerWorkgroup)); |
| |
| Int workgroupID[3] = { workgroupX, workgroupY, workgroupZ }; |
| setWorkgroupBuiltins(data, routine, workgroupID); |
| |
| For(Int i = 0, i < subgroupCount, i++) |
| { |
| auto subgroupIndex = firstSubgroup + i; |
| |
| // TODO: Replace SIMD::Int(0, 1, 2, 3) with SIMD-width equivalent |
| auto localInvocationIndex = SIMD::Int(subgroupIndex * SIMD::Width) + SIMD::Int(0, 1, 2, 3); |
| |
| // Disable lanes where (invocationIDs >= invocationsPerWorkgroup) |
| auto activeLaneMask = CmpLT(localInvocationIndex, SIMD::Int(invocationsPerWorkgroup)); |
| |
| setSubgroupBuiltins(data, routine, workgroupID, localInvocationIndex, subgroupIndex); |
| |
| shader->emit(routine, activeLaneMask, activeLaneMask, descriptorSets); |
| } |
| } |
| |
| void ComputeProgram::run( |
| const vk::DescriptorSet::Array &descriptorSetObjects, |
| const vk::DescriptorSet::Bindings &descriptorSets, |
| const vk::DescriptorSet::DynamicOffsets &descriptorDynamicOffsets, |
| const vk::Pipeline::PushConstantStorage &pushConstants, |
| uint32_t baseGroupX, uint32_t baseGroupY, uint32_t baseGroupZ, |
| uint32_t groupCountX, uint32_t groupCountY, uint32_t groupCountZ) |
| { |
| uint32_t workgroupSizeX = shader->getWorkgroupSizeX(); |
| uint32_t workgroupSizeY = shader->getWorkgroupSizeY(); |
| uint32_t workgroupSizeZ = shader->getWorkgroupSizeZ(); |
| |
| auto invocationsPerSubgroup = SIMD::Width; |
| auto invocationsPerWorkgroup = workgroupSizeX * workgroupSizeY * workgroupSizeZ; |
| auto subgroupsPerWorkgroup = (invocationsPerWorkgroup + invocationsPerSubgroup - 1) / invocationsPerSubgroup; |
| |
| Data data; |
| data.descriptorSets = descriptorSets; |
| data.descriptorDynamicOffsets = descriptorDynamicOffsets; |
| data.numWorkgroups[0] = groupCountX; |
| data.numWorkgroups[1] = groupCountY; |
| data.numWorkgroups[2] = groupCountZ; |
| data.workgroupSize[0] = workgroupSizeX; |
| data.workgroupSize[1] = workgroupSizeY; |
| data.workgroupSize[2] = workgroupSizeZ; |
| data.invocationsPerSubgroup = invocationsPerSubgroup; |
| data.invocationsPerWorkgroup = invocationsPerWorkgroup; |
| data.subgroupsPerWorkgroup = subgroupsPerWorkgroup; |
| data.pushConstants = pushConstants; |
| |
| marl::WaitGroup wg; |
| const uint32_t batchCount = 16; |
| |
| auto groupCount = groupCountX * groupCountY * groupCountZ; |
| |
| for(uint32_t batchID = 0; batchID < batchCount && batchID < groupCount; batchID++) |
| { |
| wg.add(1); |
| marl::schedule([=, &data] { |
| defer(wg.done()); |
| std::vector<uint8_t> workgroupMemory(shader->workgroupMemory.size()); |
| |
| for(uint32_t groupIndex = batchID; groupIndex < groupCount; groupIndex += batchCount) |
| { |
| auto modulo = groupIndex; |
| auto groupOffsetZ = modulo / (groupCountX * groupCountY); |
| modulo -= groupOffsetZ * (groupCountX * groupCountY); |
| auto groupOffsetY = modulo / groupCountX; |
| modulo -= groupOffsetY * groupCountX; |
| auto groupOffsetX = modulo; |
| |
| auto groupZ = baseGroupZ + groupOffsetZ; |
| auto groupY = baseGroupY + groupOffsetY; |
| auto groupX = baseGroupX + groupOffsetX; |
| MARL_SCOPED_EVENT("groupX: %d, groupY: %d, groupZ: %d", groupX, groupY, groupZ); |
| |
| using Coroutine = std::unique_ptr<rr::Stream<SpirvEmitter::YieldResult>>; |
| std::queue<Coroutine> coroutines; |
| |
| if(shader->getAnalysis().ContainsControlBarriers) |
| { |
| // Make a function call per subgroup so each subgroup |
| // can yield, bringing all subgroups to the barrier |
| // together. |
| for(uint32_t subgroupIndex = 0; subgroupIndex < subgroupsPerWorkgroup; subgroupIndex++) |
| { |
| auto coroutine = (*this)(device, &data, groupX, groupY, groupZ, workgroupMemory.data(), subgroupIndex, 1); |
| coroutines.push(std::move(coroutine)); |
| } |
| } |
| else |
| { |
| auto coroutine = (*this)(device, &data, groupX, groupY, groupZ, workgroupMemory.data(), 0, subgroupsPerWorkgroup); |
| coroutines.push(std::move(coroutine)); |
| } |
| |
| while(coroutines.size() > 0) |
| { |
| auto coroutine = std::move(coroutines.front()); |
| coroutines.pop(); |
| |
| SpirvEmitter::YieldResult result; |
| if(coroutine->await(result)) |
| { |
| // TODO: Consider result (when the enum is more than 1 entry). |
| coroutines.push(std::move(coroutine)); |
| } |
| } |
| } |
| }); |
| } |
| |
| wg.wait(); |
| |
| if(shader->containsImageWrite()) |
| { |
| vk::DescriptorSet::ContentsChanged(descriptorSetObjects, pipelineLayout, device); |
| } |
| } |
| |
| } // namespace sw |