warp-lang 1.2.2__py3-none-win_amd64.whl → 1.3.1__py3-none-win_amd64.whl

warp/dlpack.py

@@ -60,10 +60,18 @@ PyCapsule_SetName.restype = ctypes.c_int
 class _DLPackTensorHolder:
     """Class responsible for deleting DLManagedTensor memory after ownership is transferred from a capsule."""
 
+    def __new__(cls, *args, **kwargs):
+        instance = super(_DLPackTensorHolder, cls).__new__(cls)
+        instance.mem_ptr = None
+        return instance
+
     def __init__(self, mem_ptr):
         self.mem_ptr = mem_ptr
 
     def __del__(self):
+        if not self.mem_ptr:
+            return
+
         managed_tensor = DLManagedTensor.from_address(self.mem_ptr)
         if managed_tensor.deleter:
             managed_tensor.deleter(self.mem_ptr)

warp/examples/benchmarks/benchmark_cloth_warp.py

@@ -7,7 +7,7 @@
 
 import warp as wp
 
-wp.build.clear_kernel_cache()
+wp.clear_kernel_cache()
 
 
 @wp.kernel

warp/examples/benchmarks/benchmark_interop_torch.py

@@ -0,0 +1,158 @@
+# Copyright (c) 2024 NVIDIA CORPORATION.  All rights reserved.
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+import time
+
+import torch
+
+import warp as wp
+
+
+def create_simple_kernel(dtype):
+    def simple_kernel(
+        a: wp.array(dtype=dtype),
+        b: wp.array(dtype=dtype),
+        c: wp.array(dtype=dtype),
+        d: wp.array(dtype=dtype),
+        e: wp.array(dtype=dtype),
+    ):
+        pass
+
+    return wp.Kernel(simple_kernel)
+
+
+def test_from_torch(kernel, num_iters, array_size, device, warp_dtype=None):
+    warp_device = wp.get_device(device)
+    torch_device = wp.device_to_torch(warp_device)
+
+    if hasattr(warp_dtype, "_shape_"):
+        torch_shape = (array_size, *warp_dtype._shape_)
+        torch_dtype = wp.dtype_to_torch(warp_dtype._wp_scalar_type_)
+    else:
+        torch_shape = (array_size,)
+        torch_dtype = torch.float32 if warp_dtype is None else wp.dtype_to_torch(warp_dtype)
+
+    _a = torch.zeros(torch_shape, dtype=torch_dtype, device=torch_device)
+    _b = torch.zeros(torch_shape, dtype=torch_dtype, device=torch_device)
+    _c = torch.zeros(torch_shape, dtype=torch_dtype, device=torch_device)
+    _d = torch.zeros(torch_shape, dtype=torch_dtype, device=torch_device)
+    _e = torch.zeros(torch_shape, dtype=torch_dtype, device=torch_device)
+
+    wp.synchronize()
+
+    # profiler = Profiler(interval=0.000001)
+    # profiler.start()
+
+    t1 = time.time_ns()
+
+    for _ in range(num_iters):
+        a = wp.from_torch(_a, dtype=warp_dtype)
+        b = wp.from_torch(_b, dtype=warp_dtype)
+        c = wp.from_torch(_c, dtype=warp_dtype)
+        d = wp.from_torch(_d, dtype=warp_dtype)
+        e = wp.from_torch(_e, dtype=warp_dtype)
+        wp.launch(kernel, dim=array_size, inputs=[a, b, c, d, e])
+
+    t2 = time.time_ns()
+    print(f"{(t2 - t1) / 1_000_000 :8.0f} ms  from_torch(...)")
+
+    # profiler.stop()
+    # profiler.print()
+
+
+def test_array_ctype_from_torch(kernel, num_iters, array_size, device, warp_dtype=None):
+    warp_device = wp.get_device(device)
+    torch_device = wp.device_to_torch(warp_device)
+
+    if hasattr(warp_dtype, "_shape_"):
+        torch_shape = (array_size, *warp_dtype._shape_)
+        torch_dtype = wp.dtype_to_torch(warp_dtype._wp_scalar_type_)
+    else:
+        torch_shape = (array_size,)
+        torch_dtype = torch.float32 if warp_dtype is None else wp.dtype_to_torch(warp_dtype)
+
+    _a = torch.zeros(torch_shape, dtype=torch_dtype, device=torch_device)
+    _b = torch.zeros(torch_shape, dtype=torch_dtype, device=torch_device)
+    _c = torch.zeros(torch_shape, dtype=torch_dtype, device=torch_device)
+    _d = torch.zeros(torch_shape, dtype=torch_dtype, device=torch_device)
+    _e = torch.zeros(torch_shape, dtype=torch_dtype, device=torch_device)
+
+    wp.synchronize()
+
+    # profiler = Profiler(interval=0.000001)
+    # profiler.start()
+
+    t1 = time.time_ns()
+
+    for _ in range(num_iters):
+        a = wp.from_torch(_a, dtype=warp_dtype, return_ctype=True)
+        b = wp.from_torch(_b, dtype=warp_dtype, return_ctype=True)
+        c = wp.from_torch(_c, dtype=warp_dtype, return_ctype=True)
+        d = wp.from_torch(_d, dtype=warp_dtype, return_ctype=True)
+        e = wp.from_torch(_e, dtype=warp_dtype, return_ctype=True)
+        wp.launch(kernel, dim=array_size, inputs=[a, b, c, d, e])
+
+    t2 = time.time_ns()
+    print(f"{(t2 - t1) / 1_000_000 :8.0f} ms  from_torch(..., return_ctype=True)")
+
+    # profiler.stop()
+    # profiler.print()
+
+
+def test_direct_from_torch(kernel, num_iters, array_size, device, warp_dtype=None):
+    warp_device = wp.get_device(device)
+    torch_device = wp.device_to_torch(warp_device)
+
+    if hasattr(warp_dtype, "_shape_"):
+        torch_shape = (array_size, *warp_dtype._shape_)
+        torch_dtype = wp.dtype_to_torch(warp_dtype._wp_scalar_type_)
+    else:
+        torch_shape = (array_size,)
+        torch_dtype = torch.float32 if warp_dtype is None else wp.dtype_to_torch(warp_dtype)
+
+    _a = torch.zeros(torch_shape, dtype=torch_dtype, device=torch_device)
+    _b = torch.zeros(torch_shape, dtype=torch_dtype, device=torch_device)
+    _c = torch.zeros(torch_shape, dtype=torch_dtype, device=torch_device)
+    _d = torch.zeros(torch_shape, dtype=torch_dtype, device=torch_device)
+    _e = torch.zeros(torch_shape, dtype=torch_dtype, device=torch_device)
+
+    wp.synchronize()
+
+    # profiler = Profiler(interval=0.000001)
+    # profiler.start()
+
+    t1 = time.time_ns()
+
+    for _ in range(num_iters):
+        wp.launch(kernel, dim=array_size, inputs=[_a, _b, _c, _d, _e])
+
+    t2 = time.time_ns()
+    print(f"{(t2 - t1) / 1_000_000 :8.0f} ms  direct from torch")
+
+    # profiler.stop()
+    # profiler.print()
+
+
+wp.init()
+
+params = [
+    # (warp_dtype arg, kernel)
+    (None, create_simple_kernel(wp.float32)),
+    (wp.float32, create_simple_kernel(wp.float32)),
+    (wp.vec3f, create_simple_kernel(wp.vec3f)),
+    (wp.mat22f, create_simple_kernel(wp.mat22f)),
+]
+
+wp.load_module()
+
+num_iters = 100000
+
+for warp_dtype, kernel in params:
+    print(f"\ndtype={wp.context.type_str(warp_dtype)}")
+    test_from_torch(kernel, num_iters, 10, "cuda:0", warp_dtype=warp_dtype)
+    test_array_ctype_from_torch(kernel, num_iters, 10, "cuda:0", warp_dtype=warp_dtype)
+    test_direct_from_torch(kernel, num_iters, 10, "cuda:0", warp_dtype=warp_dtype)

