You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@tvm.apache.org by tq...@apache.org on 2026/10/14 15:19:18 UTC
(tvm) branch main updated: [REFACTOR] Localize lowering helpers and consolidate TIRx tests (#20719)
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 d149af9b57 [REFACTOR] Localize lowering helpers and consolidate TIRx tests (#20719)
d149af9b57 is described below
commit d149af9b57c91b4653c2c38ec3c717c5de27d0ea
Author: Tianqi Chen <tq...@users.noreply.github.com>
AuthorDate: Wed Oct 14 11:19:12 2026 -0400
[REFACTOR] Localize lowering helpers and consolidate TIRx tests (#20719)
Move temporary-allocation alignment into LLVM codegen and keep warp
pointer-scope remapping with its consumer. Remove the unused
RemoveAssume pass, consolidate TIRx tests under their owning packages,
and replace construction-heavy transform fixtures with script-defined
expected IR.
---
docs/contribute/pull_request.rst | 2 +-
python/tvm/tirx/transform/transform.py | 11 -
src/target/llvm/codegen_cpu.cc | 1 +
src/target/llvm/codegen_llvm.cc | 16 +
src/target/llvm/codegen_llvm.h | 4 +-
src/tirx/transform/ir_utils.h | 22 -
src/tirx/transform/lower_thread_allreduce.cc | 2 -
src/tirx/transform/lower_warp_memory.cc | 78 +-
src/tirx/transform/remove_assume.cc | 77 -
src/tirx/transform/update_pointer_storage_scope.cc | 99 -
src/tirx/transform/update_pointer_storage_scope.h | 54 -
tests/python/ir/test_datatype_nv_fp4.py | 56 -
tests/python/ir/test_datatype_nv_fp8.py | 128 --
tests/python/ir/test_dtype.py | 50 -
tests/python/ir/test_ir_attrs.py | 29 -
tests/python/ir/test_ir_container.py | 113 -
tests/python/ir/test_ir_json.py | 201 ++
tests/python/ir/test_ir_module.py | 45 +
tests/python/ir/test_ir_prim_expr_deep_equal.py | 13 +-
tests/python/ir/test_ir_type.py | 7 -
tests/python/ir/test_node_reflection.py | 351 ---
tests/python/ir/test_nonnullable_ir.py | 36 +-
tests/python/ir/test_pass_instrument.py | 64 -
tests/python/ir/test_unique_name_supply.py | 62 -
tests/python/relax/test_expr.py | 106 +
tests/python/relax/test_expr_functor.py | 13 +
tests/python/relax/test_pipeline.py | 11 +-
tests/python/runtime/test_runtime_tensor.py | 49 +
tests/python/sym/test_sym_analyzer_object.py | 34 +
tests/python/te/test_te_tensor.py | 23 +
.../tirx-base/test_tir_stmt_functor_substitute.py | 120 -
tests/python/tirx-base/test_tir_tensor.py | 307 ---
.../test_tir_transform_common_subexpr_elim.py | 805 -------
.../test_tir_transform_function_pass.py | 69 -
.../test_tir_transform_lower_intrin.py | 236 --
.../test_tir_transform_make_packed_api.py | 502 -----
.../test_tir_transform_narrow_datatype.py | 317 ---
.../test_tir_transform_remove_assume.py | 69 -
.../test_tir_transform_storage_rewrite.py | 755 -------
.../test_tir_transform_unroll_loop.py | 159 --
.../test_tirx_analysis_check_well_defined.py} | 0
.../analysis/test_tirx_analysis_undefined_vars.py} | 0
.../base/test_tirx_base.py} | 2 +-
.../base/test_tirx_binding_type_propagation.py} | 30 +
.../base/test_tirx_constructor.py} | 9 +
.../base/test_tirx_host_func.py} | 0
.../base/test_tirx_imm_values.py} | 0
.../base/test_tirx_index_map.py} | 0
.../base/test_tirx_intrin.py} | 0
.../base/test_tirx_nodes.py} | 6 +-
.../base/test_tirx_op_types.py} | 10 +
.../test_tir_ops.py => tirx/base/test_tirx_ops.py} | 1 -
.../base/test_tirx_ptx_cp_async.py} | 0
.../base/test_tirx_ptx_griddepcontrol.py} | 0
.../base/test_tirx_ptx_ldmatrix.py} | 0
.../base/test_tirx_ptx_mma.py} | 0
.../base/test_tirx_ptx_mma_sp.py} | 0
.../base/test_tirx_ptx_scalar_f32_math.py} | 0
.../base/test_tirx_scalable_datatype.py} | 0
.../base/test_tirx_specialize.py} | 0
tests/python/tirx/base/test_tirx_tensor_type.py | 198 ++
.../base/test_tirx_texture_scope.py} | 0
tests/python/tirx/codegen/test_tirx_codegen_fp8.py | 54 +
tests/python/{ => tirx}/test_tirx_axis_enum.py | 0
.../test_tirx_transform_bf16_legalize.py} | 147 +-
.../test_tirx_transform_common_subexpr_elim.py | 622 +++++
.../transform/test_tirx_transform_convert_ssa.py} | 391 ++--
...> test_tirx_transform_decorate_device_scope.py} | 23 +-
.../test_tirx_transform_flatten_tensor.py} | 260 ++-
...st_tirx_transform_force_narrow_index_to_i32.py} | 31 +-
.../transform/test_tirx_transform_fp8_legalize.py} | 9 +-
.../transform/test_tirx_transform_function_pass.py | 85 +
.../transform/test_tirx_transform_helpers.py} | 140 +-
...hoist_if.py => test_tirx_transform_hoist_if.py} | 0
.../transform/test_tirx_transform_lower_intrin.py | 236 ++
.../transform/test_tirx_transform_lower_tirx.py | 2370 ++++++++++++++++++++
.../test_tirx_transform_lower_tvm_builtin.py} | 183 +-
.../test_tirx_transform_make_packed_api.py | 1543 +++++++++++++
.../test_tirx_transform_naive_allocator.py | 279 +++
.../test_tirx_transform_narrow_datatype.py | 285 +++
...t_tirx_transform_pointer_value_type_rewrite.py} | 39 +-
.../transform/test_tirx_transform_remove_no_op.py} | 222 +-
...st_tirx_transform_renormalize_split_pattern.py} | 0
.../transform/test_tirx_transform_simplify.py} | 145 +-
.../test_tirx_transform_split_host_device.py} | 300 ++-
.../test_tirx_transform_storage_rewrite.py | 942 ++++++++
.../transform/test_tirx_transform_unroll_loop.py | 238 ++
.../transform/test_tirx_transform_vectorize.py} | 159 +-
.../transform/test_transform_flatten_tensor.py | 173 --
.../tirx/transform/test_transform_lower_tirx.py | 2107 -----------------
.../transform/test_transform_naive_allocator.py | 260 ---
.../test_transform_static_horizontal_fusion.py | 20 -
92 files changed, 8795 insertions(+), 7820 deletions(-)
diff --git a/docs/contribute/pull_request.rst b/docs/contribute/pull_request.rst
index 9150838e76..d349ce7d99 100644
--- a/docs/contribute/pull_request.rst
+++ b/docs/contribute/pull_request.rst
@@ -261,7 +261,7 @@ If you want to run a single test:
export PYTHONPATH=python
rm -rf python/tvm/*.pyc python/tvm/*/*.pyc python/tvm/*/*/*.pyc
- python -m pytest -v tests/python/tirx-transform/test_tir_transform_storage_rewrite.py
+ python -m pytest -v tests/python/tirx/transform/test_tirx_transform_storage_rewrite.py
# Additionally if you want to run a single test, for example test_add_one_2d inside a file.
python -m pytest -v -k "test_add_one_2d" tests/python/relax/test_frontend_tflite.py
diff --git a/python/tvm/tirx/transform/transform.py b/python/tvm/tirx/transform/transform.py
index 303e95fd7c..ffe4d1ba56 100644
--- a/python/tvm/tirx/transform/transform.py
+++ b/python/tvm/tirx/transform/transform.py
@@ -149,17 +149,6 @@ def HoistIf():
return _ffi_api.HoistIf()
-def RemoveAssume():
- """Remove all instances of prim::assume_op
-
- Returns
- -------
- fpass : tvm.ir.transform.Pass
- The result pass
- """
- return _ffi_api.RemoveAssume() # type: ignore
-
-
def BF16ComputeLegalize():
"""Legalize bf16 compute Ops.
diff --git a/src/target/llvm/codegen_cpu.cc b/src/target/llvm/codegen_cpu.cc
index d83aeefff0..3de6a46130 100644
--- a/src/target/llvm/codegen_cpu.cc
+++ b/src/target/llvm/codegen_cpu.cc
@@ -66,6 +66,7 @@
#include <unordered_set>
#include "../../backend/memory_scope.h"
+#include "../../tirx/transform/ir_utils.h"
#include "llvm_instance.h"
namespace tvm {
diff --git a/src/target/llvm/codegen_llvm.cc b/src/target/llvm/codegen_llvm.cc
index f4ec1611be..3634ff7dd1 100644
--- a/src/target/llvm/codegen_llvm.cc
+++ b/src/target/llvm/codegen_llvm.cc
@@ -109,6 +109,22 @@ namespace tvm {
namespace codegen {
using namespace tvm::prim;
+int CodeGenLLVM::GetTempAllocaAlignment(const PrimType& type, int64_t const_size) {
+ int align = runtime::kTempAllocaAlignment;
+ if (const_size > 0) {
+ int64_t element_bytes = type.StorageBytes();
+ // Only compute the total size when it can reduce the alignment. This also avoids
+ // overflowing for very large allocations.
+ if (element_bytes > 0 && const_size <= (align - 1) / element_bytes) {
+ int64_t const_s = const_size * element_bytes;
+ while (align > const_s) {
+ align = align / 2;
+ }
+ }
+ }
+ return align;
+}
+
namespace {
int GetLanesOrVScaleFactor(const PrimType& dtype) {
diff --git a/src/target/llvm/codegen_llvm.h b/src/target/llvm/codegen_llvm.h
index 91095dee8e..29cac98c50 100644
--- a/src/target/llvm/codegen_llvm.h
+++ b/src/target/llvm/codegen_llvm.h
@@ -62,7 +62,6 @@
#include <utility>
#include <vector>
-#include "../../tirx/transform/ir_utils.h"
#include "codegen_params.h"
#include "llvm_instance.h"
@@ -261,6 +260,9 @@ class CodeGenLLVM : public tvm::ExprFunctor<llvm::Value*(const Expr&)>,
llvm::GlobalValue::LinkageTypes linkage_type = llvm::GlobalValue::InternalLinkage);
protected:
+ /*! \brief Allocation alignment for constant-size temporary arrays. */
+ static int GetTempAllocaAlignment(const PrimType& type, int64_t const_size);
+
/*!
* \brief Address and type pair to assist in handling opaque pointers.
*/
diff --git a/src/tirx/transform/ir_utils.h b/src/tirx/transform/ir_utils.h
index 63acf59af9..7a6680f86e 100644
--- a/src/tirx/transform/ir_utils.h
+++ b/src/tirx/transform/ir_utils.h
@@ -60,28 +60,6 @@ inline PrimType APIType(const PrimType& t) {
return PrimType::Float(64);
}
-/*!
- * \brief Rule to get allocation alignment requirement for a given const array.
- * \param type The type of allocation.
- * \param const_size The constant size of the array.
- * \return the alignment
- */
-inline int GetTempAllocaAlignment(const PrimType& type, int64_t const_size) {
- int align = runtime::kTempAllocaAlignment;
- if (const_size > 0) {
- int64_t element_bytes = type.StorageBytes();
- // Only compute the total size when it can reduce the alignment. This also avoids
- // overflowing for very large allocations.
- if (element_bytes > 0 && const_size <= (align - 1) / element_bytes) {
- int64_t const_s = const_size * element_bytes;
- while (align > const_s) {
- align = align / 2;
- }
- }
- }
- return align;
-}
-
/*!
* \brief Convert a IR node to be SSA form.
* \param stmt The source statement to be converted.
diff --git a/src/tirx/transform/lower_thread_allreduce.cc b/src/tirx/transform/lower_thread_allreduce.cc
index 8d94ddc5ad..825fb5c816 100644
--- a/src/tirx/transform/lower_thread_allreduce.cc
+++ b/src/tirx/transform/lower_thread_allreduce.cc
@@ -36,8 +36,6 @@
#include <map>
#include <unordered_set>
-#include "update_pointer_storage_scope.h"
-
namespace tvm {
namespace tirx {
namespace detail {
diff --git a/src/tirx/transform/lower_warp_memory.cc b/src/tirx/transform/lower_warp_memory.cc
index 7f491daf30..cb495ee04c 100644
--- a/src/tirx/transform/lower_warp_memory.cc
+++ b/src/tirx/transform/lower_warp_memory.cc
@@ -44,16 +44,92 @@
#include <tvm/tirx/stmt_functor.h>
#include <tvm/tirx/transform.h>
+#include <unordered_map>
#include <unordered_set>
+#include <utility>
#include "../../sym/pattern_match.h"
#include "ir_utils.h"
-#include "update_pointer_storage_scope.h"
namespace tvm {
namespace tirx {
using namespace tvm::prim;
+namespace {
+
+class UpdatePointerStorageScope : public StmtExprMutator {
+ public:
+ using StmtExprMutator::Mutate_;
+
+ UnchangedOr<Stmt> Mutate_(const BindNode* op, InplaceMode inplace_mode) final;
+
+ UnchangedOr<Expr> Mutate_(const CallNode* op, InplaceMode inplace_mode) final;
+
+ explicit UpdatePointerStorageScope(
+ const std::unordered_map<Var, runtime::MemoryScope, ffi::ObjectPtrHash, ffi::ObjectPtrEqual>&
+ new_storage_scopes);
+
+ private:
+ std::unordered_map<const CallNode*, runtime::MemoryScope> tensor_scopes_;
+};
+
+Var WithStorageScope(const VarNode* buffer_var, runtime::MemoryScope storage_scope) {
+ auto* ptr_type = buffer_var->ty.as<PtrTypeNode>();
+ TVM_FFI_ICHECK(ptr_type) << "The provided variable is not of pointer type";
+ return Var(buffer_var->name, PtrType(ptr_type->element_type, storage_scope), buffer_var->loc);
+}
+
+UpdatePointerStorageScope::UpdatePointerStorageScope(
+ const std::unordered_map<Var, runtime::MemoryScope, ffi::ObjectPtrHash, ffi::ObjectPtrEqual>&
+ new_storage_scopes) {
+ for (auto& kv : new_storage_scopes) {
+ if (kv.first->ty.as<TensorTypeNode>()) {
+ TensorVar tensor = ffi::GetRef<Var>(kv.first.get()).as_or_throw<TensorVar>();
+ auto type = ffi::make_object<TensorTypeNode>(*tensor.ty().get());
+ type->scope = kv.second;
+ TensorVar replacement = TensorVar(tensor->name, TensorType(std::move(type)), tensor->loc);
+ VarRemapSet(kv.first, replacement);
+ } else {
+ VarRemapSet(kv.first, WithStorageScope(kv.first.get(), kv.second));
+ }
+ }
+}
+
+UnchangedOr<Stmt> UpdatePointerStorageScope::Mutate_(const BindNode* op, InplaceMode inplace_mode) {
+ const auto* call = op->value.as<CallNode>();
+ if (call &&
+ (call->op.same_as(tirx::alloc_tensor_op()) || call->op.same_as(tirx::decl_tensor_op()))) {
+ if (auto mapped = VarRemapGet(op->var); mapped != nullptr) {
+ tensor_scopes_.emplace(call, mapped.as_or_throw<TensorVar>().scope());
+ auto result = StmtExprMutator::Mutate_(op, inplace_mode);
+ tensor_scopes_.erase(call);
+ return result;
+ }
+ }
+ return StmtExprMutator::Mutate_(op, inplace_mode);
+}
+
+UnchangedOr<Expr> UpdatePointerStorageScope::Mutate_(const CallNode* op, InplaceMode inplace_mode) {
+ auto result = StmtExprMutator::Mutate_(op, inplace_mode);
+ if (auto it = tensor_scopes_.find(op); it != tensor_scopes_.end()) {
+ Expr value = std::move(result).ValueOrUnchanged(ffi::GetRef<Expr>(op));
+ auto call = value.as_or_throw<Call>();
+ auto type = call->ty_args[0].as_or_throw<TensorType>();
+ if (type->scope != it->second) {
+ auto updated = ffi::make_object<TensorTypeNode>(*type.get());
+ updated->scope = it->second;
+ auto copy = ffi::make_object<CallNode>(*call.get());
+ copy->ty_args.Set(0, TensorType(std::move(updated)));
+ return ReinferMutatedCallType(Expr(std::move(copy)), op, inplace_mode);
+ }
+ return value;
+ }
+ if (!op->op.same_as(tirx::tensor_data_ptr_op())) return result;
+ return ReinferMutatedCallType(std::move(result), op, inplace_mode);
+}
+
+} // namespace
+
// Rewrite Rule
//
// There is no special warp memory in most GPUs.
diff --git a/src/tirx/transform/remove_assume.cc b/src/tirx/transform/remove_assume.cc
deleted file mode 100644
index 8ce7f75254..0000000000
--- a/src/tirx/transform/remove_assume.cc
+++ /dev/null
@@ -1,77 +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 remove_store_undef.cc
- * \brief Remove stores of tirx::undef_op
- */
-#include <tvm/ffi/function.h>
-#include <tvm/ffi/reflection/registry.h>
-#include <tvm/ir/prim/op.h>
-#include <tvm/tirx/analysis.h>
-#include <tvm/tirx/op/memory.h>
-#include <tvm/tirx/stmt.h>
-#include <tvm/tirx/stmt_functor.h>
-#include <tvm/tirx/transform.h>
-
-namespace tvm {
-namespace tirx {
-
-// Remove any prim::assume_op calls
-class AssumeRemover : public StmtExprMutator {
- public:
- using StmtExprMutator::Mutate;
- using StmtExprMutator::Mutate_;
- using Parent = StmtExprMutator;
-
- UnchangedOr<Stmt> Mutate_(const EvaluateNode* op, InplaceMode inplace_mode) final {
- if (auto* call = op->value.as<CallNode>()) {
- if (call->op.same_as(prim::assume_op())) {
- return Evaluate(0);
- }
- }
- return StmtExprMutator::Mutate_(op, inplace_mode);
- }
-};
-
-namespace transform {
-Pass RemoveAssumeInternal() {
- auto pass_func = [](Function f, IRModule m, PassContext ctx) {
- auto* n = f.CopyOnWrite();
- n->body = ffi::make_object<AssumeRemover>()
- ->Mutate(n->body, InplaceMode::kAllow)
- .ValueOrUnchanged(n->body);
- return f;
- };
- return CreateFunctionPass(pass_func, 0, "tirx.RemoveAssumeInternal");
-}
-
-Pass RemoveAssume() {
- return Sequential({RemoveAssumeInternal(), RemoveNoOp()}, "tirx.RemoveAssume");
-}
-
-TVM_FFI_STATIC_INIT_BLOCK() {
- namespace refl = tvm::ffi::reflection;
- refl::GlobalDef().def("tirx.transform.RemoveAssume", RemoveAssume);
-}
-
-} // namespace transform
-
-} // namespace tirx
-} // namespace tvm
diff --git a/src/tirx/transform/update_pointer_storage_scope.cc b/src/tirx/transform/update_pointer_storage_scope.cc
deleted file mode 100644
index 170fae2ca8..0000000000
--- a/src/tirx/transform/update_pointer_storage_scope.cc
+++ /dev/null
@@ -1,99 +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 update_pointer_storage_scope.cc
- * \brief A pass to update storage scopes for tensor and pointer variables.
- */
-#include "update_pointer_storage_scope.h"
-
-#include <tvm/ffi/cast.h>
-#include <tvm/ffi/extra/structural_equal.h>
-#include <tvm/ir/prim/expr.h>
-#include <tvm/ir/prim/op.h>
-#include <tvm/runtime/memory_scope.h>
-#include <tvm/tirx/op/memory.h>
-#include <tvm/tirx/stmt_functor.h>
-#include <tvm/tirx/transform.h>
-
-#include <unordered_map>
-#include <utility>
-
-#include "ir_utils.h"
-
-namespace tvm {
-namespace tirx {
-
-Var WithStorageScope(const VarNode* buffer_var, runtime::MemoryScope storage_scope) {
- auto* ptr_type = buffer_var->ty.as<PtrTypeNode>();
- TVM_FFI_ICHECK(ptr_type) << "The provided variable is not of pointer type";
- return Var(buffer_var->name, PtrType(ptr_type->element_type, storage_scope), buffer_var->loc);
-}
-
-UpdatePointerStorageScope::UpdatePointerStorageScope(
- const std::unordered_map<Var, runtime::MemoryScope, ffi::ObjectPtrHash, ffi::ObjectPtrEqual>&
- new_storage_scopes) {
- for (auto& kv : new_storage_scopes) {
- if (kv.first->ty.as<TensorTypeNode>()) {
- TensorVar tensor = ffi::GetRef<Var>(kv.first.get()).as_or_throw<TensorVar>();
- auto type = ffi::make_object<TensorTypeNode>(*tensor.ty().get());
- type->scope = kv.second;
- TensorVar replacement = TensorVar(tensor->name, TensorType(std::move(type)), tensor->loc);
- VarRemapSet(kv.first, replacement);
- } else {
- VarRemapSet(kv.first, WithStorageScope(kv.first.get(), kv.second));
- }
- }
-}
-
-UnchangedOr<Stmt> UpdatePointerStorageScope::Mutate_(const BindNode* op, InplaceMode inplace_mode) {
- const auto* call = op->value.as<CallNode>();
- if (call &&
- (call->op.same_as(tirx::alloc_tensor_op()) || call->op.same_as(tirx::decl_tensor_op()))) {
- if (auto mapped = VarRemapGet(op->var); mapped != nullptr) {
- tensor_scopes_.emplace(call, mapped.as_or_throw<TensorVar>().scope());
- auto result = StmtExprMutator::Mutate_(op, inplace_mode);
- tensor_scopes_.erase(call);
- return result;
- }
- }
- return StmtExprMutator::Mutate_(op, inplace_mode);
-}
-
-UnchangedOr<Expr> UpdatePointerStorageScope::Mutate_(const CallNode* op, InplaceMode inplace_mode) {
- auto result = StmtExprMutator::Mutate_(op, inplace_mode);
- if (auto it = tensor_scopes_.find(op); it != tensor_scopes_.end()) {
- Expr value = std::move(result).ValueOrUnchanged(ffi::GetRef<Expr>(op));
- auto call = value.as_or_throw<Call>();
- auto type = call->ty_args[0].as_or_throw<TensorType>();
- if (type->scope != it->second) {
- auto updated = ffi::make_object<TensorTypeNode>(*type.get());
- updated->scope = it->second;
- auto copy = ffi::make_object<CallNode>(*call.get());
- copy->ty_args.Set(0, TensorType(std::move(updated)));
- return ReinferMutatedCallType(Expr(std::move(copy)), op, inplace_mode);
- }
- return value;
- }
- if (!op->op.same_as(tirx::tensor_data_ptr_op())) return result;
- return ReinferMutatedCallType(std::move(result), op, inplace_mode);
-}
-
-} // namespace tirx
-} // namespace tvm
diff --git a/src/tirx/transform/update_pointer_storage_scope.h b/src/tirx/transform/update_pointer_storage_scope.h
deleted file mode 100644
index ede20534e0..0000000000
--- a/src/tirx/transform/update_pointer_storage_scope.h
+++ /dev/null
@@ -1,54 +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 update_pointer_storage_scope.h
- * \brief A pass to update storage scopes for tensor and pointer variables.
- */
-#ifndef TVM_TIR_TRANSFORM_UPDATE_POINTER_STORAGE_SCOPE_H_
-#define TVM_TIR_TRANSFORM_UPDATE_POINTER_STORAGE_SCOPE_H_
-
-#include <tvm/ir/prim/expr.h>
-#include <tvm/ir/prim/op.h>
-#include <tvm/tirx/stmt_functor.h>
-
-#include <unordered_map>
-
-namespace tvm {
-namespace tirx {
-
-class UpdatePointerStorageScope : public StmtExprMutator {
- public:
- using StmtExprMutator::Mutate_;
-
- UnchangedOr<Stmt> Mutate_(const BindNode* op, InplaceMode inplace_mode) final;
-
- UnchangedOr<Expr> Mutate_(const CallNode* op, InplaceMode inplace_mode) final;
-
- explicit UpdatePointerStorageScope(
- const std::unordered_map<Var, runtime::MemoryScope, ffi::ObjectPtrHash, ffi::ObjectPtrEqual>&
- new_storage_scopes);
-
- private:
- std::unordered_map<const CallNode*, runtime::MemoryScope> tensor_scopes_;
-};
-
-} // namespace tirx
-} // namespace tvm
-#endif // TVM_TIR_TRANSFORM_UPDATE_POINTER_STORAGE_SCOPE_H_
diff --git a/tests/python/ir/test_datatype_nv_fp4.py b/tests/python/ir/test_datatype_nv_fp4.py
deleted file mode 100644
index af6bd72035..0000000000
--- a/tests/python/ir/test_datatype_nv_fp4.py
+++ /dev/null
@@ -1,56 +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: F401
-import numpy as np
-import pytest
-
-import tvm
-import tvm.testing
-import tvm.tirx as tirx
-from tvm import te
-from tvm.script import tirx as T
-
-try:
- from ml_dtypes import float4_e2m1fn
-except ImportError:
- float4_e2m1fn = None
-
-
-nv_fp4_dtypes = [(float4_e2m1fn, "float4_e2m1fn")]
-
-
-@pytest.mark.parametrize("np_dtype,dtype_str", nv_fp4_dtypes)
-def test_create_nv_fp4_nd_array(np_dtype, dtype_str):
- if np_dtype is None:
- """Skip test if ml_dtypes is not installed"""
- return
- x = np.random.rand(128, 128).astype(np_dtype)
- x_nd = tvm.runtime.tensor(x)
- assert x_nd.dtype == dtype_str
- np.testing.assert_equal(x_nd.numpy(), x)
-
-
-@pytest.mark.parametrize("np_dtype,dtype_str", nv_fp4_dtypes)
-def test_nv_fp4_buffer(np_dtype, dtype_str):
- m = te.var("m")
- n = te.var("n")
- A = tvm.tirx.decl_tensor((m, n), dtype_str)
- assert A.dtype == dtype_str
-
-
-if __name__ == "__main__":
- tvm.testing.main()
diff --git a/tests/python/ir/test_datatype_nv_fp8.py b/tests/python/ir/test_datatype_nv_fp8.py
deleted file mode 100644
index eea5701cfa..0000000000
--- a/tests/python/ir/test_datatype_nv_fp8.py
+++ /dev/null
@@ -1,128 +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: E501, F401
-import numpy as np
-import pytest
-
-import tvm
-import tvm.testing
-import tvm.tirx as tirx
-from tvm import te
-from tvm.script import s_tir as Ts
-from tvm.script import tirx as T
-
-try:
- from ml_dtypes import (
- float8_e3m4,
- float8_e4m3,
- float8_e4m3b11fnuz,
- float8_e4m3fn,
- float8_e4m3fnuz,
- float8_e5m2,
- float8_e5m2fnuz,
- float8_e8m0fnu,
- )
-except ImportError:
- float8_e3m4 = float8_e4m3 = float8_e4m3b11fnuz = float8_e4m3fn = None
- float8_e4m3fnuz = float8_e5m2 = float8_e5m2fnuz = float8_e8m0fnu = None
-
-
-def fp8_unary(dtype: str):
- @Ts.function
- def func(
- A: T.Tensor([128], dtype=dtype),
- B: T.Tensor([128], dtype=dtype),
- A_add_B: T.Tensor([128], dtype=dtype),
- A_sub_B: T.Tensor([128], dtype=dtype),
- A_mul_B: T.Tensor([128], dtype=dtype),
- A_fp32: T.Tensor([128], dtype="float32"),
- A_roundtrip: T.Tensor([128], dtype=dtype),
- ) -> None:
- for i in range(128):
- with Ts.sblock("fp8_unary"):
- vi = Ts.axis.spatial(128, i)
- A_add_B[vi] = A[vi] + B[vi]
- A_sub_B[vi] = A[vi] - B[vi]
- A_mul_B[vi] = A[vi] * B[vi]
- A_fp32[vi] = A[vi]
- A_roundtrip[vi] = A_fp32[vi]
-
- return func
-
-
-nv_fp8_dtypes = [
- (float8_e3m4, "float8_e3m4"),
- (float8_e4m3, "float8_e4m3"),
- (float8_e4m3b11fnuz, "float8_e4m3b11fnuz"),
- (float8_e4m3fn, "float8_e4m3fn"),
- (float8_e4m3fnuz, "float8_e4m3fnuz"),
- (float8_e5m2, "float8_e5m2"),
- (float8_e5m2fnuz, "float8_e5m2fnuz"),
- (float8_e8m0fnu, "float8_e8m0fnu"),
-]
-
-
-@pytest.mark.parametrize("np_dtype,dtype_str", nv_fp8_dtypes)
-def test_create_nv_fp8_nd_array(np_dtype, dtype_str):
- if np_dtype is None:
- """Skip test if ml_dtypes is not installed"""
- return
- x = np.random.rand(128, 128).astype(np_dtype)
- x_nd = tvm.runtime.tensor(x)
- assert x_nd.dtype == dtype_str
- np.testing.assert_equal(x_nd.numpy(), x)
-
-
-@pytest.mark.parametrize("np_dtype,dtype_str", nv_fp8_dtypes)
-def test_fp8_unary_op(np_dtype, dtype_str):
- func = fp8_unary(dtype_str)
- if not tvm.testing.device_enabled("llvm"):
- return
- if np_dtype is None:
- """Skip test if ml_dtypes is not installed"""
- return
- if dtype_str in ["float8_e8m0fnu", "float8_e4m3b11fnuz", "float8_e4m3fnuz", "float8_e5m2fnuz"]:
- # float8_e8m0fnu does not support arithmetic operations, and unsigned arithmetic is not tested here
- return
-
- f = tvm.compile(func, target="llvm")
- a = np.random.randn(128).astype(np_dtype)
- b = np.random.randn(128).astype(np_dtype)
- a_add_b = np.zeros(128).astype(np_dtype)
- a_sub_b = np.zeros(128).astype(np_dtype)
- a_mul_b = np.zeros(128).astype(np_dtype)
- a_fp32 = np.zeros(128).astype(np.float32)
- a_roundtrip = np.zeros(128).astype(np_dtype)
- args = list(
- map(lambda _: tvm.runtime.tensor(_), [a, b, a_add_b, a_sub_b, a_mul_b, a_fp32, a_roundtrip])
- )
- f(*args)
- expected_a_fp32 = a.astype(np.float32)
- expected_a_roundtrip = expected_a_fp32.astype(np_dtype)
- np.testing.assert_equal(args[6].numpy(), expected_a_roundtrip)
-
-
-@pytest.mark.parametrize("np_dtype,dtype_str", nv_fp8_dtypes)
-def test_nv_fp8_buffer(np_dtype, dtype_str):
- m = te.var("m")
- n = te.var("n")
- A = tvm.tirx.decl_tensor((m, n), dtype_str)
- assert A.dtype == dtype_str
-
-
-if __name__ == "__main__":
- tvm.testing.main()
diff --git a/tests/python/ir/test_dtype.py b/tests/python/ir/test_dtype.py
deleted file mode 100644
index 158a9373de..0000000000
--- a/tests/python/ir/test_dtype.py
+++ /dev/null
@@ -1,50 +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: F841
-"""Test data type related API"""
-
-import pytest
-
-import tvm
-import tvm.testing
-from tvm import DataType
-
-
-@pytest.mark.parametrize(
- "dtype_str, expected_size",
- [
- ("float32", 4),
- ("float32x4", 16),
- ("float8_e5m2x4", 4),
- ("float6_e2m3fnx4", 3),
- ("float4_e2m1fnx4", 2),
- ("uint8", 1),
- ],
-)
-def test_dtype_itemsize(dtype_str, expected_size):
- dtype = DataType(dtype_str)
- assert dtype.itemsize == expected_size
-
-
-@pytest.mark.parametrize("dtype_str", ["int32xvscalex4"])
-def test_dtype_itemmize_error(dtype_str):
- with pytest.raises(ValueError):
- size = DataType(dtype_str).itemsize
-
-
-if __name__ == "__main__":
- tvm.testing.main()
diff --git a/tests/python/ir/test_ir_attrs.py b/tests/python/ir/test_ir_attrs.py
index b89ce7a33a..6c55e81027 100644
--- a/tests/python/ir/test_ir_attrs.py
+++ b/tests/python/ir/test_ir_attrs.py
@@ -100,25 +100,6 @@ def test_op_set_signature_metadata_and_counts():
tvm.ir.Call(op, [x]).validate()
-def test_op_set_signature_preserves_custom_validator():
- op = tvm.ir.Op.get("relax.call_pure_packed")
- original_args = [(info.name, info.doc) for info in op.args_info]
- original_ty_args = [(info.name, info.doc) for info in op.ty_args_info]
- original_var_args = (op.var_args_info.name, op.var_args_info.doc)
- original_var_ty_args = (op.var_ty_args_info.name, op.var_ty_args_info.doc)
- try:
- op.set_signature()
- with pytest.raises(TypeError, match="call_pure_packed expects a function argument"):
- tvm.ir.Call(op, []).validate()
- finally:
- op.set_signature(
- original_args,
- ty_args=original_ty_args,
- var_args=original_var_args,
- var_ty_args=original_var_ty_args,
- )
-
-
def test_op_set_signature_checks_counts_without_types_or_attrs():
name = "test.python_set_signature.count_only"
tvm.ir._ffi_api.RegisterOp(name, "Count-only Python signature")
@@ -134,16 +115,6 @@ def test_op_set_signature_checks_counts_without_types_or_attrs():
assert call.attrs is not None
-def test_ptx_ops_have_variadic_call_signature():
- from tvm.backend.cuda.ptx import TABLE
-
- op = tvm.ir.Op.get(next(iter(TABLE.values())).op_name)
- assert not op.args_info
- assert op.var_args_info.name == "operands"
- tvm.ir.Call(op, []).validate()
- tvm.ir.Call(op, [tvm.ir.Var("x", tvm.ir.AnyType())]).validate()
-
-
if __name__ == "__main__":
test_dict_attrs()
test_attrs_equal()
diff --git a/tests/python/ir/test_ir_container.py b/tests/python/ir/test_ir_container.py
deleted file mode 100644
index 3180e980ed..0000000000
--- a/tests/python/ir/test_ir_container.py
+++ /dev/null
@@ -1,113 +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: E712, F401
-import numpy as np
-import pytest
-import tvm_ffi
-
-import tvm
-
-
-def test_array():
- a = tvm.runtime.convert([1, 2, 3])
- assert len(a) == 3
- assert a[-1] == 3
- a_slice = a[-3:-1]
- assert (a_slice[0], a_slice[1]) == (1, 2)
-
-
-def test_array_save_load_json():
- a = tvm.runtime.convert([1, 2, 3.5, True])
- json_str = tvm.ir.save_json(a)
- a_loaded = tvm.ir.load_json(json_str)
- assert a_loaded[1] == 2
- assert a_loaded[2] == 3.5
- assert a_loaded[3] == True
- assert isinstance(a_loaded[3], bool)
-
-
-def test_dir_array():
- a = tvm.runtime.convert([1, 2, 3])
- assert dir(a)
-
-
-def test_map():
- a = tvm.tirx.Var("a", "int32")
- b = tvm.tirx.Var("b", "int32")
- amap = tvm.runtime.convert({a: 2, b: 3})
- assert a in amap
- assert len(amap) == 2
- dd = dict(amap.items())
- assert a in dd
- assert b in dd
- assert a + 1 not in amap
- assert {x for x in amap} == {a, b}
- assert set(amap.keys()) == {a, b}
- assert set(amap.values()) == {2, 3}
-
-
-def test_str_map():
- amap = tvm.runtime.convert({"a": 2, "b": 3})
- assert "a" in amap
- assert len(amap) == 2
- dd = dict(amap.items())
- assert amap["a"] == 2
- assert "a" in dd
- assert "b" in dd
-
-
-def test_map_save_load_json():
- a = tvm.tirx.Var("a", "int32")
- b = tvm.tirx.Var("b", "int32")
- amap = tvm.runtime.convert({a: 2, b: 3})
- json_str = tvm.ir.save_json(amap)
- amap = tvm.ir.load_json(json_str)
- assert len(amap) == 2
- dd = {kv[0].name: kv[1] for kv in amap.items()}
- assert dd == {"a": 2, "b": 3}
-
-
-def test_dir_map():
- a = tvm.tirx.Var("a", "int32")
- b = tvm.tirx.Var("b", "int32")
- amap = tvm.runtime.convert({a: 2, b: 3})
- assert dir(amap)
-
-
-def test_getattr_map():
- a = tvm.tirx.Var("a", "int32")
- b = tvm.tirx.Var("b", "int32")
- amap = tvm.runtime.convert({a: 2, b: 3})
- assert isinstance(amap, tvm_ffi.Map)
-
-
-def test_in_container():
- arr = tvm.runtime.convert(["a", "b", "c"])
- assert "a" in arr
- assert "d" not in arr
-
-
-def test_tensor_container():
- x = tvm.runtime.tensor([1, 2, 3])
- arr = tvm.runtime.convert([x, x])
- assert arr[0].same_as(x)
- assert arr[1].same_as(x)
- assert isinstance(arr[0], tvm.runtime.Tensor)
-
-
-if __name__ == "__main__":
- tvm.testing.main()
diff --git a/tests/python/ir/test_ir_json.py b/tests/python/ir/test_ir_json.py
new file mode 100644
index 0000000000..28956b2dfb
--- /dev/null
+++ b/tests/python/ir/test_ir_json.py
@@ -0,0 +1,201 @@
+# 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.
+import copy
+import json
+
+import pytest
+import tvm_ffi
+
+import tvm
+
+_LEGACY_RELAX_VAR_JSON = """{
+ "root_index": 9,
+ "nodes": [
+ {"type": "ir.SourceName", "data": "legacy_relax.py"},
+ {"type": "ir.SourceLoc", "data": {"source_name": 0, "start_line": 3, "start_column": 3,
+ "end_line": 5, "end_column": 11}},
+ {"type": "ir.UnknownLoc", "data": null},
+ {"type": "ir.PrimType", "data": {"loc": 2, "dtype": "int64"}},
+ {"type": "ffi.Array", "data": [3, 3]},
+ {"type": "ir.TupleType", "data": {"loc": 2, "fields": 4}},
+ {"type": "ffi.String", "data": "legacy"},
+ {"type": "relax.expr.Var", "data": {"loc": 1, "ty": 3, "name_hint": 6}},
+ {"type": "ffi.Array", "data": [7, 7]},
+ {"type": "ir.Tuple", "data": {"loc": 1, "ty": 5, "fields": 8}}
+ ],
+ "metadata": {"tvm_version": "0.26.dev0"}
+}"""
+
+_LEGACY_TIRX_VAR_JSON = """{
+ "root_index": 6,
+ "nodes": [
+ {"type": "ir.SourceName", "data": "legacy_tirx.py"},
+ {"type": "ir.SourceLoc", "data": {"source_name": 0, "start_line": 7, "start_column": 2,
+ "end_line": 9, "end_column": 14}},
+ {"type": "ir.UnknownLoc", "data": null},
+ {"type": "ir.PrimType", "data": {"loc": 2, "dtype": "int64"}},
+ {"type": "ffi.String", "data": "legacy"},
+ {"type": "tirx.Var", "data": {"loc": 1, "ty": 3, "name": 4}},
+ {"type": "tirx.Add", "data": {"loc": 1, "ty": 3, "a": 5, "b": 5}}
+ ],
+ "metadata": {"tvm_version": "0.26.dev0"}
+}"""
+
+
+def _make_pre_name_field_json(type_key):
+ graph = json.loads(_LEGACY_TIRX_VAR_JSON)
+ graph["root_index"] = 5
+ var = graph["nodes"][5]
+ var["type"] = type_key
+ var["data"]["name_hint"] = var["data"].pop("name")
+ return json.dumps(graph)
+
+
+_PRE_NAME_FIELD_IR_VAR_JSON = _make_pre_name_field_json("ir.Var")
+_PRE_NAME_FIELD_DATAFLOW_VAR_JSON = _make_pre_name_field_json("relax.expr.DataflowVar")
+
+
+@pytest.mark.parametrize(
+ ("legacy_json", "expected_type", "var_index"),
+ [
+ (_LEGACY_RELAX_VAR_JSON, "ir.Var", 7),
+ (_LEGACY_TIRX_VAR_JSON, "ir.Var", 5),
+ (_PRE_NAME_FIELD_IR_VAR_JSON, "ir.Var", 5),
+ (_PRE_NAME_FIELD_DATAFLOW_VAR_JSON, "relax.expr.DataflowVar", 5),
+ ],
+)
+def test_var_name_legacy_json_graph_rewrite(legacy_json, expected_type, var_index):
+ from tvm.ir.json_compact import upgrade_json
+
+ original = json.loads(legacy_json)
+ expected = copy.deepcopy(original)
+ for node in expected["nodes"]:
+ if node["type"] == "tirx.Add":
+ node["type"] = "ir.prim.Add"
+ expected["nodes"][var_index]["type"] = expected_type
+ fields = expected["nodes"][var_index]["data"]
+ if "name_hint" in fields:
+ fields["name"] = fields.pop("name_hint")
+
+ upgraded = json.loads(upgrade_json(legacy_json))
+ assert upgraded == expected
+ assert upgraded["root_index"] == original["root_index"]
+ assert len(upgraded["nodes"]) == len(original["nodes"])
+
+
+def _check_legacy_var(
+ var, source_name, line, end_line, column, end_column, expected_type=tvm.ir.Var
+):
+ assert type(var) is expected_type
+ assert var.name == "legacy"
+ assert var.ty == tvm.ir.PrimType("int64")
+ assert var.loc.source_name.name == source_name
+ assert var.loc.start_line == line
+ assert var.loc.end_line == end_line
+ assert var.loc.start_column == column
+ assert var.loc.end_column == end_column
+
+
+def test_var_exact_base_legacy_tirx_json_load():
+ restored = tvm.ir.load_json(_LEGACY_TIRX_VAR_JSON.replace("tirx.Add", "prim.Add"))
+ assert isinstance(restored, tvm.ir.prim.Add)
+ assert restored.a.same_as(restored.b)
+ _check_legacy_var(restored.a, "legacy_tirx.py", 7, 9, 2, 14)
+ assert restored.loc.same_as(restored.a.loc)
+ assert {node["type"] for node in json.loads(tvm.ir.save_json(restored))["nodes"]}.isdisjoint(
+ {"relax.expr.Var", "tirx.Var", "tirx.Add"}
+ )
+
+
+def test_var_exact_base_legacy_relax_json_load():
+ restored = tvm.ir.load_json(_LEGACY_RELAX_VAR_JSON)
+ assert isinstance(restored, tvm.ir.Tuple)
+ assert restored.fields[0].same_as(restored.fields[1])
+ _check_legacy_var(restored.fields[0], "legacy_relax.py", 3, 5, 3, 11)
+ assert restored.loc.same_as(restored.fields[0].loc)
+ assert {node["type"] for node in json.loads(tvm.ir.save_json(restored))["nodes"]}.isdisjoint(
+ {"relax.expr.Var", "tirx.Var"}
+ )
+
+
+@pytest.mark.parametrize(
+ ("legacy_json", "expected_type"),
+ [
+ (_PRE_NAME_FIELD_IR_VAR_JSON, tvm.ir.Var),
+ ],
+)
+def test_var_name_legacy_json_load(legacy_json, expected_type):
+ restored = tvm.ir.load_json(legacy_json)
+ _check_legacy_var(restored, "legacy_tirx.py", 7, 9, 2, 14, expected_type)
+ graph = json.loads(tvm.ir.save_json(restored))
+ fields = graph["nodes"][graph["root_index"]]["data"]
+ assert "name" in fields
+ assert "name_hint" not in fields
+
+
+@pytest.mark.parametrize(
+ "value",
+ [
+ tvm.ir.prim.IntImm("int32", 1) + tvm.ir.prim.IntImm("int32", 10),
+ *[
+ tvm.ir.prim.FloatImm(dtype, value)
+ for dtype in ("float32", "float64")
+ for value in (float("inf"), float("-inf"))
+ ],
+ tvm.ir.prim.min_value("float32"),
+ tvm.ir.prim.max_value("float32"),
+ tvm_ffi.core.String("xyz"),
+ tvm_ffi.core.String("xy\x01z"),
+ ],
+)
+def test_json_primitive_roundtrip(value):
+ tvm.ir.assert_structural_equal(value, tvm.ir.load_json(tvm.ir.save_json(value)))
+
+
+def test_json_containers():
+ values = tvm.ir.load_json(tvm.ir.save_json([1, 2, 3.5, True]))
+ assert list(values) == [1, 2, 3.5, True]
+ assert type(values[3]) is bool
+ a = tvm.ir.Var("a", tvm.ir.PrimType("int32"))
+ b = tvm.ir.Var("b", tvm.ir.PrimType("int32"))
+ objects = tvm.runtime.convert({a: 2, b: 3})
+ assert a in objects and b in dict(objects.items())
+ assert len(objects) == 2
+ assert tvm.ir.prim.Add(a, tvm.ir.prim.IntImm("int32", 1)) not in objects
+ assert set(objects) == set(objects.keys()) == {a, b}
+ assert set(objects.values()) == {2, 3}
+ restored = tvm.ir.load_json(tvm.ir.save_json(objects))
+ assert {key.name: value for key, value in restored.items()} == {"a": 2, "b": 3}
+ x = tvm.ir.prim.IntImm("int32", 1)
+ z = tvm.ir.prim.Add(x, tvm.ir.prim.IntImm("int32", 10))
+ graph = [{"z": z, "x": x}]
+ saved = tvm.ir.save_json(graph)
+ restored = tvm.ir.load_json(saved)
+ assert len(restored) == 1
+ assert restored[0]["z"].a.same_as(restored[0]["x"])
+ tvm.ir.assert_structural_equal(restored, graph)
+ assert json.loads(saved)["metadata"]["tvm_version"] == tvm.__version__
+
+
+@pytest.mark.parametrize("dtype", ["int32", tvm_ffi.core.String("int32")])
+def test_make_node_default_location(dtype):
+ value = tvm.ir.make_node("ir.IntImm", ty=tvm.ir.PrimType(dtype), value=10)
+ assert isinstance(value, tvm.ir.prim.IntImm)
+ assert value.value == 10
+ assert set(dir(value.__class__)) <= set(dir(value))
+ assert value.loc.same_as(tvm.ir.UnknownLoc())
+ assert tvm.ir.load_json(tvm.ir.save_json(value)).loc.same_as(tvm.ir.UnknownLoc())
diff --git a/tests/python/ir/test_ir_module.py b/tests/python/ir/test_ir_module.py
new file mode 100644
index 0000000000..f2c98de810
--- /dev/null
+++ b/tests/python/ir/test_ir_module.py
@@ -0,0 +1,45 @@
+# 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.
+import pytest
+
+import tvm
+from tvm.ir.supply import UniqueNameSupply
+from tvm.ir.transform.instrument import PrintAfterAll, PrintBeforeAll
+
+
+@pytest.mark.parametrize(
+ "prefix,add_prefix,expected", [("", False, "v"), ("prefix", True, "prefix_v")]
+)
+def test_fresh_name_empty_string(prefix, add_prefix, expected):
+ names = UniqueNameSupply(prefix)
+ assert names.fresh_name("", add_prefix=add_prefix) == expected
+ assert names.fresh_name("", add_prefix=add_prefix) == expected + "_1"
+
+
+def test_module_pass_print_instruments(capsys):
+ @tvm.ir.transform.module_pass(opt_level=0, name="IdentityModule")
+ def identity(mod, ctx):
+ return mod
+
+ mod = tvm.IRModule()
+ with tvm.ir.transform.PassContext(instruments=[PrintBeforeAll(), PrintAfterAll()]):
+ result = identity(mod)
+ assert result.same_as(mod)
+ output = capsys.readouterr().out
+ assert "Before Running Pass:" in output
+ assert "After Running Pass:" in output
+ assert 'name="IdentityModule"' in output
diff --git a/tests/python/ir/test_ir_prim_expr_deep_equal.py b/tests/python/ir/test_ir_prim_expr_deep_equal.py
index d7fc2e04d7..b25077a65c 100644
--- a/tests/python/ir/test_ir_prim_expr_deep_equal.py
+++ b/tests/python/ir/test_ir_prim_expr_deep_equal.py
@@ -18,19 +18,26 @@ import tvm
def test_equal_expr():
- x = tvm.tirx.Var("x", "int32")
- y = tvm.tirx.Var("y", "int32")
+ x = tvm.ir.Var("x", tvm.ir.PrimType("int32"))
+ y = tvm.ir.Var("y", tvm.ir.PrimType("int32"))
def func1():
return x + y + 1
def func2():
- return tvm.tirx.exp(tvm.tirx.truncdiv((x + y + 1) * y, 4))
+ return tvm.ir.prim.truncdiv((x + y + 1) * y, 4)
assert tvm.ir.prim.expr_deep_equal(func1(), func1())
assert tvm.ir.prim.expr_deep_equal(func2(), func2())
assert not tvm.ir.prim.expr_deep_equal(func2(), func1())
+def test_structural_equal_maps_free_variables():
+ x = tvm.ir.Var("x", tvm.ir.PrimType("int32"))
+ y = tvm.ir.Var("y", tvm.ir.PrimType("int32"))
+ tvm.ir.assert_structural_equal(x + 1, y + 1, map_free_vars=True)
+ assert not tvm.ir.prim.expr_deep_equal(x + 1, y + 1)
+
+
if __name__ == "__main__":
test_equal_expr()
diff --git a/tests/python/ir/test_ir_type.py b/tests/python/ir/test_ir_type.py
index b0012b2a10..7bb2a7f4f2 100644
--- a/tests/python/ir/test_ir_type.py
+++ b/tests/python/ir/test_ir_type.py
@@ -14,7 +14,6 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-# ruff: noqa: F401, F821
"""Test type nodes in the IR"""
import uuid
@@ -23,7 +22,6 @@ import pytest
import tvm
from tvm.runtime import MemoryScope
-from tvm.script import tirx as T
def check_json_roundtrip(node):
@@ -76,13 +74,8 @@ def test_memory_scope_json_roundtrip():
custom = MemoryScope.register("test.ir.scope." + uuid.uuid4().hex)
for scope in (MemoryScope("global"), MemoryScope("shared.dyn"), custom):
assert tvm.ir.load_json(tvm.ir.save_json(scope)).same_as(scope)
- ty = tvm.tirx.TensorType((4,), "float32", scope=scope)
- restored = tvm.ir.load_json(tvm.ir.save_json(ty))
- assert restored.scope.same_as(scope)
- tvm.ir.assert_structural_equal(ty, restored)
if __name__ == "__main__":
- test_tensor_type_bad_constructor()
test_func_type()
test_tuple_type()
diff --git a/tests/python/ir/test_node_reflection.py b/tests/python/ir/test_node_reflection.py
deleted file mode 100644
index 9d63b228d3..0000000000
--- a/tests/python/ir/test_node_reflection.py
+++ /dev/null
@@ -1,351 +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: E712, F401, F841
-import copy
-import json
-import sys
-
-import numpy as np
-import pytest
-import tvm_ffi
-
-import tvm
-import tvm.testing
-from tvm import te
-
-
-def test_const_saveload_json():
- # save load json
- x = tvm.tirx.const(1, "int32")
- y = tvm.tirx.const(10, "int32")
- z = x + y
- z = z + z
- json_str = tvm.ir.save_json(z)
- zz = tvm.ir.load_json(json_str)
- tvm.ir.assert_structural_equal(zz, z, map_free_vars=True)
-
-
-def test_symbolic_analysis_legacy_json_load():
- # Historical symbolic-analysis objects must load through the existing JSON upgrader.
- graph = {
- "root_index": 2,
- "nodes": [
- {"type": "arith.ConstIntBound", "data": {"min_value": 2, "max_value": 7}},
- {"type": "arith.ModularSet", "data": {"coeff": 4, "base": 1}},
- {"type": "ffi.Array", "data": [0, 1, 0]},
- ],
- "metadata": {"tvm_version": tvm.__version__},
- }
- restored = tvm.ir.load_json(json.dumps(graph))
- bound, modular, shared_bound = restored
- assert isinstance(bound, tvm.sym.ConstIntBound)
- assert (bound.min_value, bound.max_value) == (2, 7)
- assert isinstance(modular, tvm.sym.ModularSet)
- assert (modular.coeff, modular.base) == (4, 1)
- assert bound.same_as(shared_bound)
- saved_types = {node["type"] for node in json.loads(tvm.ir.save_json(restored))["nodes"]}
- assert {"sym.ConstIntBound", "sym.ModularSet"} <= saved_types
- assert not any(key.startswith("arith.") for key in saved_types)
-
-
-def test_save_json_metadata_version():
- obj = tvm.runtime.convert([1, 2])
- json_str = tvm.ir.save_json(obj)
- assert json.loads(json_str)["metadata"]["tvm_version"] == tvm.__version__
- assert list(tvm.ir.load_json(json_str)) == [1, 2]
-
-
-_LEGACY_RELAX_VAR_JSON = """{
- "root_index": 9,
- "nodes": [
- {"type": "ir.SourceName", "data": "legacy_relax.py"},
- {"type": "ir.SourceLoc", "data": {"source_name": 0, "start_line": 3, "start_column": 3,
- "end_line": 5, "end_column": 11}},
- {"type": "ir.UnknownLoc", "data": null},
- {"type": "ir.PrimType", "data": {"loc": 2, "dtype": "int64"}},
- {"type": "ffi.Array", "data": [3, 3]},
- {"type": "ir.TupleType", "data": {"loc": 2, "fields": 4}},
- {"type": "ffi.String", "data": "legacy"},
- {"type": "relax.expr.Var", "data": {"loc": 1, "ty": 3, "name_hint": 6}},
- {"type": "ffi.Array", "data": [7, 7]},
- {"type": "ir.Tuple", "data": {"loc": 1, "ty": 5, "fields": 8}}
- ],
- "metadata": {"tvm_version": "0.26.dev0"}
-}"""
-
-_LEGACY_TIRX_VAR_JSON = """{
- "root_index": 6,
- "nodes": [
- {"type": "ir.SourceName", "data": "legacy_tirx.py"},
- {"type": "ir.SourceLoc", "data": {"source_name": 0, "start_line": 7, "start_column": 2,
- "end_line": 9, "end_column": 14}},
- {"type": "ir.UnknownLoc", "data": null},
- {"type": "ir.PrimType", "data": {"loc": 2, "dtype": "int64"}},
- {"type": "ffi.String", "data": "legacy"},
- {"type": "tirx.Var", "data": {"loc": 1, "ty": 3, "name": 4}},
- {"type": "tirx.Add", "data": {"loc": 1, "ty": 3, "a": 5, "b": 5}}
- ],
- "metadata": {"tvm_version": "0.26.dev0"}
-}"""
-
-
-def _make_pre_name_field_json(type_key):
- graph = json.loads(_LEGACY_TIRX_VAR_JSON)
- graph["root_index"] = 5
- var = graph["nodes"][5]
- var["type"] = type_key
- var["data"]["name_hint"] = var["data"].pop("name")
- return json.dumps(graph)
-
-
-_PRE_NAME_FIELD_IR_VAR_JSON = _make_pre_name_field_json("ir.Var")
-_PRE_NAME_FIELD_DATAFLOW_VAR_JSON = _make_pre_name_field_json("relax.expr.DataflowVar")
-
-
-@pytest.mark.parametrize(
- ("legacy_json", "expected_type", "var_index"),
- [
- (_LEGACY_RELAX_VAR_JSON, "ir.Var", 7),
- (_LEGACY_TIRX_VAR_JSON, "ir.Var", 5),
- (_PRE_NAME_FIELD_IR_VAR_JSON, "ir.Var", 5),
- (_PRE_NAME_FIELD_DATAFLOW_VAR_JSON, "relax.expr.DataflowVar", 5),
- ],
-)
-def test_var_name_legacy_json_graph_rewrite(legacy_json, expected_type, var_index):
- from tvm.ir.json_compact import upgrade_json
-
- original = json.loads(legacy_json)
- expected = copy.deepcopy(original)
- for node in expected["nodes"]:
- if node["type"] == "tirx.Add":
- node["type"] = "ir.prim.Add"
- expected["nodes"][var_index]["type"] = expected_type
- fields = expected["nodes"][var_index]["data"]
- if "name_hint" in fields:
- fields["name"] = fields.pop("name_hint")
-
- upgraded = json.loads(upgrade_json(legacy_json))
- assert upgraded == expected
- assert upgraded["root_index"] == original["root_index"]
- assert len(upgraded["nodes"]) == len(original["nodes"])
-
-
-def _check_legacy_var(
- var, source_name, line, end_line, column, end_column, expected_type=tvm.ir.Var
-):
- assert type(var) is expected_type
- assert var.name == "legacy"
- assert var.ty == tvm.ir.PrimType("int64")
- assert var.loc.source_name.name == source_name
- assert var.loc.start_line == line
- assert var.loc.end_line == end_line
- assert var.loc.start_column == column
- assert var.loc.end_column == end_column
-
-
-def test_var_exact_base_legacy_relax_json_load():
- restored = tvm.ir.load_json(_LEGACY_RELAX_VAR_JSON)
- assert isinstance(restored, tvm.relax.Tuple)
- assert restored.fields[0].same_as(restored.fields[1])
- _check_legacy_var(restored.fields[0], "legacy_relax.py", 3, 5, 3, 11)
- assert restored.loc.same_as(restored.fields[0].loc)
- assert {node["type"] for node in json.loads(tvm.ir.save_json(restored))["nodes"]}.isdisjoint(
- {"relax.expr.Var", "tirx.Var"}
- )
-
-
-def test_var_exact_base_legacy_tirx_json_load():
- restored = tvm.ir.load_json(_LEGACY_TIRX_VAR_JSON.replace("tirx.Add", "prim.Add"))
- assert isinstance(restored, tvm.tirx.Add)
- assert restored.a.same_as(restored.b)
- _check_legacy_var(restored.a, "legacy_tirx.py", 7, 9, 2, 14)
- assert restored.loc.same_as(restored.a.loc)
- assert {node["type"] for node in json.loads(tvm.ir.save_json(restored))["nodes"]}.isdisjoint(
- {"relax.expr.Var", "tirx.Var", "tirx.Add"}
- )
-
-
-@pytest.mark.parametrize(
- ("legacy_json", "expected_type"),
- [
- (_PRE_NAME_FIELD_IR_VAR_JSON, tvm.ir.Var),
- (_PRE_NAME_FIELD_DATAFLOW_VAR_JSON, tvm.relax.DataflowVar),
- ],
-)
-def test_var_name_legacy_json_load(legacy_json, expected_type):
- restored = tvm.ir.load_json(legacy_json)
- _check_legacy_var(restored, "legacy_tirx.py", 7, 9, 2, 14, expected_type)
- graph = json.loads(tvm.ir.save_json(restored))
- fields = graph["nodes"][graph["root_index"]]["data"]
- assert "name" in fields
- assert "name_hint" not in fields
-
-
-def test_dataflow_var_json_is_not_migrated_to_canonical_var():
- dataflow_var = tvm.relax.DataflowVar("value", tvm.ir.PrimType("int64"))
- graph = json.loads(tvm.ir.save_json(dataflow_var))
- root = graph["nodes"][graph["root_index"]]
- assert root["type"] == "relax.expr.DataflowVar"
- restored = tvm.ir.load_json(json.dumps(graph))
- assert type(restored) is tvm.relax.DataflowVar
-
-
-def _test_infinity_value(value, dtype):
- x = tvm.tirx.const(value, dtype)
- json_str = tvm.ir.save_json(x)
- tvm.ir.assert_structural_equal(x, tvm.ir.load_json(json_str))
-
-
-def test_infinity_value():
- _test_infinity_value(float("inf"), "float64")
- _test_infinity_value(float("-inf"), "float64")
- _test_infinity_value(float("inf"), "float32")
- _test_infinity_value(float("-inf"), "float32")
-
-
-def _test_minmax_value(value):
- json_str = tvm.ir.save_json(value)
- tvm.ir.assert_structural_equal(value, tvm.ir.load_json(json_str))
-
-
-def test_minmax_value():
- _test_minmax_value(tvm.tirx.min_value("float32"))
- _test_minmax_value(tvm.tirx.max_value("float32"))
-
-
-def test_make_smap():
- # save load json
- x = tvm.tirx.const(1, "int32")
- y = tvm.tirx.const(10, "int32")
- z = tvm.tirx.Add(x, y)
- smap = tvm.runtime.convert({"z": z, "x": x})
- json_str = tvm.ir.save_json(tvm.runtime.convert([smap]))
- arr = tvm.ir.load_json(json_str)
- assert len(arr) == 1
- assert arr[0]["z"].a == arr[0]["x"]
- tvm.ir.assert_structural_equal(arr, [smap], map_free_vars=True)
-
-
-def test_make_node():
- x = tvm.ir.make_node("ir.IntImm", ty=tvm.ir.PrimType("int32"), value=10)
- assert isinstance(x, tvm.tirx.IntImm)
- assert x.value == 10
- assert x.loc.same_as(tvm.ir.UnknownLoc())
- A = te.placeholder((10,), name="A")
- AA = tvm.ir.make_node(
- "te.Tensor", shape=A.shape, dtype=A.dtype, op=A.op, value_index=A.value_index
- )
- assert AA.op == A.op
- assert AA.value_index == A.value_index
-
- y = tvm.ir.make_node(
- "ir.IntImm",
- ty=tvm.ir.PrimType(tvm_ffi.core.String("int32")),
- value=10,
- loc=tvm.ir.UnknownLoc(),
- )
- assert isinstance(y, tvm.tirx.IntImm)
- assert y.value == 10
- assert x.loc.same_as(y.loc)
- assert tvm.ir.load_json(tvm.ir.save_json(x)).loc.same_as(tvm.ir.UnknownLoc())
-
-
-def test_make_sum():
- A = te.placeholder((2, 10), name="A")
- k = te.reduce_axis((0, 10), "k")
- B = te.compute((2,), lambda i: te.sum(A[i, k], axis=k), name="B")
- json_str = tvm.ir.save_json(B)
- BB = tvm.ir.load_json(json_str)
- assert B.op.body[0].combiner is not None
- assert BB.op.body[0].combiner is not None
-
-
-def test_string():
- # non printable str, need to store by b64
- s1 = tvm_ffi.core.String("xy\x01z")
- s2 = tvm.ir.load_json(tvm.ir.save_json(s1))
- tvm.ir.assert_structural_equal(s1, s2)
-
- # printable str, need to store by repr_str
- s1 = tvm_ffi.core.String("xyz")
- s2 = tvm.ir.load_json(tvm.ir.save_json(s1))
- tvm.ir.assert_structural_equal(s1, s2)
-
-
-def test_pass_config():
- cfg = tvm.ir.transform.PassContext(
- opt_level=1,
- config={
- "tirx.UnrollLoop": {
- "auto_max_step": 10,
- }
- },
- )
- cfg.opt_level == 1
-
- assert cfg.config["tirx.UnrollLoop"].auto_max_step == 10
- # default option
- assert cfg.config["tirx.UnrollLoop"].explicit_unroll == True
-
- # schema checking for specific config key
- with pytest.raises(TypeError):
- cfg = tvm.ir.transform.PassContext(config={"tirx.UnrollLoop": {"invalid": 1}})
-
- # schema check for un-registered config
- with pytest.raises(AttributeError):
- cfg = tvm.ir.transform.PassContext(config={"inavlid-opt": True})
-
- # schema check for wrong type
- with pytest.raises(AttributeError):
- cfg = tvm.ir.transform.PassContext(config={"tirx.UnrollLoop": 1})
-
-
-def test_dict():
- x = tvm.tirx.const(1) # a class that has Python-defined methods
- # instances should see the full class dict
- assert set(dir(x.__class__)) <= set(dir(x))
-
-
-def test_tensor():
- dev = tvm.cpu(0)
- tvm_arr = tvm.runtime.tensor(np.random.rand(4), device=dev)
- tvm_arr2 = tvm.ir.load_json(tvm.ir.save_json(tvm_arr))
- tvm.ir.assert_structural_equal(tvm_arr, tvm_arr2)
- np.testing.assert_array_equal(tvm_arr.numpy(), tvm_arr2.numpy())
-
-
-def test_tensor_dict():
- dev = tvm.cpu(0)
- m1 = {
- "key1": tvm.runtime.tensor(np.random.rand(4), device=dev),
- "key2": tvm.runtime.tensor(np.random.rand(4), device=dev),
- }
- m2 = tvm.ir.load_json(tvm.ir.save_json(m1))
- tvm.ir.assert_structural_equal(m1, m2)
-
-
-def test_free_var_equal():
- x = tvm.tirx.Var("x", ty="int32")
- y = tvm.tirx.Var("y", ty="int32")
- z = tvm.tirx.Var("z", ty="int32")
- v1 = x + y
- v1 = y + z
- tvm.ir.assert_structural_equal(x, z, map_free_vars=True)
-
-
-if __name__ == "__main__":
- tvm.testing.main()
diff --git a/tests/python/ir/test_nonnullable_ir.py b/tests/python/ir/test_nonnullable_ir.py
index 5330da8465..ea04a62fb6 100644
--- a/tests/python/ir/test_nonnullable_ir.py
+++ b/tests/python/ir/test_nonnullable_ir.py
@@ -45,43 +45,17 @@ def test_optional_statement_fields_and_roundtrip():
body = tvm.ir.Evaluate(0)
conditional = tvm.ir.If(True, body, None)
assert conditional.else_case is None
- loop = tvm.ir.For(tvm.tirx.Var("i", "int32"), 0, 4, tvm.ir.ForKind.DEFAULT, body, step=None)
+ loop = tvm.ir.For(
+ tvm.ir.Var("i", tvm.ir.PrimType("int32")), 0, 4, tvm.ir.ForKind.DEFAULT, body, step=None
+ )
assert loop.step is None
- declaration = tvm.tirx.Function([], None)
- assert declaration.body is None
- assert len(declaration.with_body(body).body) == 0
- for value in [conditional, loop, declaration]:
+ for value in [conditional, loop]:
restored = tvm.ir.load_json(tvm.ir.save_json(value))
tvm.ir.assert_structural_equal(value, restored, map_free_vars=True)
- assert "def" in declaration.script()
-def test_offset_default_and_missing_type_remain_values():
- tensor = tvm.tirx.decl_tensor((8,), "float32", elem_offset=None)
- assert int(tensor.ty.elem_offset) == 0
+def test_missing_type_and_zero_remain_values():
assert isinstance(tvm.ir.Type.missing(), tvm.ir.MissingType)
# False and zero are valid primitive expressions, not absence.
assert int(tvm.ir.Evaluate(False).value) == 0
assert int(tvm.ir.Evaluate(0).value) == 0
-
-
-def test_failed_iter_map_has_absent_padding_predicate():
- i = tvm.tirx.Var("i", "int32")
- result = tvm.sym.detect_iter_map([i], {i: tvm.ir.Range(0, i)})
- assert result.errors
- assert result.padding_predicate is None
- restored = tvm.ir.load_json(tvm.ir.save_json(result))
- assert restored.padding_predicate is None
-
-
-def test_unmapped_python_mutator_var_is_absent():
- @tvm.relax.expr_functor.mutator
- class Mutator(tvm.relax.PyExprMutator):
- pass
-
- mutator = Mutator()
- old = tvm.relax.Var("old", tvm.ir.PrimType("int32"))
- new = tvm.relax.Var("new", tvm.ir.PrimType("int32"))
- assert mutator.get_var_remap(old) is None
- mutator.set_var_remap(old, new)
- assert mutator.get_var_remap(old).same_as(new)
diff --git a/tests/python/ir/test_pass_instrument.py b/tests/python/ir/test_pass_instrument.py
deleted file mode 100644
index e469a75b55..0000000000
--- a/tests/python/ir/test_pass_instrument.py
+++ /dev/null
@@ -1,64 +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
-"""Instrument test cases."""
-
-import tvm
-from tvm import relax
-from tvm.ir.transform.instrument import PrintAfterAll, PrintBeforeAll
-from tvm.script import ir as I
-from tvm.script import relax as R
-from tvm.script import s_tir as Ts
-from tvm.script import tirx as T
-
-# pylint: disable=invalid-name,missing-function-docstring,no-value-for-parameter
-
-
-def test_tir_print_all_passes(capsys):
- @Ts.function
- def func(
- A: T.Tensor((128, 128, 128, 128), "float32"), B: T.Tensor((128, 128, 128, 128), "float32")
- ) -> None:
- for i, j, k, l in T.grid(128, 128, 128, 128):
- with Ts.sblock("B"):
- vi, vj, vk, vl = Ts.axis.remap("SSSS", [i, j, k, l])
- B[vi, vj, vk, vl] = A[vi, vj, vk, vl] * 2.0
-
- with tvm.ir.transform.PassContext(opt_level=3, instruments=[PrintBeforeAll(), PrintAfterAll()]):
- tvm.compile(func)
- all_passes_output = capsys.readouterr().out
- assert "Before Running Pass:" in all_passes_output
- assert "After Running Pass:" in all_passes_output
- assert 'name="tirx.' in all_passes_output
-
-
-def test_relax_print_all_passes(capsys):
- @I.ir_module
- class Module:
- @R.function
- def func(x: R.Tensor((16,), "float32"), y: R.Tensor((16,), "float32")):
- z = R.add(x, y)
- y = z
- return y
-
- pipeline = relax.get_pipeline("default_build")
- with tvm.ir.transform.PassContext(opt_level=3, instruments=[PrintBeforeAll(), PrintAfterAll()]):
- pipeline(Module)
- all_passes_output = capsys.readouterr().out
- assert "Before Running Pass:" in all_passes_output
- assert "After Running Pass:" in all_passes_output
- assert 'name="_pipeline"' in all_passes_output
diff --git a/tests/python/ir/test_unique_name_supply.py b/tests/python/ir/test_unique_name_supply.py
deleted file mode 100644
index f440301e1f..0000000000
--- a/tests/python/ir/test_unique_name_supply.py
+++ /dev/null
@@ -1,62 +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.
-import tvm
-import tvm.testing
-from tvm import relax as rx
-from tvm.ir.supply import UniqueNameSupply
-
-
-def _empty_relax_func():
- return rx.Function([], rx.Tuple([]))
-
-
-def test_fresh_name_empty_string():
- """Empty name should produce a valid variable name, not an empty string."""
- ns = UniqueNameSupply("")
- name = ns.fresh_name("", add_prefix=False)
- assert name == "v"
- name2 = ns.fresh_name("", add_prefix=False)
- assert name2 == "v_1"
-
-
-def test_fresh_name_empty_string_with_prefix():
- """Empty name with prefix should produce a valid variable name."""
- ns = UniqueNameSupply("prefix")
- name = ns.fresh_name("", add_prefix=True)
- assert name == "prefix_v"
- name2 = ns.fresh_name("", add_prefix=True)
- assert name2 == "prefix_v_1"
-
-
-def test_ir_module_from_expr_freshens_main_collision():
- main_gv = tvm.ir.GlobalVar("main")
- mod = tvm.IRModule.from_expr(_empty_relax_func(), {main_gv: _empty_relax_func()})
-
- assert sorted(gvar.name_hint for gvar in mod.get_global_vars()) == ["main", "main_1"]
-
-
-def test_ir_module_from_expr_reuses_existing_global_symbol():
- foo_gv = tvm.ir.GlobalVar("foo")
- func = _empty_relax_func().with_attr("global_symbol", "foo")
- mod = tvm.IRModule.from_expr(func, {foo_gv: _empty_relax_func()})
-
- assert mod.get_global_var("foo").same_as(foo_gv)
- assert [gvar.name_hint for gvar in mod.get_global_vars()] == ["foo"]
-
-
-if __name__ == "__main__":
- tvm.testing.main()
diff --git a/tests/python/relax/test_expr.py b/tests/python/relax/test_expr.py
index 402b0ef003..a8a2d2c80a 100644
--- a/tests/python/relax/test_expr.py
+++ b/tests/python/relax/test_expr.py
@@ -15,6 +15,8 @@
# specific language governing permissions and limitations
# under the License.
# ruff: noqa: F811
+import json
+
import numpy as np
import pytest
import tvm_ffi
@@ -481,6 +483,110 @@ def test_shared_operator_surface_calls_function_typed_values():
assert call.args[0].same_as(arg)
+def test_op_set_signature_preserves_custom_validator():
+ op = tvm.ir.Op.get("relax.call_pure_packed")
+ original_args = [(info.name, info.doc) for info in op.args_info]
+ original_ty_args = [(info.name, info.doc) for info in op.ty_args_info]
+ original_var_args = (op.var_args_info.name, op.var_args_info.doc)
+ original_var_ty_args = (op.var_ty_args_info.name, op.var_ty_args_info.doc)
+ try:
+ op.set_signature()
+ with pytest.raises(TypeError, match="call_pure_packed expects a function argument"):
+ tvm.ir.Call(op, []).validate()
+ finally:
+ op.set_signature(
+ original_args,
+ ty_args=original_ty_args,
+ var_args=original_var_args,
+ var_ty_args=original_var_ty_args,
+ )
+
+
+_LEGACY_TIRX_VAR_JSON = """{
+ "root_index": 6,
+ "nodes": [
+ {"type": "ir.SourceName", "data": "legacy_tirx.py"},
+ {"type": "ir.SourceLoc", "data": {"source_name": 0, "start_line": 7, "start_column": 2,
+ "end_line": 9, "end_column": 14}},
+ {"type": "ir.UnknownLoc", "data": null},
+ {"type": "ir.PrimType", "data": {"loc": 2, "dtype": "int64"}},
+ {"type": "ffi.String", "data": "legacy"},
+ {"type": "tirx.Var", "data": {"loc": 1, "ty": 3, "name": 4}},
+ {"type": "tirx.Add", "data": {"loc": 1, "ty": 3, "a": 5, "b": 5}}
+ ],
+ "metadata": {"tvm_version": "0.26.dev0"}
+}"""
+
+
+def _make_pre_name_field_json(type_key):
+ graph = json.loads(_LEGACY_TIRX_VAR_JSON)
+ graph["root_index"] = 5
+ var = graph["nodes"][5]
+ var["type"] = type_key
+ var["data"]["name_hint"] = var["data"].pop("name")
+ return json.dumps(graph)
+
+
+def _check_legacy_var(
+ var, source_name, line, end_line, column, end_column, expected_type=tvm.ir.Var
+):
+ assert type(var) is expected_type
+ assert var.name == "legacy"
+ assert var.ty == tvm.ir.PrimType("int64")
+ assert var.loc.source_name.name == source_name
+ assert var.loc.start_line == line
+ assert var.loc.end_line == end_line
+ assert var.loc.start_column == column
+ assert var.loc.end_column == end_column
+
+
+_PRE_NAME_FIELD_DATAFLOW_VAR_JSON = _make_pre_name_field_json("relax.expr.DataflowVar")
+
+
+@pytest.mark.parametrize(
+ ("legacy_json", "expected_type"),
+ [
+ (_PRE_NAME_FIELD_DATAFLOW_VAR_JSON, tvm.relax.DataflowVar),
+ ],
+)
+def test_var_name_legacy_json_load(legacy_json, expected_type):
+ restored = tvm.ir.load_json(legacy_json)
+ _check_legacy_var(restored, "legacy_tirx.py", 7, 9, 2, 14, expected_type)
+ graph = json.loads(tvm.ir.save_json(restored))
+ fields = graph["nodes"][graph["root_index"]]["data"]
+ assert "name" in fields
+ assert "name_hint" not in fields
+
+
+def test_dataflow_var_json_is_not_migrated_to_canonical_var():
+ dataflow_var = tvm.relax.DataflowVar("value", tvm.ir.PrimType("int64"))
+ graph = json.loads(tvm.ir.save_json(dataflow_var))
+ root = graph["nodes"][graph["root_index"]]
+ assert root["type"] == "relax.expr.DataflowVar"
+ restored = tvm.ir.load_json(json.dumps(graph))
+ assert type(restored) is tvm.relax.DataflowVar
+
+
+def _empty_relax_func():
+ return rx.Function([], rx.Tuple([]))
+
+
+def test_ir_module_from_expr_freshens_main_collision():
+ main_gv = tvm.ir.GlobalVar("main")
+ mod = tvm.IRModule.from_expr(_empty_relax_func(), {main_gv: _empty_relax_func()})
+
+ assert sorted(gvar.name_hint for gvar in mod.get_global_vars()) == ["main", "main_1"]
+
+
+def test_ir_module_from_expr_reuses_existing_global_symbol():
+ foo_gv = tvm.ir.GlobalVar("foo")
+ func = _empty_relax_func().with_attr("global_symbol", "foo")
+ mod = tvm.IRModule.from_expr(func, {foo_gv: _empty_relax_func()})
+
+ assert mod.get_global_var("foo").same_as(foo_gv)
+ assert [gvar.name_hint for gvar in mod.get_global_vars()] == ["foo"]
+
+
if __name__ == "__main__":
tvm.testing.main()
diff --git a/tests/python/relax/test_expr_functor.py b/tests/python/relax/test_expr_functor.py
index 9ce9a86ba2..2bae10069c 100644
--- a/tests/python/relax/test_expr_functor.py
+++ b/tests/python/relax/test_expr_functor.py
@@ -875,5 +875,18 @@ def test_function_parameter_mutation():
tvm.ir.assert_structural_equal(expected, after)
+def test_unmapped_python_mutator_var_is_absent():
+ @tvm.relax.expr_functor.mutator
+ class Mutator(tvm.relax.PyExprMutator):
+ pass
+
+ mutator = Mutator()
+ old = tvm.relax.Var("old", tvm.ir.PrimType("int32"))
+ new = tvm.relax.Var("new", tvm.ir.PrimType("int32"))
+ assert mutator.get_var_remap(old) is None
+ mutator.set_var_remap(old, new)
+ assert mutator.get_var_remap(old).same_as(new)
+
+
if __name__ == "__main__":
tvm.testing.main()
diff --git a/tests/python/relax/test_pipeline.py b/tests/python/relax/test_pipeline.py
index d0ffaaa2a0..6c9ca1551e 100644
--- a/tests/python/relax/test_pipeline.py
+++ b/tests/python/relax/test_pipeline.py
@@ -26,7 +26,7 @@ from tvm.script import tirx as T
from tvm.testing import env
-def test_pipeline_compile():
+def test_pipeline_compile(capsys):
target = tvm.target.Target("llvm", host="llvm")
pipeline = relax.pipeline.get_default_pipeline(target)
@@ -38,7 +38,14 @@ def test_pipeline_compile():
return lv0
mod = Mod
- mod = pipeline(mod)
+ from tvm.ir.transform.instrument import PrintAfterAll, PrintBeforeAll
+
+ with tvm.ir.transform.PassContext(instruments=[PrintBeforeAll(), PrintAfterAll()]):
+ mod = pipeline(mod)
+ output = capsys.readouterr().out
+ assert "Before Running Pass:" in output
+ assert "After Running Pass:" in output
+ assert 'name="_pipeline"' in output
ex = tvm.compile(mod, target)
x_np = np.random.rand(3, 4).astype(np.float32)
diff --git a/tests/python/runtime/test_runtime_tensor.py b/tests/python/runtime/test_runtime_tensor.py
index de2264700d..b43971983c 100644
--- a/tests/python/runtime/test_runtime_tensor.py
+++ b/tests/python/runtime/test_runtime_tensor.py
@@ -80,5 +80,54 @@ def test_from_dlpack_shape_one():
np.testing.assert_equal(view.numpy(), values.numpy().reshape(2, 8))
+@pytest.mark.parametrize(
+ "dtype",
+ [
+ "float4_e2m1fn",
+ "float8_e3m4",
+ "float8_e4m3",
+ "float8_e4m3b11fnuz",
+ "float8_e4m3fn",
+ "float8_e4m3fnuz",
+ "float8_e5m2",
+ "float8_e5m2fnuz",
+ "float8_e8m0fnu",
+ ],
+)
+def test_low_precision_numpy_roundtrip(dtype):
+ ml_dtypes = pytest.importorskip("ml_dtypes")
+ if not hasattr(ml_dtypes, dtype):
+ pytest.skip(f"ml_dtypes does not provide {dtype}")
+ values = np.random.default_rng(0).random((8, 8)).astype(getattr(ml_dtypes, dtype))
+ tensor = tvm.runtime.tensor(values)
+ assert tensor.dtype == dtype
+ np.testing.assert_equal(tensor.numpy(), values)
+
+
+@pytest.mark.parametrize("container", [False, True], ids=["tensor", "tensor-dict"])
+def test_tensor_json_roundtrip(container):
+ tensor = tvm.runtime.tensor(np.arange(4, dtype="float32"))
+ value = (
+ {"key1": tensor, "key2": tvm.runtime.tensor(np.arange(4, dtype="float32") + 1)}
+ if container
+ else tensor
+ )
+ restored = tvm.ir.load_json(tvm.ir.save_json(value))
+ tvm.ir.assert_structural_equal(value, restored)
+ if container:
+ for key in value:
+ np.testing.assert_array_equal(value[key].numpy(), restored[key].numpy())
+ else:
+ np.testing.assert_array_equal(value.numpy(), restored.numpy())
+
+
+def test_tensor_container_identity():
+ tensor = tvm.runtime.tensor([1, 2, 3])
+ values = tvm.runtime.convert([tensor, tensor])
+ assert values[0].same_as(tensor)
+ assert values[1].same_as(tensor)
+ assert isinstance(values[0], tvm.runtime.Tensor)
+
+
if __name__ == "__main__":
tvm.testing.main()
diff --git a/tests/python/sym/test_sym_analyzer_object.py b/tests/python/sym/test_sym_analyzer_object.py
index dd37a3ebd1..6dd2903f7e 100644
--- a/tests/python/sym/test_sym_analyzer_object.py
+++ b/tests/python/sym/test_sym_analyzer_object.py
@@ -15,6 +15,8 @@
# specific language governing permissions and limitations
# under the License.
+import json
+
import pytest
import tvm
@@ -265,5 +267,37 @@ def test_analyzer_object_clone_resets_rewrite_stats():
assert analyzer.rewrite_simplify_stats.rewrites_attempted == source_attempts
+def test_failed_iter_map_has_absent_padding_predicate():
+ i = tvm.tirx.Var("i", "int32")
+ result = tvm.sym.detect_iter_map([i], {i: tvm.ir.Range(0, i)})
+ assert result.errors
+ assert result.padding_predicate is None
+ restored = tvm.ir.load_json(tvm.ir.save_json(result))
+ assert restored.padding_predicate is None
+
+
+def test_symbolic_analysis_legacy_json_load():
+ # Historical symbolic-analysis objects must load through the existing JSON upgrader.
+ graph = {
+ "root_index": 2,
+ "nodes": [
+ {"type": "arith.ConstIntBound", "data": {"min_value": 2, "max_value": 7}},
+ {"type": "arith.ModularSet", "data": {"coeff": 4, "base": 1}},
+ {"type": "ffi.Array", "data": [0, 1, 0]},
+ ],
+ "metadata": {"tvm_version": tvm.__version__},
+ }
+ restored = tvm.ir.load_json(json.dumps(graph))
+ bound, modular, shared_bound = restored
+ assert isinstance(bound, tvm.sym.ConstIntBound)
+ assert (bound.min_value, bound.max_value) == (2, 7)
+ assert isinstance(modular, tvm.sym.ModularSet)
+ assert (modular.coeff, modular.base) == (4, 1)
+ assert bound.same_as(shared_bound)
+ saved_types = {node["type"] for node in json.loads(tvm.ir.save_json(restored))["nodes"]}
+ assert {"sym.ConstIntBound", "sym.ModularSet"} <= saved_types
+ assert not any(key.startswith("arith.") for key in saved_types)
+
+
if __name__ == "__main__":
tvm.testing.main()
diff --git a/tests/python/te/test_te_tensor.py b/tests/python/te/test_te_tensor.py
index 1b84edcda9..73dc5d0926 100644
--- a/tests/python/te/test_te_tensor.py
+++ b/tests/python/te/test_te_tensor.py
@@ -223,6 +223,29 @@ def test_tensor_inputs():
assert tuple(y.op.input_tensors) == (x,)
+def test_make_sum():
+ A = te.placeholder((2, 10), name="A")
+ k = te.reduce_axis((0, 10), "k")
+ B = te.compute((2,), lambda i: te.sum(A[i, k], axis=k), name="B")
+ json_str = tvm.ir.save_json(B)
+ BB = tvm.ir.load_json(json_str)
+ assert B.op.body[0].combiner is not None
+ assert BB.op.body[0].combiner is not None
+
+
+def test_tensor_make_node():
+ tensor = te.placeholder((10,), name="A")
+ restored = tvm.ir.make_node(
+ "te.Tensor",
+ shape=tensor.shape,
+ dtype=tensor.dtype,
+ op=tensor.op,
+ value_index=tensor.value_index,
+ )
+ assert restored.op == tensor.op
+ assert restored.value_index == tensor.value_index
+
+
if __name__ == "__main__":
test_tensor()
test_rank_zero()
diff --git a/tests/python/tirx-base/test_tir_stmt_functor_substitute.py b/tests/python/tirx-base/test_tir_stmt_functor_substitute.py
deleted file mode 100644
index 1c79bc9f67..0000000000
--- a/tests/python/tirx-base/test_tir_stmt_functor_substitute.py
+++ /dev/null
@@ -1,120 +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.
-
-import tvm_ffi
-
-import tvm
-import tvm.testing
-from tvm.script import ir as I
-from tvm.script import tirx as T
-
-
-def _apply_substitute(mod):
- """Apply substitute transform to replace the first parameter with 16."""
- func = mod["main"]
- vmap = {func.params[0]: T.int32(16)}
- new_func = tvm.tirx.Function(
- params=[],
- body=tvm_ffi.structural_map(
- func.body,
- (tvm.tirx.Var, lambda var: vmap.get(var, var)),
- order="post",
- ),
- ).with_attr("global_symbol", func.attrs["global_symbol"])
- return tvm.IRModule.from_expr(new_func)
-
-
-def test_basic_substitute():
- @I.ir_module
- class Before:
- @T.function
- def main(n: T.int32):
- for i in range(n):
- T.evaluate(i)
-
- @I.ir_module
- class Expected:
- @T.function
- def main():
- for i in range(16):
- T.evaluate(i)
-
- After = _apply_substitute(Before)
- tvm.ir.assert_structural_equal(After, Expected)
-
-
-def test_substitute_allocate():
- @I.ir_module
- class Before:
- @T.function
- def main(n: T.int32):
- A = T.alloc_tensor(ty_args=[T.Tensor((n,), "float32")])
- T.evaluate(A.data)
-
- @I.ir_module
- class Expected:
- @T.function
- def main():
- A = T.alloc_tensor(ty_args=[T.Tensor((16,), "float32")])
- T.evaluate(A.data)
-
- After = _apply_substitute(Before)
- tvm.ir.assert_structural_equal(After, Expected)
-
-
-def test_substitute_tensor_load():
- @I.ir_module
- class Before:
- @T.function
- def main(n: T.int32):
- A = T.alloc_tensor(ty_args=[T.Tensor((n,), "float32")])
- for i in range(n):
- T.evaluate(A[i])
-
- @I.ir_module
- class Expected:
- @T.function
- def main():
- A = T.alloc_tensor(ty_args=[T.Tensor((16,), "float32")])
- for i in range(16):
- T.evaluate(A[i])
-
- After = _apply_substitute(Before)
- tvm.ir.assert_structural_equal(After, Expected)
-
-
-def test_substitute_decl_tensor():
- @I.ir_module
- class Before:
- @T.function
- def main(n: T.int32):
- A = T.alloc_tensor(ty_args=[T.Tensor((n,), "float32")])
- T.evaluate(A.data)
-
- @I.ir_module
- class Expected:
- @T.function
- def main():
- A = T.alloc_tensor(ty_args=[T.Tensor((16,), "float32")])
- T.evaluate(A.data)
-
- After = _apply_substitute(Before)
- tvm.ir.assert_structural_equal(After, Expected)
-
-
-if __name__ == "__main__":
- tvm.testing.main()
diff --git a/tests/python/tirx-base/test_tir_tensor.py b/tests/python/tirx-base/test_tir_tensor.py
deleted file mode 100644
index 4c2eb95b4a..0000000000
--- a/tests/python/tirx-base/test_tir_tensor.py
+++ /dev/null
@@ -1,307 +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, F841
-
-import uuid
-
-import numpy as np
-import pytest
-
-import tvm
-import tvm.testing
-from tvm.runtime import MemoryScope
-from tvm.script import tirx as T
-
-
-def test_tensor():
- m = tvm.tirx.Var("m", "int32")
- n = tvm.tirx.Var("n", "int32")
- l = tvm.tirx.Var("l", "int32")
- Ab = tvm.tirx.decl_tensor((m, n), "float32")
- Bb = tvm.tirx.decl_tensor((n, l), "float32")
-
- assert type(Ab) is tvm.ir.Var
- assert tvm.tirx.is_tensor_var(Ab)
- assert isinstance(Ab.ty, tvm.tirx.TensorType)
- assert Ab.ty.dtype == tvm.ir.PrimType("float32")
- assert tuple(Ab.ty.shape) == (m, n)
- assert not tvm.tirx.is_tensor_var(m)
-
- serialized = tvm.ir.save_json(Ab.ty)
- assert '"tirx.TensorType"' in serialized
- restored = tvm.ir.load_json(serialized)
- assert isinstance(restored, tvm.tirx.TensorType)
- tvm.ir.assert_structural_equal(restored, Ab.ty, map_free_vars=True)
-
-
-def test_tensor_var_identity_and_global_var_properties():
- scalar = tvm.ir.Var("scalar", tvm.ir.PrimType("int32"))
- tensor = tvm.tirx.decl_tensor((8,), "float32")
-
- assert isinstance(tensor, tvm.ir.Var)
- assert isinstance(scalar, tvm.ir.Var)
- assert not tvm.tirx.is_tensor_var(scalar)
- assert tvm.tirx.is_tensor_var(tensor)
-
- assert tuple(tensor.shape) == (8,)
- assert tensor.dtype == tvm.DataType("float32")
- assert tensor.data.args[0].same_as(tensor)
-
- for name in ("shape", "dtype", "data"):
- assert not hasattr(scalar, name)
- with pytest.raises(AttributeError, match="has no attribute"):
- getattr(scalar, name)
-
-
-def test_tensor_data_is_typed_projection():
- tensor = tvm.tirx.decl_tensor((8,), "bool", scope="shared")
-
- assert tensor.ty.dtype == tvm.ir.PrimType("bool")
- assert tensor.data.ty == tvm.ir.PtrType(tvm.ir.PrimType("bool"), "shared")
- assert tensor.data.op.name == "tirx.tensor_data_ptr"
- assert tensor.data.args[0].same_as(tensor)
-
-
-def test_tensor_pointer_type_derived_from_dtype_and_scope():
- data = tvm.ir.Var("storage", tvm.ir.PtrType(tvm.ir.PrimType("uint8"), "local"))
- tensor = tvm.tirx.decl_tensor((8,), "float16", data=data)
-
- assert tensor.ty.dtype == tvm.ir.PrimType("float16")
- assert tensor.ty.scope == tvm.runtime.MemoryScope("local")
- assert tensor.data.ty == tvm.ir.PtrType(tvm.ir.PrimType("float16"), "local")
-
-
-@pytest.mark.parametrize("scope", ["global", "shared", tvm.runtime.MemoryScope("local")])
-def test_tensor_type_constructor_scope_order(scope):
- direct = tvm.tirx.TensorType((4, 8), "float32", scope, (16, 1), 4, 64, 2)
- scripted = T.Tensor((4, 8), "float32", scope, (16, 1), 4, 64, 2)
- declared = tvm.tirx.decl_tensor(
- (4, 8),
- "float32",
- scope=scope,
- strides=(16, 1),
- elem_offset=4,
- data_alignment=64,
- offset_factor=2,
- )
- tvm.ir.assert_structural_equal(direct, scripted)
- tvm.ir.assert_structural_equal(direct, declared.ty)
- assert direct.scope.same_as(tvm.runtime.MemoryScope(scope))
-
-
-@pytest.mark.parametrize("scope", ["", "tensor_type.unregistered"])
-def test_tensor_type_rejects_unregistered_scope(scope):
- with pytest.raises(ValueError):
- tvm.tirx.TensorType((4,), "float32", scope)
- with pytest.raises(ValueError):
- T.Tensor((4,), "float32", scope)
- with pytest.raises(ValueError):
- tvm.tirx.decl_tensor((4,), "float32", scope=scope)
-
-
-def test_decl_tensor_physical_data_binding():
- tensor = tvm.tirx.decl_tensor((8,), "float32")
- data = tvm.tirx.Var("data", tensor.data.ty)
-
- decl = tvm.ir.Bind(
- tensor,
- tvm.ir.Call(
- "tirx.decl_tensor",
- [data],
- ty_args=[tensor.ty],
- ),
- )
- assert decl.var.same_as(tensor)
- assert decl.value.args[0].same_as(data)
-
-
-def test_tensor_ptr_to():
- m = tvm.tirx.Var("m", "int32")
- n = tvm.tirx.Var("n", "int32")
- tensor = tvm.tirx.decl_tensor((m, n), "float32", strides=[n + 1, 1], elem_offset=7)
- pointer = tensor.ptr_to([2, 3])
- assert pointer.op.name == "tirx.address_of"
- assert pointer.ty == tvm.ir.PtrType(tvm.ir.PrimType("float32"))
- assert pointer.args[0].source.same_as(tensor)
- tvm.ir.assert_structural_equal(pointer.args[0].indices, [T.int32(2), T.int32(3)])
- shared = tvm.tirx.decl_tensor((m, n), "float32", scope="shared")
- assert shared.ptr_to([0, 0]).ty == tvm.ir.PtrType(tvm.ir.PrimType("float32"), "shared")
-
-
-def _flattened_load(tensor, indices):
- function = tvm.tirx.Function([tensor], [tvm.ir.Evaluate(tensor[tuple(indices)])])
- module = tvm.tirx.transform.FlattenTensor()(tvm.IRModule({"main": function}))
- return module["main"].body[-1].value
-
-
-def _flattened_indices(tensor, indices):
- return _flattened_load(tensor, indices).indices
-
-
-def test_tensor_load():
- m = tvm.tirx.Var("m", "int32")
- n = tvm.tirx.Var("n", "int32")
- Ab = tvm.tirx.decl_tensor((m, n), "float32", elem_offset=100)
- load = Ab[2, 3]
- tvm.ir.assert_structural_equal(load.indices, [T.int32(2), T.int32(3)])
-
-
-def test_tensor_flatten_offset():
- m = tvm.tirx.Var("m", "int32")
- n = tvm.tirx.Var("n", "int32")
- Ab = tvm.tirx.decl_tensor((m, n), "float32", elem_offset=tvm.tirx.IntImm("int32", 100))
- offset = _flattened_indices(Ab, [2, 3])
- tvm.ir.assert_structural_equal(offset, [n * 2 + 103])
-
-
-def test_tensor_index_merge_mult_mod():
- m = tvm.tirx.Var("m", "int32")
- n = tvm.tirx.Var("n", "int32")
- s = tvm.tirx.Var("s", "int32")
- k0 = tvm.tirx.Var("k0", "int32")
- k1 = tvm.tirx.Var("k1", "int32")
- A = tvm.tirx.decl_tensor((m, n), "float32")
- A_stride = tvm.tirx.decl_tensor((m, n), "float32", strides=(s, 1))
-
- def assert_simplified_equal(index_simplified, index_direct):
- (
- tvm.ir.assert_structural_equal(index_simplified, index_direct),
- f"index_simplified={index_simplified}, index_direct={index_direct}",
- )
-
- idxd = tvm.tirx.indexdiv
- idxm = tvm.tirx.indexmod
-
- # Test Case1
- index_simplified = _flattened_indices(
- A_stride, (idxd(idxm(k0, k1), s), idxm(idxm(k0, k1), s) + idxd(k0, k1) * k1)
- )
- index_direct = _flattened_indices(A_stride, (0, k0))
- assert_simplified_equal(index_simplified, index_direct)
-
- # Test Case2
- index_simplified = _flattened_indices(
- A, (idxd(idxm(k0, idxd(k1, s)), n), idxm(idxm(k0, idxd(k1, s)), n) + idxm(k0, k1))
- )
- index_direct = _flattened_indices(A, (0, idxm(k0, idxd(k1, s)) + idxm(k0, k1)))
- assert_simplified_equal(index_simplified, index_direct)
- # Test Case3
- index_simplified = _flattened_indices(
- A,
- (
- idxd((idxd(k0, idxd(k1, s)) * idxd(k1, s)), n) + idxd(idxm(k0, idxd(k1, s)), n),
- idxm((idxd(k0, idxd(k1, s)) * idxd(k1, s)), n) + idxm(idxm(k0, idxd(k1, s)), n),
- ),
- )
- index_direct = _flattened_indices(A, (0, k0))
- assert_simplified_equal(index_simplified, index_direct)
- # Test Case4 (not able to simplify)
- index_simplified = _flattened_indices(
- A, (idxd(idxm(k0, idxd(k1, s)), n), idxm(idxm(k0, idxd(k1, n)), n) + idxm(k0, k1))
- )
- index_direct = _flattened_indices(
- A, (0, idxd(idxm(k0, idxd(k1, s)), n) * n + (idxm(idxm(k0, idxd(k1, n)), n) + idxm(k0, k1)))
- )
- assert_simplified_equal(index_simplified, index_direct)
-
- # Test Case5
- B = tvm.tirx.decl_tensor((1, 14, 14, 1024))
- i = tvm.tirx.Var("i", "int32")
- j = tvm.tirx.Var("j", "int32")
- k = tvm.tirx.Var("k", "int32")
-
- index_simplified1 = _flattened_indices(
- B,
- (
- idxd(idxd(idxd((i * 50176 + j * 28672 + k), 1024), 14), 14),
- idxm(idxd(idxd((i * 50176 + j * 28672 + k), 1024), 14), 14),
- idxm(idxd((i * 50176 + j * 28672 + k), 1024), 14),
- idxm((i * 50176 + j * 28672 + k), 1024),
- ),
- )
- index_simplified2 = _flattened_indices(
- B,
- (
- idxd(idxd(i * 49 + j * 28 + idxd(k, 1024), 14), 14),
- idxm(idxd(i * 49 + j * 28 + idxd(k, 1024), 14), 14),
- idxm(i * 7 + idxd(k, 1024), 14),
- idxm(k, 1024),
- ),
- )
- index_direct = _flattened_indices(B, (0, 0, 0, (i * 50176 + j * 28672 + k)))
- assert_simplified_equal(index_simplified1, index_direct)
- assert_simplified_equal(index_simplified2, index_direct)
-
-
-@pytest.mark.parametrize(
- "shape,strides,expected",
- [([16, 32], None, 512), ([16], None, 16), ([], None, 1), ([16, 32], [40, 1], 640)],
-)
-def test_tensor_flatten(shape, strides, expected):
- """Flattening derives the storage span and binds a fresh physical view."""
- buf = tvm.tirx.decl_tensor(shape, strides=strides)
- flat = _flattened_load(buf, [0] * len(shape)).source
- assert not buf.same_as(flat)
- assert flat.data.args[0].same_as(flat)
- assert flat.data.op.name == "tirx.tensor_data_ptr"
- tvm.ir.assert_structural_equal(flat.ty.shape, [T.int32(expected)])
-
-
-def test_memory_scope_typed_conversion():
- custom = MemoryScope.register("test.tensor.scope." + uuid.uuid4().hex)
- for scope in (custom, custom.name):
- ptr_ty = tvm.ir.PtrType(tvm.ir.PrimType("float32"), scope)
- tensor = tvm.tirx.decl_tensor((4,), "float32", scope=scope)
- tensor_ty = tvm.tirx.TensorType((4,), "float32", scope)
- assert tensor_ty.scope.same_as(custom)
- assert ptr_ty.scope.same_as(custom)
- assert tensor.ty.scope.same_as(custom)
- assert tensor.scope().same_as(custom)
- assert tensor.data.ty.scope.same_as(custom)
-
- unknown = "test.memory_scope.unknown." + uuid.uuid4().hex
- for construct in (
- lambda: tvm.ir.PtrType(tvm.ir.PrimType("float32"), unknown),
- lambda: tvm.tirx.decl_tensor((4,), "float32", scope=unknown),
- lambda: T.Tensor((4,), "float32", scope=unknown),
- ):
- with pytest.raises(ValueError):
- construct()
- with pytest.raises(ValueError):
- MemoryScope(unknown)
-
-
-def test_memory_scope_global_defaults():
- global_scope = MemoryScope("global")
- assert tvm.ir.PtrType(tvm.ir.PrimType("float32")).scope.same_as(global_scope)
- assert tvm.tirx.decl_tensor((4,), "float32").scope().same_as(global_scope)
- for scope in ("global", global_scope):
- assert tvm.ir.PtrType(tvm.ir.PrimType("float32"), scope).scope.same_as(global_scope)
- assert tvm.tirx.decl_tensor((4,), "float32", scope=scope).scope().same_as(global_scope)
- assert T.Tensor((4,), "float32", scope=scope).scope.same_as(global_scope)
- for construct in (
- lambda: MemoryScope(""),
- lambda: tvm.ir.PtrType(tvm.ir.PrimType("float32"), ""),
- lambda: tvm.tirx.decl_tensor((4,), "float32", scope=""),
- lambda: T.Tensor((4,), "float32", scope=""),
- ):
- with pytest.raises(ValueError):
- construct()
-
-
-if __name__ == "__main__":
- tvm.testing.main()
diff --git a/tests/python/tirx-transform/test_tir_transform_common_subexpr_elim.py b/tests/python/tirx-transform/test_tir_transform_common_subexpr_elim.py
deleted file mode 100644
index 16efaf4351..0000000000
--- a/tests/python/tirx-transform/test_tir_transform_common_subexpr_elim.py
+++ /dev/null
@@ -1,805 +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.
-import hashlib
-
-import tvm
-from tvm.ir.base import save_json
-from tvm.script import tirx as T
-
-
-# =====================================================================
-# T1: Basic multi-level CSE
-# Two common subexpressions at different scoping levels.
-# =====================================================================
-def test_basic():
- @tvm.script.ir_module
- class Before:
- @T.function
- def main(B: T.Tensor((50,), "int32"), i1: T.int32, i2: T.int32, z3: T.int32):
- z1 = T.bind(1)
- z2 = T.bind(2)
- B[i1] = z1 + z2
- x = T.bind(1)
- y = T.bind(1)
- a = T.bind((x + y) + (z1 + z2))
- b = T.bind((x + y) + z3)
- B[i2] = a + b
-
- @tvm.script.ir_module
- class Expected:
- @T.function
- def main(B: T.Tensor((50,), "int32"), i1: T.int32, i2: T.int32, z3: T.int32):
- z1 = T.bind(1)
- z2 = T.bind(2)
- cse_v1 = T.bind(z1 + z2)
- B[i1] = cse_v1
- x = T.bind(1)
- y = T.bind(1)
- cse_v2 = T.bind(x + y)
- a = T.bind(cse_v2 + cse_v1)
- b = T.bind(cse_v2 + z3)
- B[i2] = a + b
-
- after = tvm.tirx.transform.CommonSubexprElim()(Before)
- tvm.ir.assert_structural_equal(after, Expected)
-
-
-# =====================================================================
-# T2: If -- single-branch CSE
-# Duplicated expression only in then-branch stays inside then-branch.
-# =====================================================================
-def test_if_single_branch():
- @tvm.script.ir_module
- class Before:
- @T.function
- def main(
- B: T.Tensor((50,), "int32"),
- i1: T.int32,
- i2: T.int32,
- i3: T.int32,
- y: T.int32,
- z: T.int32,
- ):
- b = T.bind(1)
- if b:
- B[i1] = y + z
- B[i2] = y + z
- else:
- B[i3] = y
-
- @tvm.script.ir_module
- class Expected:
- @T.function
- def main(
- B: T.Tensor((50,), "int32"),
- i1: T.int32,
- i2: T.int32,
- i3: T.int32,
- y: T.int32,
- z: T.int32,
- ):
- b = T.bind(1)
- if b:
- cse_v1 = T.bind(y + z)
- B[i1] = cse_v1
- B[i2] = cse_v1
- else:
- B[i3] = y
-
- after = tvm.tirx.transform.CommonSubexprElim()(Before)
- tvm.ir.assert_structural_equal(after, Expected)
-
-
-# =====================================================================
-# T3: If -- both-branch CSE
-# Duplicated expression in both branches is hoisted before the if.
-# =====================================================================
-def test_if_both_branches():
- @tvm.script.ir_module
- class Before:
- @T.function
- def main(
- B: T.Tensor((50,), "int32"),
- i1: T.int32,
- i2: T.int32,
- i3: T.int32,
- y: T.int32,
- z: T.int32,
- ):
- b = T.bind(1)
- if b:
- B[i1] = y + z
- B[i2] = y
- else:
- B[i3] = y + z
-
- @tvm.script.ir_module
- class Expected:
- @T.function
- def main(
- B: T.Tensor((50,), "int32"),
- i1: T.int32,
- i2: T.int32,
- i3: T.int32,
- y: T.int32,
- z: T.int32,
- ):
- b = T.bind(1)
- cse_v1 = T.bind(y + z)
- if b:
- B[i1] = cse_v1
- B[i2] = y
- else:
- B[i3] = cse_v1
-
- after = tvm.tirx.transform.CommonSubexprElim()(Before)
- tvm.ir.assert_structural_equal(after, Expected)
-
-
-# =====================================================================
-# T4: Cascade CSE
-# Introducing (x+y)+z creates opportunity for x+y.
-# =====================================================================
-def test_cascade():
- @tvm.script.ir_module
- class Before:
- @T.function
- def main(
- B: T.Tensor((50,), "int32"),
- i1: T.int32,
- i2: T.int32,
- i3: T.int32,
- x: T.int32,
- y: T.int32,
- z: T.int32,
- ):
- B[i1] = (x + y) + z
- B[i2] = (x + y) + z
- B[i3] = x + y
-
- @tvm.script.ir_module
- class Expected:
- @T.function
- def main(
- B: T.Tensor((50,), "int32"),
- i1: T.int32,
- i2: T.int32,
- i3: T.int32,
- x: T.int32,
- y: T.int32,
- z: T.int32,
- ):
- cse_v2 = T.bind(x + y)
- cse_v1 = T.bind(cse_v2 + z)
- B[i1] = cse_v1
- B[i2] = cse_v1
- B[i3] = cse_v2
-
- after = tvm.tirx.transform.CommonSubexprElim()(Before)
- tvm.ir.assert_structural_equal(after, Expected)
-
-
-# =====================================================================
-# T5: No change when no duplication
-# Single occurrence — pass is identity.
-# =====================================================================
-def test_no_duplication():
- @tvm.script.ir_module
- class Before:
- @T.function
- def main(x: T.int32, y: T.int32, z: T.int32):
- a = T.bind(x + (y + z))
- T.evaluate(a)
-
- @tvm.script.ir_module
- class Expected:
- @T.function
- def main(x: T.int32, y: T.int32, z: T.int32):
- a = T.bind(x + (y + z))
- T.evaluate(a)
-
- after = tvm.tirx.transform.CommonSubexprElim()(Before)
- tvm.ir.assert_structural_equal(after, Expected)
-
-
-# =====================================================================
-# T6: Deterministic output
-# Multiple runs on same input produce identical output (same var names,
-# ordering). Verified via JSON serialization hash.
-# =====================================================================
-def test_deterministic():
- NUM_TERMS = 10
- REPEATS = 10
-
- x = tvm.tirx.Var("x", "int32")
- result = tvm.tirx.Var("result", "int32")
-
- offsets = sorted([i + 1 for i in range(NUM_TERMS)])
- inc1 = [(x + offsets[i]) for i in range(NUM_TERMS)]
- inc2 = [(x + offsets[i]) for i in range(NUM_TERMS)]
-
- expression = x
- for add in inc1 + inc2:
- expression = expression + add
- body = tvm.ir.SeqStmt([tvm.ir.Bind(result, expression), tvm.ir.Evaluate(result)])
- mod = tvm.IRModule.from_expr(tvm.tirx.Function([x], body))
-
- initial_hash = None
- for _ in range(REPEATS):
- out = tvm.tirx.transform.CommonSubexprElim()(mod)
- func = out["main"]
- json_val = save_json(func)
- json_hash = hashlib.sha256(json_val.encode()).hexdigest()
- if initial_hash is None:
- initial_hash = json_hash
- assert json_hash == initial_hash
-
-
-# =====================================================================
-# T7: CSE inside for-loop
-# Common sub-expression inside a for-loop body.
-# =====================================================================
-def test_for_loop():
- @tvm.script.ir_module
- class Before:
- @T.function
- def main(B: T.Tensor((50,), "int32"), y: T.int32, z: T.int32):
- for i in range(10):
- B[i] = y + z
- B[i + 10] = y + z
-
- @tvm.script.ir_module
- class Expected:
- @T.function
- def main(B: T.Tensor((50,), "int32"), y: T.int32, z: T.int32):
- for i in range(10):
- cse_v1 = T.bind(y + z)
- B[i] = cse_v1
- B[i + 10] = cse_v1
-
- after = tvm.tirx.transform.CommonSubexprElim()(Before)
- tvm.ir.assert_structural_equal(after, Expected)
-
-
-# =====================================================================
-# T8: CSE across for-loop and outer scope
-# Expression appears both outside and inside a for-loop. Binding placed
-# in outer scope before the for.
-# =====================================================================
-def test_for_hoist():
- @tvm.script.ir_module
- class Before:
- @T.function
- def main(B: T.Tensor((50,), "int32"), y: T.int32, z: T.int32):
- B[0] = y + z
- for i in range(10):
- B[i + 1] = y + z
-
- @tvm.script.ir_module
- class Expected:
- @T.function
- def main(B: T.Tensor((50,), "int32"), y: T.int32, z: T.int32):
- cse_v1 = T.bind(y + z)
- B[0] = cse_v1
- for i in range(10):
- B[i + 1] = cse_v1
-
- after = tvm.tirx.transform.CommonSubexprElim()(Before)
- tvm.ir.assert_structural_equal(after, Expected)
-
-
-# =====================================================================
-# T9: Cannot-lift -- expressions containing TensorLoad
-# Expressions containing TensorLoad are ineligible even when duplicated,
-# because lifting them would change semantics (tensor may alias).
-# =====================================================================
-def test_cannot_lift_tensorload():
- @tvm.script.ir_module
- class Before:
- @T.function
- def main(A: T.Tensor((50,), "int32"), B: T.Tensor((50,), "int32")):
- B[0] = A[0] + A[0]
- B[1] = A[0] + A[0]
-
- @tvm.script.ir_module
- class Expected:
- @T.function
- def main(A: T.Tensor((50,), "int32"), B: T.Tensor((50,), "int32")):
- B[0] = A[0] + A[0]
- B[1] = A[0] + A[0]
-
- after = tvm.tirx.transform.CommonSubexprElim()(Before)
- tvm.ir.assert_structural_equal(after, Expected)
-
-
-# =====================================================================
-# T10: Nested if -- multi-level scope LCA
-# Expression in both branches of an inner if, but not the outer else.
-# Binding placed at the inner if's parent scope.
-# =====================================================================
-def test_nested_if():
- @tvm.script.ir_module
- class Before:
- @T.function
- def main(
- B: T.Tensor((50,), "int32"),
- c1: T.int32,
- c2: T.int32,
- y: T.int32,
- z: T.int32,
- ):
- if c1:
- if c2:
- B[0] = y + z
- else:
- B[1] = y + z
- else:
- B[2] = y
-
- @tvm.script.ir_module
- class Expected:
- @T.function
- def main(
- B: T.Tensor((50,), "int32"),
- c1: T.int32,
- c2: T.int32,
- y: T.int32,
- z: T.int32,
- ):
- if c1:
- cse_v1 = T.bind(y + z)
- if c2:
- B[0] = cse_v1
- else:
- B[1] = cse_v1
- else:
- B[2] = y
-
- after = tvm.tirx.transform.CommonSubexprElim()(Before)
- tvm.ir.assert_structural_equal(after, Expected)
-
-
-# =====================================================================
-# T11: Multiple independent CSE candidates
-# Several independent expressions, each duplicated.
-# =====================================================================
-def test_multi_independent():
- @tvm.script.ir_module
- class Before:
- @T.function
- def main(
- B: T.Tensor((50,), "int32"),
- a: T.int32,
- b: T.int32,
- c: T.int32,
- d: T.int32,
- ):
- B[0] = a + b
- B[1] = c + d
- B[2] = a + b
- B[3] = c + d
-
- @tvm.script.ir_module
- class Expected:
- @T.function
- def main(
- B: T.Tensor((50,), "int32"),
- a: T.int32,
- b: T.int32,
- c: T.int32,
- d: T.int32,
- ):
- cse_v1 = T.bind(a + b)
- B[0] = cse_v1
- cse_v2 = T.bind(c + d)
- B[1] = cse_v2
- B[2] = cse_v1
- B[3] = cse_v2
-
- after = tvm.tirx.transform.CommonSubexprElim()(Before)
- tvm.ir.assert_structural_equal(after, Expected)
-
-
-# =====================================================================
-# T12: Expression in if-condition and branch
-# One occurrence in the if-condition (parent scope), one in the
-# then-branch. Binding hoisted before the if.
-# =====================================================================
-def test_if_condition():
- @tvm.script.ir_module
- class Before:
- @T.function
- def main(B: T.Tensor((50,), "int32"), y: T.int32, z: T.int32):
- if y + z > 0:
- B[0] = y + z
-
- @tvm.script.ir_module
- class Expected:
- @T.function
- def main(B: T.Tensor((50,), "int32"), y: T.int32, z: T.int32):
- cse_v1 = T.bind(y + z)
- if cse_v1 > 0:
- B[0] = cse_v1
-
- after = tvm.tirx.transform.CommonSubexprElim()(Before)
- tvm.ir.assert_structural_equal(after, Expected)
-
-
-# =====================================================================
-# T13: Cannot-lift -- expression containing Call
-# Function calls cannot be lifted.
-# =====================================================================
-def test_cannot_lift_call():
- @tvm.script.ir_module
- class Before:
- @T.function
- def main(B: T.Tensor((50,), "int32"), x: T.int32):
- B[0] = T.call_extern("my_func", x, ty="int32") + 1
- B[1] = T.call_extern("my_func", x, ty="int32") + 1
-
- @tvm.script.ir_module
- class Expected:
- @T.function
- def main(B: T.Tensor((50,), "int32"), x: T.int32):
- B[0] = T.call_extern("my_func", x, ty="int32") + 1
- B[1] = T.call_extern("my_func", x, ty="int32") + 1
-
- after = tvm.tirx.transform.CommonSubexprElim()(Before)
- tvm.ir.assert_structural_equal(after, Expected)
-
-
-# =====================================================================
-# T14: No single-use binding
-# When all occurrences of a sub-expression are inside a deeper
-# expression that is also CSE'd, the sub-expression should NOT get
-# its own binding (it would be used only once in the parent's value).
-# =====================================================================
-def test_no_single_use_binding():
- @tvm.script.ir_module
- class Before:
- @T.function
- def main(
- B: T.Tensor((50,), "int32"),
- x: T.int32,
- y: T.int32,
- z: T.int32,
- ):
- B[0] = (x + y) + z
- B[1] = (x + y) + z
-
- @tvm.script.ir_module
- class Expected:
- @T.function
- def main(
- B: T.Tensor((50,), "int32"),
- x: T.int32,
- y: T.int32,
- z: T.int32,
- ):
- cse_v1 = T.bind((x + y) + z)
- B[0] = cse_v1
- B[1] = cse_v1
-
- after = tvm.tirx.transform.CommonSubexprElim()(Before)
- tvm.ir.assert_structural_equal(after, Expected)
-
-
-# =====================================================================
-# T15: For-loop extent lifting
-# Expression shared between loop extent and body is hoisted outside the
-# loop rather than being re-evaluated every iteration.
-# =====================================================================
-def test_for_extent_lift():
- @tvm.script.ir_module
- class Before:
- @T.function
- def main(B: T.Tensor((50,), "int32"), y: T.int32, z: T.int32):
- for i in range(y + z):
- B[i] = y + z
-
- @tvm.script.ir_module
- class Expected:
- @T.function
- def main(B: T.Tensor((50,), "int32"), y: T.int32, z: T.int32):
- cse_v1 = T.bind(y + z)
- for i in range(cse_v1):
- B[i] = cse_v1
-
- after = tvm.tirx.transform.CommonSubexprElim()(Before)
- tvm.ir.assert_structural_equal(after, Expected)
-
-
-# =====================================================================
-# T16: Loop-variable expression stays inside loop
-# An expression using the loop variable (i*4) that appears multiple
-# times inside the body must NOT be hoisted outside the loop.
-# =====================================================================
-def test_loop_var_expr_stays_inside():
- @tvm.script.ir_module
- class Before:
- @T.function
- def main(
- A: T.Tensor((50,), "int32"),
- B: T.Tensor((50,), "int32"),
- ):
- for i in range(10):
- A[i * 4] = B[i * 4]
-
- @tvm.script.ir_module
- class Expected:
- @T.function
- def main(
- A: T.Tensor((50,), "int32"),
- B: T.Tensor((50,), "int32"),
- ):
- for i in range(10):
- cse_v1 = T.bind(i * 4)
- A[cse_v1] = B[cse_v1]
-
- after = tvm.tirx.transform.CommonSubexprElim()(Before)
- tvm.ir.assert_structural_equal(after, Expected)
-
-
-# =====================================================================
-# T17: No normalization without commoning
-# Single-occurrence expression is not normalized or modified.
-# =====================================================================
-def test_no_normalization_without_commoning():
- @tvm.script.ir_module
- class Before:
- @T.function
- def main(x: T.int32, y: T.int32, z: T.int32):
- a = T.bind(x + (y + z))
- T.evaluate(a)
-
- @tvm.script.ir_module
- class Expected:
- @T.function
- def main(x: T.int32, y: T.int32, z: T.int32):
- a = T.bind(x + (y + z))
- T.evaluate(a)
-
- after = tvm.tirx.transform.CommonSubexprElim()(Before)
- tvm.ir.assert_structural_equal(after, Expected)
-
-
-# =====================================================================
-# T18: Let-bound variable -- no extraction from Let body
-# Expressions inside a Let body that reference the Let-bound variable
-# must NOT be extracted, because the variable is only in scope inside
-# the Let body, not at the statement level where Bind would be placed.
-# =====================================================================
-def test_let_body_no_extraction():
- """CSE must not extract expressions from Let bodies that use Let-bound vars."""
- x = tvm.tirx.Var("x", "int32")
- y = tvm.tirx.Var("y", "int32")
- # Let(x, 1, (x+y) + (x+y)) -- x+y appears twice but x is Let-bound
- let_expr = tvm.tirx.Let(x, tvm.tirx.IntImm("int32", 1), (x + y) + (x + y))
- buf = tvm.tirx.decl_tensor((10,), "int32", name="B")
- i = tvm.tirx.Var("i", "int32")
- store = tvm.ir.TensorStore(buf, [i], let_expr)
- loop = tvm.ir.For(
- i,
- tvm.tirx.const(0, "int32"),
- tvm.tirx.const(10, "int32"),
- tvm.ir.ForKind.DEFAULT,
- store,
- )
- func = tvm.tirx.Function([buf, y], loop)
- mod = tvm.IRModule({"main": func})
- mod_after = tvm.tirx.transform.CommonSubexprElim()(mod)
- # No CSE variables should be introduced
- script = mod_after["main"].script()
- assert "cse_v" not in script, f"CSE incorrectly extracted from Let body:\n{script}"
-
-
-# =====================================================================
-# T19: Let value -- CSE works for shared Let value expressions
-# The Let value is evaluated before the binding takes effect, so
-# expressions in the Let value CAN be CSE'd with expressions outside.
-# =====================================================================
-def test_let_value_cse():
- """CSE can extract from Let values (computed before binding)."""
- x = tvm.tirx.Var("x", "int32")
- y = tvm.tirx.Var("y", "int32")
- z = tvm.tirx.Var("z", "int32")
- # Let(x, y+z, x+1) with y+z also appearing outside the Let
- let_expr = tvm.tirx.Let(x, y + z, x + 1)
- buf = tvm.tirx.decl_tensor((10,), "int32", name="B")
- i = tvm.tirx.Var("i", "int32")
- store = tvm.ir.TensorStore(buf, [i], (y + z) + let_expr)
- loop = tvm.ir.For(
- i,
- tvm.tirx.const(0, "int32"),
- tvm.tirx.const(10, "int32"),
- tvm.ir.ForKind.DEFAULT,
- store,
- )
- func = tvm.tirx.Function([buf, y, z], loop)
- mod = tvm.IRModule({"main": func})
- mod_after = tvm.tirx.transform.CommonSubexprElim()(mod)
- # y+z should be extracted (appears in Let value AND outside)
- script = mod_after["main"].script()
- assert "cse_v" in script, f"CSE should extract y+z from Let value:\n{script}"
-
-
-# =====================================================================
-# T20: Nested Let -- both levels respected
-# Multiple nested Let expressions with common sub-expressions in the
-# innermost body. None should be extracted.
-# =====================================================================
-def test_nested_let_no_extraction():
- """CSE must not extract from nested Let bodies."""
- x = tvm.tirx.Var("x", "int32")
- y = tvm.tirx.Var("y", "int32")
- z = tvm.tirx.Var("z", "int32")
- # Let(x, 1, Let(y, 2, (x+y+z) + (x+y+z)))
- inner = (x + y + z) + (x + y + z)
- nested_let = tvm.tirx.Let(
- x, tvm.tirx.IntImm("int32", 1), tvm.tirx.Let(y, tvm.tirx.IntImm("int32", 2), inner)
- )
- buf = tvm.tirx.decl_tensor((10,), "int32", name="B")
- i = tvm.tirx.Var("i", "int32")
- store = tvm.ir.TensorStore(buf, [i], nested_let)
- loop = tvm.ir.For(
- i,
- tvm.tirx.const(0, "int32"),
- tvm.tirx.const(10, "int32"),
- tvm.ir.ForKind.DEFAULT,
- store,
- )
- func = tvm.tirx.Function([buf, z], loop)
- mod = tvm.IRModule({"main": func})
- mod_after = tvm.tirx.transform.CommonSubexprElim()(mod)
- script = mod_after["main"].script()
- assert "cse_v" not in script, f"CSE incorrectly extracted from nested Let body:\n{script}"
-
-
-# =====================================================================
-# T21: Let with lowered floordiv pattern
-# Simulates the pattern produced by LowerIntrin for floordiv:
-# Let(rmod, truncmod(a,b), Let(rdiv, div(a,b), Select(...)))
-# wrapped in Let(x, load, Let(y, load, ...))
-# This is the regression test for the CI failure in lower_intrin tests.
-# =====================================================================
-def test_let_floordiv_pattern():
- """CSE must handle the Let pattern from LowerIntrin's floordiv lowering."""
- x = tvm.tirx.Var("x", "int32")
- y = tvm.tirx.Var("y", "int32")
- rmod = tvm.tirx.Var("rmod", "int32")
- rdiv = tvm.tirx.Var("rdiv", "int32")
- # Simulate lowered floordiv: Let(rmod, x%y, Let(rdiv, x/y, Select(...)))
- select_cond = tvm.tirx.And(y >= 0, rmod >= 0) | tvm.tirx.And(y < 0, rmod <= 0)
- select_expr = tvm.tirx.Select(select_cond, rdiv, rdiv - 1)
- inner_let = tvm.tirx.Let(rdiv, tvm.tirx.Div(x, y), select_expr)
- outer_let = tvm.tirx.Let(rmod, tvm.tirx.Mod(x, y), inner_let)
- # Wrap in Let(x, load, Let(y, load, ...))
- buf_a = tvm.tirx.decl_tensor((10,), "int32", name="A")
- buf_b = tvm.tirx.decl_tensor((10,), "int32", name="B")
- buf_c = tvm.tirx.decl_tensor((10,), "int32", name="C")
- i = tvm.tirx.Var("i", "int32")
- full_expr = tvm.tirx.Let(
- x,
- tvm.tirx.TensorLoad(buf_a, [i]),
- tvm.tirx.Let(y, tvm.tirx.TensorLoad(buf_b, [i]), outer_let),
- )
- store = tvm.ir.TensorStore(buf_c, [i], full_expr)
- loop = tvm.ir.For(
- i,
- tvm.tirx.const(0, "int32"),
- tvm.tirx.const(10, "int32"),
- tvm.ir.ForKind.DEFAULT,
- store,
- )
- func = tvm.tirx.Function([buf_a, buf_b, buf_c], loop)
- mod = tvm.IRModule({"main": func})
- # Should not crash and should not extract Let-bound vars
- mod_after = tvm.tirx.transform.CommonSubexprElim()(mod)
- script = mod_after["main"].script()
- assert "cse_v" not in script, f"CSE incorrectly extracted from Let body:\n{script}"
-
-
-# =====================================================================
-# T22: No lifting of bool predicate (comparison expression)
-# A duplicated `i < n` feeds two if-statements. CSE must leave it
-# inline rather than hoisting a `cse_v: bool = (i < n)` binding.
-# =====================================================================
-def test_no_lift_bool_predicate():
- @tvm.script.ir_module
- class Before:
- @T.function
- def main(B: T.Tensor((50,), "int32"), n: T.int32, x: T.int32):
- for i in range(50):
- if i < n:
- B[i] = x
- if i < n:
- B[i] = x + 1
-
- after = tvm.tirx.transform.CommonSubexprElim()(Before)
- tvm.ir.assert_structural_equal(after, Before)
- assert "cse_v" not in after["main"].script()
-
-
-# =====================================================================
-# T23: No lifting of bool logical expression (And)
-# A duplicated `a && b` feeds two if-statements. CSE must leave it
-# inline rather than hoisting a `cse_v: bool = T.And(a, b)` binding.
-# =====================================================================
-def test_no_lift_bool_logical():
- @tvm.script.ir_module
- class Before:
- @T.function
- def main(B: T.Tensor((50,), "int32"), a: T.bool, b: T.bool, x: T.int32):
- if T.And(a, b):
- B[0] = x
- if T.And(a, b):
- B[1] = x + 1
-
- after = tvm.tirx.transform.CommonSubexprElim()(Before)
- tvm.ir.assert_structural_equal(after, Before)
- assert "cse_v" not in after["main"].script()
-
-
-# T24: Shared subtree stays SSA. UnrollLoop reuses one Stmt object across
-# positions; the identity-keyed rewriter must bind fresh cse vars per occurrence.
-def test_shared_subtree_stays_ssa():
- @tvm.script.ir_module
- class Payload:
- @T.function
- def main(B: T.Tensor((50,), "int32"), i1: T.int32, i2: T.int32):
- B[i1] = (i1 + i2) * 2
- B[i2] = (i1 + i2) * 3
-
- f = Payload["main"]
- shared = f.body # one Stmt object, placed at two positions below
- func = f.with_body(tvm.ir.SeqStmt([shared, shared]))
- after = tvm.tirx.transform.CommonSubexprElim()(tvm.IRModule({"main": func}))["main"]
-
- binds = [s for s in after.body if isinstance(s, tvm.ir.Bind)]
- assert len(binds) == 6, after.script() # (i1+i2, *2, *3) per occurrence
- bound_vars = [b.var for b in binds]
- for i, va in enumerate(bound_vars):
- for vb in bound_vars[i + 1 :]:
- assert not va.same_as(vb), f"duplicate var definitions:\n{after.script()}"
- assert tvm.tirx.analysis.verify_well_formed(after), after.script()
-
-
-if __name__ == "__main__":
- test_basic()
- test_if_single_branch()
- test_if_both_branches()
- test_cascade()
- test_no_duplication()
- test_deterministic()
- test_for_loop()
- test_for_hoist()
- test_cannot_lift_tensorload()
- test_nested_if()
- test_multi_independent()
- test_if_condition()
- test_cannot_lift_call()
- test_no_single_use_binding()
- test_for_extent_lift()
- test_loop_var_expr_stays_inside()
- test_no_normalization_without_commoning()
- test_let_body_no_extraction()
- test_let_value_cse()
- test_nested_let_no_extraction()
- test_let_floordiv_pattern()
- test_no_lift_bool_predicate()
- test_no_lift_bool_logical()
- test_shared_subtree_stays_ssa()
diff --git a/tests/python/tirx-transform/test_tir_transform_function_pass.py b/tests/python/tirx-transform/test_tir_transform_function_pass.py
deleted file mode 100644
index 2c52fc2bc1..0000000000
--- a/tests/python/tirx-transform/test_tir_transform_function_pass.py
+++ /dev/null
@@ -1,69 +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.
-import tvm
-import tvm.testing
-
-
-def test_function_pass():
- @tvm.tirx.transform.function_pass(opt_level=1)
- class TestReplaceFunc:
- """Simple test function to replace one argument to another."""
-
- def __init__(self, new_func):
- self.new_func = new_func
-
- def transform_function(self, func, mod, ctx):
- return self.new_func
-
- x = tvm.tirx.Var("x", "int32")
- y = tvm.tirx.Var("y", "int32")
- b = tvm.tirx.decl_tensor((x,), "float32")
- stmt = tvm.ir.SeqStmt([tvm.ir.Bind(x, 10), tvm.ir.Evaluate(x + 1)])
-
- func = tvm.tirx.Function([x, y, b], stmt)
-
- new_func = tvm.tirx.Function([x, y, b], tvm.ir.Evaluate(0))
-
- mod = tvm.IRModule({"main": func})
- mod = TestReplaceFunc(new_func)(mod)
-
- tvm.ir.assert_structural_equal(mod["main"].body, new_func.body)
-
-
-def test_cow_pass():
- def fapply(f):
- assert tvm.testing.object_use_count(f) == 1
- return f
-
- pidentity = tvm.tirx.transform.Apply(fapply)
- x = tvm.tirx.Var("x", "int32")
- func = tvm.tirx.Function([x], tvm.ir.Evaluate(x)).with_attr("target_bits", 32)
- func_hash = func.__hash__()
- mod = tvm.IRModule({"main": func})
- del func
- # copy on write
- mod_hash = mod.__hash__()
- mod = tvm.ir.transform.Sequential([pidentity, tvm.tirx.transform.NarrowDataType(32)])(
- mod._move()
- )
- assert mod_hash == mod.__hash__()
- assert func_hash == mod["main"].__hash__()
-
-
-if __name__ == "__main__":
- test_cow_pass()
- test_function_pass()
diff --git a/tests/python/tirx-transform/test_tir_transform_lower_intrin.py b/tests/python/tirx-transform/test_tir_transform_lower_intrin.py
deleted file mode 100644
index a65bf63382..0000000000
--- a/tests/python/tirx-transform/test_tir_transform_lower_intrin.py
+++ /dev/null
@@ -1,236 +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: RUF005
-import numpy as np
-import pytest
-
-import tvm
-import tvm.testing
-from tvm.testing import env
-
-
-def lower_intrin(params, stmt):
- """wrapper to call transformation in stmt"""
- lower_expr = tvm.ir.is_prim_expr(stmt)
- stmt = tvm.ir.Evaluate(stmt) if lower_expr else stmt
- mod = tvm.IRModule.from_expr(
- tvm.tirx.Function(params, stmt).with_attr("target", tvm.target.Target("llvm"))
- )
- mod = tvm.ir.transform.Sequential(
- [tvm.tirx.transform.StmtSimplify(), tvm.tirx.transform.LowerIntrin()]
- )(mod)
- func = mod["main"]
- stmt = func.body[0]
- return stmt.value if lower_expr else stmt.body
-
-
-def check_value(expr, variables, data, fref):
- """
- Check that expr evaluates to fref(*row) for each row in data.
- variables: list of TIR vars [x] or [x, y] bound to the columns of data.
- data: list of tuples, each tuple has len(variables) elements.
- """
- n = len(data)
- num_vars = len(variables)
- assert num_vars >= 1 and all(len(row) == num_vars for row in data)
-
- # Build input and output tensors
- input_bufs = [
- tvm.tirx.decl_tensor((n,), dtype=variables[i].ty, name=f"v{i}") for i in range(num_vars)
- ]
- out_buf = tvm.tirx.decl_tensor((n,), dtype=expr.ty, name="C")
-
- # Build loop body: for each i, bind variables[j] = input_bufs[j][i], then store expr to out
- loop_var = tvm.tirx.Var("i", "int32")
-
- def make_store(i_var):
- # Build the expression with each variable bound to the corresponding tensor load
- result = expr
- for j in range(num_vars - 1, -1, -1):
- result = tvm.tirx.Let(variables[j], tvm.tirx.TensorLoad(input_bufs[j], [i_var]), result)
- return tvm.ir.TensorStore(out_buf, [i_var], result)
-
- loop = tvm.ir.For(
- loop_var,
- tvm.tirx.const(0, "int32"),
- tvm.tirx.const(n, "int32"),
- tvm.ir.ForKind.DEFAULT,
- make_store(loop_var),
- )
-
- function = tvm.tirx.Function(input_bufs + [out_buf], loop)
- function = function.with_attr({"tirx.noalias": True, "global_symbol": "main"})
- f = tvm.compile(function, "llvm")
-
- arrays = [
- tvm.runtime.tensor(np.array([row[j] for row in data], dtype=str(variables[j].ty)))
- for j in range(num_vars)
- ]
- c = tvm.runtime.tensor(np.zeros(n, dtype=str(expr.ty)))
- f(*arrays, c)
- cref = np.array([fref(*row) for row in data])
- np.testing.assert_equal(c.numpy(), cref)
-
-
-def test_nested_byte_pointer_offset():
- data = tvm.tirx.Var("data", tvm.ir.PtrType(tvm.ir.PrimType("float32")))
- inner = tvm.tirx.ptr_byte_offset(data, 8, ty=data.ty)
- outer = tvm.tirx.ptr_byte_offset(inner, 12, ty=data.ty)
- body = tvm.ir.Evaluate(tvm.tirx.call_extern("void", "consume", outer))
- targets = ["c"]
- if env.has_llvm():
- targets.append("llvm")
- for target in targets:
- func = tvm.tirx.Function([data], body).with_attr("global_symbol", "main")
- tvm.tirx.compile(tvm.IRModule.from_expr(func), target=target).mod
-
-
-@pytest.mark.skipif(not env.has_llvm(), reason="need llvm")
-@pytest.mark.parametrize("dtype, byte_offset", [("float32x2", 16), ("float32x3", 24)])
-def test_vector_byte_pointer_offset(dtype, byte_offset):
- tensor = tvm.tirx.decl_tensor((8,), dtype, name="A")
- pointer = tvm.tirx.ptr_byte_offset(tensor.data, byte_offset, ty=tensor.data.ty)
- body = tvm.ir.Evaluate(tvm.tirx.call_extern("void", "consume", pointer))
- func = tvm.tirx.Function([tensor], body).with_attr("global_symbol", "main")
- tvm.tirx.compile(tvm.IRModule.from_expr(func), target="llvm").mod
-
-
-@pytest.mark.parametrize("shape, indices", [((), []), ((16,), [3]), ((2, 4), [1, 2])])
-def test_logical_pointer_preserves_tensor_identity(shape, indices):
- tensor = tvm.tirx.decl_tensor(shape, "float32", "tensor")
- pointer = tensor.ptr_to(indices)
- func = tvm.tirx.Function([tensor], tvm.ir.Evaluate(pointer)).with_attr(
- "target", tvm.target.Target("llvm")
- )
- lowered = tvm.tirx.transform.LowerIntrin()(tvm.IRModule.from_expr(func))["main"].body[0].value
- assert lowered.op.name == "tirx.address_of"
- load = lowered.args[0]
- assert isinstance(load, tvm.ir.TensorLoad)
- assert load.source.same_as(tensor)
- assert [int(index) for index in load.indices] == indices
-
-
-def get_ref_data():
- """Get reference data for every pairs"""
- import itertools
-
- x = range(-10, 10)
- y = list(range(-10, 10))
- y.remove(0)
- return list(itertools.product(x, y))
-
-
-@pytest.mark.skipif(not env.has_llvm(), reason="need llvm")
-def test_lower_floordiv():
- data = get_ref_data()
- for dtype in ["int32", "int64", "int16"]:
- x = tvm.tirx.Var("x", dtype)
- y = tvm.tirx.Var("y", dtype)
- zero = tvm.tirx.const(0, dtype)
- # no constraints
- res = lower_intrin([x, y], tvm.tirx.floordiv(x, y))
- check_value(res, [x, y], data, lambda a, b: a // b)
- # rhs >= 0
- res = lower_intrin([x, y], tvm.tirx.Select(y >= 0, tvm.tirx.floordiv(x, y), zero))
- check_value(res, [x, y], data, lambda a, b: a // b if b > 0 else 0)
- # involves max
- res = lower_intrin(
- [x, y], tvm.tirx.Select(y >= 0, tvm.tirx.max(tvm.tirx.floordiv(x, y), zero), zero)
- )
- check_value(res, [x, y], data, lambda a, b: max(a // b, 0) if b > 0 else 0)
- # lhs >= 0
- res = lower_intrin(
- [x, y], tvm.tirx.Select(tvm.tirx.all(y >= 0, x >= 0), tvm.tirx.floordiv(x, y), zero)
- )
- check_value(res, [x, y], data, lambda a, b: a // b if b > 0 and a >= 0 else 0)
- # const power of two
- res = lower_intrin([x, y], tvm.tirx.floordiv(x, tvm.tirx.const(8, dtype=dtype)))
- check_value(res, [x, y], [(a, b) for a, b in data if b == 8], lambda a, b: a // b)
- # floordiv(x + m, k), m and k are positive constant. 2 <= m <= k-1.
- res = lower_intrin(
- [x, y],
- tvm.tirx.floordiv(x + tvm.tirx.const(4, dtype=dtype), tvm.tirx.const(5, dtype=dtype)),
- )
- check_value(res, [x, y], [(a, b) for a, b in data if b == 5], lambda a, b: (a + 4) // b)
-
-
-@pytest.mark.skipif(not env.has_llvm(), reason="need llvm")
-def test_lower_floormod():
- data = get_ref_data()
- for dtype in ["int32", "int64", "int16"]:
- x = tvm.tirx.Var("x", dtype)
- y = tvm.tirx.Var("y", dtype)
- zero = tvm.tirx.const(0, dtype)
- # no constraints
- res = lower_intrin([x, y], tvm.tirx.floormod(x, y))
- check_value(res, [x, y], data, lambda a, b: a % b)
- # rhs >= 0
- res = lower_intrin([x, y], tvm.tirx.Select(y >= 0, tvm.tirx.floormod(x, y), zero))
- check_value(res, [x, y], data, lambda a, b: a % b if b > 0 else 0)
- # lhs >= 0
- res = lower_intrin(
- [x, y], tvm.tirx.Select(tvm.tirx.all(y >= 0, x >= 0), tvm.tirx.floormod(x, y), zero)
- )
- check_value(res, [x, y], data, lambda a, b: a % b if b > 0 and a >= 0 else 0)
- # const power of two
- res = lower_intrin([x, y], tvm.tirx.floormod(x, tvm.tirx.const(8, dtype=dtype)))
- check_value(res, [x, y], [(a, b) for a, b in data if b == 8], lambda a, b: a % b)
- # floormod(x + m, k), m and k are positive constant. 2 <= m <= k-1.
- res = lower_intrin(
- [x, y],
- tvm.tirx.floormod(x + tvm.tirx.const(4, dtype=dtype), tvm.tirx.const(5, dtype=dtype)),
- )
- check_value(res, [x, y], [(a, b) for a, b in data if b == 5], lambda a, b: (a + 4) % b)
-
-
-@pytest.mark.skipif(not env.has_llvm(), reason="need llvm")
-def test_lower_floordiv_overflow_checks():
- """
- Regression tests for overflow checks in TryFindShiftCoefficientForPositiveRange.
- Divisor is constant 3 (not 1 to avoid CSE, not power-of-two so we don't take the shift path).
- Reuses lower_intrin and check_value; overflow tests use one var [x].
- """
- # Check 3: (b-1) - a_min must not overflow (numerator and C++ int64).
- # x (int64) full range -> min_value = -2^63. With b = 3: numerator = 2 - (-2^63) > LLONG_MAX.
- x = tvm.tirx.Var("x", "int64")
- res = lower_intrin([x], tvm.tirx.floordiv(x, tvm.tirx.const(3, "int64")))
- data_check3 = [(-(2**63),), (0,), (100,)]
- check_value(res, [x], data_check3, lambda a: a // 3)
-
- # Check 4: c_value * b_value must not overflow dtype.
- # x (int16) full range -> min_value = -32768, c = ceil(32770/3) = 10923; 10923*3 > 32767.
- x = tvm.tirx.Var("x", "int16")
- res = lower_intrin([x], tvm.tirx.floordiv(x, tvm.tirx.const(3, "int16")))
- data_check4 = [(-32768,), (0,), (100,)]
- check_value(res, [x], data_check4, lambda a: a // 3)
-
- # Check 5: a_max + b*c must not overflow (offset numerator).
- # tirx.min(tirx.max(x, -10), 32758) can give bounds [-10, 32758]; b=3, c=4; a_max + 12 > 32767.
- # In practice this path may not be triggered. This test still validates correct lowering.
- x = tvm.tirx.Var("x", "int16")
- clamped = tvm.tirx.min(
- tvm.tirx.max(x, tvm.tirx.const(-10, "int16")), tvm.tirx.const(32758, "int16")
- )
- res = lower_intrin([x], tvm.tirx.floordiv(clamped, tvm.tirx.const(3, "int16")))
- data_check5 = [(-10,), (0,), (32758,), (32757,)]
- check_value(res, [x], data_check5, lambda a: (min(max(a, -10), 32758)) // 3)
-
-
-if __name__ == "__main__":
- test_lower_floordiv()
- test_lower_floormod()
- test_lower_floordiv_overflow_checks()
diff --git a/tests/python/tirx-transform/test_tir_transform_make_packed_api.py b/tests/python/tirx-transform/test_tir_transform_make_packed_api.py
deleted file mode 100644
index 097baf1d08..0000000000
--- a/tests/python/tirx-transform/test_tir_transform_make_packed_api.py
+++ /dev/null
@@ -1,502 +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.
-"""Tests for tirx.transform.MakePackedAPI TIR transform.
-
-Tests verify the transform output using TVMScript before/after patterns.
-Runtime error tests are in tests/python/codegen/test_codegen_error_handling.py.
-"""
-
-import pytest
-import tvm_ffi
-
-import tvm
-import tvm.testing
-from tvm import tirx
-from tvm.script import ir as I
-from tvm.script import tirx as T
-
-
-def _is_tensor_binding(node, *op_names):
- return (
- isinstance(node, tvm.ir.Bind)
- and isinstance(node.value, tvm.ir.Call)
- and isinstance(node.value.op, tvm.ir.Op)
- and node.value.op.name in op_names
- )
-
-
-def _find_compute_scope(func):
- result = None
-
- def _visitor(stmt):
- if isinstance(stmt, tvm.ir.RegionStmt) and stmt.op.same_as(
- tvm.ir.Op.get("tirx.compute_scope")
- ):
- nonlocal result
- result = stmt
-
- tvm_ffi.structural_walk(func.body, _visitor)
-
- return result
-
-
-@pytest.mark.parametrize("use_global_symbol", [True, False])
-def test_no_op_when_global_symbol_is_absent(use_global_symbol):
- func_attr = {"target": tvm.target.Target("llvm", host="llvm")}
-
- @T.function(private=True)
- def before():
- T.func_attr(func_attr)
- T.evaluate(0)
-
- if use_global_symbol:
- before = before.with_attr("global_symbol", "main")
-
- after = tvm.tirx.transform.MakePackedAPI()(tvm.IRModule.from_expr(before))["main"]
- if use_global_symbol:
- assert len(after.params) == 4
- else:
- tvm.ir.assert_structural_equal(before, after)
-
-
-def test_target_host_removed():
- """After MakePackedAPI, host-side target should be the host
-
- MakePackedAPI is the last transform that requires both the device
- and the host. After MakePackedAPI, the target attribute should
- only contain the host-side target.
- """
-
- host = tvm.target.Target("llvm")
-
- @I.ir_module
- class before:
- @T.function
- def main(A: T.Tensor(1, "float32")):
- T.func_attr({"global_symbol": "main", "target": T.target("cuda", host=host)})
- T.evaluate(0)
-
- after = tvm.tirx.transform.MakePackedAPI()(before)
- target_attr = after["main"].attrs["target"]
- assert str(host) == str(target_attr)
-
-
-def test_internal_subroutine_call():
- """Internal subroutines should not use the PackedFunc API
-
- A subroutine without the "global_symbol" attribute is an internal
- subroutine, and is not directly exposed to a user of the generated
- `runtime.Module`. Therefore, it doesn't need to follow the
- PackedFunc API.
- """
-
- @I.ir_module
- class before:
- @T.function
- def main(A: T.Tensor(1, "float32")):
- T.func_attr({"target": T.target("llvm", host="llvm")})
- before.subroutine(A.data)
-
- # this test fails if it's made public
- @T.function(private=True)
- def subroutine(A_data: T.handle("float32")):
- T.func_attr({"target": T.target("llvm")})
- T.evaluate(A_data)
-
- after = tvm.tirx.transform.MakePackedAPI()(before)
- tvm.ir.assert_structural_equal(before["subroutine"], after["subroutine"])
-
- compute_scope = _find_compute_scope(after["main"])
- subroutine_call_op = compute_scope.body[0].value.op
- assert isinstance(subroutine_call_op, tvm.ir.GlobalVar), (
- f"The main function's CallNode should use the subroutine's GlobalVar as the operation, "
- f"but instead has an operation of type {subroutine_call_op}"
- )
-
-
-def test_subroutine_call_to_externally_visible_subroutine():
- """Externally-visible subroutines should use the PackedFunc API
-
- Because the subroutine may be called directly by a user, it must
- use the PackedFunc API. Its signature should be updated to the
- PackedFunc signature, and call sites should be updated to use
- `T.call_cpacked`.
- """
-
- @I.ir_module
- class before:
- @T.function
- def main(A: T.Tensor(1, "float32")):
- T.func_attr({"global_symbol": "main", "target": T.target("llvm", host="llvm")})
- before.subroutine(A.data)
-
- @T.function
- def subroutine(A_data: T.handle("float32")):
- T.func_attr({"global_symbol": "subroutine", "target": T.target("llvm", host="llvm")})
- T.evaluate(A_data)
-
- after = tvm.tirx.transform.MakePackedAPI()(before)
-
- main_compute_scope = _find_compute_scope(after["main"])
- assert main_compute_scope is not None
- subroutine_compute_scope = _find_compute_scope(after["subroutine"])
- assert subroutine_compute_scope is not None
-
- subroutine_call_op = main_compute_scope.body[0].value.op
- assert (
- isinstance(subroutine_call_op, tvm.ir.Op) and subroutine_call_op.name == "tirx.call_cpacked"
- ), (
- f"The main function's CallNode should be lowered to the builtin 'tirx.call_cpacked', "
- f"but instead has an operation of type {subroutine_call_op}"
- )
-
-
-def test_zero_arg_function():
- """Zero-arg function emits num_args check but no null-pointer check."""
-
- @I.ir_module
- class Before:
- @T.function
- def func_without_arg() -> T.int64:
- T.func_attr({"target": T.target("llvm", host="llvm")})
- return T.int64(42)
-
- @I.ir_module
- class Expected:
- @T.function
- def func_without_arg(
- self_handle: T.handle,
- args: T.handle,
- num_args: T.int32,
- result: T.handle("void", "global"),
- ) -> T.int32:
- T.func_attr(
- {
- "calling_conv": 1,
- "global_symbol": "__tvm_ffi_func_without_arg",
- "target": T.target("llvm"),
- }
- )
- assert num_args == 0, (
- "TypeError",
- ["Expected ", "0", " arguments", " when calling:\n `", "func_without_arg()", "`"],
- )
- with T.compute_scope("func_without_arg_compute_"):
- T.abi_field_set(result, 0, 13, 1)
- T.abi_field_set(result, 0, 14, 0)
- T.abi_field_set(result, 0, 15, T.Cast(T.int64(42), "int64"))
- return 0
- return 0
-
- After = tvm.tirx.transform.MakePackedAPI()(Before)
- tvm.ir.assert_structural_equal(After, Expected)
-
-
-def test_pointer_return():
- """Pointer returns are stored as an opaque pointer in TVMFFIAny."""
-
- @I.ir_module
- class Before:
- @T.function
- def main(arg: T.handle) -> T.handle:
- T.func_attr({"target": T.target("llvm", host="llvm")})
- return arg
-
- after = tvm.tirx.transform.MakePackedAPI()(Before)["main"]
- return_type_indices = []
-
- def collect(node):
- if not isinstance(node, tvm.ir.Call) or node.op.name != "tirx.abi_field_set":
- return
- field = node.args[2]
- if isinstance(field, tvm.tirx.IntImm) and int(field) == 13:
- return_type_indices.append(int(node.args[3]))
-
- tvm_ffi.structural_walk(after.body, collect)
- assert 4 in return_type_indices # ffi::TypeIndex::kTVMFFIOpaquePtr
-
-
-def test_return_from_parallel_scope_is_rejected():
- """A parallel loop cannot return from its enclosing function."""
-
- i = tirx.Var("i", "int32")
- body = tvm.ir.For(i, 0, 2, tvm.ir.ForKind.PARALLEL, tvm.ir.Return(i))
- func = tirx.Function([], body, tvm.ir.PrimType("int32"))
- func = func.with_attr("global_symbol", "main")
- func = func.with_attr("target", tvm.target.Target("llvm", host="llvm"))
-
- with pytest.raises(tvm.error.InternalError, match="Return cannot be used in parallel scope"):
- tvm.tirx.transform.MakePackedAPI()(tvm.IRModule({"main": func}))
-
-
-def test_int_parameter():
- """Int parameter emits type check accepting int or bool."""
-
- @I.ir_module
- class Before:
- @T.function
- def main(arg: T.int32) -> T.int32:
- T.func_attr({"target": T.target("llvm", host="llvm")})
- if arg > 0:
- return 10
- else:
- return 20
-
- @I.ir_module
- class Expected:
- @T.function
- def main(
- self_handle: T.handle,
- args: T.handle,
- num_args: T.int32,
- result: T.handle("void", "global"),
- ) -> T.int32:
- T.func_attr(
- {
- "calling_conv": 1,
- "global_symbol": "__tvm_ffi_main",
- "target": T.target("llvm"),
- }
- )
- assert num_args == 1, (
- "TypeError",
- ["Expected ", "1", " arguments", " when calling:\n `", "main(arg: int32)", "`"],
- )
- assert not T.isnullptr(args), (
- "TypeError",
- ["args pointer is NULL", " when calling:\n `", "main(arg: int32)", "`"],
- )
- arg_type_index: T.let[T.int32] = T.abi_field_get(args, 0, 13, ty="int32")
- assert arg_type_index == 1 or arg_type_index == 2, (
- "TypeError",
- [
- "Mismatched type on argument #",
- "0",
- " when calling:\n `",
- "main(arg: int32)",
- "`,\n expected ",
- "int",
- ],
- )
- arg: T.let[T.int32] = T.Cast(T.abi_field_get(args, 0, 15, ty="int64"), "int32")
- with T.compute_scope("main_compute_"):
- if arg > 0:
- T.abi_field_set(result, 0, 13, 1)
- T.abi_field_set(result, 0, 14, 0)
- T.abi_field_set(result, 0, 15, T.Cast(10, "int64"))
- return 0
- else:
- T.abi_field_set(result, 0, 13, 1)
- T.abi_field_set(result, 0, 14, 0)
- T.abi_field_set(result, 0, 15, T.Cast(20, "int64"))
- return 0
- return 0
-
- After = tvm.tirx.transform.MakePackedAPI()(Before)
- tvm.ir.assert_structural_equal(After, Expected)
-
-
-def test_bool_parameter():
- """Bool parameter emits type check accepting bool or int."""
-
- @I.ir_module
- class Before:
- @T.function
- def main(arg: T.bool) -> T.int32:
- T.func_attr({"target": T.target("llvm", host="llvm")})
- if arg:
- return 10
- else:
- return 20
-
- @I.ir_module
- class Expected:
- @T.function
- def main(
- self_handle: T.handle,
- args: T.handle,
- num_args: T.int32,
- result: T.handle("void", "global"),
- ) -> T.int32:
- T.func_attr(
- {
- "calling_conv": 1,
- "global_symbol": "__tvm_ffi_main",
- "target": T.target("llvm"),
- }
- )
- assert num_args == 1, (
- "TypeError",
- ["Expected ", "1", " arguments", " when calling:\n `", "main(arg: bool)", "`"],
- )
- assert not T.isnullptr(args), (
- "TypeError",
- ["args pointer is NULL", " when calling:\n `", "main(arg: bool)", "`"],
- )
- arg_type_index: T.let[T.int32] = T.abi_field_get(args, 0, 13, ty="int32")
- assert arg_type_index == 2 or arg_type_index == 1, (
- "TypeError",
- [
- "Mismatched type on argument #",
- "0",
- " when calling:\n `",
- "main(arg: bool)",
- "`,\n expected ",
- "boolean",
- ],
- )
- arg: T.let[T.bool] = T.Cast(T.abi_field_get(args, 0, 15, ty="int64"), "bool")
- with T.compute_scope("main_compute_"):
- if arg:
- T.abi_field_set(result, 0, 13, 1)
- T.abi_field_set(result, 0, 14, 0)
- T.abi_field_set(result, 0, 15, T.Cast(10, "int64"))
- return 0
- else:
- T.abi_field_set(result, 0, 13, 1)
- T.abi_field_set(result, 0, 14, 0)
- T.abi_field_set(result, 0, 15, T.Cast(20, "int64"))
- return 0
- return 0
-
- After = tvm.tirx.transform.MakePackedAPI()(Before)
- tvm.ir.assert_structural_equal(After, Expected)
-
-
-def test_float_parameter():
- """Float parameter emits type check accepting float, int, or bool."""
-
- @I.ir_module
- class Before:
- @T.function
- def main(arg: T.float32) -> T.int32:
- T.func_attr({"target": T.target("llvm", host="llvm")})
- if arg > T.float32(0):
- return 10
- else:
- return 20
-
- @I.ir_module
- class Expected:
- @T.function
- def main(
- self_handle: T.handle,
- args: T.handle,
- num_args: T.int32,
- result: T.handle("void", "global"),
- ) -> T.int32:
- T.func_attr(
- {
- "calling_conv": 1,
- "global_symbol": "__tvm_ffi_main",
- "target": T.target("llvm"),
- }
- )
- assert num_args == 1, (
- "TypeError",
- ["Expected ", "1", " arguments", " when calling:\n `", "main(arg: float32)", "`"],
- )
- assert not T.isnullptr(args), (
- "TypeError",
- ["args pointer is NULL", " when calling:\n `", "main(arg: float32)", "`"],
- )
- arg_type_index: T.let[T.int32] = T.abi_field_get(args, 0, 13, ty="int32")
- assert arg_type_index == 3 or arg_type_index == 1 or arg_type_index == 2, (
- "TypeError",
- [
- "Mismatched type on argument #",
- "0",
- " when calling:\n `",
- "main(arg: float32)",
- "`,\n expected ",
- "float",
- ],
- )
- arg: T.let[T.float32] = T.Select(
- arg_type_index == 3,
- T.Cast(T.abi_field_get(args, 0, 15, ty="float64"), "float32"),
- T.Cast(T.abi_field_get(args, 0, 15, ty="int64"), "float32"),
- )
- with T.compute_scope("main_compute_"):
- if arg > T.float32(0.0):
- T.abi_field_set(result, 0, 13, 1)
- T.abi_field_set(result, 0, 14, 0)
- T.abi_field_set(result, 0, 15, T.Cast(10, "int64"))
- return 0
- else:
- T.abi_field_set(result, 0, 13, 1)
- T.abi_field_set(result, 0, 14, 0)
- T.abi_field_set(result, 0, 15, T.Cast(20, "int64"))
- return 0
- return 0
-
- After = tvm.tirx.transform.MakePackedAPI()(Before)
- tvm.ir.assert_structural_equal(After, Expected)
-
-
-def test_forward_reference_symbolic_variable():
- """MakePackedAPI succeeds when a symbolic variable is used before it is defined.
-
- When tensor A has shape (batch_size+1,) and tensor B has shape (batch_size,),
- batch_size is referenced (in A's shape check) before it is defined (from B's
- shape). The three-sequence separation (init_nest, asserts, decl_tensors)
- ensures all variable definitions precede all assertions.
- """
-
- batch_size = T.dynamic("batch_size")
-
- @I.ir_module
- class Before:
- @T.function
- def main(A: T.Tensor((batch_size + 1,), "int32"), B: T.Tensor((batch_size,), "int32")):
- T.func_attr({"target": T.target("llvm", host="llvm")})
-
- for i in range(batch_size):
- B[i] = A[i] + A[i + 1]
-
- # Should not raise "variable batch_size has been used before definition"
- After = tvm.tirx.transform.MakePackedAPI()(Before)
- assert len(After["main"].params) == 4
-
-
-def test_tensor_alignment_attached_to_tensor_var():
- """Packed ABI alignment metadata remains keyed by the logical TensorVar."""
-
- @I.ir_module
- class Before:
- @T.function
- def main(A: T.Tensor((16,), "float32", align=64)):
- T.func_attr({"global_symbol": "main", "target": T.target("llvm", host="llvm")})
- T.evaluate(A[0])
-
- after = tvm.tirx.transform.MakePackedAPI()(Before)["main"]
- alignment_nodes = []
- declared_tensors = []
-
- def collect(node):
- if isinstance(node, tvm.ir.Call) and node.op == tvm.ir.Op.get("tirx.assume_aligned"):
- alignment_nodes.append(node.args[0])
- if _is_tensor_binding(node, "tirx.decl_tensor"):
- declared_tensors.append(node.var)
-
- tvm_ffi.structural_walk(after.body, collect)
- assert len(alignment_nodes) == 1
- assert any(alignment_nodes[0].same_as(tensor) for tensor in declared_tensors)
-
-
-if __name__ == "__main__":
- tvm.testing.main()
diff --git a/tests/python/tirx-transform/test_tir_transform_narrow_datatype.py b/tests/python/tirx-transform/test_tir_transform_narrow_datatype.py
deleted file mode 100644
index e3d58a0701..0000000000
--- a/tests/python/tirx-transform/test_tir_transform_narrow_datatype.py
+++ /dev/null
@@ -1,317 +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.
-import tvm
-import tvm.testing
-from tvm.script import tirx as T
-from tvm.tirx import const
-
-
-def lower_stmt(params, stmt, target_bits):
- func = tvm.tirx.Function(params, stmt)
- func = tvm.tirx.transform.NarrowDataType(target_bits)(tvm.IRModule.from_expr(func))["main"]
- stmt = func.body[0]
- return stmt
-
-
-def lower_func_body(func, target_bits):
- """Lower a TVMScript function and return the first For loop in the body."""
- mod = tvm.IRModule.from_expr(func)
- gvar = next(iter(mod.functions.keys()))
- func = tvm.tirx.transform.NarrowDataType(target_bits)(mod)[gvar]
- body = func.body
- # With flat tensor semantics, navigate to the first For node
- if isinstance(body, tvm.ir.SeqStmt):
- for stmt in body:
- if isinstance(stmt, tvm.ir.For):
- return stmt
- while hasattr(body, "body") and not isinstance(body, tvm.ir.For):
- body = body.body[0]
- return body
-
-
-def test_basic():
- def check_const(m, n, target_bits, target_dtype):
- """Check with constant values using TVMScript closure."""
-
- @T.function
- def func(A: T.Tensor((m * n,), "float32"), B: T.Tensor((m * n,), "float32")):
- for i in T.serial(m):
- for j in T.serial(n):
- B[i * n + j] = A[i * n + j] + T.float32(1)
-
- stmt = lower_func_body(func, target_bits)
- assert stmt.loop_var.ty.dtype == target_dtype
- assert stmt.body[0].loop_var.ty.dtype == target_dtype
-
- def check_symbolic(m_dtype, n_dtype, target_bits, target_dtype):
- """Check with symbolic shapes as function parameters."""
- if m_dtype == "int32":
-
- @T.function
- def func(A: T.handle("float32"), B: T.handle("float32"), m: T.int32, n: T.int32):
- A_buf = T.decl_tensor(A, ty_args=[T.Tensor((m * n,), "float32", data=A)])
- B_buf = T.decl_tensor(B, ty_args=[T.Tensor((m * n,), "float32", data=B)])
- for i in T.serial(m):
- for j in T.serial(n):
- B_buf[i * n + j] = A_buf[i * n + j] + T.float32(1)
-
- else:
-
- @T.function
- def func(A: T.handle("float32"), B: T.handle("float32"), m: T.int64, n: T.int64):
- A_buf = T.decl_tensor(A, ty_args=[T.Tensor((m * n,), "float32", data=A)])
- B_buf = T.decl_tensor(B, ty_args=[T.Tensor((m * n,), "float32", data=B)])
- for i in T.serial(m):
- for j in T.serial(n):
- B_buf[i * n + j] = A_buf[i * n + j] + T.float32(1)
-
- stmt = lower_func_body(func, target_bits)
- assert stmt.loop_var.ty.dtype == target_dtype
- assert stmt.body[0].loop_var.ty.dtype == target_dtype
-
- # const shape
- # i32 -> i32
- check_const(2, 2, 32, "int32")
- # i64 -> i32
- check_const(const(2, dtype="int64"), const(2, dtype="int64"), 32, "int32")
- check_const(const(2**16, dtype="int64"), const(2**16, dtype="int64"), 32, "int64")
- # i32 -> i16
- check_const(2, 2, 16, "int16")
- check_const(2**10, 2**10, 16, "int32")
-
- # symbolic shape
- check_symbolic("int32", "int32", 32, "int32")
- check_symbolic("int64", "int64", 32, "int64")
-
-
-def test_thread_axis():
- # This test uses launch_thread to create RegionStmt nodes
- # and checks the dtype of thread axis variables after narrowing.
- def check_const(m, n, target_bits, target_dtype):
- @T.function
- def func(A: T.Tensor((m * n,), "float32"), B: T.Tensor((m * n,), "float32")):
- bx = T.launch_thread("blockIdx.x", m)
- tx = T.launch_thread("threadIdx.x", n)
- B[bx * n + tx] = A[bx * n + tx] + T.float32(1)
-
- mod = tvm.IRModule.from_expr(func)
- gvar = next(iter(mod.functions.keys()))
- func_narrowed = tvm.tirx.transform.NarrowDataType(target_bits)(mod)[gvar]
- stmt = func_narrowed.body[0]
- assert stmt.body_params[0].ty.dtype == target_dtype
- assert stmt.body[0].body_params[0].ty.dtype == target_dtype
-
- # i32 -> i32
- check_const(2, 32, target_bits=32, target_dtype="int32")
- # i64 -> i32
- check_const(
- const(2, dtype="int64"), const(32, dtype="int64"), target_bits=32, target_dtype="int32"
- )
- check_const(
- const(2**30, dtype="int64"),
- const(32, dtype="int64"),
- target_bits=32,
- target_dtype="int64",
- )
- # i32 -> i16
- check_const(2, 32, target_bits=16, target_dtype="int16")
- check_const(2**14, 32, target_bits=16, target_dtype="int32")
-
-
-def test_multilanes():
- # Test narrowing with vector types. Uses closure capture for m and lanes.
- def check(m, lanes, target_bits, target_dtype):
- vec_dtype = f"float32x{lanes}"
-
- @T.function
- def func(
- A: T.Tensor((m,), vec_dtype),
- B: T.Tensor((m,), vec_dtype),
- ):
- for i in T.serial(m):
- B[i] = A[i] + T.Broadcast(T.float32(1), lanes)
- A[0] = B[1]
-
- mod = tvm.IRModule.from_expr(func)
- gvar = next(iter(mod.functions.keys()))
- func_narrowed = tvm.tirx.transform.NarrowDataType(target_bits)(mod)[gvar]
- stmt = func_narrowed.body
- assert stmt.seq[0].loop_var.ty.dtype == target_dtype
-
- # i32 -> i32
- check(const(2**10, dtype="int32"), 2, target_bits=32, target_dtype="int32")
- # i64 -> i32
- check(const(2**10, dtype="int64"), 2, target_bits=32, target_dtype="int32")
- check(const(2**32, dtype="int64"), 2, target_bits=32, target_dtype="int64")
- # i32 -> i16
- check(const(2**10, dtype="int32"), 2, target_bits=16, target_dtype="int16")
- check(const(2**16, dtype="int32"), 2, target_bits=16, target_dtype="int32")
-
-
-def test_slice():
- # Test narrowing with slice indexing where tensor B has different index ranges.
- def check(m, n, target_bits, target_dtype):
- # The index may overflow in B, while not in A
- @T.function
- def func(
- A: T.Tensor((m * n,), "float32"),
- B: T.Tensor((m * n * 2,), "float32"),
- ):
- for i in T.serial(m):
- for j in T.serial(n):
- A[i * n + j] = B[i * 2 * n + 2 * j] + T.float32(1)
-
- stmt = lower_func_body(func, target_bits)
- assert stmt.loop_var.ty.dtype == target_dtype
- assert stmt.body[0].loop_var.ty.dtype == target_dtype
-
- # The maximum index is (2**15 * 2**15 - 1) * 2 <= 2**31 - 1
- check(const(2**15, "int64"), const(2**15, "int64"), target_bits=32, target_dtype="int32")
- # The maximum index is (2**15 * 2**15 - 1 + 2**15) * 2 > 2**31 - 1
- check(const(2**15, "int64"), const((2**15 + 1), "int64"), target_bits=32, target_dtype="int64")
-
-
-def test_condition():
- @T.function
- def before(A: T.Tensor((128,), "float32"), B: T.Tensor((130,), "float32")):
- for i, j in T.grid(T.int64(2), T.int64(65)):
- if i * T.int64(65) + j >= T.int64(0) and i * T.int64(65) + j < T.int64(128):
- A[i * T.int64(65) + j] = 0.0
- for i, j in T.grid(T.int64(2), T.int64(65)):
- B[i * T.int64(65) + j] = T.if_then_else(
- i * T.int64(65) + j >= T.int64(0) and i * T.int64(65) + j < T.int64(128),
- A[i * T.int64(65) + j],
- 0.0,
- )
-
- @T.function
- def expected_after(A: T.Tensor(128, "float32"), B: T.Tensor(130, "float32")):
- for i, j in T.grid(2, 65):
- if i * 65 + j >= 0 and i * 65 + j < 128:
- A[i * 65 + j] = T.float32(0)
- for i, j in T.grid(2, 65):
- B[i * 65 + j] = T.if_then_else(
- i * 65 + j >= 0 and i * 65 + j < 128, A[i * 65 + j], T.float32(0)
- )
-
- after = tvm.tirx.transform.NarrowDataType(32)(
- tvm.IRModule.from_expr(before.with_attr("global_symbol", "main"))
- )["main"]
- tvm.ir.assert_structural_equal(after, expected_after.with_attr("global_symbol", "main"))
-
-
-def test_block():
- @T.function
- def before(A: T.Tensor((128,), "float32"), B: T.Tensor((128,), "float32")):
- for i in T.serial(0, T.int64(16)):
- for j in T.serial(0, T.int64(8)):
- B[i * T.int64(8) + j] = A[i * T.int64(8) + j] + T.float32(1)
-
- @T.function
- def expected_after(A: T.Tensor((128,), "float32"), B: T.Tensor((128,), "float32")):
- for i in T.serial(0, T.int32(16)):
- for j in T.serial(0, T.int32(8)):
- B[i * T.int32(8) + j] = A[i * T.int32(8) + j] + T.float32(1)
-
- after = tvm.tirx.transform.NarrowDataType(32)(
- tvm.IRModule.from_expr(before.with_attr("global_symbol", "main"))
- )["main"]
- tvm.ir.assert_structural_equal(after, expected_after.with_attr("global_symbol", "main"))
-
-
-def test_avg_pool2d():
- @T.function
- def before(PSUM: T.Tensor((313600,), "int32"), PAVG: T.Tensor((313600,), "int32")):
- for j in T.parallel(T.int64(0), T.int64(280)):
- for i in T.serial(T.int64(0), T.int64(35)):
- for vi in T.vectorized(T.int64(0), T.int64(32)):
- PAVG[(((j * T.int64(1120)) + (i * T.int64(32))) + vi)] = T.cast(
- T.Div(
- T.cast(PSUM[(((j * T.int64(1120)) + (i * T.int64(32))) + vi)], "int64"),
- T.max(
- (
- (
- (
- T.min(
- T.int64(1),
- (T.int64(34) - T.floormod(j, T.int64(35))),
- )
- + T.int64(2)
- )
- - T.max(
- (T.int64(1) - T.floormod(j, T.int64(35))), T.int64(0)
- )
- )
- * (
- (T.min(T.int64(1), (T.int64(34) - i)) + T.int64(2))
- - T.max((T.int64(1) - i), T.int64(0))
- )
- ),
- T.int64(1),
- ),
- ),
- "int32",
- )
-
- @T.function
- def expected_after(PSUM: T.Tensor((313600,), "int32"), PAVG: T.Tensor((313600,), "int32")):
- for j in T.parallel(T.int32(0), T.int32(280)):
- for i in T.serial(T.int32(0), T.int32(35)):
- for vi in T.vectorized(T.int32(0), T.int32(32)):
- PAVG[(((j * T.int32(1120)) + (i * T.int32(32))) + vi)] = T.Div(
- PSUM[(((j * T.int32(1120)) + (i * T.int32(32))) + vi)],
- (
- (
- (
- T.min(T.int32(1), (T.int32(34) - T.floormod(j, T.int32(35))))
- + T.int32(2)
- )
- - T.max((T.int32(1) - T.floormod(j, T.int32(35))), T.int32(0))
- )
- * (
- (T.min(T.int32(1), (T.int32(34) - i)) + T.int32(2))
- - T.max((T.int32(1) - i), T.int32(0))
- )
- ),
- )
-
- after = tvm.tirx.transform.NarrowDataType(32)(
- tvm.IRModule.from_expr(before.with_attr("global_symbol", "main"))
- )
- after = tvm.tirx.transform.StmtSimplify()(after)
- tvm.ir.assert_structural_equal(after["main"], expected_after.with_attr("global_symbol", "main"))
-
-
-def test_narrow_i64_valued_tensorload_index_to_i32():
- @T.function
- def before(A: T.Tensor((16,), "int64")):
- for i in range(T.int64(15)):
- A[i + T.int64(1)] = A[i] + T.int64(1)
-
- @T.function
- def expect(A: T.Tensor((16,), "int64")):
- for i in range(15):
- A[i + 1] = A[i] + T.int64(1)
-
- after = tvm.tirx.transform.NarrowDataType(32)(
- tvm.IRModule.from_expr(before.with_attr("global_symbol", "main"))
- )["main"]
- tvm.ir.assert_structural_equal(after, expect.with_attr("global_symbol", "main"))
-
-
-if __name__ == "__main__":
- tvm.testing.main()
diff --git a/tests/python/tirx-transform/test_tir_transform_remove_assume.py b/tests/python/tirx-transform/test_tir_transform_remove_assume.py
deleted file mode 100644
index 86bb1d51d4..0000000000
--- a/tests/python/tirx-transform/test_tir_transform_remove_assume.py
+++ /dev/null
@@ -1,69 +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.
-
-import tvm
-import tvm.testing
-from tvm.script import ir as I
-from tvm.script import tirx as T
-
-
-def test_remove_assume():
- """Remove any instance of T.assume"""
-
- @I.ir_module
- class Before:
- @T.function
- def main(A: T.Tensor(1, "int32")):
- T.evaluate(T.assume(A[0] == 5))
- A[0] = 10
-
- @I.ir_module
- class Expected:
- @T.function
- def main(A: T.Tensor(1, "int32")):
- A[0] = 10
-
- After = tvm.tirx.transform.RemoveAssume()(Before)
- tvm.ir.assert_structural_equal(After, Expected)
-
-
-def test_remove_assume_loop():
- """Loops containing only T.assume should be removed"""
-
- @I.ir_module
- class Before:
- @T.function
- def main(A: T.Tensor(16, "int32")):
- for i in T.serial(16):
- T.evaluate(T.assume(A[i] == 0))
-
- for i in T.serial(16):
- A[i] = 10
-
- @I.ir_module
- class Expected:
- @T.function
- def main(A: T.Tensor(16, "int32")):
- for i in T.serial(16):
- A[i] = 10
-
- After = tvm.tirx.transform.RemoveAssume()(Before)
- tvm.ir.assert_structural_equal(After, Expected)
-
-
-if __name__ == "__main__":
- tvm.testing.main()
diff --git a/tests/python/tirx-transform/test_tir_transform_storage_rewrite.py b/tests/python/tirx-transform/test_tir_transform_storage_rewrite.py
deleted file mode 100644
index 3a906aa62d..0000000000
--- a/tests/python/tirx-transform/test_tir_transform_storage_rewrite.py
+++ /dev/null
@@ -1,755 +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: F401, F841
-import sys
-import uuid
-
-import pytest
-import tvm_ffi
-
-import tvm
-import tvm.testing
-from tvm.script import ir as I
-from tvm.script import tirx as T
-
-for _scope_name in ("local.L0A", "local.L0A2"):
- _scope = tvm.runtime.MemoryScope.register(_scope_name)
- tvm.runtime.MemoryScope.def_attr("runtime.memory_alloc_property")[_scope] = 1
-
-
-def _is_tensor_binding(node, *op_names):
- return (
- isinstance(node, tvm.ir.Bind)
- and isinstance(node.value, tvm.ir.Call)
- and isinstance(node.value.op, tvm.ir.Op)
- and node.value.op.name in op_names
- )
-
-
-def _scope_with_property(value):
- scope = tvm.runtime.MemoryScope.register("test.storage_rewrite." + uuid.uuid4().hex)
- tvm.runtime.MemoryScope.def_attr("runtime.memory_alloc_property")[scope] = value
- return scope
-
-
-def _rewritten_allocations(func):
- result = tvm.tirx.transform.StorageRewrite()(tvm.IRModule.from_expr(func))
- allocations = []
-
- def collect(node):
- if _is_tensor_binding(node, "tirx.alloc_tensor"):
- allocations.append(node)
-
- tvm_ffi.structural_walk(result["func"].body, collect)
- return allocations
-
-
-@pytest.mark.parametrize("property_value,expected_allocations", [(0, 1), (1, 1), (2, 2)])
-def test_allocation_property_controls_reuse(property_value, expected_allocations):
- scope = _scope_with_property(property_value)
-
- @T.function
- def func(output: T.Tensor((2,), "float32")):
- A = T.alloc_tensor(ty_args=[T.Tensor((128,), "float32", scope)])
- for i in T.serial(128):
- A[i] = T.float32(i)
- output[0] = A[127]
- B = T.alloc_tensor(ty_args=[T.Tensor((128,), "float32", scope)])
- for i in T.serial(128):
- B[i] = T.float32(i + 1)
- output[1] = B[127]
-
- assert len(_rewritten_allocations(func)) == expected_allocations
-
-
-@pytest.mark.parametrize("property_value,expected_allocations", [(0, 2), (1, 1), (2, 2)])
-def test_allocation_property_controls_packing(property_value, expected_allocations):
- scope = _scope_with_property(property_value)
-
- @T.function
- def func(output: T.Tensor((1,), "float32")):
- A = T.alloc_tensor(ty_args=[T.Tensor((128,), "float32", scope)])
- B = T.alloc_tensor(ty_args=[T.Tensor((128,), "float32", scope)])
- for i in T.serial(128):
- A[i] = T.float32(i)
- B[i] = T.float32(i + 1)
- output[0] = A[127] + B[127]
-
- assert len(_rewritten_allocations(func)) == expected_allocations
-
-
-@pytest.mark.parametrize("scope_name", [None, "shared.dyn"])
-def test_single_allocation_symbolic_layout_is_deferred(scope_name):
- scope = tvm.runtime.MemoryScope(scope_name) if scope_name else _scope_with_property(1)
-
- @T.function
- def func(n: T.int32, output: T.Tensor((1,), "float32")):
- A = T.alloc_tensor(ty_args=[T.Tensor((n,), "float32", scope)])
- B = T.alloc_tensor(ty_args=[T.Tensor((n,), "float32", scope)])
- for i in T.serial(n):
- A[i] = T.float32(i)
- B[i] = T.float32(i + 1)
- output[0] = A[n - 1] + B[n - 1]
-
- allocations = _rewritten_allocations(func)
- assert len(allocations) == 2
- for allocation in allocations:
- tvm.ir.assert_structural_equal(allocation.var.shape[0], func.params[0])
- assert allocation.var.scope().same_as(scope)
-
-
-def test_single_allocation_preserves_exact_metadata():
- scope = _scope_with_property(1)
-
- @T.function
- def func(output: T.Tensor((1,), "float32")):
- A = T.alloc_tensor(
- ty_args=[T.Tensor((128,), "float32", scope)], annotations={"test.constraint": 7}
- )
- B = T.alloc_tensor(ty_args=[T.Tensor((128,), "float32", scope)])
- for i in T.serial(128):
- A[i] = T.float32(i)
- B[i] = T.float32(i + 1)
- output[0] = A[127] + B[127]
-
- allocations = _rewritten_allocations(func)
- assert len(allocations) == 2
- constrained = [alloc for alloc in allocations if "test.constraint" in alloc.value.attrs]
- assert len(constrained) == 1
- assert int(constrained[0].value.attrs["test.constraint"]) == 7
- assert int(constrained[0].var.shape[0]) == 128
-
-
-@pytest.mark.parametrize("property_value", [0, 1])
-@pytest.mark.parametrize("stride,offset", [(2, 0), (1, 4)])
-def test_allocation_constraints_prevent_lifetime_reuse(property_value, stride, offset):
- scope = _scope_with_property(property_value)
-
- @T.function
- def func(output: T.Tensor((2,), "float32")):
- A = T.alloc_tensor(
- ty_args=[T.Tensor((128,), "float32", scope, strides=[stride], elem_offset=offset)]
- )
- for i in T.serial(128):
- A[i] = T.float32(i)
- output[0] = A[127]
- B = T.alloc_tensor(ty_args=[T.Tensor((128,), "float32", scope)])
- for i in T.serial(128):
- B[i] = T.float32(i + 1)
- output[1] = B[127]
-
- allocations = _rewritten_allocations(func)
- assert len(allocations) == 2
- constrained = next(alloc for alloc in allocations if alloc.var.name == "A")
- assert int(constrained.var.ty.strides[0]) == stride
- assert int(constrained.var.ty.elem_offset) == offset
-
-
-@pytest.mark.parametrize("property_value", [0, 1])
-def test_lifetime_reuse_preserves_alignment(property_value):
- scope = _scope_with_property(property_value)
-
- @T.function
- def func(output: T.Tensor((2,), "float32")):
- A = T.alloc_tensor(ty_args=[T.Tensor((128,), "float32", scope, align=256)])
- for i in T.serial(128):
- A[i] = T.float32(i)
- output[0] = A[127]
- B = T.alloc_tensor(ty_args=[T.Tensor((256,), "float32", scope, align=256)])
- for i in T.serial(256):
- B[i] = T.float32(i + 1)
- output[1] = B[255]
-
- allocations = _rewritten_allocations(func)
- assert len(allocations) == 1
- assert allocations[0].var.ty.data_alignment == 256
-
-
-def test_packing_preserves_different_alignment_contracts():
- scope = _scope_with_property(1)
-
- @T.function
- def func(output: T.Tensor((2,), "float32")):
- A = T.alloc_tensor(ty_args=[T.Tensor((17,), "float32", scope, align=64)])
- for i in T.serial(17):
- A[i] = T.float32(i)
- output[0] = A[16]
- B = T.alloc_tensor(
- ty_args=[T.Tensor((17,), "float32", scope, elem_offset=0, align=256, offset_factor=128)]
- )
- for i in T.serial(17):
- B[i] = T.float32(i + 1)
- output[1] = B[16]
-
- allocations = _rewritten_allocations(func)
- assert len(allocations) == 1
- assert allocations[0].var.ty.data_alignment == 512
- assert int(allocations[0].var.shape[0]) == 256
-
-
-def test_packing_widens_large_combined_extent():
- scope = _scope_with_property(1)
-
- @T.function
- def func(output: T.Tensor((1,), "float32")):
- A = T.alloc_tensor(ty_args=[T.Tensor((1073741824,), "float32", scope)])
- B = T.alloc_tensor(ty_args=[T.Tensor((1073741824,), "float32", scope)])
- A[0] = T.float32(1)
- B[0] = T.float32(2)
- output[0] = A[0] + B[0]
-
- allocations = _rewritten_allocations(func)
- assert len(allocations) == 1
- assert int(allocations[0].var.shape[0]) == 2147483648
- assert allocations[0].var.shape[0].ty == tvm.ir.PrimType("int64")
-
-
-def test_alloc_seq():
- scope_tb = "local.L0A"
-
- @T.function
- def func(n: T.int32):
- for i in T.serial(n):
- for j in range(10):
- A = T.alloc_tensor(ty_args=[T.Tensor((200,), "float32", scope=scope_tb)])
- A[j] = T.float32(1.2)
- for j in range(10):
- B = T.alloc_tensor(ty_args=[T.Tensor((200,), "float32", scope=scope_tb)])
- B[j] = T.float32(1.3)
-
- mod = tvm.IRModule.from_expr(func)
- body = tvm.tirx.transform.StorageRewrite()(mod)["func"].body
-
- num_alloc = [0]
-
- def verify(n):
- if _is_tensor_binding(n, "tirx.alloc_tensor"):
- num_alloc[0] += 1
- assert n.var.ty.shape[0].value == 200
-
- tvm_ffi.structural_walk(body, verify)
- assert num_alloc[0] == 1
-
-
-def test_alloc_different_dtypes():
- # Test cross-loop tensor access with tensors allocated in parent scope
- def make_mod(dtype_list, length):
- assert len(dtype_list) == 4
-
- @T.function
- def func():
- # Allocate all tensors in parent scope (before any loops)
- A = T.alloc_tensor(ty_args=[T.Tensor((length,), dtype_list[0], scope="local.L0A")])
- B = T.alloc_tensor(ty_args=[T.Tensor((length,), dtype_list[1], scope="local.L0A")])
- C = T.alloc_tensor(ty_args=[T.Tensor((length,), dtype_list[2], scope="local.L0A")])
- D = T.alloc_tensor(ty_args=[T.Tensor((length,), dtype_list[3], scope="local.L0A")])
- E = T.alloc_tensor(ty_args=[T.Tensor((length,), "int8", scope="local.L0A")])
-
- for j in range(length):
- A[j] = T.Cast(1, dtype_list[0])
- for j in range(length):
- B[j] = T.Cast(1, dtype_list[1])
- for j in range(length):
- C[j] = T.Cast(1, dtype_list[2])
- for j in range(length):
- D[j] = T.Cast(1, dtype_list[3])
- for j in range(length):
- E[j] = (
- T.Cast(A[j], "int8")
- + T.Cast(B[j], "int8")
- + T.Cast(C[j], "int8")
- + T.Cast(D[j], "int8")
- )
-
- return tvm.IRModule.from_expr(func)
-
- def dtype_bit_len(dtype):
- index = 0
- for i in dtype:
- if i.isdigit():
- break
- index += 1
- return int(dtype[index:])
-
- def offset_generater(dtype_list, length):
- dtype_len_list = [dtype_bit_len(i) for i in dtype_list]
- base_len = dtype_len_list[0]
- return sum([i * length / base_len for i in dtype_len_list])
-
- def dtype_test(dtype_list, length):
- def verify(n):
- if _is_tensor_binding(n, "tirx.alloc_tensor"):
- assert n.var.ty.shape[0].value == offset
-
- mod = make_mod(dtype_list, length)
- offset = offset_generater(dtype_list, length)
-
- body = tvm.tirx.transform.StorageRewrite()(mod)["func"].body
- tvm_ffi.structural_walk(body, verify)
-
- length = 1024
- dtype_list = ["float16", "int32", "uint16", "int8"]
- dtype_test(dtype_list, length)
-
- dtype_list = ["float32", "int32", "uint16", "int8"]
- dtype_test(dtype_list, length)
-
- dtype_list = ["float64", "int32", "uint16", "int8"]
- dtype_test(dtype_list, length)
-
- dtype_list = ["int8", "int32", "uint16", "uint8"]
- dtype_test(dtype_list, length)
-
-
-def test_address_of():
- # In this test, the storage rewrite pass is allowed to
- # combine tensors B and D, but not C
- @T.function
- def before(A: T.Tensor(8, "float32"), E: T.Tensor(8, "float32")):
- B = T.alloc_tensor(ty_args=[T.Tensor((8,), "float32")])
- for i in range(8):
- B[i] = (
- T.call_extern("deref", T.address_of(A[i]), ty="float32")
- + T.call_extern("deref", T.address_of(A[0]), ty="float32")
- + T.float32(1)
- )
- C = T.alloc_tensor(ty_args=[T.Tensor((8,), "float32")])
- for i in range(8):
- C[i] = (
- T.call_extern("deref", T.address_of(B[i]), ty="float32")
- + T.call_extern("deref", T.address_of(B[0]), ty="float32")
- + T.float32(2)
- )
- D = T.alloc_tensor(ty_args=[T.Tensor((8,), "float32")])
- for i in range(8):
- D[i] = (
- T.call_extern("deref", T.address_of(C[i]), ty="float32")
- + T.call_extern("deref", T.address_of(C[0]), ty="float32")
- + T.float32(2)
- )
- for i in range(8):
- E[i] = (
- T.call_extern("deref", T.address_of(D[i]), ty="float32")
- + T.call_extern("deref", T.address_of(D[0]), ty="float32")
- + T.float32(3)
- )
-
- def verify(n):
- if _is_tensor_binding(n, "tirx.alloc_tensor"):
- total_alloc[0] += n.var.ty.shape[0].value
-
- total_alloc = [0]
- mod = tvm.IRModule.from_expr(before.with_attr("global_symbol", "main"))
- tvm_ffi.structural_walk(mod["main"].body, verify)
- assert total_alloc[0] == 24
-
- total_alloc[0] = 0
- mod = tvm.tirx.transform.StorageRewrite()(mod)
- tvm_ffi.structural_walk(mod["main"].body, verify)
- assert total_alloc[0] == 16
-
-
-def test_parallel_alloc():
- @T.function
- def func1(n: T.int32):
- for i in T.parallel(n):
- for j in range(10):
- A = T.alloc_tensor(ty_args=[T.Tensor((n,), "float32")])
- A[j] = A[j] + T.float32(2)
-
- mod = tvm.IRModule.from_expr(func1)
- body = tvm.tirx.transform.StorageRewrite()(mod)["func1"]
-
- # With flat AllocTensor, the for body is a SeqStmt; first element is AllocTensor
- assert _is_tensor_binding(body.body[0].body[0], "tirx.alloc_tensor")
-
- @T.function
- def func2(n: T.int32):
- for t in T.serial(n):
- with T.parallel_launch():
- for i in T.parallel(n):
- for j in range(10):
- A = T.alloc_tensor(ty_args=[T.Tensor((n,), "float32")])
- A[j] = A[j] + T.float32(2)
-
- mod = tvm.IRModule.from_expr(func2)
- body = tvm.tirx.transform.StorageRewrite()(mod)["func2"]
-
- # The launch body owns the allocation executed by each worker.
- assert _is_tensor_binding(body.body[0].body[0].body[0], "tirx.alloc_tensor")
-
-
-def test_while_alloc():
- @T.function
- def func_parallel(n: T.int32):
- for i in T.parallel(n):
- j = T.alloc_tensor(ty_args=[T.Tensor((1,), "int32")])
- j[0] = 0
- while j[0] < 10:
- A = T.alloc_tensor(ty_args=[T.Tensor((n,), "float32")])
- A[j[0]] = A[j[0]] + T.float32(2)
- j[0] = j[0] + j[0] + 1
-
- @T.function
- def func_serial(n: T.int32):
- for i in T.serial(n):
- j = T.alloc_tensor(ty_args=[T.Tensor((1,), "int32")])
- j[0] = 0
- while j[0] < 10:
- A = T.alloc_tensor(ty_args=[T.Tensor((n,), "float32")])
- A[j[0]] = A[j[0]] + T.float32(2)
- j[0] = j[0] + j[0] + 1
-
- mod = tvm.IRModule.from_expr(func_parallel)
- # parallel (i, 0, n) {
- # allocate j[int32 * 1]
- # j[0] = 0
- # while((j[0] < 10)){
- # // attr [A] storage_scope = "global"
- # allocate A[float32 * n]
- # A[j[0]] = (A[j[0]] + 2f)
- # j[0] = (j[0] + (j[0] + 1))
- # }
- # }
- body = tvm.tirx.transform.StorageRewrite()(mod)["func_parallel"]
- # Navigate to inside the for loop, then check that allocations exist
- # The structure with DeclTensor is:
- # parallel (i, 0, n) { DeclTensor(j, DeclTensor(A, ...)) }
- # or with Allocate+DeclTensor pairs
- inner = body.body[0].body # inside For
- # Skip DeclTensor nodes to find Allocate
- num_alloc = [0]
-
- def count_alloc(n):
- if _is_tensor_binding(n, "tirx.alloc_tensor"):
- num_alloc[0] += 1
-
- tvm_ffi.structural_walk(inner, count_alloc)
- assert num_alloc[0] == 2 # j and A allocations
-
- mod = tvm.IRModule.from_expr(func_serial)
- body = tvm.tirx.transform.StorageRewrite()(mod)["func_serial"]
- num_alloc[0] = 0
- tvm_ffi.structural_walk(body.body, count_alloc)
- assert num_alloc[0] == 2 # j and A allocations
-
-
-@pytest.mark.parametrize("alignment,expected_size", [(4, 500), (64, 528)])
-def test_alloc_seq_type(alignment, expected_size):
- @T.function
- def func(n: T.int32):
- for i in T.serial(n):
- for j in range(10):
- A = T.alloc_tensor(
- ty_args=[T.Tensor((200,), "float32", scope="local.L0A", align=alignment)]
- )
- A1 = T.alloc_tensor(
- ty_args=[T.Tensor((200,), "float32", scope="local.L0A", align=alignment)]
- )
- A[j] = T.float32(1.2)
- A1[j] = T.float32(1.3)
- B = T.alloc_tensor(
- ty_args=[T.Tensor((200,), "int16", scope="local.L0A", align=alignment)]
- )
- B[j] = T.int16(1)
- C = T.alloc_tensor(
- ty_args=[T.Tensor((200,), "int16", scope="local.L0A", align=alignment)]
- )
- C[j] = T.int16(1)
- D = T.alloc_tensor(
- ty_args=[T.Tensor((200,), "int16", scope="local.L0A", align=alignment)]
- )
- D[j] = B[j] + C[j]
- A2 = T.alloc_tensor(
- ty_args=[T.Tensor((200,), "float32", scope="local.L0A", align=alignment)]
- )
- A2[j] = A[j]
-
- mod = tvm.IRModule.from_expr(func)
- body = tvm.tirx.transform.StorageRewrite()(mod)["func"].body
-
- num_alloc = [0]
-
- def verify(n):
- if _is_tensor_binding(n, "tirx.alloc_tensor"):
- num_alloc[0] += 1
- # Each packed region starts at a multiple of its declared byte alignment.
- assert n.var.ty.shape[0].value == expected_size
- assert n.var.ty.data_alignment == alignment
-
- tvm_ffi.structural_walk(body, verify)
- assert num_alloc[0] == 1
-
-
-def test_alloc_seq_type2():
- scope_tb = "local.L0A2"
-
- @T.function
- def func(n: T.int32):
- for i in T.serial(n):
- for j in range(10):
- A = T.alloc_tensor(ty_args=[T.Tensor((200,), "float32", scope=scope_tb)])
- A[j] = T.float32(1.2)
- for j in range(20):
- B = T.alloc_tensor(ty_args=[T.Tensor((400,), "int16", scope=scope_tb)])
- B[j] = T.int16(1)
- for j in range(10):
- C = T.alloc_tensor(ty_args=[T.Tensor((200,), "float32", scope=scope_tb)])
- C[j] = T.float32(1.2)
-
- mod = tvm.IRModule.from_expr(func)
- body = tvm.tirx.transform.StorageRewrite()(mod)["func"].body
-
- num_alloc = [0]
-
- def verify(n):
- if _is_tensor_binding(n, "tirx.alloc_tensor"):
- num_alloc[0] += 1
- assert n.var.ty.shape[0].value == 200
-
- tvm_ffi.structural_walk(body, verify)
- assert num_alloc[0] == 1
-
-
-@pytest.mark.parametrize("alignment,expected_size", [(4, 800), (64, 832)])
-def test_reuse_small_tensor(alignment, expected_size):
- @T.function
- def func(n: T.int32):
- for i in T.serial(n):
- for j in range(10):
- A = T.alloc_tensor(
- ty_args=[T.Tensor((200,), "int16", scope="local.L0A", align=alignment)]
- )
- A[j] = T.int16(1)
- B = T.alloc_tensor(
- ty_args=[T.Tensor((200,), "int16", scope="local.L0A", align=alignment)]
- )
- B[j] = T.int16(1)
- B1 = T.alloc_tensor(
- ty_args=[T.Tensor((200,), "int16", scope="local.L0A", align=alignment)]
- )
- B1[j] = A[j] + B[j]
- C = T.alloc_tensor(
- ty_args=[T.Tensor((400,), "int16", scope="local.L0A", align=alignment)]
- )
- C[j] = T.int16(1)
- D = T.alloc_tensor(
- ty_args=[T.Tensor((400,), "int16", scope="local.L0A", align=alignment)]
- )
- D[j] = T.int16(1)
- E = T.alloc_tensor(
- ty_args=[T.Tensor((400,), "int16", scope="local.L0A", align=alignment)]
- )
- E[j] = C[j]
-
- mod = tvm.IRModule.from_expr(func)
- body = tvm.tirx.transform.StorageRewrite()(mod)["func"].body
-
- num_alloc = [0]
-
- def verify(n):
- if _is_tensor_binding(n, "tirx.alloc_tensor"):
- num_alloc[0] += 1
- # Reused entries hold 400 int16 values each, padded to their alignment.
- assert n.var.ty.shape[0].value == expected_size
- assert n.var.ty.data_alignment == alignment
-
- tvm_ffi.structural_walk(body, verify)
- assert num_alloc[0] == 1
-
-
-def test_access_in_let_value():
- @T.function
- def func(A: T.Tensor((8,), "float32")):
- for i in range(8):
- B = T.alloc_tensor(ty_args=[T.Tensor((1,), "float32")])
- B[0] = 3.14
- x: T.let[T.float32] = T.exp(B[0])
- A[i] = (x + 1.0) / (x - 1.0)
-
- @T.function
- def func_rewritten(A: T.Tensor((8,), "float32")) -> None:
- B = T.alloc_tensor(ty_args=[T.Tensor((1,), "float32")])
- for i in range(8):
- B[0] = 3.14
- x: T.let[T.float32] = T.exp(B[0])
- A[i] = (x + 1.0) / (x - 1.0)
-
- mod = tvm.tirx.transform.StorageRewrite()(
- tvm.IRModule.from_expr(func.with_attr("global_symbol", "main"))
- )
- tvm.ir.assert_structural_equal(mod["main"], func_rewritten.with_attr("global_symbol", "main"))
-
-
-def test_decl_tensor_is_not_vectorized():
- """StorageRewrite leaves explicit DeclTensor views unchanged.
-
- Vectorization of DeclTensor views was dropped because the rewritten result
- violates the immutable TensorVar binding invariants.
- """
-
- @I.ir_module
- class Before:
- @T.function
- def main() -> None:
- A_data: T.let[T.handle("int32")] = T.call_extern("dummy_func", ty=T.handle("int32").ty)
- A = T.decl_tensor(A_data, ty_args=[T.Tensor([8], "int32", data=A_data)])
- A[T.ramp(0, 1, 8)] = T.broadcast(42, 8)
-
- After = tvm.tirx.transform.StorageRewrite()(Before)
- tvm.ir.assert_structural_equal(After, Before)
-
-
-def test_rewrite_decl_tensor():
- """A DeclTensor node may appear in StorageRewrite's input"""
-
- @I.ir_module
- class Before:
- @T.function
- def main(A: T.Tensor(16, "float32"), D: T.Tensor(16, "float32")):
- B = T.alloc_tensor(ty_args=[T.Tensor(16, "float32")])
- C = T.alloc_tensor(ty_args=[T.Tensor(16, "float32")])
-
- for i in range(16):
- B[i] = A[i]
-
- for i in range(16):
- C[i] = 2.0 * B[i]
-
- for i in range(16):
- D[i] = C[i]
-
- @I.ir_module
- class Expected:
- @T.function
- def main(A: T.Tensor(16, "float32"), D: T.Tensor(16, "float32")):
- B = T.alloc_tensor(ty_args=[T.Tensor(16, "float32")])
- C = T.decl_tensor(B.data, ty_args=[T.Tensor(16, "float32", data=B.data)])
-
- for i in range(16):
- B[i] = A[i]
-
- for i in range(16):
- C[i] = 2.0 * B[i]
-
- for i in range(16):
- D[i] = C[i]
-
- After = tvm.tirx.transform.StorageRewrite()(Before)
- assert tvm.tirx.analysis.verify_well_formed(After)
- tvm.ir.assert_structural_equal(After, Expected)
-
-
-def test_decl_tensor_alias_chain_uses_flat_root():
- """StorageRewrite resolves every alias in a chain to the same flat root."""
-
- @I.ir_module
- class Before:
- @T.function
- def main(D: T.Tensor(1, "float32")):
- A = T.alloc_tensor(ty_args=[T.Tensor(16, "float32")])
- B = T.decl_tensor(A.data, ty_args=[T.Tensor(16, "float32", data=A.data)])
- C = T.decl_tensor(B.data, ty_args=[T.Tensor(16, "float32", data=B.data)])
- A[0] = 1.0
- D[0] = C[0]
-
- @I.ir_module
- class Expected:
- @T.function
- def main(D: T.Tensor(1, "float32")):
- A = T.alloc_tensor(ty_args=[T.Tensor(16, "float32")])
- B = T.decl_tensor(A.data, ty_args=[T.Tensor(16, "float32", data=A.data)])
- C = T.decl_tensor(A.data, ty_args=[T.Tensor(16, "float32", data=A.data)])
- A[0] = 1.0
- D[0] = C[0]
-
- After = tvm.tirx.transform.StorageRewrite()(Before)
- assert tvm.tirx.analysis.verify_well_formed(After)
- tvm.ir.assert_structural_equal(After, Expected)
-
-
-def test_decl_tensor_alias_extends_source_lifetime():
- """An access through an alias prevents reuse of its source allocation."""
-
- @T.function
- def func(D: T.Tensor(1, "float32")):
- A = T.alloc_tensor(ty_args=[T.Tensor(16, "float32")])
- B = T.decl_tensor(A.data, ty_args=[T.Tensor(16, "float32", data=A.data)])
- A[0] = 1.0
-
- C = T.alloc_tensor(ty_args=[T.Tensor(16, "float32")])
- C[0] = 2.0
- D[0] = B[0] + C[0]
-
- after = tvm.tirx.transform.StorageRewrite()(tvm.IRModule.from_expr(func))["func"]
- allocations = []
- tvm_ffi.structural_walk(
- after.body,
- lambda node: (
- allocations.append(node) if _is_tensor_binding(node, "tirx.alloc_tensor") else None
- ),
- )
- assert len(allocations) == 2
-
-
-def test_no_orphaned_decl_tensor():
- """A DeclTensor of an unused Allocate should be removed
-
- StorageRewrite removes any allocations that are unused. When it
- does so, any DeclTensor that refers to that allocation should also
- be removed.
- """
-
- @I.ir_module
- class Before:
- @T.function
- def main(A: T.Tensor(16, "float32"), D: T.Tensor(16, "float32")):
- B = T.alloc_tensor(ty_args=[T.Tensor(16, "float32")])
- C = T.alloc_tensor(ty_args=[T.Tensor(16, "float32")])
- Unused = T.alloc_tensor(ty_args=[T.Tensor(16, "float32")])
-
- for i in range(16):
- B[i] = A[i]
-
- for i in range(16):
- C[i] = 2.0 * B[i]
-
- for i in range(16):
- D[i] = C[i]
-
- @I.ir_module
- class Expected:
- @T.function
- def main(A: T.Tensor(16, "float32"), D: T.Tensor(16, "float32")):
- B = T.alloc_tensor(ty_args=[T.Tensor(16, "float32")])
- C = T.decl_tensor(B.data, ty_args=[T.Tensor(16, "float32", data=B.data)])
-
- for i in range(16):
- B[i] = A[i]
-
- for i in range(16):
- C[i] = 2.0 * B[i]
-
- for i in range(16):
- D[i] = C[i]
-
- After = tvm.tirx.transform.StorageRewrite()(Before)
- tvm.ir.assert_structural_equal(After, Expected)
-
-
-if __name__ == "__main__":
- tvm.testing.main()
diff --git a/tests/python/tirx-transform/test_tir_transform_unroll_loop.py b/tests/python/tirx-transform/test_tir_transform_unroll_loop.py
deleted file mode 100644
index 9c3b0cfdf6..0000000000
--- a/tests/python/tirx-transform/test_tir_transform_unroll_loop.py
+++ /dev/null
@@ -1,159 +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.
-from __future__ import annotations
-
-import tvm
-from tvm.script import ir as I
-from tvm.script import tirx as T
-
-
-def test_unroll_loop():
- @I.ir_module
- class Module:
- @T.function
- def main(Ab: T.Tensor((n,), "int64"), n: T.int64): # noqa: F821
- for i in T.serial(n, n + 2):
- for j in T.unroll(8):
- Ab[j + 1] = Ab[i] + T.int64(1)
-
- mod = Module
- stmt = mod["main"].body[0]
-
- assert isinstance(stmt, tvm.ir.For)
-
- with tvm.ir.transform.PassContext(config={"tirx.UnrollLoop": {"auto_max_step": 16}}):
- ret = tvm.tirx.transform.UnrollLoop()(mod)["main"].body
- assert not any(isinstance(stmt, tvm.ir.For) for stmt in ret)
-
- with tvm.ir.transform.PassContext(config={"tirx.UnrollLoop": {"auto_max_step": 15}}):
- ret = tvm.tirx.transform.UnrollLoop()(mod)["main"].body
- assert isinstance(ret[0], tvm.ir.For)
-
- with tvm.ir.transform.PassContext(
- config={"tirx.UnrollLoop": {"auto_max_step": 16, "explicit_unroll": False}}
- ):
- ret = tvm.tirx.transform.UnrollLoop()(mod)["main"].body
- assert isinstance(ret[0], tvm.ir.For)
- assert ret[0].kind == tvm.ir.ForKind.UNROLLED
-
- @I.ir_module
- class ModuleWithPolicy:
- @T.function
- def main(Ab: T.Tensor((n,), "int64"), n: T.int64): # noqa: F821
- for i in T.serial(n, n + 2, annotations={"auto_unroll_max_step": 16}):
- for j in T.unroll(8):
- Ab[j + 1] = Ab[i] + T.int64(1)
- for i in T.serial(n, n + 2):
- for j in T.unroll(8):
- Ab[j + 1] = Ab[i] + T.int64(1)
-
- with tvm.ir.transform.PassContext(
- config={"tirx.UnrollLoop": {"auto_max_depth": 8, "explicit_unroll": False}}
- ):
- ret = tvm.tirx.transform.UnrollLoop()(ModuleWithPolicy)["main"].body
- assert isinstance(ret[0], tvm.ir.For)
- assert ret[0].kind == tvm.ir.ForKind.UNROLLED
- assert isinstance(ret[1], tvm.ir.For)
- assert ret[1].kind != tvm.ir.ForKind.UNROLLED
-
-
-def test_unroll_fake_loop():
- @I.ir_module
- class Module:
- @T.function
- def main(Ab: T.Tensor((n,), "int32"), n: T.int64): # noqa: F821
- for i in T.serial(1):
- Ab[i * 2] = 3
- for j in T.serial(10):
- Ab[j + 1] = Ab[i] + 1
-
- with tvm.ir.transform.PassContext(
- config={
- "tirx.UnrollLoop": {"auto_max_depth": 8, "auto_max_extent": 1, "explicit_unroll": False}
- }
- ):
- ret = tvm.tirx.transform.UnrollLoop()(Module)["main"].body
- assert isinstance(ret[0], tvm.ir.TensorStore)
-
-
-def test_unroll_allocations():
- @I.ir_module
- class Before:
- @T.function
- def main():
- for i in T.unroll(2):
- buf = T.alloc_tensor(ty_args=[T.Tensor([16], "float32")])
- buf[0] = 0.0
-
- @I.ir_module
- class Expected:
- @T.function
- def main():
- buf1 = T.alloc_tensor(ty_args=[T.Tensor([16], "float32")])
- buf1[0] = 0.0
- buf2 = T.alloc_tensor(ty_args=[T.Tensor([16], "float32")])
- buf2[0] = 0.0
-
- after = tvm.tirx.transform.UnrollLoop()(Before)
-
- tvm.ir.assert_structural_equal(after, Expected)
-
-
-def test_unroll_local_access():
- @I.ir_module
- class Before:
- @T.function
- def main(B: T.Tensor((64,), "float32")):
- with T.launch_thread("blockIdx.x", 4) as bx: # noqa: F841
- with T.launch_thread("threadIdx.x", 4) as tx: # noqa: F841
- A_local = T.alloc_tensor(ty_args=[T.Tensor((4,), "float32", scope="local")])
- for i in T.serial(4):
- A_local[i] = T.float32(i)
-
- @I.ir_module
- class Expected:
- @T.function
- def main(B: T.Tensor((64,), "float32")):
- with T.launch_thread("blockIdx.x", 4) as bx: # noqa: F841
- with T.launch_thread("threadIdx.x", 4) as tx: # noqa: F841
- A_local = T.alloc_tensor(ty_args=[T.Tensor((4,), "float32", scope="local")])
- A_local[0] = T.float32(0)
- A_local[1] = T.float32(1)
- A_local[2] = T.float32(2)
- A_local[3] = T.float32(3)
-
- with tvm.ir.transform.PassContext(
- config={
- "tirx.UnrollLoop": {
- "auto_max_depth": 0,
- "auto_max_extent": 1,
- "explicit_unroll": True,
- "unroll_local_access": True,
- }
- }
- ):
- after = tvm.tirx.transform.UnrollLoop()(Before)
- after = tvm.tirx.transform.StmtSimplify()(after)
-
- tvm.ir.assert_structural_equal(after, Expected)
-
-
-if __name__ == "__main__":
- test_unroll_local_access()
- test_unroll_loop()
- test_unroll_fake_loop()
- test_unroll_allocations()
diff --git a/tests/python/tirx-analysis/test_tir_analysis_verify_well_formed.py b/tests/python/tirx/analysis/test_tirx_analysis_check_well_defined.py
similarity index 100%
rename from tests/python/tirx-analysis/test_tir_analysis_verify_well_formed.py
rename to tests/python/tirx/analysis/test_tirx_analysis_check_well_defined.py
diff --git a/tests/python/tirx-analysis/test_tir_analysis_undefined_vars.py b/tests/python/tirx/analysis/test_tirx_analysis_undefined_vars.py
similarity index 100%
rename from tests/python/tirx-analysis/test_tir_analysis_undefined_vars.py
rename to tests/python/tirx/analysis/test_tirx_analysis_undefined_vars.py
diff --git a/tests/python/tirx-base/test_tir_base.py b/tests/python/tirx/base/test_tirx_base.py
similarity index 99%
rename from tests/python/tirx-base/test_tir_base.py
rename to tests/python/tirx/base/test_tirx_base.py
index 2f6fcfeb7e..e2aa19cd76 100644
--- a/tests/python/tirx-base/test_tir_base.py
+++ b/tests/python/tirx/base/test_tirx_base.py
@@ -14,7 +14,7 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-# ruff: noqa: E711, F841
+# ruff: noqa: E711
import itertools
import numpy as np
diff --git a/tests/python/tirx-base/test_binding_type_propagation.py b/tests/python/tirx/base/test_tirx_binding_type_propagation.py
similarity index 86%
rename from tests/python/tirx-base/test_binding_type_propagation.py
rename to tests/python/tirx/base/test_tirx_binding_type_propagation.py
index e2e87cc296..c805ca2331 100644
--- a/tests/python/tirx-base/test_binding_type_propagation.py
+++ b/tests/python/tirx/base/test_tirx_binding_type_propagation.py
@@ -21,6 +21,8 @@ import tvm_ffi
import tvm
from tvm import ir, relax, tirx
from tvm.ir import prim
+from tvm.script import ir as I
+from tvm.script import tirx as T
@pytest.mark.parametrize("inplace", [False, True])
@@ -162,5 +164,33 @@ def test_symbolic_call_type_rewrite_remaps_later_use(inplace):
ir.assert_structural_equal(ir.reinfer_type(new_binding.value), tensor_ty)
+def test_structural_map_tensor_shape_and_body():
+ @I.ir_module
+ class Before:
+ @T.function
+ def main(n: T.int32):
+ A = T.alloc_tensor(ty_args=[T.Tensor((n,), "float32")])
+ for i in range(n):
+ T.evaluate(A[i])
+
+ @I.ir_module
+ class Expected:
+ @T.function
+ def main():
+ A = T.alloc_tensor(ty_args=[T.Tensor((16,), "float32")])
+ for i in range(16):
+ T.evaluate(A[i])
+
+ func = Before["main"]
+ shape_var = func.params[0]
+ body = tvm_ffi.structural_map(
+ func.body,
+ (tvm.ir.Var, lambda var: T.int32(16) if var.same_as(shape_var) else var),
+ order="post",
+ )
+ After = tvm.IRModule.from_expr(tvm.tirx.Function([], body).with_attr("global_symbol", "main"))
+ tvm.ir.assert_structural_equal(After, Expected)
+
+
if __name__ == "__main__":
tvm.testing.main()
diff --git a/tests/python/tirx-base/test_tir_constructor.py b/tests/python/tirx/base/test_tirx_constructor.py
similarity index 97%
rename from tests/python/tirx-base/test_tir_constructor.py
rename to tests/python/tirx/base/test_tirx_constructor.py
index 8e871cb284..83ef15cda9 100644
--- a/tests/python/tirx-base/test_tir_constructor.py
+++ b/tests/python/tirx/base/test_tirx_constructor.py
@@ -379,5 +379,14 @@ def test_math_unary_constructor_preserves_bfloat16():
assert y.ty == tvm.ir.PrimType("bfloat16")
+def test_function_declaration_optional_body_roundtrip():
+ declaration = tvm.tirx.Function([], None)
+ assert declaration.body is None
+ assert len(declaration.with_body(tvm.ir.Evaluate(0)).body) == 0
+ restored = tvm.ir.load_json(tvm.ir.save_json(declaration))
+ tvm.ir.assert_structural_equal(declaration, restored, map_free_vars=True)
+ assert "def" in declaration.script()
+
+
if __name__ == "__main__":
tvm.testing.main()
diff --git a/tests/python/tirx-base/test_tir_host_func.py b/tests/python/tirx/base/test_tirx_host_func.py
similarity index 100%
rename from tests/python/tirx-base/test_tir_host_func.py
rename to tests/python/tirx/base/test_tirx_host_func.py
diff --git a/tests/python/tirx-base/test_tir_imm_values.py b/tests/python/tirx/base/test_tirx_imm_values.py
similarity index 100%
rename from tests/python/tirx-base/test_tir_imm_values.py
rename to tests/python/tirx/base/test_tirx_imm_values.py
diff --git a/tests/python/tirx-base/test_tir_index_map.py b/tests/python/tirx/base/test_tirx_index_map.py
similarity index 100%
rename from tests/python/tirx-base/test_tir_index_map.py
rename to tests/python/tirx/base/test_tirx_index_map.py
diff --git a/tests/python/tirx-base/test_tir_intrin.py b/tests/python/tirx/base/test_tirx_intrin.py
similarity index 100%
rename from tests/python/tirx-base/test_tir_intrin.py
rename to tests/python/tirx/base/test_tirx_intrin.py
diff --git a/tests/python/tirx-base/test_tir_nodes.py b/tests/python/tirx/base/test_tirx_nodes.py
similarity index 98%
rename from tests/python/tirx-base/test_tir_nodes.py
rename to tests/python/tirx/base/test_tirx_nodes.py
index db0faa6956..7f76d84577 100644
--- a/tests/python/tirx-base/test_tir_nodes.py
+++ b/tests/python/tirx/base/test_tirx_nodes.py
@@ -14,7 +14,6 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-# ruff: noqa: F841
import numpy as np
import pytest
@@ -354,6 +353,11 @@ def test_tensor_load_store():
s = tvm.ir.TensorStore(b, [0], 0.1)
assert isinstance(s, tvm.ir.TensorStore)
+ tensor = tvm.tirx.decl_tensor((4, 8), "float32", elem_offset=100)
+ tvm.ir.assert_structural_equal(
+ tensor[2, 3].indices, [tvm.ir.prim.IntImm("int32", 2), tvm.ir.prim.IntImm("int32", 3)]
+ )
+
@pytest.mark.parametrize("constructor", [tvm.ir.TensorLoad, tvm.tirx.TensorLoad])
def test_tensor_load_validation(constructor):
diff --git a/tests/python/tirx-base/test_tir_op_types.py b/tests/python/tirx/base/test_tirx_op_types.py
similarity index 96%
rename from tests/python/tirx-base/test_tir_op_types.py
rename to tests/python/tirx/base/test_tirx_op_types.py
index 05899e0379..1e0752aca1 100644
--- a/tests/python/tirx-base/test_tir_op_types.py
+++ b/tests/python/tirx/base/test_tirx_op_types.py
@@ -292,5 +292,15 @@ def test_tir_op_TVMBackendFreeWorkspace():
assert expr.op.name == "tirx.free_workspace"
+def test_ptx_ops_have_variadic_call_signature():
+ from tvm.backend.cuda.ptx import TABLE
+
+ op = tvm.ir.Op.get(next(iter(TABLE.values())).op_name)
+ assert not op.args_info
+ assert op.var_args_info.name == "operands"
+ tvm.ir.Call(op, []).validate()
+ tvm.ir.Call(op, [tvm.ir.Var("x", tvm.ir.AnyType())]).validate()
+
+
if __name__ == "__main__":
tvm.testing.main()
diff --git a/tests/python/tirx-base/test_tir_ops.py b/tests/python/tirx/base/test_tirx_ops.py
similarity index 99%
rename from tests/python/tirx-base/test_tir_ops.py
rename to tests/python/tirx/base/test_tirx_ops.py
index 80b1f571b9..22ec10cef7 100644
--- a/tests/python/tirx-base/test_tir_ops.py
+++ b/tests/python/tirx/base/test_tirx_ops.py
@@ -14,7 +14,6 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-# ruff: noqa: F841
import pytest
import tvm
diff --git a/tests/python/tirx-base/test_tir_ptx_cp_async.py b/tests/python/tirx/base/test_tirx_ptx_cp_async.py
similarity index 100%
rename from tests/python/tirx-base/test_tir_ptx_cp_async.py
rename to tests/python/tirx/base/test_tirx_ptx_cp_async.py
diff --git a/tests/python/tirx-base/test_tir_ptx_griddepcontrol.py b/tests/python/tirx/base/test_tirx_ptx_griddepcontrol.py
similarity index 100%
rename from tests/python/tirx-base/test_tir_ptx_griddepcontrol.py
rename to tests/python/tirx/base/test_tirx_ptx_griddepcontrol.py
diff --git a/tests/python/tirx-base/test_tir_ptx_ldmatrix.py b/tests/python/tirx/base/test_tirx_ptx_ldmatrix.py
similarity index 100%
rename from tests/python/tirx-base/test_tir_ptx_ldmatrix.py
rename to tests/python/tirx/base/test_tirx_ptx_ldmatrix.py
diff --git a/tests/python/tirx-base/test_tir_ptx_mma.py b/tests/python/tirx/base/test_tirx_ptx_mma.py
similarity index 100%
rename from tests/python/tirx-base/test_tir_ptx_mma.py
rename to tests/python/tirx/base/test_tirx_ptx_mma.py
diff --git a/tests/python/tirx-base/test_tir_ptx_mma_sp.py b/tests/python/tirx/base/test_tirx_ptx_mma_sp.py
similarity index 100%
rename from tests/python/tirx-base/test_tir_ptx_mma_sp.py
rename to tests/python/tirx/base/test_tirx_ptx_mma_sp.py
diff --git a/tests/python/tirx-base/test_tir_ptx_scalar_f32_math.py b/tests/python/tirx/base/test_tirx_ptx_scalar_f32_math.py
similarity index 100%
rename from tests/python/tirx-base/test_tir_ptx_scalar_f32_math.py
rename to tests/python/tirx/base/test_tirx_ptx_scalar_f32_math.py
diff --git a/tests/python/tirx-base/test_tir_scalable_datatype.py b/tests/python/tirx/base/test_tirx_scalable_datatype.py
similarity index 100%
rename from tests/python/tirx-base/test_tir_scalable_datatype.py
rename to tests/python/tirx/base/test_tirx_scalable_datatype.py
diff --git a/tests/python/tirx-base/test_tir_specialize.py b/tests/python/tirx/base/test_tirx_specialize.py
similarity index 100%
rename from tests/python/tirx-base/test_tir_specialize.py
rename to tests/python/tirx/base/test_tirx_specialize.py
diff --git a/tests/python/tirx/base/test_tirx_tensor_type.py b/tests/python/tirx/base/test_tirx_tensor_type.py
new file mode 100644
index 0000000000..f39324d3ca
--- /dev/null
+++ b/tests/python/tirx/base/test_tirx_tensor_type.py
@@ -0,0 +1,198 @@
+# 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.
+
+import uuid
+
+import pytest
+
+import tvm
+import tvm.testing
+from tvm.runtime import MemoryScope
+from tvm.script import tirx as T
+
+
+def test_tensor_pointer_type_derived_from_dtype_and_scope():
+ data = tvm.ir.Var("storage", tvm.ir.PtrType(tvm.ir.PrimType("uint8"), "local"))
+ tensor = tvm.tirx.decl_tensor((8,), "float16", data=data)
+
+ assert tensor.ty.dtype == tvm.ir.PrimType("float16")
+ assert tensor.ty.scope == tvm.runtime.MemoryScope("local")
+ assert tensor.data.ty == tvm.ir.PtrType(tvm.ir.PrimType("float16"), "local")
+
+
+@pytest.mark.parametrize("scope", ["global", "shared", tvm.runtime.MemoryScope("local")])
+def test_tensor_type_constructor_scope_order(scope):
+ direct = tvm.tirx.TensorType((4, 8), "float32", scope, (16, 1), 4, 64, 2)
+ scripted = T.Tensor((4, 8), "float32", scope, (16, 1), 4, 64, 2)
+ declared = tvm.tirx.decl_tensor(
+ (4, 8),
+ "float32",
+ scope=scope,
+ strides=(16, 1),
+ elem_offset=4,
+ data_alignment=64,
+ offset_factor=2,
+ )
+ tvm.ir.assert_structural_equal(direct, scripted)
+ tvm.ir.assert_structural_equal(direct, declared.ty)
+ assert direct.scope.same_as(tvm.runtime.MemoryScope(scope))
+
+
+@pytest.mark.parametrize("scope", ["", "tensor_type.unregistered"])
+def test_tensor_type_rejects_unregistered_scope(scope):
+ with pytest.raises(ValueError):
+ tvm.tirx.TensorType((4,), "float32", scope)
+ with pytest.raises(ValueError):
+ T.Tensor((4,), "float32", scope)
+ with pytest.raises(ValueError):
+ tvm.tirx.decl_tensor((4,), "float32", scope=scope)
+
+
+def test_decl_tensor_physical_data_binding():
+ tensor = tvm.tirx.decl_tensor((8,), "float32")
+ data = tvm.tirx.Var("data", tensor.data.ty)
+
+ decl = tvm.ir.Bind(
+ tensor,
+ tvm.ir.Call(
+ "tirx.decl_tensor",
+ [data],
+ ty_args=[tensor.ty],
+ ),
+ )
+ assert decl.var.same_as(tensor)
+ assert decl.value.args[0].same_as(data)
+
+
+def test_tensor_ptr_to():
+ m = tvm.tirx.Var("m", "int32")
+ n = tvm.tirx.Var("n", "int32")
+ tensor = tvm.tirx.decl_tensor((m, n), "float32", strides=[n + 1, 1], elem_offset=7)
+ pointer = tensor.ptr_to([2, 3])
+ assert pointer.op.name == "tirx.address_of"
+ assert pointer.ty == tvm.ir.PtrType(tvm.ir.PrimType("float32"))
+ assert pointer.args[0].source.same_as(tensor)
+ tvm.ir.assert_structural_equal(pointer.args[0].indices, [T.int32(2), T.int32(3)])
+ shared = tvm.tirx.decl_tensor((m, n), "float32", scope="shared")
+ assert shared.ptr_to([0, 0]).ty == tvm.ir.PtrType(tvm.ir.PrimType("float32"), "shared")
+
+
+@pytest.mark.parametrize("dtype,scope", [("float32", "global"), ("bool", "shared")])
+def test_tensor_type_and_projection(dtype, scope):
+ m, n = tvm.ir.Var("m", tvm.ir.PrimType("int32")), tvm.ir.Var("n", tvm.ir.PrimType("int32"))
+ tensor = tvm.tirx.decl_tensor((m, n), dtype, scope=scope)
+ assert type(tensor) is tvm.ir.Var
+ assert tvm.tirx.is_tensor_var(tensor)
+ assert not tvm.tirx.is_tensor_var(m)
+ assert isinstance(tensor.ty, tvm.tirx.TensorType)
+ assert tuple(tensor.ty.shape) == (m, n)
+ assert tuple(tensor.shape) == (m, n)
+ assert tensor.ty.dtype == tvm.ir.PrimType(dtype)
+ assert tensor.dtype == tvm.DataType(dtype)
+ assert tensor.data.ty == tvm.ir.PtrType(tvm.ir.PrimType(dtype), scope)
+ assert tensor.data.op.name == "tirx.tensor_data_ptr"
+ assert tensor.data.args[0].same_as(tensor)
+ serialized = tvm.ir.save_json(tensor.ty)
+ assert '"tirx.TensorType"' in serialized
+ restored = tvm.ir.load_json(serialized)
+ assert isinstance(restored, tvm.tirx.TensorType)
+ tvm.ir.assert_structural_equal(restored, tensor.ty, map_free_vars=True)
+ for name in ("shape", "dtype", "data"):
+ assert not hasattr(m, name)
+ with pytest.raises(AttributeError, match="has no attribute"):
+ getattr(m, name)
+
+
+def test_tensor_type_default_offset():
+ tensor = tvm.tirx.decl_tensor((8,), "float32", elem_offset=None)
+ assert int(tensor.ty.elem_offset) == 0
+
+
+@pytest.mark.parametrize(
+ "dtype",
+ [
+ "float4_e2m1fn",
+ "float8_e3m4",
+ "float8_e4m3",
+ "float8_e4m3b11fnuz",
+ "float8_e4m3fn",
+ "float8_e4m3fnuz",
+ "float8_e5m2",
+ "float8_e5m2fnuz",
+ "float8_e8m0fnu",
+ ],
+)
+def test_tensor_type_low_precision_dtype(dtype):
+ m = tvm.ir.Var("m", tvm.ir.PrimType("int32"))
+ n = tvm.ir.Var("n", tvm.ir.PrimType("int32"))
+ tensor = tvm.tirx.decl_tensor((m, n), dtype)
+ assert tensor.dtype == dtype
+
+
+def test_memory_scope_typed_conversion():
+ custom = MemoryScope.register("test.tensor.scope." + uuid.uuid4().hex)
+ for scope in (custom, custom.name):
+ ptr_ty = tvm.ir.PtrType(tvm.ir.PrimType("float32"), scope)
+ tensor = tvm.tirx.decl_tensor((4,), "float32", scope=scope)
+ tensor_ty = tvm.tirx.TensorType((4,), "float32", scope)
+ assert tensor_ty.scope.same_as(custom)
+ assert ptr_ty.scope.same_as(custom)
+ assert tensor.ty.scope.same_as(custom)
+ assert tensor.scope().same_as(custom)
+ assert tensor.data.ty.scope.same_as(custom)
+
+ unknown = "test.memory_scope.unknown." + uuid.uuid4().hex
+ for construct in (
+ lambda: tvm.ir.PtrType(tvm.ir.PrimType("float32"), unknown),
+ lambda: tvm.tirx.decl_tensor((4,), "float32", scope=unknown),
+ lambda: T.Tensor((4,), "float32", scope=unknown),
+ ):
+ with pytest.raises(ValueError):
+ construct()
+ with pytest.raises(ValueError):
+ MemoryScope(unknown)
+
+
+def test_memory_scope_global_defaults():
+ global_scope = MemoryScope("global")
+ assert tvm.ir.PtrType(tvm.ir.PrimType("float32")).scope.same_as(global_scope)
+ assert tvm.tirx.decl_tensor((4,), "float32").scope().same_as(global_scope)
+ for scope in ("global", global_scope):
+ assert tvm.ir.PtrType(tvm.ir.PrimType("float32"), scope).scope.same_as(global_scope)
+ assert tvm.tirx.decl_tensor((4,), "float32", scope=scope).scope().same_as(global_scope)
+ assert T.Tensor((4,), "float32", scope=scope).scope.same_as(global_scope)
+ for construct in (
+ lambda: MemoryScope(""),
+ lambda: tvm.ir.PtrType(tvm.ir.PrimType("float32"), ""),
+ lambda: tvm.tirx.decl_tensor((4,), "float32", scope=""),
+ lambda: T.Tensor((4,), "float32", scope=""),
+ ):
+ with pytest.raises(ValueError):
+ construct()
+
+
+def test_tensor_type_memory_scope_json_roundtrip():
+ custom = MemoryScope.register("test.ir.scope." + uuid.uuid4().hex)
+ for scope in (MemoryScope("global"), MemoryScope("shared.dyn"), custom):
+ ty = tvm.tirx.TensorType((4,), "float32", scope=scope)
+ restored = tvm.ir.load_json(tvm.ir.save_json(ty))
+ assert restored.scope.same_as(scope)
+ tvm.ir.assert_structural_equal(ty, restored)
+
+
+if __name__ == "__main__":
+ tvm.testing.main()
diff --git a/tests/python/tirx-base/test_tir_texture_scope.py b/tests/python/tirx/base/test_tirx_texture_scope.py
similarity index 100%
rename from tests/python/tirx-base/test_tir_texture_scope.py
rename to tests/python/tirx/base/test_tirx_texture_scope.py
diff --git a/tests/python/tirx/codegen/test_tirx_codegen_fp8.py b/tests/python/tirx/codegen/test_tirx_codegen_fp8.py
new file mode 100644
index 0000000000..ecac99933f
--- /dev/null
+++ b/tests/python/tirx/codegen/test_tirx_codegen_fp8.py
@@ -0,0 +1,54 @@
+# 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.
+import numpy as np
+import pytest
+
+import tvm
+import tvm.testing
+from tvm.script import s_tir as Ts
+from tvm.script import tirx as T
+
+
+@pytest.mark.parametrize("dtype", ["float8_e3m4", "float8_e4m3", "float8_e4m3fn", "float8_e5m2"])
+def test_fp8_float32_roundtrip(dtype):
+ if not tvm.testing.device_enabled("llvm"):
+ pytest.skip("LLVM is not enabled")
+ ml_dtypes = pytest.importorskip("ml_dtypes")
+ if not hasattr(ml_dtypes, dtype):
+ pytest.skip(f"ml_dtypes does not provide {dtype}")
+ np_dtype = getattr(ml_dtypes, dtype)
+
+ @Ts.function
+ def roundtrip(
+ A: T.Tensor((128,), dtype),
+ intermediate: T.Tensor((128,), "float32"),
+ B: T.Tensor((128,), dtype),
+ ):
+ for i in range(128):
+ with Ts.sblock("roundtrip"):
+ vi = Ts.axis.spatial(128, i)
+ intermediate[vi] = A[vi]
+ B[vi] = intermediate[vi]
+
+ source = np.random.default_rng(0).standard_normal(128).astype(np_dtype)
+ tensors = [
+ tvm.runtime.tensor(source),
+ tvm.runtime.tensor(np.zeros(128, dtype="float32")),
+ tvm.runtime.tensor(np.zeros(128, dtype=np_dtype)),
+ ]
+ tvm.compile(roundtrip, target="llvm")(*tensors)
+ np.testing.assert_equal(tensors[2].numpy(), source.astype("float32").astype(np_dtype))
diff --git a/tests/python/test_tirx_axis_enum.py b/tests/python/tirx/test_tirx_axis_enum.py
similarity index 100%
rename from tests/python/test_tirx_axis_enum.py
rename to tests/python/tirx/test_tirx_axis_enum.py
diff --git a/tests/python/tirx-transform/test_tir_transform_bf16_legalize.py b/tests/python/tirx/transform/test_tirx_transform_bf16_legalize.py
similarity index 83%
rename from tests/python/tirx-transform/test_tir_transform_bf16_legalize.py
rename to tests/python/tirx/transform/test_tirx_transform_bf16_legalize.py
index 25884a299c..81d72fc123 100644
--- a/tests/python/tirx-transform/test_tir_transform_bf16_legalize.py
+++ b/tests/python/tirx/transform/test_tirx_transform_bf16_legalize.py
@@ -14,24 +14,15 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-import tvm_ffi
import tvm
import tvm.script
+from tvm.script import ir as I
from tvm.script import tirx as T
from tvm.target import Target
from tvm.tirx.transform.transform import BindTarget
-def _is_tensor_binding(node, *op_names):
- return (
- isinstance(node, tvm.ir.Bind)
- and isinstance(node.value, tvm.ir.Call)
- and isinstance(node.value.op, tvm.ir.Op)
- and node.value.op.name in op_names
- )
-
-
def u16tof32(v):
uint32_v = v.astype("uint32")
uint32_v = uint32_v << tvm.tirx.const(16, "uint32")
@@ -161,53 +152,95 @@ def test_bf16_masked_load_store_will_legalize():
after_compute = tvm.tirx.transform.BF16ComputeLegalize()(before)
after_storage = tvm.tirx.transform.BF16StorageLegalize()(after_compute)
- def collect(mod):
- nodes = []
-
- def collect_once(node):
- if not any(node.same_as(existing) for existing in nodes):
- nodes.append(node)
-
- tvm_ffi.structural_walk(
- mod["main"].body,
- ((tvm.ir.Bind, tvm.ir.Call), collect_once),
- )
- tensors = {
- node.var.name: str(node.var.dtype)
- for node in nodes
- if _is_tensor_binding(node, "tirx.alloc_tensor", "tirx.decl_tensor")
- }
- masked_loads = [
- node
- for node in nodes
- if isinstance(node, tvm.ir.Call) and node.op.name == "tirx.masked_load"
- ]
- masked_stores = [
- node
- for node in nodes
- if isinstance(node, tvm.ir.Call) and node.op.name == "tirx.masked_store"
- ]
- return tensors, masked_loads, masked_stores
-
- compute_tensors, compute_loads, compute_stores = collect(after_compute)
- assert compute_tensors == {"A": "bfloat16", "B": "float32", "C": "bfloat16", "mask": "boolx4"}
- assert sorted(str(load.ty) for load in compute_loads) == ["bfloat16x4", "float32x4"]
- assert sorted(str(store.args[1].ty) for store in compute_stores) == [
- "bfloat16x4",
- "float32x4",
- ]
-
- storage_tensors, storage_loads, storage_stores = collect(after_storage)
- assert storage_tensors == {"A": "uint16", "B": "float32", "C": "uint16", "mask": "boolx4"}
- assert sorted(str(load.ty) for load in storage_loads) == [
- "float32x4",
- "float32x4",
- "uint16x4",
- ]
- assert sorted(str(store.args[1].ty) for store in storage_stores) == [
- "float32x4",
- "uint16x4",
- ]
+ @I.ir_module
+ class ExpectedCompute:
+ @T.function
+ def main(Aptr: T.Ptr(T.PrimType("bfloat16")), Cptr: T.Ptr(T.PrimType("bfloat16"))):
+ T.func_attr({"target": target})
+ A = T.decl_tensor(Aptr, ty_args=[T.Tensor((16,), T.bfloat16)])
+ B = T.alloc_tensor(ty_args=[T.Tensor((16,), T.float32)])
+ C = T.decl_tensor(Cptr, ty_args=[T.Tensor((16,), T.bfloat16)])
+ mask = T.alloc_tensor(ty_args=[T.Tensor((1,), T.PrimType("boolx4"), scope="local")])
+ mask[0] = T.Broadcast(T.bool(True), 4)
+ T.masked_store(
+ B,
+ T.reinterpret(
+ T.Cast(
+ T.reinterpret(T.masked_load(A, T.Ramp(0, 1, 4), mask[0]), ty="uint16x4"),
+ "uint32x4",
+ )
+ << T.Broadcast(T.uint32(16), 4),
+ ty="float32x4",
+ ),
+ T.Ramp(0, 1, 4),
+ mask[0],
+ )
+ T.masked_store(
+ C,
+ T.reinterpret(
+ T.Cast(
+ T.reinterpret(T.masked_load(B, T.Ramp(0, 1, 4), mask[0]), ty="uint32x4")
+ + (
+ (
+ T.reinterpret(
+ T.masked_load(B, T.Ramp(0, 1, 4), mask[0]), ty="uint32x4"
+ )
+ >> T.Broadcast(T.uint32(16), 4)
+ & T.Broadcast(T.uint32(1), 4)
+ )
+ + T.Broadcast(T.uint32(32767), 4)
+ )
+ >> T.Broadcast(T.uint32(16), 4),
+ "uint16x4",
+ ),
+ ty="bfloat16x4",
+ ),
+ T.Ramp(0, 1, 4),
+ mask[0],
+ )
+
+ tvm.ir.assert_structural_equal(after_compute, ExpectedCompute)
+
+ @I.ir_module
+ class ExpectedStorage:
+ @T.function
+ def main(Aptr: T.Ptr(T.PrimType("uint16")), Cptr: T.Ptr(T.PrimType("uint16"))):
+ T.func_attr({"target": target})
+ A = T.decl_tensor(Aptr, ty_args=[T.Tensor((16,), T.uint16)])
+ B = T.alloc_tensor(ty_args=[T.Tensor((16,), T.float32)])
+ C = T.decl_tensor(Cptr, ty_args=[T.Tensor((16,), T.uint16)])
+ mask = T.alloc_tensor(ty_args=[T.Tensor((1,), T.PrimType("boolx4"), scope="local")])
+ mask[0] = T.Broadcast(T.bool(True), 4)
+ T.masked_store(
+ B,
+ T.reinterpret(
+ T.Cast(T.masked_load(A, T.Ramp(0, 1, 4), mask[0]), "uint32x4")
+ << T.Broadcast(T.uint32(16), 4),
+ ty="float32x4",
+ ),
+ T.Ramp(0, 1, 4),
+ mask[0],
+ )
+ T.masked_store(
+ C,
+ T.Cast(
+ T.reinterpret(T.masked_load(B, T.Ramp(0, 1, 4), mask[0]), ty="uint32x4")
+ + (
+ (
+ T.reinterpret(T.masked_load(B, T.Ramp(0, 1, 4), mask[0]), ty="uint32x4")
+ >> T.Broadcast(T.uint32(16), 4)
+ & T.Broadcast(T.uint32(1), 4)
+ )
+ + T.Broadcast(T.uint32(32767), 4)
+ )
+ >> T.Broadcast(T.uint32(16), 4),
+ "uint16x4",
+ ),
+ T.Ramp(0, 1, 4),
+ mask[0],
+ )
+
+ tvm.ir.assert_structural_equal(after_storage, ExpectedStorage)
def test_bf16_storage_compute_scope_will_legalize():
diff --git a/tests/python/tirx/transform/test_tirx_transform_common_subexpr_elim.py b/tests/python/tirx/transform/test_tirx_transform_common_subexpr_elim.py
new file mode 100644
index 0000000000..7442b2f9b2
--- /dev/null
+++ b/tests/python/tirx/transform/test_tirx_transform_common_subexpr_elim.py
@@ -0,0 +1,622 @@
+# 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.
+import hashlib
+
+import tvm
+import tvm.testing
+from tvm.ir.base import save_json
+from tvm.script import tirx as T
+
+
+def test_basic():
+ @tvm.script.ir_module
+ class Before:
+ @T.function
+ def main(B: T.Tensor((50,), "int32"), i1: T.int32, i2: T.int32, z3: T.int32):
+ z1 = T.bind(1)
+ z2 = T.bind(2)
+ B[i1] = z1 + z2
+ x = T.bind(1)
+ y = T.bind(1)
+ a = T.bind((x + y) + (z1 + z2))
+ b = T.bind((x + y) + z3)
+ B[i2] = a + b
+
+ @tvm.script.ir_module
+ class Expected:
+ @T.function
+ def main(B: T.Tensor((50,), "int32"), i1: T.int32, i2: T.int32, z3: T.int32):
+ z1 = T.bind(1)
+ z2 = T.bind(2)
+ cse_v1 = T.bind(z1 + z2)
+ B[i1] = cse_v1
+ x = T.bind(1)
+ y = T.bind(1)
+ cse_v2 = T.bind(x + y)
+ a = T.bind(cse_v2 + cse_v1)
+ b = T.bind(cse_v2 + z3)
+ B[i2] = a + b
+
+ after = tvm.tirx.transform.CommonSubexprElim()(Before)
+ tvm.ir.assert_structural_equal(after, Expected)
+
+
+def test_if_single_branch():
+ @tvm.script.ir_module
+ class Before:
+ @T.function
+ def main(
+ B: T.Tensor((50,), "int32"),
+ i1: T.int32,
+ i2: T.int32,
+ i3: T.int32,
+ y: T.int32,
+ z: T.int32,
+ ):
+ b = T.bind(1)
+ if b:
+ B[i1] = y + z
+ B[i2] = y + z
+ else:
+ B[i3] = y
+
+ @tvm.script.ir_module
+ class Expected:
+ @T.function
+ def main(
+ B: T.Tensor((50,), "int32"),
+ i1: T.int32,
+ i2: T.int32,
+ i3: T.int32,
+ y: T.int32,
+ z: T.int32,
+ ):
+ b = T.bind(1)
+ if b:
+ cse_v1 = T.bind(y + z)
+ B[i1] = cse_v1
+ B[i2] = cse_v1
+ else:
+ B[i3] = y
+
+ after = tvm.tirx.transform.CommonSubexprElim()(Before)
+ tvm.ir.assert_structural_equal(after, Expected)
+
+
+def test_if_both_branches():
+ @tvm.script.ir_module
+ class Before:
+ @T.function
+ def main(
+ B: T.Tensor((50,), "int32"),
+ i1: T.int32,
+ i2: T.int32,
+ i3: T.int32,
+ y: T.int32,
+ z: T.int32,
+ ):
+ b = T.bind(1)
+ if b:
+ B[i1] = y + z
+ B[i2] = y
+ else:
+ B[i3] = y + z
+
+ @tvm.script.ir_module
+ class Expected:
+ @T.function
+ def main(
+ B: T.Tensor((50,), "int32"),
+ i1: T.int32,
+ i2: T.int32,
+ i3: T.int32,
+ y: T.int32,
+ z: T.int32,
+ ):
+ b = T.bind(1)
+ cse_v1 = T.bind(y + z)
+ if b:
+ B[i1] = cse_v1
+ B[i2] = y
+ else:
+ B[i3] = cse_v1
+
+ after = tvm.tirx.transform.CommonSubexprElim()(Before)
+ tvm.ir.assert_structural_equal(after, Expected)
+
+
+def test_cascade():
+ @tvm.script.ir_module
+ class Before:
+ @T.function
+ def main(
+ B: T.Tensor((50,), "int32"),
+ i1: T.int32,
+ i2: T.int32,
+ i3: T.int32,
+ x: T.int32,
+ y: T.int32,
+ z: T.int32,
+ ):
+ B[i1] = (x + y) + z
+ B[i2] = (x + y) + z
+ B[i3] = x + y
+
+ @tvm.script.ir_module
+ class Expected:
+ @T.function
+ def main(
+ B: T.Tensor((50,), "int32"),
+ i1: T.int32,
+ i2: T.int32,
+ i3: T.int32,
+ x: T.int32,
+ y: T.int32,
+ z: T.int32,
+ ):
+ cse_v2 = T.bind(x + y)
+ cse_v1 = T.bind(cse_v2 + z)
+ B[i1] = cse_v1
+ B[i2] = cse_v1
+ B[i3] = cse_v2
+
+ after = tvm.tirx.transform.CommonSubexprElim()(Before)
+ tvm.ir.assert_structural_equal(after, Expected)
+
+
+def test_no_duplication():
+ @tvm.script.ir_module
+ class Before:
+ @T.function
+ def main(x: T.int32, y: T.int32, z: T.int32):
+ a = T.bind(x + (y + z))
+ T.evaluate(a)
+
+ Expected = Before
+
+ after = tvm.tirx.transform.CommonSubexprElim()(Before)
+ tvm.ir.assert_structural_equal(after, Expected)
+
+
+def test_deterministic():
+ @tvm.script.ir_module
+ class Before:
+ @T.function
+ def main(B: T.Tensor((4,), "int32"), x: T.int32):
+ B[0] = x + 1
+ B[1] = x + 2
+ B[2] = x + 1
+ B[3] = x + 2
+
+ @tvm.script.ir_module
+ class Expected:
+ @T.function
+ def main(B: T.Tensor((4,), "int32"), x: T.int32):
+ cse_v1 = T.bind(x + 1)
+ B[0] = cse_v1
+ cse_v2 = T.bind(x + 2)
+ B[1] = cse_v2
+ B[2] = cse_v1
+ B[3] = cse_v2
+
+ hashes = set()
+ for _ in range(10):
+ After = tvm.tirx.transform.CommonSubexprElim()(Before)
+ tvm.ir.assert_structural_equal(After, Expected)
+ hashes.add(hashlib.sha256(save_json(After).encode()).hexdigest())
+ assert len(hashes) == 1
+
+
+def test_for_loop():
+ @tvm.script.ir_module
+ class Before:
+ @T.function
+ def main(B: T.Tensor((50,), "int32"), y: T.int32, z: T.int32):
+ for i in range(10):
+ B[i] = y + z
+ B[i + 10] = y + z
+
+ @tvm.script.ir_module
+ class Expected:
+ @T.function
+ def main(B: T.Tensor((50,), "int32"), y: T.int32, z: T.int32):
+ for i in range(10):
+ cse_v1 = T.bind(y + z)
+ B[i] = cse_v1
+ B[i + 10] = cse_v1
+
+ after = tvm.tirx.transform.CommonSubexprElim()(Before)
+ tvm.ir.assert_structural_equal(after, Expected)
+
+
+def test_for_hoist():
+ @tvm.script.ir_module
+ class Before:
+ @T.function
+ def main(B: T.Tensor((50,), "int32"), y: T.int32, z: T.int32):
+ B[0] = y + z
+ for i in range(10):
+ B[i + 1] = y + z
+
+ @tvm.script.ir_module
+ class Expected:
+ @T.function
+ def main(B: T.Tensor((50,), "int32"), y: T.int32, z: T.int32):
+ cse_v1 = T.bind(y + z)
+ B[0] = cse_v1
+ for i in range(10):
+ B[i + 1] = cse_v1
+
+ after = tvm.tirx.transform.CommonSubexprElim()(Before)
+ tvm.ir.assert_structural_equal(after, Expected)
+
+
+def test_cannot_lift_tensorload():
+ @tvm.script.ir_module
+ class Before:
+ @T.function
+ def main(A: T.Tensor((50,), "int32"), B: T.Tensor((50,), "int32")):
+ B[0] = A[0] + A[0]
+ B[1] = A[0] + A[0]
+
+ Expected = Before
+
+ after = tvm.tirx.transform.CommonSubexprElim()(Before)
+ tvm.ir.assert_structural_equal(after, Expected)
+
+
+def test_nested_if():
+ @tvm.script.ir_module
+ class Before:
+ @T.function
+ def main(
+ B: T.Tensor((50,), "int32"),
+ c1: T.int32,
+ c2: T.int32,
+ y: T.int32,
+ z: T.int32,
+ ):
+ if c1:
+ if c2:
+ B[0] = y + z
+ else:
+ B[1] = y + z
+ else:
+ B[2] = y
+
+ @tvm.script.ir_module
+ class Expected:
+ @T.function
+ def main(
+ B: T.Tensor((50,), "int32"),
+ c1: T.int32,
+ c2: T.int32,
+ y: T.int32,
+ z: T.int32,
+ ):
+ if c1:
+ cse_v1 = T.bind(y + z)
+ if c2:
+ B[0] = cse_v1
+ else:
+ B[1] = cse_v1
+ else:
+ B[2] = y
+
+ after = tvm.tirx.transform.CommonSubexprElim()(Before)
+ tvm.ir.assert_structural_equal(after, Expected)
+
+
+def test_multi_independent():
+ @tvm.script.ir_module
+ class Before:
+ @T.function
+ def main(
+ B: T.Tensor((50,), "int32"),
+ a: T.int32,
+ b: T.int32,
+ c: T.int32,
+ d: T.int32,
+ ):
+ B[0] = a + b
+ B[1] = c + d
+ B[2] = a + b
+ B[3] = c + d
+
+ @tvm.script.ir_module
+ class Expected:
+ @T.function
+ def main(
+ B: T.Tensor((50,), "int32"),
+ a: T.int32,
+ b: T.int32,
+ c: T.int32,
+ d: T.int32,
+ ):
+ cse_v1 = T.bind(a + b)
+ B[0] = cse_v1
+ cse_v2 = T.bind(c + d)
+ B[1] = cse_v2
+ B[2] = cse_v1
+ B[3] = cse_v2
+
+ after = tvm.tirx.transform.CommonSubexprElim()(Before)
+ tvm.ir.assert_structural_equal(after, Expected)
+
+
+def test_if_condition():
+ @tvm.script.ir_module
+ class Before:
+ @T.function
+ def main(B: T.Tensor((50,), "int32"), y: T.int32, z: T.int32):
+ if y + z > 0:
+ B[0] = y + z
+
+ @tvm.script.ir_module
+ class Expected:
+ @T.function
+ def main(B: T.Tensor((50,), "int32"), y: T.int32, z: T.int32):
+ cse_v1 = T.bind(y + z)
+ if cse_v1 > 0:
+ B[0] = cse_v1
+
+ after = tvm.tirx.transform.CommonSubexprElim()(Before)
+ tvm.ir.assert_structural_equal(after, Expected)
+
+
+def test_cannot_lift_call():
+ @tvm.script.ir_module
+ class Before:
+ @T.function
+ def main(B: T.Tensor((50,), "int32"), x: T.int32):
+ B[0] = T.call_extern("my_func", x, ty="int32") + 1
+ B[1] = T.call_extern("my_func", x, ty="int32") + 1
+
+ Expected = Before
+
+ after = tvm.tirx.transform.CommonSubexprElim()(Before)
+ tvm.ir.assert_structural_equal(after, Expected)
+
+
+def test_no_single_use_binding():
+ @tvm.script.ir_module
+ class Before:
+ @T.function
+ def main(
+ B: T.Tensor((50,), "int32"),
+ x: T.int32,
+ y: T.int32,
+ z: T.int32,
+ ):
+ B[0] = (x + y) + z
+ B[1] = (x + y) + z
+
+ @tvm.script.ir_module
+ class Expected:
+ @T.function
+ def main(
+ B: T.Tensor((50,), "int32"),
+ x: T.int32,
+ y: T.int32,
+ z: T.int32,
+ ):
+ cse_v1 = T.bind((x + y) + z)
+ B[0] = cse_v1
+ B[1] = cse_v1
+
+ after = tvm.tirx.transform.CommonSubexprElim()(Before)
+ tvm.ir.assert_structural_equal(after, Expected)
+
+
+def test_for_extent_lift():
+ @tvm.script.ir_module
+ class Before:
+ @T.function
+ def main(B: T.Tensor((50,), "int32"), y: T.int32, z: T.int32):
+ for i in range(y + z):
+ B[i] = y + z
+
+ @tvm.script.ir_module
+ class Expected:
+ @T.function
+ def main(B: T.Tensor((50,), "int32"), y: T.int32, z: T.int32):
+ cse_v1 = T.bind(y + z)
+ for i in range(cse_v1):
+ B[i] = cse_v1
+
+ after = tvm.tirx.transform.CommonSubexprElim()(Before)
+ tvm.ir.assert_structural_equal(after, Expected)
+
+
+def test_loop_var_expr_stays_inside():
+ @tvm.script.ir_module
+ class Before:
+ @T.function
+ def main(
+ A: T.Tensor((50,), "int32"),
+ B: T.Tensor((50,), "int32"),
+ ):
+ for i in range(10):
+ A[i * 4] = B[i * 4]
+
+ @tvm.script.ir_module
+ class Expected:
+ @T.function
+ def main(
+ A: T.Tensor((50,), "int32"),
+ B: T.Tensor((50,), "int32"),
+ ):
+ for i in range(10):
+ cse_v1 = T.bind(i * 4)
+ A[cse_v1] = B[cse_v1]
+
+ after = tvm.tirx.transform.CommonSubexprElim()(Before)
+ tvm.ir.assert_structural_equal(after, Expected)
+
+
+def test_let_body_no_extraction():
+ x = T.dynamic("x", "int32")
+
+ @tvm.script.ir_module
+ class Before:
+ @T.function
+ def main(B: T.Tensor((10,), "int32"), y: T.int32):
+ for i in range(10):
+ B[i] = T.Let((x + y) + (x + y), where={x: 1})
+
+ After = tvm.tirx.transform.CommonSubexprElim()(Before)
+ tvm.ir.assert_structural_equal(After, Before)
+
+
+def test_let_value_cse():
+ x = T.dynamic("x", "int32")
+
+ @tvm.script.ir_module
+ class Before:
+ @T.function
+ def main(B: T.Tensor((10,), "int32"), y: T.int32, z: T.int32):
+ for i in range(10):
+ B[i] = (y + z) + T.Let(x + 1, where={x: y + z})
+
+ @tvm.script.ir_module
+ class Expected:
+ @T.function
+ def main(B: T.Tensor((10,), "int32"), y: T.int32, z: T.int32):
+ for i in range(10):
+ cse_v1 = T.bind(y + z)
+ B[i] = cse_v1 + T.Let(x + 1, where={x: cse_v1})
+
+ After = tvm.tirx.transform.CommonSubexprElim()(Before)
+ tvm.ir.assert_structural_equal(After, Expected)
+
+
+def test_nested_let_no_extraction():
+ x = T.dynamic("x", "int32")
+ y = T.dynamic("y", "int32")
+
+ @tvm.script.ir_module
+ class Before:
+ @T.function
+ def main(B: T.Tensor((10,), "int32"), z: T.int32):
+ for i in range(10):
+ B[i] = T.Let(T.Let((x + y + z) + (x + y + z), where={y: 2}), where={x: 1})
+
+ After = tvm.tirx.transform.CommonSubexprElim()(Before)
+ tvm.ir.assert_structural_equal(After, Before)
+
+
+def test_let_floordiv_pattern():
+ x = T.dynamic("x", "int32")
+ y = T.dynamic("y", "int32")
+ rmod = T.dynamic("rmod", "int32")
+ rdiv = T.dynamic("rdiv", "int32")
+
+ @tvm.script.ir_module
+ class Before:
+ @T.function
+ def main(
+ A: T.Tensor((10,), "int32"), B: T.Tensor((10,), "int32"), C: T.Tensor((10,), "int32")
+ ):
+ for i in range(10):
+ C[i] = T.Let(
+ T.Let(
+ T.Let(
+ T.Let(
+ T.Select(
+ (y >= 0 and rmod >= 0) or (y < 0 and rmod <= 0), rdiv, rdiv - 1
+ ),
+ where={rdiv: T.Div(x, y)},
+ ),
+ where={rmod: T.Mod(x, y)},
+ ),
+ where={y: B[i]},
+ ),
+ where={x: A[i]},
+ )
+
+ After = tvm.tirx.transform.CommonSubexprElim()(Before)
+ tvm.ir.assert_structural_equal(After, Before)
+
+
+def test_no_lift_bool_predicate():
+ @tvm.script.ir_module
+ class Before:
+ @T.function
+ def main(B: T.Tensor((50,), "int32"), n: T.int32, x: T.int32):
+ for i in range(50):
+ if i < n:
+ B[i] = x
+ if i < n:
+ B[i] = x + 1
+
+ after = tvm.tirx.transform.CommonSubexprElim()(Before)
+ tvm.ir.assert_structural_equal(after, Before)
+ assert "cse_v" not in after["main"].script()
+
+
+def test_no_lift_bool_logical():
+ @tvm.script.ir_module
+ class Before:
+ @T.function
+ def main(B: T.Tensor((50,), "int32"), a: T.bool, b: T.bool, x: T.int32):
+ if T.And(a, b):
+ B[0] = x
+ if T.And(a, b):
+ B[1] = x + 1
+
+ after = tvm.tirx.transform.CommonSubexprElim()(Before)
+ tvm.ir.assert_structural_equal(after, Before)
+ assert "cse_v" not in after["main"].script()
+
+
+# T24: Shared subtree stays SSA. UnrollLoop reuses one Stmt object across
+# positions; the identity-keyed rewriter must bind fresh cse vars per occurrence.
+def test_shared_subtree_stays_ssa():
+ @tvm.script.ir_module
+ class Before:
+ @T.function
+ def main(B: T.Tensor((50,), "int32"), i1: T.int32, i2: T.int32):
+ for unused in T.unroll(2):
+ B[i1] = (i1 + i2) * 2
+ B[i2] = (i1 + i2) * 3
+
+ @tvm.script.ir_module
+ class Expected:
+ @T.function
+ def main(B: T.Tensor((50,), "int32"), i1: T.int32, i2: T.int32):
+ value = T.bind(i1 + i2)
+ twice = T.bind(value * 2)
+ B[i1] = twice
+ thrice = T.bind(value * 3)
+ B[i2] = thrice
+ value2 = T.bind(i1 + i2)
+ twice2 = T.bind(value2 * 2)
+ B[i1] = twice2
+ thrice2 = T.bind(value * 3)
+ B[i2] = thrice2
+
+ unrolled = tvm.tirx.transform.UnrollLoop()(Before)
+ body = unrolled["main"].body
+ assert body[0].same_as(body[2])
+ assert body[1].same_as(body[3])
+ After = tvm.tirx.transform.CommonSubexprElim()(unrolled)
+ tvm.ir.assert_structural_equal(After, Expected)
+ assert tvm.tirx.analysis.verify_well_formed(After["main"])
+
+
+if __name__ == "__main__":
+ tvm.testing.main()
diff --git a/tests/python/tirx-transform/test_tir_transform_convert_ssa.py b/tests/python/tirx/transform/test_tirx_transform_convert_ssa.py
similarity index 56%
rename from tests/python/tirx-transform/test_tir_transform_convert_ssa.py
rename to tests/python/tirx/transform/test_tirx_transform_convert_ssa.py
index b07fb89220..2dbaf65124 100644
--- a/tests/python/tirx-transform/test_tir_transform_convert_ssa.py
+++ b/tests/python/tirx/transform/test_tirx_transform_convert_ssa.py
@@ -26,19 +26,18 @@ from tvm.script import tirx as T
def test_reuse_in_sequential_bind():
"""De-dup sequential variable bindings"""
- # Manually construct the Function body, as SSA violations are
- # not valid TIR, and may not be expressible in future versions
- # of TVMSCript.
- var = tirx.Var("var", "int32")
- sequential_bindings = tvm.ir.SeqStmt(
- [
- tvm.ir.Bind(var, 16),
- tvm.ir.Evaluate(var),
- tvm.ir.Bind(var, 32),
- tvm.ir.Evaluate(var),
- ]
- )
- before = tirx.Function([], sequential_bindings)
+ var = T.dynamic("var", "int32")
+
+ @I.ir_module(check_well_formed=False)
+ class Before:
+ @T.function(private=True)
+ def main():
+ T.bind(16, var=var)
+ T.evaluate(var)
+ T.bind(32, var=var)
+ T.evaluate(var)
+
+ assert Before["main"].body[0].var.same_as(Before["main"].body[2].var)
@T.function(private=True)
def expected():
@@ -47,60 +46,37 @@ def test_reuse_in_sequential_bind():
var2 = T.bind(T.int32(32))
T.evaluate(var2)
- mod = tvm.IRModule.from_expr(before)
- mod = tvm.tirx.transform.ConvertSSA()(mod)
- tvm.ir.assert_structural_equal(mod["main"], expected)
+ After = tvm.tirx.transform.ConvertSSA()(Before)
+ tvm.ir.assert_structural_equal(After["main"], expected)
def test_reuse_in_nested_bind():
- """De-dup sequential bindings of the same variable.
+ """Redefinition changes all subsequent uses, including the final sibling."""
+ var = T.dynamic("var", "int32")
- In the flat Bind model, all Binds are siblings in a SeqStmt. A second
- Bind of the same variable redefines it for all subsequent siblings.
- ConvertSSA should create a new variable for the second binding and
- update all subsequent uses to refer to the new variable.
- """
+ @I.ir_module(check_well_formed=False)
+ class Before:
+ @T.function(private=True)
+ def main():
+ T.bind(32, var=var)
+ T.evaluate(var)
+ T.bind(16, var=var)
+ T.evaluate(var)
+ T.evaluate(var)
+
+ @I.ir_module
+ class Expected:
+ @T.function(private=True)
+ def main():
+ var1 = T.bind(32)
+ T.evaluate(var1)
+ var2 = T.bind(16)
+ T.evaluate(var2)
+ T.evaluate(var2)
- # Manually construct the Function body, as SSA violations are
- # not valid TIR, and may not be expressible in future versions
- # of TVMScript.
- var = tirx.Var("var", "int32")
- # Note: nested SeqStmt is flattened by the IR builder, so the input
- # is actually a flat SeqStmt with 5 elements.
- inner_seq = tvm.ir.SeqStmt(
- [
- tvm.ir.Bind(var, 16),
- tvm.ir.Evaluate(var),
- ]
- )
- outer_seq = tvm.ir.SeqStmt(
- [
- tvm.ir.Bind(var, 32),
- tvm.ir.Evaluate(var),
- inner_seq,
- tvm.ir.Evaluate(var),
- ]
- )
- before = tirx.Function([], outer_seq)
-
- # In the flat model, the second Bind(var, 16) redefines var for
- # ALL subsequent siblings including the last Evaluate.
- var1 = tirx.Var("var", "int32")
- var2 = tirx.Var("var", "int32")
- expected_body = tvm.ir.SeqStmt(
- [
- tvm.ir.Bind(var1, 32),
- tvm.ir.Evaluate(var1),
- tvm.ir.Bind(var2, 16),
- tvm.ir.Evaluate(var2),
- tvm.ir.Evaluate(var2),
- ]
- )
- expected = tirx.Function([], expected_body)
-
- mod = tvm.IRModule.from_expr(before)
- mod = tvm.tirx.transform.ConvertSSA()(mod)
- tvm.ir.assert_structural_equal(mod["main"], expected)
+ assert Before["main"].body[0].var.same_as(Before["main"].body[2].var)
+ After = tvm.tirx.transform.ConvertSSA()(Before)
+ tvm.ir.assert_structural_equal(After, Expected)
def test_reused_var_across_module():
@@ -111,12 +87,12 @@ def test_reused_var_across_module():
var = T.bind(10)
T.evaluate(var)
- before = tvm.IRModule(
- {
- "func_a": func.with_attr("global_symbol", "func_a"),
- "func_b": func.with_attr("global_symbol", "func_b"),
- }
- )
+ @I.ir_module(check_well_formed=False)
+ class before:
+ func_a = func.with_attr("global_symbol", "func_a")
+ func_b = func.with_attr("global_symbol", "func_b")
+
+ assert before["func_a"].body.same_as(before["func_b"].body)
@I.ir_module
class expected:
@@ -145,12 +121,12 @@ def test_reused_parameter():
def func(n: T.int32):
T.evaluate(n)
- before = tvm.IRModule(
- {
- "func_a": func.with_attr("global_symbol", "func_a"),
- "func_b": func.with_attr("global_symbol", "func_b"),
- }
- )
+ @I.ir_module(check_well_formed=False)
+ class before:
+ func_a = func.with_attr("global_symbol", "func_a")
+ func_b = func.with_attr("global_symbol", "func_b")
+
+ assert before["func_a"].body.same_as(before["func_b"].body)
@I.ir_module
class expected:
@@ -174,12 +150,12 @@ def test_reused_tensor_obj():
A = T.decl_tensor(a, ty_args=[T.Tensor(1, "float32", data=a)])
T.evaluate(A[0])
- before = tvm.IRModule(
- {
- "func_a": func.with_attr("global_symbol", "func_a"),
- "func_b": func.with_attr("global_symbol", "func_b"),
- }
- )
+ @I.ir_module(check_well_formed=False)
+ class before:
+ func_a = func.with_attr("global_symbol", "func_a")
+ func_b = func.with_attr("global_symbol", "func_b")
+
+ assert before["func_a"].body.same_as(before["func_b"].body)
@I.ir_module
class expected:
@@ -204,12 +180,12 @@ def test_reused_tensor_parameter():
def func(A: T.Tensor(1, "float32")):
T.evaluate(A[0])
- before = tvm.IRModule(
- {
- "func_a": func.with_attr("global_symbol", "func_a"),
- "func_b": func.with_attr("global_symbol", "func_b"),
- }
- )
+ @I.ir_module(check_well_formed=False)
+ class before:
+ func_a = func.with_attr("global_symbol", "func_a")
+ func_b = func.with_attr("global_symbol", "func_b")
+
+ assert before["func_a"].body.same_as(before["func_b"].body)
@I.ir_module
class expected:
@@ -227,17 +203,34 @@ def test_reused_tensor_parameter():
def test_reused_compound_tensor_shape_var():
"""De-duplicate implicit Vars nested in tensor parameter shapes."""
- n = tirx.Var("n", "int32")
- A = tirx.decl_tensor((tirx.max(n, 1),), layout=None)
- func = tirx.Function([A], tvm.ir.Evaluate(n))
- before = tvm.IRModule(
- {
- "func_a": func.with_attr("global_symbol", "func_a"),
- "func_b": func.with_attr("global_symbol", "func_b"),
- }
- )
+ n = T.dynamic("n", "int32")
+
+ @T.function(private=True)
+ def func(A: T.Tensor((T.max(n, 1),), "float32", layout=None)):
+ T.evaluate(n)
+
+ @I.ir_module(check_well_formed=False)
+ class before:
+ func_a = func.with_attr("global_symbol", "func_a")
+ func_b = func.with_attr("global_symbol", "func_b")
+
+ assert before["func_a"].body.same_as(before["func_b"].body)
+
+ n_a = T.dynamic("n", "int32")
+ n_b = T.dynamic("n", "int32")
+
+ @I.ir_module
+ class Expected:
+ @T.function
+ def func_a(A: T.Tensor((T.max(n_a, 1),), "float32", layout=None)):
+ T.evaluate(n_a)
+
+ @T.function
+ def func_b(A: T.Tensor((T.max(n_b, 1),), "float32", layout=None)):
+ T.evaluate(n_b)
after = tvm.tirx.transform.ConvertSSA()(before)
+ tvm.ir.assert_structural_equal(after, Expected)
func_a = after["func_a"]
func_b = after["func_b"]
n_a = func_a.params[0].shape[0].a
@@ -293,7 +286,7 @@ def test_de_duplicate_thread_idx_across_multiple_functions():
def kernel_1(A: T.Tensor([256], "float32")):
T.region(
"tirx.launch_thread",
- [tvm.ir.StringImm("threadIdx.x"), 256],
+ ["threadIdx.x", 256],
body_params=[threadIdx_x],
)
A[threadIdx_x] = A[threadIdx_x] + T.float32(1)
@@ -302,7 +295,7 @@ def test_de_duplicate_thread_idx_across_multiple_functions():
def kernel_2(A: T.Tensor([256], "float32")):
T.region(
"tirx.launch_thread",
- [tvm.ir.StringImm("threadIdx.x"), 256],
+ ["threadIdx.x", 256],
body_params=[threadIdx_x],
)
A[threadIdx_x] = A[threadIdx_x] + T.float32(1)
@@ -316,64 +309,7 @@ def test_de_duplicate_thread_idx_across_multiple_functions():
def kernel_1(A: T.Tensor([256], "float32")):
T.region(
"tirx.launch_thread",
- [tvm.ir.StringImm("threadIdx.x"), 256],
- body_params=[kernel_1_threadIdx_x],
- )
- A[kernel_1_threadIdx_x] = A[kernel_1_threadIdx_x] + T.float32(1)
-
- @T.function
- def kernel_2(A: T.Tensor([256], "float32")):
- T.region(
- "tirx.launch_thread",
- [tvm.ir.StringImm("threadIdx.x"), 256],
- body_params=[kernel_2_threadIdx_x],
- )
- A[kernel_2_threadIdx_x] = A[kernel_2_threadIdx_x] + T.float32(1)
-
- after = tvm.tirx.transform.ConvertSSA()(before)
- tvm.ir.assert_structural_equal(after, expected)
-
-
-def test_de_duplicate_thread_idx_iter_var_across_multiple_functions():
- """ConvertSSA separates a shared parameter list reused across functions.
-
- Launch regions no longer wrap their bindings in IterVars; this retains the
- existing shared-wrapper fixture using the explicit body-parameter list."""
-
- threadIdx_x = T.dynamic("threadIdx_x", "int32")
- body_params = [threadIdx_x]
-
- # complaints of multiple definitions for threadIdx_x
- @I.ir_module(check_well_formed=False)
- class before:
- @T.function
- def kernel_1(A: T.Tensor([256], "float32")):
- T.region(
- "tirx.launch_thread",
- [tvm.ir.StringImm("threadIdx.x"), 256],
- body_params=body_params,
- )
- A[threadIdx_x] = A[threadIdx_x] + T.float32(1)
-
- @T.function
- def kernel_2(A: T.Tensor([256], "float32")):
- T.region(
- "tirx.launch_thread",
- [tvm.ir.StringImm("threadIdx.x"), 256],
- body_params=body_params,
- )
- A[threadIdx_x] = A[threadIdx_x] + T.float32(1)
-
- kernel_1_threadIdx_x = T.dynamic("threadIdx_x", "int32")
- kernel_2_threadIdx_x = T.dynamic("threadIdx_x", "int32")
-
- @I.ir_module(check_well_formed=False)
- class expected:
- @T.function
- def kernel_1(A: T.Tensor([256], "float32")):
- T.region(
- "tirx.launch_thread",
- [tvm.ir.StringImm("threadIdx.x"), 256],
+ ["threadIdx.x", 256],
body_params=[kernel_1_threadIdx_x],
)
A[kernel_1_threadIdx_x] = A[kernel_1_threadIdx_x] + T.float32(1)
@@ -382,7 +318,7 @@ def test_de_duplicate_thread_idx_iter_var_across_multiple_functions():
def kernel_2(A: T.Tensor([256], "float32")):
T.region(
"tirx.launch_thread",
- [tvm.ir.StringImm("threadIdx.x"), 256],
+ ["threadIdx.x", 256],
body_params=[kernel_2_threadIdx_x],
)
A[kernel_2_threadIdx_x] = A[kernel_2_threadIdx_x] + T.float32(1)
@@ -407,13 +343,13 @@ def test_thread_idx_reused_within_and_across_functions():
def kernel_1(A: T.Tensor([256], "float32")):
with T.region(
"tirx.launch_thread",
- [tvm.ir.StringImm("threadIdx.x"), 256],
+ ["threadIdx.x", 256],
body_params=body_params,
):
A[threadIdx_x] = A[threadIdx_x] + 1.0
with T.region(
"tirx.launch_thread",
- [tvm.ir.StringImm("threadIdx.x"), 256],
+ ["threadIdx.x", 256],
body_params=body_params,
):
A[threadIdx_x] = A[threadIdx_x] + 2.0
@@ -422,13 +358,13 @@ def test_thread_idx_reused_within_and_across_functions():
def kernel_2(A: T.Tensor([256], "float32")):
with T.region(
"tirx.launch_thread",
- [tvm.ir.StringImm("threadIdx.x"), 256],
+ ["threadIdx.x", 256],
body_params=body_params,
):
A[threadIdx_x] = A[threadIdx_x] + 1.0
with T.region(
"tirx.launch_thread",
- [tvm.ir.StringImm("threadIdx.x"), 256],
+ ["threadIdx.x", 256],
body_params=body_params,
):
A[threadIdx_x] = A[threadIdx_x] + 2.0
@@ -462,70 +398,85 @@ def test_shared_shape_var_in_tensor_params_and_alloc_tensor():
function body (including AllocTensor shapes) must remain the same
Var object so that MakePackedAPI can bind it from the DLTensor shape.
"""
- n = tirx.Var("n", "int32")
- A = tirx.decl_tensor((n,), "float32", "A")
- B = tirx.decl_tensor((n,), "float32", "B")
-
- # AllocTensor with shape [n] in the body (flat, no body)
- C = tirx.decl_tensor((n,), "float32", "C")
- body = tvm.ir.SeqStmt(
- [
- tvm.ir.Bind(
- C,
- tvm.ir.Call(
- "tirx.alloc_tensor",
- [],
- attrs=tvm.ir.DictAttrs({}),
- ty_args=[C.ty],
- ),
- ),
- tvm.ir.Evaluate(1),
- ]
- )
-
- before = tirx.Function([A, B], body)
-
- mod = tvm.IRModule.from_expr(before)
- after = tvm.tirx.transform.ConvertSSA()(mod)
- # The function is already SSA — ConvertSSA should not change it.
- tvm.ir.assert_structural_equal(after["main"], before)
+ n = T.dynamic("n", "int32")
+
+ @I.ir_module
+ class Before:
+ @T.function(private=True)
+ def main(A: T.Tensor((n,), "float32"), B: T.Tensor((n,), "float32")):
+ C = T.alloc_tensor(ty_args=[T.Tensor((n,), "float32")])
+ T.evaluate(1)
+
+ After = tvm.tirx.transform.ConvertSSA()(Before)
+ tvm.ir.assert_structural_equal(After, Before)
+ assert After.same_as(Before)
def test_reused_loop_var_in_decl_tensor_elem_offset():
"""Remap a tensor whose elem_offset depends on an SSA-renamed loop var."""
- loop_var = tirx.Var("loop_var", "int32")
- tensor = tirx.decl_tensor(
- (128,),
- "float32",
- "tensor",
- elem_offset=loop_var * 128,
- scope="shared.dyn",
- )
- buffer_data = tirx.Var("buffer_data", tensor.data.ty)
- loop = tvm.ir.For(
- loop_var,
- 0,
- 128,
- tvm.ir.ForKind.DEFAULT,
- tvm.ir.SeqStmt(
- [
- tvm.ir.Bind(
- tensor,
- tvm.ir.Call(
- "tirx.decl_tensor",
- [buffer_data],
- ty_args=[tensor.ty],
- ),
- ),
- tvm.ir.Evaluate(tirx.TensorLoad(tensor, [0])),
- ]
- ),
- )
- func = tirx.Function([buffer_data], tvm.ir.SeqStmt([loop, loop, loop]))
-
- after = tvm.tirx.transform.ConvertSSA()(tvm.IRModule.from_expr(func))
-
- tvm.tirx.analysis.verify_well_formed(after["main"], assert_mode=True)
+
+ @T.function(private=True)
+ def template(data: T.handle("float32", "shared.dyn")):
+ for i in range(128):
+ tensor = T.decl_tensor(
+ data,
+ ty_args=[
+ T.Tensor((128,), "float32", data=data, elem_offset=i * 128, scope="shared.dyn")
+ ],
+ )
+ T.evaluate(tensor[0])
+
+ shared_loop = template.body[0]
+ shared_data = template.params[0]
+
+ @I.ir_module(check_well_formed=False)
+ class Before:
+ @T.function(private=True)
+ def main(data: shared_data):
+ T.add_to_parent(shared_loop)
+ T.add_to_parent(shared_loop)
+ T.add_to_parent(shared_loop)
+
+ @I.ir_module
+ class Expected:
+ @T.function(private=True)
+ def main(data: T.handle("float32", "shared.dyn")):
+ for i in range(128):
+ tensor = T.decl_tensor(
+ data,
+ ty_args=[
+ T.Tensor(
+ (128,), "float32", data=data, elem_offset=i * 128, scope="shared.dyn"
+ )
+ ],
+ )
+ T.evaluate(tensor[0])
+ for j in range(128):
+ tensor = T.decl_tensor(
+ data,
+ ty_args=[
+ T.Tensor(
+ (128,), "float32", data=data, elem_offset=j * 128, scope="shared.dyn"
+ )
+ ],
+ )
+ T.evaluate(tensor[0])
+ for k in range(128):
+ tensor = T.decl_tensor(
+ data,
+ ty_args=[
+ T.Tensor(
+ (128,), "float32", data=data, elem_offset=k * 128, scope="shared.dyn"
+ )
+ ],
+ )
+ T.evaluate(tensor[0])
+
+ assert Before["main"].body[0].same_as(Before["main"].body[1])
+ assert Before["main"].body[1].same_as(Before["main"].body[2])
+ After = tvm.tirx.transform.ConvertSSA()(Before)
+ tvm.ir.assert_structural_equal(After, Expected)
+ tvm.tirx.analysis.verify_well_formed(After["main"], assert_mode=True)
if __name__ == "__main__":
diff --git a/tests/python/tirx/transform/test_transform_decorate_device_scope.py b/tests/python/tirx/transform/test_tirx_transform_decorate_device_scope.py
similarity index 61%
rename from tests/python/tirx/transform/test_transform_decorate_device_scope.py
rename to tests/python/tirx/transform/test_tirx_transform_decorate_device_scope.py
index 9e794cae5d..6a3bc1c8ad 100644
--- a/tests/python/tirx/transform/test_transform_decorate_device_scope.py
+++ b/tests/python/tirx/transform/test_tirx_transform_decorate_device_scope.py
@@ -15,15 +15,28 @@
# specific language governing permissions and limitations
# under the License.
import tvm
+import tvm.testing
+from tvm.script import ir as I
+from tvm.script import tirx as T
def test_decorate_device():
- x = tvm.tirx.Var("x", "int32")
- mod = tvm.IRModule.from_expr(tvm.tirx.Function([x], tvm.ir.Evaluate(x)))
+ @I.ir_module
+ class Before:
+ @T.function(private=True)
+ def main(x: T.int32):
+ T.evaluate(x)
- stmt = tvm.tirx.transform.DecorateDeviceScope()(mod)["main"].body[0]
- assert stmt.op.same_as(tvm.ir.Op.get("tirx.device_scope"))
+ @I.ir_module
+ class Expected:
+ @T.function(private=True)
+ def main(x: T.int32):
+ with T.region("tirx.device_scope", []):
+ T.evaluate(x)
+
+ After = tvm.tirx.transform.DecorateDeviceScope()(Before)
+ tvm.ir.assert_structural_equal(After, Expected)
if __name__ == "__main__":
- test_decorate_device()
+ tvm.testing.main()
diff --git a/tests/python/tirx-transform/test_tir_transform_flatten_tensor.py b/tests/python/tirx/transform/test_tirx_transform_flatten_tensor.py
similarity index 63%
rename from tests/python/tirx-transform/test_tir_transform_flatten_tensor.py
rename to tests/python/tirx/transform/test_tirx_transform_flatten_tensor.py
index 14fc6f8865..298cf1b034 100644
--- a/tests/python/tirx-transform/test_tir_transform_flatten_tensor.py
+++ b/tests/python/tirx/transform/test_tirx_transform_flatten_tensor.py
@@ -16,10 +16,13 @@
# under the License.
from __future__ import annotations
+import pytest
+
import tvm
import tvm.testing
from tvm.script import ir as I
from tvm.script import tirx as T
+from tvm.tirx.transform import FlattenTensor
def _transform():
@@ -419,20 +422,253 @@ def test_flatten_inside_block():
tvm.ir.assert_structural_equal(After, Expected)
-def test_build_with_optional_unroll_explicit():
- def check(value):
- @I.ir_module
- class Module:
- @T.function
- def main(A: T.Tensor((4, 5, 6), "int16"), B: T.Tensor((4, 5, 6), "int16")):
- for ax0 in T.serial(4, annotations={"unroll_explicit": value}):
- for ax1, ax2 in T.grid(5, 6):
- B[ax0, ax1, ax2] = A[ax0, ax1, ax2]
+@pytest.mark.parametrize("value", [None, True, False])
+def test_build_with_optional_unroll_explicit(value):
+ @I.ir_module
+ class Before:
+ @T.function
+ def main(A: T.Tensor((4, 5, 6), "int16"), B: T.Tensor((4, 5, 6), "int16")):
+ for ax0 in T.serial(4, annotations={"unroll_explicit": value}):
+ for ax1, ax2 in T.grid(5, 6):
+ B[ax0, ax1, ax2] = A[ax0, ax1, ax2]
+
+ @I.ir_module
+ class Expected:
+ @T.function
+ def main(A: T.Tensor((4, 5, 6), "int16"), B: T.Tensor((4, 5, 6), "int16")):
+ A_1 = T.decl_tensor(
+ A.data, ty_args=[T.Tensor((120,), "int16", layout=None, data=A.data)]
+ )
+ B_1 = T.decl_tensor(
+ B.data, ty_args=[T.Tensor((120,), "int16", layout=None, data=B.data)]
+ )
+ for ax0 in T.serial(4, annotations={"unroll_explicit": value}):
+ for ax1, ax2 in T.grid(5, 6):
+ B_1[ax0 * 30 + ax1 * 6 + ax2] = A_1[ax0 * 30 + ax1 * 6 + ax2]
+
+ After = FlattenTensor()(Before)
+ tvm.ir.assert_structural_equal(After, Expected)
+ tvm.tirx.compile(Before, target="llvm").mod
+
+
+def test_flatten_remaps_loads_in_view_shape():
+ @I.ir_module
+ class Before:
+ @T.function(private=True)
+ def main():
+ n = T.alloc_local([1], "int32")
+ n[0] = 8
+ data = T.alloc_shared([64], "float16")
+ view = T.decl_tensor(data.data, ty_args=[T.Tensor((n[0],), "float16", scope="shared")])
+ view[0] = T.float16(0)
+
+ @I.ir_module
+ class Expected:
+ @T.function(private=True)
+ def main():
+ n = T.alloc_local([1], "int32", layout=None)
+ n[0] = 8
+ data = T.alloc_shared([64], "float16", layout=None)
+ view = T.decl_tensor(
+ data.data, ty_args=[T.Tensor((n[0],), "float16", scope="shared", layout=None)]
+ )
+ view[0] = T.float16(0)
+
+ After = FlattenTensor()(Before)
+ tvm.ir.assert_structural_equal(After, Expected)
+ body = After["main"].body
+ assert body[3].var.shape[0].source.same_as(body[0].var)
+ assert not tvm.tirx.analysis.undefined_vars(body, [])
+
+
+def test_flatten_remaps_loads_in_folded_elem_offset():
+ @I.ir_module
+ class Before:
+ @T.function(private=True)
+ def main():
+ n = T.alloc_local([1], "int32")
+ n[0] = 4
+ base = T.alloc_shared([128], "uint64")
+ mbar = T.decl_tensor(
+ base.data, ty_args=[T.Tensor((1,), "uint64", elem_offset=n[0], scope="shared")]
+ )
+ mbar[0] = T.uint64(1)
+
+ @I.ir_module
+ class Expected:
+ @T.function(private=True)
+ def main():
+ n = T.alloc_local([1], "int32", layout=None)
+ n[0] = 4
+ base = T.alloc_shared([128], "uint64", layout=None)
+ mbar = T.decl_tensor(
+ base.data,
+ ty_args=[
+ T.Tensor((1,), "uint64", elem_offset=T.int32(0), scope="shared", layout=None)
+ ],
+ )
+ mbar[n[0]] = T.uint64(1)
+
+ After = FlattenTensor()(Before)
+ tvm.ir.assert_structural_equal(After, Expected)
+ body = After["main"].body
+ assert body[-1].indices[0].source.same_as(body[0].var)
+ assert not tvm.tirx.analysis.undefined_vars(body, [])
+
+
+def test_flatten_keeps_identity_of_already_flat_tensors():
+ @I.ir_module
+ class Before:
+ @T.function(private=True)
+ def main():
+ flat = T.alloc_shared([32], "float32", layout=None)
+ flat[0] = T.float32(0)
+
+ Expected = Before
+ After = FlattenTensor()(Before)
+ tvm.ir.assert_structural_equal(After, Expected)
+ assert After["main"].body[0].var.same_as(Before["main"].body[0].var)
+
+
+def test_tensor_flatten_offset():
+ m = T.dynamic("m", "int32")
+ n = T.dynamic("n", "int32")
+
+ @I.ir_module
+ class Before:
+ @T.function(private=True)
+ def main(m: T.int32, n: T.int32, A: T.Tensor((m, n), "float32", elem_offset=T.int32(100))):
+ T.evaluate(A[2, 3])
+
+ @I.ir_module
+ class Expected:
+ @T.function(private=True)
+ def main(m: T.int32, n: T.int32, A: T.Tensor((m, n), "float32", elem_offset=T.int32(100))):
+ flat = T.decl_tensor(
+ A.data, ty_args=[T.Tensor((m * n,), "float32", elem_offset=T.int32(0), layout=None)]
+ )
+ T.evaluate(flat[n * 2 + 103])
+
+ After = FlattenTensor()(Before)
+ tvm.ir.assert_structural_equal(After, Expected)
+
+
+@pytest.mark.parametrize(
+ "shape,strides,span",
+ [([16, 32], None, 512), ([16], None, 16), ([], None, 1), ([16, 32], [40, 1], 640)],
+ ids=["matrix", "vector", "scalar", "strided"],
+)
+def test_tensor_flatten(shape, strides, span):
+ @I.ir_module
+ class Before:
+ @T.function(private=True)
+ def main(A: T.Tensor(shape, "float32", strides=strides)):
+ T.evaluate(A[tuple(0 for _ in shape)])
+
+ @I.ir_module
+ class Expected:
+ @T.function(private=True)
+ def main(A: T.Tensor(shape, "float32", strides=strides)):
+ flat = T.decl_tensor(A.data, ty_args=[T.Tensor((span,), "float32", layout=None)])
+ T.evaluate(flat[0])
+
+ After = FlattenTensor()(Before)
+ tvm.ir.assert_structural_equal(After, Expected)
+ source = After["main"].params[0]
+ flat = After["main"].body[0].var
+ assert not source.same_as(flat)
+ assert flat.data.args[0].same_as(flat)
+ assert flat.data.op.name == "tirx.tensor_data_ptr"
+
+
+@pytest.mark.parametrize(
+ "pattern", [1, 2, 3, 4], ids=["strided", "nested-mod", "split-sum", "unmerged"]
+)
+def test_tensor_index_merge_mult_mod(pattern):
+ m = T.dynamic("m", "int32")
+ n = T.dynamic("n", "int32")
+ s = T.dynamic("s", "int32")
+
+ def tensor_strides(s):
+ return (s, 1) if pattern == 1 else None
+
+ def storage_span(m, n, s):
+ return T.max(s * m, (m - 1) * s + n) if pattern == 1 else m * n
+
+ def indices(k0, k1, s, n):
+ if pattern == 1:
+ return ((k0 % k1) // s, (k0 % k1) % s + (k0 // k1) * k1)
+ if pattern == 2:
+ return ((k0 % (k1 // s)) // n, (k0 % (k1 // s)) % n + k0 % k1)
+ if pattern == 3:
+ q = (k0 // (k1 // s)) * (k1 // s)
+ return (q // n + (k0 % (k1 // s)) // n, q % n + (k0 % (k1 // s)) % n)
+ return ((k0 % (k1 // s)) // n, (k0 % (k1 // n)) % n + k0 % k1)
+
+ def expected_index(k0, k1, s, n):
+ if pattern in (1, 3):
+ return k0
+ if pattern == 2:
+ return k0 % (k1 // s) + k0 % k1
+ return ((k0 % (k1 // s)) // n) * n + (k0 % (k1 // n)) % n + k0 % k1
+
+ @I.ir_module
+ class Before:
+ @T.function(private=True)
+ def main(
+ m: T.int32,
+ n: T.int32,
+ s: T.int32,
+ k0: T.int32,
+ k1: T.int32,
+ A: T.Tensor((m, n), "float32", strides=tensor_strides(s)),
+ ):
+ T.evaluate(A[indices(k0, k1, s, n)])
+
+ @I.ir_module
+ class Expected:
+ @T.function(private=True)
+ def main(
+ m: T.int32,
+ n: T.int32,
+ s: T.int32,
+ k0: T.int32,
+ k1: T.int32,
+ A: T.Tensor((m, n), "float32", strides=tensor_strides(s)),
+ ):
+ flat = T.decl_tensor(
+ A.data, ty_args=[T.Tensor((storage_span(m, n, s),), "float32", layout=None)]
+ )
+ T.evaluate(flat[expected_index(k0, k1, s, n)])
+
+ After = FlattenTensor()(Before)
+ tvm.ir.assert_structural_equal(After, Expected)
+
+
+@pytest.mark.parametrize("factored", [False, True], ids=["linear", "factored"])
+def test_tensor_index_merge_rank_four(factored):
+ def indices(i, j, k):
+ if factored:
+ q = i * 49 + j * 28 + k // 1024
+ return (q // 14 // 14, (q // 14) % 14, (i * 7 + k // 1024) % 14, k % 1024)
+ q = i * 50176 + j * 28672 + k
+ return (q // 1024 // 14 // 14, (q // 1024 // 14) % 14, (q // 1024) % 14, q % 1024)
- tvm.tirx.compile(Module, target="llvm").mod
+ @I.ir_module
+ class Before:
+ @T.function(private=True)
+ def main(i: T.int32, j: T.int32, k: T.int32, A: T.Tensor((1, 14, 14, 1024), "float32")):
+ T.evaluate(A[indices(i, j, k)])
- for value in [None, True, False]:
- check(value)
+ @I.ir_module
+ class Expected:
+ @T.function(private=True)
+ def main(i: T.int32, j: T.int32, k: T.int32, A: T.Tensor((1, 14, 14, 1024), "float32")):
+ flat = T.decl_tensor(A.data, ty_args=[T.Tensor((200704,), "float32", layout=None)])
+ T.evaluate(flat[i * 50176 + j * 28672 + k])
+
+ After = FlattenTensor()(Before)
+ tvm.ir.assert_structural_equal(After, Expected)
if __name__ == "__main__":
diff --git a/tests/python/tirx-transform/test_tir_transform_force_narrow_index_to_i32.py b/tests/python/tirx/transform/test_tirx_transform_force_narrow_index_to_i32.py
similarity index 98%
rename from tests/python/tirx-transform/test_tir_transform_force_narrow_index_to_i32.py
rename to tests/python/tirx/transform/test_tirx_transform_force_narrow_index_to_i32.py
index 1f4e12d30b..df8837a264 100644
--- a/tests/python/tirx-transform/test_tir_transform_force_narrow_index_to_i32.py
+++ b/tests/python/tirx/transform/test_tirx_transform_force_narrow_index_to_i32.py
@@ -18,6 +18,7 @@ import pytest
import tvm
import tvm.testing
+from tvm.script import ir as I
from tvm.script import tirx as T
@@ -39,7 +40,10 @@ def test_thread_axis1():
threadIdx_x = T.launch_thread("threadIdx.x", 32)
B[blockIdx_x * 32 + threadIdx_x] = A[blockIdx_x * 32 + threadIdx_x] + T.float32(1)
- mod = tvm.IRModule.from_expr(before)
+ @I.ir_module
+ class mod:
+ main = before
+
func = tvm.tirx.transform.ForceNarrowIndexToInt32()(mod)["main"]
tvm.ir.assert_structural_equal(func, expected)
@@ -226,7 +230,10 @@ def test_thread_axis2():
],
)
- mod = tvm.IRModule.from_expr(before)
+ @I.ir_module
+ class mod:
+ main = before
+
func = tvm.tirx.transform.ForceNarrowIndexToInt32()(mod)["main"]
tvm.ir.assert_structural_equal(func, expected)
@@ -247,7 +254,10 @@ def test_block():
for j in T.serial(0, T.int32(8)):
B[i * T.int32(8) + j] = A[i * T.int32(8) + j] + T.float32(1)
- mod = tvm.IRModule.from_expr(before)
+ @I.ir_module
+ class mod:
+ main = before
+
func = tvm.tirx.transform.ForceNarrowIndexToInt32()(mod)["main"]
tvm.ir.assert_structural_equal(func, expected)
@@ -268,7 +278,10 @@ def test_i16_tensor():
for j in T.serial(0, 16):
B[i * 8 + j] = A[i * 8 + j] + T.int16(1)
- mod = tvm.IRModule.from_expr(before)
+ @I.ir_module
+ class mod:
+ main = before
+
after = tvm.tirx.transform.ForceNarrowIndexToInt32()(mod)["main"]
tvm.ir.assert_structural_equal(after, expected)
@@ -280,7 +293,10 @@ def test_fail_on_tensor_param():
for j in T.serial(0, 8):
B[i * 8 + j] = A[i * 8 + j] + T.int64(1)
- mod = tvm.IRModule.from_expr(func)
+ @I.ir_module
+ class mod:
+ main = func
+
with pytest.raises(RuntimeError):
tvm.tirx.transform.ForceNarrowIndexToInt32()(mod)["main"]
@@ -296,7 +312,10 @@ def test_fail_on_internal_tensor():
for j in T.serial(0, 8):
B[i * 8 + j] = T.cast(C[i * 8 + j] + T.int64(1), "int32")
- mod = tvm.IRModule.from_expr(func)
+ @I.ir_module
+ class mod:
+ main = func
+
with pytest.raises(RuntimeError):
tvm.tirx.transform.ForceNarrowIndexToInt32()(mod)["main"]
diff --git a/tests/python/tirx-transform/test_tir_transform_fp8_legalize.py b/tests/python/tirx/transform/test_tirx_transform_fp8_legalize.py
similarity index 98%
rename from tests/python/tirx-transform/test_tir_transform_fp8_legalize.py
rename to tests/python/tirx/transform/test_tirx_transform_fp8_legalize.py
index c6cc3e64cc..8577659a34 100644
--- a/tests/python/tirx-transform/test_tir_transform_fp8_legalize.py
+++ b/tests/python/tirx/transform/test_tirx_transform_fp8_legalize.py
@@ -17,6 +17,7 @@
import tvm
import tvm.script
import tvm.testing
+from tvm.script import ir as I
from tvm.script import tirx as T
from tvm.target import Target
from tvm.tirx.transform.transform import BindTarget
@@ -222,7 +223,13 @@ def test_fp8_compute_legalize_preserves_opaque_tensor_access(dtype, promote_dtyp
tensor = T.alloc_tensor(ty_args=[T.Tensor((16,), dtype)])
T.evaluate(T.call_extern("consume", tensor.data, ty="void"))
- before_mod = tvm.IRModule.from_expr(before)
+ input_func = before
+
+ @I.ir_module(check_well_formed=False)
+ class InputModule:
+ main = input_func
+
+ before_mod = InputModule
after = tvm.tirx.transform.FP8ComputeLegalize(promote_dtype)(before_mod)
tvm.ir.assert_structural_equal(after, before_mod)
diff --git a/tests/python/tirx/transform/test_tirx_transform_function_pass.py b/tests/python/tirx/transform/test_tirx_transform_function_pass.py
new file mode 100644
index 0000000000..3951b8b698
--- /dev/null
+++ b/tests/python/tirx/transform/test_tirx_transform_function_pass.py
@@ -0,0 +1,85 @@
+# 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.
+import tvm
+import tvm.testing
+from tvm.ir.transform.instrument import PrintAfterAll, PrintBeforeAll
+from tvm.script import ir as I
+from tvm.script import tirx as T
+
+
+def test_function_pass(capsys):
+ @I.ir_module
+ class Before:
+ @T.function(private=True)
+ def main(x: T.int32, y: T.int32, B: T.Tensor((16,), "float32")):
+ value = T.bind(10)
+ T.evaluate(value + 1)
+
+ @I.ir_module
+ class Expected:
+ @T.function(private=True)
+ def main(x: T.int32, y: T.int32, B: T.Tensor((16,), "float32")):
+ T.evaluate(0)
+
+ @tvm.tirx.transform.function_pass(opt_level=1)
+ class TestReplaceFunc:
+ def __init__(self, new_func):
+ self.new_func = new_func
+
+ def transform_function(self, func, mod, ctx):
+ return self.new_func
+
+ with tvm.ir.transform.PassContext(instruments=[PrintBeforeAll(), PrintAfterAll()]):
+ After = TestReplaceFunc(Expected["main"])(Before)
+ output = capsys.readouterr().out
+ assert "Before Running Pass:" in output
+ assert "After Running Pass:" in output
+ assert 'name="TestReplaceFunc"' in output
+ tvm.ir.assert_structural_equal(After, Expected)
+
+
+def test_cow_pass():
+ @I.ir_module
+ class Before:
+ @T.function(private=True)
+ def main(x: T.int32):
+ T.func_attr({"target_bits": 32})
+ T.evaluate(x)
+
+ @I.ir_module
+ class Expected:
+ @T.function(private=True)
+ def main(x: T.int32):
+ T.func_attr({"target_bits": 32})
+ T.evaluate(x)
+
+ def fapply(func):
+ assert tvm.testing.object_use_count(func) == 1
+ return func
+
+ func_hash = hash(Before["main"])
+ mod_hash = hash(Before)
+ After = tvm.ir.transform.Sequential(
+ [tvm.tirx.transform.Apply(fapply), tvm.tirx.transform.NarrowDataType(32)]
+ )(Before._move())
+ tvm.ir.assert_structural_equal(After, Expected)
+ assert mod_hash == hash(After)
+ assert func_hash == hash(After["main"])
+
+
+if __name__ == "__main__":
+ tvm.testing.main()
diff --git a/tests/python/tirx-transform/test_tir_transform_helpers.py b/tests/python/tirx/transform/test_tirx_transform_helpers.py
similarity index 76%
rename from tests/python/tirx-transform/test_tir_transform_helpers.py
rename to tests/python/tirx/transform/test_tirx_transform_helpers.py
index 7b14e02713..b0c6a337d1 100644
--- a/tests/python/tirx-transform/test_tir_transform_helpers.py
+++ b/tests/python/tirx/transform/test_tirx_transform_helpers.py
@@ -14,8 +14,6 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-# ruff: noqa: F841
-import pytest
import tvm
import tvm.testing
@@ -24,67 +22,37 @@ from tvm.script import tirx as T
def test_annotate_entry_func_single_function():
- @tvm.script.ir_module
- class MockModule:
+ @I.ir_module
+ class Before:
@T.function(private=True)
- def func1(A: T.Tensor((16,), "float32")):
- for i in T.serial(16):
- if i == 5:
- if i == 5:
- A[i] = 0.0
-
- mod = MockModule
- assert mod
- assert "tirx.is_entry_func" not in (mod["func1"].attrs or {})
- after = tvm.tirx.transform.AnnotateEntryFunc()(mod)
- assert (
- after["func1"].attrs
- and "tirx.is_entry_func" in after["func1"].attrs
- and after["func1"].attrs["tirx.is_entry_func"]
- )
-
+ def main():
+ T.evaluate(0)
-# Test module
-@tvm.script.ir_module
-class MockModule:
- @T.function(private=True)
- def func1(A: T.Tensor((16,), "float32")):
- for i in T.serial(16):
- if i == 5:
- if i == 5:
- A[i] = 0.0
+ @I.ir_module
+ class Expected:
+ @T.function(private=True)
+ def main():
+ T.func_attr({"tirx.is_entry_func": True})
+ T.evaluate(0)
- @T.function(private=True)
- def func2(A: T.Tensor((32,), "float32")):
- for i in T.serial(32):
- if i == 15:
- if i == 15:
- A[i] = 0.0
+ assert "tirx.is_entry_func" not in (Before["main"].attrs or {})
+ After = tvm.tirx.transform.AnnotateEntryFunc()(Before)
+ tvm.ir.assert_structural_equal(After, Expected)
-@pytest.mark.xfail
def test_annotate_entry_func_multiple_function():
- mod = MockModule
- assert mod
- assert "target" not in (mod["func1"].attrs or {})
- assert "target" not in (mod["func2"].attrs or {})
- # This should fail
- after = tvm.tirx.transform.AnnotateEntryFunc()(mod)
-
-
-def test_bind_target():
- mod = MockModule
- assert mod
+ @I.ir_module
+ class Before:
+ @T.function(private=True)
+ def func1():
+ T.evaluate(0)
- target = tvm.target.Target("cuda")
- assert "target" not in (mod["func1"].attrs or {})
- assert "target" not in (mod["func2"].attrs or {})
- after = tvm.tirx.transform.BindTarget(target)(mod)
+ @T.function(private=True)
+ def func2():
+ T.evaluate(0)
- assert "target" in after["func1"].attrs
- assert after["func1"].attrs["target"] == target
- assert "target" in after["func2"].attrs
- assert after["func2"].attrs["target"] == target
+ After = tvm.tirx.transform.AnnotateEntryFunc()(Before)
+ tvm.ir.assert_structural_equal(After, Before)
def test_bind_target_adds_attribute():
@@ -224,6 +192,7 @@ def test_bind_target_multiple_functions():
T.func_attr({"target": T.target("cuda")})
T.evaluate(0)
+ assert all("target" not in (func.attrs or {}) for func in Before.functions.values())
After = tvm.tirx.transform.BindTarget(tvm.target.Target("cuda"))(Before)
tvm.ir.assert_structural_equal(After, Expected)
@@ -340,37 +309,36 @@ def test_bind_target_with_tirx_device_entry():
def test_filter_function():
- mod = MockModule
- assert mod
- # Annotate each function for testing
- mod["func1"] = mod["func1"].with_attr("temp", "test1")
- mod["func2"] = mod["func2"].with_attr("temp", "test2")
-
- # Test condition that does not filter out anything
- def checker_filter_out_none(func: tvm.tirx.Function):
- return "temp" in func.attrs
-
- after = tvm.tirx.transform.Filter(checker_filter_out_none)(mod)
- assert len(after.functions) == 2
- # Filtered functions should satisfy the given condition.
- assert checker_filter_out_none(after["func1"])
- assert checker_filter_out_none(after["func2"])
-
- # Test condition that selectively filters out functions
- def checker_filter_out_one(func: tvm.tirx.Function):
- return ("temp" in func.attrs) and func.attrs["temp"] == "test1"
-
- after = tvm.tirx.transform.Filter(checker_filter_out_one)(mod)
- assert len(after.functions) == 1
- # Filtered functions should satisfy the given condition.
- assert checker_filter_out_one(after["func1"])
-
- # Test condition that filters out everything
- def checker_filter_out_both(func: tvm.tirx.Function):
- return "invalid_attr" in func.attrs
-
- after = tvm.tirx.transform.Filter(checker_filter_out_both)(mod)
- assert len(after.functions) == 0
+ @I.ir_module
+ class Before:
+ @T.function(private=True)
+ def func1():
+ T.func_attr({"temp": "test1"})
+ T.evaluate(0)
+
+ @T.function(private=True)
+ def func2():
+ T.func_attr({"temp": "test2"})
+ T.evaluate(0)
+
+ @I.ir_module
+ class One:
+ @T.function(private=True)
+ def func1():
+ T.func_attr({"temp": "test1"})
+ T.evaluate(0)
+
+ @I.ir_module
+ class Empty:
+ pass
+
+ for predicate, Expected in [
+ (lambda func: "temp" in func.attrs, Before),
+ (lambda func: func.attrs["temp"] == "test1", One),
+ (lambda func: "invalid_attr" in func.attrs, Empty),
+ ]:
+ After = tvm.tirx.transform.Filter(predicate)(Before)
+ tvm.ir.assert_structural_equal(After, Expected)
def test_filter_removes_global_var_map():
diff --git a/tests/python/tirx/transform/test_transform_hoist_if.py b/tests/python/tirx/transform/test_tirx_transform_hoist_if.py
similarity index 100%
rename from tests/python/tirx/transform/test_transform_hoist_if.py
rename to tests/python/tirx/transform/test_tirx_transform_hoist_if.py
diff --git a/tests/python/tirx/transform/test_tirx_transform_lower_intrin.py b/tests/python/tirx/transform/test_tirx_transform_lower_intrin.py
new file mode 100644
index 0000000000..33c7bb5e97
--- /dev/null
+++ b/tests/python/tirx/transform/test_tirx_transform_lower_intrin.py
@@ -0,0 +1,236 @@
+# 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.
+import itertools
+
+import pytest
+
+import tvm
+import tvm.testing
+from tvm.script import ir as I
+from tvm.script import tirx as T
+from tvm.testing import env
+
+
+def test_nested_byte_pointer_offset():
+ @I.ir_module
+ class Before:
+ @T.function
+ def main(data: T.handle("float32")):
+ T.func_attr({"target": T.target("llvm")})
+ T.call_extern(
+ "consume",
+ T.ptr_byte_offset(T.ptr_byte_offset(data, 8, ty=data.ty), 12, ty=data.ty),
+ ty="void",
+ )
+
+ After = tvm.tirx.transform.LowerIntrin()(Before)
+ tvm.ir.assert_structural_equal(After, Before)
+ for target in ["c"] + (["llvm"] if env.has_llvm() else []):
+ tvm.tirx.compile(Before, target=target).mod
+
+
+@pytest.mark.skipif(not env.has_llvm(), reason="need llvm")
+@pytest.mark.parametrize("dtype, byte_offset", [("float32x2", 16), ("float32x3", 24)])
+def test_vector_byte_pointer_offset(dtype, byte_offset):
+ @I.ir_module
+ class Before:
+ @T.function
+ def main(A: T.Tensor((8,), dtype)):
+ T.func_attr({"target": T.target("llvm")})
+ T.call_extern(
+ "consume", T.ptr_byte_offset(A.data, byte_offset, ty=A.data.ty), ty="void"
+ )
+
+ After = tvm.tirx.transform.LowerIntrin()(Before)
+ tvm.ir.assert_structural_equal(After, Before)
+ tvm.tirx.compile(Before, target="llvm").mod
+
+
+@pytest.mark.parametrize("shape, indices", [((), []), ((16,), [3]), ((2, 4), [1, 2])])
+def test_logical_pointer_preserves_tensor_identity(shape, indices):
+ @I.ir_module
+ class Before:
+ @T.function
+ def main(A: T.Tensor(shape, "float32")):
+ T.func_attr({"target": T.target("llvm")})
+ T.evaluate(A.ptr_to(indices))
+
+ @I.ir_module
+ class Expected:
+ @T.function
+ def main(A: T.Tensor(shape, "float32")):
+ T.func_attr({"target": T.target("llvm")})
+ T.evaluate(T.address_of(A[tuple(indices)]))
+
+ After = tvm.tirx.transform.LowerIntrin()(Before)
+ tvm.ir.assert_structural_equal(After, Expected)
+ assert After["main"].body[0].value.args[0].source.same_as(Before["main"].params[0])
+
+
+def _divmod_expression(x, y, dtype, kind, variant, lowered):
+ """The signed remainder correction shared by floor division and modulo."""
+ scalar = getattr(T, dtype)
+ zero = scalar(0)
+ operation = T.floordiv if kind == "div" else T.floormod
+ if not lowered:
+ if variant == "pow2":
+ return operation(x, scalar(8))
+ if variant == "offset":
+ return operation(x + scalar(4), scalar(5))
+ value = operation(x, y)
+ if variant == "max":
+ value = T.max(value, zero)
+ if variant in ["positive_rhs", "max"]:
+ return T.Select(y >= zero, value, zero)
+ if variant == "positive_both":
+ return T.Select(T.And(y >= zero, x >= zero), value, zero)
+ return value
+
+ bits = int(dtype[3:])
+ if variant == "pow2":
+ return x >> scalar(3) if kind == "div" else x & scalar(7)
+ if variant == "plain":
+ rmod = T.dynamic("rmod", dtype)
+ same_sign = T.Or(T.And(y >= zero, rmod >= zero), T.And(y < zero, rmod <= zero))
+ if kind == "mod":
+ return T.Let(T.Select(same_sign, rmod, rmod + y), where={rmod: T.truncmod(x, y)})
+ rdiv = T.dynamic("rdiv", dtype)
+ return T.Let(
+ T.Let(T.Select(same_sign, rdiv, rdiv - scalar(1)), where={rdiv: T.Div(x, y)}),
+ where={rmod: T.truncmod(x, y)},
+ )
+ if variant == "offset":
+ x, y = x + scalar(4), scalar(5)
+ quotient, remainder = T.Div(x, y), T.truncmod(x, y)
+ if variant == "max":
+ value = T.max(quotient, zero)
+ elif variant == "positive_both":
+ value = quotient if kind == "div" else remainder
+ elif dtype == "int16":
+ value = (
+ T.Select(remainder >= zero, quotient, quotient - scalar(1))
+ if kind == "div"
+ else T.Select(remainder >= zero, remainder, remainder + y)
+ )
+ elif kind == "div":
+ value = quotient + (remainder >> scalar(bits - 1))
+ else:
+ value = remainder + (y & (remainder >> scalar(bits - 1)))
+ if variant == "offset":
+ return value
+ condition = T.And(zero <= y, zero <= x) if variant == "positive_both" else zero <= y
+ return T.Select(condition, value, zero)
+
+
+@pytest.mark.skipif(not env.has_llvm(), reason="need llvm")
+@pytest.mark.parametrize("dtype", ["int32", "int64", "int16"])
+@pytest.mark.parametrize(
+ "kind,variant",
+ [
+ (kind, variant)
+ for kind in ["div", "mod"]
+ for variant in ["plain", "positive_rhs", "positive_both", "pow2", "offset"]
+ ]
+ + [("div", "max")],
+)
+def test_lower_floordiv_and_floormod(dtype, kind, variant):
+ scalar = getattr(T, dtype)
+
+ @I.ir_module
+ class Before:
+ @T.function
+ def main(x: scalar, y: scalar) -> scalar:
+ T.func_attr({"target": T.target("llvm")})
+ return _divmod_expression(x, y, dtype, kind, variant, False)
+
+ @I.ir_module
+ class Expected:
+ @T.function
+ def main(x: scalar, y: scalar) -> scalar:
+ T.func_attr({"target": T.target("llvm")})
+ return _divmod_expression(x, y, dtype, kind, variant, True)
+
+ After = tvm.ir.transform.Sequential(
+ [tvm.tirx.transform.StmtSimplify(), tvm.tirx.transform.LowerIntrin()]
+ )(Before)
+ tvm.ir.assert_structural_equal(After, Expected)
+ compiled = tvm.compile(After, "llvm")
+ data = itertools.product(range(-10, 10), [y for y in range(-10, 10) if y])
+ for x, y in data:
+ if (variant == "pow2" and y != 8) or (variant == "offset" and y != 5):
+ continue
+ numerator = x + 4 if variant == "offset" else x
+ expected = numerator // y if kind == "div" else numerator % y
+ if variant == "max":
+ expected = max(expected, 0)
+ if (variant in ["positive_rhs", "max"] and y < 0) or (
+ variant == "positive_both" and (y < 0 or x < 0)
+ ):
+ expected = 0
+ assert compiled(x, y) == expected
+
+
+@pytest.mark.skipif(not env.has_llvm(), reason="need llvm")
+@pytest.mark.parametrize(
+ "dtype,clamp,data",
+ [
+ ("int64", False, [-(2**63), 0, 100]),
+ ("int16", False, [-32768, 0, 100]),
+ ("int16", True, [-10, 0, 32758, 32757]),
+ ],
+)
+def test_lower_floordiv_overflow_checks(dtype, clamp, data):
+ scalar = getattr(T, dtype)
+ bits = int(dtype[3:])
+
+ def numerator(x):
+ return T.min(T.max(x, scalar(-10)), scalar(32758)) if clamp else x
+
+ @I.ir_module
+ class Before:
+ @T.function
+ def main(x: scalar) -> scalar:
+ T.func_attr({"target": T.target("llvm")})
+ return T.floordiv(numerator(x), scalar(3))
+
+ def corrected(x):
+ q, r = T.Div(numerator(x), scalar(3)), T.truncmod(numerator(x), scalar(3))
+ return (
+ T.Select(r >= scalar(0), q, q - scalar(1))
+ if dtype == "int16"
+ else q + (r >> scalar(bits - 1))
+ )
+
+ @I.ir_module
+ class Expected:
+ @T.function
+ def main(x: scalar) -> scalar:
+ T.func_attr({"target": T.target("llvm")})
+ return corrected(x)
+
+ After = tvm.ir.transform.Sequential(
+ [tvm.tirx.transform.StmtSimplify(), tvm.tirx.transform.LowerIntrin()]
+ )(Before)
+ tvm.ir.assert_structural_equal(After, Expected)
+ compiled = tvm.compile(After, "llvm")
+ for value in data:
+ expected = min(max(value, -10), 32758) if clamp else value
+ assert compiled(value) == expected // 3
+
+
+if __name__ == "__main__":
+ tvm.testing.main()
diff --git a/tests/python/tirx/transform/test_tirx_transform_lower_tirx.py b/tests/python/tirx/transform/test_tirx_transform_lower_tirx.py
new file mode 100644
index 0000000000..188e1ec0d9
--- /dev/null
+++ b/tests/python/tirx/transform/test_tirx_transform_lower_tirx.py
@@ -0,0 +1,2370 @@
+# 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.
+
+from contextlib import nullcontext
+from types import SimpleNamespace
+
+import pytest
+import tvm_ffi
+
+import tvm
+import tvm.testing
+from tvm.backend.cuda.axis import Axis as CUDAAxis
+from tvm.backend.cuda.axis import laneid, warpid
+from tvm.script import ir as I
+from tvm.script import tirx as T
+from tvm.script.parser import register_namespace
+from tvm.script.parser.protocol_registry import register_mutable_decl
+from tvm.tirx.layout import S, TileLayout
+from tvm.tirx.transform import LowerTIRx, StmtSimplify
+
+_CTA_LAUNCH_ATTRS = {
+ "cuda.kernel_attrs": {},
+ "cuda.launch_fields": ["grid.x", "grid.y", "grid.z", "block.x", "block.y", "block.z"],
+}
+
+_CLUSTER_LAUNCH_ATTRS = {
+ "cuda.kernel_attrs": {},
+ "cuda.launch_fields": [
+ "grid.x",
+ "grid.y",
+ "grid.z",
+ "block.x",
+ "block.y",
+ "block.z",
+ "cluster.x",
+ "cluster.y",
+ "cluster.z",
+ ],
+}
+
+_PREFERRED_CLUSTER_LAUNCH_ATTRS = {
+ "cuda.kernel_attrs": {},
+ "cuda.launch_fields": [
+ "grid.x",
+ "grid.y",
+ "grid.z",
+ "block.x",
+ "block.y",
+ "block.z",
+ "cluster.x",
+ "cluster.y",
+ "cluster.z",
+ "preferred_cluster.x",
+ "preferred_cluster.y",
+ "preferred_cluster.z",
+ ],
+}
+
+
+def compare(Before, Expected, transform):
+ with tvm.target.Target("cuda"):
+ After = transform()(Before)
+ tvm.ir.assert_structural_equal(After, Expected)
+
+
+def _int_range(side, axis):
+ return tuple(int(x) for x in side[axis])
+
+
+@T.function(private=True)
+def _probe_impl():
+ T.evaluate(0)
+
+
+L_LANE = T.TileLayout(T.S[32 : 1 @ laneid])
+
+
+def test_unroll_loop_optional_annotations():
+ @I.ir_module
+ class Before:
+ @T.function(private=True)
+ def main(
+ A: T.Tensor(8, "float32", layout="default"), B: T.Tensor(8, "float32", layout="default")
+ ):
+ for i in T.serial(8, annotations={"pragma_unroll": None}):
+ B[i] = A[i] + 1.0
+ for i in T.serial(8, annotations={"unroll_explicit": None}):
+ B[i] = A[i] + 2.0
+ for i in T.serial(8, annotations={"pragma_unroll": False}):
+ B[i] = A[i] + 3.0
+ for i in T.serial(8, annotations={"auto_unroll_max_step": 16, "unroll_explicit": 0}):
+ B[i] = A[i] + 4.0
+
+ @I.ir_module
+ class Expected:
+ @T.function(private=True)
+ def main(
+ A: T.Tensor(8, "float32", layout="default"), B: T.Tensor(8, "float32", layout="default")
+ ):
+ for i in T.serial(8, annotations={"pragma_unroll": None}):
+ B[i] = A[i] + 1.0
+ for i in T.serial(8):
+ B[i] = A[i] + 2.0
+ for i in T.serial(8, annotations={"pragma_unroll": False}):
+ B[i] = A[i] + 3.0
+ for i in T.unroll(8):
+ B[i] = A[i] + 4.0
+
+ lowered = tvm.tirx.transform.UnrollLoop()(Before)
+ tvm.ir.assert_structural_equal(lowered, Expected)
+
+
+def test_lower_view_get():
+ @I.ir_module
+ class Before1:
+ @T.function(private=True)
+ def main(
+ in_buf: T.Tensor(64, "float32", layout="default"),
+ out: T.Tensor(64, "float32", layout="default"),
+ ) -> None:
+ T.device_entry(launch=T.cuda.LaunchConfig(grid=(1, 1, 1), block=1 * 32))
+ bx, by, bz = (T.cuda.block_idx("x"), T.cuda.block_idx("y"), T.cuda.block_idx("z"))
+ lane_id = T.cuda.lane_id()
+ A = T.alloc_tensor(
+ ty_args=[T.Tensor([2], "float16", scope="local", layout=T.TileLayout(T.S[2:1]))]
+ )
+ B_layout = A.layout.tile(L_LANE, (32,), (2,))
+ B = A.view(64, layout=B_layout)
+ A_local = B.view(2, layout="default")
+ for i in T.vectorized(2):
+ A_local[i] = T.float32(in_buf[lane_id * 2 + i])
+ B_1 = A.view(64, layout=B_layout)
+ A_local_1 = B_1.view(2, layout="default")
+ for i in T.vectorized(2):
+ out[lane_id * 2 + i] = T.float32(A_local_1[i])
+
+ @I.ir_module
+ class Expected1:
+ @T.function(private=True)
+ def main(
+ in_buf: T.Tensor((64,), "float32", layout=None),
+ out: T.Tensor((64,), "float32", layout=None),
+ ):
+ out_1 = T.decl_tensor(
+ out.data, ty_args=[T.Tensor((64,), "float32", layout=None, data=out.data)]
+ )
+ in_buf_1 = T.decl_tensor(
+ in_buf.data, ty_args=[T.Tensor((64,), "float32", layout=None, data=in_buf.data)]
+ )
+ with T.region(
+ "tirx.device_scope",
+ [1, 1, 1, 32, 1, 1],
+ attrs=_CTA_LAUNCH_ATTRS,
+ ):
+ bx: T.let[T.int32] = 0
+ by: T.let[T.int32] = 0
+ bz: T.let[T.int32] = 0
+ lane_id: T.let[T.int32] = T.cuda.thread_idx("x") % 32
+ A = T.alloc_local((2,), "float16", layout=None)
+ B = T.decl_tensor(
+ A.data, ty_args=[T.Tensor((64,), "float16", scope="local", layout=None)]
+ )
+ A_local = T.decl_tensor(
+ A.data, ty_args=[T.Tensor((2,), "float16", scope="local", layout=None)]
+ )
+ for i in T.vectorized(2):
+ A_local[i] = T.Cast(in_buf_1[T.cuda.thread_idx("x") * 2 + i], "float16")
+ B_1 = T.decl_tensor(
+ A.data, ty_args=[T.Tensor((64,), "float16", scope="local", layout=None)]
+ )
+ A_local_1 = T.decl_tensor(
+ A.data, ty_args=[T.Tensor((2,), "float16", scope="local", layout=None)]
+ )
+ for i in T.vectorized(2):
+ out_1[T.cuda.thread_idx("x") * 2 + i] = T.Cast(A_local_1[i], "float32")
+
+ compare(Before1, Expected1, LowerTIRx)
+
+ @I.ir_module
+ class Before2:
+ @T.function(private=True)
+ def main(
+ in_buf: T.Tensor((16, 16), "float32", layout="default"),
+ out: T.Tensor((16, 16), "float32", layout="default"),
+ ) -> None:
+ T.device_entry(launch=T.cuda.LaunchConfig(grid=(1, 1, 1), block=1 * 32))
+ bx, by, bz = (T.cuda.block_idx("x"), T.cuda.block_idx("y"), T.cuda.block_idx("z"))
+ lane_id = T.cuda.lane_id()
+ atom = T.TileLayout(T.S[(1, 2) : (2, 1)])
+ tile = T.TileLayout(T.S[(2, 2) : (2, 1)])
+ warp_atom = atom.tile(L_LANE, (8, 4), (1, 2))
+ A = T.alloc_tensor(
+ ty_args=[
+ T.Tensor(
+ [4, 2], "float32", scope="local", layout=atom.tile(tile, (2, 2), (1, 2))
+ )
+ ]
+ )
+ B_layout = warp_atom.tile(tile, (2, 2), (8, 8))
+ B = A.view(16, 16, layout=B_layout)
+ A_local = B.view(2, 2, 2, layout="default")
+ for i in T.unroll(4):
+ for j in T.vectorized(2):
+ A_local[i // 2, i % 2, j] = in_buf[
+ i // 2 * 8 + lane_id // 4, i % 2 * 8 + lane_id % 4 + j
+ ]
+ B_1 = A.view(16, 16, layout=B_layout)
+ A_local_1 = B_1.view(8, layout="default")
+ for i in T.vectorized(2):
+ out[lane_id // 4 * 8 + i // 2 * 8 + lane_id % 4, lane_id % 4 * 2 + i % 2] = (
+ A_local_1[i]
+ )
+
+ @I.ir_module
+ class Expected2:
+ @T.function(private=True)
+ def main(
+ in_buf: T.Tensor((16, 16), "float32", layout=None),
+ out: T.Tensor((16, 16), "float32", layout=None),
+ ):
+ out_1 = T.decl_tensor(
+ out.data, ty_args=[T.Tensor((256,), "float32", layout=None, data=out.data)]
+ )
+ in_buf_1 = T.decl_tensor(
+ in_buf.data, ty_args=[T.Tensor((256,), "float32", layout=None, data=in_buf.data)]
+ )
+ with T.region(
+ "tirx.device_scope",
+ [1, 1, 1, 32, 1, 1],
+ attrs=_CTA_LAUNCH_ATTRS,
+ ):
+ bx: T.let[T.int32] = 0
+ by: T.let[T.int32] = 0
+ bz: T.let[T.int32] = 0
+ lane_id: T.let[T.int32] = T.cuda.thread_idx("x") % 32
+ A = T.alloc_local((8,), layout=None)
+ B = T.decl_tensor(
+ A.data, ty_args=[T.Tensor((256,), "float32", scope="local", layout=None)]
+ )
+ A_local = T.decl_tensor(
+ A.data, ty_args=[T.Tensor((8,), "float32", scope="local", layout=None)]
+ )
+ for i in T.unroll(4):
+ for j in T.vectorized(2):
+ A_local[i * 2 + j] = in_buf_1[
+ i // 2 * 128
+ + T.cuda.thread_idx("x") // 4 * 16
+ + i % 2 * 8
+ + j
+ + T.cuda.thread_idx("x") % 4
+ ]
+ B_1 = T.decl_tensor(
+ A.data, ty_args=[T.Tensor((256,), "float32", scope="local", layout=None)]
+ )
+ A_local_1 = T.decl_tensor(
+ A.data, ty_args=[T.Tensor((8,), "float32", scope="local", layout=None)]
+ )
+ for i in T.vectorized(2):
+ out_1[
+ T.cuda.thread_idx("x") // 4 * 128 + T.cuda.thread_idx("x") % 4 * 18 + i
+ ] = A_local_1[i]
+
+ compare(Before2, Expected2, LowerTIRx)
+
+ @I.ir_module
+ class Before3_wgmma_layout:
+ @T.function(private=True)
+ def main(
+ in_buf: T.Tensor((128, 128), "float32", layout="default"),
+ out: T.Tensor((128, 128), "float32", layout="default"),
+ ) -> None:
+ T.device_entry(launch=T.cuda.LaunchConfig(grid=(1, 1, 1), block=2 * 128))
+ bx, by, bz = (T.cuda.block_idx("x"), T.cuda.block_idx("y"), T.cuda.block_idx("z"))
+ wg_id = T.cuda.warpgroup_id()
+ warp_id_in_wg = T.cuda.warp_in_warpgroup()
+ lane_id = T.cuda.lane_id()
+ atom = T.TileLayout(T.S[1, 2])
+ warp_atom = atom.tile(L_LANE, (8, 4), (1, 2))
+ tile = T.TileLayout(T.S[(2, 128 // 8) : (1, 2)])
+ warp_layout = warp_atom.tile(tile, (2, 128 // 8), (8, 8))
+ L_warp = T.TileLayout(T.S[8 : 1 @ warpid])
+ layout = warp_layout.tile(L_warp, (8, 1), (16, 128))
+ acc = T.alloc_tensor(
+ ty_args=[
+ T.Tensor(
+ [64],
+ "float32",
+ scope="local",
+ layout=atom.tile(tile, (2, 128 // 8), (1, 2)),
+ )
+ ]
+ )
+ A = acc.view(128, 128, layout=layout)
+ acc_local = A.view(16, 2, 2, layout=atom.tile(tile, (2, 128 // 8), (1, 2)))
+ for i in T.serial(128 // 8):
+ for j in T.unroll(2):
+ for vec in T.vectorized(2):
+ acc_local[i, j, vec] = in_buf[
+ wg_id * 64 + warp_id_in_wg * 16 + j * 8 + lane_id // 4,
+ i * 8 + lane_id % 4 * 2 + vec,
+ ]
+ A_1 = acc.view(128, 128, layout=layout)
+ acc_local_1 = A_1.view(64, layout=atom.tile(tile, (2, 128 // 8), (1, 2)))
+ for i in T.serial(128 // 8):
+ for j in T.unroll(2):
+ for vec in T.vectorized(2):
+ out[
+ wg_id * 64 + warp_id_in_wg * 16 + j * 8 + lane_id // 4,
+ i * 8 + lane_id % 4 * 2 + vec,
+ ] = acc_local_1[i * 4 + j * 2 + vec]
+
+ @I.ir_module
+ class Expected3_wgmma_layout:
+ @T.function(private=True)
+ def main(
+ in_buf: T.Tensor((128, 128), "float32", layout=None),
+ out: T.Tensor((128, 128), "float32", layout=None),
+ ):
+ out_1 = T.decl_tensor(
+ out.data, ty_args=[T.Tensor((16384,), "float32", layout=None, data=out.data)]
+ )
+ in_buf_1 = T.decl_tensor(
+ in_buf.data, ty_args=[T.Tensor((16384,), "float32", layout=None, data=in_buf.data)]
+ )
+ with T.region(
+ "tirx.device_scope",
+ [1, 1, 1, 256, 1, 1],
+ attrs=_CTA_LAUNCH_ATTRS,
+ ):
+ warp_id_in_cta: T.let[T.int32] = T.gpu_warp_shuffle(
+ T.uint32(4294967295), T.cuda.thread_idx("x") // 32, 0, 32, 32
+ )
+ bx: T.let[T.int32] = 0
+ by: T.let[T.int32] = 0
+ bz: T.let[T.int32] = 0
+ wg_id: T.let[T.int32] = warp_id_in_cta // 4
+ warp_id_in_wg: T.let[T.int32] = warp_id_in_cta % 4
+ lane_id: T.let[T.int32] = T.cuda.thread_idx("x") % 32
+ acc = T.alloc_local((64,), layout=None)
+ B = T.decl_tensor(
+ acc.data, ty_args=[T.Tensor((16384,), "float32", scope="local", layout=None)]
+ )
+ acc_local = T.decl_tensor(
+ acc.data, ty_args=[T.Tensor((64,), "float32", scope="local", layout=None)]
+ )
+ for i in range(16):
+ for j in T.unroll(2):
+ for vec in T.vectorized(2):
+ acc_local[i % 8 * 8 + j * 4 + i // 8 * 2 + vec] = in_buf_1[
+ warp_id_in_cta * 2048
+ + j * 1024
+ + T.cuda.thread_idx("x") % 32 // 4 * 128
+ + i * 8
+ + T.cuda.thread_idx("x") % 4 * 2
+ + vec
+ ]
+ B_1 = T.decl_tensor(
+ acc.data, ty_args=[T.Tensor((16384,), "float32", scope="local", layout=None)]
+ )
+ acc_local_1 = T.decl_tensor(
+ acc.data, ty_args=[T.Tensor((64,), "float32", scope="local", layout=None)]
+ )
+ for i in range(16):
+ for j in T.unroll(2):
+ for vec in T.vectorized(2):
+ out_1[
+ warp_id_in_cta * 2048
+ + j * 1024
+ + T.cuda.thread_idx("x") % 32 // 4 * 128
+ + i * 8
+ + T.cuda.thread_idx("x") % 4 * 2
+ + vec
+ ] = acc_local_1[i % 8 * 8 + j * 4 + i // 8 * 2 + vec]
+
+ compare(Before3_wgmma_layout, Expected3_wgmma_layout, LowerTIRx)
+
+ @I.ir_module
+ class Before4_multi_view_get:
+ @T.function(private=True)
+ def main(
+ in_buf: T.Tensor(64, "float32", layout="default"),
+ out: T.Tensor(64, "float32", layout="default"),
+ ) -> None:
+ T.device_entry(launch=T.cuda.LaunchConfig(grid=(1, 1, 1), block=1 * 32))
+ bx, by, bz = (T.cuda.block_idx("x"), T.cuda.block_idx("y"), T.cuda.block_idx("z"))
+ lane_id = T.cuda.lane_id()
+ A = T.alloc_tensor(
+ ty_args=[T.Tensor([2], "float16", scope="local", layout=T.TileLayout(T.S[2:1]))]
+ )
+ B_layout = A.layout.tile(L_LANE, (32,), (2,))
+ B = A.view(64, layout=B_layout)
+ B_1 = A.view(64, layout=B_layout)
+ A_local = B.view(2, layout="default")
+ A_local[0] = T.float32(in_buf[lane_id * 2])
+ A_local_1 = B_1.view(2, layout="default")
+ A_local_1[1] = T.float32(in_buf[lane_id * 2 + 1])
+ "\n write A into out\n "
+ B_2 = A.view(64, layout=B_layout)
+ B_3 = A.view(64, layout=B_layout)
+ A_local_2 = B_2.view(2, layout="default")
+ out[lane_id * 2] = T.float32(A_local_2[0])
+ A_local_3 = B_3.view(2, layout="default")
+ out[lane_id * 2 + 1] = T.float32(A_local_3[1])
+
+ @I.ir_module
+ class Expected4_multi_view_get:
+ @T.function(private=True)
+ def main(
+ in_buf: T.Tensor((64,), "float32", layout=None),
+ out: T.Tensor((64,), "float32", layout=None),
+ ):
+ out_1 = T.decl_tensor(
+ out.data, ty_args=[T.Tensor((64,), "float32", layout=None, data=out.data)]
+ )
+ in_buf_1 = T.decl_tensor(
+ in_buf.data, ty_args=[T.Tensor((64,), "float32", layout=None, data=in_buf.data)]
+ )
+ with T.region(
+ "tirx.device_scope",
+ [1, 1, 1, 32, 1, 1],
+ attrs=_CTA_LAUNCH_ATTRS,
+ ):
+ bx: T.let[T.int32] = 0
+ by: T.let[T.int32] = 0
+ bz: T.let[T.int32] = 0
+ lane_id: T.let[T.int32] = T.cuda.thread_idx("x") % 32
+ A = T.alloc_local((2,), "float16", layout=None)
+ B = T.decl_tensor(
+ A.data, ty_args=[T.Tensor((64,), "float16", scope="local", layout=None)]
+ )
+ B_1 = T.decl_tensor(
+ A.data, ty_args=[T.Tensor((64,), "float16", scope="local", layout=None)]
+ )
+ A_local = T.decl_tensor(
+ A.data, ty_args=[T.Tensor((2,), "float16", scope="local", layout=None)]
+ )
+ A_local[0] = T.Cast(in_buf_1[T.cuda.thread_idx("x") * 2], "float16")
+ A_local_1 = T.decl_tensor(
+ A.data, ty_args=[T.Tensor((2,), "float16", scope="local", layout=None)]
+ )
+ A_local_1[1] = T.Cast(in_buf_1[T.cuda.thread_idx("x") * 2 + 1], "float16")
+ B_2 = T.decl_tensor(
+ A.data, ty_args=[T.Tensor((64,), "float16", scope="local", layout=None)]
+ )
+ B_3 = T.decl_tensor(
+ A.data, ty_args=[T.Tensor((64,), "float16", scope="local", layout=None)]
+ )
+ A_local_2 = T.decl_tensor(
+ A.data, ty_args=[T.Tensor((2,), "float16", scope="local", layout=None)]
+ )
+ out_1[T.cuda.thread_idx("x") * 2] = T.Cast(A_local_2[0], "float32")
+ A_local_3 = T.decl_tensor(
+ A.data, ty_args=[T.Tensor((2,), "float16", scope="local", layout=None)]
+ )
+ out_1[T.cuda.thread_idx("x") * 2 + 1] = T.Cast(A_local_3[1], "float32")
+
+ compare(Before4_multi_view_get, Expected4_multi_view_get, LowerTIRx)
+
+
+def test_lower_scope_id():
+ @I.ir_module
+ class Before1:
+ @T.function(private=True)
+ def main() -> None:
+ T.device_entry(launch=T.cuda.LaunchConfig(grid=(3, 4, 5), block=(32,)))
+ bx, by, bz = (T.cuda.block_idx("x"), T.cuda.block_idx("y"), T.cuda.block_idx("z"))
+ tx = T.cuda.thread_idx("x")
+ T.evaluate(bx + by + bz + tx)
+
+ @I.ir_module
+ class Expected1:
+ @T.function(private=True)
+ def main() -> None:
+ with T.region(
+ "tirx.device_scope",
+ [3, 4, 5, 32, 1, 1],
+ attrs=_CTA_LAUNCH_ATTRS,
+ ):
+ bx: T.let[T.int32] = T.cuda.block_idx("x")
+ by: T.let[T.int32] = T.cuda.block_idx("y")
+ bz: T.let[T.int32] = T.cuda.block_idx("z")
+ tx: T.let[T.int32] = T.cuda.thread_idx("x")
+ T.evaluate(bx + by + bz + tx)
+
+ compare(Before1, Expected1, LowerTIRx)
+
+ @I.ir_module
+ class Before2:
+ @T.function(private=True)
+ def main() -> None:
+ T.device_entry(
+ launch=T.cuda.LaunchConfig(grid=(8, 8, 8), block=4 * 32, cluster=(2, 2, 2))
+ )
+ cbx, cby, cbz = (
+ T.cuda.cluster_cta_id("x"),
+ T.cuda.cluster_cta_id("y"),
+ T.cuda.cluster_cta_id("z"),
+ )
+ bx, by, bz = (T.cuda.block_idx("x"), T.cuda.block_idx("y"), T.cuda.block_idx("z"))
+ warp_id = T.cuda.warp_id()
+ lane_id = T.cuda.lane_id()
+ T.evaluate(bx + by + bz + warp_id + lane_id + cbx + cby + cbz)
+
+ @I.ir_module
+ class Expected2:
+ @T.function(private=True)
+ def main() -> None:
+ with T.region(
+ "tirx.device_scope",
+ [8, 8, 8, 128, 1, 1, 2, 2, 2],
+ attrs=_CLUSTER_LAUNCH_ATTRS,
+ ):
+ warp_id_in_cta: T.let[T.int32] = T.gpu_warp_shuffle(
+ T.uint32(4294967295), T.cuda.thread_idx("x") // 32, 0, 32, 32
+ )
+ cbx: T.let[T.int32] = T.cuda.mov_sreg(32, "cluster_ctaid.x")
+ cby: T.let[T.int32] = T.cuda.mov_sreg(32, "cluster_ctaid.y")
+ cbz: T.let[T.int32] = T.cuda.mov_sreg(32, "cluster_ctaid.z")
+ bx: T.let[T.int32] = T.cuda.block_idx("x")
+ by: T.let[T.int32] = T.cuda.block_idx("y")
+ bz: T.let[T.int32] = T.cuda.block_idx("z")
+ warp_id: T.let[T.int32] = warp_id_in_cta
+ lane_id: T.let[T.int32] = T.cuda.thread_idx("x") % 32
+ T.evaluate(bx + by + bz + warp_id + lane_id + cbx + cby + cbz)
+
+ compare(Before2, Expected2, LowerTIRx)
+
+ @I.ir_module
+ class Before3:
+ @T.function(private=True)
+ def main() -> None:
+ T.device_entry(
+ launch=T.cuda.LaunchConfig(grid=(8, 10, 12), block=3 * 128, cluster=(2, 2, 1))
+ )
+ bx, by, bz = (T.cuda.block_idx("x"), T.cuda.block_idx("y"), T.cuda.block_idx("z"))
+ cbx, cby, cbz = (
+ T.cuda.cluster_cta_id("x"),
+ T.cuda.cluster_cta_id("y"),
+ T.cuda.cluster_cta_id("z"),
+ )
+ clx, cly, clz = (T.cuda.cluster_id("x"), T.cuda.cluster_id("y"), T.cuda.cluster_id("z"))
+ wg_id = T.cuda.warpgroup_id()
+ warp_id_in_wg = T.cuda.warp_in_warpgroup()
+ lane_id = T.cuda.lane_id()
+ tid_in_wg = T.cuda.thread_in_warpgroup()
+ T.evaluate(bx + by + bz)
+ T.evaluate(cbx + cby + cbz)
+ T.evaluate(clx + cly + clz)
+ T.evaluate(wg_id + warp_id_in_wg + lane_id + tid_in_wg)
+
+ @I.ir_module
+ class Expected3:
+ @T.function(private=True)
+ def main() -> None:
+ with T.region(
+ "tirx.device_scope",
+ [8, 10, 12, 384, 1, 1, 2, 2, 1],
+ attrs=_CLUSTER_LAUNCH_ATTRS,
+ ):
+ warp_id_in_cta: T.let[T.int32] = T.gpu_warp_shuffle(
+ T.uint32(4294967295), T.cuda.thread_idx("x") // 32, 0, 32, 32
+ )
+ bx: T.let[T.int32] = T.cuda.block_idx("x")
+ by: T.let[T.int32] = T.cuda.block_idx("y")
+ bz: T.let[T.int32] = T.cuda.block_idx("z")
+ cbx: T.let[T.int32] = T.cuda.mov_sreg(32, "cluster_ctaid.x")
+ cby: T.let[T.int32] = T.cuda.mov_sreg(32, "cluster_ctaid.y")
+ cbz: T.let[T.int32] = T.cuda.mov_sreg(32, "cluster_ctaid.z")
+ clx: T.let[T.int32] = T.cuda.mov_sreg(32, "clusterid.x")
+ cly: T.let[T.int32] = T.cuda.mov_sreg(32, "clusterid.y")
+ clz: T.let[T.int32] = T.cuda.mov_sreg(32, "clusterid.z")
+ wg_id: T.let[T.int32] = warp_id_in_cta // 4
+ warp_id_in_wg: T.let[T.int32] = warp_id_in_cta % 4
+ lane_id: T.let[T.int32] = T.cuda.thread_idx("x") % 32
+ tid_in_wg: T.let[T.int32] = T.cuda.thread_idx("x") % 128
+ T.evaluate(bx + by + bz)
+ T.evaluate(cbx + cby + cbz)
+ T.evaluate(clx + cly + clz)
+ T.evaluate(wg_id + warp_id_in_wg + lane_id + tid_in_wg)
+
+ compare(Before3, Expected3, LowerTIRx)
+
+
+def test_lower_ordinary_cta_has_no_cluster_launch_tags():
+ @I.ir_module
+ class Before:
+ @T.function(private=True)
+ def main() -> None:
+ T.device_entry(launch=T.cuda.LaunchConfig(grid=(1,), block=(32,)))
+
+ @I.ir_module
+ class Expected:
+ @T.function(private=True)
+ def main():
+ with T.device_scope(1, 1, 1, 32, 1, 1, attrs=_CTA_LAUNCH_ATTRS):
+ pass
+
+ with tvm.target.Target("cuda"):
+ after = LowerTIRx()(Before)["main"]
+ tvm.ir.assert_structural_equal(after, Expected["main"])
+
+
+def test_lower_explicit_singleton_cluster_launch_tags_survive_when_unused():
+ @I.ir_module
+ class Before:
+ @T.function(private=True)
+ def cluster_2d() -> None:
+ T.device_entry(launch=T.cuda.LaunchConfig(grid=(1, 1), block=(32,), cluster=(1, 1)))
+ unused_cbx, unused_cby = (T.cuda.cluster_cta_id("x"), T.cuda.cluster_cta_id("y"))
+ unused_bx, unused_by = (T.cuda.block_idx("x"), T.cuda.block_idx("y"))
+
+ @T.function(private=True)
+ def cluster_3d() -> None:
+ T.device_entry(
+ launch=T.cuda.LaunchConfig(grid=(1, 1, 1), block=(32,), cluster=(1, 1, 1))
+ )
+ unused_cbx, unused_cby, unused_cbz = (
+ T.cuda.cluster_cta_id("x"),
+ T.cuda.cluster_cta_id("y"),
+ T.cuda.cluster_cta_id("z"),
+ )
+ unused_bx, unused_by, unused_bz = (
+ T.cuda.block_idx("x"),
+ T.cuda.block_idx("y"),
+ T.cuda.block_idx("z"),
+ )
+
+ @I.ir_module
+ class Expected:
+ @T.function(private=True)
+ def cluster_2d():
+ with T.device_scope(1, 1, 1, 32, 1, 1, 1, 1, 1, attrs=_CLUSTER_LAUNCH_ATTRS):
+ unused_cbx: T.let[T.int32] = T.cuda.mov_sreg(32, "cluster_ctaid.x")
+ unused_cby: T.let[T.int32] = T.cuda.mov_sreg(32, "cluster_ctaid.y")
+ unused_bx: T.let[T.int32] = 0
+ unused_by: T.let[T.int32] = 0
+
+ @T.function(private=True)
+ def cluster_3d():
+ with T.device_scope(1, 1, 1, 32, 1, 1, 1, 1, 1, attrs=_CLUSTER_LAUNCH_ATTRS):
+ unused_cbx: T.let[T.int32] = T.cuda.mov_sreg(32, "cluster_ctaid.x")
+ unused_cby: T.let[T.int32] = T.cuda.mov_sreg(32, "cluster_ctaid.y")
+ unused_cbz: T.let[T.int32] = T.cuda.mov_sreg(32, "cluster_ctaid.z")
+ unused_bx: T.let[T.int32] = 0
+ unused_by: T.let[T.int32] = 0
+ unused_bz: T.let[T.int32] = 0
+
+ with tvm.target.Target("cuda"):
+ after = LowerTIRx()(Before)
+ tvm.ir.assert_structural_equal(after, Expected)
+
+
+def test_lower_multi_cta_cluster_launch_tags_remain_unchanged():
+ @I.ir_module
+ class Before:
+ @T.function(private=True)
+ def main() -> None:
+ T.device_entry(launch=T.cuda.LaunchConfig(grid=(2, 1), block=(32,), cluster=(2, 1)))
+ unused_cbx, unused_cby = (T.cuda.cluster_cta_id("x"), T.cuda.cluster_cta_id("y"))
+ unused_bx, unused_by = (T.cuda.block_idx("x"), T.cuda.block_idx("y"))
+
+ @I.ir_module
+ class Expected:
+ @T.function(private=True)
+ def main():
+ with T.device_scope(2, 1, 1, 32, 1, 1, 2, 1, 1, attrs=_CLUSTER_LAUNCH_ATTRS):
+ unused_cbx: T.let[T.int32] = T.cuda.mov_sreg(32, "cluster_ctaid.x")
+ unused_cby: T.let[T.int32] = T.cuda.mov_sreg(32, "cluster_ctaid.y")
+ unused_bx: T.let[T.int32] = T.cuda.block_idx("x")
+ unused_by: T.let[T.int32] = 0
+
+ with tvm.target.Target("cuda"):
+ after = LowerTIRx()(Before)["main"]
+ tvm.ir.assert_structural_equal(after, Expected["main"])
+
+
+def test_lower_singleton_cluster_preserves_preferred_cluster_tags():
+ @I.ir_module
+ class Before:
+ @T.function(private=True)
+ def main() -> None:
+ T.device_entry(
+ launch=T.cuda.LaunchConfig(
+ grid=(2, 2), block=(32,), cluster=(1, 1), preferred_cluster=[2, 2]
+ )
+ )
+ unused_cbx, unused_cby = (T.cuda.cluster_cta_id("x"), T.cuda.cluster_cta_id("y"))
+ unused_bx, unused_by = (T.cuda.block_idx("x"), T.cuda.block_idx("y"))
+
+ @I.ir_module
+ class Expected:
+ @T.function(private=True)
+ def main():
+ with T.device_scope(
+ 2, 2, 1, 32, 1, 1, 1, 1, 1, 2, 2, 1, attrs=_PREFERRED_CLUSTER_LAUNCH_ATTRS
+ ):
+ unused_cbx: T.let[T.int32] = T.cuda.mov_sreg(32, "cluster_ctaid.x")
+ unused_cby: T.let[T.int32] = T.cuda.mov_sreg(32, "cluster_ctaid.y")
+ unused_bx: T.let[T.int32] = T.cuda.block_idx("x")
+ unused_by: T.let[T.int32] = T.cuda.block_idx("y")
+
+ with tvm.target.Target("cuda"):
+ after = LowerTIRx()(Before)["main"]
+ tvm.ir.assert_structural_equal(after, Expected["main"])
+
+
+def test_lower_scope_id2():
+ @T.inline
+ def func(warp_id, tx):
+ wg_id = T.cuda.warpgroup_id()
+ T.evaluate(wg_id + warp_id + tx)
+
+ @I.ir_module
+ class Before:
+ @T.function(private=True)
+ def main():
+ T.device_entry(launch=T.cuda.LaunchConfig(grid=(3, 4, 5), block=(256,)))
+ bx, by, bz = (T.cuda.block_idx("x"), T.cuda.block_idx("y"), T.cuda.block_idx("z"))
+ warp_id = T.cuda.warp_id()
+ tx = T.cuda.thread_idx("x")
+ func(warp_id, tx)
+
+ @I.ir_module
+ class Expected:
+ @T.function(private=True)
+ def main():
+ with T.region(
+ "tirx.device_scope",
+ [3, 4, 5, 256, 1, 1],
+ attrs=_CTA_LAUNCH_ATTRS,
+ ):
+ warp_id_in_cta: T.let[T.int32] = T.gpu_warp_shuffle(
+ T.uint32(4294967295), T.cuda.thread_idx("x") // 32, 0, 32, 32
+ )
+ bx: T.let[T.int32] = T.cuda.block_idx("x")
+ by: T.let[T.int32] = T.cuda.block_idx("y")
+ bz: T.let[T.int32] = T.cuda.block_idx("z")
+ warp_id: T.let[T.int32] = warp_id_in_cta
+ tx: T.let[T.int32] = T.cuda.thread_idx("x")
+ wg_id: T.let[T.int32] = warp_id_in_cta // 4
+ T.evaluate(wg_id + warp_id + tx)
+
+ compare(Before, Expected, LowerTIRx)
+
+
+@pytest.mark.skip(
+ reason=(
+ "Tested multi-kernel-per-Function behavior where a second sibling "
+ "`with T.thread():` would redefine scope-ids and produce a second "
+ "launch. The T.device_entry() refactor allows only one device-region "
+ "marker per Function; this case is out of scope."
+ )
+)
+def test_lower_scope_id3():
+ @I.ir_module
+ class Before:
+ @T.function(private=True)
+ def main():
+ T.device_entry(launch=T.cuda.LaunchConfig(grid=(3, 4, 5), block=(128,)))
+ bx, by, bz = (T.cuda.block_idx("x"), T.cuda.block_idx("y"), T.cuda.block_idx("z"))
+ warp_id = T.cuda.warp_id()
+ tx = T.cuda.thread_idx("x")
+ T.evaluate(bx + by + bz + warp_id + tx)
+ bx, by, bz = (T.cuda.block_idx("x"), T.cuda.block_idx("y"), T.cuda.block_idx("z"))
+ warp_id = T.cuda.warp_id()
+ tx = T.cuda.thread_idx("x")
+ T.evaluate(bx + by + bz + warp_id + tx)
+
+ @I.ir_module
+ class Expected:
+ @T.function(private=True)
+ def main():
+ with T.launch_thread("blockIdx.x", 3) as blockIdx_x:
+ threadIdx_x = T.launch_thread("threadIdx.x", 128)
+ blockIdx_y = T.launch_thread("blockIdx.y", 4)
+ blockIdx_z = T.launch_thread("blockIdx.z", 5)
+ warp_id_in_cta: T.let[T.int32] = T.gpu_warp_shuffle(
+ T.uint32(4294967295), threadIdx_x // 32, 0, 32, 32
+ )
+ bx: T.let[T.int32] = blockIdx_x
+ by: T.let[T.int32] = blockIdx_y
+ bz: T.let[T.int32] = blockIdx_z
+ warp_id: T.let[T.int32] = warp_id_in_cta
+ tx: T.let[T.int32] = threadIdx_x
+ T.evaluate(bx + by + bz + warp_id + tx)
+ blockIdx_x = T.launch_thread("blockIdx.x", 6)
+ threadIdx_x = T.launch_thread("threadIdx.x", 256)
+ blockIdx_y = T.launch_thread("blockIdx.y", 7)
+ blockIdx_z = T.launch_thread("blockIdx.z", 8)
+ warp_id_in_cta: T.let[T.int32] = T.gpu_warp_shuffle(
+ T.uint32(4294967295), threadIdx_x // 32, 0, 32, 32
+ )
+ bx: T.let[T.int32] = blockIdx_x
+ by: T.let[T.int32] = blockIdx_y
+ bz: T.let[T.int32] = blockIdx_z
+ warp_id: T.let[T.int32] = warp_id_in_cta
+ tx: T.let[T.int32] = threadIdx_x
+ T.evaluate(bx + by + bz + warp_id + tx)
+
+ compare(Before, Expected, LowerTIRx)
+
+
+def test_lower_layout():
+ @I.ir_module
+ class Before:
+ @T.function(private=True)
+ def main(A: T.Tensor((128, 32), "float16", layout="default")) -> None:
+ T.device_entry(launch=T.cuda.LaunchConfig(grid=(1, 1, 1), block=(128,)))
+ bx, by, bz = (T.cuda.block_idx("x"), T.cuda.block_idx("y"), T.cuda.block_idx("z"))
+ _lane = T.cuda.lane_id()
+ tid = T.cuda.thread_idx("x")
+ A_smem = T.alloc_tensor(
+ ty_args=[
+ T.Tensor(
+ [128, 32],
+ "float16",
+ scope="shared",
+ layout=T.ComposeLayout(3, 3, 3, T.TileLayout(T.S[(512,)])),
+ )
+ ]
+ )
+ thread_col = T.meta_var(4)
+ thread_row = T.meta_var(32)
+ for tile in T.serial(128 // thread_row):
+ row = T.meta_var(tile * thread_row + tid // thread_col)
+ col = T.meta_var(tid % thread_col * 8)
+ for vec in T.vectorized(8):
+ A_smem[row, col + vec] = A[bx * 128 + row, col + vec]
+
+ compose_m = T.dynamic("compose_m", "int32")
+ compose_q = T.dynamic("compose_q", "int32")
+
+ @I.ir_module
+ class Expected:
+ @T.function(private=True)
+ def main(A: T.Tensor((128, 32), "float16", layout=None)) -> None:
+ A_1 = T.decl_tensor(
+ A.data, ty_args=[T.Tensor((4096,), "float16", layout=None, data=A.data)]
+ )
+ with T.region(
+ "tirx.device_scope",
+ [1, 1, 1, 128, 1, 1],
+ attrs=_CTA_LAUNCH_ATTRS,
+ ):
+ bx: T.let[T.int32] = 0
+ by: T.let[T.int32] = 0
+ bz: T.let[T.int32] = 0
+ _lane: T.let[T.int32] = T.cuda.thread_idx("x") % 32
+ tid: T.let[T.int32] = T.cuda.thread_idx("x")
+ A_smem = T.alloc_shared((4096,), "float16", layout=None)
+ thread_col = 4
+ thread_row = 32
+ for tile in range(128 // thread_row):
+ row = tile * thread_row + tid // thread_col
+ col = tid % thread_col * 8
+ for vec in T.vectorized(8):
+ # The swizzle lowers to its composition bindings rather than a
+ # folded closed form: compose_m is the flat element index, so
+ # compose_m // 8 is the row and compose_m % 8 the lane, which
+ # substituted back gives the same address.
+ A_smem[
+ T.Let(
+ T.Let(
+ T.shift_left(
+ T.bitwise_xor(
+ compose_q,
+ T.shift_right(T.bitwise_and(compose_q, 56), 3),
+ ),
+ 3,
+ )
+ + compose_m % 8,
+ where={compose_q: compose_m // 8},
+ ),
+ where={compose_m: tile * 1024 + T.cuda.thread_idx("x") * 8 + vec},
+ )
+ ] = A_1[tile * 1024 + T.cuda.thread_idx("x") * 8 + vec]
+
+ compare(Before, Expected, LowerTIRx)
+
+
+@pytest.mark.parametrize(
+ "target,error,message",
+ [
+ (None, tvm.error.InternalError, "Target context required"),
+ ("cuda", RuntimeError, "An explicit layout is required for shared tensors"),
+ ],
+ ids=["missing-target", "missing-shared-layout"],
+)
+def test_lower_opcall_fail(target, error, message):
+ @I.ir_module
+ class Before:
+ @T.function(private=True)
+ def main():
+ T.device_entry(launch=T.cuda.LaunchConfig(grid=1, block=32))
+ shared = T.alloc_shared((64,), "float32")
+ local = T.alloc_local(
+ (1, 32, 1), "float32", layout=T.TileLayout(T.S[(1, 32, 1) : (32, 1 @ laneid, 1)])
+ )
+ T.cuda.tile.apply(T.ptx["ld.shared.b32"], local, shared, scope="cta")
+
+ with (
+ tvm.target.Target(target) if target else nullcontext(),
+ pytest.raises(error, match=message),
+ ):
+ LowerTIRx()(Before)
+
+
+def test_lower_decl_tensor_pointer():
+ @I.ir_module
+ class Before:
+ @T.function(private=True)
+ def main():
+ T.device_entry(launch=T.cuda.LaunchConfig(grid=(1,), block=(128,)))
+ buf = T.alloc_tensor(ty_args=[T.Tensor([1024], "uint8", scope="shared.dyn")])
+ A = T.decl_tensor(
+ buf.data, ty_args=[T.Tensor([128], "float16", elem_offset=32, data=buf.data)]
+ )
+ T.evaluate(A.ptr_to([64]))
+
+ @I.ir_module
+ class Expected:
+ @T.function(private=True)
+ def main():
+ with T.region(
+ "tirx.device_scope",
+ [1, 1, 1, 128, 1, 1],
+ attrs=_CTA_LAUNCH_ATTRS,
+ ):
+ buf = T.alloc_tensor(
+ ty_args=[T.Tensor((1024,), "uint8", scope="shared.dyn", layout=None)]
+ )
+ A = T.decl_tensor(
+ buf.data,
+ ty_args=[
+ T.Tensor((128,), "float16", elem_offset=32, scope="shared.dyn", layout=None)
+ ],
+ )
+ T.address_of(A[64])
+
+ compare(Before, Expected, LowerTIRx)
+
+
+def test_lower_separate_scope_id_def():
+ @I.ir_module
+ class Before:
+ @T.function(private=True)
+ def main():
+ T.device_entry(launch=T.cuda.LaunchConfig(grid=(1,), block=(128,)))
+ tx = T.cuda.thread_idx("x")
+ if tx == 0:
+ T.evaluate(tx)
+
+ @I.ir_module
+ class Expected:
+ @T.function(private=True)
+ def main():
+ with T.region(
+ "tirx.device_scope",
+ [1, 1, 1, 128, 1, 1],
+ attrs=_CTA_LAUNCH_ATTRS,
+ ):
+ tx: T.let[T.int32] = T.cuda.thread_idx("x")
+ if tx == 0:
+ T.evaluate(tx)
+
+ compare(Before, Expected, LowerTIRx)
+
+
+def test_lower_uint32_scope_id_casts_at_bind():
+ """uint32 scope ids get a Cast at the bind; launch params stay int32."""
+
+ @I.ir_module
+ class Before:
+ @T.function(private=True)
+ def main():
+ T.device_entry(launch=T.cuda.LaunchConfig(grid=(1,), block=(128,)))
+ tx = T.cast(T.cuda.thread_idx("x"), "uint32")
+ for k in T.serial(4, dtype="uint32"):
+ T.evaluate(tx + k)
+
+ @I.ir_module
+ class Expected:
+ @T.function(private=True)
+ def main():
+ with T.region(
+ "tirx.device_scope",
+ [1, 1, 1, 128, 1, 1],
+ attrs=_CTA_LAUNCH_ATTRS,
+ ):
+ tx: T.let[T.uint32] = T.Cast(T.cuda.thread_idx("x"), "uint32")
+ for k in T.serial(T.uint32(4)):
+ T.evaluate(tx + k)
+
+ compare(Before, Expected, LowerTIRx)
+
+
+def test_lower_exec_context_infers_plain_predicate_for_dispatch():
+ seen = []
+ variant = "__probe_exec_context_plain_predicate__"
+
+ @_register_probe(variant)
+ def _probe(op_call, sctx):
+ seen.append({"scope_kind": sctx.scope_kind, "inter": sctx.inter, "intra": sctx.intra})
+
+ return _probe_impl
+
+ @I.ir_module
+ class Before:
+ @T.function(private=True)
+ def main(
+ A: T.Tensor((1,), "float32", scope="global", layout="default"),
+ B: T.Tensor((1,), "float32", scope="global", layout="default"),
+ ):
+ T.device_entry(launch=T.cuda.LaunchConfig(grid=(1,), block=4 * 32))
+ warp_id = T.cuda.warp_id()
+ lane_id = T.cuda.lane_id()
+ if (warp_id == 0) & (lane_id == 0):
+ getattr(T.cuda.tile, variant)(B[0:1], A[0:1])
+
+ @I.ir_module
+ class Expected:
+ @T.function(private=True)
+ def main(A: T.Tensor((1,), T.float32), B: T.Tensor((1,), T.float32)):
+ B_1 = T.decl_tensor(B.data, ty_args=[T.Tensor((1,), T.float32)])
+ A_1 = T.decl_tensor(A.data, ty_args=[T.Tensor((1,), T.float32)])
+ with T.device_scope(1, 1, 1, 128, 1, 1, attrs=_CTA_LAUNCH_ATTRS):
+ warp_id_in_cta: T.let[T.int32] = T.gpu_warp_shuffle(
+ T.uint32(4294967295), T.cuda.thread_idx("x") // 32, 0, 32, 32
+ )
+ warp_id: T.let[T.int32] = warp_id_in_cta
+ lane_id: T.let[T.int32] = T.cuda.thread_idx("x") % 32
+ if (warp_id == 0) & (lane_id == 0):
+ pass
+
+ with tvm.target.Target("cuda"):
+ After = LowerTIRx()(Before)
+ tvm.ir.assert_structural_equal(After, Expected)
+
+ assert len(seen) == 1
+ assert seen[0]["scope_kind"] == "thread"
+ assert _int_range(seen[0]["inter"], "laneid") == (1, 0)
+ assert _int_range(seen[0]["inter"], "warpid") == (1, 0)
+ assert _int_range(seen[0]["inter"], "cta_id") == (1, 0)
+ assert len(seen[0]["intra"]) == 0
+
+
+def test_lower_exec_context_tracks_let_aliases_and_hoists_warp_shuffle():
+ seen = []
+ variant = "__probe_let_alias_scope__"
+
+ @_register_probe(variant)
+ def _probe(op_call, sctx):
+ seen.append(sctx.inter)
+
+ return _probe_impl
+
+ @I.ir_module
+ class Before:
+ @T.function(private=True)
+ def main(
+ A: T.Tensor((1,), "float32", layout="default"),
+ B: T.Tensor((1,), "float32", layout="default"),
+ ):
+ T.device_entry(launch=T.cuda.LaunchConfig(grid=1, block=128))
+ lane: T.let[T.int32] = T.cuda.lane_id()
+ alias_lane: T.let[T.int32] = lane
+ if alias_lane == 0:
+ warp: T.let[T.int32] = T.cuda.warp_id()
+ alias_warp: T.let[T.int32] = warp
+ if alias_warp == 1:
+ getattr(T.cuda.tile, variant)(B[0:1], A[0:1])
+
+ @I.ir_module
+ class Expected:
+ @T.function(private=True)
+ def main(A: T.Tensor((1,), T.float32), B: T.Tensor((1,), T.float32)):
+ B_1 = T.decl_tensor(B.data, ty_args=[T.Tensor((1,), T.float32)])
+ A_1 = T.decl_tensor(A.data, ty_args=[T.Tensor((1,), T.float32)])
+ with T.device_scope(1, 1, 1, 128, 1, 1, attrs=_CTA_LAUNCH_ATTRS):
+ warp_id_in_cta: T.let[T.int32] = T.gpu_warp_shuffle(
+ T.uint32(4294967295), T.cuda.thread_idx("x") // 32, 0, 32, 32
+ )
+ lane: T.let[T.int32] = T.cuda.thread_idx("x") % 32
+ alias_lane: T.let[T.int32] = lane
+ if alias_lane == 0:
+ warp: T.let[T.int32] = warp_id_in_cta
+ alias_warp: T.let[T.int32] = warp
+ if alias_warp == 1:
+ pass
+
+ with tvm.target.Target("cuda"):
+ lowered = LowerTIRx()(Before)
+ tvm.ir.assert_structural_equal(lowered, Expected)
+
+ assert len(seen) == 1
+ assert _int_range(seen[0], "laneid") == (1, 0)
+ assert _int_range(seen[0], "warpid") == (1, 1)
+ code = lowered.script()
+ assert code.index("gpu_warp_shuffle") < code.index("if ")
+
+
+def test_lower_exec_context_respects_preferred_cluster_bounds():
+ seen = []
+ variant = "__probe_preferred_cluster_bounds__"
+
+ @_register_probe(variant)
+ def _probe(op_call, sctx):
+ seen.append(sctx.inter)
+
+ return _probe_impl
+
+ @I.ir_module
+ class Before:
+ @T.function(private=True)
+ def main(
+ A: T.Tensor((1,), "float32", layout="default"),
+ B: T.Tensor((1,), "float32", layout="default"),
+ ):
+ T.device_entry(
+ launch=T.cuda.LaunchConfig(
+ grid=(2, 2), block=32, cluster=1, preferred_cluster=(2, 2)
+ )
+ )
+ cby = T.cuda.cluster_cta_id("y")
+ if cby == 1:
+ getattr(T.cuda.tile, variant)(B[0:1], A[0:1])
+
+ @I.ir_module
+ class Expected:
+ @T.function(private=True)
+ def main(A: T.Tensor((1,), T.float32), B: T.Tensor((1,), T.float32)):
+ B_1 = T.decl_tensor(B.data, ty_args=[T.Tensor((1,), T.float32)])
+ A_1 = T.decl_tensor(A.data, ty_args=[T.Tensor((1,), T.float32)])
+ with T.device_scope(
+ 2, 2, 1, 32, 1, 1, 1, 1, 1, 2, 2, 1, attrs=_PREFERRED_CLUSTER_LAUNCH_ATTRS
+ ):
+ cby: T.let[T.int32] = T.cuda.mov_sreg(32, "cluster_ctaid.y")
+ if cby == 1:
+ pass
+
+ with tvm.target.Target("cuda"):
+ After = LowerTIRx()(Before)
+ tvm.ir.assert_structural_equal(After, Expected)
+
+ assert len(seen) == 1
+ assert _int_range(seen[0], "cbx") == (2, 0)
+ assert _int_range(seen[0], "cby") == (1, 1)
+
+ @I.ir_module
+ class fixed_shape_instruction:
+ @T.function(private=True)
+ def main(
+ A: T.Tensor((1,), "float32", layout="default"),
+ B: T.Tensor((1,), "float32", layout="default"),
+ ):
+ T.device_entry(
+ launch=T.cuda.LaunchConfig(
+ grid=(2, 2), block=32, cluster=1, preferred_cluster=(2, 2)
+ )
+ )
+ getattr(T.cuda.tile, variant)(B[0:1], A[0:1], scope="cluster")
+
+ with tvm.target.Target("cuda"), pytest.raises(ValueError, match="one static cluster shape"):
+ LowerTIRx()(fixed_shape_instruction)
+
+
+def test_lower_exec_context_infers_warpgroup_range_predicate_for_dispatch():
+ seen = []
+ variant = "__probe_exec_context_warpgroup_range_predicate__"
+
+ @_register_probe(variant)
+ def _probe(op_call, sctx):
+ seen.append({"scope_kind": sctx.scope_kind, "inter": sctx.inter, "intra": sctx.intra})
+
+ return _probe_impl
+
+ @I.ir_module
+ class Before:
+ @T.function(private=True)
+ def main(
+ A: T.Tensor((1,), "float32", scope="global", layout="default"),
+ B: T.Tensor((1,), "float32", scope="global", layout="default"),
+ ):
+ T.device_entry(launch=T.cuda.LaunchConfig(grid=(1,), block=2 * 128))
+ wg_id = T.cuda.warpgroup_id()
+ _lane = T.cuda.lane_id()
+ if wg_id == 0:
+ getattr(T.cuda.tile, variant)(B[0:1], A[0:1], scope="warpgroup")
+ if (0 <= wg_id) & (wg_id < 1):
+ getattr(T.cuda.tile, variant)(B[0:1], A[0:1], scope="warpgroup")
+ if (0 <= wg_id) & (wg_id < 1):
+ getattr(T.cuda.tile, variant)(B[0:1], A[0:1], scope="warpgroup")
+
+ @I.ir_module
+ class Expected:
+ @T.function(private=True)
+ def main(A: T.Tensor((1,), T.float32), B: T.Tensor((1,), T.float32)):
+ B_1 = T.decl_tensor(B.data, ty_args=[T.Tensor((1,), T.float32)])
+ A_1 = T.decl_tensor(A.data, ty_args=[T.Tensor((1,), T.float32)])
+ with T.device_scope(1, 1, 1, 256, 1, 1, attrs=_CTA_LAUNCH_ATTRS):
+ warp_id_in_cta: T.let[T.int32] = T.gpu_warp_shuffle(
+ T.uint32(4294967295), T.cuda.thread_idx("x") // 32, 0, 32, 32
+ )
+ wg_id: T.let[T.int32] = warp_id_in_cta // 4
+ _lane: T.let[T.int32] = T.cuda.thread_idx("x") % 32
+ if wg_id == 0:
+ pass
+ if (0 <= wg_id) & (wg_id < 1):
+ pass
+ if (0 <= wg_id) & (wg_id < 1):
+ pass
+
+ with tvm.target.Target("cuda"):
+ After = LowerTIRx()(Before)
+ tvm.ir.assert_structural_equal(After, Expected)
+
+ assert len(seen) == 3
+ for item in seen:
+ assert item["scope_kind"] == "warpgroup"
+ assert _int_range(item["inter"], "wgid") == (1, 0)
+ assert _int_range(item["inter"], "cta_id") == (1, 0)
+ assert _int_range(item["intra"], "laneid") == (32, 0)
+ assert _int_range(item["intra"], "wid_in_wg") == (4, 0)
+
+
+def test_lower_exec_context_tracks_cta_thread_range_predicate_for_dispatch():
+ seen = []
+ variant = "__probe_exec_context_cta_thread_range_predicate__"
+
+ @_register_probe(variant)
+ def _probe(op_call, sctx):
+ seen.append({"scope_kind": sctx.scope_kind, "inter": sctx.inter, "intra": sctx.intra})
+
+ return _probe_impl
+
+ @I.ir_module
+ class Before:
+ @T.function(private=True)
+ def main(
+ A: T.Tensor((1,), "float32", scope="global", layout="default"),
+ B: T.Tensor((1,), "float32", scope="global", layout="default"),
+ ):
+ T.device_entry(launch=T.cuda.LaunchConfig(grid=(1,), block=(256,)))
+ tid = T.cuda.thread_idx("x")
+ if (0 <= tid) & (tid < 128):
+ getattr(T.cuda.tile, variant)(B[0:1], A[0:1])
+
+ @I.ir_module
+ class Expected:
+ @T.function(private=True)
+ def main(A: T.Tensor((1,), T.float32), B: T.Tensor((1,), T.float32)):
+ B_1 = T.decl_tensor(B.data, ty_args=[T.Tensor((1,), T.float32)])
+ A_1 = T.decl_tensor(A.data, ty_args=[T.Tensor((1,), T.float32)])
+ with T.device_scope(1, 1, 1, 256, 1, 1, attrs=_CTA_LAUNCH_ATTRS):
+ tid: T.let[T.int32] = T.cuda.thread_idx("x")
+ if (0 <= tid) & (tid < 128):
+ pass
+
+ with tvm.target.Target("cuda"):
+ After = LowerTIRx()(Before)
+ tvm.ir.assert_structural_equal(After, Expected)
+
+ assert len(seen) == 1
+ assert seen[0]["scope_kind"] == "thread"
+ assert _int_range(seen[0]["inter"], "laneid") == (32, 0)
+ assert _int_range(seen[0]["inter"], "warpid") == (4, 0)
+ assert _int_range(seen[0]["inter"], "cta_id") == (1, 0)
+ assert len(seen[0]["intra"]) == 0
+
+
+def test_lower_exec_context_tracks_cta_thread_single_warp_range_predicate():
+ seen = []
+ variant = "__probe_exec_context_cta_thread_single_warp_range_predicate__"
+
+ @_register_probe(variant)
+ def _probe(op_call, sctx):
+ seen.append({"scope_kind": sctx.scope_kind, "inter": sctx.inter, "intra": sctx.intra})
+
+ return _probe_impl
+
+ @I.ir_module
+ class Before:
+ @T.function(private=True)
+ def main(
+ A: T.Tensor((1,), "float32", scope="global", layout="default"),
+ B: T.Tensor((1,), "float32", scope="global", layout="default"),
+ ):
+ T.device_entry(launch=T.cuda.LaunchConfig(grid=(1,), block=(256,)))
+ tid = T.cuda.thread_idx("x")
+ if (34 <= tid) & (tid < 40):
+ getattr(T.cuda.tile, variant)(B[0:1], A[0:1])
+
+ @I.ir_module
+ class Expected:
+ @T.function(private=True)
+ def main(A: T.Tensor((1,), T.float32), B: T.Tensor((1,), T.float32)):
+ B_1 = T.decl_tensor(B.data, ty_args=[T.Tensor((1,), T.float32)])
+ A_1 = T.decl_tensor(A.data, ty_args=[T.Tensor((1,), T.float32)])
+ with T.device_scope(1, 1, 1, 256, 1, 1, attrs=_CTA_LAUNCH_ATTRS):
+ tid: T.let[T.int32] = T.cuda.thread_idx("x")
+ if (34 <= tid) & (tid < 40):
+ pass
+
+ with tvm.target.Target("cuda"):
+ After = LowerTIRx()(Before)
+ tvm.ir.assert_structural_equal(After, Expected)
+
+ assert len(seen) == 1
+ assert seen[0]["scope_kind"] == "thread"
+ assert _int_range(seen[0]["inter"], "laneid") == (6, 2)
+ assert _int_range(seen[0]["inter"], "warpid") == (1, 1)
+ assert _int_range(seen[0]["inter"], "cta_id") == (1, 0)
+ assert len(seen[0]["intra"]) == 0
+
+
+def test_lower_exec_context_tracks_warpgroup_thread_range_predicate():
+ seen = []
+ variant = "__probe_exec_context_warpgroup_thread_range_predicate__"
+
+ @_register_probe(variant)
+ def _probe(op_call, sctx):
+ seen.append({"scope_kind": sctx.scope_kind, "inter": sctx.inter, "intra": sctx.intra})
+
+ return _probe_impl
+
+ @I.ir_module
+ class Before:
+ @T.function(private=True)
+ def main(
+ A: T.Tensor((1,), "float32", scope="global", layout="default"),
+ B: T.Tensor((1,), "float32", scope="global", layout="default"),
+ ):
+ T.device_entry(launch=T.cuda.LaunchConfig(grid=(1,), block=2 * 128))
+ wg_id = T.cuda.warpgroup_id()
+ tid_in_wg = T.cuda.thread_in_warpgroup()
+ if wg_id == 1:
+ if (32 <= tid_in_wg) & (tid_in_wg < 64):
+ getattr(T.cuda.tile, variant)(B[0:1], A[0:1], scope="warpgroup")
+
+ @I.ir_module
+ class Expected:
+ @T.function(private=True)
+ def main(A: T.Tensor((1,), T.float32), B: T.Tensor((1,), T.float32)):
+ B_1 = T.decl_tensor(B.data, ty_args=[T.Tensor((1,), T.float32)])
+ A_1 = T.decl_tensor(A.data, ty_args=[T.Tensor((1,), T.float32)])
+ with T.device_scope(1, 1, 1, 256, 1, 1, attrs=_CTA_LAUNCH_ATTRS):
+ warp_id_in_cta: T.let[T.int32] = T.gpu_warp_shuffle(
+ T.uint32(4294967295), T.cuda.thread_idx("x") // 32, 0, 32, 32
+ )
+ wg_id: T.let[T.int32] = warp_id_in_cta // 4
+ tid_in_wg: T.let[T.int32] = T.cuda.thread_idx("x") % 128
+ if wg_id == 1:
+ if (32 <= tid_in_wg) & (tid_in_wg < 64):
+ pass
+
+ with tvm.target.Target("cuda"):
+ After = LowerTIRx()(Before)
+ tvm.ir.assert_structural_equal(After, Expected)
+
+ assert len(seen) == 1
+ assert seen[0]["scope_kind"] == "warpgroup"
+ assert _int_range(seen[0]["inter"], "wgid") == (1, 1)
+ assert _int_range(seen[0]["inter"], "cta_id") == (1, 0)
+ assert _int_range(seen[0]["intra"], "laneid") == (32, 0)
+ assert _int_range(seen[0]["intra"], "wid_in_wg") == (1, 1)
+
+
+def test_lower_exec_context_tracks_dependent_conjunctive_predicate():
+ seen = []
+ variant = "__probe_exec_context_dependent_conjunctive_predicate__"
+
+ @_register_probe(variant)
+ def _probe(op_call, sctx):
+ seen.append({"scope_kind": sctx.scope_kind, "inter": sctx.inter, "intra": sctx.intra})
+
+ return _probe_impl
+
+ @I.ir_module
+ class Before:
+ @T.function(private=True)
+ def main(
+ A: T.Tensor((1,), "float32", scope="global", layout="default"),
+ B: T.Tensor((1,), "float32", scope="global", layout="default"),
+ ):
+ T.device_entry(launch=T.cuda.LaunchConfig(grid=(1,), block=2 * 128))
+ wg_id = T.cuda.warpgroup_id()
+ tid_in_wg = T.cuda.thread_in_warpgroup()
+ if ((32 <= tid_in_wg) & (tid_in_wg < 64)) & (wg_id == 1):
+ getattr(T.cuda.tile, variant)(B[0:1], A[0:1], scope="warpgroup")
+
+ @I.ir_module
+ class Expected:
+ @T.function(private=True)
+ def main(A: T.Tensor((1,), T.float32), B: T.Tensor((1,), T.float32)):
+ B_1 = T.decl_tensor(B.data, ty_args=[T.Tensor((1,), T.float32)])
+ A_1 = T.decl_tensor(A.data, ty_args=[T.Tensor((1,), T.float32)])
+ with T.device_scope(1, 1, 1, 256, 1, 1, attrs=_CTA_LAUNCH_ATTRS):
+ warp_id_in_cta: T.let[T.int32] = T.gpu_warp_shuffle(
+ T.uint32(4294967295), T.cuda.thread_idx("x") // 32, 0, 32, 32
+ )
+ wg_id: T.let[T.int32] = warp_id_in_cta // 4
+ tid_in_wg: T.let[T.int32] = T.cuda.thread_idx("x") % 128
+ if (32 <= tid_in_wg) & (tid_in_wg < 64) & (wg_id == 1):
+ pass
+
+ with tvm.target.Target("cuda"):
+ After = LowerTIRx()(Before)
+ tvm.ir.assert_structural_equal(After, Expected)
+
+ assert len(seen) == 1
+ assert seen[0]["scope_kind"] == "warpgroup"
+ assert _int_range(seen[0]["inter"], "wgid") == (1, 1)
+ assert _int_range(seen[0]["inter"], "cta_id") == (1, 0)
+ assert _int_range(seen[0]["intra"], "laneid") == (32, 0)
+ assert _int_range(seen[0]["intra"], "wid_in_wg") == (1, 1)
+
+
+def test_lower_exec_context_keeps_plain_predicate_condition():
+ @I.ir_module
+ class Before:
+ @T.function(private=True)
+ def main(A: T.Tensor((1,), "float32", scope="global", layout="default")):
+ T.device_entry(launch=T.cuda.LaunchConfig(grid=(1,), block=2 * 128))
+ wg_id = T.cuda.warpgroup_id()
+ _lane = T.cuda.lane_id()
+ if wg_id == 0:
+ T.evaluate(A[0])
+
+ @I.ir_module
+ class Expected:
+ @T.function(private=True)
+ def main(A: T.Tensor((1,), T.float32)):
+ A_1 = T.decl_tensor(A.data, ty_args=[T.Tensor((1,), T.float32)])
+ with T.device_scope(1, 1, 1, 256, 1, 1, attrs=_CTA_LAUNCH_ATTRS):
+ warp_id_in_cta: T.let[T.int32] = T.gpu_warp_shuffle(
+ T.uint32(4294967295), T.cuda.thread_idx("x") // 32, 0, 32, 32
+ )
+ wg_id: T.let[T.int32] = warp_id_in_cta // 4
+ _lane: T.let[T.int32] = T.cuda.thread_idx("x") % 32
+ if wg_id == 0:
+ T.evaluate(A_1[0])
+
+ with tvm.target.Target("cuda"):
+ lowered = LowerTIRx()(Before)
+ tvm.ir.assert_structural_equal(lowered, Expected)
+
+ script = lowered.script(extra_config={"tirx.prefix": "T"})
+ assert "if wg_id == 0:" in script
+ assert "0 <= v_1" not in script
+ assert "v_1 < 1" not in script
+
+
+def test_lower_exec_context_keeps_plain_scope_predicate_condition():
+ @I.ir_module
+ class Before:
+ @T.function(private=True)
+ def main(A: T.Tensor((1,), "float32", scope="global", layout="default")):
+ T.device_entry(launch=T.cuda.LaunchConfig(grid=(1,), block=2 * 128))
+ wg_id = T.cuda.warpgroup_id()
+ _lane = T.cuda.lane_id()
+ if wg_id == 0:
+ A[0] = T.float32(1)
+
+ @I.ir_module
+ class Expected:
+ @T.function(private=True)
+ def main(A: T.Tensor((1,), T.float32)):
+ A_1 = T.decl_tensor(A.data, ty_args=[T.Tensor((1,), T.float32)])
+ with T.device_scope(1, 1, 1, 256, 1, 1, attrs=_CTA_LAUNCH_ATTRS):
+ warp_id_in_cta: T.let[T.int32] = T.gpu_warp_shuffle(
+ T.uint32(4294967295), T.cuda.thread_idx("x") // 32, 0, 32, 32
+ )
+ wg_id: T.let[T.int32] = warp_id_in_cta // 4
+ _lane: T.let[T.int32] = T.cuda.thread_idx("x") % 32
+ if wg_id == 0:
+ A_1[0] = T.float32(1.0)
+
+ with tvm.target.Target("cuda"):
+ lowered = LowerTIRx()(Before)
+ tvm.ir.assert_structural_equal(lowered, Expected)
+
+ script = lowered.script(extra_config={"tirx.prefix": "T"})
+ assert "if wg_id == 0:" in script
+ assert "0 <= v_1" not in script
+ assert "v_1 < 1" not in script
+
+
+def test_simplify_uses_floor_div_scope_predicate_as_context_fact():
+ @I.ir_module
+ class Before:
+ @T.function(private=True)
+ def main(A: T.Tensor((16,), "float32", scope="global", layout="default")):
+ T.device_entry(launch=T.cuda.LaunchConfig(grid=(1,), block=2 * 128))
+ wg_id = T.cuda.warpgroup_id()
+ warp_id = T.cuda.warp_in_warpgroup()
+ lane_id = T.cuda.lane_id()
+ if wg_id == 0:
+ A[warp_id] = T.float32(lane_id)
+
+ @I.ir_module
+ class Expected:
+ @T.function(private=True)
+ def main(A: T.Tensor((16,), T.float32)):
+ A_1 = T.decl_tensor(A.data, ty_args=[T.Tensor((16,), T.float32)])
+ with T.device_scope(1, 1, 1, 256, 1, 1, attrs=_CTA_LAUNCH_ATTRS):
+ warp_id_in_cta: T.let[T.int32] = T.gpu_warp_shuffle(
+ T.uint32(4294967295), T.cuda.thread_idx("x") // 32, 0, 32, 32
+ )
+ wg_id: T.let[T.int32] = warp_id_in_cta // 4
+ warp_id: T.let[T.int32] = warp_id_in_cta % 4
+ lane_id: T.let[T.int32] = T.cuda.thread_idx("x") % 32
+ if wg_id == 0:
+ A_1[warp_id_in_cta % 4] = T.Cast(lane_id, "float32")
+
+ @I.ir_module
+ class ExpectedSimplified:
+ @T.function(private=True)
+ def main(A: T.Tensor((16,), T.float32)):
+ A_1 = T.decl_tensor(A.data, ty_args=[T.Tensor((16,), T.float32)])
+ with T.device_scope(1, 1, 1, 256, 1, 1, attrs=_CTA_LAUNCH_ATTRS):
+ warp_id_in_cta: T.let[T.int32] = T.gpu_warp_shuffle(
+ T.uint32(4294967295), T.cuda.thread_idx("x") // 32, 0, 32, 32
+ )
+ wg_id: T.let[T.int32] = warp_id_in_cta // 4
+ warp_id: T.let[T.int32] = warp_id_in_cta % 4
+ lane_id: T.let[T.int32] = T.cuda.thread_idx("x") % 32
+ if warp_id_in_cta // 4 == 0:
+ A_1[warp_id_in_cta] = T.Cast(T.cuda.thread_idx("x") % 32, "float32")
+
+ with tvm.target.Target("cuda"):
+ lowered = LowerTIRx()(Before)
+ simplified = StmtSimplify()(lowered)
+ tvm.ir.assert_structural_equal(lowered, Expected)
+ tvm.ir.assert_structural_equal(simplified, ExpectedSimplified)
+
+ script = simplified.script(extra_config={"tirx.prefix": "T"})
+ assert "if warp_id_in_cta // 4 == 0:" in script
+ assert "if 0 <= warp_id_in_cta" not in script
+ assert "A_1[warp_id_in_cta] = T.Cast" in script
+ assert "A_1[warp_id_in_cta % 4]" not in script
+
+
+def test_lower_exec_context_selector_filter_for_elect_sync():
+ seen = []
+ variant = "__probe_exec_context_elect_selector__"
+
+ @_register_probe(variant)
+ def _probe(op_call, sctx):
+ seen.append(sctx.inter["laneid"][1].script(extra_config={"tirx.prefix": "T"}))
+
+ return _probe_impl
+
+ @I.ir_module
+ class Before:
+ @T.function(private=True)
+ def main(
+ A: T.Tensor((1,), "float32", scope="global", layout="default"),
+ B: T.Tensor((1,), "float32", scope="global", layout="default"),
+ ):
+ T.device_entry(launch=T.cuda.LaunchConfig(grid=(1,), block=1 * 32))
+ lane_id = T.cuda.lane_id()
+ if T.cuda.elect_sync():
+ getattr(T.cuda.tile, variant)(B[0:1], A[0:1])
+ if T.cuda.elect_sync() != 0:
+ getattr(T.cuda.tile, variant)(B[0:1], A[0:1])
+ if T.cuda.elect_sync():
+ getattr(T.cuda.tile, variant)(B[0:1], A[0:1])
+
+ @I.ir_module
+ class Expected:
+ @T.function(private=True)
+ def main(A: T.Tensor((1,), T.float32), B: T.Tensor((1,), T.float32)):
+ B_1 = T.decl_tensor(B.data, ty_args=[T.Tensor((1,), T.float32)])
+ A_1 = T.decl_tensor(A.data, ty_args=[T.Tensor((1,), T.float32)])
+ with T.device_scope(1, 1, 1, 32, 1, 1, attrs=_CTA_LAUNCH_ATTRS):
+ lane_id: T.let[T.int32] = T.cuda.thread_idx("x") % 32
+ if T.cuda.elect_sync():
+ pass
+ if T.cuda.elect_sync() != T.uint32(0):
+ pass
+ if T.cuda.elect_sync():
+ pass
+
+ with tvm.target.Target("cuda"):
+ After = LowerTIRx()(Before)
+ tvm.ir.assert_structural_equal(After, Expected)
+
+ assert len(seen) == 3
+ assert any(
+ "T.gpu_active_thread_selector(lane_id, T.cuda.elect_sync())" in item for item in seen
+ )
+ assert any(
+ "T.gpu_active_thread_selector(lane_id, T.cuda.elect_sync() != T.uint32(0))" in item
+ for item in seen
+ )
+
+
+def test_lower_cleanup_accepts_bool_elect_sync_else_path():
+ @I.ir_module
+ class Before:
+ @T.function(private=True)
+ def main(A: T.Tensor((32,), "int32", scope="global", layout="default")):
+ T.device_entry(launch=T.cuda.LaunchConfig(grid=(1,), block=1 * 32))
+ lane_id = T.cuda.lane_id()
... 8242 lines suppressed ...