Refactor Vulkan Descriptor Type and Shader Stage set up.

luigifcruz · May 26, 2024 · c9237a8 · c9237a8
1 parent 8f8c510
commit c9237a8
Show file tree

Hide file tree

Showing 12 changed files with 70 additions and 67 deletions.
diff --git a/include/jetstream/render/base/buffer.hh b/include/jetstream/render/base/buffer.hh
@@ -61,6 +61,14 @@ class Buffer {
     Config config;
 };
 
+static constexpr Buffer::Target operator&(Buffer::Target a, Buffer::Target b) {
+    return static_cast<Buffer::Target>(static_cast<U64>(a) & static_cast<U64>(b));
+}
+
+static constexpr Buffer::Target operator|(Buffer::Target a, Buffer::Target b) {
+    return static_cast<Buffer::Target>(static_cast<U64>(a) | static_cast<U64>(b));
+}
+
 }  // namespace Jetstream::Render
 
 #endif
diff --git a/include/jetstream/render/metal/kernel.hh b/include/jetstream/render/metal/kernel.hh
@@ -21,9 +21,6 @@ class KernelImp<Device::Metal> : public Kernel {
 
     std::vector<std::shared_ptr<BufferImp<Device::Metal>>> buffers;
 
-    static Result checkShaderCompilation(U64);
-    static Result checkProgramCompilation(U64);
-
     friend class SurfaceImp<Device::Metal>; 
 };
 

diff --git a/include/jetstream/render/metal/program.hh b/include/jetstream/render/metal/program.hh
@@ -23,9 +23,6 @@ class ProgramImp<Device::Metal> : public Program {
     std::vector<std::shared_ptr<TextureImp<Device::Metal>>> textures;
     std::vector<std::pair<std::shared_ptr<BufferImp<Device::Metal>>, Program::Target>> buffers;
 
-    static Result checkShaderCompilation(U64);
-    static Result checkProgramCompilation(U64);
-
     friend class SurfaceImp<Device::Metal>; 
 };
 

diff --git a/include/jetstream/render/vulkan/buffer.hh b/include/jetstream/render/vulkan/buffer.hh
@@ -25,14 +25,9 @@ class BufferImp<Device::Vulkan> : public Buffer {
         return buffer;
     }
 
-    constexpr const VkDescriptorType& getDescriptorType() const {
-        return descriptorType;
-    }
-
  private:
     VkBuffer buffer;
     VkDeviceMemory memory;
-    VkDescriptorType descriptorType;
 
     friend class SurfaceImp<Device::Vulkan>;
     friend class ProgramImp<Device::Vulkan>;

diff --git a/include/jetstream/render/vulkan/kernel.hh b/include/jetstream/render/vulkan/kernel.hh
@@ -26,8 +26,7 @@ class KernelImp<Device::Vulkan> : public Kernel {
 
     std::vector<std::shared_ptr<BufferImp<Device::Vulkan>>> buffers;
 
-    static Result checkShaderCompilation(U64);
-    static Result checkProgramCompilation(U64);
+    static VkDescriptorType BufferDescriptorType(const std::shared_ptr<Buffer>& buffer);
 
     friend class SurfaceImp<Device::Vulkan>; 
 };

diff --git a/include/jetstream/render/vulkan/program.hh b/include/jetstream/render/vulkan/program.hh
@@ -29,7 +29,8 @@ class ProgramImp<Device::Vulkan> : public Program {
     std::vector<std::shared_ptr<TextureImp<Device::Vulkan>>> textures;
     std::vector<std::pair<std::shared_ptr<BufferImp<Device::Vulkan>>, Program::Target>> buffers;
 
-    static VkShaderStageFlags TargetToVulkan(const Program::Target& target);
+    static VkShaderStageFlags TargetToShaderStage(const Program::Target& target);
+    static VkDescriptorType BufferDescriptorType(const std::shared_ptr<Buffer>& buffer);
 
     friend class SurfaceImp<Device::Vulkan>; 
 };

