Acceleration of BLAST Hydra Code on GPU
Tingxing(Tim) Dong
Lawrence Livermore National Laboratory
September 9th, 2011
Mentors: Tzanio Kolev, Robert Rieben
This work performed under the auspices of the U.S. Department of Energy by Lawrence Livermore National Laboratory under Contract DE-AC52-07NA27344
Outline
Introduction of BLAST
Motivation
Details
Optimization and Restriction
Examples and Results
Conclusion
T. Dong (LLNL)
Acceleration of BLAST Hydra Code on GPU
Seminar
2 / 34
Outline
Introduction of BLAST
Motivation
Details
Optimization and Restriction
Examples and Results
Conclusion
T. Dong (LLNL)
Acceleration of BLAST Hydra Code on GPU
Seminar
3 / 34
Introduction
BLAST
Solve equations of compressible hydrodynamics with Finite Element Method(FEM)
Based on Lagrangian frame (moving mesh)
C++ code, parallelized by MPI
BLAST’s features
Curvilinear zone geometries
Higher order ﬁeld representations
Exact discrete energy conservation by construction
Reduces to classical SGH under simplifying assumptions
Support for 2D/3D meshes
Multiple options for basis functions / quadrature order
Q1Q0, Q2Q1, and Q3Q2 cases
velocity,position
density,energy,pressure
reference
random
T. Dong (LLNL)
Acceleration of BLAST Hydra Code on GPU
Seminar
4 / 34
Euler’s Equations in a Lagrangian Frame
Euler’s Equations
Momentum Conservation:
ρ d⃗v
dt = ∇· σ
Mass Conservation:
1
ρ
dρ
dt = −∇· ⃗v
Energy Conservation:
ρ de
dt = σ : ∇⃗v
Equation of State:
p = EOS(e, ρ)
Equation of Motion:
d⃗x
dt = ⃗v
Semi-discrete ﬁnite element method in BLAST
Momentum Conservation:
dv
dt = −M−1
v
F · 1
Energy Conservation:
de
dt = M−1
e
FT · v
Equation of Motion:
dx
dt = v
T. Dong (LLNL)
Acceleration of BLAST Hydra Code on GPU
Seminar
5 / 34
Outline
Introduction of BLAST
Motivation
Details
Optimization and Restriction
Examples and Results
Conclusion
T. Dong (LLNL)
Acceleration of BLAST Hydra Code on GPU
Seminar
6 / 34
Google Proﬁler of BLAST
Consider three cases of Lagrange hydro problems with diﬀerent computational workloads:
2D q2q1: least expensive, most common
2D q3q2: more expensive, greater robustness and accuracy
3D q2q1: most expensive
2D q2q1
blast
Total samples: 29855
Focusing on: 29855
Dropped nodes with <= 149 abs(samples)
Dropped edges with <= 29 samples
__libc_start_main
0 (0.0%)
of 29220 (97.9%)
main
0 (0.0%)
of 29220 (97.9%)
29220
_start
0 (0.0%)
of 29220 (97.9%)
29220
RK2AvgIntegrator
TimeStep
0 (0.0%)
of 20523 (68.7%)
20523
HydroStatePU
DeltaTEstimate
0 (0.0%)
of 8242 (27.6%)
8242
HydroStatePU
KineticEnergy
0 (0.0%)
of 341 (1.1%)
341
HydroStatePU
Eval_dv_dt
8 (0.0%)
of 20053 (67.2%)
20053
hypre_ParCSRMatrixMatvec
1 (0.0%)
of 7527 (25.2%)
99
HydroStatePU
Eval_de_dt
12 (0.0%)
of 295 (1.0%)
294
HydroStatePU
ComputeCornerForces
2918 (9.8%)
of 16502 (55.3%)
8260
HyprePCG
Mult
0 (0.0%)
of 11161 (37.4%)
11161
hypre_ParCSRMatrixMatvecT
0 (0.0%)
of 219 (0.7%)
85
DenseMatrix
Mult
166 (0.6%)
162
FiniteElementSpace
GetElementVDofs
93 (0.3%)
of 162 (0.5%)
41
2918
MultABt
4548 (15.2%)
4480
MultAtB
2510 (8.4%)
2447
DenseMatrix
CalcEigenvalues
1589 (5.3%)
1571
DenseMatrix
FNorm
1128 (3.8%)
1120
DenseMatrix
CalcSingularvalue
815 (2.7%)
793
DenseMatrix
operator*=
801 (2.7%)
768
Mult
785 (2.6%)
of 790 (2.6%)
718
hypre_PCGSolve
6 (0.0%)
of 11161 (37.4%)
11161
7015
HYPRE_ParCSRDiagScale
2288 (7.7%)
2288
hypre_SeqVectorAxpy
820 (2.7%)
820
hypre_ParVectorInnerProd
1 (0.0%)
of 580 (1.9%)
572
hypre_SeqVectorSetConstantValues
241 (0.8%)
241
hypre_SeqVectorScale
202 (0.7%)
202
8242
hypre_CSRMatrixMatvec
7515 (25.2%)
7513
7515
4521
2482
2288
1576
1120
820
806
768
754
hypre_SeqVectorInnerProd
572 (1.9%)
572
572
DenseMatrix
MultTranspose
532 (1.8%)
506
DenseMatrix
operator=
470 (1.6%)
443
267
132
Vector
GetSubVector
187 (0.6%)
47
31
241
hypre_CSRMatrixMatvecT
219 (0.7%)
219
219
202
Vector
operator*
198 (0.7%)
198
185
164
91
IdealGasEOS
SoundSpeed
159 (0.5%)
159
T. Dong (LLNL)
Acceleration of BLAST Hydra Code on GPU
Seminar
7 / 34
Google Proﬁler of BLAST
2D q3q2
blast
Total samples: 255527
Focusing on: 255527
Dropped nodes with <= 1277 abs(samples)
Dropped edges with <= 255 samples
__libc_start_main
0 (0.0%)
of 251170 (98.3%)
main
1 (0.0%)
of 251170 (98.3%)
251170
_start
0 (0.0%)
of 251170 (98.3%)
251170
RK4Integrator
TimeStep
1 (0.0%)
of 214472 (83.9%)
214472
HydroStatePU
DeltaTEstimate
1 (0.0%)
of 34758 (13.6%)
34758
HydroStatePU
KineticEnergy
0 (0.0%)
of 1351 (0.5%)
1351
HydroStatePU
Eval
0 (0.0%)
of 213707 (83.6%)
213707
HydroStatePU
Eval_dv_dt
36 (0.0%)
of 210349 (82.3%)
210349
hypre_ParCSRMatrixMatvec
19 (0.0%)
of 80662 (31.6%)
558
HydroStatePU
Eval_de_dt
24 (0.0%)
of 2697 (1.1%)
2697
HydroStatePU
ComputeCornerForces
18660 (7.3%)
of 137915 (54.0%)
103158
HyprePCG
Mult
2 (0.0%)
of 103015 (40.3%)
103015
DenseMatrix
Mult
1492 (0.6%)
1481
18660
MultABt
63910 (25.0%)
62845
MultAtB
20013 (7.8%)
19573
DenseMatrix
CalcEigenvalues
8264 (3.2%)
8137
DenseMatrix
FNorm
5629 (2.2%)
5585
DenseMatrix
CalcSingularvalue
4874 (1.9%)
4647
Mult
4640 (1.8%)
of 4691 (1.8%)
4272
DenseMatrix
operator*=
4269 (1.7%)
4077
hypre_PCGSolve
19 (0.0%)
of 103015 (40.3%)
103013
78201
HYPRE_ParCSRDiagScale
16982 (6.6%)
16982
hypre_SeqVectorAxpy
3207 (1.3%)
3206
hypre_ParVectorInnerProd
1 (0.0%)
of 2902 (1.1%)
2884
hypre_CSRMatrixMatvec
80560 (31.5%)
80558
80560
63768
34757
19866
16982
8225
5591
4833
4518
DenseMatrix
MultTranspose
4332 (1.7%)
4232
4089
3206
hypre_SeqVectorInnerProd
2878 (1.1%)
2878
2878
1490
DenseMatrix
operator=
2543 (1.0%)
2389
1490
1124
T. Dong (LLNL)
Acceleration of BLAST Hydra Code on GPU
Seminar
8 / 34
Google Proﬁler of BLAST
3D q2q1
blast
Total samples: 195012
Focusing on: 195012
Dropped nodes with <= 975 abs(samples)
Dropped edges with <= 195 samples
__libc_start_main
0 (0.0%)
of 191550 (98.2%)
main
0 (0.0%)
of 191550 (98.2%)
191550
_start
0 (0.0%)
of 191550 (98.2%)
191550
RK2AvgIntegrator
TimeStep
1 (0.0%)
of 112516 (57.7%)
112516
HydroStatePU
DeltaTEstimate
0 (0.0%)
of 77426 (39.7%)
77426
HydroStatePU
KineticEnergy
0 (0.0%)
of 983 (0.5%)
983
HydroStatePU
ComputeCornerForces
8223 (4.2%)
of 154644 (79.3%)
8223
MultABt
61576 (31.6%)
61244
MultAtB
28754 (14.7%)
28524
DenseMatrix
CalcEigenvalues
17807 (9.1%)
of 23690 (12.1%)
23526
DenseMatrix
CalcSingularvalue
7887 (4.0%)
of 13772 (7.1%)
13558
DenseMatrix
FNorm
5468 (2.8%)
5415
Mult
4789 (2.5%)
of 4825 (2.5%)
4567
DenseMatrix
operator*=
2668 (1.4%)
2531
HydroStatePU
Eval_dv_dt
8 (0.0%)
of 111675 (57.3%)
111675
77218
HyprePCG
Mult
0 (0.0%)
of 33087 (17.0%)
33087
77426
61528
hypre_PCGSolve
16 (0.0%)
of 33087 (17.0%)
33087
hypre_ParCSRMatrixMatvec
8 (0.0%)
of 29284 (15.0%)
27786
HYPRE_ParCSRDiagScale
2897 (1.5%)
2897
hypre_SeqVectorAxpy
1121 (0.6%)
1121
hypre_CSRMatrixMatvec
29248 (15.0%)
29246
29247
28683
17779
acos
905 (0.5%)
of 7681 (3.9%)
3483
cos
4613 (2.4%)
of 4654 (2.4%)
2400
7855
3925
1960
667
__ieee754_acos
6536 (3.4%)
of 6644 (3.4%)
6408
6467
5450
4720
4545
2897
2566
DenseMatrix
operator=
2049 (1.1%)
1968
DenseMatrix
MultTranspose
1941 (1.0%)
1886
1121
CalcAdjugate
992 (0.5%)
956
907
T. Dong (LLNL)
Acceleration of BLAST Hydra Code on GPU
Seminar
9 / 34
GPU vs CPU
GFLOP/s
Bandwidth
T. Dong (LLNL)
Acceleration of BLAST Hydra Code on GPU
Seminar
10 / 34
My background before this work
Fairly familiar with Finite Diﬀerence Method (FDM)
Limited understanding about Finite Element Method (FEM)
Starts from June 1st ends at August 13th 2.5months
T. Dong (LLNL)
Acceleration of BLAST Hydra Code on GPU
Seminar
11 / 34
Outline
Introduction of BLAST
Motivation
Details
Optimization and Restriction
Examples and Results
Conclusion
T. Dong (LLNL)
Acceleration of BLAST Hydra Code on GPU
Seminar
12 / 34
Generalized corner forces on the GPU
Semi-discrete ﬁnite element method in BLAST
Momentum Conservation:
dv
dt = −M−1
v
F · 1
Energy Conservation:
de
dt = M−1
e
FT · v
Equation of Motion:
dx
dt = v
Matrix F is highly ﬂoating point operation intensive and thread independent
F is constructed by two loops:
- Loop over zones in the domain(in each processor)
-Loop over quadrature points in this zone
Compute hydro forces associated with this quadrature point
on each point we compute this value absoutely independently
varies with basis functions, dimension, etc. F can be abitrarily expensive
T. Dong (LLNL)
Acceleration of BLAST Hydra Code on GPU
Seminar
13 / 34
Generalized corner forces on the GPU
Semi-discrete ﬁnite element method in BLAST
Momentum Conservation:
dv
dt = −M−1
v
F · 1
Energy Conservation:
de
dt = M−1
e
FT · v
Equation of Motion:
dx
dt = v
CUDA kernel 1: Loop over qudrature points. Compute part of F based on v, e, x (transferred
from CPU) and allocated work space (on GPU)
CUDA kernel 2: Loop over zones. each zone does a Matrix Matrix TransposeMultiplication and
assemble the F (stay on GPU)
CUDA kernel 3 (in Momentum Equaton): Compute F · 1 and either return result to the CPU or
keep on the GPU depending on the CG solver settings.
CUDA kernel 4 (in Energy Equaton): Compute FT · v based on v (results stay on GPU)
T. Dong (LLNL)
Acceleration of BLAST Hydra Code on GPU
Seminar
14 / 34
Mass matrix solve on the GPU
Semi-discrete ﬁnite element method in BLAST
Momentum Conservation:
dv
dt = −M−1
v
F · 1
Energy Conservation:
de
dt = M−1
e
FT · v
Equation of Motion:
dx
dt = v
CUDA kernel 5 (in Momentum Equaton): Custom CG solver (Provided by Stan) for M−1
v
F · 1
based on CUBLAS/CUSPARSE, with a diagonal preconditioner (We add later).
CUDA kernel 6 (in Energy Equaton): Sparse Matrix(CSR) Multiplication to solve M−1
e
FT · v by
calling CUSPARSE
Notice:
Mv and M−1
e
are computed once and read only thereafter (stay on GPU)
M−1
v
is dense, so we did not use it directly
Me is a diagonal local dense matrix, so M−1
e
is a sparse one, can be used directly
T. Dong (LLNL)
Acceleration of BLAST Hydra Code on GPU
Seminar
15 / 34
Map to CUDA Thread Hierachy
CUDA kernel 1: loop points
Each thread < −−> one quadrature point
Each thread block < −−> one or more
zones (tunable)
In fact more ﬂexsbile: one zone can be
splited into two thread blocks
CUDA kernel 2: loop zone
Each thread block < −−> one zone
Each block(zone) do MMtMult(ABt=C)
Each thread < −−> one row of Matix C
CUDA kernel 3, 4
Each thread < −−> one zone
Each thread block is composed of 32 or 64
threads(tunable)
CUDA kernel 5, 6
Call CUBLAS/CUSPARSE/MAGMA library
rountes
T. Dong (LLNL)
Acceleration of BLAST Hydra Code on GPU
Seminar
16 / 34
Kernel2: ABt = C
Each thread block(zone) does a ABt=C
A,B are not big and generally can be ﬁtted in shared and constant memory on Fermi
A is varying to each thread block and updated in each iteration; B is read only and the same
to every block
Matrix are stored in column major
Accessing global memory is coalesced
T. Dong (LLNL)
Acceleration of BLAST Hydra Code on GPU
Seminar
17 / 34
Technical details: Memory management
CUDA code can be integrated into the previous C++ code very well
Malloc GPU memory in C++ Constructor
Free GPU memory in C++ Destructor
Add a new method CUDA ComputerCornerForce
Constructor and Destructor
// in .hpp files
// declare the varibles
double *d_vec;
// in .cu files
HydroState::CUDA_Constructor//called by HydroState()
{
// malloc variables on GPU and copy initilized
// and read only data from CPU to GPU
cudaMalloc(&d_vec);
cudaMemcpy(ToDevice);
}
HydroState::CUDA_Destructor//called by ~HydroState()
{
cudaFree(d_vec);
}
Corner Force Method
// still in .cu files,
HydroState::CUDA_CornerForce
{
// copy updated hydro states(v,e,x) to GP
cudaMemcpy(ToDevice);
// compute on GPU
kernel<<< , >>>(d_vec, ...);
// copy outputs to CPU
cudaMemcpy(ToHost);
}
T. Dong (LLNL)
Acceleration of BLAST Hydra Code on GPU
Seminar
18 / 34
Technical details: GPU Class
Porting code, not developing algorithms
Maximize use of previous C++ code to avoid develope new code
In BLAST, almost everything is class
CUDA4.0 support C++ class in device code (althought not fully)
Class on CPU
class Vector
{
private:
int size;
double *data;
public:
Vector(int a)
double *GetData()
{return data;}
void Operation()
Vector()
}
Class on GPU
class Vector_GPU
{
private:
int size;
double *data;
public:
__device__ Vector_GPU(int a)
__device__ double *GetData()
{return data;}
__device__ void Operation()
__device__ Vector_GPU()
}
T. Dong (LLNL)
Acceleration of BLAST Hydra Code on GPU
Seminar
19 / 34
An example of using class in GPU kernel
#deﬁne num threads 32
#deﬁne vector length 4
global
void kernel(double* d data)
{
Vector GPU vec(vector length);
//threadIdx.x is the thread’s id, from 0-31, each thread grab its own data via pointer cast
vec.GetData()=d data + threadIdx.x * vector length;
vec.Operation();
}
int main()
{
double *d data;
//malloc a space on device memory
cudaMalloc(&d data, sizeof(double) * num threads * vector length);
//in the kernel, each thread grab its own portion of data and execute parallelly
kernel<<<1, num threads,>>>(d data);
cudaFree(d data);
return 1;
}
T. Dong (LLNL)
Acceleration of BLAST Hydra Code on GPU
Seminar
20 / 34
Technical details: Transfer of arguments
hydro state e, v, x (in the form of class objects) needs to be transferred into CUDA kernel
can not transfer C++ class objects directly like scalar
grab data (pointer) from class objects and stored into double array or structs (see
programming guide 4.0 page)
deﬁne a struct
typedef struct
{
int height;
int width;
int chunk;
int size;
double *data;
} d_Matrix;
transfer struct as argument
void configureMatrix(d_Matrix &dm);
{
// malloc memory
// grab data from C++ class objects
and copy to dm.data
// initilize height, width, chunk, size
}
int main()
{
d_Matrix dm;
configureMatrix(dm,..);
kernel<<<, >>> (d_Matirx dm);
}
T. Dong (LLNL)
Acceleration of BLAST Hydra Code on GPU
Seminar
21 / 34
Outline
Introduction of BLAST
Motivation
Details
Optimization and Restriction
Examples and Results
Conclusion
T. Dong (LLNL)
Acceleration of BLAST Hydra Code on GPU
Seminar
22 / 34
Optimization
Use cuda proﬁler to identify the hot spot (pics of proﬁler shown later)
Constant memory store static read only coeﬃcients, like basis function, weight
parameters,etc (in kenel 1-4)
Use shared (store A) and constant (store B) memory to accelerate in cuda kenel 2 ABt=C
(memory refer O(n3)), hinted by proﬁler
Implement PCG Solver which uses CUBLAS, CUSPARSE routines instead of coding myself
Hand code Eigenvalue/vector, SVD(by Veselin) (for very small matrix 2*2(2D) or 3*3(3D),
just speciﬁc to this application) in kernel 1, as can’t call LAPACK as used in C++ BLAST
T. Dong (LLNL)
Acceleration of BLAST Hydra Code on GPU
Seminar
23 / 34
Restriction
Code is developed on Telsa C1060(my local PC), tests run on Fermi
Tesla(1.3): not support dynamic malloc and free inside kernel
memory has to be pre-allocated outside kernel, even temporary variables
Fermi(2.0): supports dynamic malloc and free inside kernel
Tesla: do not support virtual function, so some codes has to be rewritten
PCG solver and kenel 6 only work on one processor at present
Fairly tuned, but not fully optimized for Fermi
Kenel 1 (also the most complicated and expensive one) reading global memory is
uncoalesced: each thread(quardarature points) access small matrix(2*2,3*3) althought
consecutively
T. Dong (LLNL)
Acceleration of BLAST Hydra Code on GPU
Seminar
24 / 34
Outline
Introduction of BLAST
Motivation
Details
Optimization and Restriction
Examples and Results
Conclusion
T. Dong (LLNL)
Acceleration of BLAST Hydra Code on GPU
Seminar
25 / 34
Tests
q2q1 2D Triple-pt
9 * 2 velocity unkowns dofs(degree of freedoms). 4 engery dofs (zones) 16 points per zones
q3q2 2D Triple-pt
16 * 2 velocity unkowns dofs. 9 engery dofs (zones) 36 points per zones
q2q1 3D Sedov wave
27 * 3 velocity unkowns dofs. 8 engery dofs (zones) 64 points per zones
T. Dong (LLNL)
Acceleration of BLAST Hydra Code on GPU
Seminar
26 / 34
Performance
2550 lines of CUDA code + 2 months
Tesla C1060
CPU: Xeon E5520 at
2.27GHz
Tesla C2050
CPU: Xeon Westmere-EP
X5660 (on Edge)
Quadro 5000
CPU: Xeon Westmere-EP
X5660
T. Dong (LLNL)
Acceleration of BLAST Hydra Code on GPU
Seminar
27 / 34
Performance
On Edge
2D q2q1 triple-pt problem
MPI+CUDA
Result of 4 MPI
Result of 4 MPI+CUDA
T. Dong (LLNL)
Acceleration of BLAST Hydra Code on GPU
Seminar
28 / 34
CPU vs GPU results
2D q3q2 triple-pt on CPU vs GPU
total energy = kinetic energy + internal energy
CPU energy change:-4.04832e-12 with iteration step 38910 at t=2.5
GPU energy change: 2.99486e-09 with iteration step 38748 at t=2.5 (below, 3x on C2050)
T. Dong (LLNL)
Acceleration of BLAST Hydra Code on GPU
Seminar
29 / 34
CPU vs GPU results
3D q2q1 Sedov on CPU
total energy = kinetic energy + internal energy
CPU: energy change:+1.89570e-13 with iteration step 848 at t=0.3
GPU: energy change:+1.26013e-11 with iteration step 848 at t=0.3 (right, 4x on C2050)
T. Dong (LLNL)
Acceleration of BLAST Hydra Code on GPU
Seminar
30 / 34
CUDA Proﬁler
Proﬁler of 2D q3q2 triple-pt
T. Dong (LLNL)
Acceleration of BLAST Hydra Code on GPU
Seminar
31 / 34
Outline
Introduction of BLAST
Motivation
Details
Optimization and Restriction
Examples and Results
Conclusion
T. Dong (LLNL)
Acceleration of BLAST Hydra Code on GPU
Seminar
32 / 34
Conclusion
GPU is well suited for computational heavy kernel
Floating points operation should accomodate the penalty of transferring data between CPU
and GPU
Optimization is a procedure of discovering
Proﬁler does help to identify the bottleneck
Use existing library instead of coding yourself (not necessarily to be the best but the most
stable)
T. Dong (LLNL)
Acceleration of BLAST Hydra Code on GPU
Seminar
33 / 34
Thanks
T. Dong (LLNL)
Acceleration of BLAST Hydra Code on GPU
Seminar
34 / 34


