Support numpy as result type when slang function returns matrix (#250)

* Support numpy as result type when slang function returns matrix

close #205

When calling a slang function which returns a matrix type, we cannot
use "numpy" type bound to the result.

This change add the support to that by
*implementing store/load method for `{RW}NDBuffer` type in slang shader, 
* add type resolve logic for ndbuffer python type when we bind nbbuffer to slang matrix, it will known what type to bind with.
* For every test in test_transform.py, also add matrix type support. So this will cover NDBuffer and Tensor type.
* For test_numpy.py add test to cover the support of binding a numpy result to slang matrix return type.

Author: kaizhangNV

slangpy/builtin/ndbuffer.py

@@ -26,6 +26,7 @@
     SlangProgramLayout,
     SlangType,
     VectorType,
+    MatrixType,
     StructuredBufferType,
     is_matching_array_type,
 )
@@ -111,9 +112,8 @@ def ndbuffer_resolve_type(
     # if implicit tensor casts enabled, allow conversion from vector to element type
     if context.options["implicit_tensor_casts"]:
         if (
-            isinstance(bound_type, VectorType)
-            and self.slang_element_type == bound_type.element_type
-        ):
+            isinstance(bound_type, VectorType) or isinstance(bound_type, MatrixType)
+        ) and self.slang_element_type == bound_type.scalar_type:
             return bound_type
 
     # Default to just casting to itself (i.e. no implicit cast)

slangpy/builtin/valueref.py

@@ -62,7 +62,7 @@ def numpy_to_slang_matrix_remove_padding(
             (slang_type.rows, slang_type.cols),
             dtype=kfr.SCALAR_TYPE_TO_NUMPY_TYPE[slang_type.slang_scalar_type],
         )
-        mat_remove_padding[:, : slang_type.cols] = mat_aligned
+        mat_remove_padding = mat_aligned[:, : slang_type.cols]
         return python_type(mat_remove_padding)
 
 

slangpy/slang/core.slang

@@ -345,7 +345,23 @@ public struct NDBuffer<T, let N : int>
         }
     }
 
-    public void store<let VD : int>(ContextND<N - 1> context, out vector<T, VD> value) {}
+    public void store<let VD : int>(ContextND<N - 1> context, in vector<T, VD> value) {}
+
+    public void load<let R : int, let C: int>(ContextND<N - 2> context, out matrix<T, R, C> value) {
+        int call_id[N];
+        for (int i = 0; i < N - 2; i++) { call_id[i] = context.call_id[i]; }
+
+        for (int r = 0; r < R; r++)
+        {
+            call_id[N - 2] = r;
+            for (int vi = 0; vi < C; vi++) {
+                call_id[N - 1] = vi;
+                value[r][vi] = get(call_id);
+            }
+        }
+    }
+
+    public void store<let R : int, let C: int>(ContextND<N - 2> context, in matrix<T, R, C> value) {}
 }
 
 public struct RWNDBuffer<T, let N : int>
@@ -405,6 +421,34 @@ public struct RWNDBuffer<T, let N : int>
             set(call_id, value[vi]);
         }
     }
+
+    public void load<let R : int, let C: int>(ContextND<N - 2> context, out matrix<T, R, C> value) {
+        int call_id[N];
+        for (int i = 0; i < N - 2; i++) { call_id[i] = context.call_id[i]; }
+
+        for (int r = 0; r < R; r++)
+        {
+            call_id[N - 2] = r;
+            for (int vi = 0; vi < C; vi++) {
+                call_id[N - 1] = vi;
+                value[r][vi] = get(call_id);
+            }
+        }
+    }
+
+    public void store<let R : int, let C: int>(ContextND<N - 2> context, in matrix<T, R, C> value) {
+        int call_id[N];
+        for (int i = 0; i < N - 2; i++) { call_id[i] = context.call_id[i]; }
+
+        for (int r = 0; r < R; r++)
+        {
+            call_id[N - 2] = r;
+            for (int vi = 0; vi < C; vi++) {
+                call_id[N - 1] = vi;
+                set(call_id, value[r][vi]);
+            }
+        }
+    }
 }
 
 namespace impl

