# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.
# pylint: disable=redefined-builtin, invalid-name, too-many-arguments, too-many-locals, too-many-positional-arguments
"""PTX MMA / ldmatrix / stmatrix intrinsics.

mma.sync.aligned has 7 form_kinds per the PTX docs (f16 / tf32 / bf16 / fp64
/ int8 / fp8 / subbyte). Each form_kind is one ``device_intrinsic`` registration;
the (shape, layouts, dtypes) modifier slots are attrs. Body computes the per-
fragment register counts at codegen time from M*N*bits/threads/frag_size and
hand-builds the asm constraint list.

ldmatrix / stmatrix each have a single PTX form (the .m8n8 .b16/.b8 variant
that TIRx uses); ``num`` and ``trans`` are modifier attrs.
"""

import re
from dataclasses import dataclass

from tvm import DataType

from ._schema import device_intrinsic
from .registry import CODEGEN_REGISTRY, register_codegen
from .types import PTXDataType
from .utils import parse_str


@dataclass
class FragAttrs:
    reg_type: str  # asm constraint letter (r / f / d)
    size: int  # bit width per register slot (32 or 64)
    ptr_type: str  # C type for the cast


_FRAG_ATTRS_MAP = {
    PTXDataType.BIT1: FragAttrs("r", 32, "uint32_t"),
    PTXDataType.INT4: FragAttrs("r", 32, "uint32_t"),
    PTXDataType.UINT4: FragAttrs("r", 32, "uint32_t"),
    PTXDataType.INT8: FragAttrs("r", 32, "uint32_t"),
    PTXDataType.UINT8: FragAttrs("r", 32, "uint32_t"),
    PTXDataType.FLOAT8_E4M3FN: FragAttrs("r", 32, "uint32_t"),
    PTXDataType.FLOAT8_E5M2: FragAttrs("r", 32, "uint32_t"),
    PTXDataType.BIT16: FragAttrs("r", 32, "uint32_t"),
    PTXDataType.FLOAT16: FragAttrs("r", 32, "uint32_t"),
    PTXDataType.BFLOAT16: FragAttrs("r", 32, "uint32_t"),
    PTXDataType.TENSOR_FLOAT32: FragAttrs("r", 32, "uint32_t"),
    PTXDataType.INT32: FragAttrs("r", 32, "int32_t"),
    PTXDataType.FLOAT32: FragAttrs("f", 32, "float"),
    PTXDataType.FLOAT64: FragAttrs("d", 64, "double"),
}


def _parse_mma_shape(shape_str):
    match = re.search(r"m(\d+)n(\d+)k(\d+)", shape_str)
    if not match:
        raise ValueError(f"Cannot parse MMA shape: {shape_str!r}")
    return tuple(map(int, match.groups()))


def _classify_mma_form(d_type, a_type, b_type):
    """Map (d, a, b) dtype triple to one of the 7 PTX form_kind tags."""
    fp16 = {"float16", "fp16"}
    tf32 = {"tensor_float32", "tf32"}
    bf16 = {"bfloat16", "bf16"}
    fp64 = {"float64", "fp64"}
    int_a = {"int8", "uint8", "s8", "u8"}
    fp8 = {"e4m3", "e5m2", "float8_e4m3fn", "float8_e4m3fnuz", "float8_e5m2"}
    subbyte = {"int4", "uint4", "bit1", "s4", "u4", "b1", "int1", "uint1"}
    if a_type in fp16 and b_type in fp16:
        return "f16"
    if a_type in tf32 and b_type in tf32:
        return "tf32"
    if a_type in bf16 and b_type in bf16:
        return "bf16"
    if a_type in fp64 and b_type in fp64:
        return "fp64"
    if a_type in int_a and b_type in int_a:
        return "int8"
    if a_type in fp8 and b_type in fp8:
        return "fp8"
    if a_type in subbyte and b_type in subbyte:
        return "subbyte"
    raise ValueError(
        f"Unknown ptx.mma form for d_type={d_type!r}, a_type={a_type!r}, b_type={b_type!r}"
    )


def _frag(dtype_str):
    return _FRAG_ATTRS_MAP[PTXDataType.from_string(dtype_str)]


def _mma_threads(shape, a_type):
    """Special case: m8n8k4 with f16 a/b uses 8 threads per fragment."""
    m, n, k = _parse_mma_shape(shape)
    if m == 8 and n == 8 and k == 4 and a_type == "float16":
        return 8
    return 32


