Repository: pygfx/pygfx (python)

Issue:
Bug when dynamically applying or removing env_map
While testing the example in #1004, I found that dynamically applying or removing the env_map was not working correctly. Initially, I thought it was a small issue, assuming we weren't tracking the env_map in MeshStandardMaterial. However, after modifying the code to enable tracking, the program threw the following exception:

```

  File "F:\pygfx\pygfx\renderers\wgpu\engine\pipeline.py", line 135, in get_cached_render_pipeline
    result = device.create_render_pipeline(
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "F:\wgpu-py\wgpu\backends\wgpu_native\_api.py", line 1717, in create_render_pipeline
    id = libf.wgpuDeviceCreateRenderPipeline(self._internal, descriptor)
         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "F:\wgpu-py\wgpu\backends\wgpu_native\_helpers.py", line 457, in proxy_func
    raise wgpu_error
  File "F:\wgpu-py\wgpu\backends\wgpu_native\_api.py", line 1717, in create_render_pipeline
    id = libf.wgpuDeviceCreateRenderPipeline(self._internal, descriptor)
         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
wgpu._classes.GPUValidationError: Validation Error

Caused by:
  In wgpuDeviceCreateRenderPipeline
    Error matching ShaderStages(FRAGMENT) shader requirements against the pipeline
      Shader global ResourceBinding { group: 0, binding: 9 } is not available in the pipeline layout
        Binding is missing from the pipeline layout
```

Upon further investigation, I realized the issue was more complicated than expected. Here are my findings:

- When bindings change, the `update_shader_data` method of `PipelineContainer` recollects the bindings, but the `_binding_codes` in `BindingDefinitions` were not cleared (which caused the above exception). 
- After I added a clear operation for `_binding_codes` in the shader's `get_bindings()` method, I found that the shader's `kwargs` were not updated too, which led to some incorrect "macro switches" in the shader. This seems a bit tricky, as we cannot clear all `kwargs` as we did with `_binding_codes`.

Pull request description:
Improve tracking of bindings
Fixes #1005

* [x] Rename "kwargs"  to "template-vars".
* [x] Split into two groups, one for init, and one for bindings, so that the latter can be cleared, so that when values are *not* set, this is detected, changing the hash, and thus triggering a recompile.
* [x] Check clearing of `_binding_codes` in `BindingDefinitions` -> ~Not necessary, since on a recompile all definitions are recreated.~ done, use `clear()`.
* [x] Have a look at the performance of the hashing. -> with the mesh in `gltf_unlit.py` re-hasing takes 12us 🤷 
* [x] Documentation.

Patch:
diff --git a/docs/advanced_shaders.rst b/docs/advanced_shaders.rst
index 62b0de0..1e90806 100644
--- a/docs/advanced_shaders.rst
+++ b/docs/advanced_shaders.rst
@@ -33,7 +33,24 @@ The shader must implement a few methods. A typical shader is shown below:
 
         type = "render"  # must be "render" or "compute"
 
+        def ___init__(self, wobject):
+            super().__init__(wobject)
+
+            # The __init__ is a good place to examine the material and geometry and set any template-variables that
+            # affect the final wgsl. By accessing `material.has_some_value` here, the value is tracked, so that when
+            # `material.has_some_value` changes later, the shader is re-compiled.
+            if material.has_some_value:
+                self["some_template_variable"] = True
+
         def get_bindings(self, wobject, shared):
+
+            # You can also set template-variables here. Again, when things that are used here change later, this
+            # is detected, and this method will be called again. When a binding has changed (e.g. a colormap is replaced
+            # with another) while the formats etc. match, the shader code is not re-composed / re-compiled, making
+            # such actions very efficient.
+            if getattr(geometry, "colors"):
+                self["use_color_buffer"] = True
+
             # Collect bindings. We must return a dict mapping slot
             # indices to Binding objects. But it's sometimes easier to
             # collect bindings in a list and then convert to a dict.
diff --git a/pygfx/renderers/wgpu/engine/pipeline.py b/pygfx/renderers/wgpu/engine/pipeline.py
index a6d7f4e..576cc89 100644
--- a/pygfx/renderers/wgpu/engine/pipeline.py
+++ b/pygfx/renderers/wgpu/engine/pipeline.py
@@ -216,6 +216,7 @@ class PipelineContainerGroup:
         # Get what has changed
         tracker = self.tracker
         changed = tracker.pop_changed()
+
         if not self.initialized:
             self.initialized = True
             changed.add("create")
@@ -356,10 +357,10 @@ class PipelineContainer:
             self.wgpu_shaders = {}
 
         if "bindings" in changed:
-            self.shader.unlock_hash()
             with tracker.track_usage("!bindings"):
