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

cpp_source = """
#include <torch/extension.h>

torch::Tensor cholesky_cuda_forward(const torch::Tensor& input);
"""

cuda_source = """
#include <torch/extension.h>
#include <cuda_runtime.h>
#include <cmath> // For std::sqrt

#define BLOCK_SIZE 256

template <typename T>
__global__ void cholesky_kernel(T* L, const T* A, int n) {
    // Each block processes one matrix in the batch.
    const int matrix_idx = blockIdx.x;
    const T* A_matrix = A + matrix_idx * n * n;
    T* L_matrix = L + matrix_idx * n * n;

    __shared__ T sdata[BLOCK_SIZE];

    // Main loop over columns
    for (int j = 0; j < n; ++j) {
        // --- Stage 1: Compute diagonal element L[j, j] ---
        T sum_sq = 0.0; // Use double literal for double precision
        
        // Parallel reduction for the sum part
        if (threadIdx.x < j) {
            T val = L_matrix[j * n + threadIdx.x];
            sum_sq = val * val;
        }
        sdata[threadIdx.x] = sum_sq;
        __syncthreads();

        // Reduce in shared memory
        for (int s = blockDim.x / 2; s > 0; s >>= 1) {
            if (threadIdx.x < s && threadIdx.x + s < j) {
                sdata[threadIdx.x] += sdata[threadIdx.x + s];
            }
            __syncthreads();
        }
        
        // Thread 0 computes the final L[j, j]
        if (threadIdx.x == 0) {
            T ajj = A_matrix[j * n + j];
            T final_sum = (j > 0) ? sdata[0] : 0.0;
            // Use sqrt, which works for both float and double
            L_matrix[j * n + j] = sqrt(ajj - final_sum);
        }
        __syncthreads(); // Sync to make L[j, j] visible to all threads

        // --- Stage 2: Compute off-diagonal elements L[i, j] for i > j ---
        T ljj = L_matrix[j * n + j]; 

        // Parallel update for the rest of the column using grid-stride loop
        for (int i = j + 1 + threadIdx.x; i < n; i += blockDim.x) {
            T sum_dot = 0.0;
            for (int k = 0; k < j; ++k) {
                sum_dot += L_matrix[i * n + k] * L_matrix[j * n + k];
            }
            L_matrix[i * n + j] = (A_matrix[i * n + j] - sum_dot) / ljj;
        }
        __syncthreads(); 
    }
}

torch::Tensor cholesky_cuda_forward(const torch::Tensor& input) {
    TORCH_CHECK(input.is_cuda(), "Input must be a CUDA tensor");
    TORCH_CHECK(input.dim() >= 2 && input.size(-1) == input.size(-2), "Input must be batches of square matrices");
    TORCH_CHECK(input.is_contiguous(), "Input must be contiguous");

    auto L = torch::zeros_like(input);
    
    const auto batch_dims = input.sizes().slice(0, input.dim() - 2);
    const int64_t batch_size = c10::size_to_dim_(input.dim() - 2, batch_dims);
    const int n = input.size(-1);

    TORCH_CHECK(n <= BLOCK_SIZE, "Matrix dimension exceeds block size for this simple kernel");
    const dim3 grid(batch_size);
    const dim3 block(BLOCK_SIZE);

    // AT_DISPATCH_FLOATING_TYPES handles both float and double
    AT_DISPATCH_FLOATING_TYPES(input.scalar_type(), "cholesky_kernel", ([&] {
        cholesky_kernel<scalar_t><<<grid, block>>>(
            L.data_ptr<scalar_t>(),
            input.data_ptr<scalar_t>(),
            n
        );
    }));

    return L;
}

"""

class ModelNew(nn.Module):
    def __init__(self):
        super(ModelNew, self).__init__()
        self.op = load_inline(
            name='cholesky_op_double',
            cpp_sources=cpp_source,
            cuda_sources=cuda_source,
            functions=['cholesky_cuda_forward'],
            verbose=False
        )

    def forward(self, input_tensor: torch.Tensor) -> torch.Tensor:
        original_shape = input_tensor.shape
        n = original_shape[-1]
        input_flat = input_tensor.reshape(-1, n, n)
        
        output_flat = self.op.cholesky_cuda_forward(input_flat)
        
        return output_flat.reshape(original_shape)