Mistranslation of loop with continuing block

[Patryk27]

Hi,

Converting this wgsl:

var<private> global: vec3<u32>;
@group(0) @binding(2) 
var global_1: texture_storage_2d<rgba16float,read_write>;

fn function() {
    var phi_79_: i32;
    var phi_89_: i32;

    let _e17 = global;
    phi_79_ = 0;
    loop {
        let _e22 = phi_79_;
        let _e23 = (_e22 < 8);
        if _e23 {
            phi_89_ = (_e22 + 1);
        } else {
            textureStore(global_1, vec2<u32>(_e17.x, _e17.y), vec4<f32>(0.10000000149011612, 0.20000000298023224, 0.30000001192092896, 1.0));
            phi_89_ = 0;
        }
        let _e26 = phi_89_;
        continue;
        continuing {
            _ = !(select(false, true, _e23));
            phi_79_ = _e26;
            break if !(select(false, true, _e23));
        }
    }
    return;
}

@compute @workgroup_size(8, 8, 1) 
fn main(@builtin(global_invocation_id) param: vec3<u32>) {
    global = param;
    function();
}

... yields an invalid msl:

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

using metal::uint;

constant metal::float4 const_type_9_ = {0.10000000149011612, 0.20000000298023224, 0.30000001192092896, 1.0};

void function(
    thread metal::uint3& global,
    metal::texture2d<float, metal::access::read_write> global_1
) {
    int phi_79_ = {};
    int phi_89_ = {};
    metal::uint3 _e17 = global;
    phi_79_ = 0;
    bool loop_init = true;
    while(true) {
        if (!loop_init) {
            bool unnamed = !(((phi_79_ < 8) ? true : false));
            phi_79_ = phi_89_;
            if (!(((phi_79_ < 8) ? true : false))) {
                break;
            }
        }
        loop_init = false;
        int _e22 = phi_79_;
        bool _e23 = _e22 < 8;
        if (_e23) {
            phi_89_ = _e22 + 1;
        } else {
            global_1.write(const_type_9_, metal::uint2(metal::uint2(_e17.x, _e17.y)));
            phi_89_ = 0;
        }
        int _e26 = phi_89_;
        continue;
    }
    return;
}

struct main_Input {
};
kernel void main_(
  metal::uint3 param [[thread_position_in_grid]]
, metal::texture2d<float, metal::access::read_write> global_1 [[user(fake0)]]
) {
    metal::uint3 global = {};
    global = param;
    function(global, global_1);
}

... with a pretty subtle mistake related to the if (!loop_init) part:

