[REFACTOR] Simplify canonical script operation adapters

docs/arch/fusion.rst

@@ -195,9 +195,9 @@ via ``cls.func_name``:
    @R.function
    def main(x: R.Tensor((10, 20), "float32")):
        with R.dataflow():
-           lv0 = R.call_tir(add, (x, const_1), out_ty=R.Tensor((10, 20), "float32"))
-           lv1 = R.call_tir(exp, (lv0,), out_ty=R.Tensor((10, 20), "float32"))
-           gv = R.call_tir(squeeze, (lv1,), out_ty=R.Tensor((10, 20), "float32"))
+           lv0 = R.call_tir(add, (x, const_1), ty_args=[R.Tensor((10, 20), "float32")])
+           lv1 = R.call_tir(exp, (lv0,), ty_args=[R.Tensor((10, 20), "float32")])
+           gv = R.call_tir(squeeze, (lv1,), ty_args=[R.Tensor((10, 20), "float32")])
            R.output(gv)
        return gv
 

docs/deep_dive/relax/learning.rst

@@ -173,9 +173,9 @@ for the end-to-end model execution. The code block below shows a TVMScript imple
             cls = Module
             n = T.int64()
             with R.dataflow():
-                lv = R.call_tir(cls.linear, (x, w0, b0), out_ty=R.Tensor((n, 256), dtype="float32"))
-                lv1 = R.call_tir(cls.relu, (lv,), out_ty=R.Tensor((n, 256), dtype="float32"))
-                lv2 = R.call_tir(cls.linear, (lv1, w1, b1), out_ty=R.Tensor((n, 10), dtype="float32"))
+                lv = R.call_tir(cls.linear, (x, w0, b0), ty_args=[R.Tensor((n, 256), dtype="float32")])
+                lv1 = R.call_tir(cls.relu, (lv,), ty_args=[R.Tensor((n, 256), dtype="float32")])
+                lv2 = R.call_tir(cls.linear, (lv1, w1, b1), ty_args=[R.Tensor((n, 10), dtype="float32")])
                 R.output(lv2)
             return lv2
 
@@ -213,7 +213,7 @@ Taking one line from the above code as an example:
 
 .. code:: python
 
-    lv = R.call_tir(cls.linear, (x, w0, b0), out_ty=R.Tensor((n, 256), dtype="float32"))
+    lv = R.call_tir(cls.linear, (x, w0, b0), ty_args=[R.Tensor((n, 256), dtype="float32")])
 
 To explain what does ``R.call_tir`` work, let us review an equivalent low-level numpy
 implementation of the operation, as follows:
@@ -241,9 +241,9 @@ Another important element in a relax function is the R.dataflow() scope annotati
 .. code:: python
 
     with R.dataflow():
-        lv = R.call_tir(cls.linear, (x, w0, b0), out_ty=R.Tensor((n, 256), dtype="float32"))
-        lv1 = R.call_tir(cls.relu, (lv,), out_ty=R.Tensor((n, 256), dtype="float32"))
-        lv2 = R.call_tir(cls.linear, (lv1, w1, b1), out_ty=R.Tensor((n, 10), dtype="float32"))
+        lv = R.call_tir(cls.linear, (x, w0, b0), ty_args=[R.Tensor((n, 256), dtype="float32")])
+        lv1 = R.call_tir(cls.relu, (lv,), ty_args=[R.Tensor((n, 256), dtype="float32")])
+        lv2 = R.call_tir(cls.linear, (lv1, w1, b1), ty_args=[R.Tensor((n, 10), dtype="float32")])
         R.output(lv2)
 
 Before we talk about the dataflow block, let us first introduce the concept of **pure** and

docs/deep_dive/relax/tutorials/relax_creation.py

@@ -37,6 +37,7 @@
 # high-level Relax operators using TVMScript.
 
 from tvm import relax, topi
+from tvm.relax import ExternFunc as _ExternFunc
 from tvm.script import ir as I
 from tvm.script import relax as R
 from tvm.script import s_tir as Ts
