[REFACTOR][S-TIR] Lift dlight into s_tir namespace

[tqchen]

This PR migrates dlight into s_tir namespace, so s_tir related components are closely grouped together.

[gemini-code-assist]

Summary of Changes

Hello @tqchen, I'm Gemini Code Assist¹! 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 undertakes a significant refactoring effort by relocating the dlight module into the tvm.s_tir namespace. This change is designed to enhance the logical grouping of components within the s_tir ecosystem, making the project structure more intuitive and maintainable. The primary impact is seen in updated import paths across the codebase, ensuring all modules correctly reference the new location of dlight.

Highlights

• Module Relocation: The dlight module has been moved from tvm.dlight to tvm.s_tir.dlight to better group s_tir related components.
• Import Path Updates: All references and import statements for dlight across documentation, backend pipelines, and test files have been updated to reflect its new namespace.
• Lazy Loading for dlight: A __getattr__ implementation has been added to tvm.s_tir.__init__.py to lazily load the dlight submodule, optimizing initial module loading.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Changelog

• docs/get_started/tutorials/ir_module.py
  • Updated dlight import from tvm.dlight to tvm.s_tir.dlight.
• docs/how_to/tutorials/cross_compilation_and_rpc.py
  • Updated dlight import from tvm.dlight to tvm.s_tir.dlight.
• docs/how_to/tutorials/customize_opt.py
  • Updated dlight import from tvm.dlight to tvm.s_tir.dlight.
• docs/how_to/tutorials/optimize_llm.py
  • Modified dlight import to be from tvm.s_tir import dlight instead of from tvm import dlight.
• python/tvm/dlight/init.py
  • Renamed to python/tvm/s_tir/dlight/__init__.py.
• python/tvm/dlight/adreno/init.py
  • Renamed to python/tvm/s_tir/dlight/adreno/__init__.py.
• python/tvm/dlight/adreno/base.py
  • Renamed to python/tvm/s_tir/dlight/adreno/base.py.
• python/tvm/dlight/adreno/convolution.py
  • Renamed to python/tvm/s_tir/dlight/adreno/convolution.py.
• python/tvm/dlight/adreno/fallback.py
  • Renamed to python/tvm/s_tir/dlight/adreno/fallback.py.
• python/tvm/dlight/adreno/layout_transform.py
  • Renamed to python/tvm/s_tir/dlight/adreno/layout_transform.py.
• python/tvm/dlight/adreno/pool.py
  • Renamed to python/tvm/s_tir/dlight/adreno/pool.py.
• python/tvm/dlight/adreno/utils.py
  • Renamed to python/tvm/s_tir/dlight/adreno/utils.py.
• python/tvm/dlight/analysis/init.py
  • Renamed to python/tvm/s_tir/dlight/analysis/__init__.py.
• python/tvm/dlight/analysis/common_analysis.py
  • Renamed to python/tvm/s_tir/dlight/analysis/common_analysis.py.
• python/tvm/dlight/analysis/gemv.py
  • Renamed to python/tvm/s_tir/dlight/analysis/gemv.py.
• python/tvm/dlight/base/init.py
  • Renamed to python/tvm/s_tir/dlight/base/__init__.py.
• python/tvm/dlight/base/common_schedules.py
  • Renamed to python/tvm/s_tir/dlight/base/common_schedules.py.
• python/tvm/dlight/base/schedule_rule.py
  • Renamed to python/tvm/s_tir/dlight/base/schedule_rule.py.
• python/tvm/dlight/base/transform.py
  • Renamed to python/tvm/s_tir/dlight/base/transform.py.
• python/tvm/dlight/base/utils.py
  • Renamed to python/tvm/s_tir/dlight/base/utils.py.
• python/tvm/dlight/benchmark/init.py
  • Renamed to python/tvm/s_tir/dlight/benchmark/__init__.py.
• python/tvm/dlight/benchmark/bench.py
  • Renamed to python/tvm/s_tir/dlight/benchmark/bench.py.
• python/tvm/dlight/benchmark/extract.py
  • Renamed to python/tvm/s_tir/dlight/benchmark/extract.py.
  • Updated import path for benchmark_prim_func.
• python/tvm/dlight/benchmark/utils.py
  • Renamed to python/tvm/s_tir/dlight/benchmark/utils.py.
