Skip to content

Avoid redundant tensor conversion in eager bindings#815

Draft
baominghelly wants to merge 1 commit into
masterfrom
perf/fix-eager-implementation-index-latest
Draft

Avoid redundant tensor conversion in eager bindings#815
baominghelly wants to merge 1 commit into
masterfrom
perf/fix-eager-implementation-index-latest

Conversation

@baominghelly

Copy link
Copy Markdown
Contributor

Summary

  • evaluate the optional implementation_index lazily so an explicitly supplied index does not also compute the default
  • convert the first required Tensor argument once in functional pybind bindings
  • reuse that converted Tensor both to select the default implementation and to call the generated dispatcher
  • support first-argument std::vector<Tensor> bindings and collision-safe generated local names
  • add scalar and vector Tensor code-generation regression coverage

Root cause

The generated functional binding used optional::value_or(DefaultImplementationIndexFor...). Because C++ evaluates function arguments eagerly, the default implementation lookup ran even when Python supplied an explicit index.

For calls without an explicit index, the generated wrapper first called DeviceFromPybind11Handle to select the default implementation and later called TensorFromPybind11Handle for dispatch. TensorFromPybind11Handle calls DeviceFromPybind11Handle internally, so the first Python Tensor's device attributes were parsed twice.

The generated wrapper now converts the first Tensor once, reads the device type from the converted C++ Tensor, and moves the same Tensor into Call<Op>.

Class constructor bindings are unchanged; this optimization targets the functional/eager call path.

Impact

This removes redundant Python/pybind11 attribute access from the common no-index path while preserving the ability to select an implementation explicitly. The explicit and default paths now have effectively equivalent host overhead after accounting for measurement noise.

Performance

Measured on an NVIDIA A100-SXM4-80GB with FP32 16x16 GEMM, 200 warmup calls, 1,000 calls per sample, 15 samples per case, and five interleaved ABBA rounds.

The before build already contained the lazy explicit-index branch, so these numbers isolate the Tensor conversion reuse:

Metric Before After Raw improvement Control-adjusted improvement
Default-index host submit 28.211 us 23.699 us 4.512 us (16.0%) 3.928 us
Default-index amortized end-to-end 28.437 us 23.722 us 4.715 us (16.6%) 4.035 us

The explicit-index path was used as the control for host drift. Every benchmark case produced the expected GEMM result (check = 1.0).

Validation

  • pytest -q tests/test_generate_wrappers.py: 14 passed
  • git diff --check
  • built and loaded before/after GEMM wheels on A100
  • verified generated scalar and vector Tensor bindings reuse the first converted Tensor

@baominghelly
baominghelly force-pushed the perf/fix-eager-implementation-index-latest branch from 0c00204 to df35b86 Compare July 24, 2026 08:55
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