submission 824647

import torch
from torch.utils.cpp_extension import load_inline
from task import input_t, output_t

# -----------------------------------------------------------------------------
# C++ / CUDA Inline Source Code for Optimized Fused Batched QR
# -----------------------------------------------------------------------------
cuda_source = """
#include <torch/extension.h>
#include <cuda_runtime.h>
#include <math.h>

__inline__ __device__ float blockReduceSum(float val) {
    __shared__ float shared[32];
    int lane = threadIdx.x & 31;
    int wid = threadIdx.x >> 5;

    // Warp-level reduction
    for (int offset = 16; offset > 0; offset /= 2) {
        val += __shfl_down_sync(0xffffffff, val, offset);
    }

    if (lane == 0) shared[wid] = val;
    __syncthreads();

    // Read from shared memory only if that warp existed
    val = (threadIdx.x < (blockDim.x >> 5))? shared[lane] : 0.0f;

    if (wid == 0) {
        for (int offset = 16; offset > 0; offset /= 2) {
            val += __shfl_down_sync(0xffffffff, val, offset);
        }
    }
    return val;
}

__global__ void unblocked_geqrf_kernel(
    const float* __restrict__ A,
    float* __restrict__ H,
    float* __restrict__ tau,
    int batch_size,
    int n
) {
    int b = blockIdx.x;
    if (b >= batch_size) return;

    extern __shared__ float sA[]; // size n * n

    // Load entire A to sA (coalesced)
    for (int i = threadIdx.x; i < n * n; i += blockDim.x) {
        sA[i] = A[b * n * n + i];
    }
    __syncthreads();

    // Perform QR factorization in shared memory
    for (int k = 0; k < n; k++) {
        // Compute norm-squared of column k below diagonal
        float local_sum = 0.0f;
        for (int i = k + 1 + threadIdx.x; i < n; i += blockDim.x) {
            float val = sA[i * n + k];
            local_sum += val * val;
        }
        float S = blockReduceSum(local_sum);

        __shared__ float s_tau, s_vscale, s_beta;
        if (threadIdx.x == 0) {
            float x_k = sA[k * n + k];
            if (S > 0.0f) {
                float g = sqrtf(x_k * x_k + S);
                float beta = (x_k >= 0.0f)? -g : g;
                s_beta = beta;
                s_tau = (beta - x_k) / beta;
                s_vscale = 1.0f / (x_k - beta);
            } else {
                s_beta = x_k;
                s_tau = 0.0f;
                s_vscale = 0.0f;
            }
            tau[b * n + k] = s_tau;
        }
        __syncthreads();

        float tau_k = s_tau;
        float vscale = s_vscale;

        if (threadIdx.x == 0) {
            sA[k * n + k] = s_beta;
        }
        for (int i = k + 1 + threadIdx.x; i < n; i += blockDim.x) {
            sA[i * n + k] *= vscale;
        }
        __syncthreads();

        // Apply Householder reflector to trailing columns (k+1 to n-1)
        if (tau_k != 0.0f) {
            for (int j = k + 1 + threadIdx.x; j < n; j += blockDim.x) {
                float d = sA[k * n + j];
                for (int i = k + 1; i < n; i++) {
                    d += sA[i * n + k] * sA[i * n + j];
                }
                sA[k * n + j] -= tau_k * d;
                for (int i = k + 1; i < n; i++) {
                    sA[i * n + j] -= tau_k * sA[i * n + k] * d;
                }
            }
        }
        __syncthreads();
    }

    // Write back sA to H (coalesced)
    for (int i = threadIdx.x; i < n * n; i += blockDim.x) {
        H[b * n * n + i] = sA[i];
    }
}

__global__ void blocked_geqrf_kernel(
    const float* __restrict__ A,
    float* __restrict__ H,
    float* __restrict__ tau,
    int batch_size,
    int n
) {
    int b = blockIdx.x;
    if (b >= batch_size) return;

    // Load entire A to H first (coalesced)
    for (int i = threadIdx.x; i < n * n; i += blockDim.x) {
        H[b * n * n + i] = A[b * n * n + i];
    }
    __syncthreads();

    // Allocate shared memory layout
    extern __shared__ float smem[];
    float* s_V = smem; // size n * 32
    float* s_W = smem + n * 32; // size 32 * n
    float* s_T = smem + n * 32 + 32 * n; // size 32 * 32

    // Pointer to this batch element's H and tau
    float* Hb = H + b * n * n;
    float* taub = tau + b * n;

    int nb = 32;
    int num_blocks = (n + nb - 1) / nb;

    for (int p = 0; p < num_blocks; p++) {
        int k_start = p * nb;
        int nb_actual = min(nb, n - k_start);
        int m = n - k_start;

        // 1. Factor the panel Hb[k_start:n, k_start:k_start+nb_actual-1]
        // Load panel to s_V
        for (int i = threadIdx.x; i < m * nb_actual; i += blockDim.x) {
            int r = i / nb_actual;
            int c = i % nb_actual;
            s_V[r * nb + c] = Hb[(k_start + r) * n + k_start + c];
        }
        __syncthreads();

        // Perform panel factorization in shared memory s_V
        for (int k = 0; k < nb_actual; k++) {
            float local_sum = 0.0f;
            for (int i = k + 1 + threadIdx.x; i < m; i += blockDim.x) {
                float val = s_V[i * nb + k];
                local_sum += val * val;
            }
            float S = blockReduceSum(local_sum);
            
            __shared__ float s_tau, s_vscale, s_beta;
            if (threadIdx.x == 0) {
                float x_k = s_V[k * nb + k];
                if (S > 0.0f) {
                    float g = sqrtf(x_k * x_k + S);
                    float beta = (x_k >= 0.0f)? -g : g;
                    s_beta = beta;
                    s_tau = (beta - x_k) / beta;
                    s_vscale = 1.0f / (x_k - beta);
                } else {
                    s_beta = x_k;
                    s_tau = 0.0f;
                    s_vscale = 0.0f;
                }
                taub[k_start + k] = s_tau;
            }
            __syncthreads();

            float tau_k = s_tau;
            float vscale = s_vscale;

            if (threadIdx.x == 0) {
                s_V[k * nb + k] = s_beta;
            }
            for (int i = k + 1 + threadIdx.x; i < m; i += blockDim.x) {
                s_V[i * nb + k] *= vscale;
            }
            __syncthreads();

            // Apply to trailing columns of the panel (columns k+1 to nb_actual-1)
            if (tau_k != 0.0f) {
                for (int j = k + 1 + threadIdx.x; j < nb_actual; j += blockDim.x) {
                    float d = s_V[k * nb + j];
                    for (int i = k + 1; i < m; i++) {
                        d += s_V[i * nb + k] * s_V[i * nb + j];
                    }
                    s_V[k * nb + j] -= tau_k * d;
                    for (int i = k + 1; i < m; i++) {
                        s_V[i * nb + j] -= tau_k * s_V[i * nb + k] * d;
                    }
                }
            }
            __syncthreads();
        }

        // Write the panel back to Hb
        for (int i = threadIdx.x; i < m * nb_actual; i += blockDim.x) {
            int r = i / nb_actual;
            int c = i % nb_actual;
            Hb[(k_start + r) * n + k_start + c] = s_V[r * nb + c];
        }
        __syncthreads();

        // 2. Compute T matrix in shared memory s_T
        for (int i = threadIdx.x; i < 1024; i += blockDim.x) {
            s_T[i] = 0.0f;
        }
        __syncthreads();

        for (int j = 0; j < nb_actual; j++) {
            float tau_j = taub[k_start + j];
            if (threadIdx.x == 0) {
                s_T[j * nb + j] = tau_j;
            }
            __syncthreads();

            if (j > 0 && tau_j != 0.0f) {
                __shared__ float s_w[32];
                if (threadIdx.x < j) {
                    int c = threadIdx.x;
                    float dot = s_V[j * nb + c]; // V[j, c]
                    for (int i = j + 1; i < m; i++) {
                        dot += s_V[i * nb + c] * s_V[i * nb + j];
                    }
                    s_w[c] = dot;
                }
                __syncthreads();

                if (threadIdx.x < j) {
                    int r = threadIdx.x;
                    float sum = 0.0f;
                    for (int c = r; c < j; c++) {
                        sum += s_T[r * nb + c] * s_w[c];
                    }
                    s_T[r * nb + j] = -tau_j * sum;
                }
                __syncthreads();
            }
        }

        // 3. Update the trailing matrix Hb[k_start:n, k_start+nb_actual:n]
        int trailing_cols = n - k_start - nb_actual;
        if (trailing_cols > 0) {
            for (int col = threadIdx.x; col < trailing_cols; col += blockDim.x) {
                float reg_W[32] = {0.0f};
                for (int i = 0; i < m; i++) {
                    float H_val = Hb[(k_start + i) * n + k_start + nb_actual + col];
                    for (int r = 0; r < nb_actual; r++) {
                        if (i == r) {
                            reg_W[r] += H_val;
                        } else if (i > r) {
                            reg_W[r] += s_V[i * nb + r] * H_val;
                        }
                    }
                }
                for (int r = 0; r < nb_actual; r++) {
                    s_W[r * trailing_cols + col] = reg_W[r];
                }
            }
            __syncthreads();

            for (int col = threadIdx.x; col < trailing_cols; col += blockDim.x) {
                float reg_Y[32] = {0.0f};
                for (int r = 0; r < nb_actual; r++) {
                    float sum = 0.0f;
                    for (int c = 0; c <= r; c++) {
                        sum += s_T[c * nb + r] * s_W[c * trailing_cols + col];
                    }
                    reg_Y[r] = sum;
                }
                for (int r = 0; r < nb_actual; r++) {
                    s_W[r * trailing_cols + col] = reg_Y[r];
                }
            }
            __syncthreads();

            for (int col = threadIdx.x; col < trailing_cols; col += blockDim.x) {
                for (int i = 0; i < m; i++) {
                    float sum = 0.0f;
                    for (int r = 0; r < nb_actual; r++) {
                        if (i == r) {
                            sum += s_W[r * trailing_cols + col];
                        } else if (i > r) {
                            sum += s_V[i * nb + r] * s_W[r * trailing_cols + col];
                        }
                    }
                    Hb[(k_start + i) * n + k_start + nb_actual + col] -= sum;
                }
            }
            __syncthreads();
        }
    }
}

void run_unblocked_geqrf(torch::Tensor A, torch::Tensor H, torch::Tensor tau) {
    int batch_size = A.size(0);
    int n = A.size(1);
    int threads = 256;
    int smem_size = n * n * sizeof(float);

    auto kernel = unblocked_geqrf_kernel;
    cudaFuncSetAttribute((const void*)kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size);

    kernel<<<batch_size, threads, smem_size>>>(
        A.data_ptr<float>(),
        H.data_ptr<float>(),
        tau.data_ptr<float>(),
        batch_size, n
    );
}

void run_blocked_geqrf(torch::Tensor A, torch::Tensor H, torch::Tensor tau) {
    int batch_size = A.size(0);
    int n = A.size(1);
    int threads = 512;
    int smem_size = (n * 32 + 32 * n + 32 * 32) * sizeof(float);

    auto kernel = blocked_geqrf_kernel;
    cudaFuncSetAttribute((const void*)kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size);

    kernel<<<batch_size, threads, smem_size>>>(
        A.data_ptr<float>(),
        H.data_ptr<float>(),
        tau.data_ptr<float>(),
        batch_size, n
    );
}
"""

