Part 2: CUDA Implementation

Fig. 24 Inline and lookaside in the gNB DU processing pipeline. The figure is from [Kundu2023B]. Note that the DGX Spark and Jetson AGX Thor platform use a unified memory architecture, which can be seen as a hybrid of the inline and lookaside memory management.

The second part of this tutorial focuses on the implementation of the LDPC decoder using CUDA and explains the common pitfalls when offloading compute-intensive functions to the GPU. Further, we show how GPU acceleration offloading of the LDPC decoding can be integrated into the OAI stack.

Function offloading to dedicated accelerators requires careful consideration of the data flow with respect to the underlying memory architecture. This is especially critical for wireless communications applications, where strict latency requirements in the receiver processing pipeline necessitate efficient memory management when utilizing dedicated signal processing accelerators.

As shown in Fig. 24, one needs to distinguish between inline and lookaside acceleration. While lookaside acceleration moves data between the CPU and the hardware accelerator, inline acceleration avoids such data movement by applying the entire processing pipeline on the hardware accelerator.

Strictly speaking, DGX Spark still moves data between the CPU and GPU. However, the overhead is significantly reduced compared to traditional split-memory architectures as DGX Spark shares the same physical memory between CPU and GPU. This has implications for the caching behavior and requires a careful implementation of the CUDA kernels to avoid performance degradation. Nevertheless, we will consider it as inline acceleration for the following discussion.

For further details on CUDA, we refer to the NVIDIA CUDA Programming Guide [CUDA2024] and for the Jetson platform to the CUDA for Tegra Memory Model [Tegra2024].

Overview

Fig. 25 Overview of CUDA implementation of the LDPC BP decoding algorithm.

The CUDA implementation can be found in plugins/ldpc_cuda/src/runtime/ldpc_decoder.cu. The core decoding algorithm is implemented in the update_cn_kernel(.) and update_vn_kernel(.) functions. Both kernels are iteratively called and perform the check node (CN) and variable node (VN) updates, respectively. The decoder stops when the maximum number of iterations is reached. An additional early stopping condition could also be applied to reduce the average number of iterations.

The pack_bits_kernel(.) kernel maps the soft-values to hard-decided bits and packs them into a more compact byte-representation which is required for the OAI processing pipeline.

CUDA Integration in OAI

The LDPC decoder is implemented as shared library that can be loaded using the OAI shared library loader. The implementation is in plugins/ldpc_cuda/src/runtime/ldpc_decoder.cu. After modifying it, rebuild the Docker images:

./scripts/build-oai-images.sh

Running the Decoder

The CUDA-based decoder can be used as a drop-in replacement for the existing decoder implementations. It can be loaded when running the gNB via the following GNB_EXTRA_OPTIONS in the .env file of the config folder.

GNB_EXTRA_OPTIONS=--loader.ldpc.shlibversion _cuda

We strongly recommend to additionally assign dedicated CPU cores to PHY-layer processing via the thread-pool option. This assigns the cores 5-9 to the PHY layer thread pool. Note that the lower CPU cores are assigned to handle the USRP-related tasks such as time synchronization.

You can now start the gNB with the CUDA-based decoder by running

scripts/start_system.sh rfsim # replace with b200 or other config folder

The GPU load can be monitored via

nvidia-smi  # or jtop on Jetson devices

Congratulations! You have now successfully accelerated the LDPC decoder using CUDA.

Check the gNB logs for the CUDA decoder:

docker logs oai-gnb

You should see the CUDA decoder being used (requires iperf3 traffic to trigger the decoder):

> ...
> [NR_PHY] I {L1_rx_thread} CUDA LDPC decoder:    123.55 us (    96.54 us / seg)
> ...

Implementation Aspects

The following sections focus on various technical aspects of the CUDA implementation and the performance implications of different memory transfer patterns.

For debugging and profiling, please refer to the tutorial on Debugging & Troubleshooting.

Memory Management

In order to offload compute-intensive processing to the GPU, data needs to be shared between the host (CPU) and the device (GPU). We can leverage the shared system memory architecture of DGX Spark and Jetson platforms to avoid the bottleneck of costly memory transfers on traditional split-memory platforms.

In fact, we can avoid the overhead of any complex API calls and memory transition operations by allocating page-locked memory on the host using cudaHostAlloc. To make input LLRs visible to the GPU, we then use a simple memcpy() operation to copy inputs from the 5G stack over into such a page-locked buffer, where unified memory addressing allows direct reads and writes both on the host and the device. For output bits, we first synchronize the parallel CUDA command stream and then simply use memcpy() to copy packed bits from the shared page-locked buffer into the 5G stack output buffer.

This simplicity is enabled by the unified memory architecture. On DGX Spark (Grace Blackwell), hardware cache coherency via NVLink-C2C ensures efficient data sharing. On Jetson platforms, the cache semantics [Tegra2024] require host caches to be active while device memory caching is disabled on page-locked buffers. For optimal performance on both platforms, we ensure maximally coalesced read-once and write-once memory access patterns (addressing consecutive memory in consecutive threads).

