onnx-runtime-ep-cuda 0.1.0-dev.5

//! Additive pointwise ops — **unary math**, **logical**, and **comparison** —
//! on the GPU via runtime-compiled (NVRTC) `extern "C"` kernels. This is CUDA
//! Wave 3 (`docs/CUDA_COVERAGE.md`), extending the [`super::elementwise`] slice
//! with the remaining CPU-EP pointwise coverage (RULES.md #4 — pointwise chains
//! have no NVIDIA library op and stay ours so they can later fuse into a GEMM
//! epilogue or a producer→activation→add chain).
//!
//! ## Scope of this slice (all limits are actionable errors, never panics)
//!
//! * **Unary math** (`Abs`, `Neg`, `Reciprocal`, `Exp`, `Log`, `Sign`, `Floor`,
//!   `Ceil`, `Round`, `Sin`, `Cos`, `Softplus`): one input, one output, identical
//!   shape, **f32/f16/bf16** (half storage computes in f32). Each formula is matched **exactly** to the CPU EP
//!   (`crates/onnx-runtime-ep-cpu/src/kernels/unary_math.rs`) so the untestable-
//!   on-this-host kernels stay numerically identical to the reference path.
//! * **Not** (`Not`): boolean element negation, matched to the CPU EP
//!   (`logical.rs` — non-zero byte is `true`, output is canonical `1`/`0`).
//! * **Comparison** (`Equal`, `Greater`, `Less`, `GreaterOrEqual`,
//!   `LessOrEqual`): two broadcast-compatible **f32/i32/i64** inputs → **Bool**
//!   output; `Equal` also accepts **Bool** inputs.
//! * **Logical** (`And`, `Or`, `Xor`): two broadcast-compatible **Bool** inputs →
//!   **Bool** output (non-zero byte is `true`, canonical `1`/`0` out).
//!
//! `dtype`: f32/f16/bf16 for unary math, f32/i32/i64 for comparison (plus bool
//! for `Equal`), and bool for
//! logical/`Not`; other dtypes return an actionable error naming the dtype/op.
//!
//! **Broadcasting:** binary comparison/logical ops reuse the same right-aligned,
//! zero-stride metadata as [`super::elementwise`].
//!
//! Each op is one thread-per-element grid-stride kernel (bandwidth-bound,
//! PyTorch-pointwise shaped).

use std::ffi::c_void;
use std::sync::{Arc, Mutex};

use cudarc::driver::{LaunchConfig, PushKernelArg};

use onnx_runtime_ep_api::{EpError, Kernel, KernelFactory, Result, TensorMut, TensorView};
use onnx_runtime_ir::{DataType, Node};

use super::elementwise::{
    BroadcastMetadataCache, BroadcastMetadataKey, is_fixed_decode_shape,
    require_matching_capture_signature,
};
use crate::error::{driver_err, not_implemented};
use crate::runtime::{CudaRuntime, cuptr};

/// Threads per block for the 1-D pointwise grids (a full warp-multiple block).
const BLOCK: u32 = 256;

/// Grid dimension for `n` elements at [`BLOCK`] threads, capped so a huge tensor
/// still fits the grid limit (the kernels are grid-stride, so a capped grid
/// still covers every element).
fn grid_for(n: usize) -> u32 {
    const MAX_BLOCKS: usize = 65_535;
    n.div_ceil(BLOCK as usize).clamp(1, MAX_BLOCKS) as u32
}

/// Reject a dtype other than the expected one with an actionable, op-named error.
fn require_dtype(op: &str, name: &str, dt: DataType, want: DataType) -> Result<()> {
    if dt != want {
        return Err(not_implemented(format!(
            "{op} with {name} dtype {dt:?} (this slice supports {want:?} only; \
             f16/bf16 pending — see docs/CUDA_COVERAGE.md)"
        )));
    }

    Ok(())
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
enum FloatDtype {
    F32,
    F16,
    Bf16,
}

impl FloatDtype {
    fn from_onnx(op: &str, name: &str, dtype: DataType) -> Result<Self> {
        match dtype {
            DataType::Float32 => Ok(Self::F32),
            DataType::Float16 => Ok(Self::F16),
            DataType::BFloat16 => Ok(Self::Bf16),
            other => Err(not_implemented(format!(
                "{op} with {name} dtype {other:?} (supported: Float32, Float16, BFloat16)"
            ))),
        }
    }

    fn suffix(self) -> &'static str {
        match self {
            Self::F32 => "f32",
            Self::F16 => "f16",
            Self::Bf16 => "bf16",
        }
    }
}

/// Reject a strided (non-contiguous) view with a "materialise first" error.
fn require_contiguous(op: &str, name: &str, contiguous: bool) -> Result<()> {
    if !contiguous {
        return Err(not_implemented(format!(
            "{op} with a non-contiguous (strided) {name}; \
             insert an explicit copy to materialise it before the op"
        )));
    }
    Ok(())
}

/// `n` as `u64`, matching the kernels' `unsigned long long` count parameter.
fn count_u64(op: &str, n: usize) -> Result<u64> {
    u64::try_from(n)
        .map_err(|_| EpError::KernelFailed(format!("cuda_ep {op}: {n} elements exceed u64")))
}

// ===========================================================================
// Unary math (f32 → f32)
// ===========================================================================