if (!(((phi_79_ < 8) ? true : false))) {

Note that it uses phi_79_ with value after the assignment, instead of using the value from-before the assignment, which causes this shader to be effectively an infinite loop (later optimized out by Metal Shader Compiler into a no-op, apparently).

I've narrowed down the issue to this place:

https://github.com/gfx-rs/naga/blob/b7f4006e46313da063f8f2f930230767c9740239/src/back/msl/writer.rs#L2960

... which I'm not sure what is supposed to do, but just commenting-out self.named_expressions.remove(&handle); generates correct code 👀

[Patryk27]

Ah, I think I see what the logic is supposed to do - maybe a low-hanging fruit would be something like this?

diff --git a/src/back/msl/writer.rs b/src/back/msl/writer.rs
index cfef6645..30df421e 100644
--- a/src/back/msl/writer.rs
+++ b/src/back/msl/writer.rs
@@ -2536,6 +2536,7 @@ impl<W: Write> Writer<W> {
         level: back::Level,
         statements: &[crate::Statement],
         context: &StatementContext,
+        with_scope: impl FnOnce(&mut Self) -> BackendResult
     ) -> BackendResult {
         // Add to the set in order to track the stack size.
         #[cfg(test)]
@@ -2685,14 +2686,19 @@ impl<W: Write> Writer<W> {
                         let lif = level.next();
                         let lcontinuing = lif.next();
                         writeln!(self.out, "{lif}if (!{gate_name}) {{")?;
-                        self.put_block(lcontinuing, continuing, context)?;
-                        if let Some(condition) = break_if {
-                            write!(self.out, "{lcontinuing}if (")?;
-                            self.put_expression(condition, &context.expression, true)?;
-                            writeln!(self.out, ") {{")?;
-                            writeln!(self.out, "{}break;", lcontinuing.next())?;
-                            writeln!(self.out, "{lcontinuing}}}")?;
-                        }
+
+                        self.put_block(lcontinuing, continuing, context, |this| {
+                            if let Some(condition) = break_if {
+                                write!(this.out, "{lcontinuing}if (")?;
+                                this.put_expression(condition, &context.expression, true)?;
+                                writeln!(this.out, ") {{")?;
+                                writeln!(this.out, "{}break;", lcontinuing.next())?;
+                                writeln!(this.out, "{lcontinuing}}}")?;
+                            }
+
+                            Ok(())
+                        })?;
+
                         writeln!(self.out, "{lif}}}")?;
                         writeln!(self.out, "{lif}{gate_name} = false;")?;
                     } else {
@@ -2952,6 +2958,8 @@ impl<W: Write> Writer<W> {
             }
         }
 
+        with_scope()?;
+
         // un-emit expressions
         //TODO: take care of loop/continuing?
         for statement in statements {

[teoxoy]

Thanks for the report!

I've narrowed down the issue to this place:

https://github.com/gfx-rs/naga/blob/b7f4006e46313da063f8f2f930230767c9740239/src/back/msl/writer.rs#L2960

... which I'm not sure what is supposed to do, but just commenting-out self.named_expressions.remove(&handle); generates correct code 👀

It doesn't seem so.

diff --git a/t.metal b/t2.metal
index 14e5c6b0..672badfc 100644
--- a/t.metal
+++ b/t2.metal
@@ -19,7 +19,7 @@ void function(
             bool unnamed = !(((phi_79_ < 8) ? true : false));
             phi_79_ = phi_89_;
             bool unnamed_1 = !(((phi_79_ < 8) ? true : false));
-            if (!(((phi_79_ < 8) ? true : false))) {
+            if (unnamed_1) {
                 break;
             }
         }

[teoxoy]

It looks like the HLSL and GLSL backends have the same issue since they all use the same strategy. The SPIR-V backend seems to be behaving properly.

[Patryk27]

It doesn't seem so.

Whoops, my bad - I've been actually checking conversion from spv (attached below) into msl which does get fixed by commenting that line - but yeah, going spv -> wgsl -> msl remains broken 🤔

; SPIR-V
; Version: 1.3
; Generator: Google rspirv; 0
; Bound: 96
; Schema: 0
               OpCapability Shader
               OpCapability Int8
               OpMemoryModel Logical Simple
               OpEntryPoint GLCompute %1 "main" %gl_GlobalInvocationID
               OpExecutionMode %1 LocalSize 8 8 1
               OpDecorate %gl_GlobalInvocationID BuiltIn GlobalInvocationId
               OpDecorate %8 Binding 2
               OpDecorate %8 DescriptorSet 0
       %uint = OpTypeInt 32 0
     %v3uint = OpTypeVector %uint 3
%_ptr_Input_v3uint = OpTypePointer Input %v3uint
       %void = OpTypeVoid
         %16 = OpTypeFunction %void
%gl_GlobalInvocationID = OpVariable %_ptr_Input_v3uint Input
      %float = OpTypeFloat 32
         %20 = OpTypeImage %float 2D 2 0 0 2 Rgba16f
%_ptr_UniformConstant_20 = OpTypePointer UniformConstant %20
     %v2uint = OpTypeVector %uint 2
        %int = OpTypeInt 32 1
      %int_0 = OpConstant %int 0
       %bool = OpTypeBool
      %int_8 = OpConstant %int 8
      %int_1 = OpConstant %int 1
       %true = OpConstantTrue %bool
    %v4float = OpTypeVector %float 4
%float_0_100000001 = OpConstant %float 0.100000001
%float_0_200000003 = OpConstant %float 0.200000003
%float_0_300000012 = OpConstant %float 0.300000012
    %float_1 = OpConstant %float 1
      %false = OpConstantFalse %bool
         %38 = OpUndef %int
          %8 = OpVariable %_ptr_UniformConstant_20 UniformConstant
         %94 = OpConstantComposite %v4float %float_0_100000001 %float_0_200000003 %float_0_300000012 %float_1
          %1 = OpFunction %void None %16
         %39 = OpLabel
         %40 = OpLoad %v3uint %gl_GlobalInvocationID
         %73 = OpCompositeExtract %uint %40 0
         %74 = OpCompositeExtract %uint %40 1
         %75 = OpCompositeConstruct %v2uint %73 %74
               OpBranch %78
         %78 = OpLabel
         %79 = OpPhi %int %int_0 %39 %89 %90
               OpLoopMerge %91 %90 None
               OpBranch %80
         %80 = OpLabel
         %81 = OpSLessThan %bool %79 %int_8
               OpSelectionMerge %87 None
               OpBranchConditional %81 %82 %84
         %82 = OpLabel
         %83 = OpIAdd %int %79 %int_1
               OpBranch %87
         %84 = OpLabel
         %93 = OpLoad %20 %8
               OpImageWrite %93 %75 %94
               OpBranch %87
         %87 = OpLabel
         %89 = OpPhi %int %83 %82 %38 %84
         %95 = OpSelect %bool %81 %true %false
               OpBranch %90
         %90 = OpLabel
               OpBranchConditional %95 %78 %91
         %91 = OpLabel
               OpReturn
               OpFunctionEnd

fwiw, I've been able to temporarily fix this issue by doing Patryk27/naga@fdcdb4f - though it feels more like a band-aid (plus it should probably call self.namer.call("loop_break");) and so I didn't try to upstream it

[jimblandy]

@Patryk27 If you're feeding Naga SPIR-V as input, we recently did some work on the translation of SPIR-V with continuing blocks in gfx-rs/naga#2290. It might be worth dumping the Naga IR (use the cli and ask for a .txt output file) and making sure that you're getting the Loop statement you'd expect.

[Patryk27]

@jimblandy: yeah, I saw that - when it comes to bugs, I'm always testing on masters 😄

That being said, it takes a moment to decode how those emits and loops come together, so I'm not really sure whether the output is correct or not - if you'd like to take a look, here it is:

output.txt

[eddyb]

I do not understand why break if turns into that kind of conditional break.

AIUI it was only added to WGSL to represent SPIR-V conditional backedges.
And SPIR-V itself only has conditional backedges to encode do-while loops.
So that's the natural translation.

That is, loop { ... continuing { ... break if !(select(false, true, _e23)); } } should turn into
do { ... } while(!!(_e23 ? true : false)); - without any overcomplicated propagation across the backedge.

[jimblandy]

@eddyb Agreed - could you file an issue to that effect?

[eddyb]

could you file an issue to that effect?

• Naga should map break if 1:1 to the backend (i.e. do-while for GLSL/HLSL/MSL). #4560

[Patryk27]

Simpler repro:

fn main() {
    var a = 0;

    loop {
        let cmp = a < 8;

        if cmp {
            a += 1;
        }

        continuing {
            break if cmp;
        }
    }
}

... gets converted into (invalid):

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

using metal::uint;


void main_(
) {
    int a = 0;
    bool loop_init = true;
    while(true) {
        if (!loop_init) {
            if (a < 8) { // aii aii
                break;
            }
        }
        loop_init = false;
        int _e2 = a;
        bool cmp = _e2 < 8;
        if (cmp) {
            int _e6 = a;
            a = _e6 + 1;
        }
    }
    return;
}

[Patryk27]

Curiously, using var cmp seems to generate the correct code!

fn main() {
    var a = 0;

    loop {
        var cmp = a < 8;

        a += 1;

        continuing {
            break if cmp;
        }
    }
}

... gets transformed into:

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

using metal::uint;


void main_(
) {
    int a = 0;
    bool cmp = {};
    bool loop_init = true;
    while(true) {
        if (!loop_init) {
            bool _e9 = cmp;
            if (cmp) {
                break;
            }
        }
        loop_init = false;
        int _e2 = a;
        cmp = _e2 < 8;
        int _e7 = a;
        a = _e7 + 1;
    }
    return;
}

[Patryk27]

Also, it's possible to crash the codegen as well 👀

fn main() {
    loop {
        let cmp = magic();

        continuing {
            break if cmp;
        }
    }
}

fn magic() -> bool {
    return true;
}

thread 'main' panicked at 'internal error: entered unreachable code', naga/src/back/msl/writer.rs:1965:17

[teoxoy]

Following the discussion in #4982 which seems to be a duplicate of this: This is actually an issue with the WGSL frontend.

[Patryk27]

Following the discussion in #4982 which seems to be a duplicate of this: This is actually an issue with the WGSL frontend.

Note that my original bug happens when going directly from spv to msl as well (I just posted wgsl, 'cause I thought it might be easier to debug) - right now I think there's two bugs, actually:

1. Inside wgsl frontend, when parsing break_if (which is fixed by my suggestion in that other thread, for which I'll prepare a PR),
2. Inside msl & hlsl & glsl backends, when outputting the loop_init idiom (this one is not fixed by my suggestion there, there's just something fundamentally wrong when break if refers to let instead of a var variable).

[teoxoy]

I see, reopening.

[Patryk27]

Alright, the main culprit seems to be this part:

wgpu/naga/src/front/wgsl/lower/mod.rs

Line 1158 in 376d901

ctx.local_table.insert(l.handle, Typed::Plain(value));

... which paired with:

wgpu/naga/src/front/wgsl/lower/mod.rs

Line 1545 in 376d901

return Ok(rctx.local_table[&local]);

... causes it to essentially copy-paste the let's expression instead of treating it as a read-only variable.

Perhaps a solution here could be to treat both lets and vars as variables, but just make sure let cannot be overwritten? (plus, of course, handle correctly logic such as var x; being correct, but let x; not)

Currently it looks like there's some special logic in a few places, which seems to treat lets as sort of const expressions, which doesn't seem correct.

edit: on a second thought, this pertains the wgsl frontend only, while the issue is present when going directly from spv as well, so the actual bug must be somewhere else.