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 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 #[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 #[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 #[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 #[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}