[Bug][Relax][CUDA] Default scheduling drops Take indices after fusion and produces incorrect output

[pinpinpoo]

Expected behavior

For x: float16[4,4] and ids: int32[1], the fused graph should compute:

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:

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:

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:

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.

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