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).


class Model(nn.Module):
    """
    **Fused probability-flow ODE integration** (flow-matching style) in one ``forward``.

    Models the ODE ``dx/dt = v_θ(x, t)`` on ``t ∈ [0, 1]`` with a learned velocity
    field. A **single fused pass** unrolls **fixed-step Euler** integration for
    ``num_steps`` substeps—typical for **real-time trajectory / action generation**
    where latency is dominated by repeated ``v_θ`` evaluations.

    This is the continuous-time analogue of a rectified-flow / conditional-flow
    sampler: one ``forward`` returns an approximate **x(1)** from **x(0)** without a
    Python loop across separate module calls (fusion target for the integrator +
    field).
    """

    def __init__(self, dim: int, num_steps: int):
        super().__init__()
        assert dim >= 1 and num_steps >= 1
        self.dim = dim
        self.num_steps = num_steps
        self.register_buffer("dt", torch.tensor(1.0 / float(num_steps)))

        self.velocity = nn.Sequential(
            nn.Linear(dim + 1, dim * 4),
            nn.SiLU(),
            nn.Linear(dim * 4, dim * 4),
            nn.SiLU(),
            nn.Linear(dim * 4, dim),
        )

    def forward(self, x0: torch.Tensor) -> torch.Tensor:
        """
        Args:
            x0: (B, dim) initial state at ``t = 0`` (e.g. noise or anchor trajectory).
        Returns:
            (B, dim) approximate state at ``t = 1`` after fused Euler integration.
        """
        B, d = x0.shape
        assert d == self.dim
        x = x0
        dt = self.dt.to(dtype=x0.dtype, device=x0.device)
        for k in range(self.num_steps):
            t = (float(k) + 0.5) / float(self.num_steps)
            t_b = torch.full((B, 1), t, device=x0.device, dtype=x0.dtype)
            inp = torch.cat([x, t_b], dim=-1)
            v = self.velocity(inp)
            x = x + dt * v
        return x


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

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

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


batch_size = 8
dim = 64
num_steps = 32


def get_inputs():
    return [torch.randn(popcorn_pri.jitter_int(batch_size), dim)]


def get_init_inputs():
    return [dim, num_steps]

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

_cpp_src = r"""
#include <torch/extension.h>
torch::Tensor flow_forward(torch::Tensor x0, torch::Tensor w1, torch::Tensor b1, torch::Tensor w2, torch::Tensor b2, torch::Tensor w3, torch::Tensor b3, int64_t num_steps);
"""

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

__device__ __forceinline__ float silu_f(float x){ return x / (1.0f + expf(-x)); }

__device__ __forceinline__ float warp_sum(float v){
#pragma unroll
    for(int o = 16; o > 0; o >>= 1) v += __shfl_down_sync(0xffffffff, v, o);
    return v;
}

__device__ __forceinline__ float block_sum(float v){
    __shared__ float sm[8];
    int tid = threadIdx.x;
    int lane = tid & 31;
    int warp = tid >> 5;
    v = warp_sum(v);
    if(lane == 0) sm[warp] = v;
    __syncthreads();
    v = (warp == 0 && lane < 8) ? sm[lane] : 0.0f;
    if(warp == 0) v = warp_sum(v);
    return v;
}

__global__ __launch_bounds__(256)
void l1_kernel(const float* __restrict__ x, const float* __restrict__ w,
               const float* __restrict__ b, float* __restrict__ h, float t){
    int bid = blockIdx.x, o = blockIdx.y, tid = threadIdx.x;
    float acc = 0.0f;
    const float* row = w + o * 65;
    if(tid < 64) acc = row[tid] * x[bid * 64 + tid];
    if(tid == 64) acc = row[64] * t;
    acc = block_sum(acc);
    if(tid == 0) h[bid * 256 + o] = silu_f(acc + b[o]);
}