Traditional memory transfer patterns designed for split-memory architectures can also be used on DGX Spark, but can lead to higher overheads depending on the exact architecture and generation. Explicit memory transfer calls such as cudaMemcpyAsync(inputs..., cudaMemcpyHostToDevice, stream) and cudaMemcpyAsync(outputs..., cudaMemcpyDeviceToHost, stream) incur additional API overheads and may depend on availability of additional copy engines; explicit transitioning of managed pageable memory allocated by cudaMallocManaged() by cudaStreamAttachMemAsync() may also require excessive API overhead for small buffer sizes. Therefore, we instead use the patterns described above to optimize for shared system memory and low latency, particularly since input and output buffers are typically small in the real-time 5G stack.

Note

On the DGX Spark architecture, cudaMallocManaged() can be even more efficient for some use cases. For an example implementation, see the 5G NR PUSCH Neural Receiver tutorial.

For comparison, we show both variants side-by-side in the following code, where the latency-optimized code path is the one with USE_UNIFIED_MEMORY defined. We copy the input LLRs from the host to the device memory in the ldpc_decoder.cu file:

    // copy input data to device-visible memory
#ifdef USE_UNIFIED_MEMORY
    memcpy(const_cast<int8_t*>(mapped_llr_in), llr_in, num_llrs * sizeof(*llr_in));
#else
    CHECK_CUDA(cudaMemcpyAsync(const_cast<int8_t*>(mapped_llr_in), llr_in, num_llrs * sizeof(*llr_in), cudaMemcpyHostToDevice, stream));
#endif
    // END marker-copy-input

#ifdef PRINT_TIMES
    clock_gettime( TIMESTAMP_CLOCK_SOURCE, &ts_cursor );

    time_ns = ts_cursor.tv_nsec - ts_begin.tv_nsec + 1000000000ll * (ts_cursor.tv_sec - ts_begin.tv_sec);
    message_count = add_measurement(&input_copy_time, time_ns, 500);
    if (message_count % 500 == 0) {
      time_ns = input_copy_time.avg_ns;
      printf("Input copy time: %llu us %llu ns\n", time_ns / 1000, time_ns - time_ns / 1000 * 1000);
      fflush(stdout);
    }
#endif

    int8_t const* llr_total = mapped_llr_in;

    for (uint32_t i = 0; i < num_iter; ++i) {
        dim3 threads(NODE_KERNEL_BLOCK, UNROLL_NODES);

        // check node update
        dim3 blocks_cn(blocks_for(bg.num_rows * Z, threads.x));
        // note: llr_msg not not read, only written to in first iteration; will be filled with outputs of this function
        update_cn_kernel<<<blocks_cn, threads, 0, stream>>>(
            llr_total, context.llr_msg_buffer,
            Z, bg.cn, bg.cn_degree, bg.cn_stride, bg.num_rows, i==0);

        // variable node update
        dim3 blocks_vn(blocks_for(bg.num_cols * Z, threads.x));
        // note: llr_total only written to
        update_vn_kernel<<<blocks_vn, threads, 0, stream>>>(
            context.llr_msg_buffer, mapped_llr_in, context.llr_total_buffer,
            Z, bg.vn, bg.vn_degree, bg.vn_stride, bg.num_cols, bg.num_rows);
        llr_total = context.llr_total_buffer;
    }

    uint8_t *mapped_llr_bits_out = context.llr_bits_out_buffer;

    // pack bits
    dim3 threads_pack(PACK_BITS_KERNEL_THREADS);
    dim3 blocks_pack(blocks_for(block_length, threads_pack.x));
    pack_bits_kernel<<<blocks_pack, threads_pack, 0, stream>>>(
        llr_total, mapped_llr_bits_out, block_length);
#ifndef USE_UNIFIED_MEMORY
    CHECK_CUDA(cudaMemcpyAsync(llr_bits, mapped_llr_bits_out, num_out_bytes, cudaMemcpyDeviceToHost, stream));
#endif

    // allow CPU access of output bits while computing syndrome
#if defined(USE_UNIFIED_MEMORY) && !defined(USE_GRAPHS)
    cudaStreamSynchronize(stream);
#endif

#ifdef PRINT_TIMES
    clock_gettime( TIMESTAMP_CLOCK_SOURCE, &ts_cursor );
#endif

    // check syndrome if additional testing is requested
    if (perform_syndrome_check) {
        dim3 threads(512);
        dim3 blocks_cn(blocks_for(num_cn, threads.x));
        compute_syndrome_kernel<<<blocks_cn, threads, 0, stream>>>(
            context.llr_total_buffer, context.syndrome_buffer,
            Z, bg.cn, bg.cn_degree, bg.cn_stride, bg.num_rows);
#ifndef USE_UNIFIED_MEMORY
      CHECK_CUDA(cudaMemcpyAsync(context.host_syndrome_buffer, context.syndrome_buffer, num_cn * sizeof(*context.syndrome_buffer) / 32, cudaMemcpyDeviceToHost, stream));
#endif
    }

