CUDA Tensor Writes#

The write-flavors examples show CUDA input followed by CPU read-back. Both synchronize the producer before the write and keep its buffer alive until the write completes. A CUDA-capable GPU and working NVIDIA driver are required.

Python#

The Python source is examples/python/write-flavors/main.py. It runs the CPU workflow first and attempts the CUDA section when Warp and a CUDA device are available. Install warp-lang alongside the example’s dependencies and run main.py. If CUDA is unavailable, the program reports that GPU ingest was skipped; a successful CPU run does not validate the GPU section.

This excerpt writes a Warp CUDA array into the selected prims:

# Finish producing GPU data, then retain it until the write completes.
device_values = wp.array(np.array([5.0, 6.0, 7.0], np.float32), dtype=wp.float32, device="cuda:0")
wp.synchronize()
write = stage.write_attribute(
    query, gpu_samples, ordinal=8, tensors=device_values, is_array=False
)
# Keep device_values alive and unchanged while the write is in progress.
write.wait()
del device_values  # Safe to release our reference after completion.

C API#

The C++ source is examples/c/write-flavors/cuda_ingest.cpp. Enable its optional CMake target with the CUDA Toolkit installed:

cmake -S examples/c/write-flavors -B build/cuda-ingest -DOVSTAGE_EXAMPLE_CUDA=ON
cmake --build build/cuda-ingest --target cuda-ingest
./build/cuda-ingest/cuda-ingest

Run these commands from the public repository root. The executable returns 0 when CPU read-back matches the input, 1 for a value mismatch, and 77 when no CUDA device can be used. checkCuda() throws on CUDA errors; waitOp() checks ovstage submission and execution errors.

This excerpt constructs the CUDA tensor descriptor and keeps its allocation alive through the write:

float values[] = {5.0f, 6.0f, 7.0f};
float* deviceValues = nullptr;
checkCuda(cudaMalloc(reinterpret_cast<void**>(&deviceValues), sizeof(values)));
checkCuda(cudaMemcpyAsync(deviceValues, values, sizeof(values), cudaMemcpyHostToDevice));
// Finish producing the GPU data before submitting it to ovstage.
checkCuda(cudaDeviceSynchronize());

int64_t shape[] = {3};
int64_t strides[] = {1};
DLTensor tensor{};
tensor.data = deviceValues;
tensor.device = {kDLCUDA, 0};
tensor.ndim = 1;
tensor.dtype = {kDLFloat, 32, 1};
tensor.shape = shape;
tensor.strides = strides;
ovstage_write_data_t write{};
write.tensors = &tensor;
write.tensor_count = 1;
write.is_array = false;
write.cuda_sync = {}; // Producer already synchronized: no additional wait.
waitOp(stage, ovstage_write_attribute(stage, query, {attribute, {}}, 8,
                                     write, OVSTAGE_PRIM_MODE_UPSERT), "CUDA write");
// Keep the GPU allocation alive until the write completes.
checkCuda(cudaFree(deviceValues));

See Device Placement (Experimental) for the experimental GPU-to-GPU workflow, extension access, event ordering, and buffer lifetimes.

On this page