Optimize For GPU

GPU Optimization for Python with NVIDIA

Treat GPU acceleration as an evidence-driven optimization, not an automatic rewrite. Preserve the user's numerical and algorithmic contract, measure with representative data, and keep the GPU version only when synchronized end-to-end benchmarks show a useful improvement.

When This Skill Applies

• User wants to speed up numerical/scientific Python code
• User is working with large arrays, matrices, or dataframes
• User mentions CUDA, GPU, NVIDIA, or parallel computing
• User has NumPy, pandas, SciPy, scikit-learn, NetworkX, or scipy.sparse.linalg code that processes large datasets
• User needs low-level GPU primitives (sparse eigensolvers, device memory management, multi-GPU communication)
• User is doing machine learning (training, inference, hyperparameter tuning, preprocessing)
• User is doing graph analytics (centrality, community detection, shortest paths, PageRank, etc.)
• User is doing vector search, nearest neighbor search, similarity search, or building a RAG pipeline
• User has Faiss, Annoy, ScaNN, or sklearn NearestNeighbors code that could be GPU-accelerated
• User wants GPU-accelerated interactive dashboards, cross-filtering, or exploratory data analysis on large datasets
• User is doing geospatial analysis (point-in-polygon, spatial joins, trajectory analysis, distance calculations) with GeoPandas or shapely
• User is doing image processing, computer vision, or medical imaging (filtering, segmentation, morphology, feature detection) with scikit-image or OpenCV
• User is working with whole-slide images (WSI), digital pathology, microscopy, or remote sensing imagery
• User is loading large binary data files into GPU memory (numpy.fromfile → cupy, or Python open() → GPU array)
• User needs to read files from S3, HTTP, or WebHDFS directly into GPU memory
• User mentions GPUDirect Storage (GDS) or wants to bypass CPU-memory staging for file IO
• User is doing physics simulation (particles, cloth, fluids, rigid bodies) or differentiable simulation
• User needs mesh operations (ray casting, closest-point queries, signed distance fields) or geometry processing on GPU
• User is doing robotics (kinematics, dynamics, control) with transforms and quaternions
• User has Python simulation loops that could be JIT-compiled to GPU kernels
• User mentions NVIDIA Warp or wants differentiable GPU simulation integrated with PyTorch/JAX
• User is doing simulations, signal processing, financial modeling, bioinformatics, physics, or any compute-intensive work
• User wants to optimize existing code and GPU acceleration is the right answer

Choose the Smallest Suitable Layer

Prefer a maintained library implementation over a custom kernel:

Do not move code out of PyTorch, JAX, TensorFlow, or another GPU-native framework merely to use one of these libraries. First remove CPU round trips and use the framework's compiler, profiler, mixed-precision, and batching facilities.

Treat these as legacy-only:

Full per-library guidance, including when each is the wrong choice and how to combine them, is in references/decision_framework.md. Install commands and CUDA version selection are in references/installation.md. Before/after conversions for every library are in references/code_transformation_patterns.md.

Optimization Workflow

1. Define the contract and baseline

• Capture a representative input, expected output, and acceptable numerical tolerance.
• Measure the current end-to-end path, including input, transfers, compute, and output.
• Profile before changing code. Use CPU profilers for CPU code and identify whether the real limit is compute, memory bandwidth, allocation, transfer, synchronization, or storage.
• Record hardware, package versions, dtypes, shapes, batch size, and warm-up policy with results.

2. Check suitability before porting

GPU execution is promising when the hot path exposes substantial independent work, runs often enough to amortize initialization and transfer, and has a working set that fits available device memory with room for temporaries. Keep a CPU path when the workload is small, mostly sequential, dominated by unsupported operations, or requires frequent host-device round trips.

Do not use fixed row-count thresholds as proof. Benchmark the user's actual shapes and hardware. For out-of-core data, estimate peak working memory and choose chunking, Dask, or a streaming design before allocating.

3. Try the least disruptive implementation

1. If the code already uses a GPU-native framework, optimize within that framework.
2. Try accelerator or backend modes (cudf.pandas, cuml.accel, nx-cugraph).
3. Move to a native GPU API only where accelerator coverage or performance is insufficient.
4. Write a custom kernel only when profiling shows an operation without a suitable library implementation.

Read the relevant library reference before writing code; compatible names can still differ in defaults, dtypes, output types, and supported arguments.

4. Keep a coherent GPU data path

• Transfer inputs once and keep intermediates device-resident.
• Reuse allocations and prefer out= or in-place forms when semantics allow.
• Batch small operations; fuse elementwise work when it removes intermediate arrays.
• Use pinned host memory and non-default streams only after profiling shows transfer overlap matters.
• Choose float32, mixed precision, or reduced-precision storage only when the contract permits it.

5. Validate semantics before speed

• Compare CPU and GPU outputs on small deterministic fixtures and representative data.
• Use explicit tolerances for floating-point results and test edge cases, NaNs, ordering, and dtypes.
• For approximate nearest-neighbor indexes, report recall@k against exact search; do not compare an exact CPU algorithm with an approximate GPU algorithm as if they were equivalent.
• Check accelerator warnings and logs for CPU fallback.

6. Benchmark GPU code correctly

GPU work is asynchronous, so a CPU timer around an unsynchronized call measures enqueue time. Warm up context creation and JIT compilation, then use CUDA events or a library-aware timer:

from cupyx.profiler import benchmark

print(benchmark(gpu_function, (arg1, arg2), n_warmup=10, n_repeat=100))