#ifdef USE_GRAPHS
    cudaGraph_t graphUpdate = {};
    CHECK_CUDA(cudaStreamEndCapture(stream, &graphUpdate));
    if (context.graphCtx) {
        cudaGraphNode_t errorNode;
        cudaGraphExecUpdateResult updateResult;
        CHECK_CUDA(cudaGraphExecUpdate(context.graphCtx, graphUpdate, &errorNode, &updateResult));
    }
    else
         CHECK_CUDA(cudaGraphInstantiate(&context.graphCtx, graphUpdate, 0));
    cudaGraphDestroy(graphUpdate);
    CHECK_CUDA(cudaGraphLaunch(context.graphCtx, stream));
#endif

#if !defined(USE_UNIFIED_MEMORY) || defined(USE_GRAPHS)
    // allow CPU access of output bits and syndrome
    cudaStreamSynchronize(stream);
#endif

#ifdef USE_UNIFIED_MEMORY
    // note: GPU synchronized before async syndrome check
    memcpy(llr_bits, mapped_llr_bits_out, num_out_bytes);
#endif

#ifdef PRINT_TIMES
    clock_gettime( TIMESTAMP_CLOCK_SOURCE, &ts_end );

    time_ns = ts_end.tv_nsec - ts_cursor.tv_nsec + 1000000000ll * (ts_end.tv_sec - ts_cursor.tv_sec);
    message_count = add_measurement(&output_copy_time, time_ns, 500);
    if (message_count % 500 == 0) {
      time_ns = output_copy_time.avg_ns;
      printf("Output copy time: %llu us %llu ns\n", time_ns / 1000, time_ns - time_ns / 1000 * 1000);
      fflush(stdout);
    }
#endif

    if (perform_syndrome_check) {
      uint32_t* p_syndrome;
#ifdef USE_UNIFIED_MEMORY
    #ifndef USE_GRAPHS
      // allow reading syndrome
      cudaStreamSynchronize(stream);
    #endif
      p_syndrome = context.syndrome_buffer;
#else
      // note: already synchronized above
      p_syndrome = context.host_syndrome_buffer;
#endif

      // check any errors indicated by syndrome
      for (uint32_t i = 0; i < num_cn / 32; i++) {
          if (p_syndrome[i] != 0) {
              return num_iter+1;
          }
      }
    }

#ifdef PRINT_TIMES
    clock_gettime( TIMESTAMP_CLOCK_SOURCE, &ts_end );

    time_ns = ts_end.tv_nsec - ts_begin.tv_nsec + 1000000000ll * (ts_end.tv_sec - ts_begin.tv_sec);
    message_count = add_measurement(&decoding_time, time_ns, 500);
    if (message_count % 500 == 0) {
      time_ns = decoding_time.avg_ns;
      printf("CUDA sync runtime: %llu us %llu ns\n", time_ns / 1000, time_ns - time_ns / 1000 * 1000);
      fflush(stdout);
    }
#endif

    return num_iter-1; // note: now index of successful iteration
}

ThreadContext& ldpc_decoder_init_context(int make_stream) {
    auto& context = thread_context;
    if (context.llr_in_buffer) // lazy
        return context;

    printf("Initializing LDPC context (TID %d)\n", (int) gettid());

    if (make_stream) {
        int highPriority = 0;
        if (cudaDeviceGetStreamPriorityRange(NULL, &highPriority))
            printf("CUDA stream priorities unsupported, %s:%d", __FILE__, __LINE__);
        CHECK_CUDA(cudaStreamCreateWithPriority(&context.stream, cudaStreamNonBlocking, highPriority));

        cudaStreamAttrValue attr = {};
        attr.syncPolicy = cudaSyncPolicyYield;
        cudaStreamSetAttribute(context.stream, cudaStreamAttributeSynchronizationPolicy, &attr);
    }

    CHECK_CUDA(cudaMallocStaging(&context.llr_in_buffer, MAX_BG_COLS * MAX_Z * sizeof(int8_t), cudaHostAllocMapped | cudaHostAllocWriteCombined));
    CHECK_CUDA(cudaMallocStaging(&context.llr_bits_out_buffer, (MAX_BLOCK_LENGTH + 7) / 8 * sizeof(uint8_t), cudaHostAllocMapped));
    CHECK_CUDA(cudaMallocStaging(&context.syndrome_buffer, MAX_BG_ROWS * MAX_Z * sizeof(uint32_t) / 32, cudaHostAllocMapped));
#ifndef USE_UNIFIED_MEMORY
    context.host_syndrome_buffer = (uint8_t*) malloc(MAX_BG_ROWS * MAX_Z * sizeof(uint8_t));
#endif
    CHECK_CUDA(cudaMalloc(&context.llr_msg_buffer, MAX_BG_ROWS * MAX_BG_COLS * MAX_Z * sizeof(llr_msg_t)));
    CHECK_CUDA(cudaMalloc(&context.llr_total_buffer, MAX_BG_COLS * MAX_Z * sizeof(llr_accumulator_t)));

    // keep track of active thread contexts for shutdown
    ThreadContext* self = &context;
    __atomic_exchange(&initialized_thread_contexts, &self, &self->next_initialized_context, __ATOMIC_ACQ_REL);

    return context;
}

