Compiler Explorer

Source code

#define STC_DEBUG // disable bounds check for better performance

#ifndef STC_CSPAN_H_INCLUDED
#define STC_CSPAN_H_INCLUDED

#include <assert.h>
#include <stdint.h>
#include <stdbool.h>
#include <stddef.h>
#include <time.h>
#include <stdio.h>

#define c_ASSERT(expr) assert(expr)
#define STC_INLINE static inline
#define STC_API static inline
#define STC_DEF static inline

#define c_MACRO_OVERLOAD(name, ...) \
    c_PASTE(c_CONCAT(name,_), c_NUMARGS(__VA_ARGS__))(__VA_ARGS__)
#define c_CONCAT(a, b) a ## b
#define c_PASTE(a, b) c_CONCAT(a, b)
#define c_EXPAND(...) __VA_ARGS__
#define c_NUMARGS(...) _c_APPLY_ARG_N((__VA_ARGS__, _c_RSEQ_N))
#define _c_APPLY_ARG_N(args) c_EXPAND(_c_ARG_N args)
#define _c_RSEQ_N 16, 15, 14, 13, 12, 11, 10, 9, 8, 7, 6, 5, 4, 3, 2, 1, 0
#define _c_ARG_N(_1, _2, _3, _4, _5, _6, _7, _8, _9, _10, _11, _12, _13, \
                 _14, _15, _16, N, ...) N

#define c_make(C, ...) \
  C##_from_n((C##_raw[])__VA_ARGS__, c_sizeof((C##_raw[])__VA_ARGS__)/c_sizeof(C##_raw))
#define c_arraylen(a) \
  (intptr_t)(sizeof(a)/sizeof 0[a])
#define c_static_assert(b) \
  ((int)(0*sizeof(int[(b) ? 1 : -1])))
#define c_LTu(a, b) \
  ((size_t)(a) < (size_t)(b))

// CSPAN

#define using_cspan(...) c_MACRO_OVERLOAD(using_cspan, __VA_ARGS__)
#define using_cspan_2(Self, T) using_cspan_3(Self, T, 1)
#define using_cspan_3(Self, T, RANK) \
    typedef T Self##_value; typedef T Self##_raw; \
    typedef struct { \
        Self##_value *data; \
        int32_t shape[RANK]; \
        cspan_idx##RANK stride; \
    } Self; \
    typedef struct { Self##_value *ref; int32_t pos[RANK]; const Self *_s; } Self##_iter; \
    \
    STC_INLINE Self Self##_from_n(Self##_raw* raw, const intptr_t n) { \
        return (Self){.data=raw, .shape={(int32_t)n}}; \
    } \
    STC_INLINE Self Self##_slice_(Self##_value* v, const int32_t shape[], const int32_t stri[], \
                                     const int rank, const int32_t a[][2]) { \
        Self s = {.data=v}; int outrank; \
        s.data += _cspan_slice(s.shape, s.stride.d, &outrank, shape, stri, rank, a); \
        c_ASSERT(outrank == RANK); \
        return s; \
    } \
    STC_INLINE Self##_iter Self##_begin(const Self* self) { \
        Self##_iter it = {.ref=self->data, .pos={0}, ._s=self}; \
        return it; \
    } \
    STC_INLINE Self##_iter Self##_end(const Self* self) { \
        Self##_iter it = {.ref=NULL}; \
        return it; \
    } \
    STC_INLINE void Self##_next(Self##_iter* it) { \
        it->ref += _cspan_next##RANK(RANK, it->pos, it->_s->shape, it->_s->stride.d); \
        if (it->pos[0] == it->_s->shape[0]) it->ref = NULL; \
    } \
    struct stc_nostruct

#define using_cspan2(Self, T) using_cspan_3(Self, T, 1); using_cspan_3(Self##2, T, 2)
#define using_cspan3(Self, T) using_cspan2(Self, T); using_cspan_3(Self##3, T, 3)
#define using_cspan4(Self, T) using_cspan3(Self, T); using_cspan_3(Self##4, T, 4)
typedef struct { int32_t d[1]; } cspan_idx1;
typedef struct { int32_t d[2]; } cspan_idx2;
typedef struct { int32_t d[3]; } cspan_idx3;
typedef struct { int32_t d[4]; } cspan_idx4;
typedef struct { int32_t d[5]; } cspan_idx5;
#define c_END -1
#define c_ALL 0,-1

