runmat_runtime/builtins/common/
gpu_helpers.rs

1use runmat_accelerate_api::{AccelProvider, GpuTensorHandle, GpuTensorStorage, HostTensorView};
2use runmat_builtins::{ComplexTensor, Tensor, Value};
3
4use crate::build_runtime_error;
5
6/// Download a GPU tensor handle to host memory, returning a dense `Tensor`.
7///
8/// This helper routes through the dispatcher so residency hooks and provider
9/// semantics stay consistent with the rest of the runtime.
10pub async fn gather_tensor_async(
11    handle: &runmat_accelerate_api::GpuTensorHandle,
12) -> crate::BuiltinResult<Tensor> {
13    let value = Value::GpuTensor(handle.clone());
14    let gathered = crate::dispatcher::gather_if_needed_async(&value).await?;
15    match gathered {
16        Value::Tensor(t) => Ok(t),
17        Value::Num(n) => Tensor::new(vec![n], vec![1, 1])
18            .map_err(|e| build_runtime_error(format!("gather: {e}")).build()),
19        Value::LogicalArray(la) => {
20            let data: Vec<f64> = la
21                .data
22                .iter()
23                .map(|&b| if b != 0 { 1.0 } else { 0.0 })
24                .collect();
25            Tensor::new(data, la.shape.clone())
26                .map_err(|e| build_runtime_error(format!("gather: {e}")).build())
27        }
28        other => {
29            Err(build_runtime_error(format!("gather: unexpected value kind {other:?}")).build())
30        }
31    }
32}
33
34/// Gather an arbitrary value, returning a host-side `Value`.
35pub async fn gather_value_async(value: &Value) -> crate::BuiltinResult<Value> {
36    crate::dispatcher::gather_if_needed_async(value).await
37}
38
39/// Upload a host complex tensor as an interleaved GPU buffer and record complex
40/// storage metadata on the returned handle.
41pub fn upload_complex_tensor(
42    provider: &dyn AccelProvider,
43    tensor: &ComplexTensor,
44) -> crate::BuiltinResult<GpuTensorHandle> {
45    let mut interleaved = Vec::with_capacity(tensor.data.len() * 2);
46    for &(re, im) in &tensor.data {
47        interleaved.push(re);
48        interleaved.push(im);
49    }
50    let view = HostTensorView {
51        data: &interleaved,
52        shape: &tensor.shape,
53    };
54    let handle = provider
55        .upload(&view)
56        .map_err(|e| build_runtime_error(format!("gpu upload: {e}")).build())?;
57    runmat_accelerate_api::set_handle_logical(&handle, false);
58    runmat_accelerate_api::set_handle_storage(&handle, GpuTensorStorage::ComplexInterleaved);
59    runmat_accelerate_api::set_handle_precision(&handle, provider.precision());
60    Ok(handle)
61}
62
63/// Wrap a GPU tensor handle, marking it as resident for downstream fusion-aware
64/// consumers and tests.
65pub fn resident_gpu_value(handle: GpuTensorHandle) -> Value {
66    runmat_accelerate_api::mark_residency(&handle);
67    Value::GpuTensor(handle)
68}
69
70/// Wrap a GPU tensor handle as a logical gpuArray value, recording metadata so that
71/// predicates like `islogical` can inspect the handle without downloading it.
72pub fn logical_gpu_value(handle: GpuTensorHandle) -> Value {
73    runmat_accelerate_api::set_handle_logical(&handle, true);
74    resident_gpu_value(handle)
75}
76
77/// Wrap a GPU tensor handle as a complex gpuArray value.
78pub fn complex_gpu_value(handle: GpuTensorHandle) -> Value {
79    runmat_accelerate_api::set_handle_logical(&handle, false);
80    runmat_accelerate_api::set_handle_storage(&handle, GpuTensorStorage::ComplexInterleaved);
81    resident_gpu_value(handle)
82}
runmat_runtime/builtins/common/gpu_helpers.rs

runmat_runtime/builtins/common/
gpu_helpers.rs