extern "C" ThreadContext* ldpc_decoder_init(int make_stream) {
    if (bg_cn[0][0])  // lazy, global
        return &ldpc_decoder_init_context(make_stream);

    printf("Initializing LDPC runtime %d\n", (int) gettid());

    const uint32_t* table_bg_cn_degree[2][8] = { { BG1_CN_DEGREE_TABLE() }, { BG2_CN_DEGREE_TABLE() } };
    const uint32_t* table_bg_vn_degree[2][8] = { { BG1_VN_DEGREE_TABLE() }, { BG2_VN_DEGREE_TABLE() } };
    const uint32_t table_bg_cn_degree_size[2][8] = { { BG1_CN_DEGREE_TABLE(sizeof) }, { BG2_CN_DEGREE_TABLE(sizeof) } };
    const uint32_t table_bg_vn_degree_size[2][8] = { { BG1_VN_DEGREE_TABLE(sizeof) }, { BG2_VN_DEGREE_TABLE(sizeof) } };
    const void* table_bg_cn[2][8] = { { BG1_CN_TABLE() }, { BG2_CN_TABLE() } };
    const void* table_bg_vn[2][8] = { { BG1_VN_TABLE() }, { BG2_VN_TABLE() } };
    const uint32_t table_bg_cn_size[2][8] = { { BG1_CN_TABLE(sizeof) }, { BG2_CN_TABLE(sizeof) } };
    const uint32_t table_bg_vn_size[2][8] = { { BG1_VN_TABLE(sizeof) }, { BG2_VN_TABLE(sizeof) } };

    for (int b = 0; b < 2; ++b) {
        for (int ils = 0; ils < 8; ++ils) {
            CHECK_CUDA(cudaMalloc(&bg_cn_degree[b][ils], table_bg_cn_degree_size[b][ils]));
            CHECK_CUDA(cudaMemcpy(const_cast<uint32_t*>(bg_cn_degree[b][ils]), table_bg_cn_degree[b][ils], table_bg_cn_degree_size[b][ils], cudaMemcpyHostToDevice));
            CHECK_CUDA(cudaMalloc(&bg_vn_degree[b][ils], table_bg_vn_degree_size[b][ils]));
            CHECK_CUDA(cudaMemcpy(const_cast<uint32_t*>(bg_vn_degree[b][ils]), table_bg_vn_degree[b][ils], table_bg_vn_degree_size[b][ils], cudaMemcpyHostToDevice));

            CHECK_CUDA(cudaMalloc(&bg_cn[b][ils], table_bg_cn_size[b][ils]));
            CHECK_CUDA(cudaMemcpy(const_cast<uint32_t*>(bg_cn[b][ils]), table_bg_cn[b][ils], table_bg_cn_size[b][ils], cudaMemcpyHostToDevice));
            CHECK_CUDA(cudaMalloc(&bg_vn[b][ils], table_bg_vn_size[b][ils]));
            CHECK_CUDA(cudaMemcpy(const_cast<uint32_t*>(bg_vn[b][ils]), table_bg_vn[b][ils], table_bg_vn_size[b][ils], cudaMemcpyHostToDevice));

            bg_cn_size[b][ils] = table_bg_cn_size[b][ils];
            bg_vn_size[b][ils] = table_bg_vn_size[b][ils];
        }
    }

    return &ldpc_decoder_init_context(make_stream);
}

extern "C" void ldpc_decoder_shutdown() {
    cudaDeviceSynchronize();

    ThreadContext* active_context = nullptr;
    __atomic_exchange(&initialized_thread_contexts, &active_context, &active_context, __ATOMIC_ACQ_REL);
    while (active_context) {
        cudaFreeStaging(active_context->llr_in_buffer);
        cudaFree(active_context->llr_msg_buffer);
        cudaFreeStaging(active_context->llr_bits_out_buffer);
        cudaFree(active_context->llr_total_buffer);
        cudaFreeStaging(active_context->syndrome_buffer);
#ifndef USE_UNIFIED_MEMORY
        free(active_context->host_syndrome_buffer);
#endif
        if (active_context->stream)
            cudaStreamDestroy(active_context->stream);

#ifdef USE_GRAPHS
        cudaGraphExecDestroy(active_context->graphCtx);
#endif

        active_context = active_context->next_initialized_context;
    }

    for (int b = 0; b < 2; ++b) {
        for (int ils = 0; ils < 8; ++ils) {
            cudaFree(&bg_cn_degree[b][ils]);
            cudaFree(&bg_vn_degree[b][ils]);
            cudaFree(&bg_cn[b][ils]);
            cudaFree(&bg_vn[b][ils]);
        }
    }
}

#ifdef ENABLE_NANOBIND

