conflict-set/InterleavingTest.cpp

#include <alloca.h>
#include <cassert>
#ifdef __x86_64__
#include <immintrin.h>
#endif

#include "third_party/nanobench.h"

struct Job {
  int *input;
  // Returned void* is a function pointer to the next continuation. We have to
  // use void* because otherwise the type would be recursive.
  typedef void *(*continuation)(Job *);
  continuation next;
};

void *stepJob(Job *j) {
  auto done = --(*j->input) == 0;
#ifdef __x86_64__
  _mm_clflush(j->input);
#endif
  return done ? nullptr : (void *)stepJob;
}

void sequential(Job **jobs, int count) {
  for (int i = 0; i < count; ++i) {
    do {
      jobs[i]->next = (Job::continuation)jobs[i]->next(jobs[i]);
    } while (jobs[i]->next);
  }
}

void sequentialNoFuncPtr(Job **jobs, int count) {
  for (int i = 0; i < count; ++i) {
    while (stepJob(jobs[i]))
      ;
  }
}

void interleaveSwapping(Job **jobs, int remaining) {
  int current = 0;
  while (remaining > 0) {
    auto next = (Job::continuation)jobs[current]->next(jobs[current]);
    jobs[current]->next = next;
    if (next == nullptr) {
      jobs[current] = jobs[remaining - 1];
      --remaining;
    } else {
      ++current;
    }
    if (current == remaining) {
      current = 0;
    }
  }
}

void interleaveBoundedCyclicList(Job **jobs, int count) {
  if (count == 0) {
    return;
  }

  constexpr int kConcurrent = 32;
  Job *inProgress[kConcurrent];
  int nextJob[kConcurrent];

  int started = std::min(kConcurrent, count);
  for (int i = 0; i < started; i++) {
    inProgress[i] = jobs[i];
    nextJob[i] = i + 1;
  }
  nextJob[started - 1] = 0;

  int prevJob = started - 1;
  int job = 0;
  for (;;) {
    auto next = (Job::continuation)inProgress[job]->next(inProgress[job]);
    inProgress[job]->next = next;
    if (next == nullptr) {
      if (started == count) {
        if (prevJob == job)
          break;
        nextJob[prevJob] = nextJob[job];
        job = prevJob;
      } else {
        int temp = started++;
        inProgress[job] = jobs[temp];
      }
    }
    prevJob = job;
    job = nextJob[job];
  }
}

#ifndef __has_attribute
#define __has_attribute(x) 0
#endif

#if __has_attribute(musttail)
#define MUSTTAIL __attribute__((musttail))
#else
#define MUSTTAIL
#endif

struct Context {
  constexpr static int kConcurrent = 32;
  Job **jobs;
  Job *inProgress[kConcurrent];
  void (*continuation[kConcurrent])(Context *, int64_t prevJob, int64_t job,
                                    int64_t started, int64_t count);
  int nextJob[kConcurrent];
};

void keepGoing(Context *context, int64_t prevJob, int64_t job, int64_t started,
               int64_t count) {
  prevJob = job;
  job = context->nextJob[job];
  MUSTTAIL return context->continuation[job](context, prevJob, job, started,
                                             count);
}

void stepJobTailCall(Context *context, int64_t prevJob, int64_t job,
                     int64_t started, int64_t count);

void complete(Context *context, int64_t prevJob, int64_t job, int64_t started,
              int64_t count) {
  if (started == count) {
    if (prevJob == job) {
      return;
    }
    context->nextJob[prevJob] = context->nextJob[job];
    job = prevJob;
  } else {
    context->inProgress[job] = context->jobs[started++];
    context->continuation[job] = stepJobTailCall;
  }
  prevJob = job;
  job = context->nextJob[job];
  MUSTTAIL return context->continuation[job](context, prevJob, job, started,
                                             count);
}

void stepJobTailCall(Context *context, int64_t prevJob, int64_t job,
                     int64_t started, int64_t count) {
  auto *j = context->inProgress[job];
  auto done = --(*j->input) == 0;
#ifdef __x86_64__
  _mm_clflush(j->input);
#endif
  if (done) {
    MUSTTAIL return complete(context, prevJob, job, started, count);
  } else {
    context->continuation[job] = stepJobTailCall;
    MUSTTAIL return keepGoing(context, prevJob, job, started, count);
  }
}

void useTailCalls(Job **jobs, int count) {
  if (count == 0) {
    return;
  }
  Context context;
  context.jobs = jobs;
  int64_t started = std::min(Context::kConcurrent, count);
  for (int i = 0; i < started; i++) {
    context.inProgress[i] = jobs[i];
    context.nextJob[i] = i + 1;
    context.continuation[i] = stepJobTailCall;
  }
  context.nextJob[started - 1] = 0;
  int prevJob = started - 1;
  int job = 0;
  return context.continuation[job](&context, prevJob, job, started, count);
}

void interleaveCyclicList(Job **jobs, int count) {
  auto *nextJob = (int *)alloca(sizeof(int) * count);

  for (int i = 0; i < count - 1; ++i) {
    nextJob[i] = i + 1;
  }
  nextJob[count - 1] = 0;

  int prevJob = count - 1;
  int job = 0;
  for (;;) {
    auto next = (Job::continuation)jobs[job]->next(jobs[job]);
    jobs[job]->next = next;
    if (next == nullptr) {
      if (prevJob == job)
        break;
      nextJob[prevJob] = nextJob[job];
      job = prevJob;
    }
    prevJob = job;
    job = nextJob[job];
  }
}

int main() {
  ankerl::nanobench::Bench bench;

  constexpr int kNumJobs = 10000;
  bench.relative(true);

  Job jobs[kNumJobs];
  Job jobsCopy[kNumJobs];
  int iters = 0;
  int originalInput[kNumJobs];
  for (int i = 0; i < kNumJobs; ++i) {
    originalInput[i] = rand() % 5 + 3;
    jobs[i].input = new int{originalInput[i]};
    jobs[i].next = stepJob;
    iters += *jobs[i].input;
  }
  bench.batch(iters);

  for (auto [scheduler, name] :
       {std::make_pair(sequentialNoFuncPtr, "sequentialNoFuncPtr"),
        std::make_pair(sequential, "sequential"),
        std::make_pair(useTailCalls, "useTailCalls"),
        std::make_pair(interleaveSwapping, "interleavingSwapping"),
        std::make_pair(interleaveBoundedCyclicList,
                       "interleaveBoundedCyclicList"),
        std::make_pair(interleaveCyclicList, "interleaveCyclicList")}) {
    for (int i = 0; i < kNumJobs; ++i) {
      *jobs[i].input = originalInput[i];
    }
    memcpy(jobsCopy, jobs, sizeof(jobs));
    Job *ps[kNumJobs];
    for (int i = 0; i < kNumJobs; ++i) {
      ps[i] = jobsCopy + i;
    }
    scheduler(ps, kNumJobs);
    for (int i = 0; i < kNumJobs; ++i) {
      if (*jobsCopy[i].input != 0) {
        fprintf(stderr, "%s failed\n", name);
        abort();
      }
    }

    bench.run(name, [&]() {
      for (int i = 0; i < kNumJobs; ++i) {
        *jobs[i].input = originalInput[i];
      }
      memcpy(jobsCopy, jobs, sizeof(jobs));
      Job *ps[kNumJobs];
      for (int i = 0; i < kNumJobs; ++i) {
        ps[i] = jobsCopy + i;
      }
      scheduler(ps, kNumJobs);
    });
  }
  for (int i = 0; i < kNumJobs; ++i) {
    delete jobs[i].input;
  }
}