import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Simple model that performs a matrix multiplication (C = A * B) where A and B are lower triangular matrices. 
    """
    def __init__(self):
        super(Model, self).__init__()
    
    def forward(self, A, B):
        """
        Performs matrix multiplication of lower triangular matrices A and B.

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

        Returns:
            torch.Tensor: The result of matrix multiplication C of shape (N, N).
        """
        return torch.tril(torch.matmul(A, B))

M = 4096

def get_inputs():
    A = torch.rand(M, M)
    B = torch.rand(M, M)
    A = torch.tril(A)
    B = torch.tril(B)
    return [A, B]

def get_init_inputs():
    return []  # No special initialization inputs needed

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

cpp_src = r"""
#include <torch/extension.h>
torch::Tensor tri_matmul_cuda(torch::Tensor A, torch::Tensor B);
"""

cuda_src = r"""
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <c10/cuda/CUDAGuard.h>
#include <c10/cuda/CUDAException.h>
#define TILE 64
__global__ __launch_bounds__(256, 2) void tri_sgemm64_kernel(const float* __restrict__ A, const float* __restrict__ B, float* __restrict__ C, int N) {
    int tj=blockIdx.x, ti=blockIdx.y;
    int tx=threadIdx.x, ty=threadIdx.y;
    int linear = ty*16 + tx;
    if (ti < tj) {
        #pragma unroll
        for (int off=0; off<4096; off+=256) {
            int l=linear+off; int rr=l>>6; int cc=l&63;
            C[(ti*TILE+rr)*N + (tj*TILE+cc)] = 0.0f;
        }
        return;
    }
    int rbase = ti*TILE + ty;
    int cbase = tj*TILE + tx;
    __shared__ float As[TILE][TILE+1];
    __shared__ float Bs[TILE][TILE+1];
    float c00=0,c01=0,c02=0,c03=0;
    float c10=0,c11=0,c12=0,c13=0;
    float c20=0,c21=0,c22=0,c23=0;
    float c30=0,c31=0,c32=0,c33=0;
    for (int kt=tj; kt<=ti; ++kt) {
        int kbase=kt*TILE;
        #pragma unroll
        for (int off=0; off<4096; off+=256) {
            int l=linear+off; int rr=l>>6; int cc=l&63;
            As[rr][cc] = A[(ti*TILE + rr)*N + (kbase + cc)];
            Bs[rr][cc] = B[(kbase + rr)*N + (tj*TILE + cc)];
        }
        __syncthreads();
        #pragma unroll
        for (int k=0;k<TILE;++k){
            float a0=As[ty][k], a1=As[ty+16][k], a2=As[ty+32][k], a3=As[ty+48][k];
            float b0=Bs[k][tx], b1=Bs[k][tx+16], b2=Bs[k][tx+32], b3=Bs[k][tx+48];
            c00=fmaf(a0,b0,c00); c01=fmaf(a0,b1,c01); c02=fmaf(a0,b2,c02); c03=fmaf(a0,b3,c03);
            c10=fmaf(a1,b0,c10); c11=fmaf(a1,b1,c11); c12=fmaf(a1,b2,c12); c13=fmaf(a1,b3,c13);
            c20=fmaf(a2,b0,c20); c21=fmaf(a2,b1,c21); c22=fmaf(a2,b2,c22); c23=fmaf(a2,b3,c23);
            c30=fmaf(a3,b0,c30); c31=fmaf(a3,b1,c31); c32=fmaf(a3,b2,c32); c33=fmaf(a3,b3,c33);
        }
        __syncthreads();
    }
    int r0=rbase, r1=rbase+16, r2=rbase+32, r3=rbase+48;
    int c0=cbase, c1=cbase+16, c2=cbase+32, c3=cbase+48;
    C[r0*N+c0]=c00; C[r0*N+c1]=c01; C[r0*N+c2]=c02; C[r0*N+c3]=c03;
    C[r1*N+c0]=c10; C[r1*N+c1]=c11; C[r1*N+c2]=c12; C[r1*N+c3]=c13;
    C[r2*N+c0]=c20; C[r2*N+c1]=c21; C[r2*N+c2]=c22; C[r2*N+c3]=c23;
    C[r3*N+c0]=c30; C[r3*N+c1]=c31; C[r3*N+c2]=c32; C[r3*N+c3]=c33;
}

torch::Tensor tri_matmul_cuda(torch::Tensor A, torch::Tensor B) {
    const at::cuda::OptionalCUDAGuard device_guard(device_of(A));
    int N=A.size(0);
    auto C=torch::empty_like(A);
    dim3 block(16,16); dim3 grid((N+TILE-1)/TILE,(N+TILE-1)/TILE);
    tri_sgemm64_kernel<<<grid,block,0,at::cuda::getCurrentCUDAStream()>>>(A.data_ptr<float>(),B.data_ptr<float>(),C.data_ptr<float>(),N);
    C10_CUDA_KERNEL_LAUNCH_CHECK(); return C;
}
"""

_ext=load_inline(name="tri_sgemm64_final_ext", cpp_sources=cpp_src, cuda_sources=cuda_src, functions=["tri_matmul_cuda"], extra_cuda_cflags=["-O3","--use_fast_math"], extra_cflags=["-O3"], verbose=False)

class ModelNew(nn.Module):
    def __init__(self):
        super(ModelNew,self).__init__()
    def forward(self,A,B):
        return _ext.tri_matmul_cuda(A,B)