NvidiaWarp-GarmentCode/warp/tests/test_fabricarray.py

# Copyright (c) 2023 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 math
import unittest
from typing import Any

import numpy as np

import warp as wp
from warp.tests.unittest_utils import *

wp.init()


# types to test fabric arrays
_fabric_types = [
    *wp.types.scalar_types,
    *[wp.types.vector(2, T) for T in wp.types.scalar_types],
    *[wp.types.vector(3, T) for T in wp.types.scalar_types],
    *[wp.types.vector(4, T) for T in wp.types.scalar_types],
    *[wp.types.matrix((2, 2), T) for T in wp.types.scalar_types],
    *[wp.types.matrix((3, 3), T) for T in wp.types.scalar_types],
    *[wp.types.matrix((4, 4), T) for T in wp.types.scalar_types],
    *[wp.types.quaternion(T) for T in wp.types.float_types],
]


def _warp_type_to_fabric(dtype, is_array=False):
    scalar_map = {
        wp.bool: "b",
        wp.int8: "i1",
        wp.int16: "i2",
        wp.int32: "i4",
        wp.int64: "i8",
        wp.uint8: "u1",
        wp.uint16: "u2",
        wp.uint32: "u4",
        wp.uint64: "u8",
        wp.float16: "f2",
        wp.float32: "f4",
        wp.float64: "f8",
    }

    if hasattr(dtype, "_wp_scalar_type_"):
        type_str = scalar_map[dtype._wp_scalar_type_]
        if len(dtype._shape_) == 1:
            role = "vector"
        else:
            role = "matrix"
    else:
        type_str = scalar_map[dtype]
        role = ""

    if is_array:
        array_depth = 1
    else:
        array_depth = 0

    return (True, type_str, dtype._length_, array_depth, role)


# returns a fabric array interface constructed from a regular array
def _create_fabric_array_interface(data: wp.array, attrib: str, bucket_sizes: list = None, copy=False):
    assert isinstance(data, wp.array)
    assert data.ndim == 1

    assert isinstance(attrib, str)

    if copy:
        data = wp.clone(data)

    if bucket_sizes is not None:
        assert hasattr(bucket_sizes, "__len__")

        # verify total size
        total_size = 0
        for bucket_size in bucket_sizes:
            total_size += bucket_size

        if total_size != data.size:
            raise RuntimeError("Bucket sizes don't add up to the size of data array")

    elif data.size > 0:
        rng = np.random.default_rng(123)

        # generate random bucket sizes
        bucket_min = 1
        bucket_max = math.ceil(0.5 * data.size)
        total_size = data.size
        size_remaining = total_size

        bucket_sizes = []
        while size_remaining >= bucket_max:
            bucket_size = rng.integers(bucket_min, high=bucket_max, dtype=int)
            bucket_sizes.append(bucket_size)
            size_remaining -= bucket_size

        if size_remaining > 0:
            bucket_sizes.append(size_remaining)

    else:
        # empty data array
        bucket_sizes = []

    dtype_size = wp.types.type_size_in_bytes(data.dtype)
    p = int(data.ptr) if data.ptr else 0
    pointers = []
    counts = []
    for bucket_size in bucket_sizes:
        pointers.append(p)
        counts.append(bucket_size)
        p += bucket_size * dtype_size

    attrib_info = {}

    attrib_info["type"] = _warp_type_to_fabric(data.dtype)
    attrib_info["access"] = 2  # ReadWrite
    attrib_info["pointers"] = pointers
    attrib_info["counts"] = counts

    iface = {}
    iface["version"] = 1
    iface["device"] = str(data.device)
    iface["attribs"] = {attrib: attrib_info}
    iface["_ref"] = data  # backref to keep the array alive

    return iface


