[naga wgsl-in] Do not eagerly concretize local const declarations of abstract types

jamienicol · jimblandy · commit 7df6e477cf0b · 2025-03-05T20:39:48.000-08:00
Instead allow the const to be converted each time it is const
evaluated as part of another expression. This allows an abstract const
to be used as a different type depending on the context.

As a result, abstract types may now find their way in to the IR, which
we don't want. This occurs because the compact pass treats named
expressions as used, mostly so that our snapshot tests are more
useful, and therefore does not remove them. To prevent this, we avoid
adding abstract-typed local consts to the named expressions list. This
will have no functional effect on any shaders produced by the
backends, but some unused local const declarations will no longer be
present.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -198,7 +198,7 @@ By @Vecvec in [#6905](https://github.com/gfx-rs/wgpu/pull/6905), [#7086](https:/
 - Error if structs have two fields with the same name. By @SparkyPotato in [#7088](https://github.com/gfx-rs/wgpu/pull/7088).
 - Forward '--keep-coordinate-space' flag to GLSL backend in naga-cli. By @cloone8 in [#7206](https://github.com/gfx-rs/wgpu/pull/7206).
 - Allow template lists to have a trailing comma. By @KentSlaney in [#7142](https://github.com/gfx-rs/wgpu/pull/7142).
-- Allow WGSL const declarations to have abstract types. By @jamienicol in [#7055](https://github.com/gfx-rs/wgpu/pull/7055).
+- Allow WGSL const declarations to have abstract types. By @jamienicol in [#7055](https://github.com/gfx-rs/wgpu/pull/7055) and [#7222](https://github.com/gfx-rs/wgpu/pull/7222).
 
 #### General
 
diff --git a/naga/src/front/wgsl/lower/mod.rs b/naga/src/front/wgsl/lower/mod.rs
@@ -1579,21 +1579,25 @@ impl<'source, 'temp> Lowerer<'source, 'temp> {
                         c.ty.map(|ast| self.resolve_ast_type(ast, &mut ectx.as_const()))
                             .transpose()?;
 
-                    let (_ty, init) = self.type_and_init(
+                    let (ty, init) = self.type_and_init(
                         c.name,
                         Some(c.init),
                         explicit_ty,
-                        AbstractRule::Concretize,
+                        AbstractRule::Allow,
                         &mut ectx.as_const(),
                     )?;
                     let init = init.expect("Local const must have init");
 
                     block.extend(emitter.finish(&ctx.function.expressions));
                     ctx.local_table
                         .insert(c.handle, Declared::Const(Typed::Plain(init)));
-                    ctx.named_expressions
-                        .insert(init, (c.name.name.to_string(), c.name.span));
-
+                    // Only add constants of non-abstract types to the named expressions
+                    // to prevent abstract types ending up in the IR.
+                    let is_abstract = ctx.module.types[ty].inner.is_abstract(&ctx.module.types);
+                    if !is_abstract {
+                        ctx.named_expressions
+                            .insert(init, (c.name.name.to_string(), c.name.span));
+                    }
                     return Ok(());
                 }
             },
diff --git a/naga/tests/in/wgsl/abstract-types-return.wgsl b/naga/tests/in/wgsl/abstract-types-return.wgsl
@@ -29,3 +29,8 @@ const one = 1;
 fn return_const_f32_const_ai() -> f32 {
     return one;
 }
+
+fn return_vec2f32_const_ai() -> vec2<f32> {
+    const vec_one = vec2(1);
+    return vec_one;
+}
diff --git a/naga/tests/in/wgsl/const_assert.wgsl b/naga/tests/in/wgsl/const_assert.wgsl
@@ -13,4 +13,8 @@ fn foo() {
   const z = x + y - 2;
   const_assert z > 0; // valid in functions.
   const_assert(z > 0);
+
+  const_assert z == 1i;
+  const_assert z > 0u;
+  const_assert z < 2.0f;
 }
diff --git a/naga/tests/out/glsl/abstract-types-return.main.Compute.glsl b/naga/tests/out/glsl/abstract-types-return.main.Compute.glsl
@@ -34,6 +34,10 @@ float return_const_f32_const_ai() {
     return 1.0;
 }
 
+vec2 return_vec2f32_const_ai() {
+    return vec2(1.0);
+}
+
 void main() {
     return;
 }
diff --git a/naga/tests/out/hlsl/abstract-types-return.hlsl b/naga/tests/out/hlsl/abstract-types-return.hlsl
@@ -40,6 +40,11 @@ float return_const_f32_const_ai()
     return 1.0;
 }
 
+float2 return_vec2f32_const_ai()
+{
+    return (1.0).xx;
+}
+
 [numthreads(1, 1, 1)]
 void main()
 {
diff --git a/naga/tests/out/ir/const_assert.compact.ron b/naga/tests/out/ir/const_assert.compact.ron
@@ -16,12 +16,8 @@
             arguments: [],
             result: None,
             local_variables: [],
-            expressions: [
-                Literal(I32(1)),
-            ],
-            named_expressions: {
-                0: "z",
-            },
+            expressions: [],
+            named_expressions: {},
             body: [
                 Return(
                     value: None,
diff --git a/naga/tests/out/ir/const_assert.ron b/naga/tests/out/ir/const_assert.ron
@@ -16,12 +16,8 @@
             arguments: [],
             result: None,
             local_variables: [],
-            expressions: [
-                Literal(I32(1)),
-            ],
-            named_expressions: {
-                0: "z",
-            },
+            expressions: [],
+            named_expressions: {},
             body: [
                 Return(
                     value: None,
diff --git a/naga/tests/out/ir/local-const.compact.ron b/naga/tests/out/ir/local-const.compact.ron
@@ -21,16 +21,6 @@
                 width: 4,
             )),
         ),
-        (
-            name: None,
-            inner: Vector(
-                size: Tri,
-                scalar: (
-                    kind: Sint,
-                    width: 4,
-                ),
-            ),
-        ),
     ],
     special_types: (
         ray_desc: None,
@@ -69,63 +59,26 @@
             result: None,
             local_variables: [],
             expressions: [
-                Literal(I32(4)),
                 Literal(I32(4)),
                 Literal(U32(4)),
                 Literal(F32(4.0)),
-                Compose(
-                    ty: 3,
-                    components: [
-                        0,
-                        0,
-                        0,
-                    ],
-                ),
-                Literal(F32(2.0)),
-                Literal(I32(4)),
                 Constant(0),
                 Constant(1),
                 Constant(2),
-                Literal(I32(4)),
-                Literal(I32(4)),
-                Literal(I32(4)),
-                Compose(
-                    ty: 3,
-                    components: [
-                        10,
-                        11,
-                        12,
-                    ],
-                ),
-                Literal(F32(2.0)),
             ],
             named_expressions: {
-                0: "a",
-                1: "b",
-                2: "c",
-                3: "d",
-                4: "e",
-                5: "f",
-                6: "ag",
-                7: "bg",
-                8: "cg",
-                9: "dg",
-                13: "eg",
-                14: "fg",
+                0: "b",
+                1: "c",
+                2: "d",
+                3: "bg",
+                4: "cg",
+                5: "dg",
             },
             body: [
-                Emit((
-                    start: 4,
-                    end: 5,
-                )),
                 Emit((
                     start: 0,
                     end: 0,
                 )),
-                Emit((
-                    start: 13,
-                    end: 14,
-                )),
                 Return(
                     value: None,
                 ),
diff --git a/naga/tests/out/ir/local-const.ron b/naga/tests/out/ir/local-const.ron
@@ -21,16 +21,6 @@
                 width: 4,
             )),
         ),
-        (
-            name: None,
-            inner: Vector(
-                size: Tri,
-                scalar: (
-                    kind: Sint,
-                    width: 4,
-                ),
-            ),
-        ),
     ],
     special_types: (
         ray_desc: None,
@@ -69,63 +59,26 @@
             result: None,
             local_variables: [],
             expressions: [
-                Literal(I32(4)),
                 Literal(I32(4)),
                 Literal(U32(4)),
                 Literal(F32(4.0)),
-                Compose(
-                    ty: 3,
-                    components: [
-                        0,
-                        0,
-                        0,
-                    ],
-                ),
-                Literal(F32(2.0)),
-                Literal(I32(4)),
                 Constant(0),
                 Constant(1),
                 Constant(2),
-                Literal(I32(4)),
-                Literal(I32(4)),
-                Literal(I32(4)),
-                Compose(
-                    ty: 3,
-                    components: [
-                        10,
-                        11,
-                        12,
-                    ],
-                ),
-                Literal(F32(2.0)),
             ],
             named_expressions: {
-                0: "a",
-                1: "b",
-                2: "c",
-                3: "d",
-                4: "e",
-                5: "f",
-                6: "ag",
-                7: "bg",
-                8: "cg",
-                9: "dg",
-                13: "eg",
-                14: "fg",
+                0: "b",
+                1: "c",
+                2: "d",
+                3: "bg",
+                4: "cg",
+                5: "dg",
             },
             body: [
-                Emit((
-                    start: 4,
-                    end: 5,
-                )),
                 Emit((
                     start: 0,
                     end: 0,
                 )),
-                Emit((
-                    start: 13,
-                    end: 14,
-                )),
                 Return(
                     value: None,
                 ),
diff --git a/naga/tests/out/msl/abstract-types-return.msl b/naga/tests/out/msl/abstract-types-return.msl
@@ -43,6 +43,11 @@ float return_const_f32_const_ai(
     return 1.0;
 }
 
+metal::float2 return_vec2f32_const_ai(
+) {
+    return metal::float2(1.0);
+}
+
 kernel void main_(
 ) {
     return;
diff --git a/naga/tests/out/spv/abstract-types-return.spvasm b/naga/tests/out/spv/abstract-types-return.spvasm
@@ -1,12 +1,12 @@
 ; SPIR-V
 ; Version: 1.1
 ; Generator: rspirv
-; Bound: 44
+; Bound: 47
 OpCapability Shader
 %1 = OpExtInstImport "GLSL.std.450"
 OpMemoryModel Logical GLSL450
-OpEntryPoint GLCompute %41 "main"
-OpExecutionMode %41 LocalSize 1 1 1
+OpEntryPoint GLCompute %44 "main"
+OpExecutionMode %44 LocalSize 1 1 1
 OpDecorate %7 ArrayStride 4
 %2 = OpTypeVoid
 %3 = OpTypeInt 32 1
@@ -25,7 +25,7 @@ OpDecorate %7 ArrayStride 4
 %30 = OpConstantComposite  %6  %22 %22
 %34 = OpTypeFunction %7
 %35 = OpConstantComposite  %7  %22 %22 %22 %22
-%42 = OpTypeFunction %2
+%45 = OpTypeFunction %2
 %10 = OpFunction  %3  None %11
 %9 = OpLabel
 OpBranch %13
@@ -68,9 +68,15 @@ OpBranch %39
 %39 = OpLabel
 OpReturnValue %22
 OpFunctionEnd
-%41 = OpFunction  %2  None %42
+%41 = OpFunction  %6  None %29
 %40 = OpLabel
-OpBranch %43
+OpBranch %42
+%42 = OpLabel
+OpReturnValue %30
+OpFunctionEnd
+%44 = OpFunction  %2  None %45
 %43 = OpLabel
+OpBranch %46
+%46 = OpLabel
 OpReturn
 OpFunctionEnd
diff --git a/naga/tests/out/wgsl/abstract-types-return.wgsl b/naga/tests/out/wgsl/abstract-types-return.wgsl
@@ -26,6 +26,10 @@ fn return_const_f32_const_ai() -> f32 {
     return 1f;
 }
 
+fn return_vec2f32_const_ai() -> vec2<f32> {
+    return vec2(1f);
+}
+
 @compute @workgroup_size(1, 1, 1) 
 fn main() {
     return;
diff --git a/naga/tests/out/wgsl/local-const.wgsl b/naga/tests/out/wgsl/local-const.wgsl
@@ -3,8 +3,6 @@ const gc: u32 = 4u;
 const gd: f32 = 4f;
 
 fn const_in_fn() {
-    const e = vec3<i32>(4i, 4i, 4i);
-    const eg = vec3<i32>(4i, 4i, 4i);
     return;
 }
 

Original file line number	Diff line number	Diff line change
`@@ -29,3 +29,8 @@ const one = 1;`
`29`	`29`	`fn return_const_f32_const_ai() -> f32 {`
`30`	`30`	`return one;`
`31`	`31`	`}`
	`32`	`+`
	`33`	`+fn return_vec2f32_const_ai() -> vec2<f32> {`
	`34`	`+ const vec_one = vec2(1);`
	`35`	`+ return vec_one;`
	`36`	`+}`
Original file line number	Diff line number	Diff line change
`@@ -13,4 +13,8 @@ fn foo() {`
`13`	`13`	`const z = x + y - 2;`
`14`	`14`	`const_assert z > 0; // valid in functions.`
`15`	`15`	`const_assert(z > 0);`
	`16`	`+`
	`17`	`+ const_assert z == 1i;`
	`18`	`+ const_assert z > 0u;`
	`19`	`+ const_assert z < 2.0f;`
`16`	`20`	`}`
Original file line number	Diff line number	Diff line change
`@@ -34,6 +34,10 @@ float return_const_f32_const_ai() {`
`34`	`34`	`return 1.0;`
`35`	`35`	`}`
`36`	`36`
	`37`	`+vec2 return_vec2f32_const_ai() {`
	`38`	`+ return vec2(1.0);`
	`39`	`+}`
	`40`	`+`
`37`	`41`	`void main() {`
`38`	`42`	`return;`
`39`	`43`	`}`
Original file line number	Diff line number	Diff line change
`@@ -40,6 +40,11 @@ float return_const_f32_const_ai()`
`40`	`40`	`return 1.0;`
`41`	`41`	`}`
`42`	`42`
	`43`	`+float2 return_vec2f32_const_ai()`
	`44`	`+{`
	`45`	`+ return (1.0).xx;`
	`46`	`+}`
	`47`	`+`
`43`	`48`	`[numthreads(1, 1, 1)]`
`44`	`49`	`void main()`
`45`	`50`	`{`