An Engineering Deep Dive

High-performance AI engineering requires bridging two historically separated domains: mathematical computational lowering (translating deep learning algorithms into silicon-friendly linear algebra) and systems infrastructure (exposing bare-metal hardware hardware to isolated container runtimes).

Understanding these domains independently is insufficient for modern performance engineers. Optimizing inference latency or cluster throughput demands tracing a high-level torch.nn.Conv2d operation all the way from tensor unrolling down to Linux device nodes, OCI runtime hooks, and PCIe execution paths.

1. Mathematical Computational Lowering: Spatial Math to GEMM

In high-level frameworks, convolutions operate on multi-dimensional tensors. A standard 2D convolution maps an input tensor $X \in \mathbb{R}^{N \times C_{in} \times H \times W}$ to an output tensor $Y \in \mathbb{R}^{N \times C_{out} \times H_{out} \times W_{out}}$ using weight tensor $W \in \mathbb{R}^{C_{out} \times C_{in} \times K_H \times K_W}$.

The Spatial-to-Matrix Translation

To eliminate non-contiguous memory access and exploit hardware matrix-multiply-accumulate (MAC) units, the 4D tensors are lowered via im2col:

1. Input Matrix ($X_{col}$): Every spatial window of shape $C_{in} \times K_H \times K_W$ scanned by the kernel is flattened into a single row vector of length $C_{in} \cdot K_H \cdot K_W$. Stacking these rows yields:
   $$X_{col} \in \mathbb{R}^{(N \cdot H_{out} \cdot W_{out}) \times (C_{in} \cdot K_H \cdot K_W)}$$
2. Weight Matrix ($W_{row}$): Each of the $C_{out}$ kernels is flattened into a row vector of length $C_{in} \cdot K_H \cdot K_W$, yielding:
   $$W_{row} \in \mathbb{R}^{C_{out} \times (C_{in} \cdot K_H \cdot K_W)}$$
3. General Matrix Multiplication (GEMM): The spatial operation lowers to a single matrix multiplication:
   $$Y_{col} = X_{col} W_{row}^T$$

$$\text{Dimensions: } \left[(N \cdot H_{out} \cdot W_{out}) \times (C_{in} \cdot K_H \cdot K_W)\right] \times \left[(C_{in} \cdot K_H \cdot K_W) \times C_{out}\right] = \left[(N \cdot H_{out} \cdot W_{out}) \times C_{out}\right]$$

4. Spatial Folding:$Y_{col}$ is transposed and reshaped (folded) back into the 4D spatial representation $Y \in \mathbb{R}^{N \times C_{out} \times H_{out} \times W_{out}}$.

Hardware Impact of Lowering

• Memory Locality:$X_{col}$ guarantees that data required for sequential compute blocks resides in contiguous cache lines (L1/L2), preventing GPU memory stall cycles.
• Tensor Core Saturation: Modern accelerators (e.g., NVIDIA Tensor Cores or AMD Matrix Core Engines) require inputs tiled in fixed matrix dimensions (e.g., $16 \times 16$). GEMM structures permit warp-level cooperative matrix load/store instructions (mma.sync).

2. The Host-Container Hardware Boundary

Once lowered GEMM calls are issued by the framework, they must reach physical hardware. In containerized environments, Docker relies on Linux namespaces (pid, net, mnt, ipc) to isolate processes. However, isolated namespaces cannot interact with hardware peripherals by default.

+-------------------------------------------------------------------------+
| DOCKER CONTAINER (Isolated User-Space Namespace)                        |
|                                                                         |
|  +-------------------------------------------------------------------+  |
|  | PyTorch Application (Python Execution Layer)                      |  |
|  +-------------------------------------------------------------------+  |
|                                    |                                    |
|                                    v                                    |
|  +-------------------------------------------------------------------+  |
|  | CUDA Runtimes & Math Libraries (cuBLAS / cuDNN / libcudart.so)     |  |
|  +-------------------------------------------------------------------+  |
+-------------------------------------------------------------------------+
                                     |
                                     | Dynamic Link / System Call
                                     v
+-------------------------------------------------------------------------+
| INJECTED BY NVIDIA CONTAINER TOOLKIT (OCI Runtime Hook)                 |
|                                                                         |
|  +-------------------------------------------------------------------+  |
|  | Host User-Space Driver API (libcuda.so / libnvidia-ml.so)          |  |
|  +-------------------------------------------------------------------+  |
|                                    |                                    |
|                                    v ioctl system calls                 |
|  +-------------------------------------------------------------------+  |
|  | Injected Character Device Node (/dev/nvidia0, /dev/nvidia-ctl)       |  |
|  +-------------------------------------------------------------------+  |
+-------------------------------------------------------------------------+
                                     |
                                     | Hardware Interconnect (PCIe Bus)
                                     v
+-------------------------------------------------------------------------+
| HOST OS & PHYSICAL SILICON                                              |
|                                                                         |
|  +-------------------------------------------------------------------+  |
|  | Kernel Module (nvidia.ko) -> Physical GPU (PCIe Bus Execution)    |  |
|  +-------------------------------------------------------------------+  |
+-------------------------------------------------------------------------+

Driver-Runtime Separation

NVIDIA drivers operate across a strict architectural boundary:

• Host Kernel Module (nvidia.ko): Interfaces directly with physical PCIe registers, managing memory allocation, DMA transfers, and hardware interrupts on the host OS.
• User-Space Driver (libcuda.so): Translates high-level API commands into low-level ioctl system calls aimed at device files.
• User-Space Runtimes (libcudart.so, cuBLAS, cuDNN): Compiled directly into application containers or environment wheels.

