Avoid redundant tensor conversion in eager bindings#815
Draft
baominghelly wants to merge 1 commit into
Draft
Conversation
baominghelly
force-pushed
the
perf/fix-eager-implementation-index-latest
branch
from
July 24, 2026 08:55
0c00204 to
df35b86
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
implementation_indexlazily so an explicitly supplied index does not also compute the defaultstd::vector<Tensor>bindings and collision-safe generated local namesRoot 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
DeviceFromPybind11Handleto select the default implementation and later calledTensorFromPybind11Handlefor dispatch.TensorFromPybind11HandlecallsDeviceFromPybind11Handleinternally, 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:
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 passedgit diff --check