Tensara Logo

tensara

Back

Solution: L2 Normalization

Kirill Barinov

ยท

Jul 13, 2026


Solution: L2 Normalization

Problem: L2 Normalization

Submitted on: 7/14/2026, 12:15:26 AM

  • Runtime: 1.29 ms
  • Performance: 2.80 GFLOPS
  • Tests: 5/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);

}

Comments