Update on "[ET-VK] Add convolution cases to codegen"

TSIA

Differential Revision: [D55829466](https://our.internmc.facebook.com/intern/diff/D55829466/)

[ghstack-poisoned]

Author: jorgep31415

backends/vulkan/runtime/graph/ops/PrepackNode.cpp

@@ -41,7 +41,7 @@ api::StorageBuffer PrepackNode::create_staging_buffer(ComputeGraph* graph) {
     size_t numel = api::utils::multiply_integers(packed.sizes());
     api::StorageBuffer staging(graph->context(), packed.dtype(), numel);
     size_t nbytes = numel * api::element_size(packed.dtype());
-    copy_zeros_to_staging(staging, nbytes);
+    set_staging_zeros(staging, nbytes);
     return staging;
   }
 

backends/vulkan/runtime/graph/ops/utils/StagingUtils.cpp

@@ -89,11 +89,10 @@ void copy_staging_to_ptr(
   memcpy_from_mapping(mapping, dst, nbytes, staging.dtype());
 }
 
-void copy_zeros_to_staging(api::StorageBuffer& staging, const size_t nbytes) {
-  void* data = malloc(nbytes);
-  memset(data, 0, nbytes);
-  copy_ptr_to_staging(data, staging, nbytes);
-  free(data);
+void set_staging_zeros(api::StorageBuffer& staging, const size_t nbytes) {
+  api::MemoryMap mapping(staging.buffer(), api::MemoryAccessType::WRITE);
+  uint8_t* data_ptr = mapping.template data<uint8_t>();
+  memset(data_ptr, 0, staging.nbytes());
 }
 
 api::ShaderInfo get_nchw_to_image_shader(const vTensor& v_dst) {

backends/vulkan/runtime/graph/ops/utils/StagingUtils.h

@@ -25,7 +25,7 @@ void copy_staging_to_ptr(
     void* dst,
     const size_t nbytes);
 
-void copy_zeros_to_staging(api::StorageBuffer& staging, const size_t nbytes);
+void set_staging_zeros(api::StorageBuffer& staging, const size_t nbytes);
 
 //
 // Functions to get shaders