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

cuda_src = r'''
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda.h>
#include <cuda_runtime.h>

#define BS 64
#define TT 16

__global__ void tri_matmul_fp32_64_kernel(const float* __restrict__ A,
                                         const float* __restrict__ B,
                                         float* __restrict__ C,
                                         int N, int tiles) {
    int tj = blockIdx.x;
    int ti = blockIdx.y;
    if (ti >= tiles || tj >= tiles) return;

    int tx = threadIdx.x;
    int ty = threadIdx.y;
    int tid = ty * TT + tx;
    int row_base = ti * BS + ty;
    int col_base = tj * BS + tx;

    // Upper triangular output tiles are exactly zero.  Writing them here avoids a full-matrix memset.
    if (ti < tj) {
        float* p = C + row_base * N + col_base;
        p[0] = 0.0f;      p[16] = 0.0f;      p[32] = 0.0f;      p[48] = 0.0f;
        p[16*N] = 0.0f;   p[16*N+16] = 0.0f; p[16*N+32] = 0.0f; p[16*N+48] = 0.0f;
        p[32*N] = 0.0f;   p[32*N+16] = 0.0f; p[32*N+32] = 0.0f; p[32*N+48] = 0.0f;
        p[48*N] = 0.0f;   p[48*N+16] = 0.0f; p[48*N+32] = 0.0f; p[48*N+48] = 0.0f;
        return;
    }

    __shared__ float As[BS][BS + 1];
    __shared__ float Bs[BS][BS + 1];

    float c00=0.0f,c01=0.0f,c02=0.0f,c03=0.0f;
    float c10=0.0f,c11=0.0f,c12=0.0f,c13=0.0f;
    float c20=0.0f,c21=0.0f,c22=0.0f,c23=0.0f;
    float c30=0.0f,c31=0.0f,c32=0.0f,c33=0.0f;

    // For output tile (ti,tj), only k in [tj*64, (ti+1)*64) can contribute.
    int k0 = tj * BS;
    int k1 = (ti + 1) * BS;

    for (int kk = k0; kk < k1; kk += BS) {
        #pragma unroll
        for (int m = 0; m < 16; ++m) {
            int idx = tid + m * 256;
            int r = idx >> 6;
            int c = idx & 63;
            As[r][c] = A[(ti * BS + r) * N + (kk + c)];
            Bs[r][c] = B[(kk + r) * N + (tj * BS + c)];
        }
        __syncthreads();

        #pragma unroll
        for (int k = 0; k < BS; ++k) {
            float a0 = As[ty     ][k];
            float a1 = As[ty + 16][k];
            float a2 = As[ty + 32][k];
            float a3 = As[ty + 48][k];
            float b0 = Bs[k][tx     ];
            float b1 = Bs[k][tx + 16];
            float b2 = Bs[k][tx + 32];
            float 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();
    }

    float* p = C + row_base * N + col_base;
    p[0]       = c00; p[16]       = c01; p[32]       = c02; p[48]       = c03;
    p[16*N]    = c10; p[16*N+16]  = c11; p[16*N+32]  = c12; p[16*N+48]  = c13;
    p[32*N]    = c20; p[32*N+16]  = c21; p[32*N+32]  = c22; p[32*N+48]  = c23;
    p[48*N]    = c30; p[48*N+16]  = c31; p[48*N+32]  = c32; p[48*N+48]  = c33;
}

torch::Tensor tri_matmul_cuda(torch::Tensor A, torch::Tensor B) {
    TORCH_CHECK(A.is_cuda() && B.is_cuda(), "A and B must be CUDA tensors");
    TORCH_CHECK(A.scalar_type() == torch::kFloat32 && B.scalar_type() == torch::kFloat32, "float32 only");
    TORCH_CHECK(A.is_contiguous() && B.is_contiguous(), "contiguous tensors required");
    int64_t N64 = A.size(0);
    auto C = torch::empty({N64, N64}, A.options());
    int N = (int)N64;
    int tiles = (N + BS - 1) / BS;
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    dim3 grid(tiles, tiles);
    dim3 block(TT, TT);
    tri_matmul_fp32_64_kernel<<<grid, block, 0, stream>>>(A.data_ptr<float>(), B.data_ptr<float>(), C.data_ptr<float>(), N, tiles);
    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return C;
}
'''

_ext = load_inline(
    name='tri_matmul_fp32_ext_final_submit',
    cpp_sources=cpp_src,
    cuda_sources=cuda_src,
    functions=['tri_matmul'],
    extra_cflags=['-O3'],
    extra_cuda_cflags=['-O3'],
    verbose=False,
)

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