AMD: `tritonamdgpu-canonicalize-pointers` asserts on `tl.where` between two base pointers (gfx1150)
gfx1150, ROCm 7.2, triton 3.7.0, torch 2.12. `tritonamdgpu-canonicalize-pointers` asserts when `tl.where` selects between two distinct base pointers and the result is then offset and loaded:
```
llvm/Support/Casting.h:560: cast<mlir::ShapedType>: Assertion `isa<To>(Val) && "cast<Ty>() argument of incompatible type!"' failed.
TritonAMDGPUCanonicalizePointers on tt.func @k -> PassManager::run failed
```
Minimal repro:
```python
import torch, triton, triton.language as tl
@triton.jit
def k(out, a, b, D: tl.constexpr, BL: tl.constexpr):
i = tl.program_id(0).to(tl.int64)
o_l = tl.arange(0, BL); o_d = tl.arange(0, D)
p = a + o_l * 0
p = tl.where(o_l == 1, b, p) # pick between two base ptrs per row
v = tl.load(p[:, None] + (i * D + o_d[None, :]))
tl.store(out + o_l[:, None] * D + o_d[None, :], v)
a = torch.randn(4, 64, device='cuda'); b = torch.randn(4, 64, device='cuda')
ptrs = torch.tensor([a.data_ptr(), b.data_ptr()], device='cuda')
out = torch.empty(4, 64, device='cuda')
k[(1,)](out, ptrs[0], ptrs[1], D=64, BL=4)
```
The base pointer is a scalar selected via `arith.select`, so the canonicalizer hits a scalar `tt.ptr` where it expects a `ShapedType`. Forcing `int64` offsets doesn't help. Looks related to the empty-uniformSum path; `amd-canonicalize-pointers-empty-uniformsum.mlir` covers one null case but not the select-between-bases one. Compiles fine on NVIDIA. Real workload is a pointer-table gather over a tuple of tensors, so the bases are genuinely different.
1 条评论