[ABI] Introduce ShapeView Minimize TensorObj exposure#67
Conversation
Summary of ChangesHello @tqchen, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request significantly refines the Application Binary Interface (ABI) for Tensor objects by introducing a new Highlights
Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
There was a problem hiding this comment.
Code Review
This pull request introduces ShapeView as a lightweight, non-owning view for tensor shapes to reduce memory allocations and copies. It also refactors TensorObj to minimize its ABI exposure by removing internal caching and data members, making the C++ API more stable. Additionally, several convenient helper functions (data_ptr, ndim, numel) have been added to ffi::Tensor. The changes are well-motivated and the implementation is solid. I've identified a couple of areas for improvement regarding type safety and code duplication, which are detailed in the comments.
This PR minimizes TensorObj ABI exposure so C++ api only depends on behavior of the DLTensor field. We also introduce ShapeView to reduce managed copy of shape. The change will make future dependencies on C++ side more stable. We also added a few helper functions such as data_ptr(), ndim(), numel() to the ffi::Tensor.
<!-- .github/pull_request_template.md --> ## 📌 Description The codegen logic for pytorch and tvm should unify after #1641 , and this PR cleans up the related codegen functions in tvm_bindings. Other changes: 1. update tvm-ffi to 0.1.0b11 to incorporate apache/tvm-ffi#67 and apache/tvm-ffi#68 2. rename of source files: `_ops.cu` and `_pybind.cu` renamed to `_binding.cu` 3. remove torch related header include/library linking in ninja files (#1642 (comment)) 4. remove the use of `use_torch_stream` in unittests, they are no longer required after apache/tvm-ffi#68 ## 🔍 Related Issues #1641 ## 🚀 Pull Request Checklist Thank you for contributing to FlashInfer! Before we review your pull request, please make sure the following items are complete. ### ✅ Pre-commit Checks - [ ] I have installed `pre-commit` by running `pip install pre-commit` (or used your preferred method). - [ ] I have installed the hooks with `pre-commit install`. - [ ] I have run the hooks manually with `pre-commit run --all-files` and fixed any reported issues. > If you are unsure about how to set up `pre-commit`, see [the pre-commit documentation](https://pre-commit.com/). ## 🧪 Tests - [ ] Tests have been added or updated as needed. - [ ] All tests are passing (`unittest`, etc.). ## Reviewer Notes cc @MasterJH5574 please let us know what changes do we need to make to help you bump to the latest version of flashinfer in MLC.
Adds a pure tree-walking interpreter for `tvm_ffi.pyast` expression trees. The
evaluator never calls `compile`/`eval`/`exec`, never synthesizes source text,
and never round-trips through the stdlib `ast` module.
- Class-level dispatch table `_EXPR_EVALUATOR_DISPATCH: {pyast-node-type: handler}`
drives `ExprEvaluator.eval`.
- `OperatorDispatch` keyed by `(op_kind, operand_type, operand_index)`; `lookup`
walks operand MRO left-to-right; `invoke` falls back to `_NATIVE_HANDLERS`
(a `{OperationKind: operator-module-callable}` table).
- `And`/`Or`/`IfThenElse`/`Parens`/`ChainedCompare` bypass user dispatch entirely
for short-circuit correctness; unary ops go direct to native; binary ops route
through `dispatch.invoke`. `ChainedCompare` is flattened to pairwise
comparisons combined with `and` semantics.
- Scope is `Mapping[str, Any]`; internally wrapped in a `ChainMap` so
`WalrusExpr` and per-lambda/per-comprehension bindings layer over but do not
mutate the caller's mapping.
- `_build_generator` evaluates the outermost iterable eagerly (matching Python
genexp semantics) and iterates remaining clauses lazily inside a generator.
- `Yield`/`YieldFrom`/`AwaitExpr` are recognized but raise `EvaluationError` on
evaluation — valid tree shapes the evaluator refuses to execute.
- `eval_assign` rejects `Attr`/`Index` targets (evaluator only produces new
local bindings) and supports at most one `StarredExpr` per sequence with
matching trailing/middle capture.
Re-exported from `tvm_ffi` (purely additive, no renames or removals):
- `eval_expr`
- `eval_assign`
- `OperatorDispatch`
- `DEFAULT_DISPATCH`
- `EvaluationError`
- `UndefinedNameError`
None — runtime library additions only.
None. No existing public API is renamed, modified, or removed.
Design doc lives outside the repo; no Sphinx docs added in this commit.
- `tests/python/test_pyast_evaluator.py`: 112 tests covering every handler,
dispatch MRO, short-circuits (`and`/`or`/chained compare), comprehension
laziness (eager outermost iter, lazy body for generator expressions),
starred/double-starred unpacking in calls/lists/dicts, walrus, f-strings
(conversion flags + `format_spec`), `eval_assign` (including starred
middle/end, Attr/Index rejection), and parity with native `eval` across
representative sources. Includes a public-API re-export sanity check.
- `uv run pytest tests/python/test_pyast_evaluator.py` → 112 passed.
- `uv run pytest tests/python/test_pyast_evaluator.py tests/python/test_pyast.py
tests/python/test_pyast_from_py.py` → 548 passed earlier in the session.
- `pre-commit run --files python/tvm_ffi/_pyast_evaluator.py
tests/python/test_pyast_evaluator.py python/tvm_ffi/__init__.py` → ruff-check
and ruff-format pass; ty has pre-existing `pytest` unresolved-import warnings
affecting every test file in the repo; one `invalid-argument-type` for
`ChainMap(dict, Mapping)` at 4 sites suppressed with `ty: ignore`; one
`invalid-assignment` at the walrus-into-ChainMap line suppressed similarly.
No new ty regressions introduced.
- Lambdas with defaults / `*args` / `**kwargs` (evaluator only supports plain
positional `Id` params; the helper strips `:` and `=` suffixes but rejects
`StarredExpr` params).
- Comprehension targets that are `Attr`/`Index`.
- Side-effectful `ChainMap` interaction from external callers.
- No fuzz testing.
- F-string debug specifier (`{x=}`) not explicitly tested.
This reverts commit 563cfd5.
This PR minimizes TensorObj ABI exposure so C++ api only depends on behavior of the DLTensor field.
We also introduce ShapeView to reduce managed copy of shape. The change will make future dependencies on C++ side more stable.
We also added a few helper functions such as data_ptr(), ndim(), numel() to the ffi::Tensor.