Move S-TIR specific node and TensorIntrin out of TIRX - #20378
Merged
Merged
Conversation
tqchen
force-pushed
the
refactor-s-tir-node-ownership
branch
from
September 17, 2026 15:00
a36d52f to
d983615
Compare
tlopex
approved these changes
Sep 17, 2026
spectrometerHBH
approved these changes
Sep 17, 2026
Own block, realization, and matching-region nodes in S-TIR, including reflection, structural callbacks, Python APIs, and serialized identities. Extend inherited traversal and dialect-aware helpers while retaining generic TIRX pass behavior and existing script syntax.
Use tensor_intrin names for the native declaration and registry implementation, with a dedicated Python class module that preserves the existing backend package and public TensorIntrin export.
Keep alignment annotations and their collector with the S-TIR compaction and lowering consumers. Visit block and allocation annotations directly without a TIRX registration bridge.
Replace dialect callback registries with inherited S-TIR analyzer, SSA, path verification and simplification paths. Keep generic structural traversal definition-aware and preserve generic specialization through buffer remapping. Normalize schedulable TE and tensor intrinsic indices locally, keep ordinary dtype transforms after block lowering, and preserve scalar index temporaries when narrowing allocations.
Place StmtSimplifier and its configuration implementation in the tirx namespace and use that shared base from S-TIR. Keep arithmetic analyzer dependencies explicit and preserve existing runtime and pass configuration identities.
Visit match-buffer definitions before read/write regions, changing the former native read/write-before-match order to agree with structural definition-before-use traversal. Generalize the private-function inline exclusion from SBlockRealize roots to every root absent from native TIRX statement dispatch. This also excludes future non-native roots whose binder and naming rules the inliner cannot preserve.
Centralize virtual-thread and asynchronous-wait keys in TIRX and retain S-TIR aliases so both schedulable and lowered consumers use the same protocol names.
Supply default region type and span metadata for the older two-field BufferRegion schema without shifting graph indices or replacing existing typed metadata.
Document dtype-changing statement substitution as a post-lowering TIRX operation while preserving structural type-preserving mappings. Share verifier fallback diagnostics, retain protected path-dispatch thunks, and make the normalizer hooks overridable without adding dialect collector or binding callbacks.
Keep the legacy parse dialect argument as documented compatibility syntax while internal callers rely on per-function attributes. Organize block-aware verification and simplification coverage under S-TIR while retaining core and mixed-dialect checks.
TRN tile implementations and initialization bodies capture buffers and indices from the scope where dispatch inserts them. Defer standalone well-formedness checks for these partial functions while retaining verification of complete parsed functions.
Represent scalar inputs as explicit parameters so intrinsic namespace roundtrips form complete functions under shared definition checks.
Follow the production lowering order when checking auto-copy padding and inspect the resulting core allocation shape.
Remove standalone coverage added for the ownership split and keep existing verification cases in their original suite with S-TIR entry points. Retain the existing API and lowering-phase adaptations.
tqchen
force-pushed
the
refactor-s-tir-node-ownership
branch
from
September 17, 2026 17:12
a4257b5 to
92328fe
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
3 participants
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.
Move SBlock, SBlockRealize, MatchBufferRegion and TensorIntrin into S-TIR ownership, including native definitions, reflection, Python APIs and serialized identities. Preserve shared TensorRegion compatibility for both typed and older object-form buffer-region graphs.
Reuse inherited traversal through explicit S-TIR analyzer, SSA, verification and simplification classes. Generic TIRX traversal uses structural definitions and remapping; storage alignment and schedulable dtype normalization stay in S-TIR. The shared statement simplifier and its configuration implementation belong to TIRX, retaining their existing runtime identities.
Adapt scheduling, TE, Relax and script consumers while preserving current script syntax and pipeline order. Define matched buffers before read/write regions, and limit private-function inlining to natively dispatchable TIRX roots. Type-preserving substitution remains structural before lowering; dtype-changing statement substitution targets core TIRX after block lowering.
Apply common variable, buffer and type well-formedness checks across parsed modules, then apply TIRX-specific execution checks to functions without the S-TIR attribute. This also checks ordinary TIRX functions in mixed modules; the legacy parse flag remains accepted without selecting verification. Separate script ownership remains a follow-up.