# returns a fabric array array interface constructed from a list of regular arrays
def _create_fabric_array_array_interface(data: list, attrib: str, bucket_sizes: list = None):
    # data should be a list of arrays
    assert isinstance(data, list)

    num_arrays = len(data)
    assert num_arrays > 0

    device = data[0].device
    dtype = data[0].dtype

    assert isinstance(attrib, str)

    if bucket_sizes is not None:
        assert hasattr(bucket_sizes, "__len__")

        # verify total size
        total_size = 0
        for bucket_size in bucket_sizes:
            total_size += bucket_size

        if total_size != num_arrays:
            raise RuntimeError("Bucket sizes don't add up to the number of given arrays")

    else:
        rng = np.random.default_rng(123)

        # generate random bucket sizes
        bucket_min = 1
        bucket_max = math.ceil(0.5 * num_arrays)
        total_size = num_arrays
        size_remaining = total_size

        bucket_sizes = []
        while size_remaining >= bucket_max:
            bucket_size = rng.integers(bucket_min, high=bucket_max, dtype=int)
            bucket_sizes.append(bucket_size)
            size_remaining -= bucket_size

        if size_remaining > 0:
            bucket_sizes.append(size_remaining)

    # initialize array of pointers to arrays and their lengths
    _array_pointers = []
    _array_lengths = []
    for i in range(num_arrays):
        _array_pointers.append(data[i].ptr)
        _array_lengths.append(data[i].size)

    array_pointers = wp.array(_array_pointers, dtype=wp.uint64, device=device)
    pointer_size = wp.types.type_size_in_bytes(array_pointers.dtype)

    lengths = wp.array(_array_lengths, dtype=wp.uint64, device=device)
    length_size = wp.types.type_size_in_bytes(lengths.dtype)

    p_pointers = int(array_pointers.ptr)
    p_lengths = int(lengths.ptr)
    pointers = []
    counts = []
    array_lengths = []
    for bucket_size in bucket_sizes:
        pointers.append(p_pointers)
        counts.append(bucket_size)
        array_lengths.append(p_lengths)
        p_pointers += bucket_size * pointer_size
        p_lengths += bucket_size * length_size

    attrib_info = {}

    attrib_info["type"] = _warp_type_to_fabric(dtype, is_array=True)
    attrib_info["access"] = 2  # ReadWrite
    attrib_info["pointers"] = pointers
    attrib_info["counts"] = counts
    attrib_info["array_lengths"] = array_lengths

    iface = {}
    iface["version"] = 1
    iface["device"] = str(device)
    iface["attribs"] = {attrib: attrib_info}
    iface["_ref"] = data  # backref to keep the data arrays alive
    iface["_ref_pointers"] = array_pointers  # backref to keep the array pointers alive
    iface["_ref_lengths"] = lengths  # backref to keep the lengths array alive

    return iface


@wp.kernel
def fa_kernel(a: wp.fabricarray(dtype=float), expected: wp.array(dtype=float)):
    i = wp.tid()

    wp.expect_eq(a[i], expected[i])

    a[i] = 2.0 * a[i]

    wp.atomic_add(a, i, 1.0)

    wp.expect_eq(a[i], 2.0 * expected[i] + 1.0)


@wp.kernel
def fa_kernel_indexed(a: wp.indexedfabricarray(dtype=float), expected: wp.indexedarray(dtype=float)):
    i = wp.tid()

    wp.expect_eq(a[i], expected[i])

    a[i] = 2.0 * a[i]

    wp.atomic_add(a, i, 1.0)

    wp.expect_eq(a[i], 2.0 * expected[i] + 1.0)


def test_fabricarray_kernel(test, device):
    data = wp.array(data=np.arange(100, dtype=np.float32), device=device)
    iface = _create_fabric_array_interface(data, "foo", copy=True)
    fa = wp.fabricarray(data=iface, attrib="foo")

    test.assertEqual(fa.dtype, data.dtype)
    test.assertEqual(fa.ndim, 1)
    test.assertEqual(fa.shape, data.shape)
    test.assertEqual(fa.size, data.size)

    wp.launch(fa_kernel, dim=fa.size, inputs=[fa, data], device=device)

    # reset data
    wp.copy(fa, data)

    # test indexed
    indices = wp.array(data=np.arange(1, data.size, 2, dtype=np.int32), device=device)
    ifa = fa[indices]
    idata = data[indices]

    test.assertEqual(ifa.dtype, idata.dtype)
    test.assertEqual(ifa.ndim, 1)
    test.assertEqual(ifa.shape, idata.shape)
    test.assertEqual(ifa.size, idata.size)

    wp.launch(fa_kernel_indexed, dim=ifa.size, inputs=[ifa, idata], device=device)

    wp.synchronize_device(device)


@wp.kernel
def fa_generic_dtype_kernel(a: wp.fabricarray(dtype=Any), b: wp.fabricarray(dtype=Any)):
    i = wp.tid()
    b[i] = a[i] + a[i]


