Sign up

Sign in

Sign up

Sign in

An almost pointless exercise in GPU optimization

Andrew Innes

Follow

--

2

Listen

Share

Unlike various of my colleagues, I’m not an ML expert who is comfortable writing funky fused operators to make big ML models run faster on GPUs in conjunction with clever quantisation tricks. Even so, I want to better understand the potential for us to make more effective use of GPUs in our products, and so I decided to revisit the first problem I ever considered trying to accelerate with a GPU. My attempt then was foiled by looking only at shader languages, which are too primitive to be effective for my problem, but this time I was aware of CUDA programming in C++ which is much more general purpose and accessible.

What is unusual about my chosen problem is that it is officially pointless. You ought not to be able to find any library that will accelerate this algorithm, because it isn’t worth writing one! But as you will see, it feels like it should be possible to make it run significantly faster by making use of the many cores in a GPU compared with the relatively few cores in a CPU. That’s what I thought then, and wanted to test now as a learning exercise.

That also makes it an interesting proxy for lots of specific algorithms you might have in your own programs, which aren’t catered for by high-performance libraries written by experts, but feel like they should be possible to accelerate because GPUs will run thousands of threads in parallel. If that is your situation, this article might help you think about what to expect and what might be involved.

TL;DR

Getting an existing C++ algorithm running on GPU is pretty easy, so it is a low bar to get started. What I learned is the importance of minimizing thread divergence and maximizing effective memory access speed. To do that effectively, I had to transform my algorithm into a state machine structure so that every thread is operating mostly in lock-step, just with different data values.

My starting, interim and final code are open to see, along with a summary of the steps I took, and the corresponding improvements or regressions at each stage. I want to focus in this article on the thought process for deciding each step, mostly by explaining the Nvidia Nsight Compute analysis which helped guide me.

In the end I managed to make my program run about 30x faster on my laptop using its GeForce GTX 1650 GPU, compared with its Core i7–9750H CPU. Only in the last two steps did it get meaningfully better than with CPU though, so be prepared for early and frequent disappointment.

If you want just the summary of what worked, jump to Progression History.

A Pointless Program

Years ago, a colleague invited me to take on his Christmas programming challenge, which was to write the fastest program he could to continuously play the card game Beggar My Neighbour. The aim, noted by John Conway as definitely not worth solving, is to try to find the longest game — with the possibility that there might be a game which never ends. The longest game found so far has 7972 turns — a rainy afternoon diversion perhaps, if you can sustain playing a card every 2.7095 seconds for six hours straight! You can see a history of new records, with a Python program that verifies them, here: https://github.com/matthewmayer/beggarmypython

The game play algorithm is almost trivial, but it has some notable features which turn out to be relevant to the challenge of effectively leveraging a GPU.

CPU Starting Point

My initial C++ program to search for long games is here. It is just a port of the Python program with a search loop that runs continuously, shuffling the deck randomly between games. I implemented two simple optimizations that seemed obvious:

The program can use multiple CPU cores by just running separate copies of the search loop in different threads, starting with different RNG seeds. Each search loop does its own random walk through potential deals.

On my laptop, throughput peaks at 2.9 million deals per second, using 12 threads to match the number of logical CPU cores.

Initial GPU Port

The beauty of General Purpose GPU programming is that you can often get started quickly with the code you already have. There is a good chance that only a few adaptations are needed to convert a program that uses multiple threads on CPU to one that runs with even more threads on GPU, as long as you don’t have extreme memory requirements (for code or data or both) which force breaking up the work into smaller pieces.

GPU cores have a similar range of machine instructions to those of CPUs, so plain algorithms will compile readily enough to give reasonable single core efficiency; you don’t need to transform your program to use a subset of variable types or data structures or code constructs, or special parallelization primitives or libraries. You mainly just need to cope with some changes to library functions (such as the random number generator in my case), and finessing of class structures to graft in the global functions needed to launch work on the GPU.

You do need to be ready for disappointment though.

If your code is similar to mine, founded on nested branching logic, you may find that the GPU won’t even be able to match the CPU performance at first. This is to be expected: CPUs are designed for running unrelated complex branching logic in each thread, dedicating significant chip area on circuitry to predict branches, speculatively execute multiple paths, and do lots of caching — 1 MB per logical CPU core in my laptop for instance. CPU cores also run fast, with 3–4GHz max speed being fairly typical. By comparison, GPUs are much lighter on hardware per core for caching, branch prediction etc, and top out at perhaps 1.5–2GHz (unless you have an extreme cooling gaming GPU). They come with faster memory though, to help compensate.

But the net effect is that an algorithm will probably run 2–3 times faster on a single CPU core than it will on a single GPU core, from a combination of the clock speed ratio and more aggressive single thread acceleration. You need to figure out how to make good use of the thousands of slow GPU cores in order to outperform a few fast CPU cores.

My initial port to GPU ran at about 1.4M deals per second (once thermal throttling kicks in), using 2048 threads (128 blocks of 16 threads each). Not an encouraging start, but in hindsight about what I should have expected.

Learn to use Nsight Compute

Early on, I decided to get comfortable using the Nsight Compute tool to analyse the GPU portion of my code in spectacular detail, rather than trying to rely on intuition and the high-level utilization figures from nvidia-smi. The approach I settled on was to step through the execution of the first couple of GPU kernel launches (the kernel here being my global function which runs the core search and game play loop for a predefined number of iterations), and then profile the next kernel run to see how much of the hardware capability was being used effectively.

The first, unflattering, report that the tool shows is the “Speed of Light” summary: ie. the percentage of theoretical peak performance on compute and memory bandwidth utilization. My first program scored around 12% for compute, and 28% for memory. (For the same program, nvidia-smi would report 88% and 28% utilization respectively, underscoring how misleading its information can be when you are trying to optimize your algorithm in the early stages.)

There were many detailed metrics underneath the headline figures which could be examined, and various analysis warnings that point out potential problem areas. Many of these sounded esoteric, but there were two reasonably clear actionable warnings:

The first one was easy to address: assign at least 32 threads per block (not 16), so we don’t leave warp capacity unused for not even trying. I ended up settling on 32 blocks of 32 threads, which increased throughput to 2.3M deals per second. nvidia-smi now reported 95% and 13% for compute and memory utilization. Nsight Compute reports we have improved average active threads per warp from 2.4 to 3.6. That left Thread Divergence as the key warning to address.

Thread Divergence

If you have learned about CUDA programming before, you may recall that thread divergence occurs when not all threads in the same warp are executing the same instruction. (Each group of 32 consecutive threads in a block will always execute together as a unit, known as a warp, and GPU hardware is heavily optimized around running these warps of threads very efficiently on the assumption that the threads are normally executing instructions in unison, working on closely related pieces of the same computation.)

It is part of the beauty of General Purpose GPU programming that thread divergence is allowed, and is handled automatically by the hardware as well as it can be, by simply letting the threads in the warp take turns to run their next instruction, when they are different. Subsets of the warp that do share the next instruction will execute together, so degradation in performance is proportional to how much the threads have diverged from each other.

In the worst case, where every thread is at its own unique point in the code, the threads are effectively time sliced on the hardware, and so running at 1/32 of the hardware’s potential. From the Nsight Compute warning, I could see we were quite close to that point:

[Warning] Instructions are executed in warps, which are groups of 32 threads. Optimal instruction throughput is achieved if all 32 threads of a warp execute the same instruction. The chosen launch configuration, early thread completion, and divergent flow control can significantly lower the number of active threads in a warp per cycle. This kernel achieves an average of 3.6 threads being active per cycle. This is further reduced to 3.1 threads per warp due to predication. The compiler may use predication to avoid an actual branch. Instead, all instructions are scheduled, but a per-thread condition code or predicate controls which threads execute the instructions. Try to avoid different execution paths within a warp when possible. In addition, ensure your kernel makes use of Independent Thread Scheduling, which allows a warp to reconverge after a data-dependent conditional block by explicitly calling __syncwarp().

Getting about 3 active threads per warp means I’m only tapping into about 1/10th of the GPU’s compute capacity, at best. Not exactly what you would assume from the 95% compute utilization figure reported by nvidia-smi — that figure really just means I’m doing something on all of the GPU’s Streaming Multiprocessor units about 95% of the time — even if that something is very inefficient at using each SM, as in this case.

I highlighted the most relevant part of the remediation advice above, which is essentially to remove the nested branching logic as much as I can. To do that in my program, where each thread is working on a different game, I realised I would have to rewrite the core game play function. The original version uses position within the code’s nested branching structure to encode part of the game state — I needed to replace that with an explicit data representation of all pieces of game state. Then every step of the inner loop would be executed in unison across all threads playing their own games, just with different data values.

Rewrite to use a lookup table

To make the inner loop purely data driven, I chose to introduce a state transition table. Game state is now tracked explicitly in variables which are treated as inputs to the transition lookup table, followed by a set of actions that is always performed on every iteration using data values from the lookup table.

A critical realization for implementing this cleanly in my case was to notice that the discard pile can be treated as a third player in the game, with some extra states to track when the discard pile is “playing” its cards to the player who won the last trick. With that mental twist in place, a fairly straightforward lookup table became possible to write by hand. The code for this version is here; it works on both CPU and GPU, but it is slower on both — about half the speed of the baseline version on CPU, and two-thirds on GPU. ☹

