import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Simple model that performs matrix multiplication (C = A * B) for upper triangular matrices.
    """
    def __init__(self):
        super(Model, self).__init__()
    
    def forward(self, A, B):
        """
        Performs matrix multiplication for upper triangular matrices.

        Args:
            A (torch.Tensor): Upper triangular matrix of shape (N, N).
            B (torch.Tensor): Upper triangular matrix of shape (N, N).

        Returns:
            torch.Tensor: The product of A and B, also an upper triangular matrix of shape (N, N).
        """
        return torch.triu(torch.matmul(A, B))

N = 4096

def get_inputs():
    """
    Generates upper triangular matrices for testing.

    Returns:
        list: A list containing two upper triangular matrices of shape (N, N).
    """
    A = torch.triu(torch.rand(N, N))
    B = torch.triu(torch.rand(N, N))
    return [A, B]

def get_init_inputs():
    """
    No specific initialization inputs are needed for this model.

    Returns:
        list: An empty list.
    """
    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 triumm_forward(torch::Tensor A, torch::Tensor B);
"""

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

#define TILE_M 64
#define TILE_N 64
#define TILE_K 32

__global__ __launch_bounds__(256, 2)
void triumm64_kernel(const float* __restrict__ A,
                     const float* __restrict__ B,
                     float* __restrict__ C,
                     int n) {
    const int blockCol = blockIdx.x * TILE_N;
    const int blockRow = blockIdx.y * TILE_M;
    const int tx = threadIdx.x; // 0..15
    const int ty = threadIdx.y; // 0..15

    // Tiles strictly below the diagonal are exactly zero.  Write them here so
    // the returned empty tensor has the same lower triangle as torch.triu.
    if (blockCol < blockRow) {
        #pragma unroll
        for (int rr = 0; rr < 4; ++rr) {
            int row = blockRow + ty + rr * 16;
            if (row < n) {
                #pragma unroll
                for (int cc = 0; cc < 4; ++cc) {
                    int col = blockCol + tx + cc * 16;
                    if (col < n) C[(long long)row * n + col] = 0.0f;
                }
            }
        }
        return;
    }

    __shared__ float As[TILE_M][TILE_K + 1];
    __shared__ float Bs[TILE_K][TILE_N + 1];

    float acc[4][4];
    #pragma unroll
    for (int i = 0; i < 4; ++i) {
        #pragma unroll
        for (int j = 0; j < 4; ++j) acc[i][j] = 0.0f;
    }

    const int tid = ty * 16 + tx;
    int kEnd = blockCol + TILE_N;
    if (kEnd > n) kEnd = n;

    // Because A and B are upper triangular, only k in [blockRow, blockCol+63]
    // can contribute to this 64x64 output tile.  Values inside the boundary
    // tiles that are outside the true per-element range are already stored as
    // zeros in A or B.
    for (int kk = blockRow; kk < kEnd; kk += TILE_K) {
        for (int idx = tid; idx < TILE_M * TILE_K; idx += 256) {
            int r = idx / TILE_K;
            int c = idx - r * TILE_K;
            int ar = blockRow + r;
            int ac = kk + c;
            As[r][c] = (ar < n && ac < n) ? A[(long long)ar * n + ac] : 0.0f;
        }
        for (int idx = tid; idx < TILE_K * TILE_N; idx += 256) {
            int r = idx / TILE_N;
            int c = idx - r * TILE_N;
            int br = kk + r;
            int bc = blockCol + c;
            Bs[r][c] = (br < n && bc < n) ? B[(long long)br * n + bc] : 0.0f;
        }
        __syncthreads();

        #pragma unroll
        for (int k = 0; k < TILE_K; ++k) {
            float a0 = As[ty +  0][k];
            float a1 = As[ty + 16][k];
            float a2 = As[ty + 32][k];
            float a3 = As[ty + 48][k];
            float b0 = Bs[k][tx +  0];
            float b1 = Bs[k][tx + 16];
            float b2 = Bs[k][tx + 32];
            float b3 = Bs[k][tx + 48];

            acc[0][0] = fmaf(a0, b0, acc[0][0]);
            acc[0][1] = fmaf(a0, b1, acc[0][1]);
            acc[0][2] = fmaf(a0, b2, acc[0][2]);
            acc[0][3] = fmaf(a0, b3, acc[0][3]);
            acc[1][0] = fmaf(a1, b0, acc[1][0]);
            acc[1][1] = fmaf(a1, b1, acc[1][1]);
            acc[1][2] = fmaf(a1, b2, acc[1][2]);
            acc[1][3] = fmaf(a1, b3, acc[1][3]);
            acc[2][0] = fmaf(a2, b0, acc[2][0]);
            acc[2][1] = fmaf(a2, b1, acc[2][1]);
            acc[2][2] = fmaf(a2, b2, acc[2][2]);
            acc[2][3] = fmaf(a2, b3, acc[2][3]);
            acc[3][0] = fmaf(a3, b0, acc[3][0]);
            acc[3][1] = fmaf(a3, b1, acc[3][1]);
            acc[3][2] = fmaf(a3, b2, acc[3][2]);
            acc[3][3] = fmaf(a3, b3, acc[3][3]);
        }
        __syncthreads();
    }

    #pragma unroll
    for (int rr = 0; rr < 4; ++rr) {
        int row = blockRow + ty + rr * 16;
        if (row < n) {
            #pragma unroll
            for (int cc = 0; cc < 4; ++cc) {
                int col = blockCol + tx + cc * 16;
                if (col < n) {
                    C[(long long)row * n + col] = (col >= row) ? acc[rr][cc] : 0.0f;
                }
            }
        }
    }
}

torch::Tensor triumm_forward(torch::Tensor A, torch::Tensor B) {
    const int n = (int)A.size(0);
    auto C = torch::empty_like(A);
    dim3 block(16, 16);
    dim3 grid((n + TILE_N - 1) / TILE_N, (n + TILE_M - 1) / TILE_M);
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    triumm64_kernel<<<grid, block, 0, stream>>>(A.data_ptr<float>(), B.data_ptr<float>(), C.data_ptr<float>(), n);
    return C;
}
"""

_triumm_ext = load_inline(
    name="triumm_upper64_ext",
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=["triumm_forward"],
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3", "--use_fast_math"],
    verbose=False,
)

class ModelNew(nn.Module):
    def __init__(self):
        super(ModelNew, self).__init__()

    def forward(self, A, B):
        return _triumm_ext.triumm_forward(A, B)