[OMP] @par(schedule='static') lowers to chunked static scheduling

[BI71317]

Summary

Codon lowers @par(schedule='static') to OpenMP chunked static scheduling, effectively behaving like schedule(static, 1) rather than no-chunk schedule(static).

This causes round-robin iteration assignment by default. For workloads where contiguous iteration ownership matters, this hurts locality, reduces throughput, and makes default static scheduling behave differently from the equivalent C/OpenMP no-chunk form.

MRE 1. Perf test

Details

from time import time

N = 16777216
THREADS = 4
REPS = 8

def value_at(i: int, r: int) -> int:
    return ((i * 1103515245 + 12345 + r) ^ (i >> 7)) & 0x7fffffff

def checksum(a: Ptr[int]) -> int:
    total = 0
    for i in range(N):
        total += a[i] & 255
    return total

def bench_default(a: Ptr[int]) -> float:
    t0 = time()
    for r in range(REPS):
        @par(schedule='static', num_threads=THREADS)
        for i in range(N):
            a[i] = value_at(i, r)
    t1 = time()
    return t1 - t0

def bench_chunk1(a: Ptr[int]) -> float:
    t0 = time()
    for r in range(REPS):
        @par(schedule='static', num_threads=THREADS, chunk_size=1)
        for i in range(N):
            a[i] = value_at(i, r)
    t1 = time()
    return t1 - t0