@wp.kernel
def fa_generic_dtype_kernel_indexed(a: wp.indexedfabricarray(dtype=Any), b: wp.indexedfabricarray(dtype=Any)):
    i = wp.tid()
    b[i] = a[i] + a[i]


def test_fabricarray_generic_dtype(test, device):
    for T in _fabric_types:
        if hasattr(T, "_wp_scalar_type_"):
            nptype = wp.types.warp_type_to_np_dtype[T._wp_scalar_type_]
        else:
            nptype = wp.types.warp_type_to_np_dtype[T]

        data = wp.array(data=np.arange(10, dtype=nptype), device=device)
        data_iface = _create_fabric_array_interface(data, "foo", copy=True)
        fa = wp.fabricarray(data=data_iface, attrib="foo")

        result = wp.zeros_like(data)
        result_iface = _create_fabric_array_interface(result, "foo", copy=True)
        fb = wp.fabricarray(data=result_iface, attrib="foo")

        test.assertEqual(fa.dtype, fb.dtype)
        test.assertEqual(fa.ndim, fb.ndim)
        test.assertEqual(fa.shape, fb.shape)
        test.assertEqual(fa.size, fb.size)

        wp.launch(fa_generic_dtype_kernel, dim=fa.size, inputs=[fa, fb], device=device)

        assert_np_equal(fb.numpy(), 2 * fa.numpy())

        # reset data
        wp.copy(fa, data)
        wp.copy(fb, result)

        # test indexed
        indices = wp.array(data=np.arange(1, data.size, 2, dtype=np.int32), device=device)
        ifa = fa[indices]
        ifb = fb[indices]

        test.assertEqual(ifa.dtype, ifb.dtype)
        test.assertEqual(ifa.ndim, ifb.ndim)
        test.assertEqual(ifa.shape, ifb.shape)
        test.assertEqual(ifa.size, ifb.size)

        wp.launch(fa_generic_dtype_kernel_indexed, dim=ifa.size, inputs=[ifa, ifb], device=device)

        assert_np_equal(ifb.numpy(), 2 * ifa.numpy())


@wp.kernel
def fa_generic_array_kernel(a: Any, b: Any):
    i = wp.tid()
    b[i] = a[i] + a[i]


def test_fabricarray_generic_array(test, device):
    for T in _fabric_types:
        if hasattr(T, "_wp_scalar_type_"):
            nptype = wp.types.warp_type_to_np_dtype[T._wp_scalar_type_]
        else:
            nptype = wp.types.warp_type_to_np_dtype[T]

        data = wp.array(data=np.arange(100, dtype=nptype), device=device)
        data_iface = _create_fabric_array_interface(data, "foo", copy=True)
        fa = wp.fabricarray(data=data_iface, attrib="foo")

        result = wp.zeros_like(data)
        result_iface = _create_fabric_array_interface(result, "foo", copy=True)
        fb = wp.fabricarray(data=result_iface, attrib="foo")

        test.assertEqual(fa.dtype, fb.dtype)
        test.assertEqual(fa.ndim, fb.ndim)
        test.assertEqual(fa.shape, fb.shape)
        test.assertEqual(fa.size, fb.size)

        wp.launch(fa_generic_array_kernel, dim=fa.size, inputs=[fa, fb], device=device)

        assert_np_equal(fb.numpy(), 2 * fa.numpy())

        # reset data
        wp.copy(fa, data)
        wp.copy(fb, result)

        # test indexed
        indices = wp.array(data=np.arange(1, data.size, 2, dtype=np.int32), device=device)
        ifa = fa[indices]
        ifb = fb[indices]

        test.assertEqual(ifa.dtype, ifb.dtype)
        test.assertEqual(ifa.ndim, ifb.ndim)
        test.assertEqual(ifa.shape, ifb.shape)
        test.assertEqual(ifa.size, ifb.size)

        wp.launch(fa_generic_array_kernel, dim=ifa.size, inputs=[ifa, ifb], device=device)

        assert_np_equal(ifb.numpy(), 2 * ifa.numpy())


