Skip to main content

weir/fixed_point_batch/ifds/
direct.rs

1#![allow(clippy::too_many_arguments)]
2
3use crate::fixed_point_batch::FixedPointBatch;
4
5impl<'a, F> FixedPointBatch<'a, F>
6where
7    F: Fn(&vyre::ir::Program, &[&[u8]], Option<[u32; 3]>, &mut Vec<Vec<u8>>) -> Result<(), String>,
8{
9    /// Prepare IFDS CSR buffers for repeated batch solves.
10    pub fn prepare_ifds(
11        &mut self,
12        num_procs: u32,
13        blocks_per_proc: u32,
14        facts_per_proc: u32,
15        intra_edges: &[(u32, u32, u32)],
16        inter_edges: &[(u32, u32, u32, u32)],
17        flow_gen: &[(u32, u32, u32)],
18        flow_kill: &[(u32, u32, u32)],
19    ) -> Result<crate::ifds_gpu::PreparedIfdsCsr, String> {
20        crate::ifds_gpu::prepare_ifds_csr_borrowed_with_scratch_via(
21            self.dispatch,
22            num_procs,
23            blocks_per_proc,
24            facts_per_proc,
25            intra_edges,
26            inter_edges,
27            flow_gen,
28            flow_kill,
29            &mut self.ifds_prepare_scratch,
30        )
31    }
32
33    /// Run IFDS reachability while reusing retained batch preparation and solve scratch.
34    #[allow(clippy::too_many_arguments)]
35    pub fn ifds(
36        &mut self,
37        num_procs: u32,
38        blocks_per_proc: u32,
39        facts_per_proc: u32,
40        intra_edges: &[(u32, u32, u32)],
41        inter_edges: &[(u32, u32, u32, u32)],
42        flow_gen: &[(u32, u32, u32)],
43        flow_kill: &[(u32, u32, u32)],
44        seed_facts: &[(u32, u32, u32)],
45        max_iterations: u32,
46    ) -> Result<Vec<u32>, String> {
47        if seed_facts.is_empty() {
48            crate::dispatch_decode::require_positive_iterations(
49                "weir fixed-point batch IFDS solve",
50                max_iterations,
51            )?;
52            return Ok(Vec::new());
53        }
54        let prepared = crate::ifds_gpu::prepare_ifds_csr_borrowed_with_scratch_via(
55            self.dispatch,
56            num_procs,
57            blocks_per_proc,
58            facts_per_proc,
59            intra_edges,
60            inter_edges,
61            flow_gen,
62            flow_kill,
63            &mut self.ifds_prepare_scratch,
64        )?;
65        crate::ifds_gpu::solve_prepared_borrowed_with_scratch_via(
66            self.dispatch,
67            &prepared,
68            seed_facts,
69            max_iterations,
70            &mut self.ifds_scratch,
71        )
72    }
73
74    /// Run IFDS reachability into caller-owned result storage while reusing
75    /// retained batch preparation and solve scratch.
76    #[allow(clippy::too_many_arguments)]
77    pub fn ifds_into(
78        &mut self,
79        num_procs: u32,
80        blocks_per_proc: u32,
81        facts_per_proc: u32,
82        intra_edges: &[(u32, u32, u32)],
83        inter_edges: &[(u32, u32, u32, u32)],
84        flow_gen: &[(u32, u32, u32)],
85        flow_kill: &[(u32, u32, u32)],
86        seed_facts: &[(u32, u32, u32)],
87        max_iterations: u32,
88        result: &mut Vec<u32>,
89    ) -> Result<(), String> {
90        if seed_facts.is_empty() {
91            crate::dispatch_decode::require_positive_iterations(
92                "weir fixed-point batch IFDS solve into",
93                max_iterations,
94            )?;
95            result.clear();
96            return Ok(());
97        }
98        let prepared = crate::ifds_gpu::prepare_ifds_csr_borrowed_with_scratch_via(
99            self.dispatch,
100            num_procs,
101            blocks_per_proc,
102            facts_per_proc,
103            intra_edges,
104            inter_edges,
105            flow_gen,
106            flow_kill,
107            &mut self.ifds_prepare_scratch,
108        )?;
109        crate::ifds_gpu::solve_prepared_borrowed_with_scratch_into_via(
110            self.dispatch,
111            &prepared,
112            seed_facts,
113            max_iterations,
114            &mut self.ifds_scratch,
115            result,
116        )
117    }
118
119    /// Run several IFDS reachability queries while preparing the IFDS CSR once
120    /// and reusing retained batch scratch across every seed set.
121    #[allow(clippy::too_many_arguments)]
122    pub fn ifds_many(
123        &mut self,
124        num_procs: u32,
125        blocks_per_proc: u32,
126        facts_per_proc: u32,
127        intra_edges: &[(u32, u32, u32)],
128        inter_edges: &[(u32, u32, u32, u32)],
129        flow_gen: &[(u32, u32, u32)],
130        flow_kill: &[(u32, u32, u32)],
131        seed_sets: &[&[(u32, u32, u32)]],
132        max_iterations: u32,
133    ) -> Result<Vec<Vec<u32>>, String> {
134        if seed_sets.is_empty() {
135            return Ok(Vec::new());
136        }
137        if seed_sets.iter().all(|seed_facts| seed_facts.is_empty()) {
138            crate::dispatch_decode::require_positive_iterations(
139                "weir fixed-point batch many IFDS solve",
140                max_iterations,
141            )?;
142            return crate::staging_reserve::reserved_result_rows(
143                seed_sets.len(),
144                "fixed-point batch empty IFDS result row",
145            );
146        }
147        let prepared = self.prepare_ifds(
148            num_procs,
149            blocks_per_proc,
150            facts_per_proc,
151            intra_edges,
152            inter_edges,
153            flow_gen,
154            flow_kill,
155        )?;
156        self.ifds_prepared_many(&prepared, seed_sets, max_iterations)
157    }
158
159    /// Run several IFDS reachability queries while preparing the IFDS CSR once
160    /// and writing into caller-owned result rows.
161    #[allow(clippy::too_many_arguments)]
162    pub fn ifds_many_into(
163        &mut self,
164        num_procs: u32,
165        blocks_per_proc: u32,
166        facts_per_proc: u32,
167        intra_edges: &[(u32, u32, u32)],
168        inter_edges: &[(u32, u32, u32, u32)],
169        flow_gen: &[(u32, u32, u32)],
170        flow_kill: &[(u32, u32, u32)],
171        seed_sets: &[&[(u32, u32, u32)]],
172        max_iterations: u32,
173        results: &mut Vec<Vec<u32>>,
174    ) -> Result<(), String> {
175        if seed_sets.is_empty() {
176            results.clear();
177            return Ok(());
178        }
179        if seed_sets.iter().all(|seed_facts| seed_facts.is_empty()) {
180            crate::dispatch_decode::require_positive_iterations(
181                "weir fixed-point batch many IFDS solve into",
182                max_iterations,
183            )?;
184            crate::staging_reserve::resize_result_rows(
185                results,
186                seed_sets.len(),
187                "fixed-point batch IFDS result row",
188            )?;
189            crate::staging_reserve::clear_result_rows(results);
190            return Ok(());
191        }
192        let prepared = self.prepare_ifds(
193            num_procs,
194            blocks_per_proc,
195            facts_per_proc,
196            intra_edges,
197            inter_edges,
198            flow_gen,
199            flow_kill,
200        )?;
201        self.ifds_prepared_many_into(&prepared, seed_sets, max_iterations, results)
202    }
203}