"""
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.h>
#include <cuda_runtime.h>

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

    float h1[64];
    float h2[40];

    for (int j = 0; j < H1; ++j) {
        float sum = 0.0f;
        for (int k = 0; k < H0; ++k) {
            sum += h0[row * H0 + k] * W1[j * H0 + k];
        }
        h1[j] = fmaxf(sum, 0.0f);
    }

    for (int j = 0; j < H2; ++j) {
        float sum = 0.0f;
        for (int k = 0; k < H1; ++k) {
            sum += h1[k] * W2[j * H1 + k];
        }
        h2[j] = fmaxf(sum, 0.0f);
    }

    for (int j = 0; j < H3; ++j) {
        float sum = 0.0f;
        for (int k = 0; k < H2; ++k) {
            sum += h2[k] * W3[j * H2 + k];
        }
        out[row * H3 + j] = 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::zeros({B, H3}, h0.options());

    int threads = 256;
    int blocks = (B + 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)
```