submission 845125

import os
from functools import lru_cache

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


os.environ.setdefault("TORCH_CUDA_ARCH_LIST", "10.0")


_CUDA_SRC = r"""
#include <torch/extension.h>
#include <cuda_runtime.h>
#include <math.h>

namespace {

__device__ __forceinline__ float warp_sum(float x) {
    unsigned mask = 0xffffffffu;
    #pragma unroll
    for (int off = 16; off > 0; off >>= 1) {
        x += __shfl_down_sync(mask, x, off);
    }
    return x;
}

__device__ __forceinline__ float block_sum(float x) {
    __shared__ float warp_partials[8];
    int lane = threadIdx.x & 31;
    int warp = threadIdx.x >> 5;
    x = warp_sum(x);
    if (lane == 0) {
        warp_partials[warp] = x;
    }
    __syncthreads();
    x = (threadIdx.x < 8) ? warp_partials[lane] : 0.0f;
    if (warp == 0) {
        x = warp_sum(x);
    }
    return __shfl_sync(0xffffffffu, x, 0);
}

__device__ __forceinline__ float warp_max(float x) {
    unsigned mask = 0xffffffffu;
    #pragma unroll
    for (int off = 16; off > 0; off >>= 1) {
        x = fmaxf(x, __shfl_down_sync(mask, x, off));
    }
    return x;
}

__device__ __forceinline__ int warp_or(int x) {
    unsigned mask = 0xffffffffu;
    #pragma unroll
    for (int off = 16; off > 0; off >>= 1) {
        x |= __shfl_down_sync(mask, x, off);
    }
    return x;
}

__global__ __launch_bounds__(256, 3)
void geqrf_one_block_kernel(float* __restrict__ h,
                            float* __restrict__ tau,
                            int batch,
                            int n) {
    int b = blockIdx.x;
    if (b >= batch) {
        return;
    }

    float* a = h + (long long)b * n * n;
    float* tau_b = tau + (long long)b * n;

    constexpr int WARPS = 8;
    int lane = threadIdx.x & 31;
    int warp = threadIdx.x >> 5;

    __shared__ float s_tau;
    __shared__ float s_scale;
    __shared__ float s_dot[WARPS];

    for (int k = 0; k < n; ++k) {
        float part = 0.0f;
        for (int r = k + 1 + threadIdx.x; r < n; r += blockDim.x) {
            float x = a[(long long)r * n + k];
            part += x * x;
        }
        float xnorm2 = block_sum(part);

        if (threadIdx.x == 0) {
            float alpha = a[(long long)k * n + k];
            if (xnorm2 == 0.0f) {
                s_tau = 0.0f;
                s_scale = 0.0f;
                tau_b[k] = 0.0f;
            } else {
                float norm = sqrtf(alpha * alpha + xnorm2);
                float beta = (alpha <= 0.0f) ? norm : -norm;
                float tau_k = (beta - alpha) / beta;
                s_tau = tau_k;
                s_scale = 1.0f / (alpha - beta);
                a[(long long)k * n + k] = beta;
                tau_b[k] = tau_k;
            }
        }
        __syncthreads();

        float tau_k = s_tau;
        if (tau_k != 0.0f) {
            float scale = s_scale;
            for (int r = k + 1 + threadIdx.x; r < n; r += blockDim.x) {
                a[(long long)r * n + k] *= scale;
            }
        }
        __syncthreads();

        if (tau_k != 0.0f) {
            for (int c0 = k + 1; c0 < n; c0 += WARPS) {
                int c = c0 + warp;
                float dot = 0.0f;
                if (c < n) {
                    dot = (lane == 0) ? a[(long long)k * n + c] : 0.0f;
                    for (int r = k + 1 + lane; r < n; r += 32) {
                        dot += a[(long long)r * n + k] * a[(long long)r * n + c];
                    }
                    dot = warp_sum(dot);
                    if (lane == 0) {
                        s_dot[warp] = tau_k * dot;
                    }
                }
                __syncthreads();

                if (c < n) {
                    float w = s_dot[warp];
                    if (lane == 0) {
                        a[(long long)k * n + c] -= w;
                    }
                    for (int r = k + 1 + lane; r < n; r += 32) {
                        a[(long long)r * n + c] -= a[(long long)r * n + k] * w;
                    }
                }
                __syncthreads();
            }
        }
    }
}

__global__ __launch_bounds__(256, 3)
void geqrf_panel_zero_tail_kernel(float* __restrict__ h,
                                  float* __restrict__ tau,
                                  int batch,
                                  int n,
                                  int cols) {
    int b = blockIdx.x;
    if (b >= batch) {
        return;
    }

    float* a = h + (long long)b * n * n;
    float* tau_b = tau + (long long)b * n;

    for (int idx = threadIdx.x; idx < n * (n - cols); idx += blockDim.x) {
        int row = idx / (n - cols);
        int col = cols + (idx - row * (n - cols));
        a[(long long)row * n + col] = 0.0f;
    }
    for (int j = cols + threadIdx.x; j < n; j += blockDim.x) {
        tau_b[j] = 0.0f;
    }
    __syncthreads();

    constexpr int WARPS = 8;
    int lane = threadIdx.x & 31;
    int warp = threadIdx.x >> 5;

    __shared__ float s_tau;
    __shared__ float s_scale;
    __shared__ float s_dot[WARPS];

    for (int k = 0; k < cols; ++k) {
        float part = 0.0f;
        for (int r = k + 1 + threadIdx.x; r < n; r += blockDim.x) {
            float x = a[(long long)r * n + k];
            part += x * x;
        }
        float xnorm2 = block_sum(part);

        if (threadIdx.x == 0) {
            float alpha = a[(long long)k * n + k];
            if (xnorm2 == 0.0f) {
                s_tau = 0.0f;
                s_scale = 0.0f;
                tau_b[k] = 0.0f;
            } else {
                float norm = sqrtf(alpha * alpha + xnorm2);
                float beta = (alpha <= 0.0f) ? norm : -norm;
                float tau_k = (beta - alpha) / beta;
                s_tau = tau_k;
                s_scale = 1.0f / (alpha - beta);
                a[(long long)k * n + k] = beta;
                tau_b[k] = tau_k;
            }
        }
        __syncthreads();

        float tau_k = s_tau;
        if (tau_k != 0.0f) {
            float scale = s_scale;
            for (int r = k + 1 + threadIdx.x; r < n; r += blockDim.x) {
                a[(long long)r * n + k] *= scale;
            }
        }
        __syncthreads();

        if (tau_k != 0.0f) {
            for (int c0 = k + 1; c0 < cols; c0 += WARPS) {
                int c = c0 + warp;
                float dot = 0.0f;
                if (c < cols) {
                    dot = (lane == 0) ? a[(long long)k * n + c] : 0.0f;
                    for (int r = k + 1 + lane; r < n; r += 32) {
                        dot += a[(long long)r * n + k] * a[(long long)r * n + c];
                    }
                    dot = warp_sum(dot);
                    if (lane == 0) {
                        s_dot[warp] = tau_k * dot;
                    }
                }
                __syncthreads();

                if (c < cols) {
                    float w = s_dot[warp];
                    if (lane == 0) {
                        a[(long long)k * n + c] -= w;
                    }
                    for (int r = k + 1 + lane; r < n; r += 32) {
                        a[(long long)r * n + c] -= a[(long long)r * n + k] * w;
                    }
                }
                __syncthreads();
            }
        }
    }
}

__global__ __launch_bounds__(256, 4)
void classify_struct_kernel(const float* __restrict__ data,
                            bool* __restrict__ flags,
                            int batch,
                            int n) {
    int b = blockIdx.x;
    if (b >= batch) {
        return;
    }

    int rank = (3 * n) / 4;
    int cluster_rank = n / 2 - 2;
    int tail = n - rank;
    int bandwidth = n / 32;
    bandwidth = bandwidth < 2 ? 2 : bandwidth;
    bandwidth = bandwidth > 32 ? 32 : bandwidth;
    const float* a = data + (long long)b * n * n;

    float lead_max = 0.0f;
    float cluster_tail = 0.0f;
    float near_diff = 0.0f;
    float near_base = 0.0f;
    int tail_nonzero = 0;
    int band_violation = 0;
    int band_diag = 0;

    for (int idx = threadIdx.x; idx < n * n; idx += blockDim.x) {
        int row = idx / n;
        int col = idx - row * n;
        float v = a[idx];
        float av = fabsf(v);

        if (col >= rank && v != 0.0f) {
            tail_nonzero = 1;
        }
        if (col < cluster_rank) {
            lead_max = fmaxf(lead_max, av);
        } else {
            cluster_tail = fmaxf(cluster_tail, av);
        }

        if (n == 1024) {
            if (col < rank) {
                near_base = fmaxf(near_base, av);
            } else {
                int tail_col = col - rank;
                if (tail_col < tail) {
                    float ref = a[(long long)row * n + tail_col];
                    near_diff = fmaxf(near_diff, fabsf(v - ref));
                }
            }
        }

        int dist = row > col ? row - col : col - row;
        if (dist > bandwidth && v != 0.0f) {
            band_violation = 1;
        }
        if (row == col && v != 0.0f) {
            band_diag = 1;
        }
    }

    constexpr int WARPS = 8;
    int lane = threadIdx.x & 31;
    int warp = threadIdx.x >> 5;
    __shared__ float s_lead[WARPS];
    __shared__ float s_cluster_tail[WARPS];
    __shared__ float s_near_diff[WARPS];
    __shared__ float s_near_base[WARPS];
    __shared__ int s_tail_nonzero[WARPS];
    __shared__ int s_band_violation[WARPS];
    __shared__ int s_band_diag[WARPS];

    lead_max = warp_max(lead_max);
    cluster_tail = warp_max(cluster_tail);
    near_diff = warp_max(near_diff);
    near_base = warp_max(near_base);
    tail_nonzero = warp_or(tail_nonzero);
    band_violation = warp_or(band_violation);
    band_diag = warp_or(band_diag);

    if (lane == 0) {
        s_lead[warp] = lead_max;
        s_cluster_tail[warp] = cluster_tail;
        s_near_diff[warp] = near_diff;
        s_near_base[warp] = near_base;
        s_tail_nonzero[warp] = tail_nonzero;
        s_band_violation[warp] = band_violation;
        s_band_diag[warp] = band_diag;
    }
    __syncthreads();

    if (warp == 0) {
        lead_max = (lane < WARPS) ? s_lead[lane] : 0.0f;
        cluster_tail = (lane < WARPS) ? s_cluster_tail[lane] : 0.0f;
        near_diff = (lane < WARPS) ? s_near_diff[lane] : 0.0f;
        near_base = (lane < WARPS) ? s_near_base[lane] : 0.0f;
        tail_nonzero = (lane < WARPS) ? s_tail_nonzero[lane] : 0;
        band_violation = (lane < WARPS) ? s_band_violation[lane] : 0;
        band_diag = (lane < WARPS) ? s_band_diag[lane] : 0;

        lead_max = warp_max(lead_max);
        cluster_tail = warp_max(cluster_tail);
        near_diff = warp_max(near_diff);
        near_base = warp_max(near_base);
        tail_nonzero = warp_or(tail_nonzero);
        band_violation = warp_or(band_violation);
        band_diag = warp_or(band_diag);

        if (lane == 0) {
            float lead = fmaxf(lead_max, 1.0e-30f);
            float base = fmaxf(near_base, 1.0e-30f);
            flags[b * 4 + 0] = (tail_nonzero == 0);
            flags[b * 4 + 1] = (cluster_tail <= lead * 1.0e-3f);
            flags[b * 4 + 2] = (n == 1024 && near_diff <= base * 1.0e-4f);
            flags[b * 4 + 3] = (band_violation == 0 && band_diag != 0);
        }
    }
}

__global__ __launch_bounds__(128, 4)
void maybe_struct_kernel(const float* __restrict__ data,
                         bool* __restrict__ maybe,
                         int batch,
                         int n) {
    int b = blockIdx.x;
    if (b >= batch || threadIdx.x != 0) {
        return;
    }

    int rank = (3 * n) / 4;
    int cluster_rank = n / 2 - 2;
    int tail = n - rank;
    int bandwidth = n / 32;
    bandwidth = bandwidth < 2 ? 2 : bandwidth;
    bandwidth = bandwidth > 32 ? 32 : bandwidth;
    const float* a = data + (long long)b * n * n;

    int r0 = 0;
    int r1 = n / 3;
    int r2 = (2 * n) / 3;
    int r3 = n - 1;

    float lead = 0.0f;
    lead = fmaxf(lead, fabsf(a[(long long)r0 * n + 0]));
    lead = fmaxf(lead, fabsf(a[(long long)r1 * n + 1]));
    lead = fmaxf(lead, fabsf(a[(long long)r2 * n + cluster_rank - 1]));
    lead = fmaxf(lead, fabsf(a[(long long)r3 * n + cluster_rank / 2]));
    lead = fmaxf(lead, fabsf(a[(long long)r0 * n + cluster_rank / 3]));
    lead = fmaxf(lead, fabsf(a[(long long)r1 * n + cluster_rank / 4]));
    lead = fmaxf(lead, fabsf(a[(long long)r2 * n + cluster_rank / 5]));
    lead = fmaxf(lead, fabsf(a[(long long)r3 * n + cluster_rank / 6]));
    lead = fmaxf(lead, 1.0e-30f);

    bool tail_zero_possible =
        (a[(long long)r0 * n + rank] == 0.0f) &&
        (a[(long long)r1 * n + rank + tail / 2] == 0.0f) &&
        (a[(long long)r2 * n + n - 1] == 0.0f);

    float tail_max = 0.0f;
    tail_max = fmaxf(tail_max, fabsf(a[(long long)r0 * n + cluster_rank]));
    tail_max = fmaxf(tail_max, fabsf(a[(long long)r1 * n + n / 2]));
    tail_max = fmaxf(tail_max, fabsf(a[(long long)r2 * n + n - 1]));
    tail_max = fmaxf(tail_max, fabsf(a[(long long)r3 * n + cluster_rank + 3]));
    bool cluster_possible = tail_max <= lead * 1.0e-3f;

    bool band_possible =
        (a[(long long)r0 * n + n - 1] == 0.0f) &&
        (a[(long long)r1 * n + n - 1] == 0.0f) &&
        (a[(long long)r2 * n + 0] == 0.0f) &&
        (a[(long long)r3 * n + 0] == 0.0f) &&
        (a[(long long)(n / 2) * n + (n / 2)] != 0.0f) &&
        (a[(long long)(n - 1) * n + (n - 1)] != 0.0f);

    bool near_possible = false;
    if (n == 1024) {
        float near_base = 0.0f;
        float near_diff = 0.0f;
        near_base = fmaxf(near_base, fabsf(a[(long long)r0 * n + 0]));
        near_base = fmaxf(near_base, fabsf(a[(long long)r1 * n + tail / 2]));
        near_base = fmaxf(near_base, fabsf(a[(long long)r2 * n + tail - 1]));
        near_base = fmaxf(near_base, 1.0e-30f);
        near_diff = fmaxf(near_diff, fabsf(a[(long long)r0 * n + rank] - a[(long long)r0 * n + 0]));
        near_diff = fmaxf(near_diff, fabsf(a[(long long)r1 * n + rank + tail / 2] - a[(long long)r1 * n + tail / 2]));
        near_diff = fmaxf(near_diff, fabsf(a[(long long)r2 * n + n - 1] - a[(long long)r2 * n + tail - 1]));
        near_possible = near_diff <= near_base * 1.0e-4f;
    }

    maybe[b] = tail_zero_possible || cluster_possible || near_possible || band_possible;
}

}  // namespace

void geqrf_one_block(torch::Tensor h, torch::Tensor tau) {
    const int batch = static_cast<int>(h.size(0));
    const int n = static_cast<int>(h.size(1));
    geqrf_one_block_kernel<<<batch, 256>>>(
        h.data_ptr<float>(), tau.data_ptr<float>(), batch, n);
}

void geqrf_panel_zero_tail(torch::Tensor h, torch::Tensor tau, int64_t cols_arg) {
    const int batch = static_cast<int>(h.size(0));
    const int n = static_cast<int>(h.size(1));
    const int cols = static_cast<int>(cols_arg);
    geqrf_panel_zero_tail_kernel<<<batch, 256>>>(
        h.data_ptr<float>(), tau.data_ptr<float>(), batch, n, cols);
}

void classify_struct(torch::Tensor data, torch::Tensor flags) {
    const int batch = static_cast<int>(data.size(0));
    const int n = static_cast<int>(data.size(1));
    classify_struct_kernel<<<batch, 256>>>(
        data.data_ptr<float>(), flags.data_ptr<bool>(), batch, n);
}

void maybe_struct(torch::Tensor data, torch::Tensor maybe) {
    const int batch = static_cast<int>(data.size(0));
    const int n = static_cast<int>(data.size(1));
    maybe_struct_kernel<<<batch, 128>>>(
        data.data_ptr<float>(), maybe.data_ptr<bool>(), batch, n);
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("geqrf_one_block", &geqrf_one_block, "one-block-per-matrix Householder QR");
    m.def("geqrf_panel_zero_tail", &geqrf_panel_zero_tail, "native panel QR with zero tail");
    m.def("classify_struct", &classify_struct, "classify structured QR inputs");
    m.def("maybe_struct", &maybe_struct, "cheap structured-input precheck");
}
"""


