warp-lang 1.4.2__py3-none-manylinux2014_aarch64.whl → 1.5.1__py3-none-manylinux2014_aarch64.whl

warp/tests/test_tile_reduce.py

@@ -0,0 +1,374 @@
+# 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 unittest
+
+import numpy as np
+
+import warp as wp
+from warp.tests.unittest_utils import *
+
+TILE_M = wp.constant(8)
+TILE_N = wp.constant(4)
+TILE_K = wp.constant(8)
+
+# num threads per-tile
+TILE_DIM = 64
+
+
+@wp.kernel
+def tile_sum_kernel(input: wp.array2d(dtype=float), output: wp.array(dtype=float)):
+    # output tile index
+    i = wp.tid()
+
+    n = input.shape[1]
+    count = int(n / TILE_DIM)
+
+    s = wp.tile_zeros(m=1, n=1, dtype=float)
+
+    for j in range(count):
+        a = wp.tile_load(input, i, j, m=1, n=TILE_DIM)
+        s += wp.tile_sum(a) * 0.5
+
+    wp.tile_store(output, i, s)
+
+
+def test_tile_reduce_sum(test, device):
+    batch_count = 56
+
+    N = TILE_DIM * 3
+
+    rng = np.random.default_rng(42)
+    input = rng.random((batch_count, N), dtype=np.float32)
+
+    input_wp = wp.array(input, requires_grad=True, device=device)
+    output_wp = wp.zeros(batch_count, requires_grad=True, device=device)
+
+    with wp.Tape() as tape:
+        wp.launch_tiled(
+            tile_sum_kernel, dim=[batch_count], inputs=[input_wp, output_wp], block_dim=TILE_DIM, device=device
+        )
+
+    sum_wp = output_wp.numpy()
+    for i in range(batch_count):
+        sum_np = np.sum(input[i]) * 0.5
+        test.assertAlmostEqual(sum_wp[i], sum_np, places=4)
+
+    output_wp.grad.fill_(1.0)
+
+    tape.backward()
+
+    assert_np_equal(input_wp.grad.numpy(), np.ones_like(input) * 0.5, tol=1.0e-4)
+
+
+@wp.kernel
+def tile_min_kernel(input: wp.array2d(dtype=float), output: wp.array(dtype=float)):
+    # output tile index
+    i = wp.tid()
+
+    a = wp.tile_load(input, i, 0, m=1, n=TILE_DIM)
+    m = wp.tile_min(a)
+
+    wp.tile_store(output, i, m)
+
+
+def test_tile_reduce_min(test, device):
+    batch_count = 56
+
+    N = TILE_DIM
+
+    rng = np.random.default_rng(42)
+    input = rng.random((batch_count, N), dtype=np.float32)
+
+    input_wp = wp.array(input, requires_grad=True, device=device)
+    output_wp = wp.zeros(batch_count, requires_grad=True, device=device)
+
+    with wp.Tape() as tape:
+        wp.launch_tiled(
+            tile_min_kernel, dim=[batch_count], inputs=[input_wp, output_wp], block_dim=TILE_DIM, device=device
+        )
+
+    min_wp = output_wp.numpy()
+    for i in range(batch_count):
+        min_np = np.min(input[i])
+        test.assertAlmostEqual(min_wp[i], min_np, places=4)
+
+
+@wp.kernel
+def tile_max_kernel(input: wp.array2d(dtype=float), output: wp.array(dtype=float)):
+    # output tile index
+    i = wp.tid()
+
+    a = wp.tile_load(input, i, 0, m=1, n=TILE_DIM)
+    m = wp.tile_max(a)
+
+    wp.tile_store(output, i, m)
+
+
+def test_tile_reduce_max(test, device):
+    batch_count = 56
+
+    N = TILE_DIM
+
+    rng = np.random.default_rng(42)
+    input = rng.random((batch_count, N), dtype=np.float32)
+
+    input_wp = wp.array(input, requires_grad=True, device=device)
+    output_wp = wp.zeros(batch_count, requires_grad=True, device=device)
+
+    with wp.Tape() as tape:
+        wp.launch_tiled(
+            tile_max_kernel, dim=[batch_count], inputs=[input_wp, output_wp], block_dim=TILE_DIM, device=device
+        )
+
+    max_wp = output_wp.numpy()
+    for i in range(batch_count):
+        max_np = np.max(input[i])
+        test.assertAlmostEqual(max_wp[i], max_np, places=4)
+
+
+@wp.kernel
+def tile_reduce_custom_kernel(input: wp.array2d(dtype=float), output: wp.array(dtype=float)):
+    # output tile index
+    i = wp.tid()
+
+    a = wp.tile_load(input, i, 0, m=1, n=TILE_DIM)
+    m = wp.tile_reduce(wp.mul, a)
+
+    wp.tile_store(output, i, m)
+
+
+def test_tile_reduce_custom(test, device):
+    batch_count = 56
+
+    N = TILE_DIM
+
+    rng = np.random.default_rng(42)
+    input = rng.random((batch_count, N), dtype=np.float32)
+
+    input_wp = wp.array(input, requires_grad=True, device=device)
+    output_wp = wp.zeros(batch_count, requires_grad=True, device=device)
+
+    with wp.Tape() as tape:
+        wp.launch_tiled(
+            tile_reduce_custom_kernel,
+            dim=[batch_count],
+            inputs=[input_wp, output_wp],
+            block_dim=TILE_DIM,
+            device=device,
+        )
+
+    prod_wp = output_wp.numpy()
+    for i in range(batch_count):
+        prod_np = np.prod(input[i])
+        test.assertAlmostEqual(prod_wp[i], prod_np, places=4)
+
+
+@wp.kernel
+def tile_grouped_sum_kernel(input: wp.array3d(dtype=float), output: wp.array(dtype=float)):
+    # output tile index
+    i = wp.tid()
+
+    a = wp.tile_load(input[i], 0, 0, m=TILE_M, n=TILE_N)
+    s = wp.tile_sum(a) * 0.5
+
+    wp.tile_store(output, i, s)
+
+
+def test_tile_reduce_grouped_sum(test, device):
+    batch_count = 56
+
+    M = TILE_M
+    N = TILE_N
+
+    rng = np.random.default_rng(42)
+    input = rng.random((batch_count, M, N), dtype=np.float32)
+
+    input_wp = wp.array(input, requires_grad=True, device=device)
+    output_wp = wp.zeros(batch_count, requires_grad=True, device=device)
+
+    with wp.Tape() as tape:
+        wp.launch_tiled(
+            tile_sum_kernel, dim=[batch_count], inputs=[input_wp, output_wp], block_dim=TILE_DIM, device=device
+        )
+
+    sum_wp = output_wp.numpy()
+    for i in range(batch_count):
+        sum_np = np.sum(input[i]) * 0.5
+        test.assertAlmostEqual(sum_wp[i], sum_np, places=4)
+
+    output_wp.grad.fill_(1.0)
+
+    tape.backward()
+
+    assert_np_equal(input_wp.grad.numpy(), np.ones_like(input) * 0.5, tol=1.0e-4)
+
+
+@wp.kernel
+def tile_reduce_simt_kernel(output: wp.array(dtype=int)):
+    # thread index
+    i = wp.tid()
+
+    t = wp.tile(i)  # convert to block wide tile
+    s = wp.tile_sum(t)  # sum over block
+
+    # update global sum
+    wp.tile_atomic_add(output, 0, 0, s)
+
+
+def test_tile_reduce_simt(test, device):
+    # use an unaligned grid dimension
+    N = TILE_DIM * 4 + 5
+
+    output = wp.zeros(shape=1, dtype=int, requires_grad=True, device=device)
+
+    with wp.Tape() as tape:
+        wp.launch(tile_reduce_simt_kernel, dim=N, inputs=[output], block_dim=TILE_DIM, device=device)
+
+    test.assertEqual(output.numpy()[0], np.sum(np.arange(N)))
+
+
+@wp.kernel
+def tile_untile_kernel(output: wp.array(dtype=int)):
+    # thread index
+    i = wp.tid()
+
+    # convert to block wide tile
+    t = wp.tile(i) * 2
+    s = wp.untile(t)
+
+    output[i] = s
+
+
+def test_tile_untile(test, device):
+    # use an unaligned grid dimension
+    N = TILE_DIM * 4 + 5
+
+    output = wp.zeros(shape=N, dtype=int, requires_grad=True, device=device)
+
+    with wp.Tape() as tape:
+        wp.launch(tile_untile_kernel, dim=N, inputs=[output], block_dim=TILE_DIM, device=device)
+
+    assert_np_equal(output.numpy(), np.arange(N) * 2)
+
+
+@wp.kernel
+def tile_untile_scalar_kernel(output: wp.array(dtype=int)):
+    # thread index
+    i = wp.tid()
+
+    # convert to block wide tile
+    t = wp.tile(i) * 2
+    s = wp.untile(t)
+
+    output[i] = s
+
+
+def test_tile_untile_scalar(test, device):
+    # use an unaligned grid dimension
+    N = TILE_DIM * 4 + 5
+
+    output = wp.zeros(shape=N, dtype=int, requires_grad=True, device=device)
+
+    with wp.Tape() as tape:
+        wp.launch(tile_untile_kernel, dim=N, inputs=[output], block_dim=TILE_DIM, device=device)
+
+    assert_np_equal(output.numpy(), np.arange(N) * 2)
+
+
+@wp.kernel
+def test_untile_vector_kernel(input: wp.array(dtype=wp.vec3), output: wp.array(dtype=wp.vec3)):
+    i = wp.tid()
+
+    v = input[i] * 0.5
+
+    t = wp.tile(v)
+    u = wp.untile(t)
+
+    output[i] = u * 2.0
+
+
+def test_tile_untile_vector(test, device):
+    input = wp.full(16, wp.vec3(1.0, 2.0, 3.0), requires_grad=True, device=device)
+    output = wp.zeros_like(input, device=device)
+
+    with wp.Tape() as tape:
+        wp.launch(test_untile_vector_kernel, dim=16, inputs=[input, output], block_dim=16, device=device)
+
+    output.grad = wp.ones_like(output, device=device)
+    tape.backward()
+
+    assert_np_equal(output.numpy(), input.numpy())
+    assert_np_equal(input.grad.numpy(), np.ones((16, 3)))
+
+
+@wp.kernel
+def tile_ones_kernel(out: wp.array(dtype=float)):
+    i = wp.tid()
+
+    t = wp.tile_ones(dtype=float, m=16, n=16)
+    s = wp.tile_sum(t)
+
+    wp.tile_store(out, 0, s)
+
+
+def test_tile_ones(test, device):
+    output = wp.zeros(1, dtype=float, device=device)
+
+    with wp.Tape() as tape:
+        wp.launch_tiled(tile_ones_kernel, dim=[1], inputs=[output], block_dim=TILE_DIM, device=device)
+
+    test.assertAlmostEqual(output.numpy()[0], 256.0)
+
+
+@wp.kernel
+def tile_arange_kernel(out: wp.array2d(dtype=int)):
+    i = wp.tid()
+
+    a = wp.tile_arange(17, dtype=int)
+    b = wp.tile_arange(5, 23, dtype=int)
+    c = wp.tile_arange(0, 34, 2, dtype=int)
+
+    wp.tile_store(out, 0, 0, a)
+    wp.tile_store(out, 1, 0, b)
+    wp.tile_store(out, 2, 0, c)
+
+
+def test_tile_arange(test, device):
+    N = 17
+
+    output = wp.zeros(shape=(3, N), dtype=int, device=device)
+
+    with wp.Tape() as tape:
+        wp.launch_tiled(tile_arange_kernel, dim=[1], inputs=[output], block_dim=TILE_DIM, device=device)
+
+    assert_np_equal(output.numpy()[0], np.arange(17))
+    assert_np_equal(output.numpy()[1], np.arange(5, 22))
+    assert_np_equal(output.numpy()[2], np.arange(0, 34, 2))
+
+
+devices = get_cuda_test_devices()
+
+
+class TestTileReduce(unittest.TestCase):
+    pass
+
+
+add_function_test(TestTileReduce, "test_tile_reduce_sum", test_tile_reduce_sum, devices=devices)
+add_function_test(TestTileReduce, "test_tile_reduce_min", test_tile_reduce_min, devices=devices)
+add_function_test(TestTileReduce, "test_tile_reduce_max", test_tile_reduce_max, devices=devices)
+add_function_test(TestTileReduce, "test_tile_reduce_custom", test_tile_reduce_custom, devices=devices)
+add_function_test(TestTileReduce, "test_tile_reduce_grouped_sum", test_tile_reduce_sum, devices=devices)
+add_function_test(TestTileReduce, "test_tile_reduce_simt", test_tile_reduce_simt, devices=devices)
+add_function_test(TestTileReduce, "test_tile_ones", test_tile_ones, devices=devices)
+add_function_test(TestTileReduce, "test_tile_arange", test_tile_arange, devices=devices)
+add_function_test(TestTileReduce, "test_tile_untile_scalar", test_tile_untile_scalar, devices=devices)
+add_function_test(TestTileReduce, "test_tile_untile_vector", test_tile_untile_vector, devices=devices)
+
+if __name__ == "__main__":
+    wp.clear_kernel_cache()
+    unittest.main(verbosity=2, failfast=True)