# PTX dtype abbreviation -> element bit width. Used by _frag_count so that
# callers passing the PTX abbreviation (e.g. "fp32") don't blow up in
# ``DataType("fp32")``.
_PTX_BITS = {
    "fp16": 16,
    "fp32": 32,
    "fp64": 64,
    "bf16": 16,
    "tf32": 32,  # tensor-float32 packs 19 significant bits into a 32-bit slot
    "s8": 8,
    "u8": 8,
    "s32": 32,
    "s4": 4,
    "u4": 4,
    "b1": 1,
    "b16": 16,
    "e4m3": 8,
    "e5m2": 8,
}


def _frag_count(dtype, dim_a, dim_b, threads):
    if dtype in _PTX_BITS:
        bits = _PTX_BITS[dtype]
    else:
        bits = DataType(dtype).bits
    size = _frag(dtype).size
    return dim_a * dim_b * bits // threads // size


# =============================================================================
# Shared helpers for the 7 mma form_kinds.
# Args layout for each form:
#   (d_ptr_in, a_ptr_in, b_ptr_in [, c_ptr_in], shape, a_layout, b_layout,
#    d_type, a_type, b_type, c_type, no_c_ptr [, saturate or bit_op])
# n_attrs = 8 for f16/tf32/bf16/fp64/fp8 (last 8 = shape, layouts, 4 dtypes, no_c_ptr)
# n_attrs = 9 for int8 (+ saturate) and subbyte (+ bit_op)
# =============================================================================


def _mma_form_parts(args, *, has_saturate=False, has_bit_op=False):
    """Compute (helper_name, c_signature, body) for one mma form invocation.

    ``args`` is the full positional arg tuple as received by codegen.
    The trailing ``n_attrs`` (8 or 9) entries are attrs.
    """
    n_extra = (1 if has_saturate else 0) + (1 if has_bit_op else 0)
    n_attrs = 8 + n_extra
    # Split off attr args from the tail (operand args are ahead).
    attrs = args[-n_attrs:]
    shape = parse_str(attrs[0])
    a_layout = parse_str(attrs[1])
    b_layout = parse_str(attrs[2])
    d_type = parse_str(attrs[3])
    a_type = parse_str(attrs[4])
    b_type = parse_str(attrs[5])
    c_type = parse_str(attrs[6])
    no_c_ptr_raw = attrs[7]
    no_c_ptr = bool(int(no_c_ptr_raw)) if hasattr(no_c_ptr_raw, "value") else bool(no_c_ptr_raw)
    saturate = False
    bit_op = ""
    if has_saturate:
        s = attrs[8]
        saturate = bool(int(s)) if hasattr(s, "value") else bool(s)
    if has_bit_op:
        bit_op = parse_str(attrs[8])

    # Fragment counts (same derivation as the contiguous form).
    m, n, k = _parse_mma_shape(shape)
    threads = _mma_threads(shape, a_type)
    d_cnt = _frag_count(d_type, m, n, threads)
    a_cnt = _frag_count(a_type, m, k, threads)
    b_cnt = _frag_count(b_type, k, n, threads)
    c_cnt = _frag_count(c_type, m, n, threads)

    # C signature: one void* per register, ordered D regs, A regs, B regs,
    # then C regs (only when the accumulator is used).
    sig_parts = (
        [f"void* d_ptr{i}" for i in range(d_cnt)]
        + [f"void* a_ptr{i}" for i in range(a_cnt)]
        + [f"void* b_ptr{i}" for i in range(b_cnt)]
    )
    if not no_c_ptr:
        sig_parts += [f"void* c_ptr{i}" for i in range(c_cnt)]
    sig = "(" + ", ".join(sig_parts) + ")"

    def _safe(s):
        return s.replace("::", "_").replace(".", "_")

    name = (
        f"ptx_mma_{shape}_{a_layout}_{b_layout}"
        f"_{_safe(d_type)}_{_safe(a_type)}_{_safe(b_type)}_{_safe(c_type)}"
        f"{'_no_c_ptr' if no_c_ptr else ''}"
        f"{'_saturate' if saturate else ''}"
    )

    d_frag = _frag(d_type)
    a_frag = _frag(a_type)
    b_frag = _frag(b_type)
    c_frag = _frag(c_type)

    saturate_inst = ".satfinite" if saturate else ""
    # PTX b1 mma requires a `.popc` suffix after the bit op (e.g. `.xor.popc`).
    bit_op_inst = f".{bit_op}.popc" if bit_op else ""

    d_type_inst = PTXDataType.from_string(d_type).to_string()
    c_type_inst = PTXDataType.from_string(c_type).to_string()
    a_type_inst = PTXDataType.from_string(a_type).to_string()
    b_type_inst = PTXDataType.from_string(b_type).to_string()

    def _slot_arr(start, cnt):
        return "{" + ", ".join(f"%{start + i}" for i in range(cnt)) + "}"

    args_template = (
        f"{_slot_arr(0, d_cnt)}, {_slot_arr(d_cnt, a_cnt)}, "
        f"{_slot_arr(d_cnt + a_cnt, b_cnt)}, {_slot_arr(d_cnt + a_cnt + b_cnt, c_cnt)}"
    )

    # Each register binds to its OWN pointer via *(T*)X_ptrN (scatter).
    d_outs = ", ".join(
        f'"=r"(*({d_frag.ptr_type}*)d_ptr{i})'
        if d_frag.reg_type == "r"
        else f'"={d_frag.reg_type}"(*({d_frag.ptr_type}*)d_ptr{i})'
        for i in range(d_cnt)
    )
    a_inputs = ", ".join(
        f'"{a_frag.reg_type}"(*({a_frag.ptr_type}*)a_ptr{i})' for i in range(a_cnt)
    )
    b_inputs = ", ".join(
        f'"{b_frag.reg_type}"(*({b_frag.ptr_type}*)b_ptr{i})' for i in range(b_cnt)
    )
    if no_c_ptr:
        c_value = "0.f" if c_frag.reg_type == "f" else "0"
        c_inputs = ", ".join(f'"{c_frag.reg_type}"({c_value})' for _ in range(c_cnt))
    else:
        c_inputs = ", ".join(
            f'"{c_frag.reg_type}"(*({c_frag.ptr_type}*)c_ptr{i})' for i in range(c_cnt)
        )

    body = (
        "    asm volatile(\n"
        f'        "mma.sync.aligned.{shape}.{a_layout}.{b_layout}{saturate_inst}'
        f'{d_type_inst}{a_type_inst}{b_type_inst}{c_type_inst}{bit_op_inst} "\n'
        f'        "{args_template};\\n"\n'
        f"        : {d_outs}\n"
        f"        : {a_inputs}, {b_inputs}, {c_inputs}\n"
        "    );"
    )
    return name, sig, body


