Kernel Style Guide#

Introduction#

This style guide establishes conventions for developing NVIDIA Warp GPU/CPU kernels in nvalchemiops. The intention is to ensure consistency in naming, design patterns, unit testing, and performance evaluation for both human and agentic developers.

Naming Conventions#

Consult this section before naming variables to ensure consistency across modules. Try and reuse one of the names below if they are semantically identical to what you are developing for; in the event that they are different enough, consider using these as inspiration and make a best effort to align to these conventions.

Batching#

Variable	Shape	Dtype	Description
`batch_idx`	`[total_atoms]`	`int32`	System index for each atom
`batch_ptr`	`[num_systems+1]`	`int32`	Cumulative atom counts

Neighbor Lists#

Variable	Shape	Dtype	Description
`neighbor_matrix`	`[total_atoms, max_neighbors]`	`int32`	Neighbor indices
`neighbor_shift_matrix`	`[total_atoms, max_neighbors, 3]`	`int32`	Integer shifts
`cartesian_shifts`	`[total_atoms, max_neighbors, 3]`	`vec3*`	Cartesian shifts
`num_neighbors`	`[total_atoms]`	`int32`	Neighbor count
`neighbor_list`	`[2, num_pairs]`	`int32`	COO format
`max_neighbors`	scalar	`int`	Max neighbors
`fill_value`	scalar	`int`	Padding value
`cutoff`, `cutoff_sq`	scalar	`float*`	Distance cutoff

Geometry#

Variable	Shape	Dtype	Description
`positions`	`[total_atoms, 3]`	`vec3*`	Atomic Cartesian coordinates
`cell`	`[num_systems, 3, 3]` or `[3, 3]`	`mat33*`	Lattice vectors (row format)
`pbc`	`[num_systems, 3]` or `[3]`	`bool`	Periodic boundary flags
`numbers`	`[total_atoms]`	`int32`	Atomic numbers (Z)

Units: Document in docstrings; if relevant to correctness then Bohr or Angstrom must be specified. In instances where units just have to be consistent (e.g. neighbor lists) then that should be stated.

Thread Indices#

Variable	Type	Description
`tid`	`int`	Single thread ID from `wp.tid()` (1D launch)
`tid_i`, `tid_j`	`int`	Thread IDs from `wp.tid()` (N-D launch)
`atom_i`, `atom_j`	`int`	Semantic atom indices (use for clarity)
`i`, `j`, `k`	`int`	Generic loop indices

Convention: Use tid for thread IDs; use atom_i/atom_j for semantic clarity; i.e. if the threads map onto atom indices, then use atom_i/atom_j over tid_i. See nvalchemiops/interactions/dispersion/dftd3.py for examples.

Warp Array Objects#

The preferred way of instantiating warp arrays is to convert from PyTorch tensors, rather than using Warp constructor methods like wp.zeros:

output = torch.zeros(..., device=..., dtype=torch.float32)
# optionally, `return_ctype=True` when it makes sense for performance
output_wp = wp.from_torch(output)

Updates to output_wp will reflect in the PyTorch tensor, output and eliminates the need for re-converting with overhead, e.g.:

# anti-pattern
output_wp = wp.zeros(..., device=...)
# run kernel...
output = wp.to_torch(output_wp)

Some additional tips for performance:

Use appropriate layouts for PyTorch tensors before constructing Warp references, including torch.Tensor.contiguous() and torch.Tensor.to(memory_format=torch.channels_last). See this blog post for details.
Use Warp vector and matrix datatypes when possible, e.g. wp.array(..., dtype=wp.vec3f) for atomic coordinates, instead of wp.array(..., dtype=wp.float32). This encourages optimal memory access patterns.
When dtype s are known ahead of time, or preferably if wrappers can be written in such a way that they are known ahead of time, reduce Python overhead with wp.from_torch(..., return_ctype=True), which avoids the need for a wp.array Python object altogether.
Consider handling gradients manually, i.e. kernels for backward passes, when possible to do so. Ensure tensors are detached from the PyTorch computational graph to avoid duplication.

Warp Array Suffix#

Suggested (not mandatory): Use _wp suffix when converting PyTorch tensors to Warp arrays in wrapper functions.

positions_wp = wp.from_torch(positions.contiguous(), dtype=wp.vec3f)
numbers_wp = wp.from_torch(numbers.contiguous(), dtype=wp.int32)

When to use: In PyTorch wrappers with both tensor types in scope. Skip: Inside pure Warp code or when no ambiguity exists.

3. Kernel Design Patterns#

Naming Conventions#

Kernel wrappers are two-tiered:

Low-level functions are decorated with @torch.library.custom_op with mutates_args, that take pre-allocated tensors and return None. These private functions dispatch the kernels and are registered so that torch.compile/backwards passes will recognize them.
Higher-level wrapper functions handle tensor allocations and call the low-level wrapper.

@torch.library.custom_op(
    "nvalchemiops::kernel_name",
    mutates_args=(...),
)
def _low_level_wrapper(...):
    """Infer devices, dtypes, dispatch correct warp kernel"""

def high_level_wrapper(...):
    """Handles tensor allocations; main entry point for users"""

Type	Convention	Example
Private kernel	Leading underscore	`_geom_cn_kernel`
Helper (`wp.func`)	Leading underscore	`_valid_neighbor`, `_s5_switch`
Public API	No underscore	`neighbor_list`, `dftd3`

