segplane.c

#include <stdio.h>
#include <stdlib.h>
#include <stdint.h>
#include <stddef.h>
#include <math.h>
#include <time.h>
#include <string.h>

#if defined(__GNUC__) || defined(__clang__)
#  define RESTRICT __restrict__
#  define INLINE static inline __attribute__((always_inline))
#else
#  define RESTRICT restrict
#  define INLINE static inline
#endif

#ifndef SMALL_NUM
#define SMALL_NUM 1.0e-8f
#endif

typedef struct {
    double best_sec;
    double mpairs_per_sec;
} bench_result;

INLINE double now_sec(void) {
    struct timespec ts;
    clock_gettime(CLOCK_MONOTONIC, &ts);
    return (double)ts.tv_sec + 1e-9 * (double)ts.tv_nsec;
}

INLINE float frand_range(float lo, float hi) {
    return lo + (hi - lo) * ((float)rand() / (float)RAND_MAX);
}

static void* aligned_malloc64(size_t bytes) {
    void* p = NULL;
    if (posix_memalign(&p, 64, bytes) != 0) return NULL;
    return p;
}

static void fill_segments_and_planes(
    float *ax0, float *ay0, float *az0,
    float *ax1, float *ay1, float *az1,
    float *pnx, float *pny, float *pnz, float *pd,
    float *rad,
    size_t n,
    float coord_range,
    float seg_len_max)
{
    for (size_t i = 0; i < n; ++i) {
        const float x0 = frand_range(-coord_range, coord_range);
        const float y0 = frand_range(-coord_range, coord_range);
        const float z0 = frand_range(-coord_range, coord_range);
        const float dx = frand_range(-seg_len_max, seg_len_max);
        const float dy = frand_range(-seg_len_max, seg_len_max);
        const float dz = frand_range(-seg_len_max, seg_len_max);
        ax0[i] = x0; ay0[i] = y0; az0[i] = z0;
        ax1[i] = x0 + dx; ay1[i] = y0 + dy; az1[i] = z0 + dz;

        float nx = frand_range(-1.0f, 1.0f);
        float ny = frand_range(-1.0f, 1.0f);
        float nz = frand_range(-1.0f, 1.0f);
        float nn = sqrtf(nx*nx + ny*ny + nz*nz);
        if (nn < SMALL_NUM) { nx = 1.0f; ny = 0.0f; nz = 0.0f; nn = 1.0f; }
        nx /= nn; ny /= nn; nz /= nn;
        pnx[i] = nx; pny[i] = ny; pnz[i] = nz;

        // Plane offset so some segments cross, some are far.
        const float anchor = frand_range(-coord_range, coord_range);
        pd[i] = anchor;
        rad[i] = frand_range(0.01f, 0.25f);
    }
}

// Kernel 1: exact segment-plane distance and crossing mask.
static void segplane3d_dist_soa(
    const float * RESTRICT ax0,
    const float * RESTRICT ay0,
    const float * RESTRICT az0,
    const float * RESTRICT ax1,
    const float * RESTRICT ay1,
    const float * RESTRICT az1,
    const float * RESTRICT pnx,
    const float * RESTRICT pny,
    const float * RESTRICT pnz,
    const float * RESTRICT pd,
    float * RESTRICT sA_out,
    float * RESTRICT sB_out,
    float * RESTRICT dist_out,
    uint8_t * RESTRICT cross_out,
    size_t n)
{
#if defined(__GNUC__) && !defined(__clang__)
#pragma GCC ivdep
#endif
    for (size_t i = 0; i < n; ++i) {
        const float nx = pnx[i], ny = pny[i], nz = pnz[i], d = pd[i];
        const float sA = nx * ax0[i] + ny * ay0[i] + nz * az0[i] + d;
        const float sB = nx * ax1[i] + ny * ay1[i] + nz * az1[i] + d;
        const float cross = (sA * sB <= 0.0f) ? 1.0f : 0.0f;
        const float aA = fabsf(sA);
        const float aB = fabsf(sB);
        const float minab = (aA < aB) ? aA : aB;
        sA_out[i] = sA;
        sB_out[i] = sB;
        dist_out[i] = cross ? 0.0f : minab;
        cross_out[i] = (uint8_t)cross;
    }
}

