Skip to content

operators.md

Latest commit

 

History

History
99 lines (69 loc) · 2.95 KB

operators.md

File metadata and controls

99 lines (69 loc) · 2.95 KB
Raw

Operators

InfiniOps operators are C++ classes with generated Python bindings. They share a common dispatch model across devices and backend implementations.

Public Call Shape

Python users call generated functions from infini.ops:

import infini.ops

infini.ops.gemm(a, b, out)

C++ development code calls operator classes:

infini::ops::Gemm::Call(a, b, out);

For calls that need a stream, workspace, or implementation selection, pass Handle and Config explicitly:

infini::ops::Handle handle;
infini::ops::Config config;

config.set_implementation_index(1);
infini::ops::Gemm::Call(handle, config, a, b, out);

Dispatch Model

Each operator has:

  • a base class under src/base/<op>.h
  • zero or more native implementations under src/native/.../ops/<op>/
  • optional PyTorch C++ implementations under src/torch/ops/<op>/
  • generated wrappers and Python bindings under generated/
  • tests under tests/test_<op>.py

Operator<Key, Device, Index> specializations provide concrete implementations. Device selects the backend and Index selects the implementation slot for that operator on that backend.

Implementation Indexes

Implementation indexes are local to an operator and device. Use an explicit index only when the backend exposes multiple implementations for the same operator.

Generated ATen-backed wrappers reserve implementation index 8. Hand-written backend implementations must avoid colliding with existing implementations for the same operator.

Operator Cache

Operator::Call(...) caches constructed operator instances per thread. The cache key includes the config implementation index and tensor/scalar geometry. Tests can call clear_cache() on generated Python operator classes for module isolation.

Code that changes tensor geometry, backend implementation selection, or workspace assumptions should account for this caching behavior.

Adding an Operator

The standard path for a native operator is:

  1. Add the base class in src/base/<op>.h.
  2. Add one or more backend implementations under src/native/.../ops/<op>/.
  3. Add or update generated wrapper inputs if needed.
  4. Add focused tests under tests/test_<op>.py.
  5. Validate a smoke build plus the focused test on every affected backend.

For PyTorch ATen-backed operators, see Adding ATen-backed operators. That page explains the generated backend path and the hand-written ATen backend path.

Smoke Coverage

Smoke builds use an operator allowlist to keep routine validation short:

python -m pip install .[dev] --no-build-isolation --no-deps \
  --config-settings=cmake.define.INFINI_RT_ROOT=/path/to/infini-rt-prefix \
  --config-settings=cmake.define.WITH_CPU=ON \
  --config-settings=cmake.define.INFINI_OPS_SMOKE_BUILD=ON

Then run:

python -m pytest tests -m smoke -q

Use full builds and broader tests for shared dispatch, wrapper generation, backend infrastructure, or high-risk operator changes.