/// NVRTC source: dtype-templated pointwise kernels for each unary math op. NVRTC
/// resolves the CUDA device intrinsics (`expf`, `logf`, `sinf`, `rintf`,
/// `log1pf`, …) with no header include. Each formula is annotated with the exact
/// CPU-EP expression it mirrors (`unary_math.rs`).
const UNARY_MATH_SRC: &str = r#"
#if __has_include(<cuda_fp16.h>) && __has_include(<cuda_bf16.h>)
#define NXRT_HAS_CUDA_HALF_HEADERS 1
#include <cuda_fp16.h>
#include <cuda_bf16.h>
#endif

template <typename T> __device__ float load_float(T value);
template <> __device__ float load_float<float>(float value) { return value; }
#ifdef NXRT_HAS_CUDA_HALF_HEADERS
template <> __device__ float load_float<__half>(__half value) { return __half2float(value); }
template <> __device__ float load_float<__nv_bfloat16>(__nv_bfloat16 value) { return __bfloat162float(value); }
#endif
template <typename T> __device__ T store_float(float value);
template <> __device__ float store_float<float>(float value) { return value; }
#ifdef NXRT_HAS_CUDA_HALF_HEADERS
template <> __device__ __half store_float<__half>(float value) { return __float2half_rn(value); }
template <> __device__ __nv_bfloat16 store_float<__nv_bfloat16>(float value) { return __float2bfloat16_rn(value); }
#endif

__device__ float op_abs(float x) { return fabsf(x); }
__device__ float op_neg(float x) { return -x; }
__device__ float op_reciprocal(float x) { return 1.0f / x; }
__device__ float op_exp(float x) { return expf(x); }
__device__ float op_log(float x) { return logf(x); }
__device__ float op_sign(float x) {
    return (x != x) ? x : ((x > 0.0f) ? 1.0f : ((x < 0.0f) ? -1.0f : 0.0f));
}
__device__ float op_floor(float x) { return floorf(x); }
__device__ float op_ceil(float x) { return ceilf(x); }
__device__ float op_round(float x) { return rintf(x); }
__device__ float op_sin(float x) { return sinf(x); }
__device__ float op_cos(float x) { return cosf(x); }
__device__ float op_softplus(float x) { return fmaxf(x, 0.0f) + log1pf(expf(-fabsf(x))); }

#define DEFINE_UNARY(NAME, TYPE, SUFFIX) \
extern "C" __global__ void NAME##_##SUFFIX(const TYPE* x, TYPE* y, const unsigned long long n) { \
    for (unsigned long long i = blockIdx.x*blockDim.x + threadIdx.x; i < n; \
         i += (unsigned long long)gridDim.x * blockDim.x) \
        y[i] = store_float<TYPE>(op_##NAME(load_float<TYPE>(x[i]))); \
}
#define DEFINE_FOR_TYPE(TYPE, SUFFIX) \
DEFINE_UNARY(abs, TYPE, SUFFIX) \
DEFINE_UNARY(neg, TYPE, SUFFIX) \
DEFINE_UNARY(reciprocal, TYPE, SUFFIX) \
DEFINE_UNARY(exp, TYPE, SUFFIX) \
DEFINE_UNARY(log, TYPE, SUFFIX) \
DEFINE_UNARY(sign, TYPE, SUFFIX) \
DEFINE_UNARY(floor, TYPE, SUFFIX) \
DEFINE_UNARY(ceil, TYPE, SUFFIX) \
DEFINE_UNARY(round, TYPE, SUFFIX) \
DEFINE_UNARY(sin, TYPE, SUFFIX) \
DEFINE_UNARY(cos, TYPE, SUFFIX) \
DEFINE_UNARY(softplus, TYPE, SUFFIX)
DEFINE_FOR_TYPE(float, f32)
#ifdef NXRT_HAS_CUDA_HALF_HEADERS
DEFINE_FOR_TYPE(__half, f16)
DEFINE_FOR_TYPE(__nv_bfloat16, bf16)
#endif
"#;

const UNARY_MATH_MODULE: &str = "pointwise_unary_math_float_v2";

/// A supported unary math op and its NVRTC entry point.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum UnaryMathOp {
    Abs,
    Neg,
    Reciprocal,
    Exp,
    Log,
    Sign,
    Floor,
    Ceil,
    Round,
    Sin,
    Cos,
    Softplus,
}

impl UnaryMathOp {
    fn stem(self) -> &'static str {
        match self {
            UnaryMathOp::Abs => "abs",
            UnaryMathOp::Neg => "neg",
            UnaryMathOp::Reciprocal => "reciprocal",
            UnaryMathOp::Exp => "exp",
            UnaryMathOp::Log => "log",
            UnaryMathOp::Sign => "sign",
            UnaryMathOp::Floor => "floor",
            UnaryMathOp::Ceil => "ceil",
            UnaryMathOp::Round => "round",
            UnaryMathOp::Sin => "sin",
            UnaryMathOp::Cos => "cos",
            UnaryMathOp::Softplus => "softplus",
        }
    }

    fn entry(self, dtype: FloatDtype) -> String {
        format!("{}_{}", self.stem(), dtype.suffix())
    }

    fn op_name(self) -> &'static str {
        match self {
            UnaryMathOp::Abs => "Abs",
            UnaryMathOp::Neg => "Neg",
            UnaryMathOp::Reciprocal => "Reciprocal",
            UnaryMathOp::Exp => "Exp",
            UnaryMathOp::Log => "Log",
            UnaryMathOp::Sign => "Sign",
            UnaryMathOp::Floor => "Floor",
            UnaryMathOp::Ceil => "Ceil",
            UnaryMathOp::Round => "Round",
            UnaryMathOp::Sin => "Sin",
            UnaryMathOp::Cos => "Cos",
            UnaryMathOp::Softplus => "Softplus",
        }
    }
}

