[kolkov]

Summary

On the gogpu/wgpu software adapter, Sum returns wrong values for tensors of 1024 elements or more. The GPU reduction (globalSumShader, globalSumShaderInt32) relies on workgroup memory shared across workgroupBarrier, and the software backend doesn't implement that yet:

• up to wgpu v0.34.6, each invocation sees its own copy, so a sum of 4096 ones comes out as 16;
• v0.34.7 refuses the pipeline, so Sum would panic instead;
• the full fix is planned for v0.34.8.

Software isn't only for CI. New() falls back to it when no hardware adapter comes up, so on a GPU-less server this silently corrupts any loss or metric that sums a large tensor.

Changes

• Sum now checks the sum kernels once per backend, on the first Sum that would use them. It runs them on a known input that spans several workgroups. If the pipeline is refused or the result is wrong, Sum runs on the CPU, which is correct on every device. The check is keyed on the result rather than a wgpu version or error value, so it works with released wgpu (silently wrong), v0.34.7 (refused) and v0.34.8 (fixed, GPU path used again).
• compileShader and getOrCreatePipeline get error-returning variants (tryCompileShader, tryGetOrCreatePipeline) that wrap the device error with %w. The panicking versions call them, and their messages are unchanged.
• Shader selection for the eager and lazy sum paths moves into one sumKernel helper.
• New TestSumAcrossWorkgroups: sizes 1, 1023, 1024, 1025, 4096, 100000 and 100001, f32 and i32, eager and lazy, compared exactly against the CPU sum. Against the old code on the software adapter, 21 of its cases fail.

The check also caught DX12 on wgpu v0.34.4 here, where the kernel read back 0, and fell back to the CPU there too.

Test plan

• TestSumAcrossWorkgroups on software (fallback taken: "globalSum kernel returned -2 for a sum of -6"), Vulkan and DX12 locally
• The same test against main fails on software (21 cases)
• golangci-lint v2.15.0 clean on the package
• CI green (CI runs -short, where GPU tests skip)

[codecov]

Codecov Report

❌ Patch coverage is 56.57895% with 33 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
internal/backend/webgpu/compute.go	56.94%	21 Missing and 10 partials ⚠️
internal/backend/webgpu/lazy_compute.go	50.00%	1 Missing and 1 partial ⚠️

📢 Thoughts on this report? Let us know!