warp/tests/test_tile_shared_memory.py

@@ -0,0 +1,190 @@
+# 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 unittest
+
+import numpy as np
+
+import warp as wp
+from warp.tests.unittest_utils import *
+
+
+# checks that we can configure shared memory to the expected size
+def test_tile_shared_mem_size(test, device):
+    DIM_M = 32
+    DIM_N = 32
+
+    BLOCK_DIM = 256
+
+    @wp.kernel
+    def compute(out: wp.array2d(dtype=float)):
+        a = wp.tile_ones(DIM_M, DIM_N, dtype=float, storage="shared")
+        b = wp.tile_ones(DIM_M, DIM_N, dtype=float, storage="shared") * 2.0
+
+        c = a + b
+        wp.tile_store(out, 0, 0, c)
+
+    out = wp.empty((DIM_M, DIM_N), dtype=float, device=device)
+
+    wp.launch_tiled(compute, dim=[1], inputs=[out], block_dim=BLOCK_DIM, device=device)
+
+    # check output
+    assert_np_equal(out.numpy(), np.ones((DIM_M, DIM_N)) * 3.0)
+
+    # check required shared memory
+    expected_forward_bytes = DIM_M * DIM_N * 4 * 2
+    expected_backward_bytes = expected_forward_bytes * 2
+
+    # check shared memory for kernel on the device
+    module_exec = compute.module.load(device, BLOCK_DIM)
+    hooks = module_exec.get_kernel_hooks(compute)
+
+    assert hooks.forward_smem_bytes == expected_forward_bytes
+    assert hooks.backward_smem_bytes == expected_backward_bytes
+
+
+# checks that we can configure shared memory > 48kb default
+def test_tile_shared_mem_large(test, device):
+    # set dimensions that require 64kb for the forward kernel
+    DIM_M = 64
+    DIM_N = 128
+
+    BLOCK_DIM = 256
+
+    # we disable backward kernel gen since 128k is not supported on most architectures
+    @wp.kernel(enable_backward=False)
+    def compute(out: wp.array2d(dtype=float)):
+        a = wp.tile_ones(DIM_M, DIM_N, dtype=float, storage="shared")
+        b = wp.tile_ones(DIM_M, DIM_N, dtype=float, storage="shared") * 2.0
+
+        c = a + b
+        wp.tile_store(out, 0, 0, c)
+
+    out = wp.empty((DIM_M, DIM_N), dtype=float, device=device)
+
+    wp.launch_tiled(compute, dim=[1], inputs=[out], block_dim=BLOCK_DIM, device=device)
+
+    # check output
+    assert_np_equal(out.numpy(), np.ones((DIM_M, DIM_N)) * 3.0)
+
+    # check required shared memory
+    expected_forward_bytes = DIM_M * DIM_N * 4 * 2
+    expected_backward_bytes = expected_forward_bytes * 2
+
+    assert expected_forward_bytes == 2**16
+
+    # check shared memory for kernel on the device
+    module_exec = compute.module.load(device, BLOCK_DIM)
+    hooks = module_exec.get_kernel_hooks(compute)
+
+    assert hooks.forward_smem_bytes == expected_forward_bytes
+    assert hooks.backward_smem_bytes == expected_backward_bytes
+
+
+# checks that we can configure dynamic shared memory during graph capture
+def test_tile_shared_mem_graph(test, device):
+    DIM_M = 32
+    DIM_N = 32
+
+    BLOCK_DIM = 256
+
+    @wp.kernel
+    def compute(out: wp.array2d(dtype=float)):
+        a = wp.tile_ones(DIM_M, DIM_N, dtype=float, storage="shared")
+        b = wp.tile_ones(DIM_M, DIM_N, dtype=float, storage="shared") * 2.0
+
+        c = a + b
+        wp.tile_store(out, 0, 0, c)
+
+    out = wp.empty((DIM_M, DIM_N), dtype=float, device=device)
+
+    wp.load_module(device=device)
+
+    wp.capture_begin(device, force_module_load=False)
+    wp.launch_tiled(compute, dim=[1], inputs=[out], block_dim=BLOCK_DIM, device=device)
+    graph = wp.capture_end(device)
+
+    wp.capture_launch(graph)
+
+    # check output
+    assert_np_equal(out.numpy(), np.ones((DIM_M, DIM_N)) * 3.0)
+
+    # check required shared memory
+    expected_forward_bytes = DIM_M * DIM_N * 4 * 2
+    expected_backward_bytes = expected_forward_bytes * 2
+
+    # check shared memory for kernel on the device
+    module_exec = compute.module.load(device, BLOCK_DIM)
+    hooks = module_exec.get_kernel_hooks(compute)
+
+    assert hooks.forward_smem_bytes == expected_forward_bytes
+    assert hooks.backward_smem_bytes == expected_backward_bytes
+
+
+# checks that stack allocations work for user functions
+def test_tile_shared_mem_func(test, device):
+    DIM_M = 32
+    DIM_N = 32
+
+    BLOCK_DIM = 256
+
+    @wp.func
+    def add_tile_small():
+        a = wp.tile_ones(16, 16, dtype=float, storage="shared")
+        b = wp.tile_ones(16, 16, dtype=float, storage="shared") * 2.0
+
+        return a + b
+
+    @wp.func
+    def add_tile_big():
+        a = wp.tile_ones(64, 64, dtype=float, storage="shared")
+        b = wp.tile_ones(64, 64, dtype=float, storage="shared") * 2.0
+
+        return a + b
+
+    @wp.kernel
+    def compute(out: wp.array2d(dtype=float)):
+        s = add_tile_small()
+        b = add_tile_big()
+
+        wp.tile_store(out, 0, 0, b)
+
+    out = wp.empty((DIM_M, DIM_N), dtype=float, device=device)
+
+    wp.launch_tiled(compute, dim=[1], inputs=[out], block_dim=BLOCK_DIM, device=device)
+
+    # check shared memory for kernel on the device
+    module_exec = compute.module.load(device, BLOCK_DIM)
+    hooks = module_exec.get_kernel_hooks(compute)
+
+    # ensure that total required dynamic shared is the larger of the two tiles
+    expected_required_shared = 64 * 64 * 4 * 2
+
+    assert hooks.forward_smem_bytes == expected_required_shared
+    assert hooks.backward_smem_bytes == expected_required_shared * 2
+
+
+devices = get_cuda_test_devices()
+
+
+class TestTileSharedMemory(unittest.TestCase):
+    pass
+
+
+add_function_test(
+    TestTileSharedMemory, "test_tile_shared_mem_size", test_tile_shared_mem_size, devices=devices, check_output=False
+)
+add_function_test(
+    TestTileSharedMemory, "test_tile_shared_mem_large", test_tile_shared_mem_large, devices=devices, check_output=False
+)
+add_function_test(TestTileSharedMemory, "test_tile_shared_mem_graph", test_tile_shared_mem_graph, devices=devices)
+add_function_test(TestTileSharedMemory, "test_tile_shared_mem_func", test_tile_shared_mem_func, devices=devices)
+
+
+if __name__ == "__main__":
+    wp.clear_kernel_cache()
+    unittest.main(verbosity=2, failfast=True)

