Skip to content

Commit

Permalink
igl | vulkan | Use SPIR-V reflection to construct descriptor sets for…
Browse files Browse the repository at this point in the history 
… compute pipelines

Summary:
Use SPIR-V reflection to construct descriptor sets for compute pipelines.

This is a part of upcoming major performance optimization for descriptor set updates.

Reviewed By: mmaurer

Differential Revision: D52265914

fbshipit-source-id: 41699d6824040f720f7d6eea4d157d9da341592b
  • Loading branch information
@corporateshark @facebook-github-bot
corporateshark authored and facebook-github-bot committed Dec 22, 2023
1 parent 06505c1 commit 63b69dc
Show file tree
Hide file tree
Showing 8 changed files with 313 additions and 163 deletions.
27 changes: 19 additions & 8 deletions src/igl/vulkan/ComputeCommandEncoder.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -28,7 +28,6 @@ ComputeCommandEncoder::ComputeCommandEncoder(const std::shared_ptr<CommandBuffer
IGL_ASSERT(commandBuffer);

ctx_.checkAndUpdateDescriptorSets();
ctx_.bindBindlessDescriptorSet(cmdBuffer_, VK_PIPELINE_BIND_POINT_COMPUTE);

isEncoding_ = true;
}
Expand Down Expand Up @@ -64,11 +63,23 @@ void ComputeCommandEncoder::bindComputePipelineState(

cps_ = static_cast<igl::vulkan::ComputePipelineState*>(pipelineState.get());

const ComputePipelineDesc& desc = cps_->getComputePipelineDesc();

ensureShaderModule(desc.shaderStages->getComputeModule().get());

binder_.bindPipeline(cps_->getVkPipeline());
binder_.bindPipeline(cps_->getVkPipeline(), &cps_->getSpvModuleInfo());

if (ctx_.config_.enableDescriptorIndexing) {
VkDescriptorSet dset = ctx_.getBindlessVkDescriptorSet();

#if IGL_VULKAN_PRINT_COMMANDS
IGL_LOG_INFO("%p vkCmdBindDescriptorSets(COMPUTE) - bindless\n", cmdBuffer_);
#endif // IGL_VULKAN_PRINT_COMMANDS
ctx_.vf_.vkCmdBindDescriptorSets(cmdBuffer_,
VK_PIPELINE_BIND_POINT_COMPUTE,
cps_->getVkPipelineLayout(),
kBindPoint_Bindless,
1,
&dset,
0,
nullptr);
}
}

