"""W4A16 weight-only quantized GEMM (CUTLASS mixed-input, scale+zero-point mode).

Scheme: x (M,K) bf16 @ dequant(w_q, scales, zeros) -> (M,N) bf16.
  w_q:  (K//2, N) uint8, low nibble = even-K row, high nibble = odd-K row
  scales/zeros: (K//128, N) bf16, per-group along K

The kernel is a CUTLASS 3.x Hopper mixed-input GEMM (example 55 pattern):
quantized weight is the register-resident A operand, TMA epilogues enabled via
the explicit swap formulation. Weights are repacked once (lazily, cached) into
the CUTLASS layout: w_q.T (N, K//2) then nibble-shuffled via cutlass::reorder_tensor.
Because CUTLASS computes `w*scale + zero`, we pass zero_eff = -zeros*scales so the
result matches the reference `(w - zeros)*scales`. Scales/zeros are upcast to fp32
so the dequant rounding matches the bf16 reference closely enough for the numeric
stress tolerance.
"""
from __future__ import annotations

import os
import threading

import torch
import torch.nn as nn

GROUP_SIZE = 128

_CUDA_SOURCE = r"""
/* CUTLASS W4A16 int4-bf16 GEMM (scale+zero-point mode). */
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>

#include "cutlass/cutlass.h"
#include "cute/tensor.hpp"
#include "cutlass/tensor_ref.h"
#include "cutlass/epilogue/collective/default_epilogue.hpp"
#include "cutlass/epilogue/thread/linear_combination.h"
#include "cutlass/gemm/dispatch_policy.hpp"
#include "cutlass/gemm/collective/collective_builder.hpp"
#include "cutlass/epilogue/collective/collective_builder.hpp"
#include "cutlass/gemm/device/gemm_universal_adapter.h"
#include "cutlass/gemm/kernel/gemm_universal.hpp"
#include "cutlass/util/packed_stride.hpp"
#include "cutlass/util/mixed_dtype_utils.hpp"

using namespace cute;

using MmaType = cutlass::bfloat16_t;
using QuantType = cutlass::uint4b_t;

using ElementA = MmaType;
using LayoutA  = cutlass::layout::RowMajor;
constexpr int AlignmentA = 128 / cutlass::sizeof_bits<ElementA>::value;

using ElementB = QuantType;
using LayoutB  = cutlass::layout::ColumnMajor;
constexpr int AlignmentB = 128 / cutlass::sizeof_bits<ElementB>::value;

using LayoutA_Transpose = typename cutlass::layout::LayoutTranspose<LayoutA>::type;
using LayoutB_Transpose = typename cutlass::layout::LayoutTranspose<LayoutB>::type;

using StrideA = cutlass::detail::TagToStrideA_t<LayoutA>;
using StrideB = cutlass::detail::TagToStrideB_t<LayoutB>;

using ValueShuffle = Layout<Shape<_2,_4>, Stride<_4,_1>>;
using MmaAtomShape = Layout<Shape<_1, Int<1>>>;
using LayoutAtomQuant = decltype(cutlass::compute_memory_reordering_atom<MmaType, MmaAtomShape, ValueShuffle>());
using LayoutB_Reordered = decltype(cute::tile_to_shape(LayoutAtomQuant{}, Layout<Shape<int,int,int>, StrideB>{}));

using ElementScale = float;
using ElementZero = float;
using LayoutScale = cutlass::layout::RowMajor;

using ElementC = cutlass::bfloat16_t;
using LayoutC  = cutlass::layout::RowMajor;
constexpr int AlignmentC = 128 / cutlass::sizeof_bits<ElementC>::value;
using ElementD = ElementC;
using LayoutD  = LayoutC;
constexpr int AlignmentD = AlignmentC;

using ElementAccumulator = float;
using ArchTag = cutlass::arch::Sm90;
using OperatorClass = cutlass::arch::OpClassTensorOp;

using StrideC = cutlass::detail::TagToStrideC_t<cutlass::layout::ColumnMajor>;
using StrideD = cutlass::detail::TagToStrideC_t<cutlass::layout::ColumnMajor>;

template <class TileShape_, class ClusterShape_, class KernelSchedule_, class BLayout_ = LayoutB_Reordered>
struct Cfg {
  using TileShape = TileShape_;
  using ClusterShape = ClusterShape_;
  using KernelSchedule = KernelSchedule_;
  using BLayout = BLayout_;
  using EpilogueSchedule = cutlass::epilogue::TmaWarpSpecializedCooperative;
  using EpilogueTileType = cutlass::epilogue::collective::EpilogueTileAuto;

  using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
      cutlass::arch::Sm90, cutlass::arch::OpClassTensorOp,
      TileShape, ClusterShape, EpilogueTileType,
      ElementAccumulator, ElementAccumulator,
      ElementC, typename cutlass::layout::LayoutTranspose<LayoutC>::type, AlignmentC,
      ElementD, typename cutlass::layout::LayoutTranspose<LayoutD>::type, AlignmentD,
      EpilogueSchedule>::CollectiveOp;

  using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
      ArchTag, OperatorClass,
      cute::tuple<ElementB, ElementScale, ElementZero>, BLayout, AlignmentB,
      ElementA, LayoutA_Transpose, AlignmentA,
      ElementAccumulator,
      TileShape, ClusterShape,
      cutlass::gemm::collective::StageCountAutoCarveout<
        static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))
      >,
      KernelSchedule>::CollectiveOp;

  using Gemm = cutlass::gemm::device::GemmUniversalAdapter<
      cutlass::gemm::kernel::GemmUniversal<
        Shape<int,int,int,int>, CollectiveMainloop, CollectiveEpilogue>>;

  using StrideScale = typename CollectiveMainloop::StrideScale;
};

namespace {

template <class G, class StrideScale>
void launch_gemm_t(int m, int n, int k, int g,
                   const ElementB* w, const ElementA* x, const ElementScale* s, const ElementScale* z,
                   ElementD* d, cudaStream_t stream) {
  int l = 1;
  int scale_k = k / g;
  auto shape_B = cute::make_shape(n, k, l);
  StrideB stride_B = cutlass::make_cute_packed_stride(StrideB{}, shape_B);
  StrideA stride_A = cutlass::make_cute_packed_stride(StrideA{}, cute::make_shape(m, k, l));
  StrideC stride_C = cutlass::make_cute_packed_stride(StrideC{}, cute::make_shape(n, m, l));
  StrideD stride_D = cutlass::make_cute_packed_stride(StrideD{}, cute::make_shape(n, m, l));
  StrideScale stride_S = cutlass::make_cute_packed_stride(StrideScale{}, cute::make_shape(n, scale_k, l));
  auto layout_B = make_layout(shape_B, stride_B);
  LayoutB_Reordered layout_B_reordered = cute::tile_to_shape(LayoutAtomQuant{}, shape_B);

  typename G::Arguments args{
    cutlass::gemm::GemmUniversalMode::kGemm,
    {n, m, k, l},
    {w, layout_B_reordered, x, stride_A, s, stride_S, g, z},
    {{1.0f, 0.0f}, d, stride_D, d, stride_D}
  };
  thread_local std::unique_ptr<uint8_t[]> workspace_cache;
  thread_local size_t workspace_size = 0;
  constexpr size_t kWorkspace = 4 * 1024 * 1024;
  if (workspace_size < kWorkspace) {
    workspace_cache.reset(new uint8_t[kWorkspace]);
    workspace_size = kWorkspace;
  }
  if (G::can_implement(args) != cutlass::Status::kSuccess) {
    throw std::runtime_error("CUTLASS W4A16: cannot implement");
  }
  G gemm;
  gemm.initialize(args, workspace_cache.get());
  gemm.run(stream);
}

}  // namespace

#define DEF_CFG(N, ...) \
  using N = Cfg<__VA_ARGS__>;

using S1 = cutlass::gemm::KernelTmaWarpSpecializedCooperative;
using S2 = cutlass::gemm::KernelTmaWarpSpecializedPingpong;
using C1 = Shape<_1,_1,_1>;
using C2 = Shape<_2,_1,_1>;

DEF_CFG(CFG0, Shape<_128,_128,Int<128>>, C2, S1)
DEF_CFG(CFG1, Shape<_256,_128,Int<128>>, C1, S1)
DEF_CFG(CFG2, Shape<_64,_128,Int<128>>, C1, S2)

void w4a16_gemm(int64_t variant,
                torch::Tensor x, torch::Tensor w, torch::Tensor scales, torch::Tensor zeros,
                torch::Tensor out, int64_t M, int64_t N, int64_t K, int64_t group) {
  cudaStream_t stream = at::cuda::getCurrentCUDAStream();
  const ElementA* xp = reinterpret_cast<const ElementA*>(x.data_ptr());
  const ElementB* wp = reinterpret_cast<const ElementB*>(w.data_ptr());
  const ElementScale* sp = reinterpret_cast<const ElementScale*>(scales.data_ptr());
  const ElementScale* zp = reinterpret_cast<const ElementScale*>(zeros.data_ptr());
  ElementD* dp = reinterpret_cast<ElementD*>(out.data_ptr());
  int m = (int)M, n = (int)N, k = (int)K, g = (int)group;
  switch (variant) {
    case 0: launch_gemm_t<typename CFG0::Gemm, typename CFG0::StrideScale>(m,n,k,g,wp,xp,sp,zp,dp,stream); break;
    case 1: launch_gemm_t<typename CFG1::Gemm, typename CFG1::StrideScale>(m,n,k,g,wp,xp,sp,zp,dp,stream); break;
    case 2: launch_gemm_t<typename CFG2::Gemm, typename CFG2::StrideScale>(m,n,k,g,wp,xp,sp,zp,dp,stream); break;
    default: throw std::runtime_error("bad variant");
  }
}

void reorder_w(torch::Tensor w_nat, torch::Tensor w_out, int64_t K, int64_t N) {
  int n = (int)N, k = (int)K, l = 1;
  auto shape_B = cute::make_shape(n, k, l);
  StrideB stride_B = cutlass::make_cute_packed_stride(StrideB{}, shape_B);
  auto layout_B = make_layout(shape_B, stride_B);
  LayoutB_Reordered layout_B_reordered = cute::tile_to_shape(LayoutAtomQuant{}, shape_B);
  cutlass::reorder_tensor(
    reinterpret_cast<const ElementB*>(w_nat.data_ptr()), layout_B,
    reinterpret_cast<ElementB*>(w_out.data_ptr()), layout_B_reordered);
}
"""