@@ -94,7 +95,7 @@ def forward(
         cls = RelaxModuleWithTIR
         with R.dataflow():
             lv0 = R.matmul(data, R.permute_dims(w0)) + b0
-            lv1 = R.call_tir(cls.relu, lv0, R.Tensor((n, 128), dtype="float32"))
+            lv1 = R.call_tir(cls.relu, lv0, ty_args=[R.Tensor((n, 128), dtype="float32")])
             lv2 = R.matmul(lv1, R.permute_dims(w1)) + b1
             R.output(lv2)
         return lv2
@@ -253,9 +254,9 @@ def forward(self, x):
     with bb.dataflow():
         lv0 = bb.emit(
             relax.call_dps_packed(
-                "env.linear",
+                _ExternFunc("env.linear"),
                 [x, fc1_weight, fc1_bias],
-                out_ty=relax.TensorType((n, 128), "float32"),
+                ty_args=[relax.TensorType((n, 128), "float32")],
             )
         )
         lv1 = bb.emit_te(topi.nn.relu, lv0)
@@ -264,7 +265,7 @@ def forward(self, x):
             relax.call_tir(
                 tir_gv,
                 [lv1, fc2_weight, fc2_bias],
-                out_ty=relax.TensorType((n, 10), "float32"),
+                ty_args=[relax.TensorType((n, 10), "float32")],
             )
         )
         bb.emit_output(gv)

docs/how_to/tutorials/mix_python_and_tvm_with_pymodule.py

@@ -292,7 +292,7 @@ def main(
             h_bias = R.call_tir(
                 cls.bias_add_tir,
                 (h, b),
-                out_ty=R.Tensor((2, 4), "float32"),
+                ty_args=[R.Tensor((2, 4), "float32")],
             )
             return R.nn.relu(h_bias)
 
@@ -361,8 +361,8 @@ def main(
             x: R.Tensor((4, 8), "float32"),
         ) -> R.Tensor((4, 8), "float32"):
             # The VM calls back into Python for these two ops
-            h = R.call_py_func("layer_norm", (x,), out_ty=R.Tensor((4, 8), "float32"))
-            out = R.call_py_func("silu", (h,), out_ty=R.Tensor((4, 8), "float32"))
+            h = R.call_py_func("layer_norm", (x,), ty_args=[R.Tensor((4, 8), "float32")])
+            out = R.call_py_func("silu", (h,), ty_args=[R.Tensor((4, 8), "float32")])
             return out
 
     mod = HybridVMModule(device=tvm.cpu(0))

docs/tirx/api/cuda.rst

@@ -61,7 +61,7 @@ The current ``Tx.cuda`` surface includes:
        ``get_tmem_addr``, ``cvta_generic_to_shared``,
        ``smem_addr_from_uint64``, ``sm100_2sm_leader_smem_addr``, and ``mov_sreg``
    * - Loads, calls, and diagnostics
-     - ``ldg``, ``func_call``, ``printf``, ``trap_when_assert_failed``,
+     - ``func_call``, ``printf``, ``trap_when_assert_failed``,
        ``nano_sleep``, ``clock64``, and ``ffs_u32``
    * - Numeric conversion and packed math
      - ``half2float``, ``bfloat162float``, ``float22half2``,

docs/tirx/native_basics/cuda/buffers.rst

@@ -444,7 +444,7 @@ to an intrinsic or inline function; ``data`` is the base pointer:
 
 .. code-block:: python
 
-    B[tx] = Tx.cuda.func_call("ld", A.ptr_to([tx]), source_code=SRC, return_type="float32")
+    B[tx] = Tx.cuda.func_call("ld", A.ptr_to([tx]), SRC, ty="float32")
 
 .. code-block:: c++
 

docs/tirx/native_basics/cuda/threads_sync.rst

@@ -76,7 +76,7 @@ Inlining raw CUDA
 -----------------
 
 For something with no intrinsic at all, inject a ``__device__`` function from a
-source string with ``Tx.cuda.func_call(name, *args, source_code=..., return_type=...)``:
+source string with ``Tx.cuda.func_call(name, *args, ..., ty=...)``:
 
 .. code-block:: python
 
@@ -91,7 +91,7 @@ source string with ``Tx.cuda.func_call(name, *args, source_code=..., return_type
         Tx.device_entry()
         bx = Tx.cta_id([1])
         tx = Tx.thread_id([256])
-        B[tx] = Tx.cuda.func_call("my_relu", A[tx], source_code=SRC, return_type="float32")
+        B[tx] = Tx.cuda.func_call("my_relu", A[tx], SRC, ty="float32")
 
 The source is emitted verbatim and the call is wired in:
 

python/tvm/backend/cuda/codegen/schema.py

@@ -132,14 +132,12 @@ def codegen(*args):
         source_code = (
             f"\n{prefix}__forceinline__ __device__ {ret_type} {name}{sig} {{\n{body_str}\n}}\n"
         )
-        kwargs = {"source_code": source_code}
+        kwargs = {}
         if tvm_return_type is not None:
-            kwargs["return_type"] = (
-                tvm_return_type(*args) if callable(tvm_return_type) else tvm_return_type
-            )
+            kwargs["ty"] = tvm_return_type(*args) if callable(tvm_return_type) else tvm_return_type
         elif ret_type != "void":
-            kwargs["return_type"] = _C_TO_TVM_DTYPE.get(ret_type, ret_type)
-        result = cuda_func_call(name, *forward, **kwargs)
+            kwargs["ty"] = _C_TO_TVM_DTYPE.get(ret_type, ret_type)
+        result = cuda_func_call(name, *forward, source_code, **kwargs)
         return (result, list(extra_deps)) if extra_deps else result
 
     codegen.__name__ = f"codegen_{op_name}"

python/tvm/backend/cuda/cpp/asm.py

@@ -260,7 +260,7 @@ def cuda_wait_until(dst, ptr, condition, scope, space, ptx_type, backoff_ns):
             f"{closing} }} while (0)\n"
         )
         operands = (condition, backoff_ns)
-    return cuda_func_call(name, *load_call.args[1:-1], *operands, source_code=source), tags
+    return cuda_func_call(name, *load_call.args[1:-1], *operands, source), tags
 
 
 # =============================================================================
@@ -525,7 +525,7 @@ def _scale(n):
                 src_ptr_in,
                 src_offset,
                 predicate,
-                source_code=source_code,
+                source_code,
             )
         # No predicate — plain cp.async.
         func_name = f"ptx_cp_async_legacy_{ca_or_cg}_{cp_size_v}_{dst_elem_bytes}_{src_elem_bytes}"