/// Factory for [`UnaryMathKernel`]; carries the op identity and shared runtime.
pub struct UnaryMathFactory {
    pub op: UnaryMathOp,
    pub runtime: Arc<CudaRuntime>,
}

impl KernelFactory for UnaryMathFactory {
    fn create(&self, _node: &Node, _input_shapes: &[Vec<usize>]) -> Result<Box<dyn Kernel>> {
        Ok(Box::new(UnaryMathKernel {
            op: self.op,
            runtime: self.runtime.clone(),
        }))
    }
}

/// NVRTC-backed f32/f16/bf16 unary math kernel.
#[derive(Debug)]
pub struct UnaryMathKernel {
    op: UnaryMathOp,
    runtime: Arc<CudaRuntime>,
}

impl UnaryMathKernel {
    fn run(&self, inputs: &[TensorView], outputs: &mut [TensorMut]) -> Result<()> {
        let op = self.op.op_name();
        if inputs.len() != 1 || outputs.len() != 1 {
            return Err(EpError::KernelFailed(format!(
                "cuda_ep {op}: expected 1 input and 1 output, got {} and {}",
                inputs.len(),
                outputs.len()
            )));
        }
        let x = &inputs[0];
        let dtype = FloatDtype::from_onnx(op, "input", x.dtype)?;
        if dtype != FloatDtype::F32 {
            self.runtime.require_nvrtc_half_headers(op)?;
        }
        if outputs[0].dtype != x.dtype {
            return Err(EpError::KernelFailed(format!(
                "cuda_ep {op}: output dtype {:?} must equal input dtype {:?}",
                outputs[0].dtype, x.dtype
            )));
        }
        require_contiguous(op, "input", x.is_contiguous())?;
        require_contiguous(op, "output", outputs[0].is_contiguous())?;

        if outputs[0].numel() != x.numel() {
            return Err(EpError::KernelFailed(format!(
                "cuda_ep {op}: output has {} elements, expected {} (same shape as input)",
                outputs[0].numel(),
                x.numel()
            )));
        }

        let n = x.numel();
        let n_u64 = count_u64(op, n)?;
        let x_ptr = cuptr(x.data_ptr::<u8>() as *const c_void);
        let y_ptr = cuptr(outputs[0].data_ptr_mut::<u8>() as *const c_void);

        let entry = self.op.entry(dtype);
        let func = self
            .runtime
            .nvrtc_function(UNARY_MATH_MODULE, UNARY_MATH_SRC, &entry)?;
        let cfg = LaunchConfig {
            grid_dim: (grid_for(n), 1, 1),
            block_dim: (BLOCK, 1, 1),
            shared_mem_bytes: 0,
        };
        let stream = self.runtime.stream();
        let mut builder = stream.launch_builder(&func);
        builder.arg(&x_ptr).arg(&y_ptr).arg(&n_u64);
        // SAFETY: `func` is the compiled unary-math entry; the (const float*,
        // float*, unsigned long long) argument list matches its signature;
        // `x_ptr`/`y_ptr` are live device allocations of `n` f32 elements, and
        // the u64 count and indexing cover their validated bounds without overflow.
        unsafe { builder.launch(cfg) }.map_err(|e| driver_err(&format!("launch {entry}"), e))?;
        if self.runtime.is_capturing()? {
            // A stream synchronize is illegal mid-capture. The launch is
            // recorded into the segment graph and replayed, so skip the sync
            // instead of erroring inside the captured segment.
            return Ok(());
        }
        self.runtime.synchronize()
    }
}

impl Kernel for UnaryMathKernel {
    fn execute(&self, inputs: &[TensorView], outputs: &mut [TensorMut]) -> Result<()> {
        self.run(inputs, outputs)
    }

    fn supports_strided_input(&self, _idx: usize) -> bool {
        false
    }

    fn capture_support(&self) -> onnx_runtime_ep_api::CaptureSupport {
        onnx_runtime_ep_api::CaptureSupport::Supported
    }
}

// ===========================================================================
// Not (bool → bool)
// ===========================================================================

/// NVRTC source: boolean negation over raw bytes. Matches the CPU EP
/// (`logical.rs`): a non-zero byte is `true`, output is canonical `1`/`0`.
const NOT_SRC: &str = r#"
extern "C" __global__ void not_bool(const unsigned char* x, unsigned char* y, const unsigned long long n) {
    // CPU: u8::from(b == 0)
    for (unsigned long long i = blockIdx.x*blockDim.x + threadIdx.x; i < n; i += (unsigned long long)gridDim.x * blockDim.x)
        y[i] = (x[i] == 0) ? 1 : 0;
}
"#;

const NOT_MODULE: &str = "pointwise_not_bool";

/// Factory for [`NotKernel`] (no attributes).
pub struct NotFactory {
    pub runtime: Arc<CudaRuntime>,
}

