This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new d68740c178 [FEAT][Script] Emit script-facing void operators through
IRBuilder (#20729)
d68740c178 is described below
commit d68740c178910557c7dcaa95328beb62c303e265
Author: Tianqi Chen <[email protected]>
AuthorDate: Wed Oct 14 15:36:46 2026 -0400
[FEAT][Script] Emit script-facing void operators through IRBuilder (#20729)
Script-facing void operators emit one Evaluate(call) under an active
IRBuilder and return AlreadyEmitted[Call] containing the exact original
call. A shared auto_emit_void_call decorator invokes one native
script-builder hook; statement receipts and explicit evaluate calls
remain single-emission, with source context on both Call and Evaluate.
Raw operator and PTX APIs remain Call constructors even inside builders.
Explicit script exports and generated Call APIs use the decorator while
classes, region frames and callable namespaces retain their interfaces.
Script PTX chains preserve instruction metadata for tile operands; CUDA
patterns, codegen and printer reconstruction use ordinary raw APIs.
---
python/tvm/backend/cuda/__init__.py | 2 +-
python/tvm/backend/cuda/script.py | 153 ++++++++-----
python/tvm/backend/metal/script.py | 31 +--
python/tvm/backend/opencl/__init__.py | 5 +-
python/tvm/backend/opencl/script.py | 22 ++
python/tvm/backend/trn/script.py | 34 +--
python/tvm/ir/op.py | 12 +-
python/tvm/relax/script/ir_builder/op.py | 3 +-
python/tvm/s_tir/script/ir_builder/ir.py | 4 +-
python/tvm/script/ir_builder/base.py | 58 ++++-
python/tvm/script/ir_builder/stmt.py | 19 +-
python/tvm/tirx/script/ir_builder/op.py | 251 ++++++++++++----------
src/script/ir_builder/ir.cc | 15 ++
tests/python/script/test_builder_void_emission.py | 78 +++++++
14 files changed, 459 insertions(+), 228 deletions(-)
diff --git a/python/tvm/backend/cuda/__init__.py
b/python/tvm/backend/cuda/__init__.py
index 244a82809c..98ad35d5a9 100644
--- a/python/tvm/backend/cuda/__init__.py
+++ b/python/tvm/backend/cuda/__init__.py
@@ -95,10 +95,10 @@ def register_backend():
def script_namespaces(**_):
"""Return CUDA-owned TVMScript namespaces."""
from . import script # pylint: disable=import-outside-toplevel
- from .ptx import PTXNamespace # pylint: disable=import-outside-toplevel
from .script import ( # pylint: disable=import-outside-toplevel
NVSHMEMNamespace,
PTXLegacyNamespace,
+ PTXNamespace,
STIRNamespace,
)
diff --git a/python/tvm/backend/cuda/script.py
b/python/tvm/backend/cuda/script.py
index 320b896abe..871da7af12 100644
--- a/python/tvm/backend/cuda/script.py
+++ b/python/tvm/backend/cuda/script.py
@@ -20,6 +20,7 @@ from __future__ import annotations
from tvm import ir as _ir
from tvm.backend.cuda import op as _cuda_op
+from tvm.script.ir_builder.base import auto_emit_void_call as
_auto_emit_void_call
from . import layout as layout
from .axis import Axis as Axis
@@ -31,6 +32,8 @@ from .launch import LaunchCompletionEvent as
LaunchCompletionEvent
from .launch import LaunchConfig as LaunchConfig
from .launch import MemSyncDomainMap as MemSyncDomainMap
from .launch import ProgrammaticEvent as ProgrammaticEvent
+from .ptx.engine import PTXNamespace as _PTXNamespace
+from .ptx.engine import _InstrChain
from .tile.instructions import make_namespace as _make_tensor_namespace
# pylint: disable=protected-access
@@ -43,9 +46,9 @@ class _CpAsyncRaw:
# Legacy variant: takes (dst_ptr, dst_offset, src_ptr, src_offset,
# cp_size). Offsets are folded into the pointers; lowers through the
# raw op below.
- self.legacy = _cuda_op.ptx_cp_async_legacy
+ self.legacy = _auto_emit_void_call(_cuda_op.ptx_cp_async_legacy)
op = _ir.Op.get("tirx.s_tir.cp_async_raw")
- self._construct = _ir.op._make_op_api(op, __name__)
+ self._construct = _auto_emit_void_call(_ir.op._make_op_api(op,
__name__))
def __call__(self, *args, **kwds):
return self._construct(*args, **kwds)
@@ -75,29 +78,35 @@ class PTXLegacyNamespace:
def __init__(self):
# Same lowered asm as T.ptx.mma, but the accumulator doubles as the
# destination and the offsets are explicit.
- self.mma = _cuda_op.ptx_legacy_mma
+ self.mma = _auto_emit_void_call(_cuda_op.ptx_legacy_mma)
# (trans, num, dtype, local_ptr, local_offset, smem_ptr, smem_offset)
- self.ldmatrix = _cuda_op.ptx_legacy_ldmatrix
+ self.ldmatrix = _auto_emit_void_call(_cuda_op.ptx_legacy_ldmatrix)
class CudaWgmmaNamespace:
"""WGMMA companions that are not PTX instructions (pure-C / empty-asm
helpers)."""
def __init__(self):
- self.noop_barrier = _cuda_op.cuda_wgmma_noop_barrier
- self.encode_matrix_descriptor =
_cuda_op.cuda_wgmma_encode_matrix_descriptor
+ self.noop_barrier =
_auto_emit_void_call(_cuda_op.cuda_wgmma_noop_barrier)
+ self.encode_matrix_descriptor = _auto_emit_void_call(
+ _cuda_op.cuda_wgmma_encode_matrix_descriptor
+ )
class CudaTcgen05Namespace:
"""tcgen05 companions that are not PTX instructions (pure-C descriptor
packers)."""
def __init__(self):
- self.encode_matrix_descriptor =
_cuda_op.cuda_tcgen05_encode_matrix_descriptor
- self.encode_instr_descriptor = _ir.op._make_op_api(
- _ir.Op.get("tirx.cuda.tcgen05_encode_instr_descriptor"), __name__
+ self.encode_matrix_descriptor = _auto_emit_void_call(
+ _cuda_op.cuda_tcgen05_encode_matrix_descriptor
+ )
+ self.encode_instr_descriptor = _auto_emit_void_call(
+
_ir.op._make_op_api(_ir.Op.get("tirx.cuda.tcgen05_encode_instr_descriptor"),
__name__)
)
- self.encode_instr_descriptor_block_scaled = _ir.op._make_op_api(
-
_ir.Op.get("tirx.cuda.tcgen05_encode_instr_descriptor_block_scaled"), __name__
+ self.encode_instr_descriptor_block_scaled = _auto_emit_void_call(
+ _ir.op._make_op_api(
+
_ir.Op.get("tirx.cuda.tcgen05_encode_instr_descriptor_block_scaled"), __name__
+ )
)
@@ -105,13 +114,13 @@ class IketNamespace:
"""Frontend-only NVIDIA IKET annotations."""
def __init__(self):
- self.mark = _cuda_op.cuda_iket_mark
- self.range_start = _cuda_op.cuda_iket_range_start
- self.range_end = _cuda_op.cuda_iket_range_end
- self.range_push = _cuda_op.cuda_iket_range_push
- self.range_pop = _cuda_op.cuda_iket_range_pop
- self.sentinel_token = _cuda_op.cuda_iket_sentinel_token
- self.official_event = _cuda_op.cuda_iket_official_event
+ self.mark = _auto_emit_void_call(_cuda_op.cuda_iket_mark)
+ self.range_start = _auto_emit_void_call(_cuda_op.cuda_iket_range_start)
+ self.range_end = _auto_emit_void_call(_cuda_op.cuda_iket_range_end)
+ self.range_push = _auto_emit_void_call(_cuda_op.cuda_iket_range_push)
+ self.range_pop = _auto_emit_void_call(_cuda_op.cuda_iket_range_pop)
+ self.sentinel_token =
_auto_emit_void_call(_cuda_op.cuda_iket_sentinel_token)
+ self.official_event =
_auto_emit_void_call(_cuda_op.cuda_iket_official_event)
TensorMapEncodeTiledAttr = _cuda_op.TensorMapEncodeTiledAttr
@@ -121,65 +130,70 @@ TCGen05InstrDescriptorBlockScaledAttrs =
_cuda_op.TCGen05InstrDescriptorBlockSca
iket = IketNamespace()
wgmma = CudaWgmmaNamespace()
tcgen05 = CudaTcgen05Namespace()
-warp_sum = _cuda_op.cuda_warp_sum
-warp_max = _cuda_op.cuda_warp_max
-warp_min = _cuda_op.cuda_warp_min
-cta_sum = _cuda_op.cuda_cta_sum
-cta_max = _cuda_op.cuda_cta_max
-cta_min = _cuda_op.cuda_cta_min
-func_call = _cuda_op.cuda_func_call
-sm100_2sm_leader_smem_addr = _cuda_op.cuda_sm100_2sm_leader_smem_addr
-timer_init = _cuda_op.timer_init
-timer_start = _cuda_op.timer_start
-timer_end = _cuda_op.timer_end
-timer_finalize = _cuda_op.timer_finalize
-mma_store = _ir.op._make_op_api(_ir.Op.get("tirx.cuda.mma_store"), __name__)
-mma_fill = _ir.op._make_op_api(_ir.Op.get("tirx.cuda.mma_fill"), __name__)
-bmma_sync = _cuda_op.bmma_sync
-atomic_add = _cuda_op.cuda_atomic_add
-atomic_cas = _cuda_op.cuda_atomic_cas
-cta_reduce = _cuda_op.cuda_cta_reduce
-mov_sreg = _cuda_op.cuda_mov_sreg
-wait_until = _cuda_op.cuda_wait_until
-warp_reduce = _cuda_op.cuda_warp_reduce
-__shfl_sync = _cuda_op.__shfl_sync
-__shfl_up_sync = _cuda_op.__shfl_up_sync
-__shfl_down_sync = _cuda_op.__shfl_down_sync
-__shfl_xor_sync = _cuda_op.__shfl_xor_sync
-__activemask = _cuda_op.__activemask
+warp_sum = _auto_emit_void_call(_cuda_op.cuda_warp_sum)
+warp_max = _auto_emit_void_call(_cuda_op.cuda_warp_max)
+warp_min = _auto_emit_void_call(_cuda_op.cuda_warp_min)
+cta_sum = _auto_emit_void_call(_cuda_op.cuda_cta_sum)
+cta_max = _auto_emit_void_call(_cuda_op.cuda_cta_max)
+cta_min = _auto_emit_void_call(_cuda_op.cuda_cta_min)
+func_call = _auto_emit_void_call(_cuda_op.cuda_func_call)
+sm100_2sm_leader_smem_addr =
_auto_emit_void_call(_cuda_op.cuda_sm100_2sm_leader_smem_addr)
+timer_init = _auto_emit_void_call(_cuda_op.timer_init)
+timer_start = _auto_emit_void_call(_cuda_op.timer_start)
+timer_end = _auto_emit_void_call(_cuda_op.timer_end)
+timer_finalize = _auto_emit_void_call(_cuda_op.timer_finalize)
+mma_store =
_auto_emit_void_call(_ir.op._make_op_api(_ir.Op.get("tirx.cuda.mma_store"),
__name__))
+mma_fill =
_auto_emit_void_call(_ir.op._make_op_api(_ir.Op.get("tirx.cuda.mma_fill"),
__name__))
+bmma_sync = _auto_emit_void_call(_cuda_op.bmma_sync)
+atomic_add = _auto_emit_void_call(_cuda_op.cuda_atomic_add)
+atomic_cas = _auto_emit_void_call(_cuda_op.cuda_atomic_cas)
+cta_reduce = _auto_emit_void_call(_cuda_op.cuda_cta_reduce)
+mov_sreg = _auto_emit_void_call(_cuda_op.cuda_mov_sreg)
+wait_until = _auto_emit_void_call(_cuda_op.cuda_wait_until)
+warp_reduce = _auto_emit_void_call(_cuda_op.cuda_warp_reduce)
+__shfl_sync = _auto_emit_void_call(_cuda_op.__shfl_sync)
+__shfl_up_sync = _auto_emit_void_call(_cuda_op.__shfl_up_sync)
+__shfl_down_sync = _auto_emit_void_call(_cuda_op.__shfl_down_sync)
+__shfl_xor_sync = _auto_emit_void_call(_cuda_op.__shfl_xor_sync)
+__activemask = _auto_emit_void_call(_cuda_op.__activemask)
tile = _make_tensor_namespace()
-_ir.op._init_op_api("tirx.cuda", __name__)
+_ir.op._init_op_api("tirx.cuda", __name__, wrapper=_auto_emit_void_call)
class NVSHMEMNamespace:
"""The NVSHMEM intrinsics submodule."""
def __init__(self):
- self.my_pe = _cuda_op.nvshmem_my_pe
- self.n_pes = _cuda_op.nvshmem_n_pes
- self.signal_op = _cuda_op.nvshmem_signal_op
- self.wait_until = _cuda_op.nvshmem_wait_until
- self.quiet = _cuda_op.nvshmem_quiet
- self.fence = _cuda_op.nvshmem_fence
- self.barrier_all = _cuda_op.nvshmem_barrier_all
+ self.my_pe = _auto_emit_void_call(_cuda_op.nvshmem_my_pe)
+ self.n_pes = _auto_emit_void_call(_cuda_op.nvshmem_n_pes)
+ self.signal_op = _auto_emit_void_call(_cuda_op.nvshmem_signal_op)
+ self.wait_until = _auto_emit_void_call(_cuda_op.nvshmem_wait_until)
+ self.quiet = _auto_emit_void_call(_cuda_op.nvshmem_quiet)
+ self.fence = _auto_emit_void_call(_cuda_op.nvshmem_fence)
+ self.barrier_all = _auto_emit_void_call(_cuda_op.nvshmem_barrier_all)
self.getmem_nbi = NVSHMEMGetMemNBINamespace()
self.putmem_nbi = NVSHMEMPutMemNBINamespace()
self.putmem_signal_nbi = NVSHMEMPutMemSignalNBINamespace()
# Canonical Op names are also exposed directly for reconstruction.
for name in ("getmem_nbi", "putmem_nbi", "putmem_signal_nbi"):
for scope in ("warp", "block"):
- setattr(self, f"{name}_{scope}", getattr(_cuda_op,
f"nvshmem_{name}_{scope}"))
+ setattr(
+ self,
+ f"{name}_{scope}",
+ getattr(getattr(self, name), scope),
+ )
class NVSHMEMGetMemNBINamespace:
"""The NVSHMEM GetMemNBI intrinsics submodule."""
def __init__(self):
- self.warp = _cuda_op.nvshmem_getmem_nbi_warp
- self.block = _cuda_op.nvshmem_getmem_nbi_block
+ self.warp = _auto_emit_void_call(_cuda_op.nvshmem_getmem_nbi_warp)
+ self.block = _auto_emit_void_call(_cuda_op.nvshmem_getmem_nbi_block)
+ @_auto_emit_void_call
def __call__(self, *args, **kwds):
return _cuda_op.nvshmem_getmem_nbi(*args, **kwds)
@@ -190,9 +204,10 @@ class NVSHMEMPutMemNBINamespace:
"""The NVSHMEM PutMemNBI intrinsics submodule."""
def __init__(self):
- self.warp = _cuda_op.nvshmem_putmem_nbi_warp
- self.block = _cuda_op.nvshmem_putmem_nbi_block
+ self.warp = _auto_emit_void_call(_cuda_op.nvshmem_putmem_nbi_warp)
+ self.block = _auto_emit_void_call(_cuda_op.nvshmem_putmem_nbi_block)
+ @_auto_emit_void_call
def __call__(self, *args, **kwds):
return _cuda_op.nvshmem_putmem_nbi(*args, **kwds)
@@ -203,10 +218,30 @@ class NVSHMEMPutMemSignalNBINamespace:
"""The NVSHMEM PutMemSignalNBI intrinsics submodule."""
def __init__(self):
- self.warp = _cuda_op.nvshmem_putmem_signal_nbi_warp
- self.block = _cuda_op.nvshmem_putmem_signal_nbi_block
+ self.warp =
_auto_emit_void_call(_cuda_op.nvshmem_putmem_signal_nbi_warp)
+ self.block =
_auto_emit_void_call(_cuda_op.nvshmem_putmem_signal_nbi_block)
+ @_auto_emit_void_call
def __call__(self, *args, **kwds):
return _cuda_op.nvshmem_putmem_signal_nbi(*args, **kwds)
# __call__ corresponds to nvshmem_putmem_signal_nbi
+
+
+class _ScriptInstrChain(_InstrChain):
+ """Preserve PTX instruction metadata while emitting script calls."""
+
+ __slots__ = ()
+
+ def __getattr__(self, name):
+ return _ScriptInstrChain(super().__getattr__(name)._cands)
+
+ __call__ = _auto_emit_void_call(_InstrChain.__call__)
+
+
+class PTXNamespace(_PTXNamespace):
+ """Script-facing PTX chains, also valid as tile instruction operands."""
+
+ def _family(self, token):
+ chain = super()._family(token)
+ return None if chain is None else _ScriptInstrChain(chain._cands)
diff --git a/python/tvm/backend/metal/script.py
b/python/tvm/backend/metal/script.py
index f3f652f4f1..d1aa0e082f 100644
--- a/python/tvm/backend/metal/script.py
+++ b/python/tvm/backend/metal/script.py
@@ -19,27 +19,34 @@
from __future__ import annotations
from tvm.backend.metal import op as _metal_op
+from tvm.script.ir_builder.base import auto_emit_void_call as
_auto_emit_void_call
class MetalNamespace:
"""The Metal intrinsics submodule."""
def __init__(self):
- self.make_filled_simdgroup_matrix =
_metal_op.make_filled_simdgroup_matrix
- self.simdgroup_load = _metal_op.simdgroup_load
- self.simdgroup_store = _metal_op.simdgroup_store
- self.simdgroup_multiply_accumulate =
_metal_op.simdgroup_multiply_accumulate
-
- cooperative_tensor_fill = staticmethod(_metal_op.cooperative_tensor_fill)
- cooperative_tensor_load = staticmethod(_metal_op.cooperative_tensor_load)
- cooperative_tensor_store = staticmethod(_metal_op.cooperative_tensor_store)
+ self.make_filled_simdgroup_matrix = _auto_emit_void_call(
+ _metal_op.make_filled_simdgroup_matrix
+ )
+ self.simdgroup_load = _auto_emit_void_call(_metal_op.simdgroup_load)
+ self.simdgroup_store = _auto_emit_void_call(_metal_op.simdgroup_store)
+ self.simdgroup_multiply_accumulate = _auto_emit_void_call(
+ _metal_op.simdgroup_multiply_accumulate
+ )
+
+ cooperative_tensor_fill =
staticmethod(_auto_emit_void_call(_metal_op.cooperative_tensor_fill))
+ cooperative_tensor_load =
staticmethod(_auto_emit_void_call(_metal_op.cooperative_tensor_load))
+ cooperative_tensor_store = staticmethod(
+ _auto_emit_void_call(_metal_op.cooperative_tensor_store)
+ )
cooperative_tensor_multiply_accumulate = staticmethod(
- _metal_op.cooperative_tensor_multiply_accumulate
+ _auto_emit_void_call(_metal_op.cooperative_tensor_multiply_accumulate)
)
- simd_shuffle = staticmethod(_metal_op.simd_shuffle)
- simd_shuffle_up = staticmethod(_metal_op.simd_shuffle_up)
- simd_shuffle_down = staticmethod(_metal_op.simd_shuffle_down)
+ simd_shuffle = staticmethod(_auto_emit_void_call(_metal_op.simd_shuffle))
+ simd_shuffle_up =
staticmethod(_auto_emit_void_call(_metal_op.simd_shuffle_up))
+ simd_shuffle_down =
staticmethod(_auto_emit_void_call(_metal_op.simd_shuffle_down))
__all__ = ["MetalNamespace"]
diff --git a/python/tvm/backend/opencl/__init__.py
b/python/tvm/backend/opencl/__init__.py
index 97c1adcd2f..5ef3b9897b 100644
--- a/python/tvm/backend/opencl/__init__.py
+++ b/python/tvm/backend/opencl/__init__.py
@@ -16,7 +16,6 @@
# under the License.
"""OpenCL-owned backend hooks."""
-import sys
from pathlib import Path
from tvm_ffi.libinfo import load_lib_ctypes
@@ -55,8 +54,10 @@ def register_backend():
except (OSError, FileNotFoundError, RuntimeError):
pass
register_device_target_detector("opencl", _detect_target_from_device)
+ from . import script
+
_init_op_api("tirx.opencl", __name__)
- builder_op.register_script_namespace("opencl", sys.modules[__name__])
+ builder_op.register_script_namespace("opencl", script)
return None
diff --git a/python/tvm/backend/opencl/script.py
b/python/tvm/backend/opencl/script.py
new file mode 100644
index 0000000000..410e73c557
--- /dev/null
+++ b/python/tvm/backend/opencl/script.py
@@ -0,0 +1,22 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""OpenCL TVMScript operator exports."""
+
+from tvm.ir.op import _init_op_api
+from tvm.script.ir_builder.base import auto_emit_void_call as
_auto_emit_void_call
+
+_init_op_api("tirx.opencl", __name__, wrapper=_auto_emit_void_call)
diff --git a/python/tvm/backend/trn/script.py b/python/tvm/backend/trn/script.py
index 11d335d23b..f16799573c 100644
--- a/python/tvm/backend/trn/script.py
+++ b/python/tvm/backend/trn/script.py
@@ -18,6 +18,8 @@
from __future__ import annotations
+from tvm.script.ir_builder.base import auto_emit_void_call as
_auto_emit_void_call
+
from . import op as _trn_op
@@ -25,22 +27,22 @@ class NKINamespace:
"""The NKI instructions submodule."""
def __init__(self):
- self.load = _trn_op.nki_load
- self.store = _trn_op.nki_store
- self.tensor_copy = _trn_op.nki_tensor_copy
- self.matmul = _trn_op.nki_matmul
- self.activation = _trn_op.nki_activation
- self.activation_reduce = _trn_op.nki_activation_reduce
- self.reciprocal = _trn_op.nki_reciprocal
- self.tensorreduce = _trn_op.nki_tensorreduce
- self.tensortensor = _trn_op.nki_tensortensor
- self.tensorscalar = _trn_op.nki_tensorscalar
- self.tensorscalar_reduce = _trn_op.nki_tensorscalar_reduce
- self.scalar_tensor_tensor = _trn_op.nki_scalar_tensor_tensor
- self.scalar_tensor_scalar = _trn_op.nki_scalar_tensor_scalar
- self.memset = _trn_op.nki_memset
- self.identity = _trn_op.nki_identity
- self.affine_select = _trn_op.nki_affine_select
+ self.load = _auto_emit_void_call(_trn_op.nki_load)
+ self.store = _auto_emit_void_call(_trn_op.nki_store)
+ self.tensor_copy = _auto_emit_void_call(_trn_op.nki_tensor_copy)
+ self.matmul = _auto_emit_void_call(_trn_op.nki_matmul)
+ self.activation = _auto_emit_void_call(_trn_op.nki_activation)
+ self.activation_reduce =
_auto_emit_void_call(_trn_op.nki_activation_reduce)
+ self.reciprocal = _auto_emit_void_call(_trn_op.nki_reciprocal)
+ self.tensorreduce = _auto_emit_void_call(_trn_op.nki_tensorreduce)
+ self.tensortensor = _auto_emit_void_call(_trn_op.nki_tensortensor)
+ self.tensorscalar = _auto_emit_void_call(_trn_op.nki_tensorscalar)
+ self.tensorscalar_reduce =
_auto_emit_void_call(_trn_op.nki_tensorscalar_reduce)
+ self.scalar_tensor_tensor =
_auto_emit_void_call(_trn_op.nki_scalar_tensor_tensor)
+ self.scalar_tensor_scalar =
_auto_emit_void_call(_trn_op.nki_scalar_tensor_scalar)
+ self.memset = _auto_emit_void_call(_trn_op.nki_memset)
+ self.identity = _auto_emit_void_call(_trn_op.nki_identity)
+ self.affine_select = _auto_emit_void_call(_trn_op.nki_affine_select)
@staticmethod
def tensorized_instruction():
diff --git a/python/tvm/ir/op.py b/python/tvm/ir/op.py
index 9fc21dcb27..ffad9de4b9 100644
--- a/python/tvm/ir/op.py
+++ b/python/tvm/ir/op.py
@@ -83,7 +83,7 @@ def _make_op_api(op, module_name):
return call
-def _init_op_api(namespace, target_module_name=None):
+def _init_op_api(namespace, target_module_name=None, *, wrapper=None):
"""Initialize registered Op callables in an already loaded Python module.
Like :func:`tvm_ffi.init_ffi_api`, the registry prefix comes first and the
@@ -104,7 +104,10 @@ def _init_op_api(namespace, target_module_name=None):
Existing Python callables retain ownership of their names; only missing
exposures are generated. Existing functions must accept printed calls or
have an appropriate exceptional printer hook. Non-callable collisions fail
- before any functions are installed. Returns None.
+ before any functions are installed. An optional ``wrapper`` adapts newly
+ generated Call constructors at export time; existing exposures and region
+ builders are untouched.
+ Returns None.
"""
if not namespace or any(
not part.isidentifier() or keyword.iskeyword(part) for part in
namespace.split(".")
@@ -134,7 +137,10 @@ def _init_op_api(namespace, target_module_name=None):
if not callable(current):
raise ValueError(f"Op {name!r} conflicts with an existing
Python attribute")
else:
- pending.append((container, parts[-1], _make_op_api(op,
target.__name__)))
+ call = _make_op_api(op, target.__name__)
+ if wrapper is not None and op.get_attr("FRegionGetBodyParams") is
None:
+ call = wrapper(call)
+ pending.append((container, parts[-1], call))
for container, name, call in pending:
setattr(container, name, call)
diff --git a/python/tvm/relax/script/ir_builder/op.py
b/python/tvm/relax/script/ir_builder/op.py
index 0c7370501b..b2fcae0a0c 100644
--- a/python/tvm/relax/script/ir_builder/op.py
+++ b/python/tvm/relax/script/ir_builder/op.py
@@ -235,6 +235,7 @@ from tvm.relax.op.builtin import stop_lift_params
from tvm.relax.type import Type
from tvm.relax.utils import convert_to_expr
from tvm.script.ir_builder import base as _base
+from tvm.script.ir_builder.base import auto_emit_void_call as
_auto_emit_void_call
from . import distributed as dist # noqa: F401
from .ir import _value, lookup_vdevice
@@ -582,7 +583,7 @@ def ne_(lhs: Any, rhs: Any, *, loc: _Loc = None) ->
_ir.Expr:
inspect = _Namespace()
-_init_op_api("relax", __name__)
+_init_op_api("relax", __name__, wrapper=_auto_emit_void_call)
__all__ = [
"ArgmaxArgminAttrs",
diff --git a/python/tvm/s_tir/script/ir_builder/ir.py
b/python/tvm/s_tir/script/ir_builder/ir.py
index d8ccdd1d29..583a71231f 100644
--- a/python/tvm/s_tir/script/ir_builder/ir.py
+++ b/python/tvm/s_tir/script/ir_builder/ir.py
@@ -29,7 +29,7 @@ from tvm.ir.type import _normalize_type
from tvm.runtime import MemoryScope
from tvm.s_tir.iter_var import IterVar
from tvm.s_tir.op import comm_reducer
-from tvm.script.ir_builder.base import MISSING
+from tvm.script.ir_builder.base import MISSING, auto_emit_void_call
from tvm.script.parser.protocol_registry import register_mutable_decl as
_register_mutable_decl
from tvm.te import CommReducer, Reduce
from tvm.tirx import Expr, IntImm, Var
@@ -71,11 +71,13 @@ def async_copy_scope():
return region("s_tir.async_copy_scope", [])
+@auto_emit_void_call
def async_commit(queue_id, *, ty=None, loc: Location = UNKNOWN_LOC):
"""Commit the preceding asynchronous copies to a queue."""
return ir.Call("s_tir.async_commit", [queue_id], ty=ty, loc=loc)
+@auto_emit_void_call
def async_wait(queue_id, inflight_count, *, ty=None, loc: Location =
UNKNOWN_LOC):
"""Wait until at most inflight_count committed groups remain in flight."""
return ir.Call("s_tir.async_wait", [queue_id, inflight_count], ty=ty,
loc=loc)
diff --git a/python/tvm/script/ir_builder/base.py
b/python/tvm/script/ir_builder/base.py
index 6593e5baf0..52ac16dcbc 100644
--- a/python/tvm/script/ir_builder/base.py
+++ b/python/tvm/script/ir_builder/base.py
@@ -20,7 +20,7 @@ from collections.abc import Callable
from contextlib import contextmanager, nullcontext
from functools import wraps
from inspect import signature
-from typing import Any, Generic, TypeVar
+from typing import Any, Generic, ParamSpec, TypeVar
from tvm_ffi import register_object as _register_object
from tvm_ffi.dataclasses import MISSING as MISSING
@@ -277,27 +277,62 @@ _T = TypeVar("_T")
class AlreadyEmitted(Generic[_T]):
- """Hold the exact emitted value; location handling preserves this receipt.
+ """Hold the exact native result of an already emitted builder operation.
Parameters
----------
- value : Any
- Object already emitted by a language variant's builder. The receipt
- retains this object without copying it. Binding or emitting the receipt
- must not emit the object again.
+ value : T
+ Native result retained without copying or emitting it. Operator APIs
+ wrap their original Call; statement APIs wrap the emitted statement.
Attributes
----------
- value : Any
- The same emitted object, available for identity checks and
source-location
- attachment.
+ value : T
+ The same object, available for identity checks and source-location
+ attachment. Builder consumers must check for this receipt before FFI
+ conversion, so they do not emit the object again.
+
+ Notes
+ -----
+ FFI conversion exposes only the wrapped object, not its emission history.
"""
__slots__ = ("value",)
+ value: _T
def __init__(self, value: _T) -> None:
self.value = value
+ def __tvm_ffi_object__(self) -> _T:
+ """Return the exact wrapped native object without copying or
emitting."""
+ return self.value
+
+
+_P = ParamSpec("_P")
+
+
+def auto_emit_void_call(func: Callable[_P, _T]) -> Callable[_P, _T |
AlreadyEmitted[ir.Call]]:
+ """Adapt an operator constructor to the script builder boundary.
+
+ Construct once and preserve the callable's signature and metadata. A void
+ Call under an active IRBuilder emits one Evaluate(call) and returns a
+ receipt containing the exact original Call. Both primitive void and the
+ empty tuple type count as void. Non-Call results, non-void calls and calls
+ outside a builder pass through unchanged.
+
+ The native hook supplies source context to both Call and Evaluate, handles
+ frameless builders, and propagates normal errors from invalid frames.
+ """
+
+ @wraps(func)
+ def wrapped(*args: _P.args, **kwargs: _P.kwargs) -> _T |
AlreadyEmitted[ir.Call]:
+ result = func(*args, **kwargs)
+ if isinstance(result, ir.Call) and _ffi_api.MaybeEmitVoidCall(result):
+ return AlreadyEmitted(result)
+ return result
+
+ return wrapped
+
class LocationEntry:
"""A materialized source range shared by generated builder operations.
@@ -351,8 +386,9 @@ def at(loc: LocationEntry | ir.Location | None, value: _T)
-> _T:
Notes
-----
- Native mutation annotates the statement held by the builder itself. Keep
- the original Python facade as well, including callable objects and frames.
+ Native mutation annotates the wrapped object in the builder IR. For an
+ operator receipt this is the Call inside its Evaluate statement. Keep the
+ original Python facade as well, including callable objects and frames.
Unsupported objects and ordinary Python values pass through unchanged.
"""
if loc is None:
diff --git a/python/tvm/script/ir_builder/stmt.py
b/python/tvm/script/ir_builder/stmt.py
index ca1265b41d..67e49faf27 100644
--- a/python/tvm/script/ir_builder/stmt.py
+++ b/python/tvm/script/ir_builder/stmt.py
@@ -104,20 +104,25 @@ def bind( # pylint: disable=invalid-name
return _ffi_api.Bind(value, type_annotation, var) # type:
ignore[attr-defined] # pylint: disable=no-member
-def evaluate(value: Expr) -> AlreadyEmitted[tvm.ir.Stmt]:
- """Emit an evaluation and return a reference to its stored statement.
+def evaluate(
+ value: Expr | AlreadyEmitted[tvm.ir.Call] | AlreadyEmitted[tvm.ir.Stmt],
+) -> AlreadyEmitted[tvm.ir.Call] | AlreadyEmitted[tvm.ir.Stmt]:
+ """Emit an expression, or pass through a call or statement emission
receipt.
Parameters
----------
- value : Expr
- The input expression to evaluate.
+ value : Expr, AlreadyEmitted[Call] or AlreadyEmitted[Stmt]
+ The input expression, or a receipt returned unchanged without emission.
Returns
-------
- result : AlreadyEmitted[Stmt]
- A receipt containing the emitted statement, so expression-statement
- handling does not emit it again.
+ result : AlreadyEmitted[Call] or AlreadyEmitted[Stmt]
+ The existing receipt unchanged, or a receipt containing the newly
+ emitted Evaluate statement, so expression-statement handling does not
+ emit it again.
"""
+ if isinstance(value, AlreadyEmitted):
+ return value
if isinstance(value, str):
value = _StringImm(value)
if isinstance(value, bool):
diff --git a/python/tvm/tirx/script/ir_builder/op.py
b/python/tvm/tirx/script/ir_builder/op.py
index a8b906b7f3..852ef38deb 100644
--- a/python/tvm/tirx/script/ir_builder/op.py
+++ b/python/tvm/tirx/script/ir_builder/op.py
@@ -34,6 +34,7 @@ from tvm.ir.prim import _ffi_api as _prim_ffi
from tvm.ir.prim import _ffi_api as _prim_ffi_api
from tvm.script.ir_builder import base as _base
from tvm.script.ir_builder.base import IRBuilder
+from tvm.script.ir_builder.base import auto_emit_void_call as
_auto_emit_void_call
from tvm.script.ir_builder.frame import IRModuleFrame
from tvm.script.ir_builder.stmt import _as_expr as _as_expr
from tvm.tirx import Expr, is_tensor_var
@@ -75,11 +76,14 @@ from tvm.tirx.expr import (
from tvm.tirx.op import CallFFIKernelAttr
from . import _ffi_api
-from .external_kernel import call_kernel
+from .external_kernel import call_kernel as _call_kernel
+
+call_kernel = _auto_emit_void_call(_call_kernel)
# pylint: enable=unused-import
+@_auto_emit_void_call
def _call_global(func: ir.GlobalVar, *args: Expr) -> Call:
"""Build a TIRX call using the declared function's exact result type."""
if IRBuilder.is_in_scope():
@@ -162,6 +166,7 @@ class WebGPUNamespace:
"""The WebGPU intrinsics submodule."""
@staticmethod
+ @_auto_emit_void_call
def subgroup_shuffle(var, lane, *, ty=None, loc: Location = UNKNOWN_LOC):
if is_tensor_var(var):
var = var[0]
@@ -173,6 +178,7 @@ class WebGPUNamespace:
)
@staticmethod
+ @_auto_emit_void_call
def subgroup_shuffle_up(var, delta, *, ty=None, loc: Location =
UNKNOWN_LOC):
if is_tensor_var(var):
var = var[0]
@@ -184,6 +190,7 @@ class WebGPUNamespace:
)
@staticmethod
+ @_auto_emit_void_call
def subgroup_shuffle_down(var, delta, *, ty=None, loc: Location =
UNKNOWN_LOC):
if is_tensor_var(var):
var = var[0]
@@ -241,226 +248,231 @@ def register_script_namespace(name: str, namespace:
object) -> object:
return namespace
-abs = _tir_op.abs # pylint: disable=redefined-builtin
+abs = _auto_emit_void_call(_tir_op.abs) # pylint: disable=redefined-builtin
-acos = _tir_op.acos
+acos = _auto_emit_void_call(_tir_op.acos)
-acosh = _tir_op.acosh
+acosh = _auto_emit_void_call(_tir_op.acosh)
-address_of = _tir_op.address_of
+address_of = _auto_emit_void_call(_tir_op.address_of)
-atomic_add = _ir.op._make_op_api(_ir.Op.get("tirx.atomic_add"), __name__)
+atomic_add =
_auto_emit_void_call(_ir.op._make_op_api(_ir.Op.get("tirx.atomic_add"),
__name__))
-call_spirv_pure_glsl450 =
_ir.op._make_op_api(_ir.Op.get("tirx.call_spirv_pure_glsl450"), __name__)
+call_spirv_pure_glsl450 = _auto_emit_void_call(
+ _ir.op._make_op_api(_ir.Op.get("tirx.call_spirv_pure_glsl450"), __name__)
+)
device_scope = _ir.op._make_op_api(_ir.Op.get("tirx.device_scope"), __name__)
-prefetch = _ir.op._make_op_api(_ir.Op.get("tirx.prefetch"), __name__)
+prefetch =
_auto_emit_void_call(_ir.op._make_op_api(_ir.Op.get("tirx.prefetch"), __name__))
-static_handle = _ir.op._make_op_api(_ir.Op.get("tirx.static_handle"), __name__)
+static_handle = _auto_emit_void_call(
+ _ir.op._make_op_api(_ir.Op.get("tirx.static_handle"), __name__)
+)
-asin = _tir_op.asin
+asin = _auto_emit_void_call(_tir_op.asin)
-asinh = _tir_op.asinh
+asinh = _auto_emit_void_call(_tir_op.asinh)
-atan = _tir_op.atan
+atan = _auto_emit_void_call(_tir_op.atan)
-atan2 = _tir_op.atan2
+atan2 = _auto_emit_void_call(_tir_op.atan2)
-atanh = _tir_op.atanh
+atanh = _auto_emit_void_call(_tir_op.atanh)
-bitwise_and = _tir_op.bitwise_and
+bitwise_and = _auto_emit_void_call(_tir_op.bitwise_and)
-bitwise_not = _tir_op.bitwise_not
+bitwise_not = _auto_emit_void_call(_tir_op.bitwise_not)
-bitwise_or = _tir_op.bitwise_or
+bitwise_or = _auto_emit_void_call(_tir_op.bitwise_or)
-bitwise_xor = _tir_op.bitwise_xor
+bitwise_xor = _auto_emit_void_call(_tir_op.bitwise_xor)
-ceil = _ir.op._make_op_api(_ir.Op.get("prim.ceil"), __name__)
+ceil = _auto_emit_void_call(_ir.op._make_op_api(_ir.Op.get("prim.ceil"),
__name__))
-clz = _ir.op._make_op_api(_ir.Op.get("prim.clz"), __name__)
+clz = _auto_emit_void_call(_ir.op._make_op_api(_ir.Op.get("prim.clz"),
__name__))
-copysign = _tir_op.copysign
+copysign = _auto_emit_void_call(_tir_op.copysign)
-cos = _tir_op.cos
+cos = _auto_emit_void_call(_tir_op.cos)
-cosh = _tir_op.cosh
+cosh = _auto_emit_void_call(_tir_op.cosh)
-erf = _tir_op.erf
+erf = _auto_emit_void_call(_tir_op.erf)
-exp = _tir_op.exp
+exp = _auto_emit_void_call(_tir_op.exp)
-exp2 = _tir_op.exp2
+exp2 = _auto_emit_void_call(_tir_op.exp2)
-exp10 = _tir_op.exp10
+exp10 = _auto_emit_void_call(_tir_op.exp10)
-gpu_thread_filter = _tir_op.gpu_thread_filter
+gpu_thread_filter = _auto_emit_void_call(_tir_op.gpu_thread_filter)
-gpu_active_thread_selector = _tir_op.gpu_active_thread_selector
+gpu_active_thread_selector =
_auto_emit_void_call(_tir_op.gpu_active_thread_selector)
-floor = _ir.op._make_op_api(_ir.Op.get("prim.floor"), __name__)
+floor = _auto_emit_void_call(_ir.op._make_op_api(_ir.Op.get("prim.floor"),
__name__))
-ceildiv = _tir_op.ceildiv
+ceildiv = _auto_emit_void_call(_tir_op.ceildiv)
-floordiv = _tir_op.floordiv
+floordiv = _auto_emit_void_call(_tir_op.floordiv)
-floormod = _tir_op.floormod
+floormod = _auto_emit_void_call(_tir_op.floormod)
-fmod = _tir_op.fmod
+fmod = _auto_emit_void_call(_tir_op.fmod)
-fma = _tir_op.fma
+fma = _auto_emit_void_call(_tir_op.fma)
-hypot = _tir_op.hypot
+hypot = _auto_emit_void_call(_tir_op.hypot)
-if_then_else = _ir.op._make_op_api(_ir.Op.get("prim.if_then_else"), __name__)
+if_then_else =
_auto_emit_void_call(_ir.op._make_op_api(_ir.Op.get("prim.if_then_else"),
__name__))
-infinity = _tir_op.infinity
+infinity = _auto_emit_void_call(_tir_op.infinity)
-isfinite = _tir_op.isfinite
+isfinite = _auto_emit_void_call(_tir_op.isfinite)
-isinf = _tir_op.isinf
+isinf = _auto_emit_void_call(_tir_op.isinf)
-isnan = _tir_op.isnan
+isnan = _auto_emit_void_call(_tir_op.isnan)
-isnullptr = _tir_op.isnullptr
+isnullptr = _auto_emit_void_call(_tir_op.isnullptr)
-ldexp = _tir_op.ldexp
+ldexp = _auto_emit_void_call(_tir_op.ldexp)
-likely = _ir.op._make_op_api(_ir.Op.get("prim.likely"), __name__)
+likely = _auto_emit_void_call(_ir.op._make_op_api(_ir.Op.get("prim.likely"),
__name__))
-log = _tir_op.log
+log = _auto_emit_void_call(_tir_op.log)
-log1p = _tir_op.log1p
+log1p = _auto_emit_void_call(_tir_op.log1p)
-log2 = _tir_op.log2
+log2 = _auto_emit_void_call(_tir_op.log2)
-log10 = _tir_op.log10
+log10 = _auto_emit_void_call(_tir_op.log10)
-max_value = _tir_op.max_value
+max_value = _auto_emit_void_call(_tir_op.max_value)
-min_value = _tir_op.min_value
+min_value = _auto_emit_void_call(_tir_op.min_value)
-nearbyint = _ir.op._make_op_api(_ir.Op.get("prim.nearbyint"), __name__)
+nearbyint =
_auto_emit_void_call(_ir.op._make_op_api(_ir.Op.get("prim.nearbyint"),
__name__))
-nextafter = _tir_op.nextafter
+nextafter = _auto_emit_void_call(_tir_op.nextafter)
-popcount = _tir_op.popcount
+popcount = _auto_emit_void_call(_tir_op.popcount)
-pow = _ir.op._make_op_api(_ir.Op.get("prim.pow"), __name__)
+pow = _auto_emit_void_call(_ir.op._make_op_api(_ir.Op.get("prim.pow"),
__name__))
-round = _ir.op._make_op_api(_ir.Op.get("prim.round"), __name__)
+round = _auto_emit_void_call(_ir.op._make_op_api(_ir.Op.get("prim.round"),
__name__))
-rsqrt = _tir_op.rsqrt
+rsqrt = _auto_emit_void_call(_tir_op.rsqrt)
-shift_left = _tir_op.shift_left
+shift_left = _auto_emit_void_call(_tir_op.shift_left)
-shift_right = _tir_op.shift_right
+shift_right = _auto_emit_void_call(_tir_op.shift_right)
-sigmoid = _tir_op.sigmoid
+sigmoid = _auto_emit_void_call(_tir_op.sigmoid)
-sin = _tir_op.sin
+sin = _auto_emit_void_call(_tir_op.sin)
-sinh = _tir_op.sinh
+sinh = _auto_emit_void_call(_tir_op.sinh)
-sqrt = _tir_op.sqrt
+sqrt = _auto_emit_void_call(_tir_op.sqrt)
-tan = _tir_op.tan
+tan = _auto_emit_void_call(_tir_op.tan)
-tanh = _tir_op.tanh
+tanh = _auto_emit_void_call(_tir_op.tanh)
-gpu_thread_return = _tir_op.gpu_thread_return
+gpu_thread_return = _auto_emit_void_call(_tir_op.gpu_thread_return)
-trunc = _ir.op._make_op_api(_ir.Op.get("prim.trunc"), __name__)
+trunc = _auto_emit_void_call(_ir.op._make_op_api(_ir.Op.get("prim.trunc"),
__name__))
-truncdiv = _tir_op.truncdiv
+truncdiv = _auto_emit_void_call(_tir_op.truncdiv)
-truncmod = _tir_op.truncmod
+truncmod = _auto_emit_void_call(_tir_op.truncmod)
-ptr_byte_offset = _tir_op.ptr_byte_offset
+ptr_byte_offset = _auto_emit_void_call(_tir_op.ptr_byte_offset)
-throw_last_error = _tir_op.throw_last_error
+throw_last_error = _auto_emit_void_call(_tir_op.throw_last_error)
-stack_alloca = _tir_op.stack_alloca
+stack_alloca = _auto_emit_void_call(_tir_op.stack_alloca)
-stack_make_shape = _tir_op.stack_make_shape
+stack_make_shape = _auto_emit_void_call(_tir_op.stack_make_shape)
-stack_make_dltensor = _tir_op.stack_make_dltensor
+stack_make_dltensor = _auto_emit_void_call(_tir_op.stack_make_dltensor)
-call_packed = _tir_op.call_packed
+call_packed = _auto_emit_void_call(_tir_op.call_packed)
_call_ffi_kernel_raw = _ir.op._make_op_api(_ir.Op.get("tirx.call_ffi_kernel"),
__name__)
+@_auto_emit_void_call
def call_ffi_kernel(*args, launch=None, **kwargs):
"""Call a kernel with a CUDA LaunchConfig, or read a printed low-level
call."""
if launch is not None:
@@ -468,125 +480,134 @@ def call_ffi_kernel(*args, launch=None, **kwargs):
return _call_ffi_kernel_raw(*args, **kwargs)
-call_cpacked = _tir_op.call_cpacked
+call_cpacked = _auto_emit_void_call(_tir_op.call_cpacked)
-call_packed_lowered = _tir_op.call_packed_lowered
+call_packed_lowered = _auto_emit_void_call(_tir_op.call_packed_lowered)
-call_cpacked_lowered = _tir_op.call_cpacked_lowered
+call_cpacked_lowered = _auto_emit_void_call(_tir_op.call_cpacked_lowered)
-tensor_data_ptr = _tir_op.tensor_data_ptr
+tensor_data_ptr = _auto_emit_void_call(_tir_op.tensor_data_ptr)
-handle_add_byte_offset = _tir_op.handle_add_byte_offset
+handle_add_byte_offset = _auto_emit_void_call(_tir_op.handle_add_byte_offset)
-abi_field_set = _tir_op.abi_field_set
+abi_field_set = _auto_emit_void_call(_tir_op.abi_field_set)
-abi_field_get = _ir.op._make_op_api(_ir.Op.get("tirx.abi_field_get"), __name__)
+abi_field_get = _auto_emit_void_call(
+ _ir.op._make_op_api(_ir.Op.get("tirx.abi_field_get"), __name__)
+)
-gpu_thread_invariant = _tir_op.gpu_thread_invariant
+gpu_thread_invariant = _auto_emit_void_call(_tir_op.gpu_thread_invariant)
-gpu_thread_allreduce = _tir_op.gpu_thread_allreduce
+gpu_thread_allreduce = _auto_emit_void_call(_tir_op.gpu_thread_allreduce)
-gpu_load_matrix_sync = _tir_op.gpu_load_matrix_sync
+gpu_load_matrix_sync = _auto_emit_void_call(_tir_op.gpu_load_matrix_sync)
-gpu_mma_sync = _tir_op.gpu_mma_sync
+gpu_mma_sync = _auto_emit_void_call(_tir_op.gpu_mma_sync)
-gpu_fill_fragment = _tir_op.gpu_fill_fragment
+gpu_fill_fragment = _auto_emit_void_call(_tir_op.gpu_fill_fragment)
-gpu_store_matrix_sync = _tir_op.gpu_store_matrix_sync
+gpu_store_matrix_sync = _auto_emit_void_call(_tir_op.gpu_store_matrix_sync)
-gpu_storage_sync = _tir_op.gpu_storage_sync
-cpu_parallel_barrier = _tir_op.cpu_parallel_barrier
+gpu_storage_sync = _auto_emit_void_call(_tir_op.gpu_storage_sync)
+cpu_parallel_barrier = _auto_emit_void_call(_tir_op.cpu_parallel_barrier)
-kernel_replace_point = _tir_op.kernel_replace_point
+kernel_replace_point = _auto_emit_void_call(_tir_op.kernel_replace_point)
-gpu_warp_shuffle = _tir_op.gpu_warp_shuffle
+gpu_warp_shuffle = _auto_emit_void_call(_tir_op.gpu_warp_shuffle)
-gpu_warp_shuffle_up = _tir_op.gpu_warp_shuffle_up
+gpu_warp_shuffle_up = _auto_emit_void_call(_tir_op.gpu_warp_shuffle_up)
-gpu_warp_shuffle_down = _tir_op.gpu_warp_shuffle_down
+gpu_warp_shuffle_down = _auto_emit_void_call(_tir_op.gpu_warp_shuffle_down)
-gpu_warp_shuffle_xor = _tir_op.gpu_warp_shuffle_xor
+gpu_warp_shuffle_xor = _auto_emit_void_call(_tir_op.gpu_warp_shuffle_xor)
-gpu_warp_activemask = _tir_op.gpu_warp_activemask
+gpu_warp_activemask = _auto_emit_void_call(_tir_op.gpu_warp_activemask)
-assume = _tir_op.assume
-assume_aligned = _tir_op.assume_aligned
+assume = _auto_emit_void_call(_tir_op.assume)
+assume_aligned = _auto_emit_void_call(_tir_op.assume_aligned)
-undef = _tir_op.undef
+undef = _auto_emit_void_call(_tir_op.undef)
-alloc_workspace = _tir_op.alloc_workspace
+alloc_workspace = _auto_emit_void_call(_tir_op.alloc_workspace)
-free_workspace = _tir_op.free_workspace
+free_workspace = _auto_emit_void_call(_tir_op.free_workspace)
-vscale = _tir_op.vscale
+vscale = _auto_emit_void_call(_tir_op.vscale)
-ignore_loop_partition = _tir_op.ignore_loop_partition
+ignore_loop_partition = _auto_emit_void_call(_tir_op.ignore_loop_partition)
-reinterpret = _ir.op._make_op_api(_ir.Op.get("tirx.reinterpret"), __name__)
+reinterpret =
_auto_emit_void_call(_ir.op._make_op_api(_ir.Op.get("tirx.reinterpret"),
__name__))
-call_extern = _ir.op._make_op_api(_ir.Op.get("tirx.call_extern"), __name__)
+call_extern =
_auto_emit_void_call(_ir.op._make_op_api(_ir.Op.get("tirx.call_extern"),
__name__))
+@_auto_emit_void_call
def call_intrin(func_name, *args, ty, attrs=None, loc: Location = UNKNOWN_LOC):
"""Call an intrinsic with an explicit result type."""
return _tir_op.call_intrin(ty, func_name, *args, attrs=attrs, loc=loc)
-call_llvm_intrin = _llvm_result_type(_tir_op.call_llvm_intrin)
+call_llvm_intrin =
_auto_emit_void_call(_llvm_result_type(_tir_op.call_llvm_intrin))
-call_llvm_pure_intrin = _llvm_result_type(_tir_op.call_llvm_pure_intrin)
+call_llvm_pure_intrin =
_auto_emit_void_call(_llvm_result_type(_tir_op.call_llvm_pure_intrin))
-call_pure_extern = _ir.op._make_op_api(_ir.Op.get("tirx.call_pure_extern"),
__name__)
+call_pure_extern = _auto_emit_void_call(
+ _ir.op._make_op_api(_ir.Op.get("tirx.call_pure_extern"), __name__)
+)
-vector_low = _ir.op._make_op_api(_ir.Op.get("tirx.vector_low"), __name__)
+vector_low =
_auto_emit_void_call(_ir.op._make_op_api(_ir.Op.get("tirx.vector_low"),
__name__))
-vector_high = _ir.op._make_op_api(_ir.Op.get("tirx.vector_high"), __name__)
+vector_high =
_auto_emit_void_call(_ir.op._make_op_api(_ir.Op.get("tirx.vector_high"),
__name__))
-vector_combine = _ir.op._make_op_api(_ir.Op.get("tirx.vector_combine"),
__name__)
+vector_combine = _auto_emit_void_call(
+ _ir.op._make_op_api(_ir.Op.get("tirx.vector_combine"), __name__)
+)
-get_active_lane_mask =
_ir.op._make_op_api(_ir.Op.get("tirx.get_active_lane_mask"), __name__)
+get_active_lane_mask = _auto_emit_void_call(
+ _ir.op._make_op_api(_ir.Op.get("tirx.get_active_lane_mask"), __name__)
+)
-masked_load = _ir.op._make_op_api(_ir.Op.get("tirx.masked_load"), __name__)
+masked_load =
_auto_emit_void_call(_ir.op._make_op_api(_ir.Op.get("tirx.masked_load"),
__name__))
-masked_store = _tir_op.masked_store
+masked_store = _auto_emit_void_call(_tir_op.masked_store)
-gpu_dp4a = _tir_op.gpu_dp4a
+gpu_dp4a = _auto_emit_void_call(_tir_op.gpu_dp4a)
broadcast = Broadcast
@@ -595,7 +616,7 @@ broadcast = Broadcast
ramp = Ramp
-fabs = _ir.op._make_op_api(_ir.Op.get("prim.fabs"), __name__)
+fabs = _auto_emit_void_call(_ir.op._make_op_api(_ir.Op.get("prim.fabs"),
__name__))
def logical_and(*values):
diff --git a/src/script/ir_builder/ir.cc b/src/script/ir_builder/ir.cc
index a935bc407a..84b2580ed0 100644
--- a/src/script/ir_builder/ir.cc
+++ b/src/script/ir_builder/ir.cc
@@ -362,6 +362,21 @@ tvm::Stmt Evaluate(Expr value) {
return stmt;
}
+// Keep operator construction independent of builder state. Python API exposure
+// invokes this hook only after constructing and selecting its call.
+bool MaybeEmitVoidCall(Call call) {
+ const auto* prim_type = call->ty.as<PrimTypeNode>();
+ bool is_void = IsVoidType(call->ty) || (prim_type &&
ffi::GetRef<PrimType>(prim_type).IsVoid());
+ if (!is_void || !IRBuilder::IsInScope()) return false;
+ IRBuilder::Current()->AttachLoc(call);
+ Evaluate(std::move(call));
+ return true;
+}
+
+TVM_FFI_STATIC_INIT_BLOCK() {
+ ffi::reflection::GlobalDef().def("script.ir_builder.MaybeEmitVoidCall",
MaybeEmitVoidCall);
+}
+
// Preserve the historical FFI entry points with one shared implementation.
TVM_FFI_STATIC_INIT_BLOCK() {
ffi::reflection::GlobalDef()
diff --git a/tests/python/script/test_builder_void_emission.py
b/tests/python/script/test_builder_void_emission.py
new file mode 100644
index 0000000000..446e3c3059
--- /dev/null
+++ b/tests/python/script/test_builder_void_emission.py
@@ -0,0 +1,78 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Core behavior of script-facing void operators in an IRBuilder."""
+
+import pytest
+import tvm_ffi
+
+from tvm import ir, tirx
+from tvm.backend import cuda
+from tvm.script import ir_builder as I
+from tvm.script import tirx as T
+from tvm.script.ir_builder.base import AlreadyEmitted
+
+
[email protected](scope="module", autouse=True)
+def register_cuda():
+ cuda.register_backend()
+
+
[email protected]("ty", [ir.PrimType("void"), ir.TupleType([])])
+def test_void_emission_and_fallback(ty):
+ assert isinstance(T.cpu_parallel_barrier(ty=ty), ir.Call)
+ with I.IRBuilder() as builder:
+ value = T.cpu_parallel_barrier(ty="int32")
+ assert isinstance(value, ir.Call)
+ assert builder.result is None
+ receipt = T.cpu_parallel_barrier(ty=ty)
+ assert isinstance(receipt, AlreadyEmitted)
+ assert T.evaluate(receipt) is receipt
+ assert builder.get().value.same_as(receipt.value)
+ assert receipt.value.ty.same_as(ty)
+ assert receipt.__tvm_ffi_object__() is receipt.value
+ assert tvm_ffi.convert(receipt).same_as(receipt.value)
+
+
+def test_parser_emits_once_with_source_locations():
+ @T.function
+ def function():
+ T.ptx.barrier.sync.aligned(0)
+ T.evaluate(T.ptx.barrier.sync.aligned(0))
+ T.cpu_parallel_barrier()
+ T.evaluate(T.cpu_parallel_barrier())
+
+ assert len(function.body.seq) == 4
+ for stmt in function.body.seq:
+ assert isinstance(stmt, ir.Evaluate)
+ assert not isinstance(stmt.loc, ir.UnknownLoc)
+ assert not isinstance(stmt.value.loc, ir.UnknownLoc)
+
+
+def test_raw_construction_and_codegen_do_not_emit():
+ from tvm.backend.cuda.codegen.registry import CODEGEN_REGISTRY
+ from tvm.backend.cuda.ptx import PTXNamespace
+
+ ptx = T.ptx.barrier.sync.aligned(0)
+ assert isinstance(ptx, ir.Call)
+ with I.IRBuilder() as builder:
+ raw = tirx.cpu_parallel_barrier()
+ assert isinstance(raw, ir.Call)
+ assert isinstance(PTXNamespace()["barrier.sync.aligned"](0), ir.Call)
+ helper, _ = CODEGEN_REGISTRY[ptx.op.name](ptx.args)
+ assert isinstance(helper, ir.Call)
+ ptx.script()
+ assert builder.result is None
