submission 876544

import torch
import torch.utils.cpp_extension as ce

if ce.CUDA_HOME is None:
    ce.CUDA_HOME = torch.__path__[0]

_code = r'''
#define TILE_THREADS 512

extern "C" __global__ void jacobi32(
    const float* input,
    float* vectors,
    float* values,
    int batch) {
    const int matrix = (int)blockIdx.x;
    const int lane = (int)threadIdx.x;
    if (matrix >= batch) return;

    __shared__ float a[1056];
    __shared__ float q[1056];
    __shared__ float cosine[16];
    __shared__ float sine[16];
    __shared__ int left[16];
    __shared__ int right[16];
    __shared__ float work[64];
    __shared__ float limits[2];
    __shared__ int order[32];

    const float* src = input + matrix * 1024;
    float* dst_q = vectors + matrix * 1024;
    float* dst_l = values + matrix * 32;
    float local_scale = 0.0f;
    float local_off = 0.0f;
    for (int index = lane; index < 1024; index += TILE_THREADS) {
        const int row = index >> 5;
        const int col = index & 31;
        const float x = src[index];
        a[row * 33 + col] = x;
        q[row * 33 + col] = row == col ? 1.0f : 0.0f;
        local_scale = fmaxf(local_scale, fabsf(x));
        if (row != col) local_off = fmaxf(local_off, fabsf(x));
    }
    for (int gap = 16; gap > 0; gap >>= 1) {
        local_scale = fmaxf(local_scale, __shfl_down_sync(0xffffffff, local_scale, gap));
        local_off = fmaxf(local_off, __shfl_down_sync(0xffffffff, local_off, gap));
    }
    const int warp = lane >> 5;
    if ((lane & 31) == 0) {
        work[warp] = local_scale;
        work[32 + warp] = local_off;
    }
    __syncthreads();
    if (lane < 32) {
        const int warps = TILE_THREADS >> 5;
        float x = lane < warps ? work[lane] : 0.0f;
        float y = lane < warps ? work[32 + lane] : 0.0f;
        for (int gap = 16; gap > 0; gap >>= 1) {
            x = fmaxf(x, __shfl_down_sync(0xffffffff, x, gap));
            y = fmaxf(y, __shfl_down_sync(0xffffffff, y, gap));
        }
        if (lane == 0) {
            limits[0] = x;
            limits[1] = y;
        }
    }
    __syncthreads();
    const float scale = limits[0];
    const float off = limits[1];

    if (off > 1.0e-7f * scale) {
        for (int sweep = 0; sweep < 8; ++sweep) {
            for (int round = 0; round < 31; ++round) {
                if (lane < 16) {
                    const int p = lane;
                    int x = p - 1 + round;
                    if (x >= 31) x -= 31;
                    const int i = p == 0 ? 0 : 1 + x;
                    int y = 30 - p + round;
                    if (y >= 31) y -= 31;
                    const int j = 1 + y;
                    left[p] = i;
                    right[p] = j;
                    const float app = a[i * 33 + i];
                    const float aqq = a[j * 33 + j];
                    const float apq = a[i * 33 + j];
                    const float cut = 1.0e-7f * (fabsf(app) + fabsf(aqq) + scale);
                    float c = 1.0f;
                    float s = 0.0f;
                    if (fabsf(apq) > cut) {
                        const float delta = aqq - app;
                        const float span = hypotf(delta, 2.0f * apq);
                        const float den = delta >= 0.0f ? delta + span : delta - span;
                        const float t = (2.0f * apq) / den;
                        c = rsqrtf(1.0f + t * t);
                        s = t * c;
                    }
                    cosine[p] = c;
                    sine[p] = s;
                }
                __syncthreads();

                for (int task = lane; task < 512; task += TILE_THREADS) {
                    const int pair = task >> 5;
                    const int row = task & 31;
                    const int i = left[pair];
                    const int j = right[pair];
                    const float c = cosine[pair];
                    const float s = sine[pair];
                    const float ai = a[row * 33 + i];
                    const float aj = a[row * 33 + j];
                    const float qi = q[row * 33 + i];
                    const float qj = q[row * 33 + j];
                    a[row * 33 + i] = c * ai - s * aj;
                    a[row * 33 + j] = s * ai + c * aj;
                    q[row * 33 + i] = c * qi - s * qj;
                    q[row * 33 + j] = s * qi + c * qj;
                }
                __syncthreads();

                for (int task = lane; task < 512; task += TILE_THREADS) {
                    const int pair = task >> 5;
                    const int col = task & 31;
                    const int i = left[pair];
                    const int j = right[pair];
                    const float c = cosine[pair];
                    const float s = sine[pair];
                    const float ai = a[i * 33 + col];
                    const float aj = a[j * 33 + col];
                    a[i * 33 + col] = col == j ? 0.0f : c * ai - s * aj;
                    a[j * 33 + col] = col == i ? 0.0f : s * ai + c * aj;
                }
                __syncthreads();
            }

        }
    }

    if (lane < 32) {
        float value = a[lane * 33 + lane];
        int source = lane;
        for (int width = 2; width <= 32; width <<= 1) {
            for (int gap = width >> 1; gap > 0; gap >>= 1) {
                const float peer_value = __shfl_xor_sync(0xffffffff, value, gap);
                const int peer_source = __shfl_xor_sync(0xffffffff, source, gap);
                const bool want_min = ((lane & width) == 0) == ((lane & gap) == 0);
                const bool take_peer = want_min ? value > peer_value : value < peer_value;
                if (take_peer) {
                    value = peer_value;
                    source = peer_source;
                }
            }
        }
        dst_l[lane] = value;
        order[lane] = source;
    }
    __syncthreads();
    for (int index = lane; index < 1024; index += TILE_THREADS) {
        const int row = index >> 5;
        const int col = index & 31;
        dst_q[index] = q[row * 33 + order[col]];
    }
}
'''

_make = getattr(torch.cuda, "compile_kernel", torch.cuda._compile_kernel)
_jacobi32 = _make(_code, "jacobi32")


def custom_kernel(data):
    if data.shape[-1] != 32:
        values, vectors = torch.linalg.eigh(data)
        return vectors, values
    batch = data.shape[0]
    flat = torch.empty(data.numel() + batch * 32, device=data.device, dtype=torch.float32)
    q = flat[: data.numel()].view_as(data)
    values = flat[data.numel() :].view(batch, 32)
    _jacobi32(grid=(batch, 1, 1), block=(512, 1, 1), args=[data, q, values, batch])
    return q, values