cpp_source = """
void run_unblocked_geqrf(torch::Tensor A, torch::Tensor H, torch::Tensor tau);
void run_blocked_geqrf(torch::Tensor A, torch::Tensor H, torch::Tensor tau);
"""

# -----------------------------------------------------------------------------
# JIT Compiler & Module Cache
# -----------------------------------------------------------------------------
_custom_geqrf_module = None

def get_geqrf_module():
    global _custom_geqrf_module
    if _custom_geqrf_module is None:
        try:
            _custom_geqrf_module = load_inline(
                name="custom_geqrf_fused_v4",
                cpp_sources=cpp_source,
                cuda_sources=cuda_source,
                functions=["run_unblocked_geqrf", "run_blocked_geqrf"],
                with_cuda=True,
                extra_cuda_cflags=["-O3", "--use_fast_math"]
            )
        except Exception as e:
            _custom_geqrf_module = False
    return _custom_geqrf_module

# -----------------------------------------------------------------------------
# Main Custom Kernel Dispatch
# -----------------------------------------------------------------------------
def custom_kernel(data: input_t) -> output_t:
    """
    Highly optimized batched geqrf solver.
    Uses unblocked shared memory QR for small matrices (n <= 176)
    and blocked shared memory QR for medium matrices (176 < n <= 512).
    """
    if data.is_cuda:
        batch_size, n, _ = data.shape
        if n <= 512:
            module = get_geqrf_module()
            if module:
                A_contiguous = data.contiguous()
                H = torch.empty_like(A_contiguous)
                tau = torch.empty((batch_size, n), dtype=data.dtype, device=data.device)
                
                if n <= 176:
                    module.run_unblocked_geqrf(A_contiguous, H, tau)
                else:
                    module.run_blocked_geqrf(A_contiguous, H, tau)
                return H, tau

    # Fallback to PyTorch native geqrf
    return torch.geqrf(data)