I have two compute kernels. The first kernel pre_initization_0 initializes some buffers using a parallel random number generator initialize some buffers. The second kernel pre_sum_weights_0 effectively sums weighted values buffer of a much smaller dimension.
#include <metal_stdlib>
#include <metal_simdgroup>
using namespace metal;
struct mt_state {
array<uint32_t, 624> array;
uint16_t index;
};
float random(device mt_state &state) {
uint16_t k = state.index;
uint16_t j = (k + 1) % 624;
uint32_t x = (state.array[k] & 0x80000000U) |
(state.array[j] & 0x7fffffffU);
uint32_t xA = x >> 1;
if (x & 0x00000001U) {
xA ^= 0x9908b0dfU;
}
j = (k + 397) % 624;
x = state.array[j]^xA;
state.array[k] = x;
state.index = (k + 1) % 624;
uint32_t y = x^(x >> 11);
y = y^((y << 7) & 0x9d2c5680U);
y = y^((y << 15) & 0xefc60000U);
return static_cast<float> (y^(y >> 18));
}
kernel void pre_initization_0(
device float *vc30c09b98 [[buffer(0)]], // x used 0
device float *vc30c09c38 [[buffer(1)]], // v_{||} used 0
device float *vc30c09cd8 [[buffer(2)]], // v_{\perp} used 0
device mt_state *sc30c50c18 [[buffer(3)]],
constant uint32_t &offset [[buffer(4)]],
uint index [[thread_position_in_grid]]) {
if (offset + index < 3000000) {
device mt_state &rc30c50c18 = sc30c50c18[index]; // used 4
const float rc30c06218 = 2.17689351e-08; // used 1
const float rc30c06318 = -46.7484322; // used 1
const float rc30c0a458 = fma(rc30c06218, random(rc30c50c18), rc30c06318); // used 1
const float rc30c07718 = 5.18059896e-05; // used 1
const float rc30c06798 = -1; // used 1
const float rc30c06618 = 2.32830644e-10; // used 2
const float rc30c06718 = 1.17549435e-38; // used 2
const float rc30c0a4f8 = fma(rc30c06618, random(rc30c50c18), rc30c06718); // used 1
const float r1034f02f8 = log(rc30c0a4f8); // used 1
const float rc30c06818 = rc30c06798*r1034f02f8; // used 1
const float r1034eedf8 = sqrt(rc30c06818); // used 1
const float rc30c06418 = 1.46291812e-09; // used 1
const float rc30c06498 = rc30c06418*random(rc30c50c18); // used 1
const float r1034eee68 = sin(rc30c06498); // used 1
const float rc30c06a18 = r1034eedf8*r1034eee68; // used 1
const float rc30c07698 = rc30c07718*rc30c06a18; // used 1
const float rc30c07598 = 3.3356411e-09; // used 1
const float rc30c07318 = -241213328; // used 1
const float rc30c0a598 = fma(rc30c06618, random(rc30c50c18), rc30c06718); // used 1
const float r1034ef238 = log(rc30c0a598); // used 1
const float rc30c06c98 = rc30c07318*r1034ef238; // used 1
const float rc310b0018 = sqrt(rc30c06c98); // used 1
const float rc30c07618 = rc30c07598*rc310b0018; // used 1
vc30c09b98[offset + index] = rc30c0a458;
vc30c09c38[offset + index] = rc30c07698;
vc30c09cd8[offset + index] = rc30c07618;
}
}
kernel void pre_sum_weights_0(
constant float *vc30c09b98 [[buffer(0)]], // x used 7
device atomic_float *vc30c0a098 [[buffer(1)]],
const texture1d<float, access::read> ac310a9000 [[texture(0)]],
const texture1d<float, access::read> ac310d8000 [[texture(1)]],
const texture1d<float, access::read> ac310d9000 [[texture(2)]],
const texture1d<float, access::read> ac310da000 [[texture(3)]],
uint index [[thread_position_in_grid]]) {
if (index < 3000000) {
const float rc30c09b98 = vc30c09b98[index]; // x used 7
const float rc30c07798 = 0.0935904533; // used 2
const float rc30d08318 = rc30c09b98 - rc30c07798; // used 1
const ushort ic30d08398 = (ushort)min(max((rc30d08318 - -46.7484322)/0.0935904533,(float)0),(float)999); // used 1
const float rc30c07c98 = -5.34242535; // used 1
const ushort ic30d08018 = (ushort)min(max((rc30c09b98 - -46.7484322)/0.0935904533,(float)0),(float)999); // used 5
const float rc30c51018 = ac310da000.read(ic30d08018).r; // used 1
const float rc30c0a6d8 = fma(rc30c07c98, rc30c09b98, rc30c51018); // used 1
const float rc30c07a18 = -10.6848507; // used 1
const float rc30c50e98 = ac310d8000.read(ic30d08018).r; // used 1
const float rc30c0a778 = fma(rc30c07a18, rc30c09b98, rc30c50e98); // used 1
const float rc30c07b98 = rc30c0a6d8*rc30c0a778; // used 1
atomic_fetch_add_explicit(&vc30c0a098[ic30d08398], rc30c07b98, memory_order_relaxed); // used 1
const float rc30c07c18 = 0.75; // used 1
const float rc30c07e18 = 114.166031; // used 1
const float rc30c50d18 = ac310a9000.read(ic30d08018).r; // used 1
const float rc30d08098 = rc30c50d18 - rc30c09b98; // used 1
const float rc30c07d98 = rc30d08098*rc30d08098; // used 1
const float rc30c07e98 = rc30c07e18*rc30c07d98; // used 1
const float rc30d08218 = rc30c07c18 - rc30c07e98; // used 1
atomic_fetch_add_explicit(&vc30c0a098[ic30d08018], rc30d08218, memory_order_relaxed); // used 1
const float rc30d08498 = rc30c07798 + rc30c09b98; // used 1
const ushort ic30d08418 = (ushort)min(max((rc30d08498 - -46.7484322)/0.0935904533,(float)0),(float)999); // used 1
const float rc30c07b18 = 0.5; // used 1
const float rc30c07918 = 1.5; // used 1
const float rc30c07818 = 10.6848507; // used 1
const float rc30c50f98 = ac310d9000.read(ic30d08018).r; // used 1
const float rc30d08298 = rc30c50f98 - rc30c09b98; // used 1
const float rc30c07a98 = rc30c07818*rc30d08298; // used 1
const float rc30d08118 = rc30c07918 - rc30c07a98; // used 1
const float rc30d34018 = rc30d08118*rc30d08118; // used 1
const float rc30d34098 = rc30c07b18*rc30d34018; // used 1
atomic_fetch_add_explicit(&vc30c0a098[ic30d08418], rc30d34098, memory_order_relaxed); // used 1
}
}
Running some timing experiments, If I initialize [[buffer(0)]] of the pre_sum_weights_0 kernel on the CPU, and run the kernel, it will run in 0.0489688 s measured using GPUEndTime - GPUStartTime in a completedHander for the command buffer.
If I run the initialization kernel pre_initization_0 before this, the kernel execution time of pre_sum_weights_0 slows to 7.20396 s.
Note I am launching these kernels from different command buffers.
Thanks for including both kernels in full. It looks like you have a 1D particle-in-cell plasma simulation here. Having the actual shader source meant I could measure this rather than reason about it. Everything below is from an M5 Max running macOS 26.5.1, so read the absolute numbers as relative rather than as predictions for your hardware.
Two different things can make a kernel look slower when another one runs ahead of it, and they are worth separating before anything else:
- The kernel's own execution time changed, because the data it reads changed.
- The kernel's execution time did not change. It spent time waiting for the earlier work to finish, and that wait got counted against it.
Both turn out to be in play here, and neither one alone accounts neatly for a factor of 147.
On the first: pre_sum_weights_0 has an execution time that depends on its input. The index it writes to is computed from the value it reads out of vc30c09b98, your buffer(0):
const ushort ic30d08018 = (ushort)min(max((rc30c09b98 - -46.7484322)/0.0935904533,(float)0),(float)999);
...
atomic_fetch_add_explicit(&vc30c0a098[ic30d08018], rc30d08218, memory_order_relaxed);
Atomic adds to the same address have to serialize, one after another. The kernel issues 9,000,000 of them, three per thread across 3,000,000 threads, into the 1000 entries of vc30c0a098. When the values in vc30c09b98 are spread across the range, those adds resolve to many different addresses and proceed concurrently. When the values are similar, the adds resolve to the same few addresses and serialize. So the execution time varies with the contents of vc30c09b98, independently of what ran before.
In each test I ran pre_initization_0 first. The only thing I changed was whether the generator state (sc30c50c18, buffer(3)) had been given a starting seed. With a seed, the values written to vc30c09b98 covered the full range and pre_sum_weights_0 was fast. Without one, every element got the same value, every add resolved to the same address, and it was slow. I also tried a third kernel that wrote just as much data but with values deliberately spread across the range, and it was fast again. So the earlier kernel is not what slows the later one. The distribution of values in vc30c09b98 is.
Worth checking on the seed, since the posted code never seeds sc30c50c18: an all-zero MT19937 state emits zero forever, and then
x = fma(2.17689351e-08, 0.0, -46.7484322) = -46.7484322
holds for every element. ic30d08398 and ic30d08018 then both evaluate to 0 and ic30d08418 to 1, so all 9,000,000 adds land on two addresses. That exact value is quick to look for in vc30c09b98.
But this only accounts for about 6x, not 147x. I tested the atomic adds in isolation. Collapsing all 9,000,000 onto one address cost roughly 6 to 7x against spreading them over the full 1000. The penalty flattens out by about eight distinct addresses.
Which brings up the second possibility. pre_initization_0 cost 31.5 ms of GPU time here against 0.085 ms for pre_sum_weights_0, so it is roughly 370 times the more expensive of the two. If any part of that wait were counted against the second kernel, it would swamp the second kernel's own time completely. I committed the pre_initization_0 command buffers without waiting and then committed pre_sum_weights_0, and its GPUStartTime and GPUEndTime still excluded the wait on this setup. That is worth confirming on yours, because the arithmetic makes it the natural candidate for the remaining gap.
One more sign points that way. Your faster figure, 0.0489688 s, is already about 575 times the 0.085 ms that pre_sum_weights_0 costs here. Even allowing for different hardware, both of your numbers are large next to what the kernel itself does. The sharper question may be what is being counted inside the faster one.
Separately, the largest saving I found is in the other kernel. Each mt_state is about 2.5 KB per thread. Covering all 3,000,000 threads at once would need roughly 7 GB, which is presumably why buffer(4)'s offset exists. Replacing the generator with a counter-based one took that 31.5 ms to 0.053 ms, with your output math for vc30c09b98, vc30c09c38, and vc30c09cd8 unchanged. It runs as a single dispatch with no state buffer at all. The values came out distributed the same way, spanning -46.7484 to 46.7458 and covering 999 of the 1000 indices:
inline uint32_t mix32(uint32_t z) {
z ^= z >> 16; z *= 0x7feb352du;
z ^= z >> 15; z *= 0x846ca68bu;
z ^= z >> 16; return z;
}
inline float random_counter(uint32_t key, thread uint32_t &ctr) {
return static_cast<float>(mix32(key ^ mix32(ctr++ * 0x9e3779b9u)));
}
// per thread: uint32_t ctr = 0; uint32_t key = mix32(index ^ seed);
pre_sum_weights_0 itself already looks well structured. Replicating vc30c0a098 across copies and accumulating into threadgroup memory both measured slower here than what you have. The threadgroup route needs a compare-exchange loop, since atomic_float is device address space only, and that loop is what costs. Reading buffer(0) as device rather than constant made no difference either.
Here are two items you can check that would likely provide answers:
pre_initization_0's own time, measured the same way you measuredpre_sum_weights_0. If it dominates on your hardware the way it does here, that locates the 7.2 s.- a histogram of
ic30d08018over all ofvc30c09b98, for both initialization paths. If the GPU-written case lands on one or two indices, that is the first mechanism above.
One thing to look for in your host code: whether the pre_sum_weights_0 command buffer is committed while the pre_initization_0 command buffers are still executing.
If you post those results, along with your GPU model and the thread count and threadgroup size you dispatch, I am happy to keep digging.