Spaces:
Running on Zero
Running on Zero
| """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() | |
| 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 | |