PyTorch demystified 3 -
Implementing and optimizing the CUDA backend

"Nobody wants just to go on. One wants to go on in an elegant way." — Alan Watts

Introduction

Now that we have built our basic infrastructure, and we have a full pipeline, it is time for our long awaited part where we get to speed things up. This part will consist of two subparts. In part one, this post, we implement a CUDA backend and optimize it. We will use NSys and NCU, two tools provided by NVIDIA, to profile our CUDA implementation, make deductions about bottlenecks, and do fixes.

The structure of this post is as follows: We first have a brief introduction into the main architecture and how we incorporate the CUDA backend. Because most of the CUDA implementation is straight forward I will outline how to do most, and we will exemplify with a few kernels that are a little more interesting. We will end up with a version of the backend that has a full CUDA backend and can be run end-to-end on the same MNIST instance we wrapped up in the previous post. For convenience I have made this a special release on my GitHub.

In part two we will optimize those kernels. Each optimization step is clearly marked by its headline, and we will see profiling outputs both before and after the fix, showing how much each step actually improved on our pipeline, and most interestingly to me personally: I give an explanation to each fix of why it actually worked. Additionally for those who want to follow along or maybe do even better fixes I have left every optimization step on a branch of the repository. You can find it here. I give a commit hash for every optimization step I perform below too.

In the end we will have improved our code execution on our benchmark from 2.399s down to 0.206s, a speedup with a factor of 11.6x. I note here that the initial naive CUDA version before I started to optimize took way longer on the benchmark, a whopping 997.691s. However, here I forgot to implement the optimizers in CUDA, resulting in unnecessary memcopies across the devices. So there is a speedup of 4800x, but that is noise that had to be debugged, rather than optimized. Anyway, we're gonna cook the full MNIST example from 25s training time down to 5 seconds (includes non-CUDA parts including data loading etc., parts that are not measured by the benchmark). And on larger models with more training data the gap is only to grow. That's a good prospect I say, so let's get started building something we can optimize.

Building a basic CUDA backend

We build our CUDA backend in two steps. At first we have to set up the basic infrastructure needed to include everything. Following that we implement the missing kernels step by step.

Setting up a basic structure

At first we do have to include CUDA code against our project. We have to tell CMake to find CUDA and where its libraries are. Additionally, because not everyone has a CUDA capable GPU or simply may not want to use it we give it a flag that is enabled by default and switches off automatically if CUDA could not be found on the machine. We do so by

# CMakeLists.txt

if (CUDA)  
  include(CheckLanguage)
  check_language(CUDA)

  if(CMAKE_CUDA_COMPILER)
    add_definitions(-D__CUDA)
    enable_language(CUDA)
    
    set(CMAKE_CUDA_STANDARD 20)
    set(CMAKE_CUDA_STANDARD_REQUIRED ON)
  else()
    message(WARNING "Could not find CUDA on system. Compiling without CUDA enabled")
  endif()
endif()

Note that we introduced a definition here as well through add_definitions(-D__CUDA). This will help us identify whether CUDA has been compile time enabled, and we will use it for compile time dependent branches in code. We have to make the CUDA headers accessible to the backend, and link CUDA against it. We are a little lazy here, so we go a little unorthodox in CMake and simply look for all .cpp files and all .cu files and append those if desired. We do get for the backend