Optimizing CUDA kernel with atomics

Please help,

I have been working on optimizing a CUDA kernel that utilizes atomic adds and comparing performance. I have been able to achieve approximately 2x faster but once the size of the input goes beyond 1025, e.g., 1026, I get the following output:

i (0): original (0.000000,0.000000), modified (16416.000000,8208.000000)

Atomic and Optimized kernels do not match.

Which is incorrect as the modified/optimized CUDA kernel should match the output from the original atomic add CUDA kernel - just faster execution. I am hoping that it is just something stupid simple I am doing wrong. Can anyone help?

The code is posted below, and I am using CUDA 12.2 (driver 535.54.03) with A100-SXM4-80GB device:

#include <cuda.h>
#include <chrono>
#include <iostream>
#include <stdio.h>
#include <string>
#include <sstream>
#include <vector>
#include <stdexcept>
#include <cstdlib>
#include <cmath>

// for easy gpu error checking
#define GPU_ERROR_CHECK(ans) do{gpuAssert((ans),__FILE__,__LINE__);}while(0)
inline void gpuAssert(cudaError_t code, const char *file, int line, bool abort=true)
{
   if (code != cudaSuccess)
   {
      fprintf(stderr,"GPUassert: %s %s %d\n", cudaGetErrorString(code), file, line);
      printf("\nCUDA KERNEL ERROR: CUDA Kernel reports error: %s\n",cudaGetErrorString(code));
      if (abort) exit(code);
   }
}