// Kernel 2: exact intersection points for crossing segment-plane pairs.
static void segplane3d_intersections_soa(
    const float * RESTRICT ax0,
    const float * RESTRICT ay0,
    const float * RESTRICT az0,
    const float * RESTRICT ax1,
    const float * RESTRICT ay1,
    const float * RESTRICT az1,
    const float * RESTRICT sA,
    const float * RESTRICT sB,
    float * RESTRICT t_out,
    float * RESTRICT px_out,
    float * RESTRICT py_out,
    float * RESTRICT pz_out,
    size_t n)
{
#if defined(__GNUC__) && !defined(__clang__)
#pragma GCC ivdep
#endif
    for (size_t i = 0; i < n; ++i) {
        const float ux = ax1[i] - ax0[i];
        const float uy = ay1[i] - ay0[i];
        const float uz = az1[i] - az0[i];
        const float denom = sA[i] - sB[i];
        const float t = sA[i] / denom;
        t_out[i] = t;
        px_out[i] = ax0[i] + t * ux;
        py_out[i] = ay0[i] + t * uy;
        pz_out[i] = az0[i] + t * uz;
    }
}

// Kernel 3: exact capsule-plane contact using deepest endpoint.
static void capsuleplane3d_contact_soa(
    const float * RESTRICT ax0,
    const float * RESTRICT ay0,
    const float * RESTRICT az0,
    const float * RESTRICT ax1,
    const float * RESTRICT ay1,
    const float * RESTRICT az1,
    const float * RESTRICT pnx,
    const float * RESTRICT pny,
    const float * RESTRICT pnz,
    const float * RESTRICT pd,
    const float * RESTRICT rad,
    float * RESTRICT sep_out,
    float * RESTRICT pen_out,
    float * RESTRICT cx_out,
    float * RESTRICT cy_out,
    float * RESTRICT cz_out,
    float * RESTRICT sx_out,
    float * RESTRICT sy_out,
    float * RESTRICT sz_out,
    uint8_t * RESTRICT hit_out,
    size_t n)
{
#if defined(__GNUC__) && !defined(__clang__)
#pragma GCC ivdep
#endif
    for (size_t i = 0; i < n; ++i) {
        const float nx = pnx[i], ny = pny[i], nz = pnz[i], d = pd[i];
        const float sA = nx * ax0[i] + ny * ay0[i] + nz * az0[i] + d;
        const float sB = nx * ax1[i] + ny * ay1[i] + nz * az1[i] + d;
        const float useA = (sA < sB) ? 1.0f : 0.0f;
        const float useB = 1.0f - useA;
        const float smin = useA * sA + useB * sB;
        const float r = rad[i];
        const float sep = smin - r;
        const float pen = (sep < 0.0f) ? -sep : 0.0f;
        const float cx = useA * ax0[i] + useB * ax1[i];
        const float cy = useA * ay0[i] + useB * ay1[i];
        const float cz = useA * az0[i] + useB * az1[i];
        sep_out[i] = sep;
        pen_out[i] = pen;
        cx_out[i] = cx;
        cy_out[i] = cy;
        cz_out[i] = cz;
        sx_out[i] = cx - r * nx;
        sy_out[i] = cy - r * ny;
        sz_out[i] = cz - r * nz;
        hit_out[i] = (uint8_t)(sep <= 0.0f);
    }
}

static size_t compact_crossings(const uint8_t *cross, size_t *idx, size_t n) {
    size_t m = 0;
    for (size_t i = 0; i < n; ++i) if (cross[i]) idx[m++] = i;
    return m;
}