• python/tvm/dlight/cpu/init.py
  • Renamed to python/tvm/s_tir/dlight/cpu/__init__.py.
• python/tvm/dlight/cpu/base.py
  • Renamed to python/tvm/s_tir/dlight/cpu/base.py.
• python/tvm/dlight/cpu/gemv.py
  • Renamed to python/tvm/s_tir/dlight/cpu/gemv.py.
• python/tvm/dlight/gpu/init.py
  • Renamed to python/tvm/s_tir/dlight/gpu/__init__.py.
  • Updated comment to reflect new module path tvm.s_tir.dlight.cuda/rocm/vulkan/metal.
• python/tvm/dlight/gpu/base.py
  • Renamed to python/tvm/s_tir/dlight/gpu/base.py.
• python/tvm/dlight/gpu/fallback.py
  • Renamed to python/tvm/s_tir/dlight/gpu/fallback.py.
• python/tvm/dlight/gpu/gemv.py
  • Renamed to python/tvm/s_tir/dlight/gpu/gemv.py.
• python/tvm/dlight/gpu/general_reduction.py
  • Renamed to python/tvm/s_tir/dlight/gpu/general_reduction.py.
• python/tvm/dlight/gpu/low_batch_gemv.py
  • Renamed to python/tvm/s_tir/dlight/gpu/low_batch_gemv.py.
• python/tvm/dlight/gpu/matmul.py
  • Renamed to python/tvm/s_tir/dlight/gpu/matmul.py.
• python/tvm/dlight/gpu/reduction.py
  • Renamed to python/tvm/s_tir/dlight/gpu/reduction.py.
• python/tvm/dlight/gpu/rmsnorm.py
  • Renamed to python/tvm/s_tir/dlight/gpu/rmsnorm.py.
• python/tvm/dlight/gpu/transpose.py
  • Renamed to python/tvm/s_tir/dlight/gpu/transpose.py.
• python/tvm/relax/backend/adreno/pipeline.py
  • Removed direct tvm.dlight import.
  • Added from tvm.s_tir import dlight as dl as an import-outside-toplevel within legalize_passes.
• python/tvm/relax/backend/cuda/pipeline.py
  • Removed direct tvm.dlight import.
  • Added from tvm.s_tir import dlight as dl as an import-outside-toplevel within legalize_passes.
• python/tvm/relax/backend/dispatch_sort_scan.py
  • Removed dlight from tvm import.
  • Added from tvm.s_tir import dlight as an import-outside-toplevel within apply_dlight_gpu_fallback.
• python/tvm/relax/backend/gpu_generic/pipeline.py
  • Removed direct tvm.dlight import.
  • Added from tvm.s_tir import dlight as dl as an import-outside-toplevel within legalize_passes.
• python/tvm/relax/backend/rocm/pipeline.py
  • Removed direct tvm.dlight import.
  • Added from tvm.s_tir import dlight as dl as an import-outside-toplevel within legalize_passes.
• python/tvm/s_tir/init.py
  • Implemented __getattr__ to lazily load the dlight submodule from .dlight.
• tests/python/codegen/test_target_codegen_cuda_fp8.py
  • Updated dlight import from tvm.dlight to tvm.s_tir.dlight.
  • Updated tvm.dlight.ApplyDefaultSchedule and its arguments to use tvm.s_tir.dlight.
• tests/python/disco/test_callback.py
  • Updated tvm.dlight.ApplyDefaultSchedule to tvm.s_tir.dlight.ApplyDefaultSchedule.
• tests/python/disco/test_ccl.py
  • Updated dlight import from tvm.dlight to tvm.s_tir.dlight.
• tests/python/disco/test_loader.py
  • Updated dlight import from tvm.dlight to tvm.s_tir.dlight.
• tests/python/dlight/test_benchmark.py
  • Renamed to tests/python/s_tir/dlight/test_benchmark.py.
  • Updated import path for benchmark functions.
• tests/python/dlight/test_cpu_gemv.py
  • Renamed to tests/python/s_tir/dlight/test_cpu_gemv.py.
  • Updated dlight import from tvm.dlight to tvm.s_tir.dlight.
• tests/python/dlight/test_gpu_conv.py
  • Renamed to tests/python/s_tir/dlight/test_gpu_conv.py.
  • Updated dlight import from tvm.dlight to tvm.s_tir.dlight.
