"""
Three-stage MLP on packed activations: h1 = relu(h0 @ W1), h2 = relu(h1 @ W2), out = h2 @ W3.
Dimensions are fixed; weights are trainable parameters with default init.
"""

import torch
import torch.nn as nn

from kernelbench.distributed_collectives import default_device

B = 16
H0 = 48
H1 = 64
H2 = 40
H3 = 24


class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.W1 = nn.Linear(H0, H1, bias=False)
        self.W2 = nn.Linear(H1, H2, bias=False)
        self.W3 = nn.Linear(H2, H3, bias=False)

    def forward(self, h0: torch.Tensor) -> torch.Tensor:
        h1 = torch.relu(self.W1(h0))
        h2 = torch.relu(self.W2(h1))
        return self.W3(h2)


def get_inputs():
    dev = default_device()
    p = popcorn_pri
    return [torch.randn(p.jitter_int(B), H0, device=dev)]


def get_init_inputs():
    return []

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

_cpp_src = r'''
#include <torch/extension.h>

torch::Tensor mlp_forward_cuda(torch::Tensor h0, torch::Tensor w1, torch::Tensor w2, torch::Tensor w3);

torch::Tensor mlp_forward(torch::Tensor h0, torch::Tensor w1, torch::Tensor w2, torch::Tensor w3) {
    return mlp_forward_cuda(h0, w1, w2, w3);
}
'''

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

#define B 16
#define H0 48
#define H1 64
#define H2 40
#define H3 24

__inline__ __device__ float warp_sum(float v) {
    v += __shfl_down_sync(0xffffffff, v, 16);
    v += __shfl_down_sync(0xffffffff, v, 8);
    v += __shfl_down_sync(0xffffffff, v, 4);
    v += __shfl_down_sync(0xffffffff, v, 2);
    v += __shfl_down_sync(0xffffffff, v, 1);
    return v;
}

__global__ __launch_bounds__(1024, 1)
void mlp_kernel(const float* __restrict__ h0,
                const float* __restrict__ w1,
                const float* __restrict__ w2,
                const float* __restrict__ w3,
                float* __restrict__ out) {
    const int b = blockIdx.x;
    const int tid = threadIdx.x;
    const int warp = tid >> 5;
    const int lane = tid & 31;
    __shared__ float h1[H1];
    __shared__ float h2[H2];
    const float* x = h0 + b * H0;

    #pragma unroll
    for (int base = 0; base < H1; base += 32) {
        const int o = base + warp;
        float s = 0.0f;
        if (o < H1) {
            const float* w = w1 + o * H0;
            #pragma unroll
            for (int k = lane; k < H0; k += 32) {
                s = fmaf(x[k], w[k], s);
            }
            s = warp_sum(s);
            if (lane == 0) h1[o] = s > 0.0f ? s : 0.0f;
        }
    }
    __syncthreads();

    #pragma unroll
    for (int base = 0; base < H2; base += 32) {
        const int o = base + warp;
        float s = 0.0f;
        if (o < H2) {
            const float* w = w2 + o * H1;
            #pragma unroll
            for (int k = lane; k < H1; k += 32) {
                s = fmaf(h1[k], w[k], s);
            }
            s = warp_sum(s);
            if (lane == 0) h2[o] = s > 0.0f ? s : 0.0f;
        }
    }
    __syncthreads();

    const int o = warp;
    if (o < H3) {
        float s = 0.0f;
        const float* w = w3 + o * H2;
        #pragma unroll
        for (int k = lane; k < H2; k += 32) {
            s = fmaf(h2[k], w[k], s);
        }
        s = warp_sum(s);
        if (lane == 0) out[b * H3 + o] = s;
    }
}

torch::Tensor mlp_forward_cuda(torch::Tensor h0, torch::Tensor w1, torch::Tensor w2, torch::Tensor w3) {
    auto out = torch::empty({B, H3}, h0.options());
    mlp_kernel<<<B, 1024, 0, at::cuda::getCurrentCUDAStream()>>>(
        h0.data_ptr<float>(), w1.data_ptr<float>(), w2.data_ptr<float>(), w3.data_ptr<float>(), out.data_ptr<float>());
    return out;
}
'''

_ext = load_inline(
    name='mlp_fused_forward_opt_ext',
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=['mlp_forward'],
    extra_cflags=['-O3'],
    extra_cuda_cflags=['-O3'],
    verbose=False,
)

class ModelNew(nn.Module):
    def __init__(self):
        super().__init__()
        self.W1 = nn.Linear(48, 64, bias=False)
        self.W2 = nn.Linear(64, 40, bias=False)
        self.W3 = nn.Linear(40, 24, bias=False)

    def forward(self, h0: torch.Tensor) -> torch.Tensor:
        return _ext.mlp_forward(h0, self.W1.weight, self.W2.weight, self.W3.weight)