static void gather_crossings(
    const float *ax0, const float *ay0, const float *az0,
    const float *ax1, const float *ay1, const float *az1,
    const float *sA, const float *sB,
    const size_t *idx, size_t m,
    float *gax0, float *gay0, float *gaz0,
    float *gax1, float *gay1, float *gaz1,
    float *gsA, float *gsB)
{
    for (size_t j = 0; j < m; ++j) {
        const size_t i = idx[j];
        gax0[j] = ax0[i]; gay0[j] = ay0[i]; gaz0[j] = az0[i];
        gax1[j] = ax1[i]; gay1[j] = ay1[i]; gaz1[j] = az1[i];
        gsA[j] = sA[i]; gsB[j] = sB[i];
    }
}

static bench_result bench_dist(
    const float *ax0, const float *ay0, const float *az0,
    const float *ax1, const float *ay1, const float *az1,
    const float *pnx, const float *pny, const float *pnz, const float *pd,
    float *sA, float *sB, float *dist, uint8_t *cross,
    size_t n, int reps)
{
    bench_result r = {1e30, 0.0};
    for (int k = 0; k < reps; ++k) {
        double t0 = now_sec();
        segplane3d_dist_soa(ax0, ay0, az0, ax1, ay1, az1, pnx, pny, pnz, pd, sA, sB, dist, cross, n);
        double t1 = now_sec();
        double dt = t1 - t0;
        if (dt < r.best_sec) r.best_sec = dt;
    }
    r.mpairs_per_sec = (double)n / r.best_sec / 1e6;
    return r;
}

static bench_result bench_intersections(
    const float *ax0, const float *ay0, const float *az0,
    const float *ax1, const float *ay1, const float *az1,
    const float *sA, const float *sB,
    float *t, float *px, float *py, float *pz,
    size_t n, int reps)
{
    bench_result r = {1e30, 0.0};
    for (int k = 0; k < reps; ++k) {
        double t0 = now_sec();
        segplane3d_intersections_soa(ax0, ay0, az0, ax1, ay1, az1, sA, sB, t, px, py, pz, n);
        double t1 = now_sec();
        double dt = t1 - t0;
        if (dt < r.best_sec) r.best_sec = dt;
    }
    r.mpairs_per_sec = (double)n / r.best_sec / 1e6;
    return r;
}

static bench_result bench_capsule(
    const float *ax0, const float *ay0, const float *az0,
    const float *ax1, const float *ay1, const float *az1,
    const float *pnx, const float *pny, const float *pnz, const float *pd,
    const float *rad,
    float *sep, float *pen, float *cx, float *cy, float *cz,
    float *sx, float *sy, float *sz, uint8_t *hit,
    size_t n, int reps)
{
    bench_result r = {1e30, 0.0};
    for (int k = 0; k < reps; ++k) {
        double t0 = now_sec();
        capsuleplane3d_contact_soa(ax0, ay0, az0, ax1, ay1, az1, pnx, pny, pnz, pd, rad,
                                   sep, pen, cx, cy, cz, sx, sy, sz, hit, n);
        double t1 = now_sec();
        double dt = t1 - t0;
        if (dt < r.best_sec) r.best_sec = dt;
    }
    r.mpairs_per_sec = (double)n / r.best_sec / 1e6;
    return r;
}

