localcorr

PyTorch local correlation kernels for optical flow networks such as RAFT and PWC-Net.

The implementations are drop-in compatible with RAFT's alt_cuda_corr and spatial_correlation_sampler, and run on NVIDIA and Intel GPUs as well as any ONNX-supported runtime (see Components).

You can check PTLFlow to see examples of optical flow models using localcorr.

Quick start

pip install localcorr

That's enough to train and run inference eagerly on Linux — the default Triton kernels are JIT-compiled on first use with no build step required (see Installation for other platforms and optional AOT plugins).

RAFT-style correlation (feature maps + sampled coordinates):

import torch
from localcorr import get_localcorr_block

# fmap1/fmap2: [B, C, H, W], coords: [B, 2, H, W] (x, y)
corr_block = get_localcorr_block(
    fmap1, fmap2, num_levels=4, radius=4,
    corr_mode="localcorr_triton",
)
corr = corr_block(coords)  # [B, levels*(2r+1)^2, H, W]

PWCNet-style spatial correlation:

from localcorr import get_localcorr_sampler

# same result as SpatialCorrelationSampler(kernel_size=1, patch_size=9, padding=0)
sampler = get_localcorr_sampler(patch_size=9, corr_mode="localcorr_triton")
corr = sampler(input1, input2)  # [B, 9, 9, H, W]

Choosing a corr mode

Everything in localcorr is selected through a single corr_mode argument (accepted by both get_localcorr_block and get_localcorr_sampler):

Notes:

• Strict vs. automatic dispatch: Each specific mode (triton, cuda, sycl) uses its own eager engine and fails with a clear error if the extension is missing or tensors are on the wrong device, rather than silently falling back. localcorr_auto is the exception: it inspects the tensor device and falls back to the Triton kernel when the native AOT extension is not built.
• Export behavior: The modes only differ in eager execution and ONNX export; see ONNX export for how each one exports.

Installation

pip install localcorr

triton (the engine behind localcorr_triton and localcorr_onnx) is installed automatically on Linux. On platforms where Triton does not provide wheels, localcorr still imports and the pure-PyTorch / ONNX-export paths work, while calling a Triton kernel raises an ImportError on first use.

TensorRT and OpenVINO are deliberately not required dependencies:

• To build or run TensorRT engines with the cuda_pyramid plugin, install TensorRT yourself (the plugin build only needs NvInfer.h and libnvinfer discoverable).
• Likewise, OpenVINO is only needed to build the sycl extension and run models using the custom op.

Building the native components (optional)

The AOT components (cuda_pyramid, sycl_pyramid, xpu) provide the fastest eager engines and the TensorRT/OpenVINO plugins listed above. They are optional and can be compiled on demand with a single command from any directory (both the pip wheel and sdist ship the native sources):

python -m localcorr.build_native cuda   # NVIDIA: AOT kernel + TensorRT plugin
python -m localcorr.build_native sycl   # Intel: OpenVINO custom-op extension
python -m localcorr.build_native xpu    # Intel: torch.xpu SYCL extension
python -m localcorr.build_native all

When pip-installed, you can also use the console script (localcorr-build cuda) or build programmatically (from localcorr import build_native; build_native.build("cuda")). From a repository checkout, the direct scripts (bash localcorr/<component>/build.sh) also work.

Requirements per target:

