Submitted on: 7/14/2026, 12:15:26 AM
1.29 ms2.80 GFLOPS5/5#include <cuda_runtime.h>
#include <math.h>
constexpr float EPSILON = 1e-7;
constexpr int NUM_THREADS = 32;
int ceil_div(int n1, int n2) {
return (n1 + n2 - 1) / n2;
}
__global__ void NormalizationKernel(const float* X, float* Y, size_t B, size_t D) {
extern __shared__ float toDivide[];
auto row = blockDim.x * blockIdx.x + threadIdx.x;
auto col = blockDim.y * blockIdx.y + threadIdx.y;
if (col % NUM_THREADS == 0 && row < B) {
float sum = 0;
float elem = 0;
for (int i = 0; i < D; ++i) {
elem = X[row * D + i];
sum += elem * elem;
}
auto sqrt_sum = sqrtf(sum);
if (!isfinite(sum) || sum < EPSILON) {
toDivide[row] = EPSILON;
} else {
toDivide[row] = sqrtf(sum);
}
}
__syncthreads();
if (row < B && col < D) {
Y[row * D + col] = X[row * D + col] / toDivide[row];
}
}
// Note: X, Y are device pointers
extern "C" void solution(const float* X, float* Y, size_t B, size_t D) {
dim3 blockDim(NUM_THREADS, NUM_THREADS);
dim3 gridDim(ceil_div(B, NUM_THREADS), ceil_div(D, NUM_THREADS));
NormalizationKernel<<<gridDim, blockDim, B * sizeof(float)>>>(X, Y, B, D);
}