[TIRX] Represent buffers as typed variables - #20063
Open
tqchen wants to merge 3 commits into
Open
Conversation
Contributor
|
Caution The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased. |
Make buffer identity an ordinary ir.Var carrying BufferType. The previous parallel Buffer object and pointer-typed data Var could diverge under generic IR substitution and rebuilding; BufferType is now the single source of truth, and buffer_data(buffer) is the explicit physical-pointer projection. This lets standard Var scoping, substitution, undefined-variable analysis, and identity maps handle buffers without a second object hierarchy. It also removes duplicate constructor state and representation-specific remapping from the common path. C++ migration: use BufferVar as a checked Var view, read logical metadata through buffer->..., use buffer.var() for identity, and use buffer.data() or builtin::buffer_data for the pointer projection. Metadata changes rebuild the BufferType and rebind the new Var through AllocBuffer or DeclBuffer. Python migration: constructors return tvm.ir.Var, value.ty carries BufferType, and is_buffer_var(value) is the runtime discriminator. Buffer remains a source-compatibility alias for imports and annotations, so isinstance(value, Buffer) must not be used to distinguish buffer variables.
Keep physical-buffer ownership explicit after Buffer becomes a typed Var. Storage planning, vector type rewriting, cleanup, and all-reduce lowering now maintain pass-local alias-to-root maps only where liveness or allocation remapping requires them; simpler passes decode buffer_data projections locally. Replace superseded DeclBuffer bindings instead of retaining stale declarations, reject aliases whose physical source has not been registered, and repair merged tagged-memory children after their parent allocation is rebuilt. Restore the alias-lifetime and original-output regressions that exercise these invariants. Preserve the Python compatibility surface deliberately: Buffer remains an annotation/import alias, is_buffer_var is the runtime discriminator, legacy metadata properties reject non-buffer Vars with AttributeError, and dtype continues to return DataType. Validated with a clean C++ build, 239 focused passing tests (3 skipped), and the complete TIRx/S-TIR Python suite; remaining broad-suite failures require unavailable LLVM, CUDA, NKI, Torch, or target-feature support.
tqchen
force-pushed
the
tvm-buffer-type-var-refactor
branch
4 times, most recently
from
July 29, 2026 07:20
997887d to
83a9b86
Compare
tqchen
force-pushed
the
tvm-buffer-type-var-refactor
branch
from
July 29, 2026 09:31
83a9b86 to
60ed387
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.
This refactors TIRx buffers into ordinary
ir.Varidentities carrying animmutable
BufferType.Rationale
The old representation kept a standalone Buffer object, its logical metadata,
and a separate pointer-typed data Var. Generic Var substitution and
buffer-specific rebuilding could therefore let identity, metadata, and pointer
type drift apart. In this change,
BufferTypeis the single source of thelogical contract and
buffer_data(buffer)is the explicit physical-pointerprojection;
DeclBufferrecords a view's binding.Simplification
maps now apply directly to buffers.
BufferType, so the physical pointer andbuffer metadata cannot disagree.
constructor state, or a parallel buffer-only substitution mechanism.
at the replacing
AllocBufferorDeclBuffer.Physical-root tracking remains deliberately local to the passes whose
correctness depends on allocation ownership:
allocation.
Each is a plain pass-local
Map<Var, Var>flattened when aDeclBufferisvisited. Parameters,
buffer_map, andAllocBufferseed roots; an alias to anunseen typed source is rejected rather than silently treated as a new root.
Simpler passes decode the one-step
buffer_dataprojection locally.Migration guide
C++
BufferVaras a checked view over an ordinaryVar.buffer->...; usebuffer.var()for identity.buffer.data()orbuiltin::buffer_data()only when a physical pointerexpression is required.
BufferType, create a fresh Var, and bind it atthe replacement definition site.
DeclBufferreplaces its superseded declaration consistently;it must not leave a stale declaration beside the new binding.
Python
tvm.ir.Var;.tycontainsBufferType.tvm.tirx.is_buffer_var(value)for runtime discrimination.tvm.tirx.Bufferremains an import and annotation compatibility alias fortvm.ir.Var, soisinstance(value, Buffer)matches every Var.tvm.tirxinstalls the legacy metadata properties on Var.Non-buffer Vars raise
AttributeError; buffer.dtypekeeps its historicalDataTyperesult.Behavior preservation
The behavior-preserving follow-up restores established LowerTIRx,
FlattenBuffer, BF16, alias-lifetime, and declaration-replacement behavior while
keeping ownership logic pass-local.
The final focused correction completes typed-buffer integration at structural
comparison, retyping, specialization, annotation, external-output, access
detection, projection lowering, printer, and target-codegen boundaries. It
also preserves the physical source handle type for declared aliases and
retains allocation-backed versus parameter-backed identity where lowering
requires that distinction.
Validation
pointer, and rank-0/rank-2 access-pointer regressions: 11 passed.
PrimTypeconversionboundaries; the corrected runtime-trace and minimal-LLVM files pass
8 tests.
commit-3 regression suite: 467 passed, 8 xfailed, 1 xpassed.
nvrtcVersion() == 13.1in the testprocess and loaded
libnvrtc.so.13.1.115. After the representation-safeTensorMap
address_ofmigration, the two complete CUDA-codegen files passed376 tests with 159 skipped.
git diff --check:passed.