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#[derive(Clone, Copy, Debug, PartialEq, Eq)]
684pub struct Engine {
685 pub backend: Backend,
687 pub dual_grain: bool,
689 pub chunk_size: Option<usize>,
691}
692
693impl Engine {
694 pub const PLAIN: Engine = Engine { backend: Backend::Auto, dual_grain: false, chunk_size: None };
696
697 #[must_use]
700 pub fn is_plain(&self) -> bool {
701 self.backend == Backend::Auto && !self.dual_grain && self.chunk_size.is_none()
702 }
703}
704
705impl Default for Engine {
706 fn default() -> Self {
707 Engine::PLAIN
708 }
709}
710
711#[derive(Clone, Copy, Debug, PartialEq, Eq)]
713pub enum Ran {
714 Device,
716 DeviceDeclined,
720 DualGrain(crate::dual_grain::GrainTiming),
722 Chunked,
724 Routed(BackendUsed),
726}
727
728#[must_use]
737pub fn scan_engine(pattern: &Pattern, input: &[u8], engine: &Engine) -> (Vec<Span>, Ran) {
738 if engine.backend == Backend::Gpu {
739 let (spans, used) = scan_with_backend(pattern, input, Backend::Gpu);
740 (spans, if used == BackendUsed::Gpu { Ran::Device } else { Ran::DeviceDeclined })
741 } else if engine.dual_grain {
742 let (spans, timing) = crate::dual_grain::scan_dual_grain(pattern, input);
743 (spans, Ran::DualGrain(timing))
744 } else if let Some(size) = engine.chunk_size {
745 assert!(size > 0, "a chunk size of zero names no chunk");
746 (crate::streaming::scan_chunked(pattern, input.chunks(size)), Ran::Chunked)
747 } else {
748 let (spans, used) = scan_with_backend(pattern, input, engine.backend);
749 (spans, Ran::Routed(used))
750 }
751}
752
753#[must_use]
767pub fn significant_stream(toks: &[Token]) -> (Vec<u32>, Vec<(u32, u32)>) {
768 let mut kinds = Vec::with_capacity(toks.len());
769 let mut spans = Vec::with_capacity(toks.len());
770 for t in toks {
771 if t.kind != TokenKind::Whitespace {
772 kinds.push(t.kind.code());
773 spans.push((t.start, t.end));
774 }
775 }
776 (kinds, spans)
777}
778
779#[must_use]
783pub fn significant_magnitudes(toks: &[Token], input: &[u8]) -> Vec<f32> {
784 toks.iter()
785 .filter(|t| t.kind != TokenKind::Whitespace)
786 .map(|t| crate::magnitude::token_magnitude(t.kind, &input[t.start()..t.end()]) * 100.0)
787 .collect()
788}
789
790pub(crate) struct ClassIds {
795 group: crate::orbit::OrbitGroup,
796 ids: std::collections::HashMap<Vec<u8>, u32>,
797}
798
799impl ClassIds {
800 pub(crate) fn new(group: crate::orbit::OrbitGroup) -> Self {
801 Self { group, ids: std::collections::HashMap::new() }
802 }
803
804 #[cfg(feature = "gpu")]
806 pub(crate) fn group(&self) -> crate::orbit::OrbitGroup {
807 self.group
808 }
809
810 pub(crate) fn id_of(&mut self, bytes: &[u8]) -> u32 {
813 let next = u32::try_from(self.ids.len()).expect("fewer than 2^32 distinct classes");
814 match self.group {
815 crate::orbit::OrbitGroup::Identity => match self.ids.get(bytes) {
816 Some(&id) => id,
817 None => {
818 self.ids.insert(bytes.to_vec(), next);
819 next
820 }
821 },
822 g => *self.ids.entry(crate::orbit::canonical(bytes, g).into_bytes()).or_insert(next),
823 }
824 }
825
826 #[cfg(feature = "gpu")]
828 pub(crate) fn find(&self, bytes: &[u8]) -> Option<u32> {
829 match self.group {
830 crate::orbit::OrbitGroup::Identity => self.ids.get(bytes).copied(),
831 g => self.ids.get(crate::orbit::canonical(bytes, g).as_bytes()).copied(),
832 }
833 }
834}
835
836#[must_use]
842pub fn significant_classes(
843 toks: &[Token],
844 input: &[u8],
845 group: crate::orbit::OrbitGroup,
846 literals: &[&[u8]],
847) -> (Vec<u32>, Vec<u32>) {
848 let mut ids = ClassIds::new(group);
849 let literal_ids: Vec<u32> = literals.iter().map(|&lit| ids.id_of(lit)).collect();
850 let token_ids: Vec<u32> = toks
851 .iter()
852 .filter(|t| t.kind != TokenKind::Whitespace)
853 .map(|t| ids.id_of(&input[t.start()..t.end()]))
854 .collect();
855 (token_ids, literal_ids)
856}
857
858#[must_use]
865pub fn significant_spectral(toks: &[Token], field: &crate::spectral::SpectralField) -> (Vec<f32>, Vec<u32>) {
866 let mut entropy = Vec::with_capacity(toks.len());
867 let mut packed = Vec::with_capacity(toks.len());
868 for t in toks.iter().filter(|t| t.kind != TokenKind::Whitespace) {
869 let sig = field.signature(t.start(), t.end());
870 entropy.push(sig.entropy * 100.0);
871 let texture = crate::nfa::texture_id(crate::spectral::texture_of(&sig));
872 let onset = u32::from(field.boundary_in(t.start(), t.end()));
873 packed.push(u32::from(sig.period) | (texture << 16) | (onset << 20));
874 }
875 (entropy, packed)
876}
877
878#[must_use]
883pub fn significant_byte_masks(toks: &[Token], input: &[u8]) -> Vec<u32> {
884 use crate::ast::ByteClass;
885 const CLASSES: [ByteClass; 6] =
886 [ByteClass::Digit, ByteClass::Word, ByteClass::Hex, ByteClass::Alpha, ByteClass::Upper, ByteClass::Lower];
887 toks.iter()
888 .filter(|t| t.kind != TokenKind::Whitespace)
889 .map(|t| {
890 let bytes = &input[t.start()..t.end()];
891 CLASSES
892 .iter()
893 .filter(|&&bc| crate::engine::byte_class_matches(bc, bytes))
894 .map(|&bc| crate::nfa::byte_class_bit(bc).expect("every class but whitespace has a bit"))
895 .fold(0, |mask, bit| mask | bit)
896 })
897 .collect()
898}
899
900#[cfg(feature = "gpu")]
904#[allow(unsafe_code)]
905mod cuda {
906 use std::sync::{Arc, OnceLock};
907
908 use cudarc::driver::{CudaContext, CudaFunction, LaunchConfig, PushKernelArg};
909 use cudarc::nvrtc::Ptx;
910
911 use super::Span;
912 use crate::ast::Pattern;
913 use crate::ends_simd::{next_match_from, next_match_from_scalar, takes_vector};
914 use crate::token::Token;
915
916 const PTX: &str = include_str!(env!("TREX_SCAN_PTX"));
920
921 struct Gpu {
923 ctx: Arc<CudaContext>,
924 func: CudaFunction,
925 }
926
927 pub(super) fn device_present() -> bool {
930 gpu().is_some()
931 }
932
933 pub(super) fn load_or_cpu<T>(
953 kernel: &str,
954 load: fn() -> Result<T, cudarc::driver::DriverError>,
955 ) -> Option<T> {
956 let prior = std::panic::take_hook();
957 std::panic::set_hook(Box::new(|_| {}));
958 let probed = std::panic::catch_unwind(load);
959 std::panic::set_hook(prior);
960 match probed {
961 Ok(Ok(loaded)) => Some(loaded),
962 Ok(Err(driver)) => {
963 crate::trace::rung("gpu", &format!("{kernel}: no usable device ({driver:?})"), 0);
964 None
965 }
966 Err(panicked) => {
967 let why = panicked
968 .downcast_ref::<String>()
969 .map(String::as_str)
970 .or_else(|| panicked.downcast_ref::<&str>().copied())
971 .unwrap_or("the CUDA library could not be loaded");
972 crate::trace::rung("gpu", &format!("{kernel}: no CUDA library ({why})"), 0);
973 None
974 }
975 }
976 }
977
978 fn gpu() -> Option<&'static Gpu> {
981 static G: OnceLock<Option<Gpu>> = OnceLock::new();
982 G.get_or_init(|| load_or_cpu("scan", load_scan_kernel)).as_ref()
983 }
984
985 fn load_scan_kernel() -> Result<Gpu, cudarc::driver::DriverError> {
987 let ctx = CudaContext::new(0)?;
988 let module = ctx.load_module(Ptx::from_src(PTX))?;
989 let func = module.load_function("trex_scan")?;
990 Ok(Gpu { ctx, func })
991 }
992
993 pub(super) fn scan_phases(
1008 pattern: &Pattern,
1009 input: &[u8],
1010 ) -> Option<(Vec<crate::engine::Span>, super::GpuPhases)> {
1011 use std::time::Instant;
1012
1013 let nfa = crate::nfa::compile_for_gpu(pattern)?;
1014 let g = gpu()?;
1015 if nfa.reads_magnitude
1016 || nfa.binds_registers
1017 || nfa.class_group.is_some()
1018 || nfa.reads_bytes
1019 || nfa.reads_spectral
1020 {
1021 eprintln!(
1022 "trex gpu phase probe: phases are timed for kind-only patterns, and this one reads a per-token property"
1023 );
1024 return None;
1025 }
1026 let mut ph = super::GpuPhases::default();
1027
1028 let t = Instant::now();
1031 crate::parallel_lex::lex_significant_parallel_held(input, |kinds, spans| {
1032 ph.lex_us = t.elapsed().as_secs_f64() * 1e6;
1033 let n = kinds.len();
1034 if n == 0 {
1035 return Some((Vec::new(), ph));
1036 }
1037 ph.tokens = n;
1038 phases_lexed(g, &nfa, kinds, spans, ph)
1039 })
1040 }
1041
1042 fn phases_lexed(
1044 g: &Gpu,
1045 nfa: &crate::nfa::GpuNfa,
1046 kinds: &[u32],
1047 spans: &[(u32, u32)],
1048 mut ph: super::GpuPhases,
1049 ) -> Option<(Vec<Span>, super::GpuPhases)> {
1050 use std::time::Instant;
1051
1052 let n = kinds.len();
1053 let stream = g.ctx.default_stream();
1054 let t = Instant::now();
1055 let staged = (|| {
1056 let d_kind = stream.clone_htod(kinds)?;
1057 let d_atom = stream.clone_htod(&nfa.atom_kind)?;
1058 let d_next = stream.clone_htod(&nfa.next_closure)?;
1059 let d_isatom = stream.clone_htod(&nfa.is_atom)?;
1060 let d_result = stream.alloc_zeros::<i32>(n)?;
1061 stream.synchronize()?;
1065 Ok::<_, cudarc::driver::DriverError>((d_kind, d_atom, d_next, d_isatom, d_result))
1066 })();
1067 let (d_kind, d_atom, d_next, d_isatom, d_result) = match staged {
1068 Ok(v) => v,
1069 Err(e) => {
1070 eprintln!("trex gpu phase probe: staging {n} tokens to the device failed: {e:?}");
1074 return None;
1075 }
1076 };
1077 ph.upload_us = t.elapsed().as_secs_f64() * 1e6;
1078
1079 let n_i = n as i32;
1080 let nstates_i = nfa.nstates as i32;
1081 let start_closure = nfa.start_closure;
1082 let match_mask = nfa.match_mask;
1083 let block = 256u32;
1084 let cfg = LaunchConfig {
1085 grid_dim: (n.div_ceil(block as usize) as u32, 1, 1),
1086 block_dim: (block, 1, 1),
1087 shared_mem_bytes: 0,
1088 };
1089
1090 let t = Instant::now();
1091 let ran = (|| {
1092 let mut builder = stream.launch_builder(&g.func);
1093 builder.arg(&d_kind);
1094 builder.arg(&n_i);
1095 builder.arg(&d_atom);
1096 builder.arg(&d_next);
1097 builder.arg(&d_isatom);
1098 builder.arg(&start_closure);
1099 builder.arg(&match_mask);
1100 builder.arg(&nstates_i);
1101 builder.arg(&d_result);
1102 unsafe { builder.launch(cfg)? };
1106 let result: Vec<i32> = stream.clone_dtoh(&d_result)?;
1107 stream.synchronize()?;
1108 Ok::<_, cudarc::driver::DriverError>(result)
1109 })();
1110 let result = match ran {
1111 Ok(r) => r,
1112 Err(e) => {
1113 eprintln!("trex gpu phase probe: the kernel over {n} tokens failed: {e:?}");
1114 return None;
1115 }
1116 };
1117 ph.kernel_and_download_us = t.elapsed().as_secs_f64() * 1e6;
1118
1119 let t = Instant::now();
1120 let out = select(spans, &result);
1121 ph.select_us = t.elapsed().as_secs_f64() * 1e6;
1122 Some((out, ph))
1123 }
1124
1125 pub(super) fn scan_pipelined(
1139 pattern: &Pattern,
1140 input: &[u8],
1141 partitions: usize,
1142 ) -> Option<(Vec<Span>, super::PipelinePhases)> {
1143 use std::sync::mpsc::{TryRecvError, channel, sync_channel};
1144 use std::time::Instant;
1145
1146 use crate::parallel_lex::SignificantWorkspace;
1147
1148 let nfa = crate::nfa::compile_for_gpu(pattern)?;
1149 if nfa.reads_magnitude
1150 || nfa.binds_registers
1151 || nfa.class_group.is_some()
1152 || nfa.reads_bytes
1153 || nfa.reads_spectral
1154 {
1155 return None;
1156 }
1157 let max_len = crate::nfa::bounded_max_len(pattern)?;
1158 let g = gpu()?;
1159 let carry = max_len.saturating_sub(1);
1160
1161 let wall = Instant::now();
1162 let mut ph = super::PipelinePhases::default();
1163 let t = Instant::now();
1164 let blobs = crate::lexer::blob_runs_parallel(input);
1165 ph.blobs_us = t.elapsed().as_secs_f64() * 1e6;
1166 let bounds = crate::parallel_lex::safe_boundaries(input, partitions.max(1));
1167 ph.partitions = bounds.len().saturating_sub(1);
1168
1169 let (to_device, from_lexer) = sync_channel::<SignificantWorkspace>(1);
1172 let (to_lexer, from_device) = channel::<SignificantWorkspace>();
1173 let nfa_ref = &nfa;
1174 let (lex_us, staged) = std::thread::scope(|sc| {
1175 let device = sc.spawn(move || device_stage(g, nfa_ref, carry, from_lexer, to_lexer));
1176 let mut lex_us = 0.0;
1177 let mut spare: Vec<SignificantWorkspace> =
1182 (0..3).map(|_| SignificantWorkspace::default()).collect();
1183 for w in bounds.windows(2) {
1184 let mut ws = match from_device.try_recv() {
1195 Ok(ws) => ws,
1196 Err(TryRecvError::Empty) => match spare.pop() {
1197 Some(ws) => ws,
1198 None => match from_device.recv() {
1199 Ok(ws) => ws,
1200 Err(std::sync::mpsc::RecvError) => break,
1201 },
1202 },
1203 Err(TryRecvError::Disconnected) => break,
1204 };
1205 let t = Instant::now();
1211 ws.kinds.clear();
1212 ws.spans.clear();
1213 ws.kinds.resize(carry, 0);
1214 ws.spans.resize(carry, (0, 0));
1215 crate::parallel_lex::lex_significant_range_onto(input, w[0], w[1], &blobs, &mut ws);
1216 lex_us += t.elapsed().as_secs_f64() * 1e6;
1217 if to_device.send(ws).is_err() {
1218 break;
1221 }
1222 }
1223 drop(to_device);
1224 (lex_us, device.join().expect("the pipeline's device thread returns"))
1225 });
1226 let staged = match staged {
1227 Ok(s) => s,
1228 Err(e) => {
1229 eprintln!("trex gpu pipeline: a window on the device failed: {e:?}");
1230 return None;
1231 }
1232 };
1233 ph.lex_us = lex_us;
1234 ph.tokens = staged.tokens;
1235 ph.device_us = staged.device_us;
1236 ph.select_us = staged.select_us;
1237 ph.carry_us = staged.carry_us;
1238 ph.wall_us = wall.elapsed().as_secs_f64() * 1e6;
1239 Some((staged.spans, ph))
1240 }
1241
1242 struct Staged {
1245 spans: Vec<Span>,
1246 tokens: usize,
1247 device_us: f64,
1248 select_us: f64,
1249 carry_us: f64,
1250 }
1251
1252 fn device_stage(
1258 g: &Gpu,
1259 nfa: &crate::nfa::GpuNfa,
1260 carry: usize,
1261 ranges: std::sync::mpsc::Receiver<crate::parallel_lex::SignificantWorkspace>,
1262 spent: std::sync::mpsc::Sender<crate::parallel_lex::SignificantWorkspace>,
1263 ) -> Result<Staged, cudarc::driver::DriverError> {
1264 use std::time::Instant;
1265
1266 let stream = g.ctx.default_stream();
1267 let d_atom = stream.clone_htod(&nfa.atom_kind)?;
1268 let d_next = stream.clone_htod(&nfa.next_closure)?;
1269 let d_isatom = stream.clone_htod(&nfa.is_atom)?;
1270 let mut staged =
1271 Staged { spans: Vec::new(), tokens: 0, device_us: 0.0, select_us: 0.0, carry_us: 0.0 };
1272 let mut kinds: Vec<u32> = Vec::new();
1276 let mut spans: Vec<(u32, u32)> = Vec::new();
1277 let mut next = 0usize;
1279 let mut last = false;
1280 while !last {
1281 let (mut window, start) = match ranges.recv() {
1287 Ok(ws) => {
1288 let t = Instant::now();
1289 let start = carry - kinds.len();
1290 staged.tokens += ws.kinds.len() - carry;
1291 let mut ws = ws;
1292 ws.kinds[start..carry].copy_from_slice(&kinds);
1293 ws.spans[start..carry].copy_from_slice(&spans);
1294 staged.carry_us += t.elapsed().as_secs_f64() * 1e6;
1295 (Some(ws), start)
1296 }
1297 Err(std::sync::mpsc::RecvError) => {
1300 last = true;
1301 (None, 0)
1302 }
1303 };
1304 let (wk, wsp): (&[u32], &[(u32, u32)]) = match &window {
1305 Some(ws) => (&ws.kinds[start..], &ws.spans[start..]),
1306 None => (&kinds, &spans),
1307 };
1308 let n = wk.len();
1309 if n == 0 {
1310 if let Some(ws) = window.take()
1314 && let Err(returned) = spent.send(ws)
1315 {
1316 drop(returned);
1317 }
1318 continue;
1319 }
1320 let t = Instant::now();
1321 let ends = window_ends(g, &stream, nfa, wk, &d_atom, &d_next, &d_isatom)?;
1322 staged.device_us += t.elapsed().as_secs_f64() * 1e6;
1323
1324 let t = Instant::now();
1325 let horizon = if last { n } else { n.saturating_sub(carry) };
1326 while next < horizon {
1327 let best = ends[next];
1328 if best > next as i32 {
1329 let e = best as usize;
1330 staged.spans.push(Span { start: wsp[next].0, end: wsp[e - 1].1 });
1331 next = e;
1332 } else {
1333 next += 1;
1334 }
1335 }
1336 staged.select_us += t.elapsed().as_secs_f64() * 1e6;
1337 next -= horizon;
1338
1339 let t = Instant::now();
1340 let (tail_kinds, tail_spans) = (wk[horizon..].to_vec(), wsp[horizon..].to_vec());
1341 kinds = tail_kinds;
1342 spans = tail_spans;
1343 staged.carry_us += t.elapsed().as_secs_f64() * 1e6;
1344 if let Some(ws) = window.take()
1348 && let Err(returned) = spent.send(ws)
1349 {
1350 drop(returned);
1351 }
1352 }
1353 Ok(staged)
1354 }
1355
1356 #[allow(clippy::too_many_arguments)]
1360 fn window_ends(
1361 g: &Gpu,
1362 stream: &Arc<cudarc::driver::CudaStream>,
1363 nfa: &crate::nfa::GpuNfa,
1364 kinds: &[u32],
1365 d_atom: &cudarc::driver::CudaSlice<u32>,
1366 d_next: &cudarc::driver::CudaSlice<u64>,
1367 d_isatom: &cudarc::driver::CudaSlice<u32>,
1368 ) -> Result<Vec<i32>, cudarc::driver::DriverError> {
1369 let d_kind = stream.clone_htod(kinds)?;
1370 device_ends(g, stream, nfa, &d_kind, kinds.len(), d_atom, d_next, d_isatom)
1371 }
1372
1373 #[allow(clippy::too_many_arguments)]
1377 fn device_view_ends(
1378 g: &Gpu,
1379 stream: &Arc<cudarc::driver::CudaStream>,
1380 nfa: &crate::nfa::GpuNfa,
1381 d_kind: &cudarc::driver::CudaView<'_, u32>,
1382 out: &mut [i32],
1383 d_atom: &cudarc::driver::CudaSlice<u32>,
1384 d_next: &cudarc::driver::CudaSlice<u64>,
1385 d_isatom: &cudarc::driver::CudaSlice<u32>,
1386 ) -> Result<(), cudarc::driver::DriverError> {
1387 let n = out.len();
1388 let d_result = stream.alloc_zeros::<i32>(n)?;
1389 let n_i = n as i32;
1390 let nstates_i = nfa.nstates as i32;
1391 let start_closure = nfa.start_closure;
1392 let match_mask = nfa.match_mask;
1393 let block = 256u32;
1394 let cfg = LaunchConfig {
1395 grid_dim: (n.div_ceil(block as usize) as u32, 1, 1),
1396 block_dim: (block, 1, 1),
1397 shared_mem_bytes: 0,
1398 };
1399 let mut builder = stream.launch_builder(&g.func);
1400 builder.arg(d_kind);
1401 builder.arg(&n_i);
1402 builder.arg(d_atom);
1403 builder.arg(d_next);
1404 builder.arg(d_isatom);
1405 builder.arg(&start_closure);
1406 builder.arg(&match_mask);
1407 builder.arg(&nstates_i);
1408 builder.arg(&d_result);
1409 unsafe { builder.launch(cfg)? };
1413 stream.memcpy_dtoh(&d_result, out)?;
1414 stream.synchronize()?;
1415 Ok(())
1416 }
1417
1418 #[allow(clippy::too_many_arguments)]
1422 fn device_ends(
1423 g: &Gpu,
1424 stream: &Arc<cudarc::driver::CudaStream>,
1425 nfa: &crate::nfa::GpuNfa,
1426 d_kind: &cudarc::driver::CudaSlice<u32>,
1427 n: usize,
1428 d_atom: &cudarc::driver::CudaSlice<u32>,
1429 d_next: &cudarc::driver::CudaSlice<u64>,
1430 d_isatom: &cudarc::driver::CudaSlice<u32>,
1431 ) -> Result<Vec<i32>, cudarc::driver::DriverError> {
1432 let d_result = stream.alloc_zeros::<i32>(n)?;
1433 let n_i = n as i32;
1434 let nstates_i = nfa.nstates as i32;
1435 let start_closure = nfa.start_closure;
1436 let match_mask = nfa.match_mask;
1437 let block = 256u32;
1438 let cfg = LaunchConfig {
1439 grid_dim: (n.div_ceil(block as usize) as u32, 1, 1),
1440 block_dim: (block, 1, 1),
1441 shared_mem_bytes: 0,
1442 };
1443 let mut builder = stream.launch_builder(&g.func);
1444 builder.arg(d_kind);
1445 builder.arg(&n_i);
1446 builder.arg(d_atom);
1447 builder.arg(d_next);
1448 builder.arg(d_isatom);
1449 builder.arg(&start_closure);
1450 builder.arg(&match_mask);
1451 builder.arg(&nstates_i);
1452 builder.arg(&d_result);
1453 unsafe { builder.launch(cfg)? };
1457 let ends: Vec<i32> = stream.clone_dtoh(&d_result)?;
1458 stream.synchronize()?;
1459 Ok(ends)
1460 }
1461
1462 pub(super) fn scan(pattern: &Pattern, input: &[u8]) -> Option<Vec<Span>> {
1463 let nfa = crate::nfa::compile_for_gpu(pattern)?;
1464 let g = gpu()?;
1465 if nfa.reads_magnitude
1466 || nfa.binds_registers
1467 || nfa.class_group.is_some()
1468 || nfa.reads_bytes
1469 || nfa.reads_spectral
1470 {
1471 return scan_props(g, &nfa, input);
1472 }
1473
1474 crate::parallel_lex::lex_significant_parallel_held(input, |kinds, spans| match scan_lexed(g, &nfa, kinds, spans) {
1477 Ok(found) => Some(found),
1478 Err(e) => {
1479 eprintln!("trex gpu scan: a device call failed: {e:?}");
1480 None
1481 }
1482 })
1483 }
1484
1485 fn scan_lexed(
1487 g: &Gpu,
1488 nfa: &crate::nfa::GpuNfa,
1489 kinds: &[u32],
1490 spans: &[(u32, u32)],
1491 ) -> Result<Vec<Span>, cudarc::driver::DriverError> {
1492 let n = kinds.len();
1493 if n == 0 {
1494 return Ok(Vec::new());
1495 }
1496
1497 let stream = g.ctx.default_stream();
1498 let d_kind = stream.clone_htod(kinds)?;
1499 let d_atom = stream.clone_htod(&nfa.atom_kind)?;
1500 let d_next = stream.clone_htod(&nfa.next_closure)?;
1501 let d_isatom = stream.clone_htod(&nfa.is_atom)?;
1502 let result = device_ends(g, &stream, nfa, &d_kind, n, &d_atom, &d_next, &d_isatom)?;
1503 Ok(select(spans, &result))
1504 }
1505
1506 pub(super) fn scan_split_now(pattern: &Pattern, input: &[u8]) -> Option<Vec<Span>> {
1509 let nfa = crate::nfa::compile_for_gpu(pattern)?;
1510 if nfa.reads_magnitude
1511 || nfa.binds_registers
1512 || nfa.class_group.is_some()
1513 || nfa.reads_bytes
1514 || nfa.reads_spectral
1515 {
1516 return None;
1517 }
1518 let g = gpu()?;
1519 Some(scan_split(g, &nfa, input))
1520 }
1521
1522 pub(super) fn scan_parts(pattern: &Pattern, input: &[u8]) -> Option<Vec<Span>> {
1529 let nfa = crate::nfa::compile_for_gpu(pattern)?;
1530 if nfa.reads_magnitude
1531 || nfa.binds_registers
1532 || nfa.class_group.is_some()
1533 || nfa.reads_bytes
1534 || nfa.reads_spectral
1535 {
1536 return None;
1537 }
1538 let g = gpu()?;
1539 crate::parallel_lex::lex_significant_parts_held(input, |parts| match scan_lexed_parts(g, &nfa, parts) {
1540 Ok(spans) => Some(spans),
1541 Err(e) => {
1542 eprintln!("trex gpu parts scan: a device call failed: {e:?}");
1543 None
1544 }
1545 })
1546 }
1547
1548 fn scan_lexed_parts(
1551 g: &Gpu,
1552 nfa: &crate::nfa::GpuNfa,
1553 parts: &[crate::lexer::Significant],
1554 ) -> Result<Vec<Span>, cudarc::driver::DriverError> {
1555 let n: usize = parts.iter().map(|p| p.kinds.len()).sum();
1556 if n == 0 {
1557 return Ok(Vec::new());
1558 }
1559 let stream = g.ctx.default_stream();
1560 let mut d_kind = unsafe { stream.alloc::<u32>(n) }?;
1564 let mut off = 0usize;
1565 for part in parts {
1566 let len = part.kinds.len();
1567 if len > 0 {
1568 stream.memcpy_htod(part.kinds.as_slice(), &mut d_kind.slice_mut(off..off + len))?;
1569 }
1570 off += len;
1571 }
1572 let d_atom = stream.clone_htod(&nfa.atom_kind)?;
1573 let d_next = stream.clone_htod(&nfa.next_closure)?;
1574 let d_isatom = stream.clone_htod(&nfa.is_atom)?;
1575 let ends = device_ends(g, &stream, nfa, &d_kind, n, &d_atom, &d_next, &d_isatom)?;
1576 Ok(select_parts(parts, &ends))
1577 }
1578
1579 fn select_parts(parts: &[crate::lexer::Significant], ends: &[i32]) -> Vec<Span> {
1583 if takes_vector(ends) {
1586 select_parts_with(parts, ends, next_match_from)
1587 } else {
1588 select_parts_with(parts, ends, next_match_from_scalar)
1589 }
1590 }
1591
1592 fn select_parts_with(
1594 parts: &[crate::lexer::Significant],
1595 ends: &[i32],
1596 scan: impl Fn(&[i32], usize, usize) -> Option<usize>,
1597 ) -> Vec<Span> {
1598 let mut out = Vec::with_capacity(ends.len() / 2 + 1);
1599 let mut a = 0usize;
1602 let mut base = 0usize;
1603 for (p, part) in parts.iter().enumerate() {
1604 let end = base + part.spans.len();
1605 while let Some(m) = scan(ends, a, end) {
1606 let e = ends[m] as usize; let last = if e <= end {
1608 part.spans[e - 1 - base].1
1609 } else {
1610 span_end_in(&parts[p + 1..], end, e - 1)
1611 };
1612 out.push(Span { start: part.spans[m - base].0, end: last });
1613 a = e;
1614 }
1615 a = a.max(end);
1619 base = end;
1620 }
1621 out
1622 }
1623
1624 fn span_end_in(rest: &[crate::lexer::Significant], mut base: usize, i: usize) -> u32 {
1627 for part in rest {
1628 let len = part.spans.len();
1629 if i < base + len {
1630 return part.spans[i - base].1;
1631 }
1632 base += len;
1633 }
1634 panic!("token {i} is past the lexed chunks, whose tokens end at {base}");
1635 }
1636
1637 pub(super) struct ResidentKinds {
1641 d_kind: Option<cudarc::driver::CudaSlice<u32>>,
1642 d_mag: Option<cudarc::driver::CudaSlice<f32>>,
1643 d_bytes: Option<cudarc::driver::CudaSlice<u32>>,
1644 d_class: Option<cudarc::driver::CudaSlice<u32>>,
1645 d_ent: Option<cudarc::driver::CudaSlice<f32>>,
1646 d_spec: Option<cudarc::driver::CudaSlice<u32>>,
1647 magnitudes: bool,
1648 spectral: bool,
1649 }
1650
1651 pub(super) fn upload_kinds(kinds: &[u32], mags: Option<&[f32]>) -> Option<ResidentKinds> {
1654 let g = gpu()?;
1655 if let Some(m) = mags {
1656 assert_eq!(m.len(), kinds.len(), "one magnitude per significant token");
1657 }
1658 let magnitudes = mags.is_some();
1659 if kinds.is_empty() {
1660 return Some(ResidentKinds {
1661 d_kind: None,
1662 d_mag: None,
1663 d_bytes: None,
1664 d_class: None,
1665 d_ent: None,
1666 d_spec: None,
1667 magnitudes,
1668 spectral: false,
1669 });
1670 }
1671 let stream = g.ctx.default_stream();
1672 let held = (|| {
1673 let d_kind = stream.clone_htod(kinds)?;
1674 let d_mag = match mags {
1675 Some(m) => Some(stream.clone_htod(m)?),
1676 None => None,
1677 };
1678 Ok::<_, cudarc::driver::DriverError>((d_kind, d_mag))
1679 })();
1680 match held {
1681 Ok((d_kind, d_mag)) => Some(ResidentKinds {
1682 d_kind: Some(d_kind),
1683 d_mag,
1684 d_bytes: None,
1685 d_class: None,
1686 d_ent: None,
1687 d_spec: None,
1688 magnitudes,
1689 spectral: false,
1690 }),
1691 Err(e) => {
1692 eprintln!("trex gpu: holding {} tokens on the device failed: {e:?}", kinds.len());
1693 None
1694 }
1695 }
1696 }
1697
1698 pub(super) fn upload_properties(
1703 kinds: &[u32],
1704 mags: &[f32],
1705 masks: &[u32],
1706 ids: &[u32],
1707 spectral: Option<(&[f32], &[u32])>,
1708 ) -> Option<ResidentKinds> {
1709 let g = gpu()?;
1710 assert!(
1711 mags.len() == kinds.len() && masks.len() == kinds.len() && ids.len() == kinds.len(),
1712 "one magnitude, mask and class id per significant token"
1713 );
1714 if let Some((e, s)) = spectral {
1715 assert!(e.len() == kinds.len() && s.len() == kinds.len(), "one spectral reading per significant token");
1716 }
1717 if kinds.is_empty() {
1718 return Some(ResidentKinds {
1719 d_kind: None,
1720 d_mag: None,
1721 d_bytes: None,
1722 d_class: None,
1723 d_ent: None,
1724 d_spec: None,
1725 magnitudes: true,
1726 spectral: spectral.is_some(),
1727 });
1728 }
1729 let stream = g.ctx.default_stream();
1730 let held = (|| {
1731 let d_kind = stream.clone_htod(kinds)?;
1732 let d_mag = stream.clone_htod(mags)?;
1733 let d_bytes = stream.clone_htod(masks)?;
1734 let d_class = stream.clone_htod(ids)?;
1735 let (d_ent, d_spec) = match spectral {
1736 Some((e, s)) => (Some(stream.clone_htod(e)?), Some(stream.clone_htod(s)?)),
1737 None => (None, None),
1738 };
1739 Ok::<_, cudarc::driver::DriverError>((d_kind, d_mag, d_bytes, d_class, d_ent, d_spec))
1740 })();
1741 match held {
1742 Ok((d_kind, d_mag, d_bytes, d_class, d_ent, d_spec)) => Some(ResidentKinds {
1743 d_kind: Some(d_kind),
1744 d_mag: Some(d_mag),
1745 d_bytes: Some(d_bytes),
1746 d_class: Some(d_class),
1747 d_ent,
1748 d_spec,
1749 magnitudes: true,
1750 spectral: spectral.is_some(),
1751 }),
1752 Err(e) => {
1753 eprintln!("trex gpu: holding {} tokens and their properties on the device failed: {e:?}", kinds.len());
1754 None
1755 }
1756 }
1757 }
1758
1759 fn scan_props_kernel() -> Option<&'static CudaFunction> {
1762 static F: OnceLock<Option<CudaFunction>> = OnceLock::new();
1763 F.get_or_init(|| {
1764 let g = gpu()?;
1765 let module = match g.ctx.load_module(Ptx::from_src(PTX)) {
1766 Ok(m) => m,
1767 Err(e) => {
1768 eprintln!("trex gpu: loading the scan module for trex_scan_props failed: {e:?}");
1769 return None;
1770 }
1771 };
1772 match module.load_function("trex_scan_props") {
1773 Ok(f) => Some(f),
1774 Err(e) => {
1775 eprintln!("trex gpu: loading trex_scan_props failed: {e:?}");
1776 None
1777 }
1778 }
1779 })
1780 .as_ref()
1781 }
1782
1783 #[allow(clippy::too_many_arguments)]
1793 fn run_scan(
1794 g: &Gpu,
1795 d_kind: &cudarc::driver::CudaSlice<u32>,
1796 d_mag: Option<&cudarc::driver::CudaSlice<f32>>,
1797 d_class: Option<&cudarc::driver::CudaSlice<u32>>,
1798 d_bytes: Option<&cudarc::driver::CudaSlice<u32>>,
1799 d_spectral: Option<(&cudarc::driver::CudaSlice<f32>, &cudarc::driver::CudaSlice<u32>)>,
1800 lit_class: Option<&[u32]>,
1801 nfa: &crate::nfa::GpuNfa,
1802 n: usize,
1803 ) -> Option<Vec<i32>> {
1804 assert_eq!(
1805 nfa.reads_spectral,
1806 d_spectral.is_some(),
1807 "a pattern that tests the spectral field scans with per-token readings, and only such a pattern does"
1808 );
1809 assert_eq!(
1810 nfa.reads_magnitude,
1811 d_mag.is_some(),
1812 "a pattern that tests a magnitude scans with per-token magnitudes, and only such a pattern does"
1813 );
1814 assert_eq!(
1815 nfa.class_group.is_some(),
1816 d_class.is_some(),
1817 "a pattern that compares a literal or a register scans with per-token classes, and only such a pattern does"
1818 );
1819 assert_eq!(
1820 nfa.reads_bytes,
1821 d_bytes.is_some(),
1822 "a pattern that tests a byte class scans with per-token masks, and only such a pattern does"
1823 );
1824 let literal_classes: Vec<u32> = match lit_class {
1825 Some(l) => l.to_vec(),
1826 None => vec![u32::MAX; nfa.nstates],
1827 };
1828 assert_eq!(literal_classes.len(), nfa.nstates, "one literal class per state");
1829 assert!(
1830 nfa.literals.iter().all(Option::is_none) || lit_class.is_some(),
1831 "a pattern with a literal scans with its literal classes resolved"
1832 );
1833 let n_i = n as i32;
1834 let nstates_i = nfa.nstates as i32;
1835 let start_closure = nfa.start_closure;
1836 let match_mask = nfa.match_mask;
1837 let reads_mag_i = i32::from(nfa.reads_magnitude);
1838 let block = 256u32;
1839 let cfg = LaunchConfig {
1840 grid_dim: (n.div_ceil(block as usize) as u32, 1, 1),
1841 block_dim: (block, 1, 1),
1842 shared_mem_bytes: 0,
1843 };
1844 let reads_spec_i = i32::from(nfa.reads_spectral);
1845 let props_func = if d_mag.is_some() || d_class.is_some() || d_bytes.is_some() || d_spectral.is_some() {
1846 Some(scan_props_kernel()?)
1847 } else {
1848 None
1849 };
1850 let stream = g.ctx.default_stream();
1851 let ran = (|| {
1852 let d_atom = stream.clone_htod(&nfa.atom_kind)?;
1853 let d_next = stream.clone_htod(&nfa.next_closure)?;
1854 let d_isatom = stream.clone_htod(&nfa.is_atom)?;
1855 let d_result = stream.alloc_zeros::<i32>(n)?;
1856 match props_func {
1857 Some(func) => {
1858 let d_lo = stream.clone_htod(&nfa.mag_lo)?;
1859 let d_hi = stream.clone_htod(&nfa.mag_hi)?;
1860 let d_back = stream.clone_htod(&nfa.ref_back)?;
1861 let d_need = stream.clone_htod(&nfa.need_mask)?;
1862 let d_lit = stream.clone_htod(&literal_classes)?;
1863 let d_ent_lo = stream.clone_htod(&nfa.ent_lo)?;
1864 let d_ent_hi = stream.clone_htod(&nfa.ent_hi)?;
1865 let d_period_need = stream.clone_htod(&nfa.period_need)?;
1866 let d_period_val = stream.clone_htod(&nfa.period_val)?;
1867 let d_texture_need = stream.clone_htod(&nfa.texture_need)?;
1868 let d_onset_need = stream.clone_htod(&nfa.onset_need)?;
1869 let spare_mag = stream.clone_htod(&[0.0f32])?;
1870 let spare_class = stream.clone_htod(&[0u32])?;
1871 let spare_bytes = stream.clone_htod(&[0u32])?;
1872 let spare_ent = stream.clone_htod(&[0.0f32])?;
1873 let spare_spec = stream.clone_htod(&[0u32])?;
1874 let mut builder = stream.launch_builder(func);
1875 builder.arg(d_kind);
1876 builder.arg(&n_i);
1877 if let Some(d) = d_mag {
1878 builder.arg(d);
1879 } else {
1880 builder.arg(&spare_mag);
1881 }
1882 builder.arg(&reads_mag_i);
1883 if let Some(d) = d_class {
1884 builder.arg(d);
1885 } else {
1886 builder.arg(&spare_class);
1887 }
1888 if let Some(d) = d_bytes {
1889 builder.arg(d);
1890 } else {
1891 builder.arg(&spare_bytes);
1892 }
1893 match d_spectral {
1894 Some((e, s)) => {
1895 builder.arg(e);
1896 builder.arg(s);
1897 }
1898 None => {
1899 builder.arg(&spare_ent);
1900 builder.arg(&spare_spec);
1901 }
1902 }
1903 builder.arg(&reads_spec_i);
1904 builder.arg(&d_atom);
1905 builder.arg(&d_lo);
1906 builder.arg(&d_hi);
1907 builder.arg(&d_back);
1908 builder.arg(&d_need);
1909 builder.arg(&d_lit);
1910 builder.arg(&d_ent_lo);
1911 builder.arg(&d_ent_hi);
1912 builder.arg(&d_period_need);
1913 builder.arg(&d_period_val);
1914 builder.arg(&d_texture_need);
1915 builder.arg(&d_onset_need);
1916 builder.arg(&d_next);
1917 builder.arg(&d_isatom);
1918 builder.arg(&start_closure);
1919 builder.arg(&match_mask);
1920 builder.arg(&nstates_i);
1921 builder.arg(&d_result);
1922 unsafe { builder.launch(cfg)? };
1928 }
1929 _ => {
1930 let mut builder = stream.launch_builder(&g.func);
1931 builder.arg(d_kind);
1932 builder.arg(&n_i);
1933 builder.arg(&d_atom);
1934 builder.arg(&d_next);
1935 builder.arg(&d_isatom);
1936 builder.arg(&start_closure);
1937 builder.arg(&match_mask);
1938 builder.arg(&nstates_i);
1939 builder.arg(&d_result);
1940 unsafe { builder.launch(cfg)? };
1945 }
1946 }
1947 let result: Vec<i32> = stream.clone_dtoh(&d_result)?;
1948 stream.synchronize()?;
1949 Ok::<_, cudarc::driver::DriverError>(result)
1950 })();
1951 match ran {
1952 Ok(result) => Some(result),
1953 Err(e) => {
1954 eprintln!("trex gpu: a scan over {n} device-held tokens failed: {e:?}");
1955 None
1956 }
1957 }
1958 }
1959
1960 fn scan_props(g: &Gpu, nfa: &crate::nfa::GpuNfa, input: &[u8]) -> Option<Vec<Span>> {
1964 let toks = crate::parallel_lex::lex_parallel(input);
1965 let (kinds, spans) = significant(&toks);
1966 let n = spans.len();
1967 if n == 0 {
1968 return Some(Vec::new());
1969 }
1970 let mags = nfa.reads_magnitude.then(|| super::significant_magnitudes(&toks, input));
1971 let masks = nfa.reads_bytes.then(|| super::significant_byte_masks(&toks, input));
1972 let spectral = nfa
1973 .reads_spectral
1974 .then(|| super::significant_spectral(&toks, &crate::spectral::analyze(input)));
1975 let (classes, lit_class) = match nfa.class_group {
1976 Some(group) => {
1977 let texts: Vec<&[u8]> = nfa.literals.iter().flatten().map(Vec::as_slice).collect();
1978 let (token_ids, literal_ids) = super::significant_classes(&toks, input, group, &texts);
1979 let mut per_state = vec![u32::MAX; nfa.nstates];
1980 let mut ids = literal_ids.into_iter();
1981 for (slot, lit) in per_state.iter_mut().zip(&nfa.literals) {
1982 if lit.is_some() {
1983 *slot = ids.next().expect("one class id per literal");
1984 }
1985 }
1986 (Some(token_ids), Some(per_state))
1987 }
1988 None => (None, None),
1989 };
1990 let stream = g.ctx.default_stream();
1991 let held = (|| {
1992 let d_kind = stream.clone_htod(&kinds)?;
1993 let d_mag = match &mags {
1994 Some(m) => Some(stream.clone_htod(m)?),
1995 None => None,
1996 };
1997 let d_class = match &classes {
1998 Some(c) => Some(stream.clone_htod(c)?),
1999 None => None,
2000 };
2001 let d_bytes = match &masks {
2002 Some(b) => Some(stream.clone_htod(b)?),
2003 None => None,
2004 };
2005 let d_spectral = match &spectral {
2006 Some((e, s)) => Some((stream.clone_htod(e)?, stream.clone_htod(s)?)),
2007 None => None,
2008 };
2009 Ok::<_, cudarc::driver::DriverError>((d_kind, d_mag, d_class, d_bytes, d_spectral))
2010 })();
2011 let (d_kind, d_mag, d_class, d_bytes, d_spectral) = match held {
2012 Ok(v) => v,
2013 Err(e) => {
2014 eprintln!("trex gpu: uploading {n} tokens and their properties failed: {e:?}");
2015 return None;
2016 }
2017 };
2018 let result = run_scan(
2019 g,
2020 &d_kind,
2021 d_mag.as_ref(),
2022 d_class.as_ref(),
2023 d_bytes.as_ref(),
2024 d_spectral.as_ref().map(|(e, s)| (e, s)),
2025 lit_class.as_deref(),
2026 nfa,
2027 n,
2028 )?;
2029 Some(select(&spans, &result))
2030 }
2031
2032 pub(super) fn scan_resident(
2039 kinds: &ResidentKinds,
2040 spans: &[(u32, u32)],
2041 text: Option<&super::HeldText>,
2042 pattern: &Pattern,
2043 ) -> Option<Vec<Span>> {
2044 let nfa = crate::nfa::compile_for_gpu(pattern)?;
2045 let g = gpu()?;
2046 if (nfa.reads_magnitude && !kinds.magnitudes) || (nfa.reads_spectral && !kinds.spectral) {
2047 return None;
2048 }
2049 let needs_text = nfa.binds_registers || nfa.class_group.is_some() || nfa.reads_bytes;
2050 if needs_text && text.is_none() {
2051 return None;
2052 }
2053 if let (Some(group), Some(t)) = (nfa.class_group, text)
2054 && group != t.classes.group()
2055 {
2056 return None;
2057 }
2058 let Some(d_kind) = &kinds.d_kind else {
2059 return Some(Vec::new());
2060 };
2061 let lit_class: Option<Vec<u32>> = match (nfa.class_group, text) {
2064 (Some(_), Some(t)) => Some(
2065 nfa.literals
2066 .iter()
2067 .map(|lit| match lit {
2068 Some(bytes) => t.classes.find(bytes).unwrap_or(u32::MAX - 1),
2069 None => u32::MAX,
2070 })
2071 .collect(),
2072 ),
2073 _ => None,
2074 };
2075 let d_mag = if nfa.reads_magnitude { kinds.d_mag.as_ref() } else { None };
2076 let d_class = if nfa.class_group.is_some() { kinds.d_class.as_ref() } else { None };
2077 let d_bytes = if nfa.reads_bytes { kinds.d_bytes.as_ref() } else { None };
2078 let d_spectral = if nfa.reads_spectral { kinds.d_ent.as_ref().zip(kinds.d_spec.as_ref()) } else { None };
2079 let result =
2080 run_scan(g, d_kind, d_mag, d_class, d_bytes, d_spectral, lit_class.as_deref(), &nfa, spans.len())?;
2081 Some(select(spans, &result))
2082 }
2083
2084 pub(super) fn scan_resident_split(
2093 kinds: &ResidentKinds,
2094 host_kinds: &[u32],
2095 spans: &[(u32, u32)],
2096 pattern: &Pattern,
2097 host_per_mille: u32,
2098 ) -> Option<Vec<Span>> {
2099 scan_resident_split_timed(kinds, host_kinds, spans, pattern, host_per_mille)
2100 .map(|(out, _)| out)
2101 }
2102
2103 pub(super) fn scan_resident_split_timed(
2107 kinds: &ResidentKinds,
2108 host_kinds: &[u32],
2109 spans: &[(u32, u32)],
2110 pattern: &Pattern,
2111 host_per_mille: u32,
2112 ) -> Option<(Vec<Span>, super::SplitPhases)> {
2113 use std::time::Instant;
2114
2115 let nfa = crate::nfa::compile_for_gpu(pattern)?;
2116 if nfa.reads_magnitude
2117 || nfa.binds_registers
2118 || nfa.class_group.is_some()
2119 || nfa.reads_bytes
2120 || nfa.reads_spectral
2121 {
2122 return None;
2123 }
2124 let g = gpu()?;
2125 let n = host_kinds.len();
2126 let Some(d_kind) = &kinds.d_kind else {
2127 return Some((Vec::new(), super::SplitPhases::default()));
2128 };
2129 let wall = Instant::now();
2130 let mut ph = super::SplitPhases::default();
2131 let mut ends = vec![0i32; n];
2136 if n < 2 {
2137 let t = Instant::now();
2138 host_anchor_ends(&nfa, host_kinds, 0, &mut ends);
2139 ph.ends_us = t.elapsed().as_secs_f64() * 1e6;
2140 let t = Instant::now();
2141 let out = select(spans, &ends);
2142 ph.select_us = t.elapsed().as_secs_f64() * 1e6;
2143 ph.wall_us = wall.elapsed().as_secs_f64() * 1e6;
2144 return Some((out, ph));
2145 }
2146 let mid = (n * host_per_mille as usize / 1000).clamp(1, n - 1);
2147 let ends_at = Instant::now();
2148 let (host_part, device_part) = ends.split_at_mut(mid);
2149 std::thread::scope(|s| {
2150 s.spawn(|| {
2151 match resident_suffix_ends(g, &nfa, d_kind, mid, device_part) {
2152 Ok(()) => {
2153 let base = i32::try_from(mid).expect("a token index within the kernel's i32 width");
2154 for slot in device_part.iter_mut() {
2155 if *slot >= 0 {
2156 *slot += base;
2157 }
2158 }
2159 }
2160 Err(e) => {
2161 eprintln!(
2162 "trex gpu: the device's share of a resident split failed and the host scanned it: {e:?}"
2163 );
2164 host_anchor_ends(&nfa, host_kinds, mid, device_part);
2165 }
2166 }
2167 });
2168 host_anchor_ends(&nfa, host_kinds, 0, host_part);
2169 });
2170 ph.ends_us = ends_at.elapsed().as_secs_f64() * 1e6;
2171 let t = Instant::now();
2172 let out = select(spans, &ends);
2173 ph.select_us = t.elapsed().as_secs_f64() * 1e6;
2174 ph.wall_us = wall.elapsed().as_secs_f64() * 1e6;
2175 Some((out, ph))
2176 }
2177
2178 fn resident_suffix_ends(
2183 g: &Gpu,
2184 nfa: &crate::nfa::GpuNfa,
2185 d_kind: &cudarc::driver::CudaSlice<u32>,
2186 mid: usize,
2187 out: &mut [i32],
2188 ) -> Result<(), cudarc::driver::DriverError> {
2189 let stream = g.ctx.default_stream();
2190 let d_atom = stream.clone_htod(&nfa.atom_kind)?;
2191 let d_next = stream.clone_htod(&nfa.next_closure)?;
2192 let d_isatom = stream.clone_htod(&nfa.is_atom)?;
2193 let view = d_kind.slice(mid..mid + out.len());
2194 device_view_ends(g, &stream, nfa, &view, out, &d_atom, &d_next, &d_isatom)
2195 }
2196
2197 fn significant(toks: &[Token]) -> (Vec<u32>, Vec<(u32, u32)>) {
2218 super::significant_stream(toks)
2219 }
2220
2221 fn select(spans: &[(u32, u32)], result: &[i32]) -> Vec<Span> {
2224 if takes_vector(result) {
2227 select_with(spans, result, next_match_from)
2228 } else {
2229 select_with(spans, result, next_match_from_scalar)
2230 }
2231 }
2232
2233 fn select_with(
2235 spans: &[(u32, u32)],
2236 result: &[i32],
2237 scan: impl Fn(&[i32], usize, usize) -> Option<usize>,
2238 ) -> Vec<Span> {
2239 let n = spans.len();
2240 let mut out = Vec::with_capacity(n / 2 + 1);
2246 let mut a = 0usize;
2247 while let Some(m) = scan(result, a, n) {
2248 let e = result[m] as usize; out.push(Span { start: spans[m].0, end: spans[e - 1].1 });
2250 a = e;
2251 }
2252 out
2253 }
2254
2255 pub(super) fn scan_placed(
2262 pattern: &Pattern,
2263 input: &[u8],
2264 ) -> Option<(Vec<Span>, super::BackendUsed)> {
2265 use flynnel::sched::call_site::Placement;
2266 let nfa = crate::nfa::compile_for_gpu(pattern)?;
2267 if nfa.reads_magnitude
2268 || nfa.binds_registers
2269 || nfa.class_group.is_some()
2270 || nfa.reads_bytes
2271 || nfa.reads_spectral
2272 {
2273 return None;
2274 }
2275 let g = gpu()?;
2276 let site = flynnel::sched::call_site::caller_site().get();
2277 let size = input.len().min(u32::MAX as usize) as u32;
2278 let cpu = || {
2279 let t = std::time::Instant::now();
2280 (crate::engine::scan(pattern, input), elapsed_ns(t))
2281 };
2282 let split = || {
2283 let t = std::time::Instant::now();
2284 (scan_split(g, &nfa, input), elapsed_ns(t))
2285 };
2286 match site.choose_placement(size) {
2287 Placement::Cpu => {
2288 let (m, ns) = cpu();
2289 site.record_placement(size, Some(ns), None);
2290 Some((m, super::BackendUsed::Cpu))
2291 }
2292 Placement::Backend => {
2293 let (m, ns) = split();
2294 site.record_placement(size, None, Some(ns));
2295 crate::trace::rung("scan", "split across the cores and the device", input.len());
2296 Some((m, super::BackendUsed::Split))
2297 }
2298 Placement::Race => {
2299 let cpu_first =
2303 RACES.fetch_add(1, std::sync::atomic::Ordering::Relaxed).is_multiple_of(2);
2304 let ((m, cpu_ns), split_ns) = if cpu_first {
2305 let c = cpu();
2306 (c, split().1)
2307 } else {
2308 let s = split().1;
2309 (cpu(), s)
2310 };
2311 site.record_placement(size, Some(cpu_ns), Some(split_ns));
2312 Some((m, super::BackendUsed::Cpu))
2313 }
2314 }
2315 }
2316
2317 static RACES: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
2319
2320 fn elapsed_ns(t: std::time::Instant) -> u64 {
2321 t.elapsed().as_nanos().min(u128::from(u64::MAX)) as u64
2322 }
2323
2324 fn scan_split(g: &'static Gpu, nfa: &crate::nfa::GpuNfa, input: &[u8]) -> Vec<Span> {
2333 crate::parallel_lex::lex_significant_parallel_held(input, |kinds, spans| {
2334 let n = kinds.len();
2335 let mut ends = vec![-1i32; n];
2336 if n < 2 {
2337 host_anchor_ends(nfa, kinds, 0, &mut ends);
2338 return select(spans, &ends);
2339 }
2340 let site = flynnel::sched::call_site::caller_site().get();
2341 let key = n.min(u32::MAX as usize) as u32;
2342 let share = site.split_cpu_share_per_mille_for(key) as usize;
2343 let mid = (n * share / 1000).clamp(1, n - 1);
2344 let (host_part, device_part) = ends.split_at_mut(mid);
2345 let (host_ns, device_ns) = std::thread::scope(|s| {
2346 let device = s.spawn(|| {
2347 let t = std::time::Instant::now();
2348 if let Err(e) = device_anchor_ends(g, nfa, &kinds[mid..], mid, device_part) {
2349 eprintln!(
2350 "trex gpu: the device's share of a split scan failed and the host scanned it: {e:?}"
2351 );
2352 host_anchor_ends(nfa, kinds, mid, device_part);
2353 }
2354 elapsed_ns(t)
2355 });
2356 let t = std::time::Instant::now();
2357 host_anchor_ends(nfa, kinds, 0, host_part);
2358 let host_ns = elapsed_ns(t);
2359 (host_ns, device.join().expect("the device share's thread returns"))
2360 });
2361 site.record_split_for(key, mid, host_ns, n - mid, device_ns);
2362 select(spans, &ends)
2363 })
2364 }
2365
2366 fn device_anchor_ends(
2370 g: &Gpu,
2371 nfa: &crate::nfa::GpuNfa,
2372 suffix: &[u32],
2373 mid: usize,
2374 out: &mut [i32],
2375 ) -> Result<(), cudarc::driver::DriverError> {
2376 let n = suffix.len();
2377 let stream = g.ctx.default_stream();
2378 let d_kind = stream.clone_htod(suffix)?;
2379 let d_atom = stream.clone_htod(&nfa.atom_kind)?;
2380 let d_next = stream.clone_htod(&nfa.next_closure)?;
2381 let d_isatom = stream.clone_htod(&nfa.is_atom)?;
2382 let d_result = stream.alloc_zeros::<i32>(n)?;
2383 let n_i = n as i32;
2384 let nstates_i = nfa.nstates as i32;
2385 let start_closure = nfa.start_closure;
2386 let match_mask = nfa.match_mask;
2387 let block = 256u32;
2388 let cfg = LaunchConfig {
2389 grid_dim: (n.div_ceil(block as usize) as u32, 1, 1),
2390 block_dim: (block, 1, 1),
2391 shared_mem_bytes: 0,
2392 };
2393 let mut builder = stream.launch_builder(&g.func);
2394 builder.arg(&d_kind);
2395 builder.arg(&n_i);
2396 builder.arg(&d_atom);
2397 builder.arg(&d_next);
2398 builder.arg(&d_isatom);
2399 builder.arg(&start_closure);
2400 builder.arg(&match_mask);
2401 builder.arg(&nstates_i);
2402 builder.arg(&d_result);
2403 unsafe { builder.launch(cfg)? };
2407 let ends: Vec<i32> = stream.clone_dtoh(&d_result)?;
2408 stream.synchronize()?;
2409 let base = i32::try_from(mid).expect("a token index within the kernel's i32 width");
2410 for (slot, e) in out.iter_mut().zip(ends) {
2411 *slot = if e < 0 { -1 } else { e + base };
2412 }
2413 Ok(())
2414 }
2415
2416 pub(super) fn host_anchor_ends(nfa: &crate::nfa::GpuNfa, kinds: &[u32], first: usize, out: &mut [i32]) {
2421 use flynnel::JobPlan;
2422 use flynnel::sched::par_iter::for_each_chunk_indexed_min_leaf;
2423 let cores = std::thread::available_parallelism().map_or(1, std::num::NonZero::get);
2424 let min_leaf = out.len().div_ceil(cores * 4).max(64);
2425 let plan = JobPlan::new(0, out.len().min(u32::MAX as usize) as u32)
2426 .with_leaf_shape(flynnel::LeafShape::PortCompute);
2427 for_each_chunk_indexed_min_leaf(&plan, out, min_leaf, |start, slots| {
2428 for (i, slot) in slots.iter_mut().enumerate() {
2429 *slot = anchor_end(nfa, kinds, first + start + i);
2430 }
2431 });
2432 }
2433
2434 #[inline]
2437 fn anchor_end(nfa: &crate::nfa::GpuNfa, kinds: &[u32], a: usize) -> i32 {
2438 let n = kinds.len();
2439 let mut active = nfa.start_closure;
2440 let mut best = -1i32;
2441 for (k, &tk) in kinds.iter().enumerate().skip(a) {
2442 if active & nfa.match_mask != 0 {
2443 best = k as i32;
2444 }
2445 let mut next = 0u64;
2446 let mut m = active;
2447 while m != 0 {
2448 let pc = m.trailing_zeros() as usize;
2449 m &= m - 1;
2450 if nfa.is_atom[pc] != 0 {
2451 let ak = nfa.atom_kind[pc];
2452 if ak == u32::MAX || ak == tk {
2453 next |= nfa.next_closure[pc];
2454 }
2455 }
2456 }
2457 if next == 0 {
2458 return best;
2459 }
2460 active = next;
2461 }
2462 if active & nfa.match_mask != 0 { n as i32 } else { best }
2463 }
2464
2465 #[cfg(test)]
2466 mod tests {
2467 use super::*;
2468
2469 fn split_corpus() -> Vec<u8> {
2472 let mut s = String::new();
2473 for i in 0..3000u32 {
2474 s.push_str(&format!("tag {} {} word {}\n", i % 7, i % 13, i));
2475 if i % 5 == 0 {
2476 s.push_str("1 2 3 4 5 6 7 8\n");
2477 }
2478 }
2479 s.into_bytes()
2480 }
2481
2482 #[test]
2485 fn host_anchor_ends_select_the_cpu_engines_spans() {
2486 let input = split_corpus();
2487 let (kinds, spans) = crate::parallel_lex::lex_significant_parallel(&input);
2488 for src in ["\\W \\N", "\\N+", "\\N{2,4}", "(\\N \\W)+", ". .", "\\W \\W \\W"] {
2489 let pat = crate::parser::parse(src).expect("pattern parses");
2490 let nfa = crate::nfa::compile_for_gpu(&pat).expect("a kind-only pattern compiles for the kernel");
2491 let mut ends = vec![-1i32; kinds.len()];
2492 host_anchor_ends(&nfa, &kinds, 0, &mut ends);
2493 assert_eq!(select(&spans, &ends), crate::engine::scan(&pat, &input), "{src}");
2494 }
2495 }
2496
2497 #[test]
2500 fn a_split_at_any_share_selects_the_cpu_engines_spans() {
2501 let Some(g) = gpu() else {
2502 eprintln!("no CUDA device present; the split agreement test did not run");
2503 return;
2504 };
2505 let input = split_corpus();
2506 let (kinds, spans) = crate::parallel_lex::lex_significant_parallel(&input);
2507 let n = kinds.len();
2508 for src in ["\\W \\N", "\\N+", "(\\N \\W)+", "\\W \\W \\W"] {
2509 let pat = crate::parser::parse(src).expect("pattern parses");
2510 let nfa = crate::nfa::compile_for_gpu(&pat).expect("a kind-only pattern compiles for the kernel");
2511 let want = crate::engine::scan(&pat, &input);
2512 for mid in [1, n / 3, n / 2, n - 1] {
2513 let mut ends = vec![-1i32; n];
2514 let (host_part, device_part) = ends.split_at_mut(mid);
2515 host_anchor_ends(&nfa, &kinds, 0, host_part);
2516 device_anchor_ends(g, &nfa, &kinds[mid..], mid, device_part)
2517 .expect("the device ran its share");
2518 assert_eq!(select(&spans, &ends), want, "{src} split at {mid} of {n}");
2519 }
2520 }
2521 }
2522
2523 #[test]
2527 fn a_pipelined_scan_selects_the_cpu_engines_spans() {
2528 if gpu().is_none() {
2529 eprintln!("no CUDA device present; the pipelined scan test did not run");
2530 return;
2531 }
2532 let input = split_corpus().repeat(4);
2533 for src in ["\\W \\N", "\\N{2,4}", "\\W \\N \\W", ". .", "\\N"] {
2536 let pat = crate::parser::parse(src).expect("pattern parses");
2537 let want = crate::engine::scan(&pat, &input);
2538 assert!(!want.is_empty(), "{src} matches nothing, so nothing would be compared");
2539 for partitions in [1, 2, 3, 7, 64] {
2540 let (got, ph) = scan_pipelined(&pat, &input, partitions).expect("the device ran the pipeline");
2541 assert_eq!(got, want, "{src} over {partitions} ranges");
2542 assert!(ph.partitions >= 1, "{src} over {partitions} ranges lexed no range");
2543 }
2544 }
2545 let unbounded = crate::parser::parse("\\N+").expect("pattern parses");
2546 assert!(scan_pipelined(&unbounded, &input, 4).is_none(), "an unbounded pattern has no carry length");
2547 }
2548
2549 #[test]
2554 fn a_parts_scan_selects_the_cpu_engines_spans() {
2555 if gpu().is_none() {
2556 eprintln!("no CUDA device present; the parts scan test did not run");
2557 return;
2558 }
2559 let many = split_corpus().repeat(4);
2560 let one = b"tag 1 2 word 3\ntag 4 5 word 6\n".to_vec();
2561 for src in ["\\W \\N", "\\N{2,4}", "\\W \\N \\W", ". .", "\\N", "\\N+"] {
2564 let pat = crate::parser::parse(src).expect("pattern parses");
2565 for input in [&many, &one] {
2566 let want = crate::engine::scan(&pat, input);
2567 assert!(!want.is_empty(), "{src} matches nothing in {} bytes", input.len());
2568 let got = scan_parts(&pat, input).expect("the device ran the parts scan");
2569 assert_eq!(got, want, "{src} over {} bytes", input.len());
2570 }
2571 assert_eq!(scan_parts(&pat, b""), Some(Vec::new()), "{src} over an empty input");
2572 }
2573 }
2574
2575 #[test]
2579 fn a_resident_split_selects_the_cpu_engines_spans() {
2580 if gpu().is_none() {
2581 eprintln!("no CUDA device present; the resident split test did not run");
2582 return;
2583 }
2584 let input = split_corpus();
2585 let held = crate::gpu::GpuTokens::upload(&input).expect("the device holds the corpus");
2586 for src in ["\\W \\N", "\\N+", "\\N{2,4}", "\\W \\N \\W", ". .", "\\N"] {
2587 let pat = crate::parser::parse(src).expect("pattern parses");
2588 let want = crate::engine::scan(&pat, &input);
2589 assert!(!want.is_empty(), "{src} matches nothing, so nothing would be compared");
2590 for share in [1u32, 100, 500, 900, 999] {
2591 let got = held.scan_split(&pat, share).expect("the device ran the resident split");
2592 assert_eq!(got, want, "{src} at a host share of {share} per mille");
2593 }
2594 }
2595 let empty = crate::gpu::GpuTokens::upload(b"").expect("the device holds an empty corpus");
2596 let pat = crate::parser::parse("\\N").expect("pattern parses");
2597 assert_eq!(empty.scan_split(&pat, 500), Some(Vec::new()), "an empty corpus has no match");
2598 }
2599 }
2600}
2601
2602#[cfg(test)]
2603mod tests {
2604 use super::*;
2605
2606 use crate::parser::parse;
2607
2608 #[test]
2609 fn eligibility_matches_the_documented_subset() {
2610 for src in [
2613 "\\N",
2614 "\\W \\N",
2615 "\\N+",
2616 "\\W*",
2617 ".",
2618 ". .",
2619 "\\N{2,4}",
2620 "(\\N \\W)+",
2621 "\\M{>6}",
2622 "\\N{mag<3}",
2623 "\\W \\N{mag>=2}",
2624 "\\W:t",
2625 "\\W:x =x",
2626 "\\W:x \\N =case x",
2627 "\"lit\"",
2628 "\\d",
2629 "\\W:x \"lit\"",
2630 ] {
2631 assert!(gpu_eligible(&parse(src).unwrap()), "{src} should be eligible");
2632 }
2633 for src in [
2637 "\\N{>+1}",
2638 "\\W+:x =x",
2639 "\\W:x \\N* =x",
2640 "\\W:x =case x \"lit\"",
2641 "\\N | \\W",
2642 "\\W\\B(.*)",
2643 "@2 \\W",
2644 ". ~\"END\"",
2645 "`[a-z]+`",
2646 ] {
2647 assert!(!gpu_eligible(&parse(src).unwrap()), "{src} should be ineligible");
2648 }
2649 }
2650
2651 #[cfg(not(feature = "gpu"))]
2652 #[test]
2653 fn scan_gpu_is_none_without_the_feature() {
2654 assert!(scan_gpu(&parse("\\N \\W").unwrap(), b"12 kg").is_none());
2657 }
2658
2659 #[cfg(feature = "gpu")]
2660 #[test]
2661 fn the_device_agrees_with_the_cpu_over_many_tokens() {
2662 let pat = parse("\\W \\N").expect("pattern parses");
2668 let mut input = String::new();
2669 for i in 0..20_000 {
2670 input.push_str(match i % 4 {
2674 0 => "tag ",
2675 1 => "tag\t\t",
2676 2 => "tag \n ",
2677 _ => "tag ",
2678 });
2679 input.push_str(&(i % 997).to_string());
2680 input.push(' ');
2681 }
2682 let bytes = input.as_bytes();
2683 let Some(device) = scan_gpu(&pat, bytes) else {
2684 eprintln!("no CUDA device present; device-agreement test did not run");
2687 return;
2688 };
2689 let cpu = crate::engine::scan(&pat, bytes);
2690 assert_eq!(device.len(), cpu.len(), "device and cpu must find the same number of matches");
2691 assert_eq!(device, cpu, "device and cpu must agree on every span");
2692 assert!(!cpu.is_empty(), "the corpus must actually match, or this asserts nothing");
2693 }
2694
2695 #[cfg(feature = "gpu")]
2696 #[test]
2697 fn the_device_agrees_with_the_cpu_on_magnitude_tests() {
2698 let mut input = String::new();
2702 for i in 0..20_000u64 {
2703 input.push_str(match i % 3 {
2704 0 => "size ",
2705 1 => "tiny\t",
2706 _ => "enormousword \n ",
2707 });
2708 input.push_str(&(i * i * 37 % 10_000_019).to_string());
2709 input.push(' ');
2710 if i % 5 == 0 {
2711 input.push_str("0.004 ");
2712 }
2713 }
2714 let bytes = input.as_bytes();
2715 for src in [
2716 "\\N{mag>6}",
2717 "\\N{mag<3}",
2718 "\\M{>1}",
2719 "\\M{<=2} \\N{mag>=5}",
2720 "(\\W \\N{mag>4})+",
2721 ] {
2722 let pat = parse(src).expect("pattern parses");
2723 assert!(gpu_eligible(&pat), "{src} is in the device subset");
2724 let Some(device) = scan_gpu(&pat, bytes) else {
2725 eprintln!("no CUDA device present; magnitude agreement test did not run");
2726 return;
2727 };
2728 let cpu = crate::engine::scan(&pat, bytes);
2729 assert!(!cpu.is_empty(), "{src} must match the corpus, or this asserts nothing");
2730 assert_eq!(device, cpu, "{src}: device and cpu must agree on every span");
2731 let held = GpuTokens::upload_with_magnitudes(bytes).expect("the device ran the scan");
2732 assert_eq!(held.scan(&pat), Some(cpu), "{src}: held tokens and cpu must agree");
2733 let kinds_only = GpuTokens::upload(bytes).expect("the device ran the scan");
2734 assert_eq!(kinds_only.scan(&pat), None, "{src}: tokens held without magnitudes decline");
2735 }
2736 }
2737
2738 #[cfg(feature = "gpu")]
2739 #[test]
2740 fn the_device_agrees_with_the_cpu_on_binds() {
2741 let words = ["alpha", "Alpha", "beta", "BETA", "gamma"];
2745 let mut input = String::new();
2746 for i in 0..20_000usize {
2747 input.push_str(words[i % 5]);
2748 input.push(' ');
2749 input.push_str(words[(i * 7 + i / 3) % 5]);
2750 input.push_str(if i % 4 == 0 { "\t" } else { " " });
2751 input.push_str(&(i % 13).to_string());
2752 input.push(' ');
2753 if i % 3 == 0 {
2754 input.push_str("Alpha 7 alpha gamma 7 gamma delta x DELTA beta beta ");
2755 }
2756 }
2757 let bytes = input.as_bytes();
2758 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"] {
2759 let pat = parse(src).expect("pattern parses");
2760 assert!(gpu_eligible(&pat), "{src} is in the device subset");
2761 let Some(device) = scan_gpu(&pat, bytes) else {
2762 eprintln!("no CUDA device present; bind agreement test did not run");
2763 return;
2764 };
2765 let cpu = crate::engine::scan(&pat, bytes);
2766 assert!(!cpu.is_empty(), "{src} must match the corpus, or this asserts nothing");
2767 assert_eq!(device, cpu, "{src}: device and cpu must agree on every span");
2768 let identity = GpuTokens::upload_with_properties(bytes, crate::orbit::OrbitGroup::Identity)
2769 .expect("the device ran the scan");
2770 let case = GpuTokens::upload_with_properties(bytes, crate::orbit::OrbitGroup::Case)
2771 .expect("the device ran the scan");
2772 let held = identity.scan(&pat).or_else(|| case.scan(&pat));
2773 assert_eq!(held, Some(cpu), "{src}: tokens held with their properties agree with the cpu");
2774 }
2775 }
2776
2777 #[cfg(feature = "gpu")]
2778 #[test]
2779 fn the_device_agrees_with_the_cpu_on_literals_and_byte_classes() {
2780 let mut input = String::new();
2784 for i in 0..20_000usize {
2785 input.push_str(["the", "The", "THE", "cat", "x1f"][i % 5]);
2786 input.push(' ');
2787 input.push_str(&(i % 97).to_string());
2788 input.push_str(if i % 3 == 0 { " , " } else { " ; " });
2789 input.push_str(["ABC", "abc", "a_b", "Zz9"][i % 4]);
2790 input.push('\n');
2791 }
2792 let bytes = input.as_bytes();
2793 for src in [
2794 "\"the\" \\N",
2795 "(?orbit:case \"the\") \\N",
2796 "\",\" \\w",
2797 "\\d \";\"",
2798 "\\u",
2799 "\\l \\W",
2800 "\\a+",
2801 "\\W:x \"the\"",
2802 ] {
2803 let pat = parse(src).expect("pattern parses");
2804 assert!(gpu_eligible(&pat), "{src} is in the device subset");
2805 let Some(device) = scan_gpu(&pat, bytes) else {
2806 eprintln!("no CUDA device present; literal and byte-class agreement test did not run");
2807 return;
2808 };
2809 let cpu = crate::engine::scan(&pat, bytes);
2810 assert!(!cpu.is_empty(), "{src} must match the corpus, or this asserts nothing");
2811 assert_eq!(device, cpu, "{src}: device and cpu must agree on every span");
2812 let identity = GpuTokens::upload_with_properties(bytes, crate::orbit::OrbitGroup::Identity)
2813 .expect("the device ran the scan");
2814 let case = GpuTokens::upload_with_properties(bytes, crate::orbit::OrbitGroup::Case)
2815 .expect("the device ran the scan");
2816 let held = identity.scan(&pat).or_else(|| case.scan(&pat));
2817 assert_eq!(held, Some(cpu), "{src}: tokens held with their properties agree with the cpu");
2818 }
2819 }
2820
2821 #[cfg(feature = "gpu")]
2822 #[test]
2823 fn the_device_agrees_with_the_cpu_on_hex_runs() {
2824 let digits = b"0123456789abcdef";
2829 let mut input = String::new();
2830 for i in 0..20_000usize {
2831 let len = [33usize, 39, 41, 52, 63, 65, 96, 32, 40, 64][i % 10];
2832 let run: String = (0..len).map(|j| char::from(digits[(i * 7 + j * 3) % 16])).collect();
2833 input.push_str(["key ", "sum ", "id\t"][i % 3]);
2834 input.push_str(&run);
2835 input.push_str(if i % 4 == 0 { " = " } else { " ; " });
2836 input.push_str(&"7".repeat(40));
2837 input.push_str(" aB3dEfGhIjKlMnOp\n");
2838 }
2839 let bytes = input.as_bytes();
2840 for src in ["\\{hex}", "\\W \\{hex}", "\\{hex} \";\"", "\\{hex} \\P \\N", "\\{hash}"] {
2841 let pat = parse(src).expect("pattern parses");
2842 assert!(gpu_eligible(&pat), "{src} is in the device subset");
2843 let Some(device) = scan_gpu(&pat, bytes) else {
2844 eprintln!("no CUDA device present; hex agreement test did not run");
2845 return;
2846 };
2847 let cpu = crate::engine::scan(&pat, bytes);
2848 assert!(!cpu.is_empty(), "{src} must match the corpus, or this asserts nothing");
2849 assert_eq!(device, cpu, "{src}: device and cpu must agree on every span");
2850 let held = GpuTokens::upload_with_properties(bytes, crate::orbit::OrbitGroup::Identity)
2851 .expect("the device ran the scan");
2852 assert_eq!(held.scan(&pat), Some(cpu), "{src}: tokens held with their properties agree with the cpu");
2853 }
2854 }
2855
2856 #[cfg(feature = "gpu")]
2857 #[test]
2858 fn the_device_agrees_with_the_cpu_on_spectral_tests() {
2859 let mut input = String::new();
2863 for i in 0..400usize {
2864 input.push_str("the quick brown fox jumps over the lazy dog and the cat sat on the mat ");
2865 if i % 3 == 0 {
2866 let mut x = 0x2545_f491_4f6c_dd1du64 ^ (i as u64);
2867 for _ in 0..200 {
2868 x = x.wrapping_mul(6_364_136_223_846_793_005).wrapping_add(1_442_695_040_888_963_407);
2869 input.push(b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/"[((x >> 58) % 64) as usize] as char);
2870 }
2871 input.push(' ');
2872 }
2873 if i % 5 == 0 {
2874 for r in 0..12 {
2875 input.push_str(&format!("{:03},{:03},{:03}\n", r, (r * 7) % 100, (r * 13) % 100));
2876 }
2877 }
2878 if i % 7 == 0 {
2879 input.push_str(&"x".repeat(150));
2880 input.push(' ');
2881 }
2882 }
2883 let bytes = input.as_bytes();
2884 for src in [
2885 "\\F{entropy>0.8}",
2886 "\\F{entropy<0.3} \\W",
2887 "\\F{period:any}",
2888 "\\F{texture:prose} \\W",
2889 "\\F{texture:data}",
2890 "\\F{onset}",
2891 "\\W \\F{entropy>0.8}",
2892 ] {
2893 let pat = parse(src).expect("pattern parses");
2894 assert!(gpu_eligible(&pat), "{src} is in the device subset");
2895 let Some(device) = scan_gpu(&pat, bytes) else {
2896 eprintln!("no CUDA device present; spectral agreement test did not run");
2897 return;
2898 };
2899 let cpu = crate::engine::scan(&pat, bytes);
2900 assert!(!cpu.is_empty(), "{src} must match the corpus, or this asserts nothing");
2901 assert_eq!(device, cpu, "{src}: device and cpu must agree on every span");
2902 let held = GpuTokens::upload_with_spectral(bytes, crate::orbit::OrbitGroup::Identity)
2903 .expect("the device ran the scan");
2904 assert_eq!(held.scan(&pat), Some(cpu), "{src}: tokens held with the spectral reading agree");
2905 let without = GpuTokens::upload_with_properties(bytes, crate::orbit::OrbitGroup::Identity)
2906 .expect("the device ran the scan");
2907 assert_eq!(without.scan(&pat), None, "{src}: tokens held without the reading decline");
2908 }
2909 assert!(!gpu_eligible(&parse("\\W:x \\F{onset}").unwrap()), "a bind beside a spectral test stays on the CPU");
2910 }
2911
2912 #[test]
2913 fn every_backend_returns_the_cpu_match_set() {
2914 let pat = parse("\\W \\N").unwrap();
2918 let input = b"tag 12 tag 34 tag 56";
2919 let baseline = crate::engine::scan(&pat, input);
2920 for b in [Backend::Auto, Backend::Gpu, Backend::Cpu] {
2921 let (m, _used) = scan_with_backend(&pat, input, b);
2922 assert_eq!(m, baseline, "{b:?} must match the CPU baseline");
2923 }
2924 }
2925
2926 #[test]
2927 fn auto_matches_the_cpu_engine_while_its_placement_learns() {
2928 let pat = parse("\\W \\N").unwrap();
2931 let mut input = String::new();
2932 for i in 0..40_000u32 {
2933 input.push_str(&format!("tag {}\n", i % 1000));
2934 }
2935 let want = crate::engine::scan(&pat, input.as_bytes());
2936 for _ in 0..40 {
2937 let (m, _used) = scan_with_backend(&pat, input.as_bytes(), Backend::Auto);
2938 assert_eq!(m, want, "Auto returns the CPU engine's spans whichever route it took");
2939 }
2940 }
2941
2942 #[test]
2943 fn auto_routes_only_kind_patterns_to_the_device() {
2944 for src in ["\\W \\N", "\\N+", ". ."] {
2945 assert!(gpu_auto_routes(&parse(src).unwrap()), "{src} reads only token kinds");
2946 }
2947 let input = "tag 12 the 34 ".repeat(2_000);
2948 for src in ["\\W:x =x", "\\N{mag>6}", "\"the\" \\W", "\\d", "\\W:t \\N"] {
2949 let pat = parse(src).unwrap();
2950 assert!(gpu_eligible(&pat), "{src} is in the device subset");
2951 assert!(!gpu_auto_routes(&pat), "{src} reads a token property");
2952 let (m, used) = scan_with_backend(&pat, input.as_bytes(), Backend::Auto);
2953 assert_eq!(used, BackendUsed::Cpu, "{src}: auto keeps a property-reading pattern on the cores");
2954 assert_eq!(m, crate::engine::scan(&pat, input.as_bytes()), "{src}: auto returns the CPU matches");
2955 }
2956 }
2957
2958 #[cfg(not(feature = "gpu"))]
2959 #[test]
2960 fn without_a_device_every_backend_runs_on_the_cpu() {
2961 assert!(!device_available());
2964 let pat = parse("\\W \\N").unwrap();
2965 let input = b"tag 12 tag 34";
2966 for b in [Backend::Auto, Backend::Gpu, Backend::Cpu] {
2967 let (_m, used) = scan_with_backend(&pat, input, b);
2968 assert_eq!(used, BackendUsed::Cpu, "no device: {b:?} must use the CPU");
2969 }
2970 }
2971}