(style-guide)=

# 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

Framework bindings own their tensors, temporary storage, dtype and layout
normalization, and conversion to Warp array views. Core launchers accept Warp
arrays and run on the caller's current Warp stream; they must not select a
PyTorch or JAX stream by inspecting an array.

For a PyTorch binding, allocate with PyTorch and create a Warp view:

```python
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.:

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

`wp.from_torch(..., dtype=...)` selects a storage-compatible Warp scalar,
vector, or matrix view; it does not perform a numerical dtype conversion.
Perform required casts with PyTorch. Preserve supported strides and materialize
a contiguous tensor only when the selected Warp view or kernel layout requires
it. Casts and layout normalization may allocate, so keep their results alive
through the final Warp launch that uses them.

#### Framework stream ownership

Framework-owned storage and Warp work must share one execution stream:

- PyTorch launch leaves establish the current PyTorch CUDA stream with the
  shared `scoped_warp_stream` context or `scoped_torch_warp_stream` decorator
  before converting storage or launching kernels.
- JAX bindings use `jax_kernel` for individual launches and `jax_callable` for
  multi-launch callbacks. These adapters supply the XLA execution stream; a
  callback must not replace it.
- Direct Warp callers retain ownership of the current Warp stream and of the
  lifetime of every input, output, and scratch array.
- Core launchers never infer which framework owns an array and never switch to
  a framework stream.

The binding operation owns materialized casts and layout conversions, temporary
allocation, and Warp views. Enter the Torch stream scope before Warp conversion
or launch and retain it through the final dependent Warp launch. Enqueue
dependent Torch work on that same Torch stream before observing results on the
host. Do not use device synchronization, `CUDA_LAUNCH_BLOCKING`, or
`record_stream` as substitutes for correct stream ownership.

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](https://pytorch.org/blog/tensor-memory-format-matters/)
  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.

```python
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.

---

## Kernel Design Patterns

### Naming Conventions

Kernel wrappers are layered by ownership:

- Core functions take Warp arrays, launch on the current Warp stream, and do
  not import a framework.
- PyTorch launch leaves establish the PyTorch stream and dispatch core kernels.
  Use `@torch.library.custom_op` with accurate `mutates_args` when compiler or
  autograd integration requires a custom operator; direct binding functions use
  the same stream contract.
- JAX bindings dispatch through `jax_kernel` or `jax_callable`. New or modified
  bindings keep inputs, outputs, materialized conversions, and temporary
  operands owned by JAX/XLA; callback-local Warp allocation in existing code is
  deferred migration work, not a pattern to extend.
- Higher-level framework wrappers handle allocation and normalization before
  calling the low-level binding.

When a custom operator is required:

```python
@torch.library.custom_op(
    "nvalchemiops::kernel_name",
    mutates_args=(...),
)
@scoped_torch_warp_stream
def _low_level_wrapper(...):
    """Establish the current PyTorch stream and dispatch the 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):

```python
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:

```python
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 `dtype`s as keys:

```python
@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:

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

    def test_kernel_basic(): ...

    def test_kernel_with_fixture(category_specific_fixture): ...
```

### Test Patterns

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

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

1. 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:

   ```python
   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=...)
   ```

1. Framework-stream regressions: on CUDA, warm compilation first, queue a
   pending framework producer, call the public API, and immediately consume its
   output in the same framework. Synchronize only when observing the final
   numerical assertion. Use a non-default PyTorch stream; for JAX, keep the
   producer, operation, and consumer in one compiled computation or in
   dependency-linked compiled calls, without an intermediate
   `block_until_ready()`.

1. 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 `dtype`s for
   `warp` kernels.

1. Edge cases: Empty systems, potentially odd input/outputs such as
   negative/positive values.

1. Helper function tests: Wrap `wp.func` in test kernel for isolated
   testing

## Performance Benchmarking

### GPU Profiling Requirements

**Critical for accurate timing**:

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

**General pattern for collecting performance data**:

```python
# 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:

```python
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.