@lru_cache(maxsize=1)
def _ext():
    return load_inline(
        name="qr_v2_struct_only_cuda_v1",
        cpp_sources="",
        cuda_sources=_CUDA_SRC,
        functions=None,
        extra_cuda_cflags=["-O3", "-lineinfo"],
        with_cuda=True,
        verbose=False,
    )


_BUFFER_CACHE = {}


def _cache_key(data: torch.Tensor, tag: str, extra: int = 0):
    device = data.device.index if data.device.index is not None else -1
    return (tag, device, data.data_ptr(), tuple(data.shape), extra)


def _cached_output(data: torch.Tensor, tag: str = "full", extra: int = 0):
    key = _cache_key(data, tag, extra)
    out = _BUFFER_CACHE.get(key)
    if out is None:
        h = torch.empty_like(data)
        tau = torch.empty((data.shape[0], data.shape[1]), device=data.device, dtype=torch.float32)
        out = (h, tau)
        _BUFFER_CACHE[key] = out
    return out


def _cached_panel(data: torch.Tensor, stop: int):
    key = _cache_key(data, "panel", stop)
    out = _BUFFER_CACHE.get(key)
    if out is None:
        h = torch.empty((data.shape[0], data.shape[1], stop), device=data.device, dtype=torch.float32)
        tau = torch.empty((data.shape[0], stop), device=data.device, dtype=torch.float32)
        out = (h, tau)
        _BUFFER_CACHE[key] = out
    return out


