Merge ArithmeticPrepack into PrepackNode (#2261)

Summary:
bypass-github-export-checks

Pull Request resolved: #2261

There's a lot of shared logic between
- `add_staging_to_tensor_node()`, which handles I/O data on execute(), and
- `ArithmeticPrepack`'s simple prepacking, on prepack().

Both just copy data to and from GPU, without any manipulation. Hence, I've decided to consolidate shared logic in this diff as well. Here are the final results:
+ Make `PrepackNode` a final class.
+ Remove all references of `impl/Packing.h`.
- Extract shared util functions to new `StagingUtils.h/cpp`.
ghstack-source-id: 217439331
exported-using-ghexport

Reviewed By: SS-JIA

Differential Revision: D54504449

fbshipit-source-id: 358f2f5acb396a05bf7758cce1f3314d3a85ba55

Author: jorgep31415

backends/vulkan/runtime/graph/ComputeGraph.cpp

@@ -8,6 +8,8 @@
 
 #include <executorch/backends/vulkan/runtime/graph/ComputeGraph.h>
 
+#include <executorch/backends/vulkan/runtime/graph/ops/StagingUtils.h>
+
 #include <executorch/backends/vulkan/runtime/graph/ops/impl/Staging.h>
 
 namespace at {

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

@@ -0,0 +1,55 @@
+/*
+ * 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 <executorch/backends/vulkan/runtime/graph/ops/ExecuteNode.h>
+
+#include <executorch/backends/vulkan/runtime/graph/ComputeGraph.h>
+
+#include <executorch/backends/vulkan/runtime/graph/ops/StagingUtils.h>
+#include <executorch/backends/vulkan/runtime/graph/ops/Utils.h>
+
+namespace at {
+namespace native {
+namespace vulkan {
+
+void PrepackNode::encode(ComputeGraph* graph) {
+  api::Context* const context = graph->context();
+  api::PipelineBarrier pipeline_barrier{};
+
+  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);
+  size_t nbytes = numel * api::element_size(tref.dtype);
+  copy_ptr_to_staging(tref.data, staging, nbytes);
+
+  std::unique_lock<std::mutex> cmd_lock = context->dispatch_lock();
+
+  api::DescriptorSet descriptor_set =
+      context->get_descriptor_set(shader_, local_workgroup_size_);
+
+  uint32_t idx = 0;
+  bind_tensor_to_descriptor_set(
+      packed,
+      pipeline_barrier,
+      api::MemoryAccessType::WRITE,
+      descriptor_set,
+      idx++);
+  bind_staging_to_descriptor_set(staging, descriptor_set, idx++);
+  descriptor_set.bind(idx, params_.buffer());
+
+  context->register_shader_dispatch(
+      descriptor_set, pipeline_barrier, shader_, global_workgroup_size_);
+}
+
+} // namespace vulkan
+} // namespace native
+} // namespace at

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

@@ -28,20 +28,37 @@ class ComputeGraph;
  * encoding of shaders transferring necessary data (such as weights and biases)
  * to the GPU.
  */
-class PrepackNode {
+class PrepackNode final {
   friend class ComputeGraph;
 
  public:
-  PrepackNode(ValueRef tref, ValueRef packed) : tref_{tref}, packed_{packed} {}
+  PrepackNode(
+      const api::ShaderInfo& shader,
+      const api::utils::uvec3& global_workgroup_size,
+      const api::utils::uvec3& local_workgroup_size,
+      const ValueRef tref,
+      const ValueRef packed,
+      api::UniformParamsBuffer&& params)
+      : shader_(shader),
+        global_workgroup_size_(global_workgroup_size),
+        local_workgroup_size_(local_workgroup_size),
+        tref_(tref),
+        packed_(packed),
+        params_(std::move(params)) {}
 
-  virtual ~PrepackNode() = default;
+  ~PrepackNode() = default;
 
- protected:
-  ValueRef tref_;
-  ValueRef packed_;
+  void encode(ComputeGraph* graph);
 
- public:
-  virtual void encode(ComputeGraph* graph) const = 0;
+ protected:
+  const api::ShaderInfo shader_;
+  const api::utils::uvec3 global_workgroup_size_;
+  const api::utils::uvec3 local_workgroup_size_;
+  const ValueRef tref_;
+  const ValueRef packed_;
+  // TODO(T180906086): pass multiple buffers and index with ValueRef.
+  // TODO(T180906457): allow re-computing param buffers.
+  api::UniformParamsBuffer params_;
 };
 
 } // namespace vulkan

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

@@ -0,0 +1,174 @@
+/*
+ * 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 <executorch/backends/vulkan/runtime/graph/ops/StagingUtils.h>
+
+#include <executorch/backends/vulkan/runtime/graph/ops/Utils.h>
+
+#include <ATen/native/vulkan/impl/Common.h>
+
+namespace at {
+namespace native {
+namespace vulkan {
+
+void memcpy_to_mapping(
+    const void* src,
+    api::MemoryMap& dst_mapping,
+    const size_t nbytes,
+    const api::ScalarType dtype) {
+#define DTYPE_CASE(ctype, vkformat, name)                    \
+  case api::ScalarType::name:                                \
+    memcpy_to_mapping_impl<ctype>(src, dst_mapping, nbytes); \
+    break;
+
+  switch (dtype) {
+    VK_FORALL_SCALAR_TYPES(DTYPE_CASE)
+    default:
+      VK_THROW("Unrecognized dtype!");
+  }
+#undef DTYPE_CASE
+}
+
+void memcpy_from_mapping(
+    api::MemoryMap& src_mapping,
+    void* dst,
+    const size_t nbytes,
+    const api::ScalarType dtype) {
+#define DTYPE_CASE(ctype, vkformat, name)                      \
+  case api::ScalarType::name:                                  \
+    memcpy_from_mapping_impl<ctype>(src_mapping, dst, nbytes); \
+    break;
+
+  switch (dtype) {
+    VK_FORALL_SCALAR_TYPES(DTYPE_CASE)
+    default:
+      VK_THROW("Unrecognized dtype!");
+  }
+#undef DTYPE_CASE
+}
+
+void copy_ptr_to_staging(
+    const void* src,
+    api::StorageBuffer& staging,
+    const size_t nbytes) {
+  api::MemoryMap mapping(staging.buffer(), api::MemoryAccessType::WRITE);
+  mapping.invalidate();
+  memcpy_to_mapping(src, mapping, nbytes, staging.dtype());
+}
+
+void copy_staging_to_ptr(
+    api::StorageBuffer& staging,
+    void* dst,
+    const size_t nbytes) {
+  api::MemoryMap mapping(staging.buffer(), api::MemoryAccessType::READ);
+  mapping.invalidate();
+  memcpy_from_mapping(mapping, dst, nbytes, staging.dtype());
+}
+
+api::ShaderInfo get_nchw_to_image_shader(const vTensor& v_dst) {
+  if (v_dst.is_quantized()) {
+    switch (v_dst.storage_type()) {
+      case api::StorageType::TEXTURE_3D:
+        switch (v_dst.dtype()) {
+          case api::ScalarType::QUInt8:
+            return VK_KERNEL(nchw_to_image_uint8);
+          case api::ScalarType::QInt8:
+            return VK_KERNEL(nchw_to_image_int8);
+          case api::ScalarType::QInt32:
+            return VK_KERNEL(nchw_to_image_int32);
+          default:
+            VK_THROW(
+                "Vulkan quantization currently not supported for dtype ",
+                v_dst.dtype());
+        }
+      case api::StorageType::TEXTURE_2D:
+        switch (v_dst.dtype()) {
+          case api::ScalarType::QUInt8:
+            return VK_KERNEL(nchw_to_image2d_uint8);
+          case api::ScalarType::QInt8:
+            return VK_KERNEL(nchw_to_image2d_int8);
+          case api::ScalarType::QInt32:
+            return VK_KERNEL(nchw_to_image2d_int32);
+          default:
+            VK_THROW(
+                "Vulkan quantization currently not supported for dtype ",
+                v_dst.dtype());
+        }
+      default:
+        VK_THROW("No kernel available!");
+      case api::StorageType::BUFFER:
+      case api::StorageType::UNKNOWN:
+        VK_THROW("Requested storage type must be a texture type.");
+    }
+  }
+
+  if (v_dst.dtype() == api::kFloat) {
+    switch (v_dst.storage_type()) {
+      case api::StorageType::TEXTURE_3D:
+        return VK_KERNEL(nchw_to_image);
+      case api::StorageType::TEXTURE_2D:
+        return VK_KERNEL(nchw_to_image2d);
+      default:
+        VK_THROW("No kernel available!");
+    }
+  } else if (v_dst.dtype() == api::kBool) {
+    switch (v_dst.storage_type()) {
+      case api::StorageType::TEXTURE_3D:
+        return VK_KERNEL(nchw_to_image_bool);
+      default:
+        VK_THROW("No kernel available!");
+    }
+  } else {
+    VK_THROW("Unsupported dtype!");
+  }
+}
+
+api::ShaderInfo get_image_to_nchw_shader(const vTensor& v_src) {
+  if (v_src.is_quantized() || v_src.dtype() == api::kBool) {
+    auto plane_size =
+        dim_at<Dim4D::Height>(v_src) * dim_at<Dim4D::Width>(v_src);
+    switch (v_src.storage_type()) {
+      case api::StorageType::TEXTURE_3D:
+        switch (v_src.dtype()) {
+          case api::ScalarType::QUInt8:
+          case api::ScalarType::QInt8:
+          case api::kBool:
+            return plane_size % 4 == 0 ? VK_KERNEL(image_to_nchw_quantized_mul4)
+                                       : VK_KERNEL(image_to_nchw_uint);
+          case api::ScalarType::QInt32:
+            return VK_KERNEL(image_to_nchw_int32);
+          default:
+            VK_THROW(
+                "Vulkan quantization currently not supported for dtype ",
+                v_src.dtype());
+        }
+      default:
+        VK_THROW("No kernel available!");
+      case api::StorageType::BUFFER:
+      case api::StorageType::UNKNOWN:
+        VK_THROW("Requested storage type must be a texture type.");
+    }
+  }
+
+  if (v_src.dtype() == api::kFloat) {
+    switch (v_src.storage_type()) {
+      case api::StorageType::TEXTURE_3D:
+        return VK_KERNEL(image_to_nchw);
+      case api::StorageType::TEXTURE_2D:
+        return VK_KERNEL(image2d_to_nchw);
+      default:
+        VK_THROW("No kernel available!");
+    }
+  } else {
+    VK_THROW("Unsupported dtype!");
+  }
+}
+
+} // namespace vulkan
+} // namespace native
+} // namespace at

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

@@ -0,0 +1,82 @@
+/*
+ * 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
+
+#ifdef USE_VULKAN_API
+
+#include <executorch/backends/vulkan/runtime/graph/ComputeGraph.h>
+
+#include <cstring>
+
+namespace at {
+namespace native {
+namespace vulkan {
+
+//
+// Functions to memcpy data into staging buffer
+//
+
+void memcpy_to_mapping(
+    const void* src,
+    api::MemoryMap& dst_mapping,
+    const size_t nbytes,
+    const api::ScalarType dtype);
+void memcpy_from_mapping(
+    const api::MemoryMap& src_mapping,
+    void* dst,
+    const size_t nbytes,
+    const api::ScalarType dtype);
+
+//
+// Utility functions for memcpy
+//
+
+template <typename T>
+void memcpy_to_mapping_impl(
+    const void* src,
+    api::MemoryMap& dst_mapping,
+    const size_t nbytes) {
+  T* data_ptr = dst_mapping.template data<T>();
+  memcpy(data_ptr, reinterpret_cast<const T*>(src), nbytes);
+}
+
+template <typename T>
+void memcpy_from_mapping_impl(
+    api::MemoryMap& src_mapping,
+    void* dst,
+    const size_t nbytes) {
+  T* data_ptr = src_mapping.template data<T>();
+  memcpy(reinterpret_cast<T*>(dst), data_ptr, nbytes);
+}
+
+//
+// Functions to copy data into and out of a staging buffer
+//
+
+void copy_ptr_to_staging(
+    const void* src,
+    api::StorageBuffer& staging,
+    const size_t nbytes);
+void copy_staging_to_ptr(
+    api::StorageBuffer& staging,
+    void* dst,
+    const size_t nbytes);
+
+//
+// Functions to get shaders
+//
+
+api::ShaderInfo get_nchw_to_image_shader(const vTensor& v_dst);
+api::ShaderInfo get_image_to_nchw_shader(const vTensor& v_src);
+
+} // namespace vulkan
+} // namespace native
+} // namespace at
+
+#endif /* USE_VULKAN_API */