1use crate::ast::{Atom, Greed, Pattern};
19use crate::engine::Span;
20use crate::token::{Token, TokenKind};
21
22#[cfg(feature = "compress")]
26mod coder;
27#[cfg(feature = "compress")]
28pub use coder::*;
29
30#[must_use]
39pub fn gpu_eligible(pattern: &Pattern) -> bool {
40 eligible_shape(pattern)
41 && (!compares_classes(pattern) || crate::nfa::device_binds_fit(pattern))
42 && !(pattern.has_spectral() && pattern.binds())
46}
47
48fn compares_classes(pattern: &Pattern) -> bool {
52 match pattern {
53 Pattern::Bind(..) | Pattern::Atom(Atom::RegisterEq(..) | Atom::Literal(..)) => true,
54 Pattern::Concat(v) | Pattern::Alt(v, _) => v.iter().any(compares_classes),
55 Pattern::Star(p, _)
56 | Pattern::Plus(p, _)
57 | Pattern::Opt(p, _)
58 | Pattern::Repeat(p, _, _, _)
59 | Pattern::Atomic(p) => compares_classes(p),
60 _ => false,
61 }
62}
63
64#[must_use]
72pub fn gpu_auto_routes(pattern: &Pattern) -> bool {
73 gpu_eligible(pattern) && !reads_token_properties(pattern)
74}
75
76fn reads_token_properties(pattern: &Pattern) -> bool {
78 match pattern {
79 Pattern::Bind(..)
80 | Pattern::Atom(
81 Atom::Magnitude(_)
82 | Atom::KindMag(..)
83 | Atom::KindPred(..)
84 | Atom::Byte(_)
85 | Atom::Literal(..)
86 | Atom::RegisterEq(..)
87 | Atom::RegisterRelated(..)
88 | Atom::LiteralWithin(..)
89 | Atom::RegisterWithin(..)
90 | Atom::RegisterKin(..)
91 | Atom::Since(..)
92 | Atom::Spectral(_),
93 ) => true,
94 Pattern::Concat(v) | Pattern::Alt(v, _) => v.iter().any(reads_token_properties),
95 Pattern::Star(p, _)
96 | Pattern::Plus(p, _)
97 | Pattern::Opt(p, _)
98 | Pattern::Repeat(p, _, _, _)
99 | Pattern::Atomic(p) => reads_token_properties(p),
100 _ => false,
101 }
102}
103
104fn eligible_shape(pattern: &Pattern) -> bool {
107 match pattern {
108 Pattern::Empty => true,
109 Pattern::Atom(Atom::Kind(crate::token::TokenKind::Whitespace)) => false,
112 Pattern::Atom(Atom::Kind(crate::token::TokenKind::Custom(_))) => false,
117 Pattern::Atom(Atom::Kind(_) | Atom::Any) => true,
118 Pattern::Atom(Atom::Magnitude(p)) => p.scope().is_none(),
122 Pattern::Atom(Atom::KindMag(k, p)) => {
123 p.scope().is_none()
124 && !matches!(k, crate::token::TokenKind::Whitespace | crate::token::TokenKind::Custom(_))
125 }
126 Pattern::Atom(Atom::RegisterEq(..) | Atom::Literal(..) | Atom::Spectral(_)) => true,
127 Pattern::Atom(Atom::Byte(bc)) => crate::nfa::byte_class_bit(*bc).is_some(),
128 Pattern::Atom(_) => false,
129 Pattern::Atomic(_) | Pattern::Within(..) => false,
134 Pattern::Concat(v) => v.iter().all(eligible_shape),
135 Pattern::Bind(_, _, p) => eligible_shape(p),
136 Pattern::Star(p, Greed::Greedy)
140 | Pattern::Plus(p, Greed::Greedy)
141 | Pattern::Opt(p, Greed::Greedy)
142 | Pattern::Repeat(p, _, _, Greed::Greedy) => {
143 eligible_shape(p)
144 }
145 Pattern::Star(_, Greed::Lazy)
146 | Pattern::Plus(_, Greed::Lazy)
147 | Pattern::Opt(_, Greed::Lazy)
148 | Pattern::Repeat(_, _, _, Greed::Lazy)
149 | Pattern::Alt(..)
150 | Pattern::Balanced(..)
151 | Pattern::Guard(..)
152 | Pattern::Assert(..)
153 | Pattern::Field(..)
154 | Pattern::Anchor(_) => false,
155 }
156}
157
158#[must_use]
163#[cfg_attr(not(feature = "gpu"), allow(unused_variables))]
164pub fn scan_gpu(pattern: &Pattern, input: &[u8]) -> Option<Vec<Span>> {
165 if !gpu_eligible(pattern) {
166 return None;
167 }
168 #[cfg(feature = "gpu")]
169 {
170 cuda::scan(pattern, input)
171 }
172 #[cfg(not(feature = "gpu"))]
173 {
174 None
175 }
176}
177
178#[must_use]
186#[cfg_attr(not(feature = "gpu"), allow(unused_variables))]
187pub fn port_anchor_ends(pattern: &Pattern, kinds: &[u32], lo: usize, out: &mut [i32]) -> bool {
188 #[cfg(feature = "gpu")]
189 {
190 let Some(nfa) = crate::nfa::compile_for_gpu(pattern) else {
191 return false;
192 };
193 if nfa.reads_magnitude
194 || nfa.binds_registers
195 || nfa.class_group.is_some()
196 || nfa.reads_bytes
197 || nfa.reads_spectral
198 {
199 return false;
200 }
201 cuda::host_anchor_ends(&nfa, kinds, lo, out);
202 true
203 }
204 #[cfg(not(feature = "gpu"))]
205 {
206 false
207 }
208}
209
210pub struct GpuTokens {
214 #[cfg(feature = "gpu")]
215 kinds: cuda::ResidentKinds,
216 #[cfg(feature = "gpu")]
219 host_kinds: Vec<u32>,
220 spans: Vec<(u32, u32)>,
221 #[cfg(feature = "gpu")]
222 text: Option<HeldText>,
223}
224
225#[cfg(feature = "gpu")]
228struct HeldText {
229 classes: ClassIds,
230}
231
232impl GpuTokens {
233 #[must_use]
239 #[cfg_attr(not(feature = "gpu"), allow(unused_variables))]
240 pub fn upload_with_properties(input: &[u8], group: crate::orbit::OrbitGroup) -> Option<Self> {
241 #[cfg(feature = "gpu")]
242 {
243 let toks = crate::parallel_lex::lex_parallel(input);
244 let (kinds, spans) = significant_stream(&toks);
245 let mags = significant_magnitudes(&toks, input);
246 let masks = significant_byte_masks(&toks, input);
247 let mut classes = ClassIds::new(group);
248 let ids: Vec<u32> = toks
249 .iter()
250 .filter(|t| t.kind != TokenKind::Whitespace)
251 .map(|t| classes.id_of(&input[t.start()..t.end()]))
252 .collect();
253 let held = cuda::upload_properties(&kinds, &mags, &masks, &ids, None)?;
254 Some(Self { kinds: held, host_kinds: kinds, spans, text: Some(HeldText { classes }) })
255 }
256 #[cfg(not(feature = "gpu"))]
257 {
258 None
259 }
260 }
261
262 #[must_use]
266 #[cfg_attr(not(feature = "gpu"), allow(unused_variables))]
267 pub fn upload_with_spectral(input: &[u8], group: crate::orbit::OrbitGroup) -> Option<Self> {
268 #[cfg(feature = "gpu")]
269 {
270 let toks = crate::parallel_lex::lex_parallel(input);
271 let (kinds, spans) = significant_stream(&toks);
272 let mags = significant_magnitudes(&toks, input);
273 let masks = significant_byte_masks(&toks, input);
274 let mut classes = ClassIds::new(group);
275 let ids: Vec<u32> = toks
276 .iter()
277 .filter(|t| t.kind != TokenKind::Whitespace)
278 .map(|t| classes.id_of(&input[t.start()..t.end()]))
279 .collect();
280 let (entropy, packed) = significant_spectral(&toks, &crate::spectral::analyze(input));
281 let held = cuda::upload_properties(&kinds, &mags, &masks, &ids, Some((&entropy, &packed)))?;
282 Some(Self { kinds: held, host_kinds: kinds, spans, text: Some(HeldText { classes }) })
283 }
284 #[cfg(not(feature = "gpu"))]
285 {
286 None
287 }
288 }
289
290 #[must_use]
294 #[cfg_attr(not(feature = "gpu"), allow(unused_variables))]
295 pub fn upload(input: &[u8]) -> Option<Self> {
296 #[cfg(feature = "gpu")]
297 {
298 let (kinds, spans) = crate::parallel_lex::lex_significant_parallel(input);
299 Some(Self { kinds: cuda::upload_kinds(&kinds, None)?, host_kinds: kinds, spans, text: None })
300 }
301 #[cfg(not(feature = "gpu"))]
302 {
303 None
304 }
305 }
306
307 #[must_use]
310 pub fn upload_with_magnitudes(input: &[u8]) -> Option<Self> {
311 Self::hold(&crate::parallel_lex::lex_parallel(input), Some(input))
312 }
313
314 #[must_use]
317 pub fn from_tokens(toks: &[Token]) -> Option<Self> {
318 Self::hold(toks, None)
319 }
320
321 #[cfg_attr(not(feature = "gpu"), allow(unused_variables))]
324 fn hold(toks: &[Token], input: Option<&[u8]>) -> Option<Self> {
325 #[cfg(feature = "gpu")]
326 {
327 let (kinds, spans) = significant_stream(toks);
328 let mags = input.map(|bytes| significant_magnitudes(toks, bytes));
329 Some(Self { kinds: cuda::upload_kinds(&kinds, mags.as_deref())?, host_kinds: kinds, spans, text: None })
330 }
331 #[cfg(not(feature = "gpu"))]
332 {
333 None
334 }
335 }
336
337 #[must_use]
339 pub fn len(&self) -> usize {
340 self.spans.len()
341 }
342
343 #[must_use]
345 pub fn is_empty(&self) -> bool {
346 self.spans.is_empty()
347 }
348
349 #[must_use]
355 #[cfg_attr(not(feature = "gpu"), allow(unused_variables))]
356 pub fn scan(&self, pattern: &Pattern) -> Option<Vec<Span>> {
357 if !gpu_eligible(pattern) {
358 return None;
359 }
360 #[cfg(feature = "gpu")]
361 {
362 cuda::scan_resident(&self.kinds, &self.spans, self.text.as_ref(), pattern)
363 }
364 #[cfg(not(feature = "gpu"))]
365 {
366 None
367 }
368 }
369
370 #[must_use]
380 #[cfg_attr(not(feature = "gpu"), allow(unused_variables))]
381 pub fn scan_split(&self, pattern: &Pattern, host_per_mille: u32) -> Option<Vec<Span>> {
382 if !gpu_eligible(pattern) {
383 return None;
384 }
385 #[cfg(feature = "gpu")]
386 {
387 cuda::scan_resident_split(&self.kinds, &self.host_kinds, &self.spans, pattern, host_per_mille)
388 }
389 #[cfg(not(feature = "gpu"))]
390 {
391 None
392 }
393 }
394
395 #[must_use]
400 #[cfg_attr(not(feature = "gpu"), allow(unused_variables))]
401 pub fn scan_split_phases(
402 &self,
403 pattern: &Pattern,
404 host_per_mille: u32,
405 ) -> Option<(Vec<Span>, SplitPhases)> {
406 if !gpu_eligible(pattern) {
407 return None;
408 }
409 #[cfg(feature = "gpu")]
410 {
411 cuda::scan_resident_split_timed(
412 &self.kinds,
413 &self.host_kinds,
414 &self.spans,
415 pattern,
416 host_per_mille,
417 )
418 }
419 #[cfg(not(feature = "gpu"))]
420 {
421 None
422 }
423 }
424}
425
426#[derive(Clone, Copy, Debug, PartialEq, Eq)]
431pub enum Backend {
432 Auto,
434 Gpu,
437 Cpu,
439}
440
441#[derive(Clone, Copy, Debug, PartialEq, Eq)]
444pub enum BackendUsed {
445 Gpu,
447 Cpu,
449 Split,
452}
453
454#[derive(Clone, Copy, Debug, Default, PartialEq)]
463pub struct GpuPhases {
464 pub tokens: usize,
466 pub lex_us: f64,
468 pub host_prep_us: f64,
471 pub upload_us: f64,
473 pub kernel_and_download_us: f64,
477 pub select_us: f64,
480}
481
482#[derive(Clone, Copy, Debug, Default, PartialEq)]
491pub struct SplitPhases {
492 pub ends_us: f64,
496 pub select_us: f64,
499 pub wall_us: f64,
501}
502
503#[derive(Clone, Copy, Debug, Default, PartialEq)]
507pub struct PipelinePhases {
508 pub partitions: usize,
510 pub tokens: usize,
512 pub wall_us: f64,
514 pub blobs_us: f64,
517 pub lex_us: f64,
519 pub device_us: f64,
522 pub select_us: f64,
524 pub carry_us: f64,
529}
530
531impl PipelinePhases {
532 #[must_use]
535 pub fn overlap_us(&self) -> f64 {
536 (self.blobs_us + self.lex_us + self.device_us + self.select_us + self.carry_us
537 - self.wall_us)
538 .max(0.0)
539 }
540}
541
542#[must_use]
549#[cfg_attr(not(feature = "gpu"), allow(unused_variables))]
550pub fn scan_gpu_phases(
551 pattern: &Pattern,
552 input: &[u8],
553) -> Option<(Vec<crate::engine::Span>, GpuPhases)> {
554 #[cfg(feature = "gpu")]
555 {
556 cuda::scan_phases(pattern, input)
557 }
558 #[cfg(not(feature = "gpu"))]
559 {
560 None
561 }
562}
563
564#[must_use]
572#[cfg_attr(not(feature = "gpu"), allow(unused_variables))]
573pub fn scan_gpu_pipelined(
574 pattern: &Pattern,
575 input: &[u8],
576 partitions: usize,
577) -> Option<(Vec<crate::engine::Span>, PipelinePhases)> {
578 #[cfg(feature = "gpu")]
579 {
580 cuda::scan_pipelined(pattern, input, partitions)
581 }
582 #[cfg(not(feature = "gpu"))]
583 {
584 None
585 }
586}
587
588#[must_use]
595#[cfg_attr(not(feature = "gpu"), allow(unused_variables))]
596pub fn scan_gpu_parts(pattern: &Pattern, input: &[u8]) -> Option<Vec<crate::engine::Span>> {
597 #[cfg(feature = "gpu")]
598 {
599 cuda::scan_parts(pattern, input)
600 }
601 #[cfg(not(feature = "gpu"))]
602 {
603 None
604 }
605}
606
607#[must_use]
613#[cfg_attr(not(feature = "gpu"), allow(unused_variables))]
614pub fn scan_gpu_split(pattern: &Pattern, input: &[u8]) -> Option<Vec<crate::engine::Span>> {
615 #[cfg(feature = "gpu")]
616 {
617 cuda::scan_split_now(pattern, input)
618 }
619 #[cfg(not(feature = "gpu"))]
620 {
621 None
622 }
623}
624
625#[must_use]
630pub fn device_available() -> bool {
631 #[cfg(feature = "gpu")]
632 {
633 cuda::device_present()
634 }
635 #[cfg(not(feature = "gpu"))]
636 {
637 false
638 }
639}
640
641#[must_use]
650pub fn scan_with_backend(
651 pattern: &Pattern,
652 input: &[u8],
653 backend: Backend,
654) -> (Vec<Span>, BackendUsed) {
655 match backend {
656 Backend::Cpu => (crate::engine::scan(pattern, input), BackendUsed::Cpu),
657 Backend::Gpu => match scan_gpu(pattern, input) {
658 Some(m) => {
659 crate::trace::rung("scan", "the device kernel", input.len());
660 (m, BackendUsed::Gpu)
661 }
662 None => (crate::engine::scan(pattern, input), BackendUsed::Cpu),
663 },
664 Backend::Auto => {
665 #[cfg(feature = "gpu")]
666 {
667 if gpu_auto_routes(pattern)
668 && device_available()
669 && let Some(placed) = cuda::scan_placed(pattern, input)
670 {
671 return placed;
672 }
673 }
674 (crate::engine::scan(pattern, input), BackendUsed::Cpu)
675 }
676 }
677}
678
679#[must_use]
693pub fn significant_stream(toks: &[Token]) -> (Vec<u32>, Vec<(u32, u32)>) {
694 let mut kinds = Vec::with_capacity(toks.len());
695 let mut spans = Vec::with_capacity(toks.len());
696 for t in toks {
697 if t.kind != TokenKind::Whitespace {
698 kinds.push(t.kind.code());
699 spans.push((t.start, t.end));
700 }
701 }
702 (kinds, spans)
703}
704
705#[must_use]
709pub fn significant_magnitudes(toks: &[Token], input: &[u8]) -> Vec<f32> {
710 toks.iter()
711 .filter(|t| t.kind != TokenKind::Whitespace)
712 .map(|t| crate::magnitude::token_magnitude(t.kind, &input[t.start()..t.end()]) * 100.0)
713 .collect()
714}
715
716pub(crate) struct ClassIds {
721 group: crate::orbit::OrbitGroup,
722 ids: std::collections::HashMap<Vec<u8>, u32>,
723}
724
725impl ClassIds {
726 pub(crate) fn new(group: crate::orbit::OrbitGroup) -> Self {
727 Self { group, ids: std::collections::HashMap::new() }
728 }
729
730 #[cfg(feature = "gpu")]
732 pub(crate) fn group(&self) -> crate::orbit::OrbitGroup {
733 self.group
734 }
735
736 pub(crate) fn id_of(&mut self, bytes: &[u8]) -> u32 {
739 let next = u32::try_from(self.ids.len()).expect("fewer than 2^32 distinct classes");
740 match self.group {
741 crate::orbit::OrbitGroup::Identity => match self.ids.get(bytes) {
742 Some(&id) => id,
743 None => {
744 self.ids.insert(bytes.to_vec(), next);
745 next
746 }
747 },
748 g => *self.ids.entry(crate::orbit::canonical(bytes, g).into_bytes()).or_insert(next),
749 }
750 }
751
752 #[cfg(feature = "gpu")]
754 pub(crate) fn find(&self, bytes: &[u8]) -> Option<u32> {
755 match self.group {
756 crate::orbit::OrbitGroup::Identity => self.ids.get(bytes).copied(),
757 g => self.ids.get(crate::orbit::canonical(bytes, g).as_bytes()).copied(),
758 }
759 }
760}
761
762#[must_use]
768pub fn significant_classes(
769 toks: &[Token],
770 input: &[u8],
771 group: crate::orbit::OrbitGroup,
772 literals: &[&[u8]],
773) -> (Vec<u32>, Vec<u32>) {
774 let mut ids = ClassIds::new(group);
775 let literal_ids: Vec<u32> = literals.iter().map(|&lit| ids.id_of(lit)).collect();
776 let token_ids: Vec<u32> = toks
777 .iter()
778 .filter(|t| t.kind != TokenKind::Whitespace)
779 .map(|t| ids.id_of(&input[t.start()..t.end()]))
780 .collect();
781 (token_ids, literal_ids)
782}
783
784#[must_use]
791pub fn significant_spectral(toks: &[Token], field: &crate::spectral::SpectralField) -> (Vec<f32>, Vec<u32>) {
792 let mut entropy = Vec::with_capacity(toks.len());
793 let mut packed = Vec::with_capacity(toks.len());
794 for t in toks.iter().filter(|t| t.kind != TokenKind::Whitespace) {
795 let sig = field.signature(t.start(), t.end());
796 entropy.push(sig.entropy * 100.0);
797 let texture = crate::nfa::texture_id(crate::spectral::texture_of(&sig));
798 let onset = u32::from(field.boundary_in(t.start(), t.end()));
799 packed.push(u32::from(sig.period) | (texture << 16) | (onset << 20));
800 }
801 (entropy, packed)
802}
803
804#[must_use]
809pub fn significant_byte_masks(toks: &[Token], input: &[u8]) -> Vec<u32> {
810 use crate::ast::ByteClass;
811 const CLASSES: [ByteClass; 6] =
812 [ByteClass::Digit, ByteClass::Word, ByteClass::Hex, ByteClass::Alpha, ByteClass::Upper, ByteClass::Lower];
813 toks.iter()
814 .filter(|t| t.kind != TokenKind::Whitespace)
815 .map(|t| {
816 let bytes = &input[t.start()..t.end()];
817 CLASSES
818 .iter()
819 .filter(|&&bc| crate::engine::byte_class_matches(bc, bytes))
820 .map(|&bc| crate::nfa::byte_class_bit(bc).expect("every class but whitespace has a bit"))
821 .fold(0, |mask, bit| mask | bit)
822 })
823 .collect()
824}
825
826#[cfg(feature = "gpu")]
830#[allow(unsafe_code)]
831mod cuda {
832 use std::sync::{Arc, OnceLock};
833
834 use cudarc::driver::{CudaContext, CudaFunction, LaunchConfig, PushKernelArg};
835 use cudarc::nvrtc::Ptx;
836
837 use super::Span;
838 use crate::ast::Pattern;
839 use crate::ends_simd::{next_match_from, next_match_from_scalar, takes_vector};
840 use crate::token::Token;
841
842 const PTX: &str = include_str!(env!("TREX_SCAN_PTX"));
846
847 struct Gpu {
849 ctx: Arc<CudaContext>,
850 func: CudaFunction,
851 }
852
853 pub(super) fn device_present() -> bool {
856 gpu().is_some()
857 }
858
859 pub(super) fn load_or_cpu<T>(
879 kernel: &str,
880 load: fn() -> Result<T, cudarc::driver::DriverError>,
881 ) -> Option<T> {
882 let prior = std::panic::take_hook();
883 std::panic::set_hook(Box::new(|_| {}));
884 let probed = std::panic::catch_unwind(load);
885 std::panic::set_hook(prior);
886 match probed {
887 Ok(Ok(loaded)) => Some(loaded),
888 Ok(Err(driver)) => {
889 crate::trace::rung("gpu", &format!("{kernel}: no usable device ({driver:?})"), 0);
890 None
891 }
892 Err(panicked) => {
893 let why = panicked
894 .downcast_ref::<String>()
895 .map(String::as_str)
896 .or_else(|| panicked.downcast_ref::<&str>().copied())
897 .unwrap_or("the CUDA library could not be loaded");
898 crate::trace::rung("gpu", &format!("{kernel}: no CUDA library ({why})"), 0);
899 None
900 }
901 }
902 }
903
904 fn gpu() -> Option<&'static Gpu> {
907 static G: OnceLock<Option<Gpu>> = OnceLock::new();
908 G.get_or_init(|| load_or_cpu("scan", load_scan_kernel)).as_ref()
909 }
910
911 fn load_scan_kernel() -> Result<Gpu, cudarc::driver::DriverError> {
913 let ctx = CudaContext::new(0)?;
914 let module = ctx.load_module(Ptx::from_src(PTX))?;
915 let func = module.load_function("trex_scan")?;
916 Ok(Gpu { ctx, func })
917 }
918
919 pub(super) fn scan_phases(
934 pattern: &Pattern,
935 input: &[u8],
936 ) -> Option<(Vec<crate::engine::Span>, super::GpuPhases)> {
937 use std::time::Instant;
938
939 let nfa = crate::nfa::compile_for_gpu(pattern)?;
940 let g = gpu()?;
941 if nfa.reads_magnitude
942 || nfa.binds_registers
943 || nfa.class_group.is_some()
944 || nfa.reads_bytes
945 || nfa.reads_spectral
946 {
947 eprintln!(
948 "trex gpu phase probe: phases are timed for kind-only patterns, and this one reads a per-token property"
949 );
950 return None;
951 }
952 let mut ph = super::GpuPhases::default();
953
954 let t = Instant::now();
957 crate::parallel_lex::lex_significant_parallel_held(input, |kinds, spans| {
958 ph.lex_us = t.elapsed().as_secs_f64() * 1e6;
959 let n = kinds.len();
960 if n == 0 {
961 return Some((Vec::new(), ph));
962 }
963 ph.tokens = n;
964 phases_lexed(g, &nfa, kinds, spans, ph)
965 })
966 }
967
968 fn phases_lexed(
970 g: &Gpu,
971 nfa: &crate::nfa::GpuNfa,
972 kinds: &[u32],
973 spans: &[(u32, u32)],
974 mut ph: super::GpuPhases,
975 ) -> Option<(Vec<Span>, super::GpuPhases)> {
976 use std::time::Instant;
977
978 let n = kinds.len();
979 let stream = g.ctx.default_stream();
980 let t = Instant::now();
981 let staged = (|| {
982 let d_kind = stream.clone_htod(kinds)?;
983 let d_atom = stream.clone_htod(&nfa.atom_kind)?;
984 let d_next = stream.clone_htod(&nfa.next_closure)?;
985 let d_isatom = stream.clone_htod(&nfa.is_atom)?;
986 let d_result = stream.alloc_zeros::<i32>(n)?;
987 stream.synchronize()?;
991 Ok::<_, cudarc::driver::DriverError>((d_kind, d_atom, d_next, d_isatom, d_result))
992 })();
993 let (d_kind, d_atom, d_next, d_isatom, d_result) = match staged {
994 Ok(v) => v,
995 Err(e) => {
996 eprintln!("trex gpu phase probe: staging {n} tokens to the device failed: {e:?}");
1000 return None;
1001 }
1002 };
1003 ph.upload_us = t.elapsed().as_secs_f64() * 1e6;
1004
1005 let n_i = n as i32;
1006 let nstates_i = nfa.nstates as i32;
1007 let start_closure = nfa.start_closure;
1008 let match_mask = nfa.match_mask;
1009 let block = 256u32;
1010 let cfg = LaunchConfig {
1011 grid_dim: (n.div_ceil(block as usize) as u32, 1, 1),
1012 block_dim: (block, 1, 1),
1013 shared_mem_bytes: 0,
1014 };
1015
1016 let t = Instant::now();
1017 let ran = (|| {
1018 let mut builder = stream.launch_builder(&g.func);
1019 builder.arg(&d_kind);
1020 builder.arg(&n_i);
1021 builder.arg(&d_atom);
1022 builder.arg(&d_next);
1023 builder.arg(&d_isatom);
1024 builder.arg(&start_closure);
1025 builder.arg(&match_mask);
1026 builder.arg(&nstates_i);
1027 builder.arg(&d_result);
1028 unsafe { builder.launch(cfg)? };
1032 let result: Vec<i32> = stream.clone_dtoh(&d_result)?;
1033 stream.synchronize()?;
1034 Ok::<_, cudarc::driver::DriverError>(result)
1035 })();
1036 let result = match ran {
1037 Ok(r) => r,
1038 Err(e) => {
1039 eprintln!("trex gpu phase probe: the kernel over {n} tokens failed: {e:?}");
1040 return None;
1041 }
1042 };
1043 ph.kernel_and_download_us = t.elapsed().as_secs_f64() * 1e6;
1044
1045 let t = Instant::now();
1046 let out = select(spans, &result);
1047 ph.select_us = t.elapsed().as_secs_f64() * 1e6;
1048 Some((out, ph))
1049 }
1050
1051 pub(super) fn scan_pipelined(
1065 pattern: &Pattern,
1066 input: &[u8],
1067 partitions: usize,
1068 ) -> Option<(Vec<Span>, super::PipelinePhases)> {
1069 use std::sync::mpsc::{TryRecvError, channel, sync_channel};
1070 use std::time::Instant;
1071
1072 use crate::parallel_lex::SignificantWorkspace;
1073
1074 let nfa = crate::nfa::compile_for_gpu(pattern)?;
1075 if nfa.reads_magnitude
1076 || nfa.binds_registers
1077 || nfa.class_group.is_some()
1078 || nfa.reads_bytes
1079 || nfa.reads_spectral
1080 {
1081 return None;
1082 }
1083 let max_len = crate::nfa::bounded_max_len(pattern)?;
1084 let g = gpu()?;
1085 let carry = max_len.saturating_sub(1);
1086
1087 let wall = Instant::now();
1088 let mut ph = super::PipelinePhases::default();
1089 let t = Instant::now();
1090 let blobs = crate::lexer::blob_runs_parallel(input);
1091 ph.blobs_us = t.elapsed().as_secs_f64() * 1e6;
1092 let bounds = crate::parallel_lex::safe_boundaries(input, partitions.max(1));
1093 ph.partitions = bounds.len().saturating_sub(1);
1094
1095 let (to_device, from_lexer) = sync_channel::<SignificantWorkspace>(1);
1098 let (to_lexer, from_device) = channel::<SignificantWorkspace>();
1099 let nfa_ref = &nfa;
1100 let (lex_us, staged) = std::thread::scope(|sc| {
1101 let device = sc.spawn(move || device_stage(g, nfa_ref, carry, from_lexer, to_lexer));
1102 let mut lex_us = 0.0;
1103 let mut spare: Vec<SignificantWorkspace> =
1108 (0..3).map(|_| SignificantWorkspace::default()).collect();
1109 for w in bounds.windows(2) {
1110 let mut ws = match from_device.try_recv() {
1121 Ok(ws) => ws,
1122 Err(TryRecvError::Empty) => match spare.pop() {
1123 Some(ws) => ws,
1124 None => match from_device.recv() {
1125 Ok(ws) => ws,
1126 Err(std::sync::mpsc::RecvError) => break,
1127 },
1128 },
1129 Err(TryRecvError::Disconnected) => break,
1130 };
1131 let t = Instant::now();
1137 ws.kinds.clear();
1138 ws.spans.clear();
1139 ws.kinds.resize(carry, 0);
1140 ws.spans.resize(carry, (0, 0));
1141 crate::parallel_lex::lex_significant_range_onto(input, w[0], w[1], &blobs, &mut ws);
1142 lex_us += t.elapsed().as_secs_f64() * 1e6;
1143 if to_device.send(ws).is_err() {
1144 break;
1147 }
1148 }
1149 drop(to_device);
1150 (lex_us, device.join().expect("the pipeline's device thread returns"))
1151 });
1152 let staged = match staged {
1153 Ok(s) => s,
1154 Err(e) => {
1155 eprintln!("trex gpu pipeline: a window on the device failed: {e:?}");
1156 return None;
1157 }
1158 };
1159 ph.lex_us = lex_us;
1160 ph.tokens = staged.tokens;
1161 ph.device_us = staged.device_us;
1162 ph.select_us = staged.select_us;
1163 ph.carry_us = staged.carry_us;
1164 ph.wall_us = wall.elapsed().as_secs_f64() * 1e6;
1165 Some((staged.spans, ph))
1166 }
1167
1168 struct Staged {
1171 spans: Vec<Span>,
1172 tokens: usize,
1173 device_us: f64,
1174 select_us: f64,
1175 carry_us: f64,
1176 }
1177
1178 fn device_stage(
1184 g: &Gpu,
1185 nfa: &crate::nfa::GpuNfa,
1186 carry: usize,
1187 ranges: std::sync::mpsc::Receiver<crate::parallel_lex::SignificantWorkspace>,
1188 spent: std::sync::mpsc::Sender<crate::parallel_lex::SignificantWorkspace>,
1189 ) -> Result<Staged, cudarc::driver::DriverError> {
1190 use std::time::Instant;
1191
1192 let stream = g.ctx.default_stream();
1193 let d_atom = stream.clone_htod(&nfa.atom_kind)?;
1194 let d_next = stream.clone_htod(&nfa.next_closure)?;
1195 let d_isatom = stream.clone_htod(&nfa.is_atom)?;
1196 let mut staged =
1197 Staged { spans: Vec::new(), tokens: 0, device_us: 0.0, select_us: 0.0, carry_us: 0.0 };
1198 let mut kinds: Vec<u32> = Vec::new();
1202 let mut spans: Vec<(u32, u32)> = Vec::new();
1203 let mut next = 0usize;
1205 let mut last = false;
1206 while !last {
1207 let (mut window, start) = match ranges.recv() {
1213 Ok(ws) => {
1214 let t = Instant::now();
1215 let start = carry - kinds.len();
1216 staged.tokens += ws.kinds.len() - carry;
1217 let mut ws = ws;
1218 ws.kinds[start..carry].copy_from_slice(&kinds);
1219 ws.spans[start..carry].copy_from_slice(&spans);
1220 staged.carry_us += t.elapsed().as_secs_f64() * 1e6;
1221 (Some(ws), start)
1222 }
1223 Err(std::sync::mpsc::RecvError) => {
1226 last = true;
1227 (None, 0)
1228 }
1229 };
1230 let (wk, wsp): (&[u32], &[(u32, u32)]) = match &window {
1231 Some(ws) => (&ws.kinds[start..], &ws.spans[start..]),
1232 None => (&kinds, &spans),
1233 };
1234 let n = wk.len();
1235 if n == 0 {
1236 if let Some(ws) = window.take()
1240 && let Err(returned) = spent.send(ws)
1241 {
1242 drop(returned);
1243 }
1244 continue;
1245 }
1246 let t = Instant::now();
1247 let ends = window_ends(g, &stream, nfa, wk, &d_atom, &d_next, &d_isatom)?;
1248 staged.device_us += t.elapsed().as_secs_f64() * 1e6;
1249
1250 let t = Instant::now();
1251 let horizon = if last { n } else { n.saturating_sub(carry) };
1252 while next < horizon {
1253 let best = ends[next];
1254 if best > next as i32 {
1255 let e = best as usize;
1256 staged.spans.push(Span { start: wsp[next].0, end: wsp[e - 1].1 });
1257 next = e;
1258 } else {
1259 next += 1;
1260 }
1261 }
1262 staged.select_us += t.elapsed().as_secs_f64() * 1e6;
1263 next -= horizon;
1264
1265 let t = Instant::now();
1266 let (tail_kinds, tail_spans) = (wk[horizon..].to_vec(), wsp[horizon..].to_vec());
1267 kinds = tail_kinds;
1268 spans = tail_spans;
1269 staged.carry_us += t.elapsed().as_secs_f64() * 1e6;
1270 if let Some(ws) = window.take()
1274 && let Err(returned) = spent.send(ws)
1275 {
1276 drop(returned);
1277 }
1278 }
1279 Ok(staged)
1280 }
1281
1282 #[allow(clippy::too_many_arguments)]
1286 fn window_ends(
1287 g: &Gpu,
1288 stream: &Arc<cudarc::driver::CudaStream>,
1289 nfa: &crate::nfa::GpuNfa,
1290 kinds: &[u32],
1291 d_atom: &cudarc::driver::CudaSlice<u32>,
1292 d_next: &cudarc::driver::CudaSlice<u64>,
1293 d_isatom: &cudarc::driver::CudaSlice<u32>,
1294 ) -> Result<Vec<i32>, cudarc::driver::DriverError> {
1295 let d_kind = stream.clone_htod(kinds)?;
1296 device_ends(g, stream, nfa, &d_kind, kinds.len(), d_atom, d_next, d_isatom)
1297 }
1298
1299 #[allow(clippy::too_many_arguments)]
1303 fn device_view_ends(
1304 g: &Gpu,
1305 stream: &Arc<cudarc::driver::CudaStream>,
1306 nfa: &crate::nfa::GpuNfa,
1307 d_kind: &cudarc::driver::CudaView<'_, u32>,
1308 out: &mut [i32],
1309 d_atom: &cudarc::driver::CudaSlice<u32>,
1310 d_next: &cudarc::driver::CudaSlice<u64>,
1311 d_isatom: &cudarc::driver::CudaSlice<u32>,
1312 ) -> Result<(), cudarc::driver::DriverError> {
1313 let n = out.len();
1314 let d_result = stream.alloc_zeros::<i32>(n)?;
1315 let n_i = n as i32;
1316 let nstates_i = nfa.nstates as i32;
1317 let start_closure = nfa.start_closure;
1318 let match_mask = nfa.match_mask;
1319 let block = 256u32;
1320 let cfg = LaunchConfig {
1321 grid_dim: (n.div_ceil(block as usize) as u32, 1, 1),
1322 block_dim: (block, 1, 1),
1323 shared_mem_bytes: 0,
1324 };
1325 let mut builder = stream.launch_builder(&g.func);
1326 builder.arg(d_kind);
1327 builder.arg(&n_i);
1328 builder.arg(d_atom);
1329 builder.arg(d_next);
1330 builder.arg(d_isatom);
1331 builder.arg(&start_closure);
1332 builder.arg(&match_mask);
1333 builder.arg(&nstates_i);
1334 builder.arg(&d_result);
1335 unsafe { builder.launch(cfg)? };
1339 stream.memcpy_dtoh(&d_result, out)?;
1340 stream.synchronize()?;
1341 Ok(())
1342 }
1343
1344 #[allow(clippy::too_many_arguments)]
1348 fn device_ends(
1349 g: &Gpu,
1350 stream: &Arc<cudarc::driver::CudaStream>,
1351 nfa: &crate::nfa::GpuNfa,
1352 d_kind: &cudarc::driver::CudaSlice<u32>,
1353 n: usize,
1354 d_atom: &cudarc::driver::CudaSlice<u32>,
1355 d_next: &cudarc::driver::CudaSlice<u64>,
1356 d_isatom: &cudarc::driver::CudaSlice<u32>,
1357 ) -> Result<Vec<i32>, cudarc::driver::DriverError> {
1358 let d_result = stream.alloc_zeros::<i32>(n)?;
1359 let n_i = n as i32;
1360 let nstates_i = nfa.nstates as i32;
1361 let start_closure = nfa.start_closure;
1362 let match_mask = nfa.match_mask;
1363 let block = 256u32;
1364 let cfg = LaunchConfig {
1365 grid_dim: (n.div_ceil(block as usize) as u32, 1, 1),
1366 block_dim: (block, 1, 1),
1367 shared_mem_bytes: 0,
1368 };
1369 let mut builder = stream.launch_builder(&g.func);
1370 builder.arg(d_kind);
1371 builder.arg(&n_i);
1372 builder.arg(d_atom);
1373 builder.arg(d_next);
1374 builder.arg(d_isatom);
1375 builder.arg(&start_closure);
1376 builder.arg(&match_mask);
1377 builder.arg(&nstates_i);
1378 builder.arg(&d_result);
1379 unsafe { builder.launch(cfg)? };
1383 let ends: Vec<i32> = stream.clone_dtoh(&d_result)?;
1384 stream.synchronize()?;
1385 Ok(ends)
1386 }
1387
1388 pub(super) fn scan(pattern: &Pattern, input: &[u8]) -> Option<Vec<Span>> {
1389 let nfa = crate::nfa::compile_for_gpu(pattern)?;
1390 let g = gpu()?;
1391 if nfa.reads_magnitude
1392 || nfa.binds_registers
1393 || nfa.class_group.is_some()
1394 || nfa.reads_bytes
1395 || nfa.reads_spectral
1396 {
1397 return scan_props(g, &nfa, input);
1398 }
1399
1400 crate::parallel_lex::lex_significant_parallel_held(input, |kinds, spans| match scan_lexed(g, &nfa, kinds, spans) {
1403 Ok(found) => Some(found),
1404 Err(e) => {
1405 eprintln!("trex gpu scan: a device call failed: {e:?}");
1406 None
1407 }
1408 })
1409 }
1410
1411 fn scan_lexed(
1413 g: &Gpu,
1414 nfa: &crate::nfa::GpuNfa,
1415 kinds: &[u32],
1416 spans: &[(u32, u32)],
1417 ) -> Result<Vec<Span>, cudarc::driver::DriverError> {
1418 let n = kinds.len();
1419 if n == 0 {
1420 return Ok(Vec::new());
1421 }
1422
1423 let stream = g.ctx.default_stream();
1424 let d_kind = stream.clone_htod(kinds)?;
1425 let d_atom = stream.clone_htod(&nfa.atom_kind)?;
1426 let d_next = stream.clone_htod(&nfa.next_closure)?;
1427 let d_isatom = stream.clone_htod(&nfa.is_atom)?;
1428 let result = device_ends(g, &stream, nfa, &d_kind, n, &d_atom, &d_next, &d_isatom)?;
1429 Ok(select(spans, &result))
1430 }
1431
1432 pub(super) fn scan_split_now(pattern: &Pattern, input: &[u8]) -> Option<Vec<Span>> {
1435 let nfa = crate::nfa::compile_for_gpu(pattern)?;
1436 if nfa.reads_magnitude
1437 || nfa.binds_registers
1438 || nfa.class_group.is_some()
1439 || nfa.reads_bytes
1440 || nfa.reads_spectral
1441 {
1442 return None;
1443 }
1444 let g = gpu()?;
1445 Some(scan_split(g, &nfa, input))
1446 }
1447
1448 pub(super) fn scan_parts(pattern: &Pattern, input: &[u8]) -> Option<Vec<Span>> {
1455 let nfa = crate::nfa::compile_for_gpu(pattern)?;
1456 if nfa.reads_magnitude
1457 || nfa.binds_registers
1458 || nfa.class_group.is_some()
1459 || nfa.reads_bytes
1460 || nfa.reads_spectral
1461 {
1462 return None;
1463 }
1464 let g = gpu()?;
1465 crate::parallel_lex::lex_significant_parts_held(input, |parts| match scan_lexed_parts(g, &nfa, parts) {
1466 Ok(spans) => Some(spans),
1467 Err(e) => {
1468 eprintln!("trex gpu parts scan: a device call failed: {e:?}");
1469 None
1470 }
1471 })
1472 }
1473
1474 fn scan_lexed_parts(
1477 g: &Gpu,
1478 nfa: &crate::nfa::GpuNfa,
1479 parts: &[crate::lexer::Significant],
1480 ) -> Result<Vec<Span>, cudarc::driver::DriverError> {
1481 let n: usize = parts.iter().map(|p| p.kinds.len()).sum();
1482 if n == 0 {
1483 return Ok(Vec::new());
1484 }
1485 let stream = g.ctx.default_stream();
1486 let mut d_kind = unsafe { stream.alloc::<u32>(n) }?;
1490 let mut off = 0usize;
1491 for part in parts {
1492 let len = part.kinds.len();
1493 if len > 0 {
1494 stream.memcpy_htod(part.kinds.as_slice(), &mut d_kind.slice_mut(off..off + len))?;
1495 }
1496 off += len;
1497 }
1498 let d_atom = stream.clone_htod(&nfa.atom_kind)?;
1499 let d_next = stream.clone_htod(&nfa.next_closure)?;
1500 let d_isatom = stream.clone_htod(&nfa.is_atom)?;
1501 let ends = device_ends(g, &stream, nfa, &d_kind, n, &d_atom, &d_next, &d_isatom)?;
1502 Ok(select_parts(parts, &ends))
1503 }
1504
1505 fn select_parts(parts: &[crate::lexer::Significant], ends: &[i32]) -> Vec<Span> {
1509 if takes_vector(ends) {
1512 select_parts_with(parts, ends, next_match_from)
1513 } else {
1514 select_parts_with(parts, ends, next_match_from_scalar)
1515 }
1516 }
1517
1518 fn select_parts_with(
1520 parts: &[crate::lexer::Significant],
1521 ends: &[i32],
1522 scan: impl Fn(&[i32], usize, usize) -> Option<usize>,
1523 ) -> Vec<Span> {
1524 let mut out = Vec::with_capacity(ends.len() / 2 + 1);
1525 let mut a = 0usize;
1528 let mut base = 0usize;
1529 for (p, part) in parts.iter().enumerate() {
1530 let end = base + part.spans.len();
1531 while let Some(m) = scan(ends, a, end) {
1532 let e = ends[m] as usize; let last = if e <= end {
1534 part.spans[e - 1 - base].1
1535 } else {
1536 span_end_in(&parts[p + 1..], end, e - 1)
1537 };
1538 out.push(Span { start: part.spans[m - base].0, end: last });
1539 a = e;
1540 }
1541 a = a.max(end);
1545 base = end;
1546 }
1547 out
1548 }
1549
1550 fn span_end_in(rest: &[crate::lexer::Significant], mut base: usize, i: usize) -> u32 {
1553 for part in rest {
1554 let len = part.spans.len();
1555 if i < base + len {
1556 return part.spans[i - base].1;
1557 }
1558 base += len;
1559 }
1560 panic!("token {i} lies past the lexed chunks, whose tokens end at {base}");
1561 }
1562
1563 pub(super) struct ResidentKinds {
1567 d_kind: Option<cudarc::driver::CudaSlice<u32>>,
1568 d_mag: Option<cudarc::driver::CudaSlice<f32>>,
1569 d_bytes: Option<cudarc::driver::CudaSlice<u32>>,
1570 d_class: Option<cudarc::driver::CudaSlice<u32>>,
1571 d_ent: Option<cudarc::driver::CudaSlice<f32>>,
1572 d_spec: Option<cudarc::driver::CudaSlice<u32>>,
1573 magnitudes: bool,
1574 spectral: bool,
1575 }
1576
1577 pub(super) fn upload_kinds(kinds: &[u32], mags: Option<&[f32]>) -> Option<ResidentKinds> {
1580 let g = gpu()?;
1581 if let Some(m) = mags {
1582 assert_eq!(m.len(), kinds.len(), "one magnitude per significant token");
1583 }
1584 let magnitudes = mags.is_some();
1585 if kinds.is_empty() {
1586 return Some(ResidentKinds {
1587 d_kind: None,
1588 d_mag: None,
1589 d_bytes: None,
1590 d_class: None,
1591 d_ent: None,
1592 d_spec: None,
1593 magnitudes,
1594 spectral: false,
1595 });
1596 }
1597 let stream = g.ctx.default_stream();
1598 let held = (|| {
1599 let d_kind = stream.clone_htod(kinds)?;
1600 let d_mag = match mags {
1601 Some(m) => Some(stream.clone_htod(m)?),
1602 None => None,
1603 };
1604 Ok::<_, cudarc::driver::DriverError>((d_kind, d_mag))
1605 })();
1606 match held {
1607 Ok((d_kind, d_mag)) => Some(ResidentKinds {
1608 d_kind: Some(d_kind),
1609 d_mag,
1610 d_bytes: None,
1611 d_class: None,
1612 d_ent: None,
1613 d_spec: None,
1614 magnitudes,
1615 spectral: false,
1616 }),
1617 Err(e) => {
1618 eprintln!("trex gpu: holding {} tokens on the device failed: {e:?}", kinds.len());
1619 None
1620 }
1621 }
1622 }
1623
1624 pub(super) fn upload_properties(
1629 kinds: &[u32],
1630 mags: &[f32],
1631 masks: &[u32],
1632 ids: &[u32],
1633 spectral: Option<(&[f32], &[u32])>,
1634 ) -> Option<ResidentKinds> {
1635 let g = gpu()?;
1636 assert!(
1637 mags.len() == kinds.len() && masks.len() == kinds.len() && ids.len() == kinds.len(),
1638 "one magnitude, mask and class id per significant token"
1639 );
1640 if let Some((e, s)) = spectral {
1641 assert!(e.len() == kinds.len() && s.len() == kinds.len(), "one spectral reading per significant token");
1642 }
1643 if kinds.is_empty() {
1644 return Some(ResidentKinds {
1645 d_kind: None,
1646 d_mag: None,
1647 d_bytes: None,
1648 d_class: None,
1649 d_ent: None,
1650 d_spec: None,
1651 magnitudes: true,
1652 spectral: spectral.is_some(),
1653 });
1654 }
1655 let stream = g.ctx.default_stream();
1656 let held = (|| {
1657 let d_kind = stream.clone_htod(kinds)?;
1658 let d_mag = stream.clone_htod(mags)?;
1659 let d_bytes = stream.clone_htod(masks)?;
1660 let d_class = stream.clone_htod(ids)?;
1661 let (d_ent, d_spec) = match spectral {
1662 Some((e, s)) => (Some(stream.clone_htod(e)?), Some(stream.clone_htod(s)?)),
1663 None => (None, None),
1664 };
1665 Ok::<_, cudarc::driver::DriverError>((d_kind, d_mag, d_bytes, d_class, d_ent, d_spec))
1666 })();
1667 match held {
1668 Ok((d_kind, d_mag, d_bytes, d_class, d_ent, d_spec)) => Some(ResidentKinds {
1669 d_kind: Some(d_kind),
1670 d_mag: Some(d_mag),
1671 d_bytes: Some(d_bytes),
1672 d_class: Some(d_class),
1673 d_ent,
1674 d_spec,
1675 magnitudes: true,
1676 spectral: spectral.is_some(),
1677 }),
1678 Err(e) => {
1679 eprintln!("trex gpu: holding {} tokens and their properties on the device failed: {e:?}", kinds.len());
1680 None
1681 }
1682 }
1683 }
1684
1685 fn scan_props_kernel() -> Option<&'static CudaFunction> {
1688 static F: OnceLock<Option<CudaFunction>> = OnceLock::new();
1689 F.get_or_init(|| {
1690 let g = gpu()?;
1691 let module = match g.ctx.load_module(Ptx::from_src(PTX)) {
1692 Ok(m) => m,
1693 Err(e) => {
1694 eprintln!("trex gpu: loading the scan module for trex_scan_props failed: {e:?}");
1695 return None;
1696 }
1697 };
1698 match module.load_function("trex_scan_props") {
1699 Ok(f) => Some(f),
1700 Err(e) => {
1701 eprintln!("trex gpu: loading trex_scan_props failed: {e:?}");
1702 None
1703 }
1704 }
1705 })
1706 .as_ref()
1707 }
1708
1709 #[allow(clippy::too_many_arguments)]
1719 fn run_scan(
1720 g: &Gpu,
1721 d_kind: &cudarc::driver::CudaSlice<u32>,
1722 d_mag: Option<&cudarc::driver::CudaSlice<f32>>,
1723 d_class: Option<&cudarc::driver::CudaSlice<u32>>,
1724 d_bytes: Option<&cudarc::driver::CudaSlice<u32>>,
1725 d_spectral: Option<(&cudarc::driver::CudaSlice<f32>, &cudarc::driver::CudaSlice<u32>)>,
1726 lit_class: Option<&[u32]>,
1727 nfa: &crate::nfa::GpuNfa,
1728 n: usize,
1729 ) -> Option<Vec<i32>> {
1730 assert_eq!(
1731 nfa.reads_spectral,
1732 d_spectral.is_some(),
1733 "a pattern that tests the spectral field scans with per-token readings, and only such a pattern does"
1734 );
1735 assert_eq!(
1736 nfa.reads_magnitude,
1737 d_mag.is_some(),
1738 "a pattern that tests a magnitude scans with per-token magnitudes, and only such a pattern does"
1739 );
1740 assert_eq!(
1741 nfa.class_group.is_some(),
1742 d_class.is_some(),
1743 "a pattern that compares a literal or a register scans with per-token classes, and only such a pattern does"
1744 );
1745 assert_eq!(
1746 nfa.reads_bytes,
1747 d_bytes.is_some(),
1748 "a pattern that tests a byte class scans with per-token masks, and only such a pattern does"
1749 );
1750 let literal_classes: Vec<u32> = match lit_class {
1751 Some(l) => l.to_vec(),
1752 None => vec![u32::MAX; nfa.nstates],
1753 };
1754 assert_eq!(literal_classes.len(), nfa.nstates, "one literal class per state");
1755 assert!(
1756 nfa.literals.iter().all(Option::is_none) || lit_class.is_some(),
1757 "a pattern with a literal scans with its literal classes resolved"
1758 );
1759 let n_i = n as i32;
1760 let nstates_i = nfa.nstates as i32;
1761 let start_closure = nfa.start_closure;
1762 let match_mask = nfa.match_mask;
1763 let reads_mag_i = i32::from(nfa.reads_magnitude);
1764 let block = 256u32;
1765 let cfg = LaunchConfig {
1766 grid_dim: (n.div_ceil(block as usize) as u32, 1, 1),
1767 block_dim: (block, 1, 1),
1768 shared_mem_bytes: 0,
1769 };
1770 let reads_spec_i = i32::from(nfa.reads_spectral);
1771 let props_func = if d_mag.is_some() || d_class.is_some() || d_bytes.is_some() || d_spectral.is_some() {
1772 Some(scan_props_kernel()?)
1773 } else {
1774 None
1775 };
1776 let stream = g.ctx.default_stream();
1777 let ran = (|| {
1778 let d_atom = stream.clone_htod(&nfa.atom_kind)?;
1779 let d_next = stream.clone_htod(&nfa.next_closure)?;
1780 let d_isatom = stream.clone_htod(&nfa.is_atom)?;
1781 let d_result = stream.alloc_zeros::<i32>(n)?;
1782 match props_func {
1783 Some(func) => {
1784 let d_lo = stream.clone_htod(&nfa.mag_lo)?;
1785 let d_hi = stream.clone_htod(&nfa.mag_hi)?;
1786 let d_back = stream.clone_htod(&nfa.ref_back)?;
1787 let d_need = stream.clone_htod(&nfa.need_mask)?;
1788 let d_lit = stream.clone_htod(&literal_classes)?;
1789 let d_ent_lo = stream.clone_htod(&nfa.ent_lo)?;
1790 let d_ent_hi = stream.clone_htod(&nfa.ent_hi)?;
1791 let d_period_need = stream.clone_htod(&nfa.period_need)?;
1792 let d_period_val = stream.clone_htod(&nfa.period_val)?;
1793 let d_texture_need = stream.clone_htod(&nfa.texture_need)?;
1794 let d_onset_need = stream.clone_htod(&nfa.onset_need)?;
1795 let spare_mag = stream.clone_htod(&[0.0f32])?;
1796 let spare_class = stream.clone_htod(&[0u32])?;
1797 let spare_bytes = stream.clone_htod(&[0u32])?;
1798 let spare_ent = stream.clone_htod(&[0.0f32])?;
1799 let spare_spec = stream.clone_htod(&[0u32])?;
1800 let mut builder = stream.launch_builder(func);
1801 builder.arg(d_kind);
1802 builder.arg(&n_i);
1803 if let Some(d) = d_mag {
1804 builder.arg(d);
1805 } else {
1806 builder.arg(&spare_mag);
1807 }
1808 builder.arg(&reads_mag_i);
1809 if let Some(d) = d_class {
1810 builder.arg(d);
1811 } else {
1812 builder.arg(&spare_class);
1813 }
1814 if let Some(d) = d_bytes {
1815 builder.arg(d);
1816 } else {
1817 builder.arg(&spare_bytes);
1818 }
1819 match d_spectral {
1820 Some((e, s)) => {
1821 builder.arg(e);
1822 builder.arg(s);
1823 }
1824 None => {
1825 builder.arg(&spare_ent);
1826 builder.arg(&spare_spec);
1827 }
1828 }
1829 builder.arg(&reads_spec_i);
1830 builder.arg(&d_atom);
1831 builder.arg(&d_lo);
1832 builder.arg(&d_hi);
1833 builder.arg(&d_back);
1834 builder.arg(&d_need);
1835 builder.arg(&d_lit);
1836 builder.arg(&d_ent_lo);
1837 builder.arg(&d_ent_hi);
1838 builder.arg(&d_period_need);
1839 builder.arg(&d_period_val);
1840 builder.arg(&d_texture_need);
1841 builder.arg(&d_onset_need);
1842 builder.arg(&d_next);
1843 builder.arg(&d_isatom);
1844 builder.arg(&start_closure);
1845 builder.arg(&match_mask);
1846 builder.arg(&nstates_i);
1847 builder.arg(&d_result);
1848 unsafe { builder.launch(cfg)? };
1854 }
1855 _ => {
1856 let mut builder = stream.launch_builder(&g.func);
1857 builder.arg(d_kind);
1858 builder.arg(&n_i);
1859 builder.arg(&d_atom);
1860 builder.arg(&d_next);
1861 builder.arg(&d_isatom);
1862 builder.arg(&start_closure);
1863 builder.arg(&match_mask);
1864 builder.arg(&nstates_i);
1865 builder.arg(&d_result);
1866 unsafe { builder.launch(cfg)? };
1871 }
1872 }
1873 let result: Vec<i32> = stream.clone_dtoh(&d_result)?;
1874 stream.synchronize()?;
1875 Ok::<_, cudarc::driver::DriverError>(result)
1876 })();
1877 match ran {
1878 Ok(result) => Some(result),
1879 Err(e) => {
1880 eprintln!("trex gpu: a scan over {n} device-held tokens failed: {e:?}");
1881 None
1882 }
1883 }
1884 }
1885
1886 fn scan_props(g: &Gpu, nfa: &crate::nfa::GpuNfa, input: &[u8]) -> Option<Vec<Span>> {
1890 let toks = crate::parallel_lex::lex_parallel(input);
1891 let (kinds, spans) = significant(&toks);
1892 let n = spans.len();
1893 if n == 0 {
1894 return Some(Vec::new());
1895 }
1896 let mags = nfa.reads_magnitude.then(|| super::significant_magnitudes(&toks, input));
1897 let masks = nfa.reads_bytes.then(|| super::significant_byte_masks(&toks, input));
1898 let spectral = nfa
1899 .reads_spectral
1900 .then(|| super::significant_spectral(&toks, &crate::spectral::analyze(input)));
1901 let (classes, lit_class) = match nfa.class_group {
1902 Some(group) => {
1903 let texts: Vec<&[u8]> = nfa.literals.iter().flatten().map(Vec::as_slice).collect();
1904 let (token_ids, literal_ids) = super::significant_classes(&toks, input, group, &texts);
1905 let mut per_state = vec![u32::MAX; nfa.nstates];
1906 let mut ids = literal_ids.into_iter();
1907 for (slot, lit) in per_state.iter_mut().zip(&nfa.literals) {
1908 if lit.is_some() {
1909 *slot = ids.next().expect("one class id per literal");
1910 }
1911 }
1912 (Some(token_ids), Some(per_state))
1913 }
1914 None => (None, None),
1915 };
1916 let stream = g.ctx.default_stream();
1917 let held = (|| {
1918 let d_kind = stream.clone_htod(&kinds)?;
1919 let d_mag = match &mags {
1920 Some(m) => Some(stream.clone_htod(m)?),
1921 None => None,
1922 };
1923 let d_class = match &classes {
1924 Some(c) => Some(stream.clone_htod(c)?),
1925 None => None,
1926 };
1927 let d_bytes = match &masks {
1928 Some(b) => Some(stream.clone_htod(b)?),
1929 None => None,
1930 };
1931 let d_spectral = match &spectral {
1932 Some((e, s)) => Some((stream.clone_htod(e)?, stream.clone_htod(s)?)),
1933 None => None,
1934 };
1935 Ok::<_, cudarc::driver::DriverError>((d_kind, d_mag, d_class, d_bytes, d_spectral))
1936 })();
1937 let (d_kind, d_mag, d_class, d_bytes, d_spectral) = match held {
1938 Ok(v) => v,
1939 Err(e) => {
1940 eprintln!("trex gpu: uploading {n} tokens and their properties failed: {e:?}");
1941 return None;
1942 }
1943 };
1944 let result = run_scan(
1945 g,
1946 &d_kind,
1947 d_mag.as_ref(),
1948 d_class.as_ref(),
1949 d_bytes.as_ref(),
1950 d_spectral.as_ref().map(|(e, s)| (e, s)),
1951 lit_class.as_deref(),
1952 nfa,
1953 n,
1954 )?;
1955 Some(select(&spans, &result))
1956 }
1957
1958 pub(super) fn scan_resident(
1965 kinds: &ResidentKinds,
1966 spans: &[(u32, u32)],
1967 text: Option<&super::HeldText>,
1968 pattern: &Pattern,
1969 ) -> Option<Vec<Span>> {
1970 let nfa = crate::nfa::compile_for_gpu(pattern)?;
1971 let g = gpu()?;
1972 if (nfa.reads_magnitude && !kinds.magnitudes) || (nfa.reads_spectral && !kinds.spectral) {
1973 return None;
1974 }
1975 let needs_text = nfa.binds_registers || nfa.class_group.is_some() || nfa.reads_bytes;
1976 if needs_text && text.is_none() {
1977 return None;
1978 }
1979 if let (Some(group), Some(t)) = (nfa.class_group, text)
1980 && group != t.classes.group()
1981 {
1982 return None;
1983 }
1984 let Some(d_kind) = &kinds.d_kind else {
1985 return Some(Vec::new());
1986 };
1987 let lit_class: Option<Vec<u32>> = match (nfa.class_group, text) {
1990 (Some(_), Some(t)) => Some(
1991 nfa.literals
1992 .iter()
1993 .map(|lit| match lit {
1994 Some(bytes) => t.classes.find(bytes).unwrap_or(u32::MAX - 1),
1995 None => u32::MAX,
1996 })
1997 .collect(),
1998 ),
1999 _ => None,
2000 };
2001 let d_mag = if nfa.reads_magnitude { kinds.d_mag.as_ref() } else { None };
2002 let d_class = if nfa.class_group.is_some() { kinds.d_class.as_ref() } else { None };
2003 let d_bytes = if nfa.reads_bytes { kinds.d_bytes.as_ref() } else { None };
2004 let d_spectral = if nfa.reads_spectral { kinds.d_ent.as_ref().zip(kinds.d_spec.as_ref()) } else { None };
2005 let result =
2006 run_scan(g, d_kind, d_mag, d_class, d_bytes, d_spectral, lit_class.as_deref(), &nfa, spans.len())?;
2007 Some(select(spans, &result))
2008 }
2009
2010 pub(super) fn scan_resident_split(
2019 kinds: &ResidentKinds,
2020 host_kinds: &[u32],
2021 spans: &[(u32, u32)],
2022 pattern: &Pattern,
2023 host_per_mille: u32,
2024 ) -> Option<Vec<Span>> {
2025 scan_resident_split_timed(kinds, host_kinds, spans, pattern, host_per_mille)
2026 .map(|(out, _)| out)
2027 }
2028
2029 pub(super) fn scan_resident_split_timed(
2033 kinds: &ResidentKinds,
2034 host_kinds: &[u32],
2035 spans: &[(u32, u32)],
2036 pattern: &Pattern,
2037 host_per_mille: u32,
2038 ) -> Option<(Vec<Span>, super::SplitPhases)> {
2039 use std::time::Instant;
2040
2041 let nfa = crate::nfa::compile_for_gpu(pattern)?;
2042 if nfa.reads_magnitude
2043 || nfa.binds_registers
2044 || nfa.class_group.is_some()
2045 || nfa.reads_bytes
2046 || nfa.reads_spectral
2047 {
2048 return None;
2049 }
2050 let g = gpu()?;
2051 let n = host_kinds.len();
2052 let Some(d_kind) = &kinds.d_kind else {
2053 return Some((Vec::new(), super::SplitPhases::default()));
2054 };
2055 let wall = Instant::now();
2056 let mut ph = super::SplitPhases::default();
2057 let mut ends = vec![0i32; n];
2062 if n < 2 {
2063 let t = Instant::now();
2064 host_anchor_ends(&nfa, host_kinds, 0, &mut ends);
2065 ph.ends_us = t.elapsed().as_secs_f64() * 1e6;
2066 let t = Instant::now();
2067 let out = select(spans, &ends);
2068 ph.select_us = t.elapsed().as_secs_f64() * 1e6;
2069 ph.wall_us = wall.elapsed().as_secs_f64() * 1e6;
2070 return Some((out, ph));
2071 }
2072 let mid = (n * host_per_mille as usize / 1000).clamp(1, n - 1);
2073 let ends_at = Instant::now();
2074 let (host_part, device_part) = ends.split_at_mut(mid);
2075 std::thread::scope(|s| {
2076 s.spawn(|| {
2077 match resident_suffix_ends(g, &nfa, d_kind, mid, device_part) {
2078 Ok(()) => {
2079 let base = i32::try_from(mid).expect("a token index within the kernel's i32 width");
2080 for slot in device_part.iter_mut() {
2081 if *slot >= 0 {
2082 *slot += base;
2083 }
2084 }
2085 }
2086 Err(e) => {
2087 eprintln!(
2088 "trex gpu: the device's share of a resident split failed and the host scanned it: {e:?}"
2089 );
2090 host_anchor_ends(&nfa, host_kinds, mid, device_part);
2091 }
2092 }
2093 });
2094 host_anchor_ends(&nfa, host_kinds, 0, host_part);
2095 });
2096 ph.ends_us = ends_at.elapsed().as_secs_f64() * 1e6;
2097 let t = Instant::now();
2098 let out = select(spans, &ends);
2099 ph.select_us = t.elapsed().as_secs_f64() * 1e6;
2100 ph.wall_us = wall.elapsed().as_secs_f64() * 1e6;
2101 Some((out, ph))
2102 }
2103
2104 fn resident_suffix_ends(
2109 g: &Gpu,
2110 nfa: &crate::nfa::GpuNfa,
2111 d_kind: &cudarc::driver::CudaSlice<u32>,
2112 mid: usize,
2113 out: &mut [i32],
2114 ) -> Result<(), cudarc::driver::DriverError> {
2115 let stream = g.ctx.default_stream();
2116 let d_atom = stream.clone_htod(&nfa.atom_kind)?;
2117 let d_next = stream.clone_htod(&nfa.next_closure)?;
2118 let d_isatom = stream.clone_htod(&nfa.is_atom)?;
2119 let view = d_kind.slice(mid..mid + out.len());
2120 device_view_ends(g, &stream, nfa, &view, out, &d_atom, &d_next, &d_isatom)
2121 }
2122
2123 fn significant(toks: &[Token]) -> (Vec<u32>, Vec<(u32, u32)>) {
2144 super::significant_stream(toks)
2145 }
2146
2147 fn select(spans: &[(u32, u32)], result: &[i32]) -> Vec<Span> {
2150 if takes_vector(result) {
2153 select_with(spans, result, next_match_from)
2154 } else {
2155 select_with(spans, result, next_match_from_scalar)
2156 }
2157 }
2158
2159 fn select_with(
2161 spans: &[(u32, u32)],
2162 result: &[i32],
2163 scan: impl Fn(&[i32], usize, usize) -> Option<usize>,
2164 ) -> Vec<Span> {
2165 let n = spans.len();
2166 let mut out = Vec::with_capacity(n / 2 + 1);
2172 let mut a = 0usize;
2173 while let Some(m) = scan(result, a, n) {
2174 let e = result[m] as usize; out.push(Span { start: spans[m].0, end: spans[e - 1].1 });
2176 a = e;
2177 }
2178 out
2179 }
2180
2181 pub(super) fn scan_placed(
2188 pattern: &Pattern,
2189 input: &[u8],
2190 ) -> Option<(Vec<Span>, super::BackendUsed)> {
2191 use flynnel::sched::call_site::Placement;
2192 let nfa = crate::nfa::compile_for_gpu(pattern)?;
2193 if nfa.reads_magnitude
2194 || nfa.binds_registers
2195 || nfa.class_group.is_some()
2196 || nfa.reads_bytes
2197 || nfa.reads_spectral
2198 {
2199 return None;
2200 }
2201 let g = gpu()?;
2202 let site = flynnel::sched::call_site::caller_site().get();
2203 let size = input.len().min(u32::MAX as usize) as u32;
2204 let cpu = || {
2205 let t = std::time::Instant::now();
2206 (crate::engine::scan(pattern, input), elapsed_ns(t))
2207 };
2208 let split = || {
2209 let t = std::time::Instant::now();
2210 (scan_split(g, &nfa, input), elapsed_ns(t))
2211 };
2212 match site.choose_placement(size) {
2213 Placement::Cpu => {
2214 let (m, ns) = cpu();
2215 site.record_placement(size, Some(ns), None);
2216 Some((m, super::BackendUsed::Cpu))
2217 }
2218 Placement::Backend => {
2219 let (m, ns) = split();
2220 site.record_placement(size, None, Some(ns));
2221 crate::trace::rung("scan", "split across the cores and the device", input.len());
2222 Some((m, super::BackendUsed::Split))
2223 }
2224 Placement::Race => {
2225 let cpu_first =
2229 RACES.fetch_add(1, std::sync::atomic::Ordering::Relaxed).is_multiple_of(2);
2230 let ((m, cpu_ns), split_ns) = if cpu_first {
2231 let c = cpu();
2232 (c, split().1)
2233 } else {
2234 let s = split().1;
2235 (cpu(), s)
2236 };
2237 site.record_placement(size, Some(cpu_ns), Some(split_ns));
2238 Some((m, super::BackendUsed::Cpu))
2239 }
2240 }
2241 }
2242
2243 static RACES: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
2245
2246 fn elapsed_ns(t: std::time::Instant) -> u64 {
2247 t.elapsed().as_nanos().min(u128::from(u64::MAX)) as u64
2248 }
2249
2250 fn scan_split(g: &'static Gpu, nfa: &crate::nfa::GpuNfa, input: &[u8]) -> Vec<Span> {
2259 crate::parallel_lex::lex_significant_parallel_held(input, |kinds, spans| {
2260 let n = kinds.len();
2261 let mut ends = vec![-1i32; n];
2262 if n < 2 {
2263 host_anchor_ends(nfa, kinds, 0, &mut ends);
2264 return select(spans, &ends);
2265 }
2266 let site = flynnel::sched::call_site::caller_site().get();
2267 let key = n.min(u32::MAX as usize) as u32;
2268 let share = site.split_cpu_share_per_mille_for(key) as usize;
2269 let mid = (n * share / 1000).clamp(1, n - 1);
2270 let (host_part, device_part) = ends.split_at_mut(mid);
2271 let (host_ns, device_ns) = std::thread::scope(|s| {
2272 let device = s.spawn(|| {
2273 let t = std::time::Instant::now();
2274 if let Err(e) = device_anchor_ends(g, nfa, &kinds[mid..], mid, device_part) {
2275 eprintln!(
2276 "trex gpu: the device's share of a split scan failed and the host scanned it: {e:?}"
2277 );
2278 host_anchor_ends(nfa, kinds, mid, device_part);
2279 }
2280 elapsed_ns(t)
2281 });
2282 let t = std::time::Instant::now();
2283 host_anchor_ends(nfa, kinds, 0, host_part);
2284 let host_ns = elapsed_ns(t);
2285 (host_ns, device.join().expect("the device share's thread returns"))
2286 });
2287 site.record_split_for(key, mid, host_ns, n - mid, device_ns);
2288 select(spans, &ends)
2289 })
2290 }
2291
2292 fn device_anchor_ends(
2296 g: &Gpu,
2297 nfa: &crate::nfa::GpuNfa,
2298 suffix: &[u32],
2299 mid: usize,
2300 out: &mut [i32],
2301 ) -> Result<(), cudarc::driver::DriverError> {
2302 let n = suffix.len();
2303 let stream = g.ctx.default_stream();
2304 let d_kind = stream.clone_htod(suffix)?;
2305 let d_atom = stream.clone_htod(&nfa.atom_kind)?;
2306 let d_next = stream.clone_htod(&nfa.next_closure)?;
2307 let d_isatom = stream.clone_htod(&nfa.is_atom)?;
2308 let d_result = stream.alloc_zeros::<i32>(n)?;
2309 let n_i = n as i32;
2310 let nstates_i = nfa.nstates as i32;
2311 let start_closure = nfa.start_closure;
2312 let match_mask = nfa.match_mask;
2313 let block = 256u32;
2314 let cfg = LaunchConfig {
2315 grid_dim: (n.div_ceil(block as usize) as u32, 1, 1),
2316 block_dim: (block, 1, 1),
2317 shared_mem_bytes: 0,
2318 };
2319 let mut builder = stream.launch_builder(&g.func);
2320 builder.arg(&d_kind);
2321 builder.arg(&n_i);
2322 builder.arg(&d_atom);
2323 builder.arg(&d_next);
2324 builder.arg(&d_isatom);
2325 builder.arg(&start_closure);
2326 builder.arg(&match_mask);
2327 builder.arg(&nstates_i);
2328 builder.arg(&d_result);
2329 unsafe { builder.launch(cfg)? };
2333 let ends: Vec<i32> = stream.clone_dtoh(&d_result)?;
2334 stream.synchronize()?;
2335 let base = i32::try_from(mid).expect("a token index within the kernel's i32 width");
2336 for (slot, e) in out.iter_mut().zip(ends) {
2337 *slot = if e < 0 { -1 } else { e + base };
2338 }
2339 Ok(())
2340 }
2341
2342 pub(super) fn host_anchor_ends(nfa: &crate::nfa::GpuNfa, kinds: &[u32], first: usize, out: &mut [i32]) {
2347 use flynnel::JobPlan;
2348 use flynnel::sched::par_iter::for_each_chunk_indexed_min_leaf;
2349 let cores = std::thread::available_parallelism().map_or(1, std::num::NonZero::get);
2350 let min_leaf = out.len().div_ceil(cores * 4).max(64);
2351 let plan = JobPlan::new(0, out.len().min(u32::MAX as usize) as u32)
2352 .with_leaf_shape(flynnel::LeafShape::PortCompute);
2353 for_each_chunk_indexed_min_leaf(&plan, out, min_leaf, |start, slots| {
2354 for (i, slot) in slots.iter_mut().enumerate() {
2355 *slot = anchor_end(nfa, kinds, first + start + i);
2356 }
2357 });
2358 }
2359
2360 #[inline]
2363 fn anchor_end(nfa: &crate::nfa::GpuNfa, kinds: &[u32], a: usize) -> i32 {
2364 let n = kinds.len();
2365 let mut active = nfa.start_closure;
2366 let mut best = -1i32;
2367 for (k, &tk) in kinds.iter().enumerate().skip(a) {
2368 if active & nfa.match_mask != 0 {
2369 best = k as i32;
2370 }
2371 let mut next = 0u64;
2372 let mut m = active;
2373 while m != 0 {
2374 let pc = m.trailing_zeros() as usize;
2375 m &= m - 1;
2376 if nfa.is_atom[pc] != 0 {
2377 let ak = nfa.atom_kind[pc];
2378 if ak == u32::MAX || ak == tk {
2379 next |= nfa.next_closure[pc];
2380 }
2381 }
2382 }
2383 if next == 0 {
2384 return best;
2385 }
2386 active = next;
2387 }
2388 if active & nfa.match_mask != 0 { n as i32 } else { best }
2389 }
2390
2391 #[cfg(test)]
2392 mod tests {
2393 use super::*;
2394
2395 fn split_corpus() -> Vec<u8> {
2398 let mut s = String::new();
2399 for i in 0..3000u32 {
2400 s.push_str(&format!("tag {} {} word {}\n", i % 7, i % 13, i));
2401 if i % 5 == 0 {
2402 s.push_str("1 2 3 4 5 6 7 8\n");
2403 }
2404 }
2405 s.into_bytes()
2406 }
2407
2408 #[test]
2411 fn host_anchor_ends_select_the_cpu_engines_spans() {
2412 let input = split_corpus();
2413 let (kinds, spans) = crate::parallel_lex::lex_significant_parallel(&input);
2414 for src in ["\\W \\N", "\\N+", "\\N{2,4}", "(\\N \\W)+", ". .", "\\W \\W \\W"] {
2415 let pat = crate::parser::parse(src).expect("pattern parses");
2416 let nfa = crate::nfa::compile_for_gpu(&pat).expect("a kind-only pattern compiles for the kernel");
2417 let mut ends = vec![-1i32; kinds.len()];
2418 host_anchor_ends(&nfa, &kinds, 0, &mut ends);
2419 assert_eq!(select(&spans, &ends), crate::engine::scan(&pat, &input), "{src}");
2420 }
2421 }
2422
2423 #[test]
2426 fn a_split_at_any_share_selects_the_cpu_engines_spans() {
2427 let Some(g) = gpu() else {
2428 eprintln!("no CUDA device present; the split agreement test did not run");
2429 return;
2430 };
2431 let input = split_corpus();
2432 let (kinds, spans) = crate::parallel_lex::lex_significant_parallel(&input);
2433 let n = kinds.len();
2434 for src in ["\\W \\N", "\\N+", "(\\N \\W)+", "\\W \\W \\W"] {
2435 let pat = crate::parser::parse(src).expect("pattern parses");
2436 let nfa = crate::nfa::compile_for_gpu(&pat).expect("a kind-only pattern compiles for the kernel");
2437 let want = crate::engine::scan(&pat, &input);
2438 for mid in [1, n / 3, n / 2, n - 1] {
2439 let mut ends = vec![-1i32; n];
2440 let (host_part, device_part) = ends.split_at_mut(mid);
2441 host_anchor_ends(&nfa, &kinds, 0, host_part);
2442 device_anchor_ends(g, &nfa, &kinds[mid..], mid, device_part)
2443 .expect("the device ran its share");
2444 assert_eq!(select(&spans, &ends), want, "{src} split at {mid} of {n}");
2445 }
2446 }
2447 }
2448
2449 #[test]
2453 fn a_pipelined_scan_selects_the_cpu_engines_spans() {
2454 if gpu().is_none() {
2455 eprintln!("no CUDA device present; the pipelined scan test did not run");
2456 return;
2457 }
2458 let input = split_corpus().repeat(4);
2459 for src in ["\\W \\N", "\\N{2,4}", "\\W \\N \\W", ". .", "\\N"] {
2462 let pat = crate::parser::parse(src).expect("pattern parses");
2463 let want = crate::engine::scan(&pat, &input);
2464 assert!(!want.is_empty(), "{src} matches nothing, so nothing would be compared");
2465 for partitions in [1, 2, 3, 7, 64] {
2466 let (got, ph) = scan_pipelined(&pat, &input, partitions).expect("the device ran the pipeline");
2467 assert_eq!(got, want, "{src} over {partitions} ranges");
2468 assert!(ph.partitions >= 1, "{src} over {partitions} ranges lexed no range");
2469 }
2470 }
2471 let unbounded = crate::parser::parse("\\N+").expect("pattern parses");
2472 assert!(scan_pipelined(&unbounded, &input, 4).is_none(), "an unbounded pattern has no carry length");
2473 }
2474
2475 #[test]
2480 fn a_parts_scan_selects_the_cpu_engines_spans() {
2481 if gpu().is_none() {
2482 eprintln!("no CUDA device present; the parts scan test did not run");
2483 return;
2484 }
2485 let many = split_corpus().repeat(4);
2486 let one = b"tag 1 2 word 3\ntag 4 5 word 6\n".to_vec();
2487 for src in ["\\W \\N", "\\N{2,4}", "\\W \\N \\W", ". .", "\\N", "\\N+"] {
2490 let pat = crate::parser::parse(src).expect("pattern parses");
2491 for input in [&many, &one] {
2492 let want = crate::engine::scan(&pat, input);
2493 assert!(!want.is_empty(), "{src} matches nothing in {} bytes", input.len());
2494 let got = scan_parts(&pat, input).expect("the device ran the parts scan");
2495 assert_eq!(got, want, "{src} over {} bytes", input.len());
2496 }
2497 assert_eq!(scan_parts(&pat, b""), Some(Vec::new()), "{src} over an empty input");
2498 }
2499 }
2500
2501 #[test]
2505 fn a_resident_split_selects_the_cpu_engines_spans() {
2506 if gpu().is_none() {
2507 eprintln!("no CUDA device present; the resident split test did not run");
2508 return;
2509 }
2510 let input = split_corpus();
2511 let held = crate::gpu::GpuTokens::upload(&input).expect("the device holds the corpus");
2512 for src in ["\\W \\N", "\\N+", "\\N{2,4}", "\\W \\N \\W", ". .", "\\N"] {
2513 let pat = crate::parser::parse(src).expect("pattern parses");
2514 let want = crate::engine::scan(&pat, &input);
2515 assert!(!want.is_empty(), "{src} matches nothing, so nothing would be compared");
2516 for share in [1u32, 100, 500, 900, 999] {
2517 let got = held.scan_split(&pat, share).expect("the device ran the resident split");
2518 assert_eq!(got, want, "{src} at a host share of {share} per mille");
2519 }
2520 }
2521 let empty = crate::gpu::GpuTokens::upload(b"").expect("the device holds an empty corpus");
2522 let pat = crate::parser::parse("\\N").expect("pattern parses");
2523 assert_eq!(empty.scan_split(&pat, 500), Some(Vec::new()), "an empty corpus has no match");
2524 }
2525 }
2526}
2527
2528#[cfg(test)]
2529mod tests {
2530 use super::*;
2531
2532 use crate::parser::parse;
2533
2534 #[test]
2535 fn eligibility_matches_the_documented_subset() {
2536 for src in [
2539 "\\N",
2540 "\\W \\N",
2541 "\\N+",
2542 "\\W*",
2543 ".",
2544 ". .",
2545 "\\N{2,4}",
2546 "(\\N \\W)+",
2547 "\\M{>6}",
2548 "\\N{mag<3}",
2549 "\\W \\N{mag>=2}",
2550 "\\W:t",
2551 "\\W:x =x",
2552 "\\W:x \\N =case x",
2553 "\"lit\"",
2554 "\\d",
2555 "\\W:x \"lit\"",
2556 ] {
2557 assert!(gpu_eligible(&parse(src).unwrap()), "{src} should be eligible");
2558 }
2559 for src in [
2563 "\\N{>+1}",
2564 "\\W+:x =x",
2565 "\\W:x \\N* =x",
2566 "\\W:x =case x \"lit\"",
2567 "\\N | \\W",
2568 "\\W\\B(.*)",
2569 "@2 \\W",
2570 ". ~\"END\"",
2571 "`[a-z]+`",
2572 ] {
2573 assert!(!gpu_eligible(&parse(src).unwrap()), "{src} should be ineligible");
2574 }
2575 }
2576
2577 #[cfg(not(feature = "gpu"))]
2578 #[test]
2579 fn scan_gpu_is_none_without_the_feature() {
2580 assert!(scan_gpu(&parse("\\N \\W").unwrap(), b"12 kg").is_none());
2583 }
2584
2585 #[cfg(feature = "gpu")]
2586 #[test]
2587 fn the_device_agrees_with_the_cpu_over_many_tokens() {
2588 let pat = parse("\\W \\N").expect("pattern parses");
2594 let mut input = String::new();
2595 for i in 0..20_000 {
2596 input.push_str(match i % 4 {
2600 0 => "tag ",
2601 1 => "tag\t\t",
2602 2 => "tag \n ",
2603 _ => "tag ",
2604 });
2605 input.push_str(&(i % 997).to_string());
2606 input.push(' ');
2607 }
2608 let bytes = input.as_bytes();
2609 let Some(device) = scan_gpu(&pat, bytes) else {
2610 eprintln!("no CUDA device present; device-agreement test did not run");
2613 return;
2614 };
2615 let cpu = crate::engine::scan(&pat, bytes);
2616 assert_eq!(device.len(), cpu.len(), "device and cpu must find the same number of matches");
2617 assert_eq!(device, cpu, "device and cpu must agree on every span");
2618 assert!(!cpu.is_empty(), "the corpus must actually match, or this asserts nothing");
2619 }
2620
2621 #[cfg(feature = "gpu")]
2622 #[test]
2623 fn the_device_agrees_with_the_cpu_on_magnitude_tests() {
2624 let mut input = String::new();
2628 for i in 0..20_000u64 {
2629 input.push_str(match i % 3 {
2630 0 => "size ",
2631 1 => "tiny\t",
2632 _ => "enormousword \n ",
2633 });
2634 input.push_str(&(i * i * 37 % 10_000_019).to_string());
2635 input.push(' ');
2636 if i % 5 == 0 {
2637 input.push_str("0.004 ");
2638 }
2639 }
2640 let bytes = input.as_bytes();
2641 for src in [
2642 "\\N{mag>6}",
2643 "\\N{mag<3}",
2644 "\\M{>1}",
2645 "\\M{<=2} \\N{mag>=5}",
2646 "(\\W \\N{mag>4})+",
2647 ] {
2648 let pat = parse(src).expect("pattern parses");
2649 assert!(gpu_eligible(&pat), "{src} is in the device subset");
2650 let Some(device) = scan_gpu(&pat, bytes) else {
2651 eprintln!("no CUDA device present; magnitude agreement test did not run");
2652 return;
2653 };
2654 let cpu = crate::engine::scan(&pat, bytes);
2655 assert!(!cpu.is_empty(), "{src} must match the corpus, or this asserts nothing");
2656 assert_eq!(device, cpu, "{src}: device and cpu must agree on every span");
2657 let held = GpuTokens::upload_with_magnitudes(bytes).expect("the device ran the scan");
2658 assert_eq!(held.scan(&pat), Some(cpu), "{src}: held tokens and cpu must agree");
2659 let kinds_only = GpuTokens::upload(bytes).expect("the device ran the scan");
2660 assert_eq!(kinds_only.scan(&pat), None, "{src}: tokens held without magnitudes decline");
2661 }
2662 }
2663
2664 #[cfg(feature = "gpu")]
2665 #[test]
2666 fn the_device_agrees_with_the_cpu_on_binds() {
2667 let words = ["alpha", "Alpha", "beta", "BETA", "gamma"];
2671 let mut input = String::new();
2672 for i in 0..20_000usize {
2673 input.push_str(words[i % 5]);
2674 input.push(' ');
2675 input.push_str(words[(i * 7 + i / 3) % 5]);
2676 input.push_str(if i % 4 == 0 { "\t" } else { " " });
2677 input.push_str(&(i % 13).to_string());
2678 input.push(' ');
2679 if i % 3 == 0 {
2680 input.push_str("Alpha 7 alpha gamma 7 gamma delta x DELTA beta beta ");
2681 }
2682 }
2683 let bytes = input.as_bytes();
2684 for src in ["\\W:x =x", "\\W:x =case x", "\\W:x \\N =x", "\\W:x \\W =case x", "(\\W:x =case x)+", "\\W:t \\N"] {
2685 let pat = parse(src).expect("pattern parses");
2686 assert!(gpu_eligible(&pat), "{src} is in the device subset");
2687 let Some(device) = scan_gpu(&pat, bytes) else {
2688 eprintln!("no CUDA device present; bind agreement test did not run");
2689 return;
2690 };
2691 let cpu = crate::engine::scan(&pat, bytes);
2692 assert!(!cpu.is_empty(), "{src} must match the corpus, or this asserts nothing");
2693 assert_eq!(device, cpu, "{src}: device and cpu must agree on every span");
2694 let identity = GpuTokens::upload_with_properties(bytes, crate::orbit::OrbitGroup::Identity)
2695 .expect("the device ran the scan");
2696 let case = GpuTokens::upload_with_properties(bytes, crate::orbit::OrbitGroup::Case)
2697 .expect("the device ran the scan");
2698 let held = identity.scan(&pat).or_else(|| case.scan(&pat));
2699 assert_eq!(held, Some(cpu), "{src}: tokens held with their properties agree with the cpu");
2700 }
2701 }
2702
2703 #[cfg(feature = "gpu")]
2704 #[test]
2705 fn the_device_agrees_with_the_cpu_on_literals_and_byte_classes() {
2706 let mut input = String::new();
2710 for i in 0..20_000usize {
2711 input.push_str(["the", "The", "THE", "cat", "x1f"][i % 5]);
2712 input.push(' ');
2713 input.push_str(&(i % 97).to_string());
2714 input.push_str(if i % 3 == 0 { " , " } else { " ; " });
2715 input.push_str(["ABC", "abc", "a_b", "Zz9"][i % 4]);
2716 input.push('\n');
2717 }
2718 let bytes = input.as_bytes();
2719 for src in [
2720 "\"the\" \\N",
2721 "(?orbit:case \"the\") \\N",
2722 "\",\" \\w",
2723 "\\d \";\"",
2724 "\\u",
2725 "\\l \\W",
2726 "\\a+",
2727 "\\W:x \"the\"",
2728 ] {
2729 let pat = parse(src).expect("pattern parses");
2730 assert!(gpu_eligible(&pat), "{src} is in the device subset");
2731 let Some(device) = scan_gpu(&pat, bytes) else {
2732 eprintln!("no CUDA device present; literal and byte-class agreement test did not run");
2733 return;
2734 };
2735 let cpu = crate::engine::scan(&pat, bytes);
2736 assert!(!cpu.is_empty(), "{src} must match the corpus, or this asserts nothing");
2737 assert_eq!(device, cpu, "{src}: device and cpu must agree on every span");
2738 let identity = GpuTokens::upload_with_properties(bytes, crate::orbit::OrbitGroup::Identity)
2739 .expect("the device ran the scan");
2740 let case = GpuTokens::upload_with_properties(bytes, crate::orbit::OrbitGroup::Case)
2741 .expect("the device ran the scan");
2742 let held = identity.scan(&pat).or_else(|| case.scan(&pat));
2743 assert_eq!(held, Some(cpu), "{src}: tokens held with their properties agree with the cpu");
2744 }
2745 }
2746
2747 #[cfg(feature = "gpu")]
2748 #[test]
2749 fn the_device_agrees_with_the_cpu_on_spectral_tests() {
2750 let mut input = String::new();
2754 for i in 0..400usize {
2755 input.push_str("the quick brown fox jumps over the lazy dog and the cat sat on the mat ");
2756 if i % 3 == 0 {
2757 let mut x = 0x2545_f491_4f6c_dd1du64 ^ (i as u64);
2758 for _ in 0..200 {
2759 x = x.wrapping_mul(6_364_136_223_846_793_005).wrapping_add(1_442_695_040_888_963_407);
2760 input.push(b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/"[((x >> 58) % 64) as usize] as char);
2761 }
2762 input.push(' ');
2763 }
2764 if i % 5 == 0 {
2765 for r in 0..12 {
2766 input.push_str(&format!("{:03},{:03},{:03}\n", r, (r * 7) % 100, (r * 13) % 100));
2767 }
2768 }
2769 if i % 7 == 0 {
2770 input.push_str(&"x".repeat(150));
2771 input.push(' ');
2772 }
2773 }
2774 let bytes = input.as_bytes();
2775 for src in [
2776 "\\F{entropy>0.8}",
2777 "\\F{entropy<0.3} \\W",
2778 "\\F{period:any}",
2779 "\\F{texture:prose} \\W",
2780 "\\F{texture:data}",
2781 "\\F{onset}",
2782 "\\W \\F{entropy>0.8}",
2783 ] {
2784 let pat = parse(src).expect("pattern parses");
2785 assert!(gpu_eligible(&pat), "{src} is in the device subset");
2786 let Some(device) = scan_gpu(&pat, bytes) else {
2787 eprintln!("no CUDA device present; spectral agreement test did not run");
2788 return;
2789 };
2790 let cpu = crate::engine::scan(&pat, bytes);
2791 assert!(!cpu.is_empty(), "{src} must match the corpus, or this asserts nothing");
2792 assert_eq!(device, cpu, "{src}: device and cpu must agree on every span");
2793 let held = GpuTokens::upload_with_spectral(bytes, crate::orbit::OrbitGroup::Identity)
2794 .expect("the device ran the scan");
2795 assert_eq!(held.scan(&pat), Some(cpu), "{src}: tokens held with the spectral reading agree");
2796 let without = GpuTokens::upload_with_properties(bytes, crate::orbit::OrbitGroup::Identity)
2797 .expect("the device ran the scan");
2798 assert_eq!(without.scan(&pat), None, "{src}: tokens held without the reading decline");
2799 }
2800 assert!(!gpu_eligible(&parse("\\W:x \\F{onset}").unwrap()), "a bind beside a spectral test stays on the CPU");
2801 }
2802
2803 #[test]
2804 fn every_backend_returns_the_cpu_match_set() {
2805 let pat = parse("\\W \\N").unwrap();
2809 let input = b"tag 12 tag 34 tag 56";
2810 let baseline = crate::engine::scan(&pat, input);
2811 for b in [Backend::Auto, Backend::Gpu, Backend::Cpu] {
2812 let (m, _used) = scan_with_backend(&pat, input, b);
2813 assert_eq!(m, baseline, "{b:?} must match the CPU baseline");
2814 }
2815 }
2816
2817 #[test]
2818 fn auto_matches_the_cpu_engine_while_its_placement_learns() {
2819 let pat = parse("\\W \\N").unwrap();
2822 let mut input = String::new();
2823 for i in 0..40_000u32 {
2824 input.push_str(&format!("tag {}\n", i % 1000));
2825 }
2826 let want = crate::engine::scan(&pat, input.as_bytes());
2827 for _ in 0..40 {
2828 let (m, _used) = scan_with_backend(&pat, input.as_bytes(), Backend::Auto);
2829 assert_eq!(m, want, "Auto returns the CPU engine's spans whichever route it took");
2830 }
2831 }
2832
2833 #[test]
2834 fn auto_routes_only_kind_patterns_to_the_device() {
2835 for src in ["\\W \\N", "\\N+", ". ."] {
2836 assert!(gpu_auto_routes(&parse(src).unwrap()), "{src} reads only token kinds");
2837 }
2838 let input = "tag 12 the 34 ".repeat(2_000);
2839 for src in ["\\W:x =x", "\\N{mag>6}", "\"the\" \\W", "\\d", "\\W:t \\N"] {
2840 let pat = parse(src).unwrap();
2841 assert!(gpu_eligible(&pat), "{src} is in the device subset");
2842 assert!(!gpu_auto_routes(&pat), "{src} reads a token property");
2843 let (m, used) = scan_with_backend(&pat, input.as_bytes(), Backend::Auto);
2844 assert_eq!(used, BackendUsed::Cpu, "{src}: auto keeps a property-reading pattern on the cores");
2845 assert_eq!(m, crate::engine::scan(&pat, input.as_bytes()), "{src}: auto returns the CPU matches");
2846 }
2847 }
2848
2849 #[cfg(not(feature = "gpu"))]
2850 #[test]
2851 fn without_a_device_every_backend_runs_on_the_cpu() {
2852 assert!(!device_available());
2855 let pat = parse("\\W \\N").unwrap();
2856 let input = b"tag 12 tag 34";
2857 for b in [Backend::Auto, Backend::Gpu, Backend::Cpu] {
2858 let (_m, used) = scan_with_backend(&pat, input, b);
2859 assert_eq!(used, BackendUsed::Cpu, "no device: {b:?} must use the CPU");
2860 }
2861 }
2862}