-                self.bindings_dicts = self.shader.get_bindings(wobject, self.shared)
-            self.shader.lock_hash()
+                self.bindings_dicts = self.shader.get_bindings_info(
+                    wobject, self.shared
+                )
             self._check_bindings()
             self.update_shader_hash()
             self.update_bind_groups()
diff --git a/pygfx/renderers/wgpu/shader/base.py b/pygfx/renderers/wgpu/shader/base.py
index a3fb0de..5b5f703 100644
--- a/pygfx/renderers/wgpu/shader/base.py
+++ b/pygfx/renderers/wgpu/shader/base.py
@@ -19,17 +19,11 @@ from .templating import apply_templating
 
 
 class ShaderInterface:
-    """Define what a shader object must look like from the pov from the pipeline."""
+    """Define what a shader object must look like from the pov of the pipeline."""
 
     def __init__(self):
         self._hash = None
 
-    def lock_hash(self):
-        self._hash = self._get_hash()
-
-    def unlock_hash(self):
-        self._hash = None
-
     @property
     def hash(self):
         """A hash of the current state of the shader. If the hash changed, it's likely that the shader changed."""
@@ -41,10 +35,10 @@ class ShaderInterface:
     def _get_hash(self):
         raise NotImplementedError()
 
-    def generate_wgsl(self, **kwargs):
+    def generate_wgsl(self, **template_vars):
         raise NotImplementedError()
 
-    def get_bindings(self, wobject, shared):
+    def get_bindings_info(self, wobject, shared):
         """Subclasses must return a dict describing the buffers and
         textures used by this shader.
 
@@ -88,23 +82,26 @@ class ShaderInterface:
 class BaseShader(ShaderInterface):
     """Base shader object to compose and template shaders using jinja2.
 