def test_fabricarray_empty(test, device):
    # Test whether common operations work with empty (zero-sized) indexed arrays
    # without throwing exceptions.

    def test_empty_ops(nrows, ncols, wptype, nptype):
        # scalar, vector, or matrix
        if ncols > 0:
            if nrows > 0:
                wptype = wp.types.matrix((nrows, ncols), wptype)
            else:
                wptype = wp.types.vector(ncols, wptype)
            dtype_shape = wptype._shape_
        else:
            dtype_shape = ()

        fill_value = wptype(42)

        # create an empty data array
        data = wp.empty(0, dtype=wptype, device=device)
        iface = _create_fabric_array_interface(data, "foo", copy=True)
        fa = wp.fabricarray(data=iface, attrib="foo")

        test.assertEqual(fa.size, 0)
        test.assertEqual(fa.shape, (0,))

        # all of these methods should succeed with zero-sized arrays
        fa.zero_()
        fa.fill_(fill_value)
        fb = fa.contiguous()

        fb = wp.empty_like(fa)
        fb = wp.zeros_like(fa)
        fb = wp.full_like(fa, fill_value)
        fb = wp.clone(fa)

        wp.copy(fa, fb)
        fa.assign(fb)

        na = fa.numpy()
        test.assertEqual(na.size, 0)
        test.assertEqual(na.shape, (0, *dtype_shape))
        test.assertEqual(na.dtype, nptype)

        test.assertEqual(fa.list(), [])

        # test indexed

        # create a zero-sized array of indices
        indices = wp.empty(0, dtype=int, device=device)

        ifa = fa[indices]

        test.assertEqual(ifa.size, 0)
        test.assertEqual(ifa.shape, (0,))

        # all of these methods should succeed with zero-sized arrays
        ifa.zero_()
        ifa.fill_(fill_value)
        ifb = ifa.contiguous()

        ifb = wp.empty_like(ifa)
        ifb = wp.zeros_like(ifa)
        ifb = wp.full_like(ifa, fill_value)
        ifb = wp.clone(ifa)

        wp.copy(ifa, ifb)
        ifa.assign(ifb)

        na = ifa.numpy()
        test.assertEqual(na.size, 0)
        test.assertEqual(na.shape, (0, *dtype_shape))
        test.assertEqual(na.dtype, nptype)

        test.assertEqual(ifa.list(), [])

    # test with scalars, vectors, and matrices
    for nptype, wptype in wp.types.np_dtype_to_warp_type.items():
        # scalars
        test_empty_ops(0, 0, wptype, nptype)

        for ncols in [2, 3, 4, 5]:
            # vectors
            test_empty_ops(0, ncols, wptype, nptype)
            # square matrices (the Fabric interface only supports square matrices right now)
            test_empty_ops(ncols, ncols, wptype, nptype)


def test_fabricarray_fill_scalar(test, device):
    for nptype, wptype in wp.types.np_dtype_to_warp_type.items():
        # create a data array
        data = wp.zeros(100, dtype=wptype, device=device)
        iface = _create_fabric_array_interface(data, "foo", copy=True)
        fa = wp.fabricarray(data=iface, attrib="foo")

        assert_np_equal(fa.numpy(), np.zeros(fa.shape, dtype=nptype))

        # fill with int value
        fill_value = 42
        fa.fill_(fill_value)
        assert_np_equal(fa.numpy(), np.full(fa.shape, fill_value, dtype=nptype))

        fa.zero_()
        assert_np_equal(fa.numpy(), np.zeros(fa.shape, dtype=nptype))

        if wptype in wp.types.float_types:
            # fill with float value
            fill_value = 13.37
            fa.fill_(fill_value)
            assert_np_equal(fa.numpy(), np.full(fa.shape, fill_value, dtype=nptype))

        # fill with Warp scalar value
        fill_value = wptype(17)
        fa.fill_(fill_value)
        assert_np_equal(fa.numpy(), np.full(fa.shape, fill_value.value, dtype=nptype))

        # reset data
        wp.copy(fa, data)

        # test indexed
        indices1 = wp.array(data=np.arange(1, data.size, 2, dtype=np.int32), device=device)
        ifa = fa[indices1]

        # ensure that the other indices remain unchanged
        indices2 = wp.array(data=np.arange(0, data.size, 2, dtype=np.int32), device=device)
        ifb = fa[indices2]

        assert_np_equal(ifa.numpy(), np.zeros(ifa.shape, dtype=nptype))
        assert_np_equal(ifb.numpy(), np.zeros(ifb.shape, dtype=nptype))

        # fill with int value
        fill_value = 42
        ifa.fill_(fill_value)
        assert_np_equal(ifa.numpy(), np.full(ifa.shape, fill_value, dtype=nptype))
        assert_np_equal(ifb.numpy(), np.zeros(ifb.shape, dtype=nptype))

        ifa.zero_()
        assert_np_equal(ifa.numpy(), np.zeros(ifa.shape, dtype=nptype))
        assert_np_equal(ifb.numpy(), np.zeros(ifb.shape, dtype=nptype))

        if wptype in wp.types.float_types:
            # fill with float value
            fill_value = 13.37
            ifa.fill_(fill_value)
            assert_np_equal(ifa.numpy(), np.full(ifa.shape, fill_value, dtype=nptype))
            assert_np_equal(ifb.numpy(), np.zeros(ifb.shape, dtype=nptype))

        # fill with Warp scalar value
        fill_value = wptype(17)
        ifa.fill_(fill_value)
        assert_np_equal(ifa.numpy(), np.full(ifa.shape, fill_value.value, dtype=nptype))
        assert_np_equal(ifb.numpy(), np.zeros(ifb.shape, dtype=nptype))


