WGSL backend writes invalid WGSL in some cases

[schell]

Given this SPV input
(fragment_equirectangular.spv.zip) naga will write invalid wgsl output. You can see in this generated wgsl source that there is an undefined variable inf:

var<private> global: vec3<f32>;
@group(0) @binding(2) 
var global_1: sampler;
@group(0) @binding(1) 
var global_2: texture_2d<f32>;
var<private> global_3: vec4<f32>;

fn function() {
    var phi_81_: bool;
    var phi_86_: vec3<f32>;

    let _e13 = global;
    let _e21 = (1.0 / sqrt(fma(_e13.z, _e13.z, fma(_e13.x, _e13.x, (_e13.y * _e13.y)))));
    if (bitcast<f32>((bitcast<u32>(_e21) & 2147483647u)) < inf) {
        phi_81_ = (_e21 > 0.0);
    } else {
        phi_81_ = false;
    }
    let _e28 = phi_81_;
    if _e28 {
        phi_86_ = (_e13 * _e21);
    } else {
        phi_86_ = vec3<f32>(0.0, 0.0, 0.0);
    }
    let _e31 = phi_86_;
    let _e40 = textureSample(global_2, global_1, vec2<f32>(fma(atan2(_e31.z, _e31.x), 0.1591, 0.5), fma(asin(_e31.y), 0.3183, 0.5)));
    global_3 = vec4<f32>(_e40.x, _e40.y, _e40.z, 1.0);
    return;
}

@fragment 
fn fragment_equirectangular(@location(0) param: vec3<f32>) -> @location(0) vec4<f32> {
    global = param;
    function();
    let _e3 = global_3;
    return _e3;
}

This source fails to parse - it errs with InvalidHeader.

Here's my test case:

        let opts = naga::front::spv::Options::default();
        let module =
            naga::front::spv::parse_u8_slice(include_bytes!("fragment_equirectangular.spv"), &opts)
            .unwrap();
        let mut validator =
            naga::valid::Validator::new(Default::default(), naga::valid::Capabilities::empty());
        let info = validator.validate(&module).unwrap();
        let wgsl =
            naga::back::wgsl::write_string(&module, &info, naga::back::wgsl::WriterFlags::empty())
            .unwrap();
        println!("wgsl: {wgsl}");
        // fails to parse, InvalidHeader
        let module = naga::front::spv::parse_u8_slice(wgsl.as_bytes(), &opts)
            .unwrap();
        let mut validator =
            naga::valid::Validator::new(Default::default(), naga::valid::Capabilities::empty());
        let _info = validator.validate(&module).unwrap();

[teoxoy]

The InvalidHeader error is due to wgsl text being passed to the spv frontend. You should pass the wgsl text to the wgsl frontend instead.

As for the undefined inf variable, is it undefined only when outputting wgsl?

[schell]

🤦 . Silly me!

After fixing that bumble the WGSL parser spits out an error that's much more on-topic:

called `Result::unwrap()` on an `Err` value: ParseError { message: "no definition in scope for identifier: 'inf'", labels: [(Span { start: 407, end: 410 }, "unknown identifier")], notes: [] }

When compiled to MSL it looks like the identifier is INFINITY so maybe the WGSL backend is assuming there is a constant inf in scope?

// language: metal2.0
#include <metal_stdlib>
#include <simd/simd.h>

using metal::uint;


void function(
    thread metal::float3& global,
    metal::sampler global_1,
    metal::texture2d<float, metal::access::sample> global_2,
    thread metal::float4& global_3
) {
    bool phi_81_ = {};
    metal::float3 phi_86_ = {};
    metal::float3 _e13 = global;
    float _e21 = 1.0 / metal::sqrt(metal::fma(_e13.z, _e13.z, metal::fma(_e13.x, _e13.x, _e13.y * _e13.y)));
    if (as_type<float>(as_type<uint>(_e21) & 2147483647u) < INFINITY) {
        phi_81_ = _e21 > 0.0;
    } else {
        phi_81_ = false;
    }
    bool _e28 = phi_81_;
    if (_e28) {
        phi_86_ = _e13 * _e21;
    } else {
        phi_86_ = metal::float3(0.0, 0.0, 0.0);
    }
    metal::float3 _e31 = phi_86_;
    metal::float4 _e40 = global_2.sample(global_1, metal::float2(metal::fma(metal::atan2(_e31.z, _e31.x), 0.1591, 0.5), metal::fma(metal::asin(_e31.y), 0.3183, 0.5)));
    global_3 = metal::float4(_e40.x, _e40.y, _e40.z, 1.0);
    return;
}

struct fragment_equirectangularInput {
    metal::float3 param [[user(loc0), center_perspective]];
};
struct fragment_equirectangularOutput {
    metal::float4 member [[color(0)]];
};
fragment fragment_equirectangularOutput fragment_equirectangular(
  fragment_equirectangularInput varyings [[stage_in]]
, metal::sampler global_1 [[user(fake0)]]
, metal::texture2d<float, metal::access::sample> global_2 [[user(fake0)]]
) {
    metal::float3 global = {};
    metal::float4 global_3 = {};
    const auto param = varyings.param;
    global = param;
    function(global, global_1, global_2, global_3);
    metal::float4 _e3 = global_3;
    return fragment_equirectangularOutput { _e3 };
}

[teoxoy]

Ah, interesting; inf is actually an f32 literal. We currently use rust's debug implementation for printing floats.

This will be fixed by gfx-rs/wgpu#4568.

However, WGSL doesn't currently support infinities or NaN. It seems we don't have validation for that yet but will in the future.

[schell]

Ok, thanks @teoxoy.

This must be happening because of my use of glam::Vec3::normalize_or_zero, which compares against infinity somewhere in its call tree (I'm using rust-gpu to compile Rust to SPIR-V for those following along).

So the short-term fix is to write a new implentation of normalize_or_zero that doesn't use infinity. What's the long term solution @teoxoy? Will the changes you mention make it "just work" or will comparing against infinities never be supported?

[teoxoy]

Will the changes you mention make it "just work"

No, it will only fix the "no definition in scope for identifier: 'inf'" error. But in the future validation will complain instead.

The long term solution to fully support infinities and NaNs will most likely be an extension that would map to SPIR-V's SPV_KHR_float_controls (gpuweb/gpuweb#2259 tracks the standardization of such a feature). I haven't looked into what the other backends need to support those.

[schell]

Thanks @teoxoy. Indeed it seems the same thing happens with any references to NaN that get produced.

[teoxoy]

With #2508 you should now get (for the SPIR-V shader in the description):

Constant expression [6] is invalid: 
        Float literal is infinite
Generating wgsl output requires validation to succeed, and it failed in a previous step