int main(void) {
    const size_t n = 4000000;
    const int reps = 6;
    srand(12345);

#define ALLOC(name) float *name = (float*)aligned_malloc64(n * sizeof(float)); if (!name) return 1
#define ALLOCU8(name) uint8_t *name = (uint8_t*)aligned_malloc64(n * sizeof(uint8_t)); if (!name) return 1
#define ALLOCSZ(name) size_t *name = (size_t*)aligned_malloc64(n * sizeof(size_t)); if (!name) return 1

    ALLOC(ax0); ALLOC(ay0); ALLOC(az0); ALLOC(ax1); ALLOC(ay1); ALLOC(az1);
    ALLOC(pnx); ALLOC(pny); ALLOC(pnz); ALLOC(pd); ALLOC(rad);
    ALLOC(sA); ALLOC(sB); ALLOC(dist); ALLOCU8(cross);
    ALLOC(t); ALLOC(px); ALLOC(py); ALLOC(pz);
    ALLOC(sep); ALLOC(pen); ALLOC(cx); ALLOC(cy); ALLOC(cz); ALLOC(sx); ALLOC(sy); ALLOC(sz); ALLOCU8(hit);
    ALLOCSZ(idx);

    fill_segments_and_planes(ax0, ay0, az0, ax1, ay1, az1, pnx, pny, pnz, pd, rad, n, 10.0f, 0.5f);

    bench_result rd = bench_dist(ax0, ay0, az0, ax1, ay1, az1, pnx, pny, pnz, pd, sA, sB, dist, cross, n, reps);
    size_t m = compact_crossings(cross, idx, n);

    float *gax0=(float*)aligned_malloc64(m*sizeof(float)); float *gay0=(float*)aligned_malloc64(m*sizeof(float)); float *gaz0=(float*)aligned_malloc64(m*sizeof(float));
    float *gax1=(float*)aligned_malloc64(m*sizeof(float)); float *gay1=(float*)aligned_malloc64(m*sizeof(float)); float *gaz1=(float*)aligned_malloc64(m*sizeof(float));
    float *gsA =(float*)aligned_malloc64(m*sizeof(float)); float *gsB =(float*)aligned_malloc64(m*sizeof(float));
    float *gt  =(float*)aligned_malloc64(m*sizeof(float)); float *gpx =(float*)aligned_malloc64(m*sizeof(float)); float *gpy =(float*)aligned_malloc64(m*sizeof(float)); float *gpz =(float*)aligned_malloc64(m*sizeof(float));
    if (!gax0||!gay0||!gaz0||!gax1||!gay1||!gaz1||!gsA||!gsB||!gt||!gpx||!gpy||!gpz) return 1;

    gather_crossings(ax0, ay0, az0, ax1, ay1, az1, sA, sB, idx, m, gax0, gay0, gaz0, gax1, gay1, gaz1, gsA, gsB);
    bench_result ri = bench_intersections(gax0, gay0, gaz0, gax1, gay1, gaz1, gsA, gsB, gt, gpx, gpy, gpz, m, reps);
    bench_result rc = bench_capsule(ax0, ay0, az0, ax1, ay1, az1, pnx, pny, pnz, pd, rad, sep, pen, cx, cy, cz, sx, sy, sz, hit, n, reps);

    size_t hits = 0;
    for (size_t i = 0; i < n; ++i) hits += (size_t)hit[i];

    printf("segment-plane 3D batched kernels benchmark\n");
    printf("n = %zu\n", n);
    printf("crossings compacted = %zu (%.2f%%)\n", m, 100.0 * (double)m / (double)n);
    printf("capsule-plane hits = %zu (%.2f%%)\n", hits, 100.0 * (double)hits / (double)n);
    printf("\n");
    printf("dist kernel:         %.6f s  %.2f Mpairs/s\n", rd.best_sec, rd.mpairs_per_sec);
    printf("intersection kernel: %.6f s  %.2f Mpairs/s  (on compacted crossings)\n", ri.best_sec, ri.mpairs_per_sec);
    printf("capsule kernel:      %.6f s  %.2f Mpairs/s\n", rc.best_sec, rc.mpairs_per_sec);

    // basic checksum so compiler keeps work
    double checksum = 0.0;
    checksum += dist[123] + sA[456] + sB[789];
    if (m > 0) checksum += gt[m/2] + gpx[m/3] + gpy[m/4] + gpz[m/5];
    checksum += sep[111] + pen[222] + cx[333] + cy[444] + cz[555] + sx[666] + sy[777] + sz[888];
    printf("checksum = %.9f\n", checksum);
    return 0;
}