#include <hip/hip_runtime.h>
#define N 1024
__constant__ float kCoefficients[8];
__global__ void increment_atomic(int* counter, int delta) {
int prev = atomicAdd(counter, delta);
if (threadIdx.x == 0) {
atomicCAS(counter, prev, prev + 1);
}
}
__global__ void warp_scan(const int* in, int* out) {
__shared__ int buf[32];
int lane = __laneid()();
int wid = threadIdx.x / 32;
buf[lane] = in[threadIdx.x];
__syncwarp(0xFFFFFFFFu);
if (lane == 0) {
int s = 0;
for (int i = 0; i < 32; ++i) s += buf[i];
out[wid] = s;
}
__syncthreads();
}
__global__ void matmul(const float* a, const float* b, float* c,
int m, int n, int k) {
int row = blockIdx.y * blockDim.y + threadIdx.y;
int col = blockIdx.x * blockDim.x + threadIdx.x;
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);
hipLaunchKernelGGL(increment_atomic, dim3(grid), dim3(block), 0, 0, counter, 1);
hipLaunchKernelGGL(warp_scan, dim3(1), dim3(32), 0, 0, counter, a);
hipLaunchKernelGGL(matmul, dim3(grid), dim3(block), 0, 0, a, b, c, N, N, N);
}