warp/tests/test_vbd.py

@@ -5,6 +5,8 @@
 # distribution of this software and related documentation without an express
 # license agreement from NVIDIA CORPORATION is strictly prohibited.
 
+import contextlib
+import io
 import unittest
 
 import warp as wp
@@ -287,14 +289,6 @@ class VBDClothSim:
             89, 99, 100
         ]
 
-        self.coloring = [
-            [9, 12, 17, 24, 31, 38, 43, 46, 50, 62, 65, 68, 80, 84, 89, 92],
-            [6, 20, 25, 32, 37, 44, 51, 56, 59, 63, 70, 75, 82, 88, 90, 94, 96],
-            [2, 8, 10, 14, 26, 29, 33, 40, 48, 52, 55, 67, 73, 79, 86, 91, 98],
-            [4, 11, 16, 23, 28, 30, 35, 42, 49, 54, 57, 71, 74, 76, 78, 93, 97],
-            [3, 15, 18, 22, 34, 36, 39, 41, 53, 58, 60, 66, 72, 85, 99, 0, 87],
-            [7, 21, 27, 45, 47, 61, 64, 69, 77, 81, 83, 95, 1, 5, 13, 19],
-        ]
         # fmt: on
 
         self.dt = 1 / 60
@@ -323,6 +317,7 @@ class VBDClothSim:
             tri_ka=stiffness,
             tri_kd=kd,
         )