Being slower on CPU was expected, since it requires more instructions to manipulate more state variables, and there is more memory access to lookup state transitions. We also have forced the adding of the discard pile to the winning player’s hand on every trick to be done slowly, as individual turns with all the overhead needed for that. However, on GPU I had hoped to see gains because now more threads in each warp should be able to execute in unison.

In fact Speed-of-Light compute and memory utilization figures did improve, to 17% and 38% respectively, but Thread Divergence was still listed as a remediation warning. We are at least up from an average of 3.6 active threads per warp to 5.3, but that is small comfort given the actual speed is now about back to where I first started on GPU, which is well behind the CPU performance.

Ah, be careful about function exits

What I had forgotten is that thread divergence will also occur when threads are exiting from a nested function inside the kernel. My game play loop would always play one game (inside the function play), then return to the search loop to swap two cards before calling play again. It basically didn’t matter that inside the play inner loop thread divergence has been largely solved, because each thread in a warp will still finish at different times, depending on how many turns are needed for their respective games. So they usually diverge at game end, and the whole warp of threads must wait for the longest game amongst them to finish. Stats on the range of game lengths shows a minimum of about 33 turns, an average of about 250 turns, and a max somewhere in the many thousands of turns at least. That’s a lot of variation.

With that realization, my next program refactoring was to include the logic for game completion book-keeping and switching to a new game as another (necessarily conditional) step in the inner game playing loop. This slows down the inner loop even further, but allows the threads to stay mostly converged across multiple games. To support this, I pre-create a backlog of games to play, ready for the next game to be picked very quickly from inside the inner loop by whichever thread is available first. (This introduces the only inter-thread synchronization mechanism needed so far, which is the use of a CUDA atomic operation to read and increment the next_deal_to_play index, which barely affects speed but solves the race condition.)

In order to allow the threads to run as long as possible in this synchronized inner loop, I decided to use a big chunk of main GPU memory to hold a large backlog, which so far has barely been used (just for the search loop’s best-game-so-far records).

We now get to a near-final version of the program, which can be recreated from the latest version using some conditional compilation definitions. This version fills a large (eg. 1M entry) backlog of deals to play in GPU memory (using another kernel of cooperating GPU threads working in parallel), which is then processed in parallel by 1024 threads (32 blocks of 32 threads). Performance did indeed improve over the initial version with the lookup table, but only to about the same speed as the baseline version achieved once I had tweaked the number of blocks and threads. What gives?!

Ah, memory speed really matters

Nsight Compute reveals Speed-of-Light is about the same as before (12% compute and 37% memory utilization), and again lists a number of warnings. This time, some other warnings now seem more relevant and actionable (as I’ve highlighted below):

“[Warning] All pipelines are under-utilized. Either this kernel is very small or it doesn’t issue enough warps per scheduler. Check the Launch Statistics and Scheduler Statistics sections for further details.”

“[Warning] Every scheduler is capable of issuing one instruction per cycle, but for this kernel each scheduler only issues an instruction every 6.0 cycles. This might leave hardware resources underutilized and may lead to less optimal performance. Out of the maximum of 8 warps per scheduler, this kernel allocates an average of 1.00 active warps per scheduler, but only an average of 0.17 warps were eligible per cycle. Eligible warps are the subset of active warps that are ready to issue their next instruction. Every cycle with no eligible warp results in no instruction being issued and the issue slot remains unused. To increase the number of eligible warps either increase the number of active warps or reduce the time the active warps are stalled.”

“[Warning] On average each warp of this kernel spends 2.4 cycles being stalled waiting for a scoreboard dependency on a L1TEX (local, global, surface, texture) operation. This represents about 39.7% of the total average of 6.0 cycles between issuing two instructions. To reduce the number of cycles waiting on L1TEX data accesses verify the memory access patterns are optimal for the target architecture, attempt to increase cache hit rates by increasing data locality or by changing the cache configuration, and consider moving frequently used data to shared memory.”

Thread Divergence is still listed as a warning, though the average active threads per warp is now at about 9. Looking at the “source” view of the kernel profile, which shows per-instruction info such as the execution counts, average active threads, distribution of instruction stall reasons etc, it looks like there are lots of times when the inner loop instructions are stalled waiting for access to results from memory.

So it looks like there are two clear problems we can address:

· Ensure there are more (warps of) threads available to schedule, so they can hide the micro-latencies that are naturally experienced by any one warp.

· Reduce the times when the inner loop is stalled accessing GPU memory.

The first issue can in theory be addressed by playing with the blocks and threads configuration. However, changing that doesn’t really make a difference: we are primarily bottlenecked on accessing data from GPU memory.

Use Shared Memory as much as possible

The best improvement we can make now is ironically to stop using main GPU memory as much as we can. In my case, that means replacing the single large backlog of deals with a very small backlog per block, held in Shared Memory.

As you may recall from CUDA Programming primers, Shared Memory is the fastest memory area (other than registers allocated by the compiler) available to programmers, but the most limited in size. On my laptop’s GPU, it is limited to 48KB — the same amount of memory as my first home computer (from the Apple II era). This memory is per-block and only available while each specific block is executing — its contents will be lost when that block ends, so relevant info needs to be copied to main memory as final output.

Thankfully for my program, using this only involves a modest code change, and basically means the inner game loop runs for fewer iterations before the local backlog is exhausted and has to be replenished. Also each small backlog is shared by fewer threads, so there is more opportunity for some threads to finish early and increase thread divergence.

With this change (which included changing blocks and threads to 16 and 128, in order to try and provide more warps to the scheduler as noted above), performance on GPU finally moved ahead of CPU for the first time, and by a fairly impressive margin, reaching about 40M deals per second (vs. 3M on CPU)!

Nsight Compute is still saying we are memory bottlenecked though. Final chance to squeeze further!

Make Shared Memory stretch as far as possible

My final change is to recognise that the core data structures are using more memory than they need to, since I’m using an enum to represent 5 possible card values, and enum by default is represented as an int in C++, which is treated as a 32-bit word in CUDA. Similarly, the lookup table data values are all defined as int, but far fewer bits would suffice. I had initially wondered whether native word sized values were more efficient at the individual instruction level, but actually GPUs (like most CPUs) efficiently support sub-word data sizes and even arbitrary bitfield operations at the instruction level, so that isn’t an important concern.

With a change to specify uint8 as my enum base type, add appropriate bitfield declarations in the lookup table struct, and use a more compact representation of deals in the backlog (not the 64 entry circular buffer representation used for playing fast), I am able to squeeze a longer backlog into the 48KB of shared memory, and also reduce the memory bandwidth needed for lookup operations and other game play steps.

The effect is rather gratifying: my final version now hits over 100M deals per second, at least until thermal throttling brings it back down to 95M or so. 😊

Nsight Compute is still saying I’m memory bound, but at this stage I’m thinking that may be just how it is. The algorithm is so light on computation that it will typically be waiting on memory (even the super-fast Shared Memory) no matter what. The next step, if it were possible to see how to restructure the game playing loop in a suitable way, would be to try and ensure we created coalesced memory access patterns (where threads in a warp all reference directly adjacent memory locations that allow the hardware to make them one read/write operation). But that seems very unlikely to be possible with each thread still working on its own deal.

Maybe there is still some potential for improvement from tweaking block and thread counts, and paying attention to memory layout and other micro-details. I’m not holding my breath!

· Final Speed-of-Light score: 18% compute, 70% memory utilization.

· nvidia-smi reports 90% and 1% for compute and memory utilization respectively.

· Thread divergence is reduced, with average active threads per warp of 20 (vs. ideal of 32).

Progression History

From initial CPU version to final decent GPU version, here was the history of changes and results. All performance numbers are in millions of deals per second.

--

--

2

Written by Andrew Innes

Chief Architect at Speechmatics

Responses (2)

Help

Status

About

Careers

Press

Blog

Privacy

Terms

Text to speech

Teams


Sign up

Sign in

Sign up

Sign in

Introduction to CUDA Optimization with Practical Examples

FreaxRuby

Follow

--

Listen

Share

Creating a CUDA Project (first_cuda.cu)

To start, you need to include the necessary libraries. For using the runtime API, include cuda_runtime.h.

#include <stdio.h>#include <cuda_runtime.h>

Main Function:

First, write a main function to check if CUDA is available.

int main() {    if (!InitCUDA()) {        return 0;    }        printf("CUDA initialized.\n");    return 0;}

Checking CUDA Availability:

bool InitCUDA() {    int count;        cudaGetDeviceCount(&count); // Get the number of available devices    if (count == 0) {        fprintf(stderr, "There is no device.\n");        return false;    }    int i;        for (i = 0; i < count; i++) {        cudaDeviceProp prop;        if (cudaGetDeviceProperties(&prop, i) == cudaSuccess) {            if (prop.major >= 1) {                break;            }        }    }        if (i == count) {        fprintf(stderr, "There is no device supporting CUDA 1.x.\n");        return false;    }    cudaSetDevice(i);    return true;}

After completing these steps, you can compile the file using nvcc, the CUDA compiler tool that separates the GPU and host parts of the code. GPU parts are compiled into intermediate code using NVIDIA's compiler, while host parts are compiled with the system's C++ compiler.

Simple Addition with CUDA (cuda_sum1.cu)

This section builds upon the initial CUDA project to calculate the sum of squares of a set of numbers.

Initial Setup:

Modify the beginning of your code to include necessary headers and define the size of the data.

#include <stdio.h>#include <stdlib.h>#include <cuda_runtime.h>#define DATA_SIZE 1048576int data[DATA_SIZE];

Generating Random Numbers:

Implement a function to generate random numbers.

void GenerateNumbers(int *number, int size) {    for (int i = 0; i < size; i++) {        number[i] = rand() % 10;    }}

Preparing Data for CUDA:

Before performing computations with CUDA, you need to copy the data from the host memory to the device memory.

GenerateNumbers(data, DATA_SIZE);int* gpudata, *result;cudaMalloc((void**) &gpudata, sizeof(int) * DATA_SIZE);cudaMalloc((void**) &result, sizeof(int));// Copy from host to device memorycudaMemcpy(gpudata, data, sizeof(int) * DATA_SIZE, cudaMemcpyHostToDevice);

Writing a Kernel Function:

In CUDA, functions executed on the device are annotated with __global__. Let's write a kernel function to compute the sum of squares.

__global__ static void sumOfSquares(int *num, int* result) {    int sum = 0;    int i;    for (i = 0; i < DATA_SIZE; i++) {        sum += num[i] * num[i];    }    *result = sum;}

Device functions have restrictions, such as not being able to return values directly.

Executing the Kernel:

Execute the function on CUDA using the following syntax:

sumOfSquares<<<1, 1, 0>>>(gpudata, result);

This configuration uses a single thread, so both the block and thread counts are set to 1, and no shared memory is used.

After execution, copy the result back to host memory, free device memory, and print the result:

int sum;cudaMemcpy(&sum, result, sizeof(int), cudaMemcpyDeviceToHost);cudaFree(gpudata);cudaFree(result);printf("sum: %d\n", sum);

To validate the correctness, you might compare the GPU result with a CPU computation.

Optimizing CUDA Addition (Parallelization — Optimization 1.0)

The initial implementation of the sum of squares did not utilize parallelization, with the entire program running on a single thread. This is not ideal for GPU architecture, which is designed for parallel execution. To improve performance, let’s parallelize the computation.

Parallelization Approach:

Split the computation into multiple parts, with each part calculating a segment of the sum of squares. Initially, sum these partial results on the CPU.

Modifications:

Define constants for data size and the number of threads.

#define DATA_SIZE 1048576#define THREAD_NUM 256

Optimized Kernel Function:

Modify the sumOfSquares function to support parallel execution.

__global__ static void sumOfSquares(int *num, int* result, clock_t* time) {    const int tid = threadIdx.x;    const int size = DATA_SIZE / THREAD_NUM;    int sum = 0;    int i;    clock_t start;    if (tid == 0) start = clock();    for (i = tid * size; i < (tid + 1) * size; i++) {        sum += num[i] * num[i];    }    result[tid] = sum;    if (tid == 0) *time = clock() - start;}

This version divides the data among multiple threads, each calculating a portion of the sum. The start time is recorded for performance measurement.

Main Function Adjustments:

Allocate memory for results and timing, then execute the optimized kernel.

int* gpudata, *result;clock_t* time;cudaMalloc((void**) &gpudata, sizeof(int) * DATA_SIZE);cudaMalloc((void**) &result, sizeof(int) * THREAD_NUM);cudaMalloc((void**) &time, sizeof(clock_t));cudaMemcpy(gpudata, data, sizeof(int) * DATA_SIZE, cudaMemcpyHostToDevice);sumOfSquares<<<1, THREAD_NUM, 0>>>(gpudata, result, time);int sum[THREAD_NUM];clock_t time_used;cudaMemcpy(&sum, result, sizeof(int) * THREAD_NUM, cudaMemcpyDeviceToHost);cudaMemcpy(&time_used, time, sizeof(clock_t), cudaMemcpyDeviceToHost);cudaFree(gpudata);cudaFree(result);cudaFree(time);

After execution, sum the partial results on the CPU for the final sum.

int final_sum = 0;for (int i = 0; i < THREAD_NUM; i++) {    final_sum += sum[i];}printf("sum: %d  time: %d\n", final_sum, time_used);

This approach significantly improves performance by utilizing parallel execution on the GPU.

Memory Access Pattern Optimization (Optimization 2.0)

The efficiency of memory access on GPUs greatly impacts performance. GPUs use DRAM, which is most efficient when accessed in a continuous manner. The previous version’s access pattern did not fully exploit this due to the way threads accessed memory.

Optimizing Memory Access:

Adjust the sumOfSquares function to ensure that threads access memory in a more efficient, continuous pattern.

__global__ static void sumOfSquares(int *num, int* result, clock_t* time) {    const int tid = threadIdx.x;    int sum = 0;    int i;    clock_t start;    if (tid == 0) start = clock();    for (i = tid; i < DATA_SIZE; i += THREAD_NUM) {        sum += num[i] * num[i];    }    result[tid] = sum;    if (tid == 0) *time = clock() - start;}

By adjusting the loop to increment by THREAD_NUM, each thread accesses continuous memory locations, improving the efficiency of memory reads and significantly boosting performance.

These optimizations demonstrate how leveraging CUDA’s parallel processing capabilities and understanding the GPU’s memory access patterns can lead to substantial performance gains in computational tasks.

Further Parallelization with Blocks (Optimization 3.0)

In CUDA, threads can be organized into blocks, where threads within the same block can share memory and synchronize their execution. This structure allows for more sophisticated parallelization strategies.

Adjusting for Block Usage:

Let’s refine our approach by using multiple blocks and threads to increase the number of parallel computations.

Define Block and Thread Numbers:

#define DATA_SIZE 1048576#define BLOCK_NUM 32#define THREAD_NUM 256

This configuration uses 32 blocks, each with 256 threads, totaling 8192 threads for computation.

Kernel Function with Blocks:

Modify the sumOfSquares kernel to utilize both blocks and threads.

__global__ static void sumOfSquares(int *num, int* result, clock_t* time) {    const int tid = threadIdx.x;    const int bid = blockIdx.x;    int sum = 0;    int i;    if (tid == 0) time[bid] = clock();    for (i = bid * THREAD_NUM + tid; i < DATA_SIZE; i += BLOCK_NUM * THREAD_NUM) {        sum += num[i] * num[i];    }

result[bid * THREAD_NUM + tid] = sum;    if (tid == 0) time[bid + BLOCK_NUM] = clock();}

This version calculates a segment of the sum per thread, with each block covering different segments. The start and end times are recorded for each block to measure performance.

Main Function Adjustments for Blocks:

Allocate memory and launch the kernel with blocks and threads.

int* gpudata, *result;clock_t* time;cudaMalloc((void**) &gpudata, sizeof(int) * DATA_SIZE);cudaMalloc((void**) &result, sizeof(int) * THREAD_NUM * BLOCK_NUM);cudaMalloc((void**) &time, sizeof(clock_t) * BLOCK_NUM * 2);cudaMemcpy(gpudata, data, sizeof(int) * DATA_SIZE, cudaMemcpyHostToDevice);sumOfSquares<<<BLOCK_NUM, THREAD_NUM, 0>>>(gpudata, result, time);int sum[THREAD_NUM * BLOCK_NUM];clock_t time_used[BLOCK_NUM * 2];cudaMemcpy(&sum, result, sizeof(int) * THREAD_NUM * BLOCK_NUM, cudaMemcpyDeviceToHost);cudaMemcpy(&time_used, time, sizeof(clock_t) * BLOCK_NUM * 2, cudaMemcpyDeviceToHost);cudaFree(gpudata);cudaFree(result);cudaFree(time);

Sum the partial results on the CPU and calculate the total execution time by finding the earliest start and the latest end time across all blocks.

int final_sum = 0;for (int i = 0; i < THREAD_NUM * BLOCK_NUM; i++) {    final_sum += sum[i];}clock_t min_start = time_used[0], max_end = time_used[BLOCK_NUM];for (int i = 1; i < BLOCK_NUM; i++) {    if (min_start > time_used[i]) min_start = time_used[i];    if (max_end < time_used[i + BLOCK_NUM]) max_end = time_used[i + BLOCK_NUM];}printf("sum: %d  time: %d\n", final_sum, max_end - min_start);

This approach further improves the performance by utilizing more threads and blocks, significantly increasing the parallelization of the computation.

Synchronization with Threads (Optimization 4.0)

To optimize further, let’s have each block calculate the sum of its threads’ results, reducing the workload on the CPU.

Using Shared Memory for Block-Wide Summation:

__global__ static void sumOfSquares(int *num, int* result, clock_t* time) {    extern __shared__ int shared[];    const int tid = threadIdx.x;    const int bid = blockIdx.x;    int i;    if (tid == 0) time[bid] = clock();    shared[tid] = 0;    for (i = bid * THREAD_NUM + tid; i < DATA_SIZE; i += BLOCK_NUM * THREAD_NUM) {        shared[tid] += num[i] * num[i];    }    __syncthreads(); // Synchronize threads in the block    // Sum the results within the block    if (tid == 0) {        for (i = 1; i < THREAD_NUM; i++) {            shared[0] += shared[i];        }        result[bid] = shared[0];    }    if (tid == 0) time[bid + BLOCK_NUM] = clock();}

This version uses shared memory within each block for intermediate results, reducing the need for global memory access and improving performance.

Main Function Adjustments for Synchronization:

Update the main function to allocate

shared memory for the kernel and adjust memory allocations and kernel invocation accordingly.

This method significantly reduces the amount of data the CPU needs to sum, leading to further performance gains.

Further Optimization with Tree-Based Reduction (Optimization 5.0)

Tree-based reduction is a technique that can improve the efficiency of parallel summation within blocks by structuring the addition in a tree-like manner, reducing the depth of dependency and improving parallelism.

Implementing Tree-Based Reduction:

Modify the kernel to use a tree-based approach for summing the results within a block, which involves halving the number of threads involved in each step of the summation.

This approach minimizes synchronization and memory access overhead, further accelerating the computation.

Conclusion

Through these optimizations, starting from simple parallelization to more advanced techniques like using blocks, shared memory, synchronization, and tree-based reduction, we’ve demonstrated how to leverage CUDA’s capabilities to achieve substantial performance improvements in computations. Each step introduces new CUDA concepts and optimization strategies, showing the potential for significant speedups in GPU-accelerated applications.

--

--

Written by FreaxRuby

🐱🐶🐯🦁🦋🤖

No responses yet

Help

Status

About

Careers

Press

Blog

Privacy

Terms

Text to speech

Teams


Sign up

Sign in

Sign up

Sign in

TDS Archive

Home

Newsletter

About

An archive of data science, data analytics, data engineering, machine learning, and artificial…

Accelerating AI/ML Model Training with Custom Operators

On the potential benefits of creating model-specific GPU kernels and their application to optimizing the use of dynamically shaped tensors

Chaim Rand

Follow

TDS Archive

--

1

Listen

Share

This post continues a long series of posts on the topic of analyzing and optimizing the runtime performance of training AI/ML models. The post could easily have been titled “PyTorch Model Performance Analysis and Optimization — Part 7”, but due to the weight of the topic at hand, we decided that a dedicated post (or series of posts) was warranted. In our previous posts, we have spoken at length about the importance of analyzing and optimizing your AI/ML workloads and the potentially significant impact it can have on the speed and costs of AI/ML model development. We have advocated for having multiple tools and techniques for profiling and optimizing training performance and have demonstrated many of these in practice. In this post we will discuss one of the more advanced optimization techniques — one that sets apart the true rock stars from the simple amateurs — creating a custom PyTorch operator in C++ and CUDA.

Popular ML frameworks, such as PyTorch, TensorFlow, and JAX are typically built using SW components that are optimized for the underlying hardware that the AI/ML workload is run on, be it a CPU, a GPU, or an AI-specific ASIC such as a Google TPU. However, inevitably, you may find the performance of certain computation blocks that comprise your model to be unsatisfactory or in-optimal. Oftentimes, tuning the low-level code blocks — often referred to as kernels — to the specific needs of the AI/ML model, can result in significant speed-ups to the runtime performance of model training and inference. Such speed-ups can be accomplished by implementing functionalities that were previously unsupported (e.g., an advanced attention block), fusing together individual operations (e.g., as in PyTorch’s tutorial on multiply-add fusion), and/or optimizing existing kernels based on the specific properties of the model at hand. Importantly, the ability to perform such customization depends on the support of both the AI HW and the ML framework. Although our focus on this post will be on NVIDIA GPUs and the PyTorch framework, it should be noted that other AI ASICs and ML frameworks enable similar capabilities for custom kernel customization. NVIDIA enables the development of custom kernels for its GPUs through its CUDA toolkit. And PyTorch includes dedicated APIs and tutorials for exposing this functionality and integrating it into the design of your model.

Our intention in this post is to draw attention to the power and potential of kernel customization and demonstrate its application to the unique challenge of training models with dynamically shaped tensors. Our intention is not — by any means — to replace the official documentation on developing custom operations. Furthermore, the examples we will share were chosen for demonstrative purposes only. We have made no effort to optimize these or verify their robustness, durability, or accuracy. If, based on this post, you choose to invest in AI/ML optimization via custom CUDA kernel development, you should be sure to undergo the appropriate training.

Toy Model — The Challenge of Dynamically Shaped Tensors

The prevalence of tensors with dynamic shapes in AI models can pose unique and exciting challenges with regards to performance optimization. We have already seen one example of this in a previous post in which we demonstrated how the use of boolean masks can trigger an undesired CPU-GPU sync event and advocated against their use. Generally speaking, AI accelerators tend to prefer tensors with fixed shapes over ones with dynamic shapes. Not only does it simplify the management of memory resources, but it also enables greater opportunity for performance optimization (e.g., using torch.compile). The toy example that follows demonstrates this challenge.

Suppose we are tasked with creating a face detection model for a next-generation digital camera. To train this model we are provided with a dataset of one million 256x256 grayscale images and associated ground-truth bounding boxes for each image. Naturally, the number of faces in each image can vary greatly, with the vast majority of images containing five or fewer faces, and just a few containing dozens or even hundreds. The requirement from our model is to support all variations. Specifically, our model needs to support the detection of up to 256 faces in an image.

To address this challenge, we define the following naïve model that generates bounding boxes and an accompanying loss function. In particular, we naïvely truncate the model outputs based on the number of target boxes rather than perform some form of assignment algorithm for matching between the bounding box predictions and ground truth targets. We (somewhat arbitrarily) choose the Generalized Intersection Over Union (GIOU) loss. A real-world solution would likely be far more sophisticated (e.g., it would include a loss component that includes a penalty for false positives).

import torchimport torch.nn as nnimport torch.nn.functional as Fclass Net(nn.Module):    def __init__(self):        super().__init__()        conv_layers = []        for i in range(4):            conv_layers.append(nn.Conv2d(4 ** i, 4 ** (i + 1), 3,                                         padding='same'))            conv_layers.append(nn.MaxPool2d(2, 2))            conv_layers.append(nn.ReLU())        self.conv_layers = nn.Sequential(*conv_layers)        self.lin1 = nn.Linear(256 * 256, 256 * 64)        self.lin2 = nn.Linear(256 * 64, 256 * 4)    def forward(self, x):        x = self.conv_layers(x.float())        x = self.lin2(F.relu(self.lin1(x.view((-1, 256 * 256)))))        return x.view((-1, 256, 4))def generalized_box_iou(boxes1, boxes2):    # loosly based on torchvision generalized_box_iou_loss code    epsilon = 1e-5    area1 = (boxes1[..., 2]-boxes1[..., 0])*(boxes1[..., 3]-boxes1[..., 1])    area2 = (boxes2[..., 2]-boxes2[..., 0])*(boxes2[..., 3]-boxes2[..., 1])    lt = torch.max(boxes1[..., :2], boxes2[..., :2])    rb = torch.min(boxes1[..., 2:], boxes2[..., 2:])    wh = (rb - lt).clamp(min=0)    inter = wh[..., 0] * wh[..., 1]    union = area1 + area2 - inter    iou = inter / union.clamp(epsilon)    lti = torch.min(boxes1[..., :2], boxes2[..., :2])    rbi = torch.max(boxes1[..., 2:], boxes2[..., 2:])    whi = (rbi - lti).clamp(min=0)    areai = (whi[..., 0] * whi[..., 1]).clamp(epsilon)    return iou - (areai - union) / areaidef loss_fn(pred, targets_list):    batch_size = len(targets_list)    total_boxes = 0    loss_sum = 0.    for i in range(batch_size):        targets = targets_list[i]        num_targets = targets.shape[0]        if num_targets > 0:            sample_preds = pred[i, :num_targets]            total_boxes += num_targets            loss_sum += generalized_box_iou(sample_preds, targets).sum()    return loss_sum / max(total_boxes, 1)

Due to the varying number of faces per image, the loss is calculated separately for each individual sample rather than a single time (for the entire batch). In particular, the CPU will launch each of the GPU kernels associated with the loss function B times, where B is the chosen batch size. Depending on the size of the batch, this could entail a significant overhead, as we will see below.

In the following block we define a dataset that generates random images and associated bounding boxes. Since the number of faces varies per image, we require a custom collate function for grouping samples into batches:

from torch.utils.data import Dataset, DataLoaderimport numpy as np# A dataset with random images and gt boxesclass FakeDataset(Dataset):    def __init__(self):        super().__init__()        self.size = 256        self.img_size = [256, 256]    def __len__(self):        return 1000000    def __getitem__(self, index):        rand_image = torch.randint(low=0, high=256,                                    size=[1]+self.img_size,                                   dtype=torch.uint8)                # set the distribution over the number of boxes to reflect the fact        # that the vast majority of images have fewer than 10 faces        n_boxes = np.clip(np.floor(np.abs(np.random.normal(0, 3)))                                      .astype(np.int32), 0, 255)                box_sizes = torch.randint(low=1, high=self.size, size=(n_boxes,2))        top_left = torch.randint(low=0, high=self.size-1, size=(n_boxes, 2))        bottom_right = torch.clamp(top_left + box_sizes, 0, self.size -1)        rand_boxes = torch.concat((top_left,bottom_right), dim = 1)        return rand_image, rand_boxes.to(torch.uint8)def collate_fn(batch):    images = torch.stack([b[0] for b in batch],dim=0)    boxes = [b[1] for b in batch]    return images, boxestrain_loader = DataLoader(    dataset = FakeDataset(),    batch_size=1024,    pin_memory=True,    num_workers=16,    collate_fn=collate_fn)

Typically, each training step starts with copying the training batch from the host (CPU) to the device (GPU). When our data samples are of fixed size, they are copied in batches. However, one of the implications of the varying number of faces per image is that the bounding box targets of each sample is copied separately requiring many more individual copy operations.

def data_to_device(data, device):    if isinstance(data, (list, tuple)):        return type(data)(            data_to_device(val, device) for val in data        )    elif isinstance(data, torch.Tensor):        return data.to(device=device, non_blocking=True)

Lastly, we define our training/evaluation loop. For the purposes of our discussion, we have chosen to focus just on the forward pass of our training loop. Note the inclusion of a PyTorch profiler object and our use of explicit synchronization events (to facilitate performance evaluation of different portions of the forward pass).

device = torch.device("cuda:0")model = torch.compile(Net()).to(device).train()# forward portion of training loop wrapped with profiler objectwith torch.profiler.profile(   schedule=torch.profiler.schedule(wait=5, warmup=5, active=10, repeat=1),   on_trace_ready=torch.profiler.tensorboard_trace_handler('/tmp/perf/'),   profile_memory=True) as prof:    for step, data in enumerate(train_loader):        with torch.profiler.record_function('copy data'):            images, boxes = data_to_device(data, device)            torch.cuda.synchronize(device)        with torch.profiler.record_function('forward'):            with torch.autocast(device_type='cuda', dtype=torch.bfloat16):                outputs = model(images)            torch.cuda.synchronize(device)        with torch.profiler.record_function('calc loss'):            loss = loss_fn(outputs, boxes)            torch.cuda.synchronize(device)        prof.step()        if step > 30:            break    # filter and print profiler results    event_list = prof.key_averages()    for i in range(len(event_list) - 1, -1, -1):        if event_list[i].key not in ['forward', 'calc loss', 'copy data']:            del event_list[i]    print(event_list.table())

Performance Analysis

Running our script on a Google Cloud g2-standard-16 VM (with a single L4 GPU), a dedicated deep learning VM image, and PyTorch 2.4.0, generates the output below (which we trimmed for readability).

-------------  ------------  ------------         Name     CPU total  CPU time avg-------------  ------------  ------------    copy data     288.164ms      28.816ms      forward        1.192s     119.221ms    calc loss        9.381s     938.067ms-------------  ------------  ------------Self CPU time total: 4.018sSelf CUDA time total: 10.107s

Despite the fact that the loss function contains far fewer operations, it completely dominates the overall step time. The overhead of the repeated invocations of the underlying GPU kernels (for each sample in the batch) is clearly evident in the Trace view in TensorBoard:

Optimization Through Concatenation

One way to reduce the number of calls to the loss function is to combine together all of the valid boxes each batch using concatenation, as shown in the following block.

def loss_with_concat(pred, targets_list):    bs = len(targets_list)    all_targets = torch.concat(targets_list, dim = 0)    num_boxes = [targets_list[i].shape[0] for i in range(bs)]    all_preds = torch.concat([pred[i,: num_boxes[i]] for i in range(bs)],                              dim=0)    total_boxes = sum(num_boxes)    loss_sum = generalized_box_iou(all_targets, all_preds).sum()    return loss_sum/max(total_boxes, 1)

The results of this optimization are captured below.

-------------  ------------  ------------         Name     CPU total  CPU time avg-------------  ------------  ------------    copy data     522.326ms      52.233ms      forward        1.187s     118.715ms    calc loss     254.047ms      25.405ms-------------  ------------  ------------Self CPU time total: 396.674msSelf CUDA time total: 1.871s

The concatenation optimization resulted in a 37X (!!) speed-up of the loss function. Note, however, that it did not address the overhead of the individual host-to-device copies of the sample ground-truth data. This overhead is captured in the screenshot below from TensorBoard’s Trace view:

Optimization Through Padding

A common approach to avoiding the use of dynamically shaped tensors is padding. In the following code block, we modify the collate function to pad (with zeros) the ground-truth bounding-boxes of each data sample to the maximum number of supported boxes, 256. (Note that the padding could also have been performed in the Dataset class.)

def collate_with_padding(batch):    images = torch.stack([b[0] for b in batch],dim=0)    padded_boxes = []    for b in batch:        p = torch.nn.functional.pad(                       b[1], (0, 0, 0, 256 - b[1].shape[0]), value = 0)        padded_boxes.append(p)    boxes = torch.stack(padded_boxes,dim=0)    return images, boxes

Padding the samples to fixed sized tensors enables us to copy the ground truth of the batch with a single call. It also allows us to compute the loss with a single invocation of the loss function. Note, that this method requires masking the resultant loss, as shown below, so that only the valid boxes are taken into consideration.

def loss_with_padding(pred, targets):    mask = (targets[...,3] > 0).to(pred.dtype)    total_boxes = mask.sum()    loss = generalized_box_iou(targets, pred)    masked_loss = loss*mask    loss_sum = masked_loss.sum()    return loss_sum/torch.clamp(total_boxes, 1)

The resultant runtime performance is captured below:

-------------  ------------  ------------         Name     CPU total  CPU time avg-------------  ------------  ------------    copy data      57.125ms       5.713ms      forward        1.315s     131.503ms    calc loss      18.438ms       1.844ms-------------  ------------  ------------Self CPU time total: 11.723msSelf CUDA time total: 1.378s

Note the nearly 10X boost in the data copy and the additional 14X boost in the loss function performance. Keep in mind that padding may increase the use of the GPU memory. In our case, this increase is less than 1%.

While the runtime of our loss function has improved dramatically, we note that the vast majority of the calculations that are performed in the loss functions are immediately masked away. We can’t help but wonder whether there is a way to further improve the performance by avoiding these redundant operations. In the next section, we will explore the opportunities provided by using custom CUDA kernels.

Creating a Custom CUDA Kernel

Many tutorials will highlight the difficulty of creating CUDA kernels and the high entrance barrier. While mastering CUDA development and tuning kernels to maximize the utilization of the GPU could indeed require years of experience as well as an intimate understanding of the GPU architecture, we strongly believe that even a novice (but ambitious) CUDA enthusiast/ML developer can succeed at — and greatly benefit from — building custom CUDA kernels. In this section we will take PyTorch’s (relatively simple) example of a C++/CUDA extension for PyTorch and enhance it with a GIOU kernel. We will do this in two stages: First we will naïvely carry over all of the GIOU logic to C++/CUDA to assess the performance impact of kernel fusion. Then, we will take advantage of our new-found low-level control to add conditional logic and reduce unneeded arithmetic operations.

Developing CUDA kernels allows you to determine the core logic that is performed in each of the GPU threads and how these are distributed onto the underlying GPU streaming multiprocessors (SMs). Doing this in the most optimal manner requires an expert understanding of the GPU architecture including the different levels of GPU memory, memory bandwidth, the on-chip acceleration engines (e.g., TensorCores), the supported number of concurrent threads per SM and how they are scheduled, and much much more. What makes things even more complicated is that these properties can vary between GPU generations and flavors. See this blog for a very basic, but very easy, introduction to CUDA.

Step 1 — Kernel Fusion

Looking back at the Trace view of our last experiment, you may notice that the forward pass of our loss calculation includes roughly thirty independent arithmetic operations which translate to launching and running an independent CUDA kernel (as can be seen by simply counting the number of cudaLaunchKernel events). This can negatively impact performance in a number of ways. For example:

Optimization through kernel fusion attempts to reduce this overhead by combining these operations into a lower number of kernels so as to reduce the overhead of multiple kernels.

In the code block below, we define a kernel that performs our GIOU on a single bounding-box prediction-target pair. We use a 1-D grid to allocate thread blocks of size 256 each where each block corresponds to one sample in the training batch and each thread corresponds to one bounding box in the sample. Thus, each thread — uniquely identified by a combination of the block and thread IDs — receives the predictions (boxes1) and targets (boxes2) and performs the GIOU calculation on the single bounding box determined by the IDs. As before, the “validity” of the bounding box is controlled by the value of the target boxes. In particular, the GIOU is explicitly zeroed wherever the corresponding box is invalid.

#include <torch/extension.h>#include <cuda.h>#include <cuda_runtime.h>namespace extension_cpp {__global__ void giou_kernel(const float* boxes1,                            const float* boxes2,                             float* giou,                             bool* mask) {  int idx = blockIdx.x * blockDim.x + threadIdx.x;  bool valid = boxes2[4*idx+3] != 0;  mask[idx] = valid;  const float epsilon = 1e-5;  const float* box1 = &boxes1[idx * 4];  const float* box2 = &boxes2[idx * 4];  // Compute area of each box  float area1 = (box1[2] - box1[0]) * (box1[3] - box1[1]);  float area2 = (box2[2] - box2[0]) * (box2[3] - box2[1]);  // Compute the intersection  float left = max(box1[0], box2[0]);  float top = max(box1[1], box2[1]);  float right = min(box1[2], box2[2]);  float bottom = min(box1[3], box2[3]);  float inter_w = max(right - left, 0);  float inter_h = max(bottom - top, 0);  float inter_area = inter_w * inter_h;  // Compute the union area  float union_area = area1 + area2 - inter_area;  // IoU  float iou_val = inter_area / max(union_area, epsilon);  // Compute the smallest enclosing box  float enclose_left = min(box1[0], box2[0]);  float enclose_top = min(box1[1], box2[1]);  float enclose_right = max(box1[2], box2[2]);  float enclose_bottom = max(box1[3], box2[3]);  float enclose_w = max(enclose_right - enclose_left, 0);  float enclose_h = max(enclose_bottom - enclose_top, 0);  float enclose_area = enclose_w * enclose_h;  float result = iou_val - (enclose_area-union_area)/max(enclose_area, epsilon);  // Generalized IoU  giou[idx] = result * valid;}at::Tensor giou_loss_cuda(const at::Tensor& a, const at::Tensor& b) {  TORCH_CHECK(a.sizes() == b.sizes());  TORCH_CHECK(a.dtype() == at::kFloat);  TORCH_CHECK(b.dtype() == at::kFloat);  TORCH_INTERNAL_ASSERT(a.device().type() == at::DeviceType::CUDA);  TORCH_INTERNAL_ASSERT(b.device().type() == at::DeviceType::CUDA);  int bs = a.sizes()[0];  at::Tensor a_contig = a.contiguous();  at::Tensor b_contig = b.contiguous();  at::Tensor giou = torch::empty({a_contig.sizes()[0], a_contig.sizes()[1]},                                  a_contig.options());  at::Tensor mask = torch::empty({a_contig.sizes()[0], a_contig.sizes()[1]},                                  a_contig.options().dtype(at::kBool));  const float* a_ptr = a_contig.data_ptr<float>();  const float* b_ptr = b_contig.data_ptr<float>();  float* giou_ptr = giou.data_ptr<float>();  bool* mask_ptr = mask.data_ptr<bool>();  // Launch the kernel  // The number of blocks is set according to the batch size.  // Each block has 256 threads corresponding to the number of boxes per sample  giou_kernel<<<bs, 256>>>(a_ptr, b_ptr, giou_ptr, mask_ptr);   at::Tensor total_boxes = torch::clamp(mask.sum(), 1);  torch::Tensor loss_sum = giou.sum();  return loss_sum/total_boxes;}// Registers CUDA implementations for giou_lossTORCH_LIBRARY_IMPL(extension_cpp, CUDA, m) {  m.impl("giou_loss", &giou_loss_cuda);}}

To complete the kernel creation, we need to add the appropriate C++ and Python operator definitions (see muladd.cpp and ops.py)

// Add the C++ definitionm.def(“giou_loss(Tensor a, Tensor b) -> Tensor”);

# define the Python operatordef giou_loss(a: Tensor, b: Tensor) -> Tensor:    return torch.ops.extension_cpp.giou_loss.default(a, b)

To compile our kernel, run the installation script (pip install .) from the base directory.

The following block uses our newly defined GIOU CUDA kernel:

def loss_with_kernel(pred, targets):    pred = pred.to(torch.float32)    targets = targets.to(torch.float32)    import extension_cpp    return extension_cpp.ops.giou_loss(pred, targets)

Note the explicit casting to torch.float32. This is a rather expensive operation that could be easily avoided by enhancing our CUDA kernel support. We leave this as an exercise to the reader :).

The results of running our script with our custom kernel are displayed below.

-------------  ------------  ------------         Name     CPU total  CPU time avg   -------------  ------------  ------------    copy data      56.901ms       5.690ms          forward        1.327s     132.704ms          calc loss       6.287ms     628.743us     -------------  ------------  ------------Self CPU time total: 6.907msSelf CUDA time total: 1.380s

Despite the naïveté of our kernel (and our inexperience at CUDA), we have boosted the loss function performance by an additional ~3X over our previous experiment (628 microseconds compare to 1.8 milliseconds). As noted above, this can be improved even further without much effort.

Step 2 — Conditional Execution

The thread-level control that CUDA provides us allows us to add a conditional statement that avoids computation on the invalid bounding boxes:

__global__ void giou_kernel(const float* boxes1,                            const float* boxes2,                            float* giou,                            bool* mask) {  int idx = blockIdx.x * blockDim.x + threadIdx.x;  bool valid = boxes2[4*idx+3] != 0;  mask[idx] = valid;  if (valid)  {    const float* box1 = &boxes1[idx * 4];    const float* box2 = &boxes2[idx * 4];    giou[idx] = compute_giou(box1, box2);  }  else  {    giou[idx] = 0;  }}

In the case of our kernel, the impact on runtime performance is negligible. The reason for this (presumably) is that our kernel is relatively small to the point that its runtime is negligible compared to the time required to load and instantiate it. The impact of our conditional execution might become apparent for larger kernels only. (The impact, as a function of the kernel size can be assessed by making our GIOU output dependent on a for loop that we run for a varying number of fixed steps. This, too, we leave as an exercise :).) It is also important to take into consideration how a conditional execution flow behaves on CUDA’s SIMT architecture, particularly, the potential performance penalty when threads belonging to the same warp diverge.

