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 6c76643d18 [REFACTOR][TIRx] Consolidate analysis and split-pattern
normalization (#20629)
6c76643d18 is described below
commit 6c76643d181a1f9f83bc83a1d3e115a5187201ee
Author: Tianqi Chen <[email protected]>
AuthorDate: Fri Oct 9 12:38:52 2026 -0400
[REFACTOR][TIRx] Consolidate analysis and split-pattern normalization
(#20629)
Keep structural undefined-variable discovery in shared IR with one
visited-variable set. Consolidate TIRx well-formedness checks under one
entry point and remove the separate SSA interfaces, while retaining
S-TIR-specific block checks.
Move split-pattern normalization into TIRx so it is available without
S-TIR dependencies. Preserve compound-expression rewrites, leave
plain-variable comparisons unchanged, and simplify the existing
normalization fixtures to serial TIRx programs.
---
include/tvm/ir/analysis.h | 13 +-
include/tvm/s_tir/transform.h | 6 -
include/tvm/tirx/analysis.h | 53 +-----
include/tvm/tirx/transform.h | 6 +
python/tvm/s_tir/backend/adreno/pipeline.py | 2 +-
python/tvm/s_tir/pipeline.py | 2 +-
python/tvm/s_tir/transform/transform.py | 11 --
python/tvm/tirx/analysis/analysis.py | 42 +----
.../tvm/tirx/script/ir_builder/parser_protocol.py | 1 -
python/tvm/tirx/transform/transform.py | 11 ++
src/tirx/analysis/verify_ssa.cc | 150 ----------------
src/tirx/analysis/verify_tirx_well_formed.cc | 195 ---------------------
src/tirx/analysis/verify_well_formed.cc | 2 +-
src/tirx/analysis/verify_well_formed.h | 95 +++++++++-
.../transform/renormalize_split_pattern.cc | 19 +-
tests/python/relax/test_relax_operators.py | 2 +-
...st_s_tir_transform_renormalize_split_pattern.py | 176 -------------------
.../tirx-analysis/test_tir_analysis_verify_ssa.py | 42 -----
.../test_tir_analysis_verify_well_formed.py | 26 +++
.../test_tir_transform_common_subexpr_elim.py | 2 +-
...test_tir_transform_renormalize_split_pattern.py | 66 +++++++
.../script/test_tirx_script_source_locations.py | 2 +-
tests/python/tirx/test_verifier.py | 6 +-
23 files changed, 231 insertions(+), 699 deletions(-)
diff --git a/include/tvm/ir/analysis.h b/include/tvm/ir/analysis.h
index b40a418bbf..e79c738860 100644
--- a/include/tvm/ir/analysis.h
+++ b/include/tvm/ir/analysis.h
@@ -42,18 +42,17 @@ namespace tvm {
* \return Undefined variables, uniquely ordered by their first use.
*/
inline ffi::Array<Var> UndefinedVars(ffi::AnyView object, const
ffi::Array<Var>& defs = {}) {
- std::unordered_set<const VarNode*> defined;
- std::unordered_set<const VarNode*> used;
- for (const Var& var : defs) defined.insert(var.get());
+ std::unordered_set<const VarNode*> visited_vars;
+ for (const Var& var : defs) visited_vars.insert(var.get());
ffi::Array<Var> undefined;
ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>(
- object, [&](const Var& var, TVMFFIDefRegionKind kind) ->
ffi::Expected<ffi::WalkResult> {
+ object, [&](const VarNode* var, TVMFFIDefRegionKind kind) ->
ffi::Expected<ffi::WalkResult> {
if (kind != kTVMFFIDefRegionKindNone) {
- defined.insert(var.get());
+ visited_vars.insert(var);
return ffi::WalkResult::Advance();
}
- if (!defined.count(var.get()) && used.insert(var.get()).second) {
- undefined.push_back(var);
+ if (visited_vars.insert(var).second) {
+ undefined.push_back(ffi::GetRef<Var>(var));
}
return ffi::WalkResult::Skip();
});
diff --git a/include/tvm/s_tir/transform.h b/include/tvm/s_tir/transform.h
index 1b67ce5084..cc22793fc3 100644
--- a/include/tvm/s_tir/transform.h
+++ b/include/tvm/s_tir/transform.h
@@ -251,12 +251,6 @@ TVM_DLL Pass HoistIfThenElse(tvm::ffi::String variant =
"");
*/
TVM_DLL Pass HoistExpression();
-/*!
- * \brief Renormalize the split pattern from floordiv(floormod()) to
floormod(floordiv()).
- * \return The pass.
- */
-TVM_DLL Pass RenormalizeSplitPattern();
-
/*!
* \brief Detect and rewrite unsafe select that contains memory access.
* \return The pass.
diff --git a/include/tvm/tirx/analysis.h b/include/tvm/tirx/analysis.h
index c2952b7613..4445d5227a 100644
--- a/include/tvm/tirx/analysis.h
+++ b/include/tvm/tirx/analysis.h
@@ -38,17 +38,6 @@ namespace tvm {
namespace tirx {
-/*!
- * \brief Verifies whether the IR stmt or Expr is in SSA form.
- * That is: each Var is defined and assigned once(in Let/For)
- *
- * \param func The function to be verified.
- * \return Whether IR is in SSA form.
- *
- * \note All passes in TIR consume and produce SSA form.
- */
-TVM_DLL bool VerifySSA(const Function& func);
-
/*!
* \brief Verify if memory accesses are legal for a specific target device
type.
*
@@ -91,6 +80,10 @@ TVM_DLL size_t CalculateWorkspaceBytes(const Function& func,
int64_t workspace_b
*
* - Each variable has a single point of definition.
*
+ * - Buffer loads have consistent source, indices, and result types.
+ *
+ * - Loop control, tile-operation categories, and scope IDs are valid.
+ *
* Dialect statements require their dialect-specific verifier.
*
* \param func The Function to be verified.
@@ -129,14 +122,6 @@ namespace transform {
using tvm::transform::Pass;
using tvm::transform::PassContext;
-/*!
- * \brief Pass variant of VerifySSA.
- *
- * \returns The pass.
- * \sa tvm::tirx::VerifySSA
- */
-TVM_DLL Pass VerifySSA();
-
/*!
* \brief Pass variant of VerifyMemory.
*
@@ -145,36 +130,6 @@ TVM_DLL Pass VerifySSA();
*/
TVM_DLL Pass VerifyMemory();
-/*!
- * \brief Pass variant of VerifyGPUCode.
- *
- * \param constraints The dict to specify constraints to check.
- *
- * \returns The pass.
- * \sa tvm::tir::VerifyGPUCode
- */
-/******** TIRx analysis helpers ********/
-
-/*!
- * \brief Verify if the given TIRX is well-formed.
- * \param func The Function to be verified.
- * \param assert_mode The indicator if it raises an error when the function is
not well-formed.
- * \param device_func The indicator if it is a device function.
- * \return Whether it is a well-formed TIRX function.
- */
-TVM_DLL bool VerifyTIRxWellFormed(const Function& func, bool assert_mode =
true,
- bool device_func = false);
-
-/*!
- * \brief Verify if the TIRX in the given IRMOdule is well-formed.
- * \param mod The IRModule to be verified.
- * \param assert_mode The indicator if it raises an error when the function is
not well-formed.
- * \param device_func The indicator if it is a device function.
- * \return Whether it is a well-formed TIRX module.
- */
-TVM_DLL bool VerifyTIRxWellFormed(const IRModule& mod, bool assert_mode = true,
- bool device_func = false);
-
} // namespace transform
} // namespace tirx
} // namespace tvm
diff --git a/include/tvm/tirx/transform.h b/include/tvm/tirx/transform.h
index 38f03f94cf..d1464701dc 100644
--- a/include/tvm/tirx/transform.h
+++ b/include/tvm/tirx/transform.h
@@ -98,6 +98,12 @@ TVM_DLL Pass UnrollLoop();
*/
TVM_DLL Pass RemoveNoOp();
+/*!
+ * \brief Renormalize the split pattern from floordiv(floormod()) to
floormod(floordiv()).
+ * \return The pass.
+ */
+TVM_DLL Pass RenormalizeSplitPattern();
+
/*!
* \brief Run statement-level arithmetic simplifications on the TIR Function.
*
diff --git a/python/tvm/s_tir/backend/adreno/pipeline.py
b/python/tvm/s_tir/backend/adreno/pipeline.py
index b7910477be..1a2decf054 100644
--- a/python/tvm/s_tir/backend/adreno/pipeline.py
+++ b/python/tvm/s_tir/backend/adreno/pipeline.py
@@ -64,7 +64,7 @@ def default_tir_pipeline(*, prepare_only=False):
[
s_tir.transform.HoistIfThenElse(),
tirx.transform.UnrollLoop(),
- s_tir.transform.RenormalizeSplitPattern(),
+ tirx.transform.RenormalizeSplitPattern(),
s_tir.transform.StmtSimplify(),
tirx.transform.RemoveNoOp(),
s_tir.transform.RewriteUnsafeSelect(),
diff --git a/python/tvm/s_tir/pipeline.py b/python/tvm/s_tir/pipeline.py
index 08cfee0c83..01d31dd7a1 100644
--- a/python/tvm/s_tir/pipeline.py
+++ b/python/tvm/s_tir/pipeline.py
@@ -64,7 +64,7 @@ def default_s_tir_pipeline(*, prepare_only=False):
[
s_tir.transform.HoistIfThenElse(),
tirx.transform.UnrollLoop(),
- s_tir.transform.RenormalizeSplitPattern(),
+ tirx.transform.RenormalizeSplitPattern(),
s_tir.transform.StmtSimplify(),
tirx.transform.RemoveNoOp(),
s_tir.transform.RewriteUnsafeSelect(),
diff --git a/python/tvm/s_tir/transform/transform.py
b/python/tvm/s_tir/transform/transform.py
index 8bf9ff3d62..1b4fc61f48 100644
--- a/python/tvm/s_tir/transform/transform.py
+++ b/python/tvm/s_tir/transform/transform.py
@@ -290,17 +290,6 @@ def HoistExpression():
return _ffi_api.HoistExpression() # type: ignore
-def RenormalizeSplitPattern():
- """Renormalize the split pattern from floordiv(floormod()) to
floormod(floordiv())
-
- Returns
- -------
- fpass : tvm.transform.Pass
- The result pass
- """
- return _ffi_api.RenormalizeSplitPattern() # type: ignore
-
-
def RewriteUnsafeSelect():
"""Detect and rewrite unsafe select that contains memory access.
diff --git a/python/tvm/tirx/analysis/analysis.py
b/python/tvm/tirx/analysis/analysis.py
index 6debad5c38..c5fa88143f 100644
--- a/python/tvm/tirx/analysis/analysis.py
+++ b/python/tvm/tirx/analysis/analysis.py
@@ -25,22 +25,6 @@ from ..function import Function
from . import _ffi_api
-def verify_ssa(func: Function) -> bool:
- """Verify if the func is in SSA form.
-
- Parameters
- ----------
- func: tvm.tirx.Function
- The module to be verified.
-
- Returns
- -------
- result : bool
- The result of verification.
- """
- return _ffi_api.verify_ssa(func) # type: ignore
-
-
def verify_memory(func: Function) -> bool:
"""Verify if func contains illegal host side direct memory access.
@@ -78,7 +62,7 @@ def undefined_vars(node: Stmt | Expr, defs: list[Var] | None
= None) -> list[Var
def verify_well_formed(obj: Function | IRModule, assert_mode: bool = True) ->
bool:
- """Verify definitions and buffer-load types in ordinary TIRX.
+ """Verify definitions, buffer-load types, loop control, tile ops, and
scope IDs.
Use ``tvm.s_tir.analysis.verify_well_formed`` for schedulable blocks.
@@ -96,27 +80,3 @@ def verify_well_formed(obj: Function | IRModule,
assert_mode: bool = True) -> bo
Whether it is a well-formed TIR function.
"""
return _ffi_api.VerifyWellFormed(obj, assert_mode) # type: ignore #
pylint: disable=no-member
-
-
-def verify_tirx_well_formed(
- obj: Function | IRModule, assert_mode: bool = True, device_func: bool =
False
-) -> bool:
- """Verify if the given TIRX is well-formed.
-
- Parameters
- ----------
- obj: Union[tvm.tirx.Function, tvm.ir.IRModule]
- The function or module to be verified.
-
- assert_mode: bool
- The indicator if it raises an error when the function is not
well-formed.
-
- device_func: bool
- The indicator if it is a device function.
-
- Returns
- -------
- result: bool
- Whether it is a well-formed TIRX function.
- """
- return _ffi_api.VerifyTIRxWellFormed(obj, assert_mode, device_func) #
type: ignore # pylint: disable=no-member
diff --git a/python/tvm/tirx/script/ir_builder/parser_protocol.py
b/python/tvm/tirx/script/ir_builder/parser_protocol.py
index 949734d98a..0369929183 100644
--- a/python/tvm/tirx/script/ir_builder/parser_protocol.py
+++ b/python/tvm/tirx/script/ir_builder/parser_protocol.py
@@ -154,7 +154,6 @@ def check_well_formed_(function: _tir.Function) -> None:
"""Validate a completed TIRx function."""
try:
_tir.analysis.verify_well_formed(function)
- _tir.analysis.verify_tirx_well_formed(function)
except Exception as error:
raise ValueError(
"Program is not well-formed. If this is deliberate, set "
diff --git a/python/tvm/tirx/transform/transform.py
b/python/tvm/tirx/transform/transform.py
index d40d5d9b15..c0b5eeea23 100644
--- a/python/tvm/tirx/transform/transform.py
+++ b/python/tvm/tirx/transform/transform.py
@@ -215,6 +215,17 @@ def CommonSubexprElim():
return _ffi_api.CommonSubexprElim() # type: ignore
+def RenormalizeSplitPattern():
+ """Renormalize the split pattern from floordiv(floormod()) to
floormod(floordiv())
+
+ Returns
+ -------
+ fpass : tvm.transform.Pass
+ The result pass
+ """
+ return _ffi_api.RenormalizeSplitPattern() # type: ignore
+
+
@_ffi.register_object("tirx.transform.StmtSimplifyConfig")
class StmtSimplifyConfig(_ffi.Object):
"""Config for stmt simplify pass"""
diff --git a/src/tirx/analysis/verify_ssa.cc b/src/tirx/analysis/verify_ssa.cc
deleted file mode 100644
index fad6db5e65..0000000000
--- a/src/tirx/analysis/verify_ssa.cc
+++ /dev/null
@@ -1,150 +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.
- */
-
-/*!
- * SSA related checks and pass.
- *
- * SSA requires each varaible to be only defined once.
- * \file verify_ssa.cc
- */
-#include <tvm/ffi/cast.h>
-#include <tvm/ffi/function.h>
-#include <tvm/ffi/reflection/registry.h>
-#include <tvm/ir/prim/expr.h>
-#include <tvm/tirx/analysis.h>
-#include <tvm/tirx/stmt_functor.h>
-
-#include <unordered_map>
-#include <unordered_set>
-#include <vector>
-
-namespace tvm {
-namespace tirx {
-
-class SSAVerifier final : public StmtExprVisitor {
- public:
- bool is_ssa_{true};
-
- ffi::Optional<VisitInterrupt> Visit(ffi::AnyView n) final {
- if (!is_ssa_) return std::nullopt;
- return StmtExprVisitor::Visit(n);
- }
- ffi::Optional<VisitInterrupt> Visit_(const prim::LetNode* op) final {
- // Weaker SSA condition
- // A single var can be binded in multiple lets
- // but they have to bind to the same value.
- // This is used to enable cases when we reuse a single let
- // expression to cosntruct a nested expr.
- // (let x = 1 in x + 1) * (let x = 1 in x + 1)
- auto it = def_map_.find(op->var);
- if (it != def_map_.end()) {
- if (!deep_equal_(it->second.as_or_throw<PrimExpr>(), op->value)) {
- is_ssa_ = false;
- return std::nullopt;
- }
- } else {
- MarkDef(op->var, op->value);
- }
- return StmtExprVisitor::Visit_(op);
- }
-
- ffi::Optional<VisitInterrupt> Visit_(const BindNode* op) final {
- MarkDef(op->var, op->value);
- return StmtExprVisitor::Visit_(op);
- }
- ffi::Optional<VisitInterrupt> Visit_(const ForNode* op) final {
- MarkDef(op->loop_var, op->loop_var);
- return StmtExprVisitor::Visit_(op);
- }
-
- ffi::Optional<VisitInterrupt> Visit_(const VarNode* node) final {
- auto var = ffi::GetRef<Var>(node);
- if (match_scope_) {
- MarkDef(var, var, true);
- }
- return StmtExprVisitor::Visit_(node);
- }
-
- void Run(const Function& func) {
- for (auto param : func->params) {
- MarkDef(param, param);
- }
-
- match_scope_ = true;
- for (const Var& param : func->params) {
- WithDefRegionKind(kTVMFFIDefRegionKindPattern, [&] { return
Visit(param->ty); });
- }
- match_scope_ = false;
- this->Visit(func->body);
- }
-
- private:
- void MarkDef(const Var& var, Expr value, bool allow_dup = false) {
- if (def_map_.count(var) != 0) {
- if (!allow_dup) {
- is_ssa_ = false;
- return;
- }
- } else {
- def_map_.insert_or_assign(var, value);
- }
- }
- // whether we are in match scope, where a var can occur multiple times.
- bool match_scope_{false};
- // deep equal
- prim::ExprDeepEqual deep_equal_;
- // def map, for let, maps to the bind value, for others maps to self.
- std::unordered_map<Var, Expr> def_map_;
-};
-
-bool VerifySSA(const Function& func) {
- auto visitor = ffi::make_object<SSAVerifier>();
- visitor->Run(func);
- return visitor->is_ssa_;
-}
-
-TVM_FFI_STATIC_INIT_BLOCK() {
- namespace refl = tvm::ffi::reflection;
- refl::GlobalDef().def("tirx.analysis.verify_ssa", VerifySSA);
-}
-
-namespace transform {
-
-Pass VerifySSA() {
- auto pass_func = [=](IRModule mod, PassContext ctx) {
- for (auto kv : mod->functions) {
- if (auto func = kv.second.as<Function>()) {
- TVM_FFI_CHECK(VerifySSA(func.value()), RuntimeError)
- << "IR is not in SSA form" << func.value();
- }
- }
- return mod;
- };
- return tvm::transform::CreateModulePass(pass_func, 0, "tirx.VerifySSA", {});
-}
-
-TVM_FFI_STATIC_INIT_BLOCK() {
- namespace refl = tvm::ffi::reflection;
- refl::GlobalDef().def("tirx.transform.VerifySSA", VerifySSA);
-}
-
-} // namespace transform
-
-} // namespace tirx
-} // namespace tvm
diff --git a/src/tirx/analysis/verify_tirx_well_formed.cc
b/src/tirx/analysis/verify_tirx_well_formed.cc
deleted file mode 100644
index 1f11f1b41e..0000000000
--- a/src/tirx/analysis/verify_tirx_well_formed.cc
+++ /dev/null
@@ -1,195 +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 tir/analysis/verify_tirx_well_formed.cc
- * \brief Check if the TIRX program is well-formed.
- */
-
-#include <tvm/ir/op.h>
-#include <tvm/ir/prim/op.h>
-#include <tvm/runtime/logging.h>
-#include <tvm/sym/analyzer.h>
-#include <tvm/tirx/analysis.h>
-#include <tvm/tirx/exec_scope.h>
-#include <tvm/tirx/op/region.h>
-#include <tvm/tirx/op_attr_types.h>
-#include <tvm/tirx/stmt.h>
-#include <tvm/tirx/stmt_functor.h>
-
-#include <exception>
-#include <optional>
-#include <tuple>
-#include <variant>
-
-#include "../ir/tir_visitor_with_path.h"
-#include "tvm/ir/module.h"
-#include "verify_well_formed.h"
-
-namespace tvm {
-namespace tirx {
-
-class ExecScopeVerifier : public Verifier<ExecScopeVerifier> {
- public:
- using Verifier::Verifier;
-
- private:
- using Verifier::Visit;
-
- void Dispatch_(const tirx::TileOpCallNode* op, ffi::reflection::AccessPath
path) override {
- static const auto& category_map =
Op::GetAttrMap<tirx::TIRxOpCategory>("TIRxOpCategory");
- Verify(category_map.get(op->op, ffi::String("")) == "tile_primitive")
- << "TIRxError: TileOpCall at " << path << " has non-tile op " <<
op->op;
- }
-};
-
-class ScopeIdVerifier : public Verifier<ScopeIdVerifier> {
- public:
- using Verifier::Verifier;
-
- private:
- using Verifier::Visit;
-
- void Dispatch_(const RegionStmtNode* op, ffi::reflection::AccessPath path)
override {
- if (op->op.same_as(tirx::device_entry_op())) {
- // Device-region marker: defs gathered from the body are verified when
- // the region exits, with launch-param sanity enforced as ``is_root``.
- size_t baseline = scope_id_def_.size();
- Verifier::Dispatch_(op, path);
- size_t total = scope_id_def_.size();
- if (total > baseline) {
- RunScopeIdVerify(path, baseline, /*is_root=*/true);
- }
- while (scope_id_def_.size() > baseline) {
- scope_id_def_.pop_back();
- }
- return;
- }
- Verifier::Dispatch_(op, path);
- }
-
- void RunScopeIdVerify(ffi::reflection::AccessPath path, size_t baseline,
bool is_root) {
- ScopeIdDefVerifier verifier;
- Verify(verifier.Verify(scope_id_def_, ScopeIdDefVerifier::Mode::kRelaxed))
- << "TIRxError: Scope at " << path << " has invalid scope_id_def";
- if (is_root) {
- // Enforce launch-parameter sanity at the device-region root.
- auto cta_thread_it = verifier.id_set.find(ScopeBinding::kCtaThread);
- if (cta_thread_it != verifier.id_set.end() &&
!(*cta_thread_it).second.is_deferred()) {
- PrimExpr ext = (*cta_thread_it).second.fused_extent();
- if (const auto* imm = ext.as<IntImmNode>()) {
- Verify(imm->value > 0) << "TIRxError: kernel at " << path
- << " has non-positive thread count " <<
imm->value;
- bool needs_warp_align =
verifier.id_set.count(ScopeBinding::kCtaWarp) ||
-
verifier.id_set.count(ScopeBinding::kWarpThread) ||
-
verifier.id_set.count(ScopeBinding::kCtaWarpgroup) ||
-
verifier.id_set.count(ScopeBinding::kWarpgroupWarp);
- if (needs_warp_align) {
- Verify(imm->value % 32 == 0)
- << "TIRxError: kernel at " << path << " uses warp-granular
bindings"
- << " but has thread count " << imm->value << " not a multiple
of 32";
- }
- }
- }
- }
- }
-
- void Dispatch_(const ScopeIdDefStmtNode* op, ffi::reflection::AccessPath
path) override {
- scope_id_def_.push_back(op->def);
- Verifier::Dispatch_(op, path);
- }
-
- Array<ScopeIdDef> scope_id_def_;
- sym::Analyzer ana_;
-};
-
-class LayoutVerifier : public Verifier<LayoutVerifier> {
- public:
- using Verifier::Verifier;
-
- private:
- using Verifier::Visit;
-};
-
-class AsyncStructsVerifier : public Verifier<AsyncStructsVerifier> {
- public:
- using Verifier::Verifier;
-
- private:
- using Verifier::Visit;
-};
-
-class DeviceFuncVerifier : public Verifier<DeviceFuncVerifier> {
- public:
- using Verifier::Verifier;
-
- private:
- using Verifier::Visit;
-};
-
-bool VerifyTIRxWellFormed(const Function& func, bool assert_mode, bool
device_func) {
- if (!ExecScopeVerifier::Verify(func, assert_mode)) {
- return false;
- }
- if (!ScopeIdVerifier::Verify(func, assert_mode)) {
- return false;
- }
- if (!LayoutVerifier::Verify(func, assert_mode)) {
- return false;
- }
- if (!AsyncStructsVerifier::Verify(func, assert_mode)) {
- return false;
- }
- if (device_func) {
- if (!DeviceFuncVerifier::Verify(func, assert_mode)) {
- return false;
- }
- }
- return true;
-}
-
-bool VerifyTIRxWellFormed(const IRModule& mod, bool assert_mode, bool
device_func) {
- for (const auto& [gvar, base_func] : mod->functions) {
- if (auto function = base_func.as<Function>()) {
- bool res = VerifyTIRxWellFormed(function.value(), assert_mode,
device_func);
- if (!res) {
- return false;
- }
- }
- }
- return true;
-}
-
-TVM_FFI_STATIC_INIT_BLOCK() {
- namespace refl = tvm::ffi::reflection;
- refl::GlobalDef().def("tirx.analysis.VerifyTIRxWellFormed",
- [](const ffi::ObjectRef& obj, bool assert_mode, bool
device_func) {
- if (auto n = obj.as<Function>()) {
- return VerifyTIRxWellFormed(n.value(),
assert_mode, device_func);
- } else if (auto n = obj.as<IRModule>()) {
- return VerifyTIRxWellFormed(n.value(),
assert_mode, device_func);
- } else {
- LOG(FATAL) << "Expects Function or IRModule, but
get "
- << obj->GetTypeKey() << " instead.";
- return false;
- }
- });
-}
-} // namespace tirx
-} // namespace tvm
diff --git a/src/tirx/analysis/verify_well_formed.cc
b/src/tirx/analysis/verify_well_formed.cc
index f64c513202..2ce62ba761 100644
--- a/src/tirx/analysis/verify_well_formed.cc
+++ b/src/tirx/analysis/verify_well_formed.cc
@@ -19,7 +19,7 @@
/*!
* \file tirx/analysis/verify_well_formed.cc
- * \brief Check if schedulable tirx is well-formed.
+ * \brief Check if TIRx is well-formed.
*/
#include "verify_well_formed.h"
diff --git a/src/tirx/analysis/verify_well_formed.h
b/src/tirx/analysis/verify_well_formed.h
index 83511377d8..1547352c72 100644
--- a/src/tirx/analysis/verify_well_formed.h
+++ b/src/tirx/analysis/verify_well_formed.h
@@ -22,6 +22,10 @@
#include <tvm/ir/prim/op.h>
#include <tvm/tirx/analysis.h>
+#include <tvm/tirx/exec_scope.h>
+#include <tvm/tirx/op/region.h>
+
+#include <utility>
#include "../ir/tir_visitor_with_path.h"
namespace tvm {
@@ -301,12 +305,101 @@ class LoopControlVerifier : public
Verifier<LoopControlVerifier<PathVisitor>, Pa
int loop_depth_{0};
};
+template <typename PathVisitor>
+class ExecScopeVerifier : public Verifier<ExecScopeVerifier<PathVisitor>,
PathVisitor> {
+ using Verifier = tirx::Verifier<ExecScopeVerifier<PathVisitor>, PathVisitor>;
+
+ public:
+ using Verifier::Verifier;
+ using Verifier::Verify;
+
+ private:
+ using Verifier::Visit;
+
+ void Dispatch_(const tirx::TileOpCallNode* op, ffi::reflection::AccessPath
path) override {
+ static const auto& category_map =
Op::GetAttrMap<tirx::TIRxOpCategory>("TIRxOpCategory");
+ Verify(category_map.get(op->op, ffi::String("")) == "tile_primitive")
+ << "TIRxError: TileOpCall at " << path << " has non-tile op " <<
op->op;
+ }
+};
+
+template <typename PathVisitor>
+class ScopeIdVerifier : public Verifier<ScopeIdVerifier<PathVisitor>,
PathVisitor> {
+ using Verifier = tirx::Verifier<ScopeIdVerifier<PathVisitor>, PathVisitor>;
+
+ public:
+ using Verifier::Verifier;
+ using Verifier::Verify;
+
+ private:
+ using Verifier::Visit;
+
+ void Visit(const Function& function, AccessPath path) override {
+ Array<ScopeIdDef> enclosing_defs = std::exchange(scope_id_def_, {});
+ Verifier::Visit(function, path);
+ scope_id_def_ = std::move(enclosing_defs);
+ }
+
+ void Dispatch_(const RegionStmtNode* op, ffi::reflection::AccessPath path)
override {
+ if (op->op.same_as(tirx::device_entry_op())) {
+ // Device-region marker: defs gathered from the body are verified when
+ // the region exits, with launch-param sanity enforced as ``is_root``.
+ size_t baseline = scope_id_def_.size();
+ Verifier::Dispatch_(op, path);
+ size_t total = scope_id_def_.size();
+ if (total > baseline) {
+ RunScopeIdVerify(path, baseline, /*is_root=*/true);
+ }
+ while (scope_id_def_.size() > baseline) {
+ scope_id_def_.pop_back();
+ }
+ return;
+ }
+ Verifier::Dispatch_(op, path);
+ }
+
+ void RunScopeIdVerify(ffi::reflection::AccessPath path, size_t baseline,
bool is_root) {
+ ScopeIdDefVerifier verifier;
+ Verify(verifier.Verify(scope_id_def_, ScopeIdDefVerifier::Mode::kRelaxed))
+ << "TIRxError: Scope at " << path << " has invalid scope_id_def";
+ if (is_root) {
+ // Enforce launch-parameter sanity at the device-region root.
+ auto cta_thread_it = verifier.id_set.find(ScopeBinding::kCtaThread);
+ if (cta_thread_it != verifier.id_set.end() &&
!(*cta_thread_it).second.is_deferred()) {
+ PrimExpr ext = (*cta_thread_it).second.fused_extent();
+ if (const auto* imm = ext.as<IntImmNode>()) {
+ Verify(imm->value > 0) << "TIRxError: kernel at " << path
+ << " has non-positive thread count " <<
imm->value;
+ bool needs_warp_align =
verifier.id_set.count(ScopeBinding::kCtaWarp) ||
+
verifier.id_set.count(ScopeBinding::kWarpThread) ||
+
verifier.id_set.count(ScopeBinding::kCtaWarpgroup) ||
+
verifier.id_set.count(ScopeBinding::kWarpgroupWarp);
+ if (needs_warp_align) {
+ Verify(imm->value % 32 == 0)
+ << "TIRxError: kernel at " << path << " uses warp-granular
bindings"
+ << " but has thread count " << imm->value << " not a multiple
of 32";
+ }
+ }
+ }
+ }
+ }
+
+ void Dispatch_(const ScopeIdDefStmtNode* op, ffi::reflection::AccessPath
path) override {
+ scope_id_def_.push_back(op->def);
+ Verifier::Dispatch_(op, path);
+ }
+
+ Array<ScopeIdDef> scope_id_def_;
+};
+
template <typename PathVisitor, typename NodeRef>
bool VerifyWellFormedCommon(const NodeRef& node, bool assert_mode) {
return UndefinedVarVerifier<PathVisitor>::Verify(node, assert_mode) &&
UndefinedBufferVerifier<PathVisitor>::Verify(node, assert_mode) &&
TensorLoadTypeVerifier<PathVisitor>::Verify(node, assert_mode) &&
- LoopControlVerifier<PathVisitor>::Verify(node, assert_mode);
+ LoopControlVerifier<PathVisitor>::Verify(node, assert_mode) &&
+ ExecScopeVerifier<PathVisitor>::Verify(node, assert_mode) &&
+ ScopeIdVerifier<PathVisitor>::Verify(node, assert_mode);
}
} // namespace tirx
} // namespace tvm
diff --git a/src/s_tir/transform/renormalize_split_pattern.cc
b/src/tirx/transform/renormalize_split_pattern.cc
similarity index 94%
rename from src/s_tir/transform/renormalize_split_pattern.cc
rename to src/tirx/transform/renormalize_split_pattern.cc
index b275741906..3d99c314b0 100644
--- a/src/s_tir/transform/renormalize_split_pattern.cc
+++ b/src/tirx/transform/renormalize_split_pattern.cc
@@ -24,18 +24,14 @@
#include <tvm/ffi/function.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ir/prim/op.h>
-#include <tvm/s_tir/analysis.h>
-#include <tvm/s_tir/stmt_functor.h>
-#include <tvm/s_tir/transform.h>
-#include <tvm/tirx/analysis.h>
#include <tvm/tirx/stmt.h>
+#include <tvm/tirx/transform.h>
-#include "../../s_tir/ir/ir_mutator_with_analyzer.h"
#include "../../sym/pattern_match.h"
+#include "../ir/ir_mutator_with_analyzer.h"
namespace tvm {
-namespace s_tir {
-using namespace tvm::tirx;
+namespace tirx {
using namespace sym;
@@ -170,7 +166,8 @@ class SplitPatternReNormalizer : public
IRMutatorWithAnalyzer {
// Pattern var match IntImm
PVar<IntImm> c1, c2;
// x < c2 <=> x/c2 < 1 <=> floor(x / c2) < 1
- TRY_RECURSIVE_REWRITE_IF(x < c2, floordiv(x, c2) < 1, c2.Eval()->value >
0); // NOLINT
+ TRY_RECURSIVE_REWRITE_IF(x < c2, floordiv(x, c2) < 1,
+ c2.Eval()->value > 0 && !x.Eval().as<VarNode>());
// NOLINT
return ret;
}
@@ -224,15 +221,15 @@ Pass RenormalizeSplitPattern() {
.ValueOrUnchanged(std::move(n->body));
return f;
};
- return CreateFunctionPass(pass_func, 0, "s_tir.RenormalizeSplitPattern", {});
+ return CreateFunctionPass(pass_func, 0, "tirx.RenormalizeSplitPattern", {});
}
TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
- refl::GlobalDef().def("s_tir.transform.RenormalizeSplitPattern",
RenormalizeSplitPattern);
+ refl::GlobalDef().def("tirx.transform.RenormalizeSplitPattern",
RenormalizeSplitPattern);
}
} // namespace transform
-} // namespace s_tir
+} // namespace tirx
} // namespace tvm
diff --git a/tests/python/relax/test_relax_operators.py
b/tests/python/relax/test_relax_operators.py
index 0008ec1f64..2570b3cc6d 100644
--- a/tests/python/relax/test_relax_operators.py
+++ b/tests/python/relax/test_relax_operators.py
@@ -424,7 +424,7 @@ def test_call_tir_canonical_roundtrip(name, tuple_result,
result_kind, namespace
tvm.tirx.Var(f"p{i}", T.Tensor((3,), "float32"))
for i in range(1 + output_count - int(inplace))
]
- func = tvm.tirx.Function(params, [tvm.tirx.Evaluate(0)],
ret_type=tvm.ir.TupleType([]))
+ func = tvm.tirx.Function(params, [tvm.ir.Evaluate(0)],
ret_type=tvm.ir.TupleType([]))
callee = relax.BlockBuilder().add_func(func, "callee")
x = relax.Var("x", tensor_type)
kwargs = {}
diff --git
a/tests/python/s_tir/transform/test_s_tir_transform_renormalize_split_pattern.py
b/tests/python/s_tir/transform/test_s_tir_transform_renormalize_split_pattern.py
deleted file mode 100644
index 50d2e453fa..0000000000
---
a/tests/python/s_tir/transform/test_s_tir_transform_renormalize_split_pattern.py
+++ /dev/null
@@ -1,176 +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 tvm
-import tvm.testing
-from tvm import s_tir
-from tvm.script import s_tir as Ts
-from tvm.script import tirx as T
-
-# fmt: off
-# pylint:
disable=no-member,invalid-name,unused-variable,line-too-long,redefined-outer-name,redundant-keyword-arg
-
[email protected]_module
-class Before:
- @Ts.function
- def main(inputs: T.Tensor((1, 4, 4, 512), "float32"), weight: T.Tensor((4,
4, 512, 256), "float32"), conv2d_transpose_nhwc: T.Tensor((1, 8, 8, 256),
"float32")) -> None:
- # function attr dict
- T.func_attr({"global_symbol": "main", "tirx.noalias": True})
- inputs_flat = T.decl_tensor([8192], dtype="float32", data=inputs.data)
- weight_flat = T.decl_tensor([2097152], dtype="float32",
data=weight.data)
- conv2d_transpose_nhwc_flat = T.decl_tensor([16384], dtype="float32",
data=conv2d_transpose_nhwc.data)
- # var definition
- # body
- blockIdx_x = T.launch_thread("blockIdx.x", 64)
- conv2d_transpose_nhwc_local = T.decl_tensor([8], "float32",
scope="local")
- PadInput_shared = T.decl_tensor([768], "float32", scope="shared")
- weight_shared = T.decl_tensor([4096], "float32", scope="shared")
- threadIdx_x = T.launch_thread("threadIdx.x", 32)
- for i2_3_init, i1_4_init, i2_4_init in T.grid(2, 2, 2):
- conv2d_transpose_nhwc_local[i1_4_init * 4 + i2_3_init * 2 +
i2_4_init] = T.float32(0)
- for i6_0 in T.serial(16):
- for ax0_ax1_ax2_ax3_fused_0 in T.serial(24):
- PadInput_shared[ax0_ax1_ax2_ax3_fused_0 * 32 + threadIdx_x] =
T.if_then_else(128 <= ax0_ax1_ax2_ax3_fused_0 * 32 + threadIdx_x and
ax0_ax1_ax2_ax3_fused_0 * 32 + threadIdx_x < 640 and 1 <= blockIdx_x // 32 * 2
+ (ax0_ax1_ax2_ax3_fused_0 * 32 + threadIdx_x) % 128 // 32 and blockIdx_x // 32
* 2 + (ax0_ax1_ax2_ax3_fused_0 * 32 + threadIdx_x) % 128 // 32 < 5,
inputs_flat[blockIdx_x // 32 * 1024 + ax0_ax1_ax2_ax3_fused_0 * 512 + i6_0 * 32
+ threadIdx_x - 2560], T.float32(0))
- for ax0_ax1_ax2_ax3_fused_0 in T.serial(32):
- weight_shared[T.ramp(ax0_ax1_ax2_ax3_fused_0 * 128 +
threadIdx_x * 4, 1, 4)] = weight_flat[T.ramp((ax0_ax1_ax2_ax3_fused_0 * 128 +
threadIdx_x * 4) // 256 * 131072 + i6_0 * 8192 + (ax0_ax1_ax2_ax3_fused_0 * 128
+ threadIdx_x * 4) % 256 // 8 * 256 + blockIdx_x % 32 * 8 + threadIdx_x % 2 *
4, 1, 4)]
- for i6_1, i2_3, i4_2, i5_2, i6_2, i1_4, i2_4 in T.grid(4, 2, 4, 4,
8, 2, 2):
- conv2d_transpose_nhwc_local[i1_4 * 4 + i2_3 * 2 + i2_4] =
conv2d_transpose_nhwc_local[i1_4 * 4 + i2_3 * 2 + i2_4] + T.if_then_else((i1_4
+ i4_2) % 2 == 0 and (i2_4 + i5_2) % 2 == 0, PadInput_shared[threadIdx_x // 8 *
128 + (i1_4 + i4_2) // 2 * 128 + (i2_4 + i5_2) // 2 * 32 + i2_3 * 32 + i6_1 * 8
+ i6_2], T.float32(0)) * weight_shared[i6_1 * 64 + i6_2 * 8 + threadIdx_x % 8 +
3840 - i5_2 * 256 - i4_2 * 1024]
- for ax1, ax2 in T.grid(2, 4):
- conv2d_transpose_nhwc_flat[threadIdx_x // 8 * 4096 + ax1 * 2048 +
blockIdx_x // 32 * 1024 + ax2 * 256 + blockIdx_x % 32 * 8 + threadIdx_x % 8] =
conv2d_transpose_nhwc_local[ax1 * 4 + ax2]
-
-
[email protected]_module
-class After:
- @Ts.function
- def main(inputs: T.Tensor((1, 4, 4, 512), "float32"), weight: T.Tensor((4,
4, 512, 256), "float32"), conv2d_transpose_nhwc: T.Tensor((1, 8, 8, 256),
"float32")) -> None:
- # function attr dict
- T.func_attr({"global_symbol": "main", "tirx.noalias": True})
- inputs_flat = T.decl_tensor([8192], dtype="float32", data=inputs.data)
- weight_flat = T.decl_tensor([2097152], dtype="float32",
data=weight.data)
- conv2d_transpose_nhwc_flat = T.decl_tensor([16384], dtype="float32",
data=conv2d_transpose_nhwc.data)
- # var definition
- # body
- blockIdx_x = T.launch_thread("blockIdx.x", 64)
- conv2d_transpose_nhwc_local = T.decl_tensor([8], "float32",
scope="local")
- PadInput_shared = T.decl_tensor([768], "float32", scope="shared")
- weight_shared = T.decl_tensor([4096], "float32", scope="shared")
- threadIdx_x = T.launch_thread("threadIdx.x", 32)
- for i2_3_init, i1_4_init, i2_4_init in T.grid(2, 2, 2):
- conv2d_transpose_nhwc_local[i1_4_init * 4 + i2_3_init * 2 +
i2_4_init] = T.float32(0)
- for i6_0 in T.serial(16):
- for ax0_ax1_ax2_ax3_fused_0 in T.serial(24):
- PadInput_shared[ax0_ax1_ax2_ax3_fused_0 * 32 + threadIdx_x] =
T.if_then_else(1 <= (ax0_ax1_ax2_ax3_fused_0 + threadIdx_x // 32) // 4 and
(ax0_ax1_ax2_ax3_fused_0 + threadIdx_x // 32) // 20 < 1 and 1 <= blockIdx_x //
32 * 2 + (ax0_ax1_ax2_ax3_fused_0 + threadIdx_x // 32) % 4 and (blockIdx_x //
32 * 2 + (ax0_ax1_ax2_ax3_fused_0 + threadIdx_x // 32) % 4) // 5 < 1,
inputs_flat[blockIdx_x // 32 * 1024 + ax0_ax1_ax2_ax3_fused_0 * 512 + i6_0 * 32
+ threadIdx_x - 2560], T.float32(0))
- for ax0_ax1_ax2_ax3_fused_0 in T.serial(32):
- weight_shared[T.ramp(ax0_ax1_ax2_ax3_fused_0 * 128 +
threadIdx_x * 4, 1, 4)] = weight_flat[T.ramp((ax0_ax1_ax2_ax3_fused_0 +
threadIdx_x * 4 // 128) // 2 * 131072 + i6_0 * 8192 + (ax0_ax1_ax2_ax3_fused_0
* 16 + threadIdx_x * 4 // 8) % 32 * 256 + blockIdx_x % 32 * 8 + threadIdx_x % 2
* 4, 1, 4)]
- for i6_1, i2_3, i4_2, i5_2, i6_2, i1_4, i2_4 in T.grid(4, 2, 4, 4,
8, 2, 2):
- conv2d_transpose_nhwc_local[i1_4 * 4 + i2_3 * 2 + i2_4] =
conv2d_transpose_nhwc_local[i1_4 * 4 + i2_3 * 2 + i2_4] + T.if_then_else((i1_4
+ i4_2) % 2 == 0 and (i2_4 + i5_2) % 2 == 0, PadInput_shared[threadIdx_x // 8 *
128 + (i1_4 + i4_2) // 2 * 128 + (i2_4 + i5_2) // 2 * 32 + i2_3 * 32 + i6_1 * 8
+ i6_2], T.float32(0)) * weight_shared[i6_1 * 64 + i6_2 * 8 + threadIdx_x % 8 +
3840 - i5_2 * 256 - i4_2 * 1024]
- for ax1, ax2 in T.grid(2, 4):
- conv2d_transpose_nhwc_flat[threadIdx_x // 8 * 4096 + ax1 * 2048 +
blockIdx_x // 32 * 1024 + ax2 * 256 + blockIdx_x % 32 * 8 + threadIdx_x % 8] =
conv2d_transpose_nhwc_local[ax1 * 4 + ax2]
-
-
[email protected]_module
-class After_simplified:
- @Ts.function
- def main(inputs: T.Tensor((1, 4, 4, 512), "float32"), weight: T.Tensor((4,
4, 512, 256), "float32"), conv2d_transpose_nhwc: T.Tensor((1, 8, 8, 256),
"float32")) -> None:
- # function attr dict
- T.func_attr({"global_symbol": "main", "tirx.noalias": True})
- # var definition
- inputs_flat = T.decl_tensor([8192], dtype="float32", data=inputs.data)
- weight_flat = T.decl_tensor([2097152], dtype="float32",
data=weight.data)
- conv2d_transpose_nhwc_flat = T.decl_tensor([16384], dtype="float32",
data=conv2d_transpose_nhwc.data)
- # body
- blockIdx_x = T.launch_thread("blockIdx.x", 64)
- conv2d_transpose_nhwc_local = T.decl_tensor([8], "float32",
scope="local")
- PadInput_shared = T.decl_tensor([768], "float32", scope="shared")
- weight_shared = T.decl_tensor([4096], "float32", scope="shared")
- threadIdx_x = T.launch_thread("threadIdx.x", 32)
- for i2_3_init, i1_4_init, i2_4_init in T.grid(2, 2, 2):
- conv2d_transpose_nhwc_local[i1_4_init * 4 + i2_3_init * 2 +
i2_4_init] = T.float32(0)
- for i6_0 in T.serial(16):
- for ax0_ax1_ax2_ax3_fused_0 in T.serial(24):
- PadInput_shared[ax0_ax1_ax2_ax3_fused_0 * 32 + threadIdx_x] =
T.if_then_else(4 <= ax0_ax1_ax2_ax3_fused_0 and ax0_ax1_ax2_ax3_fused_0 < 20
and 1 <= blockIdx_x // 32 * 2 + ax0_ax1_ax2_ax3_fused_0 % 4 and blockIdx_x //
32 * 2 + ax0_ax1_ax2_ax3_fused_0 % 4 < 5, inputs_flat[blockIdx_x // 32 * 1024 +
ax0_ax1_ax2_ax3_fused_0 * 512 + i6_0 * 32 + threadIdx_x - 2560], T.float32(0))
- for ax0_ax1_ax2_ax3_fused_0 in T.serial(32):
- weight_shared[T.ramp(ax0_ax1_ax2_ax3_fused_0 * 128 +
threadIdx_x * 4, 1, 4)] = weight_flat[T.ramp(ax0_ax1_ax2_ax3_fused_0 // 2 *
131072 + i6_0 * 8192 + ax0_ax1_ax2_ax3_fused_0 % 2 * 4096 + threadIdx_x // 2 *
256 + blockIdx_x % 32 * 8 + threadIdx_x % 2 * 4, 1, 4)]
- for i6_1, i2_3, i4_2, i5_2, i6_2, i1_4, i2_4 in T.grid(4, 2, 4, 4,
8, 2, 2):
- conv2d_transpose_nhwc_local[i1_4 * 4 + i2_3 * 2 + i2_4] =
conv2d_transpose_nhwc_local[i1_4 * 4 + i2_3 * 2 + i2_4] + T.if_then_else((i1_4
+ i4_2) % 2 == 0 and (i2_4 + i5_2) % 2 == 0, PadInput_shared[threadIdx_x // 8 *
128 + (i1_4 + i4_2) // 2 * 128 + (i2_4 + i5_2) // 2 * 32 + i2_3 * 32 + i6_1 * 8
+ i6_2], T.float32(0)) * weight_shared[i6_1 * 64 + i6_2 * 8 + threadIdx_x % 8 +
3840 - i5_2 * 256 - i4_2 * 1024]
- for ax1, ax2 in T.grid(2, 4):
- conv2d_transpose_nhwc_flat[threadIdx_x // 8 * 4096 + ax1 * 2048 +
blockIdx_x // 32 * 1024 + ax2 * 256 + blockIdx_x % 32 * 8 + threadIdx_x % 8] =
conv2d_transpose_nhwc_local[ax1 * 4 + ax2]
-
-# pylint:
enable=no-member,invalid-name,unused-variable,line-too-long,redefined-outer-name,redundant-keyword-arg
-# fmt: on
-
-
-def test_renormalize_split_pattern():
- after = tvm.s_tir.transform.RenormalizeSplitPattern()(Before)
- tvm.ir.assert_structural_equal(after, After)
- after = tvm.tirx.transform.StmtSimplify()(after)
- tvm.ir.assert_structural_equal(after, After_simplified)
-
-
[email protected]
-def impossible_equality(n: T.int32):
- # Prior to bugfix, this conditional defined the expression "2" as
- # equal to zero within the then_case. [min_value=2, max_value=0]
- if 2 == 0:
- # Then this expression evaluates n/2, using the min/max values
- # of "2", which is caught as a divide by zero error.
- if n // 2 >= 16:
- T.evaluate(0)
-
-
[email protected]
-def impossible_inequality(n: T.int32):
- # Prior to bugfix, this conditional set up a range of possible
- # values for the expression "-2" as [0, kPosInf].
- if -1 < -2:
- if n // (-2) >= 16:
- T.evaluate(0)
-
-
-integer_condition = tvm.testing.parameter(
- impossible_equality,
- impossible_inequality,
-)
-
-
-def test_analyze_inside_integer_conditional(integer_condition):
- """Avoid crash occurring in ConstIntBoundAnalyzer.
-
- Crash occurred when simplifying some expressions with provably
- false integer expressions. If the expressions were renormalized
- before calling Simplify, conditional statements could assign a
- range of possible values to integers, as if they were variables.
- This would result in divide by zero throwing an exception,
- followed by a second exception during stack unwinding causing the
- program to crash.
- """
-
- # Similar issue would occur in most transformations that subclass
- # IRMutatorWithAnalyzer. tirx.transform.StmtSimplify() is an
- # exception, as it rewrites the integer conditionals first. These
- # tests are written using RenormalizeSplitPattern as it is the
- # first case identified.
- transform = tvm.s_tir.transform.RenormalizeSplitPattern()
-
- # Issue would result in an error through while applying the transformation.
- mod = tvm.IRModule.from_expr(integer_condition)
- transform(mod)
-
-
-if __name__ == "__main__":
- tvm.testing.main()
diff --git a/tests/python/tirx-analysis/test_tir_analysis_verify_ssa.py
b/tests/python/tirx-analysis/test_tir_analysis_verify_ssa.py
deleted file mode 100644
index 6b883a0c17..0000000000
--- a/tests/python/tirx-analysis/test_tir_analysis_verify_ssa.py
+++ /dev/null
@@ -1,42 +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
-
-
-def test_verify_ssa():
- x = tvm.tirx.Var("x", "int32")
- y = tvm.tirx.Var("tindex", "int32")
- z = tvm.ir.Evaluate(x + y)
- assert tvm.tirx.analysis.verify_ssa(tvm.tirx.Function([x, y], z))
-
- assert not tvm.tirx.analysis.verify_ssa(
- tvm.tirx.Function([x, y], tvm.ir.SeqStmt([tvm.ir.Bind(x, 1), z]))
- )
-
-
-def test_verify_weak_let_ssa():
- x = tvm.tirx.Var("x", "int32")
- z1 = tvm.tirx.Let(x, 1, x + 1)
- z2 = tvm.tirx.Let(x, 2, x + 2)
-
- assert tvm.tirx.analysis.verify_ssa(tvm.tirx.Function([],
tvm.ir.Evaluate(z1 + z1)))
- assert not tvm.tirx.analysis.verify_ssa(tvm.tirx.Function([],
tvm.ir.Evaluate(z1 * z2)))
-
-
-if __name__ == "__main__":
- test_verify_ssa()
- test_verify_weak_let_ssa()
diff --git a/tests/python/tirx-analysis/test_tir_analysis_verify_well_formed.py
b/tests/python/tirx-analysis/test_tir_analysis_verify_well_formed.py
index d0d4880e19..4b8ff4a771 100644
--- a/tests/python/tirx-analysis/test_tir_analysis_verify_well_formed.py
+++ b/tests/python/tirx-analysis/test_tir_analysis_verify_well_formed.py
@@ -376,5 +376,31 @@ def
test_tensor_load_malformed_indices_return_false_without_asserting():
tvm.tirx.analysis.verify_well_formed(non_final_vector)
+def test_verify_unique_definitions():
+ x = tvm.tirx.Var("x", "int32")
+ y = tvm.tirx.Var("tindex", "int32")
+ z = tvm.ir.Evaluate(x + y)
+ assert tvm.tirx.analysis.verify_well_formed(tvm.tirx.Function([x, y], z))
+
+ assert not tvm.tirx.analysis.verify_well_formed(
+ tvm.tirx.Function([x, y], tvm.ir.SeqStmt([tvm.ir.Bind(x, 1), z])),
+ assert_mode=False,
+ )
+
+
+def test_verify_repeated_let_definitions():
+ x = tvm.tirx.Var("x", "int32")
+ z1 = tvm.tirx.Let(x, 1, x + 1)
+ z2 = tvm.tirx.Let(x, 2, x + 2)
+
+ # Each occurrence must define a distinct variable, even for a shared Let.
+ assert not tvm.tirx.analysis.verify_well_formed(
+ tvm.tirx.Function([], tvm.ir.Evaluate(z1 + z1)), assert_mode=False
+ )
+ assert not tvm.tirx.analysis.verify_well_formed(
+ tvm.tirx.Function([], tvm.ir.Evaluate(z1 * z2)), assert_mode=False
+ )
+
+
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
index 5ac77f933e..133b80529c 100644
--- a/tests/python/tirx-transform/test_tir_transform_common_subexpr_elim.py
+++ b/tests/python/tirx-transform/test_tir_transform_common_subexpr_elim.py
@@ -775,7 +775,7 @@ def test_shared_subtree_stays_ssa():
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_ssa(after), after.script()
+ assert tvm.tirx.analysis.verify_well_formed(after), after.script()
if __name__ == "__main__":
diff --git
a/tests/python/tirx-transform/test_tir_transform_renormalize_split_pattern.py
b/tests/python/tirx-transform/test_tir_transform_renormalize_split_pattern.py
new file mode 100644
index 0000000000..ab42b7f5ed
--- /dev/null
+++
b/tests/python/tirx-transform/test_tir_transform_renormalize_split_pattern.py
@@ -0,0 +1,66 @@
+# 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
+
+
+def test_renormalize_split_pattern():
+ @tvm.script.ir_module
+ class Before:
+ @T.function
+ def main(A: T.Tensor((768,), "int32"), B: T.Tensor((96,), "int32")):
+ for i in T.serial(24):
+ for j in T.serial(32):
+ A[i * 32 + j] = T.if_then_else(
+ 128 <= i * 32 + j and i * 32 + j < 640,
+ (i * 32 + j) % 128 // 32,
+ 0,
+ )
+ B[T.ramp(i * 4, 1, 4)] = T.ramp(i * 128, 1, 4) % 256 // 8
+
+ @tvm.script.ir_module
+ class After:
+ @T.function
+ def main(A: T.Tensor((768,), "int32"), B: T.Tensor((96,), "int32")):
+ for i in T.serial(24):
+ for j in T.serial(32):
+ A[i * 32 + j] = T.if_then_else(
+ 1 <= (i + j // 32) // 4 and (i + j // 32) // 20 < 1,
+ (i + j // 32) % 4,
+ 0,
+ )
+ B[T.ramp(i * 4, 1, 4)] = T.ramp(i * 128, 1, 4) // 8 % 32
+
+ @tvm.script.ir_module
+ class AfterSimplified:
+ @T.function
+ def main(A: T.Tensor((768,), "int32"), B: T.Tensor((96,), "int32")):
+ for i in T.serial(24):
+ for j in T.serial(32):
+ A[i * 32 + j] = T.if_then_else(4 <= i and i < 20, i % 4, 0)
+ B[T.ramp(i * 4, 1, 4)] = T.broadcast(i % 2 * 16, 4)
+
+ after = tvm.tirx.transform.RenormalizeSplitPattern()(Before)
+ tvm.ir.assert_structural_equal(after, After)
+ after = tvm.tirx.transform.StmtSimplify()(after)
+ tvm.ir.assert_structural_equal(after, AfterSimplified)
+
+
+if __name__ == "__main__":
+ tvm.testing.main()
diff --git a/tests/python/tirx/script/test_tirx_script_source_locations.py
b/tests/python/tirx/script/test_tirx_script_source_locations.py
index 8a7229bd4b..43b6a629ad 100644
--- a/tests/python/tirx/script/test_tirx_script_source_locations.py
+++ b/tests/python/tirx/script/test_tirx_script_source_locations.py
@@ -243,7 +243,7 @@ def
test_native_view_keeps_producer_identity_name_and_span(monkeypatch):
# A native view must preserve its producer name/span and declare its
storage exactly once.
from functools import wraps
- from tvm import ir
+ from tvm import ir, tirx
from tvm.script.ir_builder import base
original = tirx.TensorType.view
diff --git a/tests/python/tirx/test_verifier.py
b/tests/python/tirx/test_verifier.py
index 5bddb07737..96fc79de27 100644
--- a/tests/python/tirx/test_verifier.py
+++ b/tests/python/tirx/test_verifier.py
@@ -18,7 +18,7 @@ import pytest
from tvm.script import tirx as T
from tvm.script.tirx import tile as Tx
-from tvm.tirx.analysis import verify_tirx_well_formed as verify
+from tvm.tirx.analysis import verify_well_formed as verify
def test_root_scope():
@@ -244,8 +244,8 @@ def test_device_func():
T.thread_id([128])
Tx.fill(A, 0.)
# fmt: on
- verify(test1, device_func=True)
- verify(test2, device_func=True)
+ verify(test1)
+ verify(test2)
def test_preferred_cluster_validation():
