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 827f97eb8f [REFACTOR][TIR] Remove obsolete TIRx builtins (#20463)
827f97eb8f is described below
commit 827f97eb8f6883249563b1fcf87cb5414aa84941
Author: Tianqi Chen <[email protected]>
AuthorDate: Sat Sep 26 11:05:29 2026 -0400
[REFACTOR][TIR] Remove obsolete TIRx builtins (#20463)
## Summary
- Remove 14 obsolete TIRx builtins and their dependent lowering,
codegen, and Python interfaces.
- Use native `Break` and `Continue` statements for loop control; retain
`thread_return` for GPU thread exits while function `Return` carries a
value.
- Keep independent runtime DMA APIs and live synchronization paths.
Reject unsupported global storage synchronization explicitly.
---
cmake/modules/Hexagon.cmake | 6 -
include/tvm/runtime/device_api.h | 1 -
include/tvm/s_tir/stmt.h | 6 +-
include/tvm/s_tir/transform.h | 12 -
include/tvm/tirx/builtin.h | 126 +------
include/tvm/tirx/op.h | 18 +-
include/tvm/tirx/script/ir_builder/ir.h | 6 +-
python/tvm/s_tir/backend/adreno/pipeline.py | 4 -
python/tvm/s_tir/pipeline.py | 4 -
python/tvm/s_tir/transform/transform.py | 22 --
python/tvm/te/__init__.py | 2 +-
python/tvm/tirx/__init__.py | 8 +-
python/tvm/tirx/op.py | 157 +--------
python/tvm/tirx/script/ir_builder/op.py | 28 --
.../tvm/tirx/script/ir_builder/parser_protocol.py | 4 +-
src/backend/cuda/codegen/codegen_cuda.cc | 51 +--
src/backend/cuda/codegen/codegen_cuda.h | 8 -
.../hexagon/codegen/llvm/codegen_hexagon.cc | 21 --
src/backend/hexagon/runtime/profiler/lwp_handler.S | 124 -------
src/backend/hexagon/runtime/profiler/prof_utils.cc | 78 -----
src/backend/hexagon/runtime/profiler/prof_utils.h | 30 --
.../hexagon/runtime/rpc/hexagon/rpc_server.cc | 26 +-
.../hexagon/runtime/rpc/simulator/rpc_server.cc | 26 +-
src/s_tir/transform/annotate_irregular_loop.cc | 14 +-
src/s_tir/transform/inject_virtual_thread.cc | 2 -
src/s_tir/transform/lower_async_dma.cc | 215 ------------
src/s_tir/transform/profile_instrumentation.cc | 314 ------------------
src/target/llvm/codegen_cpu.cc | 30 --
src/target/llvm/codegen_cpu.h | 2 -
src/target/llvm/codegen_llvm.cc | 40 +--
src/target/llvm/codegen_llvm.h | 3 +-
src/target/source/codegen_c.cc | 9 -
src/tirx/analysis/verify_well_formed.h | 8 -
src/tirx/ir/transform.cc | 1 -
src/tirx/op/builtin.cc | 49 ---
src/tirx/op/op.cc | 27 +-
src/tirx/script/ir_builder/ir.cc | 12 +-
src/tirx/transform/lower_tvm_builtin.cc | 75 +----
tests/python/codegen/test_target_codegen_cuda.py | 3 +-
tests/python/runtime/test_runtime_trace.py | 226 -------------
.../s_tir/script/test_s_tir_script_basic_usage.py | 18 +-
.../s_tir/script/test_s_tir_script_printer.py | 4 +-
.../test_s_tir_transform_inject_virtual_thread.py | 2 +-
.../test_s_tir_transform_profiling_instr.py | 364 ---------------------
tests/python/script/test_basic_usage.py | 100 +++++-
tests/python/tirx-base/test_tir_base.py | 4 +-
tests/python/tirx-base/test_tir_constructor.py | 1 -
tests/python/tirx-base/test_tir_op_types.py | 19 --
.../tirx/script/test_tirx_script_basic_usage.py | 35 +-
49 files changed, 214 insertions(+), 2131 deletions(-)
diff --git a/cmake/modules/Hexagon.cmake b/cmake/modules/Hexagon.cmake
index 1ef746127d..ef07195cd0 100644
--- a/cmake/modules/Hexagon.cmake
+++ b/cmake/modules/Hexagon.cmake
@@ -298,14 +298,10 @@ if(USE_HEXAGON_RPC)
# TODO(masahi): Remove rpc_local_session.cc after verifying that things
work without it
"${TVM_CORE_RUNTIME_SOURCE_DIR}/rpc/rpc_local_session.cc"
)
- set(HEXAGON_PROFILER_DIR "${TVMRT_SOURCE_DIR}/profiler")
# Add the hardware-specific RPC code into the skel library.
- set_property(SOURCE ${HEXAGON_PROFILER_DIR}/lwp_handler.S PROPERTY
LANGUAGE C)
add_library(hexagon_rpc_skel SHARED
"${TVMRT_SOURCE_DIR}/rpc/hexagon/rpc_server.cc"
"${TVMRT_SOURCE_DIR}/rpc/hexagon_rpc_skel.c"
- "${HEXAGON_PROFILER_DIR}/prof_utils.cc"
- "${HEXAGON_PROFILER_DIR}/lwp_handler.S"
)
target_include_directories(hexagon_rpc_skel
SYSTEM PRIVATE "${TVMRT_SOURCE_DIR}/rpc"
@@ -314,8 +310,6 @@ if(USE_HEXAGON_RPC)
# executed via run_main_on_sim.
add_library(hexagon_rpc_sim SHARED
"${TVMRT_SOURCE_DIR}/rpc/simulator/rpc_server.cc"
- "${HEXAGON_PROFILER_DIR}/prof_utils.cc"
- "${HEXAGON_PROFILER_DIR}/lwp_handler.S"
)
target_link_libraries(hexagon_rpc_sim
-Wl,--whole-archive tvm_runtime -Wl,--no-whole-archive
diff --git a/include/tvm/runtime/device_api.h b/include/tvm/runtime/device_api.h
index a14135f3c8..ed2d2806af 100644
--- a/include/tvm/runtime/device_api.h
+++ b/include/tvm/runtime/device_api.h
@@ -417,7 +417,6 @@ TVM_RUNTIME_DLL bool RuntimeEnabled(const ffi::String&
target);
/*! \brief namespace for constant symbols */
namespace symbol {
-constexpr const char* tvm_global_barrier_state = "__tvm_global_barrier_state";
/*! \brief global function to set device */
constexpr const char* tvm_set_device = "__tvm_set_device";
} // namespace symbol
diff --git a/include/tvm/s_tir/stmt.h b/include/tvm/s_tir/stmt.h
index f5c2ef2581..810435142d 100644
--- a/include/tvm/s_tir/stmt.h
+++ b/include/tvm/s_tir/stmt.h
@@ -336,9 +336,9 @@ constexpr const char* manifest_shared_memory_local_stage =
/*!
* \brief Mark alignment of buffer dimension
- * stmt.node is Tensor
- * stmt.value is tvm_tuple(dim, align, offset)
- * This gives hint to require stride of dim to be k * align + offset.
+ * The annotation value is an array of explicit tuples
+ * (buffer_index, axis, factor, offset).
+ * This requires the stride of an axis to be k * factor + offset.
*/
constexpr const char* buffer_dim_align = "buffer_dim_align";
diff --git a/include/tvm/s_tir/transform.h b/include/tvm/s_tir/transform.h
index 5149a5fe28..8750ea4b7d 100644
--- a/include/tvm/s_tir/transform.h
+++ b/include/tvm/s_tir/transform.h
@@ -295,12 +295,6 @@ TVM_DLL Pass InstrumentBoundCheckers();
*/
TVM_DLL Pass InjectPTXLDG32(bool enable_inject = true);
-/*!
- * \brief Insert intrinsic calls to instrument function and loop level
profiling.
- * \return The pass.
- */
-TVM_DLL Pass InstrumentProfileIntrinsics();
-
/*!
* \brief Lower VTCM allocations.
* \return The pass.
@@ -326,12 +320,6 @@ TVM_DLL Pass InferFragment();
*/
TVM_DLL Pass LowerThreadAllreduce();
-/*!
- * \brief Lower Async TIR primitives to DMA copy and wait builtins.
- * \return The pass.
- */
-TVM_DLL Pass LowerAsyncDMA();
-
/*!
* \brief Rewrite global to shared memory copy on CUDA with asynchronous copy.
* \return The pass.
diff --git a/include/tvm/tirx/builtin.h b/include/tvm/tirx/builtin.h
index 4f66ab328a..146d5af881 100644
--- a/include/tvm/tirx/builtin.h
+++ b/include/tvm/tirx/builtin.h
@@ -42,17 +42,9 @@ namespace tirx {
/*! \brief Collection of builtin intrinsics as ops */
namespace builtin {
/*!
- * \brief Return from a GPU thread.
+ * \brief Return from a GPU thread without returning a function value.
*/
TVM_DLL const Op& thread_return();
-/*!
- * \brief Loop continue.
- */
-TVM_DLL const Op& continue_loop();
-/*!
- * \brief Loop break.
- */
-TVM_DLL const Op& break_loop();
/*!
* \brief Reinterpret the value using the target type.
*/
@@ -230,21 +222,6 @@ TVM_DLL const Op& ptr_byte_offset();
*/
TVM_DLL const Op& tvm_static_handle();
-/*!
- * \brief Return a unique context id, used for hint of workspace separation.
- * Different context id ganrantees not having overlapping workspace.
- */
-TVM_DLL const Op& tvm_context_id();
-
-/*!
- * \brief tvm_tuple is not an actual function and cannot codegen.
- * It is used to represent tuple structure in value field of AttrStmt,
- * for the sake of giving hint to optimization.
- *
- * void tvm_tuple(value0, value1, ..., value_n);
- */
-TVM_DLL const Op& tvm_tuple();
-
/*!
* \brief See pesudo code
*
@@ -274,14 +251,6 @@ TVM_DLL const Op& tvm_struct_get();
*/
TVM_DLL const Op& tvm_struct_set();
-/*!
- * \brief See pseudo code
- * Type lookup_param(ffi::String param_name) {
- * return __tvm_param__param_name;
- * }
- */
-TVM_DLL const Op& lookup_param();
-
/*!
* \brief See pesudo code
*
@@ -378,19 +347,6 @@ TVM_DLL const Op& tensormap_encode_tiled();
*/
TVM_DLL const Op& tvm_call_cpacked();
-/*!
- * \brief See pesudo code
- *
- * return_type tvm_call_trace_packed(name, TVMFFIAny* args) {
- * ModuleNode* env = GetCurrentEnv();
- * const ffi::Function* f = env->GetFuncFromEnv(name);
- * (*f)(args, args, len(args));
- * // return type can be int, float, handle.
- * return cast(return_type, result);
- * }
- */
-TVM_DLL const Op& tvm_call_trace_packed();
-
/*!
* \brief Mark a condition to be thread invariant.
* This means the condition must be the same for all threads.
@@ -430,25 +386,6 @@ TVM_DLL const Op& tvm_call_packed_lowered();
*/
TVM_DLL const Op& tvm_call_cpacked_lowered();
-/*!
- * \brief Lowered version of trace intrinsic, the space of value and
- * type codes are explicitly allocated. The return value is the
- * (end - 1) value on the stack.
- *
- * return_type tvm_call_trace_packed_lowered(name,
- * TVMFFIAny* args_stack,
- * int begin,
- * int end) {
- * ModuleNode* env = GetCurrentEnv();
- * const ffi::Function* f = env->GetFuncFromEnv(name);
- * f->CallPacked(ffi::PackedArgs(args_stack[begin:end]),
- * ffi::Any(args_stack + end));
- * // return type can be int, float, handle.
- * return cast(return_type, load_return_from(args_stack + end))
- * }
- */
-TVM_DLL const Op& tvm_call_trace_packed_lowered();
-
/*!
* \brief See pseudo code
*
@@ -501,12 +438,6 @@ TVM_DLL const Op& tvm_warp_shuffle_down();
TVM_DLL const Op& tvm_warp_shuffle_xor();
TVM_DLL const Op& tvm_warp_activemask();
-/*!
- * \brief Initialize the global barrier.
- * Call this at beginning of kernel that need global barrier.
- */
-TVM_DLL const Op& tvm_global_barrier_kinit();
-
/*!
* \brief See pesudo code
*
@@ -604,51 +535,6 @@ TVM_DLL const Op& texture2d_store();
*/
TVM_DLL const Op& texture2d_load();
-/*!
- * \brief Initiate a non-blocking DMA copy from source to destination
- *
- * The copy is launched immediately.
- *
- * If a `dma_start_group()` call is active, the copy will be added
- * to the current group for tracking of in-flight group counts.
- *
- * If no `dma_start_group()` call is active, the copy will be tracked
- * individually i.e. as a group with size 1.
- */
-TVM_DLL const Op& dma_copy();
-
-/*!
- * \brief Wait until the number of DMA groups in flight is less than
- * or equal to some maximum
- *
- * Calling `dma_wait()` while a group is active is unsupported.
- */
-TVM_DLL const Op& dma_wait();
-
-/*!
- * \brief Start a group of DMA copies
- *
- * Any call to `dma_copy()` that occurs after `dma_start_group()` will
- * be added to the current group for tracking of in-flight group counts.
- *
- * Only one DMA group may be active at a given time. Calling
- * `dma_start_group()` while a group is active is unsupported.
- */
-TVM_DLL const Op& dma_start_group();
-
-/*!
- * \brief End a group of DMA copies
- *
- * Track all calls to `dma_copy()` that occurred since the preceding
- * `dma_start_group()` as a single group in-flight.
- *
- * Calling `dma_end_group()` without an active group is unsupported.
- *
- * Note: A group of DMA calls may be empty, and will still contribute
- * to the count of in-flight groups used by `dma_wait()`.
- */
-TVM_DLL const Op& dma_end_group();
-
/*!
* \brief Provide a true statement that can be used for simplifications
*
@@ -666,16 +552,6 @@ TVM_DLL const Op& assume();
*/
TVM_DLL const Op& undef();
-/*!
- * \brief Profiling intrinsic
- */
-TVM_DLL const Op& start_profile_intrinsic();
-
-/*!
- * \brief Profiling intrinsic
- */
-TVM_DLL const Op& end_profile_intrinsic();
-
/*!
* \brief Get a item from any list and return it.
*
diff --git a/include/tvm/tirx/op.h b/include/tvm/tirx/op.h
index c180e89f23..7089522496 100644
--- a/include/tvm/tirx/op.h
+++ b/include/tvm/tirx/op.h
@@ -77,27 +77,13 @@ TVM_DLL Type GetType(const PrimExpr& expr);
TVM_DLL Type GetTypeFromRuntimeDataType(DLDataType dtype);
/*!
- * \brief Return from a thread.
+ * \brief Return from a GPU thread without returning a function value.
*
* \param span The location of this operation in the source.
- * \return The return expression.
+ * \return The thread return expression.
*/
TVM_DLL PrimExpr thread_return(Span span = Span());
-/*!
- * \brief Continue current loop.
- * \param span The location of this operation in the source.
- * \return The continue loop expression.
- */
-TVM_DLL PrimExpr continue_loop(Span span = Span());
-
-/*!
- * \brief Break current loop.
- * \param span The location of this operation in the source.
- * \return The break loop expression.
- */
-TVM_DLL PrimExpr break_loop(Span span = Span());
-
/*!
* Get the value of infinity.
* \param dtype The primitive type.
diff --git a/include/tvm/tirx/script/ir_builder/ir.h
b/include/tvm/tirx/script/ir_builder/ir.h
index b3c21a2c5c..ba31166835 100644
--- a/include/tvm/tirx/script/ir_builder/ir.h
+++ b/include/tvm/tirx/script/ir_builder/ir.h
@@ -274,13 +274,15 @@ tvm::tirx::Stmt Return(Expr value);
/*!
* \brief Create a break statement.
+ * \return The same statement that was added to the parent frame.
*/
-void Break();
+tvm::tirx::Stmt Break();
/*!
* \brief Create a continue statement.
+ * \return The same statement that was added to the parent frame.
*/
-void Continue();
+tvm::tirx::Stmt Continue();
/*!
* \brief Create an if statement.
diff --git a/python/tvm/s_tir/backend/adreno/pipeline.py
b/python/tvm/s_tir/backend/adreno/pipeline.py
index 0eec795233..dbd7f28fc3 100644
--- a/python/tvm/s_tir/backend/adreno/pipeline.py
+++ b/python/tvm/s_tir/backend/adreno/pipeline.py
@@ -61,8 +61,6 @@ def default_tir_pipeline(*, prepare_only=False):
]
if not bool(config.get("tirx.disable_storage_rewrite", False)):
passes.append(tirx.transform.StorageRewrite())
- if config.get("tirx.use_async_copy", False):
- passes.append(s_tir.transform.LowerAsyncDMA())
passes.extend(
[
s_tir.transform.HoistIfThenElse(),
@@ -80,8 +78,6 @@ def default_tir_pipeline(*, prepare_only=False):
passes.append(s_tir.transform.InjectPTXLDG32(True))
if not bool(config.get("tirx.disable_cse_tir", False)):
passes.append(tirx.transform.CommonSubexprElim())
- if bool(config.get("tirx.instrument_lwp", False)):
- passes.append(s_tir.transform.InstrumentProfileIntrinsics())
passes.extend(
[
# Bind the target first so that target-specific attributes are
available.
diff --git a/python/tvm/s_tir/pipeline.py b/python/tvm/s_tir/pipeline.py
index fb9e0f02d1..e91b6756e0 100644
--- a/python/tvm/s_tir/pipeline.py
+++ b/python/tvm/s_tir/pipeline.py
@@ -61,8 +61,6 @@ def default_s_tir_pipeline(*, prepare_only=False):
]
if not bool(config.get("tirx.disable_storage_rewrite", False)):
passes.append(tirx.transform.StorageRewrite())
- if config.get("tirx.use_async_copy", False):
- passes.append(s_tir.transform.LowerAsyncDMA())
passes.extend(
[
s_tir.transform.HoistIfThenElse(),
@@ -80,8 +78,6 @@ def default_s_tir_pipeline(*, prepare_only=False):
passes.append(s_tir.transform.InjectPTXLDG32(True))
if not bool(config.get("tirx.disable_cse_tir", False)):
passes.append(tirx.transform.CommonSubexprElim())
- if bool(config.get("tirx.instrument_lwp", False)):
- passes.append(s_tir.transform.InstrumentProfileIntrinsics())
passes.extend(
[
# Bind the target first so that target-specific attributes are
available.
diff --git a/python/tvm/s_tir/transform/transform.py
b/python/tvm/s_tir/transform/transform.py
index 09e04b0a15..c6b2e94a19 100644
--- a/python/tvm/s_tir/transform/transform.py
+++ b/python/tvm/s_tir/transform/transform.py
@@ -342,17 +342,6 @@ def InjectPTXLDG32(enable_inject_ptx_intrin=True):
return _ffi_api.InjectPTXLDG32(enable_inject_ptx_intrin) # type: ignore
-def InstrumentProfileIntrinsics():
- """Insert intrinsic calls to instrument function and loop level profiling.
-
- Returns
- -------
- fpass : tvm.transform.Pass
- The result pass
- """
- return _ffi_api.InstrumentProfileIntrinsics() # type: ignore
-
-
def VerifyVTCMLimit(default_target=None):
"""Verify if the size of the allocated vtcm memory satisfies the limit.
@@ -420,17 +409,6 @@ def LowerThreadAllreduce():
return _ffi_api.LowerThreadAllreduce() # type: ignore
-def LowerAsyncDMA():
- """Lower async DMA to DMA.
-
- Returns
- -------
- fpass : tvm.transform.Pass
- The result pass
- """
- return _ffi_api.LowerAsyncDMA() # type: ignore
-
-
def InjectPTXAsyncCopy():
"""Rewrite global to shared memory copy on CUDA with asynchronous copy.
diff --git a/python/tvm/te/__init__.py b/python/tvm/te/__init__.py
index 9d9b767f6c..c8f22e5e38 100644
--- a/python/tvm/te/__init__.py
+++ b/python/tvm/te/__init__.py
@@ -19,7 +19,7 @@
"""Namespace for Tensor Expression Language"""
# expose all operators in tvm tirx.op
-from tvm.tirx import any, all, min_value, max_value, trace
+from tvm.tirx import any, all, min_value, max_value
from tvm.tirx import exp, erf, tanh, sigmoid, log, tan, cos, sin, sqrt, rsqrt,
floor, ceil
from tvm.tirx import sinh, cosh, log2, log10
from tvm.tirx import asin, asinh, acos, acosh, atan, atanh
diff --git a/python/tvm/tirx/__init__.py b/python/tvm/tirx/__init__.py
index b1d1062fec..819ab45680 100644
--- a/python/tvm/tirx/__init__.py
+++ b/python/tvm/tirx/__init__.py
@@ -62,11 +62,10 @@ from .function import PrimFunc, IndexMap
from .op import call_packed_lowered, call_cpacked_lowered,
register_intrin_lowering
from .op import call_packed, call_cpacked, call_intrin, call_pure_extern,
call_extern
from .op import CallFFIKernelAttr, call_ffi_kernel, TensorMapEncodeTiledAttr,
tensormap_encode_tiled
-from .op import call_llvm_intrin, call_llvm_pure_intrin, all, any, min_value,
max_value, trace
+from .op import call_llvm_intrin, call_llvm_pure_intrin, all, any, min_value,
max_value
from .op import tvm_stack_alloca, tvm_stack_make_shape, tvm_stack_make_array
-from .op import tvm_tuple, handle_add_byte_offset, tvm_struct_get,
tvm_struct_set
-from .op import address_of, lookup_param, assume, undef
-from .op import continue_loop, break_loop
+from .op import handle_add_byte_offset, tvm_struct_get, tvm_struct_set
+from .op import address_of, assume, undef
from .op import tvm_thread_allreduce, type_annotation, tvm_access_ptr,
ptr_byte_offset
from .op import tvm_throw_last_error
from .op import (
@@ -90,7 +89,6 @@ from .op import div, indexdiv, indexmod, truncdiv, truncmod,
floordiv, floormod,
from .op import comm_reducer, min, max, sum
from .op import q_multiply_shift, q_multiply_shift_per_axis, shift_left,
shift_right
from .op import TVMBackendAllocWorkspace, TVMBackendFreeWorkspace
-from .op import start_profile_intrinsic, end_profile_intrinsic
from .op import vscale, get_active_lane_mask, get_vscale_expr
from .op import dp4a
from .op import ignore_loop_partition
diff --git a/python/tvm/tirx/op.py b/python/tvm/tirx/op.py
index 115f739b07..c93e0da604 100644
--- a/python/tvm/tirx/op.py
+++ b/python/tvm/tirx/op.py
@@ -668,50 +668,6 @@ def undef():
return call_intrin("int32", "tirx.undef")
-def start_profile_intrinsic(id):
- """Start profile intrinsic.
- Parameters
- ----------
- id : int
- The intrinsic id.
- Returns
- -------
- call : Expr
- The call expression.
- """
- return call_intrin("void", "tirx.start_profile_intrinsic", id)
-
-
-def end_profile_intrinsic(id):
- """End profile intrinsic.
- Parameters
- ----------
- id : int
- The intrinsic id.
- Returns
- -------
- call : Expr
- The call expression.
- """
- return call_intrin("void", "tirx.end_profile_intrinsic", id)
-
-
-def tvm_tuple(*value):
- """Create a tuple structure in value field of AttrStmt
-
- Parameters
- ----------
- value : Expr
- The value in tuple.
-
- Returns
- -------
- call : Expr
- The call expression.
- """
- return call_intrin("void", "tirx.tvm_tuple", *value)
-
-
def handle_add_byte_offset(handle, offset):
"""Add offset to handle
@@ -839,25 +795,6 @@ def address_of(obj: Buffer | TensorLoad | Var, span: Span
| None = None) -> Expr
raise ValueError(f"Invalid object type: {type(obj)}")
-def lookup_param(param_name, span=None):
- """Returns the param by name
-
- Parameters
- ----------
- param_name : str
- The name of param.
-
- span : Optional[Span]
- The location of this operator in the source code.
-
- Returns
- -------
- call : Expr
- The call expression.
- """
- return call_intrin("handle", "tirx.lookup_param", param_name, span=span)
-
-
def tvm_thread_allreduce(*freduce_args):
"""Perform allreduce inside threadblock.
@@ -918,17 +855,6 @@ def tvm_kernel_replace_point():
return call_intrin("void", "tirx.tvm_kernel_replace_point")
-def tvm_global_barrier_kinit():
- """Initialize the global barrier.
-
- Returns
- -------
- call : Expr
- The call expression.
- """
- return call_intrin("void", "tirx.tvm_global_barrier_kinit")
-
-
def tvm_warp_shuffle(mask, value, warp_id, width, warp_size):
"""Exchange value between threads inside a warp.
@@ -1384,47 +1310,6 @@ def all(*args, span=None):
return val
-@tvm_ffi.register_global_func("tvm.default_trace_action")
-def _tvm_default_trace_action(*args):
- print(list(args))
-
-
-def trace(args, trace_action="tvm.default_trace_action"):
- """Trace tensor data at the runtime.
-
- The trace function allows to trace specific tensor at the
- runtime. The tracing value should come as last argument.
- The trace action should be specified, by default
- tvm.default_trace_action is used.
-
- Parameters
- ----------
- args : list of Expr or Buffers.
- Positional arguments.
-
- trace_action : str.
- The name of the trace action.
-
- Returns
- -------
- call : Expr
- The call expression.
-
- See Also
- --------
- tvm.tirx.call_packed : Creates packed function.
- """
- if not isinstance(args, list):
- raise Exception("tvm.tirx.trace consumes the args as list type")
- call_args = [
- _pack_buffer(x) if is_buffer_var(x) else _reject_buffer_region(x,
"trace") for x in args
- ]
- call_args.insert(0, tvm.ir.StringImm(trace_action))
- tracing_value = args[-1]
- ret_ty = tracing_value.ty if isinstance(tracing_value, Expr) else
tracing_value.dtype
- return tvm.ir.Call(Op.get("tirx.tvm_call_trace_packed"), call_args,
ret_ty=ret_ty)
-
-
def infinity(dtype: str, span: Span | None = None) -> Any:
"""infinity value of dtype
@@ -3384,45 +3269,5 @@ def tvm_store_matrix_sync(fragment, m, n, k, index,
buffer_ptr, stride, layout):
def thread_return():
- """TVM intrinsic to call thread_return()
-
- Returns
- -------
- call : Expr
- The call expression.
- """
+ """Return from the current GPU thread without a function value."""
return call_intrin("", "tirx.thread_return")
-
-
-def continue_loop(span=None):
- """Create a tir intrinsic call to represent continue expression
-
- Parameters
- ----------
- span : Optional[Span]
- The location of this operator in the source code.
-
- Returns
- -------
- ret : Expr
- The continue expression
- """
-
- return _ffi_api.continue_loop(span)
-
-
-def break_loop(span=None):
- """Create a tir intrinsic call to represent break expression
-
- Parameters
- ----------
- span : Optional[Span]
- The location of this operator in the source code.
-
- Returns
- -------
- ret : Expr
- The break expression
- """
-
- return _ffi_api.break_loop(span)
diff --git a/python/tvm/tirx/script/ir_builder/op.py
b/python/tvm/tirx/script/ir_builder/op.py
index 716ef281e4..97ab45c37b 100644
--- a/python/tvm/tirx/script/ir_builder/op.py
+++ b/python/tvm/tirx/script/ir_builder/op.py
@@ -464,9 +464,6 @@ log2 = _op_wrapper(_tir_op.log2)
log10 = _op_wrapper(_tir_op.log10)
-lookup_param = _op_wrapper(_tir_op.lookup_param)
-
-
max_value = _op_wrapper(_tir_op.max_value)
@@ -491,12 +488,6 @@ q_multiply_shift = _op_wrapper(_tir_op.q_multiply_shift)
q_multiply_shift_per_axis = _op_wrapper(_tir_op.q_multiply_shift_per_axis)
-continue_loop = _op_wrapper(_tir_op.continue_loop)
-
-
-break_loop = _op_wrapper(_tir_op.break_loop)
-
-
round = _op_wrapper(_tir_op.round) # pylint: disable=redefined-builtin
@@ -578,9 +569,6 @@ call_packed_lowered =
_op_wrapper(_tir_op.call_packed_lowered)
call_cpacked_lowered = _op_wrapper(_tir_op.call_cpacked_lowered)
-tvm_tuple = _op_wrapper(_tir_op.tvm_tuple)
-
-
handle_add_byte_offset = _op_wrapper(_tir_op.handle_add_byte_offset)
@@ -617,9 +605,6 @@ tvm_storage_sync = _tir_op.tvm_storage_sync
tvm_kernel_replace_point = _op_wrapper(_tir_op.tvm_kernel_replace_point)
-tvm_global_barrier_kinit = _tir_op.tvm_global_barrier_kinit
-
-
tvm_warp_shuffle = _tir_op.tvm_warp_shuffle
@@ -659,12 +644,6 @@ TVMBackendAllocWorkspace =
_op_wrapper(_tir_op.TVMBackendAllocWorkspace)
TVMBackendFreeWorkspace = _op_wrapper(_tir_op.TVMBackendFreeWorkspace)
-start_profile_intrinsic = _op_wrapper(_tir_op.start_profile_intrinsic)
-
-
-end_profile_intrinsic = _op_wrapper(_tir_op.end_profile_intrinsic)
-
-
anylist_getitem = _op_wrapper(_tir_op.anylist_getitem)
@@ -929,7 +908,6 @@ __all__ = [
"bitwise_not",
"bitwise_or",
"bitwise_xor",
- "break_loop",
"broadcast",
"call_cpacked",
"call_cpacked_lowered",
@@ -947,7 +925,6 @@ __all__ = [
"ceildiv",
"clz",
"comm_reducer",
- "continue_loop",
"cooperative_tensor_fill",
"cooperative_tensor_load",
"cooperative_tensor_multiply_accumulate",
@@ -956,7 +933,6 @@ __all__ = [
"cos",
"cosh",
"dp4a",
- "end_profile_intrinsic",
"erf",
"exp",
"exp2",
@@ -988,7 +964,6 @@ __all__ = [
"logical_and",
"logical_not",
"logical_or",
- "lookup_param",
"masked_load",
"masked_store",
"max",
@@ -1016,7 +991,6 @@ __all__ = [
"sin",
"sinh",
"sqrt",
- "start_profile_intrinsic",
"tan",
"tanh",
"tensormap_encode_tiled",
@@ -1031,7 +1005,6 @@ __all__ = [
"tvm_call_packed",
"tvm_call_packed_lowered",
"tvm_fill_fragment",
- "tvm_global_barrier_kinit",
"tvm_kernel_replace_point",
"tvm_load_matrix_sync",
"tvm_mma_sync",
@@ -1045,7 +1018,6 @@ __all__ = [
"tvm_thread_allreduce",
"tvm_thread_invariant",
"tvm_throw_last_error",
- "tvm_tuple",
"tvm_warp_activemask",
"tvm_warp_shuffle",
"tvm_warp_shuffle_down",
diff --git a/python/tvm/tirx/script/ir_builder/parser_protocol.py
b/python/tvm/tirx/script/ir_builder/parser_protocol.py
index f028f37575..90fa9184c7 100644
--- a/python/tvm/tirx/script/ir_builder/parser_protocol.py
+++ b/python/tvm/tirx/script/ir_builder/parser_protocol.py
@@ -668,7 +668,7 @@ def break_(*, span: _Span = None) ->
_base.AlreadyEmitted[_tir.Stmt]:
Legality is checked on the completed function, across loop and function
boundaries.
"""
- return _base.at_(span, evaluate(_op.break_loop()))
+ return _base.with_at_group_(span, lambda:
_base.AlreadyEmitted(_ffi_api.Break()))
def continue_(*, span: _Span = None) -> _base.AlreadyEmitted[_tir.Stmt]:
@@ -676,7 +676,7 @@ def continue_(*, span: _Span = None) ->
_base.AlreadyEmitted[_tir.Stmt]:
Legality is checked on the completed function, across loop and function
boundaries.
"""
- return _base.at_(span, evaluate(_op.continue_loop()))
+ return _base.with_at_group_(span, lambda:
_base.AlreadyEmitted(_ffi_api.Continue()))
def return_(value: Any = None, *, span: _Span = None) ->
_base.AlreadyEmitted[_tir.Stmt]:
diff --git a/src/backend/cuda/codegen/codegen_cuda.cc
b/src/backend/cuda/codegen/codegen_cuda.cc
index e7a06977c9..ce06c4f8a4 100644
--- a/src/backend/cuda/codegen/codegen_cuda.cc
+++ b/src/backend/cuda/codegen/codegen_cuda.cc
@@ -172,13 +172,6 @@ std::string GetFP4Type(const PrimType& type_ty) {
CodeGenCUDA::CodeGenCUDA(Target target) : target(target) { restrict_keyword_ =
"__restrict__"; }
-void CodeGenCUDA::Init(bool output_ssa) {
- CodeGenC::Init(output_ssa);
- vid_global_barrier_state_ =
name_supply_->FreshName(runtime::symbol::tvm_global_barrier_state);
- vid_global_barrier_expect_ = name_supply_->FreshName("__barrier_expect");
- TVM_FFI_ICHECK_EQ(vid_global_barrier_state_,
runtime::symbol::tvm_global_barrier_state);
-}
-
void CodeGenCUDA::PrintFunctionSignature(const ffi::String& function_name,
const PrimFunc& func,
std::ostream& os) {
CallingConv calling_conv =
@@ -843,35 +836,7 @@ void CodeGenCUDA::PrintStorageSync(const CallNode* op) {
this->PrintIndent();
this->stream << "__syncthreads();\n";
} else if (sync == "global") {
- if (!need_global_barrier_) {
- need_global_barrier_ = true;
- this->decl_stream << "extern \"C\" __device__ unsigned " <<
vid_global_barrier_state_
- << ";\n";
- }
- // global synchronizer
- std::string is_load = PrintExpr(op->args[1]);
- std::string num_blocks = PrintExpr(op->args[2]);
- this->PrintIndent();
- // In theory only threadfence is needed
- // but we observed problems with only threadfence
- this->stream << "__threadfence_system();\n";
- this->PrintIndent();
- this->stream << "if (" << is_load << ") {\n";
- int wb = this->BeginScope();
- this->PrintIndent();
- this->stream << "atomicAdd(&" << vid_global_barrier_state_ << ", 1);\n";
- this->PrintIndent();
- std::string ptr = name_supply_->FreshName("pf");
- this->stream << "volatile unsigned* " << ptr << " = &" <<
vid_global_barrier_state_ << ";\n";
- this->PrintIndent();
- this->stream << vid_global_barrier_expect_ << " += " << num_blocks <<
";\n";
- this->PrintIndent();
- this->stream << "while (" << ptr << "[0] < " << vid_global_barrier_expect_
<< ");\n";
- this->EndScope(wb);
- this->PrintIndent();
- this->stream << "}\n";
- this->PrintIndent();
- this->stream << "__syncthreads();\n";
+ TVM_FFI_THROW(InternalError) << "Global storage synchronization is
unsupported";
}
}
@@ -1747,19 +1712,7 @@ void CodeGenCUDA::Dispatch_(const AllocBufferNode* op) {
void CodeGenCUDA::Dispatch_(const EvaluateNode* op) {
if (auto value = op->value.as<PrimExpr>(); value &&
is_const_int(value.value())) return;
- const CallNode* call = op->value.as<CallNode>();
- if (call && call->op.same_as(tirx::builtin::tvm_global_barrier_kinit())) {
- PrintIndent();
- stream << "__shared__ unsigned " << vid_global_barrier_expect_ << ";\n";
- PrintIndent();
- stream << "if (threadIdx.x == 0) {\n";
- PrintIndent();
- stream << " " << vid_global_barrier_expect_ << " = 0;\n";
- PrintIndent();
- stream << "}\n";
- } else {
- CodeGenC::Dispatch_(op);
- }
+ CodeGenC::Dispatch_(op);
}
void CodeGenCUDA::Dispatch_(const prim::RampNode* op, std::ostream& os) {
diff --git a/src/backend/cuda/codegen/codegen_cuda.h
b/src/backend/cuda/codegen/codegen_cuda.h
index c7703f9e0b..c02e19d69a 100644
--- a/src/backend/cuda/codegen/codegen_cuda.h
+++ b/src/backend/cuda/codegen/codegen_cuda.h
@@ -39,7 +39,6 @@ namespace codegen {
class CodeGenCUDA final : public CodeGenC {
public:
CodeGenCUDA(Target target);
- void Init(bool output_ssa);
std::string Finish();
bool need_include_path() {
std::vector<std::string> tag_list{"fp16", "bf16", "int8", "fp8",
@@ -97,13 +96,6 @@ class CodeGenCUDA final : public CodeGenC {
// Whether scope such as "__shared__" or "__constant__" is part of type.
bool IsScopePartOfType() const final { return false; }
- // Whether global barrier is needed.
- bool need_global_barrier_{false};
- // Global barrier state
- std::string vid_global_barrier_state_;
- // Global barrier expected node.
- std::string vid_global_barrier_expect_;
-
// Whether clusterCtaIdx.x can be emitted as the linear cluster CTA rank.
// This is only semantics-preserving for effectively 1-D clusters where the
// y/z cluster-CTA extents are both one.
diff --git a/src/backend/hexagon/codegen/llvm/codegen_hexagon.cc
b/src/backend/hexagon/codegen/llvm/codegen_hexagon.cc
index 63312077d8..472b812ef0 100644
--- a/src/backend/hexagon/codegen/llvm/codegen_hexagon.cc
+++ b/src/backend/hexagon/codegen/llvm/codegen_hexagon.cc
@@ -211,27 +211,6 @@ llvm::Value* CodeGenHexagon::Dispatch_(const
TensorLoadNode* op) {
}
llvm::Value* CodeGenHexagon::CreateIntrinsic(const CallNode* op) {
- if (op->op.same_as(builtin::start_profile_intrinsic()) ||
- op->op.same_as(builtin::end_profile_intrinsic())) {
- llvm::Value* id = MakeValue(op->args[0].as_or_throw<PrimExpr>());
- auto instrprof_id = llvm::Intrinsic::hexagon_instrprof_custom;
-#if TVM_LLVM_VERSION >= 200
- llvm::Function* func = llvm::cast<llvm::Function>(
- llvm::Intrinsic::getOrInsertDeclaration(module_.get(), instrprof_id,
{}));
-#else
- llvm::Function* func = llvm::Intrinsic::getDeclaration(module_.get(),
instrprof_id);
-#endif
- llvm::GlobalVariable* name_var =
module_->getGlobalVariable("handler_name");
- if (!name_var) {
- llvm::StringRef init_str = "lwp_handler";
- llvm::Constant* init =
llvm::ConstantDataArray::getString(module_->getContext(), init_str);
-
- name_var = new llvm::GlobalVariable(*module_, init->getType(), true,
- llvm::GlobalValue::InternalLinkage,
init, "handler_name");
- }
- llvm::Type* t_int8_p_ = llvmGetPointerTo(t_int8_, 0);
- return builder_->CreateCall(func,
{llvm::ConstantExpr::getBitCast(name_var, t_int8_p_), id});
- }
return CodeGenCPU::CreateIntrinsic(op);
}
diff --git a/src/backend/hexagon/runtime/profiler/lwp_handler.S
b/src/backend/hexagon/runtime/profiler/lwp_handler.S
deleted file mode 100644
index 8cd02dd828..0000000000
--- a/src/backend/hexagon/runtime/profiler/lwp_handler.S
+++ /dev/null
@@ -1,124 +0,0 @@
-/*
- * 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.
- */
-
-/*
-Lightweight profiling handler to record processor cycles in a buffer
-(pointed by __lwp_buffer_ptr) for a given invocation of the handler. To keep
the
-buffer size within a resonable limit, we only recond data for the first 100
-invocation of the handler for a given loop or function ID (passed in R0
register).
-The buffer size wouldn't be a concern if the loops with only siblings are
getting
-profiled. However, since the instrumentation provides several different
profiling
-options, this approach ensures that they all function as expexted. We use
second
-buffer (pointed by __lwp_counter) to keep count of the calls made to
lwp_handler
-function for each function/loop.
-
-Brief explanation of all the global variables used:
-1) __lwp_counter : Pointer to the buffer that keeps count of the number of
times handler
- is called for a given ID. To reduce the complexity of the handler,
__lwp_counter is
- indexed using the ID itself.
-2) __lwp_buffer_ptr : Pointer to the buffer that records loop/function ID,
processor cycles
- and return addresss of the handler. Return address is used to reconstruct the
call graph
- (loop-nest) to make it easier to analyze the profiling data.
-3) __lwp_buffer_size : Size of the buffer
-4) __lwp_buffer_count : Offset into main lwp buffer where data for the current
handler
-invocation needs to be written.
-
-NOTE: The handler function saves and restores R0-R5 registers which are caller
saved registers
-on Hexagon and should be handled at the callsite. However, to reduce the
codegen impact
-of the handler calls on the caller functions, we decided to move this part
into the
-handler itself.
-
-*/
- .text
- .globl lwp_handler
- .falign
- .type lwp_handler,@function
-lwp_handler:
- {
- allocframe(#32) // Allocate 32 bytes on the
stack to save R0-R5 registers (6*4bytes) and P0-P3 (4*1byte) + 4 unused bytes
as the stack has to be 8-bytes aligned
- memd(r29+#-16) = r5:4 // Save R5,R4
- r5 = p3:0 // We will save P3:0 but we
need an intermediate usual register (R5) that has already been saved
- }
- {
- memd(r29+#16) = r3:2 // Save R3,R2
- memd(r29+#8) = r1:0 // Save R1, R0
- }
- {
- memw(r29+#0) = r5 // Save P3:0 (via R5)
- r2 = add(pc,##_GLOBAL_OFFSET_TABLE_@PCREL) // Get GOT address
- }
- {
- r5 = memw(r2+##__lwp_counter@GOT) // Get address of the pointer to
__lwp_counter
- r3 = memw(r2+##__lwp_buffer_count@GOT) // Get the address of
__lwp_buffer_count
- }
- {
- r5 = memw(r5+#0) // Get the address of
__lwp_counter (address of the main lwp buffer)
- r3 = memw(r3+#0) // Get the __lwp_buffer_count
value (offset into the main buffer)
- }
- {
- r4 = memw(r5+r0<<#2) // Get the handler invocation
count for the ID (passed in R0)
- r1 = memw(r2+##__lwp_buffer_size@GOT) // Get the address of
__lwp_buffer_size
- }
- {
- r4 = add(r4,#1) // Increment count
- memw(r5+r0<<#2) = r4.new // Update count in __lwp_counter
for a given ID
- r1 = memw(r1+#0) // Get the buffer size
- }
- {
- p0 = cmp.gtu(r4,#100) // Exit if count for a given ID
is greater than 100
- if (p0.new) jump:nt .LBB0_3
- r5 = memw(r2+##__lwp_buffer_ptr@GOT) // Get address of the pointer to
__lwp_buffer_ptr
- }
- {
- r5 = memw(r5+#0) // Get address of
__lwp_buffer_ptr
- r2 = memw(r2+##__lwp_buffer_count@GOT) // Get address of
__lwp_buffer_count
- }
- {
- r4 = add(r3,#4) // Increment the offset by 4
since 4 int32 values are stored for each invocation
- if (!cmp.gtu(r1,r4.new)) jump:t .LBB0_3 // Exit if the main lwp buffer
has run out of space
- }
- {
- r5 = addasl(r5,r3,#2) // Get the address where the
data needs to be recorded
- memw(r2+#0) = r4 // Save next offset into
__lwp_buffer_count
- }
- {
- memw(r5+#0) = r31 // Save return address of this
function
- r1:0 = C15:14 // Control registers that keep
processor cycle count (64-bits)
- memw(r5+#4) = r0 // Save loop/function ID
- }
- {
- memw(r5+#12) = r1 // Save upper 32 bits
- memw(r5+#8) = r0 // Save lower 32 bits
- }
- .falign
-.LBB0_3: // Restore the registers from
the stack
- {
- r1 = memw(r29+#0) // We will restore P3:0 but need
an intermediate usual register (R1) that hasn't already been restored
- r5:4 = memd(r29+#24) // Restore R5:4
- }
- {
- r3:2 = memd(r29+#16) // Restore R3:2
- p3:0 = r1 // Restore P3:0 (via R1, not yet
restored)
- }
- {
- r1:0 = memd(r29+#8) // Restore R1:0
- dealloc_return // Deallocate the stack and
return
- }
-.Lfunc_end0:
- .size lwp_handler, .Lfunc_end0-lwp_handler
diff --git a/src/backend/hexagon/runtime/profiler/prof_utils.cc
b/src/backend/hexagon/runtime/profiler/prof_utils.cc
deleted file mode 100644
index 45cbe60703..0000000000
--- a/src/backend/hexagon/runtime/profiler/prof_utils.cc
+++ /dev/null
@@ -1,78 +0,0 @@
-/*
- * 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.
- */
-
-#include <fstream>
-#include <iostream>
-#include <sstream>
-
-// The max loop/function id used among all lwp_handler calls. Since
-// the id is used to index into the lwp_counter buffer, the size of the
-// buffer must be equal or greater than the max possible id.
-#define LWP_COUNTER_SIZE 5000
-
-// LWP_BUFFER_SIZE needs to be at most 100 * LWP_COUNTER_SIZE since 100 is
-// the max number of entries recorded for each instrumented location.
-#define LWP_BUFFER_SIZE (LWP_COUNTER_SIZE * 100)
-
-uint32_t lwp_counter[LWP_COUNTER_SIZE] = {0};
-uint32_t lwp_buffer[LWP_BUFFER_SIZE];
-uint32_t* __lwp_counter = lwp_counter;
-uint32_t* __lwp_buffer_ptr = lwp_buffer;
-uint32_t __lwp_buffer_size = LWP_BUFFER_SIZE;
-uint32_t __lwp_enable_flag = 1;
-uint32_t __lwp_buffer_count = 0;
-
-bool WriteLWPOutput(const std::string& out_json) {
- std::ostringstream s;
- s << "{\n";
- s << "\t\"entries\":[\n";
- for (size_t i = 0; i < __lwp_buffer_count; i += 4) {
- s << "\t{\n";
- s << "\t\t\"ret\":" << std::dec << lwp_buffer[i] << ",\n";
- s << "\t\t\"id\":" << std::dec << lwp_buffer[i + 1] << ",\n";
- uint64_t pcycles = (static_cast<uint64_t>(lwp_buffer[i + 3]) << 32) +
lwp_buffer[i + 2];
- s << "\t\t\"cyc\":" << std::dec << pcycles << "\n";
- s << "\t}";
- if (i < __lwp_buffer_count - 4) {
- s << ",\n";
- }
- }
- s << "\t],\n\n";
- s << "\t\"loop_counts\":[\n";
- for (size_t i = 0; i < LWP_COUNTER_SIZE; i++) {
- s << "\t\t" << lwp_counter[i] / 2;
- if (i < LWP_COUNTER_SIZE - 1)
- s << ",\n";
- else
- s << "\n";
- }
- s << "\t]\n}\n";
- std::ofstream ofc(out_json);
- if (!ofc.is_open()) {
- return false;
- }
-
- ofc << s.str() << "\n";
-
- if (!ofc) {
- return false;
- }
- ofc.close();
- return true;
-}
diff --git a/src/backend/hexagon/runtime/profiler/prof_utils.h
b/src/backend/hexagon/runtime/profiler/prof_utils.h
deleted file mode 100644
index e086f7a26b..0000000000
--- a/src/backend/hexagon/runtime/profiler/prof_utils.h
+++ /dev/null
@@ -1,30 +0,0 @@
-/*
- * 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.
- */
-
-/*!
- * \file prof_utils.h
- */
-#ifndef TVM_RUNTIME_HEXAGON_PROFILER_PROF_UTILS_H_
-#define TVM_RUNTIME_HEXAGON_PROFILER_PROF_UTILS_H_
-
-#include <string>
-
-bool WriteLWPOutput(const std::string&);
-
-#endif // TVM_RUNTIME_HEXAGON_PROFILER_PROF_UTILS_H_
diff --git a/src/backend/hexagon/runtime/rpc/hexagon/rpc_server.cc
b/src/backend/hexagon/runtime/rpc/hexagon/rpc_server.cc
index 2e7255e170..fa68f57425 100644
--- a/src/backend/hexagon/runtime/rpc/hexagon/rpc_server.cc
+++ b/src/backend/hexagon/runtime/rpc/hexagon/rpc_server.cc
@@ -39,7 +39,6 @@ extern "C" {
#include "../../../../../runtime/rpc/minrpc/minrpc_server.h"
#include "../../hexagon_common.h"
#include "../../hexagon_device_api.h"
-#include "../../profiler/prof_utils.h"
#include "hexagon_rpc.h"
namespace tvm {
@@ -330,25 +329,12 @@ __attribute__((weak)) void _Parse_fde_instr() {}
TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
- refl::GlobalDef()
- .def_packed("tvm.hexagon.load_module",
- [](tvm::ffi::PackedArgs args, tvm::ffi::Any* rv) {
- auto soname = args[0].cast<std::string>();
- auto floader =
-
tvm::ffi::Function::GetGlobalRequired("ffi.Module.load_from_file.so");
- *rv = floader(soname, "so");
- })
- .def_packed(
- "tvm.hexagon.get_profile_output", [](tvm::ffi::PackedArgs args,
tvm::ffi::Any* rv) {
- auto profiling_mode = args[0].cast<std::string>();
- auto out_file = args[1].cast<std::string>();
- if (profiling_mode.compare("lwp") == 0) {
- *rv = WriteLWPOutput(out_file);
- } else {
- HEXAGON_PRINT(ERROR, "ERROR: Unsupported profiling mode: %s",
profiling_mode.c_str());
- *rv = false;
- }
- });
+ refl::GlobalDef().def_packed(
+ "tvm.hexagon.load_module", [](tvm::ffi::PackedArgs args, tvm::ffi::Any*
rv) {
+ auto soname = args[0].cast<std::string>();
+ auto floader =
tvm::ffi::Function::GetGlobalRequired("ffi.Module.load_from_file.so");
+ *rv = floader(soname, "so");
+ });
}
void SaveBinaryToFile(const std::string& file_name, const std::string& data) {
diff --git a/src/backend/hexagon/runtime/rpc/simulator/rpc_server.cc
b/src/backend/hexagon/runtime/rpc/simulator/rpc_server.cc
index c96ddc3fec..33a60e6bb6 100644
--- a/src/backend/hexagon/runtime/rpc/simulator/rpc_server.cc
+++ b/src/backend/hexagon/runtime/rpc/simulator/rpc_server.cc
@@ -30,7 +30,6 @@
#include "../../../../../runtime/rpc/minrpc/minrpc_server.h"
#include "../../hexagon_common.h"
-#include "../../profiler/prof_utils.h"
#include "hexagon_sim_proto.h"
#include "tvm/ffi/function.h"
@@ -333,25 +332,12 @@ __attribute__((weak)) void _Parse_fde_instr() {}
TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
- refl::GlobalDef()
- .def_packed("tvm.hexagon.load_module",
- [](tvm::ffi::PackedArgs args, tvm::ffi::Any* rv) {
- auto soname = args[0].cast<std::string>();
- auto floader =
-
tvm::ffi::Function::GetGlobalRequired("ffi.Module.load_from_file.so");
- *rv = floader(soname, "so");
- })
- .def_packed(
- "tvm.hexagon.get_profile_output", [](tvm::ffi::PackedArgs args,
tvm::ffi::Any* rv) {
- auto profiling_mode = args[0].cast<std::string>();
- auto out_file = args[1].cast<std::string>();
- if (profiling_mode.compare("lwp") == 0) {
- *rv = WriteLWPOutput(out_file);
- } else {
- HEXAGON_PRINT(ERROR, "ERROR: Unsupported profiling mode: %s",
profiling_mode.c_str());
- *rv = false;
- }
- });
+ refl::GlobalDef().def_packed(
+ "tvm.hexagon.load_module", [](tvm::ffi::PackedArgs args, tvm::ffi::Any*
rv) {
+ auto soname = args[0].cast<std::string>();
+ auto floader =
tvm::ffi::Function::GetGlobalRequired("ffi.Module.load_from_file.so");
+ *rv = floader(soname, "so");
+ });
}
void SaveBinaryToFile(const std::string& file_name, const std::string& data) {
diff --git a/src/s_tir/transform/annotate_irregular_loop.cc
b/src/s_tir/transform/annotate_irregular_loop.cc
index b7ac58630d..c5b5fc7306 100644
--- a/src/s_tir/transform/annotate_irregular_loop.cc
+++ b/src/s_tir/transform/annotate_irregular_loop.cc
@@ -77,13 +77,13 @@ class IrregularLoopAnnotator : public StmtExprMutator {
return res;
}
- UnchangedOr<Stmt> Mutate_(const EvaluateNode* op, InplaceMode inplace_mode)
final {
- if (const CallNode* call = op->value.as<CallNode>()) {
- if (call->op.same_as(tirx::builtin::continue_loop()) ||
- call->op.same_as(tirx::builtin::break_loop())) {
- has_jump_ = true;
- }
- }
+ UnchangedOr<Stmt> Mutate_(const BreakNode* op, InplaceMode inplace_mode)
final {
+ has_jump_ = true;
+ return ffi::Unchanged();
+ }
+
+ UnchangedOr<Stmt> Mutate_(const ContinueNode* op, InplaceMode inplace_mode)
final {
+ has_jump_ = true;
return ffi::Unchanged();
}
diff --git a/src/s_tir/transform/inject_virtual_thread.cc
b/src/s_tir/transform/inject_virtual_thread.cc
index 7240d164f0..678000d6c0 100644
--- a/src/s_tir/transform/inject_virtual_thread.cc
+++ b/src/s_tir/transform/inject_virtual_thread.cc
@@ -340,8 +340,6 @@ class VTInjector : public s_tir::IRMutatorWithAnalyzer {
: op->args[1];
return Call(op->ty, op->op, {op->args[0], data, offset, extent,
op->args[4]});
- } else if (op->op.same_as(tirx::builtin::tvm_context_id())) {
- return allow_share_ ? Expr(ffi::GetRef<Call>(op)) : Expr(var_);
} else {
return StmtExprMutator::Mutate_(op, inplace_mode);
}
diff --git a/src/s_tir/transform/lower_async_dma.cc
b/src/s_tir/transform/lower_async_dma.cc
deleted file mode 100644
index f2f1d9fdc5..0000000000
--- a/src/s_tir/transform/lower_async_dma.cc
+++ /dev/null
@@ -1,215 +0,0 @@
-/*
- * 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.
- */
-
-/*!
- * \file lower_async_dma.cc
- */
-
-#include <tvm/ffi/cast.h>
-#include <tvm/ffi/reflection/registry.h>
-#include <tvm/runtime/logging.h>
-#include <tvm/s_tir/analysis.h>
-#include <tvm/s_tir/stmt.h>
-#include <tvm/s_tir/stmt_functor.h>
-#include <tvm/s_tir/transform.h>
-#include <tvm/sym/analyzer.h>
-#include <tvm/sym/iter_affine_map.h>
-#include <tvm/tirx/expr.h>
-#include <tvm/tirx/stmt.h>
-
-#include <optional>
-#include <set>
-
-#include "../../s_tir/ir/ir_mutator_with_analyzer.h"
-#include "../../tirx/transform/ir_utils.h"
-
-namespace tvm {
-namespace s_tir {
-using namespace tvm::tirx;
-
-class AsyncDMALowerer : public s_tir::IRMutatorWithAnalyzer {
- public:
- using s_tir::IRMutatorWithAnalyzer::Mutate;
- using s_tir::IRMutatorWithAnalyzer::Mutate_;
-
- explicit AsyncDMALowerer(bool dma_bypass_cache, const sym::Analyzer&
analyzer)
- : IRMutatorWithAnalyzer(analyzer), dma_bypass_cache_(dma_bypass_cache) {}
-
- // TODO(leiwang1999): split lower async DMA support for CUDA and Hexagon
Backend
- UnchangedOr<Stmt> Mutate_(const ForNode* loop, InplaceMode inplace_mode)
final {
- // if for loop is not within async_commit_queue_scope
- if (!async_queue_id_.has_value()) {
- return s_tir::IRMutatorWithAnalyzer::Mutate_(loop, inplace_mode);
- }
-
- // if for loop is not a memcpy of a contiguous region, it might be a cuda
cp.async behavior
- std::optional<s_tir::MemCpyDetails> mem_copy =
- s_tir::IdentifyMemCpy(ffi::GetRef<For>(loop),
ffi::GetRef<sym::Analyzer>(analyzer_));
- if (!mem_copy.has_value() || mem_copy->dest->region.size() != 1 ||
- mem_copy->source->region.size() != 1) {
- return s_tir::IRMutatorWithAnalyzer::Mutate_(loop, inplace_mode);
- }
-
- // now that we are about to perform the `copy` transform
- // save queue ID for inspection in `wait` transform
- // and, increment the number of DMA copies in the group
- queue_ids_.insert(async_queue_id_.value());
- dmas_in_group_++;
-
- tvm::PrimExpr src_min = mem_copy->source->region[0]->min;
- tvm::PrimExpr dst_min = mem_copy->dest->region[0]->min;
- tvm::PrimExpr dst_extent = mem_copy->dest->region[0]->extent;
-
- auto src = BufferLoad(mem_copy->source->source.as_or_throw<BufferVar>(),
{src_min});
- auto dst = BufferLoad(mem_copy->dest->source.as_or_throw<BufferVar>(),
{dst_min});
- PrimExpr dst_nbytes = dst_extent *
static_cast<int>(src.ty().StorageBytes());
- return Evaluate(
- Call(PrimType::Int(32), tirx::builtin::dma_copy(),
- ffi::Array<Expr>{
- PrimExpr(async_queue_id_.value()),
-
Call(mem_copy->dest->source.as_or_throw<BufferVar>().DataPointerType(),
- tirx::builtin::address_of(), ffi::Array<Expr>{dst},
Attrs(), {}, Span()),
-
Call(mem_copy->source->source.as_or_throw<BufferVar>().DataPointerType(),
- tirx::builtin::address_of(), ffi::Array<Expr>{src},
Attrs(), {}, Span()),
- dst_nbytes, PrimExpr(dma_bypass_cache_)},
- Attrs(), {}, Span())
- .as_or_throw<PrimExpr>());
- }
-
- UnchangedOr<Stmt> Mutate_(const AttrStmtNode* op, InplaceMode inplace_mode)
final {
- // populate analyzer knowledge of loop iterators
- auto previsit = s_tir::IRMutatorWithAnalyzer::Mutate_(op, inplace_mode)
- .ValueOrUnchanged(ffi::GetRef<Stmt>(op));
- if (!op->unique()) inplace_mode = InplaceMode::kDisallow;
-
- // Convert this, for example:
- // attr [0] "async_wait_queue_scope" = 0;
- // attr [0] "async_wait_inflight_count" = 0;
- //
- // To this:
- // @tirx.dma_wait(
- // 0, /* queue id */
- // 0, /* in flight count */
- // dtype=int32
- // )
- if (op->attr_key == s_tir::attr::async_wait_queue_scope) {
- // get queue ID
- auto queue_id_node = op->value.as<IntImmNode>();
- TVM_FFI_ICHECK(queue_id_node);
- int queue_id = queue_id_node->value.as<int>().value();
-
- // abort if we have not seen this queue ID in `copy` transform
- if (queue_ids_.find(queue_id) == queue_ids_.end()) {
- DLOG(INFO) << "AsyncDMALowerer exiting because the queue ID observed
in the "
- "`async_wait_queue_scope` transform has not been
previously observed in the "
- "`async_commit_queue_scope` transform";
- return previsit;
- }
-
- auto async_wait = op->body.as<AttrStmtNode>();
- if (!async_wait || async_wait->attr_key !=
s_tir::attr::async_wait_inflight_count) {
- DLOG(INFO) << "AsyncDMALowerer exiting because the body of the
`AttrStmtNode` with key "
- "`async_wait_queue_scope` does not contain an
`AttrStmtNode` with key "
- "`async_wait_inflight_count`";
- return previsit;
- }
- auto call_dma_wait = Evaluate(Call(PrimType::Int(32),
tirx::builtin::dma_wait(),
- {PrimExpr(queue_id),
async_wait->value})
- .as_or_throw<PrimExpr>());
-
- // The nested attribute is skipped by this descent.
- InplaceMode body_mode = async_wait->unique() ? inplace_mode :
InplaceMode::kDisallow;
- // concatenate the call with the body and return
- return SeqStmt({call_dma_wait, s_tir::IRMutatorWithAnalyzer::Mutate(
- ffi::AnyView(async_wait->body),
body_mode)
- .ValueOrUnchanged(async_wait->body)
- .as_or_throw<Stmt>()});
-
- // Convert this, for example:
- // attr [0] "async_commit_queue_scope" = 0;
- // attr [0] "async_scope" = 1;
- // for (ax0: int32, 0, 128) {
- // A_global[ax0] = A[ax0]
- // }
- //
- // To this:
- // @tirx.dma_copy(
- // 0, /* queue id */
- // @tirx.address_of(A_global[0], dtype=handle),
- // @tirx.address_of(A[0], dtype=handle),
- // 128, /* size */
- // dtype=int32
- // )
- } else if (op->attr_key == s_tir::attr::async_commit_queue_scope) {
- // get queue ID
- auto queue_id_node = op->value.as<IntImmNode>();
- TVM_FFI_ICHECK(queue_id_node);
- async_queue_id_ = queue_id_node->value.as<int>().value();
- auto result = s_tir::IRMutatorWithAnalyzer::Mutate_(op, inplace_mode)
- .ValueOrUnchanged(ffi::GetRef<Stmt>(op));
- if (dmas_in_group_ > 1) {
- auto call_dma_start_group =
- Evaluate(Call(PrimType::Int(32), tirx::builtin::dma_start_group(),
- {PrimExpr(async_queue_id_.value())})
- .as_or_throw<PrimExpr>());
- auto call_dma_end_group = Evaluate(Call(PrimType::Int(32),
tirx::builtin::dma_end_group(),
-
{PrimExpr(async_queue_id_.value())})
- .as_or_throw<PrimExpr>());
- result = SeqStmt({call_dma_start_group, result, call_dma_end_group});
- }
-
- async_queue_id_ = std::nullopt;
- dmas_in_group_ = 0;
- return result;
- }
- return s_tir::IRMutatorWithAnalyzer::Mutate_(op, inplace_mode);
- }
-
- private:
- int dmas_in_group_ = 0;
- std::set<int> queue_ids_;
- std::optional<int> async_queue_id_ = std::nullopt;
- bool dma_bypass_cache_;
- ffi::Map<Var, Range> input_iters = ffi::Map<Var, Range>();
-};
-
-namespace transform {
-
-Pass LowerAsyncDMA() {
- auto pass_func = [=](PrimFunc f, IRModule m, PassContext ctx) {
- auto fptr = f.CopyOnWrite();
- sym::Analyzer analyzer;
- bool dma_bypass_cache =
- ctx->GetConfig<bool>("tirx.experimental_dma_bypass_cache",
false).value();
- fptr->body = ffi::make_object<AsyncDMALowerer>(dma_bypass_cache, analyzer)
- ->Mutate(fptr->body, InplaceMode::kAllow)
- .ValueOrUnchanged(std::move(fptr->body));
- return f;
- };
- return CreatePrimFuncPass(pass_func, 0, "s_tir.LowerAsyncDMA", {});
-}
-
-TVM_FFI_STATIC_INIT_BLOCK() {
- namespace refl = tvm::ffi::reflection;
- refl::GlobalDef().def("s_tir.transform.LowerAsyncDMA", LowerAsyncDMA);
-}
-} // namespace transform
-
-} // namespace s_tir
-} // namespace tvm
diff --git a/src/s_tir/transform/profile_instrumentation.cc
b/src/s_tir/transform/profile_instrumentation.cc
deleted file mode 100644
index 458f0e06de..0000000000
--- a/src/s_tir/transform/profile_instrumentation.cc
+++ /dev/null
@@ -1,314 +0,0 @@
-/*
- * 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.
- */
-
-/*!
- * \file profile_instrumentation.cc
- */
-// Insert profile intrinsic at loop and function level. During codegen,
-// these instruction can be replaced with a call to a target specific handler
-// and can be used to capture profiling information such as processor cycles.
-
-#include <tvm/ffi/reflection/registry.h>
-#include <tvm/ir/prim/builtin.h>
-#include <tvm/ir/prim/expr.h>
-#include <tvm/s_tir/stmt_functor.h>
-#include <tvm/s_tir/transform.h>
-#include <tvm/tirx/builtin.h>
-#include <tvm/tirx/stmt.h>
-
-namespace tvm {
-namespace s_tir {
-using namespace tvm::tirx;
-namespace lwp {
-
-TVM_REGISTER_PASS_CONFIG_OPTION("s_tir.lwp_disable_func_prof", bool);
-TVM_REGISTER_PASS_CONFIG_OPTION("s_tir.lwp_max_depth", int64_t);
-TVM_REGISTER_PASS_CONFIG_OPTION("s_tir.lwp_min_height", int64_t);
-TVM_REGISTER_PASS_CONFIG_OPTION("s_tir.instr_siblings", bool);
-TVM_REGISTER_PASS_CONFIG_OPTION("s_tir.reset_start_id", bool);
-
-static int32_t start_id = 0;
-
-struct LoopInfo {
- LoopInfo() = default;
- LoopInfo(unsigned i, unsigned d, unsigned h = 0) : id(i), depth(d),
height(h) {
- has_siblings = false;
- has_parallel = false;
- }
- unsigned id;
- int32_t depth;
- int32_t height;
- bool has_siblings;
- // Set to 'true' if ForKind::kParallel is set for the current loop or one of
its ancestor
- bool has_parallel;
-};
-
-using LoopInfoMap = std::unordered_map<const ForNode*, LoopInfo>;
-// Traverse loops depth first and assign them a unique number.
-class LoopAnalyzer : public StmtExprVisitor {
- public:
- using StmtExprVisitor::Visit_;
- LoopInfoMap Analyze(const Stmt& stmt) {
- this->Visit(stmt);
- return loops;
- }
- ffi::Optional<VisitInterrupt> Visit_(const ForNode* op) final {
- LoopInfo loop_info(start_id, 0);
- start_id++;
- loop_info.height = TraverseLoop(op->body, 0);
- loops[op] = loop_info;
- return std::nullopt;
- }
-
- unsigned TraverseLoop(const Stmt& stmt, unsigned parent_depth, bool
has_parallel = false) {
- if (stmt->IsInstance<SeqStmtNode>()) {
- std::vector<const ForNode*> siblings;
- unsigned height = 0;
- bool has_loop = false;
- const SeqStmtNode* n = stmt.as<SeqStmtNode>();
- for (Stmt s : n->seq) {
- if (s->IsInstance<ForNode>()) {
- has_loop = true;
- const ForNode* f = s.as<ForNode>();
- LoopInfo loop_info(start_id, parent_depth + 1);
- start_id++;
- bool parent_parallel = false;
- if (has_parallel) {
- loop_info.has_parallel = true;
- parent_parallel = true;
- } else if (f->kind == ForKind::kParallel) {
- // has_parallel for the current loop is being set to 'false' since
the
- // intrinsic is added outside of the loop. The instrumentation
isn't
- // allowed for the subsequent nested loops.
- loop_info.has_parallel = false;
- parent_parallel = true;
- }
- siblings.push_back(f);
- height = std::max(height, TraverseLoop(f->body, parent_depth + 1,
parent_parallel));
- loop_info.height = height;
- loops[f] = loop_info;
- }
- }
- if (siblings.size() > 1) {
- for (auto* l : siblings) {
- loops[l].has_siblings = true;
- }
- }
- height = has_loop ? height + 1 : height;
- return height; // Parent's height : max of all children's height
- } else if (stmt->IsInstance<IfThenElseNode>()) {
- const IfThenElseNode* n = stmt.as<IfThenElseNode>();
- unsigned height = TraverseLoop(n->then_case, parent_depth, has_parallel);
- if (n->else_case) {
- height = std::max(height, TraverseLoop(n->else_case.value(),
parent_depth, has_parallel));
- }
- return height;
- } else if (stmt->IsInstance<ForNode>()) {
- const ForNode* f = stmt.as<ForNode>();
- LoopInfo loop_info(start_id, parent_depth + 1);
- start_id++;
- bool parent_parallel = false;
- if (has_parallel) {
- loop_info.has_parallel = true;
- parent_parallel = true;
- } else if (f->kind == ForKind::kParallel) {
- // has_parallel for the current loop is being set to 'false' since the
- // intrinsic is added outside of the loop. The instrumentation isn't
- // allowed for the subsequent nested loops.
- loop_info.has_parallel = false;
- parent_parallel = true;
- }
- unsigned height = TraverseLoop(f->body, parent_depth + 1,
parent_parallel);
- loop_info.height = height;
- loops[f] = loop_info;
- return height + 1;
- } else if (stmt->IsInstance<BindNode>()) {
- // Bind has no body; skip it and return 0 (not a loop).
- return 0;
- } else if (stmt->IsInstance<SeqStmtNode>()) {
- // For flat sequences, traverse children looking for loops.
- const SeqStmtNode* seq = stmt.as<SeqStmtNode>();
- unsigned max_height = 0;
- for (const auto& s : seq->seq) {
- max_height = std::max(max_height, TraverseLoop(s, parent_depth,
has_parallel));
- }
- return max_height;
- } else if (stmt->IsInstance<AttrStmtNode>()) {
- const AttrStmtNode* n = stmt.as<AttrStmtNode>();
- return TraverseLoop(n->body, parent_depth, has_parallel);
- } else if (stmt->IsInstance<AllocBufferNode>()) {
- return 0;
- } else {
- return 0; // inner-most loop
- }
- }
-
- private:
- LoopInfoMap loops;
-};
-
-class InstrumentIntrin : public StmtExprMutator {
- public:
- using StmtExprMutator::Mutate;
- using StmtExprMutator::Mutate_;
- UnchangedOr<ffi::Any> Mutate(ffi::AnyView value, InplaceMode inplace_mode)
override {
- if (value.as<ExprNode>()) return ffi::Unchanged();
- return StmtExprMutator::Mutate(value, inplace_mode);
- }
-
- InstrumentIntrin(int32_t max_depth, int32_t min_height, bool instr_siblings)
- : max_instr_depth_(max_depth),
- min_instr_height_(min_height),
- instr_siblings_(instr_siblings) {}
-
- void GetLoopInfo(PrimFuncNode* op) {
- auto analzer = ffi::make_object<LoopAnalyzer>();
- loops_ = analzer->Analyze(op->body);
- }
-
- UnchangedOr<Stmt> Mutate_(const ForNode* op, InplaceMode inplace_mode) final
{
- Stmt stmt = StmtExprMutator::Mutate_(op,
inplace_mode).ValueOrUnchanged(ffi::GetRef<Stmt>(op));
- if (loops_.count(op) < 1) return stmt;
-
- LoopInfo loop_info = loops_[op];
-
- if (loop_info.has_parallel) {
- return stmt;
- }
-
- // Exclude inner-most loops from instrumentation. The inner-most loop has
- // height '0' and it increases as we move outward in the loop nest.
- if (loop_info.height < min_instr_height_) {
- return stmt;
- }
-
- // Only instrument loops with a sibling
- if (instr_siblings_ && !loop_info.has_siblings) {
- return stmt;
- }
-
- // If instr_siblings_ is set, ignore max depth for instrumentation
- if (!instr_siblings_ && loop_info.depth > max_instr_depth_) {
- return stmt;
- }
- PrimExpr id = static_cast<int32_t>(loop_info.id);
- PrimExpr start_call = Call(PrimType::Void(),
tirx::builtin::start_profile_intrinsic(), {id})
- .as_or_throw<PrimExpr>();
- PrimExpr end_call = Call(PrimType::Void(),
tirx::builtin::end_profile_intrinsic(), {id})
- .as_or_throw<PrimExpr>();
- const Stmt start_profile = Evaluate(start_call);
- const Stmt end_profile = Evaluate(end_call);
- Stmt new_stmt = SeqStmt({start_profile, stmt, end_profile});
- return new_stmt;
- }
-
- private:
- LoopInfoMap loops_;
- int32_t max_instr_depth_;
- int32_t min_instr_height_;
- bool instr_siblings_;
-};
-
-class CheckParallelLoops : public StmtExprVisitor {
- public:
- using StmtExprVisitor::Visit_;
- bool HasParallelLoops(const Stmt& stmt) {
- this->Visit(stmt);
- return has_parallel;
- }
-
- private:
- ffi::Optional<VisitInterrupt> Visit_(const ForNode* op) final {
- if (op->kind == ForKind::kParallel) {
- has_parallel = true;
- } else {
- TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(StmtExprVisitor::Visit_(op));
- }
- return std::nullopt;
- }
-
- bool has_parallel = false;
-};
-
-PrimFunc AddProfileBuiltins(PrimFunc func, int32_t max_instr_depth, int32_t
min_instr_height,
- bool instr_siblings, bool
disable_func_instrumentation) {
- auto* func_ptr = func.CopyOnWrite();
-
- PrimExpr e = start_id++;
- if (!disable_func_instrumentation) {
- PrimExpr start_call = Call(PrimType::Void(),
tirx::builtin::start_profile_intrinsic(), {e})
- .as_or_throw<PrimExpr>();
- PrimExpr end_call =
- Call(PrimType::Void(), tirx::builtin::end_profile_intrinsic(),
{e}).as_or_throw<PrimExpr>();
- const Stmt start_profile = Evaluate(start_call);
- const Stmt end_profile = Evaluate(end_call);
- func_ptr->body = SeqStmt({start_profile, std::move(func_ptr->body),
end_profile});
- }
- auto p = ffi::make_object<InstrumentIntrin>(max_instr_depth,
min_instr_height, instr_siblings);
- p->GetLoopInfo(func_ptr);
- func_ptr->body =
- p->Mutate(func_ptr->body,
InplaceMode::kAllow).ValueOrUnchanged(std::move(func_ptr->body));
- return func;
-}
-
-} // namespace lwp
-
-namespace transform {
-Pass InstrumentProfileIntrinsics() {
- auto pass_func = [](IRModule m, PassContext ctx) {
- auto* mptr = m.CopyOnWrite();
-
- // All loops with depth <= max_instr_depth are instrumented. By default,
- // only outer-most loops are instrumented which has a depth of 0.
- // In addition, loops with siblings are also instrumented provided
- // their loop depth is >= min_instr_height. This is done to avoid
- // instrumenting inner-most loops.
- auto max_instr_depth = ctx->GetConfig<int64_t>("s_tir.lwp_max_depth",
0).value();
- auto min_instr_height = ctx->GetConfig<int64_t>("s_tir.lwp_min_height",
1).value();
- bool instr_siblings = ctx->GetConfig<bool>("s_tir.instr_siblings",
true).value();
- bool disable_func_instrumentation =
- ctx->GetConfig<bool>("s_tir.lwp_disable_func_prof", false).value();
- bool reset_start_id = ctx->GetConfig<bool>("s_tir.reset_start_id",
false).value();
- if (reset_start_id) lwp::start_id = 0;
- std::vector<std::pair<GlobalVar, PrimFunc>> updates;
- for (const auto& kv : mptr->functions) {
- if (auto func = kv.second.as<PrimFunc>()) {
- auto updated_func = lwp::AddProfileBuiltins(func.value(),
max_instr_depth, min_instr_height,
- instr_siblings,
disable_func_instrumentation);
- updates.push_back({kv.first, updated_func});
- }
- }
- for (const auto& pair : updates) {
- mptr->AddUnchecked(pair.first, pair.second);
- }
- return m;
- };
-
- return tvm::transform::CreateModulePass(pass_func, 0,
"s_tir.InstrumentProfileIntrinsics", {});
-}
-
-TVM_FFI_STATIC_INIT_BLOCK() {
- namespace refl = tvm::ffi::reflection;
- refl::GlobalDef().def("s_tir.transform.InstrumentProfileIntrinsics",
InstrumentProfileIntrinsics);
-}
-
-} // namespace transform
-
-} // namespace s_tir
-} // namespace tvm
diff --git a/src/target/llvm/codegen_cpu.cc b/src/target/llvm/codegen_cpu.cc
index fe43f3726d..4c779daece 100644
--- a/src/target/llvm/codegen_cpu.cc
+++ b/src/target/llvm/codegen_cpu.cc
@@ -868,34 +868,6 @@ llvm::Value* CodeGenCPU::CreateCallPacked(const CallNode*
op) {
return pc.ret_value;
}
-llvm::Value* CodeGenCPU::CreateCallTracePacked(const CallNode* op) {
- TVM_FFI_ICHECK_EQ(op->args.size(), 5U);
- PackedCall pc = MakeCallPackedLowered(
- op->args, op->ty,
static_cast<int64_t>(op->args[2].as<IntImmNode>()->value),
- static_cast<int64_t>(op->args[3].as<IntImmNode>()->value), true);
- llvm::LLVMContext* ctx = llvm_target_->GetContext();
- // Get traced value.
- llvm::Value* traced_value = MakeValue(op->args[4]);
- // The update_block handles case when we need to update the return value.
- llvm::BasicBlock* update_block = llvm::BasicBlock::Create(*ctx,
"update_block", function_);
- // The continue_block handles case when we need to return original
- // traced value.
- llvm::BasicBlock* continue_block = llvm::BasicBlock::Create(*ctx,
"continue_block", function_);
-
- // Check the ret_type_code and create cmp instruction.
- llvm::Value* cmp = builder_->CreateICmpNE(
- pc.ret_type_index, llvm::ConstantInt::get(t_int_,
ffi::TypeIndex::kTVMFFINone));
- builder_->CreateCondBr(cmp, update_block, continue_block);
- builder_->SetInsertPoint(update_block);
- builder_->CreateBr(continue_block);
- builder_->SetInsertPoint(continue_block);
- // The return value depends on from what bb we come from.
- llvm::PHINode* phi_rvalue = builder_->CreatePHI(traced_value->getType(), 2);
- phi_rvalue->addIncoming(pc.ret_value, update_block);
- phi_rvalue->addIncoming(traced_value, pc.end_block);
- return phi_rvalue;
-}
-
llvm::Value* CodeGenCPU::RuntimeTVMFFIFunctionCall() {
if (f_tvm_ffi_func_call_ != nullptr) return f_tvm_ffi_func_call_;
return GetContextPtr(gv_tvm_ffi_func_call_);
@@ -1026,8 +998,6 @@ llvm::Value* CodeGenCPU::CreateIntrinsic(const CallNode*
op) {
const ffi::Array<Expr>& args = op->args;
if (op->op.same_as(tirx::builtin::tvm_call_packed_lowered())) {
return CreateCallPacked(op);
- } else if (op->op.same_as(tirx::builtin::tvm_call_trace_packed_lowered())) {
- return CreateCallTracePacked(op);
} else if (op->op.same_as(tirx::builtin::tvm_call_cpacked_lowered())) {
return CreateCallPacked(op);
} else if (op->op.same_as(tirx::builtin::tvm_static_handle())) {
diff --git a/src/target/llvm/codegen_cpu.h b/src/target/llvm/codegen_cpu.h
index 4d253ba38b..6e4a1d153c 100644
--- a/src/target/llvm/codegen_cpu.h
+++ b/src/target/llvm/codegen_cpu.h
@@ -143,8 +143,6 @@ class CodeGenCPU : public CodeGenLLVM {
const int64_t begin, const int64_t end,
bool use_string_lookup);
// create call into tvm packed function.
llvm::Value* CreateCallPacked(const CallNode* op);
- // Create trace call into tvm packed function.
- llvm::Value* CreateCallTracePacked(const CallNode* op);
// Create parallel launch
void CreateParallelLaunch(const Stmt& body, int num_task, std::string name =
"");
// Create a new compute scope.
diff --git a/src/target/llvm/codegen_llvm.cc b/src/target/llvm/codegen_llvm.cc
index ed727151d9..cdf5b73454 100644
--- a/src/target/llvm/codegen_llvm.cc
+++ b/src/target/llvm/codegen_llvm.cc
@@ -1457,26 +1457,6 @@ llvm::Value* CodeGenLLVM::CreateIntrinsic(const
CallNode* op) {
value->addIncoming(then_value, then_value_block);
value->addIncoming(else_value, else_value_block);
return value;
- } else if (op->op.same_as(tirx::builtin::continue_loop())) {
- TVM_FFI_ICHECK(!loop_frame_jump_tgts_.empty())
- << "the tirx.continue_loop should be inserted under at least one For
or While stmts.";
- builder_->CreateBr(loop_frame_jump_tgts_.back().first);
- // LLVM allows exactly one terminator in a single basic block
- // append a new dummy basic block to avoid error.
- llvm::BasicBlock* post_dummy =
- llvm::BasicBlock::Create(*llvm_target_->GetContext(),
"post_cont_dummy", function_);
- builder_->SetInsertPoint(post_dummy);
- return post_dummy;
- } else if (op->op.same_as(tirx::builtin::break_loop())) {
- TVM_FFI_ICHECK(!loop_frame_jump_tgts_.empty())
- << "the tirx.break_loop should be inserted under at least one For or
While stmts.";
- builder_->CreateBr(loop_frame_jump_tgts_.back().second);
- // LLVM allows exactly one terminator in a single basic block
- // append a new dummy basic block to avoid error.
- llvm::BasicBlock* post_dummy =
- llvm::BasicBlock::Create(*llvm_target_->GetContext(),
"post_break_dummy", function_);
- builder_->SetInsertPoint(post_dummy);
- return post_dummy;
} else if (op->op.same_as(tirx::builtin::reinterpret())) {
llvm::Type* target = GetLLVMType(ret_type);
llvm::Value* value = MakeValue(args[0]);
@@ -1510,10 +1490,6 @@ llvm::Value* CodeGenLLVM::CreateIntrinsic(const
CallNode* op) {
} else if (op->op.same_as(tirx::builtin::atomic_add())) {
// TODO(masahi): Support atomic for CPU backend
TVM_FFI_THROW(InternalError) << "CPU backend does not support atomic add
yet.";
- } else if (op->op.same_as(tirx::builtin::start_profile_intrinsic()) ||
- op->op.same_as(tirx::builtin::end_profile_intrinsic())) {
- LOG(INFO) << "Ignoring profile_intrinsic ... " << op->op;
- return nullptr;
} else if (op->op.same_as(tirx::builtin::assume())) {
llvm::Value* cond = MakeValue(args[0]);
return builder_->CreateAssumption(cond);
@@ -2192,6 +2168,22 @@ void CodeGenLLVM::Dispatch_(const ReturnNode* op) {
builder_->SetInsertPoint(ret_dummy);
}
+void CodeGenLLVM::Dispatch_(const BreakNode* op) {
+ TVM_FFI_ICHECK(!loop_frame_jump_tgts_.empty()) << "Break requires an
enclosing loop.";
+ builder_->CreateBr(loop_frame_jump_tgts_.back().second);
+ llvm::BasicBlock* post_dummy =
+ llvm::BasicBlock::Create(*llvm_target_->GetContext(),
"post_break_dummy", function_);
+ builder_->SetInsertPoint(post_dummy);
+}
+
+void CodeGenLLVM::Dispatch_(const ContinueNode* op) {
+ TVM_FFI_ICHECK(!loop_frame_jump_tgts_.empty()) << "Continue requires an
enclosing loop.";
+ builder_->CreateBr(loop_frame_jump_tgts_.back().first);
+ llvm::BasicBlock* post_dummy =
+ llvm::BasicBlock::Create(*llvm_target_->GetContext(), "post_cont_dummy",
function_);
+ builder_->SetInsertPoint(post_dummy);
+}
+
void CodeGenLLVM::Dispatch_(const IfThenElseNode* op) {
EmitDebugLocation(op);
llvm::Value* cond = MakeValue(op->condition);
diff --git a/src/target/llvm/codegen_llvm.h b/src/target/llvm/codegen_llvm.h
index 98f991b1ba..93a6e9fe3f 100644
--- a/src/target/llvm/codegen_llvm.h
+++ b/src/target/llvm/codegen_llvm.h
@@ -241,6 +241,8 @@ class CodeGenLLVM : public
tirx::ExprFunctor<llvm::Value*(const Expr&)>,
void Dispatch_(const ForNode* op) override;
void Dispatch_(const WhileNode* op) override;
void Dispatch_(const ReturnNode* op) override;
+ void Dispatch_(const BreakNode* op) override;
+ void Dispatch_(const ContinueNode* op) override;
void Dispatch_(const IfThenElseNode* op) override;
void Dispatch_(const AllocBufferNode* op) override;
void Dispatch_(const AttrStmtNode* op) override;
@@ -594,7 +596,6 @@ class CodeGenLLVM : public
tirx::ExprFunctor<llvm::Value*(const Expr&)>,
const Op& builtin_call_pure_extern_ = builtin::call_pure_extern();
const Op& builtin_call_llvm_intrin_ = builtin::call_llvm_intrin();
const Op& builtin_call_llvm_pure_intrin_ = builtin::call_llvm_pure_intrin();
- const Op& builtin_lookup_param_ = builtin::lookup_param();
const Op& builtin_tvm_call_cpacked_lowered_ =
builtin::tvm_call_cpacked_lowered();
void EmitDebugLocation();
diff --git a/src/target/source/codegen_c.cc b/src/target/source/codegen_c.cc
index dd59944d85..c8de443497 100644
--- a/src/target/source/codegen_c.cc
+++ b/src/target/source/codegen_c.cc
@@ -714,10 +714,6 @@ void CodeGenC::Dispatch_(const CallNode* op, std::ostream&
os) { // NOLINT(*)
TVM_FFI_ICHECK(buffer && buffer->ty.as<BufferTypeNode>())
<< "buffer_data expects a Var with BufferType";
os << GetVarID(buffer);
- } else if (op->op.same_as(tirx::builtin::continue_loop())) {
- os << "continue;";
- } else if (op->op.same_as(tirx::builtin::break_loop())) {
- os << "break;";
} else if (op->op.same_as(builtin_call_extern_) ||
op->op.same_as(builtin_call_pure_extern_)) {
TVM_FFI_ICHECK_GE(op->args.size(), 1U);
auto func = op->args[0].as_or_throw<StringImm>();
@@ -892,11 +888,6 @@ void CodeGenC::Dispatch_(const CallNode* op, std::ostream&
os) { // NOLINT(*)
os << " != ";
this->PrintExpr(op->args[0], os);
os << ")";
- } else if (op->op.same_as(tirx::builtin::lookup_param())) {
- TVM_FFI_ICHECK_EQ(op->args.size(), 1);
- const StringImmNode* str = op->args[0].as<StringImmNode>();
- TVM_FFI_ICHECK(str != nullptr);
- os << "__tvm_param__" << str->value;
} else if (op->op.same_as(tirx::builtin::tvm_thread_invariant())) {
os << "(";
this->PrintExpr(op->args[0], os);
diff --git a/src/tirx/analysis/verify_well_formed.h
b/src/tirx/analysis/verify_well_formed.h
index f65b55d5b6..42820bf2a9 100644
--- a/src/tirx/analysis/verify_well_formed.h
+++ b/src/tirx/analysis/verify_well_formed.h
@@ -317,14 +317,6 @@ class LoopControlVerifier : public
Verifier<LoopControlVerifier<PathVisitor>, Pa
<< " requires an enclosing loop in the same
function.";
}
- void Dispatch_(const CallNode* op, AccessPath path) override {
- if (op->op.same_as(builtin::break_loop()) ||
op->op.same_as(builtin::continue_loop())) {
- Verify(loop_depth_ > 0) << "ValueError: " << op->op << " at " << path
- << " requires an enclosing loop in the same
function.";
- }
- PathVisitor::Dispatch_(op, path);
- }
-
int loop_depth_{0};
};
diff --git a/src/tirx/ir/transform.cc b/src/tirx/ir/transform.cc
index 9696ea09f0..eb512dd90b 100644
--- a/src/tirx/ir/transform.cc
+++ b/src/tirx/ir/transform.cc
@@ -45,7 +45,6 @@ TVM_REGISTER_PASS_CONFIG_OPTION("tirx.add_lower_pass",
ffi::Array<ffi::Array<ffi
TVM_REGISTER_PASS_CONFIG_OPTION("tirx.debug_keep_trivial_loop", bool);
TVM_REGISTER_PASS_CONFIG_OPTION("tirx.use_async_copy", bool);
TVM_REGISTER_PASS_CONFIG_OPTION("tirx.merge_static_smem", bool);
-TVM_REGISTER_PASS_CONFIG_OPTION("tirx.instrument_lwp", bool);
TVM_REGISTER_PASS_CONFIG_OPTION("tirx.vtcm_capacity", int64_t);
TVM_REGISTER_PASS_CONFIG_OPTION("tirx.s_tir.ldg32", bool);
TVM_REGISTER_PASS_CONFIG_OPTION("tirx.enable_fast_math", bool);
diff --git a/src/tirx/op/builtin.cc b/src/tirx/op/builtin.cc
index f2856276aa..24ef6af404 100644
--- a/src/tirx/op/builtin.cc
+++ b/src/tirx/op/builtin.cc
@@ -93,16 +93,6 @@ TIR_DEFINE_BUILTIN_FUNC(thread_return)
static_cast<int64_t>(CallEffectKind::kControlJump))
.set_num_inputs(0);
-TIR_DEFINE_BUILTIN_FUNC(continue_loop)
- .set_attr<TCallEffectKind>("TCallEffectKind",
-
static_cast<int64_t>(CallEffectKind::kControlJump))
- .set_num_inputs(0);
-
-TIR_DEFINE_BUILTIN_FUNC(break_loop)
- .set_attr<TCallEffectKind>("TCallEffectKind",
-
static_cast<int64_t>(CallEffectKind::kControlJump))
- .set_num_inputs(0);
-
// tirx.filter: escape hatch for non-canonical thread-set filter predicates
// used as an IfThenElse condition. (var, cond) -- ``var`` names the
// active-set axis the compiler should collapse to a singleton if it cannot
@@ -186,13 +176,6 @@ TIR_DEFINE_BUILTIN_FUNC(tvm_static_handle)
.set_attr<TCallEffectKind>("TCallEffectKind",
static_cast<int64_t>(CallEffectKind::kSpecialCallArg));
-TIR_DEFINE_BUILTIN_FUNC(tvm_context_id)
- .set_num_inputs(0)
- .set_attr<TCallEffectKind>("TCallEffectKind",
static_cast<int64_t>(CallEffectKind::kReadState));
-
-TIR_DEFINE_BUILTIN_FUNC(tvm_tuple).set_attr<TCallEffectKind>(
- "TCallEffectKind", static_cast<int64_t>(CallEffectKind::kEmbedInfo));
-
TIR_DEFINE_BUILTIN_FUNC(handle_add_byte_offset)
.set_num_inputs(2)
.set_attr<TCallEffectKind>("TCallEffectKind",
static_cast<int64_t>(CallEffectKind::kPure));
@@ -208,11 +191,6 @@ TIR_DEFINE_BUILTIN_FUNC(tvm_struct_set)
.set_attr<TCallEffectKind>("TCallEffectKind",
static_cast<int64_t>(CallEffectKind::kUpdateState));
-TIR_DEFINE_BUILTIN_FUNC(lookup_param)
- .set_num_inputs(4)
- .set_attr<TCallEffectKind>("TCallEffectKind",
-
static_cast<int64_t>(CallEffectKind::kUpdateState));
-
TIR_DEFINE_BUILTIN_FUNC(tvm_throw_last_error)
.set_num_inputs(0)
.set_attr<TCallEffectKind>("TCallEffectKind",
static_cast<int64_t>(CallEffectKind::kOpaque));
@@ -243,9 +221,6 @@ TIR_DEFINE_BUILTIN_FUNC(tvm_call_cpacked)
.set_attr<TCallEffectKind>("TCallEffectKind",
static_cast<int64_t>(CallEffectKind::kOpaque))
.set_attr<TScriptPrinterName>("TScriptPrinterName",
ffi::String("call_cpacked"), /*plevel=*/20);
-TIR_DEFINE_BUILTIN_FUNC(tvm_call_trace_packed)
- .set_attr<TCallEffectKind>("TCallEffectKind",
static_cast<int64_t>(CallEffectKind::kOpaque));
-
TIR_DEFINE_BUILTIN_FUNC(tvm_thread_invariant)
.set_num_inputs(1)
.set_attr<TCallEffectKind>("TCallEffectKind",
static_cast<int64_t>(CallEffectKind::kPure));
@@ -260,9 +235,6 @@ TIR_DEFINE_BUILTIN_FUNC(tvm_call_cpacked_lowered)
.set_attr<TScriptPrinterName>("TScriptPrinterName",
ffi::String("call_cpacked_lowered"),
/*plevel=*/20);
-TIR_DEFINE_BUILTIN_FUNC(tvm_call_trace_packed_lowered)
- .set_attr<TCallEffectKind>("TCallEffectKind",
static_cast<int64_t>(CallEffectKind::kOpaque));
-
// TODO(tvm-team) revisit storage sync once we have a good memory hierachy
structure.
TIR_DEFINE_BUILTIN_FUNC(tvm_storage_sync)
.set_attr<TCallEffectKind>("TCallEffectKind",
static_cast<int64_t>(CallEffectKind::kOpaque));
@@ -286,9 +258,6 @@ TIR_DEFINE_BUILTIN_FUNC(tvm_warp_shuffle_xor)
TIR_DEFINE_BUILTIN_FUNC(tvm_warp_activemask)
.set_attr<TCallEffectKind>("TCallEffectKind",
static_cast<int64_t>(CallEffectKind::kOpaque));
-TIR_DEFINE_BUILTIN_FUNC(tvm_global_barrier_kinit)
- .set_attr<TCallEffectKind>("TCallEffectKind",
static_cast<int64_t>(CallEffectKind::kOpaque));
-
TIR_DEFINE_BUILTIN_FUNC(tvm_thread_allreduce)
.set_attr<TCallEffectKind>("TCallEffectKind",
static_cast<int64_t>(CallEffectKind::kOpaque));
TIR_DEFINE_BUILTIN_FUNC(cooperative_tensor_fill)
@@ -337,18 +306,6 @@ TIR_DEFINE_BUILTIN_FUNC(texture2d_load)
.set_attr<TVectorizable>("TVectorizable", true)
.set_attr<TCallEffectKind>("TCallEffectKind",
static_cast<int64_t>(CallEffectKind::kOpaque));
-TIR_DEFINE_BUILTIN_FUNC(dma_copy).set_attr<TCallEffectKind>(
- "TCallEffectKind", static_cast<int64_t>(CallEffectKind::kOpaque));
-
-TIR_DEFINE_BUILTIN_FUNC(dma_wait).set_attr<TCallEffectKind>(
- "TCallEffectKind", static_cast<int64_t>(CallEffectKind::kOpaque));
-
-TIR_DEFINE_BUILTIN_FUNC(dma_start_group)
- .set_attr<TCallEffectKind>("TCallEffectKind",
static_cast<int64_t>(CallEffectKind::kOpaque));
-
-TIR_DEFINE_BUILTIN_FUNC(dma_end_group)
- .set_attr<TCallEffectKind>("TCallEffectKind",
static_cast<int64_t>(CallEffectKind::kOpaque));
-
TIR_DEFINE_BUILTIN_FUNC(assume)
.set_attr<TCallEffectKind>("TCallEffectKind",
static_cast<int64_t>(CallEffectKind::kEmbedInfo))
.set_num_inputs(1);
@@ -357,12 +314,6 @@ TIR_DEFINE_BUILTIN_FUNC(undef)
.set_attr<TCallEffectKind>("TCallEffectKind",
static_cast<int64_t>(CallEffectKind::kReadState))
.set_num_inputs(0);
-TIR_DEFINE_BUILTIN_FUNC(start_profile_intrinsic)
- .set_attr<TCallEffectKind>("TCallEffectKind",
static_cast<int64_t>(CallEffectKind::kPure));
-
-TIR_DEFINE_BUILTIN_FUNC(end_profile_intrinsic)
- .set_attr<TCallEffectKind>("TCallEffectKind",
static_cast<int64_t>(CallEffectKind::kPure));
-
TIR_DEFINE_BUILTIN_FUNC(anylist_getitem)
.set_attr<TCallEffectKind>("TCallEffectKind",
static_cast<int64_t>(CallEffectKind::kReadState));
diff --git a/src/tirx/op/op.cc b/src/tirx/op/op.cc
index 425824798b..e6c30042ec 100644
--- a/src/tirx/op/op.cc
+++ b/src/tirx/op/op.cc
@@ -135,21 +135,6 @@ PrimExpr q_multiply_shift(PrimExpr x, PrimExpr y, PrimExpr
q, PrimExpr s, Span s
.as_or_throw<PrimExpr>();
}
-PrimExpr thread_return(Span span) {
- return Call(PrimType::Void(), tirx::builtin::thread_return(), {}, {}, {},
span)
- .as_or_throw<PrimExpr>();
-}
-
-PrimExpr continue_loop(Span span) {
- return Call(PrimType::Void(), tirx::builtin::continue_loop(), {}, {}, {},
span)
- .as_or_throw<PrimExpr>();
-}
-
-PrimExpr break_loop(Span span) {
- return Call(PrimType::Void(), tirx::builtin::break_loop(), {}, {}, {}, span)
- .as_or_throw<PrimExpr>();
-}
-
TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
refl::GlobalDef().def("tirx.RegisterOpLowerIntrinsic",
@@ -159,13 +144,15 @@ TVM_FFI_STATIC_INIT_BLOCK() {
});
}
+PrimExpr thread_return(Span span) {
+ return Call(PrimType::Void(), tirx::builtin::thread_return(), {}, {}, {},
span)
+ .as_or_throw<PrimExpr>();
+}
+
TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
- refl::GlobalDef()
- .def("tirx.thread_return", thread_return)
- .def("tirx.continue_loop", continue_loop)
- .def("tirx.break_loop", break_loop);
-};
+ refl::GlobalDef().def("tirx.thread_return", thread_return);
+}
PrimExpr logaddexp(PrimExpr a, PrimExpr b, Span span) {
TVM_FFI_ICHECK(IsFloatType(a.ty())) << a;
diff --git a/src/tirx/script/ir_builder/ir.cc b/src/tirx/script/ir_builder/ir.cc
index 2687f0bae5..675bf816d1 100644
--- a/src/tirx/script/ir_builder/ir.cc
+++ b/src/tirx/script/ir_builder/ir.cc
@@ -483,9 +483,17 @@ tvm::tirx::Stmt Return(Expr value) {
return stmt;
}
-void Break() { AddToParent(tvm::tirx::Break(Span())); }
+tvm::tirx::Stmt Break() {
+ tvm::tirx::Stmt stmt = tvm::tirx::Break(Span());
+ AddToParent(stmt);
+ return stmt;
+}
-void Continue() { AddToParent(tvm::tirx::Continue(Span())); }
+tvm::tirx::Stmt Continue() {
+ tvm::tirx::Stmt stmt = tvm::tirx::Continue(Span());
+ AddToParent(stmt);
+ return stmt;
+}
IfFrame If(PrimExpr condition) {
ffi::ObjectPtr<IfFrameNode> n = ffi::make_object<IfFrameNode>();
diff --git a/src/tirx/transform/lower_tvm_builtin.cc
b/src/tirx/transform/lower_tvm_builtin.cc
index 9f14fe4d95..1d0e240b7f 100644
--- a/src/tirx/transform/lower_tvm_builtin.cc
+++ b/src/tirx/transform/lower_tvm_builtin.cc
@@ -427,18 +427,13 @@ class BuiltinLower : public StmtExprMutator {
args.push_back(IntImm(PrimType::Int(32), attr->force_cu_dtype));
}
Call packed(op->ty, builtin::tvm_call_packed(), args);
- return MakeCallPackedGeneric(packed.get(), 0,
builtin::tvm_call_packed_lowered(), false);
+ return MakeCallPackedGeneric(packed.get(), 0,
builtin::tvm_call_packed_lowered());
}
if (op->op.same_as(builtin::tvm_call_packed()) ||
(op->op.same_as(builtin::call_ffi_kernel()) && !preserve_ffi_kernel_))
{
- return MakeCallPackedGeneric(op, 0, builtin::tvm_call_packed_lowered(),
- /* use_last_value_as_traced_value*/ false);
+ return MakeCallPackedGeneric(op, 0, builtin::tvm_call_packed_lowered());
} else if (op->op.same_as(builtin::tvm_call_cpacked())) {
- return MakeCallPackedGeneric(op, 0, builtin::tvm_call_cpacked_lowered(),
- /* use_last_value_as_traced_value*/ false);
- } else if (op->op.same_as(builtin::tvm_call_trace_packed())) {
- return MakeCallPackedGeneric(op, 0,
builtin::tvm_call_trace_packed_lowered(),
- /* use_last_value_as_traced_value*/ true);
+ return MakeCallPackedGeneric(op, 0, builtin::tvm_call_cpacked_lowered());
} else if (op->op.same_as(builtin::anylist_setitem_call_packed())) {
return MakeAnyListSetItemCallPacked(op,
builtin::tvm_call_packed_lowered());
} else if (op->op.same_as(builtin::anylist_setitem_call_cpacked())) {
@@ -447,16 +442,6 @@ class BuiltinLower : public StmtExprMutator {
return MakeShape(op);
} else if (op->op.same_as(builtin::tvm_stack_make_array())) {
return MakeArray(op);
- } else if (op->op.same_as(builtin::tvm_context_id())) {
- return IntImm(op->ty.as_or_throw<PrimType>(), 0);
- } else if (op->op.same_as(builtin::dma_copy())) {
- return MakeDMACopy(op);
- } else if (op->op.same_as(builtin::dma_wait())) {
- return MakeDMAWait(op);
- } else if (op->op.same_as(builtin::dma_start_group())) {
- return MakeDMAStartGroup(op);
- } else if (op->op.same_as(builtin::dma_end_group())) {
- return MakeDMAEndGroup(op);
} else {
return StmtExprMutator::Mutate_(op, inplace_mode);
}
@@ -476,49 +461,6 @@ class BuiltinLower : public StmtExprMutator {
return StringImm("device_api." + device_name + "." + method_name);
}
- PrimExpr MakeDMACopy(const CallNode* op) {
- PrimExpr queue_id = op->args[0].as_or_throw<PrimExpr>();
- Expr dst = op->args[1];
- Expr src = op->args[2];
- PrimExpr size = op->args[3].as_or_throw<PrimExpr>();
- PrimExpr bypass_cache = op->args[4].as_or_throw<PrimExpr>();
-
- auto method_name = GetDeviceMethodName("dma_copy");
- Call call_packed = Call(PrimType::Int(32), builtin::tvm_call_packed(),
- {method_name, queue_id, dst, src, size,
bypass_cache});
- return Mutate(call_packed.as_or_throw<PrimExpr>(), InplaceMode::kDisallow)
- .ValueOrUnchanged(call_packed.as_or_throw<PrimExpr>());
- }
-
- PrimExpr MakeDMAWait(const CallNode* op) {
- PrimExpr queue_id = op->args[0].as_or_throw<PrimExpr>();
- PrimExpr inflight = op->args[1].as_or_throw<PrimExpr>();
-
- auto method_name = GetDeviceMethodName("dma_wait");
- Call call_packed =
- Call(PrimType::Int(32), builtin::tvm_call_packed(), {method_name,
queue_id, inflight});
- return Mutate(call_packed.as_or_throw<PrimExpr>(), InplaceMode::kDisallow)
- .ValueOrUnchanged(call_packed.as_or_throw<PrimExpr>());
- }
-
- PrimExpr MakeDMAStartGroup(const CallNode* op) {
- PrimExpr queue_id = op->args[0].as_or_throw<PrimExpr>();
-
- auto method_name = GetDeviceMethodName("dma_start_group");
- Call call_packed = Call(PrimType::Int(32), builtin::tvm_call_packed(),
{method_name, queue_id});
- return Mutate(call_packed.as_or_throw<PrimExpr>(), InplaceMode::kDisallow)
- .ValueOrUnchanged(call_packed.as_or_throw<PrimExpr>());
- }
-
- PrimExpr MakeDMAEndGroup(const CallNode* op) {
- PrimExpr queue_id = op->args[0].as_or_throw<PrimExpr>();
-
- auto method_name = GetDeviceMethodName("dma_end_group");
- Call call_packed = Call(PrimType::Int(32), builtin::tvm_call_packed(),
{method_name, queue_id});
- return Mutate(call_packed.as_or_throw<PrimExpr>(), InplaceMode::kDisallow)
- .ValueOrUnchanged(call_packed.as_or_throw<PrimExpr>());
- }
-
// call shape
Expr MakeShape(const CallNode* op) {
// if args.size() == 0, it represents a scalar shape ()
@@ -656,7 +598,7 @@ class BuiltinLower : public StmtExprMutator {
Expr list_handle = op->args[0];
PrimExpr list_index = op->args[1].as_or_throw<PrimExpr>();
- Call call = MakeCallPackedGeneric(op, 2, lowered_op, false);
+ Call call = MakeCallPackedGeneric(op, 2, lowered_op);
Expr args_stack = call->args[1];
// The stack offset of return value stack_end
PrimExpr ret_offset = call->args[3].as_or_throw<PrimExpr>();
@@ -675,10 +617,8 @@ class BuiltinLower : public StmtExprMutator {
* \param op The call
* \param name_offset The beginning of function name and call packed section.
* \param lowered_packed_op The target lowered op.
- * \param pass_last_arg_as_traced_value Whether to pass last argument as
traced value
*/
- Call MakeCallPackedGeneric(const CallNode* op, size_t name_offset, const Op&
lowered_packed_op,
- bool pass_last_arg_as_traced_value) {
+ Call MakeCallPackedGeneric(const CallNode* op, size_t name_offset, const Op&
lowered_packed_op) {
auto& scope = alloca_scope_.back();
auto& prep_seq = prep_seq_stack_.back();
@@ -725,11 +665,6 @@ class BuiltinLower : public StmtExprMutator {
ffi::Array<Expr> packed_args = {op->args[name_offset], scope.stack_ffi_any,
ConstInt32(arg_stack_begin),
ConstInt32(arg_stack_begin + num_args)};
- if (pass_last_arg_as_traced_value) {
- // pass in last element as traced value
- // used by call_packed_traced
- packed_args.push_back(op->args[op->args.size() - 1]);
- }
return Call(op->ty, lowered_packed_op, packed_args);
}
diff --git a/tests/python/codegen/test_target_codegen_cuda.py
b/tests/python/codegen/test_target_codegen_cuda.py
index bf7a0d2316..7e4b31492d 100644
--- a/tests/python/codegen/test_target_codegen_cuda.py
+++ b/tests/python/codegen/test_target_codegen_cuda.py
@@ -1152,8 +1152,6 @@ def test_device_host_call_same_func():
tvm.testing.run_with_gpu_lock(run_and_check)
[email protected]
[email protected](not env.has_cuda(), reason="need cuda")
def test_thread_return():
@I.ir_module
class Module:
@@ -1168,6 +1166,7 @@ def test_thread_return():
lib = tvm.compile(Module, target="cuda")
cuda_code = lib.mod.imports[0].inspect_source()
assert "return;" in cuda_code
+ assert "return 0;" not in cuda_code
@pytest.mark.gpu
diff --git a/tests/python/runtime/test_runtime_trace.py
b/tests/python/runtime/test_runtime_trace.py
deleted file mode 100644
index c2ddb042ce..0000000000
--- a/tests/python/runtime/test_runtime_trace.py
+++ /dev/null
@@ -1,226 +0,0 @@
-# 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.
-# ruff: noqa: F811
-import numpy as np
-
-import tvm
-from tvm import te
-
-
-def test_trace_default_action():
- n = 2
- x = te.placeholder((n, n, n), name="X", dtype="float32")
- y = te.compute(x.shape, lambda i, j, k: tvm.tirx.trace([i, j, k,
x[i][j][k]]))
- f = tvm.compile(te.create_prim_func([x, y]), target="llvm")
- xnd = tvm.runtime.tensor(np.ones((n, n, n), dtype=x.dtype.dtype))
- ynd = tvm.runtime.tensor(np.zeros((n, n, n), dtype=y.dtype.dtype))
- f(xnd, ynd)
-
-
-def test_trace_expr_assign():
- @tvm.register_global_func("tvm.tirx.trace_callback2")
- def trace_buffer(x):
- return
-
- def check_assign(dtype):
- n = 4
- x = te.placeholder((n, n, n), name="X", dtype=dtype)
- y = te.compute(
- x.shape, lambda i, j, k: tvm.tirx.trace([x[i][j][k]],
"tvm.tirx.trace_callback2")
- )
- z = te.compute(
- x.shape, lambda i, j, k: tvm.tirx.trace([y[i][j][k]],
"tvm.tirx.trace_callback2")
- )
- f = tvm.compile(te.create_prim_func([x, y, z]), "llvm")
-
- xnd = tvm.runtime.tensor(np.ones((n, n, n), dtype=x.dtype.dtype))
- ynd = tvm.runtime.tensor(np.zeros((n, n, n), dtype=y.dtype.dtype))
- znd = tvm.runtime.tensor(np.zeros((n, n, n), dtype=z.dtype.dtype))
- f(xnd, ynd, znd)
-
- assert np.array_equal(xnd.numpy(), np.ones((n, n, n)))
- assert np.array_equal(ynd.numpy(), np.ones((n, n, n)))
- assert np.array_equal(znd.numpy(), np.ones((n, n, n)))
-
- for t in ["float64", "float32", "int64", "int32"]:
- check_assign(t)
-
-
-def test_trace_expr_sum_generated():
- @tvm.register_global_func("tvm.tirx.trace_callback3")
- def trace_buffer(x):
- return
-
- def check_expr_sum(dtype):
- n = 4
- a = te.placeholder((n, n, n), name="a", dtype=dtype)
- b = te.placeholder((n, n, n), name="b", dtype=dtype)
- c = te.compute(
- a.shape,
- lambda i, j, k: (
- tvm.tirx.trace([a[i][j][k]], "tvm.tirx.trace_callback3")
- + tvm.tirx.trace([b[i][j][k]], "tvm.tirx.trace_callback3")
- ),
- )
- f = tvm.compile(te.create_prim_func([a, b, c]))
- xnd = tvm.runtime.tensor(np.array(np.ones((n, n, n),
dtype=a.dtype.dtype)))
- ynd = tvm.runtime.tensor(np.array(np.ones((n, n, n),
dtype=b.dtype.dtype)))
- znd = tvm.runtime.tensor(np.zeros((n, n, n), dtype=c.dtype.dtype))
- f(xnd, ynd, znd)
- assert np.array_equal(znd.numpy(), xnd.numpy() + ynd.numpy())
-
- for t in ["float64", "float32", "int64", "int32"]:
- check_expr_sum(t)
-
-
-def test_trace_expr_sum_args():
- @tvm.register_global_func("tvm.tirx.trace_silent")
- def silent(*args):
- return
-
- def check_expr_sum(dtype):
- n = 4
- a = te.placeholder((n, n, n), name="a", dtype=dtype)
- b = te.placeholder((n, n, n), name="b", dtype=dtype)
- e = te.placeholder((n, n, n), name="e", dtype=dtype)
- d = te.placeholder((n, n, n), name="d", dtype=dtype)
-
- c = te.compute(
- a.shape,
- lambda i, j, k: (
- tvm.tirx.trace([i, j, k, a[i][j][k]], "tvm.tirx.trace_silent")
- + tvm.tirx.trace([i, j, k, b[i][j][k]],
"tvm.tirx.trace_silent")
- + tvm.tirx.trace([i, j, k, d[i][j][k]],
"tvm.tirx.trace_silent")
- + tvm.tirx.trace([i, j, k, e[i][j][k]],
"tvm.tirx.trace_silent")
- ),
- )
- f = tvm.compile(te.create_prim_func([a, b, d, e, c]))
- a_nd = tvm.runtime.tensor(np.array(np.ones((n, n, n),
dtype=a.dtype.dtype)))
- b_nd = tvm.runtime.tensor(np.array(np.ones((n, n, n),
dtype=b.dtype.dtype)))
- d_nd = tvm.runtime.tensor(np.array(np.ones((n, n, n),
dtype=d.dtype.dtype)))
- e_nd = tvm.runtime.tensor(np.array(np.ones((n, n, n),
dtype=e.dtype.dtype)))
- c_nd = tvm.runtime.tensor(np.zeros((n, n, n), dtype=c.dtype.dtype))
- f(a_nd, b_nd, d_nd, e_nd, c_nd)
- assert np.array_equal(
- c_nd.numpy(), a_nd.numpy() + b_nd.numpy() + d_nd.numpy() +
e_nd.numpy()
- )
-
- for t in ["float64", "float32", "int64", "int32"]:
- check_expr_sum(t)
-
-
-def test_trace_expr_sum_custom():
- @tvm.register_global_func("tvm.tirx.trace_callback4")
- def trace_buffer(x):
- return
-
- def check_expr_sum_custom(dtype):
- n = 4
- a = te.placeholder((n, n), name="a", dtype=dtype)
- b = te.placeholder((n, n), name="b", dtype=dtype)
- c = te.compute(
- a.shape,
- lambda i, j: (
- tvm.tirx.trace([a[i][j]], "tvm.tirx.trace_callback4")
- + tvm.tirx.trace([b[i][j]], "tvm.tirx.trace_callback4")
- ),
- )
- f = tvm.compile(te.create_prim_func([a, b, c]))
- npa = np.array(
- [[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 1, 0], [0, 0, 0, 1]],
- dtype=a.dtype.dtype,
- )
- npb = np.array(
- [[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 1, 0], [0, 0, 0, 1]],
- dtype=a.dtype.dtype,
- )
- xnd = tvm.runtime.tensor(npa)
- ynd = tvm.runtime.tensor(npb)
- znd = tvm.runtime.tensor(np.zeros((n, n), dtype=c.dtype.dtype))
- f(xnd, ynd, znd)
- assert np.array_equal(znd.numpy(), npa + npb)
-
- for t in ["float64", "float32", "int64", "int32"]:
- check_expr_sum_custom(t)
-
-
-def test_trace_can_change_traced_value_int():
- @tvm.register_global_func("tvm.tirx.trace_change_int_first")
- def trace_buffer(x):
- return 13
-
- @tvm.register_global_func("tvm.tirx.trace_change_int_second")
- def trace_buffer(x):
- return 14
-
- def check_assign(dtype):
- n = 4
- x = te.placeholder((n,), name="X", dtype=dtype)
- y = te.compute(x.shape, lambda i: tvm.tirx.trace([x[i]],
"tvm.tirx.trace_change_int_first"))
- z = te.compute(
- x.shape, lambda i: tvm.tirx.trace([y[i]],
"tvm.tirx.trace_change_int_second")
- )
- f = tvm.compile(te.create_prim_func([x, y, z]))
-
- xnd = tvm.runtime.tensor(np.ones((n,), dtype=x.dtype.dtype))
- ynd = tvm.runtime.tensor(np.zeros((n,), dtype=y.dtype.dtype))
- znd = tvm.runtime.tensor(np.zeros((n,), dtype=z.dtype.dtype))
- f(xnd, ynd, znd)
- check_array_first = np.array([13, 13, 13, 13])
- check_array_second = np.array([14, 14, 14, 14])
- assert np.array_equal(ynd.numpy(), check_array_first)
- assert np.array_equal(znd.numpy(), check_array_second)
-
- for t in ["int64", "int32"]:
- check_assign(t)
-
-
-def test_trace_can_change_traced_value_float():
- @tvm.register_global_func("tvm.tirx.trace_change_float_first")
- def trace_buffer(x):
- return 13.0
-
- @tvm.register_global_func("tvm.tirx.trace_change_float_second")
- def trace_buffer(x):
- return 14.0
-
- def check_assign(dtype):
- n = 4
- x = te.placeholder((n,), name="X", dtype=dtype)
- y = te.compute(
- x.shape, lambda i: tvm.tirx.trace([x[i]],
"tvm.tirx.trace_change_float_first")
- )
- z = te.compute(
- x.shape, lambda i: tvm.tirx.trace([y[i]],
"tvm.tirx.trace_change_float_second")
- )
- f = tvm.compile(te.create_prim_func([x, y, z]), target="llvm")
-
- xnd = tvm.runtime.tensor(np.ones((n,), dtype=x.dtype.dtype))
- ynd = tvm.runtime.tensor(np.zeros((n,), dtype=y.dtype.dtype))
- znd = tvm.runtime.tensor(np.zeros((n,), dtype=z.dtype.dtype))
- f(xnd, ynd, znd)
- check_array_first = np.array([13.0, 13.0, 13.0, 13.0])
- check_array_second = np.array([14.0, 14.0, 14.0, 14.0])
- assert np.array_equal(ynd.numpy(), check_array_first)
- assert np.array_equal(znd.numpy(), check_array_second)
-
- for t in ["float64", "float32"]:
- check_assign(t)
-
-
-if __name__ == "__main__":
- tvm.testing.main()
diff --git a/tests/python/s_tir/script/test_s_tir_script_basic_usage.py
b/tests/python/s_tir/script/test_s_tir_script_basic_usage.py
index e81aa0c2f3..d5b5d711bb 100644
--- a/tests/python/s_tir/script/test_s_tir_script_basic_usage.py
+++ b/tests/python/s_tir/script/test_s_tir_script_basic_usage.py
@@ -1018,15 +1018,7 @@ def test_return_statement():
def test_loop_jump_statement():
- """`break` and `continue` evaluates to TIR intrinsics"""
-
- @Ts.prim_func
- def explicit():
- for i in range(16):
- if i % 2 == 0:
- T.evaluate(T.continue_loop())
- if i < 15:
- T.evaluate(T.break_loop())
+ """`break` and `continue` emit native TIR statements."""
@Ts.prim_func
def implicit():
@@ -1036,7 +1028,13 @@ def test_loop_jump_statement():
if i < 15:
break
- assert_structural_equal_ignore_global_symbol(implicit, explicit)
+ jumps = []
+ tvm_ffi.structural_walk(
+ implicit.body,
+ lambda node: jumps.append(node) if isinstance(node, tirx.Break |
tirx.Continue) else None,
+ )
+ assert sum(isinstance(node, tirx.Break) for node in jumps) == 1
+ assert sum(isinstance(node, tirx.Continue) for node in jumps) == 1
@pytest.mark.parametrize(
diff --git a/tests/python/s_tir/script/test_s_tir_script_printer.py
b/tests/python/s_tir/script/test_s_tir_script_printer.py
index a312109ac6..819a6edcf4 100644
--- a/tests/python/s_tir/script/test_s_tir_script_printer.py
+++ b/tests/python/s_tir/script/test_s_tir_script_printer.py
@@ -703,9 +703,9 @@ def main(A: T.Buffer((4,), "float32"), B: T.Buffer((4,),
"float32")):
for i in range(1000):
if i % 13 == 0:
A[1] = A[1] + T.float32(1.0)
- T.continue_loop()
+ continue
if A[0] >= B[0]:
- T.break_loop()
+ break
"""
_assert_print(main, expected_output)
diff --git
a/tests/python/s_tir/transform/test_s_tir_transform_inject_virtual_thread.py
b/tests/python/s_tir/transform/test_s_tir_transform_inject_virtual_thread.py
index cc559c7c0c..fd6c433eeb 100644
--- a/tests/python/s_tir/transform/test_s_tir_transform_inject_virtual_thread.py
+++ b/tests/python/s_tir/transform/test_s_tir_transform_inject_virtual_thread.py
@@ -46,7 +46,7 @@ def test_vthread():
"int32",
"Run",
B.access_ptr("r"),
- T.call_intrin("int32", "tirx.tvm_context_id"),
+ vt_x,
)
)
C_buf[i * nthread + vt_x] = B[i] + T.float32(1)
diff --git
a/tests/python/s_tir/transform/test_s_tir_transform_profiling_instr.py
b/tests/python/s_tir/transform/test_s_tir_transform_profiling_instr.py
deleted file mode 100644
index 848390aca4..0000000000
--- a/tests/python/s_tir/transform/test_s_tir_transform_profiling_instr.py
+++ /dev/null
@@ -1,364 +0,0 @@
-# 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.
-# ruff: noqa: E741, F401
-
-import numpy
-
-import tvm
-import tvm.testing
-from tvm import s_tir
-from tvm.ir.module import IRModule
-from tvm.script import s_tir as Ts
-from tvm.script import tirx as T
-
-default_lwp_test_config = {
- "tirx.instrument_lwp": True,
- "s_tir.lwp_disable_func_prof": True,
- "s_tir.reset_start_id": True,
-}
-
-
[email protected]_func
-def input1(
- A: T.Buffer((8, 8, 128), dtype="int32"),
- B: T.Buffer((8, 8, 128), dtype="int32"),
- C: T.Buffer((8, 8, 128), dtype="int32"),
-) -> None:
- for i, j in T.grid(8, 8):
- for k, l in T.grid(8, 16):
- with Ts.sblock("B"):
- vi, vj, vk, vl = Ts.axis.remap("SSSS", [i, j, k, l])
- B[vi, vj, vk * 16 + vl] = A[vi, vj, vk * 16 + vl] * 2
- for k, l in T.grid(8, 16):
- with Ts.sblock("C"):
- vi, vj, vk, vl = Ts.axis.remap("SSSS", [i, j, k, l])
- C[vi, vj, vk * 16 + vl] = B[vi, vj, vk * 16 + vl] * 2
-
-
[email protected]_func
-def input2(
- A: T.Buffer((8, 8, 128), dtype="int32"),
- B: T.Buffer((8, 8, 128), dtype="int32"),
- C: T.Buffer((8, 8, 128), dtype="int32"),
- D: T.Buffer((8, 8, 128), dtype="int32"),
-) -> None:
- for i in T.serial(0, 8):
- for j in T.serial(0, 8):
- for k, l in T.grid(8, 16):
- with Ts.sblock("B"):
- vi, vj, vk, vl = Ts.axis.remap("SSSS", [i, j, k, l])
- B[vi, vj, vk * 16 + vl] = A[vi, vj, vk * 16 + vl] * 2
- for k, l in T.grid(8, 16):
- with Ts.sblock("B"):
- vi, vj, vk, vl = Ts.axis.remap("SSSS", [i, j, k, l])
- B[vi, vj, vk * 16 + vl] = B[vi, vj, vk * 16 + vl] * D[vi,
vj, vk * 16 + vl]
- for j in T.serial(0, 8):
- for k, l in T.grid(8, 16):
- with Ts.sblock("C"):
- vi, vj, vk, vl = Ts.axis.remap("SSSS", [i, j, k, l])
- C[vi, vj, vk * 16 + vl] = B[vi, vj, vk * 16 + vl] + 2
- for k, l in T.grid(8, 16):
- with Ts.sblock("B"):
- vi, vj, vk, vl = Ts.axis.remap("SSSS", [i, j, k, l])
- C[vi, vj, vk * 16 + vl] = C[vi, vj, vk * 16 + vl] * D[vi,
vj, vk * 16 + vl]
-
-
[email protected]_func
-def input3(
- A: T.Buffer((8, 8, 128), dtype="int32"),
- B: T.Buffer((8, 8, 128), dtype="int32"),
- C: T.Buffer((8, 8, 128), dtype="int32"),
- D: T.Buffer((8, 8, 128), dtype="int32"),
-) -> None:
- for i in T.serial(0, 8):
- for j in T.parallel(0, 8):
- for k in T.serial(0, 8):
- for l in T.serial(0, 16):
- with Ts.sblock("B"):
- vi, vj, vk, vl = Ts.axis.remap("SSSS", [i, j, k, l])
- B[vi, vj, vk * 16 + vl] = A[vi, vj, vk * 16 + vl] * 2
- for k in T.serial(0, 8):
- for l in T.serial(0, 16):
- with Ts.sblock("B"):
- vi, vj, vk, vl = Ts.axis.remap("SSSS", [i, j, k, l])
- B[vi, vj, vk * 16 + vl] = B[vi, vj, vk * 16 + vl] *
D[vi, vj, vk * 16 + vl]
- for j in T.serial(0, 8):
- for k in T.parallel(0, 8):
- for l in T.serial(0, 16):
- with Ts.sblock("C"):
- vi, vj, vk, vl = Ts.axis.remap("SSSS", [i, j, k, l])
- C[vi, vj, vk * 16 + vl] = B[vi, vj, vk * 16 + vl] + 2
- for k in T.parallel(0, 8):
- for l in T.serial(0, 16):
- with Ts.sblock("B"):
- vi, vj, vk, vl = Ts.axis.remap("SSSS", [i, j, k, l])
- C[vi, vj, vk * 16 + vl] = C[vi, vj, vk * 16 + vl] *
D[vi, vj, vk * 16 + vl]
-
-
[email protected]_func
-def test1_expected_output(
- A: T.Buffer((8, 8, 128), dtype="int32"),
- B: T.Buffer((8, 8, 128), dtype="int32"),
- C: T.Buffer((8, 8, 128), dtype="int32"),
-) -> None:
- for i, j in T.grid(8, 8):
- T.evaluate(T.start_profile_intrinsic(3, dtype="handle"))
- for k, l in T.grid(8, 16):
- with Ts.sblock("B"):
- vi, vj, vk, vl = Ts.axis.remap("SSSS", [i, j, k, l])
- B[vi, vj, vk * 16 + vl] = A[vi, vj, vk * 16 + vl] * 2
- T.evaluate(T.end_profile_intrinsic(3, dtype="handle"))
- T.evaluate(T.start_profile_intrinsic(5, dtype="handle"))
- for k, l in T.grid(8, 16):
- with Ts.sblock("C"):
- vi, vj, vk, vl = Ts.axis.remap("SSSS", [i, j, k, l])
- C[vi, vj, vk * 16 + vl] = B[vi, vj, vk * 16 + vl] * 2
- T.evaluate(T.end_profile_intrinsic(5, dtype="handle"))
-
-
[email protected]_func
-def test2_expected_output(
- A: T.Buffer((8, 8, 128), dtype="int32"),
- B: T.Buffer((8, 8, 128), dtype="int32"),
- C: T.Buffer((8, 8, 128), dtype="int32"),
-) -> None:
- T.evaluate(T.start_profile_intrinsic(1, dtype="handle"))
- for i in T.serial(0, 8):
- T.evaluate(T.start_profile_intrinsic(2, dtype="handle"))
- for j in T.serial(0, 8):
- for k in T.serial(0, 8):
- for l in T.serial(0, 16):
- with Ts.sblock("B"):
- vi, vj, vk, vl = Ts.axis.remap("SSSS", [i, j, k, l])
- B[vi, vj, vk * 16 + vl] = A[vi, vj, vk * 16 + vl] * 2
- for k in T.serial(0, 8):
- for l in T.serial(0, 16):
- with Ts.sblock("C"):
- vi, vj, vk, vl = Ts.axis.remap("SSSS", [i, j, k, l])
- C[vi, vj, vk * 16 + vl] = B[vi, vj, vk * 16 + vl] * 2
- T.evaluate(T.end_profile_intrinsic(2, dtype="handle"))
- T.evaluate(T.end_profile_intrinsic(1, dtype="handle"))
-
-
[email protected]_func
-def test3_expected_output(
- A: T.Buffer((8, 8, 128), dtype="int32"),
- B: T.Buffer((8, 8, 128), dtype="int32"),
- C: T.Buffer((8, 8, 128), dtype="int32"),
-) -> None:
- T.evaluate(T.start_profile_intrinsic(1, dtype="handle"))
- for i in T.serial(0, 8):
- T.evaluate(T.start_profile_intrinsic(2, dtype="handle"))
- for j in T.serial(0, 8):
- T.evaluate(T.start_profile_intrinsic(3, dtype="handle"))
- for k in T.serial(0, 8):
- for l in T.serial(0, 16):
- with Ts.sblock("B"):
- vi, vj, vk, vl = Ts.axis.remap("SSSS", [i, j, k, l])
- B[vi, vj, vk * 16 + vl] = A[vi, vj, vk * 16 + vl] * 2
- T.evaluate(T.end_profile_intrinsic(3, dtype="handle"))
- T.evaluate(T.start_profile_intrinsic(5, dtype="handle"))
- for k in T.serial(0, 8):
- for l in T.serial(0, 16):
- with Ts.sblock("C"):
- vi, vj, vk, vl = Ts.axis.remap("SSSS", [i, j, k, l])
- C[vi, vj, vk * 16 + vl] = B[vi, vj, vk * 16 + vl] * 2
- T.evaluate(T.end_profile_intrinsic(5, dtype="handle"))
- T.evaluate(T.end_profile_intrinsic(2, dtype="handle"))
- T.evaluate(T.end_profile_intrinsic(1, dtype="handle"))
-
-
[email protected]_func
-def test4_expected_output(
- A: T.Buffer((8, 8, 128), dtype="int32"),
- B: T.Buffer((8, 8, 128), dtype="int32"),
- C: T.Buffer((8, 8, 128), dtype="int32"),
- D: T.Buffer((8, 8, 128), dtype="int32"),
-) -> None:
- for i in T.serial(0, 8):
- T.evaluate(T.start_profile_intrinsic(2, dtype="handle"))
- for j in T.serial(0, 8):
- T.evaluate(T.start_profile_intrinsic(3, dtype="handle"))
- for k, l in T.grid(8, 16):
- with Ts.sblock("B"):
- vi, vj, vk, vl = Ts.axis.remap("SSSS", [i, j, k, l])
- B[vi, vj, vk * 16 + vl] = A[vi, vj, vk * 16 + vl] * 2
- T.evaluate(T.end_profile_intrinsic(3, dtype="handle"))
- T.evaluate(T.start_profile_intrinsic(5, dtype="handle"))
- for k, l in T.grid(8, 16):
- with Ts.sblock("B"):
- vi, vj, vk, vl = Ts.axis.remap("SSSS", [i, j, k, l])
- B[vi, vj, vk * 16 + vl] = B[vi, vj, vk * 16 + vl] * D[vi,
vj, vk * 16 + vl]
- T.evaluate(T.end_profile_intrinsic(5, dtype="handle"))
- T.evaluate(T.end_profile_intrinsic(2, dtype="handle"))
- T.evaluate(T.start_profile_intrinsic(7, dtype="handle"))
- for j in T.serial(0, 8):
- T.evaluate(T.start_profile_intrinsic(8, dtype="handle"))
- for k, l in T.grid(8, 16):
- with Ts.sblock("C"):
- vi, vj, vk, vl = Ts.axis.remap("SSSS", [i, j, k, l])
- C[vi, vj, vk * 16 + vl] = B[vi, vj, vk * 16 + vl] + 2
- T.evaluate(T.end_profile_intrinsic(8, dtype="handle"))
- T.evaluate(T.start_profile_intrinsic(10, dtype="handle"))
- for k, l in T.grid(8, 16):
- with Ts.sblock("B"):
- vi, vj, vk, vl = Ts.axis.remap("SSSS", [i, j, k, l])
- C[vi, vj, vk * 16 + vl] = C[vi, vj, vk * 16 + vl] * D[vi,
vj, vk * 16 + vl]
- T.evaluate(T.end_profile_intrinsic(10, dtype="handle"))
- T.evaluate(T.end_profile_intrinsic(7, dtype="handle"))
-
-
[email protected]_func
-def test5_expected_output(
- A: T.Buffer((8, 8, 128), dtype="int32"),
- B: T.Buffer((8, 8, 128), dtype="int32"),
- C: T.Buffer((8, 8, 128), dtype="int32"),
-) -> None:
- T.evaluate(T.start_profile_intrinsic(1, dtype="handle"))
- for i in T.serial(0, 8):
- T.evaluate(T.start_profile_intrinsic(2, dtype="handle"))
- for j in T.serial(0, 8):
- for k in T.serial(0, 8):
- for l in T.serial(0, 16):
- with Ts.sblock("B"):
- vi, vj, vk, vl = Ts.axis.remap("SSSS", [i, j, k, l])
- B[vi, vj, vk * 16 + vl] = A[vi, vj, vk * 16 + vl] * 2
- for k in T.serial(0, 8):
- for l in T.serial(0, 16):
- with Ts.sblock("C"):
- vi, vj, vk, vl = Ts.axis.remap("SSSS", [i, j, k, l])
- C[vi, vj, vk * 16 + vl] = B[vi, vj, vk * 16 + vl] * 2
- T.evaluate(T.end_profile_intrinsic(2, dtype="handle"))
- T.evaluate(T.end_profile_intrinsic(1, dtype="handle"))
-
-
[email protected]_func
-def test6_expected_output(
- A: T.Buffer((8, 8, 128), dtype="int32"),
- B: T.Buffer((8, 8, 128), dtype="int32"),
- C: T.Buffer((8, 8, 128), dtype="int32"),
- D: T.Buffer((8, 8, 128), dtype="int32"),
-) -> None:
- for i in T.serial(0, 8):
- T.evaluate(T.start_profile_intrinsic(2, dtype="handle"))
- for j in T.parallel(0, 8):
- for k in T.serial(0, 8):
- for l in T.serial(0, 16):
- with Ts.sblock("B"):
- vi, vj, vk, vl = Ts.axis.remap("SSSS", [i, j, k, l])
- B[vi, vj, vk * 16 + vl] = A[vi, vj, vk * 16 + vl] * 2
- for k in T.serial(0, 8):
- for l in T.serial(0, 16):
- with Ts.sblock("B"):
- vi, vj, vk, vl = Ts.axis.remap("SSSS", [i, j, k, l])
- B[vi, vj, vk * 16 + vl] = B[vi, vj, vk * 16 + vl] *
D[vi, vj, vk * 16 + vl]
- T.evaluate(T.end_profile_intrinsic(2, dtype="handle"))
- T.evaluate(T.start_profile_intrinsic(7, dtype="handle"))
- for j in T.serial(0, 8):
- T.evaluate(T.start_profile_intrinsic(8, dtype="handle"))
- for k in T.parallel(0, 8):
- for l in T.serial(0, 16):
- with Ts.sblock("C"):
- vi, vj, vk, vl = Ts.axis.remap("SSSS", [i, j, k, l])
- C[vi, vj, vk * 16 + vl] = B[vi, vj, vk * 16 + vl] + 2
- T.evaluate(T.end_profile_intrinsic(8, dtype="handle"))
- T.evaluate(T.start_profile_intrinsic(10, dtype="handle"))
- for k in T.parallel(0, 8):
- for l in T.serial(0, 16):
- with Ts.sblock("B"):
- vi, vj, vk, vl = Ts.axis.remap("SSSS", [i, j, k, l])
- C[vi, vj, vk * 16 + vl] = C[vi, vj, vk * 16 + vl] *
D[vi, vj, vk * 16 + vl]
- T.evaluate(T.end_profile_intrinsic(10, dtype="handle"))
- T.evaluate(T.end_profile_intrinsic(7, dtype="handle"))
-
-
-# By default, only loops with siblings are instrumented.
-def test1():
- with tvm.transform.PassContext(config=default_lwp_test_config):
- mod = tvm.IRModule.from_expr(input1.with_attr("global_symbol", "main"))
- mod = tvm.s_tir.transform.InstrumentProfileIntrinsics()(mod)
- tvm.ir.assert_structural_equal(
- mod["main"], test1_expected_output.with_attr("global_symbol", "main")
- )
-
-
-# By default, only loops with siblings are instrumented. Here, 'lwp_max_depth'
-# doesn't have any effect unless 'instr_siblings' is set to False (ex: test3).
-def test2():
- test2_config = default_lwp_test_config.copy()
- test2_config.update({"s_tir.lwp_max_depth": 3})
- with tvm.transform.PassContext(config=test2_config):
- mod = tvm.IRModule.from_expr(input1.with_attr("global_symbol", "main"))
- mod = tvm.s_tir.transform.InstrumentProfileIntrinsics()(mod)
- tvm.ir.assert_structural_equal(
- mod["main"], test1_expected_output.with_attr("global_symbol", "main")
- )
-
-
-# test3: Use 'lwp_max_depth' to instrument loops upto a certain depth. This
flag
-# is effective only when 'instr_siblings' is disabled. Also, note that
inner-most
-# loops are always excluded from instrumentation unless overwritten using
-# 'lwp_min_height' (ex: test5)
-def test3():
- test3_config = default_lwp_test_config.copy()
- test3_config.update({"s_tir.lwp_max_depth": 3, "s_tir.instr_siblings":
False})
- with tvm.transform.PassContext(config=test3_config):
- mod = tvm.IRModule.from_expr(input1.with_attr("global_symbol", "main"))
- mod = tvm.s_tir.transform.InstrumentProfileIntrinsics()(mod)
- tvm.ir.assert_structural_equal(
- mod["main"], test3_expected_output.with_attr("global_symbol", "main")
- )
-
-
-# test4: Use 'lwp_min_height' to exclude inner loops upto a certain height from
-# instrumentation.
-def test4():
- with tvm.transform.PassContext(config=default_lwp_test_config):
- mod = tvm.IRModule.from_expr(input2.with_attr("global_symbol", "main"))
- mod = tvm.s_tir.transform.InstrumentProfileIntrinsics()(mod)
- tvm.ir.assert_structural_equal(
- mod["main"], test4_expected_output.with_attr("global_symbol", "main")
- )
-
-
-# test5: Use both 'lwp_min_height' and 'lwp_max_depth'.
-# instrumentation.
-def test5():
- test5_config = default_lwp_test_config.copy()
- test5_config.update(
- {"s_tir.lwp_max_depth": 3, "s_tir.instr_siblings": False,
"s_tir.lwp_min_height": 2}
- )
- with tvm.transform.PassContext(config=test5_config):
- mod = tvm.IRModule.from_expr(input1.with_attr("global_symbol", "main"))
- mod = tvm.s_tir.transform.InstrumentProfileIntrinsics()(mod)
- tvm.ir.assert_structural_equal(
- mod["main"], test5_expected_output.with_attr("global_symbol", "main")
- )
-
-
-# test6: Tests instrumentation for the parallel loops
-def test6():
- with tvm.transform.PassContext(config=default_lwp_test_config):
- mod = tvm.IRModule.from_expr(input3.with_attr("global_symbol", "main"))
- mod = tvm.s_tir.transform.InstrumentProfileIntrinsics()(mod)
- tvm.ir.assert_structural_equal(
- mod["main"], test6_expected_output.with_attr("global_symbol", "main")
- )
-
-
-if __name__ == "__main__":
- tvm.testing.main()
diff --git a/tests/python/script/test_basic_usage.py
b/tests/python/script/test_basic_usage.py
index 75bc8208a5..2e45c50fa9 100644
--- a/tests/python/script/test_basic_usage.py
+++ b/tests/python/script/test_basic_usage.py
@@ -26,7 +26,7 @@ import pytest
# ruff: noqa: F841
from minilang import Value
-from tvm import DataType, ir
+from tvm import DataType, error, ir
from tvm.ir import prim
from tvm.ir._overload_prim_expr import EqualOp
from tvm.script import ir as I
@@ -213,6 +213,104 @@ def
test_void_branch_statements_need_no_synthetic_named_output(language):
assert [value for kind, value in main.body] == [None, None, 9]
+def test_ordinary_iterator_binding_calls_the_custom_iterator_once(language):
+ # An ordinary iterator alias calls its captured implementation exactly
once.
+ M = language.M
+ calls = []
+
+ def custom_range(extent):
+ calls.append(extent)
+ return M.grid(2)
+
+ @M.function
+ def main():
+ iterator = custom_range
+ for i in iterator(4):
+ M.record(i)
+
+ assert calls == [4]
+ variable = main.body[0][1]
+ assert variable.op == "loop" and variable.args == (2,) and variable.name
== "i"
+
+
+def test_body_annotation_reads_a_preceding_ordinary_local(language):
+ # A body annotation must read the current local value, not a stale
enclosing capture.
+ M = language.M
+ annotations = []
+
+ def annotation(shape):
+ annotations.append(shape)
+ return M.Tensor(shape)
+
+ M.annotation = annotation
+
+ @M.function
+ def main():
+ shape = (4,)
+ value: M.annotation(shape) = 1
+ M.record(value)
+ shape = (8,)
+ M.record(shape)
+
+ assert annotations == [(4,)]
+ assert main.body == [("emit", 1), ("emit", (8,))]
+
+
+def test_native_concise_scopes_unwind_with_their_parent():
+ # Nested concise thread scopes must preserve the original variables in the
constructed IR.
+ from tvm import tirx
+ from tvm.script import tirx as T
+
+ variables = []
+
+ def observe(*items):
+ variables.extend(items)
+
+ @T.prim_func
+ def main():
+ bx = T.launch_thread("blockIdx.x", 2)
+ tx = T.launch_thread("threadIdx.x", 32)
+ observe(bx, tx)
+ T.evaluate(bx + tx)
+
+ bx, tx = variables
+ body = main.body
+ assert isinstance(body, tirx.AttrStmt) and isinstance(body.body,
tirx.AttrStmt)
+ assert body.node.var.same_as(bx) and body.body.node.var.same_as(tx)
+ assert body.body.body.value.a.same_as(bx) and
body.body.body.value.b.same_as(tx)
+
+
+def test_loop_control_validation_preserves_valid_and_unchecked_ir():
+ # Invalid loop placement must be rejected, while direct IR construction
preserves the node.
+ from tvm import ir, tirx
+ from tvm.script import tirx as T
+
+ invalid = tirx.PrimFunc(params=[], body=tirx.Break())
+ ir.assert_structural_equal(invalid.body, tirx.Break())
+ assert not tirx.analysis.verify_well_formed(invalid, assert_mode=False)
+ with pytest.raises(error.InternalError, match="requires an enclosing
loop"):
+ tirx.analysis.verify_well_formed(invalid)
+
+ @T.prim_func
+ def valid():
+ for i in range(2):
+ break
+
+ assert isinstance(valid.body, tirx.For)
+ ir.assert_structural_equal(valid.body.body, invalid.body)
+
+ @I.ir_module(check_well_formed=False, extra_vars={"invalid": invalid})
+ class Unchecked:
+ bad = invalid
+
+ assert Unchecked["bad"].same_as(invalid)
+ with pytest.raises(ValueError, match="requires an enclosing loop"):
+
+ @I.ir_module(extra_vars={"invalid": invalid})
+ class Rejected:
+ bad = invalid
+
+
def test_written_comparison_order(primitive_language):
# A leading host literal must not reverse the written native comparison or
its source span.
M = primitive_language.M
diff --git a/tests/python/tirx-base/test_tir_base.py
b/tests/python/tirx-base/test_tir_base.py
index 5118105b22..7a8dcd09ab 100644
--- a/tests/python/tirx-base/test_tir_base.py
+++ b/tests/python/tirx-base/test_tir_base.py
@@ -160,7 +160,7 @@ def test_control_flow_jump():
assert out == 1.0
-def test_break_loop():
+def test_break_statement():
@T.prim_func
def func(In: T.Buffer((2,), "int32"), Out: T.Buffer((2,), "int32")):
Out[0] = 0
@@ -187,7 +187,7 @@ def test_break_loop():
assert b[1] == 9
-def test_continue_loop():
+def test_continue_statement():
@T.prim_func
def func(Out: T.Buffer((2,), "int32")):
T.func_attr({"global_symbol": "main"})
diff --git a/tests/python/tirx-base/test_tir_constructor.py
b/tests/python/tirx-base/test_tir_constructor.py
index c4f697d650..bacd9f7579 100644
--- a/tests/python/tirx-base/test_tir_constructor.py
+++ b/tests/python/tirx-base/test_tir_constructor.py
@@ -223,7 +223,6 @@ def test_buffer_region_call_wrappers_reject():
lambda: tvm.tirx.call_cpacked("consume", region, 0),
lambda: tvm.tirx.call_packed_lowered("consume", region),
lambda: tvm.tirx.call_cpacked_lowered("consume", region, 0),
- lambda: tvm.tirx.trace([region]),
lambda: T.evaluate(region),
]
for call in calls:
diff --git a/tests/python/tirx-base/test_tir_op_types.py
b/tests/python/tirx-base/test_tir_op_types.py
index 3eb41d56fc..6ff7c05e6e 100644
--- a/tests/python/tirx-base/test_tir_op_types.py
+++ b/tests/python/tirx-base/test_tir_op_types.py
@@ -23,14 +23,6 @@ from tvm import tirx
from tvm.backend.cuda import op as _cuda_op
-def test_tir_op_tvm_tuple():
- x = tirx.Var("x", ty="float32")
- y = tirx.Var("y", ty="float32")
- z = tirx.Var("z", ty="float32")
- expr = tirx.tvm_tuple(x, y, z, 1, 2, 3)
- assert expr.op.name == "tirx.tvm_tuple"
-
-
def test_tir_op_tvm_struct_get():
x = tirx.Var("x", ty="handle")
expr = tirx.tvm_struct_get(x, 1, 2, dtype="int32")
@@ -56,17 +48,6 @@ def test_tir_op_address_of():
assert scalar_address.ty == tvm.ir.PointerType(tvm.ir.PrimType("uint32"))
-def test_tir_op_trace_pointer():
- pointer = tirx.Var("pointer",
tvm.ir.PointerType(tvm.ir.PrimType("float32")))
- traced = tirx.trace([pointer])
- assert traced.ty == pointer.ty
-
-
-def test_tir_op_lookup_param():
- expr = tirx.lookup_param("p0")
- assert expr.op.name == "tirx.lookup_param"
-
-
def test_tir_op_reinterpret():
x = tirx.Var("x", ty="int32")
expr = tirx.reinterpret("float32", x)
diff --git a/tests/python/tirx/script/test_tirx_script_basic_usage.py
b/tests/python/tirx/script/test_tirx_script_basic_usage.py
index 4945bab51a..dfc2a61d3a 100644
--- a/tests/python/tirx/script/test_tirx_script_basic_usage.py
+++ b/tests/python/tirx/script/test_tirx_script_basic_usage.py
@@ -515,22 +515,33 @@ def
test_module_string_constants_keep_common_constructor_values():
assert isinstance(Module.attrs["type"], ir.StringType)
+def test_thread_return_is_distinct_from_function_return():
+ @T.prim_func
+ def thread_exit():
+ T.thread_return()
+
+ @T.prim_func
+ def function_exit():
+ return 0
+
+ assert isinstance(thread_exit.body, tirx.Evaluate)
+ assert thread_exit.body.value.op.name == "tirx.thread_return"
+ assert isinstance(function_exit.body, tirx.Return)
+ assert isinstance(function_exit.body.value, tirx.IntImm)
+ assert function_exit.body.value.value == 0
+
+
def test_loop_control_validation_preserves_valid_and_unchecked_ir():
# Invalid loop placement must be rejected, while disabled checks preserve
the original IR.
- from tvm import ir, tirx
+ from tvm import error, ir, tirx
- @T.prim_func(check_well_formed=False)
- def invalid():
- T.evaluate(tirx.break_loop())
-
- # This is exactly the native node emitted by source `break`; constructing
- # the intrinsic explicitly keeps the surrounding Python definition valid.
- ir.assert_structural_equal(invalid.body, tirx.Evaluate(tirx.break_loop()))
- with pytest.raises(ValueError, match="requires an enclosing loop"):
+ invalid = tirx.PrimFunc(params=[], body=tirx.Break())
- @T.prim_func
- def rejected():
- T.evaluate(tirx.break_loop())
+ # Direct construction retains the native statement for an explicit
verifier pass.
+ ir.assert_structural_equal(invalid.body, tirx.Break())
+ assert not tirx.analysis.verify_well_formed(invalid, assert_mode=False)
+ with pytest.raises(error.InternalError, match="requires an enclosing
loop"):
+ tirx.analysis.verify_well_formed(invalid)
@T.prim_func
def valid():
