[ET-VK] Organize utils

jorgep31415 · jorgep31415 · commit 25d279d66400 · 2024-03-08T12:10:40.000-08:00
Refactor our `OpUtils.*`, `StagingUtils.*`, and `Utils.*` which have accumulated helper functions during the redesign. We split them and move them into two new folders: 1. `ops/utils`: Used by files in `ops/` and `ops/impl/`. 2. `ops/impl/utils`: Used only by files in `ops/impl/`. In general, this helps to unflatten the `ops/` directory: ``` [jorgep31415devvm15882.vll0 /data/users/jorgep31415/fbsource/xplat/executorch/backends/vulkan/runtime/graph/ops (98b0e789c)]$ tree . ├── ExecuteNode.cpp ├── ExecuteNode.h ├── impl │ ├── Arithmetic.cpp │ ├── Arithmetic.h │ ├── Staging.cpp │ ├── Staging.h │ └── utils │ ├── DimUtils.h │ ├── ScalarUtils.h │ ├── TensorUtils.cpp │ └── TensorUtils.h ├── OperatorRegistry.cpp ├── OperatorRegistry.h ├── PrepackNode.cpp ├── PrepackNode.h └── utils ├── BindingUtils.cpp ├── BindingUtils.h ├── StagingUtils.cpp └── StagingUtils.h 3 directories, 18 files ``` Since we're heavily updating include declarations, we also introduce the convention to only include the VK-API via ``` #include <ATen/native/vulkan/api/api.h> ``` instead of including more specific files, e.g., ``` #include <ATen/native/vulkan/api/Tensor.h> ``` Differential Revision: [D54690346](https://our.internmc.facebook.com/intern/diff/D54690346/) ghstack-source-id: 217990425 Pull Request resolved: #2323
diff --git a/backends/vulkan/runtime/graph/ComputeGraph.cpp b/backends/vulkan/runtime/graph/ComputeGraph.cpp
@@ -6,12 +6,15 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-#include <executorch/backends/vulkan/runtime/graph/ComputeGraph.h>
+// @lint-ignore-every CLANGTIDY
+// facebook-security-vulnerable-integer-sign-conversion
 
-#include <executorch/backends/vulkan/runtime/graph/ops/StagingUtils.h>
+#include <executorch/backends/vulkan/runtime/graph/ComputeGraph.h>
 
 #include <executorch/backends/vulkan/runtime/graph/ops/impl/Staging.h>
 
