"""
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

cuda_source = '''
#include <torch/extension.h>
#include <cuda_runtime.h>

__global__ void mlp_kernel(const float* h0, const float* W1, const float* W2, const float* W3,
                           float* out, int B, int H0, int H1, int H2, int H3) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int total = B * H3;
    if (idx >= total) return;
    int b = idx / H3;
    int h3 = idx % H3;

    // Compute h1 = relu(h0[b] @ W1.T)  W1 is (H1, H0)
    float h1_val[64];
    for (int i = 0; i < H1; i++) {
        float sum = 0.0f;
        for (int j = 0; j < H0; j++) {
            sum += h0[b * H0 + j] * W1[i * H0 + j];
        }
        h1_val[i] = fmaxf(sum, 0.0f);
    }

    // h2 = relu(h1 @ W2.T) W2 (H2, H1)
    float h2_val[40];
    for (int i = 0; i < H2; i++) {
        float sum = 0.0f;
        for (int j = 0; j < H1; j++) {
            sum += h1_val[j] * W2[i * H1 + j];
        }
        h2_val[i] = fmaxf(sum, 0.0f);
    }

    // out = h2 @ W3.T  W3 (H3, H2)
    float sum = 0.0f;
    for (int j = 0; j < H2; j++) {
        sum += h2_val[j] * W3[h3 * H2 + j];
    }
    out[b * H3 + h3] = sum;
}

torch::Tensor mlp_forward(torch::Tensor h0, torch::Tensor W1, torch::Tensor W2, torch::Tensor W3) {
    int B = h0.size(0);
    int H0 = h0.size(1);
    int H1 = W1.size(0);
    int H2 = W2.size(0);
    int H3 = W3.size(0);
    auto out = torch::empty({B, H3}, h0.options());
    int threads = 256;
    int blocks = (B * H3 + threads - 1) / threads;
    mlp_kernel<<<blocks, threads>>>(h0.data_ptr<float>(), W1.data_ptr<float>(),
                                     W2.data_ptr<float>(), W3.data_ptr<float>(),
                                     out.data_ptr<float>(), B, H0, H1, H2, H3);
    return out;
}
'''

cpp_source = '''
torch::Tensor mlp_forward(torch::Tensor h0, torch::Tensor W1, torch::Tensor W2, torch::Tensor W3);
'''

module = load_inline(
    name='mlp_cuda',
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=['mlp_forward'],
    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 module.mlp_forward(h0, self.W1.weight, self.W2.weight, self.W3.weight)