[Bug] TIRx/CUDA: an annotated host-level scalar (`c: T.float32 = ...`) is allocated on the device and then written from the host (SIGSEGV)

[KitKyoD]

Expected behavior

In a T.prim_func compiled for cuda, a scalar assigned at the host level with a type annotation,

c: T.float32 = T.float32(0.0)

should behave like the plain binding c = T.float32(0.0) (the value reaches the kernel as a scalar argument), or the parser / lowering should reject it with a diagnostic. It should not produce a host function that writes into device memory.

Actual behavior

The annotated form is parsed as a one-element mutable cell. With target cuda the host side lowers that cell to a workspace allocation on the device and then stores the initial value into it from the host; the kernel gets an extra pointer parameter and reads the cell through it. Calling the compiled function segfaults with no diagnostic.

Generated code for the reproducer below (main aafeaadc, 2026-10-02):

[Annotated]
  parsed  : c: T.float32 = T.float32(0.0)
  kernel  : extern "C" __global__ void __launch_bounds__(8) main_kernel(float* __restrict__ B_ptr, float* __restrict__ C_ptr, float* __restrict__ c_ptr);
  host IR : allocates workspace | host store of 0.0
            %c = tail call ptr %0(i32 2, i32 %dev_id, i64 4, i32 2, i32 32)      ; %0 = load ptr, ptr @__TVMBackendAllocWorkspace  -> device_type 2 (CUDA), 4 bytes
            store float 0.000000e+00, ptr %c, align 4                            ; host store into that device pointer
[Plain]
  parsed  : c: T.let[T.float32] = T.float32(0.0)
  kernel  : extern "C" __global__ void __launch_bounds__(8) main_kernel(float* __restrict__ B_ptr, float* __restrict__ C_ptr);
  host IR : no workspace | no host store

The kernel body of the annotated variant is C_ptr[threadIdx.x] = B_ptr[threadIdx.x] + c_ptr[0];. On a GPU (RTX A6000, CUDA 13.0, a tree based on main e0ed4aad4 of 2026-09-26, built with USE_CUDA=ON) the annotated variant dies with SIGSEGV (return code -11) at the call, while the plain variant and an inlined constant run and return the expected values. The generated CUDA source and the two host-IR lines above are the same on main aafeaadc, which I checked without a GPU (see Environment).

Where it comes from, as far as I can read it:

• python/tvm/script/parser/transpile.py, visit_AnnAssign: x: X.ty = value is rewritten to X.decl_mutable_cell_(value, ty=X.ty, ...).
• python/tvm/tirx/script/ir_builder/parser_protocol.py, decl_mutable_cell_: a primitive annotation allocates local_scalar(dtype) (a one-element buffer in local scope) and stores the initializer into it.
• src/tirx/transform/lower_tvm_builtin.cc: BuiltinLower::Build takes the device type from the function's target (kind->default_device_type, 2 for cuda), and the allocation lowering keeps a stack allocation only for kDLCPU + global scope + small size; everything else becomes TVMBackendAllocWorkspace(device_type, device_id, bytes, ...). So the host-level scalar cell of a cuda function is allocated on the device, while the store that initializes it stays in host code.

Environment

• TVM main aafeaadc (2026-10-02), Python 3.12, LLVM 19.1.7, Linux x86_64. This build has USE_CUDA=OFF: the CUDA code generator is compiled unconditionally, so tvm.compile(mod, target="cuda") still produces the CUDA source and the host LLVM IR, which is all the reproducer inspects by default.
• The crash itself: USE_CUDA=ON build of a tree based on main e0ed4aad4 (2026-09-26, after the transpiler parser landed), CUDA 13.0.88, RTX A6000 (sm_86), Python 3.12. The generated code there is byte-identical to the aafeaadc output above.

Steps to reproduce

Save the script below and run python repro_host_scalar_device_workspace.py (no GPU needed; it prints the kernel signature and the host-IR lines). With a CUDA device, python repro_host_scalar_device_workspace.py --run also calls both functions; the annotated one segfaults.

