Robust Buffer Access conflicts with lazy initialization

[kvark]

WebGPU allows for RBA behavior on out-of-bound accesses into buffers. This behavior says that 0, 1, or any value from the underlying buffer (not binding!) can be returned or written.
However, if the buffer ranges are lazily initialized to zero, we can end up in a situation where shader reads values outside of the initialized range. This is not acceptable.

A quick solution would be to consider any storage buffer bindings to cover the whole buffer for the purpose of lazy initialization. cc @Wumpf

Note also that the extent of the problem depends on the bounds checking policy:

wgpu/wgpu-hal/src/vulkan/adapter.rs

Line 843 in 3afab9c

index_bounds_check_policy: naga::back::IndexBoundsCheckPolicy::Restrict,

. If it's ReadZeroSkipWrite then we can keep doing what we do today. cc @jimblandy

[jimblandy]

On Metal, Firefox always uses ReadZeroSkipWrite for buffer accesses because wgpu_server_device_create_shader_module always requests bounds checks.

On Vulkan, the spec says:

• A buffer access is considered to be out of bounds if any of the following are true: ...
  • The pointer was not formed by OpImageTexelPointer and the object pointed to is not wholly contained within the bound range.

Since that refers to "the bound range", not the buffer, and wgpu-core does initialize the bound range, I think it's impossible for shaders to see uninitialized buffer contents on Vulkan as well. I believe "the bound range" is what's established via the vkDescriptorBufferInfo passed to vkUpdateDescriptorSets, and wgpu-hal sets that up based on the bound range in the wgpu-core bind group.

On D3D12, we don't implement any bounds checks at all, I believe because the platform promises WebGPU-compatible behavior itself.

So I think this is not a bug.

[jimblandy]

In the Vulkan spec:

VK_EXT_robustness2
...
Issues:

1. Why do VkPhysicalDeviceRobustness2PropertiesEXT::robustUniformBufferAccessSizeAlignment and VkPhysicalDeviceRobustness2PropertiesEXT::robustStorageBufferAccessSizeAlignment exist?

   RESOLVED: Some implementations cannot efficiently tightly bounds-check all buffer accesses. Rather, the size of the bound range is padded to some power of two multiple, up to 256 bytes for uniform buffers and up to 4 bytes for storage buffers, and that padded size is bounds-checked. This is sufficient to implement D3D-like behavior, because D3D only allows binding whole uniform buffers or ranges that are a multiple of 256 bytes, and D3D raw and structured buffers only support 32-bit accesses.

[kvark]

So, if RBA2 is supported, we need to respect those access alignments for the purpose of range initialization.
If RBA2 is not supported, we can't even rely on a particular alignment there, and it's technically "any value within the buffer".

[jimblandy]

@kvark That seems to contradict the language in the description of vkPhysicalDeviceFeatures::robustBufferAccess that I quoted above, doesn't it?

I wonder why they didn't update that, especially with the mentions of robustBufferAccess2 in immediately subsequent paragraphs.

[jimblandy]

Ahh, later the Vulkan spec says:

• Out-of-bounds buffer loads will return any of the following values:
  ...
  • Values from anywhere within the memory range(s) bound to the buffer (possibly including bytes of memory past the end of the buffer, up to the end of the bound range).

So Vulkan agrees with WebGPU on what is considered "out of bounds", but doesn't meet its requirements for how loads of such references should behave.

[jimblandy]

@kvark You said:

A quick solution would be to consider any storage buffer bindings to cover the whole buffer for the purpose of lazy initialization.

Why would it be limited to storage buffers? Doesn't this mean that pretty much any buffer access could read uninitialized content from elsewhere in the buffer?

It looks to me like this makes lazy initialization difficult to keep.

[jimblandy]

Part of the rationale here is that it's assumed that, even if the platform's bounds-checking tools have some kind of required alignment, up to 256, for the enforcement, we can just round up [edit] binding sizes and buffer sizes to the nearest such boundary, and let platform-provided bounds checks return values from within the buffer.

VkPhysicalDeviceRobustness2PropertiesEXT includes the following two fields:

• robustStorageBufferAccessSizeAlignment is the number of bytes that the range of a storage buffer descriptor is rounded up to when used for bounds-checking when the robustBufferAccess2 feature is enabled. This value must be either 1 or 4.

• robustUniformBufferAccessSizeAlignment is the number of bytes that the range of a uniform buffer descriptor is rounded up to when used for bounds-checking when the robustBufferAccess2 feature is enabled. This value must be a power of two in the range [1, 256].