• tests/python/dlight/test_gpu_fallback.py
  • Renamed to tests/python/s_tir/dlight/test_gpu_fallback.py.
  • Updated dlight import from tvm.dlight to tvm.s_tir.dlight.
• tests/python/dlight/test_gpu_gemv.py
  • Renamed to tests/python/s_tir/dlight/test_gpu_gemv.py.
  • Updated dlight import from tvm.dlight to tvm.s_tir.dlight.
• tests/python/dlight/test_gpu_general_reduction.py
  • Renamed to tests/python/s_tir/dlight/test_gpu_general_reduction.py.
  • Updated dlight import from tvm.dlight to tvm.s_tir.dlight.
• tests/python/dlight/test_gpu_low_batch_gemv.py
  • Renamed to tests/python/s_tir/dlight/test_gpu_low_batch_gemv.py.
  • Updated dlight import from tvm.dlight to tvm.s_tir.dlight.
• tests/python/dlight/test_gpu_matmul.py
  • Renamed to tests/python/s_tir/dlight/test_gpu_matmul.py.
  • Updated dlight import from tvm.dlight to tvm.s_tir.dlight.
• tests/python/dlight/test_gpu_matmul_tensorize.py
  • Renamed to tests/python/s_tir/dlight/test_gpu_matmul_tensorize.py.
  • Updated dlight import from tvm.dlight to tvm.s_tir.dlight.
• tests/python/dlight/test_gpu_reduction.py
  • Renamed to tests/python/s_tir/dlight/test_gpu_reduction.py.
  • Updated dlight import from tvm.dlight to tvm.s_tir.dlight.
• tests/python/dlight/test_gpu_rmsnorm.py
  • Renamed to tests/python/s_tir/dlight/test_gpu_rmsnorm.py.
  • Updated dlight import from tvm.dlight to tvm.s_tir.dlight.
• tests/python/dlight/test_gpu_transpose.py
  • Renamed to tests/python/s_tir/dlight/test_gpu_transpose.py.
  • Updated dlight import from tvm.dlight to tvm.s_tir.dlight.
• tests/python/dlight/test_primitives.py
  • Renamed to tests/python/s_tir/dlight/test_primitives.py.
• tests/python/relax/backend/clml/utils.py
  • Updated dlight import from tvm.dlight to tvm.s_tir.dlight.
• tests/python/relax/nvshmem/test_runtime_builtin_kv_cache_transfer.py
  • Updated dlight import from tvm.dlight to tvm.s_tir.dlight.
• tests/python/relax/test_backend_dispatch_sort_scan.py
  • Modified dlight import to be from tvm.s_tir import dlight instead of from tvm import dlight.
• tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_cpu.py
  • Updated dlight import from tvm.dlight to tvm.s_tir.dlight.
• tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_flashinfer.py
  • Updated dlight import from tvm.dlight to tvm.s_tir.dlight.
• tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_mla_flashinfer.py
  • Updated dlight import from tvm.dlight to tvm.s_tir.dlight.
• tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_mla_tir.py
  • Updated dlight import from tvm.dlight to tvm.s_tir.dlight.
• tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_tir.py
  • Updated dlight import from tvm.dlight to tvm.s_tir.dlight.
• tests/python/relax/test_runtime_builtin_rnn_state.py
  • Updated dlight import from tvm.dlight to tvm.s_tir.dlight.
• tests/python/relax/test_vm_build.py
  • Updated tvm.dlight.ApplyDefaultSchedule to tvm.s_tir.dlight.ApplyDefaultSchedule.
• tests/scripts/release/make_notes.py
  • Added 's_tir': 'S-TIR' mapping for release notes generation.

Activity

• The pull request was created by tqchen with the title "[REFACTOR][S-TIR] Lift dlight into s_tir namespace" and a description stating its purpose to migrate dlight into the s_tir namespace for better grouping.
• No further review comments or activity have been recorded in the provided context.

Using Gemini Code Assist

The 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 /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

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 .gemini/ folder in the base of the repository. Detailed instructions can be found here.

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

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request is a large-scale refactoring that moves the dlight module into the s_tir namespace. This is a good change for better code organization. The changes mostly consist of updating import paths across the codebase, including documentation, tests, and backend pipelines. The use of __getattr__ for lazy loading in tvm.s_tir is a nice touch. The changes look solid. I've found one minor inconsistency in a test file and left a suggestion to improve it.