perf: reuse ATen tensors and optimize PyTorch fallback wrappers

[baominghelly]

Summary

• Preserve the originating at::Tensor as an opaque source handle when a Python tensor enters the bindings.
• Reuse that tensor in generated PyTorch fallback wrappers instead of rebuilding an at::Tensor with at::from_blob on every dispatch.
• Keep the existing metadata-based from_blob conversion for tensors without a source handle.
• Read dtype, device, shape, and strides directly from the C++ at::Tensor.
• Compile the hot-path PyTorch bridge as a separate optimized object target.
• Compile generated PyTorch wrappers with a configurable optimization level, defaulting to -O3 instead of the previous fixed -O0.
• Add code-generation and build-structure regression coverage for the optimized bridge path.

Motivation

Generated PyTorch fallback wrappers previously reconstructed an at::Tensor for every tensor argument on every call. Although tensor data remained zero-copy, repeatedly creating the wrapper added substantial host-side dispatch overhead for small and latency-sensitive operators.

This change carries the original at::Tensor through the InfiniOps Tensor wrapper and reuses its existing TensorImpl. The core tensor type keeps the handle type-erased, so it does not gain a PyTorch dependency.

The bridge runs once per tensor argument, so it is built as a separate optimized object target. Generated wrappers are also optimized by default. INFINI_OPS_TORCH_OPT_LEVEL accepts 0, 1, 2, 3, or s, allowing builds to override the default 3 when required.

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
• 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

The shared PyTorch fallback conversion and generated-wrapper build paths change for every backend. Local build, correctness, and performance validation was performed on CPU and NVIDIA; the other vendor environments were not locally available.

Smoke Test Result

Build:

CUDACXX=/usr/local/cuda/bin/nvcc \
CMAKE_BUILD_PARALLEL_LEVEL=64 \
python3 -m pip install --force-reinstall --no-build-isolation --no-deps -v . \
  --config-settings=cmake.define.WITH_CPU=ON \
  --config-settings=cmake.define.WITH_NVIDIA=ON \
  --config-settings=cmake.define.WITH_TORCH=ON \
  --config-settings=cmake.define.INFINI_RT_ROOT=/data/shared/baoming/workplace/new_ops/infinirt-prefix \
  --config-settings=cmake.define.INFINI_OPS_TORCH_OPT_LEVEL=3

Torch wrapper optimization: -O3
Successfully built InfiniOps
Successfully installed InfiniOps-0.1.0

Generator and build-structure unit tests:

python3 -m pytest tests/test_generate_torch_ops.py tests/test_generate_wrappers.py -q
25 passed in 1.13s

NVIDIA generated fallback suite:

CUDA_VISIBLE_DEVICES=4 python3 -m pytest tests/test_torch_ops.py \
  --devices nvidia -q --tb=short \
  -k "not mse_loss and not smooth_l1_loss and not soft_margin_loss"

1499 passed, 3145 skipped, 54 deselected, 23 warnings in 14.47s

mse_loss, smooth_l1_loss, and soft_margin_loss are excluded because of known mismatches in the generic Python reference harness. These mismatches are unrelated to this change and reproduce independently.

Test Results on Supported Platforms

Platform	Affected	Build / Smoke Result	Full Result / Notes
NVIDIA	Yes	CPU/NVIDIA/Torch O3 wheel build passed	Generated fallback suite: 1499 passed, 3145 skipped, 54 deselected
CPU	Yes	Included in local wheel build	Targeted fallback conversion checks passed
Iluvatar	Yes	Not tested locally	Shared PyTorch fallback path; platform environment unavailable
MetaX	Yes	Not tested locally	Separate host-compiler bridge path covered structurally; platform environment unavailable
Cambricon	Yes	Not tested locally	Shared PyTorch fallback path; platform environment unavailable
Moore	Yes	Not tested locally	Separate host-compiler bridge path covered structurally; platform environment unavailable
Ascend	Yes	Not tested locally	Shared PyTorch fallback path; platform environment unavailable

Benchmark / Performance Impact

Command:

CUDA_VISIBLE_DEVICES=4 python3 scripts/sweep_ntops.py

Environment and methodology:

• NVIDIA A100-SXM4-80GB
• PyTorch 2.6.0+cu124, CUDA 12.4
• FP16, standard LLM shapes
• 126 measured operator/shape configurations across 43 operators
• Same output buffers and reference convention before and after

Final result relative to the original fixed-O0, non-reuse implementation:

Metric	Before	After
Median InfiniOps latency	28.30 us	10.155 us
Median PyTorch reference latency	5.835 us	5.315 us
Median dispatch gap	20.77 us	4.14 us
Median speed ratio vs. PyTorch reference	0.210x	0.557x
Configurations below 0.6x	108 / 126	86 / 126

Median InfiniOps latency decreased by 64.12%, corresponding to approximately 2.79x throughput at the wrapper level. Normalizing for the reference-side timing difference gives an improvement of approximately 60.61%. The median dispatch gap decreased by approximately 80.07%.

Notes for Reviewers

Please focus on:

• Ownership and lifetime of the type-erased source handle in Tensor.
• The invariant that tensors not originating from pybind still use the existing metadata-based from_blob path.
• Keeping src/torch/pybind11_.cc outside the generated Torch wrapper target so the dispatch bridge remains optimized.
• The configurable generated-wrapper optimization level and its default -O3 behavior.
• The separate host-compiler object path used by MetaX and Moore builds.

The excluded loss-operator cases are generic reference-harness mismatches and are unrelated to tensor reuse or compiler optimization.