_BUILD_LOCK = threading.Lock()
_ext = None
_ext_failed = False


def _build_extension():
    global _ext, _ext_failed
    if _ext is not None or _ext_failed:
        return _ext
    with _BUILD_LOCK:
        if _ext is not None or _ext_failed:
            return _ext
        try:
            from torch.utils.cpp_extension import load_inline
            cpp_decl = (
                "void w4a16_gemm(int64_t variant, torch::Tensor x, torch::Tensor w, torch::Tensor scales, torch::Tensor zeros, "
                "torch::Tensor out, int64_t M, int64_t N, int64_t K, int64_t group);\n"
                "void reorder_w(torch::Tensor w_nat, torch::Tensor w_out, int64_t K, int64_t N);\n"
            )
            cutlass_root = os.environ.get("CUTLASS_ROOT", "/tmp/cutlass")
            extra_inc = [
                f"{cutlass_root}/include",
                f"{cutlass_root}/tools/util/include",
                f"{cutlass_root}/examples/55_hopper_mixed_dtype_gemm",
                f"{cutlass_root}/examples/common",
                "/home/ubuntu/.local/lib/python3.10/site-packages/pybind11/include",
            ]
            build_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "build", "w4a16_ext")
            os.makedirs(build_dir, exist_ok=True)
            os.environ.setdefault("TORCH_CUDA_ARCH_LIST", "9.0")
            _ext = load_inline(
                name="w4a16_ext",
                cpp_sources=cpp_decl,
                cuda_sources=_CUDA_SOURCE,
                functions=["w4a16_gemm", "reorder_w"],
                extra_include_paths=extra_inc,
                extra_cuda_cflags=["-arch=sm_90a", "--expt-relaxed-constexpr", "-std=c++17", "-O3",
                                   "-DCUTLASS_ENABLE_CUB=0"],
                verbose=False,
                build_directory=build_dir,
            )
        except Exception as e:  # noqa: BLE001
            _ext_failed = True
            print(f"[w4a16] CUTLASS extension build failed: {e}", flush=True)
            return None
        return _ext