+#include <executorch/backends/vulkan/runtime/graph/ops/utils/StagingUtils.h>
+
 namespace at {
 namespace native {
 namespace vulkan {
diff --git a/backends/vulkan/runtime/graph/ComputeGraph.h b/backends/vulkan/runtime/graph/ComputeGraph.h
@@ -12,9 +12,7 @@
 
 #ifdef USE_VULKAN_API
 
-#include <ATen/native/vulkan/api/Context.h>
-#include <ATen/native/vulkan/api/Tensor.h>
-#include <ATen/native/vulkan/api/Types.h>
+#include <ATen/native/vulkan/api/api.h>
 
 #include <executorch/backends/vulkan/runtime/graph/GraphConfig.h>
 
diff --git a/backends/vulkan/runtime/graph/GraphConfig.h b/backends/vulkan/runtime/graph/GraphConfig.h
@@ -10,7 +10,7 @@
 
 #ifdef USE_VULKAN_API
 
-#include <ATen/native/vulkan/api/Context.h>
+#include <ATen/native/vulkan/api/api.h>
 
 namespace at {
 namespace native {
diff --git a/backends/vulkan/runtime/graph/ops/ExecuteNode.cpp b/backends/vulkan/runtime/graph/ops/ExecuteNode.cpp
@@ -10,7 +10,7 @@
 
 #include <executorch/backends/vulkan/runtime/graph/ComputeGraph.h>
 
-#include <executorch/backends/vulkan/runtime/graph/ops/Utils.h>
+#include <executorch/backends/vulkan/runtime/graph/ops/utils/BindingUtils.h>
 
 namespace at {
 namespace native {
diff --git a/backends/vulkan/runtime/graph/ops/ExecuteNode.h b/backends/vulkan/runtime/graph/ops/ExecuteNode.h
@@ -10,9 +10,7 @@
 
 #ifdef USE_VULKAN_API
 
-#include <ATen/native/vulkan/api/Context.h>
-#include <ATen/native/vulkan/api/Tensor.h>
-#include <ATen/native/vulkan/api/Types.h>
+#include <ATen/native/vulkan/api/api.h>
 
 #include <executorch/backends/vulkan/runtime/graph/containers/Value.h>
 
diff --git a/backends/vulkan/runtime/graph/ops/OperatorRegistry.cpp b/backends/vulkan/runtime/graph/ops/OperatorRegistry.cpp
@@ -8,8 +8,6 @@
 
 #include <executorch/backends/vulkan/runtime/graph/ops/OperatorRegistry.h>
 
-#include <executorch/backends/vulkan/runtime/graph/ops/impl/Arithmetic.h>
-
 namespace at {
 namespace native {
 namespace vulkan {
diff --git a/backends/vulkan/runtime/graph/ops/PrepackNode.cpp b/backends/vulkan/runtime/graph/ops/PrepackNode.cpp
@@ -10,8 +10,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/Utils.h>
+#include <executorch/backends/vulkan/runtime/graph/ops/utils/BindingUtils.h>
+#include <executorch/backends/vulkan/runtime/graph/ops/utils/StagingUtils.h>
 
 namespace at {
 namespace native {
diff --git a/backends/vulkan/runtime/graph/ops/PrepackNode.h b/backends/vulkan/runtime/graph/ops/PrepackNode.h
@@ -10,9 +10,7 @@
 
 #ifdef USE_VULKAN_API
 
-#include <ATen/native/vulkan/api/Context.h>
-#include <ATen/native/vulkan/api/Tensor.h>
-#include <ATen/native/vulkan/api/Types.h>
+#include <ATen/native/vulkan/api/api.h>
 
 #include <executorch/backends/vulkan/runtime/graph/containers/Value.h>
 
diff --git a/backends/vulkan/runtime/graph/ops/impl/Arithmetic.cpp b/backends/vulkan/runtime/graph/ops/impl/Arithmetic.cpp
@@ -8,11 +8,13 @@
 
 #include <executorch/backends/vulkan/runtime/graph/ops/impl/Arithmetic.h>
 
-#include <executorch/backends/vulkan/runtime/graph/ops/OpUtils.h>
 #include <executorch/backends/vulkan/runtime/graph/ops/OperatorRegistry.h>
 
 #include <executorch/backends/vulkan/runtime/graph/ops/impl/Staging.h>
 
+#include <executorch/backends/vulkan/runtime/graph/ops/impl/utils/ScalarUtils.h>
+#include <executorch/backends/vulkan/runtime/graph/ops/impl/utils/TensorUtils.h>
+
 namespace at {
 namespace native {
 namespace vulkan {
diff --git a/backends/vulkan/runtime/graph/ops/impl/Arithmetic.h b/backends/vulkan/runtime/graph/ops/impl/Arithmetic.h
@@ -12,8 +12,6 @@
 
 #include <executorch/backends/vulkan/runtime/graph/ComputeGraph.h>
 
-#include <executorch/backends/vulkan/runtime/graph/ops/Utils.h>
-
 namespace at {
 namespace native {
 namespace vulkan {
diff --git a/backends/vulkan/runtime/graph/ops/impl/Staging.cpp b/backends/vulkan/runtime/graph/ops/impl/Staging.cpp
@@ -8,9 +8,10 @@
 
 #include <executorch/backends/vulkan/runtime/graph/ops/impl/Staging.h>
 
-#include <executorch/backends/vulkan/runtime/graph/ops/OpUtils.h>
-#include <executorch/backends/vulkan/runtime/graph/ops/StagingUtils.h>
-#include <executorch/backends/vulkan/runtime/graph/ops/Utils.h>
+#include <executorch/backends/vulkan/runtime/graph/ops/utils/StagingUtils.h>
+
+#include <executorch/backends/vulkan/runtime/graph/ops/impl/utils/DimUtils.h>
+#include <executorch/backends/vulkan/runtime/graph/ops/impl/utils/TensorUtils.h>
 
 namespace at {
 namespace native {
diff --git a/backends/vulkan/runtime/graph/ops/impl/utils/DimUtils.h b/backends/vulkan/runtime/graph/ops/impl/utils/DimUtils.h
@@ -12,8 +12,6 @@
 
 #include <ATen/native/vulkan/api/api.h>
 
-#include <executorch/backends/vulkan/runtime/graph/containers/Value.h>
-
 namespace at {
 namespace native {
 namespace vulkan {
@@ -75,27 +73,6 @@ uint32_t dim_at(const vTensor& v_in) {
   return dim_at<N>(v_in.sizes());
 }
 
-/*
- * For most global work group sizes, returns {4, 4, 4}, but adjusts the size for
- * 2D global work group sizes. Always maintains a total of 64 invocations
- */
-api::utils::uvec3 adaptive_work_group_size(
-    const api::utils::uvec3& global_work_group);
-
-template <typename T>
-T extract_scalar(const Value& value) {
-  if (value.isInt()) {
-    return static_cast<T>(value.toInt());
-  }
-  if (value.isDouble()) {
-    return static_cast<T>(value.toDouble());
-  }
-  if (value.isBool()) {
-    return static_cast<T>(value.toBool());
-  }
-  VK_THROW("Cannot extract scalar from Value with type ", value.type());
-}
-
 } // namespace vulkan
 } // namespace native
 } // namespace at
diff --git a/backends/vulkan/runtime/graph/ops/impl/utils/ScalarUtils.h b/backends/vulkan/runtime/graph/ops/impl/utils/ScalarUtils.h
@@ -0,0 +1,39 @@
+/*
+ * 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 <ATen/native/vulkan/api/api.h>
+
+#include <executorch/backends/vulkan/runtime/graph/containers/Value.h>
+
+namespace at {
+namespace native {
+namespace vulkan {
+
+template <typename T>
+T extract_scalar(const Value& value) {
+  if (value.isInt()) {
+    return static_cast<T>(value.toInt());
+  }
+  if (value.isDouble()) {
+    return static_cast<T>(value.toDouble());
+  }
+  if (value.isBool()) {
+    return static_cast<T>(value.toBool());
+  }
+  VK_THROW("Cannot extract scalar from Value with type ", value.type());
+}
+
+} // namespace vulkan
+} // namespace native
+} // namespace at
+
+#endif /* USE_VULKAN_API */
diff --git a/backends/vulkan/runtime/graph/ops/impl/utils/TensorUtils.cpp b/backends/vulkan/runtime/graph/ops/impl/utils/TensorUtils.cpp
@@ -6,7 +6,9 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-#include <executorch/backends/vulkan/runtime/graph/ops/OpUtils.h>
+#include <executorch/backends/vulkan/runtime/graph/ops/impl/utils/TensorUtils.h>
+
+#include <executorch/backends/vulkan/runtime/graph/ops/impl/utils/DimUtils.h>
 
 namespace at {
 namespace native {
@@ -29,6 +31,14 @@ api::utils::uvec3 adaptive_work_group_size(
   return local_group_size;
 }
 
+api::utils::ivec4 get_size_as_ivec4(const vTensor& t) {
+  return api::utils::make_ivec4(
+      {dim_at<Dim4D::Width>(t),
+       dim_at<Dim4D::Height>(t),
+       dim_at<Dim4D::Channel>(t),
+       dim_at<Dim4D::Batch>(t)});
+}
+
 } // namespace vulkan
 } // namespace native
 } // namespace at
diff --git a/backends/vulkan/runtime/graph/ops/impl/utils/TensorUtils.h b/backends/vulkan/runtime/graph/ops/impl/utils/TensorUtils.h
@@ -0,0 +1,28 @@
+/*
+ * 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 <ATen/native/vulkan/api/api.h>
+
+namespace at {
+namespace native {
+namespace vulkan {
+
+api::utils::uvec3 adaptive_work_group_size(
+    const api::utils::uvec3& global_work_group);
+
+api::utils::ivec4 get_size_as_ivec4(const vTensor& t);
+
+} // namespace vulkan
+} // namespace native
+} // namespace at
+
+#endif /* USE_VULKAN_API */
diff --git a/backends/vulkan/runtime/graph/ops/utils/BindingUtils.cpp b/backends/vulkan/runtime/graph/ops/utils/BindingUtils.cpp
@@ -6,22 +6,12 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-#include <executorch/backends/vulkan/runtime/graph/ops/Utils.h>
-
-#include <executorch/backends/vulkan/runtime/graph/ops/OpUtils.h>
+#include <executorch/backends/vulkan/runtime/graph/ops/utils/BindingUtils.h>
 
 namespace at {
 namespace native {
 namespace vulkan {
 
-api::utils::ivec4 get_size_as_ivec4(const vTensor& t) {
-  return api::utils::make_ivec4(
-      {dim_at<Dim4D::Width>(t),
-       dim_at<Dim4D::Height>(t),
-       dim_at<Dim4D::Channel>(t),
-       dim_at<Dim4D::Batch>(t)});
-}
-
 void bind_tensor_to_descriptor_set(
     vTensor& tensor,
     api::PipelineBarrier& pipeline_barrier,
diff --git a/backends/vulkan/runtime/graph/ops/utils/BindingUtils.h b/backends/vulkan/runtime/graph/ops/utils/BindingUtils.h
@@ -16,8 +16,6 @@ namespace at {
 namespace native {
 namespace vulkan {
 
-api::utils::ivec4 get_size_as_ivec4(const vTensor& t);
-
 void bind_tensor_to_descriptor_set(
     vTensor& tensor,
     api::PipelineBarrier& pipeline_barrier,
diff --git a/backends/vulkan/runtime/graph/ops/utils/StagingUtils.cpp b/backends/vulkan/runtime/graph/ops/utils/StagingUtils.cpp
@@ -6,15 +6,36 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-#include <executorch/backends/vulkan/runtime/graph/ops/StagingUtils.h>
+// @lint-ignore-every CLANGTIDY facebook-security-vulnerable-memcpy
 
-#include <executorch/backends/vulkan/runtime/graph/ops/OpUtils.h>
-#include <executorch/backends/vulkan/runtime/graph/ops/Utils.h>
+#include <executorch/backends/vulkan/runtime/graph/ops/utils/StagingUtils.h>
+
+#include <executorch/backends/vulkan/runtime/graph/ops/impl/utils/DimUtils.h>
+
+#include <cstring>
 
 namespace at {
 namespace native {
 namespace vulkan {
 
+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);
+}
+
 void memcpy_to_mapping(
     const void* src,
     api::MemoryMap& dst_mapping,
diff --git a/backends/vulkan/runtime/graph/ops/utils/StagingUtils.h b/backends/vulkan/runtime/graph/ops/utils/StagingUtils.h
@@ -12,49 +12,10 @@
 
 #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
 //
diff --git a/backends/vulkan/test/vulkan_compute_api_test.cpp b/backends/vulkan/test/vulkan_compute_api_test.cpp
