import torch
import torch.nn as nn

# Default ModelNew is PyTorch-only. run_and_check: use check_kernel=False with backend=cuda
# until the candidate file contains a real CUDA kernel (see kernel_static_checker).


def _fake_quant_symmetric_int8(x: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:
    """Symmetric int8 fake-quant: round(x/scale) clamped to [-127, 127], times scale."""
    q = torch.clamp(torch.round(x / scale), -127.0, 127.0)
    return q * scale


class Model(nn.Module):
    """
    **QuantVLA selective layout fusion** (reference target for a fused CUDA/Triton kernel).

    Vision–language–action stacks often use **token-wise mixed precision**: some
    timesteps or modalities tolerate aggressive quantization while others stay in
    full precision. This module fuses in **one forward**:

      1. **Layout routing** — ``sigmoid(layout_logits)`` yields a per-token mix
         coefficient in ``(0, 1)`` (quant vs FP residual).
      2. **Quantitative path** — global symmetric **int8-style fake quantization**
         (scale from the tensor max, ``/ 127``).
      3. **Selective blend** — ``α · x_quant + (1 - α) · x``.
      4. **Head** — linear projection on the mixed activations.

    This matches a “quantitative selective layout” block: the layout logits (from a
    router, modality tags, or calibration) choose how strongly each position uses the
    quantized layout vs the identity layout.
    """

    def __init__(self, dim: int):
        super().__init__()
        self.dim = dim
        self.out_proj = nn.Linear(dim, dim, bias=True)

    def forward(self, x: torch.Tensor, layout_logits: torch.Tensor) -> torch.Tensor:
        """
        Args:
            x: (B, T, D) VLA token activations (fused modalities).
            layout_logits: (B, T) unbounded logits; larger → more mass on quantized path.
        Returns:
            (B, T, D) updated activations.
        """
        B, T, D = x.shape
        assert D == self.dim
        assert layout_logits.shape == (B, T)

        alpha = torch.sigmoid(layout_logits).unsqueeze(-1)
        with torch.no_grad():
            amax = x.detach().abs().max().clamp(min=1e-8)
        scale = amax / 127.0
        x_q = _fake_quant_symmetric_int8(x, scale)
        mixed = alpha * x_q + (1.0 - alpha) * x
        return self.out_proj(mixed)


class ModelNew(nn.Module):
    """KernelBench candidate entry point; default delegates to ``Model``."""

    def __init__(self, dim: int):
        super().__init__()
        self._impl = Model(dim)

    def forward(self, x: torch.Tensor, layout_logits: torch.Tensor) -> torch.Tensor:
        return self._impl(x, layout_logits)


batch_size = 8
seq_len = 64
dim = 128


def get_inputs():
    p = popcorn_pri
    x = torch.randn(p.jitter_int(batch_size), p.jitter_int(seq_len), dim)
    layout_logits = torch.randn(p.jitter_int(batch_size), p.jitter_int(seq_len))
    return [x, layout_logits]


def get_init_inputs():
    return [dim]

import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline

cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>

// Kernel 1: Compute global max of |x| using grid-level reduction
__global__ void amax_kernel(const float* __restrict__ x, float* __restrict__ block_max, int64_t n) {
    __shared__ float smem[256];
    int tid = threadIdx.x;
    int bid = blockIdx.x;
    
    float local_max = 0.0f;
    
    // Grid-stride loop
    for (int64_t i = bid * blockDim.x + tid; i < n; i += blockDim.x * gridDim.x) {
        float v = fabsf(x[i]);
        if (v > local_max) local_max = v;
    }
    
    smem[tid] = local_max;
    __syncthreads();
    
    // Block reduction
    for (int s = blockDim.x >> 1; s > 0; s >>= 1) {
        if (tid < s && smem[tid + s] > smem[tid]) {
            smem[tid] = smem[tid + s];
        }
        __syncthreads();
    }
    
    if (tid == 0) {
        block_max[bid] = smem[0];
    }
}

// Kernel 2: Final reduction to get global max
__global__ void reduce_final(float* block_max, float* global_max, int n) {
    __shared__ float smem[256];
    int tid = threadIdx.x;
    
    float local_max = 0.0f;
    for (int i = tid; i < n; i += blockDim.x) {
        if (block_max[i] > local_max) local_max = block_max[i];
    }
    
    smem[tid] = local_max;
    __syncthreads();
    
    for (int s = blockDim.x >> 1; s > 0; s >>= 1) {
        if (tid < s && smem[tid + s] > smem[tid]) {
            smem[tid] = smem[tid + s];
        }
        __syncthreads();
    }
    
    if (tid == 0) {
        global_max[0] = smem[0];
    }
}

// Kernel 3: Precompute alphas = sigmoid(logits)
__global__ void sigmoid_kernel(const float* __restrict__ logits, float* __restrict__ alphas, int64_t n) {
    for (int64_t i = blockIdx.x * blockDim.x + threadIdx.x; i < n; i += blockDim.x * gridDim.x) {
        alphas[i] = 1.0f / (1.0f + expf(-logits[i]));
    }
}

// Kernel 4: Fused fake quant + blend using precomputed alphas
__global__ void quant_blend_kernel_opt(
    const float* __restrict__ x,
    const float* __restrict__ alphas,  // Precomputed sigmoid
    float* __restrict__ out,
    float scale,
    int64_t B, int64_t T, int64_t D
) {
    int64_t total = B * T * D;
    
    for (int64_t i = blockIdx.x * blockDim.x + threadIdx.x; i < total; i += blockDim.x * gridDim.x) {
        // Compute indices
        int64_t t = (i / D) % T;
        int64_t b = i / (T * D);
        
        // Get precomputed alpha
        float alpha = alphas[b * T + t];
        
        // Fake quant
        float xv = x[i];
        float q = roundf(xv / scale);
        q = fmaxf(-127.0f, fminf(127.0f, q));
        float xq = q * scale;
        
        // Blend
        out[i] = alpha * xq + (1.0f - alpha) * xv;
    }
}

// Host function: compute amax
torch::Tensor compute_amax(torch::Tensor x) {
    int64_t n = x.numel();
    int block_size = 256;
    int num_blocks = std::min(1024, (int)((n + block_size - 1) / block_size));
    
    // Storage for partial results
    auto block_max = torch::empty({num_blocks}, x.options());
    auto global_max = torch::empty({1}, x.options());
    
    // Launch first kernel
    amax_kernel<<<num_blocks, block_size>>>(
        x.data_ptr<float>(),
        block_max.data_ptr<float>(),
        n
    );
    
    // Launch final reduction
    reduce_final<<<1, block_size>>>(
        block_max.data_ptr<float>(),
        global_max.data_ptr<float>(),
        num_blocks
    );
    
    // Clamp to minimum
    return global_max.clamp(1e-8);
}

// Host function: compute sigmoid
torch::Tensor compute_sigmoid(torch::Tensor logits) {
    auto alphas = torch::empty_like(logits);
    int64_t n = logits.numel();
    int block_size = 256;
    int num_blocks = (n + block_size - 1) / block_size;
    
    sigmoid_kernel<<<num_blocks, block_size>>>(
        logits.data_ptr<float>(),
        alphas.data_ptr<float>(),
        n
    );
    
    return alphas;
}

// Host function: quant blend
torch::Tensor quant_blend_opt(torch::Tensor x, torch::Tensor alphas, float scale) {
    auto out = torch::empty_like(x);
    int64_t B = x.size(0), T = x.size(1), D = x.size(2);
    int64_t total = B * T * D;
    
    int block_size = 256;
    int num_blocks = (total + block_size - 1) / block_size;
    
    quant_blend_kernel_opt<<<num_blocks, block_size>>>(
        x.data_ptr<float>(),
        alphas.data_ptr<float>(),
        out.data_ptr<float>(),
        scale,
        B, T, D
    );
    
    return out;
}
"""

cpp_source = """
torch::Tensor compute_amax(torch::Tensor x);
torch::Tensor compute_sigmoid(torch::Tensor logits);
torch::Tensor quant_blend_opt(torch::Tensor x, torch::Tensor alphas, float scale);
"""

fused_module = load_inline(
    name='qvla_opt',
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=['compute_amax', 'compute_sigmoid', 'quant_blend_opt'],
    verbose=True,
    extra_cuda_cflags=['-O3']
)

class ModelNew(nn.Module):
    def __init__(self, dim: int):
        super().__init__()
        self.dim = dim
        self.out_proj = nn.Linear(dim, dim, bias=True)

    def forward(self, x: torch.Tensor, layout_logits: torch.Tensor) -> torch.Tensor:
        x = x.contiguous()
        layout_logits = layout_logits.contiguous()
        
        # Compute amax using custom kernel
        amax = fused_module.compute_amax(x)
        scale = (amax / 127.0).item()
        
        # Precompute alphas
        alphas = fused_module.compute_sigmoid(layout_logits)
        
        # Fused quant + blend
        mixed = fused_module.quant_blend_opt(x, alphas, scale)
        
        # Linear projection
        return self.out_proj(mixed)