ncca-ngl 0.3.4__py3-none-any.whl → 0.5.0__py3-none-any.whl

ncca/ngl/webgpu/base_webgpu_pipeline.py

@@ -0,0 +1,354 @@
+"""
+Abstract base classes for WebGPU rendering pipelines.
+Provides common functionality for buffer management, pipeline creation, and rendering.
+"""
+
+from abc import ABC, abstractmethod
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+import numpy as np
+import wgpu
+
+from .webgpu_constants import NGLToWebGPU
+
+
+class BaseWebGPUPipeline(ABC):
+    """
+    Abstract base class for all WebGPU rendering pipelines.
+
+    Provides common functionality for:
+    - Buffer management and creation
+    - Pipeline configuration
+    - Uniform buffer handling
+    - Resource cleanup
+    """
+
+    def __init__(
+        self,
+        device: wgpu.GPUDevice,
+        texture_format: wgpu.TextureFormat = wgpu.TextureFormat.rgba8unorm,
+        depth_format: wgpu.TextureFormat = wgpu.TextureFormat.depth24plus,
+        msaa_sample_count: int = 4,
+        data_type: str = "Vec3",
+        stride: int = 0,
+    ):
+        """
+        Initialize base pipeline.
+
+        Args:
+            device: WebGPU device
+            texture_format: Colour attachment format
+            depth_format: Depth attachment format
+            msaa_sample_count: Number of MSAA samples
+            data_type: Vertex data type (e.g., "Vec3", "Vec2")
+            stride: Vertex buffer stride. If 0, inferred from data_type
+        """
+        self.device = device
+        self.texture_format = texture_format
+        self.depth_format = depth_format
+        self.msaa_sample_count = msaa_sample_count
+        self._data_type = data_type
+
+        if stride != 0:
+            self._stride = stride
+        else:
+            self._stride = NGLToWebGPU.stride_from_type(self._data_type)
+
+        # Core pipeline resources
+        self.pipeline: Optional[wgpu.GPURenderPipeline] = None
+        self.uniform_buffer: Optional[wgpu.GPUBuffer] = None
+        self.bind_group: Optional[wgpu.GPUBindGroup] = None
+
+        # Initialize uniform data structure
+        self.uniform_data = np.zeros((), dtype=self.get_dtype())
+        self._set_default_uniforms()
+
+        # Create the pipeline
+        self._create_pipeline()
+
+    @abstractmethod
+    def get_dtype(self) -> np.dtype:
+        """Get the numpy dtype for the uniform buffer structure."""
+        pass
+
+    @abstractmethod
+    def _get_shader_code(self) -> str:
+        """Get the WGSL shader code for this pipeline."""
+        pass
+
+    @abstractmethod
+    def _get_vertex_buffer_layouts(self) -> List[Dict[str, Any]]:
+        """Get vertex buffer layout configurations for the pipeline."""
+        pass
+
+    @abstractmethod
+    def _get_primitive_topology(self) -> wgpu.PrimitiveTopology:
+        """Get the primitive topology for the pipeline."""
+        pass
+
+    @abstractmethod
+    def _set_default_uniforms(self) -> None:
+        """Set default values for uniform data."""
+        pass
+
+    @abstractmethod
+    def _get_pipeline_label(self) -> str:
+        """Get the label for the pipeline."""
+        pass
+
+    def _create_pipeline(self) -> None:
+        """Create the render pipeline and associated resources."""
+        # Load shader
+        shader_module = self.device.create_shader_module(code=self._get_shader_code())
+
+        # Create render pipeline
+        self.pipeline = self.device.create_render_pipeline(
+            label=self._get_pipeline_label(),
+            layout="auto",
+            vertex={
+                "module": shader_module,
+                "entry_point": "vertex_main",
+                "buffers": self._get_vertex_buffer_layouts(),
+            },
+            fragment={
+                "module": shader_module,
+                "entry_point": "fragment_main",
+                "targets": [{"format": self.texture_format}],
+            },
+            primitive={"topology": self._get_primitive_topology()},
+            depth_stencil={
+                "format": self.depth_format,
+                "depth_write_enabled": True,
+                "depth_compare": wgpu.CompareFunction.less,
+            },
+            multisample={"count": self.msaa_sample_count},
+        )
+
+        # Create uniform buffer
+        self.uniform_buffer = self.device.create_buffer_with_data(
+            data=self.uniform_data.tobytes(),
+            usage=int(wgpu.BufferUsage.UNIFORM | wgpu.BufferUsage.COPY_DST),
+            label=f"{self._get_pipeline_label()}_uniform_buffer",
+        )
+
+        # Create bind group
+        bind_group_layout = self.pipeline.get_bind_group_layout(0)
+        self.bind_group = self.device.create_bind_group(
+            layout=bind_group_layout,
+            entries=[
+                {
+                    "binding": 0,
+                    "resource": {"buffer": self.uniform_buffer},
+                }
+            ],
+        )
+
+    def _create_or_update_buffer(
+        self,
+        current_buffer: Optional[wgpu.GPUBuffer],
+        data: Union[np.ndarray, wgpu.GPUBuffer],
+        usage: wgpu.BufferUsage,
+        buffer_label: str,
+    ) -> Tuple[Optional[wgpu.GPUBuffer], int]:
+        """
+        Create or update a GPU buffer with new data.
+
+        Args:
+            current_buffer: Existing buffer (may be None)
+            data: New data (numpy array or GPU buffer)
+            usage: Buffer usage flags
+            buffer_label: Label for the buffer
+
+        Returns:
+            Tuple of (buffer, data_size)
+        """
+        if isinstance(data, wgpu.GPUBuffer):
+            # Use provided buffer directly
+            return data, data.size
+
+        # Handle numpy array
+        data_bytes = data.astype(np.float32).tobytes()
+        data_size = len(data_bytes)
+
+        # Create new buffer if needed or existing one is too small
+        if current_buffer is None or current_buffer.size < data_size:
+            if current_buffer:
+                current_buffer.destroy()
+            buffer = self.device.create_buffer_with_data(
+                data=data_bytes,
+                usage=usage,
+                label=buffer_label,
+            )
+            return buffer, data_size
+        else:
+            # Update existing buffer
+            self.device.queue.write_buffer(current_buffer, 0, data_bytes)
+            return current_buffer, data_size
+
+    def _process_vertex_data(
+        self,
+        data: Optional[Union[np.ndarray, wgpu.GPUBuffer]],
+        default_value: Optional[np.ndarray] = None,
+        padding_size: Optional[int] = None,
+        buffer_label: str = "vertex_buffer",
+    ) -> Optional[Union[wgpu.GPUBuffer, Tuple[wgpu.GPUBuffer, int]]]:
+        """
+        Process vertex data, handling numpy arrays, GPU buffers, and defaults.
+
+        Args:
+            data: Input data (numpy array, GPU buffer, or None)
+            default_value: Default value if data is None
+            padding_size: Size to pad arrays to (for alignment)
+            buffer_label: Label for created buffers
+
+        Returns:
+            Processed buffer(s) or None
+        """
+        if data is None and default_value is not None:
+            data = default_value
+
+        if data is None:
+            return None
+
+        if isinstance(data, wgpu.GPUBuffer):
+            return data
+
+        # Handle numpy array
+        if padding_size:
+            # Pad array to specified size
+            if data.ndim == 1:
+                padded_data = np.zeros(padding_size, dtype=np.float32)
+                padded_data[: len(data)] = data.astype(np.float32)
+            else:
+                padded_data = np.zeros((data.shape[0], padding_size), dtype=np.float32)
+                padded_data[:, : data.shape[1]] = data.astype(np.float32)
+            data = padded_data
+
+        buffer, _ = self._create_or_update_buffer(
+            None,  # Always create new for processed data
+            data,
+            wgpu.BufferUsage.VERTEX | wgpu.BufferUsage.COPY_DST,
+            buffer_label,
+        )
+        return buffer
+
+    @abstractmethod
+    def set_data(self, **kwargs) -> None:
+        """
+        Set rendering data (vertices, colours, etc.).
+
+        Args:
+            **kwargs: Pipeline-specific data parameters
+        """
+        pass
+
+    @abstractmethod
+    def update_uniforms(self, **kwargs) -> None:
+        """
+        Update uniform buffer values.
+
+        Args:
+            **kwargs: Pipeline-specific uniform parameters
+        """
+        pass
+
+    @abstractmethod
+    def render(self, render_pass: wgpu.GPURenderPassEncoder, **kwargs) -> None:
+        """
+        Render using this pipeline.
+
+        Args:
+            render_pass: Active render pass encoder
+            **kwargs: Pipeline-specific render parameters
+        """
+        pass
+
+    def cleanup(self) -> None:
+        """Release pipeline resources. Can be overridden for additional cleanup."""
+        if self.uniform_buffer:
+            self.uniform_buffer.destroy()
+
+
+class BasePointPipeline(BaseWebGPUPipeline):
+    """
+    Base class for point rendering pipelines.
+
+    Provides common functionality for:
+    - Point billboarding
+    - Quad generation
+    - Circle clipping in fragment shader
+    """
+
+    def _get_primitive_topology(self) -> wgpu.PrimitiveTopology:
+        """Points are rendered as triangle strips for quad generation."""
+        return wgpu.PrimitiveTopology.triangle_strip
+
+    def _get_default_vertex_layouts(self, has_colour_buffer: bool = False) -> List[Dict[str, Any]]:
+        """
+        Get default vertex buffer layouts for point rendering.
+
+        Args:
+            has_colour_buffer: Whether to include colour buffer layout
+
+        Returns:
+            List of vertex buffer layout configurations
+        """
+        layouts = [
+            {
+                "array_stride": self._stride,
+                "step_mode": "instance",
+                "attributes": [
+                    {
+                        "format": NGLToWebGPU.vertex_format(self._data_type),
+                        "offset": 0,
+                        "shader_location": 0,
+                    },
+                ],
+            },
+        ]
+
+        if has_colour_buffer:
+            layouts.append({
+                "array_stride": NGLToWebGPU.stride_from_type("Vec3"),
+                "step_mode": "instance",
+                "attributes": [
+                    {
+                        "format": NGLToWebGPU.vertex_format("Vec3"),
+                        "offset": 0,
+                        "shader_location": 1,
+                    },
+                ],
+            })
+
+        return layouts
+
+    def _render_points(
+        self,
+        render_pass: wgpu.GPURenderPassEncoder,
+        position_buffer: wgpu.GPUBuffer,
+        colour_buffer: Optional[wgpu.GPUBuffer] = None,
+        num_points: Optional[int] = None,
+    ) -> None:
+        """
+        Common point rendering implementation.
+
+        Args:
+            render_pass: Active render pass encoder
+            position_buffer: Buffer containing point positions
+            colour_buffer: Optional buffer containing point colours
+            num_points: Number of points to render
+        """
+        if position_buffer is None:
+            return
+
+        count = num_points if num_points is not None else getattr(self, "num_points", 0)
+
+        render_pass.set_pipeline(self.pipeline)
+        render_pass.set_bind_group(0, self.bind_group, [], 0, 999999)
+        render_pass.set_vertex_buffer(0, position_buffer)
+
+        if colour_buffer:
+            render_pass.set_vertex_buffer(1, colour_buffer)
+
+        # 4 vertices per quad for point rendering
+        render_pass.draw(4, count)

