Skip to content

[TIRx] Flatten cuda/trn backend operator folder - #20081

Merged
spectrometerHBH merged 2 commits into
apache:mainfrom
spectrometerHBH:main
Aug 1, 2026
Merged

[TIRx] Flatten cuda/trn backend operator folder#20081
spectrometerHBH merged 2 commits into
apache:mainfrom
spectrometerHBH:main

Conversation

@spectrometerHBH

Copy link
Copy Markdown
Contributor

Summary

Flatten the CUDA/TRN backend Python layout by removing the pass-through operator package layer:

  • tvm.backend.cuda.operator.intrinsicstvm.backend.cuda.intrinsics
  • tvm.backend.cuda.operator.tile_primitivetvm.backend.cuda.tile_primitive
  • tvm.backend.trn.operator.tile_primitivetvm.backend.trn.tile_primitive

The operator package was an empty shell: its __init__.py only re-exported its subpackages, it added a meaningless path segment, and it shadowed the stdlib operator name. No compatibility shim is kept; all in-repo imports (source, tests, docs) are rewritten, and the tvm.tirx.<backend>.* alias redirect picks the new paths up automatically. Cross-package references inside the moved packages are relative imports that resolve unchanged after the one-level move.

Companion kernel-side PR: mlc-ai/tirx-kernels#13.

Test

Full TIRx suite on 8x B200 (16 xdist workers): 2620 passed, 41 skipped, 3 xpassed in 386s; bench_suite --check-imports and registry --cc 10 --strict import gates green; pre-commit green.

Move intrinsics/ and tile_primitive/ out of backend/cuda/operator/
up to backend/cuda/, and tile_primitive/ out of backend/trn/operator/
up to backend/trn/. The operator package was an empty shell that only
re-exported its subpackages, added a path segment, and shadowed the
stdlib operator name.

No compatibility shim: all in-repo imports (source, tests, docs) are
rewritten to the new paths. The tvm.tirx.<backend>.* alias redirect
picks the new paths up automatically, and tirx-kernels is updated in
a companion commit.

Verified with the full TIRx suite on 8x B200: 2620 passed, 41 skipped,
3 xpassed, with both kernel import gates green.
@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.

The gpu docs precheck fails on any non-allowlisted WARNING. The new
backend.rst sections tripped it twice:

- tile_primitive/__init__.py star-imports its submodules without
  __all__, so :imported-members: leaked every transitively imported
  name (tvm.tirx expr/stmt/layout classes, arith.Analyzer,
  dataclasses.replace), producing duplicate-description warnings and
  docutils errors from non-RST docstrings.
- intrinsics/__init__.py defines __all__, and __all__-listed imported
  members are documented even without :imported-members:, so
  PTXDataType got indexed twice (duplicate-description warning).

Drop :imported-members: from the tile_primitive section (empty page,
same as the old operator section) and add :no-index: to the
intrinsics section (content still rendered, not indexed). Verified
locally with a minimal sphinx 8.1.3 build: zero warnings.
@spectrometerHBH
spectrometerHBH merged commit 4a09f4e into apache:main Aug 1, 2026
8 checks passed
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.

2 participants