def _register_mma_form(form_kind, *, has_saturate=False, has_bit_op=False):
    n_attrs = 8 + (1 if has_saturate else 0) + (1 if has_bit_op else 0)

    def _parts(*args, hs=has_saturate, hb=has_bit_op):
        return _mma_form_parts(args, has_saturate=hs, has_bit_op=hb)

    device_intrinsic(
        f"_ptx_mma_{form_kind}",
        n_attrs=n_attrs,
        helper_name=lambda *a: _parts(*a)[0],
        c_signature=lambda *a: _parts(*a)[1],
        body=lambda *a: _parts(*a)[2],
    )


# Form 1 — f16. Form 2 — tf32. Form 3 — bf16. Form 4 — fp64. Form 6 — fp8.
# All share the same 8-attr layout (no saturate / bit_op).
for _kind in ("f16", "tf32", "bf16", "fp64", "fp8"):
    _register_mma_form(_kind)
del _kind

# Form 5 — int8 (+ saturate).
_register_mma_form("int8", has_saturate=True)

# Form 7 — subbyte (+ bit_op for b1).
_register_mma_form("subbyte", has_bit_op=True)


@register_codegen("ptx_mma")
def codegen_ptx_mma(
    shape,
    a_layout,
    b_layout,
    d_type,
    a_type,
    b_type,
    c_type,
    d_cnt,
    a_cnt,
    b_cnt,
    c_cnt,
    no_c_ptr,
    *rest,
):
    """Classify (d, a, b) dtype triple to one of 7 form_kinds and forward.

    ``rest`` = flattened per-register pointers (d_cnt + a_cnt + b_cnt + c_cnt of
    them) followed by ``saturate`` and optionally ``bit_op``.
    """
    shape = parse_str(shape)
    a_layout = parse_str(a_layout)
    b_layout = parse_str(b_layout)
    d_type = parse_str(d_type)
    a_type = parse_str(a_type)
    b_type = parse_str(b_type)
    c_type = parse_str(c_type)
    d_cnt = int(d_cnt)
    a_cnt = int(a_cnt)
    b_cnt = int(b_cnt)
    c_cnt = int(c_cnt)
    no_c_ptr = bool(int(no_c_ptr)) if hasattr(no_c_ptr, "value") else bool(no_c_ptr)

    n_ptrs = d_cnt + a_cnt + b_cnt + (0 if no_c_ptr else c_cnt)
    ptrs = list(rest[:n_ptrs])
    trailing = list(rest[n_ptrs:])
    saturate = bool(trailing[0]) if trailing else False
    bit_op_v = ""
    if len(trailing) >= 2:
        bo = trailing[1]
        bit_op_v = parse_str(bo) if isinstance(bo, str) else (bo if bo is not None else "")

    kind = _classify_mma_form(d_type, a_type, b_type)

    # op_args are the flattened per-register pointers (already in PTX order:
    # D regs, A regs, B regs, then C regs unless no_c_ptr).
    op_args = ptrs

    attr_args = [shape, a_layout, b_layout, d_type, a_type, b_type, c_type, no_c_ptr]
    if kind == "int8":
        attr_args.append(saturate)
    elif kind == "subbyte":
        attr_args.append(bit_op_v)

    result = CODEGEN_REGISTRY[f"tirx._ptx_mma_{kind}"](op_args + attr_args)
    return result[0] if isinstance(result, tuple) else result


