Skip to content

[TIRX] Represent buffers as typed variables - #20063

Open
tqchen wants to merge 3 commits into
mainfrom
tvm-buffer-type-var-refactor
Open

[TIRX] Represent buffers as typed variables#20063
tqchen wants to merge 3 commits into
mainfrom
tvm-buffer-type-var-refactor

Conversation

@tqchen

@tqchen tqchen commented Jul 29, 2026

Copy link
Copy Markdown
Member

This refactors TIRx buffers into ordinary ir.Var identities carrying an
immutable 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, BufferType is the single source of the
logical contract and buffer_data(buffer) is the explicit physical-pointer
projection; DeclBuffer records a view's binding.

Simplification

  • Generic Var scoping, substitution, undefined-variable analysis, and identity
    maps now apply directly to buffers.
  • The pointer type is derived from BufferType, so the physical pointer and
    buffer metadata cannot disagree.
  • The common IR path no longer needs a second Buffer object, duplicate
    constructor state, or a parallel buffer-only substitution mechanism.
  • Metadata-changing transforms create a fresh typed Var and bind it explicitly
    at the replacing AllocBuffer or DeclBuffer.

Physical-root tracking remains deliberately local to the passes whose
correctness depends on allocation ownership:

  1. Storage liveness/reuse attributes alias accesses to the owning allocation.
  2. Pointer vector-type inference and rewriting aggregate evidence at that root.
  3. LowerTIRx cleanup projects aliases to the remapped physical allocation.
  4. Thread-allreduce's deferred remapper replaces or removes the owning
    allocation.

Each is a plain pass-local Map<Var, Var> flattened when a DeclBuffer is
visited. Parameters, buffer_map, and AllocBuffer seed roots; an alias to an
unseen typed source is rejected rather than silently treated as a new root.
Simpler passes decode the one-step buffer_data projection locally.

Migration guide

C++

  • Treat BufferVar as a checked view over an ordinary Var.
  • Read logical metadata through buffer->...; use buffer.var() for identity.
  • Use buffer.data() or builtin::buffer_data() only when a physical pointer
    expression is required.
  • To change metadata, rebuild BufferType, create a fresh Var, and bind it at
    the replacement definition site.
  • A transformed DeclBuffer replaces its superseded declaration consistently;
    it must not leave a stale declaration beside the new binding.

Python

  • Buffer constructors return tvm.ir.Var; .ty contains BufferType.
  • Use tvm.tirx.is_buffer_var(value) for runtime discrimination.
  • tvm.tirx.Buffer remains an import and annotation compatibility alias for
    tvm.ir.Var, so isinstance(value, Buffer) matches every Var.
  • Importing tvm.tirx installs the legacy metadata properties on Var.
    Non-buffer Vars raise AttributeError; buffer .dtype keeps its historical
    DataType result.

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

  • Debug build with LLVM 18.1.8 and CUDA: passed.
  • LowerIntrin, C address-generation, padded-vector LLVM, packed CUDA, raw
    pointer, and rank-0/rank-2 access-pointer regressions: 11 passed.
  • Coherent NumPy 2.3.5 reproduced both missed PrimType conversion
    boundaries; the corrected runtime-trace and minimal-LLVM files pass
    8 tests.
  • Broader affected transform, C-host, runtime-trace, minimal-LLVM, and
    commit-3 regression suite: 467 passed, 8 xfailed, 1 xpassed.
  • Matched CUDA 13.1 validation asserted nvrtcVersion() == 13.1 in the test
    process and loaded libnvrtc.so.13.1.115. After the representation-safe
    TensorMap address_of migration, the two complete CUDA-codegen files passed
    376 tests with 159 skipped.
  • Complete TIR and Relax specialization files: 12 passed.
  • Isolated schedule-state regression: 1 passed.
  • H200 paged-attention float32 prefill/decode nvcc regression: 1 passed.
  • Repository-pinned formatting, all pre-commit hooks, and git diff --check:
    passed.

@gemini-code-assist

Copy link
Copy Markdown
Contributor

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

tqchen added 2 commits July 29, 2026 00:14
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
tqchen force-pushed the tvm-buffer-type-var-refactor branch 4 times, most recently from 997887d to 83a9b86 Compare July 29, 2026 07:20
@tqchen
tqchen force-pushed the tvm-buffer-type-var-refactor branch from 83a9b86 to 60ed387 Compare July 29, 2026 09:31
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