pinpinpoo opened a new issue, #20421:
URL: https://github.com/apache/tvm/issues/20421
### Expected behavior
For `x: float16[4,4]` and `ids: int32[1]`, the fused graph should compute:
```python
y = (x + np.float16(1))[ids, :]
```
Enabling fusion should preserve the selected row.
### Actual behavior
The unfused path matches an independent NumPy reference. The fused path
produces incorrect output for nonzero indices.
With an input whose rows contain float16-rounded values `0`, `1/3`, `2/3`,
and `1`, respectively, I observed:
| ids | Expected / OFF | ON |
|---|---|---|
| [0] | [1, 1, 1, 1] | [1, 1, 1, 1] |
| [2] | [1.666015625, 1.666015625, 1.666015625, 1.666015625] | [1, 1, 1, 1] |
| [3] | [2, 2, 2, 2] | [1, 1, 1, 1] |
The 4x4 case reproduced in three independent processes. Each process tested
four inputs (one row-encoded input and three random inputs), indices 0/2/3, and
five repeated executions. OFF passed all 60 checks; ON failed all 40 checks
using nonzero indices. A 1x4 control passed on both paths.
This case was reduced from an extracted official FuseTIR test graph with a
4096x4096 input. The reduction changes only shape extents and preserves
explicit `T.int64` expressions, read/write regions, the `Primitive` function
attribute, and dataflow structure.
The saved IR points to default scheduling. Immediately after `FuseTIR`, the
index is still used:
```python
T_take_intermediate[v_ax0, v_ax1] = Out_intermediate[input_ids[v_ax0], v_ax1]
```
After `ApplyDefaultSchedule`, the body no longer reads `input_ids`. The
schedule trace records `Fallback` accepting this function. The generated CUDA
kernel is:
```cpp
extern "C" __global__ void __launch_bounds__(1024) fused_func_kernel(
half* __restrict__ T_take_intermediate_ptr,
half* __restrict__ input_embeds_ptr) {
if (((int)threadIdx.x) < 16) {
T_take_intermediate_ptr[(((int)threadIdx.x) & 3)] =
input_embeds_ptr[((int)threadIdx.x)] + __float2half_rn(1.0f);
}
}
```
The kernel has no index argument and lets four different input rows write to
each output address. The output above is the observed result; the conflicting
writes mean that the selected row is not reliable.
### Environment
- TVM: `0.26.dev0+source.8f328e8`
- Commit: `8f328e802cfe5e41fcc8f5c17e7582b1c28bfce4`
- GPU: NVIDIA RTX A6000 (`sm_86`)
- NVIDIA driver: `550.120`
- CUDA compiler: `12.4.131`
- LLVM: 15, host target `{"kind": "llvm", "mcpu": "generic"}`
- Linux x86_64, kernel `6.8.0-52-generic`, glibc 2.35
- Python: 3.11.16; NumPy: 2.4.6
### Steps to reproduce
Save the following as `repro_add_take_4x4.py` and run in the TVM CUDA
environment:
```bash
for trial in 0 1 2; do
python repro_add_take_4x4.py "$trial"
done
```
OFF skips `FuseOps` and `FuseTIR`; ON runs the CUDA pipeline with both
passes enabled. The script uses the same input arrays for both paths and saves
the IR after fusion and scheduling, plus every output comparison. The explicit
integer widths and primitive-function context are retained from the failing
graph.
```python
from pathlib import Path
import json
import sys
import numpy as np
import tvm
from tvm import relax, s_tir
from tvm.relax.backend.cuda import pipeline
from tvm.script import ir as I, relax as R, tirx as T
SOURCE = r'''
@I.ir_module
class Module:
@T.prim_func(private=True, s_tir=True)
def add(A: T.Buffer((T.int64(4), T.int64(4)), 'float16'), Out:
T.Buffer((T.int64(4), T.int64(4)), 'float16')):
for i, j in T.grid(T.int64(4), T.int64(4)):
with T.sblock('add'):
vi, vj = T.axis.remap('SS', [i, j])
T.reads(A[vi, vj])
T.writes(Out[vi, vj])
Out[vi, vj] = A[vi, vj] + T.float16(1)
@T.prim_func(private=True, s_tir=True)
def take(A: T.Buffer((T.int64(4), T.int64(4)), 'float16'), B:
T.Buffer((T.int64(1),), 'int32'), T_take: T.Buffer((T.int64(1), T.int64(4)),
'float16')):
for ax0, ax1 in T.grid(T.int64(1), T.int64(4)):
with T.sblock('T_take'):
v_ax0, v_ax1 = T.axis.remap('SS', [ax0, ax1])
T.reads(A[B[v_ax0], v_ax1], B[v_ax0])
T.writes(T_take[v_ax0, v_ax1])
T_take[v_ax0, v_ax1] = A[B[v_ax0], v_ax1]
@R.function(private=True)
def fused_func(input_ids: R.Tensor((1,), dtype='int32'), input_embeds:
R.Tensor((4, 4), dtype='float16')) -> R.Tensor((1, 4), dtype='float16'):
R.func_attr({'Primitive': 1})
cls = Module
with R.dataflow():
lv = R.call_tir(cls.add, (input_embeds,), out_ty=R.Tensor((4,
4), dtype='float16'))
gv = R.call_tir(cls.take, (lv, input_ids), out_ty=R.Tensor((1,
4), dtype='float16'))
R.output(gv)
return gv
@R.function
def main(input_ids: R.Tensor((1,), dtype='int32'), input_embeds:
R.Tensor((4, 4), dtype='float16')) -> R.Tensor((1, 4), dtype='float16'):
cls = Module
with R.dataflow():
gv: R.Tensor((1, 4), dtype='float16') =
cls.fused_func(input_ids, input_embeds)
R.output(gv)
return gv
'''
def build(source, enabled, target, directory):
directory.mkdir(parents=True, exist_ok=True)
mod = tvm.script.from_source(source, extra_vars=dict(I=I, R=R, T=T),
s_tir=True)
mod = relax.transform.ConvertToDataflow()(mod)
with target, tvm.transform.PassContext(opt_level=3):
steps = pipeline.library_dispatch_passes(target) +
pipeline.legalize_passes(target)
for i, transform in enumerate(steps):
name = str(transform.info.name).split(".")[-1]
if not enabled and name in ("FuseOps", "FuseTIR"):
continue
mod = transform(mod)
if name in ("FuseTIR", "ApplyDefaultSchedule"):
(directory / f"{i:02d}_{name}.py").write_text(mod.script())
for transform in pipeline.dataflow_lower_passes(target) +
pipeline.finalize_passes(target):
mod = transform(mod)
return relax.build(mod, target=target, relax_pipeline=None,
tir_pipeline="default")
def main():
trial = int(sys.argv[1]) if len(sys.argv) > 1 else 0
out = Path(f"add_take_repro_{trial}")
device = tvm.cuda(0)
assert device.exist
target = tvm.target.Target(
{"kind": "cuda", "arch": "sm_" +
str(device.compute_version).replace(".", "")},
host={"kind": "llvm", "mcpu": "generic"},
)
vms = {name: relax.VirtualMachine(build(SOURCE, name == "on", target,
out/name), device)
for name in ("off", "on")}
encoded = np.broadcast_to((np.arange(4, dtype=np.float32)/3)[:, None],
(4, 4)).astype("float16").copy()
datasets = [("row_encoded", encoded)] + [
(f"random_{seed}", np.random.default_rng(seed).uniform(.1, .5, (4,
4)).astype("float16"))
for seed in range(3)
]
observations = []
failures = 0
for label, x in datasets:
dx = tvm.runtime.tensor(x, device=device)
for index in (0, 2, 3):
ids = np.array([index], dtype="int32")
di = tvm.runtime.tensor(ids, device=device)
# Preserve float16 Add rounding before indexing; independent of
TVM.
expected = (x.astype("float32") +
np.float32(1)).astype("float16")[ids, :]
for repeat in range(5):
values = {}
for name in (("off", "on") if repeat % 2 == 0 else ("on",
"off")):
vm = vms[name]
vm.set_input("main", di, dx)
vm.invoke_stateful("main")
device.sync()
got = vm.get_outputs("main").numpy().copy()
assert got.shape == expected.shape and got.dtype ==
expected.dtype
values[name] = got
np.testing.assert_array_equal(values["off"], expected)
passed = bool(np.isclose(values["on"], expected, rtol=.02,
atol=.002).all())
failures += not passed
observations.append(dict(dataset=label, index=index,
repeat=repeat,
input=x.tolist(), expected=expected.tolist(),
off=values["off"].tolist(),
on=values["on"].tolist(), on_passed=passed))
if label == "row_encoded" and repeat == 0:
print("index:", index, "expected:", expected, "OFF:",
values["off"], "ON:", values["on"])
(out/"observations.json").write_text(json.dumps(observations, indent=2))
print(f"OFF: all {len(observations)} checks passed; ON:
{failures}/{len(observations)} failed")
if __name__ == "__main__":
main()
```
### Triage
- needs-triage
- backend:cuda
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
