Implementation:NVIDIA DALI KernelContext

Knowledge Sources	NVIDIA_DALI
Domains	Kernels, GPU_Computing
Last Updated	2026-02-08 16:00 GMT

Overview

Defines the kernel execution context, the abstract Scratchpad interface for temporary memory allocation, and backend-specific context structures used throughout the DALI kernel framework.

Description

The context.h header is a foundational piece of the DALI kernel infrastructure. It provides three key components: the Context template (specialized for CPU and GPU backends), the Scratchpad abstract class for temporary memory management, and the unified KernelContext struct that bundles both together. The GPU specialization of Context carries a cudaStream_t for asynchronous GPU execution.

The Scratchpad class is an abstract interface that allows kernels to obtain auxiliary working memory without managing allocation lifetimes directly. It supports multiple memory kinds (host, pinned, device, managed) and provides convenience methods for allocating raw memory, typed arrays, tensors, and tensor lists. All memory allocated through a Scratchpad is released at once when the scratchpad object goes out of scope, following a bulk-deallocation pattern.

The KernelContext struct combines a CPU context, a GPU context, and a pointer to a Scratchpad, forming the standard execution environment passed to every kernel's Setup and Run methods. This design decouples kernels from specific memory management strategies and allows different scratchpad implementations (such as DynamicScratchpad) to be injected at runtime.

Usage

Use KernelContext whenever implementing or invoking a DALI kernel. The context is passed as the first argument to both Setup and Run methods. Before calling Run, the caller should assign a valid Scratchpad implementation to the context's scratchpad pointer. For GPU kernels, the gpu.stream field must be set to a valid CUDA stream. The Scratchpad methods such as AllocateGPU, AllocateHost, AllocTensor, and ToGPU are used inside kernel implementations to request temporary working memory.

Code Reference

Source Location

• Repository: NVIDIA_DALI
• File: dali/kernels/context.h
• Lines: 1-214

Signature

template <typename ComputeBackend>
struct Context {};

template <>
struct Context<ComputeGPU> {
  cudaStream_t stream = AccessOrder::null_stream();
};

class Scratchpad {
 public:
  template <typename MemoryKind>
  inline void *Alloc(size_t bytes, size_t alignment);

  template <typename MemoryKind, typename T, int dim>
  TensorView<kind2storage_t<MemoryKind>, T, dim> AllocTensor(TensorShape<dim> shape);

  template <typename MemoryKind, typename T, int dim>
  TensorListView<kind2storage_t<MemoryKind>, T, dim>
  AllocTensorList(TensorListShape<dim> shape);

  template <typename MemoryKind, typename T>
  T *Allocate(size_t count, size_t alignment = alignof(T));

  template <typename T>
  T *AllocateGPU(size_t count, size_t alignment = alignof(T));

  template <typename T>
  T *AllocateHost(size_t count, size_t alignment = alignof(T));

  template <typename T>
  T *AllocatePinned(size_t count, size_t alignment = alignof(T));

  template <typename T>
  T *AllocateManaged(size_t count, size_t alignment = alignof(T));

  template <typename Collection, typename T = std::remove_const_t<element_t<Collection>>>
  if_array_like<Collection, T*> ToGPU(cudaStream_t stream, const Collection &c);

  template <typename Collection, typename T = std::remove_const_t<element_t<Collection>>>
  if_iterable<Collection, T*> ToHost(const Collection &c);

  virtual void *Alloc(mm::memory_kind_id kind_id, size_t bytes, size_t alignment) = 0;
};

using CPUContext = Context<ComputeCPU>;
using GPUContext = Context<ComputeGPU>;

struct KernelContext {
  CPUContext cpu;
  GPUContext gpu;
  Scratchpad *scratchpad = nullptr;
};

Import

#include "dali/kernels/context.h"

I/O Contract

Inputs

Name	Type	Required	Description
bytes	size_t	Yes	Number of bytes to allocate (for Alloc/Allocate methods)
alignment	size_t	No	Alignment requirement for the allocation, defaults to alignof(T)
count	size_t	Yes	Number of elements of type T to allocate (for Allocate* methods)
shape	TensorShape<dim> / TensorListShape<dim>	Yes	Shape of tensor/tensor list to allocate (for AllocTensor/AllocTensorList)
stream	cudaStream_t	Yes	CUDA stream for GPU context and async copy operations (for ToGPU)
c	Collection	Yes	Source collection to copy to GPU/Host/Pinned/Managed memory

Outputs

Name	Type	Description
(pointer)	T*	Pointer to allocated memory of the requested type and kind
(tensor view)	TensorView<Backend, T, dim>	View over scratchpad-allocated tensor memory
(tensor list view)	TensorListView<Backend, T, dim>	View over scratchpad-allocated tensor list memory

Usage Examples

Allocating GPU Memory in a Kernel

#include "dali/kernels/context.h"

void MyKernel::Run(KernelContext &ctx,
                   const OutListGPU<float, 3> &out,
                   const InListGPU<float, 3> &in) {
  // Allocate temporary GPU buffer for 1024 floats
  float *tmp = ctx.scratchpad->AllocateGPU<float>(1024);

  // Copy a host vector to GPU memory via scratchpad
  std::vector<int> offsets = {0, 100, 200, 300};
  int *gpu_offsets = ctx.scratchpad->ToGPU(ctx.gpu.stream, offsets);

  // Allocate a temporary tensor on the GPU
  auto tmp_tensor = ctx.scratchpad->AllocTensor<mm::memory_kind::device, float, 2>(
      TensorShape<2>{64, 64});
}

Setting Up a KernelContext

#include "dali/kernels/context.h"
#include "dali/kernels/dynamic_scratchpad.h"

cudaStream_t stream;
cudaStreamCreate(&stream);

DynamicScratchpad scratchpad(AccessOrder(stream));
KernelContext ctx;
ctx.gpu.stream = stream;
ctx.scratchpad = &scratchpad;

// Now ctx can be passed to kernel Setup and Run methods

Related Pages

• Environment:NVIDIA_DALI_CUDA_GPU_Environment