Allow storage buffer limits beyond i32 max as a native extension

[rcoreilly]

Description

vulkaninfo on my NVIDIA H100 system reports:

maxStorageBufferRange = 4294967295

but the limits in wgpu (via wgpu_native 25.0.2.1) are:

MaxStorageBufferBindingSize = 2147483647

I could really use the larger buffers for my compute shaders...

shouldn't it be getting the same limits as vulkan reports?

[Wumpf]

Size is limited to i32 max right now since indexing at least at some point was (believed to be?) limited to i32.

see

• wgpu/wgpu-hal/src/auxil/mod.rs

  Line 56 in 0c0df24

  pub const MAX_I32_BINDING_SIZE: u32 = (1 << 31) - 1;

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

  Line 1271 in 0c0df24

  .min(crate::auxil::MAX_I32_BINDING_SIZE),

Warrants a deeper look if that's still and always the case. Could be that we can do u32 max only as a native only feature; binding stack, wgsl and output spir-v need to be scrutinized for this

Also affects dx12

[rcoreilly]

This appears to be the relevant discussion of this issue: gpuweb/gpuweb#1135

It looks like the conclusion was ultimately to support unsigned u32 indexing without limits, right?

Also, a typical storage buffer contains f32 values so even a signed i32 index can access the entire space of a (1 << 32)-1 sized storage buffer, given the 4 bytes per element. I believe the hardware uses u32 underneath typically in any case, right?

For reference, I'm using the full 4294967295 buffer sizes on the mac version of my code (in wgpu) and it works fine, so apparently the metal backend is fine with this already.

[cwfitzgerald]

The comment here explains why:

wgpu/wgpu-hal/src/auxil/mod.rs

Lines 52 to 56 in 0c0df24

	/// Maximum binding size for the shaders that only support `i32` indexing.
	/// Interestingly, the index itself can't reach that high, because the minimum
	/// element size is 4 bytes, but the compiler toolchain still computes the
	/// offset at some intermediate point, internally, as i32.
	pub const MAX_I32_BINDING_SIZE: u32 = (1 << 31) - 1;

There are driver bugs where binding offsets are computed using i32 byte offsets internally in the shader. In order to protect against them, we trim it down to i32::MAX. It depends on the backend, but metal doesn't have this problem.

WebGPU only actually guarentees 128MB

[rcoreilly]

@cwfitzgerald thanks for the explanation. For compute usage, the NVIDIA on Linux platform is the de-facto standard as I'm sure you know, and only being able to access 1/2 the available storage buffer size is a significant impediment for larger-scale compute usage. I'm running neural network models on H100 chips and they have 80Gb of available ram, and it is much more efficient to access that ram in 4Gb chunks rather than 2Gb.

I'm sure the NVIDIA driver does not have the bug, as I was able to use the larger buffers through vulkan directly. And are you saying that the maxStorageBufferRange reported by vulkan is therefore subject to these bugs -- i.e., you're preventing bugs that would otherwise show up if using vulkan directly?

In any case, would it be possible to add support for an environment variable, or even better, a white-list of known-safe drivers, to override this constraint?

I understand that this is not the standard use-case, but if you want to be able to replace vulkan as a serious compute platform for native use, I think this would be an essential thing to support.

[Wumpf]

I think it's worth considering putting this in as a native-only feature that has the appropriate bells and whistles. Just like we allow disabling some validations, we should allow extending the limit here. Might be a bit more wonky since this essentially overrides the previously reported limit, but I don't see why we shouldn't allow this if a user is opting in to the risk