#include <nanobind/nanobind.h>
#include <nanobind/ndarray.h>

namespace nb = nanobind;

NB_MODULE(ldpc_decoder, m) {
    m.def("decode", [](uint32_t BG, uint32_t Z,
                       const nb::ndarray<int8_t, nb::shape<-1>, nb::device::cpu>& llrs,
                       uint32_t block_length, uint32_t num_iter) {
        auto* context = ldpc_decoder_init(1); // lazy

        size_t num_bytes = (block_length + 7) / 8 * 8;
        uint8_t *data = new uint8_t[num_bytes];
        memset(data, 0, num_bytes);
        nb::capsule owner(data, [](void *p) noexcept { delete[] (uint8_t*) p; });

        ldpc_decode(context, 0, BG, Z, llrs.data(),
                    block_length, data,
                    num_iter, true);

        return nb::ndarray<nb::numpy, uint8_t>(data, {num_bytes}, owner);
    });
}

#endif

    // copy input LLRs from the host to the device memory
    cudaMemcpyAsync(context.llr_in_buffer, p_in, memorySize_llr_in, cudaMemcpyHostToDevice, context.stream);

After decoding, we make the output bits available via a host-side copy in parallel to an asynchronous syndrome check on the device:

    // pack LDPC output bits on the device
    int pack_thread_blocks = (block_length + 127) / 128;
    pack_decoded_bit_kernel<<<pack_thread_blocks, 128, 0, context.stream>>>(context.dev_llr_out, context.dev_bits_out, block_length);

   // allow CPU access of output bits while computing syndrome
#ifdef USE_UNIFIED_MEMORY
    cudaStreamSynchronize(context.stream);

    // ... schedule syndrome or CRC check ...

    // while syndrome computations are running on the device, copy output bits to 5G stack output buffer
    memcpy(p_out, context.dev_bits_out, memorySize_bits_out);
#else
    cudaCheck(cudaMemcpyAsync(p_out, context.dev_bits_out, memorySize_bits_out, cudaMemcpyDeviceToHost, context.stream));

    // ... schedule syndrome or CRC check ...

    // allow CPU access of output bits and syndrome
    cudaStreamSynchronize(context.stream);
#endif

Kernel Optimization

The core idea of CUDA programming is to parallelize the execution of the kernel on the GPU. The kernel is executed by a grid of blocks, each containing a number of independent threads. This level of parallelism allows for significant speedups compared to the serial execution on the CPU.

For a detailed introduction to CUDA programming, we refer to the CUDA Programming Guide. Conceptually, a CUDA kernel is defined as follows

// Kernel definition: __global__ indicates that the function is a CUDA kernel
__global__ void my_kernel(float *data, int N) {
   // Each thread processes one element of the array

   // We calculate the global thread index from the block and thread indices
   // idx is unique for each thread
   int idx = threadIdx.x + blockIdx.x * blockDim.x;
   if (idx < N) {
      // process the idx-th element of the array
      data[idx] = ...;
   }
}

This kernel can now be launched by specifying its grid and block dimensions via

// Launch the kernel
// <<<1, 32>>> specifies the grid and block sizes
my_kernel<<<1, 32>>>(d_data, N);

For the case of LDPC decoding, the decoder can be parallelized over the number of variable (VN) and check node (CN) updates, respectively. An overview of the CUDA implementation is given in Fig. 25. The VN update kernel is given as

__launch_bounds__(UNROLL_NODES*NODE_KERNEL_BLOCK, 3)
static __global__ void update_vn_kernel(llr_msg_t const* __restrict__ llr_msg, int8_t const* __restrict__ llr_ch, llr_accumulator_t* __restrict__ llr_total,
                                        uint32_t Z, uint32_t const* __restrict__ bg_vn, uint32_t const* __restrict__ bg_vn_degree, uint32_t max_degree, uint32_t num_cols, uint32_t num_rows) {
    uint32_t tid = blockIdx.x * blockDim.x + threadIdx.x;

    uint32_t i = tid % Z; // for i in range(Z)
    uint32_t idx_col = tid / Z; // for idx_col in range(num_cols)

    uint32_t vn_degree = bg_vn_degree[idx_col];

    // list of tuples (idx_row = index_cn, s)
    // idx_col = idx_vn and msg_offset omitted,
    // msg spread out to idx_cn + idx_vn * num_cn
    uint32_t const* variable_nodes = &bg_vn[idx_col * max_degree]; // len(vn) = vn_degree

#if UNROLL_NODES > 1
    __shared__ int msg_sums[UNROLL_NODES][NODE_KERNEL_BLOCK+1];
    msg_sums[threadIdx.y][threadIdx.x] = 0;
#endif

    __syncwarp();

    int msg_sum = 0;
    // accumulate all incoming LLRs
    if (idx_col < num_cols) {
    for (uint32_t j = threadIdx.y; j < vn_degree; j += UNROLL_NODES) {
        uint32_t vn = variable_nodes[j];

        // see packing layout above
        uint32_t idx_row = vn & 0xffffu; // note: little endian
        uint32_t s = vn >> 16;           // ...
        uint32_t msg_offset = idx_row + idx_col * num_rows;

        // index of the msg in the LLR array
        // it is the idx_col-th variable node, and the j-th message from the idx_row-th check node
        uint32_t msg_idx = msg_offset * Z + (i-s+(Z<<8))%Z;

        // accumulate all incoming LLRs
        msg_sum += llr_msg[msg_idx];
    }
    }

    // add the channel LLRs
    __syncwarp();
    if (threadIdx.y == 0) {
        if (idx_col < num_cols)
        msg_sum += llr_ch[idx_col*Z + i];
    }

    msg_sum = min(max(msg_sum, -MAX_LLR_ACCUMULATOR_VALUE), MAX_LLR_ACCUMULATOR_VALUE);

#if UNROLL_NODES > 1
	msg_sums[threadIdx.y][threadIdx.x] = msg_sum;

    __syncthreads();

    if (threadIdx.y == 0) {
        int msg_sum = 0;
        for (int i = 0; i < UNROLL_NODES; ++i) {
            msg_sum += msg_sums[i][threadIdx.x];
        }
        msg_sums[0][threadIdx.x] = msg_sum;
    }

    __syncthreads();

    msg_sum = msg_sums[0][threadIdx.x];
#endif

    msg_sum = min(max(msg_sum, -MAX_LLR_ACCUMULATOR_VALUE), MAX_LLR_ACCUMULATOR_VALUE);

    __syncwarp();
    if (idx_col < num_cols)
    llr_total[idx_col*Z + i] = llr_accumulator_t(msg_sum);
}

