Skip to content

Commit 8306cb5

Browse files
pytorchbotSS-JIA
andauthored
[ET-VK][ez] Fix undefined behaviour in ambiguous ParamsBuffer constructor (#7482)
## 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 Co-authored-by: Stephen Jia <[email protected]>
1 parent 2600cc8 commit 8306cb5

File tree

2 files changed

+7
-6
lines changed

2 files changed

+7
-6
lines changed

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

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -31,8 +31,9 @@ class ParamsBuffer final {
3131
vulkan_buffer_(
3232
context_p_->adapter_ptr()->vma().create_params_buffer(block)) {}
3333

34-
template <typename Block>
35-
ParamsBuffer(Context* context_p, const VkDeviceSize nbytes)
34+
// The last bool argument, though unused, is required to disambiguate this
35+
// constructor from the one above.
36+
ParamsBuffer(Context* context_p, const VkDeviceSize nbytes, const bool unused)
3637
: context_p_(context_p),
3738
vulkan_buffer_(
3839
context_p_->adapter_ptr()->vma().create_uniform_buffer(nbytes)) {}

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

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -659,7 +659,7 @@ utils::GPUMemoryLayout vTensor::estimate_memory_layout() const {
659659

660660
const vkapi::BufferBindInfo vTensor::sizes_ubo() {
661661
if (!uniforms_.buffer()) {
662-
uniforms_ = ParamsBuffer(storage_.context_, kMaxUniformBufferSize);
662+
uniforms_ = ParamsBuffer(storage_.context_, kMaxUniformBufferSize, true);
663663
}
664664
if (sizes_uniform_offset_ == kUniformOffsetUnset) {
665665
VK_CHECK_COND(
@@ -674,7 +674,7 @@ const vkapi::BufferBindInfo vTensor::sizes_ubo() {
674674

675675
const vkapi::BufferBindInfo vTensor::strides_ubo() {
676676
if (!uniforms_.buffer()) {
677-
uniforms_ = ParamsBuffer(storage_.context_, kMaxUniformBufferSize);
677+
uniforms_ = ParamsBuffer(storage_.context_, kMaxUniformBufferSize, true);
678678
}
679679
if (unsqueezed_strides_offset_ == kUniformOffsetUnset) {
680680
VK_CHECK_COND(
@@ -691,7 +691,7 @@ const vkapi::BufferBindInfo vTensor::strides_ubo() {
691691

692692
const vkapi::BufferBindInfo vTensor::logical_limits_ubo() {
693693
if (!uniforms_.buffer()) {
694-
uniforms_ = ParamsBuffer(storage_.context_, kMaxUniformBufferSize);
694+
uniforms_ = ParamsBuffer(storage_.context_, kMaxUniformBufferSize, true);
695695
}
696696
if (logical_limits_uniform_offset_ == kUniformOffsetUnset) {
697697
VK_CHECK_COND(
@@ -707,7 +707,7 @@ const vkapi::BufferBindInfo vTensor::logical_limits_ubo() {
707707

708708
const vkapi::BufferBindInfo vTensor::numel_ubo() {
709709
if (!uniforms_.buffer()) {
710-
uniforms_ = ParamsBuffer(storage_.context_, kMaxUniformBufferSize);
710+
uniforms_ = ParamsBuffer(storage_.context_, kMaxUniformBufferSize, true);
711711
}
712712
if (numel_uniform_offset_ == kUniformOffsetUnset) {
713713
VK_CHECK_COND(

0 commit comments

Comments
 (0)