Skip to main content

weir/summary/
dispatch.rs

1#![allow(clippy::too_many_arguments)]
2
3use super::{summarize_function_with_count, SummarizeFunctionScratch};
4use crate::dispatch_input_cache::OwnedDispatchInputs;
5use crate::soundness::{DataflowEvidence, Soundness};
6
7pub fn summarize_function_evidence_via<F>(
8    dispatch: &F,
9    fn_ast_in: &[u32],
10    callgraph_in: &[u32],
11    cached_summary_in: &[u32],
12    bit_count: u32,
13) -> Result<DataflowEvidence<Vec<u32>>, String>
14where
15    F: Fn(&vyre::ir::Program, &[Vec<u8>], Option<[u32; 3]>) -> Result<Vec<Vec<u8>>, String>,
16{
17    summarize_function_via(
18        dispatch,
19        fn_ast_in,
20        callgraph_in,
21        cached_summary_in,
22        bit_count,
23    )
24    .map(|value| DataflowEvidence::new(value, Soundness::MayOver))
25}
26
27pub fn summarize_function_via<F>(
28    dispatch: &F,
29    fn_ast_in: &[u32],
30    callgraph_in: &[u32],
31    cached_summary_in: &[u32],
32    bit_count: u32,
33) -> Result<Vec<u32>, String>
34where
35    F: Fn(&vyre::ir::Program, &[Vec<u8>], Option<[u32; 3]>) -> Result<Vec<Vec<u8>>, String>,
36{
37    let owned_inputs = OwnedDispatchInputs::new();
38    summarize_function_borrowed_via(
39        &|program, inputs, grid| owned_inputs.dispatch(program, inputs, grid, dispatch),
40        fn_ast_in,
41        callgraph_in,
42        cached_summary_in,
43        bit_count,
44    )
45}
46
47/// GPU function-summary projection using borrowed dispatch inputs.
48///
49/// This avoids constructing an owned input vector for backends that already
50/// accept borrowed byte slices.
51pub fn summarize_function_borrowed_via<F>(
52    dispatch: &F,
53    fn_ast_in: &[u32],
54    callgraph_in: &[u32],
55    cached_summary_in: &[u32],
56    bit_count: u32,
57) -> Result<Vec<u32>, String>
58where
59    F: Fn(&vyre::ir::Program, &[&[u8]], Option<[u32; 3]>) -> Result<Vec<Vec<u8>>, String>,
60{
61    let mut scratch = SummarizeFunctionScratch::default();
62    summarize_function_borrowed_with_scratch_via(
63        dispatch,
64        fn_ast_in,
65        callgraph_in,
66        cached_summary_in,
67        bit_count,
68        &mut scratch,
69    )
70}
71
72/// GPU function-summary projection using borrowed dispatch inputs and
73/// caller-owned staging/output scratch.
74pub fn summarize_function_borrowed_with_scratch_via<F>(
75    dispatch: &F,
76    fn_ast_in: &[u32],
77    callgraph_in: &[u32],
78    cached_summary_in: &[u32],
79    bit_count: u32,
80    scratch: &mut SummarizeFunctionScratch,
81) -> Result<Vec<u32>, String>
82where
83    F: Fn(&vyre::ir::Program, &[&[u8]], Option<[u32; 3]>) -> Result<Vec<Vec<u8>>, String>,
84{
85    let words =
86        crate::dispatch_decode::bitset_word_capacity("summarize_function result", bit_count)?
87            .max(1);
88    let mut result = crate::staging_reserve::reserved_vec(words, "summary result word")?;
89    summarize_function_borrowed_into_result_with_scratch_via(
90        dispatch,
91        fn_ast_in,
92        callgraph_in,
93        cached_summary_in,
94        bit_count,
95        scratch,
96        &mut result,
97    )?;
98    Ok(result)
99}
100
101/// GPU function-summary projection using caller-owned staging/output scratch
102/// and caller-owned decoded result storage.
103pub fn summarize_function_borrowed_into_result_with_scratch_via<F>(
104    dispatch: &F,
105    fn_ast_in: &[u32],
106    callgraph_in: &[u32],
107    cached_summary_in: &[u32],
108    bit_count: u32,
109    scratch: &mut SummarizeFunctionScratch,
110    result: &mut Vec<u32>,
111) -> Result<(), String>
112where
113    F: Fn(&vyre::ir::Program, &[&[u8]], Option<[u32; 3]>) -> Result<Vec<Vec<u8>>, String>,
114{
115    summarize_function_staged_into_result_via(
116        &|program, inputs, grid, outputs| {
117            let result = dispatch(program, inputs, grid)?;
118            crate::output_scratch::replace_outputs_preserving_slots(outputs, result);
119            Ok(())
120        },
121        fn_ast_in,
122        callgraph_in,
123        cached_summary_in,
124        bit_count,
125        &mut scratch.outputs,
126        &mut scratch.ast_bytes,
127        &mut scratch.callgraph_bytes,
128        &mut scratch.cached_summary_bytes,
129        &mut scratch.out_bytes,
130        result,
131    )?;
132    crate::dispatch_decode::require_bitset_tail_clear(
133        "summarize_function output",
134        result,
135        bit_count,
136    )
137}
138
139/// GPU function-summary projection using borrowed dispatch inputs and
140/// caller-owned output storage.
141pub fn summarize_function_borrowed_into_via<F>(
142    dispatch: &F,
143    fn_ast_in: &[u32],
144    callgraph_in: &[u32],
145    cached_summary_in: &[u32],
146    bit_count: u32,
147    outputs: &mut Vec<Vec<u8>>,
148) -> Result<Vec<u32>, String>
149where
150    F: Fn(&vyre::ir::Program, &[&[u8]], Option<[u32; 3]>, &mut Vec<Vec<u8>>) -> Result<(), String>,
151{
152    let semantic_words = crate::dispatch_decode::bitset_word_capacity(
153        "summarize_function semantic width",
154        bit_count,
155    )?;
156    let words = semantic_words.max(1);
157    let byte_capacity = words
158        .checked_mul(std::mem::size_of::<u32>())
159        .ok_or_else(|| {
160            "summarize_function byte capacity overflows host usize. Fix: split the summary bitset before GPU dispatch."
161                .to_string()
162        })?;
163    let mut ast_bytes = crate::staging_reserve::reserved_vec(byte_capacity, "summary AST byte")?;
164    let mut callgraph_bytes =
165        crate::staging_reserve::reserved_vec(byte_capacity, "summary callgraph byte")?;
166    let mut cached_summary_bytes =
167        crate::staging_reserve::reserved_vec(byte_capacity, "summary cached byte")?;
168    let mut out_bytes = crate::staging_reserve::reserved_vec(byte_capacity, "summary output byte")?;
169    let mut result = crate::staging_reserve::reserved_vec(words, "summary result word")?;
170    summarize_function_staged_into_result_via(
171        dispatch,
172        fn_ast_in,
173        callgraph_in,
174        cached_summary_in,
175        bit_count,
176        outputs,
177        &mut ast_bytes,
178        &mut callgraph_bytes,
179        &mut cached_summary_bytes,
180        &mut out_bytes,
181        &mut result,
182    )?;
183    Ok(result)
184}
185
186fn summarize_function_staged_into_result_via<F>(
187    dispatch: &F,
188    fn_ast_in: &[u32],
189    callgraph_in: &[u32],
190    cached_summary_in: &[u32],
191    bit_count: u32,
192    outputs: &mut Vec<Vec<u8>>,
193    ast_bytes: &mut Vec<u8>,
194    callgraph_bytes: &mut Vec<u8>,
195    cached_summary_bytes: &mut Vec<u8>,
196    out_bytes: &mut Vec<u8>,
197    result: &mut Vec<u32>,
198) -> Result<(), String>
199where
200    F: Fn(&vyre::ir::Program, &[&[u8]], Option<[u32; 3]>, &mut Vec<Vec<u8>>) -> Result<(), String>,
201{
202    let semantic_words = crate::dispatch_decode::bitset_word_capacity(
203        "summarize_function semantic width",
204        bit_count,
205    )?;
206    let words = semantic_words.max(1);
207    let program = summarize_function_with_count("ast", "cg", "cached", "out", bit_count);
208    crate::dispatch_decode::pack_exact_u32_slots_into(
209        fn_ast_in,
210        semantic_words,
211        "summarize_function ast",
212        ast_bytes,
213    )?;
214    crate::dispatch_decode::pack_exact_u32_slots_into(
215        callgraph_in,
216        semantic_words,
217        "summarize_function callgraph",
218        callgraph_bytes,
219    )?;
220    crate::dispatch_decode::pack_exact_u32_slots_into(
221        cached_summary_in,
222        semantic_words,
223        "summarize_function cached",
224        cached_summary_bytes,
225    )?;
226    crate::dispatch_decode::require_bitset_tail_clear(
227        "summarize_function ast",
228        fn_ast_in,
229        bit_count,
230    )?;
231    crate::dispatch_decode::require_bitset_tail_clear(
232        "summarize_function callgraph",
233        callgraph_in,
234        bit_count,
235    )?;
236    crate::dispatch_decode::require_bitset_tail_clear(
237        "summarize_function cached",
238        cached_summary_in,
239        bit_count,
240    )?;
241    crate::dispatch_decode::pad_dispatch_min_words(ast_bytes, words)?;
242    crate::dispatch_decode::pad_dispatch_min_words(callgraph_bytes, words)?;
243    crate::dispatch_decode::pad_dispatch_min_words(cached_summary_bytes, words)?;
244    let out_byte_len = words.checked_mul(4).ok_or_else(|| {
245        "summarize_function output byte length overflows host usize. Fix: split the summary bitset before GPU dispatch.".to_string()
246    })?;
247    let dispatch_words = u32::try_from(words)
248        .map_err(|_| "summarize_function word count exceeds u32 dispatch metadata. Fix: split the summary bitset before GPU dispatch.".to_string())?
249        .max(1);
250    crate::dispatch_decode::try_write_zero_bytes(
251        out_bytes,
252        out_byte_len,
253        "summarize_function output zero staging",
254    )?;
255    let inputs = [
256        ast_bytes.as_slice(),
257        callgraph_bytes.as_slice(),
258        cached_summary_bytes.as_slice(),
259        out_bytes.as_slice(),
260    ];
261    dispatch(&program, &inputs, Some([dispatch_words, 1, 1]), outputs)?;
262    crate::dispatch_decode::unpack_only_exact_u32_into(
263        outputs,
264        "summarize_function",
265        "summary",
266        words,
267        result,
268    )?;
269    crate::dispatch_decode::require_bitset_tail_clear(
270        "summarize_function output",
271        result,
272        bit_count,
273    )
274}