[ET-VK] Handle scalar tensor and mutable buffer inputs in Vulkan delegate runtime

## Context

* Handle scalar tensor inputs by adding them to the graph as symbolic ints
* Add support for symint inputs in the Vulkan delegate
* Add type checking for Vulkan delegate inputs and outputs

This is needed for Transformer models, which receive a an `input_pos` integer scalar tensor as an input. `input_pos` is used in KV cache updates and determines the sizes of the cache slices.

### Why are scalar tensors added as symint?

Adding scalar tensors as symints makes more sense than adding them as real tensors, since symints are commonly used to inform tensor shapes. Adding scalar tensors as symints allow them to be easily accessible by the CPU at graph encoding and resizing time, as well as easily accesible by the GPU within compute shaders.

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

ghstack-source-id: 246578362
Pull Request resolved: #5930

Author: SS-JIA

backends/vulkan/runtime/VulkanBackend.cpp

@@ -192,6 +192,12 @@ class GraphBuilder {
     UIntVector dims_fb = tensor_fb->dims();
     const std::vector<int64_t> dims_vector(dims_fb->cbegin(), dims_fb->cend());
 
+    // For scalar tensors, add them as SymInts instead of tensors
+    if (dtype == vkapi::kInt && utils::multiply_integers(dims_vector) == 1) {
+      ref_mapping_[fb_id] = compute_graph_->add_symint(0);
+      return;
+    }
+
     utils::GPUMemoryLayout memory_layout =
         tensor_fb->memory_layout() == vkgraph::VkMemoryLayout::DEFAULT_LAYOUT
         ? compute_graph_->suggested_memory_layout(dims_vector)
@@ -312,10 +318,16 @@ class GraphBuilder {
       add_value_to_graph(fb_id, value);
     }
 
-    // Parse the inputs
+    // Parse the inputs, which will be tensors most of the time but can also be
+    // symints and tensorrefs (which will be the case if the original graph had)
+    // mutable buffers.
     for (const uint32_t fb_id : *flatbuffer_->input_ids()) {
       const ValueRef ref = get_fb_id_valueref(fb_id);
-      compute_graph_->set_input_tensor(ref);
+      if (compute_graph_->val_is_tensor(ref)) {
+        compute_graph_->set_input_tensor(ref);
+      } else {
+        compute_graph_->set_val_as_input(ref);
+      }
     }
 
     // Parse the operators
@@ -354,10 +366,15 @@ class GraphBuilder {
       }
     }
 
-    // Parse the outputs
+    // Parse the outputs, which will be mostly tensors.  For some reason,
+    // mutable buffers are shown to be returned in the fx.Graph but do not get
+    // returned by the delegate; this may be an implementation detail of how the
+    // executorch emitter handles mutable buffers.
     for (const uint32_t fb_id : *flatbuffer_->output_ids()) {
       const ValueRef ref = get_fb_id_valueref(fb_id);
-      compute_graph_->set_output_tensor(ref);
+      if (compute_graph_->val_is_tensor(ref)) {
+        compute_graph_->set_output_tensor(ref);
+      }
     }
   }
 };
@@ -508,13 +525,44 @@ class VulkanBackend final : public ::executorch::runtime::BackendInterface {
     const size_t num_inputs = compute_graph->inputs().size();
     bool should_propagate_resize = false;
     for (size_t i = 0; i < num_inputs; i++) {
-      bool was_resized =
-          maybe_resize_input(compute_graph, i, args[i]->toTensor());
-      should_propagate_resize = should_propagate_resize || was_resized;
-      compute_graph->copy_into_staging(
-          compute_graph->inputs()[i].staging,
-          args[i]->toTensor().const_data_ptr(),
-          args[i]->toTensor().numel());
+      const ValueRef iref = compute_graph->inputs()[i].value;
+      if (compute_graph->val_is_tensor(iref)) {
+        VK_CHECK_COND(args[i]->isTensor());
+        bool was_resized =
+            maybe_resize_input(compute_graph, i, args[i]->toTensor());
+        should_propagate_resize = should_propagate_resize || was_resized;
+        compute_graph->copy_into_staging(
+            compute_graph->inputs()[i].staging,
+            args[i]->toTensor().const_data_ptr(),
+            args[i]->toTensor().numel());
+      } else if (compute_graph->val_is_symint(iref)) {
+        int32_t scalar_tensor_val = 0;
+        const int32_t cur_val = compute_graph->read_symint(iref);
+        if (args[i]->isTensor()) {
+          exec_aten::Tensor& scalar_tensor_src = args[i]->toTensor();
+          exec_aten::ScalarType dtype = scalar_tensor_src.scalar_type();
+          if (dtype == exec_aten::ScalarType::Int) {
+            scalar_tensor_val = *scalar_tensor_src.const_data_ptr<int32_t>();
+          } else if (dtype == exec_aten::ScalarType::Long) {
+            scalar_tensor_val =
+                int32_t(*scalar_tensor_src.const_data_ptr<int64_t>());
+          }
+          if (scalar_tensor_val != cur_val) {
+            compute_graph->set_symint(iref, scalar_tensor_val);
+            // Since symint inputs may impact tensor's sizes, trigger a resize
+            // if any symbolic integer shapes are updated.
+            should_propagate_resize = true;
+          }
+        } else {
+          VK_THROW(
+              "Cannot handle symint arg to graph that is not derived from a "
+              "scalar tensor at the moment.");
+        }
+      } else {
+        VK_THROW(
+            "Could not handle input with type ",
+            compute_graph->get_val_type(iref));
+      }
     }
 
     if (should_propagate_resize) {
@@ -523,13 +571,21 @@ class VulkanBackend final : public ::executorch::runtime::BackendInterface {
     compute_graph->execute();
 
     for (size_t i = 0; i < compute_graph->outputs().size(); i++) {
-      maybe_resize_output(compute_graph, i, args[num_inputs + i]->toTensor());
-      // args holds inputs directly followed by outputs, so the i'th output
-      // for compute_graph corresponds to the (i + num_inputs)'th arg
-      compute_graph->copy_from_staging(
-          compute_graph->outputs()[i].staging,
-          args[num_inputs + i]->toTensor().mutable_data_ptr(),
-          args[num_inputs + i]->toTensor().numel());
+      const ValueRef oref = compute_graph->outputs()[i].value;
+      if (compute_graph->val_is_tensor(oref)) {
+        VK_CHECK_COND(args[i]->isTensor());
+        maybe_resize_output(compute_graph, i, args[num_inputs + i]->toTensor());
+        // args holds inputs directly followed by outputs, so the i'th output
+        // for compute_graph corresponds to the (i + num_inputs)'th arg
+        compute_graph->copy_from_staging(
+            compute_graph->outputs()[i].staging,
+            args[num_inputs + i]->toTensor().mutable_data_ptr(),
+            args[num_inputs + i]->toTensor().numel());
+      } else {
+        VK_THROW(
+            "Could not handle output with type ",
+            compute_graph->get_val_type(oref));
+      }
     }
 
 #ifdef ET_EVENT_TRACER_ENABLED

backends/vulkan/runtime/graph/ComputeGraph.h

@@ -555,6 +555,14 @@ class ComputeGraph final {
 
   int32_t read_symint(const ValueRef idx);
 
+  inline void set_val_as_input(const ValueRef idx) {
+    inputs_.push_back({idx, kDummyValueRef});
+  }
+
+  inline void set_val_as_output(const ValueRef idx) {
+    outputs_.push_back({idx, kDummyValueRef});
+  }
+
   /*
    * Convenience function to add an input tensor along with its staging buffer
    */