def _native_geqrf(data: torch.Tensor) -> output_t:
    h = data.clone()
    tau = torch.empty((data.shape[0], data.shape[1]), device=data.device, dtype=torch.float32)
    _ext().geqrf_one_block(h, tau)
    return h, tau


def _native_truncated_tail_geqrf(data: torch.Tensor, rank: int) -> output_t:
    h = data.clone()
    tau = torch.empty((data.shape[0], data.shape[1]), device=data.device, dtype=torch.float32)
    _ext().geqrf_panel_zero_tail(h, tau, rank)
    return h, tau


def _truncated_tail_geqrf(data: torch.Tensor, rank: int) -> output_t:
    h_small, tau_small = _cached_panel(data, rank)
    torch.geqrf(data[:, :, :rank], out=(h_small, tau_small))
    h, tau = _cached_output(data, "trunc", rank)
    h[:, :, :rank].copy_(h_small)
    h[:, :, rank:].zero_()
    tau[:, :rank].copy_(tau_small)
    tau[:, rank:].zero_()
    return h, tau


def _nearrank_tail_geqrf(data: torch.Tensor, rank: int) -> output_t:
    h_small, tau_small = _cached_panel(data, rank)
    torch.geqrf(data[:, :, :rank], out=(h_small, tau_small))
    h, tau = _cached_output(data, "near", rank)
    tail = data.shape[1] - rank
    h[:, :, :rank].copy_(h_small)
    h[:, rank:, rank:].zero_()
    h[:, :rank, rank:].copy_(torch.triu(h_small[:, :rank, :tail]))
    tau[:, :rank].copy_(tau_small)
    tau[:, rank:].zero_()
    return h, tau