warp/examples/benchmarks/benchmark_launches.py

@@ -128,7 +128,7 @@ def ksz(s: Sz):
     tid = wp.tid()  # noqa: F841
 
 
-wp.build.clear_kernel_cache()
+wp.clear_kernel_cache()
 
 devices = wp.get_devices()
 num_launches = 100000

warp/examples/core/example_cupy.py

@@ -0,0 +1,78 @@
+# Copyright (c) 2024 NVIDIA CORPORATION.  All rights reserved.
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+###########################################################################
+# Example CuPy
+#
+# The example demonstrates interoperability with CuPy on CUDA devices
+# and NumPy on CPU devices.
+###########################################################################
+
+import warp as wp
+
+
+@wp.kernel
+def saxpy(x: wp.array(dtype=float), y: wp.array(dtype=float), a: float):
+    i = wp.tid()
+    y[i] = a * x[i] + y[i]
+
+
+class Example:
+    def __init__(self):
+        device = wp.get_device()
+
+        self.n = 10
+        self.a = 1.0
+
+        if device.is_cuda:
+            # use CuPy arrays on CUDA devices
+            import cupy as cp
+
+            print(f"Using CuPy on device {device}")
+
+            # tell CuPy to use the same device
+            with cp.cuda.Device(device.ordinal):
+                self.x = cp.arange(self.n, dtype=cp.float32)
+                self.y = cp.ones(self.n, dtype=cp.float32)
+        else:
+            # use NumPy arrays on CPU
+            import numpy as np
+
+            print("Using NumPy on CPU")
+
+            self.x = np.arange(self.n, dtype=np.float32)
+            self.y = np.ones(self.n, dtype=np.float32)
+
+    def step(self):
+        # Launch a Warp kernel on the pre-allocated arrays.
+        # When running on a CUDA device, these are CuPy arrays.
+        # When running on the CPU, these are NumPy arrays.
+        #
+        # Note that the arrays can be passed to Warp kernels directly.  Under the hood,
+        # Warp uses the __cuda_array_interface__ and __array_interface__ protocols to
+        # access the data.
+        wp.launch(saxpy, dim=self.n, inputs=[self.x, self.y, self.a])
+
+    def render(self):
+        print(self.y)
+
+
+if __name__ == "__main__":
+    import argparse
+
+    parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
+    parser.add_argument("--device", type=str, default=None, help="Override the default Warp device.")
+    parser.add_argument("--num_frames", type=int, default=10, help="Total number of frames.")
+
+    args = parser.parse_known_args()[0]
+
+    with wp.ScopedDevice(args.device):
+        example = Example()
+
+        for _ in range(args.num_frames):
+            example.step()
+            example.render()

