Spaces:
Running on Zero
Running on Zero
| """Phase40 fixed 128-byte segmented AC with exact owned-ending-byte output. | |
| One instance retains original staged AC/PFAC controls and the same preparation, | |
| workspace and evaluator. Segmentation changes scheduling only; source offsets | |
| and boundary tests retain whole-row semantics. No length sorting is performed. | |
| """ | |
| from __future__ import annotations | |
| import hashlib | |
| from pathlib import Path | |
| import time | |
| import numpy as np | |
| from . import phase32a_gpu, phase32b_gpu, phase38_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, phase38_gpu) | |
| } | |
| KERNEL_SOURCE_PATH = Path(__file__).with_name("phase40_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() | |
| _BASE_KERNEL_FILE_SHA256_AT_IMPORT = hashlib.sha256(phase38_gpu.KERNEL_SOURCE_PATH.read_bytes()).hexdigest() | |
| KERNEL_INCLUDE_PATHS = (phase38_gpu.KERNEL_SOURCE_PATH,) | |
| MAX_CANDIDATES = phase38_gpu.MAX_CANDIDATES | |
| TILE_BYTES = 128 | |
| workspace_nbytes = phase38_gpu.workspace_nbytes | |
| workspace_bytes = workspace_nbytes | |
| STAT_NAMES = phase38_gpu.STAT_NAMES | |
| SEGMENTED_STAT_NAMES = ( | |
| "tile_jobs", "transition_attempts", "owned_terminal_visits", "boundary_passed_events", | |
| *STAT_NAMES[4:22], "halo_transition_attempts", "normalized_bytes", "max_depth_per_row_sum", | |
| ) | |
| def _stream_max_lengths(arrays): | |
| """Derive each stream's maximum from its actual completed-DFA output IDs. | |
| Raw key IDs may be interleaved across streams; do not infer their partition | |
| from contiguous slices of key_lengths. Repeated suffix outputs are harmless | |
| when taking this maximum. Empty streams retain zero and are not launched. | |
| """ | |
| lengths = [] | |
| metadata = arrays["dfa_meta"].reshape(3, 4) | |
| for stream, meta in enumerate(metadata): | |
| transition_start, width, output_start, pattern_count = map(int, meta) | |
| transition_end = int(metadata[stream + 1, 0]) if stream < 2 else arrays["transitions"].size | |
| state_count, remainder = divmod(transition_end - transition_start, width) | |
| if remainder or state_count < 1: | |
| raise RuntimeError("Phase40 malformed inherited DFA dimensions") | |
| first = int(arrays["output_offsets"][output_start]) | |
| last = int(arrays["output_offsets"][output_start + state_count]) | |
| keys = arrays["output_keys"][first:last] | |
| maximum = int(arrays["key_lengths"][keys].max()) if keys.size else 0 | |
| if bool(pattern_count) != bool(maximum): | |
| raise RuntimeError("Phase40 stream output lengths disagree with its patterns") | |
| lengths.append(maximum) | |
| result = np.asarray(lengths, dtype=np.int32) | |
| result.setflags(write=False) | |
| return result | |
| class ExactSegmentedFrozenUnicodeCudaClassifier(phase38_gpu.ExactStagedFrozenUnicodeCudaClassifier): | |
| def __init__(self, *args, matcher="segmented", tile_bytes=TILE_BYTES, **kwargs): | |
| if matcher not in ("ac", "pfac", "segmented"): | |
| raise ValueError("matcher must be ac, pfac or segmented") | |
| if isinstance(tile_bytes, bool) or not isinstance(tile_bytes, int) or tile_bytes != TILE_BYTES: | |
| raise ValueError("Phase40 preregistered tile size is exactly 128 UTF-8 bytes") | |
| super().__init__(*args, matcher="ac" if matcher == "segmented" else matcher, **kwargs) | |
| self.matcher = matcher | |
| self._cuda_source_sha256 = hashlib.sha256(CUDA_SOURCE.encode()).hexdigest() | |
| self._stream_length_identity = None | |
| if not self.unsupported_reason: | |
| lengths = _stream_max_lengths(self._arrays) | |
| self._arrays["segmented_max_lengths"] = lengths | |
| self._stream_length_identity = { | |
| "maximum_pattern_utf8_bytes_by_stream": lengths.tolist(), | |
| "sha256": hashlib.sha256(lengths.tobytes()).hexdigest(), "bytes": lengths.nbytes, | |
| } | |
| self._config_requested_bytes = sum(max(value.nbytes, value.dtype.itemsize) | |
| for value in self._arrays.values()) | |
| def backend_report(self): | |
| return {**super().backend_report, | |
| "backend": "cpu" if self.unsupported_reason else "experimental_segmented_frozen_unicode_cuda", | |
| "phase40_kernel_file_sha256": _KERNEL_FILE_SHA256_AT_IMPORT, | |
| "included_phase38_kernel_file_sha256": _BASE_KERNEL_FILE_SHA256_AT_IMPORT, | |
| "kernel_entrypoints": ["range38_prepare", "range38_match_ac", "range38_match_pfac", | |
| "range40_match_segmented", "range38_evaluate"], | |
| "tile_bytes": TILE_BYTES, "stream_length_identity": self._stream_length_identity, | |
| "tile_ownership": "final matching byte in the disjoint owned tile; whole-row neighbor boundary checks", | |
| "tile_warmup": "AC starts at root at max(0,tile_start-Lmax+1), before the owned ending bytes", | |
| "tile_work_bound": "n + max(0,ceil(n/128)-1)*min(max(0,Lmax-1),n) transition attempts per stream before output-overflow early exits", | |
| "segmented_work_assignment": "one128-lane block per article; lane-strided128-byte tiles; no sorting", | |
| "long_pattern_policy": "existing supported taxonomy and row-byte caps; no additional length cutoff"} | |
| # Source-matched inherited methods: only selector/loading, matching dispatch | |
| # and diagnostic labels differ. Preparation/evaluation/transfers are retained. | |
| def select_matcher(self, matcher): | |
| if matcher not in ("ac", "pfac", "segmented"): | |
| raise ValueError("matcher must be ac, pfac or segmented") | |
| 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 _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(("range40_" if name == "match_segmented" else "range38_") + name) | |
| for name in ("prepare", "match_ac", "match_pfac", "evaluate", "match_segmented")} | |
| # 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 _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() | |
| if self.matcher == "segmented": | |
| 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 ( | |
| "transitions", "classes", "dfa_meta", "output_offsets", "output_keys", "key_lengths", | |
| "segmented_max_lengths")) | |
| arguments.extend((b["codes"], b["match_events"], b["event_counts"], np.int32(self.event_capacity), | |
| b["stats"], collect)) | |
| self._kernel((rows,), (128,), tuple(arguments)) | |
| else: | |
| 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) | |
| stat_names = SEGMENTED_STAT_NAMES if self.matcher == "segmented" else STAT_NAMES | |
| 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["phase40_device_statistics"] = { | |
| "instrumented": True, "matcher": self.matcher, | |
| "depth_unit": ("transition attempts per tile including preceding warm-up bytes" if self.matcher == "segmented" else | |
| "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 | |