Skip to content

feat(ops): add PyTorch-aligned convolution operators - #882

Merged
voltjia merged 1 commit into
masterfrom
feat/add-conv-operators
Aug 6, 2026
Merged

feat(ops): add PyTorch-aligned convolution operators#882
voltjia merged 1 commit into
masterfrom
feat/add-conv-operators

Conversation

@voltjia

@voltjia voltjia commented Aug 5, 2026

Copy link
Copy Markdown
Collaborator

Summary

  • Add public PyTorch-aligned Conv1d, Conv2d, Conv3d, and Convolution operators.
  • Keep backend registrations split by operator while sharing validation, metadata storage, CUDA dispatch, and the kernel under the canonical ops/convolution path.
  • Preserve the deprecated ConvInfinilm public API and tests while reusing the same dimension-independent kernel.
  • Keep same-name std::string and std::vector<int64_t> overloads distinct in generated Python bindings, with focused regression coverage.

Motivation

ConvInfinilm was introduced for InfiniLM integration and is now deprecated. The new operators provide standard public replacements aligned with PyTorch while retaining InfiniOps' trailing explicit output convention.

The shared CUDA implementation previously lived under ops/conv, even though no corresponding public Conv operator existed. This revision adds the real PyTorch-aligned Convolution operator and places the shared implementation under ops/convolution.

API Alignment

InfiniOps API Alignment target Upstream interface InfiniOps C++ adaptation Evidence
conv1d PyTorch Python API conv1d(input, weight, bias=None, stride=1, padding=0, dilation=1, groups=1) -> Tensor conv1d(input, weight, bias, stride, padding, dilation, groups, out); padding is either vector<int64_t> or string PyTorch conv1d
conv2d PyTorch Python API conv2d(input, weight, bias=None, stride=1, padding=0, dilation=1, groups=1) -> Tensor conv2d(input, weight, bias, stride, padding, dilation, groups, out); padding is either vector<int64_t> or string PyTorch conv2d
conv3d PyTorch Python API conv3d(input, weight, bias=None, stride=1, padding=0, dilation=1, groups=1) -> Tensor conv3d(input, weight, bias, stride, padding, dilation, groups, out); padding is either vector<int64_t> or string PyTorch conv3d
convolution PyTorch Python API and ATen schema convolution(input, weight, bias, stride, padding, dilation, transposed, output_padding, groups) -> Tensor convolution(input, weight, bias, stride, padding, dilation, transposed, output_padding, groups, out) PyTorch torch.convolution, fixed ATen schema

The explicit out tensor is placed after inputs and attributes per CONTRIBUTING.md. Numeric scalar/default convenience forms are intentionally not added. String padding accepts valid and same, and same requires unit stride as in PyTorch.

The Convolution interface represents every upstream parameter. This PR intentionally reuses only the existing forward 1D/2D/3D kernel: descriptor construction asserts for transposed=true and for nonzero output_padding values. Transposed convolution implementation is outside this PR's scope.

Type of Change

  • feat - new feature / new operator / new platform
  • fix - bug fix
  • perf - performance improvement (no behavioral change)
  • refactor - code restructuring without behavior change
  • test - adding or fixing tests only
  • docs - documentation only
  • build / ci - build system or CI configuration
  • chore - tooling, formatting, or other non-code changes
  • Breaking change

Platforms Affected

  • CPU (WITH_CPU)
  • NVIDIA (WITH_NVIDIA)
  • Iluvatar (WITH_ILUVATAR)
  • MetaX (WITH_METAX)
  • Cambricon (WITH_CAMBRICON)
  • Moore (WITH_MOORE)
  • Ascend (WITH_ASCEND)
  • PyTorch C++ bindings (WITH_TORCH)
  • Build system / CMake / CI
  • Python bindings / user-facing API

Smoke Test Result

Current revision: 40caf513.

Local interface and formatting checks passed. A fresh NVIDIA build and runtime test is pending because ssh nvidia currently times out before authentication at 172.22.163.40:14735.

The preceding revision 302ffe99 passed an NVIDIA A100-SXM4-80GB smoke build and smoke tests in accelerator-dev/nvidia:latest:

Successfully built InfiniOps
Successfully installed InfiniOps-0.1.0
66 passed, 34 skipped, 8330 deselected in 5.04s

Test Results on Supported Platforms

Platform Affected Current Result Notes
NVIDIA Yes CI and fresh runtime validation pending Preceding revision passed smoke/full builds and 72 focused convolution tests.
Iluvatar Yes CI pending Backend registration uses the shared CUDA implementation; backend-owner runtime validation requested.
MetaX Yes CI pending Backend registration uses the shared CUDA implementation; preceding revision CI passed.
Cambricon No N/A Not affected.
Moore Yes CI pending Backend registration uses the shared CUDA implementation; preceding revision CI passed.
Ascend No N/A Not affected.

Current revision checks:

python -m pytest --noconftest tests/test_generate_wrappers.py -q
23 passed in 0.75s

python scripts/generate_wrappers.py \
  --devices nvidia iluvatar metax moore \
  --ops conv1d conv2d conv3d convolution \
  --strict-ops
passed; emitted all 16 operator/backend provider includes

ruff check <changed Python files>
All checks passed!

ruff format --check <changed Python files>
6 files already formatted

clang-format 21.1.8 --dry-run --Werror <changed C++ files>
passed

git diff --check
passed

Preceding revision focused NVIDIA runtime result:

python -m pytest \
  tests/test_conv1d.py \
  tests/test_conv2d.py \
  tests/test_conv3d.py \
  tests/test_conv_infinilm.py -q -rs
72 passed, 72 skipped, 1 warning in 4.14s

Benchmark / Performance Impact

N/A. The public operators reuse the existing ConvInfinilm CUDA kernel; this PR does not claim a performance change.

Notes for Reviewers

  • ConvInfinilm remains public, deprecated, and covered by its existing tests.
  • Shared CUDA code now lives in src/native/cuda/ops/convolution/kernel.{h,cuh}; there is no unmatched ops/conv directory and no forwarding header.
  • Each backend still registers Conv1d, Conv2d, and Conv3d in separate operator directories. The new Convolution registration is also separate.
  • The wrapper generator change only permits name-based vector fallback when libclang reports the argument as int. This prevents a recognized std::string padding overload from being rewritten as a vector.
  • No torch_ops.yaml or generated artifact is changed.

@voltjia
voltjia force-pushed the feat/add-conv-operators branch from 7521a86 to 302ffe9 Compare August 6, 2026 02:17
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant