#include <hip/hip_runtime.h>
#include <hip/device_functions.h>
#define WARP 32
__constant__ float c_scale[4];
__device__ __forceinline__ float fast_scale(float v, int idx) {
return v * c_scale[idx & 3];
}
__device__ __noinline__ float slow_path(float v) {
if (v < 0.0f) return 0.0f;
return __sinf(v);
}
__launch_bounds__(256, 2)
__global__ void apply(const float* in, float* out, int n) {
int i = blockIdx.x * blockDim.x + threadIdx.x;
if (i < n) {
float v = fast_scale(in[i], i);
out[i] = slow_path(v);
}
}
__global__ void warp_sum(const float* in, float* out, int n) {
int lane = __laneid()();
float v = (threadIdx.x < n) ? in[threadIdx.x] : 0.0f;
__syncwarp(0xFFFFFFFFu);
__shared__ float partial[WARP];
partial[lane] = v;
__syncthreads();
if (lane == 0) {
float s = 0.0f;
for (int i = 0; i < WARP; ++i) s += partial[i];
out[blockIdx.x] = s;
}
}
int main(void) {
const int N = 1 << 18;
float* d_in = nullptr;
float* d_out = nullptr;
float* d_warp = nullptr;
hipMalloc((void**)&d_in, N * sizeof(float));
hipMalloc((void**)&d_out, N * sizeof(float));
hipMalloc((void**)&d_warp, (N / WARP) * sizeof(float));
float scale_init[4] = {1.0f, 2.0f, 3.0f, 4.0f};
hipMemcpyToSymbol(c_scale, scale_init, sizeof(scale_init));
dim3 grid(N / 256);
dim3 block(256);
hipLaunchKernelGGL(apply, dim3(grid), dim3(block), 0, 0, d_in, d_out, N);
hipLaunchKernelGGL(warp_sum, dim3(N / WARP), dim3(WARP), 0, 0, d_in, d_warp, N);
hipDeviceSynchronize();
hipFree(d_in);
hipFree(d_out);
hipFree(d_warp);
return 0;
}