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

cpp_src = """
torch::Tensor kronecker_product_cuda(torch::Tensor a, torch::Tensor b);
"""

cuda_src = """
#include <torch/extension.h>
#include <cuda_runtime.h>

__global__ void __launch_bounds__(512) kronecker_product_kernel(
    const float* __restrict__ a,
    const float* __restrict__ b,
    float* __restrict__ output,
    int r1, int c1,
    int r2, int c2,
    int out_cols,
    int64_t total_elements
) {
    int64_t idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= total_elements) return;

    int col_out = idx % out_cols;
    int row_out = idx / out_cols;

    int row_a = row_out / r2;
    int col_a = col_out / c2;
    int row_b = row_out % r2;
    int col_b = col_out % c2;

    float val_a = a[row_a * c1 + col_a];
    float val_b = b[row_b * c2 + col_b];

    output[idx] = val_a * val_b;
}

torch::Tensor kronecker_product_cuda(torch::Tensor a, torch::Tensor b) {
    if (a.dim() != 2 || b.dim() != 2) {
        return torch::kron(a, b);
    }

    int r1 = a.size(0);
    int c1 = a.size(1);
    int r2 = b.size(0);
    int c2 = b.size(1);

    int out_rows = r1 * r2;
    int out_cols = c1 * c2;
    int64_t total_elements = (int64_t)out_rows * out_cols;

    auto output = torch::empty({out_rows, out_cols}, a.options());

    a = a.contiguous();
    b = b.contiguous();

    const int block_size = 512;
    int64_t grid_size = (total_elements + block_size - 1) / block_size;
    if (grid_size > 2147483647) grid_size = 2147483647;

    kronecker_product_kernel<<<grid_size, block_size>>>(
        a.data_ptr<float>(),
        b.data_ptr<float>(),
        output.data_ptr<float>(),
        r1, c1,
        r2, c2,
        out_cols,
        total_elements
    );

    return output;
}
"""

class ModelNew(nn.Module):
    def __init__(self):
        super().__init__()
        self.module = load_inline(
            name="kronecker_product_opt",
            cpp_sources=cpp_src,
            cuda_sources=cuda_src,
            functions=["kronecker_product_cuda"],
            verbose=False,
            extra_cuda_cflags=["-O3"]
        )

    def forward(self, a, b):
        return self.module.kronecker_product_cuda(a, b)

def get_inputs():
    a = torch.randn(64, 64, dtype=torch.float32)
    b = torch.randn(64, 64, dtype=torch.float32)
    return [a, b]

def get_init_inputs():
    return []