If we round up buffer sizes to these amounts, and then round up lazy initialization regions to these amounts, that should suffice to meet the spec's requirements.

I think we want to handle this as follows:

• Extend wgpu_hal::Alignments (available from ExposedAdapter) to provide the two alignments above.

  • On Vulkan, we obtain them from VkPhysicalDeviceRobustness2PropertiesEXT.
  • On Metal, since wgpu-hal directs Naga to inject its own bounds checks, we can set these to 1.
  • On Direct3D 12, my understanding is that what you ask for is what you get, so it can be 1 there too.

• Document wgpu-hal to promise that, when a bindgroup is created, shaders can only access the bound portion of the buffer, rounded up to the values from Alignments.

• Fix wgpu-core to round up lazy initialization regions to the next multiple of those alignments.

[jimblandy]

Testing this is a problem. By definition, a test needs to be able to detect whether the change to the lazy initialization code is present or absent. That's tricky here, for a few reasons:

• We don't know what the contents of the uninitialized portion of the buffer are. Maybe they're just naturally zero, in which case the test can't distinguish between the fixed and broken lazy initialization code.

• Platforms may or may not permit us to see beyond the end of bound ranges, introducing platform variation that is irrelevant to what we want to observe.

One approach to the second problem might be to add a wgpu_types::BufferSize field to Device whose value is added to the size of all bindings just before they're passed to wgpu_hal, which could be set by tests. That would create a mismatch between the enforcement provided by wgpu_hal (whether by relying on platform facilities or having Naga inject bounds checks),

I'm not sure how to get around the first problem. It doesn't help that Direct3D 12 lacks any way to fill buffers, short of dispatching a compute shader to write to them.

[jimblandy]

Because every approach I can think of to test this is very involved, and this issue has been on our P0 list for a very long time, I'd like to land a fix without additional tests.

[teoxoy]

• Document wgpu-hal to promise that, when a bindgroup is created, shaders can only access the bound portion of the buffer, rounded up to the values from Alignments.

One detail that's worth bringing up is that storage buffer bindings shouldn't be enlarged because they can contain runtime-sized arrays.

The WebGPU spec does say this:

• effective buffer binding size(resource) is a multiple of 4.
• resource.offset is a multiple of limits.minStorageBufferOffsetAlignment.

So I think we only need to make sure that minStorageBufferOffsetAlignment is at least as high as robustStorageBufferAccessSizeAlignment to satisfy Vulkan's requirements.

I think we should be fine here since the buffer binding size is already guaranteed to be a multiple of 4.

Now for uniform buffers:

The WebGPU spec only requires:

• resource.offset is a multiple of limits.minUniformBufferOffsetAlignment.

Given that uniform buffers can't contain runtime-sized arrays we can enlarge bindings but we should make sure that minUniformBufferOffsetAlignment is at least as high as robustUniformBufferAccessSizeAlignment so that we don't have to do heroics to pad data structures in the shader.

EDIT: crossed sections that don't apply. The robust{Storage,Uniform}BufferAccessSizeAlignments only apply to the size of the bindings, and not to the offset of the binding.

[teoxoy]

Actually, if I read this section of the Vulkan spec correctly it sounds like we don't need to do anything?

[jimblandy]

I think we should be fine here since the buffer binding size is already guaranteed to be a multiple of 4.

Okay, yes. To restate: the most lenient Vulkan VK_EXT_robustness2 implementation won't effectively round up the size of any binding that WebGPU would permit. WebGPU says, for a "storage" and "read-only-storage" binding to be valid, it must be true that:

effective buffer binding size(resource) is a multiple of 4.

Then, Vulkan says, in describing VkPhysicalDeviceRobustness2PropertiesEXT:

• robustStorageBufferAccessSizeAlignment is the number of bytes that the range of a storage buffer descriptor is rounded up to when used for bounds-checking when the robustBufferAccess2 feature is enabled. This value must be either 1 or 4.

So wgpu-hal can simply promise that storage bindings are accurately checked.

[jimblandy]

One detail that's worth bringing up is that storage buffer bindings shouldn't be enlarged because they can contain runtime-sized arrays.

Yes - we shouldn't actually change binding sizes. Rather, wgpu-hal should simply note how the requested binding size relates to what shaders can actually get at.

[teoxoy]

Ignore this comment #1813 (comment), I misunderstood what the spec was saying.

The plan in #1813 (comment) makes sense to me now 👍