python/dpt/kernels/transmission.py

1"""Strict mixed-precision transmission kernels, with no launch-time allocation.23The host boundary owns domain, shape, stream and alias validation. Every launch4must use ``record_tape=False``: the public operator supplies its explicit VJP.5Mask bits select transmission (1), counts (2), log transmission (4), removed6primary fraction (8). Beam modes are absent (0), fixed scalar (1), device scalar7(2) and per-pixel (3). Unused array arguments may be empty device placeholders.89Status is an int32 array with at least one element: bit 0 is invalid input,10bit 1 is a non-finite gradient. The caller resets it before checked execution.11Only exceptional lanes update status; numerical reductions never use atomics.12"""1314from functools import cache1516import warp as wp1718TILE_SIZE = 25619STRICT_OPTIONS = {"fast_math": False, "fuse_fp": True, "enable_backward": False}202122# region book:transmission-weak-attenuation23@wp.func_native("return -expm1f(-optical_depth);")24def removed_primary(optical_depth: wp.float32) -> wp.float32:25    """Evaluate the decrement directly, without subtracting two near-unities."""26    ...272829@wp.func_native("return -log1pf(-decrement);")30def optical_depth_from_decrement(decrement: wp.float32) -> wp.float32:31    """Invert a supplied decrement; the caller requires 0 <= decrement < 1."""32    ...333435# endregion book:transmission-weak-attenuation363738# region book:transmission-evaluate39@wp.func_native("return -value;")40def log_transmission_exact(value: wp.float32) -> wp.float32:41    """Use IEEE unary negation; Warp 1.17's generic neg is ``0 - value``."""42    ...434445@wp.func46def transmission_factor(optical_depth: wp.float32) -> wp.float64:47    """Keep exponent range until all requested physical products are formed."""48    if optical_depth <= wp.float32(64.0):49        # Here exp(-L) is normal in FP32. Promotion retains its range, while50        # keeping the common case on the single-precision exponential path.51        return wp.float64(wp.exp(-optical_depth))52    return wp.exp(-wp.float64(optical_depth))535455# endregion book:transmission-evaluate565758# region book:transmission-forward59@cache60def get_forward_kernel(mask: int, beam_mode: int):61    """Return a specialised pointwise kernel, launched with ``dim=P``.6263    Inputs: L, beam_array, beam_scalar. Outputs: T, counts, log_T, removed.64    Static selection eliminates unused array reads, stores and mathematics.65    """66    if mask < 1 or mask > 15 or beam_mode not in (0, 1, 2, 3):67        raise ValueError("Invalid output mask or beam mode")68    if mask & 2 and beam_mode == 0:69        raise ValueError("Counts require an open beam")7071    @wp.kernel(module="unique", module_options=STRICT_OPTIONS)72    def forward(73        optical_depth: wp.array(dtype=wp.float32),74        beam: wp.array(dtype=wp.float32),75        beam_scalar: wp.float32,76        transmission: wp.array(dtype=wp.float32),77        counts: wp.array(dtype=wp.float32),78        log_transmission: wp.array(dtype=wp.float32),79        removed: wp.array(dtype=wp.float32),80    ):81        p = wp.tid()82        depth = optical_depth[p]83        if wp.static(mask & 3 != 0):84            factor = transmission_factor(depth)85            if wp.static(mask & 1 != 0):86                transmission[p] = wp.float32(factor)87            if wp.static(mask & 2 != 0):88                illumination = beam_scalar89                if wp.static(beam_mode == 2):90                    illumination = beam[0]91                elif wp.static(beam_mode == 3):92                    illumination = beam[p]93                # In particular, never multiply illumination by stored T.94                count = wp.float32(wp.float64(illumination) * factor)95                if count == wp.float32(0.0):96                    count = wp.float32(0.0)97                counts[p] = count98        if wp.static(mask & 4 != 0):99            log_transmission[p] = log_transmission_exact(depth)100        if wp.static(mask & 8 != 0):101            decrement = removed_primary(depth)102            if decrement == wp.float32(0.0):103                decrement = wp.float32(0.0)104            removed[p] = decrement105106    return forward107108109# endregion book:transmission-forward110111112@wp.func113def checked_gradient(value: wp.float64, status: wp.array(dtype=wp.int32)) -> wp.float32:114    """Round once and flag overflow, including after tape accumulation."""115    rounded = wp.float32(value)116    if not wp.isfinite(rounded):117        wp.atomic_or(status, 0, 2)118    return rounded119120121# region book:transmission-vjp122@wp.func123def weighted_depth_adjoint(124    factor: wp.float64,125    illumination: wp.float32,126    seed_transmission: wp.float32,127    seed_counts: wp.float32,128    seed_log: wp.float32,129    seed_removed: wp.float32,130) -> wp.float64:131    """Differentiate the real-valued map, retaining range before rounding."""132    optical = (wp.float64(seed_removed) - wp.float64(seed_transmission)) * factor133    count = wp.float64(seed_counts) * wp.float64(illumination) * factor134    return optical - count - wp.float64(seed_log)135136137# endregion book:transmission-vjp138139140@cache141def get_vjp_kernel(142    seed_mask: int, beam_mode: int, write_depth: bool, write_beam: bool, accumulate: bool143):144    """Pointwise VJP: inputs L, beam, scalar, seed_T/counts/log/removed;145    outputs grad_L, grad_beam, status. Launch ``dim=P``.146147    ``write_beam`` is only for per-pixel beams. Broadcast beams use the tiled148    reduction below. ``accumulate`` is reserved for the tape adapter.149    """150    if seed_mask < 0 or seed_mask > 15 or beam_mode not in (0, 1, 2, 3):151        raise ValueError("Invalid seed mask or beam mode")152    if seed_mask & 2 and beam_mode == 0:153        raise ValueError("Count cotangents require an open beam")154    if write_beam and beam_mode != 3:155        raise ValueError("Pointwise beam gradients require a per-pixel beam")156157    @wp.kernel(module="unique", module_options=STRICT_OPTIONS)158    def vjp(159        optical_depth: wp.array(dtype=wp.float32),160        beam: wp.array(dtype=wp.float32),161        beam_scalar: wp.float32,162        seed_transmission: wp.array(dtype=wp.float32),163        seed_counts: wp.array(dtype=wp.float32),164        seed_log: wp.array(dtype=wp.float32),165        seed_removed: wp.array(dtype=wp.float32),166        grad_depth: wp.array(dtype=wp.float32),167        grad_beam: wp.array(dtype=wp.float32),168        status: wp.array(dtype=wp.int32),169    ):170        p = wp.tid()171        factor = wp.float64(0.0)172        if wp.static(seed_mask & 11 != 0):173            factor = transmission_factor(optical_depth[p])174        a = wp.float32(0.0)175        b = wp.float32(0.0)176        c = wp.float32(0.0)177        r = wp.float32(0.0)178        if wp.static(seed_mask & 1 != 0):179            a = seed_transmission[p]180        if wp.static(seed_mask & 2 != 0):181            b = seed_counts[p]182        if wp.static(seed_mask & 4 != 0):183            c = seed_log[p]184        if wp.static(seed_mask & 8 != 0):185            r = seed_removed[p]186        if wp.static(write_depth):187            illumination = wp.float32(0.0)188            if wp.static(seed_mask & 2 != 0):189                illumination = beam_scalar190                if wp.static(beam_mode == 2):191                    illumination = beam[0]192                elif wp.static(beam_mode == 3):193                    illumination = beam[p]194            derivative = weighted_depth_adjoint(factor, illumination, a, b, c, r)195            if wp.static(accumulate):196                derivative = derivative + wp.float64(grad_depth[p])197            grad_depth[p] = checked_gradient(derivative, status)198        if wp.static(write_beam):199            derivative_beam = wp.float64(b) * factor200            if wp.static(accumulate):201                derivative_beam = derivative_beam + wp.float64(grad_beam[p])202            grad_beam[p] = checked_gradient(derivative_beam, status)203204    return vjp205206207@wp.func208def beam_contribution(depth: wp.float32, seed: wp.float32) -> wp.float64:209    return wp.float64(seed) * transmission_factor(depth)210211212# region book:transmission-beam-reduction213@cache214def get_beam_partial_kernel():215    """Inputs L, seed_counts; output FP64 partials.216217    Use ``launch_tiled(dim=ceil(P/256), block_dim=256)``. Tile loads zero-pad218    their bounds, so padded seeds contribute zero even though exp(-0) is one.219    The separate pass deliberately avoids retaining an 8P-byte contribution220    image or contending for a global scalar. Its traffic is accounted separately.221    """222223    @wp.kernel(module="unique", module_options=STRICT_OPTIONS)224    def partials(225        optical_depth: wp.array(dtype=wp.float32),226        seed_counts: wp.array(dtype=wp.float32),227        output: wp.array(dtype=wp.float64),228    ):229        block = wp.tid()230        depths = wp.tile_load(optical_depth, shape=TILE_SIZE, offset=block * TILE_SIZE)231        seeds = wp.tile_load(seed_counts, shape=TILE_SIZE, offset=block * TILE_SIZE)232        contributions = wp.tile_map(beam_contribution, depths, seeds)233        total = wp.tile_sum(contributions)234        wp.tile_store(output, total, offset=block)235236    return partials237238239# endregion book:transmission-beam-reduction240241242@cache243def get_inverse_kernel():244    """Inputs decrement; outputs optical depth; launch ``dim=P``."""245246    @wp.kernel(module="unique", module_options=STRICT_OPTIONS)247    def inverse(delta: wp.array(dtype=wp.float32), depth: wp.array(dtype=wp.float32)):248        p = wp.tid()249        value = optical_depth_from_decrement(delta[p])250        if value == wp.float32(0.0):251            value = wp.float32(0.0)252        depth[p] = value253254    return inverse255256257@cache258def get_inverse_vjp_kernel(accumulate: bool):259    """Inputs delta, seed_L; outputs grad_delta, status; launch ``dim=P``."""260261    @wp.kernel(module="unique", module_options=STRICT_OPTIONS)262    def inverse_vjp(263        delta: wp.array(dtype=wp.float32),264        seed: wp.array(dtype=wp.float32),265        gradient: wp.array(dtype=wp.float32),266        status: wp.array(dtype=wp.int32),267    ):268        p = wp.tid()269        derivative = wp.float64(seed[p]) / (wp.float64(1.0) - wp.float64(delta[p]))270        if wp.static(accumulate):271            derivative = derivative + wp.float64(gradient[p])272        gradient[p] = checked_gradient(derivative, status)273274    return inverse_vjp275276277@wp.kernel(module="unique", module_options=STRICT_OPTIONS)278def finite_nonnegative(values: wp.array(dtype=wp.float32), status: wp.array(dtype=wp.int32)):279    """Inputs values; outputs status; launch dim=values.shape[0]."""280    value = values[wp.tid()]281    if not wp.isfinite(value) or value < wp.float32(0.0):282        wp.atomic_or(status, 0, 1)283284285@wp.kernel(module="unique", module_options=STRICT_OPTIONS)286def finite_seed(values: wp.array(dtype=wp.float32), status: wp.array(dtype=wp.int32)):287    """Inputs cotangents; outputs status; launch dim=values.shape[0]."""288    if not wp.isfinite(values[wp.tid()]):289        wp.atomic_or(status, 0, 1)290291292@wp.kernel(module="unique", module_options=STRICT_OPTIONS)293def decrement_domain(values: wp.array(dtype=wp.float32), status: wp.array(dtype=wp.int32)):294    """Inputs decrements; outputs status; launch dim=values.shape[0]."""295    value = values[wp.tid()]296    if not wp.isfinite(value) or value < wp.float32(0.0) or value >= wp.float32(1.0):297        wp.atomic_or(status, 0, 1)298299300@wp.kernel(module="unique", module_options={"enable_backward": True})301def dependency_marker(302    depth: wp.array(dtype=wp.float32),303    beam: wp.array(dtype=wp.float32),304    transmission: wp.array(dtype=wp.float32),305    counts: wp.array(dtype=wp.float32),306    log_transmission: wp.array(dtype=wp.float32),307    removed: wp.array(dtype=wp.float32),308):309    """Zero-dimensional tape bookkeeping only; never executes a CUDA thread."""310    pass311