# Per-shape kernel config selection (from the CUTLASS config sweep).
#   v8  = Tile(128,128,128), Cluster(2,1,1), Cooperative   -- best decode/small-M
#   v6  = Tile(256,128,128), Cluster(1,1,1), Cooperative   -- best large-M (compute)
#   v10 = Tile(64,128,128),  Cluster(1,1,1), Pingpong      -- best small-N decode
def _pick_variant(M, N):
    if N <= 4096 and M <= 1:
        return 2
    if M <= 1:
        return 0
    if M <= 32:
        return 0
    return 1


class Model(nn.Module):
    def __init__(self, M: int, N: int, K: int, group_size: int = GROUP_SIZE):
        super().__init__()
        self.M, self.N, self.K = M, N, K
        self.group_size = group_size
        self.register_buffer("w_q", torch.empty(K // 2, N, dtype=torch.uint8))
        self.register_buffer("scales", torch.empty(K // group_size, N, dtype=torch.bfloat16))
        self.register_buffer("zeros", torch.empty(K // group_size, N, dtype=torch.bfloat16))
        self._prepared = None
        self._out_cache = None

    def _prepare(self):
        wp = self.w_q
        sc = self.scales
        ze = self.zeros
        w_nat = wp.T.contiguous()  # (N, K//2)
        w_shuf = torch.empty_like(w_nat)
        sc_ = sc.float().contiguous()
        ze_eff = (-(ze.float()) * sc.float()).contiguous()
        ext = _build_extension()
        if ext is None:
            return None
        ext.reorder_w(w_nat, w_shuf, self.K, self.N)
        return ext, w_shuf, sc_, ze_eff, _pick_variant(self.M, self.N), self.group_size

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if self._prepared is None:
            self._prepared = self._prepare()
        if self._prepared is None:
            return self._forward_fallback(x)
        ext, w_shuf, sc_, ze_eff, variant, group = self._prepared
        M = x.shape[0]
        x = x.contiguous()
        if self._out_cache is None or self._out_cache.shape[0] != M:
            self._out_cache = torch.empty((M, self.N), dtype=torch.bfloat16, device=x.device)
        out = self._out_cache
        ext.w4a16_gemm(variant, x, w_shuf, sc_, ze_eff, out, M, self.N, self.K, group)
        return out

    def _forward_fallback(self, x: torch.Tensor) -> torch.Tensor:
        # Reference-style fallback (unpack -> dequant -> matmul).
        K = self.K
        wp, sc, ze = self.w_q, self.scales, self.zeros
        wu = torch.empty((K, self.N), dtype=torch.uint8, device=wp.device)
        wu[0::2] = wp & 0xF
        wu[1::2] = (wp >> 4) & 0xF
        wbf = (wu.to(torch.bfloat16) - ze.repeat_interleave(self.group_size, dim=0)) * sc.repeat_interleave(
            self.group_size, dim=0)
        return x.to(torch.bfloat16) @ wbf


M = 1
N = 12288
K = 4096


def get_inputs():
    x = torch.randn(M, K, dtype=torch.bfloat16)
    return [x]


def get_init_inputs():
    return [M, N, K]