-------------  ------------  ------------         Name     CPU total  CPU time avg-------------  ------------  ------------    copy data      57.008ms       5.701ms      forward        1.318s     131.850ms    calc loss       6.234ms     623.426us-------------  ------------  ------------Self CPU time total: 7.139msSelf CUDA time total: 1.371s

Results and Next Steps

We summarize the results of our experiments in the table below.

Importantly, our work is not done. Admittedly, we have taken some shortcuts in the example we have shared:

Summary

In this post we demonstrated the potential of the use of a custom CUDA kernel on the runtime performance of AI/ML applications. We attempted, in particular, to utilize the low-level control enabled by CUDA to introduce a conditional flow to limit the number of redundant arithmetic operations in the case of dynamically shaped inputs. While the performance boost resulting from the fusion of multiple kernel operations was significant, we found the size of our kernel to be too small to benefit from the conditional execution flow.

Throughout many of our posts we have emphasized the importance of having multiple tools and techniques for optimizing ML and reducing its costs. Custom kernel development is one of the most powerful techniques at our disposal. However, for many AI/ML engineers, it is also one of the most intimidating techniques. We hope that we have succeeded in convincing you that this opportunity is within reach of any ML developer and that it does not require major specialization in CUDA.

In recent years, new frameworks have been introduced with the goal of making custom kernel development and optimization more accessible to AI/ML developers. One of the most popular of these frameworks is Triton. In our next post we will continue our exploration of the topic of custom kernel development by assessing the capabilities and potential impact of developing Triton kernels.