-    Templating variables can be provided as kwargs, set (and get) as attributes,
-    or passed as kwargs to ``generate_wgsl()``.
+    Templating variables can be passed kwargs, set(and get) as attributes, or passed as kwargs to ``generate_wgsl()``.
+    They should only be set in __init__() and get_bindings_info().
+
+    The template variables affect the hash. After `pipeline.py` calls get_bindings_info(), it checks
+    the hash to see if the shader has changed, and recompiles the wgsl if it has.
     """
 
     type = "render"  # must be "compute" or "render"
     needs_bake_function = False
 
-    def __init__(self, wobject, **kwargs):
+    def __init__(self, wobject, **template_vars):
         super().__init__()
 
-        # The shader kwargs
-        self.kwargs = kwargs
+        # The shader templating variables
+        self._template_vars_init = {}
+        self._template_vars_init.update(template_vars)
+        self._template_vars_bindings = {}
 
-        # The stored hash. If set, the hash is locked, and kwargs cannot
-        # be set. This is to prevent the shader implementations setting
-        # shader kwargs *after* the pipeline has obtained the hash.
-        self._hash = None
+        # Start off with the current dict set. This is reset when get_bindings_info() gets called.
+        self._template_vars_current = self._template_vars_init
 
         # Handling binding definitions is handled by a wrapped object.
         self._binding_definitions = BindingDefinitions()
@@ -118,24 +115,27 @@ class BaseShader(ShaderInterface):
         self["colormap_dim"] = None
 
     def __setitem__(self, key, value):
-        if hasattr(self.__class__, key):
-            msg = f"Templating variable {key} causes name clash with class attribute."
-            raise KeyError(msg)
-        if self._hash is not None:
+        if self._template_vars_current is None:
             raise RuntimeError(
-                "Shader is trying to set shaders kwargs while they're locked."
+                "Attempt to set template variable outside of _init() or _get_bindings()."
             )
-        self.kwargs[key] = value
+        self._template_vars_current[key] = value
 
     def __getitem__(self, key):
-        return self.kwargs[key]
+        try:
+            return self._template_vars_bindings[key]
+        except KeyError:
+            return self._template_vars_init[key]
 
     def _get_hash(self):
+        # Getting the hash also locks template var setting
+        self._template_vars_current = None
+
         # The full name of this shader class.
         fullname = self.__class__.__module__ + "." + self.__class__.__name__
 
         # If we assume that the shader class produces the same code for
-        # a specific set of kwargs, we can use the fullname in the hash
+        # a specific set of template-vars, we can use the fullname in the hash
         # instead of the actual code. This assumption is valid in
         # general, but can break down in a few specific situations, the
         # most common one being an interactive session. In this case,
@@ -148,50 +148,88 @@ class BaseShader(ShaderInterface):
         name_probably_defines_code = fullname.count(".") >= 2
 
         if name_probably_defines_code:
-            # Faster, but assumes that the produced code only depends on kwargs.
+            # Faster, but assumes that the produced code only depends on template-vars.
             return hash_from_value(
-                [fullname, self._binding_definitions.get_code(), self.kwargs]
+                [
+                    fullname,
+                    self._binding_definitions.get_code(),
+                    self._template_vars_init,
+                    self._template_vars_bindings,
+                ]
             )
         else:
             # More reliable (e.g. in an interactive session).
             return hash_from_value(
-                [self._binding_definitions.get_code(), self.get_code(), self.kwargs]
+                [
+                    self._binding_definitions.get_code(),
+                    self.get_code(),
+                    self._template_vars_init,
+                    self._template_vars_bindings,
+                ]
             )
 
+    def get_bindings_info(self, wobject, shared):
+        # We assume that in normal usage (by the pipeline.py logic),
+        # this method is called first after initialization,
+        # which means that we can handle the _template_vars_current here.
+        self._template_vars_current = self._template_vars_bindings
+
+        # Clear the binding-based template variables, and the binding definitions, so that when get_bindings()
+        # not-sets a variable (which was previously set), this affects the hash and is detected.
+        self._template_vars_bindings.clear()
+        self._binding_definitions.clear()
+
+        try:
+            return self.get_bindings(wobject, shared)
+        finally:
+            self._template_vars_current = None
+            self._hash = None  # template vars my have changed, so force a recalculation
+
+    def get_bindings(self, wobject, shared):
+        # Default implementation returns zero bindings
+        return {0: {}}
+
     def get_code(self):
         """Implement this to compose the total (but still templated)
         shader. This method is called by ``generate_wgsl()``.
         """
         raise NotImplementedError()
 
-    def generate_wgsl(self, **kwargs):
+    def generate_wgsl(self, **more_template_vars):
         """Generate the final WGSL. Calls get_code() and then resolves
         the templating variables, varyings, and depth output.
         """
 
         # Compose shader variables
-        shader_kwargs = {}
-        shader_kwargs.update(self.kwargs)
-        shader_kwargs.update(kwargs)
-        shader_kwargs["bindings_code"] = self._binding_definitions.get_code()
-
-        # If shader kwargs are used in ``get_code()``, make sure its correct.
-        ori_kwargs, self.kwargs = self.kwargs, shader_kwargs
+        template_vars = {}
+        template_vars.update(self._template_vars_init)
+        template_vars.update(self._template_vars_bindings)
+        template_vars.update(more_template_vars)
+        template_vars["bindings_code"] = self._binding_definitions.get_code()
+
+        # If template variables are used in ``get_code()``, make sure its correct.
+        ori_template_vars, self._template_vars_bindings = (
+            self._template_vars_bindings,
+            template_vars,
+        )
 
         try:
             # Obtain base code
             code1 = self.get_code()
             # Templating and resolving includes
-            code2 = apply_templating(code1, **shader_kwargs)
+            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)
             return code3
         finally:
-            self.kwargs = ori_kwargs
+            self._template_vars_bindings = ori_template_vars
 
     def define_bindings(self, bindgroup, bindings_dict):
         """Define a collection of bindings organized in a dict."""
+        if self._template_vars_current is not self._template_vars_bindings:
+            raise RuntimeError("define_bindings can only be used from get_bindings")
+
         for index, binding in bindings_dict.items():
             self._binding_definitions.define_binding(bindgroup, index, binding)
 
@@ -201,6 +239,9 @@ class BaseShader(ShaderInterface):
         The binding must be a Binding object. The code that defines the binding
         will be inserted in ``pygfx.std.wgsl``.
         """
+        if self._template_vars_current is not self._template_vars_bindings:
+            raise RuntimeError("define_bindings can only be used from get_bindings")
+
         self._binding_definitions.define_binding(bindgroup, index, binding)
 
     # ----- Colormap stuff
diff --git a/pygfx/renderers/wgpu/shader/bindings.py b/pygfx/renderers/wgpu/shader/bindings.py
index c5a15e1..bbac9d3 100644
--- a/pygfx/renderers/wgpu/shader/bindings.py
+++ b/pygfx/renderers/wgpu/shader/bindings.py
@@ -18,6 +18,11 @@ class BindingDefinitions:
         self._binding_codes = {}
         self._uniform_struct_names = {}  # dtype -> name
 
+    def clear(self):
+        self._typedefs.clear()
+        self._binding_codes.clear()
+        self._uniform_struct_names.clear()
+
     def get_code(self):
         """Get the wgsl source code for the collected bindings."""
         code = ""


Test patch:
diff --git a/tests/renderers/test_shader_composer.py b/tests/renderers/test_shader_composer.py
index fd2f9b7..2883534 100644
--- a/tests/renderers/test_shader_composer.py
+++ b/tests/renderers/test_shader_composer.py
@@ -87,6 +87,9 @@ def test_uniform_definitions():
 
     shader = MyShader()
 
+    # Make it look like we're inside get_bindings ...
+    shader._template_vars_current = shader._template_vars_bindings
+
     # Fails
     with raises(TypeError):  # Not a valid struct type
         shader.define_binding(0, 0, Binding("zz", "buffer/uniform", "not a struct"))