impl KernelFactory for NotFactory {
    fn create(&self, _node: &Node, _input_shapes: &[Vec<usize>]) -> Result<Box<dyn Kernel>> {
        Ok(Box::new(NotKernel {
            runtime: self.runtime.clone(),
        }))
    }
}

/// NVRTC-backed boolean `Not` kernel.
#[derive(Debug)]
pub struct NotKernel {
    runtime: Arc<CudaRuntime>,
}

impl NotKernel {
    fn run(&self, inputs: &[TensorView], outputs: &mut [TensorMut]) -> Result<()> {
        if inputs.len() != 1 || outputs.len() != 1 {
            return Err(EpError::KernelFailed(format!(
                "cuda_ep Not: expected 1 input and 1 output, got {} and {}",
                inputs.len(),
                outputs.len()
            )));
        }
        let x = &inputs[0];
        require_dtype("Not", "input", x.dtype, DataType::Bool)?;
        require_dtype("Not", "output", outputs[0].dtype, DataType::Bool)?;
        require_contiguous("Not", "input", x.is_contiguous())?;
        require_contiguous("Not", "output", outputs[0].is_contiguous())?;

        if outputs[0].numel() != x.numel() {
            return Err(EpError::KernelFailed(format!(
                "cuda_ep Not: output has {} elements, expected {} (same shape as input)",
                outputs[0].numel(),
                x.numel()
            )));
        }

        let n = x.numel();
        let n_u64 = count_u64("Not", n)?;
        let x_ptr = cuptr(x.data_ptr::<u8>() as *const c_void);
        let y_ptr = cuptr(outputs[0].data_ptr_mut::<u8>() as *const c_void);

        let func = self
            .runtime
            .nvrtc_function(NOT_MODULE, NOT_SRC, "not_bool")?;
        let cfg = LaunchConfig {
            grid_dim: (grid_for(n), 1, 1),
            block_dim: (BLOCK, 1, 1),
            shared_mem_bytes: 0,
        };
        let stream = self.runtime.stream();
        let mut builder = stream.launch_builder(&func);
        builder.arg(&x_ptr).arg(&y_ptr).arg(&n_u64);
        // SAFETY: `func` is the compiled `not_bool` entry; the (const uchar*,
        // uchar*, unsigned long long) argument list matches its signature; both
        // pointers are live device allocations of `n` 1-byte bool elements, and
        // the u64 count and indexing cover their validated bounds without overflow.
        unsafe { builder.launch(cfg) }.map_err(|e| driver_err("launch not_bool", e))?;
        if self.runtime.is_capturing()? {
            // A stream synchronize is illegal mid-capture. The launch is
            // recorded into the segment graph and replayed, so skip the sync
            // instead of erroring inside the captured segment.
            return Ok(());
        }
        self.runtime.synchronize()
    }
}

impl Kernel for NotKernel {
    fn execute(&self, inputs: &[TensorView], outputs: &mut [TensorMut]) -> Result<()> {
        self.run(inputs, outputs)
    }

    fn supports_strided_input(&self, _idx: usize) -> bool {
        false
    }

    fn capture_support(&self) -> onnx_runtime_ep_api::CaptureSupport {
        onnx_runtime_ep_api::CaptureSupport::Supported
    }
}

// ===========================================================================
// Comparison (f32/i32/i64, same dtype → bool) — NumPy broadcasting
// ===========================================================================

/// NVRTC source: one `extern "C"` kernel per comparison op/dtype. Outputs are
/// canonical 1-byte bool values, per ONNX comparison semantics.
const CMP_SRC: &str = r#"
__device__ __forceinline__ void broadcast_indices(unsigned long long out, const unsigned long long* m, int rank, unsigned long long* ai, unsigned long long* bi) {
    *ai = 0; *bi = 0;
    for (int axis = rank - 1; axis >= 0; --axis) {
        unsigned long long coord = out % m[axis]; out /= m[axis];
        *ai += coord * m[rank + axis]; *bi += coord * m[2 * rank + axis];
    }
}
#define DEFINE_CMP(name, type, suffix, expr) \
extern "C" __global__ void name##_##suffix(const type* a, const type* b, unsigned char* y, const unsigned long long* m, int rank, const unsigned long long n) { \
    for (unsigned long long i = blockIdx.x*blockDim.x + threadIdx.x; i < n; i += (unsigned long long)gridDim.x * blockDim.x) { \
        unsigned long long ai, bi; broadcast_indices(i, m, rank, &ai, &bi); y[i] = (expr) ? 1 : 0; \
    } \
}
#define DEFINE_CMP_FOR_TYPE(type, suffix) \
DEFINE_CMP(equal, type, suffix, a[ai] == b[bi]) \
DEFINE_CMP(greater, type, suffix, a[ai] > b[bi]) \
DEFINE_CMP(less, type, suffix, a[ai] < b[bi]) \
DEFINE_CMP(greater_equal, type, suffix, a[ai] >= b[bi]) \
DEFINE_CMP(less_equal, type, suffix, a[ai] <= b[bi])
DEFINE_CMP_FOR_TYPE(float, f32)
DEFINE_CMP_FOR_TYPE(int, i32)
DEFINE_CMP_FOR_TYPE(long long, i64)
DEFINE_CMP(equal, unsigned char, bool, a[ai] == b[bi])
"#;

const CMP_MODULE: &str = "pointwise_compare";