--

--

1

Published in TDS Archive

An archive of data science, data analytics, data engineering, machine learning, and artificial intelligence writing from the former Towards Data Science Medium publication.

Written by Chaim Rand

I am a Machine Learning Algorithm Developer working on Autonomous Vehicle technologies at Mobileye. The views expressed in my posts are my own.

Responses (1)

Help

Status

About

Careers

Press

Blog

Privacy

Terms

Text to speech

Teams


Sign up

Sign in

Sign up

Sign in

Member-only story

CUDA 3: Your Checklist for Optimizing CUDA Kernels

Rimika Dhara

Follow

--

Share

Hi! This is the fourth article in the series I have been writing about Programming in CUDA. For this post, I want to talk about how to optimize CUDA kernels and how we can build intuition behind kernel optimizations. I am not going to provide code here, just simple concepts and strategies that we can use to optimize the kernels. If you’re catching up, feel free to check out the previous three articles I’ve written on this topic:

Let’s get started!

1. Understanding the Issue

So, you have your CUDA kernel built; the logic is sound, and it’s giving you accurate results. But, it seems to be taking wayyyy too long. Whether it’s part of a deep learning model or a standalone kernel, slow CUDA kernels can become major performance bottlenecks. Just like when wanting to fix a broken car, it’s better to know which part of the model to start fixing. When we talk about performance, we often aim to maximize the GPU’s peak performance, which hinges on three things: compute, memory bandwidth, and overhead. Let’s discuss those:

--

--

Written by Rimika Dhara

Aspiring SWE and ML Engineer. Current SWE Intern at Oracle.

No responses yet

Help

Status

About

Careers

Press

Blog

Privacy

Terms

Text to speech

Teams


Sign up

Sign in

Sign up

Sign in

Unlock Warp‑Level Performance: DeepSeek’s Practical Techniques for Specialized GPU Tasks

Amin Sedaghat

Follow

--

Listen

Share

Introduction

Amid the push toward trillion‐parameter models, the DeepSeek team has showcased astonishing engineering feats — like partitioning Streaming Multiprocessors (SMs) into specialized warp‐level communication tasks, weaving in lower‐level PTX instructions, and auto‐tuning “communication chunks” for concurrency.

Their approach elegantly circumvents conventional limitations. In fact, DeepSeek achieves a level of warp specialization that NVIDIA itself hasn’t fully tackled. Orchestrating InfiniBand (IB) ↔ NVLink communication pipelines at warp granularity, minimizing overhead and pushing performance even further.

Reflecting this focus on warp‐level allocations, the DeepSeek paper notes:

[W]e employ the warp specialization technique (Bauer et al., 2014) and partition 20 SMs into 10 communication channels… In addition, both dispatching and combining kernels overlap with the computation stream, so we also consider their impact on other SM computation kernels. Specifically, we employ customized PTX (Parallel Thread Execution) instructions and auto-tune the communication chunk size, which significantly reduces the use of the L2 cache and the interference to other SMs.

