kernel.py

import os
import torch
from torch.utils.cpp_extension import load_inline

_qr_mod = None

CUDA_SRC = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <c10/cuda/CUDAException.h>
#include <cublas_v2.h>

#include <vector>
#include <algorithm>
#include <cmath>

#define CUBLAS_CHECK(cmd)                                                     \
  do {                                                                        \
    cublasStatus_t status_ = (cmd);                                           \
    TORCH_CHECK(status_ == CUBLAS_STATUS_SUCCESS,                             \
                "cuBLAS failure, status=", static_cast<int>(status_),         \
                " at ", __FILE__, ":", __LINE__);                             \
  } while (0)

static constexpr int QR_BLOCK = 512;
static constexpr int MAX_NB = 64;

__inline__ __device__ float warp_sum(float v) {
  unsigned mask = 0xffffffffu;
  v += __shfl_down_sync(mask, v, 16);
  v += __shfl_down_sync(mask, v, 8);
  v += __shfl_down_sync(mask, v, 4);
  v += __shfl_down_sync(mask, v, 2);
  v += __shfl_down_sync(mask, v, 1);
  return v;
}

__inline__ __device__ float block_sum(float v, float* smem) {
  int tid = threadIdx.x;
  int lane = tid & 31;
  int wid = tid >> 5;
  int nwarps = (blockDim.x + 31) >> 5;

  v = warp_sum(v);
  if (lane == 0) smem[wid] = v;
  __syncthreads();

  v = (tid < nwarps) ? smem[lane] : 0.0f;
  if (wid == 0) v = warp_sum(v);
  if (tid == 0) smem[0] = v;
  __syncthreads();

  return smem[0];
}

__global__ void to_col_major_kernel(
    const float* __restrict__ A,
    float* __restrict__ C,
    long long total,
    int n) {
  long long stride = (long long)blockDim.x * gridDim.x;
  for (long long idx = (long long)blockIdx.x * blockDim.x + threadIdx.x;
       idx < total;
       idx += stride) {
    long long nn = (long long)n * n;
    long long b = idx / nn;
    long long rem = idx - b * nn;
    int r = (int)(rem / n);
    int c = (int)(rem - (long long)r * n);

    C[b * nn + r + (long long)c * n] =
        A[b * nn + (long long)r * n + c];
  }
}

__global__ void from_col_major_kernel(
    const float* __restrict__ C,
    float* __restrict__ H,
    long long total,
    int n) {
  long long stride = (long long)blockDim.x * gridDim.x;
  for (long long idx = (long long)blockIdx.x * blockDim.x + threadIdx.x;
       idx < total;
       idx += stride) {
    long long nn = (long long)n * n;
    long long b = idx / nn;
    long long rem = idx - b * nn;
    int r = (int)(rem / n);
    int c = (int)(rem - (long long)r * n);

    H[b * nn + (long long)r * n + c] =
        C[b * nn + r + (long long)c * n];
  }
}

// Unblocked Householder QR of one panel per batch item.
// Only updates columns inside the current panel; the trailing matrix is updated later
// by the compact WY block reflector through cuBLAS.
__global__ void panel_qr_kernel(
    float* __restrict__ C,
    float* __restrict__ tau,
    int n,
    int k,
    int ib) {
  extern __shared__ float sh[];
  int b = blockIdx.x;
  int tid = threadIdx.x;
  int lane = tid & 31;
  int warp = tid >> 5;
  int nwarps = blockDim.x >> 5;

  long long base = (long long)b * n * n;

  for (int jj = 0; jj < ib; ++jj) {
    int col = k + jj;

    float ss = 0.0f;
    for (int r = col + 1 + tid; r < n; r += blockDim.x) {
      float x = C[base + r + (long long)col * n];
      ss += x * x;
    }

    float sigma = block_sum(ss, sh);

    if (tid == 0) {
      float alpha = C[base + col + (long long)col * n];
      float beta = alpha;
      float tauv = 0.0f;
      float inv = 0.0f;

      if (sigma == 0.0f) {
        if (alpha < 0.0f) {
          beta = -alpha;
          tauv = 2.0f;
        }
      } else {
        float norm = sqrtf(alpha * alpha + sigma);
        beta = (alpha >= 0.0f) ? -norm : norm;
        tauv = (beta - alpha) / beta;
        inv = 1.0f / (alpha - beta);
      }

      C[base + col + (long long)col * n] = beta;
      tau[(long long)b * n + col] = tauv;
      sh[0] = tauv;
      sh[1] = inv;
    }
    __syncthreads();

    float tauv = sh[0];
    float inv = sh[1];

    if (inv != 0.0f) {
      for (int r = col + 1 + tid; r < n; r += blockDim.x) {
        C[base + r + (long long)col * n] *= inv;
      }
    }
    __syncthreads();

    // Apply current reflector to the rest of the panel.
    if (tauv != 0.0f) {
      for (int pp = jj + 1 + warp; pp < ib; pp += nwarps) {
        int j = k + pp;

        float dot = (lane == 0)
            ? C[base + col + (long long)j * n]
            : 0.0f;

        for (int r = col + 1 + lane; r < n; r += 32) {
          dot += C[base + r + (long long)col * n] *
                 C[base + r + (long long)j * n];
        }

        dot = warp_sum(dot);
        float w = tauv * dot;

        if (lane == 0) {
          C[base + col + (long long)j * n] -= w;
        }

        for (int r = col + 1 + lane; r < n; r += 32) {
          C[base + r + (long long)j * n] -=
              C[base + r + (long long)col * n] * w;
        }
      }
    }
    __syncthreads();
  }
}