def test_fabricarray_fill_vector(test, device):
    # test filling a vector array with scalar or vector values (vec_type, list, or numpy array)

    for nptype, wptype in wp.types.np_dtype_to_warp_type.items():
        # vector types
        vector_types = [
            wp.types.vector(2, wptype),
            wp.types.vector(3, wptype),
            wp.types.vector(4, wptype),
            wp.types.vector(5, wptype),
        ]

        for vec_type in vector_types:
            vec_len = vec_type._length_

            data = wp.zeros(100, dtype=vec_type, device=device)
            iface = _create_fabric_array_interface(data, "foo", copy=True)
            fa = wp.fabricarray(data=iface, attrib="foo")

            assert_np_equal(fa.numpy(), np.zeros((*fa.shape, vec_len), dtype=nptype))

            # fill with int scalar
            fill_value = 42
            fa.fill_(fill_value)
            assert_np_equal(fa.numpy(), np.full((*fa.shape, vec_len), fill_value, dtype=nptype))

            # test zeroing
            fa.zero_()
            assert_np_equal(fa.numpy(), np.zeros((*fa.shape, vec_len), dtype=nptype))

            # vector values can be passed as a list, numpy array, or Warp vector instance
            fill_list = [17, 42, 99, 101, 127][:vec_len]
            fill_arr = np.array(fill_list, dtype=nptype)
            fill_vec = vec_type(fill_list)

            expected = np.tile(fill_arr, fa.size).reshape((*fa.shape, vec_len))

            # fill with list of vector length
            fa.fill_(fill_list)
            assert_np_equal(fa.numpy(), expected)

            # clear
            fa.zero_()

            # fill with numpy array of vector length
            fa.fill_(fill_arr)
            assert_np_equal(fa.numpy(), expected)

            # clear
            fa.zero_()

            # fill with vec instance
            fa.fill_(fill_vec)
            assert_np_equal(fa.numpy(), expected)

            if wptype in wp.types.float_types:
                # fill with float scalar
                fill_value = 13.37
                fa.fill_(fill_value)
                assert_np_equal(fa.numpy(), np.full((*fa.shape, vec_len), fill_value, dtype=nptype))

                # fill with float list of vector length
                fill_list = [-2.5, -1.25, 1.25, 2.5, 5.0][:vec_len]

                fa.fill_(fill_list)

                expected = np.tile(np.array(fill_list, dtype=nptype), fa.size).reshape((*fa.shape, vec_len))

                assert_np_equal(fa.numpy(), expected)

            # reset data
            wp.copy(fa, data)

            # test indexed
            indices1 = wp.array(data=np.arange(1, data.size, 2, dtype=np.int32), device=device)
            ifa = fa[indices1]

            # ensure that the other indices remain unchanged
            indices2 = wp.array(data=np.arange(0, data.size, 2, dtype=np.int32), device=device)
            ifb = fa[indices2]

            assert_np_equal(ifa.numpy(), np.zeros((*ifa.shape, vec_len), dtype=nptype))
            assert_np_equal(ifb.numpy(), np.zeros((*ifb.shape, vec_len), dtype=nptype))

            # fill with int scalar
            fill_value = 42
            ifa.fill_(fill_value)
            assert_np_equal(ifa.numpy(), np.full((*ifa.shape, vec_len), fill_value, dtype=nptype))
            assert_np_equal(ifb.numpy(), np.zeros((*ifb.shape, vec_len), dtype=nptype))

            # test zeroing
            ifa.zero_()
            assert_np_equal(ifa.numpy(), np.zeros((*ifa.shape, vec_len), dtype=nptype))
            assert_np_equal(ifb.numpy(), np.zeros((*ifb.shape, vec_len), dtype=nptype))

            # vector values can be passed as a list, numpy array, or Warp vector instance
            fill_list = [17, 42, 99, 101, 127][:vec_len]
            fill_arr = np.array(fill_list, dtype=nptype)
            fill_vec = vec_type(fill_list)

            expected = np.tile(fill_arr, ifa.size).reshape((*ifa.shape, vec_len))

            # fill with list of vector length
            ifa.fill_(fill_list)
            assert_np_equal(ifa.numpy(), expected)
            assert_np_equal(ifb.numpy(), np.zeros((*ifb.shape, vec_len), dtype=nptype))

            # clear
            ifa.zero_()

            # fill with numpy array of vector length
            ifa.fill_(fill_arr)
            assert_np_equal(ifa.numpy(), expected)
            assert_np_equal(ifb.numpy(), np.zeros((*ifb.shape, vec_len), dtype=nptype))

            # clear
            ifa.zero_()

            # fill with vec instance
            ifa.fill_(fill_vec)
            assert_np_equal(ifa.numpy(), expected)
            assert_np_equal(ifb.numpy(), np.zeros((*ifb.shape, vec_len), dtype=nptype))

            if wptype in wp.types.float_types:
                # fill with float scalar
                fill_value = 13.37
                ifa.fill_(fill_value)
                assert_np_equal(ifa.numpy(), np.full((*ifa.shape, vec_len), fill_value, dtype=nptype))
                assert_np_equal(ifb.numpy(), np.zeros((*ifb.shape, vec_len), dtype=nptype))

                # fill with float list of vector length
                fill_list = [-2.5, -1.25, 1.25, 2.5, 5.0][:vec_len]

                ifa.fill_(fill_list)

                expected = np.tile(np.array(fill_list, dtype=nptype), ifa.size).reshape((*ifa.shape, vec_len))

                assert_np_equal(ifa.numpy(), expected)
                assert_np_equal(ifb.numpy(), np.zeros((*ifb.shape, vec_len), dtype=nptype))