We’ll begin with a simple 4×4 tile transpose, demonstrating how warp shuffle intrinsics can remove unneeded block‐wide synchronization. Then we’ll grow that idea into advanced warp‐level partitioning for real‐world tasks like all‐to‐all communication in multi‐node MoE workloads. Even if your own applications aren’t at DeepSeek’s hyper‐scale, you’ll still find these warp‐specialization ideas invaluable for high‐performance GPU work.

1. The Challenge: Rearranging Data at Warp Scope

Sometimes, you only need 32 threads (or even fewer) to rearrange data. Yet the default CUDA approach forces an entire block to synchronize (via __syncthreads()), which can stall many threads that aren’t involved.

DeepSeek’s large‐scale multi‐node MoE training pipeline demonstrates how such optimizations compound. Individual warps might handle IB sends, NVLink receives, or matrix‐multiplication tasks — enabling partial concurrency and better GPU utilization. By working at the warp scope, DeepSeek avoids the bottlenecks of block‐wide synchronization.

Here’s a simple diagram illustrating the difference between full block sync vs. focusing only on warp‐level coordination:

2. A Naive 4×4 Transpose (Block-Level Sync)

Below is a minimal kernel that uses shared memory and `__syncthreads()` to transpose a 4×4 chunk, causing an entire block to wait even though only 16 threads are actually active:

// Naive 4×4 Transpose (Block-Level Sync)__global__ void transpose_naive(float* output, const float* input){    __shared__ float tile[16];  // 4×4 tile    int tid = threadIdx.x;    // Only first 16 threads do the transfer    if (tid < 16)    {        tile[tid] = input[tid];    }    __syncthreads();  // Block-wide barrier    if (tid < 16)    {        int row = tid / 4;        int col = tid % 4;        int transposedId = col * 4 + row;        output[transposedId] = tile[tid];    }}

Explanation:

3. Warp-Level Optimization with Shuffle Intrinsics

Now let’s refine the approach: if data exchange is restricted to a single warp, we can avoid shared memory altogether by using warp shuffle intrinsics (__shfl_sync, etc.). This drastically cuts latency and removes block‐wide synchronization, key enablers in DeepSeek’s method of assigning different communication tasks to separate warps.

Below is the warp shuffle version. Notice that we only use half a warp (16 threads), so we manage “inactive” lanes with __activemask().

///////////////////////////////////////////////////////////////////////  // Warp-Level 4×4 Transpose with __shfl_sync (Partial Warp Example)  ///////////////////////////////////////////////////////////////////////__global__ void transpose_warp_shfl(float* output, const float* input){    int tid = threadIdx.x;    // Contains the set of lanes actually active in the warp    unsigned int mask = __activemask();    if (tid < 16)    {        float val = input[tid];        int row = tid / 4;        int col = tid % 4;        // Identify which lane originally held the transposed element        int srcTid = col * 4 + row;        // Use the valid lane mask to avoid undefined data from inactive lanes        float transposedVal = __shfl_sync(mask, val, srcTid, 32);        output[tid] = transposedVal;    }}

Explanation:

4. PTX and Warp Specialization: The Broader Context

DeepSeek’s large‐scale approach extends this same strategy beyond small tiles. Specialized warps handle IB sending, NVLink forwarding, or MoE gating without major stalls. This is done by carefully constructing kernels (and sometimes diving into PTX) to:

NVIDIA doesn’t provide a direct “pin warp to IB/NVLink” feature, yet DeepSeek’s team discovered and fine‐tuned these warp partitions manually. They also demonstrated that with customized PTX, they could push concurrency and reduce L2 traffic to levels not typically seen in standard CUDA pipelines.

5. Practical Pitfalls in Real Scenarios

If you only use part of a warp (say 16 out of 32 lanes), it’s crucial to manage all remaining “inactive” lanes properly. Otherwise, unintended reads or writes could occur. The CUDA intrinsic __activemask() helps you safely identify which lanes are active, ensuring data exchange remains correct. Meanwhile, if your data requirements exceed a single warp’s scope, you can’t rely on warp shuffle intrinsics alone. Larger tiles or computations typically call for shared memory, block‐level synchronization, or more intricate multi‐warp patterns.

Another important detail is memory visibility. Even though you’ve sidestepped block‐wide synchronization in warp‐only tasks, once another warp (or block) needs to read or write partial results, the proper fences and synchronization come back into the picture. Within‐warp data remains invisible to the wider block or grid unless you employ standard barriers.

Finally, specialized warps might coordinate both IB/NVLink traffic and local computation in theory, but making this overlap work in practice requires careful tuning and real‐world profiling. You’ll likely need to refine kernel launches, stream priorities, and chunk sizes, among other parameters.

6. When (and When Not) To Use Inline PTX

Even though __shfl_sync compiles into efficient instructions, sometimes you need super‐fine control , like specifying certain registers or caching modes that the compiler doesn’t expose. That’s where inline PTX steps in. But it’s a double‐edged sword:

Inline PTX should be your last resort. Nevertheless, DeepSeek’s customized instructions highlight how effective it can be. They overcame cross‐node communication bottlenecks by reducing L2 usage and carefully scheduling instructions at warp granularity, going far beyond what the default CUDA runtime typically offers.

7. Key Takeaways

When your operation fits within the confines of a single warp, shuffle intrinsics are your best friend, eliminating the need for block‐wide barriers and extra memory transfers. Just remember that partial warps demand careful handling: inactive lanes can wreak havoc if they’re not properly managed. With the help of __activemask(), you can safeguard your data against unintended reads or writes, ensuring your intra‐warp tasks stay predictable.

As you scale up toward multi‐warp processes, shared memory and block‐level sync steps inevitably rejoin the picture. This doesn’t mean you should fear them; rather, treat them as useful tools within more intricate concurrency patterns. Borrowing a page from DeepSeek’s playbook, you can dedicate entire warps to IB or NVLink communication and let other warps handle compute, resulting in better resource utilization and near‐continuous workstreams when concurrency is tuned correctly.

Lastly, if even the most aggressive compiler flags and CUDA intrinsics can’t hit your performance target, inline PTX stands as your final frontier. It grants you unrivaled control over registers, caching, and instruction scheduling, albeit at the cost of greater development complexity and maintenance. Just as DeepSeek did, tread carefully in this realm: successes here can be monumental, but so too can the engineering overhead.

Conclusion

DeepSeek’s mastery of warp‐level transposition, communication overlap, and custom PTX optimizations proves that GPUs are more malleable than most developers including myself realize. By restricting data exchange to the warp and dynamically allocating warp roles for communication or compute, the DeepSeek team achieves concurrency levels that standard tools seldom provide.

While everyone’s chasing the latest GPU that costs more than a brand new car, some crafty engineers at DeepSeek were like “Hold my coffee (or most likely tea)” and achieved what NVIDIA themselves hadn’t tackled yet. Sometimes all it takes is being clever with the threads you’ve got.

References:

https://github.com/deepseek-ai/DeepSeek-V3/blob/main/DeepSeek_V3.pdf

--

--

Written by Amin Sedaghat

http://aminsed.github.io/Portfolio/ https://github.com/Aminsed

No responses yet

Help

Status

About

Careers

Press

Blog

Privacy

Terms

Text to speech

Teams


Sign up

Sign in

Sign up

Sign in

Mirage : multi-level superoptimizer for faster tensor computation in LLMs

SACHIN KUMAR

Follow

--

Listen

Share

Current deep neural network (DNN) frameworks generally specify DNN computation using tensor programs, which are directed acyclic graphs whose nodes and edges represent tensor algebra operators and tensors. For optimizing an input tensor program, existing frameworks like Pytorch and Tensorflow ) use manually designed rules to map the tensor program to expert-written GPU kernels, which requires extensive engineering efforts to design and implement optimization rules.

In this paper [1], authors present Mirage, the first multi-level superoptimizer for tensor programs . Mirage is able to automatically discover and verify sophisticated tensor program optimizations that require joint optimization of algebraic transformations, schedule transformations, and discovery of new custom kernels.

Mirage Workflow and Key idea

i) Key idea

ii) Workflow

a) Program Partitioning

b) Expression-guided 𝜇Graph generator

c) Probabilistic equivalence verifier

d) 𝜇Graph optimizer

Mirage outperforms existing systems by up to 3.5×, by exploiting subtle custom kernels and optimizations missing in existing systems.

Multi-Level Graph Representation

a) GPU hierarchy

b) Kernel graph

c) Block graph

d) Grid dimensions

e) For-loop dimensions

f) Thread graph

g) Tensor layout

Case Study: Group-Query Attention

Expression-Guided 𝜇Graph Generator

Mirage 𝜇Graph generator automatically discovers potential 𝜇Graphs for an input tensor program, by finding 𝜇Graphs that capture optimizations at all of the kernel, block, and thread levels.

i) Kernel and Block Graph Generation

ii) Thread Graph Construction

iii) Pruning via Abstract Expressions