# =============================================================================
# ldmatrix / stmatrix — m8n8 fragment load/store. PTX docs lists 3 ldmatrix
# forms (m8n8 + m8n16 + m16n16); TIRx uses only the m8n8 form. 1
# device_intrinsic each. ``num`` (.x1/.x2/.x4) and ``trans`` are modifier
# attrs; the asm body loops over per-register constraints based on
# (num, dtype).
# =============================================================================


def _ldmatrix_parts(*args):
    # args = (smem_ptr, dst0, dst1, ..., dst{N-1}, num, dtype, trans)
    # The last 3 entries are the codegen attrs (n_attrs=3).
    num = int(args[-3])
    dtype = parse_str(args[-2])
    trans_b = bool(int(args[-1])) if hasattr(args[-1], "value") else bool(args[-1])
    if num not in (1, 2, 4):
        raise ValueError(f"ldmatrix .num must be one of {{1, 2, 4}}, got {num}")
    if dtype not in ("b16", "b8"):
        raise ValueError(f"ldmatrix dtype must be 'b16' or 'b8', got {dtype!r}")
    n_regs = num if dtype == "b16" else num // 2
    trans_inst = ".trans" if trans_b else ""
    slot_list = "{" + ", ".join(f"%{i}" for i in range(n_regs)) + "}"
    reg_decls = ", ".join(f"r{i}" for i in range(n_regs))
    out_constraints = ", ".join(f'"=r"(r{i})' for i in range(n_regs))
    dst_assigns = "\n".join(f"    *(uint32_t*)dst{i} = r{i};" for i in range(n_regs))
    name = f"ptx_ldmatrix_{num}_{dtype.replace('::', '_').replace('.', '_')}_{1 if trans_b else 0}"
    sig = "(void* smem_ptr, " + ", ".join(f"void* dst{i}" for i in range(n_regs)) + ")"
    body = (
        f"    uint32_t {reg_decls};\n"
        "    unsigned int addr = __cvta_generic_to_shared(smem_ptr);\n"
        "    asm volatile(\n"
        f'        "ldmatrix.sync.aligned.m8n8.x{num}{trans_inst}.shared.{dtype} '
        f'{slot_list}, [%{n_regs}];"\n'
        f"        : {out_constraints}\n"
        f'        : "r"(addr));\n'
        f"{dst_assigns}"
    )
    return name, sig, body


device_intrinsic(
    "_ptx_ldmatrix_impl",
    n_attrs=3,
    c_signature=lambda *a: _ldmatrix_parts(*a)[1],
    helper_name=lambda *a: _ldmatrix_parts(*a)[0],
    body=lambda *a: _ldmatrix_parts(*a)[2],
)


