Add API to read value of SymInt andParamsBuffer (pytorch#5754)

Summary:
Pull Request resolved: pytorch#5754

## Context

This diff adds an API to read the value from a `SymInt`. This functionality will be useful because `SymInt`s may be needed to set tensor sizes, in addition to being used as arguments to shaders.
ghstack-source-id: 245517196
exported-using-ghexport

Reviewed By: jorgep31415

Differential Revision: D63642093

fbshipit-source-id: 14ce9daecd5520cf078d3b3259eaafb5b654d834

Author: SS-JIA

backends/vulkan/runtime/api/containers/ParamsBuffer.h

@@ -56,12 +56,29 @@ class ParamsBuffer final {
     }
     // Fill the uniform buffer with data in block
     {
-      vkapi::MemoryMap mapping(vulkan_buffer_, vkapi::MemoryAccessType::WRITE);
+      vkapi::MemoryMap mapping(vulkan_buffer_, vkapi::kWrite);
       Block* data_ptr = mapping.template data<Block>();
 
       *data_ptr = block;
     }
   }
+
+  template <typename T>
+  T read() const {
+    T val;
+    if (sizeof(val) != nbytes_) {
+      VK_THROW(
+          "Attempted to store value from ParamsBuffer to type of different size");
+    }
+    // Read value from uniform buffer and store in val
+    {
+      vkapi::MemoryMap mapping(vulkan_buffer_, vkapi::kRead);
+      T* data_ptr = mapping.template data<T>();
+
+      val = *data_ptr;
+    }
+    return val;
+  }
 };
 
 } // namespace api

backends/vulkan/runtime/graph/ComputeGraph.cpp

@@ -416,6 +416,10 @@ void ComputeGraph::set_symint(const ValueRef idx, const int32_t val) {
   get_symint(idx)->set(val);
 }
 
+int32_t ComputeGraph::read_symint(const ValueRef idx) {
+  return get_symint(idx)->get();
+}
+
 SharedObject& ComputeGraph::get_shared_object(const int64_t idx) {
   if (idx >= shared_objects_.size()) {
     shared_objects_.resize(static_cast<size_t>(idx + 1));

backends/vulkan/runtime/graph/ComputeGraph.h

@@ -536,6 +536,8 @@ class ComputeGraph final {
 
   void set_symint(const ValueRef idx, const int32_t val);
 
+  int32_t read_symint(const ValueRef idx);
+
   /*
    * Convenience function to add an input tensor along with its staging buffer
    */

backends/vulkan/runtime/graph/containers/SymInt.cpp

@@ -17,6 +17,10 @@ void SymInt::set(const int32_t val) {
   gpu_buffer.update(val);
 }
 
+int32_t SymInt::get() {
+  return gpu_buffer.read<int32_t>();
+}
+
 void SymInt::operator=(const int32_t val) {
   gpu_buffer.update(val);
 }

backends/vulkan/runtime/graph/containers/SymInt.h

@@ -35,6 +35,8 @@ struct SymInt final {
 
   void set(const int32_t val);
 
+  int32_t get();
+
   void operator=(const int32_t val);
 };
 

backends/vulkan/test/vulkan_compute_api_test.cpp

@@ -1433,6 +1433,9 @@ TEST(VulkanComputeGraphTest, test_simple_graph_with_symint) {
     int scalar_val = i - 3.0f;
     graph.set_symint(scalar, scalar_val);
 
+    int32_t scalar_val_read = graph.read_symint(scalar);
+    EXPECT_TRUE(scalar_val_read == scalar_val);
+
     float val_a = i + 2.0f;
     float val_out = val_a + scalar_val;