void ComputeCommandEncoder::dispatchThreadGroups(const Dimensions& threadgroupCount,
Expand All @@ -77,7 +88,7 @@ void ComputeCommandEncoder::dispatchThreadGroups(const Dimensions& threadgroupCo

IGL_ASSERT_MSG(cps_, "Did you forget to call bindComputePipeline()?");

binder_.updateBindings();
binder_.updateBindings(cps_->getVkPipelineLayout(), cps_);
// threadgroupSize is controlled inside compute shaders
ctx_.vf_.vkCmdDispatch(
cmdBuffer_, threadgroupCount.width, threadgroupCount.height, threadgroupCount.depth);
Expand Down Expand Up @@ -180,7 +191,7 @@ void ComputeCommandEncoder::bindPushConstants(const void* data, size_t length, s
}

ctx_.vf_.vkCmdPushConstants(cmdBuffer_,
ctx_.pipelineLayoutCompute_->getVkPipelineLayout(),
cps_->getVkPipelineLayout(),
VK_SHADER_STAGE_COMPUTE_BIT,
(uint32_t)offset,
(uint32_t)length,
Expand Down
33 changes: 31 additions & 2 deletions src/igl/vulkan/ComputePipelineState.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@
#include <igl/vulkan/Device.h>
#include <igl/vulkan/ShaderModule.h>
#include <igl/vulkan/VulkanContext.h>
#include <igl/vulkan/VulkanDescriptorSetLayout.h>
#include <igl/vulkan/VulkanDevice.h>
#include <igl/vulkan/VulkanPipelineBuilder.h>
#include <igl/vulkan/VulkanPipelineLayout.h>
Expand All @@ -19,7 +20,7 @@ namespace vulkan {

ComputePipelineState::ComputePipelineState(const igl::vulkan::Device& device,
ComputePipelineDesc desc) :
PipelineState(device.getVulkanContext(), nullptr, desc.debugName.c_str()),
PipelineState(device.getVulkanContext(), desc.shaderStages.get(), desc.debugName.c_str()),
device_(device),
desc_(std::move(desc)) {}

Expand Down Expand Up @@ -57,6 +58,34 @@ VkPipeline ComputePipelineState::getVkPipeline() const {
return pipeline_;
}

// @fb-only
const VkDescriptorSetLayout DSLs[] = {
dslCombinedImageSamplers_->getVkDescriptorSetLayout(),
dslUniformBuffers_->getVkDescriptorSetLayout(),
dslStorageBuffers_->getVkDescriptorSetLayout(),
ctx.getBindlessVkDescriptorSetLayout(),
};

const VkPhysicalDeviceLimits& limits = ctx.getVkPhysicalDeviceProperties().limits;

constexpr uint32_t kPushConstantsSize = 128;

if (!IGL_VERIFY(kPushConstantsSize <= limits.maxPushConstantsSize)) {
IGL_LOG_ERROR("Push constants size exceeded %u (max %u bytes)",
kPushConstantsSize,
limits.maxPushConstantsSize);
}

pipelineLayout_ = std::make_unique<VulkanPipelineLayout>(
ctx,
ctx.getVkDevice(),
DSLs,
static_cast<uint32_t>(ctx.config_.enableDescriptorIndexing
? IGL_ARRAY_NUM_ELEMENTS(DSLs)
: IGL_ARRAY_NUM_ELEMENTS(DSLs) - 1u),
ivkGetPushConstantRange(VK_SHADER_STAGE_COMPUTE_BIT, 0, kPushConstantsSize),
IGL_FORMAT("Pipeline Layout: {}", desc_.debugName).c_str());

const auto& shaderModule = desc_.shaderStages->getComputeModule();

igl::vulkan::VulkanComputePipelineBuilder()
Expand All @@ -67,7 +96,7 @@ VkPipeline ComputePipelineState::getVkPipeline() const {
.build(ctx.vf_,
ctx.device_->getVkDevice(),
ctx.pipelineCache_,
ctx.pipelineLayoutCompute_->getVkPipelineLayout(),
pipelineLayout_->getVkPipelineLayout(),
&pipeline_,
desc_.debugName.c_str());

Expand Down
4 changes: 1 addition & 3 deletions src/igl/vulkan/PipelineState.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -18,9 +18,7 @@ namespace igl::vulkan {
PipelineState::PipelineState(const VulkanContext& ctx,
IShaderStages* stages,
const char* debugName) {
if (!stages) {
return;
}
IGL_ASSERT(stages);

auto* sm = static_cast<igl::vulkan::ShaderModule*>(stages->getComputeModule().get());

Expand Down
14 changes: 7 additions & 7 deletions src/igl/vulkan/RenderCommandEncoder.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -275,7 +275,7 @@ void RenderCommandEncoder::initialize(const RenderPassDesc& renderPass,
bindScissorRect(scissor);

ctx_.checkAndUpdateDescriptorSets();
ctx_.bindBindlessDescriptorSet(cmdBuffer_, VK_PIPELINE_BIND_POINT_GRAPHICS);
ctx_.bindBindlessDescriptorSet(cmdBuffer_);

ctx_.vf_.vkCmdBeginRenderPass(cmdBuffer_, &bi, VK_SUBPASS_CONTENTS_INLINE);

Expand Down Expand Up @@ -416,7 +416,7 @@ void RenderCommandEncoder::bindRenderPipelineState(
"Make sure your render pass and render pipeline both have matching depth attachments");
}

binder_.bindPipeline(VK_NULL_HANDLE);
binder_.bindPipeline(VK_NULL_HANDLE, nullptr);
}

void RenderCommandEncoder::bindDepthStencilState(
Expand Down Expand Up @@ -590,7 +590,7 @@ void RenderCommandEncoder::bindPipeline() {
return;
}

binder_.bindPipeline(rps->getVkPipeline(dynamicState_));
binder_.bindPipeline(rps->getVkPipeline(dynamicState_), nullptr);
}

void RenderCommandEncoder::draw(PrimitiveType primitiveType,
Expand All @@ -609,7 +609,7 @@ void RenderCommandEncoder::draw(PrimitiveType primitiveType,

ensureVertexBuffers();

binder_.updateBindings();
binder_.updateBindings(ctx_.pipelineLayoutGraphics_->getVkPipelineLayout(), nullptr);
dynamicState_.setTopology(primitiveTypeToVkPrimitiveTopology(primitiveType));
bindPipeline();

Expand Down Expand Up @@ -646,7 +646,7 @@ void RenderCommandEncoder::drawIndexed(PrimitiveType primitiveType,

ensureVertexBuffers();

binder_.updateBindings();
binder_.updateBindings(ctx_.pipelineLayoutGraphics_->getVkPipelineLayout(), nullptr);
dynamicState_.setTopology(primitiveTypeToVkPrimitiveTopology(primitiveType));
bindPipeline();

Expand Down Expand Up @@ -689,7 +689,7 @@ void RenderCommandEncoder::multiDrawIndirect(PrimitiveType primitiveType,

ensureVertexBuffers();

binder_.updateBindings();
binder_.updateBindings(ctx_.pipelineLayoutGraphics_->getVkPipelineLayout(), nullptr);
dynamicState_.setTopology(primitiveTypeToVkPrimitiveTopology(primitiveType));
bindPipeline();

Expand Down Expand Up @@ -717,7 +717,7 @@ void RenderCommandEncoder::multiDrawIndexedIndirect(PrimitiveType primitiveType,

ensureVertexBuffers();

binder_.updateBindings();
binder_.updateBindings(ctx_.pipelineLayoutGraphics_->getVkPipelineLayout(), nullptr);
dynamicState_.setTopology(primitiveTypeToVkPrimitiveTopology(primitiveType));
bindPipeline();

Expand Down
34 changes: 25 additions & 9 deletions src/igl/vulkan/ResourcesBinder.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@
*/

#include <igl/vulkan/Buffer.h>
#include <igl/vulkan/PipelineState.h>
#include <igl/vulkan/ResourcesBinder.h>
#include <igl/vulkan/SamplerState.h>
#include <igl/vulkan/Texture.h>
Expand Down Expand Up @@ -138,31 +139,46 @@ void ResourcesBinder::bindTexture(uint32_t index, igl::vulkan::Texture* tex) {
}
}

void ResourcesBinder::updateBindings() {
void ResourcesBinder::updateBindings(VkPipelineLayout layout, const vulkan::PipelineState* state) {
IGL_PROFILER_FUNCTION_COLOR(IGL_PROFILER_COLOR_UPDATE);

IGL_ASSERT(layout != VK_NULL_HANDLE);

if (isDirtyFlags_ & DirtyFlagBits_Textures) {
ctx_.updateBindingsTextures(cmdBuffer_, bindPoint_, bindingsTextures_);
ctx_.updateBindingsTextures(cmdBuffer_,
layout,
bindPoint_,
bindingsTextures_,
state ? state->dslCombinedImageSamplers_.get() : nullptr,
info_);
}
if (isDirtyFlags_ & DirtyFlagBits_UniformBuffers) {
ctx_.updateBindingsUniformBuffers(cmdBuffer_, bindPoint_, bindingsUniformBuffers_);
ctx_.updateBindingsUniformBuffers(cmdBuffer_,
layout,
bindPoint_,
bindingsUniformBuffers_,
state ? state->dslUniformBuffers_.get() : nullptr,
info_);
}
if (isDirtyFlags_ & DirtyFlagBits_StorageBuffers) {
ctx_.updateBindingsStorageBuffers(cmdBuffer_, bindPoint_, bindingsStorageBuffers_);
ctx_.updateBindingsStorageBuffers(cmdBuffer_,
layout,
bindPoint_,
bindingsStorageBuffers_,
state ? state->dslStorageBuffers_.get() : nullptr,
info_);
}

if (isDirtyFlags_) {
// TODO: rebind all descriptor sets here in one command
isDirtyFlags_ = 0;
}
isDirtyFlags_ = 0;
}

void ResourcesBinder::bindPipeline(VkPipeline pipeline) {
void ResourcesBinder::bindPipeline(VkPipeline pipeline, const util::SpvModuleInfo* info) {
if (lastPipelineBound_ == pipeline) {
return;
}

lastPipelineBound_ = pipeline;
info_ = info;

if (pipeline != VK_NULL_HANDLE) {
#if IGL_VULKAN_PRINT_COMMANDS
Expand Down
9 changes: 7 additions & 2 deletions src/igl/vulkan/ResourcesBinder.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -15,8 +15,12 @@

namespace igl {
namespace vulkan {
namespace util {
struct SpvModuleInfo;
} // namespace util

class Buffer;
class PipelineState;
class SamplerState;
class Texture;
class VulkanBuffer;
Expand Down Expand Up @@ -64,11 +68,11 @@ class ResourcesBinder final {

/// @brief Convenience function that updates all bindings in the context for all resource types
/// that have been modified since the last time this function was called
void updateBindings();
void updateBindings(VkPipelineLayout layout, const vulkan::PipelineState* state);

/// @brief If the pipeline passed in as a parameter is different than the last pipeline bound
/// through this class, binds it and cache it as the last pipeline bound. Does nothing otherwise
void bindPipeline(VkPipeline pipeline);
void bindPipeline(VkPipeline pipeline, const util::SpvModuleInfo* info);

private:
friend class VulkanContext;
Expand Down Expand Up @@ -96,6 +100,7 @@ class ResourcesBinder final {
BindingsBuffers bindingsUniformBuffers_;
BindingsBuffers bindingsStorageBuffers_;
VkPipelineBindPoint bindPoint_ = VK_PIPELINE_BIND_POINT_GRAPHICS;
const util::SpvModuleInfo* info_ = nullptr;
};

} // namespace vulkan
Expand Down
Loading

0 comments on commit 63b69dc

Please sign in to comment.