// Pack the panel Householder vectors into dense V, column-major, lda=m.
// V is m x ib, where m = n-k.
__global__ void pack_v_kernel(
    const float* __restrict__ C,
    float* __restrict__ V,
    long long total,
    int n,
    int k,
    int m,
    int ib,
    long long strideV) {
  long long step = (long long)blockDim.x * gridDim.x;

  for (long long idx = (long long)blockIdx.x * blockDim.x + threadIdx.x;
       idx < total;
       idx += step) {
    long long per = (long long)m * ib;
    int b = (int)(idx / per);
    long long rem = idx - (long long)b * per;
    int row = (int)(rem % m);
    int col = (int)(rem / m);

    float val = 0.0f;
    if (row == col) {
      val = 1.0f;
    } else if (row > col) {
      long long cbase = (long long)b * n * n;
      val = C[cbase + (k + row) + (long long)(k + col) * n];
    }

    V[(long long)b * strideV + row + (long long)col * m] = val;
  }
}

// Form T for compact WY: H_panel = I - V T V^T = H_0 H_1 ... H_{ib-1}.
__global__ void form_t_kernel(
    const float* __restrict__ C,
    const float* __restrict__ tau,
    float* __restrict__ T,
    int n,
    int k,
    int m,
    int ib,
    long long strideT) {
  extern __shared__ float shared[];
  float* red = shared;
  float* tmp = shared + 32;

  int b = blockIdx.x;
  int tid = threadIdx.x;
  long long cbase = (long long)b * n * n;
  float* Tb = T + (long long)b * strideT;

  for (int idx = tid; idx < ib * ib; idx += blockDim.x) {
    Tb[idx] = 0.0f;
  }
  __syncthreads();

  for (int i = 0; i < ib; ++i) {
    float taui = tau[(long long)b * n + k + i];

    if (taui != 0.0f) {
      for (int j = 0; j < i; ++j) {
        float s = 0.0f;

        // dot(v_j, v_i), starting at row i because v_i is zero above i.
        for (int row = i + tid; row < m; row += blockDim.x) {
          float vi = (row == i)
              ? 1.0f
              : C[cbase + (k + row) + (long long)(k + i) * n];

          // Since row >= i > j, this is below the diagonal of v_j.
          float vj = C[cbase + (k + row) + (long long)(k + j) * n];

          s += vj * vi;
        }

        float dot = block_sum(s, red);
        if (tid == 0) {
          Tb[j + (long long)i * ib] = -taui * dot;
        }
        __syncthreads();
      }

      if (tid == 0) {
        // T(0:i, i) = T(0:i, 0:i) * T(0:i, i)
        for (int r = 0; r < i; ++r) {
          float acc = 0.0f;
          for (int l = 0; l < i; ++l) {
            acc += Tb[r + (long long)l * ib] * Tb[l + (long long)i * ib];
          }
          tmp[r] = acc;
        }
        for (int r = 0; r < i; ++r) {
          Tb[r + (long long)i * ib] = tmp[r];
        }
        Tb[i + (long long)i * ib] = taui;
      }
    }

    __syncthreads();
  }
}

