tile-ai / tile-ai/tilelang

[BUG][Fuzzer][ice-on-valid-code] `T.q_multiply_shift` with `s=1` and a power-of-2 multiplier aborts with an internal shift-range ICHECK instead of computing the result

Open
#3,003 0 comments 0 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
7.4k
Forks
742
Avg merge
1d 2h
Merged PRs (30d)
108

Description

### Required prerequisites

- [x] I have read the documentation .
- [x] I have searched the [Issue Tracker](https://github.com/tile-ai/tilelang/issues) that this hasn't already been reported. (comment there if it has.)

### What version of TileLang are you using?

0.1.13 (latest release)

### System information

Linux, NVIDIA L40S (sm_89), CUDA 12.8, PyTorch 2.8.0. The abort happens during host-side `LowerIntrin` (before any GPU code is generated), so it is target-independent — no specific GPU is needed to reproduce.

### Problem description

`T.q_multiply_shift(x, y, q, s)` aborts the compiler with an internal shift-range check when the shift `s == 1` and the multiplier `y` is exactly `1 << 30`:

```
InternalError: Check failed: (pb->value >= 0 && pb->value < rtype.bits()) is false:
Shift amount must be non-negative and less than 32 for type int32
```

`s = 1` is a documented, ordinary shift value: the op's docstring types `s` as an unrestricted "Integer shift" (only `q > 0` is required), and the same call with any other multiplier of the same magnitude compiles and runs. The failure is specific to the power-of-2 multiplier value `1 << 30` combined with `s == 1`; nearby shifts (`s = 0, 2, 3, -1`) all compile and return correct results.

Not a regression — see Provenance.

### Reproducible example code

```python
import torch, tilelang, tilelang.language as T

N, Q = 128, 31
POW2 = 1 << 30 # power-of-2 multiplier (represents 0.5)

def build(y, s):
@T.prim_func
def qm(A: T.Tensor((N,), "int32"), O: T.Tensor((N,), "int32")):
with T.Kernel(1, threads=N) as bx:
t = T.get_thread_binding()
O[t] = T.q_multiply_shift(A[t], y, Q, s) # out = round(x*y*2^-s)
return qm

a = torch.randint(-(1 << 20), 1 << 20, (N,), dtype=torch.int32, device="cuda")

# CONTROL: same s=1, a non-power-of-2 multiplier -> compiles and computes correctly
tilelang.compile(build((1 << 30) + 1, 1), out_idx=[1])(a) # -> OK

# TRIGGER: s=1 with the power-of-2 multiplier -> aborts at legalize time
tilelang.compile(build(POW2, 1), out_idx=[1])(a) # -> InternalError (below)
```

### Traceback

```
tvm::left_shift(tvm::PrimExpr, tvm::PrimExpr, tvm::Span)
...
InternalError: Check failed: (pb->value >= 0 && pb->value < rtype.bits()) is false:
Shift amount must be non-negative and less than 32 for type int32
```

(raised from `tilelang.transform.LowerIntrin()` while legalizing the `tir.q_multiply_shift` call.)

### Expected behavior

Compile and return the same value the general (non-power-of-2) path already returns for the same `s`. `s = 1` is a legal shift and the result is computable — the control above, and the same op at `s ∈ {-1, 0, 2, 3}`, all produce correct output; only the `1 << 30, s = 1` combination aborts.

### Additional context

**Root cause.** The `q_multiply_shift` legalization rule mishandles the boundary case `s == 1` in its power-of-2 fast path: it selects a rounding-and-right-shift branch that assumes the shift amount is at least 1, but at `s == 1` that amount is 0, so it emits a left shift by `-1` that the TIR shift builder rejects.

Mechanism (the off-by-one)

In the `tirx.q_multiply_shift` `FLegalize` rule ([intrin_rule.cc L283-L297](https://github.com/tile-ai/tvm/blob/8df8ebd61659505dd7005f4820b030e983e2c005/src/target/intrin_rule.cc#L283-L297)), the multiplier `y == 1 << 30` takes the power-of-2 special case:

```cpp
PrimExpr exp = s - 1;
int exp_val = get_int_value(s) - 1;
if (exp_val > 0) {
// power of 2 is greater than 0, apply left shift.
return x << exp; // s >= 2
} else {
// power of 2 is less than 0, round and then apply right shift.
PrimExpr one = make_const(lp_dtype, 1);
exp = -exp; // s == 1 -> exp == 0
PrimExpr rounding_factor = one << (exp - 1); // 1 << -1 -> ICHECK abort
PrimExpr rounded_t = x + rounding_factor;
return rounded_t >> exp;
}
```

The branch guard is `exp_val > 0` (i.e. `s >= 2`), so the else-branch is entered for every `s <= 1`. Its comment says "power of 2 is less than 0", but that only holds for `s < 1`; at the boundary `s == 1` the exponent `exp` is exactly `0`, and `one << (exp - 1)` is `1 << -1`. The [`<<` shift builder](https://github.com/tile-ai/tvm/blob/8df8ebd61659505dd7005f4820b030e983e2c005/src/target/intrin_rule.cc#L294) const-folds that and its shift-range ICHECK aborts (`Shift amount must be non-negative and less than 32 for type int32`). For `s == 1` the mathematically correct result is `x` (0.5 · 2¹ · x = x), i.e. `x >> 0` with no rounding term; the `exp == 0` case needs no rounding factor.

Trigger boundary (which s trip vs dodge, and the correct-sibling control)

`y = 1<<30, Q = 31`, output compared against the general (non-power-of-2) semantics computed in int64 (`ls=max(s,0)`, `rs=max(-s,0)`, `((x<> (rs+q)`), all cells run this session on sm_89 / 0.1.13:

| `s` | branch | result |
|---|---|---|
| 100 | `x << exp` (exp=99) | **ICE (left shift ≥ 32)** — distinct boundary, see §Generalization |
| 33 | `x << exp` (exp=32) | **ICE (left shift ≥ 32)** — distinct boundary |
| 32 | `x << exp` (exp=31) | compiles, maxerr 0 |
| 3 | `x << exp` | compiles, maxerr 0 |
| 2 | `x << exp` | compiles, maxerr 0 |
| **1** | else, `exp==0` | **ICE (`1 << -1`)** ← this report |
| 0 | else, `exp==1` | compiles, maxerr 0 |
| -1 | else, `exp==2` | compiles, maxerr 0 |
| -5 | else, `exp==6` | compiles, maxerr 0 |

In the `else` (round-and-right-shift) branch, only `s == 1` aborts; every other `s ≤ 0` compiles and matches the int64 reference exactly. This is a **tight-guard / off-by-one boundary** bug: the guard `exp_val > 0` is off by one at the single boundary value `s == 1`. (The `s ≥ 33` aborts are a *separate* unclamped-large-left-shift defect in the sibling `x << exp` branch — same source region, distinct root; recorded in §Generalization.) The control in the repro (same `s == 1`, multiplier `(1<<30)+1`) takes the general path and returns the correct value (maxerr 0 vs the int64 reference), which proves `s == 1` is a legal, computable input rather than a harness artifact.

This is one of three distinct defects in the same `tirx.q_multiply_shift` legalization (`intrin_rule.cc`): (1) this power-of-2 `s == 1` off-by-one; (2) the shared `get_int_value` helper asserts a compile-time-constant multiplier, so a runtime `y` crashes with `broadcast_node != nullptr`; (3) the per-axis op passes an integer `is_lshift_required` straight into a boolean `Select`, tripping an `is_bool()` check. Same file and op family, three separate root causes on different lines — this report covers only (1).

**Suggested fix.** In the `q_multiply_shift` power-of-2 branch (`intrin_rule.cc` L283-L297), handle `exp == 0` so it returns `x` directly (no rounding factor, no shift) — equivalently, widen the left-shift guard from `exp_val > 0` to `exp_val >= 0` so `s == 1` takes `x << exp` with `exp == 0`. Both are small, localized changes.

§17 Generalization (two-level root + 4-axis sweep, all cells run this session on sm_89 / 0.1.13)

**Root, two levels.**
- **Source-level:** the `tirx.q_multiply_shift` `FLegalize` rule, `intrin_rule.cc` L283–L297. The fragile idiom is an unguarded `one << (exp - 1)` (L294) where `exp` can be `0`, producing a shift by `-1` that the TIR `<<` builder's range ICHECK rejects. More broadly the whole rule emits shift amounts with **no range clamp** — neither the `exp==0` underflow (this report) nor a `> 31` overflow is checked.
- **Operator-level:** the `QMultiplyShift` op family that all route through `intrin_rule.cc` — `tirx.q_multiply_shift`, `tirx.q_multiply_shift_per_axis`, and the shared static `QMultiplyShift` helper (L224). The helper computes the same `one << (total_right_shift - 1)` rounding factor (L247), so the same underflow class exists there.

**4-axis findings.**

| axis | cell tested | observed result | same-root? |
|---|---|---|---|
| similar-logic (guard boundary) | pow2 path, `s = 33` and `s = 100` | `InternalError: Shift amount must be … less than 32` — from the *sibling* `x << exp` branch (L288) with `exp = s-1 ≥ 32`; `s = 32` (exp 31) compiles | **distinct-adjacent** (over-large left-shift, no clamp; L288 not L294). Same source region, different defect — recorded, not filed here |
| related-source | shared helper `one << (total_right_shift - 1)` (L247) forced to `total_right_shift == 0` via per_axis `rs=0, q=0` | `InternalError: Shift amount must be … less than 64` (int64 here) — same underflow class | **same root**, but only reachable **off-contract** (docstring requires `q > 0`); B029's face is reachable *in-contract* (`s = 1` is documented-legal), which is why B029 is the fileable instance |
| related-operator | `T.q_multiply_shift_per_axis(x, 1<<30, 0, 0, 31, 1, 0)` (normal int flags) | `InternalError: (condition.dtype().is_bool()) is false` — integer `is_lshift_required` fed into a boolean `Select` | **distinct** root (bool-typed `Select` misuse); this is the draft's noted defect (3). Passing bool-cast flags makes per_axis compile — so it is a separate bug, not this one |
| related-type / scope (vectorized) | pow2 `s = 1` under `T.vectorized(4)` (lanes > 1) | **compiles (dodges the ICE) but silently returns `0` for all 128 inputs** — reference (int64) is in ±10⁶ range: `o[:3]=[0,0,0]` vs `ref[:3]=[-116505,-620886,-956886]`; the identical kernel at `s = 0` gives maxerr 0 (control) | **SAME root**, **silent face** — see below |
| related-type (extra shifts) | pow2 `s = -5` | compiles, maxerr 0 | control — confirms only the `exp==0` boundary is affected in the `else` branch |

**New same-root finding — silent-value face under vectorization.** The same `exp==0` off-by-one that *loudly* aborts in scalar form does **not** abort when the `q_multiply_shift(s=1)` expression is vectorized (`T.vectorized`, lanes > 1): the `1 << -1` const-folds to `0` on the vector/broadcast path instead of tripping the scalar ICHECK, so the rounding factor is poisoned and every lane's output collapses to `0`. Verified: the identical vectorized kernel at `s = 0` matches the int64 reference exactly (maxerr 0), while at `s = 1` all 128 outputs are `0` against a reference spanning ±10⁶. This means the L294 off-by-one is not merely a loud compile abort — in vectorized code it becomes a **silent wrong-value**. This strengthens the report's impact and belongs to the same root cause; folded here rather than filed separately.

**Example-run (PART 1).** There is **no shipped example or test to run** for this op: `git grep q_multiply_shift v0.1.13 -- examples testing` returns no match; the identifier exists only in the language-definition files. The docstring in `tilelang/language/tir/op.py` (v0.1.13) was read directly and types `s` as an unrestricted "Integer shift" (`q` "Needs to be > 0"), confirming `s = 1` is a documented-legal input. Reach is therefore grounded in the public op surface, not in any shipped call site.

**Provenance.** The QMultiplyShift power-of-2 lowering is vendored TVM code in the TileLang TVM fork (`3rdparty/tvm` pinned at `8df8ebd6` for 0.1.13); the boundary bug is present as long as this rule has shipped. Origin in upstream Apache TVM is unverified.

**Dedup.** I searched the open and closed tracker and found no existing report of this defect.

**Reach.** `T.q_multiply_shift` is a documented, publicly exported op (its docstring — verified in `tilelang/language/tir/op.py` at v0.1.13 — types `s` as an unrestricted "Integer shift" and only requires `q > 0`); `s = 1` and a power-of-2 multiplier are ordinary requantize parameters. `git grep q_multiply_shift v0.1.13 -- examples testing` returns **no match** (verified this session): the identifier appears only in the language-definition files (`language/tir/op.py`, `language/tir/ir.py{,i}`, `language/ast/ir.py`), never in a shipped example or test. So no test exercises this legalization path at any `s` — which is why CI is green, and there is no shipped example to run for this op. The trigger is the intersection of the power-of-2 multiplier value (`y == 1<<30`, the only value that enters the fast path) and the single shift `s == 1`; every other `s` on that path compiles (see §Generalization for the exact boundaries).

**Impact.** The trigger is narrow: a power-of-2 multiplier entering the fast path combined with the single boundary shift `s == 1`. In **scalar** form it is a compile-time abort during host-side `LowerIntrin` — loud, caught immediately, blocks that one kernel from building. But the same off-by-one is **not always loud**: when the `q_multiply_shift(s=1)` expression is **vectorized** (lanes > 1), the `1 << -1` const-folds to `0` instead of aborting, and the kernel compiles but silently returns `0` for every element (verified this session against an int64 reference — see §Generalization). So the same root cause has both a loud (scalar) and a silent-wrong-value (vectorized) face. Fixing the `exp == 0` case closes both: a legal `s == 1` requantize input then computes the correct value instead of aborting or producing zeros.

Contributor guide

Open the contributing guide

Research direction

Read 3rdparty/tvm/src/target/intrin_rule.cc around the q_multiply_shift legalization rule at lines 283–297, then review the public operation in tilelang/language/tir/op.py. Run the supplied Python reproduction and verify that the s=1 power-of-two case compiles and matches the general path, including the vectorized behavior; there is no shipped q_multiply_shift test identified in examples or testing.

Written by the indexing model from the issue text.

Assessment

Tech stack
cpp, python
Domain
compilers
Issue type
Bug
Difficulty
3/5
Estimated time
1-2 days
Activity status
Quiet
Clarity
Clearly specified
Newbie friendliness
70/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.