@@ -549,7 +549,7 @@ def _scale(n):
             dst_offset,
             src_ptr_in,
             src_offset,
-            source_code=source_code,
+            source_code,
         )
     else:
         raise ValueError(f"cp_async_raw codegen expects 5/6 args, got {len(args)}")

python/tvm/backend/cuda/cpp/builtins.py

@@ -18,18 +18,17 @@
 """Thin wrappers around CUDA builtins and library calls.
 
 Everything here compiles to a plain CUDA C++ expression -- an intrinsic
-(``__ldg``, ``__shfl_xor_sync``, ``atomicAdd``), a type-punning
+(``__shfl_xor_sync``, ``atomicAdd``), a type-punning
 reinterpret, or a small arithmetic helper. Nothing in this module emits
 inline asm; anything that needs a hand-written PTX body lives in
 :mod:`.asm`, and single PTX instructions live in the ``T.ptx`` table.
 
 * warp / CTA reductions (templated butterfly shuffle-XOR)
 * packed float2 / bf16 / fp8 conversions and arithmetic
-* ``__ldg`` scalar and vector loads, atomics, address casts
+* Atomics and address casts
 * thread queries, scheduling hints, and the CUDA-side sync helpers
 """
 
-from tvm import DataType
 from tvm.backend.cuda.op import cuda_func_call
 
 from ..codegen.registry import register_codegen
@@ -92,7 +91,7 @@ def codegen_cuda_warp_reduce(value, op, width):
 
     func_name = f"tvm_builtin_cuda_warp_reduce_{op_str}_{width_int}"
     source_code = _warp_reduce_source(func_name, width_int, step_expr)
-    return cuda_func_call(func_name, value, source_code=source_code, return_type=value.ty)
+    return cuda_func_call(func_name, value, source_code, ty=value.ty)
 
 
 @register_codegen("cuda_cta_reduce")
@@ -125,7 +124,7 @@ def codegen_cuda_cta_reduce(value, op, num_warps, scratch):
         "    return scratch[0];\n"
         "}\n"
     )
-    return cuda_func_call(func_name, value, scratch, source_code=cta_body, return_type=value.ty)
+    return cuda_func_call(func_name, value, scratch, cta_body, ty=value.ty)
 
 
 # =============================================================================
@@ -276,92 +275,6 @@ def codegen_cuda_cta_reduce(value, op, num_warps, scratch):
 )
 
 
