[ET-VK] Fix metadata UBO VVL warnings (#7484)

* [ET-VK][ez] Fix undefined behaviour in ambiguous `ParamsBuffer` constructor

## Context

I discovered this bug when trying to execute the `vulkan_compute_api_test` binary on Windows. Almost all the tests were failing, with compute shaders producing incorrect results. After bisecting the change, it turns out the culprit is #7015. The diff introduced an alternative templated constructor for `ParamsBuffer` which would initialize an empty UBO with a specified size instead of wrapping a pre-existing object.

The issue is that these constructors are ambiguous because they both are template constructors and both only accept one argument. Therefore, the original constructor would be called when certain callsites intended to call the new constructor. This results in a UBO being created with an incorrect size, and resulted in the tensor's metadata being passed incorrectly into a compute shader.

To fix, I added a dummy argument into the new constructor for disambiguation purposes. I also changed it so that it's not templated, since there's no reason for it to be templated.

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

ghstack-source-id: 260031108
Pull Request resolved: #7478

* [ET-VK] Create Pipeline layouts with push constant ranges when required

## Context

#7223 added the ability to use push constants in shaders. However, one thing the diff missed was not specifying that the compute pipeline layout needed to include a push constant upon creation. The Vulkan validation layers warns against this, and on certain GPUs such as the integrated Intel GPU on my windows laptop compute shaders will produce incorrect output.

This diff makes the change such that the compute pipeline layout will be created with a push constant block if necessary.

## Solution

Change the key of the pipeline layout cache to accept an additional push constant size field. The push constant size will be used to create the pipeline layout with a push constant block of the specified size.

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

ghstack-source-id: 260031109
Pull Request resolved: #7479

* [ET-VK] Fix metadata UBO VVL warnings

## Context

Recently #7015 was implemented so that all tensor metadata (e.g. sizes, strides) would be stored in a single UBO instead of with separate UBO objects. This helps with memory savings presumably due to defragmentation of memory allocations.

However, once the change was introduced, I noticed two new warnings produced by the Vulkan Validation Layer.

The first complains that the offset of a UBO descriptor is not a multiple of the `minUniformBufferOffsetAlignment` field reported by the physical device properties.

The second complains that the range of a UBO descriptor exceeds the offset + range of the underlying UBO object.

# Solution

To address the first one, instead of using `sizeof(utils::ivec4)` to determine the offset per metadata field, check the `minUniformBufferOffsetAlignment` field of reported by the device and use that instead.

The second warning arises because the logic in the constructor of `BufferBindInfo` had a mistake; instead of using the range of the underlying UBO object, it should use the range subtracted by the user specified offset.

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

ghstack-source-id: 260031110
Pull Request resolved: #7480

---------

Co-authored-by: Stephen Jia <[email protected]>

Author: pytorchbot

backends/vulkan/runtime/api/containers/Tensor.cpp

@@ -658,66 +658,77 @@ utils::GPUMemoryLayout vTensor::estimate_memory_layout() const {
 }
 
 const vkapi::BufferBindInfo vTensor::sizes_ubo() {
+  const size_t size_per_ubo = context()->adapter_ptr()->min_ubo_alignment();
+  const size_t max_ubo_size = kMaxMetadataFieldCount * size_per_ubo;
   if (!uniforms_.buffer()) {
-    uniforms_ = ParamsBuffer(storage_.context_, kMaxUniformBufferSize, true);
+    uniforms_ = ParamsBuffer(storage_.context_, max_ubo_size, true);
   }
   if (sizes_uniform_offset_ == kUniformOffsetUnset) {
     VK_CHECK_COND(
-        (uniforms_size_ + kSizePerUniform) <= kMaxUniformBufferSize,
+        (uniforms_size_ + size_per_ubo) <= max_ubo_size,
         "Uniform data allocation has exceeded Tensor uniform buffer size");
     sizes_uniform_offset_ = uniforms_size_;
-    uniforms_size_ += kSizePerUniform;
+    uniforms_size_ += size_per_ubo;
     uniforms_.update(utils::make_whcn_ivec4(sizes_), sizes_uniform_offset_);
   }
-  return vkapi::BufferBindInfo(uniforms_.buffer(), sizes_uniform_offset_);
+  return vkapi::BufferBindInfo(
+      uniforms_.buffer(), sizes_uniform_offset_, size_per_ubo);
 }
 
 const vkapi::BufferBindInfo vTensor::strides_ubo() {
+  const size_t size_per_ubo = context()->adapter_ptr()->min_ubo_alignment();
+  const size_t max_ubo_size = kMaxMetadataFieldCount * size_per_ubo;
   if (!uniforms_.buffer()) {
-    uniforms_ = ParamsBuffer(storage_.context_, kMaxUniformBufferSize, true);
+    uniforms_ = ParamsBuffer(storage_.context_, max_ubo_size, true);
   }
   if (unsqueezed_strides_offset_ == kUniformOffsetUnset) {
     VK_CHECK_COND(
-        (uniforms_size_ + kSizePerUniform) <= kMaxUniformBufferSize,
+        (uniforms_size_ + size_per_ubo) <= max_ubo_size,
         "Uniform data allocation has exceeded Tensor uniform buffer size");
     unsqueezed_strides_offset_ = uniforms_size_;
-    uniforms_size_ += kSizePerUniform;
+    uniforms_size_ += size_per_ubo;
     uniforms_.update(
         utils::make_whcn_ivec4(unsqueezed_strides_),
         unsqueezed_strides_offset_);
   }
-  return vkapi::BufferBindInfo(uniforms_.buffer(), unsqueezed_strides_offset_);
+  return vkapi::BufferBindInfo(
+      uniforms_.buffer(), unsqueezed_strides_offset_, size_per_ubo);
 }
 
 const vkapi::BufferBindInfo vTensor::logical_limits_ubo() {
+  const size_t size_per_ubo = context()->adapter_ptr()->min_ubo_alignment();
+  const size_t max_ubo_size = kMaxMetadataFieldCount * size_per_ubo;
   if (!uniforms_.buffer()) {
-    uniforms_ = ParamsBuffer(storage_.context_, kMaxUniformBufferSize, true);
+    uniforms_ = ParamsBuffer(storage_.context_, max_ubo_size, true);
   }
   if (logical_limits_uniform_offset_ == kUniformOffsetUnset) {
     VK_CHECK_COND(
-        (uniforms_size_ + kSizePerUniform) <= kMaxUniformBufferSize,
+        (uniforms_size_ + size_per_ubo) <= max_ubo_size,
         "Uniform data allocation has exceeded Tensor uniform buffer size");
     logical_limits_uniform_offset_ = uniforms_size_;
-    uniforms_size_ += kSizePerUniform;
+    uniforms_size_ += size_per_ubo;
     uniforms_.update(logical_limits(), logical_limits_uniform_offset_);
   }
   return vkapi::BufferBindInfo(
-      uniforms_.buffer(), logical_limits_uniform_offset_);
+      uniforms_.buffer(), logical_limits_uniform_offset_, size_per_ubo);
 }
 
 const vkapi::BufferBindInfo vTensor::numel_ubo() {
+  const size_t size_per_ubo = context()->adapter_ptr()->min_ubo_alignment();
+  const size_t max_ubo_size = kMaxMetadataFieldCount * size_per_ubo;
   if (!uniforms_.buffer()) {
-    uniforms_ = ParamsBuffer(storage_.context_, kMaxUniformBufferSize, true);
+    uniforms_ = ParamsBuffer(storage_.context_, max_ubo_size, true);
   }
   if (numel_uniform_offset_ == kUniformOffsetUnset) {
     VK_CHECK_COND(
-        (uniforms_size_ + kSizePerUniform) <= kMaxUniformBufferSize,
+        (uniforms_size_ + size_per_ubo) <= max_ubo_size,
         "Uniform data allocation has exceeded Tensor uniform buffer size");
     numel_uniform_offset_ = uniforms_size_;
-    uniforms_size_ += kSizePerUniform;
+    uniforms_size_ += size_per_ubo;
     uniforms_.update(numel(), numel_uniform_offset_);
   }
-  return vkapi::BufferBindInfo(uniforms_.buffer(), numel_uniform_offset_);
+  return vkapi::BufferBindInfo(
+      uniforms_.buffer(), numel_uniform_offset_, size_per_ubo);
 }
 
 size_t vTensor::staging_buffer_numel() const {

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

@@ -348,16 +348,13 @@ class vTensor final {
   uint32_t numel_uniform_offset_;
   uint32_t logical_limits_uniform_offset_;
 
-  // Size allocated for each uniform
-  // each uniform is assumed to be a vec of 4 ints to maintain 16 byte alignemnt
-  constexpr static size_t kSizePerUniform = sizeof(utils::ivec4);
-  // Total size of tensor's uniform buffer
-  constexpr static size_t kMaxUniformBufferSize =
-      4 * // we have 4 uniforms that are passed on to shaders
-      kSizePerUniform;
-
-  // Initial value of uniform buffer offsets
-  constexpr static uint32_t kUniformOffsetUnset = kMaxUniformBufferSize;
+  // Maximum number of metadata fields that can be stored in the metadata UBO.
+  // This is used to calculate the size of the UBO that should be allocated.
+  constexpr static size_t kMaxMetadataFieldCount = 4;
+
+  // Initial value of uniform buffer offsets. 1 is selected as it is essentially
+  // impossible for a ubo to have an offset of 1.
+  constexpr static uint32_t kUniformOffsetUnset = 1;
 
   vTensorStorage storage_;
 

backends/vulkan/runtime/vk_api/Adapter.h

@@ -207,6 +207,10 @@ class Adapter final {
     return supports_8bit_storage_buffers() && supports_int8_shader_types();
   }
 
+  inline size_t min_ubo_alignment() const {
+    return physical_device_.min_ubo_alignment;
+  }
+
   // Command Buffer Submission
 
   void

backends/vulkan/runtime/vk_api/Descriptor.cpp

@@ -28,7 +28,17 @@ BufferBindInfo::BufferBindInfo(
     const uint32_t offset_p)
     : handle(buffer_p.handle()),
       offset(buffer_p.mem_offset() + offset_p),
-      range(buffer_p.mem_range()) {}
+      range(buffer_p.mem_range() - offset_p) {}
+
+BufferBindInfo::BufferBindInfo(
+    const VulkanBuffer& buffer_p,
+    const uint32_t offset_p,
+    const uint32_t range_p)
+    : handle(buffer_p.handle()),
+      offset(buffer_p.mem_offset() + offset_p),
+      range(range_p) {
+  VK_CHECK_COND(range_p <= (buffer_p.mem_range() - offset_p));
+}
 
 //
 // ParamsBindList

backends/vulkan/runtime/vk_api/Descriptor.h

@@ -34,6 +34,10 @@ struct BufferBindInfo final {
 
   BufferBindInfo();
   BufferBindInfo(const VulkanBuffer& buffer_p, const uint32_t offset_p = 0u);
+  BufferBindInfo(
+      const VulkanBuffer& buffer_p,
+      const uint32_t offset_p,
+      const uint32_t range_p);
 };
 
 struct ParamsBindList final {

backends/vulkan/runtime/vk_api/Device.cpp

@@ -39,10 +39,17 @@ PhysicalDevice::PhysicalDevice(VkPhysicalDevice physical_device_handle)
       num_compute_queues(0),
       supports_int16_shader_types(false),
       has_unified_memory(false),
-      has_timestamps(properties.limits.timestampComputeAndGraphics),
-      timestamp_period(properties.limits.timestampPeriod) {
+      has_timestamps(false),
+      timestamp_period(0),
+      min_ubo_alignment(0) {
   // Extract physical device properties
   vkGetPhysicalDeviceProperties(handle, &properties);
+
+  // Extract fields of interest
+  has_timestamps = properties.limits.timestampComputeAndGraphics;
+  timestamp_period = properties.limits.timestampPeriod;
+  min_ubo_alignment = properties.limits.minUniformBufferOffsetAlignment;
+
   vkGetPhysicalDeviceMemoryProperties(handle, &memory_properties);
 
   VkPhysicalDeviceFeatures2 features2{

backends/vulkan/runtime/vk_api/Device.h

@@ -49,6 +49,7 @@ struct PhysicalDevice final {
   bool has_unified_memory;
   bool has_timestamps;
   float timestamp_period;
+  size_t min_ubo_alignment;
 
   explicit PhysicalDevice(VkPhysicalDevice);
 };