# mypy: allow-untyped-defs
"""Packed interval mask analysis for flex flash attention.

Flex attention can evaluate arbitrary ``mask_mod`` functions lane by lane, but
SM100 packed masks can use a cheaper path when the kept KV lanes in each
32-lane chunk are window-like. This pass recognizes mask graphs where, after
substituting ``kv_idx + mask_lane`` for the per-lane KV index, the predicate can
be represented as a small union of half-open intervals over ``mask_lane``.

This works for causal, sliding-window, block, and document-bound masks whose
bounds depend only on lane-uniform values such as ``q_idx``, ``b_idx``, or aux
loads indexed by them. It intentionally rejects masks that need arbitrary
per-lane inclusion, such as ``doc_ids[b, kv_idx] == doc_ids[b, q_idx]`` or
``kv_idx % 2 == 0``; those would degenerate into many singleton intervals and
are better handled by the generic mask path.
"""

import dataclasses
from collections.abc import Mapping, Sequence
from itertools import product

import sympy

import torch
from torch.fx import GraphModule
from torch.utils._sympy.functions import FloorDiv, Max, Min

from ...codegen.cutedsl.lane_analysis import decompose_affine_lane_expr
from ...virtualized import NullHandler, V
from .aux_vectorization import fx_aux_index_to_sympy


# Fall back to scalar mask lowering before interval expansion makes generated CuTe too large.
MAX_PACKED_MASK_INTERVALS_FOR_CODE_SIZE = 8
LANE_UNIFORM_BINARY_OPS: Mapping[object, str] = {
    torch.ops.aten.add.Tensor: "+",
    torch.ops.aten.add.Scalar: "+",
    torch.ops.aten.sub.Tensor: "-",
    torch.ops.aten.sub.Scalar: "-",
    torch.ops.aten.mul.Tensor: "*",
    torch.ops.aten.mul.Scalar: "*",
    torch.ops.aten.remainder.Tensor: "%",
    torch.ops.aten.remainder.Scalar: "%",
}


@dataclasses.dataclass(frozen=True)
class PackedMaskInterval:
    """Half-open keep interval [lower_lane, upper_lane_exclusive) for a packed mask.

    Packed interval lowering replaces per-lane boolean mask evaluation with one
    32-bit keep mask when the mask predicate can be expressed as lane bounds.
    For example, ``q_idx >= kv_idx`` over a 32-lane KV window starting at
    ``kv_idx`` keeps lanes ``[0, q_idx - kv_idx + 1)``.
    """

    lower_lane: sympy.Expr
    upper_lane_exclusive: sympy.Expr
    # CuTe code for temporary symbols that stand in for lane-uniform aux loads.
    symbol_codes: Mapping[sympy.Symbol, str] = dataclasses.field(
        default_factory=dict, compare=False, repr=False
    )

    @classmethod
    def full(cls) -> "PackedMaskInterval":
        return cls(sympy.Integer(0), sympy.Integer(32))

    def is_full(self) -> bool:
        return self.lower_lane == 0 and self.upper_lane_exclusive == 32

    def with_symbol_codes(
        self, symbol_codes: Mapping[sympy.Symbol, str]
    ) -> "PackedMaskInterval":
        return dataclasses.replace(self, symbol_codes=dict(symbol_codes))

    def render_lower(self) -> str:
        return self._render(self.lower_lane)

    def render_upper(self) -> str:
        return self._render(self.upper_lane_exclusive)

    def keep_mask_expr(self) -> str:
        """Render the 32-bit CuTe keep-mask expression for this interval."""
        lower = self.render_lower()
        upper = self.render_upper()
        return (
            "(utils.shr_u32(cutlass.Uint32(0xFFFFFFFF), "
            "cutlass.Uint32(min(max(cutlass.Int32(32) - "
            f"{upper}, cutlass.Int32(0)), cutlass.Int32(32)))) & "
            "utils.shl_u32(cutlass.Uint32(0xFFFFFFFF), "
            "cutlass.Uint32(min(max("
            f"{lower}, cutlass.Int32(0)), cutlass.Int32(32)))))"
        )

    def _render(self, expr: sympy.Expr) -> str:
        rendered = sympy_to_cute_index(expr, self.symbol_codes)
        assert rendered is not None
        return rendered