__forceinline__ __host__ __device__ float dot(float3 a, float3 b)
{
    return a.x * b.x + a.y * b.y + a.z * b.z;
}

__forceinline__ __host__ __device__ float length(float3 v)
{
    return sqrtf(dot(v, v));
}

__forceinline__ __device__ float2 myKernel(float val) {
    return make_float2(val, val / 2);
}
/**
 * @brief Works for values less than 1026 samples, that is up to and including 1025 samples
 */
__global__ void optimized_org_kernel(const float3 * __restrict__ pos_a, const float3 * __restrict__ pos_b,
                                     const uint32_t input_size, const uint32_t output_size,
                                     float2 * __restrict__ result, const int region,
                                     const uint32_t first_idx_x, const uint32_t last_idx_x)
{
    // Calculate thread indices
    uint32_t thridx_x = threadIdx.x + blockDim.x * blockIdx.x + first_idx_x;
    uint32_t stride_x = blockDim.x * gridDim.x;
    uint32_t thridx_y = threadIdx.y + blockDim.y * blockIdx.y;
    uint32_t stride_y = blockDim.y * gridDim.y;

    float3 distance3;
    float distance;
    uint32_t output_start, output_end;

    // Local accumulation variables to reduce the number of atomic operations
    float2 local_accum = make_float2(0.0f, 0.0f);

    for (uint32_t x = thridx_x; x < last_idx_x; x += stride_x) {
        // Calculate the distance between points
        distance3.x = pos_a[x].x - pos_b[x].x;
        distance3.y = pos_a[x].y - pos_b[x].y;
        distance3.z = pos_a[x].z - pos_b[x].z;
        distance = length(distance3);

        // Determine the output range for this thread
        output_start = __fdiv_rz(output_size, 2) + __fdiv_rz(distance, output_size) - region;
        output_end = output_start + region;

        // Clamp the values to ensure they stay within bounds
        output_start = max(0u, output_start);
        output_end = min(output_end, output_size);

        for (uint32_t y = thridx_y; y < output_size; y += stride_y) {
            // Only accumulate within the valid range
            if (y >= output_start && y < output_end) {
                float2 lval = myKernel(1.0f);
                local_accum.x += lval.x;
                local_accum.y += lval.y;
            }
        }
    }

    // Write back the accumulated values using atomic operations
    if (local_accum.x != 0.0f || local_accum.y != 0.0f) {
        atomicAdd(&result[thridx_y].x, local_accum.x);
        atomicAdd(&result[thridx_y].y, local_accum.y);
    }
}

__global__ void org_kernel(const float3 * pos_a, const float3 * pos_b,
                           const uint32_t input_size, const uint32_t output_size,
                           float2 * result, const int region,
                           const uint32_t first_idx_x, const uint32_t last_idx_x)
{
    uint32_t thridx_x = threadIdx.x + blockDim.x * blockIdx.x + first_idx_x;
    uint32_t thridx_y = threadIdx.y + blockDim.y * blockIdx.y;
    uint32_t stride_x = blockDim.x * gridDim.x;
    uint32_t stride_y = blockDim.y * gridDim.y;

    float3 distance3 = make_float3(0.0f, 0.0f, 0.0f);
    float distance = 0;

    uint32_t output_start, output_end;

    for(uint32_t x = thridx_x; x < last_idx_x; x += stride_x){
        // distance calcs
        distance3.x = pos_a[x].x - pos_b[x].x;
        distance3.y = pos_a[x].y - pos_b[x].y;
        distance3.z = pos_a[x].z - pos_b[x].z;

        distance = length(distance3);
        output_start = __fdiv_rz(output_size, 2) + __fdiv_rz(distance, output_size) - region;
        output_end = output_start + region;

        for(uint32_t y = thridx_y; y < output_size; y += stride_y){
            if((y < output_end) && (y >= output_start)){
                float2 lval = myKernel(1.0f);
                atomicAdd(&result[y].x, lval.x);
                atomicAdd(&result[y].y, lval.y);
            }
        }
    }
}

bool eval_arrays_equal(float2 * d_org, float2 * d_mod, uint32_t n) {
    if (d_org == nullptr || d_mod == nullptr) {
        throw std::invalid_argument("Arrays are NULL.");
    }
    if (n < 1) {
        throw std::invalid_argument("Invalid array length, less than 1.");
    }

    float2 * h_org;
    float2 * h_mod;
    size_t sz = n * sizeof(float2);
    h_org = (float2*)malloc(sz);
    h_mod  = (float2*)malloc(sz);

    GPU_ERROR_CHECK(cudaMemcpy(h_org, d_org, sz, cudaMemcpyDeviceToHost));
    GPU_ERROR_CHECK(cudaMemcpy(h_mod, d_mod, sz, cudaMemcpyDeviceToHost));

    GPU_ERROR_CHECK(cudaDeviceSynchronize());
    for (uint32_t i = 0; i < n; ++i) {
        if (h_org[i].x != h_mod[i].x || h_org[i].y != h_mod[i].y) {
            printf("\ti (%i): original (%f,%f), modified (%f,%f)\n", i, h_org[i].x, h_org[i].y, h_mod[i].x, h_mod[i].y);
            return false;
        }
    }

    free(h_org);
    free(h_mod);

    // Every element is equal
    return true;
}

void printDeviceProperties() {
    int32_t device;
    cudaError_t error = cudaGetDevice(&device);
    if (error != cudaSuccess) {
        std::cerr << "Failed to get current device: " << cudaGetErrorString(error) << std::endl;
        return;
    }

    cudaDeviceProp deviceProp;
    error = cudaGetDeviceProperties(&deviceProp, device);
    if (error != cudaSuccess) {
        std::cerr << "Failed to get device properties: " << cudaGetErrorString(error) << std::endl;
        return;
    }

    std::cout << "Device " << device << ": \"" << deviceProp.name << "\"" << std::endl;
    std::cout << "  CUDA Capability: " << deviceProp.major << "." << deviceProp.minor << std::endl;
    std::cout << "  Total Global Memory: " << deviceProp.totalGlobalMem / (1024 * 1024) << " MB" << std::endl;
    std::cout << "  Shared Memory per Block: " << deviceProp.sharedMemPerBlock / 1024 << " KB" << std::endl;
    std::cout << "  Registers per Block: " << deviceProp.regsPerBlock << std::endl;
    std::cout << "  Warp Size: " << deviceProp.warpSize << std::endl;
    std::cout << "  Max Threads per Block: " << deviceProp.maxThreadsPerBlock << std::endl;
    std::cout << "  Max Threads Dim: [" << deviceProp.maxThreadsDim[0] << ", "
              << deviceProp.maxThreadsDim[1] << ", " << deviceProp.maxThreadsDim[2] << "]" << std::endl;
    std::cout << "  Max Grid Size: [" << deviceProp.maxGridSize[0] << ", "
              << deviceProp.maxGridSize[1] << ", " << deviceProp.maxGridSize[2] << "]" << std::endl;
    std::cout << "  Clock Rate: " << deviceProp.clockRate / 1000 << " MHz" << std::endl;
    std::cout << "  Total Constant Memory: " << deviceProp.totalConstMem / 1024 << " KB" << std::endl;
    std::cout << "  Multiprocessor Count: " << deviceProp.multiProcessorCount << std::endl;
    std::cout << "  Compute Mode: " << deviceProp.computeMode << std::endl;
}

int main(int argc, char * argv[]) {
    if (argc != 2) {
        fprintf(stderr, "\nPass the number of array elements via command line as follows:\n");
        fprintf(stderr, "./xTest <num_elems>\n\n");
        return EXIT_FAILURE;
    }

    // Dimensions
    const uint32_t BLOCK_WIDTH = 512;
    dim3 nblks(BLOCK_WIDTH,1,1);
    dim3 nthreads(1,BLOCK_WIDTH,1);

    // Retrieve command-line argument
    uint32_t n_values = static_cast<uint32_t>(std::stoi(argv[1]));

    uint32_t region = 3;

    uint32_t n_float3s = n_values;
    uint32_t float3_sz = n_float3s * sizeof(float3);
    uint32_t output_sz = n_values * sizeof(float2);

    // Allocate host & device side
    float2 *d_out_org;
    float2 *d_out_mod;
    GPU_ERROR_CHECK(cudaMalloc(&d_out_org, output_sz));
    GPU_ERROR_CHECK(cudaMalloc(&d_out_mod, output_sz));
    GPU_ERROR_CHECK(cudaMemset(d_out_org, 0, output_sz));
    GPU_ERROR_CHECK(cudaMemset(d_out_mod, 0, output_sz));

    // Float3s
    float3 *pos_a, *pos_b;
    float3 *d_pos_a, *d_pos_b;
    pos_a = (float3*)malloc(float3_sz);
    pos_b = (float3*)malloc(float3_sz);
    for(size_t p = 0; p < n_float3s; ++p){
        pos_a[p] = make_float3(1,1,1);
        pos_b[p] = make_float3(0.1,0.1,0.1);
    }
    GPU_ERROR_CHECK(cudaMalloc(&d_pos_a, float3_sz));
    GPU_ERROR_CHECK(cudaMalloc(&d_pos_b, float3_sz));
    GPU_ERROR_CHECK(cudaMemcpy(d_pos_a, pos_a, float3_sz, cudaMemcpyHostToDevice));
    GPU_ERROR_CHECK(cudaMemcpy(d_pos_b, pos_b, float3_sz, cudaMemcpyHostToDevice));
    GPU_ERROR_CHECK(cudaDeviceSynchronize());

    float total_time_org = 0.0f;
    float total_time_mod = 0.0f;

    uint32_t first_idx_x = 0;
    uint32_t last_idx_x  = n_values;

    const uint32_t n_passes = 16;

    for (uint32_t pass = 0; pass < n_passes; ++pass) {
        auto start = std::chrono::high_resolution_clock::now();

        // Original atomic add kernel
        org_kernel<<<nblks,nthreads,0,0>>>(d_pos_a, d_pos_b, n_values, n_values, d_out_org, region, first_idx_x, last_idx_x);
        GPU_ERROR_CHECK(cudaDeviceSynchronize());

        auto stop = std::chrono::high_resolution_clock::now();
        total_time_org += static_cast<float>(std::chrono::duration_cast<std::chrono::nanoseconds>(stop - start).count());

        start = std::chrono::high_resolution_clock::now();

        // Optimized atomic add kernel
        optimized_org_kernel<<<nblks,nthreads,0,0>>>(d_pos_a, d_pos_b, n_values, n_values, d_out_mod, region, first_idx_x, last_idx_x);
        GPU_ERROR_CHECK(cudaDeviceSynchronize());

        stop = std::chrono::high_resolution_clock::now();
        total_time_mod += static_cast<float>(std::chrono::duration_cast<std::chrono::nanoseconds>(stop - start).count());
    }

    // Check for fidelity
    if (eval_arrays_equal(d_out_org, d_out_mod, n_values)) {
        printf("\nFidelity achieved.\n");
        printf("\tTotal number of passes: %d\n", n_passes);
        float org_time = (total_time_org / n_passes);
        float mod_time = (total_time_mod / n_passes);
        printf("\t[ORIGINAL] Time: %8.9f (us.)\n", org_time);
        printf("\t[MODIFIED] Time: %8.9f (us.)\n", mod_time);
        printf("\tSpeedup Factor: %8.9f\n", (org_time / mod_time));
    } else {
        printf("\nAtomic and Optimized kernels do not match.\n");
        return EXIT_FAILURE;
    }

    GPU_ERROR_CHECK(cudaPeekAtLastError());
    GPU_ERROR_CHECK(cudaDeviceSynchronize());
    GPU_ERROR_CHECK(cudaDeviceReset());

    return EXIT_SUCCESS;
}

Thanks to anyone who can point out any error or point a direction that might fix this issue.

The original kernel adds a number to multiple columns. The modified kernel sums up multiple numbers and adds it to a single column. How can this ever be equivalent?

Did you investigate all suggestions already given in the original thread here ?

Hi @striker159

Thank you for the reply.

Yes, I have been looking into the suggestions. I don’t think shared memory is useful in my case as I was able to get it working but during execution the performance was not actually much better, maybe 3%. So, I wanted to look in another direction hopefully without shared memory - just lowering the number of necessary atomicAdds if possible.

I think I just solved the problem, by using warp level primitives and only calling atomicAdd once per block.

The performance is somewhere between 1.5x to 2.8x faster with the mod_kernel as compared to the org_kernel.  I have posted the code below for anyone who may like to use it for their own project(s).

Thanks to all for assistance pointing me in some good direction(s) for optimizing this CUDA kernel.

#include <cuda.h>
#include <cuda_runtime.h>
#include <iostream>
#include <iomanip>
#include <stdexcept>
#include <chrono>
#include <cstdlib>

// for easy gpu error checking
#define GPU_ERROR_CHECK(ans) do{gpuAssert((ans),__FILE__,__LINE__);}while(0)
inline void gpuAssert(cudaError_t code, const char *file, int line, bool abort=true)
{
   if (code != cudaSuccess)
   {
      fprintf(stderr,"GPUassert: %s %s %d\n", cudaGetErrorString(code), file, line);
      printf("\nCUDA KERNEL ERROR: CUDA Kernel reports error: %s\n",cudaGetErrorString(code));
      if (abort) exit(code);
   }
}
/**
 * @brief CUDA DEVICE kernels executes scalar dot product.
 *
 * @param a The first float3.
 * @param b The second float3.
 * @return floating-point value that is the scalar dot product.
 */
__forceinline__ __host__ __device__ float dot(float3 a, float3 b) {
	return a.x * b.x + a.y * b.y + a.z * b.z;
}

/**
 * @brief CUDA DEVICE kernel executes Euclidean length of input float3.
 *
 * @param v The float3 whose x, y, z components length is being computed from.
 * @return floating-point value that is the Euclidean length of input float3.
 */
__forceinline__ __host__ __device__ float length(float3 v) {
	return sqrtf(dot(v, v));
}

/**
 * @brief CUDA DEVICE kernel is a toy operation for demonstration purposes, whereby the
 * input value is modified and returned as a float2 datatype.
 *
 * @param val The input value being modified.
 * @return float2 version of input float with modifications applied.
 */
__forceinline__ __host__ __device__ float2 myKernel(float val) {
	return make_float2(val, val / 2);
}

/**
 * @brief CUDA DEVICE kernel that executes a warp-level summation of input float2 value.
 * @details Allows data to be summed without the use of extra memory space, that is, shared
 * directly across threads in a single warp (32-threads).
 *
 * @param val The float2 value being summed across warp.
 * @return float2 value summed across threads in a single warp.
 */
__inline__ __device__ float2 warpReduceSum(float2 val) {
    for (int offset = warpSize / 2; offset > 0; offset /= 2) {
        val.x += __shfl_down_sync(0xffffffff, val.x, offset);
        val.y += __shfl_down_sync(0xffffffff, val.y, offset);
    }
    return val;
}

/**
 * @brief CUDA DEVICE kernel that calls CUDA intrinsic @ref atomicAdd only on the
 * first thread and warp in the block.
 *
 * @param address The resulting global address where the value is added and stored.
 * @param val The value being added to the global address.
 */
__inline__ __device__ void atomicAddWarp(float2 *address, float2 val) {
    if (threadIdx.x % warpSize == 0) {
        atomicAdd(&address->x, val.x);
        atomicAdd(&address->y, val.y);
    }
}
__global__ void org_kernel(const float3 * pos_a, const float3 * pos_b,
                           const uint32_t input_size, const uint32_t output_size,
                           float2 * result, const int32_t region,
                           const uint32_t first_idx_x, const uint32_t last_idx_x) {
    // Compute indices
    uint32_t thridx_x = threadIdx.x + blockDim.x * blockIdx.x + first_idx_x;
    uint32_t thridx_y = threadIdx.y + blockDim.y * blockIdx.y;
    uint32_t stride_x = blockDim.x * gridDim.x;
    uint32_t stride_y = blockDim.y * gridDim.y;

    float3 distance3 = make_float3(0.0f, 0.0f, 0.0f);
    float distance = 0;

    uint32_t output_start, output_end;

    for(uint32_t x = thridx_x; x < last_idx_x; x += stride_x){
        // distance calcs
        distance3.x = pos_a[x].x - pos_b[x].x;
        distance3.y = pos_a[x].y - pos_b[x].y;
        distance3.z = pos_a[x].z - pos_b[x].z;

        distance = length(distance3);
        output_start = __fdiv_rz(output_size, 2) + __fdiv_rz(distance, output_size) - region;
        output_end = output_start + region;

        for(uint32_t y = thridx_y; y < output_size; y += stride_y){
            if((y < output_end) && (y >= output_start)) {
                float2 lval = myKernel(1.0f);
                atomicAdd(&result[y].x, lval.x);
                atomicAdd(&result[y].y, lval.y);
            }
        }
    }
}
__global__ void mod_kernel(const float3 * __restrict__ pos_a, const float3 * __restrict__ pos_b,
                           const uint32_t input_size, const uint32_t output_size,
                           float2 * __restrict__ result, const int32_t region,
                           const uint32_t first_idx_x, const uint32_t last_idx_x) {
    // Compute indices
    uint32_t thridx_x = threadIdx.x + blockDim.x * blockIdx.x + first_idx_x;
    uint32_t thridx_y = threadIdx.y + blockDim.y * blockIdx.y;
    uint32_t stride_x = blockDim.x * gridDim.x;
    uint32_t stride_y = blockDim.y * gridDim.y;

    if (thridx_x >= last_idx_x) return;

    float3 distance3;
    float distance;
    uint32_t output_start, output_end;

    for (uint32_t x = thridx_x; x < last_idx_x; x += stride_x) {
        // Pre-calculate distance components
        distance3.x = pos_a[x].x - pos_b[x].x;
        distance3.y = pos_a[x].y - pos_b[x].y;
        distance3.z = pos_a[x].z - pos_b[x].z;

        // Compute the distance and the output indices range
        distance = sqrtf(distance3.x * distance3.x +
                         distance3.y * distance3.y +
                         distance3.z * distance3.z);

        output_start = __fdividef(output_size, 2) + __fdividef(distance, output_size) - region;
        output_end = output_start + region;

        // Restrict output range to valid indices
        output_start = max(output_start, 0U);
        output_end = min(output_end, output_size);

        for (uint32_t y = thridx_y; y < output_size; y += stride_y) {
            if (y >= output_start && y < output_end) {
                float2 lval = myKernel(1.0f);
                // Execute warp-level primitives then only call atomic add once per block
                float2 warp_sum = warpReduceSum(lval);
                atomicAddWarp(&result[y], warp_sum);
            }
        }
    }
}
bool eval_arrays(const float2 * d_arr1, const float2 * d_arr2, const uint32_t n) {
    if (d_arr1 == nullptr || d_arr2 == nullptr) {
	    throw std::invalid_argument("Null array(s).");
    }
    if (n < 1) {
	    throw std::invalid_argument("Invalid array length.");
    }

    float2 * h_arr1 = nullptr;
    float2 * h_arr2 = nullptr;
    h_arr1 = new float2[n];
    h_arr2 = new float2[n];
    GPU_ERROR_CHECK(cudaMemcpy(h_arr1, d_arr1, n * sizeof(float2), cudaMemcpyDeviceToHost));
    GPU_ERROR_CHECK(cudaMemcpy(h_arr2, d_arr2, n * sizeof(float2), cudaMemcpyDeviceToHost));

    for (uint32_t i = 0; i < n; ++i) {
        if (h_arr1[i].x != h_arr2[i].x || h_arr1[i].y != h_arr2[i].y) {
            std::cout << "Index: " << i << " Array 1 (" << h_arr1[i].x << "," << h_arr1[i].y << "), Array 2 ("
                      << h_arr2[i].x << "," << h_arr2[i].y << ")\n";
            delete [] h_arr1;
            delete [] h_arr2;
            return false;
	    }
    }

    delete [] h_arr1;
    delete [] h_arr2;

    // Every element in both arrays was the same
    return true;
}

int main(int argc, char * argv[]) {
    if (argc != 2) {
        std::cerr << "\nPass the number of array elements via command line as follows:\n";
        std::cerr << "./xOptimize <num_elems>\n\n";
        return EXIT_FAILURE;
    }

    // Get number of array elements from command line
    int n_values = std::stoi(argv[1]);
    if (n_values < 1) {
        std::cerr << "Invalid number of array elements: " << n_values << std::endl;
        return EXIT_FAILURE;
    }

    // Defined sizes
    const uint32_t BLOCK_WIDTH = 512;
    size_t float3_sz = n_values * sizeof(float3);
    size_t output_sz = n_values * sizeof(float2);

    // HOST-side positions
    float3 *pos_a = nullptr;
    float3 *pos_b = nullptr;
    pos_a = new float3[n_values];
    pos_b = new float3[n_values];
    for (int i = 0; i < n_values; ++i) {
        pos_a[i] = make_float3(i, i + 1, i + 2);
        pos_b[i] = make_float3(i + 0.5f, i + 1.5f, i + 2.5f);
    }
    // DEVICE-side positions
    float3 *d_pos_a = nullptr;
    float3 *d_pos_b = nullptr;
    GPU_ERROR_CHECK(cudaMalloc(&d_pos_a, float3_sz));
    GPU_ERROR_CHECK(cudaMalloc(&d_pos_b, float3_sz));
    GPU_ERROR_CHECK(cudaMemcpy(d_pos_a, pos_a, float3_sz, cudaMemcpyHostToDevice));
    GPU_ERROR_CHECK(cudaMemcpy(d_pos_b, pos_b, float3_sz, cudaMemcpyHostToDevice));

    // DEVICE-side outputs
    float2 *d_out_org = nullptr;
    float2 *d_out_mod = nullptr;
    GPU_ERROR_CHECK(cudaMalloc(&d_out_org, output_sz));
    GPU_ERROR_CHECK(cudaMalloc(&d_out_mod, output_sz));
    GPU_ERROR_CHECK(cudaMemset(d_out_org, 0, output_sz));
    GPU_ERROR_CHECK(cudaMemset(d_out_mod, 0, output_sz));

    float total_time_org = 0.0f;
    float total_time_mod = 0.0f;
    uint32_t first_idx_x = 0;
    uint32_t last_idx_x = n_values;

    int region = 3;
    dim3 nthreads(BLOCK_WIDTH, 1, 1);
    dim3 nblocks(1, BLOCK_WIDTH, 1);

    const uint32_t n_passes = 16;
    for (uint32_t pass = 0; pass < n_passes; ++pass) {

        auto start = std::chrono::high_resolution_clock::now();

        // Original atomic kernel
        org_kernel<<<nblocks, nthreads>>>(d_pos_a, d_pos_b, n_values, n_values, d_out_org, region, first_idx_x, last_idx_x);
        GPU_ERROR_CHECK(cudaDeviceSynchronize());

        auto stop = std::chrono::high_resolution_clock::now();
        total_time_org += static_cast<float>(std::chrono::duration_cast<std::chrono::nanoseconds>(stop - start).count());

        start = std::chrono::high_resolution_clock::now();

        // Modified atomic kernel
        mod_kernel<<<nblocks, nthreads>>>(d_pos_a, d_pos_b, n_values, n_values, d_out_mod, region, first_idx_x, last_idx_x);
        GPU_ERROR_CHECK(cudaDeviceSynchronize());

        stop = std::chrono::high_resolution_clock::now();
        total_time_mod += static_cast<float>(std::chrono::duration_cast<std::chrono::nanoseconds>(stop - start).count());
    }

    std::cout << std::fixed << std::setprecision(4);
    total_time_org /= n_passes;
    total_time_mod /= n_passes;

    std::cout << "\nTotal number of passes: " << n_passes << std::endl;
    std::cout << "Original CUDA Kernel Time: " << total_time_org << " (us.)\n";
    std::cout << "Modified CUDA Kernel Time: " << total_time_mod << " (us.)\n";
    std::cout << "Speedup factor: " << (total_time_org / total_time_mod) << std::endl;

    // Check fidelity
    if (eval_arrays(d_out_org, d_out_mod, n_values)) {
        std::cout << "\nFidelity achieved.\n\n";
    }else {
        std::cout << "\nFidelity not achieved.\n\n";
    }

    return EXIT_SUCCESS;
}

When trying to improve a kernel, your main focus should be correctness, not speed. If you don’t want correct results, you could remove the kernel which will give you the greatest speedup.

As already pointed out by others in the first thread, your modified kernels are not equivalent to the original kernel.

This is your original code

for(uint32_t y = thridx_y; y < output_size; y += stride_y){
            if((y < output_end) && (y >= output_start)) {
                float2 lval = myKernel(1.0f);
                atomicAdd(&result[y].x, lval.x);
            }
        }

Let’s plug in some numbers for simplicity.

for(uint32_t y = 0; y < 5; y += 1){
            if((y < 5) && (y >= 0)) {
                float2 lval = myKernel(1.0f);
                atomicAdd(&result[y].x, lval.x);
            }
        }

This means if result is initialized with [0,0,0,0,0], it will be [1,1,1,1,1] after the loop.

But if you accumulate the values for multiple y in any way, and then write to output, you will get different results.

local_accum = 0
        for(uint32_t y = 0; y < 5; y += 1){
            if((y < 5) && (y >= 0)) {
                float2 lval = myKernel(1.0f);
                local_accum.x += lval.x;
            }
        }
       atomicAdd(&result[0].x, local_accum.x);

This will output [5,0,0,0,0], not [1,1,1,1,1]

The simplest solution the reduce the number of atomics in the original kernel is to only perform atomicAdd(result.x, 1) .  Then afterwards use a second kernel which computes result.y = result.x / 2

Hi @striker159,

Thank you for the reply.

Yes, my first goal has always been to ensure the correctness of the modified code with the original code then to speed it up, if possible. This is why it was such a problem, speeding it up is not too difficult but speeding it up and maintaining consistency with the original results, that’s the issue.

I understand that the following code would give incorrect results when comparing with the original:

local_accum = 0
        for(uint32_t y = 0; y < 5; y += 1){
            if((y < 5) && (y >= 0)) {
                float2 lval = myKernel(1.0f);
                local_accum.x += lval.x;
            }
        }
       atomicAdd(&result[0].x, local_accum.x);

That is why I have since abandoned that idea.

However, is there any reason that the following snippet of code from above wouldn’t work?

...
__inline__ __device__ float2 warpReduceSum(float2 val) {
    for (int offset = warpSize / 2; offset > 0; offset /= 2) {
        val.x += __shfl_down_sync(0xffffffff, val.x, offset);
        val.y += __shfl_down_sync(0xffffffff, val.y, offset);
    }
    return val;
}
__inline__ __device__ void atomicAddWarp(float2 *address, float2 val) {
    if (threadIdx.x % warpSize == 0) {
        atomicAdd(&address->x, val.x);
        atomicAdd(&address->y, val.y);
    }
}
__global__ void mod_kernel(const float3 * __restrict__ pos_a, const float3 * __restrict__ pos_b,
                           const uint32_t input_size, const uint32_t output_size,
                           float2 * __restrict__ result, const int32_t region,
                           const uint32_t first_idx_x, const uint32_t last_idx_x) {
    // Compute indices
    uint32_t thridx_x = threadIdx.x + blockDim.x * blockIdx.x + first_idx_x;
    uint32_t thridx_y = threadIdx.y + blockDim.y * blockIdx.y;
    uint32_t stride_x = blockDim.x * gridDim.x;
    uint32_t stride_y = blockDim.y * gridDim.y;

    if (thridx_x >= last_idx_x) return;

    float3 distance3;
    float distance;
    uint32_t output_start, output_end;

    for (uint32_t x = thridx_x; x < last_idx_x; x += stride_x) {
        // Pre-calculate distance components
        distance3.x = pos_a[x].x - pos_b[x].x;
        distance3.y = pos_a[x].y - pos_b[x].y;
        distance3.z = pos_a[x].z - pos_b[x].z;

        // Compute the distance and the output indices range
        distance = sqrtf(distance3.x * distance3.x +
                         distance3.y * distance3.y +
                         distance3.z * distance3.z);

        output_start = __fdividef(output_size, 2) + __fdividef(distance, output_size) - region;
        output_end = output_start + region;

        // Restrict output range to valid indices
        output_start = max(output_start, 0U);
        output_end = min(output_end, output_size);

        for (uint32_t y = thridx_y; y < output_size; y += stride_y) {
            if (y >= output_start && y < output_end) {
                float2 lval = myKernel(1.0f);
                // Execute warp-level primitives then only call atomic add once per block
                float2 warp_sum = warpReduceSum(lval);
                atomicAddWarp(&result[y], warp_sum);
            }
        }
    }
}
...

Thanks again for the help.

First of all, there is no guarantee that all threads in the warp reach the reduction code.

That aside, the reduction approach has the same problem. Values for multiple y are combined.
For example, in thread0 y=0, in thread1 y=1. Then the warpsum will be 2, and both result[0] and result[1] will be set to 2. But the correct result would be result[0] = 1 and result[1] = 1.

Okay, but wouldn’t calling one kernel to operate on the x component and then another kernel to operate on the y component end up costing more time than just calling the original kernel once? Unless the input is very large to maybe mitigate.

I guess I am trying to understand if maybe the original kernel is as fast as it can be given the parameters of the problem itself.

After the first x calculations, store the intermediate results in shared memory, use syncthreads() to synchronize block-wise, then use one or some of the warps for y processing within the block.

Hi @Curefab,

Thank you for the reply. Using shared memory is a good idea, and I have tried various versions of this methodology. While the results are valid in that they match the original values, the performance is really no better or if it is it is negligible. This is why I am beginning to think that maybe the original kernel is as fast as it can get.

This is why I am beginning to think that maybe the original kernel is as fast as it can get.

That is quite likely.

You can make some further tests with changed parameters or theoretical calculations of the maximum speed (e.g. how many bytes you read/write vs. the bandwidth of device memory) combined with Compute Nsight to support this assumption.

Related topics

Powered by Discourse, best viewed with JavaScript enabled


Optimizing CUDA.jl performance for small array operations

Hi,

I’m trying to get better GPU performance with CUDA.jl for small array operations.

So, I’ve started to port SymbolicRegression.jl to the GPU using CUDA.jl. It seems I’ve gotten the main evaluation part of the code to use the corresponding CUDA operations (which was REALLY straightforward by the way, great job!) , but it’s slower than I would like.

Part of the problem is that during symbolic regression, you typically work on small amounts of data; maybe a matrix of size 5x1000. Without some clever fusion of tree evaluations, this means one needs to worry about the time it takes to launch kernels which makes things tricky.

As a MWE, consider the following code:

using CUDA, BenchmarkTools, Statistics
for N in [1000, 10000]
    c1 = CUDA.ones(Float32, N)
    c2 = ones(Float32, N)
    res1 = @benchmark CUDA.@sync cos.($c1);
    res2 = @benchmark cos.($c2);
    println("Size $N: CUDA=$(median(res1.times)); CPU=$(median(res2.times))")
end

On my v100, this gives me (in microseconds):

Size 1000: CUDA=26021.0; CPU=9086.0
Size 10000: CUDA=24419.5; CPU=87287.0

The GPU scales so well the array size is negligible. But the baseline time it takes to launch a kernel means I can’t exploit the full power of the GPU for evaluating trees on these small arrays.

Is there something I can do to improve the kernel launch speed here?

I also tried:

res1 = @benchmark CUDA.@sync blocking=false cos.($c1);

which was mentioned in the CUDA.jl docs to be better for profiling short executions. This lowers the the evaluation time to ~11000, but unfortunately this is still not enough.

Thanks for any advice!
Cheers,
Miles

This is the order I’d probably try things in:

I’m not at all familiar with the algorithms, but could the fusion be easier than it appears? For one, you can always add “batch” dimensions easily like:

c1a = CUDA.ones(Float32, N)
c1b = CUDA.ones(Float32, N)
c1 = cat(c1a, c1b, dims=2)
res1a, res2a = eachslice(cos.(c1), dims=2)

Also, you can “lazily” create broadcasted objects then evaluate them with a single “fused” kernel call like:

using Base: broadcasted, materialize
bc1 = broadcasted(cos, c1)
bc2 = broadcasted(sin, c2)
bc = broadcasted(+, bc1, bc2)
res = materialize(bc) # one kernel call, w/o temp arrays

You could use multiple threads. I think you may need to be on CUDA#master (might need to do some reading in Automatic task-based concurrency using local streams by maleadt · Pull Request #662 · JuliaGPU/CUDA.jl · GitHub and links therein too) but the latest versions each thread uses its own stream, which can then execute concurrently. In general it’d be good to do some profiling to make sure the GPU is doing what you think it is.

You could combine more of your code into single kernels by writing custom kernels (still pure Julia). KernelAbstractions.jl offers a pretty nice interface to it, although even CUDA.@cuda isn’t too bad.

Thanks @marius311. Replies to each point numbered below.

I should mention first that the C++ CUDA kernel launch cost is only ~10 microseconds, so I definitely think the seeming ~25,000 microseconds launch cost here can be cut down. i.e., this seem to be a software problem rather than hardware or algorithm. But these other tricks to combine kernels might be enough to account for it.

Best,
Miles

I should mention first that the C++ CUDA kernel launch cost is only ~10 microseconds, so I definitely think the seeming ~25,000 microseconds launch cost here

Quick note but I think those benchmark times are in nanoseconds, so its not 25,000 microseconds, its 25 microseconds, on the order of the kernel launch.

Anyway, thanks for the info, that helps. Are the inputs to the ~1000 equations the same at least? As in, does this reduce to applying an array of functions to some input? If so, search these foums for “array of functions” and you’ll get several links discussing this, some even mentioning GPU (this is probalby a good one to follow links from, also, this may be the identical question). This may be naive since I haven’t really read those threads, but here’s my solution messing around with it briefly:

using CUDA

fs = (sin, x->x^2, x->tan(x^2))

apply_funcs(fs, x) = ntuple(i->fs[i](x), length(fs))

apply_funcs.(Ref(fs), cu(rand(1000)))

I haven’t done a super proper profile, but I’m pretty sure that final call results in a single kernel launch. Certainly benchmarking it on a V100, its faster than 3 consecutive ones:

julia> x = cu(rand(Float32, 1000));

julia> @btime CUDA.@sync apply_funcs.($(Ref(fs)), $x);
  22.183 μs (47 allocations: 1.55 KiB)

julia> @btime CUDA.@sync (fs[1].($x); fs[2].($x); fs[3].($x));
  39.731 μs (78 allocations: 1.84 KiB)

Something else which is cool is that if you look at e.g. @code_llvm apply_funcs(fs, 2.) you’ll see that x^2 is only calculated once, I believe in this case the compiler is smart enough to eliminate the common sub-expression. I’d guess the CUDA compiler would do the same thing, although I’m not familiar enough with how to check exactly.

A limitation of this is that it stops working once your tuple of functions is longer than length 10 because of this. But maybe batching 10 functions together like this is enough to saturate the GPU, and if not, you can always write your own ntuple-like function and grow that if-statement a bit larger.

Thanks! This is very helpful.

Okay, good to know re: units. I’m not sure why I thought it was in microseconds… but yeah 25 us seems much more reasonable.

Are the inputs to the ~1000 equations the same at least?

Only at the leafs of each equation. Two equations might both use the variable x1, for instance. But equations also have constants, and these will always be slightly different due to mutations. In the branch nodes, the inputs end up being completely different unless that particular subtree is identical, which is rare. I tried using memoization for small subtrees at one point, but this didn’t help… there’s just so many different subtrees possible in a typical search.

Re: apply_funcs, good idea. Will try this. To be honest, thinking more about this, I wonder how doable it would be to completely batch evaluation of many equations…

Actually, maybe this make things easier for your point 1!

Right now my recursive evaluation essentially looks like this:

function evalTree(X, tree)
    if tree.degree == 0
        if tree.constant
            return fill(tree.value, size(X, 2))
        end
        return X[tree.feature, :]
    elseif tree.degree == 1
        x = evalTree(X, tree.left)
        op = unary_operators[tree.op]
        return op.(x)
    else
        x = evalTree(X, tree.left)
        y = evalTree(X, tree.right)
        op = binary_operators[tree.op]
        return op.(x, y)
    end
end

(though I do some operator fusing for small subtrees, and also run things behind a function barrier, per your helpful advice in the other thread  ).

I guess I am wondering if it would be efficient to: (1) pass a list of trees to evaluate, and (2) walk the binary tree for ALL trees up to the maximum depth, using an identity operator for i when tree[i]==nothing  at the current depth, and converting all unary operators to binary via (x, y) -> op(x)

What do you think?

Although, maybe the fact that the tuple of operators changes every single time means that this would be re-compiling the kernel each launch… I guess one could instead pass an array of indices for the functions to call, and manage all of that inside the kernel assuming fixed operators?

Related topics

Powered by Discourse, best viewed with JavaScript enabled


How to speed up that simple CUDA kernel?

Hi! I have a kernel:

__global__ void filter_small(unsigned arr_size, float *arr, ResultAndPos *results, float threshold, int *n_results) {
    const int itemsPerThread = 32;
    int begin = blockIdx.x * blockDim.x * itemsPerThread + threadIdx.x * itemsPerThread;
    int end = begin + itemsPerThread;
    if (end > arr_size)
      end = arr_size;

    for (int index = beg; index < end; index++) {
      if (arr[index] < threshold) {
        int oldIdx = atomicAdd(n_results, 1);
        results[oldIdx] = ResultAndPos{arr[index], index};
      }
    }
  }

So basically it's very simple filter, that leaves elements smaller that the threshold. But the array itself is very large, ~ 5-15 millions of floats.

And I launch it as follows:

int blockSize = 128;
    int itemsPerThread = 32;
    int itemsPerBlock = itemsPerThread * blockSize;
    int numBlocks = (N + itemsPerBlock - 1) / itemsPerBlock;

    filter_small<<<numBlocks, blockSize>>>(N, arr_ptr, results_ptr, threshold, d_n_results);

For ten million, it executes around ~12 ms which is quite slow. How can I speed it up?

Thank you!

OS: linux ubuntu 18.04, cuda 11.1, nvidia 3060

Create your account and connect with a world of communities.

Anyone can view, post, and comment to this community

Top Posts


This makes the project proportionally harder in my opinion because you need to be that much more efficient with moving data through the memory hierarchy. With tensor cores, to get anywhere close to cuBLAS, you need to start with something like the most efficient kernel in simon's article, and then do stuff like shared memory swizzling, async global memory copies, double buffering, and writing a really efficient kernel epilogue to accumulate the C matrix into the product.I came across this article a while ago and it inspired me to take a stab at this^, and as of now I have gotten to ~80% of the cuBLAS tensor core performance where the kernel is mostly compute bound, and I am close to giving up on the last ~20%, because I think I may need to write the inner loop in SASS to make sure the instruction mix between shared memory loads, mma instructions, and synchronizations is perfectly balanced so that none of the hardware pipelines get overloaded (see link below), and I have enough compassion for myself to not spend my free time doing stuff like that :). There are also certain things implemented in CUTLASS that seem important (look up serpentine traversal) but NVIDIA engineers wont talk about the hardware details required to understand why this helps.Article on this is forthcominghttps://github.com/NervanaSystems/maxas/wiki/SGEMM

I came across this article a while ago and it inspired me to take a stab at this^, and as of now I have gotten to ~80% of the cuBLAS tensor core performance where the kernel is mostly compute bound, and I am close to giving up on the last ~20%, because I think I may need to write the inner loop in SASS to make sure the instruction mix between shared memory loads, mma instructions, and synchronizations is perfectly balanced so that none of the hardware pipelines get overloaded (see link below), and I have enough compassion for myself to not spend my free time doing stuff like that :). There are also certain things implemented in CUTLASS that seem important (look up serpentine traversal) but NVIDIA engineers wont talk about the hardware details required to understand why this helps.Article on this is forthcominghttps://github.com/NervanaSystems/maxas/wiki/SGEMM

Article on this is forthcominghttps://github.com/NervanaSystems/maxas/wiki/SGEMM

https://github.com/NervanaSystems/maxas/wiki/SGEMM

I’d be so happy if SASS were documented and ptxas were open source, sometimes I spend entire days going through whitepapers and various sources of online documentation to get more hardware details…

My guess is that people nowadays are gradually moving away from raw CUDA programming and moving towards things like Triton etc, and you won't be focusing on pure GEMM since you tend to do some fusion.The Triton tutorial claims their performance is on par with cuBLAS.https://triton-lang.org/main/getting-started/tutorials/03-ma...

The Triton tutorial claims their performance is on par with cuBLAS.https://triton-lang.org/main/getting-started/tutorials/03-ma...

https://triton-lang.org/main/getting-started/tutorials/03-ma...

Your guess is wrong. Besides the fact that there's much more to life than matmul (for which triton is just ok), the other obvious fact is that triton has exactly 1 frontend (python) and there's much more to life than that frontend.I find that basically in every thread about low-level work there's someone making some weird comment about how triton or mojo or XYZ supplants CUDA or assembly or whatever. I can't understand how this comes about because absolutely no one working in these areas thinks XYZ is going to supplant anything. So it's invariably outsiders making these claims and I cannot fathom why any outsider would be motivated to make claims from the outside.

I find that basically in every thread about low-level work there's someone making some weird comment about how triton or mojo or XYZ supplants CUDA or assembly or whatever. I can't understand how this comes about because absolutely no one working in these areas thinks XYZ is going to supplant anything. So it's invariably outsiders making these claims and I cannot fathom why any outsider would be motivated to make claims from the outside.

As an outsider CUDA is so intimidating so the promise of Triton etc is very appearing and I wanted to get sold.

i have PRs in Triton - i'm well familiar with the fact that triton is an MLIR project.> C++ straight using MLIRthat's like saying llvm ir is usable through C++ ... or hell that's like saying NVPTX is usable through C++. it's not just not a frontend it's the exact opposite: it's emitting IR using IR builders.

> C++ straight using MLIRthat's like saying llvm ir is usable through C++ ... or hell that's like saying NVPTX is usable through C++. it's not just not a frontend it's the exact opposite: it's emitting IR using IR builders.

that's like saying llvm ir is usable through C++ ... or hell that's like saying NVPTX is usable through C++. it's not just not a frontend it's the exact opposite: it's emitting IR using IR builders.

Knowing that reaching broad devex parity is very expensive I think the real win is figuring out what specific problem you have and building community and robust software support around that.

It's the fact that AMD doesn't prioritize the reliability of its hardware and software stack. If I run llama.cpp on Vulkan I get a reasonable speedup, but if I raise the batch size to 512, the GPU is starting to make strange noises and shuts the PC down midway. Very cool. 98% of zero is still zero.

In fact cuBLAS and CUDA are kinda orthogonal in that you're either calling a pre-built cuBLAS kernel or writing your own CUDA kernel but not really combining the two.I'd say CUDA shines more because of stability, documentation, community support + examples, and ability to use modern C++ features in GPU code.

I'd say CUDA shines more because of stability, documentation, community support + examples, and ability to use modern C++ features in GPU code.

Targeting nvidia GPUs? Or in general? For whom?Building a performant BLAS library is hard but certainly not impossible. The tricks discussed in this post are hardly anything new either. Now, making a BLAS competitive with Nvidia's on its own GPUs is bound to be tough. But not technically unfeasible (after all, you can drop down to PTX if needed).

Building a performant BLAS library is hard but certainly not impossible. The tricks discussed in this post are hardly anything new either. Now, making a BLAS competitive with Nvidia's on its own GPUs is bound to be tough. But not technically unfeasible (after all, you can drop down to PTX if needed).

On average over 20 runs:CuBLAS (./sgemm 0) has 50.9 TFLOPS.My kernel has 61.8 TFLOPS, so it's actually +21% speedup in this benchmark.How do I collect my paycheck?

CuBLAS (./sgemm 0) has 50.9 TFLOPS.My kernel has 61.8 TFLOPS, so it's actually +21% speedup in this benchmark.How do I collect my paycheck?

My kernel has 61.8 TFLOPS, so it's actually +21% speedup in this benchmark.How do I collect my paycheck?

How do I collect my paycheck?

On a 4090 gpu, average of 20 runs of SGEMM_CUDA:  size    tflops_cublas  tflops_my  diff
  4096²   50.8-50.9      61.8       +21%
  8192²   56.3-56.4      67.1       +19%
  16384²  53.6           66.7       +24%

I guess the right thing to do now would be to hire a B2B salesman and figure out, which company needs it.

size    tflops_cublas  tflops_my  diff
  4096²   50.8-50.9      61.8       +21%
  8192²   56.3-56.4      67.1       +19%
  16384²  53.6           66.7       +24%

I guess the right thing to do now would be to hire a B2B salesman and figure out, which company needs it.

size    tflops_cublas  tflops_my  diff
  4096²   50.8-50.9      61.8       +21%
  8192²   56.3-56.4      67.1       +19%
  16384²  53.6           66.7       +24%

I have seen how those high-performance libraries are made and I'm still in awe at the quality and quantity of the staffing involved. Those were the smartest and most knowledgeable engineers I met in my career.

Generalizing from a micro benchmark is typically hubris.

Then there are also numerics: being fast is not enough if your implementation accumulates a lot of rounding errors doing so. Floating point arithmetic can and will mess up your results in unexpected ways. -funsafe famously is neither fun nor safe.Maybe tooling will catch up and make it easier. Think tinygrad with beamsearch, triton or halide.

Maybe tooling will catch up and make it easier. Think tinygrad with beamsearch, triton or halide.


Maharshi's blog

Home Blog

Learning CUDA by optimizing softmax: A worklog

04 Jan, 2025

The softmax operation is crucial. It is used extensively as a layer within deep learning models like transformers where it normalizes raw scores (logits) into a probability distribution. This property makes it particularly useful in classification tasks, where each output neuron represents the likelihood of a specific class. Optimizing softmax, especially in the context of GPU programming with CUDA, presents many opportunities for learning.

In this worklog, we will start by benchmarking PyTorch's softmax operation then finally we will iteratively optimize it in CUDA. The NVIDIA GPU used for this worklog is one GTX 1050Ti (that's all I have got right now).

The full code is available on my GitHub: Optimizing softmax in CUDA

Let's start.

The math

Before getting into it all, let's take a moment to understand the math behind the softmax operation. Softmax for an input vector X having N elements, produces an output vector O with N elements, where the ith element in the output vector is defined as:

Note that softmax operation depends on the current element xi and also on the sum of exponentials of all the elements of the input vector X. We will call this sum as the "normalization factor" (or, norm) henceforth.

Usually, instead of a single vector we deal with a matrix of shape (M,N) consisting of M rows where each row is a vector of N elements. Softmax is then performed along the columns of this matrix. The output here will be another matrix of the same shape.

Throughout this worklog, we will be working with a matrix of shape (1024,32768) i.e. 33,554,432 floating point numbers in total.

Example of the softmax output on a vector containing 5 elements:

import torch
import torch.nn.functional as F

vector = torch.randn(5, dtype=torch.float32)
print("Input vector:", vector)

# softmax along the last dimension
output = F.softmax(vector, dim=-1)
print("Output vector:", output)

Input vector: tensor([-1.3701,  0.7485,  0.1610, -2.0154,  1.0918])

Output vector: tensor([0.0382, 0.3176, 0.1765, 0.0200, 0.4477])

There is a problem though:

If the values of xi are very large (or very small), then the exponentials might cause overflow or underflow considering the precision limits of floating point numbers on a modern computer. We cannot represent and work with very large or very small numbers. This means for extreme values, the above version of softmax is NOT numerically stable.

But... there is a fix! We can modify the above equation in such a way that the overall operation becomes numerically stable while being correct: We subtract the maximum value xmax of the vector (a constant) from each xi before computing the exponential. This subtraction operation "shifts" the numbers to a range that can work nicely with floating point numbers. The numerically stable softmax equation becomes:

How this "shifted" equation results in the correct softmax output is left as an excersice to the reader :)

How fast is PyTorch?

We can get a baseline metric on how fast PyTorch is for computing the softmax operation, along the last dimension, on a randomly initialized matrix.

Following the above example, we can get a quick measure for the execution time of the softmax function:

import time
import torch
import torch.nn.functional as F

# Initialize the matrix on devuce
matrix = torch.randn(1024, 32768, device='cuda', dtype=torch.float32)

# Warm up
_ = torch.nn.functional.softmax(matrix, dim=-1)

# Ensure all CUDA operations are finished
torch.cuda.synchronize()

total_time = 0
n_iters = 5

for i in range(n_iters):
    # Measure time
    torch.cuda.synchronize()  # Ensure all CUDA operations are finished
    start = time.time()
    _ = torch.nn.functional.softmax(matrix, dim=-1)
    torch.cuda.synchronize()  # Synchronize again
    end = time.time()

    total_time += (end - start) * 1000
    print(total_time)

print(f"Softmax computation time (average): {(total_time/n_iters):.3f} ms")

Softmax computation time (average): 7.226 ms

From our quick test, PyTorch takes around 7.2 milliseconds to process and compute softmax on the entire matrix. Now, let's see how far can we go with implementing softmax in CUDA.

Kernel 1 - Naive softmax

In this kernel, we will assume that each thread in a block processes and computes one entire row of the input matrix. If the number of threads in one block is N_THREADS, then we need a total of ceil(M / N_THREADS) blocks to process the entire matrix.

The figure below shows this. Note that row = blockDim.x * blockIdx.x + threadIdx.x is the row which each thread within some block will process.

The actual computation is quite intuitive here. Softmax is calculated in three passes over the input array:

Pass 1 - Calculation of the maximum: The whole input row is first traversed from left (index = 0) to right (index = N - 1) to find the maximum value xmax.

Pass 2 - Calculation of the norm: The whole input row is traversed from left to right again, but this time the normalization factor is computed using the xmax value from the first pass, for each element.

Pass 3 - Softmax computation: The whole input row is traversed again from left to right and for each element the exponential of (x−xmax) is divided by the norm calculated in the second pass.

Below is the specific code snippet that does this:

int row = blockDim.x * blockIdx.x + threadIdx.x;

if (row < M) {
    // maximum of this row
    float x_max = -INFINITY;
    // norm factor of this row
    float norm = 0.0f;

    // output in 3 passes
    for (int col = 0; col < N; col++) {
        int i = row * N + col;
        x_max = max(x_max, input[i]);
    }
    for (int col = 0; col < N; col++) {
        int i = row * N + col;
        norm += expf(input[i] - x_max);
    }
    for (int col = 0; col < N; col++) {
        int i = row * N + col;
        output[i] = expf(input[i] - x_max) / norm;
    }
}

Running this kernel results in:

>> GPU allocation time: 10.727424 ms
>> Host to device transfer time: 26.176161 ms
>> Kernel execution time: 124.102112 ms
>> Device to host transfer time: 37.320896 ms

The naive kernel takes around 124.10 milliseconds to execute. This is 17.24 times slower compared to PyTorch's 7.2 milliseconds.

Can we improve it? Ofcourse we can.

Kernel 2 - Online softmax

Three passes to compute softmax is not at all optimal. Maybe there's a way to "fuse" the first pass (calculating the maximum) and the second pass (calculating the norm) together.

To do this, we will exploit the multiplication property of exponentials i.e.

To calculate the xmax and norm in just one pass, at each step we need to multiply the "current norm" with a "correction term".

For example, consider the following input vector: V=[3,2,5,1] for which we need to compute maximum and norm. We will now iterate through this input vector to see what correction term do we need and when do we need it.

Assume that the variables maxi and normi will represent maximum and norm untill the ith element.

Starting at i=0:

Note that after the first iteration, the values for maximum and norm are the correct values (but till the first index).

Next at i=1:

We add the "previous norm" value to the "current norm" value at each iteration.

Now at i=2:

Finally at i=3:

After the final iteration, we remain with:

and,

We just calculated both maximum and norm factor in only one pass by using a correction term and by exploiting the property of multiplying exponentials! The correction term is:

Now, to write this algorithm as a CUDA kernel, we simply use the naive kernel and "fuse" the first two loops into one:

int row = blockDim.x * blockIdx.x + threadIdx.x;

if (row < M) {
    float x_max = -INFINITY;
    float norm = 0.0f;

    // pass 1
    for (int col = 0; col < N; col++) {
        int i = row * N + col;
        float curr = input[i];
        if (curr > x_max) {
            // correct the global norm here
            norm = norm * expf(x_max - curr);
            x_max = curr;
        }
        norm += expf(curr - x_max);
    }
    // pass 2
    for (int col = 0; col < N; col++) {
        int i = row * N + col;
        input[i] = expf(input[i] - x_max) / norm;
    }
}

Running this kernel results in:

>> GPU allocation time: 10.431488 ms
>> Host to device transfer time: 25.897375 ms
>> Kernel execution time: 88.149567 ms
>> Device to host transfer time: 33.533314 ms

Using this simple trick (also called online softmax) we see that this kernel is 1.39 times (around 28.12%) faster than the naive kernel.

That's a clever improvement, but we can do more. We need to dive deeper into how we can use threads within one block to parallelize the computations even more by collaborating with each other.

Kernel 3 - Shared memory and reductions

The more you learn about GPU programming with CUDA, the more you will realize that memory is structured into hierarchies. The list below shows the access speeds of different memory hierarchies from fastest to slowest.

The kernels above uses only global GPU memory. Reading from and writing to global memory is expensive and time consuming, so we need to somehow reduce the access and storing time.

The idea here is to have each block (thread block) process one row of the input matrix and the threads within each block will process only a chunk of the entire row. Have a look at the figure below to understand which elements will each thread load.

Here tid = threadIdx.x loads elements spaced by blockDim.x so that the threads with different tids load consecutive elements from the input row. This helps in achieving memory coalescing where accessing consecutive addresses from the global memory is faster than accessing random addresses.

There is a problem though: To calculate the values of maximum and norm, we need to have access to all the elements of the input row. How will we do that if different threads have access to only a chunk of the input row?

This is where reductions come into play. Bear with me on this one.

Let's assume each thread has its own private set of variables called local_max and local_norm and also suppose that there are N_THREADS threads in total. Now, the thread with tid = i will compute the local max and local norm using the elements i, i + blockDim.x, i + 2*blockDim.x and so on.

After all the threads in a block complete processing their respective chunks, we will be left with N_THREADS values for local_max and local_norm. To calculate the global maximum value, we need to "reduce" these N_THREADS local maximum values to 1 global maximum value. The figure below will help you understand this.

However, to perform this "block-level" reduction we will need to store the local maximum value in the shared memory of the block. Each thread will store its local maximum as:

smem[tid] = local_max;
__syncthreads();

Note we also add a sync barrier to ensure that each thread correctly stores its local maximum into the corresponding address in the shared memory and waits for other threads before moving on to the reduction step.

We will now use the shared memory to reduce the N_THREADS local maximum values to 1 value and then store it in the first address (smem[0]) in the shared memory. The reduction step looks like:

for (int stride = blockDim.x / 2; stride > 0; stride /= 2) {
    if (tid < stride) {
        smem[tid] = max(smem[tid], smem[tid + stride]);
    }
    // sync before next iteration
    __syncthreads();
}

float global_max = smem[0];
__syncthreads();

This code block performs reduction in O(log(N)) time complexity which is faster than reducing linearly i.e. O(N) complexity. Let's see an example of this reduction with 8 threads where the shared memory will contain 8 maximum values in the start:

Initially:

smem = [3, 7, 2, 8, 6, 4, 5, 1]

First Iteration (stride = 4):

Each thread with tid < 4 compares smem[tid] with smem[tid + stride] and updates smem[tid] with the maximum.

Comparisons:

tid = 0: smem[0] = max(smem[0], smem[4]) = max(3, 6) = 6
tid = 1: smem[1] = max(smem[1], smem[5]) = max(7, 4) = 7
tid = 2: smem[2] = max(smem[2], smem[6]) = max(2, 5) = 5
tid = 3: smem[3] = max(smem[3], smem[7]) = max(8, 1) = 8

Updated smem:

smem = [6, 7, 5, 8, 6, 4, 5, 1]

Second Iteration (stride = 2):

Each thread with tid < 2 compares smem[tid] with smem[tid + stride] and updates smem[tid].

Comparisons:

tid = 0: smem[0] = max(smem[0], smem[2]) = max(6, 5) = 6
tid = 1: smem[1] = max(smem[1], smem[3]) = max(7, 8) = 8

Updated smem:

smem = [6, 8, 5, 8, 6, 4, 5, 1]

Third Iteration (stride = 1):

Each thread with tid < 1 compares smem[tid] with smem[tid + stride] and updates smem[tid].

Comparison:

tid = 0: smem[0] = max(smem[0], smem[1]) = max(6, 8) = 8

Updated smem:

smem = [8, 8, 5, 8, 6, 4, 5, 1]

Final State:

After the reduction, the maximum value is stored in smem[0], which is:

global_max = smem[0] = 8

This shows how in only 3 iterations, we performed the reduction and got access to the global maximum value from the 8 threads. We do the same reduction for local_norm as well to find the global norm value. The only difference for local norm value is that, instead of performing the max operation we perform the + operation.

Here's how the kernel looks like for reduction of the maximum value:

__shared__ float smem[1024];

int row = blockIdx.x;
int tid = threadIdx.x;

// edge condition (we don't process further)
if (row >= M) return;

float* input_row = xd + row * N;
float* output_row = resd + row * N;
float local_max = -INFINITY;
float local_norm = 0.0f;

for (int i = tid; i < N; i += blockDim.x) {
    float x = input_row[i];
    if (x > local_max) {
        local_norm *= expf(local_max - x);
        local_max = x;
    }
    local_norm += expf(x - local_max);
}
__syncthreads();

smem[tid] = local_max;
__syncthreads();

for (int stride = blockDim.x / 2; stride > 0; stride /= 2) {
    if (tid < stride) {
        smem[tid] = max(smem[tid], smem[tid + stride]);
    }
    __syncthreads();
}

float global_max = smem[0];
__syncthreads();

The output from this kernel looks like:

>> GPU allocation time: 10.464928 ms
>> Host to device transfer time: 22.674080 ms
>> Kernel execution time: 6.612160 ms
>> Device to host transfer time: 41.318016 ms

Right away we see that this kernel which uses shared memory and reductions is already around 8.33% (1.09 times) faster than PyTorch's implementation.

Can we improve this even more? Let's see.

Kernel 4 - Shuffle instructions

This kernel will be largely similar to the previous one with one difference. If you notice carefully, in the reduction operations for local maximum value and local norm value we are accessing the shared memory and syncing the threads in every iteration. Even though accessing shared memory is fast, what if we could eliminate the usage of shared memory and syncing barriers while reducing the values?

Before explaining how, we need to understand the concept of warps within thread blocks:

Warps are a fundamental unit of execution within a thread block. A warp is a group of 32 threads in a thread block that execute the same instruction simultaneously (SIMD: Single Instruction, Multiple Data). All threads in a warp execute instructions in lockstep, meaning all 32 threads execute the same instruction at the same time on different data. If a thread block contains N threads, the number of warps is ceil(N / 32). Also, when threads in a warp follow different execution paths (e.g., due to conditional statements), it leads to warp divergence, reducing performance as the threads execute sequentially instead of in parallel.

In our case, if we have blockDim.x = 1024 then each block is composed of 32 warps (each warp consisting of 32 threads).

To limit the usage of shared memory, CUDA provides us with shuffle instructions which are specialized intrinsics that allow threads within a warp to directly exchange data without the overhead of shared memory. These are warp-level primitives and are highly efficient because they use registers to exchange data which is faster than using shared memory (according to the hierarchy).

Suppose in one block we have N_THREADS threads in total. That means, we have NW = ceil(N_THREADS / warp_size) warps where warp_size is usually 32 threads. Now, instead of doing a block-level reduction using shared memory what if we first perform a warp-level reduction:

From N_THREADS values, doing a warp-level reduction for every warp available will leave us with NW values across the block that needs to be reduced further. So, the first available warp can load the values from the remaining warps, and then perform a warp-level reduction again to get the final value. Let's consider an example to ease your mind:

Suppose there are 16 threads that have already calculated their respective local maximum values. Also, assume that warp_size = 4 which means there are 4 warps in total. The values are [3, 7, 2, 9, 4, 1, 8, 5, 10, 6, 12, 11, 13, 14, 15, 16].

Step 1: Warp-level reduction

The warp size is 4, so there are 4 warps in the block (16 threads / 4 threads per warp). Each warp performs its own reduction.

Warp 0 (Threads 0 to 3: Values [3, 7, 2, 9]):

Offset = 2:

Offset = 1:

Result for Warp 0: 9 (stored in Thread 0 of the warp).

Warp 1 (Threads 4 to 7: Values [4, 1, 8, 5]):

Offset = 2:

Offset = 1:

Result for Warp 1: 8 (stored in Thread 4 of the warp).

Warp 2 (Threads 8 to 11: Values [10, 6, 12, 11]):

Offset = 2:

Offset = 1:

Result for Warp 2: 12 (stored in Thread 8 of the warp).

Warp 3 (Threads 12 to 15: Values [13, 14, 15, 16]):

Offset = 2:

Offset = 1:

Result for Warp 3: 16 (stored in Thread 12 of the warp).

Step 2 - Block-level reduction

At this point, the maximum values from each warp are stored in the first thread of each warp: [9, 8, 12, 16].

The block-level reduction begins.

Store Warp Results in Shared Memory:

Synchronize Threads:

Perform Final Reduction Using First Warp:

First Warp Reduction (smem = [9, 8, 12, 16]):

Offset = 2:

Offset = 1:

Global Block Maximum: 16 (stored in smem[0]).

At this point, we have the global maximum value for the entire block using warp-level reductions.

How to actually perform these warp-level reductions though? CUDA provides us with shuffle instructions for that. We will use the __shfl_down_sync instruction to perform reduction. Here's how it works:

It is a CUDA warp-level primitive that shifts data values down within a warp. Threads in the warp exchange data based on a specified offset, and threads that would receive data from out-of-bounds threads are assigned a default value. The syntax for __shfl_down_sync is:

T __shfl_down_sync(unsigned mask, T var, int delta, int width=warpSize);

Here:

Consider the following piece of code:

int val = threadIdx.x;
int shifted_val = __shfl_down_sync(0xFFFFFFFF, val, 1);

For delta = 1:

The reduction code for this kernel looks like:

float val = local_max;
for (int offset = warp_size / 2; offset > 0; offset /= 2) {
    val = fmaxf(val, __shfl_down_sync(0xffffffff, val, offset));
}

if (blockDim.x > warp_size) {
    if (tid % warp_size == 0) {
        // which warp are we at?
        // store the value in its first thread index
        smem[tid / warp_size] = val;
    }
    __syncthreads();

    if (tid < warp_size) {
        val = (tid < CEIL_DIV(blockDim.x, warp_size)) ? smem[tid] : -INFINITY;
        for (int offset = warp_size / 2; offset > 0; offset /= 2) {
            val = fmaxf(val, __shfl_down_sync(0xffffffff, val, offset));
        }
        if (tid == 0) smem[0] = val;
    }
} else {
    if (tid == 0) smem[0] = val;
}
__syncthreads();

float global_max = smem[0];
__syncthreads();

and the kernel outputs:

>> GPU allocation time: 10.542080 ms
>> Host to device transfer time: 25.580065 ms
>> Kernel execution time: 5.174400 ms
>> Device to host transfer time: 45.923008 ms

This kernel is around 1.29 times (or, 22.73%) faster than the shared memory kernel! Using shuffle instructions eliminated the need of using sync barriers __syncthreads in each iteration as well.

Conclusion

In this worklog, we iteratively optimized the softmax operation starting from PyTorch and then writing a custom CUDA kernel for the same. With the above improvements, our custom softmax CUDA kernel became around 1.41 times (or, 29.17%) faster than PyTorch on RTX 1050Ti.

Thank you for reading!

#CUDA
#softmax