Skip to content

[BugFix][Language] Accept integer per-axis shift flags - #3372

Draft
Metis-dot wants to merge 1 commit into
tile-ai:mainfrom
Metis-dot:fix/per-axis-integer-shift-flags
Draft

Metis-dot wants to merge 1 commit into
tile-ai:mainfrom
Metis-dot:fix/per-axis-integer-shift-flags

Conversation

@Metis-dot

Copy link
Copy Markdown

Summary

Refs #3006.

Normalize integer is_lshift_required and is_rshift_required flags to nonzero boolean expressions in TileLang's existing operator-adapter layer. The documented integer flags otherwise reach a Select condition unchanged and abort q_multiply_shift_per_axis legalization with an internal is_bool() check.

Existing boolean flags retain their IR unchanged. Add target-independent legalization/arithmetic regressions for constant and symbolic signed/unsigned flags, plus C-target CPU compilation and numerical execution regressions.

Scope

This fixes calls through the modern TileLang language facade and backend dialects that share it. Raw tvm.tirx calls and the legacy AST surface are unchanged; no TVM submodule or native C++ changes are included.

Validation

  • Official TileLang 0.1.15 wheel, Python 3.12, CPU PyTorch 2.14.1, apache-tvm-ffi 0.1.12
  • Final focused regressions: baseline 26 failed / 4 passed; patched 30 passed
  • C-target CPU kernels compiled and executed with integer and existing boolean flags, matching an independent integer reference exactly
  • Separate runtime check with int64 flags loaded from a tensor, including zero, negative, and high-bit values, matched the nonzero reference
  • Existing CPU atomic suite: 25 passed
  • Explicit SM80 CUDA source-only lowering: integer flags fail before the patch and lower successfully after it; boolean controls lower in both cases
  • Repository format.sh --files, all applicable pre-commit hooks, pinned Ruff 0.16.6, and git diff --check passed

The tests used an overlay containing the exact modified Python operator file over the wheel's unchanged native libraries, not a rebuilt checkout. The wheel's relevant intrin_rule.cc is byte-identical to the current pinned TVM source. No GPU execution, NVCC compilation, LLVM runtime, or full-suite validation is claimed. A neighboring CUDA-marked test file could not collect because existing GPU feature probing aborts without a CUDA driver.

Provenance

The implementation and tests are AI-generated and were independently checked in a separate AI review pass. This draft is available for maintainer review.

@coderabbitai

coderabbitai Bot commented Oct 2, 2026

Copy link
Copy Markdown
Contributor

Important

Draft PR not reviewed

Draft PRs are not automatically reviewed by default.

  • Trigger a manual review

To automatically review draft PRs, update your CodeRabbit configuration:

reviews:
  auto_review:
    drafts: true
  • Autopilot · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

Autopilot is currently an internal CodeRabbit preview.


Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

Comment @coderabbitai help to get the list of available commands.

@github-actions

github-actions Bot commented Oct 2, 2026

Copy link
Copy Markdown

👋 Hi! Thank you for contributing to the TileLang project.

Please remember to run pre-commit run --all-files in the root directory of the project to ensure your changes are properly linted and formatted. This will help ensure your contribution passes the format check.

We appreciate you taking this step! Our team will review your contribution, and we look forward to your awesome work! 🚀

This branch has not been deployed

No deployments
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