-# =============================================================================
-# __ldg — typed read-only cached loads, scalar and vector. Source is ``void*``
-# so callers may pass a typed pointer or handle_add_byte_offset result; the
-# helper casts per ``dtype``.
-# =============================================================================
-
-
-def _int_attr(value):
-    return int(value.value) if hasattr(value, "value") else int(value)
-
-
-_CUDA_LDG_CTYPES = {
-    "int8": "signed char",
-    "uint8": "unsigned char",
-    "int16": "short",
-    "uint16": "unsigned short",
-    "int32": "int",
-    "uint32": "unsigned int",
-    "int64": "long long",
-    "uint64": "unsigned long long",
-    "float16": "half",
-    "bfloat16": "nv_bfloat16",
-    "float32": "float",
-    "float64": "double",
-}
-
-_CUDA_LDG_VECTOR_CTYPES = {
-    "int32": "int",
-    "uint32": "unsigned int",
-    "float32": "float",
-}
-_CUDA_LDG_VECTOR_BASES = {"int32": "int", "uint32": "uint", "float32": "float"}
-
-
-def _cuda_ldg_suffix(dtype: str) -> str:
-    return dtype.replace("float", "f").replace("uint", "u").replace("int", "i")
-
-
-@register_codegen("cuda_ldg")
-def codegen_cuda_ldg(*args):
-    if len(args) == 2:
-        addr, dtype = args
-        dtype = str(DataType(parse_str(dtype)))
-        if dtype not in _CUDA_LDG_CTYPES:
-            raise ValueError(f"Unsupported CUDA __ldg dtype {dtype!r}")
-        c_type = _CUDA_LDG_CTYPES[dtype]
-        func_name = f"tvm_builtin_cuda_ldg_{_cuda_ldg_suffix(dtype)}"
-        source_code = f"""
-__forceinline__ __device__ {c_type} {func_name}(void* src) {{
-    return __ldg(reinterpret_cast<const {c_type}*>(src));
-}}
-"""
-        return cuda_func_call(func_name, addr, source_code=source_code, return_type=dtype)
-
-    if len(args) < 5:
-        raise ValueError(f"cuda_ldg expects 2 args or vector form, got {len(args)}")
-    *dsts, addr, dtype, vec, dst_count = args
-    dtype = str(DataType(parse_str(dtype)))
-    vec = parse_str(vec)
-    dst_count = _int_attr(dst_count)
-    vec_len = int(vec[1:]) if vec else 1
-    if dtype not in _CUDA_LDG_VECTOR_CTYPES:
-        raise ValueError(f"Unsupported vector CUDA __ldg dtype {dtype!r}")
-    if vec not in ("v2", "v4") or dst_count != vec_len or len(dsts) != vec_len:
-        raise ValueError(
-            f"vector CUDA __ldg expects dst_count=len(dsts)=vec_len for v2/v4, "
-            f"got vec={vec!r}, dst_count={dst_count}, len(dsts)={len(dsts)}"
-        )
-    c_type = _CUDA_LDG_VECTOR_CTYPES[dtype]
-    vec_type = f"{_CUDA_LDG_VECTOR_BASES[dtype]}{vec_len}"
-    members = ("x", "y", "z", "w")[:vec_len]
-    func_name = f"tvm_builtin_cuda_ldg_{_cuda_ldg_suffix(dtype)}_{vec}_to_dst{dst_count}"
-    params = ", ".join(f"void* dst{i}" for i in range(vec_len))
-    stores = "\n".join(
-        f"    *reinterpret_cast<{c_type}*>(dst{i}) = v.{member};"
-        for i, member in enumerate(members)
-    )
-    source_code = f"""
-__forceinline__ __device__ void {func_name}({params}, void* src) {{
-    {vec_type} v = __ldg(reinterpret_cast<const {vec_type}*>(src));
-{stores}
-}}
-"""
-    return cuda_func_call(func_name, *dsts, addr, source_code=source_code, return_type="void")
-
-
 # =============================================================================
 # Atomics — templated wrappers around CUDA's ``atomicAdd`` / ``atomicCAS``.
 # =============================================================================

python/tvm/backend/cuda/cpp/instrument.py

@@ -117,8 +117,8 @@ def codegen_cuda_iket_official_event(event_id, source_code, payload=None):
     return cuda_func_call(
         "tvm_builtin_iket_official_event",
         *args,
-        source_code=source_code,
-        return_type="uint32",
+        source_code,
+        ty="uint32",
     )
 
 
@@ -171,7 +171,7 @@ def c_type(dtype: str) -> str:
     printf({fmt_literal}{comma_call_args});
 }}
 """
-    return cuda_func_call(func_name, *args, source_code=source_code)
+    return cuda_func_call(func_name, *args, source_code)
 
 
 device_intrinsic(