warp/examples/fem/example_apic_fluid.py

@@ -15,17 +15,13 @@
 import numpy as np
 
 import warp as wp
+import warp.examples.fem.utils as fem_example_utils
 import warp.fem as fem
 import warp.sim.render
 from warp.fem import Domain, Field, Sample, at_node, div, grad, integrand
 from warp.sim import Model, State
 from warp.sparse import BsrMatrix, bsr_mm, bsr_mv, bsr_transposed
 
-try:
-    from .bsr_utils import bsr_cg
-except ImportError:
-    from bsr_utils import bsr_cg
-
 
 @wp.func
 def collision_sdf(x: wp.vec3):
@@ -130,7 +126,7 @@ def scale_transposed_divergence_mat(
     tr_divergence_mat_values: wp.array(dtype=wp.mat(shape=(3, 1), dtype=float)),
     inv_fraction_int: wp.array(dtype=float),
 ):
-    # In-place scaling of gradient operator rows wiht inverse mass
+    # In-place scaling of gradient operator rows with inverse mass
 
     u_i = wp.tid()
     block_beg = tr_divergence_mat_offsets[u_i]
@@ -140,22 +136,24 @@ def scale_transposed_divergence_mat(
         tr_divergence_mat_values[b] = tr_divergence_mat_values[b] * inv_fraction_int[u_i]
 
 
-@wp.kernel
-def compute_particle_ijk(positions: wp.array(dtype=wp.vec3), voxel_size: float, ijks: wp.array(dtype=wp.vec3i)):
-    # Index-space coordinates of grid cell containing each particle
-
-    p = wp.tid()
-    pos = positions[p] / voxel_size
-    ijks[p] = wp.vec3i(int(wp.floor(pos[0])), int(wp.floor(pos[1])), int(wp.floor(pos[2])))
-
-
-def solve_incompressibility(divergence_mat: BsrMatrix, inv_volume, pressure, velocity, quiet: bool = False):
+def solve_incompressibility(
+    divergence_mat: BsrMatrix, dirichlet_projector: BsrMatrix, inv_volume, pressure, velocity, quiet: bool = False
+):
     """Solve for divergence-free velocity delta:
 
     delta_velocity = inv_volume * transpose(divergence_mat) * pressure
     divergence_mat * (velocity + delta_velocity) = 0
+    dirichlet_projector * delta_velocity = 0
     """
 
+    # Constraint-free divergence -- computed *before* projection of divergence_mat
+    rhs = wp.empty_like(pressure)
+    bsr_mv(A=divergence_mat, x=velocity, y=rhs, alpha=-1.0)
+
+    # Project matrix to enforce boundary conditions
+    # divergence_matrix -= divergence_matrix * vel_projector
+    bsr_mm(alpha=-1.0, x=divergence_mat, y=dirichlet_projector, z=divergence_mat, beta=1.0)
+
     # Build transposed gradient matrix, scale with inverse fraction
     transposed_divergence_mat = bsr_transposed(divergence_mat)
     wp.launch(
@@ -171,9 +169,7 @@ def solve_incompressibility(divergence_mat: BsrMatrix, inv_volume, pressure, vel
     # For simplicity, assemble Schur complement and solve with CG
     schur = bsr_mm(divergence_mat, transposed_divergence_mat)
 
-    rhs = wp.zeros_like(pressure)
-    bsr_mv(A=divergence_mat, x=velocity, y=rhs, alpha=-1.0, beta=0.0)
-    bsr_cg(schur, b=rhs, x=pressure, quiet=quiet, tol=1.0e-6)
+    fem_example_utils.bsr_cg(schur, b=rhs, x=pressure, quiet=quiet, tol=1.0e-6)
 
     # Apply pressure to velocity
     bsr_mv(A=transposed_divergence_mat, x=pressure, y=velocity, alpha=1.0, beta=1.0)
@@ -249,23 +245,11 @@ class Example:
 
         self.current_frame = self.current_frame + 1
 
-        particle_ijk = wp.empty(self.state_0.particle_count, dtype=wp.vec3i)
-
         with wp.ScopedTimer(f"simulate frame {self.current_frame}", active=True):
             for _s in range(self.sim_substeps):
-                # Compute the voxel coordinates for each particle.
-                # `Volume.allocate_by_voxels` accepts world positions, but allocates
-                # the voxels with the closest origin rather than the enclosing ones
-                # (i.e, it "round"s the positions, while here we eant to "floor" it)
-                wp.launch(
-                    compute_particle_ijk,
-                    dim=particle_ijk.shape,
-                    inputs=[self.state_0.particle_q, self.voxel_size, particle_ijk],
-                )
-
                 # Allocate the voxels and create the warp.fem geometry
                 volume = wp.Volume.allocate_by_voxels(
-                    voxel_points=particle_ijk,
+                    voxel_points=self.state_0.particle_q,
                     voxel_size=self.voxel_size,
                 )
                 grid = fem.Nanogrid(volume)
@@ -354,13 +338,10 @@ class Example:
                     output_dtype=float,
                 )
 
-                # Project matrix to enforce boundary conditions
-                # divergence_matrix -= divergence_matrix * vel_projector
-                bsr_mm(alpha=-1.0, x=divergence_matrix, y=vel_projector, z=divergence_matrix, beta=1.0)
-
                 # Solve unilateral incompressibility
                 solve_incompressibility(
                     divergence_matrix,
+                    vel_projector,
                     inv_volume,
                     pressure_field.dof_values,
                     velocity_field.dof_values,

warp/examples/fem/example_burgers.py

@@ -16,17 +16,8 @@
 ###########################################################################
 
 import warp as wp
+import warp.examples.fem.utils as fem_example_utils
 import warp.fem as fem
-import warp.sparse as sp
-
-# Import example utilities
-# Make sure that works both when imported as module and run as standalone file
-try:
-    from .bsr_utils import invert_diagonal_bsr_mass_matrix
-    from .plot_utils import Plot
-except ImportError:
-    from bsr_utils import invert_diagonal_bsr_mass_matrix
-    from plot_utils import Plot
 
 
 @fem.integrand
@@ -97,7 +88,7 @@ def slope_limiter(domain: fem.Domain, s: fem.Sample, u: fem.Field, dx: wp.vec2):
     # Assumes regular grid topology
 
     center_coords = fem.Coords(0.5, 0.5, 0.0)
-    cell_center = fem.types.make_free_sample(s.element_index, center_coords)
+    cell_center = fem.make_free_sample(s.element_index, center_coords)
     center_pos = domain(cell_center)
 
     u_center = u(cell_center)
@@ -149,19 +140,19 @@ class Example:
         matrix_inertia = fem.integrate(
             vel_mass_form, fields={"u": trial, "v": self._test}, output_dtype=wp.float32, nodal=True
         )
-        self._inv_mass_matrix = sp.bsr_copy(matrix_inertia)
-        invert_diagonal_bsr_mass_matrix(self._inv_mass_matrix)
+        self._inv_mass_matrix = wp.sparse.bsr_copy(matrix_inertia)
+        fem_example_utils.invert_diagonal_bsr_matrix(self._inv_mass_matrix)
 
         # Initial condition
         self.velocity_field = vector_space.make_field()
         fem.interpolate(initial_condition, dest=self.velocity_field)
 
-        # Velocity nor field -- for visualization purposes
+        # Velocity norm field -- for visualization purposes
         self.velocity_norm_field = scalar_space.make_field()
         fem.interpolate(velocity_norm, dest=self.velocity_norm_field, fields={"u": self.velocity_field})
 
-        self.renderer = Plot()
-        self.renderer.add_surface("u_norm", self.velocity_norm_field)
+        self.renderer = fem_example_utils.Plot()
+        self.renderer.add_field("u_norm", self.velocity_norm_field)
 
     def _velocity_delta(self, trial_velocity):
         # Integration on sides
@@ -186,7 +177,7 @@ class Example:
                 alpha=1.0,
                 beta=1.0,
             )
-        return sp.bsr_mv(self._inv_mass_matrix, rhs)
+        return self._inv_mass_matrix @ rhs
 
     def step(self):
         self.current_frame += 1
@@ -223,7 +214,7 @@ class Example:
 
     def render(self):
         self.renderer.begin_frame(time=self.current_frame * self.sim_dt)
-        self.renderer.add_surface("u_norm", self.velocity_norm_field)
+        self.renderer.add_field("u_norm", self.velocity_norm_field)
         self.renderer.end_frame()
 
 

warp/examples/fem/example_convection_diffusion.py

@@ -15,19 +15,9 @@
 ###########################################################################
 
 import warp as wp
+import warp.examples.fem.utils as fem_example_utils
 import warp.fem as fem
 
-# Import example utilities
-# Make sure that works both when imported as module and run as standalone file
-try:
-    from .bsr_utils import bsr_cg
-    from .mesh_utils import gen_trimesh
-    from .plot_utils import Plot
-except ImportError:
-    from bsr_utils import bsr_cg
-    from mesh_utils import gen_trimesh
-    from plot_utils import Plot
-
 
 @fem.integrand
 def initial_condition(domain: fem.Domain, s: fem.Sample):
@@ -94,8 +84,8 @@ class Example:
         self.current_frame = 0
 
         if tri_mesh:
-            positions, tri_vidx = gen_trimesh(res=wp.vec2i(res))
-            geo = fem.Trimesh2D(tri_vertex_indices=tri_vidx, positions=positions)
+            positions, tri_vidx = fem_example_utils.gen_trimesh(res=wp.vec2i(res))
+            geo = fem.Trimesh2D(tri_vertex_indices=tri_vidx, positions=positions, build_bvh=True)
         else:
             geo = fem.Grid2D(res=wp.vec2i(res))
 
@@ -116,8 +106,8 @@ class Example:
             output_dtype=float,
         )
 
-        self.renderer = Plot()
-        self.renderer.add_surface("phi", self._phi_field)
+        self.renderer = fem_example_utils.Plot()
+        self.renderer.add_field("phi", self._phi_field)
 
     def step(self):
         self.current_frame += 1
@@ -131,11 +121,11 @@ class Example:
         )
 
         # Solve linear system
