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.