diff options
Diffstat (limited to 'tools/gfx/vulkan/vk-shader-object-layout.cpp')
| -rw-r--r-- | tools/gfx/vulkan/vk-shader-object-layout.cpp | 983 |
1 files changed, 983 insertions, 0 deletions
diff --git a/tools/gfx/vulkan/vk-shader-object-layout.cpp b/tools/gfx/vulkan/vk-shader-object-layout.cpp new file mode 100644 index 000000000..e1d06c791 --- /dev/null +++ b/tools/gfx/vulkan/vk-shader-object-layout.cpp @@ -0,0 +1,983 @@ +// vk-shader-object-layout.cpp +#include "vk-shader-object-layout.h" + +namespace gfx +{ + +using namespace Slang; + +namespace vk +{ + +Index ShaderObjectLayoutImpl::Builder::findOrAddDescriptorSet(Index space) +{ + Index index; + if (m_mapSpaceToDescriptorSetIndex.TryGetValue(space, index)) + return index; + + DescriptorSetInfo info = {}; + info.space = space; + + index = m_descriptorSetBuildInfos.getCount(); + m_descriptorSetBuildInfos.add(info); + + m_mapSpaceToDescriptorSetIndex.Add(space, index); + return index; +} + +VkDescriptorType ShaderObjectLayoutImpl::Builder::_mapDescriptorType( + slang::BindingType slangBindingType) +{ + switch (slangBindingType) + { + case slang::BindingType::PushConstant: + default: + SLANG_ASSERT("unsupported binding type"); + return VK_DESCRIPTOR_TYPE_MAX_ENUM; + + case slang::BindingType::Sampler: + return VK_DESCRIPTOR_TYPE_SAMPLER; + case slang::BindingType::CombinedTextureSampler: + return VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER; + case slang::BindingType::Texture: + return VK_DESCRIPTOR_TYPE_SAMPLED_IMAGE; + case slang::BindingType::MutableTexture: + return VK_DESCRIPTOR_TYPE_STORAGE_IMAGE; + case slang::BindingType::TypedBuffer: + return VK_DESCRIPTOR_TYPE_UNIFORM_TEXEL_BUFFER; + case slang::BindingType::MutableTypedBuffer: + return VK_DESCRIPTOR_TYPE_STORAGE_TEXEL_BUFFER; + case slang::BindingType::RawBuffer: + case slang::BindingType::MutableRawBuffer: + return VK_DESCRIPTOR_TYPE_STORAGE_BUFFER; + case slang::BindingType::InputRenderTarget: + return VK_DESCRIPTOR_TYPE_INPUT_ATTACHMENT; + case slang::BindingType::InlineUniformData: + return VK_DESCRIPTOR_TYPE_INLINE_UNIFORM_BLOCK_EXT; + case slang::BindingType::RayTracingAccelerationStructure: + return VK_DESCRIPTOR_TYPE_ACCELERATION_STRUCTURE_KHR; + case slang::BindingType::ConstantBuffer: + return VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER; + } +} + +/// Add any descriptor ranges implied by this object containing a leaf +/// sub-object described by `typeLayout`, at the given `offset`. + +void ShaderObjectLayoutImpl::Builder::_addDescriptorRangesAsValue( + slang::TypeLayoutReflection* typeLayout, BindingOffset const& offset) +{ + // First we will scan through all the descriptor sets that the Slang reflection + // information believes go into making up the given type. + // + // Note: We are initializing the sets in order so that their order in our + // internal data structures should be deterministically based on the order + // in which they are listed in Slang's reflection information. + // + Index descriptorSetCount = typeLayout->getDescriptorSetCount(); + for (Index i = 0; i < descriptorSetCount; ++i) + { + SlangInt descriptorRangeCount = typeLayout->getDescriptorSetDescriptorRangeCount(i); + if (descriptorRangeCount == 0) + continue; + auto descriptorSetIndex = + findOrAddDescriptorSet(offset.bindingSet + typeLayout->getDescriptorSetSpaceOffset(i)); + } + + // For actually populating the descriptor sets we prefer to enumerate + // the binding ranges of the type instead of the descriptor sets. + // + Index bindRangeCount = typeLayout->getBindingRangeCount(); + for (Index i = 0; i < bindRangeCount; ++i) + { + auto bindingRangeIndex = i; + auto bindingRangeType = typeLayout->getBindingRangeType(bindingRangeIndex); + switch (bindingRangeType) + { + default: + break; + + // We will skip over ranges that represent sub-objects for now, and handle + // them in a separate pass. + // + case slang::BindingType::ParameterBlock: + case slang::BindingType::ConstantBuffer: + case slang::BindingType::ExistentialValue: + case slang::BindingType::PushConstant: + continue; + } + + // Given a binding range we are interested in, we will then enumerate + // its contained descriptor ranges. + + Index descriptorRangeCount = + typeLayout->getBindingRangeDescriptorRangeCount(bindingRangeIndex); + if (descriptorRangeCount == 0) + continue; + auto slangDescriptorSetIndex = + typeLayout->getBindingRangeDescriptorSetIndex(bindingRangeIndex); + auto descriptorSetIndex = findOrAddDescriptorSet( + offset.bindingSet + typeLayout->getDescriptorSetSpaceOffset(slangDescriptorSetIndex)); + auto& descriptorSetInfo = m_descriptorSetBuildInfos[descriptorSetIndex]; + + Index firstDescriptorRangeIndex = + typeLayout->getBindingRangeFirstDescriptorRangeIndex(bindingRangeIndex); + for (Index j = 0; j < descriptorRangeCount; ++j) + { + Index descriptorRangeIndex = firstDescriptorRangeIndex + j; + auto slangDescriptorType = typeLayout->getDescriptorSetDescriptorRangeType( + slangDescriptorSetIndex, descriptorRangeIndex); + + // Certain kinds of descriptor ranges reflected by Slang do not + // manifest as descriptors at the Vulkan level, so we will skip those. + // + switch (slangDescriptorType) + { + case slang::BindingType::ExistentialValue: + case slang::BindingType::InlineUniformData: + case slang::BindingType::PushConstant: + continue; + default: + break; + } + + auto vkDescriptorType = _mapDescriptorType(slangDescriptorType); + VkDescriptorSetLayoutBinding vkBindingRangeDesc = {}; + vkBindingRangeDesc.binding = + offset.binding + (uint32_t)typeLayout->getDescriptorSetDescriptorRangeIndexOffset( + slangDescriptorSetIndex, descriptorRangeIndex); + vkBindingRangeDesc.descriptorCount = + (uint32_t)typeLayout->getDescriptorSetDescriptorRangeDescriptorCount( + slangDescriptorSetIndex, descriptorRangeIndex); + vkBindingRangeDesc.descriptorType = vkDescriptorType; + vkBindingRangeDesc.stageFlags = VK_SHADER_STAGE_ALL; + + descriptorSetInfo.vkBindings.add(vkBindingRangeDesc); + } + } + + // We skipped over the sub-object ranges when adding descriptors above, + // and now we will address that oversight by iterating over just + // the sub-object ranges. + // + Index subObjectRangeCount = typeLayout->getSubObjectRangeCount(); + for (Index subObjectRangeIndex = 0; subObjectRangeIndex < subObjectRangeCount; + ++subObjectRangeIndex) + { + auto bindingRangeIndex = + typeLayout->getSubObjectRangeBindingRangeIndex(subObjectRangeIndex); + auto bindingType = typeLayout->getBindingRangeType(bindingRangeIndex); + + auto subObjectTypeLayout = typeLayout->getBindingRangeLeafTypeLayout(bindingRangeIndex); + SLANG_ASSERT(subObjectTypeLayout); + + BindingOffset subObjectRangeOffset = offset; + subObjectRangeOffset += + BindingOffset(typeLayout->getSubObjectRangeOffset(subObjectRangeIndex)); + + switch (bindingType) + { + // A `ParameterBlock<X>` never contributes descripto ranges to the + // decriptor sets of a parent object. + // + case slang::BindingType::ParameterBlock: + default: + break; + + case slang::BindingType::ExistentialValue: + // An interest/existential-typed sub-object range will only contribute + // descriptor ranges to a parent object in the case where it has been + // specialied, which is precisely the case where the Slang reflection + // information will tell us about its "pending" layout. + // + if (auto pendingTypeLayout = subObjectTypeLayout->getPendingDataTypeLayout()) + { + BindingOffset pendingOffset = BindingOffset(subObjectRangeOffset.pending); + _addDescriptorRangesAsValue(pendingTypeLayout, pendingOffset); + } + break; + + case slang::BindingType::ConstantBuffer: + { + // A `ConstantBuffer<X>` range will contribute any nested descriptor + // ranges in `X`, along with a leading descriptor range for a + // uniform buffer to hold ordinary/uniform data, if there is any. + + SLANG_ASSERT(subObjectTypeLayout); + + auto containerVarLayout = subObjectTypeLayout->getContainerVarLayout(); + SLANG_ASSERT(containerVarLayout); + + auto elementVarLayout = subObjectTypeLayout->getElementVarLayout(); + SLANG_ASSERT(elementVarLayout); + + auto elementTypeLayout = elementVarLayout->getTypeLayout(); + SLANG_ASSERT(elementTypeLayout); + + BindingOffset containerOffset = subObjectRangeOffset; + containerOffset += BindingOffset(subObjectTypeLayout->getContainerVarLayout()); + + BindingOffset elementOffset = subObjectRangeOffset; + elementOffset += BindingOffset(elementVarLayout); + + _addDescriptorRangesAsConstantBuffer( + elementTypeLayout, containerOffset, elementOffset); + } + break; + + case slang::BindingType::PushConstant: + { + // This case indicates a `ConstantBuffer<X>` that was marked as being + // used for push constants. + // + // Much of the handling is the same as for an ordinary + // `ConstantBuffer<X>`, but of course we need to handle the ordinary + // data part differently. + + SLANG_ASSERT(subObjectTypeLayout); + + auto containerVarLayout = subObjectTypeLayout->getContainerVarLayout(); + SLANG_ASSERT(containerVarLayout); + + auto elementVarLayout = subObjectTypeLayout->getElementVarLayout(); + SLANG_ASSERT(elementVarLayout); + + auto elementTypeLayout = elementVarLayout->getTypeLayout(); + SLANG_ASSERT(elementTypeLayout); + + BindingOffset containerOffset = subObjectRangeOffset; + containerOffset += BindingOffset(subObjectTypeLayout->getContainerVarLayout()); + + BindingOffset elementOffset = subObjectRangeOffset; + elementOffset += BindingOffset(elementVarLayout); + + _addDescriptorRangesAsPushConstantBuffer( + elementTypeLayout, containerOffset, elementOffset); + } + break; + } + } +} + +/// Add the descriptor ranges implied by a `ConstantBuffer<X>` where `X` is +/// described by `elementTypeLayout`. +/// +/// The `containerOffset` and `elementOffset` are the binding offsets that +/// should apply to the buffer itself and the contents of the buffer, respectively. +/// + +void ShaderObjectLayoutImpl::Builder::_addDescriptorRangesAsConstantBuffer( + slang::TypeLayoutReflection* elementTypeLayout, + BindingOffset const& containerOffset, + BindingOffset const& elementOffset) +{ + // If the type has ordinary uniform data fields, we need to make sure to create + // a descriptor set with a constant buffer binding in the case that the shader + // object is bound as a stand alone parameter block. + if (elementTypeLayout->getSize(SLANG_PARAMETER_CATEGORY_UNIFORM) != 0) + { + auto descriptorSetIndex = findOrAddDescriptorSet(containerOffset.bindingSet); + auto& descriptorSetInfo = m_descriptorSetBuildInfos[descriptorSetIndex]; + VkDescriptorSetLayoutBinding vkBindingRangeDesc = {}; + vkBindingRangeDesc.binding = containerOffset.binding; + vkBindingRangeDesc.descriptorCount = 1; + vkBindingRangeDesc.descriptorType = VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER; + vkBindingRangeDesc.stageFlags = VK_SHADER_STAGE_ALL; + descriptorSetInfo.vkBindings.add(vkBindingRangeDesc); + } + + _addDescriptorRangesAsValue(elementTypeLayout, elementOffset); +} + +/// Add the descriptor ranges implied by a `PushConstantBuffer<X>` where `X` is +/// described by `elementTypeLayout`. +/// +/// The `containerOffset` and `elementOffset` are the binding offsets that +/// should apply to the buffer itself and the contents of the buffer, respectively. +/// + +void ShaderObjectLayoutImpl::Builder::_addDescriptorRangesAsPushConstantBuffer( + slang::TypeLayoutReflection* elementTypeLayout, + BindingOffset const& containerOffset, + BindingOffset const& elementOffset) +{ + // If the type has ordinary uniform data fields, we need to make sure to create + // a descriptor set with a constant buffer binding in the case that the shader + // object is bound as a stand alone parameter block. + auto ordinaryDataSize = (uint32_t)elementTypeLayout->getSize(SLANG_PARAMETER_CATEGORY_UNIFORM); + if (ordinaryDataSize != 0) + { + auto pushConstantRangeIndex = containerOffset.pushConstantRange; + + VkPushConstantRange vkPushConstantRange = {}; + vkPushConstantRange.size = ordinaryDataSize; + vkPushConstantRange.stageFlags = VK_SHADER_STAGE_ALL; // TODO: be more clever + + while ((uint32_t)m_ownPushConstantRanges.getCount() <= pushConstantRangeIndex) + { + VkPushConstantRange emptyRange = {0}; + m_ownPushConstantRanges.add(emptyRange); + } + + m_ownPushConstantRanges[pushConstantRangeIndex] = vkPushConstantRange; + } + + _addDescriptorRangesAsValue(elementTypeLayout, elementOffset); +} + +/// Add binding ranges to this shader object layout, as implied by the given +/// `typeLayout` + +void ShaderObjectLayoutImpl::Builder::addBindingRanges(slang::TypeLayoutReflection* typeLayout) +{ + SlangInt bindingRangeCount = typeLayout->getBindingRangeCount(); + for (SlangInt r = 0; r < bindingRangeCount; ++r) + { + slang::BindingType slangBindingType = typeLayout->getBindingRangeType(r); + uint32_t count = (uint32_t)typeLayout->getBindingRangeBindingCount(r); + slang::TypeLayoutReflection* slangLeafTypeLayout = + typeLayout->getBindingRangeLeafTypeLayout(r); + + Index baseIndex = 0; + Index subObjectIndex = 0; + switch (slangBindingType) + { + case slang::BindingType::ConstantBuffer: + case slang::BindingType::ParameterBlock: + case slang::BindingType::ExistentialValue: + baseIndex = m_subObjectCount; + subObjectIndex = baseIndex; + m_subObjectCount += count; + break; + case slang::BindingType::RawBuffer: + case slang::BindingType::MutableRawBuffer: + if (slangLeafTypeLayout->getType()->getElementType() != nullptr) + { + // A structured buffer occupies both a resource slot and + // a sub-object slot. + subObjectIndex = m_subObjectCount; + m_subObjectCount += count; + } + baseIndex = m_resourceViewCount; + m_resourceViewCount += count; + break; + case slang::BindingType::Sampler: + baseIndex = m_samplerCount; + m_samplerCount += count; + m_totalBindingCount += 1; + break; + + case slang::BindingType::CombinedTextureSampler: + baseIndex = m_combinedTextureSamplerCount; + m_combinedTextureSamplerCount += count; + m_totalBindingCount += 1; + break; + + case slang::BindingType::VaryingInput: + baseIndex = m_varyingInputCount; + m_varyingInputCount += count; + break; + + case slang::BindingType::VaryingOutput: + baseIndex = m_varyingOutputCount; + m_varyingOutputCount += count; + break; + default: + baseIndex = m_resourceViewCount; + m_resourceViewCount += count; + m_totalBindingCount += 1; + break; + } + + BindingRangeInfo bindingRangeInfo; + bindingRangeInfo.bindingType = slangBindingType; + bindingRangeInfo.count = count; + bindingRangeInfo.baseIndex = baseIndex; + bindingRangeInfo.subObjectIndex = subObjectIndex; + + // We'd like to extract the information on the GLSL/SPIR-V + // `binding` that this range should bind into (or whatever + // other specific kind of offset/index is appropriate to it). + // + // A binding range represents a logical member of the shader + // object type, and it may encompass zero or more *descriptor + // ranges* that describe how it is physically bound to pipeline + // state. + // + // If the current bindign range is backed by at least one descriptor + // range then we can query the binding offset of that descriptor + // range. We expect that in the common case there will be exactly + // one descriptor range, and we can extract the information easily. + // + if (typeLayout->getBindingRangeDescriptorRangeCount(r) != 0) + { + SlangInt descriptorSetIndex = typeLayout->getBindingRangeDescriptorSetIndex(r); + SlangInt descriptorRangeIndex = typeLayout->getBindingRangeFirstDescriptorRangeIndex(r); + + auto set = typeLayout->getDescriptorSetSpaceOffset(descriptorSetIndex); + auto bindingOffset = typeLayout->getDescriptorSetDescriptorRangeIndexOffset( + descriptorSetIndex, descriptorRangeIndex); + + bindingRangeInfo.setOffset = uint32_t(set); + bindingRangeInfo.bindingOffset = uint32_t(bindingOffset); + } + + m_bindingRanges.add(bindingRangeInfo); + } + + SlangInt subObjectRangeCount = typeLayout->getSubObjectRangeCount(); + for (SlangInt r = 0; r < subObjectRangeCount; ++r) + { + SlangInt bindingRangeIndex = typeLayout->getSubObjectRangeBindingRangeIndex(r); + auto& bindingRange = m_bindingRanges[bindingRangeIndex]; + auto slangBindingType = typeLayout->getBindingRangeType(bindingRangeIndex); + slang::TypeLayoutReflection* slangLeafTypeLayout = + typeLayout->getBindingRangeLeafTypeLayout(bindingRangeIndex); + + // A sub-object range can either represent a sub-object of a known + // type, like a `ConstantBuffer<Foo>` or `ParameterBlock<Foo>` + // (in which case we can pre-compute a layout to use, based on + // the type `Foo`) *or* it can represent a sub-object of some + // existential type (e.g., `IBar`) in which case we cannot + // know the appropraite type/layout of sub-object to allocate. + // + RefPtr<ShaderObjectLayoutImpl> subObjectLayout; + switch (slangBindingType) + { + default: + { + auto varLayout = slangLeafTypeLayout->getElementVarLayout(); + auto subTypeLayout = varLayout->getTypeLayout(); + ShaderObjectLayoutImpl::createForElementType( + m_renderer, subTypeLayout, subObjectLayout.writeRef()); + } + break; + + case slang::BindingType::ExistentialValue: + if (auto pendingTypeLayout = slangLeafTypeLayout->getPendingDataTypeLayout()) + { + ShaderObjectLayoutImpl::createForElementType( + m_renderer, pendingTypeLayout, subObjectLayout.writeRef()); + } + break; + } + + SubObjectRangeInfo subObjectRange; + subObjectRange.bindingRangeIndex = bindingRangeIndex; + subObjectRange.layout = subObjectLayout; + + // We will use Slang reflection infromation to extract the offset information + // for each sub-object range. + // + // TODO: We should also be extracting the uniform offset here. + // + subObjectRange.offset = SubObjectRangeOffset(typeLayout->getSubObjectRangeOffset(r)); + subObjectRange.stride = SubObjectRangeStride(slangLeafTypeLayout); + + switch (slangBindingType) + { + case slang::BindingType::ParameterBlock: + m_childDescriptorSetCount += subObjectLayout->getTotalDescriptorSetCount(); + m_childPushConstantRangeCount += subObjectLayout->getTotalPushConstantRangeCount(); + break; + + case slang::BindingType::ConstantBuffer: + m_childDescriptorSetCount += subObjectLayout->getChildDescriptorSetCount(); + m_totalBindingCount += subObjectLayout->getTotalBindingCount(); + m_childPushConstantRangeCount += subObjectLayout->getTotalPushConstantRangeCount(); + break; + + case slang::BindingType::ExistentialValue: + if (subObjectLayout) + { + m_childDescriptorSetCount += subObjectLayout->getChildDescriptorSetCount(); + m_totalBindingCount += subObjectLayout->getTotalBindingCount(); + m_childPushConstantRangeCount += subObjectLayout->getTotalPushConstantRangeCount(); + + // An interface-type range that includes ordinary data can + // increase the size of the ordinary data buffer we need to + // allocate for the parent object. + // + uint32_t ordinaryDataEnd = + subObjectRange.offset.pendingOrdinaryData + + (uint32_t)bindingRange.count * subObjectRange.stride.pendingOrdinaryData; + + if (ordinaryDataEnd > m_totalOrdinaryDataSize) + { + m_totalOrdinaryDataSize = ordinaryDataEnd; + } + } + break; + + default: + break; + } + + m_subObjectRanges.add(subObjectRange); + } +} + +Result ShaderObjectLayoutImpl::Builder::setElementTypeLayout( + slang::TypeLayoutReflection* typeLayout) +{ + typeLayout = _unwrapParameterGroups(typeLayout, m_containerType); + m_elementTypeLayout = typeLayout; + + m_totalOrdinaryDataSize = (uint32_t)typeLayout->getSize(); + + // Next we will compute the binding ranges that are used to store + // the logical contents of the object in memory. These will relate + // to the descriptor ranges in the various sets, but not always + // in a one-to-one fashion. + + addBindingRanges(typeLayout); + + // Note: This routine does not take responsibility for + // adding descriptor ranges at all, because the exact way + // that descriptor ranges need to be added varies between + // ordinary shader objects, root shader objects, and entry points. + + return SLANG_OK; +} + +SlangResult ShaderObjectLayoutImpl::Builder::build(ShaderObjectLayoutImpl** outLayout) +{ + auto layout = RefPtr<ShaderObjectLayoutImpl>(new ShaderObjectLayoutImpl()); + SLANG_RETURN_ON_FAIL(layout->_init(this)); + + returnRefPtrMove(outLayout, layout); + return SLANG_OK; +} + +Result ShaderObjectLayoutImpl::createForElementType( + DeviceImpl* renderer, + slang::TypeLayoutReflection* elementType, + ShaderObjectLayoutImpl** outLayout) +{ + Builder builder(renderer); + builder.setElementTypeLayout(elementType); + + // When constructing a shader object layout directly from a reflected + // type in Slang, we want to compute the descriptor sets and ranges + // that would be used if this object were bound as a parameter block. + // + // It might seem like we need to deal with the other cases for how + // the shader object might be bound, but the descriptor ranges we + // compute here will only ever be used in parameter-block case. + // + // One important wrinkle is that we know that the parameter block + // allocated for `elementType` will potentially need a buffer `binding` + // for any ordinary data it contains. + + bool needsOrdinaryDataBuffer = + builder.m_elementTypeLayout->getSize(SLANG_PARAMETER_CATEGORY_UNIFORM) != 0; + uint32_t ordinaryDataBufferCount = needsOrdinaryDataBuffer ? 1 : 0; + + // When binding the object, we know that the ordinary data buffer will + // always use a the first available `binding`, so its offset will be + // all zeroes. + // + BindingOffset containerOffset; + + // In contrast, the `binding`s used by all the other entries in the + // parameter block will need to be offset by one if there was + // an ordinary data buffer. + // + BindingOffset elementOffset; + elementOffset.binding = ordinaryDataBufferCount; + + // Furthermore, any `binding`s that arise due to "pending" data + // in the type of the object (due to specialization for existential types) + // will need to come after all the other `binding`s that were + // part of the "primary" (unspecialized) data. + // + uint32_t primaryDescriptorCount = + ordinaryDataBufferCount + (uint32_t)builder.m_elementTypeLayout->getSize( + SLANG_PARAMETER_CATEGORY_DESCRIPTOR_TABLE_SLOT); + elementOffset.pending.binding = primaryDescriptorCount; + + // Once we've computed the offset information, we simply add the + // descriptor ranges as if things were declared as a `ConstantBuffer<X>`, + // since that is how things will be laid out inside the parameter block. + // + builder._addDescriptorRangesAsConstantBuffer( + builder.m_elementTypeLayout, containerOffset, elementOffset); + return builder.build(outLayout); +} + +ShaderObjectLayoutImpl::~ShaderObjectLayoutImpl() +{ + for (auto& descSetInfo : m_descriptorSetInfos) + { + getDevice()->m_api.vkDestroyDescriptorSetLayout( + getDevice()->m_api.m_device, descSetInfo.descriptorSetLayout, nullptr); + } +} + +Result ShaderObjectLayoutImpl::_init(Builder const* builder) +{ + auto renderer = builder->m_renderer; + + initBase(renderer, builder->m_elementTypeLayout); + + m_bindingRanges = builder->m_bindingRanges; + + m_descriptorSetInfos = _Move(builder->m_descriptorSetBuildInfos); + m_ownPushConstantRanges = builder->m_ownPushConstantRanges; + m_resourceViewCount = builder->m_resourceViewCount; + m_samplerCount = builder->m_samplerCount; + m_combinedTextureSamplerCount = builder->m_combinedTextureSamplerCount; + m_childDescriptorSetCount = builder->m_childDescriptorSetCount; + m_totalBindingCount = builder->m_totalBindingCount; + m_subObjectCount = builder->m_subObjectCount; + m_subObjectRanges = builder->m_subObjectRanges; + m_totalOrdinaryDataSize = builder->m_totalOrdinaryDataSize; + + m_containerType = builder->m_containerType; + + // Create VkDescriptorSetLayout for all descriptor sets. + for (auto& descriptorSetInfo : m_descriptorSetInfos) + { + VkDescriptorSetLayoutCreateInfo createInfo = {}; + createInfo.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_LAYOUT_CREATE_INFO; + createInfo.pBindings = descriptorSetInfo.vkBindings.getBuffer(); + createInfo.bindingCount = (uint32_t)descriptorSetInfo.vkBindings.getCount(); + VkDescriptorSetLayout vkDescSetLayout; + SLANG_RETURN_ON_FAIL(renderer->m_api.vkCreateDescriptorSetLayout( + renderer->m_api.m_device, &createInfo, nullptr, &vkDescSetLayout)); + descriptorSetInfo.descriptorSetLayout = vkDescSetLayout; + } + return SLANG_OK; +} + +DeviceImpl* ShaderObjectLayoutImpl::getDevice() +{ + return static_cast<DeviceImpl*>(m_renderer); +} + +Result EntryPointLayout::Builder::build(EntryPointLayout** outLayout) +{ + RefPtr<EntryPointLayout> layout = new EntryPointLayout(); + SLANG_RETURN_ON_FAIL(layout->_init(this)); + + returnRefPtrMove(outLayout, layout); + return SLANG_OK; +} + +void EntryPointLayout::Builder::addEntryPointParams(slang::EntryPointLayout* entryPointLayout) +{ + m_slangEntryPointLayout = entryPointLayout; + setElementTypeLayout(entryPointLayout->getTypeLayout()); + m_shaderStageFlag = VulkanUtil::getShaderStage(entryPointLayout->getStage()); + + // Note: we do not bother adding any descriptor sets/ranges here, + // because the descriptor ranges of an entry point will simply + // be allocated as part of the descriptor sets for the root + // shader object. +} + +Result EntryPointLayout::_init(Builder const* builder) +{ + auto renderer = builder->m_renderer; + + SLANG_RETURN_ON_FAIL(Super::_init(builder)); + + m_slangEntryPointLayout = builder->m_slangEntryPointLayout; + m_shaderStageFlag = builder->m_shaderStageFlag; + return SLANG_OK; +} + +RootShaderObjectLayout::~RootShaderObjectLayout() +{ + if (m_pipelineLayout) + { + m_renderer->m_api.vkDestroyPipelineLayout( + m_renderer->m_api.m_device, m_pipelineLayout, nullptr); + } +} + +Index RootShaderObjectLayout::findEntryPointIndex(VkShaderStageFlags stage) +{ + auto entryPointCount = m_entryPoints.getCount(); + for (Index i = 0; i < entryPointCount; ++i) + { + auto entryPoint = m_entryPoints[i]; + if (entryPoint.layout->getShaderStageFlag() == stage) + return i; + } + return -1; +} + +Result RootShaderObjectLayout::create( + DeviceImpl* renderer, + slang::IComponentType* program, + slang::ProgramLayout* programLayout, + RootShaderObjectLayout** outLayout) +{ + RootShaderObjectLayout::Builder builder(renderer, program, programLayout); + builder.addGlobalParams(programLayout->getGlobalParamsVarLayout()); + + SlangInt entryPointCount = programLayout->getEntryPointCount(); + for (SlangInt e = 0; e < entryPointCount; ++e) + { + auto slangEntryPoint = programLayout->getEntryPointByIndex(e); + + EntryPointLayout::Builder entryPointBuilder(renderer); + entryPointBuilder.addEntryPointParams(slangEntryPoint); + + RefPtr<EntryPointLayout> entryPointLayout; + SLANG_RETURN_ON_FAIL(entryPointBuilder.build(entryPointLayout.writeRef())); + + builder.addEntryPoint(entryPointLayout); + } + + SLANG_RETURN_ON_FAIL(builder.build(outLayout)); + + return SLANG_OK; +} + +Result RootShaderObjectLayout::_init(Builder const* builder) +{ + auto renderer = builder->m_renderer; + + SLANG_RETURN_ON_FAIL(Super::_init(builder)); + + m_program = builder->m_program; + m_programLayout = builder->m_programLayout; + m_entryPoints = _Move(builder->m_entryPoints); + m_pendingDataOffset = builder->m_pendingDataOffset; + m_renderer = renderer; + + // If the program has unbound specialization parameters, + // then we will avoid creating a final Vulkan pipeline layout. + // + // TODO: We should really create the information necessary + // for binding as part of a separate object, so that we have + // a clean seperation between what is needed for writing into + // a shader object vs. what is needed for binding it to the + // pipeline. We eventually need to be able to create bindable + // state objects from unspecialized programs, in order to + // support dynamic dispatch. + // + if (m_program->getSpecializationParamCount() != 0) + return SLANG_OK; + + // Otherwise, we need to create a final (bindable) layout. + // + // We will use a recursive walk to collect all the `VkDescriptorSetLayout`s + // that are required for the global scope, sub-objects, and entry points. + // + SLANG_RETURN_ON_FAIL(addAllDescriptorSets()); + + // We will also use a recursive walk to collect all the push-constant + // ranges needed for this object, sub-objects, and entry points. + // + SLANG_RETURN_ON_FAIL(addAllPushConstantRanges()); + + // Once we've collected the information across the entire + // tree of sub-objects + + // Now call Vulkan API to create a pipeline layout. + VkPipelineLayoutCreateInfo pipelineLayoutCreateInfo = {}; + pipelineLayoutCreateInfo.sType = VK_STRUCTURE_TYPE_PIPELINE_LAYOUT_CREATE_INFO; + pipelineLayoutCreateInfo.setLayoutCount = (uint32_t)m_vkDescriptorSetLayouts.getCount(); + pipelineLayoutCreateInfo.pSetLayouts = m_vkDescriptorSetLayouts.getBuffer(); + if (m_allPushConstantRanges.getCount()) + { + pipelineLayoutCreateInfo.pushConstantRangeCount = + (uint32_t)m_allPushConstantRanges.getCount(); + pipelineLayoutCreateInfo.pPushConstantRanges = m_allPushConstantRanges.getBuffer(); + } + SLANG_RETURN_ON_FAIL(m_renderer->m_api.vkCreatePipelineLayout( + m_renderer->m_api.m_device, &pipelineLayoutCreateInfo, nullptr, &m_pipelineLayout)); + return SLANG_OK; +} + +/// Add all the descriptor sets implied by this root object and sub-objects + +Result RootShaderObjectLayout::addAllDescriptorSets() +{ + SLANG_RETURN_ON_FAIL(addAllDescriptorSetsRec(this)); + + // Note: the descriptor ranges/sets for direct entry point parameters + // were already enumerated into the ranges/sets of the root object itself, + // so we don't wnat to add them again. + // + // We do however have to deal with the possibility that an entry + // point could introduce "child" descriptor sets, e.g., because it + // has a `ParameterBlock<X>` parameter. + // + for (auto& entryPoint : getEntryPoints()) + { + SLANG_RETURN_ON_FAIL(addChildDescriptorSetsRec(entryPoint.layout)); + } + + return SLANG_OK; +} + +/// Recurisvely add descriptor sets defined by `layout` and sub-objects + +Result RootShaderObjectLayout::addAllDescriptorSetsRec(ShaderObjectLayoutImpl* layout) +{ + // TODO: This logic assumes that descriptor sets are all contiguous + // and have been allocated in a global order that matches the order + // of enumeration here. + + for (auto& descSetInfo : layout->getOwnDescriptorSets()) + { + m_vkDescriptorSetLayouts.add(descSetInfo.descriptorSetLayout); + } + + SLANG_RETURN_ON_FAIL(addChildDescriptorSetsRec(layout)); + return SLANG_OK; +} + +/// Recurisvely add descriptor sets defined by sub-objects of `layout` + +Result RootShaderObjectLayout::addChildDescriptorSetsRec(ShaderObjectLayoutImpl* layout) +{ + for (auto& subObject : layout->getSubObjectRanges()) + { + auto bindingRange = layout->getBindingRange(subObject.bindingRangeIndex); + switch (bindingRange.bindingType) + { + case slang::BindingType::ParameterBlock: + SLANG_RETURN_ON_FAIL(addAllDescriptorSetsRec(subObject.layout)); + break; + + default: + if (auto subObjectLayout = subObject.layout) + { + SLANG_RETURN_ON_FAIL(addChildDescriptorSetsRec(subObject.layout)); + } + break; + } + } + + return SLANG_OK; +} + +/// Add all the push-constant ranges implied by this root object and sub-objects + +Result RootShaderObjectLayout::addAllPushConstantRanges() +{ + SLANG_RETURN_ON_FAIL(addAllPushConstantRangesRec(this)); + + for (auto& entryPoint : getEntryPoints()) + { + SLANG_RETURN_ON_FAIL(addChildPushConstantRangesRec(entryPoint.layout)); + } + + return SLANG_OK; +} + +/// Recurisvely add push-constant ranges defined by `layout` and sub-objects + +Result RootShaderObjectLayout::addAllPushConstantRangesRec(ShaderObjectLayoutImpl* layout) +{ + // TODO: This logic assumes that push-constant ranges are all contiguous + // and have been allocated in a global order that matches the order + // of enumeration here. + + for (auto pushConstantRange : layout->getOwnPushConstantRanges()) + { + pushConstantRange.offset = m_totalPushConstantSize; + m_totalPushConstantSize += pushConstantRange.size; + + m_allPushConstantRanges.add(pushConstantRange); + } + + SLANG_RETURN_ON_FAIL(addChildPushConstantRangesRec(layout)); + return SLANG_OK; +} + +/// Recurisvely add push-constant ranges defined by sub-objects of `layout` + +Result RootShaderObjectLayout::addChildPushConstantRangesRec(ShaderObjectLayoutImpl* layout) +{ + for (auto& subObject : layout->getSubObjectRanges()) + { + if (auto subObjectLayout = subObject.layout) + { + SLANG_RETURN_ON_FAIL(addAllPushConstantRangesRec(subObject.layout)); + } + } + + return SLANG_OK; +} + +Result RootShaderObjectLayout::Builder::build(RootShaderObjectLayout** outLayout) +{ + RefPtr<RootShaderObjectLayout> layout = new RootShaderObjectLayout(); + SLANG_RETURN_ON_FAIL(layout->_init(this)); + returnRefPtrMove(outLayout, layout); + return SLANG_OK; +} + +void RootShaderObjectLayout::Builder::addGlobalParams( + slang::VariableLayoutReflection* globalsLayout) +{ + setElementTypeLayout(globalsLayout->getTypeLayout()); + + // We need to populate our descriptor sets/ranges with information + // from the layout of the global scope. + // + // While we expect that the parameter in the global scope start + // at an offset of zero, it is also worth querying the offset + // information because it could impact the locations assigned + // to "pending" data in the case of static specialization. + // + BindingOffset offset(globalsLayout); + + // Note: We are adding descriptor ranges here based directly on + // the type of the global-scope layout. The type layout for the + // global scope will either be something like a `struct GlobalParams` + // that contains all the global-scope parameters or a `ConstantBuffer<GlobalParams>` + // and in either case the `_addDescriptorRangesAsValue` can properly + // add all the ranges implied. + // + // As a result we don't require any special-case logic here to + // deal with the possibility of a "default" constant buffer allocated + // for global-scope parameters of uniform/ordinary type. + // + _addDescriptorRangesAsValue(globalsLayout->getTypeLayout(), offset); + + // We want to keep track of the offset that was applied to "pending" + // data because we will need it again later when it comes time to + // actually bind things. + // + m_pendingDataOffset = offset.pending; +} + +void RootShaderObjectLayout::Builder::addEntryPoint(EntryPointLayout* entryPointLayout) +{ + auto slangEntryPointLayout = entryPointLayout->getSlangLayout(); + auto entryPointVarLayout = slangEntryPointLayout->getVarLayout(); + + // The offset information for each entry point needs to + // be adjusted by any offset for "pending" data that + // was recorded in the global-scope layout. + // + // TODO(tfoley): Double-check that this is correct. + + BindingOffset entryPointOffset(entryPointVarLayout); + entryPointOffset.pending += m_pendingDataOffset; + + EntryPointInfo info; + info.layout = entryPointLayout; + info.offset = entryPointOffset; + + // Similar to the case for the global scope, we expect the + // type layout for the entry point parameters to be either + // a `struct EntryPointParams` or a `PushConstantBuffer<EntryPointParams>`. + // Rather than deal with the different cases here, we will + // trust the `_addDescriptorRangesAsValue` code to handle + // either case correctly. + // + _addDescriptorRangesAsValue(entryPointVarLayout->getTypeLayout(), entryPointOffset); + + m_entryPoints.add(info); +} + +} // namespace vk +} // namespace gfx |