+        builder.color()
 
         self.model = builder.finalize(device=device)
         self.model.ground = True
@@ -331,11 +326,6 @@ class VBDClothSim:
         self.model.soft_contact_ke = 1.0e4
         self.model.soft_contact_kd = 1.0e2
 
-        coloring_wp = []
-        for color in self.coloring:
-            coloring_wp.append(wp.array(color, dtype=wp.int32, device=self.model.device))
-        self.model.coloring = coloring_wp
-
         self.dt = self.dt / self.num_substeps
         self.fixed_particles = [0, 9]
 
@@ -367,19 +357,21 @@ class VBDClothSim:
             model.particle_flags = wp.array(flags, device=model.device)
 
 
-def test_vbd_cloth(test, device):
-    example = VBDClothSim(device)
-    example.run(test)
-
-
 devices = get_test_devices()
 
 
 class TestVBD(unittest.TestCase):
-    pass
+    def test_vbd_cloth(self):
+        for device in devices:
+            with contextlib.redirect_stdout(io.StringIO()) as f:
+                example = VBDClothSim(device)
+            self.assertRegex(
+                f.getvalue(),
+                r"Warp UserWarning: The graph is not optimizable anymore, terminated with a max/min ratio: 2.0 without reaching the target ratio: 1.1",
+            )
 
