wai-quantum 0.5.0

// wai-quantum fabric: the dense f64 state vector's gate kernels for OpenCL.
//
// A transcription of `quantum_sv`'s three formulas, with contraction into
// fused multiply-add off, so every operation is one IEEE binary64 operation rounded to
// nearest, in the order the CPU performs it:
//
// - a diagonal gate multiplies amplitude x by its diagonal entry, unless the
//   entry is exactly 1 (then the amplitude is left as it is: multiplying by 1
//   could flip the sign of a zero);
// - a gate on one qubit applies its 2x2 matrix to each pair, left to right:
//   m00*a - m00i*ai + m01*b - m01i*bi;
// - a gate on k >= 2 qubits computes each output of a group as a row of the
//   matrix against the group's members in local-index order, accumulated from
//   zero in that order.
//
// (`local` is an address space in OpenCL C, so the local index is `li`.)
//
// Each thread owns one amplitude, pair or group, and a launch is one gate, so
// no thread reads what another writes within a launch.

#pragma OPENCL EXTENSION cl_khr_fp64 : enable
#pragma OPENCL FP_CONTRACT OFF

typedef ulong u64;

// The maximum group size of a gate on k qubits: 2^k <= SV_MAX_D.
#define SV_MAX_D 64

__kernel void sv_diag(
    __global double* restrict re,
    __global double* restrict im,
    u64 len,
    __global const uint* restrict bitpos,
    uint k,
    __global const double* restrict diag)
{
    u64 x = (u64)get_global_id(0);
    u64 stride = (u64)get_global_size(0);
    for (; x < len; x += stride) {
        u64 li = 0;
        for (uint b = 0; b < k; b++) {
            li = (li << 1) | ((x >> bitpos[b]) & 1);
        }
        double cr = diag[2 * li];
        double ci = diag[2 * li + 1];
        if (cr == 1.0 && ci == 0.0) {
            continue;
        }
        double ar = re[x];
        double ai = im[x];
        re[x] = cr * ar - ci * ai;
        im[x] = cr * ai + ci * ar;
    }
}

__kernel void sv_one(
    __global double* restrict re,
    __global double* restrict im,
    u64 pairs,
    u64 s,
    double m00r, double m00i, double m01r, double m01i,
    double m10r, double m10i, double m11r, double m11i)
{
    u64 p = (u64)get_global_id(0);
    u64 stride = (u64)get_global_size(0);
    for (; p < pairs; p += stride) {
        u64 low = p & (s - 1);
        u64 i = ((p - low) << 1) | low;
        u64 j = i + s;
        double ar = re[i], ai = im[i], br = re[j], bi = im[j];
        re[i] = m00r * ar - m00i * ai + m01r * br - m01i * bi;
        im[i] = m00r * ai + m00i * ar + m01r * bi + m01i * br;
        re[j] = m10r * ar - m10i * ai + m11r * br - m11i * bi;
        im[j] = m10r * ai + m10i * ar + m11r * bi + m11i * br;
    }
}

__kernel void sv_many(
    __global double* restrict re,
    __global double* restrict im,
    u64 groups,
    __global const uint* restrict sorted,
    uint k,
    __global const u64* restrict offsets,
    __global const double* restrict m)
{
    u64 t = (u64)get_global_id(0);
    u64 stride = (u64)get_global_size(0);
    uint d = 1u << k;
    double vr[SV_MAX_D];
    double vi[SV_MAX_D];
    for (; t < groups; t += stride) {
        // t with a zero inserted at each target position, ascending.
        u64 base = t;
        for (uint b = 0; b < k; b++) {
            u64 pos = sorted[b];
            u64 low = base & (((u64)1 << pos) - 1);
            base = ((base >> pos) << (pos + 1)) | low;
        }
        for (uint j = 0; j < d; j++) {
            vr[j] = re[base | offsets[j]];
            vi[j] = im[base | offsets[j]];
        }
        for (uint r = 0; r < d; r++) {
            double sr = 0.0, si = 0.0;
            for (uint c = 0; c < d; c++) {
                double er = m[2 * (r * d + c)];
                double ei = m[2 * (r * d + c) + 1];
                sr += er * vr[c] - ei * vi[c];
                si += er * vi[c] + ei * vr[c];
            }
            re[base | offsets[r]] = sr;
            im[base | offsets[r]] = si;
        }
    }
}