def test_fabricarray_fill_matrix(test, device):
    # test filling a matrix array with scalar or matrix values (mat_type, nested list, or 2d numpy array)

    for nptype, wptype in wp.types.np_dtype_to_warp_type.items():
        # matrix types
        matrix_types = [
            # square matrices only
            wp.types.matrix((2, 2), wptype),
            wp.types.matrix((3, 3), wptype),
            wp.types.matrix((4, 4), wptype),
            wp.types.matrix((5, 5), wptype),
        ]

        for mat_type in matrix_types:
            mat_len = mat_type._length_
            mat_shape = mat_type._shape_

            data = wp.zeros(100, dtype=mat_type, device=device)
            iface = _create_fabric_array_interface(data, "foo", copy=True)
            fa = wp.fabricarray(data=iface, attrib="foo")

            assert_np_equal(fa.numpy(), np.zeros((*fa.shape, *mat_shape), dtype=nptype))

            # fill with scalar
            fill_value = 42
            fa.fill_(fill_value)
            assert_np_equal(fa.numpy(), np.full((*fa.shape, *mat_shape), fill_value, dtype=nptype))

            # test zeroing
            fa.zero_()
            assert_np_equal(fa.numpy(), np.zeros((*fa.shape, *mat_shape), dtype=nptype))

            # matrix values can be passed as a 1d numpy array, 2d numpy array, flat list, nested list, or Warp matrix instance
            if wptype != wp.bool:
                fill_arr1 = np.arange(mat_len, dtype=nptype)
            else:
                fill_arr1 = np.ones(mat_len, dtype=nptype)

            fill_arr2 = fill_arr1.reshape(mat_shape)
            fill_list1 = list(fill_arr1)
            fill_list2 = [list(row) for row in fill_arr2]
            fill_mat = mat_type(fill_arr1)

            expected = np.tile(fill_arr1, fa.size).reshape((*fa.shape, *mat_shape))

            # fill with 1d numpy array
            fa.fill_(fill_arr1)
            assert_np_equal(fa.numpy(), expected)

            # clear
            fa.zero_()

            # fill with 2d numpy array
            fa.fill_(fill_arr2)
            assert_np_equal(fa.numpy(), expected)

            # clear
            fa.zero_()

            # fill with flat list
            fa.fill_(fill_list1)
            assert_np_equal(fa.numpy(), expected)

            # clear
            fa.zero_()

            # fill with nested list
            fa.fill_(fill_list2)
            assert_np_equal(fa.numpy(), expected)

            # clear
            fa.zero_()

            # fill with mat instance
            fa.fill_(fill_mat)
            assert_np_equal(fa.numpy(), expected)

            # reset data
            wp.copy(fa, data)

            # test indexed
            indices1 = wp.array(data=np.arange(1, data.size, 2, dtype=np.int32), device=device)
            ifa = fa[indices1]

            # ensure that the other indices remain unchanged
            indices2 = wp.array(data=np.arange(0, data.size, 2, dtype=np.int32), device=device)
            ifb = fa[indices2]

            assert_np_equal(ifa.numpy(), np.zeros((*ifa.shape, *mat_shape), dtype=nptype))
            assert_np_equal(ifb.numpy(), np.zeros((*ifb.shape, *mat_shape), dtype=nptype))

            # fill with scalar
            fill_value = 42
            ifa.fill_(fill_value)
            assert_np_equal(ifa.numpy(), np.full((*ifa.shape, *mat_shape), fill_value, dtype=nptype))
            assert_np_equal(ifb.numpy(), np.zeros((*ifb.shape, *mat_shape), dtype=nptype))

            # test zeroing
            ifa.zero_()
            assert_np_equal(ifa.numpy(), np.zeros((*ifa.shape, *mat_shape), dtype=nptype))
            assert_np_equal(ifb.numpy(), np.zeros((*ifb.shape, *mat_shape), dtype=nptype))

            # matrix values can be passed as a 1d numpy array, 2d numpy array, flat list, nested list, or Warp matrix instance
            if wptype != wp.bool:
                fill_arr1 = np.arange(mat_len, dtype=nptype)
            else:
                fill_arr1 = np.ones(mat_len, dtype=nptype)
            fill_arr2 = fill_arr1.reshape(mat_shape)
            fill_list1 = list(fill_arr1)
            fill_list2 = [list(row) for row in fill_arr2]
            fill_mat = mat_type(fill_arr1)

            expected = np.tile(fill_arr1, ifa.size).reshape((*ifa.shape, *mat_shape))

            # fill with 1d numpy array
            ifa.fill_(fill_arr1)
            assert_np_equal(ifa.numpy(), expected)
            assert_np_equal(ifb.numpy(), np.zeros((*ifb.shape, *mat_shape), dtype=nptype))

            # clear
            ifa.zero_()

            # fill with 2d numpy array
            ifa.fill_(fill_arr2)
            assert_np_equal(ifa.numpy(), expected)
            assert_np_equal(ifb.numpy(), np.zeros((*ifb.shape, *mat_shape), dtype=nptype))

            # clear
            ifa.zero_()

            # fill with flat list
            ifa.fill_(fill_list1)
            assert_np_equal(ifa.numpy(), expected)
            assert_np_equal(ifb.numpy(), np.zeros((*ifb.shape, *mat_shape), dtype=nptype))

            # clear
            ifa.zero_()

            # fill with nested list
            ifa.fill_(fill_list2)
            assert_np_equal(ifa.numpy(), expected)
            assert_np_equal(ifb.numpy(), np.zeros((*ifb.shape, *mat_shape), dtype=nptype))

            # clear
            ifa.zero_()

            # fill with mat instance
            ifa.fill_(fill_mat)
            assert_np_equal(ifa.numpy(), expected)
            assert_np_equal(ifb.numpy(), np.zeros((*ifb.shape, *mat_shape), dtype=nptype))