As the OAI processing pipeline uses multiple threads, we need to ensure proper multi-threading support in the CUDA kernel. This is done via a thread specific context that is passed to the kernel. This ensures that each thread operates on its own CUDA stream and, thus, can be executed in parallel without interference.

struct ThreadContext {
    cudaStream_t stream = 0;

    // Device memory declarations - use raw pointers instead of device symbols
    int8_t* llr_in_buffer = nullptr;
    uint8_t* llr_bits_out_buffer = nullptr;
    uint32_t* syndrome_buffer = nullptr;
#ifndef USE_UNIFIED_MEMORY
    uint32_t* host_syndrome_buffer = nullptr;
#endif
    llr_msg_t* llr_msg_buffer = nullptr;
    llr_accumulator_t* llr_total_buffer = nullptr;

#ifdef USE_GRAPHS
    cudaGraphExec_t graphCtx = nullptr;
#endif

    // list of thread contexts for shutdown
    ThreadContext* next_initialized_context = nullptr;
};
static __thread ThreadContext thread_context = { };

Further, the decoder uses clipping values for the extrinsic messages and the VN accumulator.

    // add the channel LLRs
    __syncwarp();
    if (threadIdx.y == 0) {
        if (idx_col < num_cols)
        msg_sum += llr_ch[idx_col*Z + i];
    }

    msg_sum = min(max(msg_sum, -MAX_LLR_ACCUMULATOR_VALUE), MAX_LLR_ACCUMULATOR_VALUE);

#if UNROLL_NODES > 1
	msg_sums[threadIdx.y][threadIdx.x] = msg_sum;

    __syncthreads();

    if (threadIdx.y == 0) {
        int msg_sum = 0;
        for (int i = 0; i < UNROLL_NODES; ++i) {
            msg_sum += msg_sums[i][threadIdx.x];
        }
        msg_sums[0][threadIdx.x] = msg_sum;
    }

    __syncthreads();

    msg_sum = msg_sums[0][threadIdx.x];
#endif

    msg_sum = min(max(msg_sum, -MAX_LLR_ACCUMULATOR_VALUE), MAX_LLR_ACCUMULATOR_VALUE);

    __syncwarp();
    if (idx_col < num_cols)
    llr_total[idx_col*Z + i] = llr_accumulator_t(msg_sum);

    // clip msg magnitudes to MAX_LLR_VALUE
    min_1 = min(max(min_1, -MAX_LLR_MSG_VALUE), MAX_LLR_MSG_VALUE);
    min_2 = min(max(min_2, -MAX_LLR_MSG_VALUE), MAX_LLR_MSG_VALUE);

Following the same principles, the CN update kernel is given as

