Function create_nvshmem_array

tests/python/tirx/codegen/test_codegen_nvshmem.py:49–62

Create and optionally initialize an nvshmem-accessible DNDArray.

(sess, shape, dtype, init_data_fn=None, zero_out=True)

def create_nvshmem_array(sess, shape, dtype, init_data_fn=None, zero_out=True):
    """Create and optionally initialize an nvshmem-accessible DNDArray."""
    nvshmem_empty = sess.get_global_func("runtime.disco.nvshmem.empty")
    arr = nvshmem_empty(ShapeTuple(shape), dtype, None)

    if init_data_fn:
        for i in range(NUM_WORKERS):
            arr.debug_copy_from(i, init_data_fn(i, shape, dtype))
    elif zero_out:
        zero_data = np.zeros(shape, dtype=dtype)
        for i in range(NUM_WORKERS):
            arr.debug_copy_from(i, zero_data)

    return arr


@pytest.mark.gpu