WebGPU: TransposeProgram WGSL shader fails with type mismatch on Chrome 141+ (Tint IR rewrite) #8638
Description
Description
After Chrome 141's Tint WGSL compiler overhaul (which replaced AST transformations with a new IR), the TransposeProgram shader fails to compile with a type mismatch error.
Error Message
Error while parsing WGSL: :145:41 error: type mismatch for argument 2 in call to 'setOutputAtIndex', expected 'f32', got 'i32'
setOutputAtIndex(flatIndex, A[getIndexFromCoords3D(
^^^^^^^^^^^^^^^^^^^^^^^
vec3<i32>(coords.y,coords.x,coords.z), uniforms.aShape)]);
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
- While calling [Device].CreateShaderModule([ShaderModuleDescriptor "TransposeProgram"]).
[Invalid ShaderModule "TransposeProgram"] is invalid.
- While validating compute stage ([Invalid ShaderModule "TransposeProgram"], entryPoint: "_start").
- While calling [Device].CreateComputePipeline([ComputePipelineDescriptor "TransposeProgram"]).
Environment
- TensorFlow.js version: 4.22.0
- Browser: Chrome 141+ (with new Tint IR-based WGSL compiler)
- Backend: WebGPU
- OS: macOS (also reported on other platforms)
Root Cause Analysis
Chrome 141 completed a 2.5-year project to overhaul Tint's internals, inserting an Intermediate Representation (IR) between the AST and backend code generators. This new compiler has stricter type validation.
The issue is in transpose_webgpu.ts:
setOutputAtIndex(flatIndex, A[getIndexFromCoords3D(vec3<i32>(...), uniforms.aShape)]);When tensor A has an integer data type, A[index] returns i32, but setOutputAtIndex expects f32. The old Tint compiler was lenient about this implicit conversion; the new IR-based compiler correctly rejects it.
Suggested Fix
Add explicit f32() casting when reading from integer tensors:
setOutputAtIndex(flatIndex, f32(A[getIndexFromCoords3D(...)]));Or conditionally cast based on the tensor's dtype.
References
- Chrome 141 WebGPU - Tint IR Overhaul
- Related past fixes: PR [webgpu] Fix wrongly used u32 of WGSL #5559, PR [webgpu] Use i32 as coords type in WGSL #5628 (similar i32/f32 type issues)
Workaround
Use WebGL backend instead of WebGPU until this is fixed.