#define cspan_md(array, ...) \
    {.data=array, .shape={__VA_ARGS__}, .stride={.d={__VA_ARGS__}}}

/* For static initialization, use cspan_make(). c_make() for non-static only. */
#define cspan_make(SpanType, ...) \
    {.data=(SpanType##_value[])__VA_ARGS__, .shape={sizeof((SpanType##_value[])__VA_ARGS__)/sizeof(SpanType##_value)}}

#define cspan_slice(OutSpan, parent, ...) \
    OutSpan##_slice_((parent)->data, (parent)->shape, (parent)->stride.d, cspan_rank(parent) + \
                     c_static_assert(cspan_rank(parent) == sizeof((int32_t[][2]){__VA_ARGS__})/sizeof(int32_t[2])), \
                     (const int32_t[][2]){__VA_ARGS__})
     
/* create a cspan from a cvec, cstack, cdeq, cqueue, or cpque (heap) */
#define cspan_from(container) \
    {.data=(container)->data, .shape={(int32_t)(container)->_len}}

#define cspan_from_array(array) \
    {.data=(array) + c_static_assert(sizeof(array) != sizeof(void*)), .shape={c_arraylen(array)}}

#define cspan_size(self) _cspan_size((self)->shape, cspan_rank(self))
#define cspan_rank(self) c_arraylen((self)->shape)

#define cspan_index(self, ...) c_PASTE(cspan_idx_, c_NUMARGS(__VA_ARGS__))(self, __VA_ARGS__)
#define cspan_idx_1 cspan_idx_4
#define cspan_idx_2 cspan_idx_4
#define cspan_idx_3 cspan_idx_4
#define cspan_idx_4(self, ...) \
    c_PASTE(_cspan_idx, c_NUMARGS(__VA_ARGS__))((self)->shape, (self)->stride, __VA_ARGS__) // small/fast
#define cspan_idx_5(self, ...) \
    (_cspan_idxN(c_NUMARGS(__VA_ARGS__), (self)->shape, (self)->stride.d, (int32_t[]){__VA_ARGS__}) + \
     c_static_assert(cspan_rank(self) == c_NUMARGS(__VA_ARGS__))) // general

#define cspan_at(self, ...) ((self)->data + cspan_index(self, __VA_ARGS__))
#define cspan_front(self) ((self)->data)
#define cspan_back(self) ((self)->data + cspan_size(self) - 1)

// cspan_subspanN. (N<4) Optimized, same as e.g. cspan_slice(Span3, &ms3, {off,off+count}, {c_ALL}, {c_ALL});
#define cspan_subspan(self, offset, count) \
    {.data=cspan_at(self, offset), .shape={count}}
#define cspan_subspan2(self, offset, count) \
    {.data=cspan_at(self, offset, 0), .shape={count, (self)->shape[1]}, .stride={(self)->stride}}
#define cspan_subspan3(self, offset, count) \
    {.data=cspan_at(self, offset, 0, 0), .shape={count, (self)->shape[1], (self)->shape[2]}, .stride={(self)->stride}}

// cspan_submdN: reduce rank (N<5) Optimized, same as e.g. cspan_slice(Span2, &ms4, {x}, {y}, {c_ALL}, {c_ALL});
#define cspan_submd4(...) c_MACRO_OVERLOAD(cspan_submd4, __VA_ARGS__)
#define cspan_submd3(...) c_MACRO_OVERLOAD(cspan_submd3, __VA_ARGS__)
#define cspan_submd2(self, x) \
    {.data=cspan_at(self, x, 0), .shape={(self)->shape[1]}}
#define cspan_submd3_2(self, x) \
    {.data=cspan_at(self, x, 0, 0), .shape={(self)->shape[1], (self)->shape[2]}, \
                                    .stride={.d={0, (self)->stride.d[2]}}}
#define cspan_submd3_3(self, x, y) \
    {.data=cspan_at(self, x, y, 0), .shape={(self)->shape[2]}}
#define cspan_submd4_2(self, x) \
    {.data=cspan_at(self, x, 0, 0, 0), .shape={(self)->shape[1], (self)->shape[2], (self)->shape[3]}, \
                                       .stride={.d={0, (self)->stride.d[2], (self)->stride.d[3]}}}
#define cspan_submd4_3(self, x, y) \
    {.data=cspan_at(self, x, y, 0, 0), .shape={(self)->shape[2], (self)->shape[3]}, .stride={.d={0, (self)->stride.d[3]}}}
#define cspan_submd4_4(self, x, y, z) \
    {.data=cspan_at(self, x, y, z, 0), .shape={(self)->shape[3]}}

// private definitions:

STC_INLINE intptr_t _cspan_size(const int32_t shape[], int rank) {
    intptr_t sz = shape[0];
    while (rank-- > 1) sz *= shape[rank];
    return sz;
}

STC_INLINE intptr_t _cspan_idx1(const int32_t shape[1], const cspan_idx1 stri, int32_t x)
    { c_ASSERT(c_LTu(x, shape[0])); return x; }

STC_INLINE intptr_t _cspan_idx2(const int32_t shape[2], const cspan_idx2 stri, int32_t x, int32_t y)
    { c_ASSERT(c_LTu(x, shape[0]) && c_LTu(y, shape[1])); return (intptr_t)stri.d[1]*x + y; }

STC_INLINE intptr_t _cspan_idx3(const int32_t shape[3], const cspan_idx3 stri, int32_t x, int32_t y, int32_t z) {
    c_ASSERT(c_LTu(x, shape[0]) && c_LTu(y, shape[1]) && c_LTu(z, shape[2]));
    return (intptr_t)stri.d[2]*(stri.d[1]*x + y) + z;
}
STC_INLINE intptr_t _cspan_idx4(const int32_t shape[4], const cspan_idx4 stri, int32_t x, int32_t y,
                                                                               int32_t z, int32_t w) {
    c_ASSERT(c_LTu(x, shape[0]) && c_LTu(y, shape[1]) && c_LTu(z, shape[2]) && c_LTu(w, shape[3]));
    return (intptr_t)stri.d[3]*(stri.d[2]*(stri.d[1]*x + y) + z) + w;
}

STC_API intptr_t _cspan_idxN(int rank, const int32_t shape[], const int32_t stri[], const int32_t a[]);
STC_API intptr_t _cspan_next2(int rank, int32_t pos[], const int32_t shape[], const int32_t stride[]);
#define _cspan_next1(r, pos, d, s) (++pos[0], 1)
#define _cspan_next3 _cspan_next2
#define _cspan_next4 _cspan_next2
#define _cspan_next5 _cspan_next2

STC_API intptr_t _cspan_slice(int32_t odim[], int32_t ostri[], int* orank, 
                              const int32_t shape[], const int32_t stri[], 
                              int rank, const int32_t a[][2]);

/* -------------------------- IMPLEMENTATION ------------------------- */

STC_DEF intptr_t _cspan_idxN(int rank, const int32_t shape[], const int32_t stri[], const int32_t a[]) {
    intptr_t off = a[0];
    bool ok = c_LTu(a[0], shape[0]);
    for (int i = 1; i < rank; ++i) {
        off *= stri[i];
        off += a[i];
        ok &= c_LTu(a[i], shape[i]);
    }
    c_ASSERT(ok);
    return off;
}

STC_DEF intptr_t _cspan_next2(int rank, int32_t pos[], const int32_t shape[], const int32_t stride[]) {
    intptr_t off = 1, rs = 1;
    ++pos[rank - 1];
    while (--rank && pos[rank] == shape[rank]) {
        pos[rank] = 0, ++pos[rank - 1];
        const intptr_t ds = rs*shape[rank];
        rs *= stride[rank];
        off += rs - ds;
    }
    return off;
}

STC_DEF intptr_t _cspan_slice(int32_t odim[], int32_t ostri[], int* orank, 
                              const int32_t shape[], const int32_t stri[], 
                              int rank, const int32_t a[][2]) {
    intptr_t off = 0;
    int i = 0, j = 0, ok = true;
    int32_t t, s = 1;
    for (; i < rank; ++i) {
        off *= stri[i];
        off += a[i][0];
        switch (a[i][1]) { 
            case 0: s *= stri[i]; ok &= c_LTu(a[i][0], shape[i]); continue;
            case -1: t = shape[i]; break;
            default: t = a[i][1]; break; 
        }
        odim[j] = t - a[i][0];
        ostri[j] = s*stri[i];
        ok &= c_LTu(0, odim[j]) & !c_LTu(shape[i], t);
        s = 1; ++j;
    }
    *orank = j;
    c_ASSERT(ok);
    return off;
}
#endif

//////////////////////////////////////////////////////////////////////////////////////////

using_cspan3(MD, double);

// define the dimensions of a 3d-array
enum {nx = 50, ny = 51, nz = 100};
intptr_t n = 10000;

// define the contents of two nx x ny x nz arrays in and out
double Vout[nx*ny*nz];
double Vin[nx*ny*nz]; //, 1.23;

// define some slice indices for each dimension
int lx=10, ly=20, lz=30;
int hx=25, hy=35, hz=45;

static void TraditionalForLoop(intptr_t state) {

clock_t t = clock();
  double sum = 0;
  for (int s=0; s<state; ++s) {
    for (int x=lx; x<hx; ++x)
      for (int y=ly; y<hy; ++y)
        for (int z=lz; z<hz; ++z) {
          double d = Vin[nz*(ny*x + y) + z];
          Vout[nz*(ny*x + y) + z] += d, sum += d;
        }
  }
  t = clock() - t;
  printf("forloop: %.1f ms, %f\n", 1000.0f*t/CLOCKS_PER_SEC, sum);
}

static void MDRanges_setup(intptr_t state) {

double sum = 0;
  // Code before the loop is not measured
  clock_t t = clock();
  for (intptr_t s=0; s<state; ++s) {
    MD3 r_in = cspan_md(Vin, nx, ny, nz);
    MD3 r_out = cspan_md(Vout, nx, ny, nz);

r_in = cspan_slice(MD3, &r_in, {lx, hx}, {ly, hy}, {lz, hz});
    r_out = cspan_slice(MD3, &r_out, {lx, hx}, {ly, hy}, {lz, hz});
    MD3_iter i = MD3_begin(&r_in); // can be iterated "flat".
    MD3_iter o = MD3_begin(&r_out);
    sum += Vin[s % nx];
  }
  t = clock() - t;
  printf("setup: %.1f ms, %f\n", 1000.0f*t/CLOCKS_PER_SEC, sum);
}

static void MDRanges_loop_over_joined(intptr_t state) {

MD3 r_in = cspan_md(Vin, nx, ny, nz);
  MD3 r_out = cspan_md(Vout, nx, ny, nz);
  r_in = cspan_slice(MD3, &r_in, {lx, hx}, {ly, hy}, {lz, hz});
  r_out = cspan_slice(MD3, &r_out, {lx, hx}, {ly, hy}, {lz, hz});

// C++23: for (auto [o, i] : std::views::zip(flat(r_out), flat(r_in))) { o = i; }
  double sum = 0;
  clock_t t = clock();
  for (intptr_t s=0; s<state; ++s) {
    MD3_iter i = MD3_begin(&r_in);
    MD3_iter o = MD3_begin(&r_out);

// Code before the loop is not measured
    for (; i.ref; MD3_next(&i), MD3_next(&o)) {
        *o.ref += *i.ref; sum += *i.ref;
    }
  }
  t = clock() - t;
  printf("joined: %.1f ms, %f\n", 1000.0f*t/CLOCKS_PER_SEC, sum);
}

static void MDRanges_nested_loop(intptr_t state) {

// C++23: for (auto [o, i] : std::views::zip(flat(r_out), flat(r_in))) { o = i; }

// Code before the loop is not measured
  clock_t t = clock();
  double sum = 0;
  for (intptr_t s=0; s<state; ++s) {
    for (int x = 0; x < r_in.shape[0]; ++x)
      for (int y = 0; y < r_in.shape[1]; ++y)
          for (int z = 0; z < r_in.shape[2]; ++z) {
              double d = *cspan_at(&r_in, x,y,z);
              *cspan_at(&r_out, x,y,z) += d, sum += d;
           }

}
  t = clock() - t;
  printf("nested: %.1f ms, %f\n", 1000.0f*t/CLOCKS_PER_SEC, sum);
}

int main()
{
    for (int i=0; i<nx*ny*nz; ++i) Vin[i] = i + 0.25678;
    MDRanges_setup(n);
    TraditionalForLoop(n);
    MDRanges_nested_loop(n);
    MDRanges_loop_over_joined(n);
}