[BUG][Fuzzer][ice-on-valid-code] `T.pow` on `bfloat16` fails to compile with an `InternalError` instead of running (as `float16` does)
- 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.)
### What version of TileLang are you using?
`0.1.11` (release commit `cd37ed5`, bundling `tile-ai/tvm` submodule `ec7f7bd9`).
### System information
```
GPU: NVIDIA H100 80GB HBM3 (sm_90)
Python: 3.12.3
PyTorch: 2.12.1+cu130 (CUDA 13.0)
TileLang: 0.1.11 (tag v0.1.11 = cd37ed5fc35ae7a60a1277c8eb49028174ac51e6)
```
Arch-independent: the defect is a frontend type-guard that fires in `tvm::pow` before any target codegen, so it does not depend on the GPU. Confirmed on H100 `sm_90` (this run); other arches are INFERRED (not run).
### Problem description
`T.pow(x, y)` on `bfloat16` operands fails to compile with
```
tvm.error.InternalError: Check failed: (x.dtype().is_float()) is false: power only applies to float
```
The identical kernel with `float16` operands compiles and runs. The `bfloat16` case is the one reproduced here (run on H100 this session). By inspection of the guard, `float8` and the integer types fail the same `is_float()` check and are also rejected before lowering (source-inferred, not separately run this session), whereas `float32`/`float16` are accepted.
### Reproducible example code
```python
import tilelang
import tilelang.language as T
import torch
M, N = 8, 8
@tilelang.jit
def powk(A, B):
A: T.Tensor((M, N), "bfloat16")
B: T.Tensor((M, N), "bfloat16")
C = T.empty((M, N), "bfloat16")
with T.Kernel(M, threads=N) as bm:
for j in T.serial(N):
C[bm, j] = T.pow(A[bm, j], B[bm, j])
return C
A = torch.rand(M, N, device="cuda", dtype=torch.bfloat16) + 0.5
B = torch.rand(M, N, device="cuda", dtype=torch.bfloat16) + 0.5
out = powk(A, B) # InternalError: power only applies to float
# The identical kernel with dtype "float16" compiles and runs.
```
### Traceback
```pytb
File ".../tilelang/language/tir/op.py", line 3088, in pow
return _tvm_op.pow(x, y, span)
File ".../3rdparty/tvm/python/tvm/tirx/op.py", line 2983, in pow
return _ffi_api._OpPow(x, y, span)
File "/project/3rdparty/tvm/src/tirx/op/op.cc", line 853, in tvm::PrimExpr tvm::pow(PrimExpr, PrimExpr, Span)
tvm.error.InternalError: Check failed: (x.dtype().is_float()) is false: power only applies to float
```
### Expected behavior
`T.pow` should accept `bfloat16` consistently with `float16` — computing in `float` and converting back — instead of aborting compilation. `float16` already works, so this is a dtype-coverage gap rather than a "pow can't do this" limitation.
### Additional context
**Root cause.** The frontend `pow` builder guards its input dtype with
[`TVM_FFI_ICHECK(x.dtype().is_float()) << "power only applies to float"`](https://github.com/tile-ai/tvm/blob/ec7f7bd92c185d8c0b84a5dc17709a0860aa5714/src/tirx/op/op.cc#L853). `DataType::is_float()` returns [`code() == DataType::kFloat`](https://github.com/tile-ai/tvm/blob/ec7f7bd92c185d8c0b84a5dc17709a0860aa5714/include/tvm/runtime/data_type.h#L154) — it matches only the `kFloat` type code (`float32` and `float16`, which share that code), and is a separate predicate from `is_bfloat()`/`is_float8()`. So `bfloat16` (code `kBFloat`), the `float8` types, and the integer types all fail this guard and are rejected before lowering. The precedent for the intended behavior is 42 lines below in the same file: the `abs` builder admits bf16/tfloat with [`x.dtype().is_float() || x.dtype().is_bfloat() || x.dtype().is_tfloat()`](https://github.com/tile-ai/tvm/blob/ec7f7bd92c185d8c0b84a5dc17709a0860aa5714/src/tirx/op/op.cc#L895).
**Suggested fix (two parts — the second is required, see the runtime note below).**
1. Relax the guard at [`op.cc#L853`](https://github.com/tile-ai/tvm/blob/ec7f7bd92c185d8c0b84a5dc17709a0860aa5714/src/tirx/op/op.cc#L853) to mirror the `abs` precedent:
```cpp
// before:
TVM_FFI_ICHECK(x.dtype().is_float()) << "power only applies to float";
// after:
TVM_FFI_ICHECK(x.dtype().is_float() || x.dtype().is_bfloat() || x.dtype().is_tfloat())
<< "power only applies to float";
```
(Whether `float8`/integer inputs should additionally be admitted, or rejected with a clearer message, is the maintainers' call.) This edit lands in the bundled `tile-ai/tvm` submodule (compiled into `libtvm`), so it requires a from-source build.
2. Add a `bfloat16_t` `hpow` device overload in the JIT header [`src/tl_templates/cuda/common.h`](https://github.com/tile-ai/tilelang/blob/cd37ed5fc35ae7a60a1277c8eb49028174ac51e6/src/tl_templates/cuda/common.h#L43), which today defines `#define hpow powf` and has no bf16 overload. Follow the existing `TL_PATCH` compute-in-float pattern used for [`__habs(bfloat16_t)`](https://github.com/tile-ai/tilelang/blob/cd37ed5fc35ae7a60a1277c8eb49028174ac51e6/src/tl_templates/cuda/common.h#L98) / [`hrsqrt(half_t)`](https://github.com/tile-ai/tilelang/blob/cd37ed5fc35ae7a60a1277c8eb49028174ac51e6/src/tl_templates/cuda/common.h#L103): `bfloat16_t hpow(bfloat16_t x, bfloat16_t y) { return bfloat16_t(powf(float(x), float(y))); }`.
**Why part 2 is needed (runtime-verified).** Relaxing only the guard eliminates the frontend `InternalError` for bf16, but bf16 codegen then lowers `T.pow` to `hpow(bfloat16_t, bfloat16_t)`, which resolves to the `#define hpow powf` macro — there is no bf16 `hpow` overload — so nvcc aborts. With both parts applied, bf16 `T.pow` compiles **and** runs (max abs diff vs a torch fp32 reference = 0.0037, within bf16 tolerance); fp16 and fp32 controls are unaffected.
**Distinct from the math-intrinsic codegen gaps (#2565).** This is a *frontend type-guard* wrong-reject in `tvm::pow` that fires *before* any codegen — unlike the unary math intrinsics (`log`/`sqrt`/`tanh`/…) tracked in the open family issue **#2565**, which are accepted by the frontend and then fail at *CUDA codegen* on a missing `h…` overload. #2565 explicitly scopes `T.pow` out ("a distinct frontend type-guard issue and is out of scope here"); this report is that carved-out complement. The root layers differ (frontend `is_float()` guard here vs codegen template coverage in #2565), so they need separate fixes — though part 2 of the fix below (the bf16 `hpow` device overload) is the *same missing-`h…`-overload class* as #2565 and as the already-closed #2384 (`T.rsqrt`/bf16). Net: same broad "bf16 math coverage" theme as #2565/#2384/#2383, but a different op and a frontend-vs-codegen root that #2565 deliberately left for a separate issue.
Provenance, dedup, and regression scope
- **Provenance / regression scope.** Not a regression — no single introducing PR: the `is_float()`-only guard is long-standing and still present unchanged on the `tile-ai/tvm` default branch `tilelang_main` (verified `op.cc:853` on `tilelang_main` on 2026-07-09), with the bf16-admitting `abs` precedent still sitting at `op.cc:895` right below it. Confirmed live on H100 (`sm_90`), tilelang 0.1.11: `float16` compiles+runs (sample output 0.3857); `bfloat16` aborts at `op.cc:853` with the traceback above. Source confirmed at the pinned SHAs (guard at `op.cc:853`, precedent at `op.cc:895`, `is_float()` at `data_type.h:154`, `#define hpow powf` at `common.h:43`).
- **Dedup.** No existing `pow` / "power only applies" / `hpow` issue on the tracker, open or closed (searched `pow`, `hpow`, `power only applies to float`, `bfloat16 pow` on 2026-07-09 — 0 hits for the pow guard). Related, all distinct:
- **#2565** (OPEN) — "Many `T.*` math intrinsics … fail to compile on bfloat16/float16/float8": the broad codegen-coverage family. It explicitly excludes `T.pow` ("a distinct frontend type-guard issue and is out of scope here"), so this report is the complement, not a duplicate. Different root layer (frontend guard vs codegen overloads).
- **#2384** (CLOSED) — `T.rsqrt`/bf16: same missing-`h…`-overload *codegen* class as part 2 of the fix, but a different op and a codegen (not frontend) root.
- **#2383** (CLOSED) — `T.exp`/bf16 silent: another closed member of the same bf16-coverage family (self-referential `fast_exp`), again codegen-layer, distinct op.
Contributor guide
Research direction
Start with the pow guard in 3rdparty/tvm/src/tirx/op/op.cc and compare it with the nearby abs handling, then inspect src/tl_templates/cuda/common.h for the existing hpow definition and bf16 helper patterns. Rebuild from source and rerun the supplied bfloat16 kernel, checking that it compiles and runs while float16 and float32 controls remain unaffected.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- cpp, python, pytorch
- Domain
- backend, compilers
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Quiet
- Clarity
- Clearly specified
- Newbie friendliness
- 64/100