-        bsr_cg(self._matrix, x=self._phi_field.dof_values, b=rhs, quiet=self._quiet, tol=1.0e-12)
+        fem_example_utils.bsr_cg(self._matrix, x=self._phi_field.dof_values, b=rhs, quiet=self._quiet, tol=1.0e-12)
 
     def render(self):
         self.renderer.begin_frame(time=self.current_frame * self.sim_dt)
-        self.renderer.add_surface("phi", self._phi_field)
+        self.renderer.add_field("phi", self._phi_field)
         self.renderer.end_frame()
 
 

warp/examples/fem/example_convection_diffusion_dg.py

@@ -15,31 +15,14 @@
 ###########################################################################
 
 import warp as wp
+import warp.examples.fem.utils as fem_example_utils
 import warp.fem as fem
-from warp.sparse import bsr_axpy
-
-# Import example utilities
-# Make sure that works both when imported as module and run as standalone file
-try:
-    from .bsr_utils import bsr_cg
-    from .example_convection_diffusion import (
-        diffusion_form,
-        inertia_form,
-        initial_condition,
-        velocity,
-    )
-    from .mesh_utils import gen_quadmesh, gen_trimesh
-    from .plot_utils import Plot
-except ImportError:
-    from bsr_utils import bsr_cg
-    from example_convection_diffusion import (
-        diffusion_form,
-        inertia_form,
-        initial_condition,
-        velocity,
-    )
-    from mesh_utils import gen_quadmesh, gen_trimesh
-    from plot_utils import Plot
+from warp.examples.fem.example_convection_diffusion import (
+    diffusion_form,
+    inertia_form,
+    initial_condition,
+    velocity,
+)
 
 
 # Standard transport term, on cells' interior