def bench_block_chunk(a: Ptr[int]) -> float:
    t0 = time()
    for r in range(REPS):
        @par(schedule='static', num_threads=THREADS, chunk_size=N // THREADS)
        for i in range(N):
            a[i] = value_at(i, r)
    t1 = time()
    return t1 - t0

def run_one(name: str, fn, a: Ptr[int]):
    elapsed = fn(a)
    chk = checksum(a)
    bytes_written = float(N * 8 * REPS)
    gib_per_s = bytes_written / elapsed / (1024.0 * 1024.0 * 1024.0)
    print(name, "checksum:", chk, "time:", elapsed, "bandwidth:", gib_per_s, "GiB/s")

a = Ptr[int](N)
warm = bench_default(a)

print("Codon static round-robin penalty: N=", N, " THREADS=", THREADS, " REPS=", REPS, sep="")
run_one("@par(schedule='static')", bench_default, a)
run_one("@par(schedule='static', chunk_size=1)", bench_chunk1, a)
run_one("@par(schedule='static', chunk_size=N//THREADS)", bench_block_chunk, a)

Result

$ ./24_static_round_robin_penalty
Codon static round-robin penalty: N=16777216 THREADS=4 REPS=8
@par(schedule='static') checksum: 2139095040 time: 0.0911033 bandwidth: 10.9765 GiB/s
@par(schedule='static', chunk_size=1) checksum: 2139095040 time: 0.0951467 bandwidth: 10.5101 GiB/s
@par(schedule='static', chunk_size=N//THREADS) checksum: 2139095040 time: 0.0275688 bandwidth: 36.2729 GiB/s

As you can see, default static non-chunk works poorly as well as Chunk size with 1 RR does.

Almost same semantic code with C works differently.

Details

#include <omp.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>

#ifndef N
#define N 16777216
#endif

#ifndef THREADS
#define THREADS 4
#endif

#ifndef REPS
#define REPS 8
#endif

typedef long long i64;

static inline i64 value_at(i64 i, int r) {
  return ((i * 1103515245LL + 12345LL + r) ^ (i >> 7)) & 0x7fffffffLL;
}

static i64 checksum(i64 *a) {
  i64 total = 0;
  for (i64 i = 0; i < (i64)N; ++i)
    total += a[i] & 255;
  return total;
}

static void bench_static_default(i64 *a, double *elapsed) {
  double t0 = omp_get_wtime();
  for (int r = 0; r < REPS; ++r) {
#pragma omp parallel for schedule(static) num_threads(THREADS)
    for (i64 i = 0; i < (i64)N; ++i)
      a[i] = value_at(i, r);
  }
  *elapsed = omp_get_wtime() - t0;
}

static void bench_static1(i64 *a, double *elapsed) {
  double t0 = omp_get_wtime();
  for (int r = 0; r < REPS; ++r) {
#pragma omp parallel for schedule(static, 1) num_threads(THREADS)
    for (i64 i = 0; i < (i64)N; ++i)
      a[i] = value_at(i, r);
  }
  *elapsed = omp_get_wtime() - t0;
}

static void bench_static_block_chunk(i64 *a, double *elapsed) {
  double t0 = omp_get_wtime();
  for (int r = 0; r < REPS; ++r) {
#pragma omp parallel for schedule(static, N / THREADS) num_threads(THREADS)
    for (i64 i = 0; i < (i64)N; ++i)
      a[i] = value_at(i, r);
  }
  *elapsed = omp_get_wtime() - t0;
}

static void run_one(const char *name, void (*fn)(i64 *, double *), i64 *a) {
  double elapsed = 0.0;
  fn(a, &elapsed);
  i64 chk = checksum(a);
  double bytes = (double)N * sizeof(i64) * REPS;
  printf("%s checksum: %lld time: %.9f bandwidth: %.3f GiB/s\n", name, chk,
         elapsed, bytes / elapsed / (1024.0 * 1024.0 * 1024.0));
}

int main(void) {
  i64 *a = NULL;
  if (posix_memalign((void **)&a, 64, (size_t)N * sizeof(i64)) != 0)
    return 1;

  double warm_elapsed = 0.0;
  bench_static_default(a, &warm_elapsed);

  printf("C OpenMP static round-robin penalty: N=%d THREADS=%d REPS=%d\n", N,
         THREADS, REPS);
  run_one("schedule(static)", bench_static_default, a);
  run_one("schedule(static,1)", bench_static1, a);
  run_one("schedule(static,N/THREADS)", bench_static_block_chunk, a);

  free(a);
  return 0;
}

Build option and Result

gcc -O3 -march=native -fopenmp \
  <file_name>.c \
  -o <file_name>

OMP_DYNAMIC=FALSE OMP_PROC_BIND=close OMP_PLACES=cores \
./<file_name>

C OpenMP static round-robin penalty: N=16777216 THREADS=4 REPS=8
schedule(static) checksum: 2139095040 time: 0.030114027 bandwidth: 33.207 GiB/s
schedule(static,1) checksum: 2139095040 time: 0.077665618 bandwidth: 12.876 GiB/s
schedule(static,N/THREADS) checksum: 2139095040 time: 0.038195507 bandwidth: 26.181 GiB/s

MRE 2. check each iteration thread index

Details

import openmp as omp

N = 32
THREADS = 4

def print_owners(name: str, owners: Ptr[int]):
    print(name, "owners:", end="")
    for i in range(N):
        print(" ", owners[i], sep="", end="")
    print()

    for t in range(THREADS):
        print("  thread", t, ":", end="")
        for i in range(N):
            if owners[i] == t:
                print(" ", i, sep="", end="")
        print()

owners_default = Ptr[int](N)
owners_chunk1 = Ptr[int](N)
owners_chunk4 = Ptr[int](N)
owners_block = Ptr[int](N)

@par(schedule='static', num_threads=THREADS)
for i in range(N):
    owners_default[i] = omp.get_thread_num()

@par(schedule='static', num_threads=THREADS, chunk_size=1)
for i in range(N):
    owners_chunk1[i] = omp.get_thread_num()

@par(schedule='static', num_threads=THREADS, chunk_size=4)
for i in range(N):
    owners_chunk4[i] = omp.get_thread_num()

@par(schedule='static', num_threads=THREADS, chunk_size=N // THREADS)
for i in range(N):
    owners_block[i] = omp.get_thread_num()

print("Codon static distribution: N=", N, " THREADS=", THREADS, sep="")
print_owners("@par(schedule='static')", owners_default)
print_owners("@par(schedule='static', chunk_size=1)", owners_chunk1)
print_owners("@par(schedule='static', chunk_size=4)", owners_chunk4)
print_owners("@par(schedule='static', chunk_size=N//THREADS)", owners_block)

Result

$ ./23_static_distribution
Codon static distribution: N=32 THREADS=4
@par(schedule='static') owners: 0 1 2 3 0 1 2 3 0 1 2 3 0 1 2 3 0 1 2 3 0 1 2 3 0 1 2 3 0 1 2 3
  thread 0 : 0 4 8 12 16 20 24 28
  thread 1 : 1 5 9 13 17 21 25 29
  thread 2 : 2 6 10 14 18 22 26 30
  thread 3 : 3 7 11 15 19 23 27 31
@par(schedule='static', chunk_size=1) owners: 0 1 2 3 0 1 2 3 0 1 2 3 0 1 2 3 0 1 2 3 0 1 2 3 0 1 2 3 0 1 2 3
  thread 0 : 0 4 8 12 16 20 24 28
  thread 1 : 1 5 9 13 17 21 25 29
  thread 2 : 2 6 10 14 18 22 26 30
  thread 3 : 3 7 11 15 19 23 27 31
@par(schedule='static', chunk_size=4) owners: 0 0 0 0 1 1 1 1 2 2 2 2 3 3 3 3 0 0 0 0 1 1 1 1 2 2 2 2 3 3 3 3
  thread 0 : 0 1 2 3 16 17 18 19
  thread 1 : 4 5 6 7 20 21 22 23
  thread 2 : 8 9 10 11 24 25 26 27
  thread 3 : 12 13 14 15 28 29 30 31
@par(schedule='static', chunk_size=N//THREADS) owners: 0 0 0 0 0 0 0 0 1 1 1 1 1 1 1 1 2 2 2 2 2 2 2 2 3 3 3 3 3 3 3 3
  thread 0 : 0 1 2 3 4 5 6 7
  thread 1 : 8 9 10 11 12 13 14 15
  thread 2 : 16 17 18 19 20 21 22 23
  thread 3 : 24 25 26 27 28 29 30 31

As you can see, Yes, it really takes RR for default non chunk iteration launch.

Root Cause

Codon currently lowers default @par(schedule='static') as chunked static scheduling.

At the LLVM IR level, the generated call uses sched index 33 (sched index for static, unordered, chunk) instead of 34 (sched index for static, unordered, non-chunk):

...
  %11 = call {} @__kmpc_for_static_init_8(ptr nonnull @._STATIC_LOOP_IDENT.0, i32 %7, i32 33, ptr nonnull %3, ptr nonnull %4, ptr nonnull %5, ptr nonnull %6, i64 1, i64 %.unpack.i)
...

This is because in \codon\cir\transform\parallel\schedule.cpp\, nullIfneg func must return nullptr for non-chunk case:

Value *nullIfNeg(Value *v) {
  if (v && util::isConst<int64_t>(v) && util::getConst<int64_t>(v) <= 0)
    return nullptr;
  return v;
}

But when non-chunk case is given, chunk size variable in CIR is not Const -1 literal, but a Variable:

(var '".default.std.openmp.for_par.0:0.chunk_size.1258" (type '"Int64") (global false) (external false) (thread-local false))

so, func isConst is false, returning v even for non chunk case.