Probabilistic Equivalence Verifier

𝜇Graph Optimizer

Evaluation

i) Implementation

ii) Experimental Setup

iii) Performance Results

Conclusion

Paper discussed: https://www.cs.cmu.edu/~zhihaoj2/papers/mirage.pdf

Github link of this paper: https://github.com/mirage-project/mirage

References:

--

--

Written by SACHIN KUMAR

Staff MLE@Chegg/Applied Generative AI Researcher | https://www.linkedin.com/in/techsachinkumar/

No responses yet

Help

Status

About

Careers

Press

Blog

Privacy

Terms

Text to speech

Teams


Sign up

Sign in

Sign up

Sign in

Optimizing Vector Dot Product in CUDA: Exploring Shared Memory and Reduction Techniques

Dhanush

Follow

--

Listen

Share

Introduction

In the realm of parallel computing, CUDA (Compute Unified Device Architecture) provides a powerful platform to harness the full potential of GPUs. This blog post dives into the intricacies of implementing a vector dot product using CUDA, focusing on the efficient use of shared memory, handling race conditions, and applying reduction techniques to optimize performance.

Understanding the Vector Dot Product

The dot product of two vectors is a fundamental operation in many scientific and engineering applications. For vectors a\mathbf{a}a and b\mathbf{b}b of size nnn, the dot product is calculated as:

dot_product= ∑i=(0 to n−1)a[i]×b[i]

In this post, we’ll explore how to perform this operation on a GPU, leveraging CUDA for parallel computation.

1. Code Overview

The primary goal of this CUDA program is to compute the dot product of two vectors a\mathbf{a}a and b\mathbf{b}b using GPU parallelism. The result of the dot product is stored in an array c\mathbf{c}c, where each element corresponds to a partial dot product computed by a block of threads.

2. Code Breakdown

Here’s a step-by-step explanation of the code:

#include<iostream>#define n 10

using namespace std;const int threadperblock = 256;

__global__ void dot_product(int *a, int *b, int *c) {    __shared__ float cache[threadperblock];

int tid = threadIdx.x + blockIdx.x * blockDim.x;    int cacheIndex = threadIdx.x;

float temp = 0;    while (tid < n) {        temp += a[tid] * b[tid];        tid += blockDim.x * gridDim.x;    }

cache[cacheIndex] = temp;

__syncthreads();

int i = blockDim.x / 2;    while (i != 0) {        if (cacheIndex < i)            cache[cacheIndex] += cache[cacheIndex + i];

__syncthreads();        i /= 2;    }

if (cacheIndex == 0) {        c[blockIdx.x] = cache[0];    }}

3. Host Code

int main() {    int a[n], b[n], c[n], cpu_stored, stored_dot_val;    int *dev_a, *dev_b, *dev_c;

for (int i = 0; i < n; i++) {        a[i] = i + 2;        b[i] = i + 1;    }    cpu_stored = 0;    for (int i = 0; i < n; i++) {        cpu_stored += a[i] * b[i];    }    cout << "CPU Stored Value: " << cpu_stored << endl;

cudaMalloc((void **) &dev_a, n * sizeof(int));    cudaMalloc((void **) &dev_b, n * sizeof(int));    cudaMalloc((void **) &dev_c, n * sizeof(int));

cudaMemcpy(dev_a, a, n * sizeof(int), cudaMemcpyHostToDevice);    cudaMemcpy(dev_b, b, n * sizeof(int), cudaMemcpyHostToDevice);

int blockpergrid = (n + threadperblock - 1) / threadperblock;

dot_product<<<blockpergrid, threadperblock>>>(dev_a, dev_b, dev_c);    cudaMemcpy(&c, dev_c, blockpergrid * sizeof(int), cudaMemcpyDeviceToHost);

stored_dot_val = 0;    for (int i = 0; i < blockpergrid; i++) {        stored_dot_val += c[i];    }

cout << "Stored Value is: " << stored_dot_val << endl;    cudaFree(dev_a);    cudaFree(dev_b);    cudaFree(dev_c);    return 0;}

Conclusion

In this blog post, we’ve explored an efficient CUDA implementation for computing the dot product of vectors, emphasizing the use of shared memory and the reduction technique to enhance performance. By leveraging shared memory, we minimized global memory accesses, which significantly improves computational speed. The reduction technique, applied within each block, ensures that partial results are combined effectively, making the algorithm more efficient.

Handling race conditions and using synchronization barriers like __syncthreads() is crucial for ensuring the correctness of parallel computations. This approach allows multiple threads to safely collaborate and achieve accurate results.

For a complete view of the CUDA code and further exploration, visit my GitHub repository. There, you can find the full implementation and experiment with the code to better understand the concepts discussed.

Happy coding, and feel free to reach out with any questions or feedback!

--

--

Written by Dhanush

Graduate student from CSU- Dominguez Hills. Skilled in CUDA C/C++, parallel programming, and data structures. Proven expertise in React, ASP.NET Core & Docker.

No responses yet

Help

Status

About

Careers

Press

Blog

Privacy

Terms

Text to speech

Teams


Sign up

Sign in

Sign up

Sign in

Hardware Low-Level: CUDA Kernel Optimization from Scratch

Nova

Follow

--

Listen

Share

Reading notes for https://siboehm.com/articles/22/CUDA-MMM and

https://www.xiaohongshu.com/explore/66d6bf4a000000001d015511?app_platform=ios&app_version=8.55.8&share_from_user_hidden=true&xsec_source=app_share&type=normal&xsec_token=CBBbLBO8SY12Z54WFDil8dNRGqxlIGrO42iILYCU8gsKA=&author_share=1&xhsshare=CopyLink&shareRedId=Nz9ERDY2Ok86Ozg4PEo0Rzw7OUo3SkpB&apptime=1729026758&share_id=7eb4d12349894b4085d7b02ed3465cb0

Matrix multiplication is one of the most important algorithms in deep learning using GPUs, accounting for nearly all FLOPs. The author wrote a CUDA matrix multiplication kernel from scratch and progressively optimized it, eventually achieving performance close to cuBLAS. The goal is to understand the performance characteristics of modern GPUs through this process, including global memory access coalescing, shared memory caching, and utilization optimization.

Optimization Process

Initial Kernel

The simplest implementation involves each thread calculating a single element in matrix C. The performance of this method is low due to inefficient memory access. Performance: 309 GFLOPs, accounting for 1.3% of cuBLAS.

Global Memory Access Coalescing

This optimization step adjusts the data access pattern so that threads within the same warp can sequentially read data, achieving coalesced global memory accesses. Performance improves to 1986.5 GFLOPs, accounting for 8.5% of cuBLAS.

Shared Memory Cache

By using shared memory to cache blocks of matrices A and B in shared memory, the number of accesses to global memory is significantly reduced, which greatly decreases the overhead of memory accesses and thus improves computing efficiency. Performance reaches 2980.3 GFLOPs, accounting for 12.8% of cuBLAS.

1D Block Division

Each thread calculates multiple elements in matrix C, which reduces the frequency of accessing shared memory. The performance significantly improves to 8474.7 GFLOPs, accounting for 36.5% of cuBLAS.

2D Block Division

By further dividing the work, each thread calculates larger blocks of matrix C, increasing computational density. Performance reaches 15971.7 GFLOPs, accounting for 68.7% of cuBLAS.

Vectorized Memory Access

Vectorization is a method that leverages hardware parallelism by using data types with a width of 4 (e.g., float4) for memory read and write operations, processing more data per operation. This strategy greatly reduces the number of memory instructions, improving memory bandwidth utilization.

Through automatic tuning

To find the optimal parameters, performance further improves to 19721.0 GFLOPS, accounting for 84.8% of cuBLAS.

Warp blocking optimization

It increases parallelism and improves the locality of register caching. Performance reaches 21779.3 GFLOPS, accounting for 93.7% of cuBLAS.

Currently:

Using tensor memory accelerator (TMA) for memory I/O async operations and warp group MMA for MMA async operations.

Quoted from Nvidia’s Hopper webpage:

TMA allows applications to transfer 1D and up to 5D tensors between global memory and shared memory, in both directions, as well as between the shared memory regions of different SMs in the same cluster (refer to [Thread Block Clusters](https://docs.nvidia.com/cuda/hopper-tuning-guide/index.html#thread-block-clusters)). Additionally, for writes from shared memory to global memory, it allows specifying element wise reduction operations such as add/min/max as well as bitwise and/or for most common data types.

This has several advantages:

- Avoids using registers for moving data between the different memory spaces. - Avoids using SM instructions for moving data: a single thread can issue large data movement instructions to the TMA unit. The whole block can then continue working on other instructions while the data is in flight and only wait for the data to be consumed when actually necessary. - Enables users to write warp specialized codes, where specific warps specialize on data movement between the different memory spaces while other warps only work on local data within the SM.

Extra comments

- This implementation is not necessary in nowadays with tensor cores, when we have cuBLAS on Hopper.- Testing matrix is small, cuBL calls a relatively simple kernel while it has many more complex ones.- Kernel here used TN-layout where matrix A is transposed, but cuBLAS uses NN-layout, so the comparison may not be properly utilized.

--

--

Written by Nova

No responses yet

Help

Status

About

Careers

Press

Blog

Privacy

Terms

Text to speech

Teams