perf(torch): reuse ATen wrappers in generated operators

[voltjia]

Summary

• Add a source-neutral TargetTensorPool<Adapter> with stable entries and move-only call-scoped leases.
• Generate per-argument ATen wrapper pools for required, optional, and list tensor parameters in generated Torch operators.
• Rebind cached ATen wrappers to the current data pointer without retaining the source tensor, its storage, or a source-native handle in Tensor.
• Extract ATen metadata directly at the Python boundary while preserving the generic Tensor-like attribute protocol and keeping the Python-specific bridge out of libinfiniops.so.

Motivation

Generated Torch fallback operators currently call at::from_blob for every tensor argument on every invocation. This does not copy tensor data, but it repeatedly constructs TensorImpl/Storage wrappers. The Python binding also obtains tensor metadata through repeated Python attribute calls.

PR #798 avoids most of that cost by retaining and reusing the caller's ATen handle. This PR instead caches target-native wrappers behind a compile-time adapter. An X -> InfiniOps -> ATen path can therefore reuse ATen wrappers without putting an X-specific owner or handle into the generic Tensor/operator cache.

Related: #798

Type of Change

• feat - new feature / new operator / new platform
• fix - bug fix
• perf - performance improvement (no intended 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 (requires a ! in the Conventional Commits prefix or a BREAKING CHANGE: footer)

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 generated path also covers Hygon, which is not listed in the template.

Smoke Test Result

Final commit: 2671579bec529394d343b74fbce6bc889dfe75cb
Image: accelerator-dev/nvidia:latest (nvcr.io/nvidia/pytorch:25.12-py3)
Hardware: NVIDIA A100-SXM4-80GB

python -m pip install --force-reinstall --no-build-isolation --no-deps '.[dev]' \
  --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=/home/huangjiacheng/infiniops-pr798-perf-20260715/infinirt-prefix \
  --config-settings=cmake.define.INFINI_OPS_SMOKE_BUILD=ON \
  --config-settings=cmake.define.INFINI_OPS_OPS=add,mul,cast,cat,gemm,matmul,linear,rms_norm,swiglu,silu_and_mul,causal_softmax,causal_softmax_infinilm,abs,clamp,exp,cutlass_scaled_mm,moe_sum,flash_attn_varlen_func,scaled_dot_product_attention,stack \
  --config-settings=cmake.define.INFINI_OPS_TORCH_OPS=abs,clamp,exp,stack
python -m pytest tests -m smoke -q

90 passed, 15 skipped, 12845 deselected in 21.32s

Test Results on Supported Platforms

Platform	Affected	Build / Smoke Result	Full Result / Notes
NVIDIA	Yes	smoke passed	Final-SHA installed-wheel probes passed; full suite not run
Iluvatar	Yes	not run	Shared generated Torch path changed; requires platform CI/owner validation
MetaX	Yes	not run	Custom host-compiler bridge path requires platform CI/owner validation
Cambricon	Yes	not run	Shared generated Torch path changed; requires platform CI/owner validation
Moore	Yes	not run	Custom host-compiler bridge path requires platform CI/owner validation
Ascend	Yes	not run	Shared generated Torch path changed; requires platform CI/owner validation

The same smoke build included CPU. Additional final-SHA checks:

• python -m pytest tests/test_target_tensor_pool.py tests/test_generate_torch_ops.py tests/test_generate_wrappers.py -q: 30 passed in 1.83s.
• Ruff 0.15.22 check and format check passed for all changed Python files.
• clang-format 21.1.8 passed for all changed C++ files.
• Installed-wheel probes passed for CPU Abs; CUDA Abs, optional Clamp, and list Stack; alternating data pointers; bool attention masks; zero-element reuse; and non-contiguous tensors with nonzero storage offsets.
• Non-Torch Tensor-like objects using the original data_ptr/shape/dtype/device/stride() protocol passed required Abs, optional Clamp, and list Stack calls in the Release wheel.
• An unsupported complex Torch dtype was rejected by the existing unsupported-dtype path in a Release child process instead of being silently treated as uint8.
• import infini.ops before import torch passed.
• Allocation-retention probes returned to the same 141312 allocated-byte baseline after releasing 64 MiB (Abs), 128 MiB (Clamp), and 192 MiB (Stack) payloads.
• readelf -d confirmed that libinfiniops.so has no libtorch_python.so dependency.

The full final-SHA suite and non-NVIDIA platform smoke suites were not run. This PR remains a draft for those CI/platform checks.

Benchmark / Performance Impact

Benchmark: generated abs Torch fallback, torch.float16, shapes [1, 4096], [8, 4096], and [1, 8192] on A100-SXM4-80GB. The harness used two independent GPU/CPU placements, four Williams-balanced panels per round, seven timed blocks per shape, and alternating data-pointer correctness checks.

Build	Base	This PR	#798	This PR vs. Base	This PR vs. #798
O0	19.592 us	9.222 us	9.065 us	-52.9%	+1.7%
O3	15.201 us	7.817 us	8.206 us	-48.6%	-4.7%

The performance panel used the pre-publication implementation commit c856ae9; the tested float16 hot path is unchanged in the final commit. These numbers are post-hoc directional sensitivity results. The preregistered primary analysis failed its control-stability checks, so the exact ratios are not presented as definitive latency estimates.

Notes for Reviewers

• Pool reuse assumes an operator instance fixes shape, strides, dtype, and device; Rebind changes only the borrowed data pointer.
• Optional and list leases are deliberately declared before copied at::Tensor containers so the native wrappers are destroyed before their leases are released.
• AtenTensorAdapter installs a non-owning c10::DataPtr; no source tensor or storage owner is retained.
• The pool follows the existing mutable operator-instance concurrency contract and adds no hot-path lock.
• Hand-written Torch operators still use their existing conversion paths; this change targets generated Torch operators.
• Please focus on DataPtr ownership, nested-call lease behavior, list-pool high-water allocation, and the MetaX/Moore system-compiler bridge path.