/// A supported comparison op (same-type operands, bool output).
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum CmpOp {
    Equal,
    Greater,
    Less,
    GreaterOrEqual,
    LessOrEqual,
}

impl CmpOp {
    fn entry(self, dtype: DataType) -> Option<&'static str> {
        let suffix = match dtype {
            DataType::Float32 => "f32",
            DataType::Int32 => "i32",
            DataType::Int64 => "i64",
            DataType::Bool if self == Self::Equal => "bool",
            _ => return None,
        };
        Some(match (self, suffix) {
            (CmpOp::Equal, "f32") => "equal_f32",
            (CmpOp::Equal, "i32") => "equal_i32",
            (CmpOp::Equal, "i64") => "equal_i64",
            (CmpOp::Equal, "bool") => "equal_bool",
            (CmpOp::Greater, "f32") => "greater_f32",
            (CmpOp::Greater, "i32") => "greater_i32",
            (CmpOp::Greater, "i64") => "greater_i64",
            (CmpOp::Less, "f32") => "less_f32",
            (CmpOp::Less, "i32") => "less_i32",
            (CmpOp::Less, "i64") => "less_i64",
            (CmpOp::GreaterOrEqual, "f32") => "greater_equal_f32",
            (CmpOp::GreaterOrEqual, "i32") => "greater_equal_i32",
            (CmpOp::GreaterOrEqual, "i64") => "greater_equal_i64",
            (CmpOp::LessOrEqual, "f32") => "less_equal_f32",
            (CmpOp::LessOrEqual, "i32") => "less_equal_i32",
            (CmpOp::LessOrEqual, "i64") => "less_equal_i64",
            _ => unreachable!("unsupported comparison dtype was filtered above"),
        })
    }

    fn op_name(self) -> &'static str {
        match self {
            CmpOp::Equal => "Equal",
            CmpOp::Greater => "Greater",
            CmpOp::Less => "Less",
            CmpOp::GreaterOrEqual => "GreaterOrEqual",
            CmpOp::LessOrEqual => "LessOrEqual",
        }
    }
}

/// Returns a claim-time rejection reason for a comparison dtype contract.
pub(crate) fn comparison_unsupported_reason(op: &str, input_dtypes: &[DataType]) -> Option<String> {
    let Some(&a) = input_dtypes.first() else {
        return Some(format!("{op}: missing operand dtype for CUDA EP"));
    };
    let Some(&b) = input_dtypes.get(1) else {
        return Some(format!("{op}: missing second operand dtype for CUDA EP"));
    };
    if a != b {
        return Some(format!(
            "{op}: operands must have the same dtype on CUDA EP (got {a:?} and {b:?})"
        ));
    }
    let supported = matches!(a, DataType::Float32 | DataType::Int32 | DataType::Int64)
        || (op == "Equal" && a == DataType::Bool);
    (!supported).then(|| format!("{op}: operand dtype {a:?} not supported on CUDA EP"))
}

// ===========================================================================
// Logical (bool, bool → bool) — NumPy broadcasting
// ===========================================================================

/// NVRTC source: one `extern "C"` kernel per logical op — two bool operands (a
/// non-zero byte is `true`, matching the CPU `Not`), 1-byte bool output.
const LOGICAL_SRC: &str = r#"
__device__ __forceinline__ void broadcast_indices(unsigned long long out, const unsigned long long* m, int rank, unsigned long long* ai, unsigned long long* bi) {
    *ai = 0; *bi = 0;
    for (int axis = rank - 1; axis >= 0; --axis) {
        unsigned long long coord = out % m[axis]; out /= m[axis];
        *ai += coord * m[rank + axis]; *bi += coord * m[2 * rank + axis];
    }
}
#define DEFINE_LOGICAL(name, expr) \
extern "C" __global__ void name(const unsigned char* a, const unsigned char* b, unsigned char* y, const unsigned long long* m, int rank, const unsigned long long n) { \
    for (unsigned long long i = blockIdx.x*blockDim.x + threadIdx.x; i < n; i += (unsigned long long)gridDim.x * blockDim.x) { \
        unsigned long long ai, bi; broadcast_indices(i, m, rank, &ai, &bi); y[i] = (expr) ? 1 : 0; \
    } \
}
DEFINE_LOGICAL(and_bool, (a[ai] != 0) && (b[bi] != 0))
DEFINE_LOGICAL(or_bool, (a[ai] != 0) || (b[bi] != 0))
DEFINE_LOGICAL(xor_bool, (a[ai] != 0) != (b[bi] != 0))
"#;

const LOGICAL_MODULE: &str = "pointwise_logical_bool";

/// A supported logical op (bool operands, bool output).
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum LogicalOp {
    And,
    Or,
    Xor,
}

impl LogicalOp {
    fn entry(self) -> &'static str {
        match self {
            LogicalOp::And => "and_bool",
            LogicalOp::Or => "or_bool",
            LogicalOp::Xor => "xor_bool",
        }
    }

    fn op_name(self) -> &'static str {
        match self {
            LogicalOp::And => "And",
            LogicalOp::Or => "Or",
            LogicalOp::Xor => "Xor",
        }
    }
}

/// The dtype contract for a binary op: operand dtype and output dtype.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
enum BinaryKind {
    /// Same-dtype operands, bool output (comparison).
    Compare(CmpOp),
    /// bool operands, bool output (logical).
    LogicalBool,
}