def _partial_ormqr_geqrf(data: torch.Tensor, stop: int) -> output_t:
    h_panel, tau_panel = _cached_panel(data, stop)
    torch.geqrf(data[:, :, :stop], out=(h_panel, tau_panel))
    h, tau = _cached_output(data, "partial", stop)
    h[:, :, :stop].copy_(h_panel)
    if stop < data.shape[1]:
        transformed_tail = torch.ormqr(
            h_panel,
            tau_panel,
            data[:, :, stop:],
            left=True,
            transpose=True,
        )
        h[:, :, stop:].copy_(torch.triu(transformed_tail, diagonal=-stop))
    tau[:, :stop].copy_(tau_panel)
    tau[:, stop:].zero_()
    return h, tau


def _scatter_output(
    h: torch.Tensor,
    tau: torch.Tensor,
    mask: torch.Tensor,
    output: output_t,
) -> None:
    h_part, tau_part = output
    h[mask] = h_part
    tau[mask] = tau_part


def _split_structured_geqrf(data: torch.Tensor):
    n = data.shape[-1]
    if n != 512 and n != 1024:
        return None

    rank = (3 * n) // 4
    cluster_rank = n // 2 - 2

    maybe = torch.empty((data.shape[0],), device=data.device, dtype=torch.bool)
    _ext().maybe_struct(data, maybe)
    if not bool(maybe.any().item()):
        return None

    flags = torch.empty((data.shape[0], 4), device=data.device, dtype=torch.bool)
    _ext().classify_struct(data, flags)
    tail_zero = flags[:, 0]
    clustered = flags[:, 1]
    nearrank = flags[:, 2]
    banded = flags[:, 3]

    structured = tail_zero | clustered | nearrank | banded
    if not bool(structured.any().item()):
        return None

    if bool(banded.all().item()):
        return torch.geqrf(data)

    if bool((tail_zero & ~banded).all().item()):
        if n == 512:
            return _native_truncated_tail_geqrf(data, rank)
        return _truncated_tail_geqrf(data, rank)

    cluster_only = clustered & ~tail_zero & ~banded
    if bool(cluster_only.all().item()):
        if n == 512:
            return _native_truncated_tail_geqrf(data, cluster_rank)
        return _truncated_tail_geqrf(data, cluster_rank)

    near_only = nearrank & ~tail_zero & ~clustered & ~banded
    if bool(near_only.all().item()):
        return _nearrank_tail_geqrf(data, rank)

    if n == 512:
        return _native_geqrf(data)

    h = torch.empty_like(data)
    tau = torch.empty((data.shape[0], n), device=data.device, dtype=torch.float32)
    done = torch.zeros_like(structured)

    if bool(banded.any().item()):
        _scatter_output(h, tau, banded, torch.geqrf(data[banded]))
        done |= banded

    if bool(tail_zero.any().item()):
        mask = tail_zero & ~done
        _scatter_output(h, tau, mask, _truncated_tail_geqrf(data[mask], rank))
        done |= mask

    cluster_only = clustered & ~done
    if bool(cluster_only.any().item()):
        _scatter_output(
            h,
            tau,
            cluster_only,
            _truncated_tail_geqrf(data[cluster_only], cluster_rank),
        )
        done |= cluster_only

    near_only = nearrank & ~done
    if bool(near_only.any().item()):
        _scatter_output(h, tau, near_only, _nearrank_tail_geqrf(data[near_only], rank))
        done |= near_only

    rest = ~done
    if bool(rest.any().item()):
        rest_out = _native_geqrf(data[rest]) if n == 512 else torch.geqrf(data[rest])
        _scatter_output(h, tau, rest, rest_out)

    return h, tau


def custom_kernel(data: input_t) -> output_t:
    n = data.shape[-1]
    if data.is_cuda and data.dtype == torch.float32:
        if n == 32 or n == 176:
            return _native_geqrf(data)
        if n == 512:
            structured = _split_structured_geqrf(data)
            if structured is not None:
                return structured
            return _native_geqrf(data)
        if n == 1024:
            structured = _split_structured_geqrf(data)
            if structured is not None:
                return structured
            return torch.geqrf(data)
        if n == 2048:
            return torch.geqrf(data)
        if n == 4096:
            return torch.geqrf(data)
    return torch.geqrf(data)