Adds compute shaders support and scaling with FSR.

Fixes #6842

Comments By: phosit, Stan
Differential Revision: https://code.wildfiregames.com/D5218
This was SVN commit r28010.
This commit is contained in:
vladislavbelov
2024-01-17 19:40:27 +00:00
parent f9f798158a
commit e3f46bb809
63 changed files with 5337 additions and 118 deletions
@@ -1,4 +1,4 @@
/* Copyright (C) 2023 Wildfire Games.
/* Copyright (C) 2024 Wildfire Games.
* This file is part of 0 A.D.
*
* 0 A.D. is free software: you can redistribute it and/or modify
@@ -171,6 +171,7 @@ std::unique_ptr<CShaderProgram> CShaderProgram::Create(
#define EL(x) const int el_##x = programXeroFile.GetElementID(#x)
#define AT(x) const int at_##x = programXeroFile.GetAttributeID(#x)
EL(binding);
EL(compute);
EL(descriptor_set);
EL(descriptor_sets);
EL(fragment);
@@ -231,6 +232,10 @@ std::unique_ptr<CShaderProgram> CShaderProgram::Create(
uint32_t texturesDescriptorSetSize = 0;
std::unordered_map<CStrIntern, uint32_t> textureMapping;
VkDescriptorType storageImageDescriptorType = VK_DESCRIPTOR_TYPE_MAX_ENUM;
uint32_t storageImageDescriptorSetSize = 0;
std::unordered_map<CStrIntern, uint32_t> storageImageMapping;
auto addDescriptorSets = [&](const XMBElement& element) -> bool
{
const bool useDescriptorIndexing =
@@ -309,6 +314,23 @@ std::unique_ptr<CShaderProgram> CShaderProgram::Create(
texturesDescriptorSetSize =
std::max(texturesDescriptorSetSize, binding + 1);
}
else if (type == "storageImage" || type == "storageBuffer")
{
const CStrIntern name{attributes.GetNamedItem(at_name)};
storageImageMapping[name] = binding;
storageImageDescriptorSetSize =
std::max(storageImageDescriptorSetSize, binding + 1);
const VkDescriptorType descriptorType = type == "storageBuffer"
? VK_DESCRIPTOR_TYPE_STORAGE_BUFFER
: VK_DESCRIPTOR_TYPE_STORAGE_IMAGE;
if (storageImageDescriptorType == VK_DESCRIPTOR_TYPE_MAX_ENUM)
storageImageDescriptorType = descriptorType;
else if (storageImageDescriptorType != descriptorType)
{
LOGERROR("Shader should have storages of the same type.");
return false;
}
}
else
{
LOGERROR("Unsupported binding: '%s'", type.c_str());
@@ -325,6 +347,13 @@ std::unique_ptr<CShaderProgram> CShaderProgram::Create(
{
if (programChild.GetNodeName() == el_vertex)
{
if (shaderProgram->m_PipelineBindPoint != VK_PIPELINE_BIND_POINT_MAX_ENUM &&
shaderProgram->m_PipelineBindPoint != VK_PIPELINE_BIND_POINT_GRAPHICS)
{
LOGERROR("Shader program can't mix different pipelines: '%s'.", name.c_str());
return nullptr;
}
shaderProgram->m_PipelineBindPoint = VK_PIPELINE_BIND_POINT_GRAPHICS;
const VfsPath shaderModulePath =
L"shaders/" + programChild.GetAttributes().GetNamedItem(at_file).FromUTF8();
shaderProgram->m_FileDependencies.emplace_back(shaderModulePath);
@@ -386,6 +415,13 @@ std::unique_ptr<CShaderProgram> CShaderProgram::Create(
}
else if (programChild.GetNodeName() == el_fragment)
{
if (shaderProgram->m_PipelineBindPoint != VK_PIPELINE_BIND_POINT_MAX_ENUM &&
shaderProgram->m_PipelineBindPoint != VK_PIPELINE_BIND_POINT_GRAPHICS)
{
LOGERROR("Shader program can't mix different pipelines: '%s'.", name.c_str());
return nullptr;
}
shaderProgram->m_PipelineBindPoint = VK_PIPELINE_BIND_POINT_GRAPHICS;
const VfsPath shaderModulePath =
L"shaders/" + programChild.GetAttributes().GetNamedItem(at_file).FromUTF8();
shaderProgram->m_FileDependencies.emplace_back(shaderModulePath);
@@ -413,6 +449,42 @@ std::unique_ptr<CShaderProgram> CShaderProgram::Create(
}
}
}
else if (programChild.GetNodeName() == el_compute)
{
if (shaderProgram->m_PipelineBindPoint != VK_PIPELINE_BIND_POINT_MAX_ENUM &&
shaderProgram->m_PipelineBindPoint != VK_PIPELINE_BIND_POINT_COMPUTE)
{
LOGERROR("Shader program can't mix different pipelines: '%s'.", name.c_str());
return nullptr;
}
shaderProgram->m_PipelineBindPoint = VK_PIPELINE_BIND_POINT_COMPUTE;
const VfsPath shaderModulePath =
L"shaders/" + programChild.GetAttributes().GetNamedItem(at_file).FromUTF8();
shaderProgram->m_FileDependencies.emplace_back(shaderModulePath);
shaderProgram->m_ShaderModules.emplace_back(
CreateShaderModule(device, shaderModulePath));
if (shaderProgram->m_ShaderModules.back() == VK_NULL_HANDLE)
return nullptr;
VkPipelineShaderStageCreateInfo computeShaderStageInfo{};
computeShaderStageInfo.sType = VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO;
computeShaderStageInfo.stage = VK_SHADER_STAGE_COMPUTE_BIT;
computeShaderStageInfo.module = shaderProgram->m_ShaderModules.back();
computeShaderStageInfo.pName = "main";
shaderProgram->m_Stages.emplace_back(std::move(computeShaderStageInfo));
XERO_ITER_EL(programChild, stageChild)
{
if (stageChild.GetNodeName() == el_push_constant)
{
if (!addPushConstant(stageChild, VK_SHADER_STAGE_COMPUTE_BIT))
return nullptr;
}
else if (stageChild.GetNodeName() == el_descriptor_sets)
{
if (!addDescriptorSets(stageChild))
return nullptr;
}
}
}
}
if (shaderProgram->m_Stages.empty())
@@ -421,6 +493,8 @@ std::unique_ptr<CShaderProgram> CShaderProgram::Create(
return nullptr;
}
ENSURE(shaderProgram->m_PipelineBindPoint != VK_PIPELINE_BIND_POINT_MAX_ENUM);
for (size_t index = 0; index < shaderProgram->m_PushConstants.size(); ++index)
shaderProgram->m_PushConstantMapping[shaderProgram->m_PushConstants[index].name] = index;
std::vector<VkPushConstantRange> pushConstantRanges;
@@ -482,6 +556,12 @@ std::unique_ptr<CShaderProgram> CShaderProgram::Create(
device, VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER, texturesDescriptorSetSize, std::move(textureMapping));
layouts.emplace_back(shaderProgram->m_TextureBinding->GetDescriptorSetLayout());
}
if (storageImageDescriptorSetSize > 0)
{
shaderProgram->m_StorageImageBinding.emplace(
device, storageImageDescriptorType, storageImageDescriptorSetSize, std::move(storageImageMapping));
layouts.emplace_back(shaderProgram->m_StorageImageBinding->GetDescriptorSetLayout());
}
VkPipelineLayoutCreateInfo pipelineLayoutCreateInfo{};
pipelineLayoutCreateInfo.sType = VK_STRUCTURE_TYPE_PIPELINE_LAYOUT_CREATE_INFO;
@@ -527,6 +607,8 @@ int32_t CShaderProgram::GetBindingSlot(const CStrIntern name) const
return it->second + m_PushConstants.size();
if (const int32_t bindingSlot = m_TextureBinding.has_value() ? m_TextureBinding->GetBindingSlot(name) : -1; bindingSlot != -1)
return bindingSlot + m_PushConstants.size() + m_UniformMapping.size();
if (const int32_t bindingSlot = m_StorageImageBinding.has_value() ? m_StorageImageBinding->GetBindingSlot(name) : -1; bindingSlot != -1)
return bindingSlot + m_PushConstants.size() + m_UniformMapping.size() + (m_TextureBinding.has_value() ? m_TextureBinding->GetBoundDeviceObjects().size() : 0);
return -1;
}
@@ -551,6 +633,8 @@ void CShaderProgram::Unbind()
{
if (m_TextureBinding.has_value())
m_TextureBinding->Unbind();
if (m_StorageImageBinding.has_value())
m_StorageImageBinding->Unbind();
}
void CShaderProgram::PreDraw(CRingCommandContext& commandContext)
@@ -578,16 +662,63 @@ void CShaderProgram::PreDraw(CRingCommandContext& commandContext)
}
}
void CShaderProgram::PreDispatch(
CRingCommandContext& commandContext)
{
PreDraw(commandContext);
if (m_StorageImageBinding.has_value())
for (CTexture* texture : m_StorageImageBinding->GetBoundDeviceObjects())
if (texture)
{
if (!(texture->GetUsage() & ITexture::Usage::SAMPLED) && texture->IsInitialized())
continue;
VkImageLayout oldLayout = VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL;
if (!texture->IsInitialized())
oldLayout = VK_IMAGE_LAYOUT_UNDEFINED;
Utilities::SetTextureLayout(
commandContext.GetCommandBuffer(), texture,
oldLayout,
VK_IMAGE_LAYOUT_GENERAL,
VK_ACCESS_SHADER_READ_BIT, VK_ACCESS_SHADER_READ_BIT | VK_ACCESS_SHADER_WRITE_BIT,
VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT | VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT,
VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT);
}
}
void CShaderProgram::PostDispatch(CRingCommandContext& commandContext)
{
if (m_StorageImageBinding.has_value())
for (CTexture* texture : m_StorageImageBinding->GetBoundDeviceObjects())
if (texture)
{
if (!(texture->GetUsage() & ITexture::Usage::SAMPLED) && texture->IsInitialized())
continue;
Utilities::SetTextureLayout(
commandContext.GetCommandBuffer(), texture,
VK_IMAGE_LAYOUT_GENERAL,
VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL,
VK_ACCESS_SHADER_READ_BIT | VK_ACCESS_SHADER_WRITE_BIT, VK_ACCESS_SHADER_READ_BIT,
VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT,
VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT | VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT);
}
}
void CShaderProgram::BindOutdatedDescriptorSets(
CRingCommandContext& commandContext)
{
// TODO: combine calls after more sets to bind.
PS::StaticVector<std::tuple<uint32_t, VkDescriptorSet>, 1> descriptortSets;
PS::StaticVector<std::tuple<uint32_t, VkDescriptorSet>, 2> descriptortSets;
if (m_TextureBinding.has_value() && m_TextureBinding->IsOutdated())
{
constexpr uint32_t TEXTURE_BINDING_SET = 1u;
descriptortSets.emplace_back(TEXTURE_BINDING_SET, m_TextureBinding->UpdateAndReturnDescriptorSet());
}
if (m_StorageImageBinding.has_value() && m_StorageImageBinding->IsOutdated())
{
constexpr uint32_t STORAGE_IMAGE_BINDING_SET = 2u;
descriptortSets.emplace_back(STORAGE_IMAGE_BINDING_SET, m_StorageImageBinding->UpdateAndReturnDescriptorSet());
}
for (const auto [firstSet, descriptorSet] : descriptortSets)
{
@@ -687,6 +818,17 @@ void CShaderProgram::SetTexture(const int32_t bindingSlot, CTexture* texture)
}
}
void CShaderProgram::SetStorageTexture(const int32_t bindingSlot, CTexture* texture)
{
if (bindingSlot < 0)
return;
const int32_t offset = static_cast<int32_t>(m_PushConstants.size() + m_UniformMapping.size() + (m_TextureBinding.has_value() ? m_TextureBinding->GetBoundDeviceObjects().size() : 0));
ENSURE(bindingSlot >= offset);
ENSURE(m_StorageImageBinding.has_value());
const uint32_t index = bindingSlot - offset;
m_StorageImageBinding->SetObject(index, texture);
}
} // namespace Vulkan
} // namespace Backend