• cuda — CUDA toolkit (nvcc), CMake >= 3.18, and TensorRT (NvInfer.h + libnvinfer). Set TRT_ROOT or TENSORRT_ROOT if TensorRT is not in a default system path. Output: localcorr/cuda_pyramid/build/liblocalcorr_cuda_pyramid_plugin.so
• sycl — Intel oneAPI (icpx with SYCL) and the OpenVINO runtime (the pip package headers are used by default; override with OV_INSTALL_DIR). Output: localcorr/sycl_pyramid/build/liblocalcorr_sycl_pyramid.so (the SYCL kernel) and localcorr/sycl_pyramid/build/liblocalcorr_sycl_pyramid_ov.so (the OpenVINO extension, plus test_kernel for correctness verification).
• xpu — a PyTorch XPU build (pip install torch --index-url [https://download.pytorch.org/whl/xpu](https://download.pytorch.org/whl/xpu)), oneAPI (icpx), and (for the Triton fallback only) triton-xpu plus the Level Zero headers. Output: localcorr/xpu/build/localcorr_sycl_xpu*.so

Usage

Eager execution & training

All modes support fp16 and fp32 inputs as well as PyTorch autograd for training. Across all backends, backward passes compute gradients with respect to the feature maps (fmap1/fmap2 or input1/input2), while gradients for the sampling coordinates (coords) are returned as zeros.

If you want the same model code to run unchanged across NVIDIA and Intel GPUs, use corr_mode="localcorr_auto" (which also accepts an optional xpu_backend argument in both factories):

# RAFT-style: backend resolved once from the device of fmap1/fmap2
corr_block = get_localcorr_block(
    fmap1, fmap2, num_levels=4, radius=4, corr_mode="localcorr_auto"
)
corr = corr_block(coords)  # [B, levels*(2r+1)^2, H, W]

# PWCNet-style: backend dispatched per call from the device of input1/input2
# (uses the LocalCorrCudaPyramidSampler class under the hood)
sampler = get_localcorr_sampler(patch_size=9, corr_mode="localcorr_auto")
corr = sampler(input1, input2)  # [B, 9, 9, H, W]

How localcorr_auto dispatches:

• NVIDIA GPUs: Uses the AOT CUDA kernel (localcorr_cuda) when the extension is built, otherwise falls back to the Triton kernel.
• Intel GPUs (torch.xpu): Uses the fused SYCL kernel when the extension is built, otherwise falls back to the Triton kernel (xpu_backend="triton" forces Triton; defaults to "sycl").
• ONNX export: Emits a single LocalCorrCudaPyramid plugin node (which also sidesteps the fact that the fused Triton kernel cannot be traced for export), or standard ONNX ops when the "onnx" backend is active.

ONNX export

By default, ONNX export emits LocalCorrCudaPyramid custom nodes understood by the TensorRT ONNX parser (via the CUDA plugin) and OpenVINO (via the sycl_pyramid SYCL custom-op extension).

To export models for runtimes without custom-plugin support, switch to the standard-ops backend before exporting:

import localcorr

localcorr.set_onnx_backend("onnx")  # or "tensorrt" (default)
# ... then export with the legacy tracer:
# torch.onnx.export(model, args, "model.onnx", dynamo=False, opset_version=17)

Important: Always pass dynamo=False when calling torch.onnx.export. Recent PyTorch versions default to the Dynamo exporter, which silently drops the correlation and emits an all-zeros placeholder instead of the plugin node.

Key behaviors during export:

• Environment variable: You can also select the standard-ops backend via LOCALCORR_ONNX_BACKEND=onnx. This setting only affects ONNX export; eager execution is unchanged.
• Mode compatibility: Models configured with corr_mode="localcorr_triton" cannot be exported with the default "tensorrt" backend (the fused Triton kernel raises an error on export). Either set the backend to "onnx" or initialize the model with corr_mode="localcorr_onnx".
• localcorr_onnx mode: Always exports using standard ONNX ops without touching global state, while running the fast fused Triton kernel in eager mode. This produces plugin-free ONNX and TensorRT engines at the cost of a slower correlation subgraph than the fused plugin.
• Standard-ops design: Rather than calling grid_sample directly over multi-channel feature maps, the standard-ops path contracts channels via matmul and gathers displacement windows from single-channel correlation planes, which is significantly faster (see design notes in localcorr/onnx_generic.py).

TensorRT inference (NVIDIA)

Models exported with LocalCorrCudaPyramid nodes require the TensorRT runtime to load the compiled plugin library before deserializing the engine (after building the cuda target):

import ctypes
import tensorrt as trt

trt_logger = trt.Logger(trt.Logger.WARNING)
trt.init_libnvinfer_plugins(trt_logger, "")
ctypes.CDLL(
    "localcorr/cuda_pyramid/build/liblocalcorr_cuda_pyramid_plugin.so",
    mode=ctypes.RTLD_GLOBAL,
)

For command-line tools such as trtexec, you can pass the library via LD_PRELOAD or --staticPlugins / --dynamicPlugins.

OpenVINO inference (Intel)

After building the sycl extension, register the custom op that maps the custom_ops::LocalCorrCudaPyramid ONNX node to the SYCL kernel:

import localcorr.openvino_ext as lcorr_ov  # Import BEFORE openvino:
# preloads the oneAPI OpenCL runtime so the SYCL extension finds the
# OpenCL CPU platform.
import openvino as ov

core = ov.Core()
lcorr_ov.add_localcorr_extension(core)  # registers the custom op
model = core.read_model("model.onnx")    # LocalCorrCudaPyramid resolves

Notes:

• SYCL device selection: The extension computes on the SYCL GPU device by default (auto: GPU preferred, CPU fallback when no GPU is available). Set LOCALCORR_SYCL_DEVICE=cpu|gpu|default|auto to override.
• OpenVINO device plugins: The OpenVINO GPU and NPU plugins cannot execute the custom op directly (Reference nodes are a CPU-plugin capability), so load the OpenVINO model on CPU or AUTO (which delegates the custom op and dispatches the SYCL kernel onto the GPU). For standalone NPU deployment, export with the standard-ops backend (localcorr_onnx).
• oneAPI CPU workaround: The OpenCL CPU runtime in oneAPI 2026.1 crashes or hangs on SYCL sub-group shuffles; the CPU execution paths avoid them (see localcorr/sycl_pyramid/localcorr_cuda_pyramid_kernel.dp.cpp).
• Symbol visibility: The oneAPI runtime preload in localcorr.openvino_ext intentionally uses RTLD_LOCAL; loading those libraries with RTLD_GLOBAL breaks the OpenVINO GPU plugin's JIT compiler (see module docstring for details).

Components

localcorr_triton

Triton kernels, JIT-compiled at first use (pip install triton, no build step):

• triton/pyramid_corr.py: RAFT-style multi-level local correlation around sampled coordinates (LocalCorrTriton per-level, and LocalCorrTritonFused computing all pyramid levels in one launch).
• triton/spatial_corr.py: PWCNet-style single-level spatial correlation (TritonSpatialCorrelationSampler), equivalent to SpatialCorrelationSampler(kernel_size=1, patch_size, padding=0).

localcorr_cuda_pyramid

An AOT-compiled CUDA kernel and TensorRT IPluginV3 plugin that computes all correlation pyramid levels in one kernel launch (using templated fast paths with warp-shuffle reductions; see localcorr/cuda_pyramid/localcorr_cuda_pyramid_kernel.h). During ONNX export, each correlation call is emitted as a single LocalCorrCudaPyramid node resolved by the TensorRT ONNX parser.

The same plugin serves both correlation styles:

• RAFT-style: levels = N (pyramid), x-major displacement order.
• PWCNet-style: levels = 1, y-major displacement order (channel = dy_idx * patch_size + dx_idx).

In eager mode (corr_mode="localcorr_cuda"), the kernel is invoked through ctypes on PyTorch's current CUDA stream, paired with a fused backward kernel.

localcorr_sycl_xpu (PyTorch on Intel GPUs)

Exposes the fused SYCL kernel to PyTorch as a torch.xpu (SyclExtension) extension. Models running on torch.xpu devices compute the entire correlation pyramid in a single fused kernel launch on the Intel GPU, ordered on PyTorch's XPU stream for both forward and backward passes.

localcorr_sycl_pyramid (OpenVINO on Intel)

A SYCL port of the fused kernel paired with an OpenVINO custom-operation extension that maps custom_ops::LocalCorrCudaPyramid ONNX nodes to the SYCL kernel. See OpenVINO inference for usage.

Plugin interface

ONNX node specification (domain custom_ops):

LocalCorrCudaPyramid(fmap1, fmap2_packed, coords) -> corr
  attributes: radius (int), levels (int), dilation (int, default 1),
              ymajor (int, default 0)
  fmap1:        [B, H1, W1, C]                     (NHWC, fp16 or fp32)
  fmap2_packed: [B, C * sum_i h_i*w_i]             (NHWC levels packed,
                                                    h_i = H1>>i, w_i = W1>>i,
                                                    batch-major)
  coords:       [B, 2, H1, W1]                     (planar, x in channel 0,
                                                    level-0 scale)
  corr:         [B, levels, (2*radius+1)^2, H1, W1]

The displacement channel order is d = iy + rd*ix by default (RAFT) or d = ix + rd*iy when ymajor=1 (PWCNet / SpatialCorrelationSampler order). dilation controls the step between consecutive displacement positions (matching dilation_patch in spatial correlation samplers).

Tests

Unit tests run on CPU by default; tests requiring a CUDA device, an Intel GPU, or optional packages (onnx, onnxruntime, openvino) are skipped automatically when unavailable. From the repository root:

pytest tests

(Prefer pytest over python -m pytest to avoid any local directory shadowing with the triton package; from the parent directory, pytest localcorr/tests also works.)

Acknowledgments

Most of the code in this repository was written by GLM-5.3, a large language model trained by Z.ai, with human guidance, testing, and review. Honestly, I wouldn't have known how to write custom CUDA, SYCL, Triton, TensorRT, and OpenVINO C++ extensions all by myself.

The implementations are based on:

• alt_cuda_corr from the RAFT project — the original RAFT-style local correlation kernels, reimplemented here as a single fused pyramid kernel.
• spatial_correlation_sampler by Clément Pinard — the PWCNet-style spatial correlation package whose semantics (kernel_size=1, patch_size, padding=0) this package is compatible with.