Update on "[ET-VK] Adding push constant and ubo verison of select and slice ops to improve memory and performance."

Adding push constant and ubo verison of select and slice ops to improve memory and performance.
*   Updated `transfer_buffer.yaml` and `transfer_texture.yaml` to include `UBO_PARAMS` parameter and generate variants for `select` and `slice` ops with UBO parameters.
*   Updated `transfer.glsl` to generate ubo and push constant versions of `select` and `slice` ops with UBO parameters.

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

[ghstack-poisoned]

Author: trivedivivek

backends/vulkan/runtime/graph/ComputeGraph.cpp

@@ -549,15 +549,6 @@ vkapi::BufferBindInfo ComputeGraph::get_or_create_int_param_buffer(
   }
 }
 
-int32_t ComputeGraph::get_or_create_int(
-    const ValueRef idx,
-    const int32_t default_val) {
-  if (values_.at(idx).isNone()) {
-    return default_val;
-  }
-  return extract_scalar<int32_t>(idx);
-}
-
 void ComputeGraph::set_symint(const ValueRef idx, const int32_t val) {
   get_symint(idx)->set(val);
 }

backends/vulkan/runtime/graph/ComputeGraph.h

@@ -424,7 +424,7 @@ class ComputeGraph final {
   // Scalar Value Extraction
   //
 
-  bool is_scalar(const ValueRef idx) const {
+  bool is_scalar_or_none(const ValueRef idx) const {
     const Value& value = values_.at(idx);
     return value.isInt() || value.isDouble() || value.isBool() ||
         value.isNone();
@@ -445,6 +445,15 @@ class ComputeGraph final {
     VK_THROW("Cannot extract scalar from Value with type ", value.type());
   }
 
+  template <typename T>
+  T extract_scalar_or(const ValueRef idx, const T default_value) {
+    Value& value = values_.at(idx);
+    if (value.isNone()) {
+      return default_value;
+    }
+    return extract_scalar<T>(idx);
+  }
+
   template <typename T>
   std::optional<T> extract_optional_scalar(const ValueRef idx) {
     if (val_is_none(idx)) {
@@ -685,8 +694,6 @@ class ComputeGraph final {
       const ValueRef idx,
       const int32_t default_value);
 
-  int32_t get_or_create_int(const ValueRef idx, const int32_t default_value);
-
   void set_symint(const ValueRef idx, const int32_t val);
 
   int32_t read_symint(const ValueRef idx);

backends/vulkan/runtime/graph/ops/impl/Transfer.cpp

@@ -46,8 +46,9 @@ void add_transfer_copy_node(
     int32_t step_ref;
   } transfer_params{static_cast<int32_t>(dim_whcn), 0, 0};
 
-  const bool param_is_scalar = graph.is_scalar(index_or_start_ref) &&
-      (transfer_type == TransferType::SELECT || graph.is_scalar(step_ref));
+  const bool param_is_scalar = graph.is_scalar_or_none(index_or_start_ref) &&
+      (transfer_type == TransferType::SELECT ||
+       graph.is_scalar_or_none(step_ref));
 
   vkapi::ParamsBindList param_buffers;
   if (!param_is_scalar) {
@@ -61,9 +62,9 @@ void add_transfer_copy_node(
     }
   } else {
     transfer_params.index_or_start_ref =
-        graph.get_or_create_int(index_or_start_ref, 0);
+        graph.extract_scalar_or<int32_t>(index_or_start_ref, 0);
     if (transfer_type != TransferType::SELECT) {
-      transfer_params.step_ref = graph.get_or_create_int(step_ref, 1);
+      transfer_params.step_ref = graph.extract_scalar_or<int32_t>(step_ref, 1);
     }
   }