1use xlog_core::{Result, ScalarType, XlogError};
36
37use super::{wcoj_kernels, CudaKernelProvider, WCOJ_MODULE};
38use crate::device_runtime::StreamId;
39use crate::launch::LaunchRecorder;
40use crate::CudaBuffer;
41use crate::{LaunchAsync, LaunchConfig};
42
43const BLOCK_SIZE: u32 = 256;
44
45pub const FJ_DELTA_MAX_DOMAIN: u32 = 1 << 16;
48
49#[derive(Debug, Clone, Copy)]
56pub struct FjDeltaCols {
57 pub delta_carry: usize,
59 pub delta_key: usize,
61 pub r_carry: usize,
63 pub r_value: usize,
65}
66
67impl FjDeltaCols {
68 pub const CANONICAL: Self = Self {
70 delta_carry: 0,
71 delta_key: 1,
72 r_carry: 0,
73 r_value: 1,
74 };
75
76 fn validate(&self, ctx: &str) -> Result<()> {
77 let ok = |a: usize, b: usize| a < 2 && b < 2 && a != b;
78 if !ok(self.delta_carry, self.delta_key) || !ok(self.r_carry, self.r_value) {
79 return Err(XlogError::Kernel(format!(
80 "{ctx}: invalid column roles {self:?} (arity-2 indices, pairwise distinct)"
81 )));
82 }
83 Ok(())
84 }
85}
86
87fn require_binary_u32_class(buf: &CudaBuffer, name: &str, ctx: &str) -> Result<()> {
88 if buf.arity() != 2 {
89 return Err(XlogError::Kernel(format!(
90 "{ctx}: {name} must be arity-2, got {}",
91 buf.arity()
92 )));
93 }
94 for idx in 0..2 {
95 match buf.schema().column_type(idx) {
96 Some(ScalarType::U32) | Some(ScalarType::Symbol) => {}
97 other => {
98 return Err(XlogError::Kernel(format!(
99 "{ctx}: {name} column {idx} must be U32/Symbol, got {other:?}"
100 )));
101 }
102 }
103 }
104 Ok(())
105}
106
107impl CudaKernelProvider {
108 pub fn fj_delta_novel_u32_recorded(
118 &self,
119 delta: &CudaBuffer,
120 edge: &CudaBuffer,
121 full_r: &CudaBuffer,
122 cols: FjDeltaCols,
123 domain: u32,
124 launch_stream: StreamId,
125 ) -> Result<CudaBuffer> {
126 let ctx = "fj_delta_novel_u32_recorded";
127 let runtime = self.memory().runtime().ok_or_else(|| {
128 XlogError::Kernel(format!(
129 "{ctx} requires a runtime-backed GpuMemoryManager \
130 (constructed via with_runtime)"
131 ))
132 })?;
133 let cu_stream = runtime
134 .stream_pool()
135 .resolve(launch_stream)
136 .ok_or_else(|| {
137 XlogError::Kernel(format!(
138 "{ctx}: launch_stream StreamId({}) does not resolve",
139 launch_stream.0
140 ))
141 })?;
142
143 cols.validate(ctx)?;
144 require_binary_u32_class(delta, "delta", ctx)?;
145 require_binary_u32_class(edge, "edge", ctx)?;
146 require_binary_u32_class(full_r, "full_r", ctx)?;
147 if domain == 0 || domain > FJ_DELTA_MAX_DOMAIN {
148 return Err(XlogError::Kernel(format!(
149 "{ctx}: domain {domain} outside (0, {FJ_DELTA_MAX_DOMAIN}] \
150 (dense-domain spike bound)"
151 )));
152 }
153
154 let row_count = |buf: &CudaBuffer| -> Result<u32> {
155 match buf.cached_row_count() {
156 Some(c) => Ok(c),
157 None => self.dtoh_scalar_untracked::<u32>(buf.num_rows_device(), 0),
158 }
159 };
160 let n_delta = row_count(delta)?;
161 let n_edge = row_count(edge)?;
162 let n_r = row_count(full_r)?;
163
164 let out_schema = full_r.schema().clone();
165 if n_delta == 0 || n_edge == 0 {
166 return self.create_empty_buffer(out_schema);
167 }
168
169 let words_per_row = domain.div_ceil(32);
170 let n_words = u64::from(domain) * u64::from(words_per_row);
171 let n_words = n_words as u32;
173
174 let delta_x = delta.column(cols.delta_carry).ok_or_else(|| {
175 XlogError::Kernel(format!("{ctx}: delta column {} missing", cols.delta_carry))
176 })?;
177 let delta_y = delta.column(cols.delta_key).ok_or_else(|| {
178 XlogError::Kernel(format!("{ctx}: delta column {} missing", cols.delta_key))
179 })?;
180 let edge_y = edge
181 .column(0)
182 .ok_or_else(|| XlogError::Kernel(format!("{ctx}: edge column 0 missing")))?;
183 let edge_z = edge
184 .column(1)
185 .ok_or_else(|| XlogError::Kernel(format!("{ctx}: edge column 1 missing")))?;
186 let r_x = full_r.column(cols.r_carry).ok_or_else(|| {
187 XlogError::Kernel(format!("{ctx}: full_r column {} missing", cols.r_carry))
188 })?;
189 let r_z = full_r.column(cols.r_value).ok_or_else(|| {
190 XlogError::Kernel(format!("{ctx}: full_r column {} missing", cols.r_value))
191 })?;
192 let delta_x_v = self.column_as_u32_view(delta_x, n_delta as usize)?;
193 let delta_y_v = self.column_as_u32_view(delta_y, n_delta as usize)?;
194 let edge_y_v = self.column_as_u32_view(edge_y, n_edge as usize)?;
195 let edge_z_v = self.column_as_u32_view(edge_z, n_edge as usize)?;
196
197 let range_lo = self.memory().alloc::<u32>(n_delta as usize)?;
199 let mut wp = self.memory().alloc::<u32>(n_delta as usize + 1)?;
200 {
201 let mut scan_scratch = self.multiblock_scan_u32_scratch_for_len(n_delta + 1)?;
202 let mut rec = LaunchRecorder::new_strict(launch_stream);
203 rec.read_column(delta_y);
204 rec.read_column(edge_y);
205 rec.write(&range_lo);
206 rec.write(&wp);
207 for level in scan_scratch.levels() {
208 rec.read_write(level);
209 }
210 rec.preflight(runtime)
211 .map_err(|e| XlogError::Kernel(format!("{ctx}: range preflight failed: {e}")))?;
212 let kernel = self
213 .device()
214 .inner()
215 .get_func(WCOJ_MODULE, wcoj_kernels::FJ_DELTA_RANGE_U32)
216 .ok_or_else(|| {
217 XlogError::Kernel("fj_delta_range_u32 kernel not found".to_string())
218 })?;
219 let grid = n_delta.div_ceil(BLOCK_SIZE);
220 let enqueue = |stream: &crate::launch::CudaEnqueue<'_>| -> Result<()> {
224 unsafe {
225 kernel
226 .clone()
227 .launch_in(
228 stream,
229 LaunchConfig {
230 grid_dim: (grid, 1, 1),
231 block_dim: (BLOCK_SIZE, 1, 1),
232 shared_mem_bytes: 0,
233 },
234 (&delta_y_v, n_delta, &edge_y_v, n_edge, &range_lo, &mut wp),
235 )
236 .map_err(|e| {
237 XlogError::Kernel(format!("fj_delta_range_u32 launch failed: {e}"))
238 })?;
239 }
240 self.multiblock_scan_u32_inplace_on_stream(
241 &mut wp,
242 n_delta + 1,
243 stream,
244 &mut scan_scratch,
245 )?;
246 Ok(())
247 };
248 let rec = unsafe { rec.enqueue_prepared_with(&cu_stream, enqueue) }
250 .map_err(|error| error.into_xlog_error())?;
251 rec.commit()
252 .map_err(|e| XlogError::Kernel(format!("{ctx}: range commit failed: {e}")))?;
253 }
254 cu_stream
255 .synchronize()
256 .map_err(|e| XlogError::Kernel(format!("{ctx}: range sync failed: {e}")))?;
257 let total_work = u64::from(self.dtoh_scalar_untracked::<u32>(&wp, n_delta as usize)?);
258 if total_work == 0 {
259 return self.create_empty_buffer(out_schema);
260 }
261 if total_work > u64::from(u32::MAX - 1) {
262 return Err(XlogError::Kernel(format!(
263 "{ctx}: candidate work {total_work} exceeds the u32 work-index space"
264 )));
265 }
266 let total_work = total_work as u32;
267
268 let mut bitmap = self.memory().alloc::<u32>(n_words as usize)?;
270 self.device()
271 .inner()
272 .memset_zeros(&mut bitmap)
273 .map_err(|e| XlogError::Kernel(format!("{ctx}: zero bitmap failed: {e}")))?;
274 let mut error_flag = self.memory().alloc::<u32>(1)?;
275 self.device()
276 .inner()
277 .memset_zeros(&mut error_flag)
278 .map_err(|e| XlogError::Kernel(format!("{ctx}: zero error flag failed: {e}")))?;
279 {
280 let mut rec = LaunchRecorder::new_strict(launch_stream);
281 rec.read_column(delta_x);
282 rec.read_column(edge_z);
283 rec.read(&range_lo);
284 rec.read(&wp);
285 rec.read_write(&bitmap);
286 rec.write(&error_flag);
287 if n_r > 0 {
288 rec.read_column(r_x);
289 rec.read_column(r_z);
290 }
291 rec.preflight(runtime)
292 .map_err(|e| XlogError::Kernel(format!("{ctx}: mark preflight failed: {e}")))?;
293 let mark = self
294 .device()
295 .inner()
296 .get_func(WCOJ_MODULE, wcoj_kernels::FJ_DELTA_MARK_U32)
297 .ok_or_else(|| {
298 XlogError::Kernel("fj_delta_mark_u32 kernel not found".to_string())
299 })?;
300 let grid = total_work.div_ceil(BLOCK_SIZE);
301 let enqueue = |stream: &crate::launch::CudaEnqueue<'_>| -> Result<()> {
305 unsafe {
306 mark.clone()
307 .launch_in(
308 stream,
309 LaunchConfig {
310 grid_dim: (grid, 1, 1),
311 block_dim: (BLOCK_SIZE, 1, 1),
312 shared_mem_bytes: 0,
313 },
314 (
315 &delta_x_v,
316 n_delta,
317 &range_lo,
318 &wp,
319 total_work,
320 &edge_z_v,
321 &mut bitmap,
322 words_per_row,
323 domain,
324 &mut error_flag,
325 ),
326 )
327 .map_err(|e| {
328 XlogError::Kernel(format!("fj_delta_mark_u32 launch failed: {e}"))
329 })?;
330 }
331 if n_r > 0 {
332 let r_x_v = self.column_as_u32_view(r_x, n_r as usize)?;
333 let r_z_v = self.column_as_u32_view(r_z, n_r as usize)?;
334 let subtract = self
335 .device()
336 .inner()
337 .get_func(WCOJ_MODULE, wcoj_kernels::FJ_DELTA_SUBTRACT_U32)
338 .ok_or_else(|| {
339 XlogError::Kernel("fj_delta_subtract_u32 kernel not found".to_string())
340 })?;
341 let grid = n_r.div_ceil(BLOCK_SIZE);
342 unsafe {
346 subtract
347 .clone()
348 .launch_in(
349 stream,
350 LaunchConfig {
351 grid_dim: (grid, 1, 1),
352 block_dim: (BLOCK_SIZE, 1, 1),
353 shared_mem_bytes: 0,
354 },
355 (
356 &r_x_v,
357 &r_z_v,
358 n_r,
359 &mut bitmap,
360 words_per_row,
361 domain,
362 &mut error_flag,
363 ),
364 )
365 .map_err(|e| {
366 XlogError::Kernel(format!(
367 "fj_delta_subtract_u32 launch failed: {e}"
368 ))
369 })?;
370 }
371 }
372 Ok(())
373 };
374 let rec = unsafe { rec.enqueue_prepared_with(&cu_stream, enqueue) }
376 .map_err(|error| error.into_xlog_error())?;
377 rec.commit()
378 .map_err(|e| XlogError::Kernel(format!("{ctx}: mark commit failed: {e}")))?;
379 }
380 cu_stream
381 .synchronize()
382 .map_err(|e| XlogError::Kernel(format!("{ctx}: mark sync failed: {e}")))?;
383 if self.dtoh_scalar_untracked::<u32>(&error_flag, 0)? != 0 {
384 return Err(XlogError::Kernel(format!(
385 "{ctx}: id outside domain {domain} (fail-closed; raise domain or \
386 renumber the fixture)"
387 )));
388 }
389
390 let mut counts = self.memory().alloc::<u32>(n_words as usize + 1)?;
392 {
393 let mut scan_scratch = self.multiblock_scan_u32_scratch_for_len(n_words + 1)?;
394 let mut rec = LaunchRecorder::new_strict(launch_stream);
395 rec.read(&bitmap);
396 rec.write(&counts);
397 for level in scan_scratch.levels() {
398 rec.read_write(level);
399 }
400 rec.preflight(runtime)
401 .map_err(|e| XlogError::Kernel(format!("{ctx}: count preflight failed: {e}")))?;
402 let popcount = self
403 .device()
404 .inner()
405 .get_func(WCOJ_MODULE, wcoj_kernels::FJ_DELTA_POPCOUNT)
406 .ok_or_else(|| {
407 XlogError::Kernel("fj_delta_popcount kernel not found".to_string())
408 })?;
409 let grid = n_words.div_ceil(BLOCK_SIZE);
410 let enqueue = |stream: &crate::launch::CudaEnqueue<'_>| -> Result<()> {
412 unsafe {
413 popcount
414 .clone()
415 .launch_in(
416 stream,
417 LaunchConfig {
418 grid_dim: (grid, 1, 1),
419 block_dim: (BLOCK_SIZE, 1, 1),
420 shared_mem_bytes: 0,
421 },
422 (&bitmap, n_words, &mut counts),
423 )
424 .map_err(|e| {
425 XlogError::Kernel(format!("fj_delta_popcount launch failed: {e}"))
426 })?;
427 }
428 self.multiblock_scan_u32_inplace_on_stream(
429 &mut counts,
430 n_words + 1,
431 stream,
432 &mut scan_scratch,
433 )?;
434 Ok(())
435 };
436 let rec = unsafe { rec.enqueue_prepared_with(&cu_stream, enqueue) }
438 .map_err(|error| error.into_xlog_error())?;
439 rec.commit()
440 .map_err(|e| XlogError::Kernel(format!("{ctx}: count commit failed: {e}")))?;
441 }
442 cu_stream
443 .synchronize()
444 .map_err(|e| XlogError::Kernel(format!("{ctx}: count sync failed: {e}")))?;
445 let total_novel = self.dtoh_scalar_untracked::<u32>(&counts, n_words as usize)?;
446 if total_novel == 0 {
447 return self.create_empty_buffer(out_schema);
448 }
449
450 let out_x = self.memory().alloc::<u32>(total_novel as usize)?;
451 let out_z = self.memory().alloc::<u32>(total_novel as usize)?;
452 {
453 let mut rec = LaunchRecorder::new_strict(launch_stream);
454 rec.read(&bitmap);
455 rec.read(&counts);
456 rec.write(&out_x);
457 rec.write(&out_z);
458 rec.preflight(runtime)
459 .map_err(|e| XlogError::Kernel(format!("{ctx}: emit preflight failed: {e}")))?;
460 let emit = self
461 .device()
462 .inner()
463 .get_func(WCOJ_MODULE, wcoj_kernels::FJ_DELTA_EMIT_U32)
464 .ok_or_else(|| {
465 XlogError::Kernel("fj_delta_emit_u32 kernel not found".to_string())
466 })?;
467 let grid = n_words.div_ceil(BLOCK_SIZE);
468 let enqueue = |stream: &crate::launch::CudaEnqueue<'_>| -> Result<()> {
471 unsafe {
472 emit.clone()
473 .launch_in(
474 stream,
475 LaunchConfig {
476 grid_dim: (grid, 1, 1),
477 block_dim: (BLOCK_SIZE, 1, 1),
478 shared_mem_bytes: 0,
479 },
480 (&bitmap, words_per_row, n_words, &counts, &out_x, &out_z),
481 )
482 .map_err(|e| {
483 XlogError::Kernel(format!("fj_delta_emit_u32 launch failed: {e}"))
484 })?;
485 }
486 Ok(())
487 };
488 let rec = unsafe { rec.enqueue_prepared_with(&cu_stream, enqueue) }
490 .map_err(|error| error.into_xlog_error())?;
491 rec.commit()
492 .map_err(|e| XlogError::Kernel(format!("{ctx}: emit commit failed: {e}")))?;
493 }
494 cu_stream
495 .synchronize()
496 .map_err(|e| XlogError::Kernel(format!("{ctx}: emit sync failed: {e}")))?;
497
498 let d_nr = self.memory().alloc::<u32>(1)?;
499 let mut rec_num_rows = LaunchRecorder::new_strict(launch_stream);
500 rec_num_rows.write(&d_nr);
501 rec_num_rows.preflight(runtime).map_err(|error| {
502 XlogError::Kernel(format!("{ctx}: num_rows preflight failed: {error}"))
503 })?;
504 let rec_num_rows = unsafe {
505 rec_num_rows.enqueue_prepared_with(&cu_stream, |enqueue| {
506 self.initialize_launch_metadata_u32(
507 total_novel,
508 &d_nr,
509 enqueue,
510 &format!("{ctx}: result num_rows"),
511 )
512 })
513 }
514 .map_err(crate::launch::LaunchEnqueueError::into_xlog_error)?;
515 rec_num_rows.commit().map_err(|error| {
516 XlogError::Kernel(format!("{ctx}: num_rows commit failed: {error}"))
517 })?;
518 let columns = if cols.r_carry == 0 {
521 vec![out_x.into_bytes().into(), out_z.into_bytes().into()]
522 } else {
523 vec![out_z.into_bytes().into(), out_x.into_bytes().into()]
524 };
525 Ok(CudaBuffer::from_columns_with_host_count(
526 columns,
527 u64::from(total_novel),
528 d_nr,
529 out_schema,
530 total_novel,
531 ))
532 }
533
534 pub fn fj_delta_columns_max_u32(
539 &self,
540 inputs: &[(&CudaBuffer, &[usize])],
541 launch_stream: StreamId,
542 ) -> Result<u32> {
543 let ctx = "fj_delta_columns_max_u32";
544 let runtime = self.memory().runtime().ok_or_else(|| {
545 XlogError::Kernel(format!("{ctx} requires a runtime-backed GpuMemoryManager"))
546 })?;
547 let cu_stream = runtime
548 .stream_pool()
549 .resolve(launch_stream)
550 .ok_or_else(|| {
551 XlogError::Kernel(format!(
552 "{ctx}: launch_stream StreamId({}) does not resolve",
553 launch_stream.0
554 ))
555 })?;
556 let mut d_max = self.memory().alloc::<u32>(1)?;
557 self.device()
558 .inner()
559 .memset_zeros(&mut d_max)
560 .map_err(|e| XlogError::Kernel(format!("{ctx}: zero max cell failed: {e}")))?;
561 let kernel = self
562 .device()
563 .inner()
564 .get_func(WCOJ_MODULE, wcoj_kernels::FJ_DELTA_MAX_U32)
565 .ok_or_else(|| XlogError::Kernel("fj_delta_max_u32 kernel not found".to_string()))?;
566 for (buf, col_idxs) in inputs {
567 let n = match buf.cached_row_count() {
568 Some(c) => c,
569 None => self.dtoh_scalar_untracked::<u32>(buf.num_rows_device(), 0)?,
570 };
571 if n == 0 {
572 continue;
573 }
574 for &idx in *col_idxs {
575 let col = buf
576 .column(idx)
577 .ok_or_else(|| XlogError::Kernel(format!("{ctx}: column {idx} missing")))?;
578 let view = self.column_as_u32_view(col, n as usize)?;
579 let mut rec = LaunchRecorder::new_strict(launch_stream);
580 rec.read_column(col);
581 rec.read_write(&d_max);
582 rec.preflight(runtime)
583 .map_err(|e| XlogError::Kernel(format!("{ctx}: preflight failed: {e}")))?;
584 let grid = n.div_ceil(BLOCK_SIZE);
585 let enqueue = |stream: &crate::launch::CudaEnqueue<'_>| -> Result<()> {
588 unsafe {
589 kernel
590 .clone()
591 .launch_in(
592 stream,
593 LaunchConfig {
594 grid_dim: (grid, 1, 1),
595 block_dim: (BLOCK_SIZE, 1, 1),
596 shared_mem_bytes: 0,
597 },
598 (&view, n, &mut d_max),
599 )
600 .map_err(|e| {
601 XlogError::Kernel(format!("fj_delta_max_u32 launch failed: {e}"))
602 })?;
603 }
604 Ok(())
605 };
606 let rec = unsafe { rec.enqueue_prepared_with(&cu_stream, enqueue) }
608 .map_err(|error| error.into_xlog_error())?;
609 rec.commit()
610 .map_err(|e| XlogError::Kernel(format!("{ctx}: commit failed: {e}")))?;
611 }
612 }
613 cu_stream
614 .synchronize()
615 .map_err(|e| XlogError::Kernel(format!("{ctx}: sync failed: {e}")))?;
616 self.dtoh_scalar_untracked::<u32>(&d_max, 0)
617 }
618}