IntervalSet = tuple[PackedMaskInterval, ...]
MaybeIntervalSet = IntervalSet | None


def is_bool_full_node(node: torch.fx.Node, value: bool) -> bool:
    return (
        node.op == "call_function"
        and node.target is torch.ops.aten.full.default
        and len(node.args) >= 2
        and node.args[0] == []
        and node.args[1] is value
    )


def _simplify_expr(expr: sympy.Expr) -> sympy.Expr:
    """Simplify with Inductor sizevars when available, otherwise use SymPy."""
    if isinstance(V.graph, NullHandler):
        return sympy.simplify(expr)
    return V.graph.sizevars.simplify(expr)


def sympy_to_cute_index(
    expr: sympy.Expr, symbol_codes: Mapping[sympy.Symbol, str] | None = None
) -> str | None:
    """Render a supported SymPy integer expression as CuTe index code."""
    expr = _simplify_expr(expr)
    match expr:
        case sympy.Integer():
            return f"cutlass.Int32({int(expr)})"
        case sympy.Symbol():
            if symbol_codes is not None and expr in symbol_codes:
                return symbol_codes[expr]
            if expr.name == "q_idx":
                return "q_idx[0]"
            if expr.name == "kv_idx":
                return "kv_idx[0]"
        case FloorDiv():
            lhs, rhs = (sympy_to_cute_index(arg, symbol_codes) for arg in expr.args)
            if lhs is not None and rhs is not None:
                return f"({lhs} // {rhs})"
        case sympy.Add():
            args = []
            for arg in expr.args:
                cute_arg = sympy_to_cute_index(arg, symbol_codes)
                if cute_arg is None:
                    return None
                args.append(cute_arg)
            return "(" + " + ".join(args) + ")"
        case sympy.Mul():
            args = []
            for arg in expr.args:
                cute_arg = sympy_to_cute_index(arg, symbol_codes)
                if cute_arg is None:
                    return None
                args.append(cute_arg)
            return "(" + " * ".join(args) + ")"
        case Min() | Max():
            args = []
            for arg in expr.args:
                cute_arg = sympy_to_cute_index(arg, symbol_codes)
                if cute_arg is None:
                    return None
                args.append(cute_arg)
            op = "min" if isinstance(expr, Min) else "max"
            return f"{op}(" + ", ".join(args) + ")"
    if isinstance(expr, (sympy.Min, sympy.Max)):
        args = []
        for arg in expr.args:
            cute_arg = sympy_to_cute_index(arg, symbol_codes)
            if cute_arg is None:
                return None
            args.append(cute_arg)
        op = "min" if isinstance(expr, sympy.Min) else "max"
        return f"{op}(" + ", ".join(args) + ")"
    return None


def is_aten_index_node(node: torch.fx.Node) -> bool:
    return node.op == "call_function" and node.target is torch.ops.aten.index.Tensor


def fx_node_dtype(expr: torch.fx.Node) -> torch.dtype | None:
    """Read an FX node dtype from metadata, following aten.index bases if needed."""
    tensor_meta = expr.meta.get("tensor_meta")
    if tensor_meta is not None:
        return tensor_meta.dtype
    val = expr.meta.get("val")
    if isinstance(val, torch.Tensor):
        return val.dtype
    if is_aten_index_node(expr):
        base = expr.args[0]
        if isinstance(base, torch.fx.Node):
            return fx_node_dtype(base)
    return None


def fx_node_shape(expr: torch.fx.Node) -> tuple[int, ...] | None:
    """Read an FX node shape from tensor metadata or fake tensor value."""
    tensor_meta = expr.meta.get("tensor_meta")
    if tensor_meta is not None:
        return tuple(tensor_meta.shape)
    val = expr.meta.get("val")
    if isinstance(val, torch.Tensor):
        return tuple(val.shape)
    return None


