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 triu_mm_cuda(torch::Tensor A, torch::Tensor B);
torch::Tensor triu_mm(torch::Tensor A, torch::Tensor B) { return triu_mm_cuda(A,B); }
'''

cuda_src = r'''
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
#define TILE 64
#define TS 16
#define RPT 4
__global__ void tri64_kernel(const float* __restrict__ A,const float* __restrict__ B,float* __restrict__ C,int N){
 int bj=blockIdx.x, bi=blockIdx.y; int tx=threadIdx.x, ty=threadIdx.y; int tid=ty*TS+tx;
 int rowBase=bi*TILE+ty*RPT; int colBase=bj*TILE+tx*RPT;
 if(bi>bj){
  #pragma unroll
  for(int ii=0;ii<RPT;++ii){ int r=rowBase+ii;
   #pragma unroll
   for(int jj=0;jj<RPT;++jj) C[r*N+colBase+jj]=0.0f;
  }
  return;
 }
 __shared__ float As[TILE][TILE+1];
 __shared__ float Bs[TILE][TILE+1];
 float c[4][4];
 #pragma unroll
 for(int i=0;i<4;i++){
  #pragma unroll
  for(int j=0;j<4;j++) c[i][j]=0.0f;
 }
 #pragma unroll 1
 for(int bk=bi; bk<=bj; ++bk){
  int kbase=bk*TILE;
  #pragma unroll
  for(int q=0;q<16;++q){
   int idx=tid+q*256; int lr=idx>>6; int lc=idx&63;
   As[lr][lc]=A[(bi*TILE+lr)*N + kbase+lc];
   Bs[lr][lc]=B[(kbase+lr)*N + bj*TILE+lc];
  }
  __syncthreads();
  #pragma unroll
  for(int k=0;k<TILE;++k){
   float a0=As[ty*4+0][k]; float a1=As[ty*4+1][k]; float a2=As[ty*4+2][k]; float a3=As[ty*4+3][k];
   float b0=Bs[k][tx*4+0]; float b1=Bs[k][tx*4+1]; float b2=Bs[k][tx*4+2]; float b3=Bs[k][tx*4+3];
   c[0][0]=fmaf(a0,b0,c[0][0]); c[0][1]=fmaf(a0,b1,c[0][1]); c[0][2]=fmaf(a0,b2,c[0][2]); c[0][3]=fmaf(a0,b3,c[0][3]);
   c[1][0]=fmaf(a1,b0,c[1][0]); c[1][1]=fmaf(a1,b1,c[1][1]); c[1][2]=fmaf(a1,b2,c[1][2]); c[1][3]=fmaf(a1,b3,c[1][3]);
   c[2][0]=fmaf(a2,b0,c[2][0]); c[2][1]=fmaf(a2,b1,c[2][1]); c[2][2]=fmaf(a2,b2,c[2][2]); c[2][3]=fmaf(a2,b3,c[2][3]);
   c[3][0]=fmaf(a3,b0,c[3][0]); c[3][1]=fmaf(a3,b1,c[3][1]); c[3][2]=fmaf(a3,b2,c[3][2]); c[3][3]=fmaf(a3,b3,c[3][3]);
  }
  __syncthreads();
 }
 #pragma unroll
 for(int ii=0;ii<4;++ii){ int r=rowBase+ii;
  #pragma unroll
  for(int jj=0;jj<4;++jj) C[r*N+colBase+jj]=c[ii][jj];
 }
}
torch::Tensor triu_mm_cuda(torch::Tensor A,torch::Tensor B){int N=(int)A.size(0); auto C=torch::empty_like(A); dim3 block(16,16); dim3 grid((N+63)/64,(N+63)/64); tri64_kernel<<<grid,block>>>(A.data_ptr<float>(),B.data_ptr<float>(),C.data_ptr<float>(),N); return C;}
'''
_ext=load_inline(name='tri_mm_custom64_final',cpp_sources=cpp_src,cuda_sources=cuda_src,functions=['triu_mm'],with_cuda=True,extra_cflags=['-O3'],extra_cuda_cflags=['-O3'])
class ModelNew(nn.Module):
 def __init__(self): super(ModelNew,self).__init__()
 def forward(self,A,B): return _ext.triu_mm(A,B)