file(GLOB_RECURSE CORE_SOURCES
    computational_graph/*.cpp
    data_modeling/*.cpp
    module/*.cpp
    system/*.cpp
    training/*.cpp
    utility/*.cpp
)

if(CMAKE_CUDA_COMPILER)
    file(GLOB_RECURSE CUDA_SOURCES
        computational_graph/*.cu
        data_modeling/*.cu
        module/*.cu
        system/*.cu
        training/*.cu
        utility/*.cu
    )
    list(APPEND CORE_SOURCES ${CUDA_SOURCES})
endif()

add_library(BackendCore SHARED ${CORE_SOURCES})

target_include_directories(BackendCore PUBLIC
    ${CMAKE_CURRENT_SOURCE_DIR}
)

set_target_properties(BackendCore PROPERTIES
    LIBRARY_OUTPUT_DIRECTORY "${PYTHON_MODULE_DIR}" # make sure Python-modules see backend
)

if(CMAKE_CUDA_COMPILER)
    set_target_properties(BackendCore PROPERTIES
        CUDA_SEPARABLE_COMPILATION ON
        CMAKE_CUDA_ARCHITECTURES native # we can target alternative architectures here
    )

    find_package(CUDAToolkit REQUIRED)
    target_include_directories(BackendCore PRIVATE 
        ${CUDAToolkit_INCLUDE_DIRS}
    )
    target_link_libraries(BackendCore 
        CUDA::cudart
    )
endif()

With that we now should be able to readily include our first CUDA files into our C++ backend. Remember the definition we introduced in CMake? We will now use it and include our first kernel. The casual scalar multiplication will serve as a base example illustrating the workflow. Our Tensor class will get the following enhancements:

// src/backend/data_modeling/tensor.cpp

#ifdef __CUDA
#include "utility/cuda/cuda_common.cuh"
#include "data_modeling/cuda/tensor_ops.cuh"
#endif

// ... 

Tensor Tensor::operator*(const ftype scalar) const {
  Tensor res(dims, values->getDevice(), requiresGrad);
  switch(values->getDevice()){
    case Device::CPU:
      for (tensorSize_t i = 0; i < values->getSize(); ++i) {
        res.values->data()[i] = values->data()[i] * scalar;
      }
      break;
    case Device::CUDA:
      #ifdef __CUDA
        cuda_impl::scalarmul(res, *this, scalar);
      #else
        __throw_runtime_error("Not compiled with CUDA");
      #endif
      break;
  }

  return res;
}

// ...

What we do here is check the device of the tensor we call the operation on, and either allocate the result on the GPU or in host RAM. A switch-case is either gonna call the CPU version or the CUDA version. Additionally, we guard against errors with a conditional compilation, giving us meaningful output when CUDA has been requested, but the code has not been compiled against. We will use this structure, i.e. the switch-case with the error output in virtually every operation where we intend to branch off into CUDA code.

We also have to store the CUDA kernels somewhere. In the example we did a conditional include that only uses the CUDA headers when the compile time flag has been enabled. For ease of architecture we write kernels exactly where they belong. That is, respective CUDA kernels are in a directory named cuda in the same directory where the CPU version resides in. CUDA kernels that can be shared by multiple modules of the backend are in a folder with path src/backend/shared/cuda. Therefore the directory structure looks like the following now:

.
├── src
│   ├── CMakeLists.txt
│   ├── backend
│   │   ├── computational_graph
│   │   │   ├── activation_functions
│   │   │   │   └── cuda
│   │   │   ├── loss_functions
│   │   │   │   └── cuda
│   │   │   ├── tensor_ops
│   │   │   │   └── cuda
│   │   ├── data_modeling
│   │   │   └── cuda
│   │   ├── module
│   │   │   ├── activation_functions
│   │   │   │   └── cuda
│   │   │   ├── layers
│   │   │   │   └── cuda
│   │   ├── shared
│   │   │   └── cuda
│   │   ├── system
│   │   ├── training
│   │   │   ├── loss_functions
│   │   │   │   └── cuda
│   │   └── utility
│   │       └── cuda
...

Note that I only show directories in the tree above. An exception to the rule is the cuda folder in utility, but I'd wager this also contains what you'd expect it to contain. Mainly some wrappers for better error messaging and some generic functions checking device capability. With that we can now bring CUDA alive in our backend.

CUDA memory management and custom kernels

As previously mentioned, we don't really have to work out every kernel for this blogpost. For us now it is mostly important to look at the main points, and point out some kernels that are of more interest to us. The previous section introduced the basic pattern we will apply through our CUDA implementations. But it did not show yet how we get CUDA tensors in the first place. Given our preliminary thoughts in part 1 of this series, this is rather a plug-and-play operation. There, we came up with a tensorValues_t struct, whose purpose was to manage the underlying data of the tensors. We gave it a resize operation that every tensor constructor called, and malloc'ed memory in there. All we have to do now is to give it the device and call cudaMalloc instead, as in

// src/backend/data_modeling/tensor.cpp

void Tensor::tensorValues_t::resize(const tensorSize_t size) {
  this->size = size;
  switch (device) {
    case Device::CPU:
      values = static_cast<ftype*>(std::malloc(this->size * sizeof(ftype)));
      break;
    case Device::CUDA:
      #ifdef __CUDA
        cudaErrchk(cudaMalloc((void**) &values, this->size * sizeof(ftype)));
      #else
        std::__throw_invalid_argument("Not compiled with CUDA.");
      #endif
      break;
  }
}

After this we also have to make sure that we update the constructor of tensorValues_t accordingly, as well as all the copy operations and set-device operations. The latter has two ways: If we are on the GPU and we move to the CPU we have to malloc, copy from the device to the host, and free memory on the GPU. If we start out on the CPU and move to the GPU we have to go the other way. No surprises there essentially. Lastly we update our getters and setters of the tensorValues_t to fetch elements from either the CPU or the GPU. And that already concludes our memory management.

Because we designed our frontend with this in mind, we do not have to update the Python binding and can simply instantiate a tensor now, telling it that it is a GPU resident as in t = Tensor.ones([2, 2], Device.CUDA), and we will have a tensor that lives on the GPU.

Adding a few kernels is just what you think it is. In the previous section we took a look at how CUDA kernels are called through a switch-case statement, here is the CUDA implementation of the kernel we showed there:

// src/backend/data_modeling/cuda/tensor_ops.cuh

#ifndef __CUDA
static_assert(false, "File should not be included without CUDA enabled"); // for safety reasons
#endif // __CUDA

// includes here...

class Tensor;

namespace cuda_impl {
  void scalarmul(Tensor& res, const Tensor& src, ftype scalar);
}

// src/backend/data_modeling/cuda/tensor_ops.cu

namespace { // anonymous namespace for all non-shared CUDA kernels in this project
  __global__ void scalarmulKernel(ftype* const res, const ftype* const left, ftype scalar, tensorSize_t size) {
    int gid = blockDim.x * blockIdx.x + threadIdx.x;
    if(gid >= size)
      return;
    
    res[gid] = left[gid] * scalar;
  }
}

namespace cuda_impl {
  void scalarmul(Tensor& res, const Tensor& src, ftype scalar) {
    constexpr int threadsPerBlock = 256;
    const int blocksPerGrid = (src.getSize() + threadsPerBlock - 1) / threadsPerBlock;

    scalarmulKernel<<<blocksPerGrid, threadsPerBlock>>>(res.getData(), src.getData(), scalar, src.getSize());
    cudaErrchk(cudaDeviceSynchronize());
  }
}

We have this modular design for clean code reasons, but also just because CUDA demands it. Essentially, when implementing CUDA kernels we have to use .cu file endings, to signal NVCC that it's time for it to shine. We could now either rename all our .cpp files to .cu, which would force every user to have a CUDA toolchain on their machine, or we do what I did here: Have separate .cu files with conditional compilation, and we cleanly separate CUDA code from pure C++ code. I note here that cudaMalloc and other utility functions given by CUDA are pure C++ code included through #include "cuda_runtime.h" and can reside in .cpp files.

I also note here that the custom kernel we just implemented above is what CUDA engineers like to call "embarrassingly parallel". In other words, pretty trivial to implement functionally. Most of the kernels introduced in this new version follow this pattern, making for a quick win. A few exceptions however stand out, and I will present them as an optional section below. However, one thing I want to note is that the kernels have not been written all too terrible to start with. For instance, I used hardware intrinsics when possible. For instance, inside the sigmoid kernel we compute a single element using this kernel:

// src/backend/shared/cuda/common_kernels.cuh

namespace cuda_impl { // in cuda_impl because it is shared
  template<typename T>
  __device__ __forceinline__ ftype cudaSigmoid(const ftype x) {
    if constexpr (std::is_same_v<T, float>) {
      ftype z = expf(-fabsf(x));
      ftype s = 1.0f / (1.0f + z);
      return (x >= 0.f) ? s : z * s; // x < 0 => e^x/(e^x+1) 
    } 
    else if constexpr (std::is_same_v<T, double>) {
      ftype z = exp(-abs(x));
      ftype s = 1.0 / (1.0 + z);
      return (x >= 0.0) ? s : z * s; // x < 0 => e^x/(e^x+1) 
    }
    else {
      static_assert(always_false<T>, "Unexpected value for ftype encountered");
    }
  }
}

The kernel lives in shared because it is used both in the forward pass and in the backward pass. A forced inline ensures that we do not call kernels from within a kernel, and a templated overload distinguishes whether to use single-precision-specific instructions (expf, fabsf) or generic instructions (exp, abs). The generic instructions take both single- and double-precision values, but we can play with other functions exposed through the CUDA API. While I did not verify this myself at this point, we can later revisit this and play with these functions to squeeze out performance. In theory we would expect the expf to be faster though, since it is explicitly designed for single-precision floating-point numbers. In general for floating points there is a trade-off between accuracy and speed. The larger our representation, the better our results, but the slower the computation itself. NVIDIA provides intrinsics for this that reduce the number of operations even further, at the cost of precision.

Optional: A brief history of floating points

▲ Collapse

This optional is a brief historical breakdown of floating point numbers and how they evolved over time. A short disclaimer: This is obviously for the younger generation of engineers, including spoiled millennials like me, and I have not personally lived through all those events I describe in the following. It is based on my personal research through the wild west the internet has become, and could be wrong. I am confident the information is correct, but if you find something wrong, please let me know.

Essentially, when representing numbers we distinguish, roughly, two kinds of number representations: integer values and floating point numbers (FPs). Normally we learn the integer representation first, since that is the easier one: binary representation, conversion between decimal, octal, hexadecimal, and binary format, addition and subtraction in the respective number systems, and signed representation, nowadays the 2's complement (enabling us to push both addition and subtraction through an adder, hence saving silicon). Pretty straight forward once you get the idea.

FPs on the other hand are not that easy. I myself have never had an introduction into actual circuitry doing basic arithmetic on FPs, and the reason is that even basic operations get much more complicated than on integer values. On top of that, floating points come with a much larger set of operations we want to perform on them. Integers are normally constrained to addition and multiplication plus their inverse, unless you want to do something more fancy. Floating points, on the other hand, add trigonometric functions (sine, cosine, tangent), root operations such as the square root, exponentials, logarithms, and any other hip operation you can come up with. This makes it inherently harder to find a single representation that works well across all these operations, and for all intended purposes.

The result was that in the beginning of computing (relative to our modern standpoint that is) vendors had different ways to represent floats, mostly disagreeing on how many bits to use for representation and how to distribute them (see figure below). Intel CPUs did it differently than IBM, which did it differently than Seymour Cray, etc., leading to problems such as incompatibilities between hardware of different vendors, different user-facing results on each system, and each representation having its own set of quirks. To tackle the challenge Intel brought in Canadian mathematician William Kahan, and secretly they began working on an 'ultimate FP representation'.

At the same time, the IEEE called for a unified standard, and Intel allowed Kahan to submit his work, which led to what is now called the IEEE-754 standard. It is a standard that generally puts up the scaffolding for how FPs are computed nowadays, which is why if you work on code that is a little closer to hardware, you've probably seen the name many times already. It starts out with how to represent numbers, see the figure below. In addition, it also provides rules for how to round, how to control rounding errors, and exception handling. The two special FP types you have probably seen already, NaN and Infinity, have their root in this standard. The standard also defines a positive and negative zero FP representation, indicating whether we approached zero from the positive or the negative side.

Intel on its side developed the Intel 8087 math co-processor in 1980, 5 years before IEEE-754 officially launched. The centerpiece of it was the x87 architecture, which uses a stack-based register file. Kahan determined that an 80-bit representation of floating-point numbers is enough to remove numerical errors that arise from chaining floating point operations, which is why x86-64 processors still support 80-bit wide FP operations today. However, while mathematically sound, 80-bit sits awkwardly in a computer, since it is not a natural power of two. Computers, however, tend to fetch from addresses aligned to a power of two, resulting in alignment problems outside of the floating point unit. Modern computers that need genuine higher precision than 80 bits bank on the fact that hardware got a whole lot more efficient than it used to be back in the day, and moved directly up to 128-bit. Practically it also turned out that for most of us 32-bit and in some cases 64-bit are more than sufficient, making the 80-bit format obsolete today. Intel used later capability upgrades of its processors, starting with SSE and cemented by SSE2, to softly phase out the x87 FP unit, which it still carries only for backward compatibility reasons.

However, the story of floating point numbers does not end here. For decades this has worked just fine. Whoever follows recent developments, however, will have noticed that we all of a sudden encountered a new wave of floating point representations. For instance, the C++ standard just got enhanced with a few new FP types. Most are not a surprise, but looking closer we find an std::bfloat16_t type. CUDA supports a range of smaller types, mainly FP4, FP6, and FP8, along with the brainfloat type that C++ also has. But how come?

The answer lies in the rise of modern LLMs. Researchers found that for large models shorter representations are enough to still yield good results, leading to the notion of quantization. Here, smaller data types have been used to perform the floating point operations in LLMs, reaching good results with higher compute throughput and cutting the memory footprint. Google's brainfloat (BFloat) takes this idea a step further. Given that precision (the mantissa) matters less, but the exponent range needed in LLMs can be wide, this data type shifts the bit budget accordingly. For instance, the IEEE-754 half-precision (FP16) format assigns 10 bits to the mantissa and 5 bits to the exponent (plus 1 sign bit). BFloat16 keeps the sign bit but instead spends 8 bits on the exponent and 7 on the mantissa, reducing the mantissa by three bits and increasing the exponent by the same number. These are, however, very recent developments and can change in the future, so it's worth keeping an eye out.

Other optimizations consistently included in our first wiring up of CUDA are classic warp-level reductions by defining a volatile pointer to the beginning of the array, and then scan through the array based on thread-id, see below (we will switch to modern warp shuffle instructions in an optimization step down below). I also used shared memory where applicable, and bit-shifts for divisions by 2, which happens pretty much all the time for parallel reduction kernels.

// src/backend/module/activation_functions/cuda/activations.cu

namespace { 
  // example of a warp-level reduction kernel - pattern applied consistently throughout code
  template<int maxoffset>
  __forceinline__ __device__ void softmaxWarpSumReduce(volatile ftype* const input, const tensorSize_t stride, const int offset) {
    static_assert(maxoffset > 0 && maxoffset <= 32, "Invalid value for template");
    
    if(maxoffset == 32) {
      if(offset + 32 < stride) input[offset] += input[offset + 32];
    }
    if(maxoffset >= 16) {
      if(offset + 16 < stride) input[offset] += input[offset + 16];
    }
    if(maxoffset >= 8) {
      if(offset + 8 < stride) input[offset] += input[offset + 8];
    }
    if(maxoffset >= 4) {
      if(offset + 4 < stride) input[offset] += input[offset + 4];
    }
    if(maxoffset >= 2) {
      if(offset + 2 < stride) input[offset] += input[offset + 2];
    }
    if(maxoffset >= 1) {
      if(offset + 1 < stride) input[offset] += input[offset + 1];
    }
  }
}

To the astute reader wanting to see optimization, arguably the more interesting part of this blogpost, I recommend to skip the next optional sections, covering two example kernels that were not that trivial, and that show my train of thought when I implemented the initial design. For those who are interested in seeing more what we will be dealing with before we optimize, the optional parts are for you.

Optimizing the CUDA backend

Finally we have arrived at one of my personal highlights of this whole project, the CUDA backend optimization. We will be going through the optimization step-by-step and explain it and why it works. For everyone who wants to follow along I kept every step on a separate branch which you can find on my GitHub.

There is a script I used for benchmarking, which may or may not have been updated in between optimization steps to reflect the new setting. You can find the script here. The base loop is depicted below, and for everyone who does want the code of the benchmark explained I do that in an optional box below.

Brief disclaimer: There was an out-of-bounds indexing bug I found later in the createContiguousCopyKernel in file tensor_ops.cu, and it has been phased out by optimization step 3. It shows up rarely and thus had been undiscovered. If you follow along and your code crashes you can fix it by copying the version at the end of the branch. However, the bug has no influence on the benchmarking and optimization results.

# optimization/mnist_cuda.py
if __name__ == "__main__":
    x, y_int = load_mnist()
    y = to_one_hot(y_int)

    x_train, x_val, y_train, y_val = train_test_split(
        x, y, test_size=0.1, random_state=42
    )

    # upload once; all batching and shuffling happens on the GPU
    x_train_gpu = to_gpu(x_train)
    y_train_gpu = to_gpu(y_train)
    x_val_gpu   = to_gpu(x_val)

    net = make_net()
    loss_fn = CrossEntropyWithSoftmax()
    optim = RmsProp(net.parameters(), 0.0001, 0.999)

    # warm up GPU
    BATCH_COUNT = 10
    train_epoch(net, loss_fn, optim, x_train_gpu, y_train_gpu, max_batch_count=BATCH_COUNT)

    times = []
    n_epochs = 5
    for epoch in range(n_epochs):
        
        start = time.perf_counter()
        train_loss = train_epoch(net, loss_fn, optim, x_train_gpu, y_train_gpu, max_batch_count=BATCH_COUNT)
        elapsed = time.perf_counter() - start
        print(f"Elapsed time GPU: {elapsed:.4f}s")

        times.append(elapsed)

    median_time = sorted(times)[len(times) // 2]
    print(f"Median time GPU: {median_time:.4f}s")

# optimization/mnist_cuda.py

def train_epoch(net, loss_fn, optim, x_gpu, y_gpu, batch_size=64, max_batch_count=50):
    n = x_gpu.dims[0]

    indices = list(range(n))
    random.shuffle(indices)
    x_shuf = x_gpu.slice(indices)
    y_shuf = y_gpu.slice(indices)

    total_loss = 0.0
    n_batches = 0
    max_batches = math.ceil(n / batch_size) # not used, my bad
    for start in range(0, n, batch_size):
        end = min(start + batch_size, n)
        xb = x_shuf.slice(start, end)
        yb = y_shuf.slice(start, end)

        pred = net.forward(xb)
        loss = loss_fn(yb, pred)
        loss.backward()

        optim.clipGradients(1.0)
        optim.step()
        optim.zeroGrad()

        total_loss += loss.getitem(0)
        n_batches += 1
        if n_batches == max_batch_count:
          return total_loss / n_batches

No surprises here, just a normal training loop. Lastly we benchmark through running this training loop five times, and taking the time, and we record the median. I note here that initially in the script I made a mistake of taking the last number that was put out. The numbers I give in this blogpost in the initial steps are thus not the median, but the last output, but that should be ok, as throughout this whole run of the benchmarks did not vary significantly in between the five training loops. Rerunning will very likely give me the same numbers with minimal deviations.

Opt. 1 - Remove unnecessary memcopies

This first version is benchmarked on a small example, because it was so slow to run. We will move on from that very quickly and get to larger examples, but for now we benchmark on a mere 5 epochs with 10 batches each. The batch-size remains 64 samples throughout this whole blog entry, as well as the next. We start off by running NSys to see the execution on a continuous timeline. I note that because NSys significantly slows down the run and generates huge amounts of data then needs to load (sometimes overwhelming itself) I ran NSys on a smaller example than the benchmark itself. A screenshot of this run is in the following:

The main information is that rather embarrassingly, 99.4% of the time we spend copying data, with about half of that being device-to-host copies. This kind of copy specifically should be negligible in this script, since we copy once from the host to the device, and afterwards the only real copies back should just be the training losses and network outputs at best. We can use NSys to go deeper into where the memcopy-calls are actually made and find that I forgot to map the gradient clipping onto the GPU. That is, each time an optimizer clips the gradients the CPU would execute the associated parts of the code.

Additionally, since we copy each element individually through auto g = (*grads)[i]; and grads->set((*grads)[i] * scale, i); we have created the worst possible way to copy to and from the device as well. The fix then is straight forward: Write a kernel for gradient clipping and incorporate it into the CUDA framework.

I personally struggle with calling this an optimization step, since it is more of a debugging. The code did not do what I expected it to do, and I diagnosed and fixed it. However, I already had my naive CUDA implementation as a release out there, and this serves as a stark example of how expensive seemingly innocuous memcopies can wreck your application, so I kept it in. The speedup is magnificent. In this example I used a reduced benchmark, but going to the final benchmark I used for my last optimization steps, this step has reduced the runtime from 997.691s down to a meager 2.399s. Can't really complain about that.

Opt. 2 - Speed up slicing by using non-blocking CUDA kernel instead

For everyone looking at the before-after numbers and wondering about the contradiction with the optimization step before, I have increased the workload for the benchmark now. More specifically, the number of batches we push through per epoch has been increased from 10 to 50 now, making the before number larger than the after from the previous example.

To start out with the optimization, we once again employ NSys, giving us the following timeline:

We find ourselves still copying a lot of data, but this time we spend the bulk of it on device-to-device copies. The culprit this time is the slicing operation. Looking into the Python code we use slicing to implement the shuffling operation at the beginning of each epoch. The code for that looks like this:

def train_epoch(net, loss_fn, optim, x_gpu, y_gpu, batch_size=64, max_batch_count=50):
    n = x_gpu.dims[0]

    indices = list(range(n))
    random.shuffle(indices)
    x_shuf = x_gpu.slice(indices)
    y_shuf = y_gpu.slice(indices)
    # ...

The slicing, at it stands before optimizing, looks like this:

// src/backend/data_modeling/tensor.cpp

Tensor Tensor::getSlice(span<const tensorDim_t> indices) const {
  // .. set up result

  values->copyValues(*res.values, indices, res.getDims().getStride(0));
  return res;
}

void Tensor::tensorValues_t::copyValues(tensorValues_t& target, span<const tensorDim_t> indices, 
                                        const tensorSize_t sizeOfDim) const {
  assert(target.size >= sizeOfDim * indices.size());

  switch(device){
    case Device::CPU:
      // ...
    case Device::CUDA:
      #ifdef __CUDA
        tensorSize_t targetOffset = 0;
        for(tensorDim_t idx: indices){
          tensorSize_t thisOffset = idx * sizeOfDim;
          copyValues(target, thisOffset, thisOffset+sizeOfDim, targetOffset);
          targetOffset += sizeOfDim;
        }
        cudaErrchk(cudaDeviceSynchronize());
      #else // the casual check for any CUDA call
        __throw_invalid_argument("Not compiled with CUDA");
      #endif
      break;
  }
}

void Tensor::tensorValues_t::copyValues(tensorValues_t& target, tensorSize_t low, 
                                        tensorSize_t high, tensorSize_t targetOffset) const {
  // safety checks here

  switch(device){
    case Device::CPU:
      std::memcpy(target.values+targetOffset, values+low, (high-low) * sizeof(ftype));
      break;
    case Device::CUDA:
      #ifdef __CUDA
        cudaErrchk(cudaMemcpy(target.values+targetOffset, values+low, (high-low) * sizeof(ftype), cudaMemcpyDeviceToDevice));
      #else
        __throw_runtime_error("Not compiled with CUDA");
      #endif
      break;
  }
}

The gist is the following: Initially, getSlice sets up the result, and the first copyValues-methods called determines a range to copy from for each index to slice over. This for-loop then calls subsequent memcopies. Pretty straight forward, but the astute reader can find two issues with this design.

The first problem is that memcopies are blocking. That is, if somewhere along the code I call a standard cudaMemcpy, then the memcopy is issued on the host, and the host code blocks until all data is on the device - there is an actual reason for this design I talk about optionally below. For now it is worth knowing that we lose momentum by issuing a series of subsequent memcopies, each one waiting for the previous one to finish. For better bus utilization we should focus on having less memcopies, but each one better transport a large chunk of neighboring memory cells. I.e. we'd rather have large coalesced memory accesses than many smaller ones where applicable.

The second problem is negligible compared with the first one, but still worth mentioning: We do have two overloads of copyValues here, and each one checks the device. Given that the device does not change in between the two we could just manually get the memcopies from the inner call into the outer call, saving a conditional check. For now we will only fix the first one though, as this is where we lose the bulk, and the CUDA fix naturally does away with this. What's left on the CPU side is work for later.

To fix this issue we have multiple ways. For instance, we could use multiple CUDA streams, hence overlap the memcopies via overlapping the memcopies. The runtime overhead then lies in context creation for each stream, plus the complexity we introduce in the source code itself.

There is however an easier solution that I chose instead. Since we already have a destination allocated on the GPU we can bypass the memcopies entirely and simply write a quick manual kernel copying the data for us. The overhead here is in thread creation, but once they exist each thread will do its own indexing, giving us both coalesced memory accesses where applicable, it bypasses the additional stream initialization for now, and it gives us an easy way to synchronize, since we can just launch it on the default stream asynchronously. The kernel is depicted below.

// src/backend/data_modeling/cuda/tensor_ops.cu

namespace cuda_impl {
  void getSlice(Tensor& res, const Tensor& src, span<const tensorDim_t> idx) {
    constexpr int threadsPerBlock = 256;
    const int blocks = (res.getSize() + threadsPerBlock - 1) / threadsPerBlock;

    tensorDim_t* idx_d;
    cudaErrchk(cudaMalloc(&idx_d, idx.size() * sizeof(tensorDim_t)));
    const auto sizeOfDim = res.getDims().getStride(0);

    cudaErrchk(cudaMemcpy(idx_d, idx.data(), idx.size() * sizeof(tensorDim_t), cudaMemcpyHostToDevice));

    getSliceKernel<<<blocks, threadsPerBlock>>>(res.getData(), src.getData(), idx_d, sizeOfDim, res.getSize());
    cudaErrchk(cudaDeviceSynchronize());
    cudaErrchk(cudaFree(idx_d));
  }
}

Opt. 3 - Speed up matmul in the backward block through kernel fusion

As per usual, we start by profiling through NSys. The following timeline shows that we finally moved into kernel space:

We find ourselves 60% of the time running kernels, and out of that more than 70% in the createContiguousCopyKernel. The question is where this one gets called the most, but since this is a new kernel not explained yet I give an optional explanation below. If you are familiar with it you can skip this section and move on to the next one.

// src/backend/data_modeling/dim_type.h

class Dimension final {
  using dim_t = std::array<tensorSize_t, MAX_NDIMS>;
  struct shallowCopyToken {};

  private:
    // those two indicate the structure of the contiguous data that lies underneath
    std::shared_ptr< std::vector<tensorDim_t> > contiguousDims;
    std::shared_ptr<dim_t> contiguousStrides;

    std::vector<tensorDim_t> dims;
    dim_t strides;
    
    // ...

  public:

    // ...

    bool inOriginalState() const noexcept {
      return *contiguousStrides == strides;
    }

    const auto getStrides() const noexcept { return strides; }
    const auto getContiguousStrides() const noexcept { return contiguousStrides; }
    void makeContiguous();

    // ...
}

By digging deeper we find that the kernel is extensively called in the backward-pass of the matrix multiplication operation. Looking at the node again we find the following operation:

// src/backend/computational_graph/tensor_ops/matmul_node.cpp

vector<shared_ptr<Tensor>> MatMulNode::backward(const Tensor& upstreamGrad) {
    assert(!upstreamGrad.getRequiresGrad());
    
    return {
        make_shared<Tensor>(upstreamGrad.matmul(*parents[1])), 
        make_shared<Tensor>(parents[0]->matmul(upstreamGrad))
    };
}

I talked about this method earlier in part 2a in an optional block, but if you allow me to quote myself, "Given two matrices $A$ and $B$, and computing their matrix multiplication $C = A @ B$, we obtain the following for the derivatives: $\frac{\partial C}{\partial A} = C @ B^T$ and $\frac{\partial C}{\partial B} = A @ C^T$". The code on the other side does the following: It first transposes one of the two tensors, which induces a copy with a reshuffling of the underlying data, then multiplies. But since we need the transposition how to solve that one?

Kernel fusion to the rescue: We need to do the transposition in any case, but what we can do is bake it directly into our matrix multiplication, where we will have to think about our indexing so that we hit the right elements. Luckily this is also rather straight forward, all we have to do is virtually swap rows and columns inside the kernel, but the logic remains the same. To further avoid decision making at runtime we template our kernel, and the transposition information will be evaluated at compile time, leaning on modern constexpr-features. The kernel looks like the following after the refactor:

// src/backend/data_modeling/cuda/tensor_ops.cu

namespace {
  __global__ void matMul2DKernel(ftype* const res, 
                                 const ftype* const left, const ftype* const right,
                                 const tensorDim_t leftRows, const tensorDim_t rightRows, 
                                 const tensorDim_t leftCols, const tensorDim_t rightCols, 
                                 const tensorSize_t resSize)
  {
    const int tid = threadIdx.x;
    const int gid = blockDim.x * blockIdx.x + tid;
    if(gid >= resSize) 
      return;

    const int K = transposeLeft ? leftRows :  leftCols;
    const int N = transposeRight ? rightRows : rightCols;

    const int i = gid / N;
    const int j = gid % N;

    // C[i, j] = sum_{k=0}^{leftCols} A[i, k] * B[k, j]
    ftype cij = 0;
    for (int k = 0; k < K; k++) {

      int leftIdx = transposeLeft ? k * leftCols + i : i * leftCols + k;
      int rightIdx = transposeRight ? j * rightCols + k : k * rightCols + j;
            
      cij += left[leftIdx] * right[rightIdx];
    }

    res[gid] = cij; // smem[tid];
  }
}

Important to note here is that there's a decision to be made on whether the input parameters, i.e. the number of rows and columns, represent the physical or the logical layout, with the physical being the actual alignment in memory and the logical the alignment as it would virtually exist after transposing the respective tensor. I preferred being lazy here and chose that they represent the physical layout, for no other reason as that being easier for me to imagine and reason about.

The basic difference with the previous, unoptimized kernel is then the introduction of K and N, representing the virtual layout then, and finally moving over that in the inner loop. The loop decides whether it wants to move over rows or columns then, i.e. it does the index swap. All we have to do then is to update the call to matmul through

// src/backend/computational_graph/tensor_ops/matmul_node.cpp

vector<shared_ptr<Tensor>> MatMulNode::backward(const Tensor& upstreamGrad) {
    assert(!upstreamGrad.getRequiresGrad());
    // TODO: optimize operators
    return {
        make_shared<Tensor>(upstreamGrad.matmul(*parents[1], false, true)), 
        make_shared<Tensor>(parents[0]->matmul(upstreamGrad, true, false))
    };
}

The signature for matmul has been adjusted to pass the flags down into the actual matmul implementation as in Tensor matmul(const Tensor& other, const bool transposeLeft, const bool transposeRight) const;, and the code decides dynamically which template instance of the kernel to call.

We will see this kernel again later and optimize it further, and then the same for the CPU later as well, so do not feel obliged to dig into all the knitty-gritties of that one if you missed some detail. Getting the idea and why this optimization step works is enough for now.

Opt. 4 - Memcopies revisited: On-device initialization and memory pool

This optimization step is certainly one of my favorite ones due to the simplicity and elegance of the solutions. Because we are already in a really fast area, at least on this small dataset, we disabled the number of batches now, and an epoch trains on the whole dataset from now on. However, since we are already very quick in general we will see more and more negligible improvements only. Looking at NSys again we once find ourselves running lots of memcopy-operations, shown below. So this is another chance for us to look at memory management.

This step will, unlike the other ones, be split up in two parts. The first part won't have an impact on the training dataset, but will alter the NSys output, as it will reduce those host-to-device copies. The second one won't change the memcopies, hence not alter the NSys output, but it is something that has been burning under my fingers for a while now, and now I finally have a chance to make it shine.

Opt. 4.1 - On-device initialization

This first part is pretty straight forward. When initializing weights for networks, so far the code went the following route: First, it would allocate an array on the CPU, populate it with values, and then copy that whole chunk of memory onto the GPU. This is happening in the constructor of the layer module we implemented earlier, as in

// src/backend/module/layers/ff_layer.cpp

FfLayer::FfLayer(tensorDim_t inSize, tensorDim_t outSize, Device d, 
    bool useBias, bool requiresGrad, shared_ptr<InitializerBase> init)
  : requiresGrad{requiresGrad} 
{
  if(!init){
    init = make_shared<NormalXavierInitializer>(inSize, outSize);  
  }

  weights = make_shared<Tensor>(Dimension({inSize, outSize}), d, requiresGrad);
  weights->reset(init);
    
  if(useBias){
    bias = make_shared<Tensor>(vector<tensorDim_t>{outSize}, d, requiresGrad);
    TensorFunctions::ToZeros(*bias);
  }
}

The fix is then to initialize the weights directly on the GPU. CUDA offers CuRand as a ready solution for that, so I won't go into detail here. The initializers in their final forms for now can be found in the src/backend/shared directory.

Opt. 4.2 - Use a memory pool for tensors

With that done we finally find ourselves only doing necessary memcopies from host to device and back, set apart a few not-so-obvious cases. For now we can consider that done though. However, there is another part of the code that I know stings, and we're gonna take it down. Let's take a brief look at the construction and destruction of the tensorValues_t:

// src/backend/data_modeling/tensor.cpp

// note: I removed the CUDA safety checks here for legibility 

Tensor::tensorValues_t::~tensorValues_t() noexcept {
  switch(device){
    case Device::CPU:
      free(values);
      break;
    case Device::CUDA:
      cudaErrchk(cudaFree(values));
      break;
  }
}

// called by Tensor-class constructor
void Tensor::tensorValues_t::resize(const tensorSize_t size) {
  this->size = size;
  switch (device) {
    case Device::CPU:
      values = static_cast<ftype*>(std::malloc(this->size * sizeof(ftype)));
      break;
    case Device::CUDA:
      
      cudaErrchk(cudaMalloc((void**) &values, this->size * sizeof(ftype)));
      break;
  }
}

We see that both construction and destruction are fairly simple: They encounter a malloc for a new tensor, and when it goes down, it frees that space again. The problem is that many of those allocations and free-operations are unnecessary. Say we have a training loop. Then in each forward pass a chain of tensors is constructed in the background. All these tensors have their lifetime limited to a maximum of a forward-backward pass, and get destroyed and re-created during the next pass at the very maximum achievable lifetime. Then, in the next forward pass we create tensors with exactly the same shapes and dimensions all over again. And the same goes for the backward pass.

The problem with this is at hand. But so is the solution: This is an ideal case for a memory pool. Because sizes and dimensions are the same in every forward- and backward pass, and when requesting memory we only really care about the full size of the tensor, we can simply request and free based on the size. I give the memory pool with its main two functions in the following:

// src/backend/shared/memory_pool.h

// for legibility I removed the CUDA safety checks below 

namespace mempool_impl {
  template<typename T>
  class MemoryPool final {
    private:
      // hold pointers to unused allocated memory.
      std::unordered_map<tensorSize_t, std::vector<T*>> freeLists;
      std::unordered_map<tensorSize_t, std::vector<T*>> freeListsCuda;
      
    public:
      MemoryPool() = default;
      ~MemoryPool() noexcept;

      T* request(Device d, tensorSize_t n);
      void giveback(T* ptr, Device d, tensorSize_t n);
      void flush(Device d) noexcept;
  };

  template<typename T>
  T* MemoryPool<T>::request(const Device d, const tensorSize_t n) {
    switch(d) {
      case Device::CPU:
      {
        // CPU version mirrors CUDA version
      }
      case Device::CUDA:
      {
        auto& list = freeListsCuda[n];
        if(!list.empty()) {
          T* ptr = list.back();
          list.pop_back();
          return ptr;
        }

        T* ptr;
        auto err = cudaMalloc((void**) &ptr, n * sizeof(T));
        if(err != cudaSuccess) {
          // ran out of memory, free up cached memory and retry
          flush(Device::CUDA);
          cudaErrchk(cudaMalloc((void**) &ptr, n * sizeof(T)));
        }

        return ptr;
      }
    }
  }

  template<typename T>
  void MemoryPool<T>::giveback(T* ptr, const Device d, const tensorSize_t n) {
    switch(d) {
      case Device::CPU:
        freeLists[n].push_back(ptr); 
        break;
      case Device::CUDA:
        freeListsCuda[n].push_back(ptr); 
        break;
    }
  }
}

As we can see, we see... nothing much; well, not entirely true. I omitted a few methods in the snippet above, since they should be obvious. The interesting one is request, and giveback does the inverse, which is much simpler. Request essentially first checks for what device we request a memory allocation. The CPU version here is commented away, because it mirrors the GPU version nearly one-by-one, just with different API calls.

After requesting the device we can then perform a (cuda-)malloc. If nothing comes back we ran out of memory, in which case we have two choices. Either we let the program crash, or we decide to give up memory and continue. The latter case gets the program running, but could do so very inefficiently, since we run the risk of recycling the same memory all over again in a loop. I decided for the latter and went brute force for now, flushing the entire reserved memory the pool holds. There are surely smoother ways out there, but that's again future work.

I also note that the template is not per se necessary. One could simply run one memory pool and demand the size to allocate in bytes instead. That way the responsibility to request the right amount of size is pushed back onto the programmer. This is almost purely a matter of taste, and I went with the simpler API, that in turn demands that for other data types we need an extra pool. I have identified two pools so far, and I do not think that the amount grows very much at all from this point on, so this is it for now. We then can instantiate the pool via inline static mempool_impl::MemoryPool<ftype> tensorPool;, which introduces a rather subtle bug that I fix in the next post. For now we can rewrite our earlier tensor-creation and destruction behavior as in the next code segment and be done with this part of the optimization.

// src/backend/data_modeling/tensor.cpp

Tensor::tensorValues_t::~tensorValues_t() noexcept {
  if(values != nullptr) {
    mempool::tensorPool.giveback(values, device, size);
  }
}

void Tensor::tensorValues_t::resize(const tensorSize_t size) {
  this->size = size;

  if(values != nullptr) [[unlikely]] {
    mempool::tensorPool.giveback(values, device, size);
  }
  values = mempool::tensorPool.request(device, size);
}

Opt. 5 - Fuse matmul kernels into a single kernel call

This one is very interesting. Technically it should have improved things, but on the practical side things stagnated. Before I get into that let's take a look at the nsys output. I increased the amount of computation in the script we run to profile, therefore the amount of computation has heavily lifted from the initial offloading of data onto the GPU and toward more computation. In other words, we have a lot more computation now, and much less memcopies. We find ourselves with the following output:

As can be seen here, the vast majority of the CUDA side is happening in the matmul-kernel, so that is the one we are going to take a look at. On top of that, looking at the timeline as it happens, it is also clear why our improvements are getting smaller and smaller. The spikes that are happening in the row of the kernel calls, those are our training loop. Visually they occupy only roughly 30% of our entire timeline we are working on in our benchmarks. Given how the rest of the code is already that large compared with that segment we can see how fast we already are, since it is vastly outnumbered by a small series of Python calls surrounding everything in the beginning and the end of the whole script.

Anyway, in a previous post we already laid hands on the matmul, introducing a templated kernel that overloads in transposition of either the left or the right kernel. NSys now tells us that one of those is dominating all the others, more clearly to be seen in the command-line via an nsys stat my_file.nsys-rep, parroting back at us

 Time (%)  Total Time (ns)  Instances  Avg (ns)   Med (ns)   Min (ns)  Max (ns)   StdDev (ns)                                                  Name                                                
 --------  ---------------  ---------  ---------  ---------  --------  ---------  -----------  ----------------------------------------------------------------------------------------------------
     65.2      909,975,555      9,000  101,108.4   50,561.0       960    255,717    108,457.5  void ::matMul2DKernel<(bool)0, (bool)1>(float *, const float *, const float *,unsigned in…
     14.5      202,607,249      9,000   22,511.9    9,856.0     5,760     90,050     20,763.3  void ::matMul2DKernel<(bool)0, (bool)0>(float *, const float *, const float *, unsigned in…
     10.0      139,692,572      9,000   15,521.4    8,096.0     2,944     36,321     14,293.5  void ::matMul2DKernel<(bool)1, (bool)0>(float *, const float *, const float *, unsigned in…
      1.8       24,791,568     18,000    1,377.3      704.0       640      6,752      1,321.4  ::stepRmsPropKernel(float *, float *, const float *, float, float, unsigned int)
      1.2       17,160,733     30,009      571.9      480.0       448      1,345        236.7  void cub::_V_300304_SM_750::detail::transform::transform_kernel::elementwiseaddKernel(float *, const float *, const float *, unsigned int)

Interestingly, the NSys-UI (NSight) and its command line output disagree on the real big kernels by a few percentage points, but not on the overarching picture. The smaller kernels do fit well again. We take it as is and see that the majority of the kernel space is used by the matmul templated overload transposing the right matrix, leaving the left one intact. That is the most counterintuitive for us, since manually written we can create the best memory access pattern out of the box for that case. Interesting is also that this one has the fastest execution in the best case, which is to be expected, but the worst in the worst case.

Running NCU it tells us that this is due to partial waves, see the figure below. This means that at some point we run a wave of blocks that keep some streaming multiprocessors (SMs) empty, hence wasting compute resources that sit idle. The fix would be to tune the number of blocks, ideally in a way that they are a multiple of the SMs, and distribute the threads over them.

However, we don't always have to follow the profiler when we think there's a better way. I already hinted earlier that there are some inefficiencies in the matmul-kernels, and before proceeding to what the compiler tells us we can do away with the obvious first. Otherwise we start tuning a kernel that we will phase out anyway, and our work was for naught. To show the current problem again, here's the kernel call.

// src/backend/data_modeling/cuda/tensor_ops.cu

namespace cuda_impl {
  void matmul(Tensor& res, const Tensor& left, const Tensor& right) {
    constexpr int threadsPerBlock = 256;
    const int blocksPerGrid = (res.getSize() + threadsPerBlock - 1) / threadsPerBlock;
    
    // sizes of the 2D matrices respectively
    const tensorSize_t leftSize = left.getDims().get(-1) * left.getDims().get(-2); 
    const tensorSize_t rightSize = right.getDims().get(-1) * right.getDims().get(-2);
    const tensorSize_t resSize = left.getDims().get(-2) * right.getDims().get(-1);

    tensorSize_t leftOffset = 0;
    tensorSize_t rightOffset = 0;
    tensorSize_t resOffset = 0;

    while(leftOffset < left.getSize()){
      if(!(transposeLeft || transposeRight)) {
        matMul2DKernel<false, false><<<blocksPerGrid, threadsPerBlock>>>(res.getData() + resOffset, left.getData() + leftOffset, right.getData() + rightOffset,
                                       left.getDims().get(-2), right.getDims().get(-2), left.getDims().get(-1), right.getDims().get(-1), resSize);
      }
      else if(transposeLeft && transposeRight) [[unlikely]] {
        // another kernel call
      }
      // ... cover other template overloads

      leftOffset += leftSize;
      rightOffset += rightSize;
      resOffset += resSize;
    }

    cudaErrchk(cudaDeviceSynchronize());
  }
}

As shown here, the host code separates the batch of matrices into segments of 2D-matrices, and then multiplies them individually in a loop. We do not use CUDA streams here, meaning that instead of processing an entire batch of 2D-matrices in parallel, we take each of them and multiply them in a sequential fashion. Needless to say that this design is horrible at best. Thus, in this step we remove the outer while loop. Instead, we will use the CUDA utilities and use the additional y- and z- coordinates CUDA hands us for indexing threads and blocks. First, we have to compute how many of those multiplications we aim to perform:

const tensorSize_t nMultiplications = res.getSize() / resSize;
const int blocks = (resSize + threadsPerBlock - 1) / threadsPerBlock;
dim3 numBlocks(blocks, nMultiplications);

These will be our kernel call parameters. Next, we have to change our kernel itself to index into the right array. We give the kernel itself the size parameters that we computed earlier, and we get the offsets inside the kernels through using the block-ids, and we can then simply index inside the kernel.

const int leftOffset = blockIdx.y * leftSize;
const int rightOffset = blockIdx.y * rightSize;

// offsets computed outside of kernel before
cij += left[leftOffset + leftIdx] * right[rightOffset + rightIdx];

// same for indexing into result

This step clearly marks a trade-off. We had a series of sequential calls to the kernels, which we replaced by one single kernel call. The kernel itself however became larger by adding more input arguments, but also slightly more computations. This can have additional side effects, such as increasing register pressure, if they happened to be at the limit already. Asking NSys and NCU however gives us a rather sobering answer: Our fix didn't seem to have any impact on anything at all. However, we also did not regress and have some more elegant code in the end, so why not? For reference, below is the new NSys report before we move on to fix 6.

 Time (%)  Total Time (ns)  Instances  Avg (ns)   Med (ns)   Min (ns)  Max (ns)   StdDev (ns)                                                  Name                                                
 --------  ---------------  ---------  ---------  ---------  --------  ---------  -----------  ----------------------------------------------------------------------------------------------------
     65.3      909,392,179      9,000  101,043.6   50,529.0       960    256,197    108,385.1  void ::matMul2DKernel<(bool)0, (bool)1>(float *, const float *, const float *, unsigned in…
     14.5      202,405,897      9,000   22,489.5    9,856.0     5,728     89,986     20,742.1  void ::matMul2DKernel<(bool)0, (bool)0>(float *, const float *, const float *, unsigned in…
     10.0      139,579,312      9,000   15,508.8    8,096.0     2,912     36,385     14,277.1  void ::matMul2DKernel<(bool)1, (bool)0>(float *, const float *, const float *, unsigned in…

We can also take a look at the NSys output to compare with before and see that our fix did little here. For everyone who wants to go deeper into this fix I provide an optional section below. Everyone else can skip to the next section, as we never planned to hover here for too long and will refine the kernel further anyway.

To be on the safe side that we made no mistake we quickly take a look at the matmul kernel itself in isolation. We employ Google benchmark to write up a quick benchmark, and for completeness reasons we test the matmul in multiple settings. Because we are interested in comparing the launch overhead and blocking with the fused kernel we can focus on one instance of the matmul kernel, rather than testing every template overload. It is to be noted that one dry run should be performed on the device to remove overhead that comes from CUDA's JIT behavior and context initialization. Google Benchmarks returns us these numbers before the fix:

----------------------------------------------------------------------------------------
Benchmark                              Time             CPU   Iterations UserCounters...
----------------------------------------------------------------------------------------
BM_MatMul_CUDA/64/64/64             8.35 us         8.35 us        82625 GFLOP/s=62.8207/s
BM_MatMul_CUDA/256/256/256          55.3 us         55.3 us        12941 GFLOP/s=606.583/s
BM_MatMul_CUDA/512/512/512           444 us          444 us         1539 GFLOP/s=605.077/s
BM_MatMul_CUDA/1024/1024/1024       3738 us         3738 us          191 GFLOP/s=574.47/s
BM_MatMul_CUDA/64/784/256           57.1 us         57.1 us        11573 GFLOP/s=449.763/s
BM_MatMul_CUDA/64/256/128           15.2 us         15.2 us        45986 GFLOP/s=275.334/s
BM_MatMul_CUDA/64/128/10            11.0 us         11.0 us        63442 GFLOP/s=14.8342/s

And this is the output after the fix:

----------------------------------------------------------------------------------------
Benchmark                              Time             CPU   Iterations UserCounters...
----------------------------------------------------------------------------------------
BM_MatMul_CUDA/64/64/64             8.23 us         8.23 us        83068 GFLOP/s=63.7151/s
BM_MatMul_CUDA/256/256/256          54.8 us         54.8 us        12880 GFLOP/s=611.837/s
BM_MatMul_CUDA/512/512/512           440 us          440 us         1546 GFLOP/s=610.399/s
BM_MatMul_CUDA/1024/1024/1024       3714 us         3714 us          193 GFLOP/s=578.278/s
BM_MatMul_CUDA/64/784/256           57.0 us         56.9 us        11628 GFLOP/s=451.117/s
BM_MatMul_CUDA/64/256/128           15.2 us         15.2 us        46025 GFLOP/s=276.244/s
BM_MatMul_CUDA/64/128/10            11.0 us         11.0 us        63601 GFLOP/s=14.8537/s

Now there's good and bad news in here. The good news is that perhaps our fix indeed did do something. After all the improved kernels are slightly better across most inputs. Moreover, the larger the batch dimension, that is the larger the number of loops the pre-fix version iterates over, the larger the improvement. The bad news is that the fix is in terms of around 20 microseconds even for a batch-size of 1024, which is far beyond what our actual MNIST benchmark has. The fix thus only ever gets really visible on extremely large inputs and only matters to those who want every last drop of performance.

Good to know, we wonder if we could have spent our time better than this. But I personally like this version much better, as the code itself is just more elegant. And why not, maybe those microseconds will matter at some point. And as a small bit of extra information, the degradation in performance we measured on the benchmark is likely noise rather than a genuine regression. The next fix however is going to change a lot for our kernel.

Opt. 6 - Tiled matmul kernel

Because we just ran NSys to get the NSight visualization we do not need to do that again. We also do not need to, since we know that the matMul2DKernel absorbs all of our computational efforts. We can thus go directly into NCU and get a direct profile of our kernel. We call ncu --set full -o [output-file] python3 [benchmark-file]. Looking into the analysis and the recommendations we find from NCU we find two things: 1. It recommends us to tune our kernel again, specifically the launch grid configurations, and 2. it has a problem with our memory access pattern.

Because the grid launch configurations are a matter of both kernel and its input we will focus on the more generalizable fix, namely the memory access pattern. We can see that in general it has an issue with both the load and the store into the L1 cache. Given our kernel it is pretty obvious where that is coming from: The memory access patterns from our two source tensors, and the stores back into the resulting tensor.

We once again have two ways to solve this: We can do what I have been doing in an earlier iteration already, that is to swap loops. I did this in part 1 already, resulting in flat addressing. However, this will also only go that far. Our memory access pattern will be better, resulting in more coalesced fetches. But once a row or column, whatever we iterate through at the moment, becomes larger than our caches, we are again loading the same data over and over again into cache.

There is a better solution though, one very much worth knowing, and that is the tiled matmul. While this technique is not limited to the matmul operation, the matmul operation is the one where it is normally first encountered. The basic idea is the following: Say we have two matrices, both of size $\mathbb{R}^N$. The full matrix multiplication then costs $\mathcal{O}(N^3)$ operations (if one does not consider Strassen's algorithm). The normal matmul as one would do it by hand, walking through the rows of the left matrix and the columns of the right matrix, would load each datapoint anew from memory.

The idea of the tiled approach is then to instead of loading one row and one column for one element, to load a block (a tile) of the left matrix, a block of the right matrix, and then multiply those to get a tile of the result. If the tilesize is tuned correctly this method results in no wasted memory fetches, and since the multiplication runs on level 1 cache memory (both on CPU and GPU, the CPU version comes in part 4 of this series), we get amazing speed out of the memory hierarchy. Because this is a larger and interesting topic I personally think that it is worth it looking a little bit closer. Given that this is large enough to become a complete blogpost on its own, and also given that those already exist, I softly nudge the interested reader toward this one. From my side you get a ready kernel, but let's start with the kernel call:

// src/backend/data_modeling/cuda/tensor_ops.cu

namespace cuda_impl {
  void matmul(Tensor& res, const Tensor& left, const Tensor& right, const bool transposeLeft, const bool transposeRight) {
    constexpr int MATMUL_TILESIZE = 32; // choose 16 (threadsPerBlock=256) or 32 (threadsPerBlock=1024)
    constexpr dim3 threadsPerBlock(MATMUL_TILESIZE, MATMUL_TILESIZE);
    
    // sizes of the 2D matrices respectively
    const tensorSize_t leftSize = left.getDims().get(-1) * left.getDims().get(-2); 
    const tensorSize_t rightSize = right.getDims().get(-1) * right.getDims().get(-2);
    const tensorSize_t resSize = res.getDims().get(-2) * res.getDims().get(-1);

    const tensorSize_t nMultiplications = res.getSize() / resSize;

    // x = horizontal, y = vertical
    const int blocksX = (res.getDims().get(-1) + MATMUL_TILESIZE - 1) / MATMUL_TILESIZE;
    const int blocksY = (res.getDims().get(-2) + MATMUL_TILESIZE - 1) / MATMUL_TILESIZE;
    dim3 numBlocks(blocksX, blocksY, nMultiplications);

    if(!(transposeLeft || transposeRight)) {
      matMul2DKernel<MATMUL_TILESIZE * MATMUL_TILESIZE, false, false><<<numBlocks, threadsPerBlock>>>(
                                                      res.getData(), left.getData(), right.getData(), 
                                                      left.getDims().get(-2), left.getDims().get(-1), 
                                                      right.getDims().get(-2), right.getDims().get(-1), 
                                                      res.getDims().get(-2), res.getDims().get(-1),
                                                      leftSize, rightSize, resSize);
    }
    // .. call the other template overloads

    cudaErrchk(cudaDeviceSynchronize());
  }
}

At first we choose the size of a tile. We make the tiles quadratic for now, and sizes are either 16 or 32. I chose 32, since that is the warpsize, enhancing throughput on first glance. We will see later if it was a good choice. We make the threads be indexed two-dimensional. nMultiplications encodes how many 2D-matrices we compute later on, and we call $\left (\left \lceil \frac{nCols_{res}}{tilesize} \right \rceil, \left \lceil \frac{nRows_{res}}{tilesize} \right \rceil, nMultiplications \right)$ blocks. Each of these blocks has $(tilesize, tilesize)$ threads, and each thread computes one element of the result by sliding the tile across the two matrices. This is shown in the kernel below.

// src/backend/data_modeling/cuda/tensor_ops.cu

namespace {
  template<int tileSize, bool transposeLeft, bool transposeRight>
  __global__ void matMul2DKernel(ftype* const res, 
                                 const ftype* const left, const ftype* const right,
                                 const tensorDim_t leftRows, const tensorDim_t leftCols, 
                                 const tensorDim_t rightRows, const tensorDim_t rightCols,
                                 const tensorDim_t resRows, const tensorDim_t resCols,
                                 const tensorSize_t leftSize, const tensorSize_t rightSize, const tensorSize_t resSize)
  {
    __shared__ ftype smemA[tileSize];
    __shared__ ftype smemB[tileSize];

    const tensorSize_t K = transposeLeft ? leftRows : leftCols;
    const tensorSize_t N = transposeRight ? rightRows : rightCols;

    const tensorSize_t i = blockIdx.y * blockDim.y + threadIdx.y;
    const tensorSize_t j = blockIdx.x * blockDim.x + threadIdx.x;

    const tensorSize_t leftOffset = blockIdx.z * leftSize;
    const tensorSize_t rightOffset = blockIdx.z * rightSize;
    const tensorSize_t resOffset = blockIdx.z * resSize;

    ftype cij = 0;
    for (tensorSize_t k = 0; k < K; k += blockDim.x) {
      // load tile into smem
      const tensorSize_t leftIdx = transposeLeft ? (k + threadIdx.x) * leftCols + i : i * leftCols + k + threadIdx.x;
      const tensorSize_t rightIdx = transposeRight ? j * rightCols + (k + threadIdx.y) : (k + threadIdx.y) * rightCols + j;

      smemA[threadIdx.y * blockDim.x + threadIdx.x] = (i < resRows && (k + threadIdx.x) < K) ? left[leftOffset + leftIdx] : 0.0f;
      smemB[threadIdx.y * blockDim.x + threadIdx.x] = ((k + threadIdx.y) < K && j < resCols) ? right[rightOffset + rightIdx] : 0.0f;
      __syncthreads();

      for (int l = 0; l < blockDim.x; l++) {
        cij += smemA[threadIdx.y * blockDim.x + l] * smemB[l * blockDim.x + threadIdx.x];
      }
      __syncthreads();
    }

    if (i < resRows && j < resCols) {
      res[resOffset + i * N + j] = cij;
    }
  }
}

Now it pays off having read the previous iterations we performed on the matmul kernel, since now everything coalesces into this one, along with the tiles residing in shared memory. Kernels initially determine which element they are computing, which 2D-matrix they are in, and resolve the template overload. We then enter a loop iterating over blockDim.x. This works because blockDim.x corresponds with the tilesize. We could have put this into the loop, making it a compile time expression and saving a little overhead.

Inside the loop we first load the elements into shared memory. We need to be careful here to not cross the boundaries of the outer matrices. Values above the boundary get zero padded, one operation. If we moved the boundary checks into the inner loop that follows we would cast this one boundary check into blockDim.x boundary checks, hence the zero-padding is another micro-optimization. The inner loop then simply sums over the respective elements in shared memory, making this the actual matmul. I note here that these are coalesced accesses. Fetching from smemA is a broadcast-read, since all threads of a warp access the same memory location. And the fetches from smemB access through threadIdx.x, making neighboring threads access neighboring cells, one memory bank at a time, since we neatly organized everything in 32 32-bit floating points.

In summary, we squeezed everything into shared memory, where we reduced latency through memory accesses, and we achieved better memory access patterns.

Opt. 7 - Remove unnecessary device syncs

This will be a rather quick one again. Whenever we called a CUDA kernel, we have followed the pattern laid out by every CUDA tutorial: 1. Compute launch parameters, 2. call kernel, 3. synchronize the device, and 4. cleanup if necessary. However, on this whole project we never used an alternative stream, meaning two subsequent CUDA kernel calls will happen in a sequential fashion, even when kernel calls run asynchronously to the host code.

We can therefore remove any unnecessary cudaDeviceSynchronize() call. The host code will run through, schedule the CUDA code, and block as soon as it reaches a blocking statement, which is in between training steps, as sometimes we copy dimensionality information onto the device, but it would be latest at the end of any training loop, when we capture the loss or the host needs data from the device for its own decision making.

The fix is then pretty straight forward. Showing the improvement in a profiler would be harder to capture. We can show that we spend less time on the synchronize, or spend less time in user space. I chose to do two Nsight profiles, which we can see below. One version before the fix, one version after the fix. I capture the first 10ms of kernel calls in both cases, and we can see the kernel calls are denser on the version after the fix.

Opt. 8 - Use modern warp shuffle intrinsics in reduction kernels

This last one has been burning for a while, and if it's only for the simple reason that I grew tired of the deprecated warnings of the compiler on how we used the warp unrolled reduction pattern in our reduction kernels. That is the pattern we used in our reduction kernels, specifically exemplified by the following kernel, which I copy here again:

// src/backend/module/activation_functions/cuda/activations.cu

namespace { 
  template<int maxoffset>
  __forceinline__ __device__ void softmaxWarpSumReduce(volatile ftype* const input, const tensorSize_t stride, const int offset) {
    static_assert(maxoffset > 0 && maxoffset <= 32, "Invalid value for template");
    
    if(maxoffset == 32) {
      if(offset + 32 < stride) input[offset] += input[offset + 32];
    }
    if(maxoffset >= 16) {
      if(offset + 16 < stride) input[offset] += input[offset + 16];
    }
    if(maxoffset >= 8) {
      if(offset + 8 < stride) input[offset] += input[offset + 8];
    }
    if(maxoffset >= 4) {
      if(offset + 4 < stride) input[offset] += input[offset + 4];
    }
    if(maxoffset >= 2) {
      if(offset + 2 < stride) input[offset] += input[offset + 2];
    }
    if(maxoffset >= 1) {
      if(offset + 1 < stride) input[offset] += input[offset + 1];
    }
  }
}

This is normally applied to shared memory, i.e. input is normally a pointer to an address in device specific shared memory. The reason for this is to suppress optimizations by the kernel, making sure that each warp writes back into shared memory. That allows us to remove the synchronization instructions, which is what we actually save through this optimization.

However, given that this is such a common pattern NVIDIA stepped in and gave us warp shuffles in later iterations of its hardware. Warp shuffles are essentially a means for threads in a warp to cooperate further by sharing values stored in thread-local registers with other threads of the same warp. There are multiple instructions of how to do that, for instance cross-pattern, enabling neighboring threads to share, shuffle ups, where each thread gives its register content to a thread with a higher index. We obviously take the inverse route, meaning we shuffle down. I used the shuffle down with synchronization instructions, since we still need to synchronize in between threads for the reduction. The kernel above now looks like what I give below, and we apply this pattern to all reduction kernels.

// src/backend/shared/cuda/common_softmax.cuh

// baked directly into the kernel, bypassing the device kernel completely
for(int offset = maxoffset; offset > 0; offset >>= 1) {
  maxVal = cudaMax<ftype>(maxVal, __shfl_down_sync(0xFFFFFFFF, maxVal, offset, stride)); 
}

I went for the looped version here, one can easily unroll that loop by hand, saving even more, or simply use a #pragma unroll before the loop to tell the compiler to do it for you. In my implementation referenced by the commit hash I did not unroll the loops, which is where the benchmark number comes from. Just out of curiosity I also created a branch out of the commit to check if the loop unroll does something, and the number went down from 0.3325s without the loop unrolling to 0.3318s with loop unrolling. Given that our kernels are mostly short and registers are likely free to use the unrolled version is likely the better one in most cases, but the three numbers I got out of this experiment here are in the milliseconds and not very clean numbers. So what actually is noise and what is a genuine improvement is hard to measure here, unless we isolate our kernels and do some hard benchmarking. I for my part am pretty happy knowing at the very least it does not move the needle for now, but we stay canonical and got rid of the compiler complaining to us that we're using phone booths and carrier pigeons. And as a side, I just included the unrolled loops on my main branch.

Comparison with cuBLAS and outlook

That wraps up our optimization journey so far. While there is still a lot that can be done we have to lower the gun at some point and call it a day, or else we will never get to eat our hard-earned dinner today. As a last step, and just out of curiosity, I wanted to check out how much we can still squeeze out of it. Running NSight again we can see that we still lose the bulk in the matmul kernel, a well known problem by the way. Matrix multiplications are just expensive operations, and given that there are so many involved in every deep learning application including LLMs there is a lot of research out there just optimizing this one single operation.

Anyway, since we lose most of our fuel in that operation I decided to quickly compare it with a baseline. cuBLAS has ready made implementations of common operations used in linear algebra, and is provided by NVIDIA and hence optimized by NVIDIA engineers. I ran the benchmarks implemented in Google-Benchmark earlier (found here), parroting back the following numbers at me for the cuBLAS version:

----------------------------------------------------------------------------------------
Benchmark                              Time             CPU   Iterations UserCounters...
----------------------------------------------------------------------------------------
BM_MatMul_CUDA/64/64/64             12.3 us         12.3 us        57092 GFLOP/s=42.5191/s
BM_MatMul_CUDA/256/256/256          23.4 us         23.4 us        31298 GFLOP/s=1.4355k/s
BM_MatMul_CUDA/512/512/512          66.1 us         66.1 us        10415 GFLOP/s=4.05948k/s
BM_MatMul_CUDA/1024/1024/1024        400 us          400 us         1696 GFLOP/s=5.3725k/s
BM_MatMul_CUDA/64/784/256           19.6 us         19.6 us        34714 GFLOP/s=1.31137k/s
BM_MatMul_CUDA/64/256/128           12.3 us         12.3 us        55978 GFLOP/s=340.338/s
BM_MatMul_CUDA/64/128/10            12.3 us         12.3 us        56551 GFLOP/s=13.3226/s

Doing the same with our custom kernel does the following now:

----------------------------------------------------------------------------------------
Benchmark                              Time             CPU   Iterations UserCounters...
----------------------------------------------------------------------------------------
BM_MatMul_CUDA/64/64/64             7.51 us         7.50 us        93948 GFLOP/s=69.8632/s
BM_MatMul_CUDA/256/256/256          56.9 us         56.9 us        12746 GFLOP/s=589.361/s
BM_MatMul_CUDA/512/512/512           408 us          408 us         1713 GFLOP/s=657.182/s
BM_MatMul_CUDA/1024/1024/1024       3293 us         3293 us          218 GFLOP/s=652.158/s
BM_MatMul_CUDA/64/784/256           56.8 us         56.8 us        11748 GFLOP/s=452.521/s
BM_MatMul_CUDA/64/256/128           15.3 us         15.3 us        45443 GFLOP/s=273.79/s
BM_MatMul_CUDA/64/128/10            9.10 us         9.10 us        76300 GFLOP/s=18.0097/s

Interestingly, on smaller input sizes our custom kernel wins, but gets its booty kicked really quickly by cuBLAS on rising input sizes. The question is why. Sadly cuBLAS is not open-source, hence we cannot really look into what it does differently. But given that the library call has no immediate kernel calls and hence not really kernel launch parameters we give it chances are that there is some checking going on beneath the hood on input sizes. It will likely also check our hardware and call subsequent kernels. Hence, step 1 for us to go further would also be to tune the kernel launch parameters accordingly, which was also one of the top NCU recommendations to follow next. Other optimizations we could perform are the use of dedicated hardware, since CUDA now has TensorCores, dedicated hardware specifically designed for deep learning applications.

Other optimizations that apply to the whole framework would be to use optimized libraries where applicable, but we can do more than that. The sigmoid function screams at us to use compile time optimizations, where we can compile a precomputed table at compile time. The runtime interface then reduces to a mere lookup in the table. Other things we can do are kernel fusions, for instance a matmul and add for the feedforward layer matmul and add bias term. This saves us further fetches from memory and one kernel launch overhead.

Another interesting direction to look at is mixed precision floating points. In the beginning of this blogpost I talked about floating point numbers and mentioned that recently more floating point types have been introduced to both CUDA and C++. This increases throughput by both saving memory, but more importantly, enabling larger memory and compute throughput. For instance, CUDA supports intrinsics where one thread can perform one single operation on two FP-16 operands at the same time, similar to what Intel did with AVX. That means that one warp could have a throughput of 64 operands per cycle rather than 32 during those operations.

However, those optimizations will be part of a later blogpost, shall I find the time for it. For now it is as it is, but if you feel motivated to go deeper into the matter feel free to pick up from where I left off, and if you find good optimizations, please let me in as well.

Summary

In the previous three blogposts leading up to this one we implemented a deep learning framework akin to PyTorch in C++. In contrast, this part of the series was all about CUDA. We implemented a CUDA backend, ending up with an end-to-end training pipeline in CUDA. We then used NSys and NCU to profile both our timeline and individual kernels, optimizing our CUDA backend with a speedup of roughly 11x on our whole training-pipeline as is . To do so we used techniques such as kernel fusion, custom kernels to bypass and streamline memory-copies, a memory pool, and tiled kernels for cache optimal behavior. In the end we compared with cuBLAS, where we still got beaten, but we also now know how to continue in order to compete with it.

My next and last post of this series will go a similar path. With a satisfactory CUDA backend, at least for now, we will optimize the C++ side as well. We will use similar techniques and once again spend some time with the matmul operation, but we will also use other optimization techniques and find and fix some bugs that we silently introduced. This one will hopefully be very educational as well for you, as it surely was for me. Until then, and don't forget to reach out or give feedback if you engaged with this post. Cheers.