/// Factory for a binary comparison kernel.
pub struct CmpFactory {
    pub op: CmpOp,
    pub runtime: Arc<CudaRuntime>,
}

impl KernelFactory for CmpFactory {
    fn create(&self, _node: &Node, _input_shapes: &[Vec<usize>]) -> Result<Box<dyn Kernel>> {
        Ok(Box::new(BinaryPredKernel {
            op_name: self.op.op_name(),
            entry: "",
            module: CMP_MODULE,
            src: CMP_SRC,
            kind: BinaryKind::Compare(self.op),
            runtime: self.runtime.clone(),
            metadata: Mutex::new(BroadcastMetadataCache::new(self.runtime.clone())),
            last_capture_safe_signature: Mutex::new(None),
        }))
    }
}

/// Factory for a binary logical kernel.
pub struct LogicalFactory {
    pub op: LogicalOp,
    pub runtime: Arc<CudaRuntime>,
}

impl KernelFactory for LogicalFactory {
    fn create(&self, _node: &Node, _input_shapes: &[Vec<usize>]) -> Result<Box<dyn Kernel>> {
        Ok(Box::new(BinaryPredKernel {
            op_name: self.op.op_name(),
            entry: self.op.entry(),
            module: LOGICAL_MODULE,
            src: LOGICAL_SRC,
            kind: BinaryKind::LogicalBool,
            runtime: self.runtime.clone(),
            metadata: Mutex::new(BroadcastMetadataCache::new(self.runtime.clone())),
            last_capture_safe_signature: Mutex::new(None),
        }))
    }
}

/// NVRTC-backed binary predicate kernel producing a **Bool** output. Covers both
/// comparison (f32/i32/i64 operands, plus bool `Equal`) and logical (bool
/// operands) families via
/// [`BinaryKind`], with NumPy-style right-aligned broadcasting.
#[derive(Debug)]
pub struct BinaryPredKernel {
    op_name: &'static str,
    entry: &'static str,
    module: &'static str,
    src: &'static str,
    kind: BinaryKind,
    runtime: Arc<CudaRuntime>,
    /// Persistent broadcast metadata so a captured launch performs no per-step
    /// host allocation/upload/free/synchronize (the seam Marsten identified).
    metadata: Mutex<BroadcastMetadataCache>,
    /// The exact dtype/shape signature recorded by the most recent successful
    /// fixed-decode call. `Some` iff the op is currently capture-safe.
    last_capture_safe_signature: Mutex<Option<PredCaptureSignature>>,
}

/// The dtype + operand/broadcast shapes a captured predicate launch is pinned
/// to. Capture is admitted only while the live call matches this exactly.
#[derive(Clone, Debug, PartialEq, Eq)]
struct PredCaptureSignature {
    dtype: DataType,
    shapes: BroadcastMetadataKey,
}

impl BinaryPredKernel {
    fn run(&self, inputs: &[TensorView], outputs: &mut [TensorMut]) -> Result<()> {
        let mut last_signature = self.last_capture_safe_signature.lock().map_err(|_| {
            EpError::KernelFailed(
                "cuda_ep binary predicate capture signature lock was poisoned".into(),
            )
        })?;
        let warmed_signature = last_signature.take();
        let op = self.op_name;
        if inputs.len() != 2 || outputs.len() != 1 {
            return Err(EpError::KernelFailed(format!(
                "cuda_ep {op}: expected 2 inputs and 1 output, got {} and {}",
                inputs.len(),
                outputs.len()
            )));
        }
        let a = &inputs[0];
        let b = &inputs[1];
        let entry = match self.kind {
            BinaryKind::Compare(cmp) => {
                let Some(entry) = cmp.entry(a.dtype) else {
                    return Err(not_implemented(format!(
                        "{op}: operand dtype {:?} not supported on CUDA EP",
                        a.dtype
                    )));
                };
                require_dtype(op, "B", b.dtype, a.dtype)?;
                entry
            }
            BinaryKind::LogicalBool => {
                require_dtype(op, "A", a.dtype, DataType::Bool)?;
                require_dtype(op, "B", b.dtype, DataType::Bool)?;
                self.entry
            }
        };
        // Comparison and logical ops always emit Bool.
        require_dtype(op, "output", outputs[0].dtype, DataType::Bool)?;
        require_contiguous(op, "A", a.is_contiguous())?;
        require_contiguous(op, "B", b.is_contiguous())?;
        require_contiguous(op, "output", outputs[0].is_contiguous())?;

        let out_shape = onnx_runtime_ir::broadcast_shapes(a.shape, b.shape).map_err(EpError::Ir)?;
        if outputs[0].shape != out_shape {
            return Err(EpError::KernelFailed(format!(
                "cuda_ep {op}: output shape {:?} must equal broadcast shape {:?}",
                outputs[0].shape, out_shape
            )));
        }

        let n = outputs[0].numel();
        let n_u64 = count_u64(op, n)?;
        let a_ptr = cuptr(a.data_ptr::<u8>() as *const c_void);
        let b_ptr = cuptr(b.data_ptr::<u8>() as *const c_void);
        let y_ptr = cuptr(outputs[0].data_ptr_mut::<u8>() as *const c_void);

        // Pin capture eligibility to a shape whose broadcast metadata is stable
        // across steady decode. A predicate that drives control flow (e.g. the
        // LongRoPE `Greater(total_seq_len, threshold)` feeding an `If`) is a
        // rank-0 scalar or single-element tensor whose metadata is trivially
        // constant; a fixed-decode one-row output likewise qualifies. Only such
        // a signature, matching the warmed launch, may enter capture (mirroring
        // the elementwise binary kernel). The `If` consumer's branch flip is
        // separately guarded by the executor's control-flow seam invalidation.
        let capture_eligible =
            out_shape.iter().product::<usize>() == 1 || is_fixed_decode_shape(&out_shape);
        let current_signature = capture_eligible.then(|| PredCaptureSignature {
            dtype: a.dtype,
            shapes: BroadcastMetadataKey {
                a_shape: a.shape.to_vec(),
                b_shape: b.shape.to_vec(),
                out_shape: out_shape.clone(),
            },
        });
        require_matching_capture_signature(
            &self.runtime,
            op,
            warmed_signature.as_ref(),
            current_signature.as_ref(),
        )?;

        let func = self.runtime.nvrtc_function(self.module, self.src, entry)?;
        let mut metadata = self.metadata.lock().map_err(|_| {
            EpError::KernelFailed("cuda_ep binary predicate metadata lock was poisoned".into())
        })?;
        let metadata_ptr = metadata.prepare(a.shape, b.shape, &out_shape)?;
        let rank = i32::try_from(out_shape.len())
            .map_err(|_| EpError::KernelFailed(format!("cuda_ep {op}: rank exceeds i32")))?;
        let cfg = LaunchConfig {
            grid_dim: (grid_for(n), 1, 1),
            block_dim: (BLOCK, 1, 1),
            shared_mem_bytes: 0,
        };
        let stream = self.runtime.stream();
        let mut builder = stream.launch_builder(&func);
        builder
            .arg(&a_ptr)
            .arg(&b_ptr)
            .arg(&y_ptr)
            .arg(&metadata_ptr)
            .arg(&rank)
            .arg(&n_u64);
        // SAFETY: `func` is the compiled predicate entry; its argument list is
        // (const T*, const T*, unsigned char*, metadata, rank, count), where T
        // matches the validated same-type operands. All pointers cover their
        // respective allocations, with matching rank/count and indexing. The
        // metadata pointer is the persistent cache buffer, valid across replays.
        unsafe { builder.launch(cfg) }.map_err(|e| driver_err(&format!("launch {entry}"), e))?;
        *last_signature = current_signature;
        Ok(())
    }
}