std::vector<torch::Tensor> qr_forward(torch::Tensor A) {
  TORCH_CHECK(A.is_cuda(), "A must be CUDA");
  TORCH_CHECK(A.scalar_type() == at::kFloat, "A must be torch.float32");
  TORCH_CHECK(A.dim() == 3, "A must have shape [batch, n, n]");
  TORCH_CHECK(A.size(1) == A.size(2), "A must be square");

  c10::cuda::OptionalCUDAGuard guard(A.device());

  auto Ac = A.contiguous();
  int64_t B64 = Ac.size(0);
  int64_t n64 = Ac.size(1);
  TORCH_CHECK(n64 <= INT_MAX, "n too large");

  int B = (int)B64;
  int n = (int)n64;

  auto C = torch::empty_like(Ac);
  auto H = torch::empty_like(Ac);
  auto tau = torch::empty({B64, n64}, Ac.options());

  if (B == 0 || n == 0) {
    return {H, tau};
  }

  auto stream = at::cuda::getCurrentCUDAStream();

  long long total = (long long)B * n * n;
  int threads = 256;
  int blocks = (int)std::min<long long>((total + threads - 1) / threads, 65535LL);

  to_col_major_kernel<<<blocks, threads, 0, stream>>>(
      Ac.data_ptr<float>(), C.data_ptr<float>(), total, n);
  C10_CUDA_KERNEL_LAUNCH_CHECK();

  // Use smaller panels for 512-ish cases, larger for 1024+.
  int nb = (n >= 1024) ? 64 : 32;
  nb = std::min(nb, MAX_NB);

  long long strideV = (long long)n * MAX_NB;
  long long strideT = (long long)MAX_NB * MAX_NB;
  long long strideW = (long long)MAX_NB * n;
  long long strideCmat = (long long)n * n;

  auto V = torch::empty({B64, strideV}, Ac.options());
  auto T = torch::empty({B64, strideT}, Ac.options());
  auto W = torch::empty({B64, strideW}, Ac.options());
  auto W2 = torch::empty({B64, strideW}, Ac.options());

  cublasHandle_t handle = at::cuda::getCurrentCUDABlasHandle();
  CUBLAS_CHECK(cublasSetStream(handle, stream));

  cublasMath_t old_math;
  cublasPointerMode_t old_ptr;
  CUBLAS_CHECK(cublasGetMathMode(handle, &old_math));
  CUBLAS_CHECK(cublasGetPointerMode(handle, &old_ptr));

  // Important for QR correctness: avoid TF32 tensor-core accumulation.
  CUBLAS_CHECK(cublasSetMathMode(handle, CUBLAS_PEDANTIC_MATH));
  CUBLAS_CHECK(cublasSetPointerMode(handle, CUBLAS_POINTER_MODE_HOST));

  const float one = 1.0f;
  const float zero = 0.0f;
  const float minus_one = -1.0f;

  float* Cptr = C.data_ptr<float>();
  float* tauptr = tau.data_ptr<float>();
  float* Vptr = V.data_ptr<float>();
  float* Tptr = T.data_ptr<float>();
  float* Wptr = W.data_ptr<float>();
  float* W2ptr = W2.data_ptr<float>();

  for (int k = 0; k < n; k += nb) {
    int ib = std::min(nb, n - k);
    int m = n - k;
    int nt = n - k - ib;

    panel_qr_kernel<<<B, QR_BLOCK, 32 * sizeof(float), stream>>>(
        Cptr, tauptr, n, k, ib);
    C10_CUDA_KERNEL_LAUNCH_CHECK();

    if (nt > 0) {
      long long vtotal = (long long)B * m * ib;
      int vblocks = (int)std::min<long long>(
          (vtotal + threads - 1) / threads, 65535LL);

      pack_v_kernel<<<vblocks, threads, 0, stream>>>(
          Cptr, Vptr, vtotal, n, k, m, ib, strideV);
      C10_CUDA_KERNEL_LAUNCH_CHECK();

      form_t_kernel<<<B, QR_BLOCK, (32 + MAX_NB) * sizeof(float), stream>>>(
          Cptr, tauptr, Tptr, n, k, m, ib, strideT);
      C10_CUDA_KERNEL_LAUNCH_CHECK();

      float* A2 = Cptr + k + (long long)(k + ib) * n;

      // W = V^T A2
      CUBLAS_CHECK(cublasSgemmStridedBatched(
          handle,
          CUBLAS_OP_T, CUBLAS_OP_N,
          ib, nt, m,
          &one,
          Vptr, m, strideV,
          A2, n, strideCmat,
          &zero,
          Wptr, ib, strideW,
          B));

      // W2 = T^T W, because trailing update applies H_panel^T.
      CUBLAS_CHECK(cublasSgemmStridedBatched(
          handle,
          CUBLAS_OP_T, CUBLAS_OP_N,
          ib, nt, ib,
          &one,
          Tptr, ib, strideT,
          Wptr, ib, strideW,
          &zero,
          W2ptr, ib, strideW,
          B));

      // A2 -= V W2
      CUBLAS_CHECK(cublasSgemmStridedBatched(
          handle,
          CUBLAS_OP_N, CUBLAS_OP_N,
          m, nt, ib,
          &minus_one,
          Vptr, m, strideV,
          W2ptr, ib, strideW,
          &one,
          A2, n, strideCmat,
          B));
    }
  }

  CUBLAS_CHECK(cublasSetMathMode(handle, old_math));
  CUBLAS_CHECK(cublasSetPointerMode(handle, old_ptr));

  from_col_major_kernel<<<blocks, threads, 0, stream>>>(
      C.data_ptr<float>(), H.data_ptr<float>(), total, n);
  C10_CUDA_KERNEL_LAUNCH_CHECK();

  return {H, tau};
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
  m.def("qr_forward", &qr_forward, "Batched square compact Householder QR");
}
"""

def _load_qr_mod():
    global _qr_mod
    if _qr_mod is None:
        if torch.cuda.is_available() and "TORCH_CUDA_ARCH_LIST" not in os.environ:
            major, minor = torch.cuda.get_device_capability()
            os.environ["TORCH_CUDA_ARCH_LIST"] = f"{major}.{minor}"

        _qr_mod = load_inline(
            name="b200_compact_householder_qr_ext",
            cpp_sources="",
            cuda_sources=CUDA_SRC,
            functions=None,
            with_cuda=True,
            extra_cuda_cflags=[
                "-O3",
                "--ptxas-options=-O3",
            ],
            extra_ldflags=["-lcublas"],
            verbose=False,
        )
    return _qr_mod

def solve(A: torch.Tensor):
    return _load_qr_mod().qr_forward(A)