__global__ __launch_bounds__(256)
void l2_kernel(const float* __restrict__ h1, const float* __restrict__ w,
               const float* __restrict__ b, float* __restrict__ h2){
    int bid = blockIdx.x, o = blockIdx.y, tid = threadIdx.x;
    const float* row = w + o * 256;
    float acc = row[tid] * h1[bid * 256 + tid];
    acc = block_sum(acc);
    if(tid == 0) h2[bid * 256 + o] = silu_f(acc + b[o]);
}

__global__ __launch_bounds__(256)
void l3_kernel(const float* __restrict__ x, const float* __restrict__ h2,
               const float* __restrict__ w, const float* __restrict__ b,
               float* __restrict__ out, float dt){
    int bid = blockIdx.x, o = blockIdx.y, tid = threadIdx.x;
    const float* row = w + o * 256;
    float acc = row[tid] * h2[bid * 256 + tid];
    acc = block_sum(acc);
    if(tid == 0) out[bid * 64 + o] = fmaf(dt, acc + b[o], x[bid * 64 + o]);
}

torch::Tensor flow_forward(torch::Tensor x0, torch::Tensor w1, torch::Tensor b1,
                           torch::Tensor w2, torch::Tensor b2, torch::Tensor w3,
                           torch::Tensor b3, int64_t num_steps){
    auto opts = x0.options();
    int B = (int)x0.size(0);
    auto h1 = torch::empty({B, 256}, opts);
    auto h2 = torch::empty({B, 256}, opts);
    auto x1 = torch::empty_like(x0);
    auto x2 = torch::empty_like(x0);
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    dim3 gl1(B, 256), gl3(B, 64);
    float dt = 1.0f / (float)num_steps;
    const float* src = x0.data_ptr<float>();
    float* dst = x1.data_ptr<float>();
    const float *pw1 = w1.data_ptr<float>(), *pb1 = b1.data_ptr<float>();
    const float *pw2 = w2.data_ptr<float>(), *pb2 = b2.data_ptr<float>();
    const float *pw3 = w3.data_ptr<float>(), *pb3 = b3.data_ptr<float>();
    float *ph1 = h1.data_ptr<float>(), *ph2 = h2.data_ptr<float>();
    float *px1 = x1.data_ptr<float>(), *px2 = x2.data_ptr<float>();

    for(int k = 0; k < (int)num_steps; ++k){
        float t = ((float)k + 0.5f) * dt;
        l1_kernel<<<gl1, 256, 0, stream>>>(src, pw1, pb1, ph1, t);
        l2_kernel<<<gl1, 256, 0, stream>>>(ph1, pw2, pb2, ph2);
        dst = (k & 1) ? px2 : px1;
        l3_kernel<<<gl3, 256, 0, stream>>>(src, ph2, pw3, pb3, dst, dt);
        src = dst;
    }
    return (num_steps & 1) ? x1 : x2;
}
"""

_ext = load_inline(
    name="flow64_ext_final",
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=["flow_forward"],
    extra_cuda_cflags=["-O3", "--use_fast_math"],
    extra_cflags=["-O3"],
    verbose=False,
)

class ModelNew(nn.Module):
    def __init__(self, dim: int, num_steps: int):
        super().__init__()
        assert dim >= 1 and num_steps >= 1
        self.dim = dim
        self.num_steps = num_steps
        self.register_buffer("dt", torch.tensor(1.0 / float(num_steps)))
        self.velocity = nn.Sequential(
            nn.Linear(dim + 1, dim * 4),
            nn.SiLU(),
            nn.Linear(dim * 4, dim * 4),
            nn.SiLU(),
            nn.Linear(dim * 4, dim),
        )

    def forward(self, x0: torch.Tensor) -> torch.Tensor:
        return _ext.flow_forward(
            x0.contiguous(),
            self.velocity[0].weight.contiguous(), self.velocity[0].bias.contiguous(),
            self.velocity[2].weight.contiguous(), self.velocity[2].bias.contiguous(),
            self.velocity[4].weight.contiguous(), self.velocity[4].bias.contiguous(),
            self.num_steps,
        )