Repository: pygfx/pygfx
Pull request #1053 (2025-03-24 21:20:51)

Refactor engine's resolving of fragment shader output 
**a.ka. get rid of `get_fragment_output`**

This is part of the work in #1002, but it seemed orthogonal enough to roll into a separate PR, which makes review easier.

### In short

Earlier, (custom) shaders ended like this:
```wgsl
@fragment
fn fs_main(varyings: Varyings) -> FragmentOutput {
   ...
    var out = get_fragment_output(varyings.positions, the_final_color);
    out.picking  = ..  # maybe
    out.depth = ... # some shaders
    return out;
}
```

The `get_fragment_output()` provided a hook, allowing the blenders to provide an implementation to "do their thing".

In the new situation of this pr, the hook is implemented differently. A bit more complex, but leading to simpler shaders for users:
```wgsl
@fragment
fn fs_main(varyings: Varyings) -> FragmentOutput {
   ...
    var out = FragmentOutput;
    out.color = the_final_color;
    out.picking  = ..  # maybe
    out.depth = ... # some shaders
    return out;
}
```

### Motivation

* The first arg in `get_fragment_output()` feels arbitrary. It's only needed by the dither blending and weighted with depth. It also does not 'scale'.
* For classic blending, the code is now as one would expect, and the shader resolving indeed will not change it (after #1002).
* For special blending (mainly dither and weighted), the code also looks natural.
* It will allow users to more easily tweak the special blending, e.g. changing the seed for the dither, or the weight for weighted blending, using simply `out.weight = ...`.
* It will make it easier / more natural to support custom/additional render targets in the future.

### Compatibility

To help users who write custom shaders transition, shaders can still be written using the old approach. The wgsl is converted to the 'new approach' as part of the shader resolving. We can remove that in a later version.

### Performance

I also tweaked the performance of the shader resolving step. I made changes that made it a lot faster, but also some more checks (e.g. the resolving now correctly knows whether it's in a vertex/fragment entrypoint). In total the resolving is still more than twice as fast. Don't expect to 'feel' the difference, I don't think this was a bottleneck, but all bits help.

### Changes in a list

* [x] Refactor `resolve.py` for performance.
* [x] Refactor output resolving to do more work, to allow simpler shader.
* [x] Added some trickery to support the old usage. 
* [x] Added many tests. 
* [x] Updated blender.py accordingly ... which will be completely overhauled in #1002. 
* [x] Docs.
* [x] Update our shaders and examples that use custom shaders.



almarklein: This was harder than I thought when I started it. And all it does is change the API for writing custom shaders a bit. I still think its worth it though, especially if we formalize the API for custom-shaders more in the furture, since its cleaner/simpler, and provides an easy way to turn the knobs of e.g. weighted blending for specific use-cases.
almarklein: Another way to put it: that `get_fragment_output` was a nice trick that provides a very easy shortcut, but now let's do it properly.

Patch:
diff --git a/docs/advanced_shaders.rst b/docs/advanced_shaders.rst
index 1e90806..6cdf01c 100644
--- a/docs/advanced_shaders.rst
+++ b/docs/advanced_shaders.rst
@@ -209,9 +209,9 @@ FragmentOutput
 In a somewhat similar way, the output of the fragment shader is
 predefined. Though in this case the output is determined by the blend
 mode and render pass (opaque or transparent), and the details are hidden
-from the shader author. This way, Pygfx can support advanced handling
-of transparency without affecting individual shaders.
-All fragment functions in Pygfx are somewhat like this:
+from the shader author. This way, Pygfx can support special blend modes
+without affecting individual shaders.
+All fragment functions in Pygfx look somewhat like this:
 
 
 .. code-block:: python
@@ -223,11 +223,39 @@ All fragment functions in Pygfx are somewhat like this:
             @stage(fragment)
             fn fs_main(varyings: Varyings) -> FragmentOutput {
                 ...
-                var out = get_fragment_output(varyings.position, color);
+                var out: FragmentOutput;
+                out.color = vec4<f32>(...);
                 return out;
             }
             """
 
+For some types of blending the output struct is modified automatically,
+and users can influence this process. E.g. to explicitly set a seed for
+dithered blending:
+
+.. code-block::
+
+    ...
+    var out: FragmentOutput;
+    out.color = vec4<f32>(...);
+    $$ if blending == 'dither'
+    out.seed1 = f32(...);
+    $$ endif
+    return out;
+
+... or set the weight for weighted blending:
+
+.. code-block::
+
+    ...
+    var out: FragmentOutput;
+    out.color = vec4<f32>(...);
+    $$ if blending == 'weighted'
+    out.weight = f32(...);
+    $$ endif
+    return out;
+
+
 Picking
 -------
 
@@ -257,7 +285,8 @@ to unpack the picking info. See e.g. the picking of a mesh:
             @stage(fragment)
             fn fs_main(varyings: Varyings) -> FragmentOutput {
                 ...
-                var out = get_fragment_output(varyings.position, color);
+                var out: FragmentOutput;
+                out.color = color;
 
                 // The builtin write_pick templating variable should be used
                 // to ensure picking info is only written in the appropriate render pass
@@ -297,8 +326,8 @@ you can include the wgsl code for clipping planes in your shader like this:
                 // clipping planes
                 {$ include 'pygfx.clipping_planes.wgsl' $}
 
-                var out = get_fragment_output(varyings.position, color);
-                ...
+                var out: FragmentOutput;
+                out.color = color;
                 return out;
             }
             """
diff --git a/pygfx/renderers/wgpu/engine/blender.py b/pygfx/renderers/wgpu/engine/blender.py
index 5959706..361ecaa 100644
--- a/pygfx/renderers/wgpu/engine/blender.py
+++ b/pygfx/renderers/wgpu/engine/blender.py
@@ -13,7 +13,7 @@ from .renderer import WgpuRenderer
 
 # Notes:
 # - The user code provides color as-is in rgba.
-# - In the get_fragment_output logic defined in the shaders here, the color
+# - In the apply_virtual_fields_of_fragment_output logic defined in the shaders here, the color
 #   is pre-multiplied with the alpha.
 # - All fixed-pipeling blending options assume premultiplied alpha.
 
@@ -71,9 +71,9 @@ class BasePass:
 
         Notes:
 
-        * This code gets injected into the shader, so the material shaders
-          can use get_fragment_output.
-        * This code should define FragmentOutput and get_fragment_output.
+        * This code gets injected into the shader, defining FragmentOutput.
+        * This code may define virtual fields in FragmentOutput using a special comment.
+        * If it does, it should define apply_virtual_fields_of_fragment_output(outp: ptr<function,FragmentOutput>, *virtual_args)
         """
         return ""
 
@@ -153,14 +153,13 @@ class OpaquePass(BasePass):
     def get_shader_code(self, blender):
         return """
         struct FragmentOutput {
+            // virtualfield color : vec4<f32> = vec4<f32>(0.0)
             @location(0) color: vec4<f32>,
             @location(1) pick: vec4<u32>,
         };
-        fn get_fragment_output(position: vec4<f32>, color: vec4<f32>) -> FragmentOutput {
+        fn apply_virtual_fields_of_fragment_output(outp: ptr<function,FragmentOutput>, color: vec4<f32>) {
             if (color.a < 1.0 - ALPHA_COMPARE_EPSILON ) { discard; }
-            var out : FragmentOutput;
-            out.color = vec4<f32>(color.rgb, 1.0);
-            return out;
+            (*outp).color = vec4<f32>(color.rgb, 1.0);
         }
         """
 
@@ -174,13 +173,12 @@ class FullOpaquePass(OpaquePass):
     def get_shader_code(self, blender):
         return """
         struct FragmentOutput {
+            // virtualfield color : vec4<f32> = vec4<f32>(0.0)
             @location(0) color: vec4<f32>,
             @location(1) pick: vec4<u32>,
         };
-       fn get_fragment_output(position: vec4<f32>, color: vec4<f32>) -> FragmentOutput {
-            var out : FragmentOutput;
-            out.color = vec4<f32>(color.rgb, 1.0);  // make every fragment opaque
-            return out;
+        fn apply_virtual_fields_of_fragment_output(outp: ptr<function,FragmentOutput>, color: vec4<f32>) {
+            (*outp).color = vec4<f32>(color.rgb, 1.0);
         }
         """
 
@@ -193,30 +191,28 @@ class DitherPass(OpaquePass):
 
     def get_shader_code(self, blender):
         return """
-
+        // We want the seed for the randon function to be such that the result is
+        // deterministic, so that rendered images can be visually compared. This
+        // is why the object-id should not be used. Using the xy ndc coord is a
+        // no-brainer seed. Using only these will give an effect often observed
+        // in games, where the pattern is "stuck to the screen". We also seed with
+        // the depth, since this covers *a lot* of cases, e.g. different objects
+        // behind each-other, as well as the same object having different parts
+        // at the same screen pixel. This only does not cover cases where objects
+        // are exactly on top of each other. Therefore we use rgba as another seed.
+        // So the only case where the same pattern may be used for different
+        // fragments if an object is at the same depth and has the same color.
         struct FragmentOutput {
+            // virtualfield seed1 : f32 = varyings.position.x * varyings.position.y * varyings.position.z
+            // virtualfield seed2 : f32 = out.color.r * 0.12 + out.color.g * 0.34 + out.color.b * 0.56 + out.color.a * 0.78
             @location(0) color: vec4<f32>,
             @location(1) pick: vec4<u32>,
         };
-        fn get_fragment_output(position: vec4<f32>, color: vec4<f32>) -> FragmentOutput {
-            // We want the seed for the ramdon function to be such that the result is
-            // deterministic, so that rendered images can be visually compared. This
-            // is why the object-id should not be used. Using the xy ndc coord is a
-            // no-brainer seed. Using only these will give an effect often observed
-            // in games, where the pattern is "stuck to the screen". We also seed with
-            // the depth, since this covers *a lot* of cases, e.g. different objects
-            // behind each-other, as well as the same object having different parts
-            // at the same screen pixel. This only does not cover cases where objects
-            // are exactly on top of each other. Therefore we use rgba as another seed.
-            // So the only case where the same pattern may be used for different
-            // fragments if an object is at the same depth and has the same color.
-            var out : FragmentOutput;
-            let seed1 = position.x * position.y * position.z;
-            let seed2 = color.r * 0.12 + color.g * 0.34 + color.b * 0.56 + color.a * 0.78;
+        fn apply_virtual_fields_of_fragment_output(outp: ptr<function,FragmentOutput>, seed1:f32, seed2:f32) {
             let rand = random2(vec2<f32>(seed1, seed2));
-            if ( color.a < 1.0 - ALPHA_COMPARE_EPSILON && color.a < rand ) { discard; }
-            out.color = vec4<f32>(color.rgb, 1.0);  // fragments that pass through are opaque
-            return out;
+            let alpha = (*outp).color.a;
+            if ( alpha < 1.0 - ALPHA_COMPARE_EPSILON && alpha < rand ) { discard; }
+            (*outp).color.a = 1.0;  // fragments that pass through are opaque
         }
         """
 
@@ -249,13 +245,12 @@ class SimpleSinglePass(OpaquePass):
         # Take depth into account, but don't treat transparent fragments differently
         return """
         struct FragmentOutput {
+            // virtualfield color : vec4<f32> = vec4<f32>(0.0)
             @location(0) color: vec4<f32>,
             @location(1) pick: vec4<u32>,
         };
-        fn get_fragment_output(position: vec4<f32>, color: vec4<f32>) -> FragmentOutput {
-            var out : FragmentOutput;
-            out.color = vec4<f32>(color.rgb * color.a, color.a);
-            return out;
+        fn apply_virtual_fields_of_fragment_output(outp: ptr<function,FragmentOutput>, color: vec4<f32>) {
+            (*outp).color = vec4<f32>(color.rgb * color.a, color.a);
         }
         """
 
@@ -319,13 +314,12 @@ class SimpleTransparencyPass(BasePass):
     def get_shader_code(self, blender):
         return """
         struct FragmentOutput {
+            // virtualfield color : vec4<f32> = vec4<f32>(0.0)
             @location(0) color: vec4<f32>,
         };
-        fn get_fragment_output(position: vec4<f32>, color: vec4<f32>) -> FragmentOutput {
+        fn apply_virtual_fields_of_fragment_output(outp: ptr<function,FragmentOutput>, color: vec4<f32>) {
             if (color.a <= ALPHA_COMPARE_EPSILON) { discard; }
-            var out : FragmentOutput;
-            out.color = vec4<f32>(color.rgb * color.a, color.a);
-            return out;
+            (*outp).color = vec4<f32>(color.rgb * color.a, color.a);
         }
         """
 
@@ -420,19 +414,18 @@ class WeightedTransparencyPass(BasePass):
     def get_shader_code(self, blender):
         code = """
         struct FragmentOutput {
+            // virtualfield color : vec4<f32> = vec4<f32>(0.0)
+            // virtualfield depth_for_weight: f32 = varyings.position.z;
             @location(0) accum: vec4<f32>,
             @location(1) reveal: f32,
         };
-        fn get_fragment_output(position: vec4<f32>, color: vec4<f32>) -> FragmentOutput {
-            let depth = position.z;
+        fn apply_virtual_fields_of_fragment_output(outp: ptr<function,FragmentOutput>, color: vec4<f32>, depth: f32) {
             let alpha = color.a;
             if (alpha <= ALPHA_COMPARE_EPSILON) { discard; }
             let premultiplied = color.rgb * alpha;
             WEIGHT_CODE
-            var out : FragmentOutput;
-            out.accum = vec4<f32>(premultiplied, alpha) * weight;
-            out.reveal = alpha;  // yes, alpha, not weight
-            return out;
+            (*outp).accum = vec4<f32>(premultiplied, alpha) * weight;
+            (*outp).reveal = alpha;  // yes, alpha, not weight
             // Note 1: could also take user-specified transmittance into account.
             // Note 2: its also possible to undo a fragment contribution. For this the accum
             // and reveal buffer must be float to avoid clamping. And we'd do `abs(color.a)` above.
@@ -516,14 +509,13 @@ class FrontmostTransparencyPass(BasePass):
     def get_shader_code(self, blender):
         return """
         struct FragmentOutput {
+            // virtualfield color : vec4<f32> = vec4<f32>(0.0)
             @location(0) color: vec4<f32>,
             @location(1) pick: vec4<u32>,
         };
-        fn get_fragment_output(position: vec4<f32>, color: vec4<f32>) -> FragmentOutput {
+        fn apply_virtual_fields_of_fragment_output(outp: ptr<function,FragmentOutput>, color: vec4<f32>) {
             if (color.a <= ALPHA_COMPARE_EPSILON || color.a >= 1.0 - ALPHA_COMPARE_EPSILON) { discard; }
-            var out : FragmentOutput;
-            out.color = vec4<f32>(color.rgb * color.a, color.a);
-            return out;
+            (*outp).color = vec4<f32>(color.rgb * color.a, color.a);
         }
         """
 
@@ -646,6 +638,7 @@ class BaseFragmentBlender:
         return {
             "blending_code": self.passes[pass_index].get_shader_code(self),
             "write_pick": self.passes[pass_index].write_pick,
+            "blending": self.name,
         }
 
     def get_pass_count(self):
@@ -1044,12 +1037,6 @@ class AdditivePass(BasePass):
         struct FragmentOutput {
             @location(0) color: vec4<f32>,
         };
-        fn get_fragment_output(position: vec4<f32>, color: vec4<f32>) -> FragmentOutput {
-            let depth = position.z;
-            var out : FragmentOutput;
-            out.color = color;
-            return out;
-        }
         """
 
 
diff --git a/pygfx/renderers/wgpu/shader/base.py b/pygfx/renderers/wgpu/shader/base.py
index eb77221..1cf027e 100644
--- a/pygfx/renderers/wgpu/shader/base.py
+++ b/pygfx/renderers/wgpu/shader/base.py
@@ -13,7 +13,7 @@ from ..engine.utils import (
     hash_from_value,
 )
 from ..engine.binding import Binding
-from .resolve import resolve_varyings, resolve_depth_output
+from .resolve import resolve_shadercode
 from .bindings import BindingDefinitions
 from .templating import apply_templating
 
@@ -218,9 +218,8 @@ class BaseShader(ShaderInterface):
             code1 = self.get_code()
             # Templating and resolving includes
             code2 = apply_templating(code1, **template_vars)
-            # Auto-insert varyings struct, and Tweak FragmentOutput if depth is set in frag shader.
-            code3 = resolve_varyings(code2)
-            code3 = resolve_depth_output(code3)
+            # Resolve Varyings and FragmentOutput
+            code3 = resolve_shadercode(code2)
             return code3
         finally:
             self._template_vars_bindings = ori_template_vars
diff --git a/pygfx/renderers/wgpu/shader/resolve.py b/pygfx/renderers/wgpu/shader/resolve.py
index 54ac70f..dc1cd09 100644
--- a/pygfx/renderers/wgpu/shader/resolve.py
+++ b/pygfx/renderers/wgpu/shader/resolve.py
@@ -1,4 +1,72 @@
 import re
+import logging
+
+
+logger = logging.getLogger("pygfx")
+warned_for = set()
+
+
+def resolve_shadercode(wgsl):
+    """Apply all shader resolve opearions."""
+    return _resolve(wgsl, VaryingResolver(), OutputResolver())
+
+
+def resolve_varyings(wgsl):
+    return _resolve(wgsl, VaryingResolver())
+
+
+def resolve_output(wgsl):
+    return _resolve(wgsl, OutputResolver())
+
+
+def _resolve(wgsl, *resolvers):
+    assert isinstance(wgsl, str)
+
+    # Split into lines, which is easier to process. Ensure it ends with newline in the end.
+    lines = wgsl.splitlines()
+    while lines and not lines[0]:
+        lines.pop(0)
+    if not (lines and lines[-1] == ""):
+        lines.append("")
+
+    curly_braces_level = 0
+    current_func_kind = None
+
+    # Apply resolve functions. They modify the lines
+    for linenr, ori_line in enumerate(lines):
+        line = ori_line.split("//", 1)[0].strip()
+        if not line:
+            continue
+
+        if line.startswith("fn "):
+            prevline = lines[linenr - 1].strip() if linenr > 0 else ""
+            current_func_kind = "normal"
+            if prevline.startswith("@"):
+                if prevline.startswith("@vertex"):
+                    current_func_kind = "vertex"
+                elif prevline.startswith("@fragment"):
+                    current_func_kind = "fragment"
+                elif prevline.startswith("@compute"):
+                    current_func_kind = "compute"
+            curly_braces_level = 0
+
+        if current_func_kind:
+            curly_braces_level += line.count("{") - line.count("}")
+            if curly_braces_level <= 0:
+                current_func_kind = None
+
+        for resolver in resolvers:
+            resolver.process_line(linenr, line, current_func_kind)
+
+    for resolver in resolvers:
+        resolver.update_lines(lines)
+
+    # Return as str
+    return "\n".join(lines)
+
+
+def indent_from_line(line):
+    return line[: len(line) - len(line.lstrip())]
 
 
 varying_types = ["f32", "vec2<f32>", "vec3<f32>", "vec4<f32>"]
@@ -8,198 +76,300 @@ varying_types = (
     + [t.replace("f", "u") for t in varying_types]
 )
 
-re_varying_getter = re.compile(r"[\s,\(\[]varyings\.(\w+)", re.UNICODE)
-re_varying_setter = re.compile(r"\A\s*?varyings\.(\w+)(\.\w+)?\s*?\=")
+re_varying_getter = re.compile(r"[^\w]varyings\.(\w+)", re.UNICODE)
+re_varying_setter = re.compile(r"varyings\.(\w+)(\.\w+)?\s*\=", re.UNICODE)
 builtin_varyings = {"position": "vec4<f32>"}
 
 
-def resolve_varyings(wgsl):
+class VaryingResolver:
     """Resolve varyings in the given wgsl:
     * Detect varyings being used.
     * Check that these are also set.
     * Remove assignments of varyings that are not used.
     * Include the Varyings struct.
     """
-    assert isinstance(wgsl, str)
 
-    # Split into lines, which is easier to process. Ensure it ends with newline in the end.
-    lines = wgsl.splitlines()
-    while lines and not lines[0].strip():
-        lines.pop(0)
-    while lines and not lines[-1].strip():
-        lines.pop(-1)
-    lines.append("")
-
-    # Prepare dicts that map name to list-of-linenr. And a tupe dict.
-    assigned_varyings = {}
-    used_varyings = {}
-    types = {}  # varying types
-
-    # We try to find the function that first uses the Varyings struct.
-    struct_insert_pos = None
-
-    # Go over all lines to:
-    # - find the lines where a varying is set
-    # - collect the types of these varyings
-    for linenr, line in enumerate(lines):
-        match = re_varying_setter.match(line)
-        if match:
-            # Get parts
-            name = match.group(1)
-            attr = match.group(2)
-            # Handle builtin
-            if name in builtin_varyings:
-                used_varyings[name] = []
-                types[name] = builtin_varyings[name]
-            # Find type
-            type = line[match.end() :].split("(")[0].strip().replace(" ", "")
-            if type not in varying_types:
-                type = ""
-            # Triage
-            if attr:
-                pass  # Not actually a type but an attribute access
-            elif not type:
-                raise TypeError(
-                    f"Varying {name!r} assignment needs an explicit cast (of a correct type), e.g. `varying.{name} = f32(3.0);`:\n{line}"
-                )
-            elif name in types and type != types[name]:
-                raise TypeError(
-                    f"Varying {name!r} assignment does not match expected type {types[name]}:\n{line}"
-                )
-            else:
-                types[name] = type
-            # Store position
-            assigned_varyings.setdefault(name, []).append(linenr)
-
-    # Go over all lines to:
-    # - collect all used varyings
-    # - find where the vertex-shader starts
-    in_vertex_shader = False
-    current_func_linenr = 0
-    for linenr, line in enumerate(lines):
-        line = line.strip()
-        # Detect when we enter a new function
-        if line.startswith("fn "):
-            current_func_linenr = linenr
-            if line.startswith("fn vs_main"):
-                in_vertex_shader = True
-            else:
-                in_vertex_shader = False
-        # Remove comments (shader code has no strings that can contain slashes)
-        line = line.split("//")[0]
-        if "Varyings" in line and struct_insert_pos is None:
-            struct_insert_pos = current_func_linenr
-        # Everything we find here is a match (prepend a space to allow an easier regexp)
-        for match in re_varying_getter.finditer(" " + line):
-            name = match.group(1)
-            this_varying_is_set_on_this_line = linenr in assigned_varyings.get(name, [])
-            if this_varying_is_set_on_this_line:
-                pass
-            elif in_vertex_shader:
-                # If varyings are used in another way than setting, in the vertex shader,
-                # we should either consider them "used", or possibly break the shader if
-                # the used varying is disabled. So let's just not allow it.
-                raise TypeError(
-                    f"Varying {name!r} is read in the vertex shader, but only writing is allowed:\n{line}"
+    def __init__(self):
+        # Prepare dicts that map name to list-of-linenr. And a tupe dict.
+        self.assigned_varyings = {}
+        self.used_varyings = {}
+        self.types = {}  # varying types
+
+        # We try to find the function that first uses the Varyings struct.
+        self.struct_linenr = None
+
+        # Go over all lines to:
+        # - find the lines where a varying is set
+        # - collect the types of these varyings
+        # - collect all used varyings
+        # - find where the vertex-shader starts
+
+    def process_line(self, linenr, line, current_func_kind):
+        # Look for varyings being set (in the vertex shader)
+        # Note that this implicitly matches from the start of the stripped line.
+        in_vertex_shader = current_func_kind == "vertex"
+
+        # Put the Varyings struct right above where its first used
+        if self.struct_linenr is None:
+            if "Varyings" in line:
+                self.struct_linenr = linenr if current_func_kind else 0
+
+        # Look for varyings being set (in vertex shader)
+        if in_vertex_shader:
+            match = re_varying_setter.match(line)
+            if match:
+                # Get parts
+                name = match.group(1)
+                attr = match.group(2)
+                # Handle builtin
+                if name in builtin_varyings:
+                    self.used_varyings[name] = []
+                    self.types[name] = builtin_varyings[name]
+                # Find type
+                type = line[match.end() :].split("(")[0].strip().replace(" ", "")
+                if type not in varying_types:
+                    type = ""
+                # Triage
+                if attr:
+                    pass  # Not actually a type but an attribute access
+                elif not type:
+                    raise TypeError(
+                        f"Varying {name!r} assignment needs an explicit cast (of a correct type), e.g. `varying.{name} = f32(3.0);`:\n{line}"
+                    )
+                elif name in self.types and type != self.types[name]:
+                    raise TypeError(
+                        f"Varying {name!r} assignment does not match expected type {self.types[name]}:\n{line}"
+                    )
+                else:
+                    self.types[name] = type
+                # Store position
+                self.assigned_varyings.setdefault(name, []).append(linenr)
+
+        # Look for varyings being get (in any function)
+        if current_func_kind and "varyings." in line:
+            for match in re_varying_getter.finditer(" " + line):
+                name = match.group(1)
+                this_varying_is_set_on_this_line = linenr in self.assigned_varyings.get(
+                    name, []
                 )
-            else:
-                used_varyings.setdefault(name, []).append(linenr)
-
-    # Check if all used varyings are assigned
-    for name in used_varyings:
-        if name not in assigned_varyings:
-            line = lines[used_varyings[name][0]]
-            raise TypeError(f"Varying {name!r} is read, but not assigned:\n{line}")
-
-    # Comment-out the varying setter if its unused elsewhere in the shader
-    for name, linenrs in assigned_varyings.items():
-        if name not in used_varyings:
-            for linenr in linenrs:
-                line = lines[linenr]
-                indent = line[: len(line) - len(line.lstrip())]
-                lines[linenr] = indent + "// unused: " + line[len(indent) :]
-                # Deal with multiple-line assignments
-                line_s = line.strip()
-                while not line_s.endswith(";"):
-                    linenr += 1
-                    line_s = lines[linenr].strip()
-                    unexpected = "fn ", "struct ", "var ", "let ", "}"
-                    if line_s.startswith(unexpected) or linenr == len(lines) - 1:
-                        raise TypeError(
-                            f"Varying {name!r} assignment seems to be missing a semicolon:\n{line}"
-                        )
-                    lines[linenr] = indent + "// " + line_s
-
-    # Build and insert the struct
-    if struct_insert_pos is not None:
-        # Maybe we should move up a bit
-        if struct_insert_pos > 0:
-            if lines[struct_insert_pos - 1].lstrip().startswith("@"):
-                struct_insert_pos -= 1
-        # First divide into slot-based and builtins
-        used_varyings = set(used_varyings)
-        used_builtins = used_varyings.intersection(builtin_varyings)
-        used_slots = used_varyings.difference(used_builtins)
-        used_slots = list(sorted(used_slots))
-        # Build struct
-        struct_lines = ["struct Varyings {"]
-        for slotnr, name in enumerate(used_slots):
-            struct_lines.append(f"    @location({slotnr}) {name} : {types[name]},")
-        for name in sorted(used_builtins):
-            struct_lines.append(f"    @builtin({name}) {name} : {types[name]},")
-        struct_lines.append("};\n")
-        # Apply indentation and insert
-        line = lines[struct_insert_pos]
-        indent = line[: len(line) - len(line.lstrip())]
-        struct_lines = [indent + line for line in struct_lines]
-        lines.insert(struct_insert_pos, "\n".join(struct_lines))
-    else:
-        assert not used_varyings, "woops, did not expect used_varyings here"
-
-    # Return modified code
-    return "\n".join(lines)
+                if this_varying_is_set_on_this_line:
+                    pass
+                elif in_vertex_shader:
+                    # If varyings are used in another way than setting, in the vertex shader,
+                    # we should either consider them "used", or possibly break the shader if
+                    # the used varying is disabled. So let's just not allow it.
+                    raise TypeError(
+                        f"Varying {name!r} is read in the vertex shader, but only writing is allowed:\n{line}"
+                    )
+                else:
+                    self.used_varyings.setdefault(name, []).append(linenr)
 
+    def update_lines(self, lines):
+        assigned_varyings = self.assigned_varyings
+        used_varyings = self.used_varyings
+        types = self.types
+        struct_linenr = self.struct_linenr
 
-re_depth_setter = re.compile(r"\A\s*?out\.depth\s*?\=")
+        # Check if all used varyings are assigned
+        for name in used_varyings:
+            if name not in assigned_varyings:
+                line = lines[used_varyings[name][0]]
+                raise TypeError(f"Varying {name!r} is read, but not assigned:\n{line}")
 
+        # Comment-out the varying setter if its unused elsewhere in the shader
+        for name, linenrs in assigned_varyings.items():
+            if name not in used_varyings:
+                for linenr in linenrs:
+                    line = lines[linenr]
+                    indent = indent_from_line(line)
+                    lines[linenr] = indent + "// unused: " + line[len(indent) :]
+                    # Deal with multiple-line assignments
+                    line_s = line.strip()
+                    while not line_s.endswith(";"):
+                        linenr += 1
+                        line_s = lines[linenr].strip()
+                        unexpected = "fn ", "struct ", "var ", "let ", "}"
+                        if line_s.startswith(unexpected) or linenr == len(lines) - 1:
+                            raise TypeError(
+                                f"Varying {name!r} assignment seems to be missing a semicolon:\n{line}"
+                            )
+                        lines[linenr] = indent + "// " + line_s
 
-def resolve_depth_output(wgsl):
+        # Build and insert the struct
+        if struct_linenr is not None:
+            # Refine position
+            while struct_linenr > 0 and not lines[struct_linenr].lstrip().startswith(
+                "fn"
+            ):
+                struct_linenr -= 1
+            if struct_linenr > 0 and lines[struct_linenr - 1].lstrip().startswith("@"):
+                struct_linenr -= 1
+            # First divide into slot-based and builtins
+            used_varyings = set(used_varyings)
+            used_builtins = used_varyings.intersection(builtin_varyings)
+            used_slots = used_varyings.difference(used_builtins)
+            used_slots = list(sorted(used_slots))
+            # Build struct
+            struct_lines = ["struct Varyings {"]
+            for slotnr, name in enumerate(used_slots):
+                struct_lines.append(f"    @location({slotnr}) {name} : {types[name]},")
+            for name in sorted(used_builtins):
+                struct_lines.append(f"    @builtin({name}) {name} : {types[name]},")
+            struct_lines.append("};\n")
+            # Apply indentation and insert
+            line = lines[struct_linenr]
+            indent = indent_from_line(line)
+            new_lines = [indent + line for line in struct_lines] + [line]
+            lines[struct_linenr] = "\n".join(new_lines)
+        else:
+            assert not used_varyings, "woops, did not expect used_varyings here"
+
+
+re_out_create = re.compile(r"var\s*out\s*\:\s*FragmentOutput\s*;")
+re_out_create_old = re.compile(r"var\s*out\s*\=\s*get_fragment_output\(")
+re_out_field = re.compile(r"out\.(\w+)\s*\=", re.UNICODE)
+re_out_return = re.compile(r"return\s*out\s*;", re.MULTILINE)
+
+
+class OutputResolver:
     """When out.depth is set (in the fragment shader), adjust the FragmentOutput
     to accept depth.
     """
-    assert isinstance(wgsl, str)
 
-    # Split into lines, which is easier to process. Ensure it ends with newline in the end.
-    lines = wgsl.splitlines()
-    while lines and not lines[0].strip():
-        lines.pop(0)
-    while lines and not lines[-1].strip():
-        lines.pop(-1)
-    lines.append("")
-
-    # Detect whether the depth is set in the shader. We're going to assume
-    # this is in the fragment shader. We check for "out.depth =".
-    # Background: by default the depth is based on the geometry (set
-    # by vertex shader and interpolated). It is possible for a fragment
-    # shader to write the depth instead. If this is done, the GPU cannot
-    # do early depth testing; the fragment shader must be run for the
-    # depth to be known.
-    depth_is_set = False
-    struct_linenr = -1
-    for linenr, line in enumerate(lines):
-        if line.lstrip().startswith("struct FragmentOutput {"):
-            struct_linenr = linenr
-        elif re_depth_setter.match(line):
-            depth_is_set = True
-            if struct_linenr >= 0:
-                break
-
-    if depth_is_set:
-        if struct_linenr < 0:
-            raise TypeError("FragmentOutput definition not found.")
-        depth_field = "    @builtin(frag_depth) depth : f32,"
-        line = lines[struct_linenr]
-        indent = line[: len(line) - len(line.lstrip())]
-        lines.insert(struct_linenr + 1, indent + depth_field)
+    def __init__(self):
+        # Detect whether the depth is set in the shader. We're going to assume
+        # this is in the fragment shader. We check for "out.depth =".
+        # Background: by default the depth is based on the geometry (set
+        # by vertex shader and interpolated). It is possible for a fragment
+        # shader to write the depth instead. If this is done, the GPU cannot
+        # do early depth testing; the fragment shader must be run for the
+        # depth to be known.
 
-    return "\n".join(lines)
+        self.struct_linenr = None
+        self.assigned_fields = None  # list of (field-name, linenr)
+        self.assigned_fields_list = []  # list of assigned_fields
+
+    def process_line(self, linenr, line, current_func_kind):
+        in_fragment_shader = current_func_kind == "fragment"
+
+        if not in_fragment_shader:
+            self.assigned_fields = None
+
+        if self.struct_linenr is None:
+            if line.startswith("struct FragmentOutput {"):
+                self.struct_linenr = linenr
+
+        if in_fragment_shader:
+            if self.assigned_fields is None:
+                # Detect var out: FragmentOutput
+                if re_out_create.match(line) or re_out_create_old.match(line):
+                    self.assigned_fields = [("__create", linenr)]
+                    self.assigned_fields_list.append(self.assigned_fields)
+            else:
+                # Detect fields being set
+                match = re_out_field.match(line)
+                if match:
+                    name = match.group(1)
+                    self.assigned_fields.append((name, linenr))
+                # Detect return, because that's the best place to insert our call
+                match = re_out_return.match(line)
+                if match:
+                    self.assigned_fields.append(("return", linenr))
+                    self.assigned_fields = None
+
+    def update_lines(self, lines):
+        struct_linenr = self.struct_linenr
+        assigned_fields_list = self.assigned_fields_list
+
+        # Collect virtual fields, specified in comments inside the FragmentOutput definition.
+        # These will be the args for apply_virtual_fields_of_fragment_output()
+        # E.g. // virtualfield foo : f32 = vec4<f32>(0.0)
+        virtual_field_prefix = "// virtualfield "
+        virtual_fields = []
+        if struct_linenr is not None:
+            linenr = struct_linenr
+            while linenr < len(lines) - 1:
+                linenr += 1
+                line = lines[linenr].lstrip()
+                if line.startswith(virtual_field_prefix):
+                    name_type_value = line[len(virtual_field_prefix) :].split("//")[0]
+                    name, colon, type_value = name_type_value.partition(":")
+                    type, eq, value = type_value.partition("=")
+                    assert colon and eq, (
+                        "A virtualfield needs to be of the form 'name : type = defaultvalue'. {name!r} does not"
+                    )
+                    virtual_fields.append(
+                        (name.strip(), type.strip(), value.strip(" ;"))
+                    )
+                else:
+                    break
+
+        depth_is_set = False
+
+        # Handle each place where a FragmentOutput is returned
+        for assigned_fields in assigned_fields_list:
+            create_linenr = next(i for name, i in assigned_fields if name == "__create")
+            create_indent = indent_from_line(lines[create_linenr])
+
+            # Handle get_fragment_output deprecated usage. Not pretty, but it works and is temporary.
+            # TODO: after one or two releases, remove the deprecated get_fragment_output case.
+            create_line = lines[create_linenr]
+            if re_out_create_old.match(create_line.lstrip()):
+                if "get_fragment_output" not in warned_for:
+                    warned_for.add("get_fragment_output")
+                    logger.warning("get_fragment_output() in WGSL is deprecated")
+                new_create_line = create_indent + "var out: FragmentOutput;"
+                colorset_line = (
+                    create_indent
+                    + "out.color ="
+                    + create_line.split("=", 1)[1].replace(");", ").color;")
+                )
+                lines[create_linenr - 1] += "\n" + new_create_line
+                lines[create_linenr] = colorset_line
+                assigned_fields.append(("color", create_linenr))
+                create_linenr -= 1
+
+            # Collect virtual argument values
+            args = []
+            for name, type, default_value in virtual_fields:
+                linenrs = [i for n, i in assigned_fields if n == name]
+                if not linenrs:
+                    args.append(default_value)
+                else:
+                    prefixed_name = f"out_virtualfield_{name}"
+                    args.append(prefixed_name)
+                    lines[create_linenr] += (
+                        f"\n{create_indent}var {prefixed_name}: {type};"
+                    )
+
+                    for linenr in linenrs:
+                        lines[linenr] = lines[linenr].replace(
+                            f"out.{name}", f"{prefixed_name}", 1
+                        )
+
+            # Process special arguments
+            if any(i for name, i in assigned_fields if name == "depth"):
+                depth_is_set = True
+
+            # Apply the virtual fields. The call to the apply-function is
+            # inserted right before 'return out', or right after the last struct
+            # field is set (if we did not detect a return).
+            if args:
+                linenr = max(i for _, i in assigned_fields)
+                args.insert(0, "&out")
+                line = lines[linenr]
+                is_return_line = line.lstrip().startswith("return")
+                indent = indent_from_line(line)
+                extra_line = f"{indent}apply_virtual_fields_of_fragment_output({', '.join(args)});"
+                new_lines = [extra_line, line] if is_return_line else [line, extra_line]
+                lines[linenr] = "\n".join(new_lines)
+
+        # If depth is set, add the depth field to the struct definition.
+        if depth_is_set:
+            if struct_linenr is None:
+                raise TypeError("FragmentOutput definition not found.")
+            depth_field = "    @builtin(frag_depth) depth : f32,"
+            line = lines[struct_linenr]
+            indent = indent_from_line(line)
+            new_lines = [line, indent + depth_field]
+            lines[struct_linenr] = "\n".join(new_lines)
diff --git a/pygfx/renderers/wgpu/shaders/lineshader.py b/pygfx/renderers/wgpu/shaders/lineshader.py
index 197e7ff..f32f457 100644
--- a/pygfx/renderers/wgpu/shaders/lineshader.py
+++ b/pygfx/renderers/wgpu/shaders/lineshader.py
@@ -514,7 +514,9 @@ class ThinLineShader(LineShader):
             let opacity = color.a * u_material.opacity;
             let out_color = vec4<f32>(physical_color, opacity);
 
-            return get_fragment_output(varyings.position, out_color);
+            var out: FragmentOutput;
+            out.color = out_color;
+            return out;
         }
         """
 
diff --git a/pygfx/renderers/wgpu/wgsl/background.wgsl b/pygfx/renderers/wgpu/wgsl/background.wgsl
index 214b457..e2f86ab 100644
--- a/pygfx/renderers/wgpu/wgsl/background.wgsl
+++ b/pygfx/renderers/wgpu/wgsl/background.wgsl
@@ -10,7 +10,7 @@ struct VertexInput {
 
 @vertex
 fn vs_main(in: VertexInput) -> Varyings {
-    var varyings: Varyings;
+
     // Define positions at the four corners of the viewport, at the largest depth
     var positions = array<vec2<f32>, 4>(
         vec2<f32>(-1.0, -1.0),
@@ -43,6 +43,7 @@ fn vs_main(in: VertexInput) -> Varyings {
         let texcoord = vec3<f32>(virtual_pos * 0.5 + 0.5, 0.0);
     $$ endif
 
+    var varyings: Varyings;
     varyings.position = vec4<f32>(pos, 0.9999999, 1.0);
     varyings.texcoord = vec3<f32>(texcoord);
     //varyings.world_pos = vec3<f32>(ndc_to_world_pos(vec4<f32>(virtual_pos, 0.9999999, 1.0);));
@@ -84,11 +85,9 @@ fn fs_main(varyings: Varyings) -> FragmentOutput {
     // We can apply clipping planes, but maybe a background should not be clipped?
     // apply_clipping_planes(in.world_pos);
 
-    // This is the opaque pass.
-    // A fragment of the background could be transparent, but it should still be
-    // written in the opaque pass in order for it to really be background.
-    // So we fool the blender into thinking this fragment is opaque, even if its not.
-    var out = get_fragment_output(varyings.position, vec4<f32>(out_color.rgb, 1.0));
+    // Force opaque. We may undo this when the ordered2 mechanics are removed.
+    var out: FragmentOutput;
+    out.color = vec4<f32>(out_color.rgb, 1.0);
     $$ if write_pick
         // We omit any extra information in the pick info
         // While we figure out exactly how best to return it.
@@ -96,6 +95,5 @@ fn fs_main(varyings: Varyings) -> FragmentOutput {
         // https://github.com/pygfx/pygfx/pull/700
         out.pick = pick_pack(u32(u_wobject.id), 20);
     $$ endif
-    out.color = vec4<f32>(out_color);
     return out;
 }
\ No newline at end of file
diff --git a/pygfx/renderers/wgpu/wgsl/grid.wgsl b/pygfx/renderers/wgpu/wgsl/grid.wgsl
index 69b6b1f..2a4b80f 100644
--- a/pygfx/renderers/wgpu/wgsl/grid.wgsl
+++ b/pygfx/renderers/wgpu/wgsl/grid.wgsl
@@ -11,7 +11,6 @@ struct VertexInput {
 
 @vertex
 fn vs_main(in: VertexInput) -> Varyings {
-    var varyings: Varyings;
 
     // Calculate reference points
     let p0 = (u_wobject.world_transform * vec4<f32>(0.0, 0.0, 0.0, 1.0)).xyz;  // "origin" of the plane
@@ -151,6 +150,7 @@ fn vs_main(in: VertexInput) -> Varyings {
 
     var pos_ndc = u_stdinfo.projection_transform * u_stdinfo.cam_transform * vec4<f32>(pos, 1.0);
 
+    var varyings: Varyings;
     // Store positions and the view direction in the world
     varyings.position = vec4<f32>(pos_ndc);
     varyings.world_pos = vec3<f32>(ndc_to_world_pos(out.position));
@@ -279,7 +279,8 @@ fn fs_main(varyings: Varyings) -> FragmentOutput {
     // We can apply clipping planes, but maybe a grid should not be clipped?
     // apply_clipping_planes(in.world_pos);
 
-    var out = get_fragment_output(varyings.position, out_color);
+    var out: FragmentOutput;
+    out.color = out_color;
     $$ if write_pick
         // Just the object id for now
         out.pick = pick_pack(u32(u_wobject.id), 20);
diff --git a/pygfx/renderers/wgpu/wgsl/image.wgsl b/pygfx/renderers/wgpu/wgsl/image.wgsl
index 66a382e..2219b19 100644
--- a/pygfx/renderers/wgpu/wgsl/image.wgsl
+++ b/pygfx/renderers/wgpu/wgsl/image.wgsl
@@ -53,7 +53,8 @@ fn fs_main(varyings: Varyings) -> FragmentOutput {
     let opacity = color.a * u_material.opacity;
     let out_color = vec4<f32>(physical_color, opacity);
 
-    var out = get_fragment_output(varyings.position, out_color);
+    var out: FragmentOutput;
+    out.color = out_color;
 
     $$ if write_pick
     // The wobject-id must be 20 bits. In total it must not exceed 64 bits.
diff --git a/pygfx/renderers/wgpu/wgsl/line.wgsl b/pygfx/renderers/wgpu/wgsl/line.wgsl
index 51318db..57f0591 100644
--- a/pygfx/renderers/wgpu/wgsl/line.wgsl
+++ b/pygfx/renderers/wgpu/wgsl/line.wgsl
@@ -892,7 +892,8 @@ fn fs_main(varyings: Varyings, @builtin(front_facing) is_front: bool) -> Fragmen
     let opacity = min(1.0, color.a) * alpha * u_material.opacity;
     let out_color = vec4<f32>(physical_color, opacity);
 
-    var out = get_fragment_output(varyings.position, out_color);
+    var out: FragmentOutput;
+    out.color = out_color;
 
     // Set picking info.
     $$ if write_pick
diff --git a/pygfx/renderers/wgpu/wgsl/mesh.wgsl b/pygfx/renderers/wgpu/wgsl/mesh.wgsl
index 44b8c16..f94a976 100644
--- a/pygfx/renderers/wgpu/wgsl/mesh.wgsl
+++ b/pygfx/renderers/wgpu/wgsl/mesh.wgsl
@@ -631,7 +631,8 @@ fn fs_main(varyings: Varyings, @builtin(front_facing) is_front: bool) -> Fragmen
 
     let out_color = vec4<f32>(physical_color, diffuse_color.a);
 
-    var out = get_fragment_output(varyings.position, out_color);
+    var out: FragmentOutput;
+    out.color = out_color;
 
     $$ if write_pick
     // The wobject-id must be 20 bits. In total it must not exceed 64 bits.
diff --git a/pygfx/renderers/wgpu/wgsl/mesh_normal_lines.wgsl b/pygfx/renderers/wgpu/wgsl/mesh_normal_lines.wgsl
index 4709bc8..9efe420 100644
--- a/pygfx/renderers/wgpu/wgsl/mesh_normal_lines.wgsl
+++ b/pygfx/renderers/wgpu/wgsl/mesh_normal_lines.wgsl
@@ -71,7 +71,8 @@ fn fs_main(varyings: Varyings, @builtin(front_facing) is_front: bool) -> Fragmen
     let out_color = vec4<f32>(physical_color, opacity);
 
     // Wrap up
-    var out = get_fragment_output(varyings.position, out_color);
+    var out: FragmentOutput;
+    out.color = out_color;
 
     $$ if write_pick
     // The wobject-id must be 20 bits. In total it must not exceed 64 bits.
diff --git a/pygfx/renderers/wgpu/wgsl/mesh_slice.wgsl b/pygfx/renderers/wgpu/wgsl/mesh_slice.wgsl
index 65fb04f..64d8e8b 100644
--- a/pygfx/renderers/wgpu/wgsl/mesh_slice.wgsl
+++ b/pygfx/renderers/wgpu/wgsl/mesh_slice.wgsl
@@ -231,7 +231,9 @@ fn fs_main(varyings: Varyings) -> FragmentOutput {
     let opacity = min(1.0, color_value.a) * alpha;
     let out_color = vec4<f32>(physical_color, opacity);
 
-    var out = get_fragment_output(varyings.position, out_color);
+    var out: FragmentOutput;
+    out.color = out_color;
+
     $$ if write_pick
     // The wobject-id must be 20 bits. In total it must not exceed 64 bits.
     out.pick = (
diff --git a/pygfx/renderers/wgpu/wgsl/points.wgsl b/pygfx/renderers/wgpu/wgsl/points.wgsl
index 3cfdd13..511e1ee 100644
--- a/pygfx/renderers/wgpu/wgsl/points.wgsl
+++ b/pygfx/renderers/wgpu/wgsl/points.wgsl
@@ -370,7 +370,8 @@ fn fs_main(varyings: Varyings) -> FragmentOutput {
         let out_color = vec4<f32>(srgb2physical(the_color.rgb), the_color.a * u_material.opacity);
     $$ endif
 
-    var out = get_fragment_output(varyings.position, out_color);
+    var out: FragmentOutput;
+    out.color = out_color;
 
     $$ if write_pick
     // The wobject-id must be 20 bits. In total it must not exceed 64 bits.
diff --git a/pygfx/renderers/wgpu/wgsl/std.wgsl b/pygfx/renderers/wgpu/wgsl/std.wgsl
index 6a41cfe..d7de5e2 100644
--- a/pygfx/renderers/wgpu/wgsl/std.wgsl
+++ b/pygfx/renderers/wgpu/wgsl/std.wgsl
@@ -9,7 +9,7 @@
 //
 // Also adds auto-generated code for:
 //
-// * FragmentOutput and get_fragment_output (implementation depends on blend mode and picking)
+// * FragmentOutput (implementation depends on blend mode and picking)
 // * Bindings (as specified via shader.define_binding())
 //
 // In addition to this, the shader will also:
@@ -102,9 +102,17 @@ fn is_orthographic() -> bool {
 
 // ----- Things related to output
 
-// Implements get_fragment_output
+// Implements FragmentOutput and optionally apply_virtual_fields_of_fragment_output()
 {{ blending_code }}
-
+struct StubColorWrapper {
+    color: vec4<f32>,
+}
+// Temporary for backwards compat
+fn get_fragment_output(position: vec4<f32>, color: vec4<f32>) -> StubColorWrapper {
+    var wrapper : StubColorWrapper;
+    wrapper.color = color;
+    return wrapper;
+}
 
 fn srgb2physical(color: vec3<f32>) -> vec3<f32> {
     // In Python, the below reads as
diff --git a/pygfx/renderers/wgpu/wgsl/text.wgsl b/pygfx/renderers/wgpu/wgsl/text.wgsl
index 4012445..8a1b0d2 100644
--- a/pygfx/renderers/wgpu/wgsl/text.wgsl
+++ b/pygfx/renderers/wgpu/wgsl/text.wgsl
@@ -298,7 +298,8 @@ fn fs_main(varyings: Varyings) -> FragmentOutput {
     // Debug
     //color_out = vec4<f32>(atlas_value, 0.0, 0.0, 1.0);
 
-    var out = get_fragment_output(varyings.position, color_out);
+    var out: FragmentOutput;
+    out.color = color_out;
 
     // Move text closer to camera, since its often overlaid on something.
     // The text is moved closer than the outline so that the outline of one character
diff --git a/pygfx/renderers/wgpu/wgsl/volume_ray.wgsl b/pygfx/renderers/wgpu/wgsl/volume_ray.wgsl
index 266b64f..d842aca 100644
--- a/pygfx/renderers/wgpu/wgsl/volume_ray.wgsl
+++ b/pygfx/renderers/wgpu/wgsl/volume_ray.wgsl
@@ -117,13 +117,10 @@ fn fs_main(varyings: Varyings) -> FragmentOutput {
     let world_pos = u_wobject.world_transform * vec4<f32>(data_pos, 1.0);
     let ndc_pos = u_stdinfo.projection_transform * u_stdinfo.cam_transform * world_pos;
 
-    // Maybe we did the work for nothing
-
-    // Get fragment output. Note that the depth arg only affects the
-    // blending - setting the depth attribute actually sets the fragment depth.
-    let depth = ndc_pos.z / ndc_pos.w;
-    var out = get_fragment_output(vec4<f32>(ndc_pos.xy, depth, 1.0), render_out.color);
-    out.depth = depth;
+    // Create fragment output.
+    var out: FragmentOutput;
+    out.color = render_out.color;
+    out.depth = ndc_pos.z / ndc_pos.w;
 
     $$ if write_pick
     // The wobject-id must be 20 bits. In total it must not exceed 64 bits.
diff --git a/pygfx/renderers/wgpu/wgsl/volume_slice.wgsl b/pygfx/renderers/wgpu/wgsl/volume_slice.wgsl
index 60cb445..7b73021 100644
--- a/pygfx/renderers/wgpu/wgsl/volume_slice.wgsl
+++ b/pygfx/renderers/wgpu/wgsl/volume_slice.wgsl
@@ -184,7 +184,8 @@ fn fs_main(varyings: Varyings) -> FragmentOutput {
     let out_color = vec4<f32>(physical_color, opacity);
 
 
-    var out = get_fragment_output(varyings.position, out_color);
+    var out: FragmentOutput;
+    out.color = out_color;
 
     $$ if write_pick
     // The wobject-id must be 20 bits. In total it must not exceed 64 bits.


Test patch:
diff --git a/tests/renderers/test_resolve_fragmentoutput.py b/tests/renderers/test_resolve_fragmentoutput.py
new file mode 100644
index 0000000..ce3117a
--- /dev/null
+++ b/tests/renderers/test_resolve_fragmentoutput.py
@@ -0,0 +1,535 @@
+from pygfx.renderers.wgpu.shader.resolve import resolve_output
+
+
+def test_fragmentoutput_nothing():
+    # No fields are set, and there are no virtual fields.
+    code1 = """
+    struct FragmentOutput {
+        @location(0) color: vec4<f32>,
+    }
+    @fragment
+    fn fs_main() {
+    }
+    """.strip()
+    code2 = code1
+    assert resolve_output(code1).strip() == code2
+
+
+def test_fragmentoutput_no_virtual_fields():
+    # No virtual fields, but the real fields are set
+    code1 = """
+    struct FragmentOutput {
+        @location(0) color: vec4<f32>,
+        @location(1) foo: f32,
+    }
+    @fragment
+    fn fs_main() {
+        var out: FragmentOutput;
+        out.color = vec4<f32>(1.0);
+        out.foo = 42.0;
+        return out;
+    }
+    """.strip()
+    code2 = code1
+    assert resolve_output(code1).strip() == code2
+
+
+def test_fragmentoutput_no_output():
+    # Has virtual fields, but code does not use FragmentOutput in fragment shader
+    code1 = """
+    struct FragmentOutput {
+        // virtualfield foo : f32 = 42;
+        @location(0) color: vec4<f32>,
+    }
+    fn some function() {
+        var out: FragmentOutput;
+        out.color = vec4<f32>(1.0);
+        out.foo = 42.0;
+        return out;
+    }
+    """.strip()
+    code2 = code1
+    assert resolve_output(code1).strip() == code2
+
+
+def test_fragmentoutput_virtual_using_default():
+    # out.foo is not specified, so it uses the default
+
+    code1 = """
+    struct FragmentOutput {
+        // virtualfield foo : f32 = 42;
+        @location(0) color: vec4<f32>,
+    }
+    @fragment
+    fn fs_main() {
+        var out: FragmentOutput;
+        return out;
+    }
+    """.strip()
+
+    code2 = """
+    struct FragmentOutput {
+        // virtualfield foo : f32 = 42;
+        @location(0) color: vec4<f32>,
+    }
+    @fragment
+    fn fs_main() {
+        var out: FragmentOutput;
+        apply_virtual_fields_of_fragment_output(&out, 42);
+        return out;
+    }
+    """.strip()
+
+    code3 = resolve_output(code1).strip()
+    assert code3 == code2
+
+
+def test_fragmentoutput_virtual_no_return():
+    # test special case; does not rely on return out
+
+    code1 = """
+    struct FragmentOutput {
+        // virtualfield foo : f32 = 42;
+        @location(0) color: vec4<f32>,
+    }
+    @fragment
+    fn fs_main() {
+        var out: FragmentOutput;
+    }
+    """.strip()
+
+    code2 = """
+    struct FragmentOutput {
+        // virtualfield foo : f32 = 42;
+        @location(0) color: vec4<f32>,
+    }
+    @fragment
+    fn fs_main() {
+        var out: FragmentOutput;
+        apply_virtual_fields_of_fragment_output(&out, 42);
+    }
+    """.strip()
+
+    code3 = resolve_output(code1).strip()
+    assert code3 == code2
+
+
+def test_fragmentoutput_virtual_is_set():
+    # out.foo is set, use that
+
+    code1 = """
+    struct FragmentOutput {
+        // virtualfield foo : f32 = 42;
+        @location(0) color: vec4<f32>,
+    }
+    @fragment
+    fn fs_main() {
+        var out: FragmentOutput;
+        out.foo = 7;
+        return out;
+    }
+    """.strip()
+
+    code2 = """
+    struct FragmentOutput {
+        // virtualfield foo : f32 = 42;
+        @location(0) color: vec4<f32>,
+    }
+    @fragment
+    fn fs_main() {
+        var out: FragmentOutput;
+        var out_virtualfield_foo: f32;
+        out_virtualfield_foo = 7;
+        apply_virtual_fields_of_fragment_output(&out, out_virtualfield_foo);
+        return out;
+    }
+    """.strip()
+
+    code3 = resolve_output(code1).strip()
+    assert code3 == code2
+
+
+def test_fragmentoutput_virtual_is_set_in_branch():
+    # out.foo is set, use that
+
+    code1 = """
+    struct FragmentOutput {
+        // virtualfield foo : f32 = 42;
+        @location(0) color: vec4<f32>,
+    }
+    @fragment
+    fn fs_main() {
+        var out: FragmentOutput;
+        if (something) {
+            out.foo = 7;
+        } else {
+            out.foo = 9;
+        }
+        return out;
+    }
+    """.strip()
+
+    code2 = """
+    struct FragmentOutput {
+        // virtualfield foo : f32 = 42;
+        @location(0) color: vec4<f32>,
+    }
+    @fragment
+    fn fs_main() {
+        var out: FragmentOutput;
+        var out_virtualfield_foo: f32;
+        if (something) {
+            out_virtualfield_foo = 7;
+        } else {
+            out_virtualfield_foo = 9;
+        }
+        apply_virtual_fields_of_fragment_output(&out, out_virtualfield_foo);
+        return out;
+    }
+    """.strip()
+
+    code3 = resolve_output(code1).strip()
+    assert code3 == code2
+
+
+def test_fragmentoutput_virtual_is_also_real_field_unset():
+    # out.color is both a real and a virtual field, but not set
+
+    code1 = """
+    struct FragmentOutput {
+        // virtualfield color : vec4<f32> = vec4<f32>(0.0);
+        @location(0) color: vec4<f32>,
+    }
+    @fragment
+    fn fs_main() {
+        var out: FragmentOutput;
+        return out;
+    }
+    """.strip()
+
+    code2 = """
+    struct FragmentOutput {
+        // virtualfield color : vec4<f32> = vec4<f32>(0.0);
+        @location(0) color: vec4<f32>,
+    }
+    @fragment
+    fn fs_main() {
+        var out: FragmentOutput;
+        apply_virtual_fields_of_fragment_output(&out, vec4<f32>(0.0));
+        return out;
+    }
+    """.strip()
+
+    code3 = resolve_output(code1).strip()
+    assert code3 == code2
+
+
+def test_fragmentoutput_virtual_is_also_real_field_set():
+    # out.color is both a real and a virtual field, and set
+
+    code1 = """
+    struct FragmentOutput {
+        // virtualfield color : vec4<f32> = vec4<f32>(0.0);
+        @location(0) color: vec4<f32>,
+    }
+    @fragment
+    fn fs_main() {
+        var out: FragmentOutput;
+        out.color = some_color;
+        return out;
+    }
+    """.strip()
+
+    code2 = """
+    struct FragmentOutput {
+        // virtualfield color : vec4<f32> = vec4<f32>(0.0);
+        @location(0) color: vec4<f32>,
+    }
+    @fragment
+    fn fs_main() {
+        var out: FragmentOutput;
+        var out_virtualfield_color: vec4<f32>;
+        out_virtualfield_color = some_color;
+        apply_virtual_fields_of_fragment_output(&out, out_virtualfield_color);
+        return out;
+    }
+    """.strip()
+
+    code3 = resolve_output(code1).strip()
+    assert code3 == code2
+
+
+def test_fragmentoutput_virtual_is_multiline():
+    # out.foo is set using a multiline value
+
+    code1 = """
+    struct FragmentOutput {
+        // virtualfield foo : vec2<f32> = vec2<f32>(0.0);
+        @location(0) color: vec4<f32>,
+    }
+    @fragment
+    fn fs_main() {
+        var out: FragmentOutput;
+        out.foo = vec2<f32>(
+                1.0,
+                2.0
+        );
+        return out;
+    }
+    """.strip()
+
+    code2 = """
+    struct FragmentOutput {
+        // virtualfield foo : vec2<f32> = vec2<f32>(0.0);
+        @location(0) color: vec4<f32>,
+    }
+    @fragment
+    fn fs_main() {
+        var out: FragmentOutput;
+        var out_virtualfield_foo: vec2<f32>;
+        out_virtualfield_foo = vec2<f32>(
+                1.0,
+                2.0
+        );
+        apply_virtual_fields_of_fragment_output(&out, out_virtualfield_foo);
+        return out;
+    }
+    """.strip()
+
+    code3 = resolve_output(code1).strip()
+    assert code3 == code2
+
+
+def test_fragmentoutput_multiple_places():
+    # Capable of dealing with multiple end-points
+
+    code1 = """
+    struct FragmentOutput {
+        // virtualfield foo : f32 = 42;
+        // virtualfield bar : vec2<f32> = vec2<f32>(0.0);
+        @location(0) color: vec4<f32>,
+    }
+    @fragment
+    fn fs_main1() {
+        var out: FragmentOutput;
+        out.foo = 7;
+        return out;
+    }
+    fn another function() {
+        let out = 2;
+        return out;
+    }
+    @fragment
+    fn fs_main2() {
+        var out: FragmentOutput;
+        out.bar = vec2<f32>(1.0, 2.0);
+        return out;
+    }
+    """.strip()
+
+    code2 = """
+    struct FragmentOutput {
+        // virtualfield foo : f32 = 42;
+        // virtualfield bar : vec2<f32> = vec2<f32>(0.0);
+        @location(0) color: vec4<f32>,
+    }
+    @fragment
+    fn fs_main1() {
+        var out: FragmentOutput;
+        var out_virtualfield_foo: f32;
+        out_virtualfield_foo = 7;
+        apply_virtual_fields_of_fragment_output(&out, out_virtualfield_foo, vec2<f32>(0.0));
+        return out;
+    }
+    fn another function() {
+        let out = 2;
+        return out;
+    }
+    @fragment
+    fn fs_main2() {
+        var out: FragmentOutput;
+        var out_virtualfield_bar: vec2<f32>;
+        out_virtualfield_bar = vec2<f32>(1.0, 2.0);
+        apply_virtual_fields_of_fragment_output(&out, 42, out_virtualfield_bar);
+        return out;
+    }
+    """.strip()
+
+    code3 = resolve_output(code1).strip()
+    assert code3 == code2
+
+
+def test_fragmentoutput_add_depth_simple():
+    # depth field is set
+
+    code1 = """
+    struct FragmentOutput {
+        @location(0) color: vec4<f32>,
+    }
+    @fragment
+    fn fs_main() {
+        var out: FragmentOutput;
+        out.depth = 0.0;
+        return out;
+    }
+    """
+
+    code2 = """
+    struct FragmentOutput {
+        @builtin(frag_depth) depth : f32,
+        @location(0) color: vec4<f32>,
+    }
+    @fragment
+    fn fs_main() {
+        var out: FragmentOutput;
+        out.depth = 0.0;
+        return out;
+    }
+    """.strip()
+
+    code3 = resolve_output(code1).strip()
+    assert code3 == code2
+
+
+def test_fragmentoutput_add_depth_with_virtual():
+    # Depth field is set, and there are also virtual fields
+
+    code1 = """
+    struct FragmentOutput {
+        // virtualfield foo : f32 = 42;
+        @location(0) color: vec4<f32>,
+    }
+    @fragment
+    fn fs_main() {
+        var out: FragmentOutput;
+        out.color = vec4<f32>(1.0);
+        out.depth = 0.0;
+        return out;
+    }
+    """
+
+    code2 = """
+    struct FragmentOutput {
+        @builtin(frag_depth) depth : f32,
+        // virtualfield foo : f32 = 42;
+        @location(0) color: vec4<f32>,
+    }
+    @fragment
+    fn fs_main() {
+        var out: FragmentOutput;
+        out.color = vec4<f32>(1.0);
+        out.depth = 0.0;
+        apply_virtual_fields_of_fragment_output(&out, 42);
+        return out;
+    }
+    """.strip()
+
+    code3 = resolve_output(code1).strip()
+    assert code3 == code2
+
+
+def test_fragmentoutput_legacy_simple():
+    # using get_fragment_output(), simple
+
+    code1 = """
+    struct FragmentOutput {
+        @location(0) color: vec4<f32>,
+    }
+    @fragment
+    fn fs_main() {
+        var out = get_fragment_output(varyings.position, out_color);
+        return out;
+    }
+    """.strip()
+
+    code2 = """
+    struct FragmentOutput {
+        @location(0) color: vec4<f32>,
+    }
+    @fragment
+    fn fs_main() {
+        var out: FragmentOutput;
+        out.color = get_fragment_output(varyings.position, out_color).color;
+        return out;
+    }
+    """.strip()
+
+    code3 = resolve_output(code1).strip()
+    assert code3 == code2
+
+
+def test_fragmentoutput_legacy_with_virtual():
+    # using get_fragment_output(), and a virtual field
+
+    code1 = """
+    struct FragmentOutput {
+        // virtualfield foo : f32 = 42;
+        @location(0) color: vec4<f32>,
+    }
+    @fragment
+    fn fs_main() {
+        var out = get_fragment_output(varyings.position, out_color);
+        out.foo = 7;
+        return out;
+    }
+    """.strip()
+
+    code2 = """
+    struct FragmentOutput {
+        // virtualfield foo : f32 = 42;
+        @location(0) color: vec4<f32>,
+    }
+    @fragment
+    fn fs_main() {
+        var out: FragmentOutput;
+        var out_virtualfield_foo: f32;
+        out.color = get_fragment_output(varyings.position, out_color).color;
+        out_virtualfield_foo = 7;
+        apply_virtual_fields_of_fragment_output(&out, out_virtualfield_foo);
+        return out;
+    }
+    """.strip()
+
+    code3 = resolve_output(code1).strip()
+    assert code3 == code2
+
+
+def test_fragmentoutput_legacy_virtual_is_real():
+    # using get_fragment_output(), and using color as both a virtual and real field
+
+    code1 = """
+    struct FragmentOutput {
+        // virtualfield color : vec4<f32> = vec4<f32>(0.0);
+        @location(0) color: vec4<f32>,
+    }
+    @fragment
+    fn fs_main() {
+        var out = get_fragment_output(varyings.position, out_color);
+        return out;
+    }
+    """.strip()
+
+    code2 = """
+    struct FragmentOutput {
+        // virtualfield color : vec4<f32> = vec4<f32>(0.0);
+        @location(0) color: vec4<f32>,
+    }
+    @fragment
+    fn fs_main() {
+        var out: FragmentOutput;
+        var out_virtualfield_color: vec4<f32>;
+        out_virtualfield_color = get_fragment_output(varyings.position, out_color).color;
+        apply_virtual_fields_of_fragment_output(&out, out_virtualfield_color);
+        return out;
+    }
+    """.strip()
+
+    code3 = resolve_output(code1).strip()
+    assert code3 == code2
+
+
+if __name__ == "__main__":
+    for f in list(globals().values()):
+        if callable(f) and f.__name__.startswith("test_"):
+            print(f.__name__)
+            f()
diff --git a/tests/renderers/test_resolve_varyings.py b/tests/renderers/test_resolve_varyings.py
index 7b981ca..c1403ad 100644
--- a/tests/renderers/test_resolve_varyings.py
+++ b/tests/renderers/test_resolve_varyings.py
@@ -1,4 +1,4 @@
-from pygfx.renderers.wgpu.shader.base import resolve_varyings
+from pygfx.renderers.wgpu.shader.resolve import resolve_varyings
 from pytest import raises
 
 
@@ -14,6 +14,7 @@ def test_varyings_struct_position0():
     # If no varyings are used, the struct is not inserted
 
     code1 = """
+    @vertex
     fn vs_main() {
     }
     fn fs_main() {
@@ -30,6 +31,7 @@ def test_varyings_struct_position1():
     # But if the struct type is used, we insert it
 
     code1 = """
+    @vertex
     fn vs_main() -> Varyings {
     }
     fn fs_main(varyings : Varyings) {
@@ -40,6 +42,7 @@ def test_varyings_struct_position1():
     struct Varyings {
     };
 
+    @vertex
     fn vs_main() -> Varyings {
     }
     fn fs_main(varyings : Varyings) {
@@ -54,6 +57,7 @@ def test_varyings_struct_position2():
     # Indentation is preserved
 
     code1 = """
+            @vertex
             fn vs_main() -> Varyings {
             }
             fn fs_main(varyings : Varyings) {
@@ -64,6 +68,7 @@ def test_varyings_struct_position2():
             struct Varyings {
             };
 
+            @vertex
             fn vs_main() -> Varyings {
             }
             fn fs_main(varyings : Varyings) {
@@ -80,6 +85,7 @@ def test_varyings_struct_position3():
     code1 = """
     fn foo() = {
     }
+    @vertex
     fn vs_main() -> Varyings {
     }
     fn fs_main(varyings : Varyings) {
@@ -92,6 +98,7 @@ def test_varyings_struct_position3():
     struct Varyings {
     };
 
+    @vertex
     fn vs_main() -> Varyings {
     }
     fn fs_main(varyings : Varyings) {
@@ -128,7 +135,7 @@ def test_varyings_struct_position5():
     code1 = """
     fn foo() = {
     }
-    @blabla
+    @@vertex
     fn vs_main() -> Varyings {
     }
     fn fs_main(varyings : Varyings) {
@@ -141,7 +148,7 @@ def test_varyings_struct_position5():
     struct Varyings {
     };
 
-    @blabla
+    @@vertex
     fn vs_main() -> Varyings {
     }
     fn fs_main(varyings : Varyings) {
@@ -156,6 +163,7 @@ def test_varyings_remove1():
     # If no varyings are used, any set varyings are removed.
 
     code1 = """
+    @vertex
     fn vs_main() -> Varyings {
         var varyings : Varyings;
         varyings.foo = f32(something1);
@@ -170,6 +178,7 @@ def test_varyings_remove1():
     struct Varyings {
     };
 
+    @vertex
     fn vs_main() -> Varyings {
         var varyings : Varyings;
         // unused: varyings.foo = f32(something1);
@@ -188,6 +197,7 @@ def test_varyings_remove2():
     # Used varyings are held
 
     code1 = """
+    @vertex
     fn vs_main() -> Varyings {
         var varyings : Varyings;
         varyings.foo = f32(something1);
@@ -206,6 +216,7 @@ def test_varyings_remove2():
         @location(0) bar : vec2<f32>,
     };
 
+    @vertex
     fn vs_main() -> Varyings {
         var varyings : Varyings;
         // unused: varyings.foo = f32(something1);
@@ -227,6 +238,7 @@ def test_varyings_remove3():
     # Comments are taken into account
 
     code1 = """
+    @vertex
     fn vs_main() -> Varyings {
         var varyings : Varyings;
         varyings.foo = f32(something1);
@@ -245,6 +257,7 @@ def test_varyings_remove3():
     struct Varyings {
     };
 
+    @vertex
     fn vs_main() -> Varyings {
         var varyings : Varyings;
         // unused: varyings.foo = f32(something1);
@@ -267,6 +280,7 @@ def test_varyings_remove4():
     # Usage can actually be to the left of the assignment op
 
     code1 = """
+    @vertex
     fn vs_main() -> Varyings {
         var varyings : Varyings;
         varyings.foo = f32(something1);
@@ -283,6 +297,7 @@ def test_varyings_remove4():
         @location(0) foo : f32,
     };
 
+    @vertex
     fn vs_main() -> Varyings {
         var varyings : Varyings;
         varyings.foo = f32(something1);
@@ -302,6 +317,7 @@ def test_varyings_remove5():
     # Varying assignments can be multi-line
 
     code1 = """
+    @vertex
     fn vs_main() -> Varyings {
         var varyings : Varyings;
         varyings.spam = vec3<f32>(
@@ -315,6 +331,7 @@ def test_varyings_remove5():
     struct Varyings {
     };
 
+    @vertex
     fn vs_main() -> Varyings {
         var varyings : Varyings;
         // unused: varyings.spam = vec3<f32>(
@@ -332,6 +349,7 @@ def test_varyings_struct1():
     # Usage in other funcs counts, and varyings are sorted alphabetically
 
     code1 = """
+    @vertex
     fn vs_main() -> Varyings {
         var varyings : Varyings;
         varyings.foo = f32(something1);
@@ -356,6 +374,7 @@ def test_varyings_struct1():
         @location(2) spam : vec3<f32>,
     };
 
+    @vertex
     fn vs_main() -> Varyings {
         var varyings : Varyings;
         varyings.foo = f32(something1);
@@ -381,6 +400,7 @@ def test_varyings_attr1():
     # Can set varying attribute
 
     code1 = """
+    @vertex
     fn vs_main() -> Varyings {
         var varyings : Varyings;
         varyings.color = vec4<f32>(color);
@@ -398,6 +418,7 @@ def test_varyings_attr1():
         @location(0) color : vec4<f32>,
     };
 
+    @vertex
     fn vs_main() -> Varyings {
         var varyings : Varyings;
         varyings.color = vec4<f32>(color);
@@ -418,6 +439,7 @@ def test_varyings_attr2():
     # Can set varying attribute, and all setter lines are commented
 
     code1 = """
+    @vertex
     fn vs_main() -> Varyings {
         var varyings : Varyings;
         varyings.color = vec4<f32>(color);
@@ -433,6 +455,7 @@ def test_varyings_attr2():
     struct Varyings {
     };
 
+    @vertex
     fn vs_main() -> Varyings {
         var varyings : Varyings;
         // unused: varyings.color = vec4<f32>(color);
@@ -452,6 +475,7 @@ def test_varyings_builtin1():
     # Builtins are not slots
 
     code1 = """
+    @vertex
     fn vs_main() -> Varyings {
         var varyings : Varyings;
         varyings.position = vec4<f32>(ndc_pos);
@@ -473,6 +497,7 @@ def test_varyings_builtin1():
         @builtin(position) position : vec4<f32>,
     };
 
+    @vertex
     fn vs_main() -> Varyings {
         var varyings : Varyings;
         varyings.position = vec4<f32>(ndc_pos);
@@ -495,6 +520,7 @@ def test_varyings_builtin2():
     # Builtins are considered used when set
 
     code1 = """
+    @vertex
     fn vs_main() -> Varyings {
         var varyings : Varyings;
         varyings.position = vec4<f32>(ndc_pos);
@@ -510,6 +536,7 @@ def test_varyings_builtin2():
         @builtin(position) position : vec4<f32>,
     };
 
+    @vertex
     fn vs_main() -> Varyings {
         var varyings : Varyings;
         varyings.position = vec4<f32>(ndc_pos);
@@ -528,6 +555,7 @@ def test_varyings_error_readonly_in_vs1():
     # Cannot use varying as a getter in vs_main
 
     code1 = """
+    @vertex
     fn vs_main() -> Varyings {
         var varyings : Varyings;
         varyings.position = vec4<f32>(ndc_pos);
@@ -548,6 +576,7 @@ def test_varyings_error_need_type_when_set():
     # Assignment needs type annotation
 
     code1 = """
+    @vertex
     fn vs_main() -> Varyings {
         var varyings : Varyings;
         varyings.position = ndc_pos;
@@ -567,6 +596,7 @@ def test_varyings_error_assignments_invalid_multiline1():
     # Assignment needs type annotation
 
     code1 = """
+    @vertex
     fn vs_main() -> Varyings {
         var varyings : Varyings;
         varyings.foo = vec4<f32>(
@@ -577,6 +607,7 @@ def test_varyings_error_assignments_invalid_multiline1():
     info.match("missing a semicolon")
 
     code = """
+    @vertex
     fn vs_main() -> Varyings {
         var varyings : Varyings;
         varyings.foo = vec4<f32>(
@@ -596,6 +627,7 @@ def test_varyings_error_need_type_match():
     # When assigning multiple times, the type should match
 
     code1 = """
+    @vertex
     fn vs_main() -> Varyings {
         var varyings : Varyings;
         varyings.position = vec4<f32>(ndc_pos);
@@ -616,6 +648,7 @@ def test_varyings_error_used_but_not_assigned():
     # When assigning multiple times, the type should match
 
     code1 = """
+    @vertex
     fn vs_main() -> Varyings {
     }
 
diff --git a/tests/renderers/test_shader_composer.py b/tests/renderers/test_shader_composer.py
index 2883534..7ddfb34 100644
--- a/tests/renderers/test_shader_composer.py
+++ b/tests/renderers/test_shader_composer.py
@@ -1,5 +1,4 @@
 from pygfx.renderers.wgpu.shader import BaseShader
-from pygfx.renderers.wgpu.shader.base import resolve_depth_output
 from pygfx.renderers.wgpu import Binding
 from pygfx.utils import array_from_shadertype
 from pygfx import Buffer
@@ -188,41 +187,7 @@ def test_uniform_definitions():
         shader.define_binding(0, 0, Binding("zz", "buffer/uniform", struct))
 
 
-def test_resolve_depth_output():
-    code1 = """
-    struct FragmentOutput {
-        @location(0) color: vec4<f32>,
-    }
-    fn fs_main() {
-    }
-    """.strip()
-    code2 = code1
-    assert resolve_depth_output(code1).strip() == code2
-
-    code1 = """
-    struct FragmentOutput {
-        @location(0) color: vec4<f32>,
-    }
-    fn fs_main() {
-        out.depth = 0.0;
-    }
-    """
-
-    code2 = """
-    struct FragmentOutput {
-        @builtin(frag_depth) depth : f32,
-        @location(0) color: vec4<f32>,
-    }
-    fn fs_main() {
-        out.depth = 0.0;
-    }
-    """.strip()
-
-    assert resolve_depth_output(code1).strip() == code2
-
-
 if __name__ == "__main__":
     test_templating()
     test_logic_beyond_templating()
     test_uniform_definitions()
-    test_resolve_depth_output()