Source-linked AI summary
Interactive Debugger for Performance Portable Python HPC Kernels
Ivan Grigorik, Gabriel Kosmacher, George Biros, Milos Gligoric
TL;DR
Python HPC eDSLs lack interactive debugging that preserves target-device execution, leaving developers with prints, assertions, or altered CPU-only runs. PKDB provides pdb-style on-device debugging with live evaluation and kernel call site substitution, and its evaluation reports moderate overhead alongside faster debugging than sequential CPU-only mode. Its current implementation is a PyKokkos proof of concept, with broader applicability requiring substantial engineering.
Problem
Python eDSL developers largely rely on print statements, assertions, or altered CPU-only execution because interactive debugging for device kernels is limited.
Method
PKDB combines pdb-style debugging for on-device Python HPC kernels with live code evaluation and kernel call site substitution, using platform-specific controllers over a modular architecture.
Results
PKDB introduces moderate debug overhead and substantially outperforms PyKokkos-Debug's sequential Python-only execution across evaluated workloads and backends.
Takeaways & Limitations
PKDB provides practical interactive debugging for performance-portable Python HPC kernels while preserving native CPU and GPU execution and avoiding full application restarts for kernel changes.
Takeaways & Limitations
PKDB is demonstrated with PyKokkos, and applying its core ideas to other eDSLs or mixed-language stacks may require substantial engineering.
Abstract
from arXiv · showhide
We propose PKDB, the first interactive debugger for GPU and multithreaded low-level kernels written in Python. Python is widely used in high performance computing (HPC), with frameworks such as PyKokkos translating Python-embedded domain-specific languages to native code that runs across OpenMP-threaded CPUs and various GPUs. Yet interactive debugging support for such code is absent: developers resort to print statements, framework-specific assertions, or CPU-only execution, the last of which requires altering the program or its data and can mask device-specific bugs. PKDB enables standard interactive debugging like breakpoints, stepping, and variable inspection while preserving actual on-device execution without source modification. Beyond these fundamentals, PKDB introduces two advanced capabilities that exploit the dynamic nature of Python and PyKokkos: (i) Live code evaluation, which lets developers execute arbitrary Python expressions or entire kernels in the middle of a paused kernel without restarting the process; (ii) Kernel call site substitution, which allows an actively running kernel to be updated and reloaded on the fly, so only the kernel is recompiled and re-executed without restarting the application. Our performance evaluation on Intel, AMD, and NVIDIA CPUs, and NVIDIA and AMD GPUs shows that PKDB introduces limited overhead and is practical for everyday use while introducing critical debugging features to the Python HPC ecosystem.
I. INTRODUCTION
PKDB addresses the lack of interactive debugging for Python GPU and multithreaded HPC kernels by bringing familiar debugging workflows to actual device execution. It adds live evaluation and kernel call site substitution while evaluating support across heterogeneous CPUs and GPUs.
- Motivation and contribution: PKDB targets the limited development tooling available for Python eDSLs and provides the first interactive debugger for GPU and multithreaded low-level Python kernels.Existing practice relies on print statements, assertions, or altered CPU-only execution, which may not expose device-specific parallel behavior.
- Core debugging: PKDB mirrors pdb while debugging application kernels on CPUs and GPUs without source modification, preserving on-device execution and supporting heterogeneous accelerator types.The interface supports breakpoints, stepping, continuation, variable inspection, and context switching between Python and generated device code.
- Advanced capabilities: Live code evaluation executes arbitrary Python code or entire kernels during a paused session, allowing developers to inspect or modify program state without restarting.PKDB can execute multiple kernel versions concurrently to reduce live-evaluation latency.
- Advanced capabilities: Kernel call site substitution updates and reloads an actively running kernel so it can be recompiled and re-executed without restarting the application.This supports fixing and continuing kernel execution within the same debugger session.
- Evaluation: PKDB is evaluated on ExaMiniMD, a Boltzmann-kinetics solver, and a periodic Ewald sum across Intel, NVIDIA, and AMD CPUs and NVIDIA and AMD GPUs.The evaluation compares debugging workflows and measures time saved by call site substitution versus full application restart.
II. BACKGROUND AND EXAMPLE
Kokkos provides performance-portable execution across diverse hardware, while PyKokkos brings this model into a Python eDSL that translates kernels to Kokkos and executes them on target devices. The example illustrates a GPU kernel launched from Python, alongside the usability challenges that motivate PKDB.
- Kokkos framework: Kokkos maps a single C++ source to diverse CPU and GPU backends through execution spaces and parallel execution policies.RangePolicy distributes one-dimensional iterations across threads, while TeamPolicy provides hierarchical parallelism with teams that can synchronize and share local data.
- End-to-end example: The Figure 1 workflow launches yAx with a CUDA RangePolicy and just-in-time translates it into a Kokkos C++ functor whose results remain accessible through the original Python objects.The example uses CuPy arrays on the GPU and a parallel reduction over N threads.
- Kokkos framework: Kokkos remains difficult to use because template-heavy diagnostics, explicit memory and device management, and C++-centric development increase the workflow burden for Python-oriented developers.
- PyKokkos: PyKokkos lets developers write kernels in a Python subset, then automatically translates, compiles, and executes them on a selected device through Kokkos-compatible dispatch APIs.The eDSL uses @pk.workunit kernels with parallel_for, parallel_reduce, or parallel_scan and paired execution spaces and policies.
C. PKDB example
PKDB provides a unified debugging session spanning Python host code and generated device kernels, with breakpoints, stepping, inspection, and context switching across execution spaces. Its architecture connects pdb⋆, platform-specific controllers, and target debuggers while preserving the PyKokkos dispatch path.
- Interactive session: PKDB lets developers set source-line breakpoints in both Python host code and PyKokkos kernels through one transparent interface.CUDA-GDB handles device-kernel breakpoints, while pdb handles Python host breakpoints.
- Interactive session: The example session moves from a Python breakpoint into a device-kernel breakpoint, then continues deeper into device execution.The developer inspects y on the Python side before continuing to the kernel’s temp-variable definition.
- Interactive session: PKDB fetches requested device-array slices with standard Python range syntax, avoiding direct full-array transfer into the debugger.The displayed values match earlier Python-side output, confirming the kernel receives the expected input.
- Interactive session: Live evaluation computes device_sum directly on the GPU and cross-checks it with an equivalent Python-side expression.This demonstrates evaluation without moving the array data to the host.
- Architecture: PKDB’s architecture uses three cooperating processes—pdb⋆, a controller, and a target debugger—connected through two PTYs.Commands travel from pdb⋆ through the controller to the target debugger, with responses returning along the same path.
- Architecture: At each parallel dispatch, PKDB captures the environment, discovers the execution space, compiles with debugging information, and builds source-to-generated-code line mappings.The discovered execution space selects the controller, while disabled optimizations preserve debug information.
- Architecture: Platform-specific controllers connect common PKDB commands to OpenMP, CUDA, and HIP target debuggers while the rest of the architecture remains platform-agnostic.The controller API is shared across targets and commands are exposed through pdb⋆.
B. Live code evaluation
Live code evaluation executes arbitrary Python expressions or new PyKokkos kernels at a breakpoint using the captured environment. Device arrays are passed through shared GPU memory rather than copied to the host, so evaluated code can inspect or modify live state.
- Evaluation semantics: PKDB evaluates arbitrary Python expressions at a breakpoint using global variables, functions, and modules captured at the parallel-dispatch boundary.Assignments can remain visible after execution resumes, and evaluated expressions may dispatch new PyKokkos kernels.
- Device data handling: Avoiding device-to-host copies reduces transfer cost and prevents out-of-memory errors for datasets that already fill device memory.The evaluated expression can write through the shared allocation, making those changes visible to the rest of the program.
- Device data handling: Device-array arguments are reconstructed in the evaluation context through IPC handles that expose views of the original device allocations.The mechanism retrieves device pointers, opens IPC memory handles, and builds CuPy views without copying the underlying data.
C. Kernel call site substitution
Kernel call site substitution lets users replace kernels during a debug session without restarting the process. PKDB records global or location-specific replacements, checks signatures at dispatch, and executes the substituted kernel when applicable.
- Purpose: PKDB’s kernel call site substitution replaces kernels during a debug session without restarting the process.The mechanism is intended to avoid restart costs in long-running HPC jobs.
- Design rationale: PKDB exploits runtime kernel compilation and call-site interposition instead of modifying Python source files or using record/replay.The approach redirects the callee when each parallel operation is dispatched.
- Commands: The hotswap command supports either global replacement of every future dispatch or replacement restricted to one source-line call site.Location-specific substitutions leave other dispatches of the original kernel unchanged.
- Safety condition: The replacement must match the original kernel signature; otherwise PKDB reports an error and falls back to the original kernel.The application algorithm explicitly validates signatures before constructing the substituted dispatch.
- Mechanism: PKDB stores global and per-line substitution maps and checks them at each parallel_<op> dispatch before applying a replacement.The dispatch logic preserves the original source call site while redirecting execution when a substitution is registered.
D. Concurrent kernel launch
Concurrent kernel launch lets users run multiple kernel clauses at a breakpoint, including variants targeting different execution spaces. By default, each clause receives isolated array data, while mutable annotations can share storage but may introduce races.
- Purpose: Concurrent kernel launch runs multiple kernels at a breakpoint so their results can be compared directly on live data.Different clauses may target different execution spaces, such as GPU and CPU.
- Command syntax: The parallel_launch syntax specifies kernel names, optional execution spaces, policies, and arguments for multiple semicolon-separated clauses.Names and expressions resolve through Python eval against the live environment at the dispatch boundary.
- Data isolation: By default, array arguments use isolated copies so concurrently executing kernels do not interfere.CuPy arrays use host snapshots for GPU execution spaces, while NumPy arrays receive separate CPU buffers for OpenMP execution spaces.
- Mutable data: The mut: annotation enables shared array storage, but marking the same buffer mutable in multiple clauses can cause races and silent corruption.PKDB applies the annotations directly and leaves correctness to the user.
E. Thread-aware continue
PKDB makes multithreaded continuation deterministic by controlling GDB scheduling, while also allowing users to resume selected thread ranges.
- PKDB locks the GDB scheduler and continues threads explicitly, one by one, so all threads reach breakpoints deterministically.
- The extended continue threads [begin:end] command resumes only a specified thread range, while continue without a range resumes all threads.PKDB does not specifically handle breakpoints inside thread-dependent if branches.
IV. EVALUATION
The evaluation measures PKDB’s debugging overhead, compares it with PyKokkos Debug mode, assesses kernel call site substitution cost, and applies PKDB to two research applications.
- Evaluation questions: The experiments measure wall-time differences between PKDB debug mode and pdb across CUDA, HIP, and OpenMP backends.
- Evaluation questions: The evaluation compares PKDB’s unmodified GPU/OpenMP execution with PyKokkos Debug mode, which lowers parallel work to serial Python-only execution.
- Evaluation questions: The study measures kernel call site substitution against editing source files and restarting debugging sessions, and includes Boltzmann and Ewald case studies.
- Setup: Experiments use the machines listed in Table IV, with four runs per experiment, one discarded as a dry run, and averages over the remaining runs.
- Setup: For Section IV-C, disabling PKDB’s breakpoint-skipping optimization measures full framework overhead rather than pdb⋆ in isolation.
- Subjects: The subjects include ExaMiniMD, a PyKokkos particle-code mini-application, and Boltzmann, a particle-in-cell solver combining NumPy/CuPy with PyKokkos kernels.
C. Debug overhead
PKDB’s debug overhead remains bounded across the evaluated applications and backends, while preserving substantial speed advantages over serial Python debugging and making kernel substitution inexpensive.
- Debug overhead: 1.95×: ExaMiniMD’s mean tPKDB/tpdb ratio does not exceed 1.75× for OpenMP, 1.95× for CUDA, or 1.54× for HIP.
- Debug overhead: 2.33×: Boltzmann’s mean tPKDB/tpdb ratio does not exceed 2.18× for OpenMP, 2.33× for CUDA, or 2.14× for HIP.
- Debug overhead: 1.86×: Ewald’s mean tPKDB/tpdb ratio does not exceed 1.70× for OpenMP, 1.86× for CUDA, or 1.84× for HIP.The 5M Ewald runtime is unavailable for Figure 4c because of Local-server memory limits unrelated to PKDB.
- PyKokkos Debug comparison: PKDB debugging produces significant speedups over PyKokkos-Debug, which lowers parallel kernels to sequential Python-only execution.
- PyKokkos Debug comparison: 152.91–685.83s: PyKokkos-Debug spans this range on Local, while PKDB-OpenMP stays between 6.81 and 9.89s and PKDB-CUDA between 5.53 and 5.93s.
- Call site substitution: PKDB’s call site substitution benchmark swaps four alternate ExaMiniMD workunits and compares hotswap with automated edit-and-debug sessions at atom size 420k.
- Call site substitution: Tens of milliseconds: the hotswap command adds only this much wall time, whereas edit-and-debug launches another PKDB session and reruns the program to the call sites.
- Call site substitution: Call site substitution has constant asymptotic time under the stated infinite-process assumption, while edit-and-debug scales linearly with the number of call sites.
F. Case studies
PKDB is used to debug two PyKokkos research applications: a Boltzmann particle-in-cell kinetics code and an Ewald summation code.
- The case studies apply PKDB to the Boltzmann and Ewald PyKokkos applications.
1) Boltzmann:
PKDB helped diagnose real Boltzmann and Ewald failures during device execution by combining interactive debugging with device- and thread-level inspection. The case studies traced failures to a broken scratch-size translation and an out-of-bounds cell index.
- Boltzmann: PKDB exposed a Boltzmann CUDA illegal-address failure after stepping through the workunit to the crash.The failure occurred despite apparently consistent argument memory spaces and layouts.
- Boltzmann: A broken PyKokkos scratch-size specification left the translated C++ kernel without the required team cache allocation.The debugger showed pk.TeamMember.team_size equal to 0, after which fixing PyKokkos restored successful execution.
- Ewald: PKDB revealed Ewald out-of-bounds indexing by evaluating array lengths and inspecting the attempted index 528.The investigation used Python-side expressions during device debugging.
- Ewald: The Ewald bug was in _get_cell_fp64, which mishandled particles exactly on cell boundaries and produced an invalid counter index.A device-side parallel_print exposed per-thread cell_xyz values before pk.atomic_add accessed counter.
- Ewald: PKDB located the bug in real Ewald research code and supported device-side, per-thread inspection rather than only host-thread debugging.The case study emphasizes that PKDB led developers to the exact location requiring correction.
V. LIMITATIONS AND FUTURE WORK
PKDB currently trades performance for observability and is demonstrated primarily with PyKokkos, while its broader applicability remains prospective. Future work targets user studies and an eDSL-agnostic implementation layer.
- Limitations: PKDB's debugger-instrumented Kokkos build slows end-to-end runs relative to the paired release build, although the evaluation describes the overhead as moderate.The limitation reflects the usual performance cost of full-featured debugging instrumentation.
- Case study: The Ewald case study presents a PKDB transcript combining failure logs, broken code, and debugger stops at the relevant lines.The figure organizes the failure evidence across logs, source, and an interactive session.
- Scope: PKDB was selected as a proof-of-concept for PyKokkos, and extending its core ideas to other eDSLs or mixed-language stacks may require substantial engineering.The paper names Triton, Numba, Pallas, and Python bindings to native CUDA, HIP, or OpenMP code as possible targets.
- Future work: Future work includes user studies against print-debugging and restricted pdb workflows, plus an eDSL-agnostic PKDB layer.These directions aim to reduce dependence on a particular eDSL.
- Positioning: PKDB combines pdb-style composition with platform-specific debugging for performance-portable kernels, live code evaluation, and kernel call site substitution.Its related-work positioning distinguishes support for concurrent accelerators and dynamic kernel updates.