submission 800090

#!POPCORN leaderboard qr_v2
#!POPCORN gpu B200

# from-scratch batched blocked householder qr. panel+T fp32 (correct, from lmaf).
# heavy trailing gemms (W=V^T C, C-=V Y) on tf32x3 tensor cores via nvcuda::wmma,
# now BLOCK-TILED: 4 warps per block share one staged A operand (V) with lean
# single-buffer staging for higher resident CTA count on skinny updates.
# W is consumed inside its producer: each warp applies T^T to its 16x16 W tile
# and writes Y directly. the panel is staged row-wise in shared memory when it fits.

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

CUDA_SRC = r"""
#include <torch/extension.h>
#include <cuda_runtime.h>
#include <mma.h>
#include <math.h>
#include <stdexcept>
#include <vector>

using namespace nvcuda;

#define NB 16
#define NWARP 4                 // warps per gemm block
#define NTILE (NWARP * 16)      // cols handled per block

__device__ float block_sum(float x, float* scratch) {
    int tid = threadIdx.x;
    int lane = tid & 31;
    int wid = tid >> 5;
    #pragma unroll
    for (int o = 16; o > 0; o >>= 1) x += __shfl_down_sync(0xffffffff, x, o);
    if (lane == 0) scratch[wid] = x;
    __syncthreads();
    int nw = (blockDim.x + 31) >> 5;
    x = (tid < nw) ? scratch[tid] : 0.0f;
    if (wid == 0) {
        #pragma unroll
        for (int o = 16; o > 0; o >>= 1) x += __shfl_down_sync(0xffffffff, x, o);
    }
    if (tid == 0) scratch[0] = x;
    __syncthreads();
    return scratch[0];
}

template<int N>
__device__ void block_sums(float (&x)[N], int count, float* scratch, float* out) {
    int tid = threadIdx.x;
    int lane = tid & 31;
    int wid = tid >> 5;
    int nw = (blockDim.x + 31) >> 5;

    for (int p = 0; p < count; ++p) {
        float v = x[p];
        #pragma unroll
        for (int o = 16; o > 0; o >>= 1) v += __shfl_down_sync(0xffffffff, v, o);
        if (lane == 0) scratch[wid * N + p] = v;
    }
    __syncthreads();

    if (wid == 0) {
        for (int p = lane; p < count; p += 32) {
            float v = 0.0f;
            for (int w = 0; w < nw; ++w) v += scratch[w * N + p];
            out[p] = v;
        }
    }
    __syncthreads();
}

__global__ void panel_kernel_lwoc(float* __restrict__ A, float* __restrict__ tau,
                                  float* __restrict__ V, float* __restrict__ T,
                                  int n, int j0, int ib) {
    int b = blockIdx.x;
    float* M = A + (size_t)b * n * n;
    float* Tb = T + (size_t)b * NB * NB;
    float* Vb = V + (size_t)b * n * NB;
    float* taub = tau + (size_t)b * n;
    int tid = threadIdx.x, nt = blockDim.x;

    extern __shared__ float smem[];
    float* col = smem;
    float* scratch = col + n;
    float* ds = scratch + 32 * NB;
    float* bc = ds + NB;
    float* tmp = bc + 4;

    int m = n - j0;
    for (int idx = tid; idx < m * ib; idx += nt) Vb[idx] = 0.0f;
    for (int idx = tid; idx < NB * NB; idx += nt) Tb[idx] = 0.0f;
    __syncthreads();

    for (int kk = 0; kk < ib; ++kk) {
        int k = j0 + kk;
        float local = 0.0f;
        for (int i = k + tid; i < n; i += nt) { float x = M[(size_t)i * n + k]; col[i] = x; local += x * x; }
        float sumsq = block_sum(local, scratch);
        float alpha = col[k];
        if (tid == 0) {
            float norm = sqrtf(sumsq);
            float beta = (alpha >= 0.0f) ? -norm : norm;
            float below = sumsq - alpha * alpha;
            float tauk, scale;
            if (norm == 0.0f || below <= 0.0f) { tauk = 0.0f; scale = 1.0f; beta = alpha; }
            else { tauk = (beta - alpha) / beta; scale = alpha - beta; }
            bc[0] = tauk; bc[1] = beta; bc[2] = scale; taub[k] = tauk;
        }
        __syncthreads();
        float tauk = bc[0], beta = bc[1], scale = bc[2];
        float inv = 1.0f / scale;

        for (int i = k + 1 + tid; i < n; i += nt) {
            float vi = (tauk == 0.0f) ? 0.0f : col[i] * inv;
            col[i] = vi; M[(size_t)i * n + k] = vi; Vb[(size_t)(i - j0) * NB + kk] = vi;
        }
        if (tid == 0) { col[k] = 1.0f; M[(size_t)k * n + k] = beta; Vb[(size_t)kk * NB + kk] = 1.0f; }
        __syncthreads();

        if (tauk != 0.0f) {
            int lane = tid & 31, wid = tid >> 5, nwarps = nt >> 5;
            for (int jj = kk + 1 + wid; jj < ib; jj += nwarps) {
                int j = j0 + jj;
                float w = 0.0f;
                for (int i = k + lane; i < n; i += 32) w += col[i] * M[(size_t)i * n + j];
                #pragma unroll
                for (int o = 16; o > 0; o >>= 1) w += __shfl_down_sync(0xffffffff, w, o);
                w = __shfl_sync(0xffffffff, w, 0);
                float tw = tauk * w;
                for (int i = k + lane; i < n; i += 32) M[(size_t)i * n + j] -= col[i] * tw;
            }
        }
        __syncthreads();

        if (tid == 0) Tb[kk * NB + kk] = tauk;
        float dlocal[NB];
        #pragma unroll
        for (int p = 0; p < NB; ++p) dlocal[p] = 0.0f;
        for (int i = k + tid; i < n; i += nt) {
            float ci = col[i];
            #pragma unroll
            for (int p = 0; p < NB; ++p) {
                if (p < kk) dlocal[p] += Vb[(size_t)(i - j0) * NB + p] * ci;
            }
        }
        if (kk > 0) {
            block_sums<NB>(dlocal, kk, scratch, ds);
            for (int p = tid; p < kk; p += nt) {
                tmp[p] = -tauk * ds[p];
            }
            __syncthreads();
            for (int p = tid; p < kk; p += nt) {
                float s = 0.0f;
                for (int q = 0; q < kk; ++q) s += Tb[p * NB + q] * tmp[q];
                Tb[p * NB + kk] = s;
            }
        }
        __syncthreads();
    }
}

__global__ void panel_kernel(float* __restrict__ A, float* __restrict__ tau,
                             float* __restrict__ V, float* __restrict__ T,
                             int n, int j0, int ib) {
    int b = blockIdx.x;
    float* M = A + (size_t)b * n * n;
    float* Tb = T + (size_t)b * NB * NB;
    float* Vb = V + (size_t)b * n * NB;
    float* taub = tau + (size_t)b * n;
    int tid = threadIdx.x, nt = blockDim.x;
    int m = n - j0;

    extern __shared__ float smem[];
    float* sP = smem;
    float* scratch = sP + (size_t)m * ib;
    float* ds = scratch + 32 * NB;
    float* bc = ds + NB;
    float* tmp = bc + 4;

    for (int idx = tid; idx < m * ib; idx += nt) {
        int r = idx / ib;
        int c = idx - r * ib;
        sP[idx] = M[(size_t)(j0 + r) * n + (j0 + c)];
        Vb[(size_t)r * NB + c] = 0.0f;
    }
    for (int idx = tid; idx < NB * NB; idx += nt) Tb[idx] = 0.0f;
    __syncthreads();

    for (int kk = 0; kk < ib; ++kk) {
        float local = 0.0f;
        for (int r = kk + tid; r < m; r += nt) {
            float x = sP[(size_t)r * ib + kk];
            local += x * x;
        }
        float sumsq = block_sum(local, scratch);
        float alpha = sP[(size_t)kk * ib + kk];
        if (tid == 0) {
            float norm = sqrtf(sumsq);
            float beta = (alpha >= 0.0f) ? -norm : norm;
            float below = sumsq - alpha * alpha;
            float tauk, scale;
            if (norm == 0.0f || below <= 0.0f) { tauk = 0.0f; scale = 1.0f; beta = alpha; }
            else { tauk = (beta - alpha) / beta; scale = alpha - beta; }
            bc[0] = tauk; bc[1] = beta; bc[2] = scale; taub[j0 + kk] = tauk;
        }
        __syncthreads();
        float tauk = bc[0], beta = bc[1], scale = bc[2];
        float inv = 1.0f / scale;

        for (int r = kk + 1 + tid; r < m; r += nt) {
            float vi = (tauk == 0.0f) ? 0.0f : sP[(size_t)r * ib + kk] * inv;
            sP[(size_t)r * ib + kk] = vi;
            Vb[(size_t)r * NB + kk] = vi;
        }
        if (tid == 0) {
            sP[(size_t)kk * ib + kk] = beta;
            Vb[(size_t)kk * NB + kk] = 1.0f;
        }
        __syncthreads();

        if (tauk != 0.0f) {
            int lane = tid & 31, wid = tid >> 5, nwarps = nt >> 5;
            for (int jj = kk + 1 + wid; jj < ib; jj += nwarps) {
                float w = 0.0f;
                for (int r = kk + lane; r < m; r += 32) {
                    float v = (r == kk) ? 1.0f : sP[(size_t)r * ib + kk];
                    w += v * sP[(size_t)r * ib + jj];
                }
                #pragma unroll
                for (int o = 16; o > 0; o >>= 1) w += __shfl_down_sync(0xffffffff, w, o);
                w = __shfl_sync(0xffffffff, w, 0);
                float tw = tauk * w;
                for (int r = kk + lane; r < m; r += 32) {
                    float v = (r == kk) ? 1.0f : sP[(size_t)r * ib + kk];
                    sP[(size_t)r * ib + jj] -= v * tw;
                }
            }
        }
        __syncthreads();

        if (tid == 0) Tb[kk * NB + kk] = tauk;
        float dlocal[NB];
        #pragma unroll
        for (int p = 0; p < NB; ++p) dlocal[p] = 0.0f;
        for (int r = kk + tid; r < m; r += nt) {
            float ci = (r == kk) ? 1.0f : sP[(size_t)r * ib + kk];
            #pragma unroll
            for (int p = 0; p < NB; ++p) {
                if (p < kk) dlocal[p] += Vb[(size_t)r * NB + p] * ci;
            }
        }
        if (kk > 0) {
            block_sums<NB>(dlocal, kk, scratch, ds);
            for (int p = tid; p < kk; p += nt) tmp[p] = -tauk * ds[p];
            __syncthreads();
            for (int p = tid; p < kk; p += nt) {
                float s = 0.0f;
                for (int q = 0; q < kk; ++q) s += Tb[p * NB + q] * tmp[q];
                Tb[p * NB + kk] = s;
            }
        }
        __syncthreads();
    }

    for (int idx = tid; idx < m * ib; idx += nt) {
        int r = idx / ib;
        int c = idx - r * ib;
        M[(size_t)(j0 + r) * n + (j0 + c)] = sP[idx];
    }
}

// load+split+3mma helper: As is 16x8 (ld8), Bs slice is 8x16 inside an 8xNTILE (ld NTILE).
__device__ __forceinline__ void mma3(
    wmma::fragment<wmma::accumulator,16,16,8,float>& cf,
    const float* As, const float* Bslice) {
    wmma::fragment<wmma::matrix_a,16,16,8,wmma::precision::tf32,wmma::row_major> af, af_lo;
    wmma::fragment<wmma::matrix_b,16,16,8,wmma::precision::tf32,wmma::row_major> bf, bf_lo;
    wmma::load_matrix_sync(af, As, 8);
    wmma::load_matrix_sync(bf, Bslice, NTILE);
    #pragma unroll
    for (int t = 0; t < af.num_elements; t++) { float hi = wmma::__float_to_tf32(af.x[t]); af_lo.x[t] = wmma::__float_to_tf32(af.x[t]-hi); af.x[t]=hi; }
    #pragma unroll
    for (int t = 0; t < bf.num_elements; t++) { float hi = wmma::__float_to_tf32(bf.x[t]); bf_lo.x[t] = wmma::__float_to_tf32(bf.x[t]-hi); bf.x[t]=hi; }
    wmma::mma_sync(cf, af, bf, cf);
    wmma::mma_sync(cf, af_lo, bf, cf);
    wmma::mma_sync(cf, af, bf_lo, cf);
}

// Y = T^T @ (V^T @ C). block does 16 rows x NTILE cols; 4 warps share staged V^T tile.
__global__ void wmma_wy_kernel(const float* __restrict__ A, const float* __restrict__ V,
                               const float* __restrict__ T, float* __restrict__ Y,
                               int n, int j0, int ib, int m, int cols) {
    int tn = blockIdx.x, tm = blockIdx.y, b = blockIdx.z;
    int row0 = tm * 16, col0 = tn * NTILE;
    const float* Ab = A + (size_t)b * n * n + (size_t)j0 * n + (j0 + ib);
    const float* Vb = V + (size_t)b * n * NB;
    const float* Tb = T + (size_t)b * NB * NB;
    float* Yb = Y + (size_t)b * NB * n;

    __shared__ float As[16 * 8];
    __shared__ float Bs[8 * NTILE];
    __shared__ float Cs[NWARP * 256];
    int tid = threadIdx.x, w = tid >> 5;

    wmma::fragment<wmma::accumulator,16,16,8,float> cf;
    wmma::fill_fragment(cf, 0.0f);

    for (int k0 = 0; k0 < m; k0 += 8) {
        for (int idx = tid; idx < 16 * 8; idx += blockDim.x) {
            int r = idx >> 3, c = idx & 7, p = row0 + r, kk = k0 + c;
            As[idx] = (p < ib && kk < m) ? Vb[(size_t)kk * NB + p] : 0.0f;   // V^T[p,kk]=V[kk,p]
        }
        for (int idx = tid; idx < 8 * NTILE; idx += blockDim.x) {
            int c = idx / NTILE, r = idx % NTILE, kk = k0 + c, col = col0 + r;
            Bs[idx] = (kk < m && col < cols) ? Ab[(size_t)kk * n + col] : 0.0f;  // C[kk,col]
        }
        __syncthreads();
        mma3(cf, As, Bs + w * 16);
        __syncthreads();
    }
    wmma::store_matrix_sync(Cs + w * 256, cf, 16, wmma::mem_row_major);
    __syncwarp();
    for (int idx = (tid & 31); idx < 256; idx += 32) {
        int r = idx >> 4, c = idx & 15, p = row0 + r, col = col0 + w * 16 + c;
        if (p < ib && col < cols) {
            float y = 0.0f;
            #pragma unroll
            for (int q = 0; q < NB; ++q) {
                if (q < ib) y += Tb[(size_t)q * NB + p] * Cs[w * 256 + q * 16 + c];
            }
            Yb[(size_t)p * n + col] = y;
        }
    }
}

// C -= V @ Y. block does 16 rows x NTILE cols; 4 warps share staged V tile. K=ib.
__global__ void wmma_c_kernel(float* __restrict__ A, const float* __restrict__ V,
                              const float* __restrict__ Y, int n, int j0, int ib, int m, int cols) {
    int tn = blockIdx.x, tm = blockIdx.y, b = blockIdx.z;
    int row0 = tm * 16, col0 = tn * NTILE;
    float* Ab = A + (size_t)b * n * n + (size_t)j0 * n + (j0 + ib);
    const float* Vb = V + (size_t)b * n * NB;
    const float* Yb = Y + (size_t)b * NB * n;

    __shared__ float As[16 * 8];
    __shared__ float Bs[8 * NTILE];
    __shared__ float Cs[NWARP * 256];
    int tid = threadIdx.x, w = tid >> 5;

    wmma::fragment<wmma::accumulator,16,16,8,float> cf;
    wmma::fill_fragment(cf, 0.0f);

    for (int k0 = 0; k0 < ib; k0 += 8) {
        for (int idx = tid; idx < 16 * 8; idx += blockDim.x) {
            int r = idx >> 3, c = idx & 7, row = row0 + r, p = k0 + c;
            As[idx] = (row < m && p < ib) ? Vb[(size_t)row * NB + p] : 0.0f;   // V[row,p]
        }
        for (int idx = tid; idx < 8 * NTILE; idx += blockDim.x) {
            int c = idx / NTILE, r = idx % NTILE, p = k0 + c, col = col0 + r;
            Bs[idx] = (p < ib && col < cols) ? Yb[(size_t)p * n + col] : 0.0f;  // Y[p,col]
        }
        __syncthreads();
        mma3(cf, As, Bs + w * 16);
        __syncthreads();
    }
    wmma::store_matrix_sync(Cs + w * 256, cf, 16, wmma::mem_row_major);
    for (int idx = (tid & 31); idx < 256; idx += 32) {
        int r = idx >> 4, c = idx & 15, row = row0 + r, col = col0 + w * 16 + c;
        if (row < m && col < cols) Ab[(size_t)row * n + col] -= Cs[w * 256 + idx];
    }
}

std::vector<torch::Tensor> qr_forward(torch::Tensor A) {
    TORCH_CHECK(A.is_cuda(), "A must be cuda");
    TORCH_CHECK(A.scalar_type() == torch::kFloat32, "A must be fp32");
    auto Ac = A.contiguous();
    int64_t batch, n;
    bool was2 = (Ac.dim() == 2);
    if (Ac.dim() == 3) { batch = Ac.size(0); n = Ac.size(1); }
    else { batch = 1; n = Ac.size(0); Ac = Ac.reshape({1, n, n}); }

    auto a = Ac.clone();
    auto tau = torch::empty({batch, n}, A.options());
    auto V = torch::empty({batch, n, NB}, A.options());
    auto T = torch::empty({batch, NB, NB}, A.options());
    auto Y = torch::empty({batch, NB, n}, A.options());

    int threads = 256;
    size_t aux = 32 * NB + NB + 4 + NB;
    size_t cap = 200 * 1024;
    static bool attr_set = false;
    if (!attr_set) {
        cudaError_t attr = cudaFuncSetAttribute(
            panel_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, (int)cap);
        TORCH_CHECK(attr == cudaSuccess, "panel attr");
        attr_set = true;
    }

    for (int j0 = 0; j0 < (int)n; j0 += NB) {
        int ib = ((int)n - j0 < NB) ? ((int)n - j0) : NB;
        int m = (int)n - j0;
        size_t panel_shmem = ((size_t)m * ib + aux) * sizeof(float);
        if (panel_shmem <= cap) {
            panel_kernel<<<(int)batch, threads, panel_shmem>>>(
                a.data_ptr<float>(), tau.data_ptr<float>(), V.data_ptr<float>(),
                T.data_ptr<float>(), (int)n, j0, ib);
        } else {
            size_t old_shmem = ((size_t)n + aux) * sizeof(float);
            panel_kernel_lwoc<<<(int)batch, threads, old_shmem>>>(
                a.data_ptr<float>(), tau.data_ptr<float>(), V.data_ptr<float>(),
                T.data_ptr<float>(), (int)n, j0, ib);
        }
        TORCH_CHECK(cudaGetLastError() == cudaSuccess, "panel");

        int cols = (int)n - j0 - ib;
        if (cols <= 0) continue;

        dim3 wg((cols + NTILE - 1) / NTILE, (ib + 15) / 16, (int)batch);
        wmma_wy_kernel<<<wg, NWARP * 32>>>(a.data_ptr<float>(), V.data_ptr<float>(),
                                           T.data_ptr<float>(), Y.data_ptr<float>(),
                                           (int)n, j0, ib, m, cols);
        TORCH_CHECK(cudaGetLastError() == cudaSuccess, "wy");

        dim3 cg((cols + NTILE - 1) / NTILE, (m + 15) / 16, (int)batch);
        wmma_c_kernel<<<cg, NWARP * 32>>>(a.data_ptr<float>(), V.data_ptr<float>(),
                                          Y.data_ptr<float>(), (int)n, j0, ib, m, cols);
        TORCH_CHECK(cudaGetLastError() == cudaSuccess, "c");
    }
    if (was2) return {a.squeeze(0), tau.squeeze(0)};
    return {a, tau};
}
"""

CPP_SRC = r"""
std::vector<torch::Tensor> qr_forward(torch::Tensor A);
"""

_mod = load_inline(
    name="qr_spld_v1",
    cpp_sources=[CPP_SRC],
    cuda_sources=[CUDA_SRC],
    functions=["qr_forward"],
    verbose=False,
    extra_cuda_cflags=["-O3"],
)


def custom_kernel(data: input_t) -> output_t:
    a, tau = _mod.qr_forward(data)
    return a, tau