impl Kernel for BinaryPredKernel {
    fn execute(&self, inputs: &[TensorView], outputs: &mut [TensorMut]) -> Result<()> {
        self.run(inputs, outputs)
    }

    fn supports_strided_input(&self, _idx: usize) -> bool {
        false
    }

    fn capture_support(&self) -> onnx_runtime_ep_api::CaptureSupport {
        // Only the exact fixed-row signature recorded by the most recent
        // successful call may enter capture; the persistent metadata cache means
        // the launch itself performs no per-step host alloc/upload/free/sync.
        match self.last_capture_safe_signature.lock() {
            Ok(signature) if signature.is_some() => onnx_runtime_ep_api::CaptureSupport::Supported,
            Ok(_) => onnx_runtime_ep_api::CaptureSupport::unsupported(format!(
                "{} broadcast shape/dtype signature does not match the warmed capture signature",
                self.op_name
            )),
            Err(_) => onnx_runtime_ep_api::CaptureSupport::unsupported(format!(
                "{} capture signature is unavailable because its state lock was poisoned",
                self.op_name
            )),
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn unary_math_entry_points_are_present_in_source() {
        for op in [
            UnaryMathOp::Abs,
            UnaryMathOp::Neg,
            UnaryMathOp::Reciprocal,
            UnaryMathOp::Exp,
            UnaryMathOp::Log,
            UnaryMathOp::Sign,
            UnaryMathOp::Floor,
            UnaryMathOp::Ceil,
            UnaryMathOp::Round,
            UnaryMathOp::Sin,
            UnaryMathOp::Cos,
            UnaryMathOp::Softplus,
        ] {
            assert!(
                UNARY_MATH_SRC.contains(&format!("DEFINE_UNARY({},", op.stem())),
                "missing NVRTC generator for {}",
                op.op_name()
            );
        }
    }

    #[test]
    fn cmp_entry_points_are_present_in_source() {
        for op in [
            CmpOp::Equal,
            CmpOp::Greater,
            CmpOp::Less,
            CmpOp::GreaterOrEqual,
            CmpOp::LessOrEqual,
        ] {
            let stem = op
                .entry(DataType::Float32)
                .unwrap()
                .strip_suffix("_f32")
                .unwrap();
            assert!(
                CMP_SRC.contains(&format!("DEFINE_CMP({stem}, type, suffix,")),
                "missing NVRTC generator for {}",
                op.op_name()
            );
        }
        for suffix in ["float, f32", "int, i32", "long long, i64"] {
            assert!(CMP_SRC.contains(suffix), "missing comparison type {suffix}");
        }
        assert_eq!(CmpOp::Equal.entry(DataType::Bool), Some("equal_bool"));
        assert_eq!(CmpOp::Greater.entry(DataType::Bool), None);
    }

    #[test]
    fn logical_entry_points_are_present_in_source() {
        for op in [LogicalOp::And, LogicalOp::Or, LogicalOp::Xor] {
            assert!(
                LOGICAL_SRC.contains(&format!("DEFINE_LOGICAL({},", op.entry())),
                "missing NVRTC entry {} for {}",
                op.entry(),
                op.op_name()
            );
        }
        assert!(NOT_SRC.contains("void not_bool("), "missing not_bool entry");
    }

    #[test]
    fn round_uses_ties_to_even_intrinsic() {
        // ONNX Round is round-half-to-even; `roundf` (half-away-from-zero) would
        // be wrong, so the kernel must use `rintf`.
        assert!(
            UNARY_MATH_SRC.contains("op_round(float x) { return rintf(x); }"),
            "Round must use rintf"
        );
        assert!(
            !UNARY_MATH_SRC.contains("roundf("),
            "Round must not use half-away-from-zero roundf"
        );
    }

    #[test]
    fn sign_handles_nan_and_zero_like_cpu() {
        // NaN -> NaN (v != v guard) and the zero case falls through to 0.0f.
        assert!(
            UNARY_MATH_SRC.contains("(x != x) ? x"),
            "sign must guard NaN"
        );
    }

    #[test]
    fn entry_points_are_all_distinct() {
        let mut seen = std::collections::HashSet::new();
        let unary = [
            UnaryMathOp::Abs,
            UnaryMathOp::Neg,
            UnaryMathOp::Reciprocal,
            UnaryMathOp::Exp,
            UnaryMathOp::Log,
            UnaryMathOp::Sign,
            UnaryMathOp::Floor,
            UnaryMathOp::Ceil,
            UnaryMathOp::Round,
            UnaryMathOp::Sin,
            UnaryMathOp::Cos,
            UnaryMathOp::Softplus,
        ]
        .map(|o| o.entry(FloatDtype::F32));
        let cmp = [
            CmpOp::Equal,
            CmpOp::Greater,
            CmpOp::Less,
            CmpOp::GreaterOrEqual,
            CmpOp::LessOrEqual,
        ]
        .map(|o| o.entry(DataType::Float32).unwrap());
        let logical = [LogicalOp::And, LogicalOp::Or, LogicalOp::Xor].map(|o| o.entry());
        for e in unary
            .into_iter()
            .chain(cmp.map(str::to_owned))
            .chain(logical.map(str::to_owned))
        {
            assert!(seen.insert(e.clone()), "duplicate entry point {e}");
        }
    }

    #[test]
    fn require_dtype_rejects_actionably() {
        let e = require_dtype("Exp", "input", DataType::Int64, DataType::Float32).unwrap_err();
        let msg = format!("{e}");
        assert!(msg.contains("Int64"), "{msg}");
        assert!(msg.contains("Float32"), "{msg}");
    }

    #[test]
    fn require_contiguous_rejects_strided_actionably() {
        let e = require_contiguous("And", "A", false).unwrap_err();
        let msg = format!("{e}");
        assert!(msg.contains("non-contiguous"), "{msg}");
        assert!(msg.contains("materialise"), "{msg}");
    }

    #[test]
    fn comparison_dtype_contract_is_actionable() {
        assert_eq!(
            comparison_unsupported_reason("Equal", &[DataType::Int64, DataType::Int64]),
            None
        );
        let reason =
            comparison_unsupported_reason("Greater", &[DataType::Bool, DataType::Bool]).unwrap();
        assert!(reason.contains("Bool"), "{reason}");
        let reason =
            comparison_unsupported_reason("Equal", &[DataType::Int32, DataType::Int64]).unwrap();
        assert!(reason.contains("same dtype"), "{reason}");
    }

    #[test]
    fn grid_covers_all_elements() {
        assert_eq!(grid_for(0), 1);
        assert_eq!(grid_for(1), 1);
        assert_eq!(grid_for(BLOCK as usize), 1);
        assert_eq!(grid_for(BLOCK as usize + 1), 2);
        assert_eq!(grid_for(usize::MAX / 2), 65_535);
    }

    #[test]
    fn near_i32_max_uses_u64_count_and_indexing() {
        let near_i32_max = (i32::MAX as usize) + 1;
        let count: u64 = count_u64("Exp", near_i32_max).unwrap();
        assert_eq!(count, (i32::MAX as u64) + 1);

        const LOOP: &str = "for (unsigned long long i = blockIdx.x*blockDim.x + threadIdx.x; i < n; i += (unsigned long long)gridDim.x * blockDim.x)";
        assert!(UNARY_MATH_SRC.contains("const unsigned long long n)"));
        assert!(UNARY_MATH_SRC.contains("for (unsigned long long i ="));
        assert!(UNARY_MATH_SRC.contains("i += (unsigned long long)gridDim.x * blockDim.x"));
        for (name, source, kernel_count) in [
            ("Not", NOT_SRC, 1),
            ("comparison macro", CMP_SRC, 1),
            ("logical macro", LOGICAL_SRC, 1),
        ] {
            assert_eq!(
                source.matches("const unsigned long long n)").count(),
                kernel_count,
                "{name} count parameters must be unsigned 64-bit"
            );
            assert_eq!(
                source.matches(LOOP).count(),
                kernel_count,
                "{name} kernels must use unsigned 64-bit grid-stride indexing"
            );
            assert!(
                !source.contains("const int n)") && !source.contains("for (int i"),
                "{name} source regressed to signed 32-bit indexing"
            );
        }
    }
}