ncca/ngl/webgpu/custom_shader_pipeline.py

@@ -0,0 +1,288 @@
+"""
+Custom WebGPU pipeline that accepts user-provided shader source files.
+Provides flexibility for users to write their own WGSL shaders.
+"""
+
+import os
+from typing import Dict, List, Optional, Union
+
+import numpy as np
+import wgpu
+
+from .base_webgpu_pipeline import BaseWebGPUPipeline
+from .webgpu_constants import NGLToWebGPU
+
+
+class CustomShaderPipeline(BaseWebGPUPipeline):
+    """
+    A WebGPU pipeline that uses custom shader source provided by the user.
+
+    This pipeline allows users to provide their own WGSL shader source code
+    while handling the boilerplate for buffer management, uniform updates,
+    and rendering setup.
+    """
+
+    def __init__(
+        self,
+        device: wgpu.GPUDevice,
+        shader_source: str,
+        vertex_formats: Optional[List[Union[str, wgpu.VertexFormat]]] = None,
+        primitive_topology: wgpu.PrimitiveTopology = wgpu.PrimitiveTopology.triangle_list,
+        texture_format: wgpu.TextureFormat = wgpu.TextureFormat.rgba8unorm,
+        depth_format: wgpu.TextureFormat = wgpu.TextureFormat.depth24plus,
+        msaa_sample_count: int = 4,
+        uniform_struct_definition: Optional[str] = None,
+        pipeline_label: str = "CustomShaderPipeline",
+    ):
+        """
+        Initialize custom shader pipeline.
+
+        Args:
+            device: WebGPU device
+            shader_source: WGSL shader source code as string
+            vertex_formats: List of vertex data formats (e.g., ["Vec3", "Vec3"] for position+colour)
+            primitive_topology: Primitive topology for rendering
+            texture_format: Colour attachment format
+            depth_format: Depth attachment format
+            msaa_sample_count: Number of MSAA samples
+            uniform_struct_definition: Custom uniform struct definition (optional)
+            pipeline_label: Label for debugging
+        """
+        self._shader_source = shader_source
+        self._vertex_formats = vertex_formats or ["Vec3"]
+        self._primitive_topology = primitive_topology
+        self._uniform_struct_definition = uniform_struct_definition
+        self._pipeline_label = pipeline_label
+
+        # Calculate stride from vertex formats
+        self._stride = sum(
+            NGLToWebGPU.stride_from_type(fmt) for fmt in self._vertex_formats
+        )
+
+        # Initialize buffers storage
+        self.vertex_buffers: Dict[int, wgpu.GPUBuffer] = {}
+        self.num_vertices = 0
+
+        # Call parent constructor after setting up our attributes
+        super().__init__(
+            device=device,
+            texture_format=texture_format,
+            depth_format=depth_format,
+            msaa_sample_count=msaa_sample_count,
+            data_type=self._vertex_formats[0],  # Use first format as default
+            stride=self._stride,
+        )
+
+    def get_dtype(self) -> np.dtype:
+        """Get the numpy dtype for the uniform buffer structure."""
+        if self._uniform_struct_definition:
+            # For custom uniforms, we need to parse the struct or use a default
+            # For now, use a basic MVP + colour structure
+            return np.dtype(
+                [
+                    ("MVP", np.float32, (4, 4)),
+                    ("colour", np.float32, 4),
+                ]
+            )
+        else:
+            # Default uniform structure
+            return np.dtype(
+                [
+                    ("MVP", np.float32, (4, 4)),
+                    ("colour", np.float32, 4),
+                ]
+            )
+
+    def _get_shader_code(self) -> str:
+        """Return the custom shader source."""
+        return self._shader_source
+
+    def _get_vertex_buffer_layouts(self) -> List[Dict[str, any]]:
+        """Get vertex buffer layouts based on provided vertex formats."""
+        if len(self._vertex_formats) == 1:
+            # Single interleaved buffer
+            return [
+                {
+                    "array_stride": self._stride,
+                    "step_mode": "vertex",
+                    "attributes": [
+                        {
+                            "format": NGLToWebGPU.vertex_format(
+                                self._vertex_formats[0]
+                            ),
+                            "offset": 0,
+                            "shader_location": 0,
+                        },
+                    ],
+                },
+            ]
+        else:
+            # Multiple separate buffers or interleaved with multiple attributes
+            layouts = []
+            current_offset = 0
+
+            for i, fmt in enumerate(self._vertex_formats):
+                stride = NGLToWebGPU.stride_from_type(fmt)
+                layouts.append(
+                    {
+                        "array_stride": stride,
+                        "step_mode": "vertex",
+                        "attributes": [
+                            {
+                                "format": NGLToWebGPU.vertex_format(fmt),
+                                "offset": 0,
+                                "shader_location": i,
+                            },
+                        ],
+                    }
+                )
+                current_offset += stride
+
+            return layouts
+
+    def _get_primitive_topology(self) -> wgpu.PrimitiveTopology:
+        """Return the primitive topology."""
+        return self._primitive_topology
+
+    def _set_default_uniforms(self) -> None:
+        """Set default uniform values."""
+        self.uniform_data["MVP"] = np.eye(4, dtype=np.float32)
+        self.uniform_data["colour"] = np.array([1.0, 1.0, 1.0, 1.0], dtype=np.float32)
+
+    def _get_pipeline_label(self) -> str:
+        """Return the pipeline label."""
+        return self._pipeline_label
+
+    def set_data(
+        self,
+        positions: Optional[np.ndarray] = None,
+        colours: Optional[np.ndarray] = None,
+        interleaved_data: Optional[np.ndarray] = None,
+        **kwargs,
+    ) -> None:
+        """
+        Set vertex data for rendering.
+
+        Args:
+            positions: Vertex position data (N, 3)
+            colours: Vertex colour data (N, 3) or (N, 4)
+            interleaved_data: Pre-interleaved vertex data
+            **kwargs: Additional vertex data arrays (e.g., velocities, life, initial_position)
+        """
+        if interleaved_data is not None:
+            # Use pre-interleaved data
+            buffer = self._create_or_update_buffer(
+                self.vertex_buffers.get(0),
+                interleaved_data,
+                wgpu.BufferUsage.VERTEX | wgpu.BufferUsage.COPY_DST,
+                f"{self._pipeline_label}_vertex_buffer_0",
+            )
+            self.vertex_buffers[0] = buffer[0]
+            self.num_vertices = len(interleaved_data)
+        else:
+            # Handle separate vertex data arrays
+            binding = 0
+
+            if positions is not None:
+                buffer = self._create_or_update_buffer(
+                    self.vertex_buffers.get(binding),
+                    positions,
+                    wgpu.BufferUsage.VERTEX | wgpu.BufferUsage.COPY_DST,
+                    f"{self._pipeline_label}_vertex_buffer_{binding}",
+                )
+                self.vertex_buffers[binding] = buffer[0]
+                self.num_vertices = len(positions)
+                binding += 1
+
+            if colours is not None:
+                buffer = self._create_or_update_buffer(
+                    self.vertex_buffers.get(binding),
+                    colours,
+                    wgpu.BufferUsage.VERTEX | wgpu.BufferUsage.COPY_DST,
+                    f"{self._pipeline_label}_vertex_buffer_{binding}",
+                )
+                self.vertex_buffers[binding] = buffer[0]
+                binding += 1
+
+            # Handle additional vertex attributes (velocities, life, initial_position, etc.)
+            for attr_name, attr_data in kwargs.items():
+                if attr_data is not None:
+                    buffer = self._create_or_update_buffer(
+                        self.vertex_buffers.get(binding),
+                        attr_data,
+                        wgpu.BufferUsage.VERTEX | wgpu.BufferUsage.COPY_DST,
+                        f"{self._pipeline_label}_vertex_buffer_{binding}",
+                    )
+                    self.vertex_buffers[binding] = buffer[0]
+                    binding += 1
+
+    def update_uniforms(self, **kwargs) -> None:
+        """
+        Update uniform buffer values.
+
+        Args:
+            **kwargs: Uniform values to update (e.g., mvp=matrix, colour=array)
+        """
+        if "mvp" in kwargs:
+            self.uniform_data["MVP"] = kwargs["mvp"]
+
+        if "colour" in kwargs:
+            colour = kwargs["colour"]
+            if len(colour) == 3:
+                self.uniform_data["colour"] = np.array([*colour, 1.0], dtype=np.float32)
+            else:
+                self.uniform_data["colour"] = np.array(colour, dtype=np.float32)
+
+        # Update the GPU buffer
+        if self.uniform_buffer:
+            self.device.queue.write_buffer(
+                self.uniform_buffer, 0, self.uniform_data.tobytes()
+            )
+
+    def render(self, render_pass: wgpu.GPURenderPassEncoder, **kwargs) -> None:
+        """
+        Render using this pipeline.
+
+        Args:
+            render_pass: Active render pass encoder
+            **kwargs: Additional render parameters
+        """
+        if self.pipeline is None or self.num_vertices == 0:
+            return
+
+        render_pass.set_pipeline(self.pipeline)
+        render_pass.set_bind_group(0, self.bind_group, [], 0, 999999)
+
+        # Set vertex buffers
+        for binding, buffer in self.vertex_buffers.items():
+            render_pass.set_vertex_buffer(binding, buffer)
+
+        # Draw
+        render_pass.draw(self.num_vertices)
+
+    @classmethod
+    def from_file(
+        cls, device: wgpu.GPUDevice, shader_file: str, **kwargs
+    ) -> "CustomShaderPipeline":
+        """
+        Create a CustomShaderPipeline from a WGSL file.
+
+        Args:
+            device: WebGPU device
+            shader_file: Path to WGSL shader file
+            **kwargs: Additional arguments to pass to constructor
+
+        Returns:
+            CustomShaderPipeline instance
+        """
+        if not os.path.exists(shader_file):
+            raise FileNotFoundError(f"Shader file not found: {shader_file}")
+
+        with open(shader_file, "r", encoding="utf-8") as f:
+            shader_source = f.read()
+
+        # Use filename as default pipeline label if not provided
+        if "pipeline_label" not in kwargs:
+            kwargs["pipeline_label"] = f"Custom_{os.path.basename(shader_file)}"
+
+        return cls(device, shader_source, **kwargs)