__launch_bounds__(UNROLL_NODES*NODE_KERNEL_BLOCK, 3)
static __global__ void update_cn_kernel(llr_accumulator_t const* __restrict__ llr_total, llr_msg_t* __restrict__ llr_msg,
                                        uint32_t Z, uint32_t const* __restrict__ bg_cn, uint32_t const* __restrict__ bg_cn_degree, uint32_t max_degree, uint32_t num_rows,
                                        bool first_iter) {
    uint32_t tid = blockIdx.x * blockDim.x + threadIdx.x;

    uint32_t i = tid % Z; // for i in range(Z)
    uint32_t idx_row = tid / Z; // for idx_row in range(num_rows)

    uint32_t cn_degree = bg_cn_degree[idx_row];

    // list of tuples (idx_col = idx_vn, s),
    // idx_row = idx_cn and msg_offset omitted,
    // msg spread out to idx_cn + idx_vn * num_cn
    uint32_t const* check_nodes = &bg_cn[idx_row * max_degree]; // len(cn) = cn_degree

    // search the "extrinsic" min of all incoming LLRs
    // this means we need to find the min and the second min of all incoming LLRs
    int min_1 = INT_MAX;
    int min_2 = INT_MAX;
    int idx_min = -1;
    int node_sign = 1;
    uint32_t msg_signs = 0; // bitset, 0 == positive; max degree is 19

#if UNROLL_NODES > 1
    __shared__ int mins1[UNROLL_NODES][NODE_KERNEL_BLOCK+1];
    __shared__ int mins2[UNROLL_NODES][NODE_KERNEL_BLOCK+1];
    __shared__ int idx_mins[UNROLL_NODES][NODE_KERNEL_BLOCK+1];
    __shared__ uint32_t signs[UNROLL_NODES][NODE_KERNEL_BLOCK+1];
    mins1[threadIdx.y][threadIdx.x] = INT_MAX;
    mins2[threadIdx.y][threadIdx.x] = INT_MAX;
    signs[threadIdx.y][threadIdx.x] = 0;
#endif

    __syncwarp();

    if (idx_row < num_rows) {
    for (uint32_t ii = threadIdx.y; ii < cn_degree; ii += UNROLL_NODES) {
        uint32_t cn = check_nodes[ii];

        // see packing layout above
        uint32_t idx_col = cn & 0xffffu; // note: little endian
        uint32_t s = cn >> 16;           // ...
        uint32_t msg_offset = idx_row + idx_col * num_rows;

        uint32_t msg_idx = msg_offset * Z + i;

        // total VN message
        int t = llr_total[idx_col*Z + (i+s)%Z];

        // make extrinsic by subtracting the previous msg
        if (!first_iter)
            t -= __ldg(&llr_msg[msg_idx]);

        // store sign for 2nd recursion
        // note: could be also used for syndrome-based check or early termination
        int sign = (t >= 0 ? 1 : -1);
        node_sign *= sign;
        msg_signs |= (t < 0) << ii; // for later sign calculation

        // find min and second min
        int t_abs = abs(t);
        if (t_abs < min_1) {
            min_2 = min_1;
            min_1 = t_abs;
            idx_min = msg_idx;
        } else if (t_abs < min_2)
            min_2 = t_abs;
    }
    }

#if UNROLL_NODES > 1
	mins1[threadIdx.y][threadIdx.x] = min_1;
	mins2[threadIdx.y][threadIdx.x] = min_2;
	idx_mins[threadIdx.y][threadIdx.x] = idx_min;
	signs[threadIdx.y][threadIdx.x] = msg_signs;

    __syncthreads();

    if (threadIdx.y == 0) {
        int min_1 = INT_MAX;
        int min_2 = INT_MAX;
        int idx_min = -1;
        uint32_t msg_signs = 0; // bitset, 0 == positive; max degree is 19
        for (int i = 0; i < UNROLL_NODES; ++i) {
            int t_abs = mins1[i][threadIdx.x];
            if (t_abs < min_1) {
                min_2 = min_1;
                min_1 = t_abs;
                idx_min = idx_mins[i][threadIdx.x];
            } else if (t_abs < min_2)
                min_2 = t_abs;
            min_2 = min(min_2, mins2[i][threadIdx.x]);
            msg_signs |= signs[i][threadIdx.x];
        }
        mins1[0][threadIdx.x] = min_1;
        mins2[0][threadIdx.x] = min_2;
        idx_mins[0][threadIdx.x] = idx_min;
        signs[0][threadIdx.x] = msg_signs;
    }

    __syncthreads();

    min_1 = mins1[0][threadIdx.x];
    min_2 = mins2[0][threadIdx.x];
    idx_min = idx_mins[0][threadIdx.x];
    msg_signs = signs[0][threadIdx.x];

    node_sign = (__popc(msg_signs) & 1) ? -1 : 1;
#endif

    // START marker-cnp-damping
    // apply damping factor
    min_1 = APPLY_DAMPING_INT(min_1); // min_1 * DAMPING_FACTOR, e.g. *3/4
    min_2 = APPLY_DAMPING_INT(min_2); // min_2 * DAMPING_FACTOR, e.g. *3/4
    // END marker-cnp-damping

    // clip msg magnitudes to MAX_LLR_VALUE
    min_1 = min(max(min_1, -MAX_LLR_MSG_VALUE), MAX_LLR_MSG_VALUE);
    min_2 = min(max(min_2, -MAX_LLR_MSG_VALUE), MAX_LLR_MSG_VALUE);
    // END marker-vnp-clipping

    __syncwarp();

    // apply min and second min to the outgoing LLR
    if (idx_row < num_rows) {
    for (uint32_t ii = threadIdx.y; ii < cn_degree; ii += UNROLL_NODES) {
         uint32_t cn = check_nodes[ii];

        // see packing layout above
        uint32_t idx_col = cn & 0xffffu; // note: little endian
        uint32_t msg_offset = idx_row + idx_col * num_rows;

        uint32_t msg_idx = msg_offset * Z + i;
        int min_val;
        if (msg_idx == idx_min)
            min_val = min_2;
        else
            min_val = min_1;

        int msg_sign = (msg_signs >> ii) & 0x1 ? -1 : 1;

        // and update outgoing msg including sign
        llr_msg[msg_idx] = llr_msg_t(min_val * node_sign * msg_sign);
    }
    }
}