@wp.kernel
def fa_generic_sums_kernel(a: wp.fabricarrayarray(dtype=Any), sums: wp.array(dtype=Any)):
    i = wp.tid()

    # get sub-array using wp::view()
    row = a[i]

    # get sub-array length
    count = row.shape[0]

    # compute sub-array sum
    for j in range(count):
        sums[i] = sums[i] + row[j]


@wp.kernel
def fa_generic_sums_kernel_indexed(a: wp.indexedfabricarrayarray(dtype=Any), sums: wp.array(dtype=Any)):
    i = wp.tid()

    # get sub-array using wp::view()
    row = a[i]

    # get sub-array length
    count = row.shape[0]

    # compute sub-array sum
    for j in range(count):
        sums[i] = sums[i] + row[j]


def test_fabricarrayarray(test, device):
    for T in _fabric_types:
        if hasattr(T, "_wp_scalar_type_"):
            nptype = wp.types.warp_type_to_np_dtype[T._wp_scalar_type_]
        else:
            nptype = wp.types.warp_type_to_np_dtype[T]

        n = 100

        min_length = 1
        max_length = 10
        arrays = []
        expected_sums = []
        expected_sums_indexed = []

        # generate data arrays
        length = min_length
        for i in range(n):
            if length > max_length:
                length = min_length

            na = np.arange(1, length + 1, dtype=nptype)

            arrays.append(wp.array(data=na, device=device))
            expected_sums.append(na.sum())

            # every second index
            if i % 2 == 0:
                expected_sums_indexed.append(na.sum())

            length += 1

        data_iface = _create_fabric_array_array_interface(arrays, "foo")
        fa = wp.fabricarrayarray(data=data_iface, attrib="foo")

        sums = wp.zeros_like(fa)

        test.assertEqual(fa.dtype, sums.dtype)
        test.assertEqual(fa.ndim, 2)
        test.assertEqual(sums.ndim, 1)
        test.assertEqual(fa.shape, sums.shape)
        test.assertEqual(fa.size, sums.size)

        wp.launch(fa_generic_sums_kernel, dim=fa.size, inputs=[fa, sums], device=device)

        assert_np_equal(sums.numpy(), np.array(expected_sums, dtype=nptype))

        # test indexed
        indices = wp.array(data=np.arange(0, n, 2, dtype=np.int32), device=device)
        ifa = fa[indices]

        sums = wp.zeros_like(ifa)

        test.assertEqual(ifa.dtype, sums.dtype)
        test.assertEqual(ifa.ndim, 2)
        test.assertEqual(sums.ndim, 1)
        test.assertEqual(ifa.shape, sums.shape)
        test.assertEqual(ifa.size, sums.size)

        wp.launch(fa_generic_sums_kernel_indexed, dim=ifa.size, inputs=[ifa, sums], device=device)

        assert_np_equal(sums.numpy(), np.array(expected_sums_indexed, dtype=nptype))


