cuda/src/range_curate_cuda/_runtime/phase38_gpu.py

"""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