"""Phase38 staged frozen normalization, AC/PFAC matching and exact row evaluation. One instance owns both matcher configurations and one reusable workspace. This research adapter changes no maintained backend and loads only an explicit AOT cubin. Event order is unspecified; the inherited evaluator's output is exact. """ from __future__ import annotations import hashlib from pathlib import Path import threading import time import numpy as np from . import phase32a_gpu, phase32b_gpu _SOURCE_SHA256_AT_IMPORT = hashlib.sha256(Path(__file__).read_bytes()).hexdigest() _BASE_SOURCE_SHA256_AT_IMPORT = { module.__name__: hashlib.sha256(Path(module.__file__).read_bytes()).hexdigest() for module in (phase32a_gpu, phase32b_gpu) } KERNEL_SOURCE_PATH = Path(__file__).with_name("phase38_kernel.cu") CUDA_SOURCE = KERNEL_SOURCE_PATH.read_text(encoding="utf-8") _KERNEL_FILE_SHA256_AT_IMPORT = hashlib.sha256(KERNEL_SOURCE_PATH.read_bytes()).hexdigest() MAX_CANDIDATES = phase32b_gpu.MAX_CANDIDATES MAX_EVENT_CAPACITY = 512 MAX_BATCH_BYTES = 32 * 1024 * 1024 MAX_BATCH_ROWS = 32768 MAX_SHARED_GPU_BYTES = 2 * 1024**3 STAT_NAMES = ( "search_jobs", "transition_attempts", "terminal_visits", "boundary_passed_events", "depth_0", "depth_1", "depth_2_3", "depth_4_7", "depth_8_15", "depth_16_31", "depth_32_63", "depth_64_127", "depth_128_255", "depth_256_511", "depth_512_1023", "depth_1024_2047", "depth_2048_4095", "depth_4096_8191", "depth_8192_16383", "depth_16384_32767", "depth_32768_65535", "depth_65536_plus", "capacity_stopped_jobs", "normalized_bytes", "max_depth_per_row_sum", ) def workspace_nbytes(byte_count, row_count, event_capacity=MAX_EVENT_CAPACITY): """Exact requested array bytes, excluding configs/pool rounding/CUDA handles.""" for name, value, limit in (("byte_count", byte_count, MAX_BATCH_BYTES), ("row_count", row_count, MAX_BATCH_ROWS), ("event_capacity", event_capacity, MAX_EVENT_CAPACITY)): minimum = 1 if name == "event_capacity" else 0 if isinstance(value, bool) or not isinstance(value, int) or not minimum <= value <= limit: raise ValueError(f"{name} outside Phase38 capacity") return 15 * byte_count + (2109 + 16 * event_capacity + 8 * len(STAT_NAMES)) * row_count + 8 # Runner compatibility: both names describe requested arrays, not pool totals. workspace_bytes = workspace_nbytes def _pack_pfac_tables(specs, arrays): """Build genuine trie edges with terminal-only outputs and existing key IDs. Reconstruct the three key streams in the same specified order as _pack_specs; consumer storage and all evaluator arrays remain the inherited ones. """ patterns, _pc, _consumers, exact, *_rest = specs streams = [[(phase32a_gpu._ascii(value), index) for index, value in enumerate(patterns)], [], []] seen, next_key = set(), len(patterns) for value, mode, _kind, _index in exact: stream = 1 if mode == "case_sensitive" else 2 encoded = phase32a_gpu._ascii(value) if stream == 2: encoded = encoded.lower() identity = (stream, encoded) if identity not in seen: seen.add(identity) streams[stream].append((encoded, next_key)) next_key += 1 if next_key != arrays["key_lengths"].size: raise RuntimeError("Phase38 key order differs from the inherited taxonomy") transitions, metadata, output_offsets, output_keys = [], [], [], [] for stream, patterns in enumerate(streams): children, outputs = [{}], [set()] for pattern, key in patterns: if len(pattern) != int(arrays["key_lengths"][key]): raise RuntimeError("Phase38 key length differs from the inherited taxonomy") state = 0 for byte in pattern: if byte not in children[state]: children[state][byte] = len(children) children.append({}) outputs.append(set()) state = children[state][byte] outputs[state].add(key) width = int(arrays["dfa_meta"][stream * 4 + 1]) if (len(transitions) + len(children) * width) * 4 > phase32a_gpu.MAX_DFA_BYTES: raise phase32a_gpu.UnsupportedGpuTaxonomy("Phase38 combined PFAC tries exceed 32 MiB") dense = np.zeros((len(children), width), dtype=np.uint32) class_map = arrays["classes"][stream * 256:(stream + 1) * 256] for state, edges in enumerate(children): for byte, child in edges.items(): dense[state, int(class_map[byte])] = child metadata.extend((len(transitions), width, len(output_offsets), len(patterns))) transitions.extend(dense.reshape(-1).tolist()) for values in outputs: output_offsets.append(len(output_keys)) output_keys.extend(sorted(values)) output_offsets.append(len(output_keys)) return { "pfac_transitions": np.asarray(transitions, dtype=np.uint32), "pfac_meta": np.asarray(metadata, dtype=np.int32), "pfac_output_offsets": np.asarray(output_offsets, dtype=np.int32), "pfac_output_keys": np.asarray(output_keys, dtype=np.int32), } class ExactStagedFrozenUnicodeCudaClassifier(phase32b_gpu.ExactFrozenUnicodeCudaClassifier): def __init__(self, *args, matcher="ac", event_capacity=MAX_EVENT_CAPACITY, batch_bytes=MAX_BATCH_BYTES, batch_rows=MAX_BATCH_ROWS, max_row_bytes=65536, kernel_event_diagnostics=False, **kwargs): workspace_nbytes(batch_bytes, batch_rows, event_capacity) if matcher not in ("ac", "pfac"): raise ValueError("matcher must be ac or pfac") if not isinstance(kernel_event_diagnostics, bool): raise TypeError("kernel_event_diagnostics must be boolean") super().__init__(*args, batch_bytes=batch_bytes, batch_rows=batch_rows, max_row_bytes=max_row_bytes, **kwargs) self.matcher, self.event_capacity = matcher, event_capacity self.kernel_event_diagnostics = kernel_event_diagnostics self._cuda_source_sha256 = hashlib.sha256(CUDA_SOURCE.encode()).hexdigest() self._kernels, self._events = {}, None self._call_lock = threading.Lock() self._failed_drain = False self._active_transfer_owners = None if not self.unsupported_reason: from range_curate._core.accelerated.compact import _build_native_specs try: self._arrays.update(_pack_pfac_tables( _build_native_specs(self.cpu_engine._screening_indicator), self._arrays)) except phase32a_gpu.UnsupportedGpuTaxonomy as exc: self.unsupported_reason = str(exc) self._config_requested_bytes = sum(max(value.nbytes, value.dtype.itemsize) for value in self._arrays.values()) def select_matcher(self, matcher): if matcher not in ("ac", "pfac"): raise ValueError("matcher must be ac or pfac") if not self._call_lock.acquire(blocking=False): raise RuntimeError("cannot change the matcher during an active call") try: if self._failed_drain: raise RuntimeError("Phase38 CUDA drain failed; end this worker process") self.matcher = matcher if self._kernels: self._kernel = self._kernels["match_" + matcher] finally: self._call_lock.release() def screen_arrow(self, *args, **kwargs): if not self._call_lock.acquire(blocking=False): raise RuntimeError("Phase38 shared workspace already has an active call") try: result = super().screen_arrow(*args, **kwargs) if self.diagnostics and "diagnostics" in self.last_profile and not self.unsupported_reason: self.last_profile["diagnostics"]["device_event_timings"] = True self.last_profile["diagnostics"]["device_event_scope"] = "three separately launched stages; diagnostic calls only" return result finally: self._call_lock.release() @property def backend_report(self): return {**super().backend_report, "backend": "cpu" if self.unsupported_reason else "experimental_staged_frozen_unicode_cuda", "matcher": self.matcher, "match_event_capacity": self.event_capacity, "match_event_schema": ["key", "original_codepoint_start", "original_codepoint_end", "stream"], "match_event_order": "unspecified; every bounded event reduced before output", "match_event_overflow": "saturated cap+1; exact whole-row code3 CPU fallback", "workspace_requested_max_bytes": workspace_nbytes(self.batch_bytes, self.batch_rows, self.event_capacity), "config_requested_bytes": self._config_requested_bytes, "kernel_event_diagnostics": self.kernel_event_diagnostics, "phase38_kernel_file_sha256": _KERNEL_FILE_SHA256_AT_IMPORT, "kernel_entrypoints": ["range38_prepare", "range38_match_ac", "range38_match_pfac", "range38_evaluate"], "shared_gpu_limit_bytes": MAX_SHARED_GPU_BYTES, "gpu_pool_accounting": "default pool total includes all live/cached controls; no pool trimming", "pfac_work_assignment": "one block per article; independent lane-strided byte starts in three streams"} def _guard_allocations(self, sizes): # Conservative before-allocation bound: do not credit cached blocks or # soon-to-be-replaced arrays. Root preallocates every control to max once. projected = int(self.cp.get_default_memory_pool().total_bytes()) projected += sum(((int(size) + 511) // 512) * 512 for size in sizes) if projected > MAX_SHARED_GPU_BYTES: raise MemoryError("Phase38 live/cached pool plus requested allocations exceeds 2 GiB") def _initialize(self): if self._failed_drain: raise RuntimeError("Phase38 CUDA drain failed; end this worker process") if self._kernels: return if not self.cubin_path.is_file(): raise FileNotFoundError("Build the Phase38 AOT cubin explicitly: " + str(self.cubin_path)) import cupy as cp self.cp = cp try: self._guard_allocations(max(value.nbytes, value.dtype.itemsize) for value in self._arrays.values()) self._module = cp.RawModule(path=str(self.cubin_path)) kernels = {name: self._module.get_function("range38_" + name) for name in ("prepare", "match_ac", "match_pfac", "evaluate")} # Keep each config owner on the instance even if a later copy fails. for name, value in self._arrays.items(): self._device_config[name] = cp.asarray(value if value.size else np.zeros(1, dtype=value.dtype)) self._events = tuple(cp.cuda.Event() for _ in range(4)) cp.cuda.Stream.null.synchronize() self._cubin_sha256 = hashlib.sha256(self.cubin_path.read_bytes()).hexdigest() self._kernel_attributes = { name: {field: int(value.attributes[field]) if value.attributes.get(field) is not None else None for field in ("num_regs", "local_size_bytes", "shared_size_bytes")} for name, value in kernels.items()} self._kernels, self._kernel = kernels, kernels["match_" + self.matcher] except BaseException as error: self._drain(error) raise def _reserve(self, byte_count, row_count): workspace_nbytes(byte_count, row_count, self.event_capacity) cp = self.cp specs = {} if byte_count > self._capacity_bytes: specs.update(data=(byte_count, cp.uint8), normalized=(2 * byte_count, cp.uint8), source_offsets=(2 * byte_count, cp.uint32), raw_source_offsets=(byte_count, cp.uint32)) if row_count > self._capacity_rows: specs.update(offsets=(row_count + 1, cp.int64), timestamps=(row_count, cp.int64), candidates=(row_count * MAX_CANDIDATES * 4, cp.int32), codes=(row_count, cp.uint8), facts=((row_count, 8), cp.uint32), normalized_lengths=(row_count, cp.int32), source_lengths=(row_count, cp.int32), match_events=((row_count, self.event_capacity, 4), cp.int32), event_counts=(row_count, cp.uint32), stats=((row_count, len(STAT_NAMES)), cp.uint64)) self._guard_allocations(int(np.prod(shape)) * np.dtype(dtype).itemsize for shape, dtype in specs.values()) pending = {} try: for name, (shape, dtype) in specs.items(): pending[name] = cp.empty(shape, dtype=dtype) except BaseException: # No stream references these allocations yet; old complete workspace # and capacity metadata remain usable. Cached blocks stay accounted. raise self._workspace.update(pending) self._capacity_bytes = max(self._capacity_bytes, byte_count) self._capacity_rows = max(self._capacity_rows, row_count) def _drain(self, error): try: if self.cp is not None: self.cp.cuda.Stream.null.synchronize() except BaseException as drain_error: self._failed_drain = True # Retain host input and device owners when completion is unknown. raise error from drain_error self._active_transfer_owners = None def _gpu_stage(self, batch, *, text_column, timestamp_column, diagnostic=None): if self._failed_drain: raise RuntimeError("Phase38 CUDA drain failed; end this worker process") try: return self._gpu_stage_staged(batch, text_column=text_column, timestamp_column=timestamp_column, diagnostic=diagnostic) except BaseException as error: self._drain(error) raise def _gpu_stage_staged(self, batch, *, text_column, timestamp_column, diagnostic=None): cp, total_started = self.cp, time.perf_counter() data, offsets = phase32a_gpu._utf8_buffers(batch.column(batch.schema.get_field_index(text_column))) epochs = phase32a_gpu._timestamps(batch.column(batch.schema.get_field_index(timestamp_column))) rows, byte_count = batch.num_rows, data.nbytes if not rows: raise ValueError("Phase38 device stage requires at least one row") self._reserve(max(1, byte_count), rows) b, c = self._workspace, self._device_config self._active_transfer_owners = (batch, data, offsets, epochs) preparation = time.perf_counter() - total_started if diagnostic is not None: phase32a_gpu._record_interval(diagnostic, "host_prepare", "main", total_started, total_started + preparation) started = time.perf_counter() if byte_count: b["data"][:byte_count].set(data) b["offsets"][:rows + 1].set(offsets) b["timestamps"][:rows].set(epochs) cp.cuda.Stream.null.synchronize() h2d = time.perf_counter() - started if diagnostic is not None: phase32a_gpu._record_interval(diagnostic, "h2d", "main", started, started + h2d) collect_stats = diagnostic is not None or self.kernel_event_diagnostics count, collect = np.int32(rows), np.int32(collect_stats) row_grid = ((rows + 127) // 128,) started = time.perf_counter() if collect_stats: self._events[0].record() self._kernels["prepare"](row_grid, (128,), ( b["data"], b["offsets"], count, np.int32(self.max_row_bytes), c["unicode_index"], c["unicode_pool"], c["python_whitespace"], b["normalized"], b["source_offsets"], b["raw_source_offsets"], b["normalized_lengths"], b["source_lengths"], b["event_counts"], b["codes"], b["facts"], b["stats"], collect)) if collect_stats: self._events[1].record() pfac = self.matcher == "pfac" matcher_arrays = ("pfac_transitions", "classes", "pfac_meta", "pfac_output_offsets", "pfac_output_keys") if pfac else ( "transitions", "classes", "dfa_meta", "output_offsets", "output_keys") arguments = [b["data"], b["offsets"], count, b["normalized"], b["source_offsets"], b["raw_source_offsets"], b["normalized_lengths"]] arguments.extend(c[name] for name in matcher_arrays) arguments.extend((c["key_lengths"], b["codes"], b["match_events"], b["event_counts"], np.int32(self.event_capacity), b["stats"], collect)) self._kernel((rows,) if pfac else row_grid, (128,), tuple(arguments)) if collect_stats: self._events[2].record() arguments = [b["offsets"], b["timestamps"], count] arguments.extend(c[name] for name in ("key_offsets", "key_consumers", "consumers", "aliases", "lexical", "contexts", "clauses", "term_ids", "term_bytes", "term_offsets")) arguments.extend((np.int32(self._conjunction_count), b["normalized"], b["source_offsets"], b["normalized_lengths"], b["source_lengths"], b["match_events"], b["event_counts"], np.int32(self.event_capacity), b["candidates"], b["codes"], b["facts"])) self._kernels["evaluate"](row_grid, (128,), tuple(arguments)) if collect_stats: self._events[3].record() cp.cuda.Stream.null.synchronize() kernel = time.perf_counter() - started if diagnostic is not None: phase32a_gpu._record_interval(diagnostic, "kernel_host_interval", "main", started, started + kernel) started = time.perf_counter() codes = cp.asnumpy(b["codes"][:rows]) dense = cp.asnumpy(b["facts"][:rows]) if np.any(codes == 7): raise RuntimeError("Phase38 frozen mapping guard failed before native completion") stats = cp.asnumpy(b["stats"][:rows]) if collect_stats else None event_counts = cp.asnumpy(b["event_counts"][:rows]) if collect_stats else None d2h = time.perf_counter() - started if diagnostic is not None: phase32a_gpu._record_interval(diagnostic, "d2h", "main", started, started + d2h) self._active_transfer_owners = None for values in (codes, dense, offsets): values.setflags(write=False) profile = {"rows": rows, "bytes": byte_count, "host_prepare_seconds": preparation, "h2d_seconds": h2d, "kernel_seconds": kernel, "d2h_seconds": d2h} if stats is not None: totals = stats.sum(axis=0, dtype=np.uint64) profile.update({"device_" + name: int(value) for name, value in zip(STAT_NAMES, totals)}) profile["event_capacity_rows"] = int(np.count_nonzero(event_counts > self.event_capacity)) profile["buffered_match_events"] = int(np.minimum(event_counts, self.event_capacity).sum()) if collect_stats: for index, name in enumerate(("prepare", "matching", "evaluation")): profile["device_event_" + name + "_seconds"] = float(cp.cuda.get_elapsed_time( self._events[index], self._events[index + 1])) / 1000 if diagnostic is not None: diagnostic["phase38_device_statistics"] = { "instrumented": True, "matcher": self.matcher, "depth_unit": "transition attempts per starting position for PFAC; per stream for AC", "overflow_censors_work": True, "event_order": "unspecified", "max_depth": int(stats[:, 24].max()), "device_event_timings": True} diagnostic["transfer_payload_bytes"] = { "h2d": {"text": byte_count, "row_offsets": offsets.nbytes, "timestamps": epochs.nbytes}, "d2h": {"route_codes": codes.nbytes, "dense_facts": dense.nbytes, "statistics": stats.nbytes, "event_counts": event_counts.nbytes}} return codes, dense, offsets, profile