OCI Runtime Hook Mechanics

When passing --gpus all to Docker, the default runtime delegates container setup to the nvidia-container-runtime-hook. This process operates in three phases:

1. Pre-start Interception: The hook intercepts container startup right after the root filesystem is mounted, but before the application entrypoint executes.
2. Device Node Injection: The hook reads host GPU state and creates character device nodes inside the container’s isolated /dev directory (e.g., /dev/nvidia0, /dev/nvidia-ctl, /dev/nvidia-uvm) using mknod.
3. Library Bind-Mounting: The hook bind-mounts host user-space driver libraries (such as libcuda.so.XXX.XX) directly into the container’s /usr/lib/x86_64-linux-gnu/ path.

Because the container accesses the host kernel module (nvidia.ko) directly via these injected device nodes, there is zero virtualization overhead during GPU execution.

3. Unified Execution Pipeline: Tracing a Forward Pass

Tracing a lowered convolution execution from Python down to silicon hardware reveals the step-by-step control and data flow:

4. End-to-End Implementation & Profiling Script

The following executable Python script demonstrates this vertical unified stack. It computes a manual im2col GEMM lowering, validates correctness against native PyTorch C++ kernels, profiles memory duplication overhead, and queries hardware device metrics via injected CUDA drivers.

Python

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

def verify_gpu_passthrough():
    """Queries device driver status injected via OCI runtime hook."""
    print("=== 1. Container Hardware Passthrough Telemetry ===")
    assert torch.cuda.is_available(), "CUDA Passthrough Failed: No GPU device nodes detected!"
    
    device_count = torch.cuda.device_count()
    device_name = torch.cuda.get_device_name(0)
    capability = torch.cuda.get_device_capability(0)
    
    print(f"Detected GPU Devices : {device_count}")
    print(f"Primary Device Name  : {device_name}")
    print(f"Compute Capability   : {capability[0]}.{capability[1]}")
    print(f"Driver/CUDA Version  : {torch.version.cuda}\n")

def profile_im2col_gemm_lowering():
    """
    Executes and profiles manual im2col GEMM lowering vs native C++ cuDNN kernel.
    """
    print("=== 2. Computational Lowering & Execution Verification ===")
    
    # Define problem dimensions (NCHW)
    N, C_in, H, W = 16, 64, 56, 56
    C_out, K_H, K_W = 128, 3, 3
    stride, padding = 1, 1

    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

    # Allocate input tensors directly on GPU memory
    x = torch.randn(N, C_in, H, W, device=device)
    weights = torch.randn(C_out, C_in, K_H, K_W, device=device)
    bias = torch.randn(C_out, device=device)

    # Benchmark 1: Native cuDNN Implementation
    torch.cuda.synchronize()
    start_native = time.perf_counter()
    y_native = F.conv2d(x, weights, bias, stride=stride, padding=padding)
    torch.cuda.synchronize()
    time_native = (time.perf_counter() - start_native) * 1000

    # Benchmark 2: Manual im2col + GEMM Lowering Implementation
    torch.cuda.synchronize()
    start_manual = time.perf_counter()
    
    # Step A: im2col lowering via unfold
    # Shape: (N, C_in * K_H * K_W, H_out * W_out)
    x_unfolded = F.unfold(x, kernel_size=(K_H, K_W), padding=padding, stride=stride)
    
    # Transpose for GEMM: (N, H_out * W_out, C_in * K_H * K_W)
    X_col = x_unfolded.transpose(1, 2)
    
    # Step B: Weight matrix reshaping into (C_out, C_in * K_H * K_W)
    W_row = weights.view(C_out, -1)
    
    # Step C: GEMM Execution via batched matrix multiplication
    # (N, L, K) @ (K, C_out) -> (N, L, C_out)
    Y_col = torch.matmul(X_col, W_row.t())
    
    if bias is not None:
        Y_col += bias

    # Step D: Spatial Folding back to (N, C_out, H_out, W_out)
    H_out = (H + 2 * padding - K_H) // stride + 1
    W_out = (W + 2 * padding - K_W) // stride + 1
    y_manual = Y_col.transpose(1, 2).view(N, C_out, H_out, W_out)
    
    torch.cuda.synchronize()
    time_manual = (time.perf_counter() - start_manual) * 1000

    # Compute maximum floating-point discrepancy
    max_diff = torch.max(torch.abs(y_native - y_manual)).item()

    print(f"Native cuDNN Execution Time : {time_native:.4f} ms")
    print(f"Manual GEMM Execution Time  : {time_manual:.4f} ms")
    print(f"Absolute Pointwise Difference: {max_diff:.6e}")
    
    # Memory duplication analysis
    raw_size_mb = (x.element_size() * x.nelement()) / (1024 * 1024)
    unrolled_size_mb = (X_col.element_size() * X_col.nelement()) / (1024 * 1024)
    
    print("\n=== 3. Memory Duplication Overhead Analysis ===")
    print(f"Raw Input Tensor Size      : {raw_size_mb:.2f} MB")
    print(f"Unrolled (im2col) Size     : {unrolled_size_mb:.2f} MB")
    print(f"VRAM Memory Inflation Factor: {unrolled_size_mb / raw_size_mb:.2f}x")

if __name__ == "__main__":
    verify_gpu_passthrough()
    profile_im2col_gemm_lowering()