slangpy/tests/slangpy_tests/helpers.py

@@ -135,7 +135,10 @@ def read_ndbuffer_from_numpy(buffer: NDBuffer) -> np.ndarray:
     data = np.array([])
     shape = np.prod(np.array(buffer.shape))
     for i in range(shape):
-        data = np.append(data, cursor[i].read())
+        element = cursor[i].read()
+        if cursor.element_type_layout.kind == TypeReflection.Kind.matrix:
+            element = element.to_numpy()
+        data = np.append(data, element)
 
     return data
 
@@ -149,13 +152,19 @@ def write_ndbuffer_from_numpy(buffer: NDBuffer, data: np.ndarray, element_count:
             element_count = 1
         elif cursor.element_type_layout.kind == TypeReflection.Kind.vector:
             element_count = cursor.element_type.col_count
+        elif cursor.element_type_layout.kind == TypeReflection.Kind.matrix:
+            element_count = cursor.element_type.row_count * cursor.element_type.col_count
         else:
             raise ValueError(
                 f"element_count not set and type is not scalar or vector: {cursor.element_type_layout.kind}"
             )
 
     for i in range(shape):
         buffer_data = np.array(data[i * element_count : (i + 1) * element_count])
+        if cursor.element_type_layout.kind == TypeReflection.Kind.matrix:
+            buffer_data = buffer_data.reshape(
+                cursor.element_type.row_count, cursor.element_type.col_count
+            )
         cursor[i].write(buffer_data)
 
     cursor.apply()

slangpy/tests/slangpy_tests/test_numpy.py

@@ -37,6 +37,10 @@
     }
     return res;
 }
+
+matrix<float, R, C> matFunc1<int R, int C>(matrix<float, R, C> input){
+    return input + 2.0f;
+}
 """
 
 
@@ -161,8 +165,8 @@ def test_return_numpy_matrix(device_type: DeviceType):
 
     module = load_test_module(device_type)
 
-    for R in range(2, 4):
-        for C in range(2, 4):
+    for R in range(2, 5):
+        for C in range(2, 5):
             funName = f"matFunc<{R}, {C}>"
             func = module.find_function(funName)
             assert func is not None
@@ -177,8 +181,8 @@ def test_return_numpy_matrix(device_type: DeviceType):
 def test_setup_numpy_matrix(device_type: DeviceType):
 
     module = load_test_module(device_type)
-    for R in range(2, 4):
-        for C in range(2, 4):
+    for R in range(2, 5):
+        for C in range(2, 5):
             funName = f"flattenMatrix<{R}, {C}>"
             func = module.find_function(funName)
             assert func is not None
@@ -189,5 +193,25 @@ def test_setup_numpy_matrix(device_type: DeviceType):
             assert np.allclose(res["data"][0 : R * C], np.ones(R * C))
 
 
+# test that we can use "numpy" as the result type for a function that returns a matrix
+@pytest.mark.parametrize("device_type", helpers.DEFAULT_DEVICE_TYPES)
+def test_numpy_matrix_as_result(device_type: DeviceType):
+
+    module = load_test_module(device_type)
+
+    for R in range(2, 5):
+        for C in range(2, 5):
+            funName = f"matFunc1<{R}, {C}>"
+            func = module.find_function(funName)
+            assert func is not None
+            matType = getattr(spy, f"float{R}x{C}")
+            N = R * C
+            res = func(matType(np.arange(1, N + 1).reshape(R, C)), _result="numpy")
+
+            assert res is not None
+            assert res.shape == (R, C)
+            assert np.allclose(res, (np.arange(1, N + 1) + 2.0).reshape(R, C))
+
+
 if __name__ == "__main__":
     pytest.main([__file__, "-v", "-s"])

slangpy/tests/slangpy_tests/test_tensor_with_grads.py

@@ -9,13 +9,6 @@
 import sys
 
 
-if sys.platform == "darwin":
-    pytest.skip(
-        "Skipping on macOS: Waiting for slang-gfx fix for resource clear API https://github.com/shader-slang/slang/issues/6640",
-        allow_module_level=True,
-    )
-
-
 def get_test_tensors(device: Device, N: int = 4):
     np.random.seed(0)