# Minimal reproducer: an annotated scalar assigned in the host region of a TIRx prim_func becomes
# a one-element allocation that the CUDA host function makes on the DEVICE and then writes from
# the host. Prints the generated kernel signature and the host-side allocate/store lines; runs the
# function only when asked (needs a CUDA device; it segfaults there).
#
#   python repro_host_scalar_device_workspace.py            # prints the evidence, no device needed
#   python repro_host_scalar_device_workspace.py --run      # also calls the function (needs a GPU)
import re
import sys

import tvm
from tvm.script import ir as I
from tvm.script import tirx as T

N = 8


@I.ir_module
class Annotated:
    @T.prim_func
    def main(B: T.Buffer((N,), "float32"), C: T.Buffer((N,), "float32")):
        T.func_attr({"tirx.noalias": True})
        c: T.float32 = T.float32(0.0)                 # annotated host-level scalar
        for i in T.thread_binding(N, thread="threadIdx.x"):
            C[i] = B[i] + c


@I.ir_module
class Plain:
    @T.prim_func
    def main(B: T.Buffer((N,), "float32"), C: T.Buffer((N,), "float32")):
        T.func_attr({"tirx.noalias": True})
        c = T.float32(0.0)                            # plain assignment: parsed as a binding
        for i in T.thread_binding(N, thread="threadIdx.x"):
            C[i] = B[i] + c


def show(name, mod):
    lib = tvm.compile(mod, target="cuda")
    cuda = lib.mod.imports[0].inspect_source()
    sig = next(l.strip() for l in cuda.splitlines() if "__global__" in l)
    ll = lib.mod.inspect_source("ll")
    lines = ll.splitlines()
    # the host loads the workspace-allocator function pointer, then calls it a few lines later:
    #   %k = load ptr, ptr @__TVMBackendAllocWorkspace
    #   %c = tail call ptr %k(i32 <device_type>, i32 %dev_id, i64 <bytes>, i32 <dtype code>, i32 <bits>)
    regs = [m.group(1) for l in lines
            for m in [re.match(r"\s*(%\w+) = load ptr, ptr @__TVMBackendAllocWorkspace", l)] if m]
    alloc = [l.strip() for l in lines
             if any(re.match(rf"\s*%\w+ = tail call ptr {re.escape(r)}\(", l) for r in regs)]
    store = [l.strip() for l in lines if re.match(r"\s*store float 0\.0", l)]
    print(f"[{name}]")
    print("  parsed  :", next(l.strip() for l in mod.script().splitlines() if re.match(r"\s*c\b", l)))
    print("  kernel  :", sig)
    print("  host IR :", "allocates workspace" if alloc else "no workspace", "|",
          "host store of 0.0" if store else "no host store")
    for l in alloc + store:
        print("           ", l)
    return lib


if __name__ == "__main__":
    print("tvm", tvm.__version__)
    libs = {name: show(name, mod) for name, mod in (("Annotated", Annotated), ("Plain", Plain))}
    if "--run" in sys.argv:
        import numpy as np
        dev = tvm.cuda(0)
        for name, lib in libs.items():
            b = tvm.runtime.tensor(np.arange(N, dtype="float32"), dev)
            c = tvm.runtime.empty((N,), "float32", dev)
            print(f"[{name}] calling ...", flush=True)
            lib(b, c)
            dev.sync()
            print(f"[{name}] ran:", c.numpy().tolist())

Triage

• needs-triage
• backend:cuda
• (component: TIRx script parser / lower_tvm_builtin)

[KitKyoD]

Note for anyone running this on current main: after #20519 (2026-10-02) the buffer annotation is spelled T.Tensor; with T.Buffer replaced by T.Tensor the script prints the same kernel signature and host-IR lines on main 043b57fbe9b8 (checked 2026-10-03).