+            example.run(self)
 
-add_function_test(TestVBD, "test_vbd_cloth", test_vbd_cloth, devices=devices)
 
 if __name__ == "__main__":
     wp.clear_kernel_cache()

warp/tests/test_volume.py

@@ -5,6 +5,8 @@
 # distribution of this software and related documentation without an express
 # license agreement from NVIDIA CORPORATION is strictly prohibited.
 
+import os
+import tempfile
 import unittest
 from typing import Any
 
@@ -890,6 +892,46 @@ def test_volume_aniso_transform(test, device):
     assert_np_equal(transform, np.array(volume.get_grid_info().transform_matrix).reshape(3, 3))
 
 
+def test_volume_write(test, device):
+    codecs = ["none", "zip", "blosc"]
+    try:
+        import blosc  # noqa: F401 I001
+    except ImportError:
+        codecs.pop()
+
+    for volume_name in ("float", "vec3f", "index"):
+        for codec in codecs:
+            with test.subTest(volume_name=volume_name, codec=codec):
+                volume = volumes[volume_name][device.alias]
+                fd, file_path = tempfile.mkstemp(suffix=".nvdb")
+                os.close(fd)
+                try:
+                    volume.save_to_nvdb(file_path, codec=codec)
+                    with open(file_path, "rb") as f:
+                        volume_2 = wp.Volume.load_from_nvdb(f)
+                    next_volume = volume
+                    while next_volume:
+                        np.testing.assert_array_equal(next_volume.array().numpy(), volume_2.array().numpy())
+                        next_volume = next_volume.load_next_grid()
+                        volume_2 = volume_2.load_next_grid()
+
+                finally:
+                    os.remove(file_path)
+
+    with test.subTest(volume_write="unsupported"):
+        volume = volumes["index"][device.alias]
+        volume = volume.load_next_grid()
+
+        fd, file_path = tempfile.mkstemp(suffix=".nvdb")
+        os.close(fd)
+
+        try:
+            with test.assertRaises(RuntimeError):
+                volume.save_to_nvdb(file_path, codec=codec)
+        finally:
+            os.remove(file_path)
+
+
 class TestVolume(unittest.TestCase):
     def test_volume_new_del(self):
         # test the scenario in which a volume is created but not initialized before gc
@@ -930,6 +972,7 @@ add_function_test(
 add_function_test(TestVolume, "test_volume_multiple_grids", test_volume_multiple_grids, devices=devices)
 add_function_test(TestVolume, "test_volume_feature_array", test_volume_feature_array, devices=devices)
 add_function_test(TestVolume, "test_volume_sample_index", test_volume_sample_index, devices=devices)
+add_function_test(TestVolume, "test_volume_write", test_volume_write, devices=[wp.get_device("cpu")])
 
 points = {}
 points_jittered = {}