tile-ai / tile-ai/tilelang

[BUG][Fuzzer][wrong-code] `T.max`/`T.min` on float constants silently folds order-dependently on NaN instead of computing the commutative result

Open
#2,882 1 comment 0 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
7.4k
Forks
745
Avg merge
1d 1h
Merged PRs (30d)
104

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.)
- [x] I have tried the latest version of TileLang.

### What version of TileLang are you using?

Reproduced on 0.1.9 (tag `v0.1.9`, `441c3b06`; bundled TVM submodule `0e15b274`). The offending lines are unchanged on `main` today — the tile-ai/tvm fork's `const_fold.h` still uses `std::max`/`std::min` on floats (see Root cause) — so this is expected to reproduce on 0.1.13 as well.

### System information

NVIDIA L40S (sm_89), CUDA 13.0, PyTorch 2.13.0, Python 3.13. The wrong value is produced host-side at compile time (constant folding), so it is architecture-independent; the GPU is used only to read back the folded constant.

### Problem description

`T.max`/`T.min` on two float **compile-time constants** returns a result that depends on argument order when one operand is `NaN`: `T.max(1.0, nan)` folds to `1.0`, but `T.max(nan, 1.0)` folds to `nan` (same for `T.min`). The kernel compiles with no error and the wrong value is baked into the emitted CUDA as a constant.

The identical operation is commutative when the operands are *not* compile-time constants: with the NaN fed through a tensor, `max(1.0, n)` and `max(n, 1.0)` both return `1.0` on the device. So the compile-time fold of `T.max`/`T.min` disagrees with the runtime lowering of the same op, and contradicts itself under operand swap.

Not a regression — the folding path has used `std::max`/`std::min` since these lines existed; it fails identically wherever the two constants reach the folder.

Emitted CUDA + runtime / C++ controls (all run on 0.1.9)

Constant case — the wrong values are folded into the emitted constants:

```c++
Out[0] = 0x1p+0f/*1.000000e+00*/; // T.max(1.0, nan)
Out[1] = CUDART_NAN_F; // T.max(nan, 1.0)
```

Same op with the NaN passed through a tensor (`n = A[0]`) — lowers to `max(...)`, commutative:

```c++
Out[0] = max(0x1p+0f, n); // -> 1.0
Out[1] = max(n, 0x1p+0f); // -> 1.0
```
Device output for `max(1,NaN)`, `max(NaN,1)`, `min(1,NaN)`, `min(NaN,1)`: `[1.0, 1.0, 1.0, 1.0]` — the runtime path has one well-defined answer.

The two library calls involved:
```
std::max(1.0f, NaN) isnan=0 std::max(NaN, 1.0f) isnan=1 <- order-dependent (the fold)
fmaxf(1.0f, NaN) isnan=0 fmaxf(NaN, 1.0f) isnan=0 <- order-independent (runtime)
```
Finite-constant control: `T.max(1.0, 2.0)` and `T.max(2.0, 1.0)` both fold to `2.0` — only a NaN operand triggers it.

### Reproducible example code

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

@T.prim_func
def main(Out: T.Tensor((2,), "float32")):
with T.Kernel(1, threads=1) as bx:
nan = float('nan')
Out[0] = T.max(1.0, nan) # folds to 1.0
Out[1] = T.max(nan, 1.0) # folds to nan <- same op, operands swapped

out = tilelang.compile(main, out_idx=[0])()
print(out.tolist()) # -> [1.0, nan] (order-dependent; both should be 1.0)
```

### Traceback

No traceback — the kernel compiles and runs to completion; the result is silently wrong and order-dependent.

### Expected behavior

Both orderings should agree, matching the runtime lowering of the same op (CUDA `fmaxf`/`fminf`, i.e. IEEE-754 `maxNum`/`minNum`: with exactly one NaN operand, return the non-NaN operand; order-independent). The emitted-code control above shows the runtime path already returns `1.0` both ways, so there is a well-defined correct target — `T.max(1.0, nan)` and `T.max(nan, 1.0)` should both fold to `1.0`.

### Additional context

**Root cause.** The compile-time folder does not use IEEE min/max semantics for the float case — it calls C++ `std::max`/`std::min`, which return the *first* argument on an unordered (NaN) comparison, so the fold depends on operand order. Concretely, [`TryConstFold`](https://github.com/tile-ai/tvm/blob/0e15b274bce8b46f971abf5ac390e844aa6acee5/src/arith/const_fold.h#L343) does `return FloatImm(rtype, std::max(fa->value, fb->value));` and [`TryConstFold`](https://github.com/tile-ai/tvm/blob/0e15b274bce8b46f971abf5ac390e844aa6acee5/src/arith/const_fold.h#L332) the same with `std::min`. The runtime lowering instead emits `max(...)`/`min(...)` which resolve to CUDA `fmaxf`/`fminf` (commutative, NaN-dropping); only the constant-operand path takes the `std::max`/`std::min` branch and diverges.

**Suggested fix.** One direction is to make the float fold match the runtime `fmaxf`/`fminf` semantics — return the non-NaN operand when exactly one operand is NaN (and NaN when both are) — rather than calling `std::max`/`std::min` directly, so the folded value equals what the lowered op computes. The integer branches are unaffected. (This lives in the bundled `tile-ai/tvm` fork's `const_fold.h`, so the fix is in the vendored TVM, not tilelang proper.)

**Provenance.** The `std::max`/`std::min` float fold is long-standing code in the bundled TVM fork; present at least since the 0.1.9 submodule pin (`0e15b274`) and unchanged on the fork's `main` today. I did not trace the exact introducing PR (shallow clone), so origin before that is unverified.

**Dedup.** Searched the tracker (open+closed) for max/min/nan/const-fold/commutative. Distinct from #2638 (folds the *algebraic* identities `x*0→0` / `x−x→0` in the simplifier, a different rewrite and node type) and #2697 (the *reduce* write-back drops NaN under `clear=False`, not scalar `T.max`/`T.min` const-folding). No issue covers order-dependent NaN folding of `T.max`/`T.min`.

**Reach.** Triggering requires both operands of `T.max`/`T.min` to be compile-time-constant floats with one being NaN. `T.max`/`T.min` appear in 60 example files (softmax/clamp/relu), but a `grep` of `examples/` finds none passing a constant NaN (or `float('inf')`/literal) to them — they take runtime operands, which lower to the commutative `fmaxf`/`fminf` and are unaffected. So no shipped example trips it; it fires only when a NaN sentinel constant reaches the folder (e.g. a NaN init/clamp bound written as a literal, or one the simplifier proves constant).

Contributor guide

Open the contributing guide

Research direction

Start with the bundled TVM fork's src/arith/const_fold.h, especially TryConstFold and TryConstFold, and run the Python reproducer from the issue. Compare constant folding with the emitted runtime fmaxf/fminf behavior; done means both operand orders with one NaN produce the same non-NaN result without changing integer folding.

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
55/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.