diff --git a/include/jetstream/render/webgpu/program.hh b/include/jetstream/render/webgpu/program.hh
@@ -30,8 +30,6 @@ class ProgramImp<Device::WebGPU> : public Program {
     std::vector<std::shared_ptr<TextureImp<Device::WebGPU>>> textures;
     std::vector<std::pair<std::shared_ptr<BufferImp<Device::WebGPU>>, Program::Target>> buffers;
 
-    static Result checkShaderCompilation(U64);
-    static Result checkProgramCompilation(U64);
     static wgpu::ShaderStage TargetToWebGPU(const Program::Target& target);
 
     friend class SurfaceImp<Device::WebGPU>; 

diff --git a/src/memory/vulkan/buffer.cc b/src/memory/vulkan/buffer.cc
@@ -241,7 +241,8 @@ Implementation::TensorBuffer(std::shared_ptr<TensorStorageMetadata>&,
     // TODO: Add a global way to specify usage.
     bufferInfo.usage =  VK_BUFFER_USAGE_TRANSFER_SRC_BIT | 
                         VK_BUFFER_USAGE_TRANSFER_DST_BIT | 
-                        VK_BUFFER_USAGE_VERTEX_BUFFER_BIT;
+                        VK_BUFFER_USAGE_VERTEX_BUFFER_BIT |
+                        VK_BUFFER_USAGE_STORAGE_BUFFER_BIT;
     bufferInfo.sharingMode = VK_SHARING_MODE_EXCLUSIVE;
     bufferInfo.pNext = &extImageCreateInfo;
 

diff --git a/src/modules/lineplot/generic.cc b/src/modules/lineplot/generic.cc
@@ -96,7 +96,7 @@ Result Lineplot<D, T>::createPresent() {
         cfg.buffer = gridPoints.data();
         cfg.elementByteSize = sizeof(F32);
         cfg.size = gridPoints.size();
-        cfg.target = Render::Buffer::Target::VERTEX;
+        cfg.target = Render::Buffer::Target::STORAGE;
         cfg.enableZeroCopy = false;
         JST_CHECK(window->build(gridPointsBuffer, cfg));
     }
@@ -108,7 +108,7 @@ Result Lineplot<D, T>::createPresent() {
         cfg.buffer = buffer;
         cfg.elementByteSize = sizeof(F32);
         cfg.size = gridVertices.size();
-        cfg.target = Render::Buffer::Target::VERTEX;
+        cfg.target = Render::Buffer::Target::VERTEX | Render::Buffer::Target::STORAGE;
         cfg.enableZeroCopy = enableZeroCopy;
         JST_CHECK(window->build(gridVerticesBuffer, cfg));
     }
@@ -165,7 +165,7 @@ Result Lineplot<D, T>::createPresent() {
         cfg.buffer = buffer;
         cfg.elementByteSize = sizeof(F32);
         cfg.size = signalPoints.size();
-        cfg.target = Render::Buffer::Target::VERTEX;
+        cfg.target = Render::Buffer::Target::STORAGE;
         cfg.enableZeroCopy = enableZeroCopy;
         JST_CHECK(window->build(signalPointsBuffer, cfg));
     }
@@ -177,7 +177,7 @@ Result Lineplot<D, T>::createPresent() {
         cfg.buffer = buffer;
         cfg.elementByteSize = sizeof(F32);
         cfg.size = signalVertices.size();
-        cfg.target = Render::Buffer::Target::VERTEX;
+        cfg.target = Render::Buffer::Target::VERTEX | Render::Buffer::Target::STORAGE;
         cfg.enableZeroCopy = enableZeroCopy;
         JST_CHECK(window->build(signalVerticesBuffer, cfg));
     }

diff --git a/src/render/vulkan/buffer.cc b/src/render/vulkan/buffer.cc
@@ -14,25 +14,6 @@ Result Implementation::create() {
     auto& device = Backend::State<Device::Vulkan>()->getDevice();
     auto& physicalDevice = Backend::State<Device::Vulkan>()->getPhysicalDevice();
 
-    switch (config.target) {
-        case Target::VERTEX:
-            break;
-        case Target::VERTEX_INDICES:
-            break;
-        case Target::STORAGE:
-            descriptorType = VK_DESCRIPTOR_TYPE_STORAGE_BUFFER;
-            break;
-        case Target::UNIFORM:
-            descriptorType = VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER;
-            break;
-        case Target::STORAGE_DYNAMIC:
-            descriptorType = VK_DESCRIPTOR_TYPE_STORAGE_BUFFER_DYNAMIC;
-            break;
-        case Target::UNIFORM_DYNAMIC:
-            descriptorType = VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER_DYNAMIC;
-            break;
-    }
-
     if (config.enableZeroCopy) {
         buffer = reinterpret_cast<VkBuffer>(config.buffer);
     } else {
@@ -41,25 +22,23 @@ Result Implementation::create() {
         VkBufferUsageFlags bufferUsageFlag = 0;
         bufferUsageFlag |= VK_BUFFER_USAGE_TRANSFER_SRC_BIT;
         bufferUsageFlag |= VK_BUFFER_USAGE_TRANSFER_DST_BIT;
-        switch (config.target) {
-            case Target::VERTEX:
-                bufferUsageFlag |= VK_BUFFER_USAGE_VERTEX_BUFFER_BIT;
-                break;
-            case Target::VERTEX_INDICES:
-                bufferUsageFlag |= VK_BUFFER_USAGE_INDEX_BUFFER_BIT;
-                break;
-            case Target::STORAGE:
-                bufferUsageFlag |= VK_BUFFER_USAGE_STORAGE_BUFFER_BIT;
-                break;
-            case Target::UNIFORM:
-                bufferUsageFlag |= VK_BUFFER_USAGE_UNIFORM_BUFFER_BIT;
-                break;
-            case Target::STORAGE_DYNAMIC:
-                bufferUsageFlag |= VK_BUFFER_USAGE_STORAGE_BUFFER_BIT;
-                break;
-            case Target::UNIFORM_DYNAMIC:
-                bufferUsageFlag |= VK_BUFFER_USAGE_UNIFORM_BUFFER_BIT;
-                break;
+
+        if ((config.target & Target::VERTEX) == Target::VERTEX) {
+            bufferUsageFlag |= VK_BUFFER_USAGE_VERTEX_BUFFER_BIT;
+        }
+
+        if ((config.target & Target::VERTEX_INDICES) == Target::VERTEX_INDICES) {
+            bufferUsageFlag |= VK_BUFFER_USAGE_INDEX_BUFFER_BIT;
+        }
+
+        if ((config.target & Target::STORAGE) == Target::STORAGE ||
+            (config.target & Target::STORAGE_DYNAMIC) == Target::STORAGE_DYNAMIC) {
+            bufferUsageFlag |= VK_BUFFER_USAGE_STORAGE_BUFFER_BIT;
+        }
+
+        if ((config.target & Target::UNIFORM) == Target::UNIFORM ||
+            (config.target & Target::UNIFORM_DYNAMIC) == Target::UNIFORM_DYNAMIC) {
+            bufferUsageFlag |= VK_BUFFER_USAGE_UNIFORM_BUFFER_BIT;
         }
 
         // Create buffer.

diff --git a/src/render/vulkan/kernel.cc b/src/render/vulkan/kernel.cc
@@ -43,8 +43,7 @@ Result Implementation::create() {
 
         VkDescriptorSetLayoutBinding binding{};
         binding.binding = i;
-        // TODO: Why buffer->getDescriptorType() doesn't work?
-        binding.descriptorType = VK_DESCRIPTOR_TYPE_STORAGE_BUFFER;
+        binding.descriptorType = BufferDescriptorType(buffer);
         binding.descriptorCount = 1;
         binding.stageFlags = VK_SHADER_STAGE_COMPUTE_BIT;
 
@@ -89,8 +88,7 @@ Result Implementation::create() {
         descriptorWriteBuffer.dstSet = descriptorSet;
         descriptorWriteBuffer.dstBinding = i;
         descriptorWriteBuffer.dstArrayElement = 0;
-        // TODO: Why buffer->getDescriptorType() doesn't work?
-        descriptorWriteBuffer.descriptorType = VK_DESCRIPTOR_TYPE_STORAGE_BUFFER;
+        descriptorWriteBuffer.descriptorType = BufferDescriptorType(buffer);
         descriptorWriteBuffer.descriptorCount = 1;
         descriptorWriteBuffer.pBufferInfo = &bufferInfo;
 
@@ -169,4 +167,19 @@ Result Implementation::encode(VkCommandBuffer& commandBuffer) {
     return Result::SUCCESS;
 }
 
+VkDescriptorType Implementation::BufferDescriptorType(const std::shared_ptr<Buffer>& buffer) {
+    const auto& bufferType = buffer->getConfig().target;
+
+    if ((bufferType & Buffer::Target::UNIFORM) == Buffer::Target::UNIFORM) {
+        return VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER;
+    }
+
+    if ((bufferType & Buffer::Target::STORAGE) == Buffer::Target::STORAGE) {
+        return VK_DESCRIPTOR_TYPE_STORAGE_BUFFER;
+    }
+
+    JST_ERROR("[VULKAN] Invalid buffer usage.");
+    throw Result::ERROR;
+}
+
 }  // namespace Jetstream::Render
diff --git a/src/render/vulkan/program.cc b/src/render/vulkan/program.cc
@@ -67,9 +67,9 @@ Result Implementation::create(VkRenderPass& renderPass,
         auto& [buffer, target] = buffers[i];
 
         VkDescriptorSetLayoutBinding binding{};
-        binding.descriptorType = buffer->getDescriptorType();
+        binding.descriptorType = BufferDescriptorType(buffer);
         binding.descriptorCount = 1;
-        binding.stageFlags = TargetToVulkan(target);
+        binding.stageFlags = TargetToShaderStage(target);
         binding.binding = bindingOffset++;
 
         bindings.push_back(binding);
@@ -130,7 +130,7 @@ Result Implementation::create(VkRenderPass& renderPass,
         descriptorWriteBuffer.dstSet = descriptorSet;
         descriptorWriteBuffer.dstBinding = bindingOffset++;
         descriptorWriteBuffer.dstArrayElement = 0;
-        descriptorWriteBuffer.descriptorType = buffer->getDescriptorType();
+        descriptorWriteBuffer.descriptorType = BufferDescriptorType(buffer);
         descriptorWriteBuffer.descriptorCount = 1;
         descriptorWriteBuffer.pBufferInfo = &bufferInfo;
 
@@ -341,7 +341,7 @@ Result Implementation::encode(VkCommandBuffer& commandBuffer, VkRenderPass&) {
     return Result::SUCCESS;
 }
 
-VkShaderStageFlags Implementation::TargetToVulkan(const Program::Target& target) {
+VkShaderStageFlags Implementation::TargetToShaderStage(const Program::Target& target) {
     VkShaderStageFlags flags = 0;
 
     if (static_cast<U8>(target & Program::Target::VERTEX) > 0) {
@@ -355,4 +355,19 @@ VkShaderStageFlags Implementation::TargetToVulkan(const Program::Target& target)
     return flags;
 }
 
+VkDescriptorType Implementation::BufferDescriptorType(const std::shared_ptr<Buffer>& buffer) {
+    const auto& bufferType = buffer->getConfig().target;
+
+    if ((bufferType & Buffer::Target::UNIFORM) == Buffer::Target::UNIFORM) {
+        return VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER;
+    }
+
+    if ((bufferType & Buffer::Target::STORAGE) == Buffer::Target::STORAGE) {
+        return VK_DESCRIPTOR_TYPE_STORAGE_BUFFER;
+    }
+
+    JST_ERROR("[VULKAN] Invalid buffer usage.");
+    throw Result::ERROR;
+}
+
 }  // namespace Jetstream::Render