import warp as wp

# num threads per-tile
TILE_DIM = 32

@wp.kernel
def test_kernel(y: wp.array(dtype=wp.mat33)):
    i = wp.tid()
    I = wp.identity(3, dtype=wp.float32)
    m = wp.float32(i) * I

    t = wp.tile(m, preserve_type=True)
    sum = wp.tile_reduce(wp.add, t)

    wp.tile_store(y, sum)


y = wp.zeros(shape=1, dtype=wp.mat33)

wp.launch(test_kernel, dim=TILE_DIM, inputs=[], outputs=[y], block_dim=TILE_DIM)