# explicit kernel overloads
for T in _fabric_types:
    wp.overload(fa_generic_dtype_kernel, [wp.fabricarray(dtype=T), wp.fabricarray(dtype=T)])
    wp.overload(fa_generic_dtype_kernel_indexed, [wp.indexedfabricarray(dtype=T), wp.indexedfabricarray(dtype=T)])

    wp.overload(fa_generic_array_kernel, [wp.fabricarray(dtype=T), wp.fabricarray(dtype=T)])
    wp.overload(fa_generic_array_kernel, [wp.indexedfabricarray(dtype=T), wp.indexedfabricarray(dtype=T)])

    wp.overload(fa_generic_sums_kernel, [wp.fabricarrayarray(dtype=T), wp.array(dtype=T)])
    wp.overload(fa_generic_sums_kernel_indexed, [wp.indexedfabricarrayarray(dtype=T), wp.array(dtype=T)])


devices = get_test_devices()


class TestFabricArray(unittest.TestCase):
    pass


# fabric arrays
add_function_test(TestFabricArray, "test_fabricarray_kernel", test_fabricarray_kernel, devices=devices)
add_function_test(TestFabricArray, "test_fabricarray_empty", test_fabricarray_empty, devices=devices)
add_function_test(TestFabricArray, "test_fabricarray_generic_dtype", test_fabricarray_generic_dtype, devices=devices)
add_function_test(TestFabricArray, "test_fabricarray_generic_array", test_fabricarray_generic_array, devices=devices)
add_function_test(TestFabricArray, "test_fabricarray_fill_scalar", test_fabricarray_fill_scalar, devices=devices)
add_function_test(TestFabricArray, "test_fabricarray_fill_vector", test_fabricarray_fill_vector, devices=devices)
add_function_test(TestFabricArray, "test_fabricarray_fill_matrix", test_fabricarray_fill_matrix, devices=devices)

# fabric arrays of arrays
add_function_test(TestFabricArray, "test_fabricarrayarray", test_fabricarrayarray, devices=devices)


if __name__ == "__main__":
    wp.build.clear_kernel_cache()
    unittest.main(verbosity=2)