Note that the choice of the clipping values is critical for the performance of the decoder, in particular as the decoder uses mostly signed char variables. This keeps the memory footprint low and can be configured via the following macros

static const int MAX_LLR_ACCUMULATOR_VALUE = 127;
typedef int8_t llr_accumulator_t;
static const int MAX_LLR_MSG_VALUE = 127;
typedef int8_t llr_msg_t;

For a given lifting factor Z and a number of rows num_rows and columns num_cols given from the BG selection, the following grid and block dimensions are used

dim3 threads(256);
// check node update
dim3 blocks_cn(blocks_for(bg.num_rows * Z, threads.x));
// VN update
dim3 blocks_vn(blocks_for(bg.num_cols * Z, threads.x));

inline __host__ __device__ uint32_t blocks_for(uint32_t elements, int block_size) {
    return int( uint32_t(elements + (block_size-1)) / uint32_t(block_size) );
}

Note that the grid and block dimensions are architecture dependent and should be tuned for the specific GPU.

After decoding, the output is hard-decided and packed into a byte-array.

__launch_bounds__(PACK_BITS_KERNEL_THREADS, 6)
static __global__ void pack_bits_kernel(llr_accumulator_t const* __restrict__ llr_total, uint8_t* __restrict__ bits, uint32_t block_length) {
    uint32_t tid = blockIdx.x * blockDim.x + threadIdx.x;

    uint32_t coop_byte = 0;
    // 1 bit per thread
    if (tid < block_length)
        coop_byte = (llr_total[tid] < 0) << (7 - (threadIdx.x & 7)); // note: highest to lowest bit

    // use fast lane shuffles to assemble one byte per group of 8 adjacent threads
    coop_byte += __shfl_xor_sync(0xffffffff, coop_byte, 1); // xxyyzzww
    coop_byte += __shfl_xor_sync(0xffffffff, coop_byte, 2); // xxxxyyyy
    coop_byte += __shfl_xor_sync(0xffffffff, coop_byte, 4); // xxxxxxxx

    // share bytes across thread group to allow one coalesced write by first N threads
    __shared__ uint32_t bit_block_shared[PACK_BITS_KERNEL_THREADS / 8];
    if ((threadIdx.x & 0x7) == 0)
        bit_block_shared[threadIdx.x / 8] = coop_byte;

    __syncthreads();

    // the first (PACK_BITS_KERNEL_THREADS / 8) threads pack 8 bits each
    if (threadIdx.x < PACK_BITS_KERNEL_THREADS / 8 && blockIdx.x * PACK_BITS_KERNEL_THREADS + threadIdx.x * 8 < block_length) {
        bits[blockIdx.x * PACK_BITS_KERNEL_THREADS / 8 + threadIdx.x] = bit_block_shared[threadIdx.x];
    }
}

Note

The decoder returns the number of iterations and declares a decoding failure by returning max_num_iter if the CRC or the syndrome check fails.

Testing

Unit tests verify the CUDA decoder against Sionna’s reference implementation using nanobind for Python-CUDA binding. The tests cover multiple code rates for both base graphs (BG1 and BG2).

cd plugins/ldpc_cuda
pytest tests/unit/

The integration test runs the full 5G stack with the CUDA decoder and verifies end-to-end operation via iperf3 traffic.

cd plugins/ldpc_cuda/tests/integration
./run_integration.sh

See plugins/ldpc_cuda/tests/ for the full test suite.

Outlook - Weighted Belief Propagation

An extension of classical belief propagation is the weighted belief propagation [Nachmani2016] using gradient descent to optimize trainable parameters during message passing. An example implementation can be found in the Sionna tutorial on Weighted Belief Propagation.

When looking at the current implementation of the CN update, we can already see a damping factor of 3/4 applied to each outgoing CN message [Pretti2005].

#define APPLY_DAMPING_INT(x) (x*3/4)

    // apply damping factor
    min_1 = APPLY_DAMPING_INT(min_1); // min_1 * DAMPING_FACTOR, e.g. *3/4
    min_2 = APPLY_DAMPING_INT(min_2); // min_2 * DAMPING_FACTOR, e.g. *3/4

One could now implement a weighted CN update by simply replacing the damping factor with a trainable parameter from the weighted BP tutorial.

We hope you enjoyed this tutorial. You are now well equipped to accelerate other compute intensive parts of the 5G stack as well by following the same principles.