From d897ac376b358ee1213e2dd03db4174df9ee61b9 Mon Sep 17 00:00:00 2001 From: Jorge Pineda Date: Tue, 2 Jul 2024 09:34:15 -0700 Subject: [PATCH 1/2] [ET-VK] Move `ParamsBuffer` and `StorageBuffer` to standalone files Includes the renaming of `UniformParamsBuffer` to `ParamsBuffer` for brevity. These objects aren't tightly coupled to `Context` and hence they are better placed in standalone files. Differential Revision: [D59281543](https://our.internmc.facebook.com/intern/diff/D59281543/) [ghstack-poisoned] --- backends/vulkan/runtime/api/Context.cpp | 59 ---------- backends/vulkan/runtime/api/Context.h | 104 ------------------ backends/vulkan/runtime/api/ParamsBuffer.cpp | 67 +++++++++++ backends/vulkan/runtime/api/ParamsBuffer.h | 68 ++++++++++++ backends/vulkan/runtime/api/StorageBuffer.h | 70 ++++++++++++ backends/vulkan/runtime/api/Tensor.cpp | 9 +- backends/vulkan/runtime/api/Tensor.h | 9 +- backends/vulkan/runtime/api/api.h | 2 + backends/vulkan/runtime/graph/ComputeGraph.h | 4 +- backends/vulkan/test/utils/test_utils.cpp | 4 +- .../vulkan/test/vulkan_compute_api_test.cpp | 6 +- 11 files changed, 223 insertions(+), 179 deletions(-) create mode 100644 backends/vulkan/runtime/api/ParamsBuffer.cpp create mode 100644 backends/vulkan/runtime/api/ParamsBuffer.h create mode 100644 backends/vulkan/runtime/api/StorageBuffer.h diff --git a/backends/vulkan/runtime/api/Context.cpp b/backends/vulkan/runtime/api/Context.cpp index 9407cf76f5b..ec10354c67d 100644 --- a/backends/vulkan/runtime/api/Context.cpp +++ b/backends/vulkan/runtime/api/Context.cpp @@ -8,11 +8,6 @@ #include -#include -#include -#include -#include - #ifndef VULKAN_DESCRIPTOR_POOL_SIZE #define VULKAN_DESCRIPTOR_POOL_SIZE 1024u #endif @@ -220,59 +215,5 @@ Context* context() { return context.get(); } -// -// UniformParamsBuffer -// - -namespace { - -void memcpy_to_buffer(const VulkanBuffer& src, VulkanBuffer& dst) { - MemoryMap dst_mapping(dst, MemoryAccessType::WRITE); - - MemoryMap src_mapping(src, MemoryAccessType::READ); - src_mapping.invalidate(); - - void* dst_ptr = dst_mapping.template data(); - void* src_ptr = src_mapping.template data(); - - // @lint-ignore CLANGTIDY facebook-security-vulnerable-memcpy - memcpy(dst_ptr, src_ptr, src.mem_size()); -} - -} // namespace - -UniformParamsBuffer::UniformParamsBuffer(const UniformParamsBuffer& other) - : context_p_(other.context_p_), vulkan_buffer_{} { - if (other.vulkan_buffer_) { - vulkan_buffer_ = context_p_->adapter_ptr()->vma().create_uniform_buffer( - other.vulkan_buffer_.mem_size()); - - memcpy_to_buffer(other.vulkan_buffer_, vulkan_buffer_); - } -} - -UniformParamsBuffer& UniformParamsBuffer::operator=( - const UniformParamsBuffer& other) { - if (&other != this) { - context_p_ = other.context_p_; - - // Move vulkan_buffer_ to another VulkanBuffer for cleanup - if (vulkan_buffer_) { - VulkanBuffer temp_buffer(std::move(vulkan_buffer_)); - context_p_->register_buffer_cleanup(temp_buffer); - } - // vulkan_buffer_ should now be empty - - if (other.vulkan_buffer_) { - vulkan_buffer_ = context_p_->adapter_ptr()->vma().create_uniform_buffer( - other.vulkan_buffer_.mem_size()); - - memcpy_to_buffer(other.vulkan_buffer_, vulkan_buffer_); - } - } - - return *this; -} - } // namespace api } // namespace vkcompute diff --git a/backends/vulkan/runtime/api/Context.h b/backends/vulkan/runtime/api/Context.h index fc986b83eb0..a175a4a2677 100644 --- a/backends/vulkan/runtime/api/Context.h +++ b/backends/vulkan/runtime/api/Context.h @@ -10,19 +10,12 @@ // @lint-ignore-every CLANGTIDY facebook-hte-BadMemberName -#include - #include #include #include #include -#include #include #include -#include -#include - -#include namespace vkcompute { namespace api { @@ -218,103 +211,6 @@ class Context final { void flush(); }; -class UniformParamsBuffer final { - private: - Context* context_p_; - size_t nbytes_; - VulkanBuffer vulkan_buffer_; - - public: - UniformParamsBuffer() : context_p_{nullptr}, vulkan_buffer_{} {} - - template - UniformParamsBuffer(Context* context_p, const Block& block) - : context_p_(context_p), - nbytes_(sizeof(block)), - vulkan_buffer_( - context_p_->adapter_ptr()->vma().create_params_buffer(block)) {} - - UniformParamsBuffer(const UniformParamsBuffer&); - UniformParamsBuffer& operator=(const UniformParamsBuffer&); - - UniformParamsBuffer(UniformParamsBuffer&&) = default; - UniformParamsBuffer& operator=(UniformParamsBuffer&&) = default; - - ~UniformParamsBuffer() { - if (vulkan_buffer_) { - context_p_->register_buffer_cleanup(vulkan_buffer_); - } - } - - const VulkanBuffer& buffer() const { - return vulkan_buffer_; - } - - template - void update(const Block& block) { - if (sizeof(block) != nbytes_) { - VK_THROW( - "Attempted to update UniformParamsBuffer with data of different size"); - } - // Fill the uniform buffer with data in block - { - MemoryMap mapping(vulkan_buffer_, MemoryAccessType::WRITE); - Block* data_ptr = mapping.template data(); - - *data_ptr = block; - } - } -}; - -class StorageBuffer final { - private: - Context* context_p_; - ScalarType dtype_; - size_t numel_; - size_t nbytes_; - VulkanBuffer vulkan_buffer_; - - public: - StorageBuffer( - Context* context_p, - const ScalarType dtype, - const size_t numel, - const bool gpuonly = false) - : context_p_(context_p), - dtype_(dtype), - numel_(numel), - nbytes_(element_size(dtype_) * numel_), - vulkan_buffer_(context_p_->adapter_ptr()->vma().create_storage_buffer( - nbytes_, - gpuonly)) {} - - StorageBuffer(const StorageBuffer&) = delete; - StorageBuffer& operator=(const StorageBuffer&) = delete; - - StorageBuffer(StorageBuffer&&) = default; - StorageBuffer& operator=(StorageBuffer&&) = default; - - ~StorageBuffer() { - context_p_->register_buffer_cleanup(vulkan_buffer_); - } - - inline ScalarType dtype() { - return dtype_; - } - - inline VulkanBuffer& buffer() { - return vulkan_buffer_; - } - - inline size_t numel() { - return numel_; - } - - inline size_t nbytes() { - return nbytes_; - } -}; - bool available(); // The global runtime is retrieved using this function, where it is declared as diff --git a/backends/vulkan/runtime/api/ParamsBuffer.cpp b/backends/vulkan/runtime/api/ParamsBuffer.cpp new file mode 100644 index 00000000000..d0af87e9303 --- /dev/null +++ b/backends/vulkan/runtime/api/ParamsBuffer.cpp @@ -0,0 +1,67 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#include + +#include +#include + +namespace vkcompute { +namespace api { + +namespace { + +void memcpy_to_buffer(const VulkanBuffer& src, VulkanBuffer& dst) { + MemoryMap dst_mapping(dst, MemoryAccessType::WRITE); + + MemoryMap src_mapping(src, MemoryAccessType::READ); + src_mapping.invalidate(); + + void* dst_ptr = dst_mapping.template data(); + void* src_ptr = src_mapping.template data(); + + // @lint-ignore CLANGTIDY facebook-security-vulnerable-memcpy + memcpy(dst_ptr, src_ptr, src.mem_size()); +} + +} // namespace + +ParamsBuffer::ParamsBuffer(const ParamsBuffer& other) + : context_p_(other.context_p_), vulkan_buffer_{} { + if (other.vulkan_buffer_) { + vulkan_buffer_ = context_p_->adapter_ptr()->vma().create_uniform_buffer( + other.vulkan_buffer_.mem_size()); + + memcpy_to_buffer(other.vulkan_buffer_, vulkan_buffer_); + } +} + +ParamsBuffer& ParamsBuffer::operator=(const ParamsBuffer& other) { + if (&other != this) { + context_p_ = other.context_p_; + + // Move vulkan_buffer_ to another VulkanBuffer for cleanup + if (vulkan_buffer_) { + VulkanBuffer temp_buffer(std::move(vulkan_buffer_)); + context_p_->register_buffer_cleanup(temp_buffer); + } + // vulkan_buffer_ should now be empty + + if (other.vulkan_buffer_) { + vulkan_buffer_ = context_p_->adapter_ptr()->vma().create_uniform_buffer( + other.vulkan_buffer_.mem_size()); + + memcpy_to_buffer(other.vulkan_buffer_, vulkan_buffer_); + } + } + + return *this; +} + +} // namespace api +} // namespace vkcompute diff --git a/backends/vulkan/runtime/api/ParamsBuffer.h b/backends/vulkan/runtime/api/ParamsBuffer.h new file mode 100644 index 00000000000..2cf6452efc8 --- /dev/null +++ b/backends/vulkan/runtime/api/ParamsBuffer.h @@ -0,0 +1,68 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +// @lint-ignore-every CLANGTIDY facebook-hte-BadMemberName + +#include + +#include + +namespace vkcompute { +namespace api { + +class ParamsBuffer final { + private: + Context* context_p_; + size_t nbytes_; + VulkanBuffer vulkan_buffer_; + + public: + ParamsBuffer() : context_p_{nullptr}, vulkan_buffer_{} {} + + template + ParamsBuffer(Context* context_p, const Block& block) + : context_p_(context_p), + nbytes_(sizeof(block)), + vulkan_buffer_( + context_p_->adapter_ptr()->vma().create_params_buffer(block)) {} + + ParamsBuffer(const ParamsBuffer&); + ParamsBuffer& operator=(const ParamsBuffer&); + + ParamsBuffer(ParamsBuffer&&) = default; + ParamsBuffer& operator=(ParamsBuffer&&) = default; + + ~ParamsBuffer() { + if (vulkan_buffer_) { + context_p_->register_buffer_cleanup(vulkan_buffer_); + } + } + + const VulkanBuffer& buffer() const { + return vulkan_buffer_; + } + + template + void update(const Block& block) { + if (sizeof(block) != nbytes_) { + VK_THROW("Attempted to update ParamsBuffer with data of different size"); + } + // Fill the uniform buffer with data in block + { + MemoryMap mapping(vulkan_buffer_, MemoryAccessType::WRITE); + Block* data_ptr = mapping.template data(); + + *data_ptr = block; + } + } +}; + +} // namespace api +} // namespace vkcompute diff --git a/backends/vulkan/runtime/api/StorageBuffer.h b/backends/vulkan/runtime/api/StorageBuffer.h new file mode 100644 index 00000000000..4e8128ee7c5 --- /dev/null +++ b/backends/vulkan/runtime/api/StorageBuffer.h @@ -0,0 +1,70 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +// @lint-ignore-every CLANGTIDY facebook-hte-BadMemberName + +#include + +#include + +namespace vkcompute { +namespace api { + +class StorageBuffer final { + private: + Context* context_p_; + ScalarType dtype_; + size_t numel_; + size_t nbytes_; + VulkanBuffer vulkan_buffer_; + + public: + StorageBuffer( + Context* context_p, + const ScalarType dtype, + const size_t numel, + const bool gpuonly = false) + : context_p_(context_p), + dtype_(dtype), + numel_(numel), + nbytes_(element_size(dtype_) * numel_), + vulkan_buffer_(context_p_->adapter_ptr()->vma().create_storage_buffer( + nbytes_, + gpuonly)) {} + + StorageBuffer(const StorageBuffer&) = delete; + StorageBuffer& operator=(const StorageBuffer&) = delete; + + StorageBuffer(StorageBuffer&&) = default; + StorageBuffer& operator=(StorageBuffer&&) = default; + + ~StorageBuffer() { + context_p_->register_buffer_cleanup(vulkan_buffer_); + } + + inline ScalarType dtype() { + return dtype_; + } + + inline VulkanBuffer& buffer() { + return vulkan_buffer_; + } + + inline size_t numel() { + return numel_; + } + + inline size_t nbytes() { + return nbytes_; + } +}; + +} // namespace api +} // namespace vkcompute diff --git a/backends/vulkan/runtime/api/Tensor.cpp b/backends/vulkan/runtime/api/Tensor.cpp index 687b39dd7de..0b77a066f06 100644 --- a/backends/vulkan/runtime/api/Tensor.cpp +++ b/backends/vulkan/runtime/api/Tensor.cpp @@ -172,7 +172,7 @@ api::VulkanBuffer& vTensor::buffer( const api::BufferBindInfo vTensor::sizes_ubo() { if (!sizes_uniform_.buffer()) { - sizes_uniform_ = api::UniformParamsBuffer( + sizes_uniform_ = api::ParamsBuffer( storage_.context_, api::utils::make_whcn_ivec4(sizes_)); } return api::BufferBindInfo(sizes_uniform_.buffer()); @@ -181,14 +181,14 @@ const api::BufferBindInfo vTensor::sizes_ubo() { const api::BufferBindInfo vTensor::texture_limits_ubo() { if (!texture_limits_uniform_.buffer()) { texture_limits_uniform_ = - api::UniformParamsBuffer(storage_.context_, texture_limits_); + api::ParamsBuffer(storage_.context_, texture_limits_); } return api::BufferBindInfo(texture_limits_uniform_.buffer()); } const api::BufferBindInfo vTensor::texel_strides_ubo() { if (!texel_strides_uniform_.buffer()) { - texel_strides_uniform_ = api::UniformParamsBuffer( + texel_strides_uniform_ = api::ParamsBuffer( storage_.context_, api::utils::make_whcn_ivec4( calculate_strides(padded_sizes_, memory_layout_))); @@ -198,8 +198,7 @@ const api::BufferBindInfo vTensor::texel_strides_ubo() { const api::BufferBindInfo vTensor::ntexels_ubo() { if (!ntexels_uniform_.buffer()) { - ntexels_uniform_ = - api::UniformParamsBuffer(storage_.context_, texel_numel()); + ntexels_uniform_ = api::ParamsBuffer(storage_.context_, texel_numel()); } return api::BufferBindInfo(ntexels_uniform_.buffer()); } diff --git a/backends/vulkan/runtime/api/Tensor.h b/backends/vulkan/runtime/api/Tensor.h index 0b0abfb23e3..1d57e35ba3f 100644 --- a/backends/vulkan/runtime/api/Tensor.h +++ b/backends/vulkan/runtime/api/Tensor.h @@ -11,6 +11,7 @@ // @lint-ignore-every CLANGTIDY facebook-hte-BadMemberName #include +#include #include namespace vkcompute { @@ -180,10 +181,10 @@ class vTensor final { * Refer to the comments for the corresponding *_ubo() functions for more * context about the data contained in each buffer. */ - api::UniformParamsBuffer sizes_uniform_; - api::UniformParamsBuffer texture_limits_uniform_; - api::UniformParamsBuffer texel_strides_uniform_; - api::UniformParamsBuffer ntexels_uniform_; + api::ParamsBuffer sizes_uniform_; + api::ParamsBuffer texture_limits_uniform_; + api::ParamsBuffer texel_strides_uniform_; + api::ParamsBuffer ntexels_uniform_; vTensorStorage storage_; diff --git a/backends/vulkan/runtime/api/api.h b/backends/vulkan/runtime/api/api.h index 16e2b969871..8d78374e4eb 100644 --- a/backends/vulkan/runtime/api/api.h +++ b/backends/vulkan/runtime/api/api.h @@ -13,10 +13,12 @@ #include #include #include +#include #include #include #include #include +#include #include #include diff --git a/backends/vulkan/runtime/graph/ComputeGraph.h b/backends/vulkan/runtime/graph/ComputeGraph.h index 14965b52b5d..1c1a2b9f986 100644 --- a/backends/vulkan/runtime/graph/ComputeGraph.h +++ b/backends/vulkan/runtime/graph/ComputeGraph.h @@ -95,7 +95,7 @@ class ComputeGraph final { std::unique_ptr context_; std::vector shared_objects_; std::vector values_; - std::vector param_ubos_; + std::vector param_ubos_; std::vector> prepack_nodes_; std::vector> execute_nodes_; @@ -382,7 +382,7 @@ class ComputeGraph final { template const api::BufferBindInfo create_params_buffer(const Block& data) { - param_ubos_.emplace_back(api::UniformParamsBuffer(context_.get(), data)); + param_ubos_.emplace_back(api::ParamsBuffer(context_.get(), data)); return api::BufferBindInfo(param_ubos_.back().buffer()); } diff --git a/backends/vulkan/test/utils/test_utils.cpp b/backends/vulkan/test/utils/test_utils.cpp index 67a4790deb3..e2eecc5b835 100644 --- a/backends/vulkan/test/utils/test_utils.cpp +++ b/backends/vulkan/test/utils/test_utils.cpp @@ -126,7 +126,7 @@ void record_conv2d_prepack_weights_op( add_dtype_suffix(kernel_name, v_dst); api::ShaderInfo shader = VK_KERNEL_FROM_STR(kernel_name); - api::UniformParamsBuffer original_sizes_ubo( + api::ParamsBuffer original_sizes_ubo( context, api::utils::make_ivec4(original_sizes, /*reverse = */ true)); api::SpecVarList specialization_constants = {}; @@ -217,7 +217,7 @@ void record_index_fill_buffer(api::Context* context, vTensor& v_ten) { break; } - api::UniformParamsBuffer params(api::context(), int32_t(v_ten.numel())); + api::ParamsBuffer params(api::context(), int32_t(v_ten.numel())); { api::PipelineBarrier pipeline_barrier{}; diff --git a/backends/vulkan/test/vulkan_compute_api_test.cpp b/backends/vulkan/test/vulkan_compute_api_test.cpp index 19aa3ee0a78..fcad7184c60 100644 --- a/backends/vulkan/test/vulkan_compute_api_test.cpp +++ b/backends/vulkan/test/vulkan_compute_api_test.cpp @@ -217,7 +217,7 @@ TEST_F(VulkanComputeAPITest, spec_var_shader_test) { float offset = 1.5f; { - api::UniformParamsBuffer params(api::context(), int32_t(len)); + api::ParamsBuffer params(api::context(), int32_t(len)); uint32_t len_div4 = api::utils::div_up(uint32_t(len), uint32_t(4)); api::PipelineBarrier pipeline_barrier{}; api::context()->submit_compute_job( @@ -262,7 +262,7 @@ TEST_F(VulkanComputeAPITest, update_params_between_submit) { {5.0, 5.0, 5.0, 5.0}, }; - api::UniformParamsBuffer params(api::context(), block); + api::ParamsBuffer params(api::context(), block); { api::PipelineBarrier pipeline_barrier{}; @@ -323,7 +323,7 @@ void test_storage_buffer_type(const size_t len) { break; } - api::UniformParamsBuffer params(api::context(), int32_t(len)); + api::ParamsBuffer params(api::context(), int32_t(len)); { uint32_t len_div4 = api::utils::div_up(uint32_t(len), uint32_t(4)); From 73ab79b9958a11243abd7ea039f7c3083c64a54a Mon Sep 17 00:00:00 2001 From: Jorge Pineda Date: Tue, 2 Jul 2024 13:32:52 -0700 Subject: [PATCH 2/2] Update on "[ET-VK] Move `ParamsBuffer` and `StorageBuffer` to standalone files" Includes the renaming of `UniformParamsBuffer` to `ParamsBuffer` for brevity. These objects aren't tightly coupled to `Context` and hence they are better placed in standalone files. Differential Revision: [D59281543](https://our.internmc.facebook.com/intern/diff/D59281543/) [ghstack-poisoned] --- backends/vulkan/runtime/api/ParamsBuffer.cpp | 1 - 1 file changed, 1 deletion(-) diff --git a/backends/vulkan/runtime/api/ParamsBuffer.cpp b/backends/vulkan/runtime/api/ParamsBuffer.cpp index d0af87e9303..28ac835e1da 100644 --- a/backends/vulkan/runtime/api/ParamsBuffer.cpp +++ b/backends/vulkan/runtime/api/ParamsBuffer.cpp @@ -8,7 +8,6 @@ #include -#include #include namespace vkcompute {