cuda/src/range_curate_cuda/_runtime/phase40_gpu.py

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

    @property
    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