From 3f27222671cdece20c379e3d593de442659e2019 Mon Sep 17 00:00:00 2001 From: Jorge Pineda Date: Tue, 2 Apr 2024 12:06:27 -0700 Subject: [PATCH] [ET-VK][EZ] Fix StorageBuffer size in PrepackNode MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Fixing a bug I wrote, but in my defense the bug was in the code I copied. 😛 This change decreases the size of our staging `StorageBuffer`s which were previously bigger than necessary. There's a few numbers describing size at play: ``` numel -> CPU buffer nbytes -> numel * cpu_dtype gpu_numel -> GPU texture gpu_nbytes -> gpu_numel * gpu_dtype ``` Firstly, `StorageBuffer`'s ctor takes `numel` not `nbytes`. Secondly, we should use the CPU size not the GPU size which may be aligned up to a multiple of 4. Differential Revision: [D55619076](https://our.internmc.facebook.com/intern/diff/D55619076/) [ghstack-poisoned] --- backends/vulkan/runtime/graph/ops/PrepackNode.cpp | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/backends/vulkan/runtime/graph/ops/PrepackNode.cpp b/backends/vulkan/runtime/graph/ops/PrepackNode.cpp index bedd1f1d8fe..60d1982d97e 100644 --- a/backends/vulkan/runtime/graph/ops/PrepackNode.cpp +++ b/backends/vulkan/runtime/graph/ops/PrepackNode.cpp @@ -39,10 +39,8 @@ void PrepackNode::encode(ComputeGraph* graph) { TensorRef tref = graph->get_val(tref_).toTensorRef(); vTensor packed = graph->get_val(packed_).toTensor(); - // TODO: Extract to standalone function, to support other types of prepacking. - api::StorageBuffer staging( - graph->context(), packed.dtype(), packed.gpu_nbytes()); size_t numel = api::utils::multiply_integers(tref.sizes); + api::StorageBuffer staging(graph->context(), tref.dtype, numel); size_t nbytes = numel * api::element_size(tref.dtype); copy_ptr_to_staging(tref.data, staging, nbytes);