Linters#

In many cases, naming scientific variables compactly can go against PEP-style formatting/style guides. You can decorate lines where variables are declared to disable false positives (particularly from Sonar):

dE_dCN = ...  # NOSONAR (S125) "math formula"

Similarly, there can be branching based off numerics that are picked up by Sonar, which is not necessarily bad advice, albeit somewhat irrelevant:

if value == 0.0: ...  # NOSONAR (S1244) "warp kernel"

Precision Support#

When appropriate to do so, overload kernels programmatically and store the results in a dictionary with dtypes as keys:

@wp.kernel
def _my_kernel(positions: wp.array(dtype=Any), values: wp.array(dtype=Any)):
    # ... generic kernel code ...

# Register overloads
kernel_overloads = {}
for scalar_type, vec_type in zip([wp.float16, wp.float32, wp.float64],
                                  [wp.vec3h, wp.vec3f, wp.vec3d]):
    kernel_overloads[(scalar_type, vec_type)] = wp.overload(_my_kernel, {
        "positions": wp.array(dtype=vec_type),
        "values": wp.array(dtype=scalar_type),
    })

# runtime retrieval; determines and launches appropriate overload
def kernel_wrapper(...):
    scalar_type = scalar_data.dtype
    vec_type = vector_data.dtype
    kernel_func = kernel_overloads[(scalar_type, vec_type)]
    wp.launch(kernel_func, ...)

Not all precisions should be supported: if results are known to underflow particularly at lower precisions, do not add them to the overloads.

Documentation#

Use NumPy-style docstrings with the following specific items within each category:

Summary line
Parameters (with shape, dtype, description)
- Array outputs should be denoted with OUTPUT to indicate arrays that are expected to be pre-allocated, and are mutated in place by a kernel.
Returns
- Not generally applicable to Warp kernels, and more for PyTorch wrappers
Notes (launch patterns, caveats)
- Document the thread abstraction (e.g. per-atom, per-system)
- Known performance characteristics
See Also
- Reference related kernels, particularly those that are run immediately before or after the current kernel.

Unit Testing#

Test Organization#

Try and mirror test modules with the Python package tree: for example, a test_dftd3.py mirrors dftd3.py.
Use conftest.py for shared fixtures, such as systems to test against, devices, etc.
Group categories of tests in classes:

class TestCategory:
    @pytest.fixture
    def category_specific_fixture(): ...

    def test_kernel_basic(): ...

    def test_kernel_with_fixture(category_specific_fixture): ...

Test Patterns#

Parametrized tests for variations, particularly with device and datatypes (precision):

@pytest.mark.parametrize("dtype", ["float32", "float64"])
@pytest.mark.parametrize("device", ["cuda:0", "cpu"])
def test_kernel(dtype, device):
    # ... test implementation ...

Full pipeline tests (primary coverage): test the end-to-end workflow, making sure that the values are not just finite, but once correctness has been established, check numerical tests to make sure kernels are numerically stable between changes:

def test_dftd3_full_pipeline():
    """Test complete pipeline exercising all kernels."""
    energy, forces, cn = dftd3(...)
    assert torch.isfinite(energy).all()


def test_kernel_values():
    """Check numerical results against hardcoded reference values"""
    values = kernel(...)
    hardcoded_values = torch.tensor(...)
    assert torch.allclose(values, hardcoded_values, rtol=..., atol=...)

Known fail states: If a kernel or function is known to fail predictably with certain inputs or conditions, there must be tests that capture this using patterns such as pytest.raises(Exception). Good things to capture here are incorrect shapes and dtypes for warp kernels.
Edge cases: Empty systems, potentially odd input/outputs such as negative/positive values.
Helper function tests: Wrap wp.func in test kernel for isolated testing

Performance Benchmarking#

GPU Profiling Requirements#

Critical for accurate timing:

Warmup runs: GPU kernels compile on first launch (JIT)
Synchronization: GPU operations are asynchronous - always call torch.cuda.synchronize()/wp.synchronize()
CUDA Events: More relevant, and therefore preferable over time.perf_counter() calls
Multiple runs: following warm up runs to average statistics
Units: Report units at relevant timescales: micro/milliseconds is typical but may depend on the kernel

General pattern for collecting performance data:

# Warmup
for _ in range(warmup_runs):
    func()
    torch.cuda.synchronize()

# Timing with CUDA events
times = []
torch.cuda.memory.reset_peak_memory_stats()
for _ in range(timing_runs):
    torch.cuda.synchronize()
    start_event = torch.cuda.Event(enable_timing=True)
    end_event = torch.cuda.Event(enable_timing=True)

    start_event.record()
    func()
    end_event.record()

    torch.cuda.synchronize()
    times.append(start_event.elapsed_time(end_event))  # milliseconds
peak_memory = torch.cuda.memory.max_memory_allocated()

NVTX annotations for Nsight Systems profiling:

import nvtx

@nvtx.annotate("compute_dftd3", color="red")
def my_function(...):
    return dftd3(...)

What to Exclude from Timing#

Pre-compute separately (not part of kernel timing):

Neighbor list construction (unless benchmarking neighbor lists)
Parameter loading
Data transfers (unless benchmarking transfers)
One-time allocations

See benchmarks/interactions/dispersion/benchmark_d3_synthetic.py and benchmarks/benchmark_neighborlist.py for complete examples.