@register_codegen("ptx_ldmatrix")
def codegen_ptx_ldmatrix(trans, num, dtype, smem_ptr, *dst_handles):
    trans = bool(trans)
    num = int(num)
    dtype = parse_str(dtype)
    if dtype.startswith("."):
        dtype = dtype[1:]
    n_regs = num if dtype == "b16" else num // 2
    if len(dst_handles) != n_regs:
        raise ValueError(
            f"ldmatrix .x{num}.{dtype} codegen expects {n_regs} dst handles, got {len(dst_handles)}"
        )
    result = CODEGEN_REGISTRY["tirx._ptx_ldmatrix_impl"](
        [smem_ptr, *dst_handles, num, dtype, trans]
    )
    return result[0] if isinstance(result, tuple) else result


def _stmatrix_parts(*args):
    # args = (smem_ptr, src0, src1, ..., src{N-1}, trans, num, dtype, shape, space)
    # The last 5 entries are codegen attrs (n_attrs=5).
    trans_arg, num_arg, dtype_arg, shape_arg, space_arg = args[-5:]
    n_regs = int(num_arg)
    trans_b = bool(int(trans_arg)) if hasattr(trans_arg, "value") else bool(trans_arg)
    dtype = parse_str(dtype_arg)
    shape = parse_str(shape_arg)
    space = parse_str(space_arg)
    if dtype.startswith("."):
        dtype = dtype[1:]
    if n_regs not in (1, 2, 4):
        raise ValueError(f"stmatrix .num must be one of {{1, 2, 4}}, got {n_regs}")
    if dtype not in ("b16", "b8"):
        raise ValueError(f"stmatrix .type must be b16 or b8, got {dtype!r}")
    if shape not in ("m8n8", "m16n8"):
        raise ValueError(f"stmatrix .shape must be m8n8 or m16n8, got {shape!r}")
    if space not in ("shared", "shared::cta"):
        raise ValueError(f"stmatrix state space must be shared or shared::cta, got {space!r}")
    if shape == "m16n8" and not trans_b:
        raise ValueError("stmatrix .m16n8 requires .trans")
    trans_inst = ".trans" if trans_b else ""
    slot_list = "{" + ", ".join(f"%{i}" for i in range(n_regs)) + "}"
    src_loads = "\n".join(f"    uint32_t r{i} = *(uint32_t*)src{i};" for i in range(n_regs))
    in_constraints = ", ".join(f'"r"(r{i})' for i in range(n_regs))
    name = f"ptx_stmatrix_{shape}_{n_regs}_{1 if trans_b else 0}_{space.replace('::', '_')}_{dtype}"
    sig = "(void* smem_ptr, " + ", ".join(f"void* src{i}" for i in range(n_regs)) + ")"
    body = (
        f"{src_loads}\n"
        "    unsigned int addr = __cvta_generic_to_shared(smem_ptr);\n"
        "    asm volatile(\n"
        f'        "stmatrix.sync.aligned.{shape}.x{n_regs}{trans_inst}.{space}.{dtype} '
        f'[%{n_regs}], {slot_list};"\n'
        "        :\n"
        f'        : {in_constraints}, "r"(addr));'
    )
    return name, sig, body


device_intrinsic(
    "_ptx_stmatrix_impl",
    n_attrs=5,
    c_signature=lambda *a: _stmatrix_parts(*a)[1],
    helper_name=lambda *a: _stmatrix_parts(*a)[0],
    body=lambda *a: _stmatrix_parts(*a)[2],
)


@register_codegen("ptx_stmatrix")
def codegen_ptx_stmatrix(trans, num, dtype, shape, space, smem_ptr, *src_handles):
    trans = bool(trans)
    num = int(num)
    dtype_str = parse_str(dtype)
    if dtype_str.startswith("."):
        dtype_str = dtype_str[1:]
    n_regs = num
    if len(src_handles) != n_regs:
        raise ValueError(
            f"stmatrix .x{num}.{dtype_str} codegen expects {n_regs} src handles, "
            f"got {len(src_handles)}"
        )
    result = CODEGEN_REGISTRY["tirx._ptx_stmatrix_impl"](
        [smem_ptr, *src_handles, trans, num, dtype, shape, space]
    )
    return result[0] if isinstance(result, tuple) else result