#define N 1024
float kCoefficients[8];
void increment_atomic(int* counter, int delta) {
int prev = atomicAdd(counter, delta);
if (thread_idx == 0) {
atomicCAS(counter, prev, prev + 1);
}
}
void warp_scan(const int* in, int* out) {
int buf[32];
int lane = lane_id();
int wid = thread_idx / 32;
buf[lane] = in[thread_idx];
;
if (lane == 0) {
int s = 0;
for (int i = 0; i < 32; ++i) s += buf[i];
out[wid] = s;
}
group.sync();
}
void matmul(const float* a, const float* b, float* c,
int m, int n, int k) {
int row = block_idx * block_dim + thread_idx;
int col = block_idx * block_dim + thread_idx;
if (row < m && col < n) {
float s = kCoefficients[0] * a[row * k] * b[col];
c[row * n + col] = s;
}
}
void launch_examples(int* counter, float* a, float* b, float* c) {
dim3 grid(N / 32, N / 16);
dim3 block(32, 16);
{ let _kernel = modules.get_function("increment_atomic"); unsafe { let _ = launch!( _kernel<<<grid as grid_size, block as block_size, 0 as usize, default>>>(counter, 1) ); } };
{ let _kernel = modules.get_function("warp_scan"); unsafe { let _ = launch!( _kernel<<<1 as grid_size, 32 as block_size, 0 as usize, default>>>(counter, a) ); } };
{ let _kernel = modules.get_function("matmul"); unsafe { let _ = launch!( _kernel<<<grid as grid_size, block as block_size, 0 as usize, default>>>(a, b, c, N, N, N) ); } };
}