Slangpy numpy marshall + cleanup (#155)

* Initial work to get native slangpy running again

* All slangpy tests now working

* Start optimizations to shift slangpy to using shadercursor for uniforms

* NativeType->NativeMarshall

* Started on native value and buffer marshalling

* Native NDBuffer largely working

* Handle broadcasting+output alloc for ndbuffer

* Raw dispatch for ndbuffer

* NDBuffer has to/from numpy functions

* Pre PR cleanup

* Numpy marshall for slangpy

* Slangpy shape uses size_t

* Replace shape size_t with int64

* Revert "Replace shape size_t with int64"

This reverts commit 6740ea6.

* Revert "Slangpy shape uses size_t"

This reverts commit 83b4e12.

src/sgl/python/sgl_ext.cpp

@@ -55,6 +55,8 @@ SGL_PY_DECLARE(ui_widgets);
 
 SGL_PY_DECLARE(utils_renderdoc);
 SGL_PY_DECLARE(utils_slangpy);
+SGL_PY_DECLARE(utils_slangpy_buffer);
+SGL_PY_DECLARE(utils_slangpy_value);
 SGL_PY_DECLARE(utils_tev);
 SGL_PY_DECLARE(utils_texture_loader);
 
@@ -129,8 +131,12 @@ NB_MODULE(sgl_ext, m_)
 
     m.def_submodule("renderdoc", "RenderDoc module");
     SGL_PY_IMPORT(utils_renderdoc);
+
     m.def_submodule("slangpy", "SlangPy module");
     SGL_PY_IMPORT(utils_slangpy);
+    SGL_PY_IMPORT(utils_slangpy_buffer);
+    SGL_PY_IMPORT(utils_slangpy_value);
+
     m.def_submodule("tev", "tev image viewer module");
     SGL_PY_IMPORT(utils_tev);
     SGL_PY_IMPORT(utils_texture_loader);

src/sgl/utils/python/slangpy.cpp

@@ -39,9 +39,18 @@ void NativeMarshall::write_shader_cursor_pre_dispatch(
     if (!cd_val.is_none()) {
         ShaderCursor child_field = cursor[binding->get_variable_name()];
         write_shader_cursor(child_field, cd_val);
-        read_back.append(nb::make_tuple(binding, value, cd_val));
+        store_readback(binding, read_back, value, cd_val);
     }
 }
+void NativeMarshall::store_readback(
+    NativeBoundVariableRuntime* binding,
+    nb::list& read_back,
+    nb::object value,
+    nb::object data
+) const
+{
+    read_back.append(nb::make_tuple(binding, value, data));
+}
 
 void NativeBoundVariableRuntime::populate_call_shape(std::vector<int>& call_shape, nb::object value)
 {
@@ -612,7 +621,8 @@ SGL_PY_EXPORT(utils_slangpy)
             &NativeSlangType::get_type_reflection,
             &NativeSlangType::set_type_reflection,
             D_NA(NativeSlangType, type_reflection)
-        );
+        )
+        .def_prop_rw("shape", &NativeSlangType::get_shape, &NativeSlangType::set_shape, D_NA(NativeSlangType, shape));
 
     nb::class_<NativeMarshall, PyNativeMarshall, Object>(slangpy, "NativeMarshall") //
         .def(
@@ -648,34 +658,6 @@ SGL_PY_EXPORT(utils_slangpy)
         .def("create_output", &NativeMarshall::create_output, D_NA(NativeMarshall, create_output))
         .def("read_output", &NativeMarshall::read_output, D_NA(NativeMarshall, read_output));
 
-    nb::class_<NativeValueMarshall, NativeMarshall>(slangpy, "NativeValueMarshall") //
-        .def(
-            "__init__",
-            [](NativeValueMarshall& self) { new (&self) NativeValueMarshall(); },
-            D_NA(NativeValueMarshall, NativeValueMarshall)
-        );
-
-    nb::class_<NativeNDBufferMarshall, PyNativeNDBufferMarshall, NativeMarshall>(slangpy, "NativeNDBufferMarshall") //
-        .def(
-            "__init__",
-            [](NativeNDBufferMarshall& self,
-               int dims,
-               bool writable,
-               ref<NativeSlangType> slang_type,
-               ref<NativeSlangType> slang_element_type,
-               int element_stride)
-            { new (&self) PyNativeNDBufferMarshall(dims, writable, slang_type, slang_element_type, element_stride); },
-            "dims"_a,
-            "writable"_a,
-            "slang_type"_a,
-            "slang_element_type"_a,
-            "element_stride"_a,
-            D_NA(NativeNDBufferMarshall, NativeNDBufferMarshall)
-        )
-        .def_prop_ro("dims", &sgl::slangpy::NativeNDBufferMarshall::dims)
-        .def_prop_ro("writable", &sgl::slangpy::NativeNDBufferMarshall::writable)
-        .def_prop_ro("slang_element_type", &sgl::slangpy::NativeNDBufferMarshall::slang_element_type);
-
     nb::class_<NativeBoundVariableRuntime, Object>(slangpy, "NativeBoundVariableRuntime") //
         .def(nb::init<>(), D_NA(NativeBoundVariableRuntime, NativeBoundVariableRuntime))
         .def_prop_rw(
@@ -921,36 +903,4 @@ SGL_PY_EXPORT(utils_slangpy)
             D_NA(CallContext, call_shape)
         )
         .def_prop_ro("call_mode", &CallContext::call_mode, D_NA(CallContext, call_mode));
-
-    nb::class_<NativeNDBufferDesc>(slangpy, "NativeNDBufferDesc")
-        .def(nb::init<>())
-        .def_rw("dtype", &NativeNDBufferDesc::dtype)
-        .def_rw("element_stride", &NativeNDBufferDesc::element_stride)
-        .def_rw("shape", &NativeNDBufferDesc::shape)
-        .def_rw("strides", &NativeNDBufferDesc::strides)
-        .def_rw("usage", &NativeNDBufferDesc::usage)
-        .def_rw("memory_type", &NativeNDBufferDesc::memory_type);
-
-    nb::class_<NativeNDBuffer, Object>(slangpy, "NativeNDBuffer")
-        .def(nb::init<ref<Device>, NativeNDBufferDesc>())
-        .def_prop_ro("device", &NativeNDBuffer::device)
-        .def_prop_rw("slangpy_signature", &NativeNDBuffer::slangpy_signature, &NativeNDBuffer::set_slagpy_signature)
-        .def_prop_ro("dtype", &NativeNDBuffer::dtype)
-        .def_prop_ro("shape", &NativeNDBuffer::shape)
-        .def_prop_ro("strides", &NativeNDBuffer::strides)
-        .def_prop_ro("element_count", &NativeNDBuffer::element_count)
-        .def_prop_ro("usage", &NativeNDBuffer::usage)
-        .def_prop_ro("memory_type", &NativeNDBuffer::memory_type)
-        .def_prop_ro("storage", &NativeNDBuffer::storage)
-        .def(
-            "to_numpy",
-            [](NativeNDBuffer& self) { return buffer_to_numpy(self.storage().get()); },
-            D_NA(NativeNDBuffer, buffer_to_numpy)
-        )
-        .def(
-            "from_numpy",
-            [](NativeNDBuffer& self, nb::ndarray<nb::numpy> data) { buffer_from_numpy(self.storage().get(), data); },
-            "data"_a,
-            D_NA(NativeNDBuffer, buffer_from_numpy)
-        );
 }

src/sgl/utils/python/slangpy.h

@@ -53,8 +53,15 @@ class NativeSlangType : public Object {
     /// Set the reflection type.
     void set_type_reflection(const ref<TypeReflection>& reflection) { m_type_reflection = reflection; }
 
+    /// Get the shape of the type.
+    Shape get_shape() const { return m_shape; }
+
+    /// Set the shape of the type.
+    void set_shape(const Shape& shape) { m_shape = shape; }
+
 private:
     ref<TypeReflection> m_type_reflection;
+    Shape m_shape;
 };
 
 /// Base class for a marshal to a slangpy supported type.
@@ -145,6 +152,10 @@ class NativeMarshall : public Object {
         return nb::none();
     };
 
+protected:
+    void
+    store_readback(NativeBoundVariableRuntime* binding, nb::list& read_back, nb::object value, nb::object data) const;
+
 private:
     Shape m_concrete_shape;
     ref<NativeSlangType> m_slang_type;