runmat_runtime/execution/
capture.rs

1use std::collections::HashSet;
2
3use runmat_accelerate_api::GpuTensorHandle;
4use runmat_value::Value;
5
6use crate::RuntimeError;
7
8/// Validate values captured by a lazy future before they can cross a spawn boundary.
9pub fn validate_spawn_capture(value: &Value) -> Result<(), RuntimeError> {
10    for_each_gpu_handle(value, &mut |handle| {
11        let provider = runmat_accelerate_api::provider_for_handle(handle).ok_or_else(|| {
12            crate::runtime_error::semantic_error(
13                "SpawnProviderUnavailable",
14                format!(
15                    "spawn cannot capture GPU handle buffer {} (device {}) without an active provider",
16                    handle.buffer_id, handle.device_id
17                ),
18            )
19        })?;
20        let policy = provider.spawn_handle_concurrency();
21        if matches!(
22            policy,
23            runmat_accelerate_api::SpawnHandleConcurrency::Reject
24        ) {
25            return Err(crate::runtime_error::semantic_error(
26                "SpawnGpuHandleUnsupported",
27                format!(
28                    "spawn cannot capture GPU handle buffer {} on provider '{}' (spawn_handle_concurrency={})",
29                    handle.buffer_id,
30                    provider.device_info(),
31                    policy.as_str()
32                ),
33            ));
34        }
35        Ok(())
36    })
37}
38
39fn for_each_gpu_handle(
40    value: &Value,
41    operation: &mut impl FnMut(&GpuTensorHandle) -> Result<(), RuntimeError>,
42) -> Result<(), RuntimeError> {
43    visit_gpu_handles(value, operation, &mut HashSet::new())
44}
45
46fn visit_gpu_handles(
47    value: &Value,
48    operation: &mut impl FnMut(&GpuTensorHandle) -> Result<(), RuntimeError>,
49    visited_handles: &mut HashSet<usize>,
50) -> Result<(), RuntimeError> {
51    match value {
52        Value::GpuTensor(handle) => operation(handle),
53        Value::Cell(cell) => visit_values(&cell.data, operation, visited_handles),
54        Value::Struct(value) => {
55            for value in value.fields.values() {
56                visit_gpu_handles(value, operation, visited_handles)?;
57            }
58            Ok(())
59        }
60        Value::Object(value) => {
61            for value in value.properties.values() {
62                visit_gpu_handles(value, operation, visited_handles)?;
63            }
64            Ok(())
65        }
66        Value::ObjectArray(value) => visit_values(value.data(), operation, visited_handles),
67        Value::Closure(value) => visit_values(&value.captures, operation, visited_handles),
68        Value::OutputList(values) => visit_values(values, operation, visited_handles),
69        Value::HandleObject(handle) => {
70            let address = runmat_gc::gc_handle_addr(&handle.target);
71            if visited_handles.insert(address) {
72                runmat_gc::gc_with_value(&handle.target, |value| {
73                    visit_gpu_handles(value, operation, visited_handles)
74                })
75                .map_err(|error| RuntimeError::new(format!("invalid handle target: {error}")))??;
76            }
77            Ok(())
78        }
79        Value::Int(_)
80        | Value::Num(_)
81        | Value::Complex(_, _)
82        | Value::Bool(_)
83        | Value::LogicalArray(_)
84        | Value::String(_)
85        | Value::StringArray(_)
86        | Value::CharArray(_)
87        | Value::Symbolic(_)
88        | Value::SymbolicArray(_)
89        | Value::Tensor(_)
90        | Value::SparseTensor(_)
91        | Value::ComplexTensor(_)
92        | Value::Listener(_)
93        | Value::FunctionHandle(_)
94        | Value::ExternalFunctionHandle(_)
95        | Value::MethodFunctionHandle(_)
96        | Value::BoundFunctionHandle { .. }
97        | Value::ClassRef(_)
98        | Value::MException(_)
99        | Value::Future(_)
100        | Value::Task(_)
101        | Value::Pool(_)
102        | Value::Job(_)
103        | Value::Foreign(_) => Ok(()),
104    }
105}
106
107fn visit_values(
108    values: &[Value],
109    operation: &mut impl FnMut(&GpuTensorHandle) -> Result<(), RuntimeError>,
110    visited_handles: &mut HashSet<usize>,
111) -> Result<(), RuntimeError> {
112    for value in values {
113        visit_gpu_handles(value, operation, visited_handles)?;
114    }
115    Ok(())
116}
117
118#[cfg(test)]
119mod tests {
120    use super::*;
121    use runmat_accelerate_api::{
122        AccelDownloadFuture, AccelProvider, HostTensorView, SpawnHandleConcurrency,
123        ThreadProviderGuard,
124    };
125    use runmat_value::{CellArray, HandleRef, StructValue};
126
127    struct RejectProvider;
128    static REJECT_PROVIDER: RejectProvider = RejectProvider;
129
130    impl AccelProvider for RejectProvider {
131        fn upload(&self, _host: &HostTensorView) -> anyhow::Result<GpuTensorHandle> {
132            Err(anyhow::anyhow!("unsupported"))
133        }
134
135        fn download<'a>(&'a self, _handle: &'a GpuTensorHandle) -> AccelDownloadFuture<'a> {
136            Box::pin(async { Err(anyhow::anyhow!("unsupported")) })
137        }
138
139        fn free(&self, _handle: &GpuTensorHandle) -> anyhow::Result<()> {
140            Ok(())
141        }
142
143        fn device_info(&self) -> String {
144            "reject-provider".into()
145        }
146
147        fn device_id(&self) -> u32 {
148            41
149        }
150    }
151
152    struct ShareProvider;
153    static SHARE_PROVIDER: ShareProvider = ShareProvider;
154
155    impl AccelProvider for ShareProvider {
156        fn upload(&self, _host: &HostTensorView) -> anyhow::Result<GpuTensorHandle> {
157            Err(anyhow::anyhow!("unsupported"))
158        }
159
160        fn download<'a>(&'a self, _handle: &'a GpuTensorHandle) -> AccelDownloadFuture<'a> {
161            Box::pin(async { Err(anyhow::anyhow!("unsupported")) })
162        }
163
164        fn free(&self, _handle: &GpuTensorHandle) -> anyhow::Result<()> {
165            Ok(())
166        }
167
168        fn device_info(&self) -> String {
169            "share-provider".into()
170        }
171
172        fn device_id(&self) -> u32 {
173            42
174        }
175
176        fn spawn_handle_concurrency(&self) -> SpawnHandleConcurrency {
177            SpawnHandleConcurrency::ImmutableShare
178        }
179    }
180
181    fn gpu(device_id: u32, buffer_id: u64) -> Value {
182        Value::GpuTensor(GpuTensorHandle {
183            shape: vec![1],
184            device_id,
185            buffer_id,
186            descriptor: Default::default(),
187        })
188    }
189
190    #[test]
191    fn spawn_capture_obeys_provider_concurrency_policy() {
192        let _guard = ThreadProviderGuard::set(Some(&REJECT_PROVIDER));
193        let error = validate_spawn_capture(&gpu(41, 7)).expect_err("reject capture");
194        assert_eq!(error.identifier(), Some("RunMat:SpawnGpuHandleUnsupported"));
195        drop(_guard);
196
197        let _guard = ThreadProviderGuard::set(Some(&SHARE_PROVIDER));
198        validate_spawn_capture(&gpu(42, 9)).expect("immutable sharing is safe");
199    }
200
201    #[test]
202    fn spawn_capture_rejects_missing_provider() {
203        let _guard = ThreadProviderGuard::set(None);
204        let error = validate_spawn_capture(&gpu(99, 13)).expect_err("missing provider");
205        assert_eq!(error.identifier(), Some("RunMat:SpawnProviderUnavailable"));
206    }
207
208    #[test]
209    fn spawn_capture_recurses_through_cells_closures_and_handle_objects() {
210        let _guard = ThreadProviderGuard::set(Some(&REJECT_PROVIDER));
211        let cell = Value::Cell(
212            CellArray::new(vec![Value::Num(1.0), gpu(41, 11)], 1, 2).expect("test cell"),
213        );
214        assert_eq!(
215            validate_spawn_capture(&cell).unwrap_err().identifier(),
216            Some("RunMat:SpawnGpuHandleUnsupported")
217        );
218
219        let closure = Value::Closure(runmat_value::Closure {
220            function_name: "worker".into(),
221            bound_function: None,
222            captures: vec![gpu(41, 21)],
223        });
224        assert_eq!(
225            validate_spawn_capture(&closure).unwrap_err().identifier(),
226            Some("RunMat:SpawnGpuHandleUnsupported")
227        );
228
229        let mut payload = StructValue::new();
230        payload.fields.insert("nested".into(), gpu(41, 31));
231        let target = runmat_gc::gc_allocate(Value::Struct(payload)).expect("gc payload");
232        let object = Value::HandleObject(HandleRef {
233            class_name: "Payload".into(),
234            target,
235            valid: true,
236        });
237        assert_eq!(
238            validate_spawn_capture(&object).unwrap_err().identifier(),
239            Some("RunMat:SpawnGpuHandleUnsupported")
240        );
241    }
242}
runmat_runtime/execution/capture.rs

runmat_runtime/execution/
capture.rs
