nomic-ai/kompute
1
0
Fork 0
mirror of https://github.com/nomic-ai/kompute.git synced 2026-05-11 17:09:59 +00:00
Code Issues Packages Projects Releases Wiki Activity

Refactored descriptor sets

Browse Source
This commit is contained in:
Alejandro Saucedo
2020-08-22 10:15:34 +01:00
parent eecadbe36b
commit 9d97ca07a7
9 changed files with 56 additions and 45 deletions
Download Patch File Download Diff File Expand all files Collapse all files

63
src/Algorithm.cpp
View File

@@ -32,8 +32,6 @@ void Algorithm::init(std::string shaderFilePath,
std::vector<std::shared_ptr<Tensor>> tensorParams) {
SPDLOG_DEBUG("Kompute Algorithm init started");
spdlog::info("Loading shader with file path {}", shaderFilePath);
// TODO: Move to util function
this->createParameters(tensorParams);
this->createShaderModule(shaderFilePath);
@@ -41,29 +39,29 @@ void Algorithm::init(std::string shaderFilePath,
}
void Algorithm::createParameters(std::vector<std::shared_ptr<Tensor>>& tensorParams) {
std::vector<vk::DescriptorPoolSize> descriptorPoolSizes;
SPDLOG_DEBUG("Kompute Algorithm createParameters started");
for (std::shared_ptr<Tensor> tensorParam : tensorParams) {
descriptorPoolSizes.push_back(
// TODO: Explore design for having multiple descriptor pool sizes
std::vector<vk::DescriptorPoolSize> descriptorPoolSizes = {
vk::DescriptorPoolSize(
vk::DescriptorType::eStorageBuffer,
1 // Descriptor count
)
);
}
};
// TODO: Explore design for having more than 1 set configurable
vk::DescriptorPoolCreateInfo descriptorPoolInfo(
vk::DescriptorPoolCreateFlags(),
1, // Max sets
descriptorPoolSizes.size(),
static_cast<uint32_t>(descriptorPoolSizes.size()),
descriptorPoolSizes.data());
SPDLOG_DEBUG("Kompute Algorithm creating descriptor pool");
this->mDescriptorPool = std::make_shared<vk::DescriptorPool>();
this->mDevice->createDescriptorPool(&descriptorPoolInfo, nullptr, this->mDescriptorPool.get());
// TODO: Explore allowing descriptor set bind index
std::vector<vk::DescriptorSetLayoutBinding> descriptorSetBindings;
// TODO: Explore allowing descriptor set bind index to be configurable by user to specify which tensors woudl go on each binding
for (size_t i = 0; i < tensorParams.size(); i++) {
descriptorSetBindings.push_back(
vk::DescriptorSetLayoutBinding(
@@ -77,10 +75,11 @@ void Algorithm::createParameters(std::vector<std::shared_ptr<Tensor>>& tensorPar
// This is the component that is fed into the pipeline
vk::DescriptorSetLayoutCreateInfo descriptorSetLayoutInfo(
vk::DescriptorSetLayoutCreateFlags(),
descriptorSetBindings.size(),
static_cast<uint32_t>(descriptorSetBindings.size()),
descriptorSetBindings.data()
);
SPDLOG_DEBUG("Kompute Algorithm creating descriptor set layout");
// TODO: We createa signle descriptor set layout which would have to be extended if multiple set layouts to be supported
this->mDescriptorSetLayout = std::make_shared<vk::DescriptorSetLayout>();
this->mDevice->createDescriptorSetLayout(&descriptorSetLayoutInfo, nullptr, this->mDescriptorSetLayout.get());
@@ -90,30 +89,40 @@ void Algorithm::createParameters(std::vector<std::shared_ptr<Tensor>>& tensorPar
1, // Descriptor set layout count
this->mDescriptorSetLayout.get());
SPDLOG_DEBUG("Kompute Algorithm allocating descriptor sets");
std::vector<vk::DescriptorSet> descriptorSets =
this->mDevice->allocateDescriptorSets(descriptorSetAllocateInfo);
if (descriptorSets.size() != tensorParams.size()) {
throw std::runtime_error("Number of descriptor sets does not match number of paramters");
}
this->mDescriptorSet = std::make_shared<vk::DescriptorSet>();
this->mDevice->allocateDescriptorSets(&descriptorSetAllocateInfo, this->mDescriptorSet.get());
std::vector<vk::WriteDescriptorSet> computeWriteDescriptorSets;
for (size_t i = 0; i < descriptorSets.size(); i++) {
for (size_t i = 0; i < tensorParams.size(); i++) {
std::shared_ptr<Tensor> currTensor = tensorParams[i];
vk::DescriptorSet& currDescriptorSet = descriptorSets[i];
this->mDescriptorSets.push_back(std::make_shared<vk::DescriptorSet>(currDescriptorSet));
vk::DescriptorBufferInfo descriptorBufferInfo = currTensor->constructDescriptorBufferInfo();
vk::DescriptorBufferInfo descriptorBufferInfo =
tensorParams[i]->constructDescriptorBufferInfo();
// TODO: Explore design exposing the destination array element
computeWriteDescriptorSets.push_back(
vk::WriteDescriptorSet());
vk::WriteDescriptorSet(
*this->mDescriptorSet,
i, // Destination binding
0, // Destination array element
1, // Descriptor count
vk::DescriptorType::eStorageBuffer,
nullptr, // Descriptor image info
&descriptorBufferInfo));
}
SPDLOG_DEBUG("Kompute Algorithm updating descriptor sets");
this->mDevice->updateDescriptorSets(computeWriteDescriptorSets, nullptr);
SPDLOG_DEBUG("Kompue Algorithm successfully run init");
}
void Algorithm::createShaderModule(std::string shaderFilePath) {
SPDLOG_DEBUG("Kompute Algorithm createShaderModule started");
std::ifstream fileStream(
shaderFilePath, std::ios::binary | std::ios::in | std::ios::ate);
@@ -166,13 +175,13 @@ void Algorithm::recordDispatch(uint32_t x, uint32_t y, uint32_t z) {
this->mCommandBuffer->bindPipeline(vk::PipelineBindPoint::eCompute, *this->mPipeline);
// TODO: Simplify interaction given we store array of pointers
std::vector<vk::DescriptorSet> descriptorSetRefs(this->mDescriptorSets.size());
for (size_t i = 0; i < this->mDescriptorSets.size(); i++) {
descriptorSetRefs[i] = *this->mDescriptorSets[i];
}
this->mCommandBuffer->bindDescriptorSets(vk::PipelineBindPoint::eCompute, *this->mPipelineLayout, 0, descriptorSetRefs, nullptr);
this->mCommandBuffer->bindDescriptorSets(
vk::PipelineBindPoint::eCompute,
*this->mPipelineLayout,
0, // First set
*this->mDescriptorSet,
nullptr // Dispatcher
);
this->mCommandBuffer->dispatch(x, y, z);
}