@dataclasses.dataclass
class PackedMaskAnalyzer:
    """Lower one Flex mask_mod FX graph to packed 32-lane mask intervals.

    The analyzer treats each packed mask as a 32-lane window starting at kv_idx,
    rewrites supported mask predicates into bounds over mask_lane, and records
    renderable code for lane-uniform aux loads used by those bounds. For example,
    a document mask like ``doc_offsets[b_idx] <= kv_idx and kv_idx < q_idx + 128``
    is analyzed over the packed lane expression ``kv_idx + mask_lane`` and
    becomes the interval ``[doc_offsets[b_idx] - kv_idx, q_idx + 128 - kv_idx)``.

    placeholders: FX placeholders for the mask_mod ABI and captured inputs.
    q_idx: FX placeholder for the scalar query index.
    kv_idx: FX placeholder for the first KV index in the packed 32-lane window.
    q_symbol: SymPy symbol used for q_idx during interval analysis.
    kv_symbol: SymPy symbol used for the packed window base KV index.
    lane_symbol: SymPy symbol for the lane offset within the packed window.
    symbol_codes: CuTe renderings for temporary symbols introduced for aux loads.
    aux_load_symbols: Stable symbols assigned to supported lane-uniform aux loads.
    next_symbol_id: Counter for naming temporary aux-load symbols.
    """

    placeholders: Sequence[torch.fx.Node]
    q_idx: torch.fx.Node
    kv_idx: torch.fx.Node
    q_symbol: sympy.Symbol = dataclasses.field(
        default_factory=lambda: sympy.Symbol("q_idx", integer=True, nonnegative=True)
    )
    kv_symbol: sympy.Symbol = dataclasses.field(
        default_factory=lambda: sympy.Symbol("kv_idx", integer=True, nonnegative=True)
    )
    lane_symbol: sympy.Symbol = dataclasses.field(
        default_factory=lambda: sympy.Symbol(
            "mask_lane", integer=True, nonnegative=True
        )
    )
    symbol_codes: dict[sympy.Symbol, str] = dataclasses.field(default_factory=dict)
    aux_load_symbols: dict[torch.fx.Node, sympy.Symbol] = dataclasses.field(
        default_factory=dict
    )
    next_symbol_id: int = 0

    def node_to_intervals(self, node: torch.fx.Node) -> MaybeIntervalSet:
        """Lower one supported boolean FX node into packed mask intervals."""
        if is_bool_full_node(node, True):
            return (PackedMaskInterval.full(),)
        if is_bool_full_node(node, False):
            return ()
        if node.op != "call_function":
            return None
        match node.target:
            case torch.ops.aten.bitwise_and.Tensor | torch.ops.aten.logical_and.default:
                return self.combine_binary_intervals(node, intersect=True)
            case torch.ops.aten.bitwise_or.Tensor | torch.ops.aten.logical_or.default:
                return self.combine_binary_intervals(node, intersect=False)
            case torch.ops.aten.le.Tensor | torch.ops.aten.le.Scalar:
                return self.comparison_to_intervals(
                    node.args[0], node.args[1], strict=False
                )
            case torch.ops.aten.lt.Tensor | torch.ops.aten.lt.Scalar:
                return self.comparison_to_intervals(
                    node.args[0], node.args[1], strict=True
                )
            case torch.ops.aten.ge.Tensor | torch.ops.aten.ge.Scalar:
                return self.comparison_to_intervals(
                    node.args[1], node.args[0], strict=False
                )
            case torch.ops.aten.gt.Tensor | torch.ops.aten.gt.Scalar:
                return self.comparison_to_intervals(
                    node.args[1], node.args[0], strict=True
                )
            case torch.ops.aten.eq.Tensor | torch.ops.aten.eq.Scalar:
                return self.equality_to_intervals(node.args[0], node.args[1])
            case _:
                return None

    def combine_binary_intervals(
        self, node: torch.fx.Node, *, intersect: bool
    ) -> MaybeIntervalSet:
        lhs, rhs = node.args
        if not isinstance(lhs, torch.fx.Node) or not isinstance(rhs, torch.fx.Node):
            return None
        lhs_intervals = self.node_to_intervals(lhs)
        rhs_intervals = self.node_to_intervals(rhs)
        if lhs_intervals is None or rhs_intervals is None:
            return None
        if intersect:
            return self.merge_intersection(lhs_intervals, rhs_intervals)
        if (
            len(lhs_intervals) + len(rhs_intervals)
            > MAX_PACKED_MASK_INTERVALS_FOR_CODE_SIZE
        ):
            return None
        return lhs_intervals + rhs_intervals

    def comparison_to_intervals(
        self,
        lhs: object,
        rhs: object,
        *,
        strict: bool,
    ) -> MaybeIntervalSet:
        exprs = self.mask_operands_to_sympy(lhs, rhs)
        if exprs is None:
            return None
        return self.lane_comparison_to_intervals(*exprs, strict=strict)

    def equality_to_intervals(self, lhs: object, rhs: object) -> MaybeIntervalSet:
        exprs = self.mask_operands_to_sympy(lhs, rhs)
        if exprs is None:
            return None
        return self.lane_equality_to_intervals(*exprs)

    def mask_operands_to_sympy(
        self, lhs: object, rhs: object
    ) -> tuple[sympy.Expr, sympy.Expr] | None:
        lhs_expr = self.fx_mask_expr_to_sympy(lhs)
        rhs_expr = self.fx_mask_expr_to_sympy(rhs)
        if lhs_expr is None or rhs_expr is None:
            return None
        return lhs_expr, rhs_expr

    def lane_comparison_to_intervals(
        self,
        lhs_expr: sympy.Expr,
        rhs_expr: sympy.Expr,
        *,
        strict: bool,
    ) -> MaybeIntervalSet:
        """Convert a lane-affine comparison into one packed keep interval."""
        diff = V.graph.sizevars.simplify(lhs_expr - rhs_expr)
        affine = decompose_affine_lane_expr(diff, self.lane_symbol)
        if affine is None:
            return None
        lane_coeff, rest = affine
        if lane_coeff == 1:
            upper = V.graph.sizevars.simplify(-rest if strict else -rest + 1)
            return self.interval_if_renderable(sympy.Integer(0), upper)
        if lane_coeff == -1:
            lower = V.graph.sizevars.simplify(rest + 1 if strict else rest)
            return self.interval_if_renderable(lower, sympy.Integer(32))
        return None

    def lane_equality_to_intervals(
        self, lhs_expr: sympy.Expr, rhs_expr: sympy.Expr
    ) -> MaybeIntervalSet:
        """Convert a lane-affine equality into singleton keep intervals."""
        if isinstance(lhs_expr, FloorDiv) and isinstance(rhs_expr, FloorDiv):
            intervals = self._floor_div_equality_to_intervals(lhs_expr, rhs_expr)
            if intervals is not None:
                return intervals
        diff = V.graph.sizevars.simplify(lhs_expr - rhs_expr)
        affine = decompose_affine_lane_expr(diff, self.lane_symbol)
        if affine is None:
            return None
        lane_coeff, rest = affine
        if lane_coeff in (1, -1):
            lane_value = -rest if lane_coeff == 1 else rest
            lower = V.graph.sizevars.simplify(lane_value)
            upper = V.graph.sizevars.simplify(lane_value + 1)
            return self.interval_if_renderable(lower, upper)
        return None

    def _floor_div_equality_to_intervals(
        self, lhs_expr: FloorDiv, rhs_expr: FloorDiv
    ) -> MaybeIntervalSet:
        """Convert matching block-id equality into a contiguous lane interval."""
        lhs_base, lhs_divisor = lhs_expr.args
        rhs_base, rhs_divisor = rhs_expr.args
        if lhs_divisor != rhs_divisor:
            return None
        block_start = None
        if lhs_base == self.q_symbol and rhs_base == self.kv_symbol + self.lane_symbol:
            block_start = V.graph.sizevars.simplify(
                lhs_expr * lhs_divisor - self.kv_symbol
            )
        elif (
            rhs_base == self.q_symbol and lhs_base == self.kv_symbol + self.lane_symbol
        ):
            block_start = V.graph.sizevars.simplify(
                rhs_expr * rhs_divisor - self.kv_symbol
            )
        if block_start is None:
            return None
        lower = V.graph.sizevars.simplify(block_start)
        upper = V.graph.sizevars.simplify(block_start + lhs_divisor)
        return self.interval_if_renderable(lower, upper)

    def interval_if_renderable(
        self, lower: sympy.Expr, upper: sympy.Expr
    ) -> MaybeIntervalSet:
        if (
            sympy_to_cute_index(lower, self.symbol_codes) is not None
            and sympy_to_cute_index(upper, self.symbol_codes) is not None
        ):
            return (PackedMaskInterval(lower, upper),)
        return None

    def merge_intersection(
        self,
        intervals: IntervalSet,
        new_intervals: IntervalSet,
    ) -> MaybeIntervalSet:
        """Intersect two interval sets, falling back if code size would grow too much."""
        if (
            len(intervals) * len(new_intervals)
            > MAX_PACKED_MASK_INTERVALS_FOR_CODE_SIZE
        ):
            return None
        return tuple(
            PackedMaskInterval(
                V.graph.sizevars.simplify(Max(lhs.lower_lane, rhs.lower_lane)),
                V.graph.sizevars.simplify(
                    Min(lhs.upper_lane_exclusive, rhs.upper_lane_exclusive)
                ),
            )
            for lhs, rhs in product(intervals, new_intervals)
        )

    def fx_mask_expr_to_sympy(self, expr: object) -> sympy.Expr | None:
        """Translate a mask expression with kv_idx expanded by mask_lane."""
        return fx_aux_index_to_sympy(
            expr,
            {
                self.q_idx: self.q_symbol,
                self.kv_idx: self.kv_symbol + self.lane_symbol,
            },
            self.mask_aux_load_to_symbol,
        )

    def mask_aux_load_to_symbol(self, node: torch.fx.Node) -> sympy.Expr | None:
        """Assign a temporary SymPy symbol to a supported aux tensor load."""
        if not is_aten_index_node(node):
            return None
        if node in self.aux_load_symbols:
            return self.aux_load_symbols[node]
        lane_uniform_code = self.render_lane_uniform_scalar_expr(node)
        if lane_uniform_code is None:
            return None
        symbol = sympy.Symbol(f"mask_bound_{self.next_symbol_id}", integer=True)
        self.next_symbol_id += 1
        self.symbol_codes[symbol] = lane_uniform_code
        self.aux_load_symbols[node] = symbol
        return symbol

    def render_lane_uniform_scalar_expr(
        self,
        expr: object,
        *,
        for_index: bool = False,
        index_dim_size: int | sympy.Expr | None = None,
    ) -> str | None:
        """Render an expression that is invariant across the packed KV lanes."""
        if isinstance(expr, int | sympy.Integer):
            index = int(expr)
            if for_index and index < 0:
                if index_dim_size is None:
                    return None
                index = V.graph.sizevars.guard_int(index + index_dim_size)
            return f"cutlass.Int32({index})"
        if not isinstance(expr, torch.fx.Node):
            return None

        if expr is self.kv_idx:
            return None
        if expr is self.q_idx:
            return "q_idx[0]"
        if len(self.placeholders) >= 2:
            if expr is self.placeholders[0]:
                return "b_idx[0]"
            if expr is self.placeholders[1]:
                return "h_idx[0]"
        for aux_idx, placeholder in enumerate(self.placeholders[4:]):
            if expr is placeholder:
                return f"aux_tensors[{aux_idx}]"

        if is_aten_index_node(expr):
            return self._render_lane_uniform_index_expr(expr, for_index=for_index)
        if expr.op != "call_function":
            return None
        match expr.target:
            case torch.ops.aten.div.Tensor_mode:
                if expr.kwargs.get("rounding_mode") != "floor":
                    return None
                op = "//"
            case _:
                op = LANE_UNIFORM_BINARY_OPS.get(expr.target)
                if op is None:
                    return None
        args = expr.args
        if len(args) < 2:
            return None
        lhs = self.render_lane_uniform_scalar_expr(args[0])
        rhs = self.render_lane_uniform_scalar_expr(args[1])
        if lhs is None or rhs is None:
            return None
        return f"({lhs} {op} {rhs})"

    def _render_lane_uniform_index_expr(
        self,
        expr: torch.fx.Node,
        *,
        for_index: bool,
    ) -> str | None:
        """Render a scalar integer aten.index load that is uniform across lanes."""
        if fx_node_dtype(expr) not in (
            torch.int8,
            torch.int16,
            torch.int32,
            torch.int64,
            torch.uint8,
        ):
            return None
        result_shape = fx_node_shape(expr)
        if result_shape is not None and len(result_shape) != 0:
            return None
        base, indices = expr.args
        base_code = self.render_lane_uniform_scalar_expr(base)
        base_shape = fx_node_shape(base) if isinstance(base, torch.fx.Node) else None
        if base_code is None or not isinstance(indices, (list, tuple)):
            return None
        if base_shape is None or len(indices) != len(base_shape):
            return None
        index_codes = []
        for dim, index in enumerate(indices):
            dim_size = base_shape[dim]
            index_code = self.render_lane_uniform_scalar_expr(
                index, for_index=True, index_dim_size=dim_size
            )
            if index_code is None:
                return None
            if (
                dim_size is not None
                and isinstance(index, torch.fx.Node)
                and index
                not in (self.q_idx, self.placeholders[0], self.placeholders[1])
            ):
                if not isinstance(dim_size, int | sympy.Integer):
                    return None
                dtype = fx_node_dtype(index)
                integer_type = "Int64" if dtype == torch.int64 else "Int32"
                size_code = f"cutlass.{integer_type}({dim_size})"
                zero_code = f"cutlass.{integer_type}(0)"
                index_code = (
                    f"({index_code} + {size_code} "
                    f"if {index_code} < {zero_code} else {index_code})"
                )
            index_codes.append(index_code)
        load = f"{base_code}[{', '.join(index_codes)}]"
        if for_index and fx_node_dtype(expr) == torch.int64:
            return load
        return f"cutlass.Int32({load})"


def select_packed_mask_intervals(
    graph_module: GraphModule,
) -> MaybeIntervalSet:
    """Entry point for selecting packed mask intervals from a mask_mod graph.

    Returns ``None`` when the graph is unsupported, when interval packing would
    increase code size too much, or when analysis reduces the whole mask to the
    unconditional ``[0, 32)`` interval.
    """
    graph = graph_module.graph
    nodes = list(graph.nodes)
    placeholders = [node for node in nodes if node.op == "placeholder"]
    output = [node for node in nodes if node.op == "output"]
    if len(placeholders) < 4 or len(output) != 1:
        return None

    q_idx, kv_idx = placeholders[2], placeholders[3]
    output_val = output[0].args[0]
    if not isinstance(output_val, torch.fx.Node):
        return None

    analyzer = PackedMaskAnalyzer(
        placeholders=placeholders,
        q_idx=q_idx,
        kv_idx=kv_idx,
    )
    intervals = analyzer.node_to_intervals(output_val)
    if intervals is None:
        return None
    if len(intervals) == 1 and intervals[0].is_full():
        return None
    return tuple(
        interval.with_symbol_codes(analyzer.symbol_codes) for interval in intervals
    )