@@ -58,6 +41,10 @@ def upwind_transport_form(s: fem.Sample, domain: fem.Domain, phi: fem.Field, psi
     vel = velocity(pos, ang_vel)
     vel_n = wp.dot(vel, fem.normal(domain, s))
 
+    if wp.min(pos) <= 0.0 or wp.max(pos) >= 1.0:  # boundary side
+        return phi(s) * (-psi(s) * vel_n + 0.5 * psi(s) * wp.abs(vel_n))
+
+    # interior side
     return fem.jump(phi, s) * (-fem.average(psi, s) * vel_n + 0.5 * fem.jump(psi, s) * wp.abs(vel_n))
 
 
@@ -79,7 +66,7 @@ def sip_diffusion_form(
 
 
 class Example:
-    def __init__(self, quiet=False, degree=2, resolution=50, mesh="grid", viscosity=0.001, ang_vel=1.0):
+    def __init__(self, quiet=False, degree=2, resolution=50, mesh="grid", viscosity=0.0001, ang_vel=1.0):
         self._quiet = quiet
 
         res = resolution
@@ -87,10 +74,10 @@ class Example:
         self.current_frame = 0
 
         if mesh == "tri":
-            positions, tri_vidx = gen_trimesh(res=wp.vec2i(resolution))
+            positions, tri_vidx = fem_example_utils.gen_trimesh(res=wp.vec2i(resolution))
             geo = fem.Trimesh2D(tri_vertex_indices=tri_vidx, positions=positions)
         elif mesh == "quad":
-            positions, quad_vidx = gen_quadmesh(res=wp.vec2i(resolution))
+            positions, quad_vidx = fem_example_utils.gen_quadmesh(res=wp.vec2i(resolution))
             geo = fem.Quadmesh2D(quad_vertex_indices=quad_vidx, positions=positions)
         else:
             geo = fem.Grid2D(res=wp.vec2i(resolution))
@@ -124,37 +111,30 @@ class Example:
         side_test = fem.make_test(space=scalar_space, domain=sides)
         side_trial = fem.make_trial(space=scalar_space, domain=sides)
 
-        bsr_axpy(
-            fem.integrate(
-                upwind_transport_form,
-                fields={"phi": side_trial, "psi": side_test},
-                values={"ang_vel": ang_vel},
-            ),
-            y=matrix_transport,
+        matrix_transport += fem.integrate(
+            upwind_transport_form,
+            fields={"phi": side_trial, "psi": side_test},
+            values={"ang_vel": ang_vel},
         )
 
         matrix_diffusion = fem.integrate(
             diffusion_form,
             fields={"u": trial, "v": self._test},
         )
-        bsr_axpy(
-            fem.integrate(
-                sip_diffusion_form,
-                fields={"phi": side_trial, "psi": side_test},
-            ),
-            y=matrix_diffusion,
+
+        matrix_diffusion += fem.integrate(
+            sip_diffusion_form,
+            fields={"phi": side_trial, "psi": side_test},
         )
 
-        self._matrix = matrix_inertia
-        bsr_axpy(x=matrix_transport, y=self._matrix)
-        bsr_axpy(x=matrix_diffusion, y=self._matrix, alpha=viscosity)
+        self._matrix = matrix_inertia + matrix_transport + viscosity * matrix_diffusion
 
         # Initial condition
         self._phi_field = scalar_space.make_field()
         fem.interpolate(initial_condition, dest=self._phi_field)
 
-        self.renderer = Plot()
-        self.renderer.add_surface("phi", self._phi_field)
+        self.renderer = fem_example_utils.Plot()
+        self.renderer.add_field("phi", self._phi_field)
 
     def step(self):
         self.current_frame += 1
@@ -166,13 +146,13 @@ class Example:
         )
 
         phi = wp.zeros_like(rhs)
-        bsr_cg(self._matrix, b=rhs, x=phi, method="bicgstab", quiet=self._quiet)
+        fem_example_utils.bsr_cg(self._matrix, b=rhs, x=phi, method="bicgstab", quiet=self._quiet)
 
         wp.utils.array_cast(in_array=phi, out_array=self._phi_field.dof_values)
 
     def render(self):
         self.renderer.begin_frame(time=self.current_frame * self.sim_dt)
-        self.renderer.add_surface("phi", self._phi_field)
+        self.renderer.add_field("phi", self._phi_field)
         self.renderer.end_frame()