Add write function for binding (#893)

Author: ccummingsNV

slangpy/core/dispatchdata.py

@@ -218,12 +218,15 @@ def dispatch(
         **kwargs: dict[str, Any],
     ) -> None:
 
-        # Merge uniforms
+        # Merge uniforms and collect writer tuples
         uniforms: dict[str, Any] = {}
+        writers: list[tuple] = []
         if opts.uniforms is not None:
             for u in opts.uniforms:
                 if isinstance(u, dict):
                     uniforms.update(u)
+                elif isinstance(u, tuple):
+                    writers.append(u)
                 else:
                     uniforms.update(u(self))  # type: ignore (need to work out native dispatch)
         uniforms.update(vars)
@@ -245,6 +248,8 @@ def dispatch(
         compute_pass = command_encoder.begin_compute_pass()
         cursor = ShaderCursor(compute_pass.bind_pipeline(self.compute_pipeline))
         cursor.write(uniforms)
+        for fn, args, wkwargs in writers:
+            fn(cursor, *args, **wkwargs)
         cursor.find_entry_point(0).write(call_data)
         compute_pass.dispatch(thread_count)
         compute_pass.end()

slangpy/core/function.py

@@ -156,6 +156,16 @@ def set(self, *args: Any, **kwargs: Any):
                 "Set requires either keyword arguments or 1 dictionary / hook argument"
             )
 
+    def write(self, fn: Callable, *args: Any, **kwargs: Any):
+        """
+        Specify a writer function that receives a ShaderCursor and optional arguments
+        to write uniforms directly. The function signature should be:
+            fn(cursor: ShaderCursor, *args, **kwargs)
+        """
+        if not callable(fn):
+            raise ValueError("write() requires a callable as the first argument")
+        return FunctionNodeSet(self, (fn, args, kwargs))
+
     def cuda_stream(self, stream: NativeHandle) -> "FunctionNode":
         """
         Specify a CUDA stream to use for the function. This is useful for synchronizing with other

slangpy/tests/slangpy_tests/test_sets_and_hooks.py

@@ -1,9 +1,11 @@
 # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
+from typing import Any
+
 import pytest
 import numpy as np
 
-from slangpy import DeviceType, Module
+from slangpy import DeviceType, Module, ShaderCursor
 from slangpy.types import Tensor
 from slangpy.testing import helpers
 
@@ -65,5 +67,27 @@ def test_set_with_callback(device_type: DeviceType):
     assert np.allclose(res_data, val_data + 10)
 
 
+@pytest.mark.parametrize("device_type", helpers.DEFAULT_DEVICE_TYPES)
+def test_write(device_type: DeviceType):
+    m = load_test_module(device_type)
+    assert m is not None
+
+    add_k = m.add_k.as_func()
+
+    val = Tensor.empty(m.device, dtype=float, shape=(10,))
+    val_data = np.zeros(10, dtype=np.float32)  # np.random.rand(10).astype(np.float32)
+    val.copy_from_numpy(val_data)
+
+    def writer(cursor: ShaderCursor, *args: Any, **kwargs: Any):
+        cursor.write({"params": {"k": kwargs["myvalue"]}})
+
+    add_k = add_k.write(writer, myvalue=10)
+
+    res = add_k(val)
+
+    res_data = res.to_numpy().view(dtype=np.float32)
+    assert np.allclose(res_data, val_data + 10)
+
+
 if __name__ == "__main__":
     pytest.main([__file__, "-v", "-s"])

src/slangpy_ext/utils/slangpy.cpp

@@ -933,6 +933,13 @@ nb::object NativeCallData::exec(
             for (auto u : uniforms) {
                 if (nb::isinstance<nb::dict>(u)) {
                     write_shader_cursor(cursor, nb::cast<nb::dict>(u));
+                } else if (nb::isinstance<nb::tuple>(u)) {
+                    // Writer tuple: (fn, args, kwargs)
+                    nb::tuple t = nb::cast<nb::tuple>(u);
+                    nb::object fn = t[0];
+                    nb::tuple args = nb::cast<nb::tuple>(t[1]);
+                    nb::dict kwargs = nb::cast<nb::dict>(t[2]);
+                    fn(nb::cast(cursor), *args, **kwargs);
                 } else {
                     write_shader_cursor(cursor, nb::cast<nb::dict>(u(this)));
                 }