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 f1899b6a99 [REFACTOR] Expose analyzer-aware statement functors (#20681)
f1899b6a99 is described below
commit f1899b6a99a963b3a8a5e422e6334536ca9812c9
Author: Tianqi Chen <[email protected]>
AuthorDate: Mon Oct 12 18:02:22 2026 -0400
[REFACTOR] Expose analyzer-aware statement functors (#20681)
---
include/tvm/s_tir/stmt_functor_with_analyzer.h | 86 ++++++++++++++++++++++
.../tvm/tirx/stmt_functor_with_analyzer.h | 62 +++++++++++++---
src/backend/trn/transform/lower_trainium_layout.cc | 10 +--
src/s_tir/analysis/domain_touched.cc | 9 +--
src/s_tir/analysis/identify_memcpy.cc | 9 +--
src/s_tir/analysis/oob_checker.cc | 10 +--
src/s_tir/backend/adreno/inject_texture_alloc.cc | 10 +--
src/s_tir/backend/adreno/texture_flatten.cc | 12 +--
src/s_tir/ir/ir_mutator_with_analyzer.cc | 5 +-
src/s_tir/ir/ir_mutator_with_analyzer.h | 60 ---------------
src/s_tir/ir/ir_visitor_with_analyzer.cc | 4 +-
src/s_tir/ir/ir_visitor_with_analyzer.h | 49 ------------
.../schedule/primitive/layout_transformation.cc | 12 +--
src/s_tir/schedule/transform.cc | 6 +-
src/s_tir/schedule/transform.h | 11 ++-
src/s_tir/transform/inject_permuted_layout.cc | 20 ++---
src/s_tir/transform/inject_virtual_thread.cc | 18 ++---
.../transform/using_assume_to_reduce_branches.cc | 10 +--
src/tirx/ir/ir_mutator_with_analyzer.cc | 39 +++++-----
src/tirx/ir/ir_visitor_with_analyzer.cc | 22 +++---
src/tirx/ir/ir_visitor_with_analyzer.h | 76 -------------------
src/tirx/transform/flatten_tensor.cc | 14 ++--
src/tirx/transform/lower_intrin.cc | 26 +++----
src/tirx/transform/lower_tirx_cleanup.cc | 14 ++--
src/tirx/transform/remove_no_op.cc | 10 +--
src/tirx/transform/renormalize_split_pattern.cc | 15 ++--
src/tirx/transform/stmt_simplify.cc | 3 +-
src/tirx/transform/stmt_simplify.h | 13 ++--
28 files changed, 291 insertions(+), 344 deletions(-)
diff --git a/include/tvm/s_tir/stmt_functor_with_analyzer.h
b/include/tvm/s_tir/stmt_functor_with_analyzer.h
new file mode 100644
index 0000000000..aee173e357
--- /dev/null
+++ b/include/tvm/s_tir/stmt_functor_with_analyzer.h
@@ -0,0 +1,86 @@
+/*
+ * 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 tvm/s_tir/stmt_functor_with_analyzer.h
+ * \brief Statement and expression functors with an analyzer context.
+ */
+
+#ifndef TVM_S_TIR_STMT_FUNCTOR_WITH_ANALYZER_H_
+#define TVM_S_TIR_STMT_FUNCTOR_WITH_ANALYZER_H_
+
+#include <tvm/s_tir/stmt_functor.h>
+#include <tvm/tirx/stmt_functor_with_analyzer.h>
+
+namespace tvm {
+namespace s_tir {
+
+class TVM_DLL StmtExprVisitorWithAnalyzer : public
tirx::StmtExprVisitorWithAnalyzer {
+ public:
+ using Parent = tirx::StmtExprVisitorWithAnalyzer;
+ using Parent::Visit;
+ using Parent::Visit_;
+ TVM_DEFINE_OBJECT_FUNCTOR_DEFAULT_CONSTRUCTOR(StmtExprVisitorWithAnalyzer,
Parent)
+ virtual ffi::Optional<VisitInterrupt> Visit_(const SBlockNode* op);
+ virtual ffi::Optional<VisitInterrupt> Visit_(const SBlockRealizeNode* op) {
+ return s_tir::StmtExprVisitor::VisitBlockRealize(this, op);
+ }
+
+ protected:
+ static void InitVTable(VTable* vtable) {
+ Parent::InitVTable(vtable);
+ SetDispatch<StmtExprVisitorWithAnalyzer, SBlockNode>(vtable);
+ SetDispatch<StmtExprVisitorWithAnalyzer, SBlockRealizeNode>(vtable);
+ }
+};
+
+class TVM_DLL StmtExprMutatorWithAnalyzer : public
tirx::StmtExprMutatorWithAnalyzer {
+ public:
+ using Parent = tirx::StmtExprMutatorWithAnalyzer;
+ using Parent::Mutate;
+ using Parent::Mutate_;
+ explicit StmtExprMutatorWithAnalyzer(const sym::Analyzer& analyzer)
+ : StmtExprMutatorWithAnalyzer(analyzer.get()) {}
+ explicit StmtExprMutatorWithAnalyzer(sym::AnalyzerObj* analyzer)
+ : Parent(analyzer, GlobalVTable()) {}
+ virtual UnchangedOr<Stmt> Mutate_(const SBlockNode* op, InplaceMode
inplace_mode);
+ virtual UnchangedOr<Stmt> Mutate_(const SBlockRealizeNode* op, InplaceMode
inplace_mode) {
+ return s_tir::StmtExprMutator::MutateBlockRealize(this, op, inplace_mode);
+ }
+
+ protected:
+ static void InitVTable(VTable* vtable) {
+ Parent::InitVTable(vtable);
+ SetDispatch<StmtExprMutatorWithAnalyzer, SBlockNode>(vtable);
+ SetDispatch<StmtExprMutatorWithAnalyzer, SBlockRealizeNode>(vtable);
+ }
+ static const VTable* GlobalVTable() {
+ static const VTable table = [] {
+ VTable table;
+ InitVTable(&table);
+ table.Finalize();
+ return table;
+ }();
+ return &table;
+ }
+};
+
+} // namespace s_tir
+} // namespace tvm
+#endif // TVM_S_TIR_STMT_FUNCTOR_WITH_ANALYZER_H_
diff --git a/src/tirx/ir/ir_mutator_with_analyzer.h
b/include/tvm/tirx/stmt_functor_with_analyzer.h
similarity index 68%
rename from src/tirx/ir/ir_mutator_with_analyzer.h
rename to include/tvm/tirx/stmt_functor_with_analyzer.h
index 17343bbe62..62535ac931 100644
--- a/src/tirx/ir/ir_mutator_with_analyzer.h
+++ b/include/tvm/tirx/stmt_functor_with_analyzer.h
@@ -18,14 +18,16 @@
*/
/*!
- * \file tirx/ir/ir_mutator_with_analyzer.h
- * \brief IR mutator base-class with an analyzer context.
+ * \file tvm/tirx/stmt_functor_with_analyzer.h
+ * \brief Statement and expression functors with an analyzer context.
*/
-#ifndef TVM_TIRX_IR_IR_MUTATOR_WITH_ANALYZER_H_
-#define TVM_TIRX_IR_IR_MUTATOR_WITH_ANALYZER_H_
+
+#ifndef TVM_TIRX_STMT_FUNCTOR_WITH_ANALYZER_H_
+#define TVM_TIRX_STMT_FUNCTOR_WITH_ANALYZER_H_
#include <tvm/ffi/cast.h>
#include <tvm/ffi/extra/structural_visit.h>
+#include <tvm/ir/prim/expr.h>
#include <tvm/ir/scope_stack.h>
#include <tvm/ir/with_context.h>
#include <tvm/sym/analyzer.h>
@@ -38,6 +40,43 @@
namespace tvm {
namespace tirx {
+class TVM_TIRX_DLL StmtExprVisitorWithAnalyzer : public StmtExprVisitor {
+ public:
+ TVM_DEFINE_OBJECT_FUNCTOR_DEFAULT_CONSTRUCTOR(StmtExprVisitorWithAnalyzer,
StmtExprVisitor)
+
+ PrimExpr Simplify(const PrimExpr& expr) { return analyzer_->Simplify(expr); }
+
+ using StmtExprVisitor::Visit_;
+
+ ffi::Optional<VisitInterrupt> Visit_(const ForNode* op);
+ ffi::Optional<VisitInterrupt> Visit_(const BindNode* op);
+ ffi::Optional<VisitInterrupt> Visit_(const IfNode* op);
+ ffi::Optional<VisitInterrupt> Visit_(const RegionStmtNode* op);
+ ffi::Optional<VisitInterrupt> Visit_(const AssertStmtNode* op);
+ ffi::Optional<VisitInterrupt> Visit_(const CallNode* op);
+ ffi::Optional<VisitInterrupt> Visit_(const prim::LetNode* op);
+
+ // StmtExprVisitorWithAnalyzer deliberately does not handle Select nodes,
+ // because both sides of a Select node are visited regardless of the
+ // condition.
+
+ protected:
+ static void InitVTable(VTable* vtable);
+ explicit StmtExprVisitorWithAnalyzer(const VTable* vtable) :
StmtExprVisitor(vtable) {}
+ /*! \brief internal analyzer field. */
+ sym::Analyzer analyzer_;
+
+ /*! \brief Scope stack for accumulated assert constraints. */
+ ScopeStack<WithGroup<sym::ConstraintContext>> constraint_scope_;
+
+ /*! \brief Extract a constraint from a conditional statement
+ *
+ * Intended for preparing argument for use in
+ * `With<ConstraintContext>`.
+ */
+ PrimExpr ExtractRealCondition(PrimExpr condition) const;
+};
+
/*!
* \brief IRMutator with an analyzer context.
*
@@ -47,14 +86,14 @@ namespace tirx {
*
* \sa src/tirx/ir/ir_mutator_with_analyzer.cc
*/
-class IRMutatorWithAnalyzer : public StmtExprMutator {
+class TVM_TIRX_DLL StmtExprMutatorWithAnalyzer : public StmtExprMutator {
public:
using StmtExprMutator::Mutate;
using StmtExprMutator::Mutate_;
- explicit IRMutatorWithAnalyzer(const sym::Analyzer& analyzer)
- : IRMutatorWithAnalyzer(analyzer.get()) {}
- explicit IRMutatorWithAnalyzer(sym::AnalyzerObj* analyzer)
- : IRMutatorWithAnalyzer(analyzer, GlobalVTable()) {}
+ explicit StmtExprMutatorWithAnalyzer(const sym::Analyzer& analyzer)
+ : StmtExprMutatorWithAnalyzer(analyzer.get()) {}
+ explicit StmtExprMutatorWithAnalyzer(sym::AnalyzerObj* analyzer)
+ : StmtExprMutatorWithAnalyzer(analyzer, GlobalVTable()) {}
// override functions that need to populate the context information.
UnchangedOr<Stmt> Mutate_(const ForNode* op, InplaceMode inplace_mode)
override;
@@ -68,7 +107,7 @@ class IRMutatorWithAnalyzer : public StmtExprMutator {
protected:
static void InitVTable(VTable* vtable);
- IRMutatorWithAnalyzer(sym::AnalyzerObj* analyzer, const VTable* vtable)
+ StmtExprMutatorWithAnalyzer(sym::AnalyzerObj* analyzer, const VTable* vtable)
: StmtExprMutator(vtable), analyzer_(analyzer) {}
static const VTable* GlobalVTable();
/*!
@@ -126,6 +165,7 @@ class IRMutatorWithAnalyzer : public StmtExprMutator {
}
}
};
+
} // namespace tirx
} // namespace tvm
-#endif // TVM_TIRX_IR_IR_MUTATOR_WITH_ANALYZER_H_
+#endif // TVM_TIRX_STMT_FUNCTOR_WITH_ANALYZER_H_
diff --git a/src/backend/trn/transform/lower_trainium_layout.cc
b/src/backend/trn/transform/lower_trainium_layout.cc
index 2bc5a2fb70..8813c18d1b 100644
--- a/src/backend/trn/transform/lower_trainium_layout.cc
+++ b/src/backend/trn/transform/lower_trainium_layout.cc
@@ -30,6 +30,7 @@
#include <tvm/tirx/op/memory.h>
#include <tvm/tirx/stmt.h>
#include <tvm/tirx/stmt_functor.h>
+#include <tvm/tirx/stmt_functor_with_analyzer.h>
#include <tvm/tirx/transform.h>
#include <algorithm>
@@ -37,7 +38,6 @@
#include <utility>
#include <vector>
-#include "../../../tirx/ir/ir_mutator_with_analyzer.h"
#include "../../../tirx/transform/flattened_tensor.h"
namespace tvm {
@@ -53,7 +53,7 @@ static bool IsTrainiumLayout(const TileLayoutNode* layout) {
});
}
-class TrainiumLayoutApplier : public tirx::IRMutatorWithAnalyzer {
+class TrainiumLayoutApplier : public tirx::StmtExprMutatorWithAnalyzer {
public:
static std::pair<Stmt, ffi::Array<Var>> Lower(const Stmt& stmt, const
ffi::Array<Var>& params) {
sym::Analyzer ana;
@@ -88,10 +88,10 @@ class TrainiumLayoutApplier : public
tirx::IRMutatorWithAnalyzer {
}
explicit TrainiumLayoutApplier(const sym::Analyzer& analyzer)
- : tirx::IRMutatorWithAnalyzer(analyzer) {}
+ : tirx::StmtExprMutatorWithAnalyzer(analyzer) {}
protected:
- using IRMutatorWithAnalyzer::Mutate_;
+ using StmtExprMutatorWithAnalyzer::Mutate_;
ffi::Any MutateTileArgument(const ffi::Any& any) {
if (auto tensor = any.as<TensorVar>()) {
@@ -142,7 +142,7 @@ class TrainiumLayoutApplier : public
tirx::IRMutatorWithAnalyzer {
{tensor.type()}, call->loc),
op->loc);
}
- return IRMutatorWithAnalyzer::Mutate_(op, inplace_mode);
+ return StmtExprMutatorWithAnalyzer::Mutate_(op, inplace_mode);
}
TensorVar GetFlattenedTensor(TensorVar buf, bool is_alloc = false) {
diff --git a/src/s_tir/analysis/domain_touched.cc
b/src/s_tir/analysis/domain_touched.cc
index 466292c713..1508032613 100644
--- a/src/s_tir/analysis/domain_touched.cc
+++ b/src/s_tir/analysis/domain_touched.cc
@@ -27,6 +27,7 @@
#include <tvm/runtime/logging.h>
#include <tvm/s_tir/analysis.h>
#include <tvm/s_tir/stmt_functor.h>
+#include <tvm/s_tir/stmt_functor_with_analyzer.h>
#include <tvm/sym/int_set.h>
#include <tvm/te/tensor.h>
@@ -34,8 +35,6 @@
#include <unordered_map>
#include <unordered_set>
-#include "../../s_tir/ir/ir_visitor_with_analyzer.h"
-
namespace tvm {
namespace s_tir {
@@ -63,9 +62,9 @@ using TensorDomainAccess = std::tuple<LoadAccess,
StoreAccess, CombinedAccess>;
} // namespace
// Find Read region of the tensor in the stmt.
-class TensorTouchedDomain final : public s_tir::IRVisitorWithAnalyzer {
+class TensorTouchedDomain final : public s_tir::StmtExprVisitorWithAnalyzer {
public:
- using s_tir::IRVisitorWithAnalyzer::Visit_;
+ using s_tir::StmtExprVisitorWithAnalyzer::Visit_;
std::unordered_map<const VarNode*, TensorDomainAccess>&
GetAccessedTensorRegions() {
return tensor_access_map_;
@@ -99,7 +98,7 @@ class TensorTouchedDomain final : public
s_tir::IRVisitorWithAnalyzer {
}
private:
- using Parent = s_tir::IRVisitorWithAnalyzer;
+ using Parent = s_tir::StmtExprVisitorWithAnalyzer;
ffi::Optional<VisitInterrupt> Visit_(const TensorLoadNode* op) final {
TensorVar tensor = op->source.as_or_throw<tvm::tirx::TensorVar>();
diff --git a/src/s_tir/analysis/identify_memcpy.cc
b/src/s_tir/analysis/identify_memcpy.cc
index ef9da0e86c..1360f01cc4 100644
--- a/src/s_tir/analysis/identify_memcpy.cc
+++ b/src/s_tir/analysis/identify_memcpy.cc
@@ -27,6 +27,7 @@
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ir/prim/op.h>
#include <tvm/s_tir/analysis.h>
+#include <tvm/s_tir/stmt_functor_with_analyzer.h>
#include <tvm/sym/int_set.h>
#include <tvm/sym/iter_affine_map.h>
#include <tvm/tirx/analysis.h>
@@ -38,8 +39,6 @@
#include <string>
#include <variant>
-#include "../../s_tir/ir/ir_visitor_with_analyzer.h"
-
namespace tvm {
namespace s_tir {
using namespace tvm::tirx;
@@ -303,9 +302,9 @@ TVM_FFI_STATIC_INIT_BLOCK() {
refl::GlobalDef().def("s_tir.analysis._identify_memcpy", [](const Stmt&
stmt) {
ffi::Array<ffi::ObjectRef> output;
- struct Visitor : s_tir::IRVisitorWithAnalyzer {
+ struct Visitor : s_tir::StmtExprVisitorWithAnalyzer {
public:
- using s_tir::IRVisitorWithAnalyzer::Visit_;
+ using s_tir::StmtExprVisitorWithAnalyzer::Visit_;
explicit Visitor(ffi::Array<ffi::ObjectRef>* output) : output(output) {}
ffi::Array<ffi::ObjectRef>* output;
@@ -322,7 +321,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
TVM_FFI_THROW(InternalError) << "Internal error, unhandled
std::variant type";
}
- return IRVisitorWithAnalyzer::Visit_(op);
+ return StmtExprVisitorWithAnalyzer::Visit_(op);
}
};
diff --git a/src/s_tir/analysis/oob_checker.cc
b/src/s_tir/analysis/oob_checker.cc
index 8c3f459b74..783cb613ce 100644
--- a/src/s_tir/analysis/oob_checker.cc
+++ b/src/s_tir/analysis/oob_checker.cc
@@ -22,9 +22,9 @@
*/
#include <tvm/ffi/reflection/registry.h>
+#include <tvm/s_tir/stmt_functor_with_analyzer.h>
#include <tvm/tirx/transform.h>
-#include "../../s_tir/ir/ir_visitor_with_analyzer.h"
#include "../schedule/error.h"
namespace tvm {
@@ -69,22 +69,22 @@ class OOBError : public s_tir::ScheduleErrorContextObj {
IRModule mod_;
std::vector<OOBLocation> locations_;
};
-class OOBCheckerVisitor final : public s_tir::IRVisitorWithAnalyzer {
+class OOBCheckerVisitor final : public s_tir::StmtExprVisitorWithAnalyzer {
public:
- using s_tir::IRVisitorWithAnalyzer::Visit_;
+ using s_tir::StmtExprVisitorWithAnalyzer::Visit_;
ffi::Optional<VisitInterrupt> Visit_(const TensorStoreNode* node) final {
for (size_t i = 0; i < node->dest.as_or_throw<TensorVar>()->shape.size();
i++) {
CheckBounds(node, node->dest.as_or_throw<TensorVar>(), i);
}
- return IRVisitorWithAnalyzer::Visit_(node);
+ return StmtExprVisitorWithAnalyzer::Visit_(node);
}
ffi::Optional<VisitInterrupt> Visit_(const TensorLoadNode* node) final {
TensorVar tensor = node->source.as_or_throw<tvm::tirx::TensorVar>();
for (size_t i = 0; i < tensor->shape.size(); i++) {
CheckBounds(node, tensor, i);
}
- return IRVisitorWithAnalyzer::Visit_(node);
+ return StmtExprVisitorWithAnalyzer::Visit_(node);
}
template <class T>
diff --git a/src/s_tir/backend/adreno/inject_texture_alloc.cc
b/src/s_tir/backend/adreno/inject_texture_alloc.cc
index 48cea9be51..329c5088e7 100644
--- a/src/s_tir/backend/adreno/inject_texture_alloc.cc
+++ b/src/s_tir/backend/adreno/inject_texture_alloc.cc
@@ -26,13 +26,13 @@
#include <tvm/s_tir/analysis.h>
#include <tvm/s_tir/backend/adreno/transform.h>
#include <tvm/s_tir/stmt_functor.h>
+#include <tvm/s_tir/stmt_functor_with_analyzer.h>
#include <tvm/sym/iter_affine_map.h>
#include <tvm/tirx/analysis.h>
#include <tvm/tirx/op/abi.h>
#include <tvm/tirx/op/memory.h>
#include "../../../backend/opencl/runtime/texture.h"
-#include "../../../s_tir/ir/ir_mutator_with_analyzer.h"
#include "../../../tirx/transform/ir_utils.h"
namespace tvm {
@@ -47,10 +47,10 @@ using runtime::IsTextureStorage;
/*!
* \brief Inject Texture Alloc Intrinsic right after tensor allocations are
realized.
*/
-class TextureAllocInjector : public s_tir::IRMutatorWithAnalyzer {
+class TextureAllocInjector : public s_tir::StmtExprMutatorWithAnalyzer {
public:
- using s_tir::IRMutatorWithAnalyzer::Mutate;
- using s_tir::IRMutatorWithAnalyzer::Mutate_;
+ using s_tir::StmtExprMutatorWithAnalyzer::Mutate;
+ using s_tir::StmtExprMutatorWithAnalyzer::Mutate_;
static Function Inject(Function func) {
sym::Analyzer ana;
@@ -61,7 +61,7 @@ class TextureAllocInjector : public
s_tir::IRMutatorWithAnalyzer {
return func;
}
- explicit TextureAllocInjector(const sym::Analyzer& ana) :
IRMutatorWithAnalyzer(ana) {}
+ explicit TextureAllocInjector(const sym::Analyzer& ana) :
StmtExprMutatorWithAnalyzer(ana) {}
private:
UnchangedOr<Stmt> Mutate_(const BindNode* op, InplaceMode inplace_mode)
final {
diff --git a/src/s_tir/backend/adreno/texture_flatten.cc
b/src/s_tir/backend/adreno/texture_flatten.cc
index e1ba8c1370..3b7d0d92b2 100644
--- a/src/s_tir/backend/adreno/texture_flatten.cc
+++ b/src/s_tir/backend/adreno/texture_flatten.cc
@@ -28,6 +28,7 @@
#include <tvm/ir/prim/expr.h>
#include <tvm/ir/prim/op.h>
#include <tvm/s_tir/backend/adreno/transform.h>
+#include <tvm/s_tir/stmt_functor_with_analyzer.h>
#include <tvm/te/operation.h>
#include <tvm/tirx/op/memory.h>
#include <tvm/tirx/stmt.h>
@@ -36,7 +37,6 @@
#include "../../../backend/opencl/runtime/texture.h"
#include "../../../runtime/thread_storage_scope.h"
-#include "../../../s_tir/ir/ir_visitor_with_analyzer.h"
namespace tvm {
namespace s_tir {
@@ -52,7 +52,8 @@ class TextureLoweringBase : public StmtExprMutator {
using StmtExprMutator::Mutate;
using StmtExprMutator::Mutate_;
- explicit TextureLoweringBase(const ffi::Array<Var>& params,
IRVisitorWithAnalyzer* bound_analyzer)
+ explicit TextureLoweringBase(const ffi::Array<Var>& params,
+ StmtExprVisitorWithAnalyzer* bound_analyzer)
: bound_analyzer_{bound_analyzer} {
for (const Var& param : params) {
if (auto tensor = param.as<TensorVar>()) {
@@ -81,7 +82,7 @@ class TextureLoweringBase : public StmtExprMutator {
// Set of all external input and output tensors
std::unordered_set<TensorVar, ffi::ObjectPtrHash, ffi::ObjectPtrEqual>
extern_tensor_;
// Bound analzer
- IRVisitorWithAnalyzer* bound_analyzer_;
+ StmtExprVisitorWithAnalyzer* bound_analyzer_;
};
// Lower Nd storage access to 2d texture access using lowering convention
@@ -91,7 +92,8 @@ class TextureFlattener : public TextureLoweringBase {
using TextureLoweringBase::Mutate;
using TextureLoweringBase::Mutate_;
- explicit TextureFlattener(const ffi::Array<Var>& params,
IRVisitorWithAnalyzer* bound_analyzer)
+ explicit TextureFlattener(const ffi::Array<Var>& params,
+ StmtExprVisitorWithAnalyzer* bound_analyzer)
: TextureLoweringBase(params, bound_analyzer) {}
UnchangedOr<Stmt> Mutate_(const TensorStoreNode* op, InplaceMode
inplace_mode) final {
@@ -168,7 +170,7 @@ class TextureFlattener : public TextureLoweringBase {
Function TextureFlattenHandler(Function func) {
auto fptr = func.CopyOnWrite();
- auto bound_analyzer = ffi::make_object<IRVisitorWithAnalyzer>();
+ auto bound_analyzer = ffi::make_object<StmtExprVisitorWithAnalyzer>();
bound_analyzer->Visit(fptr->body);
fptr->body = ffi::make_object<TextureFlattener>(fptr->params,
bound_analyzer.get())
->Mutate(fptr->body, InplaceMode::kAllow)
diff --git a/src/s_tir/ir/ir_mutator_with_analyzer.cc
b/src/s_tir/ir/ir_mutator_with_analyzer.cc
index 30c527d89a..84a3a5adac 100644
--- a/src/s_tir/ir/ir_mutator_with_analyzer.cc
+++ b/src/s_tir/ir/ir_mutator_with_analyzer.cc
@@ -17,12 +17,13 @@
* under the License.
*/
-#include "ir_mutator_with_analyzer.h"
+#include <tvm/s_tir/stmt_functor_with_analyzer.h>
namespace tvm {
namespace s_tir {
using namespace tirx;
-UnchangedOr<Stmt> IRMutatorWithAnalyzer::Mutate_(const SBlockNode* op,
InplaceMode inplace_mode) {
+UnchangedOr<Stmt> StmtExprMutatorWithAnalyzer::Mutate_(const SBlockNode* op,
+ InplaceMode
inplace_mode) {
return constraint_scope_.WithNewScope([&]() -> UnchangedOr<Stmt> {
for (const auto& iter_var : op->iter_vars) {
analyzer_->Bind(iter_var->var, iter_var->dom.value());
diff --git a/src/s_tir/ir/ir_mutator_with_analyzer.h
b/src/s_tir/ir/ir_mutator_with_analyzer.h
deleted file mode 100644
index 24ad2b9260..0000000000
--- a/src/s_tir/ir/ir_mutator_with_analyzer.h
+++ /dev/null
@@ -1,60 +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.
- */
-
-#ifndef TVM_S_TIR_IR_MUTATOR_WITH_ANALYZER_H_
-#define TVM_S_TIR_IR_MUTATOR_WITH_ANALYZER_H_
-
-#include <tvm/s_tir/stmt_functor.h>
-
-#include "../../tirx/ir/ir_mutator_with_analyzer.h"
-
-namespace tvm {
-namespace s_tir {
-class IRMutatorWithAnalyzer : public tirx::IRMutatorWithAnalyzer {
- public:
- using Parent = tirx::IRMutatorWithAnalyzer;
- using Parent::Mutate;
- using Parent::Mutate_;
- explicit IRMutatorWithAnalyzer(const sym::Analyzer& analyzer)
- : IRMutatorWithAnalyzer(analyzer.get()) {}
- explicit IRMutatorWithAnalyzer(sym::AnalyzerObj* analyzer) :
Parent(analyzer, GlobalVTable()) {}
- virtual UnchangedOr<Stmt> Mutate_(const SBlockNode* op, InplaceMode
inplace_mode);
- virtual UnchangedOr<Stmt> Mutate_(const SBlockRealizeNode* op, InplaceMode
inplace_mode) {
- return s_tir::StmtExprMutator::MutateBlockRealize(this, op, inplace_mode);
- }
-
- protected:
- static void InitVTable(VTable* vtable) {
- Parent::InitVTable(vtable);
- SetDispatch<IRMutatorWithAnalyzer, SBlockNode>(vtable);
- SetDispatch<IRMutatorWithAnalyzer, SBlockRealizeNode>(vtable);
- }
- static const VTable* GlobalVTable() {
- static const VTable table = [] {
- VTable table;
- InitVTable(&table);
- table.Finalize();
- return table;
- }();
- return &table;
- }
-};
-} // namespace s_tir
-} // namespace tvm
-#endif // TVM_S_TIR_IR_MUTATOR_WITH_ANALYZER_H_
diff --git a/src/s_tir/ir/ir_visitor_with_analyzer.cc
b/src/s_tir/ir/ir_visitor_with_analyzer.cc
index d8dd7d78bc..02f0ddb40d 100644
--- a/src/s_tir/ir/ir_visitor_with_analyzer.cc
+++ b/src/s_tir/ir/ir_visitor_with_analyzer.cc
@@ -17,12 +17,12 @@
* under the License.
*/
-#include "ir_visitor_with_analyzer.h"
+#include <tvm/s_tir/stmt_functor_with_analyzer.h>
namespace tvm {
namespace s_tir {
using namespace tirx;
-ffi::Optional<VisitInterrupt> IRVisitorWithAnalyzer::Visit_(const SBlockNode*
op) {
+ffi::Optional<VisitInterrupt> StmtExprVisitorWithAnalyzer::Visit_(const
SBlockNode* op) {
return constraint_scope_.WithNewScope([&]() -> ffi::Optional<VisitInterrupt>
{
for (const auto& iter_var : op->iter_vars) {
analyzer_->Bind(iter_var->var, iter_var->dom.value());
diff --git a/src/s_tir/ir/ir_visitor_with_analyzer.h
b/src/s_tir/ir/ir_visitor_with_analyzer.h
deleted file mode 100644
index ea52dba33f..0000000000
--- a/src/s_tir/ir/ir_visitor_with_analyzer.h
+++ /dev/null
@@ -1,49 +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.
- */
-
-#ifndef TVM_S_TIR_IR_VISITOR_WITH_ANALYZER_H_
-#define TVM_S_TIR_IR_VISITOR_WITH_ANALYZER_H_
-
-#include <tvm/s_tir/stmt_functor.h>
-
-#include "../../tirx/ir/ir_visitor_with_analyzer.h"
-
-namespace tvm {
-namespace s_tir {
-class IRVisitorWithAnalyzer : public tirx::IRVisitorWithAnalyzer {
- public:
- using Parent = tirx::IRVisitorWithAnalyzer;
- using Parent::Visit;
- using Parent::Visit_;
- TVM_DEFINE_OBJECT_FUNCTOR_DEFAULT_CONSTRUCTOR(IRVisitorWithAnalyzer, Parent)
- virtual ffi::Optional<VisitInterrupt> Visit_(const SBlockNode* op);
- virtual ffi::Optional<VisitInterrupt> Visit_(const SBlockRealizeNode* op) {
- return s_tir::StmtExprVisitor::VisitBlockRealize(this, op);
- }
-
- protected:
- static void InitVTable(VTable* vtable) {
- Parent::InitVTable(vtable);
- SetDispatch<IRVisitorWithAnalyzer, SBlockNode>(vtable);
- SetDispatch<IRVisitorWithAnalyzer, SBlockRealizeNode>(vtable);
- }
-};
-} // namespace s_tir
-} // namespace tvm
-#endif // TVM_S_TIR_IR_VISITOR_WITH_ANALYZER_H_
diff --git a/src/s_tir/schedule/primitive/layout_transformation.cc
b/src/s_tir/schedule/primitive/layout_transformation.cc
index b6b55147cd..8f51c3bdac 100644
--- a/src/s_tir/schedule/primitive/layout_transformation.cc
+++ b/src/s_tir/schedule/primitive/layout_transformation.cc
@@ -24,12 +24,12 @@
#include <tvm/ir/prim/expr.h>
#include <tvm/runtime/logging.h>
#include <tvm/s_tir/stmt.h>
+#include <tvm/s_tir/stmt_functor_with_analyzer.h>
#include <tvm/sym/analyzer.h>
#include <optional>
#include <variant>
-#include "../../../s_tir/ir/ir_mutator_with_analyzer.h"
#include "../utils.h"
namespace tvm {
@@ -804,10 +804,10 @@ class ReuseBlocksCollector : public
s_tir::StmtExprVisitor {
ffi::Map<SBlock, SBlock> new_block_to_old_;
};
-class TransformLayoutRewriter : public s_tir::IRMutatorWithAnalyzer {
+class TransformLayoutRewriter : public s_tir::StmtExprMutatorWithAnalyzer {
public:
- using s_tir::IRMutatorWithAnalyzer::Mutate;
- using s_tir::IRMutatorWithAnalyzer::Mutate_;
+ using s_tir::StmtExprMutatorWithAnalyzer::Mutate;
+ using s_tir::StmtExprMutatorWithAnalyzer::Mutate_;
/*!
* \brief Rewrite the access to the tensor after the transformation
@@ -847,7 +847,7 @@ class TransformLayoutRewriter : public
s_tir::IRMutatorWithAnalyzer {
const IndexMap& index_map,
const TransformLayoutPlanner::TransformPlan& plan,
const sym::Analyzer& analyzer)
- : IRMutatorWithAnalyzer(analyzer),
+ : StmtExprMutatorWithAnalyzer(analyzer),
old_tensor_(old_tensor),
new_tensor_(new_tensor),
index_map_(index_map),
@@ -865,7 +865,7 @@ class TransformLayoutRewriter : public
s_tir::IRMutatorWithAnalyzer {
*indices = this->IterMapSimplifyWithContext(*indices, true);
}
- using Parent = s_tir::IRMutatorWithAnalyzer;
+ using Parent = s_tir::StmtExprMutatorWithAnalyzer;
UnchangedOr<ffi::Any> Mutate(ffi::AnyView value, InplaceMode inplace_mode)
final {
const auto* stmt = value.as<StmtNode>();
diff --git a/src/s_tir/schedule/transform.cc b/src/s_tir/schedule/transform.cc
index 86fee5ba52..56952901cb 100644
--- a/src/s_tir/schedule/transform.cc
+++ b/src/s_tir/schedule/transform.cc
@@ -467,7 +467,7 @@ void
BlockTensorAccessSimplifier::SimplifyTensorIndices(ffi::Array<PrimExpr>* in
UnchangedOr<Stmt> BlockTensorAccessSimplifier::Mutate_(const SBlockNode* op,
InplaceMode
inplace_mode) {
- SBlock block = s_tir::IRMutatorWithAnalyzer::Mutate_(op, inplace_mode)
+ SBlock block = s_tir::StmtExprMutatorWithAnalyzer::Mutate_(op, inplace_mode)
.ValueOrUnchanged(ffi::GetRef<Stmt>(op))
.as_or_throw<SBlock>();
auto* n = block.CopyOnWrite();
@@ -478,7 +478,7 @@ UnchangedOr<Stmt>
BlockTensorAccessSimplifier::Mutate_(const SBlockNode* op,
UnchangedOr<Stmt> BlockTensorAccessSimplifier::Mutate_(const TensorStoreNode*
op,
InplaceMode
inplace_mode) {
- TensorStore node = s_tir::IRMutatorWithAnalyzer::Mutate_(op, inplace_mode)
+ TensorStore node = s_tir::StmtExprMutatorWithAnalyzer::Mutate_(op,
inplace_mode)
.ValueOrUnchanged(ffi::GetRef<Stmt>(op))
.as_or_throw<TensorStore>();
SimplifyTensorIndices(&node.CopyOnWrite()->indices);
@@ -487,7 +487,7 @@ UnchangedOr<Stmt>
BlockTensorAccessSimplifier::Mutate_(const TensorStoreNode* op
UnchangedOr<PrimExpr> BlockTensorAccessSimplifier::Mutate_(const
TensorLoadNode* op,
InplaceMode
inplace_mode) {
- TensorLoad node = s_tir::IRMutatorWithAnalyzer::Mutate_(op, inplace_mode)
+ TensorLoad node = s_tir::StmtExprMutatorWithAnalyzer::Mutate_(op,
inplace_mode)
.ValueOrUnchanged(ffi::GetRef<PrimExpr>(op))
.as_or_throw<TensorLoad>();
SimplifyTensorIndices(&node.CopyOnWrite()->indices);
diff --git a/src/s_tir/schedule/transform.h b/src/s_tir/schedule/transform.h
index ba0fed9a36..b540f48d80 100644
--- a/src/s_tir/schedule/transform.h
+++ b/src/s_tir/schedule/transform.h
@@ -24,12 +24,11 @@
#include <tvm/s_tir/schedule/state.h>
#include <tvm/s_tir/stmt.h>
#include <tvm/s_tir/stmt_functor.h>
+#include <tvm/s_tir/stmt_functor_with_analyzer.h>
#include <unordered_map>
#include <utility>
-#include "../../s_tir/ir/ir_mutator_with_analyzer.h"
-
namespace tvm {
namespace s_tir {
using namespace tvm::tirx;
@@ -211,10 +210,10 @@ ffi::Optional<s_tir::LoopRV> TileWithTensorIntrin(const
s_tir::Schedule& sch,
/*!
* \brief Simplifier for indices of tensor access and block tensor access
regions.
*/
-class BlockTensorAccessSimplifier : public s_tir::IRMutatorWithAnalyzer {
+class BlockTensorAccessSimplifier : public s_tir::StmtExprMutatorWithAnalyzer {
public:
- using s_tir::IRMutatorWithAnalyzer::Mutate;
- using s_tir::IRMutatorWithAnalyzer::Mutate_;
+ using s_tir::StmtExprMutatorWithAnalyzer::Mutate;
+ using s_tir::StmtExprMutatorWithAnalyzer::Mutate_;
/*!
* \brief Simplify indices of tensor access and block tensor access regions
in the statement
@@ -228,7 +227,7 @@ class BlockTensorAccessSimplifier : public
s_tir::IRMutatorWithAnalyzer {
}
explicit BlockTensorAccessSimplifier(const sym::Analyzer& analyzer)
- : IRMutatorWithAnalyzer(analyzer) {}
+ : StmtExprMutatorWithAnalyzer(analyzer) {}
private:
void SimplifyAccessRegion(ffi::Array<TensorRegion>* old_access_regions);
diff --git a/src/s_tir/transform/inject_permuted_layout.cc
b/src/s_tir/transform/inject_permuted_layout.cc
index 708671b94c..875191dd8c 100644
--- a/src/s_tir/transform/inject_permuted_layout.cc
+++ b/src/s_tir/transform/inject_permuted_layout.cc
@@ -26,13 +26,13 @@
#include <tvm/ir/prim/op.h>
#include <tvm/s_tir/stmt.h>
#include <tvm/s_tir/stmt_functor.h>
+#include <tvm/s_tir/stmt_functor_with_analyzer.h>
#include <tvm/s_tir/transform.h>
#include <tvm/sym/analyzer.h>
#include <tvm/tirx/function.h>
#include <tvm/tirx/op/memory.h>
#include "../../runtime/thread_storage_scope.h"
-#include "../../s_tir/ir/ir_mutator_with_analyzer.h"
#include "../../support/utils.h"
#include "../../tirx/transform/ir_utils.h"
@@ -58,10 +58,10 @@ ffi::Optional<Var> GetBufferDataVar(const ffi::Any& data) {
} // namespace
-class PermutedLayoutInjector : public IRMutatorWithAnalyzer {
+class PermutedLayoutInjector : public StmtExprMutatorWithAnalyzer {
public:
- using IRMutatorWithAnalyzer::Mutate;
- using IRMutatorWithAnalyzer::Mutate_;
+ using StmtExprMutatorWithAnalyzer::Mutate;
+ using StmtExprMutatorWithAnalyzer::Mutate_;
static Function Transform(Function func) {
Analyzer analyzer;
@@ -75,7 +75,7 @@ class PermutedLayoutInjector : public IRMutatorWithAnalyzer {
}
explicit PermutedLayoutInjector(Function func, const Analyzer& analyzer)
- : IRMutatorWithAnalyzer(analyzer) {
+ : StmtExprMutatorWithAnalyzer(analyzer) {
for (const Var& param : func->params) {
if (auto tensor = param.as<TensorVar>()) {
tensor_map_.insert({tensor.value().var(), tensor.value()});
@@ -152,13 +152,13 @@ class PermutedLayoutInjector : public
IRMutatorWithAnalyzer {
if (op->annotations.count(tvm::s_tir::attr::kPermutedLayout) == 0 ||
!CheckAnnotation(op->annotations.at(tvm::s_tir::attr::kPermutedLayout))) {
- return IRMutatorWithAnalyzer::Mutate_(op, inplace_mode);
+ return StmtExprMutatorWithAnalyzer::Mutate_(op, inplace_mode);
}
auto prev_permute = permute_;
permute_ = true;
- SBlock block = IRMutatorWithAnalyzer::Mutate_(op, inplace_mode)
+ SBlock block = StmtExprMutatorWithAnalyzer::Mutate_(op, inplace_mode)
.ValueOrUnchanged(ffi::GetRef<Stmt>(op))
.as_or_throw<SBlock>();
@@ -210,7 +210,7 @@ class PermutedLayoutInjector : public IRMutatorWithAnalyzer
{
// Rewrite write from global to shared.dyn or shared
// We assume the shape of the shared memory is [..., row_size, col_size],
// where row_size is divisible by 64, or divisible by 32 and col_size is
divisible by 2.
- auto store = IRMutatorWithAnalyzer::Mutate_(op, inplace_mode)
+ auto store = StmtExprMutatorWithAnalyzer::Mutate_(op, inplace_mode)
.ValueOrUnchanged(ffi::GetRef<Stmt>(op))
.as_or_throw<TensorStore>();
@@ -231,7 +231,7 @@ class PermutedLayoutInjector : public IRMutatorWithAnalyzer
{
UnchangedOr<PrimExpr> Mutate_(const TensorLoadNode* op, InplaceMode
inplace_mode) final {
// Rewrite load from shared or shared.dyn to global
- auto load = IRMutatorWithAnalyzer::Mutate_(op, inplace_mode)
+ auto load = StmtExprMutatorWithAnalyzer::Mutate_(op, inplace_mode)
.ValueOrUnchanged(ffi::GetRef<PrimExpr>(op))
.as_or_throw<TensorLoad>();
@@ -308,7 +308,7 @@ class PermutedLayoutInjector : public IRMutatorWithAnalyzer
{
static const Op ptx_ldmatrix_op = Op::Get("tirx.ptx_legacy.ldmatrix");
static const Op mma_store_op = Op::Get("tirx.cuda.mma_store");
if (!permute_ || (!op->op.same_as(ptx_ldmatrix_op) &&
!op->op.same_as(mma_store_op))) {
- return IRMutatorWithAnalyzer::Mutate_(op, inplace_mode);
+ return StmtExprMutatorWithAnalyzer::Mutate_(op, inplace_mode);
}
bool is_ldmatrix = op->op.same_as(ptx_ldmatrix_op);
int pointer_index = is_ldmatrix ? 5 : 2;
diff --git a/src/s_tir/transform/inject_virtual_thread.cc
b/src/s_tir/transform/inject_virtual_thread.cc
index df0b54537a..eb8f3c9223 100644
--- a/src/s_tir/transform/inject_virtual_thread.cc
+++ b/src/s_tir/transform/inject_virtual_thread.cc
@@ -29,13 +29,13 @@
#include <tvm/ir/prim/op.h>
#include <tvm/s_tir/stmt.h>
#include <tvm/s_tir/stmt_functor.h>
+#include <tvm/s_tir/stmt_functor_with_analyzer.h>
#include <tvm/s_tir/transform.h>
#include <tvm/tirx/op/memory.h>
#include <tvm/tirx/op/region.h>
#include <unordered_set>
-#include "../../s_tir/ir/ir_mutator_with_analyzer.h"
#include "ir_utils.h"
namespace tvm {
@@ -244,15 +244,15 @@ class VarTouchedAnalysis : public StmtExprVisitor {
// Inject virtual thread loop
// rewrite the buffer access pattern when necessary.
-class VTInjector : public s_tir::IRMutatorWithAnalyzer {
+class VTInjector : public s_tir::StmtExprMutatorWithAnalyzer {
public:
- using s_tir::IRMutatorWithAnalyzer::Mutate;
- using s_tir::IRMutatorWithAnalyzer::Mutate_;
+ using s_tir::StmtExprMutatorWithAnalyzer::Mutate;
+ using s_tir::StmtExprMutatorWithAnalyzer::Mutate_;
// constructor
VTInjector(sym::AnalyzerObj* analyzer, Var var, int num_threads,
const std::unordered_set<const VarNode*>& touched_var, bool
allow_share)
- : IRMutatorWithAnalyzer(analyzer),
+ : StmtExprMutatorWithAnalyzer(analyzer),
var_(var),
num_threads_(num_threads),
touched_var_(touched_var),
@@ -699,12 +699,12 @@ class VTInjector : public s_tir::IRMutatorWithAnalyzer {
*/
};
-class VirtualThreadInjector : public s_tir::IRMutatorWithAnalyzer {
+class VirtualThreadInjector : public s_tir::StmtExprMutatorWithAnalyzer {
public:
- using s_tir::IRMutatorWithAnalyzer::Mutate;
- using s_tir::IRMutatorWithAnalyzer::Mutate_;
+ using s_tir::StmtExprMutatorWithAnalyzer::Mutate;
+ using s_tir::StmtExprMutatorWithAnalyzer::Mutate_;
- using IRMutatorWithAnalyzer::IRMutatorWithAnalyzer;
+ using StmtExprMutatorWithAnalyzer::StmtExprMutatorWithAnalyzer;
UnchangedOr<Stmt> Mutate_(const RegionStmtNode* op, InplaceMode
inplace_mode) final {
Stmt stmt = StmtExprMutator::Mutate_(op,
inplace_mode).ValueOrUnchanged(ffi::GetRef<Stmt>(op));
diff --git a/src/s_tir/transform/using_assume_to_reduce_branches.cc
b/src/s_tir/transform/using_assume_to_reduce_branches.cc
index 80475cdd75..7f083e21f9 100644
--- a/src/s_tir/transform/using_assume_to_reduce_branches.cc
+++ b/src/s_tir/transform/using_assume_to_reduce_branches.cc
@@ -41,12 +41,12 @@
#include <tvm/relax/expr.h>
#include <tvm/relax/op_attr_types.h>
#include <tvm/s_tir/stmt_functor.h>
+#include <tvm/s_tir/stmt_functor_with_analyzer.h>
#include <tvm/s_tir/transform.h>
#include <tvm/tirx/function.h>
#include <unordered_map>
-#include "../../s_tir/ir/ir_mutator_with_analyzer.h"
#include "../../sym/constraint_extract.h"
#include "tvm/ir/expr.h"
namespace tvm {
@@ -77,10 +77,10 @@ class AssumeChecker : public StmtExprVisitor {
}
};
-class ParseAssumeAndOvercompute : public IRMutatorWithAnalyzer {
+class ParseAssumeAndOvercompute : public StmtExprMutatorWithAnalyzer {
public:
- using IRMutatorWithAnalyzer::Mutate;
- using IRMutatorWithAnalyzer::Mutate_;
+ using StmtExprMutatorWithAnalyzer::Mutate;
+ using StmtExprMutatorWithAnalyzer::Mutate_;
/* This class analyzes the complete function.
It parses the tensor assumptions and eliminates the redundant branch
@@ -119,7 +119,7 @@ class ParseAssumeAndOvercompute : public
IRMutatorWithAnalyzer {
else_condition_context.
This class handles all these scenarios.*/
- using Parent = IRMutatorWithAnalyzer;
+ using Parent = StmtExprMutatorWithAnalyzer;
explicit ParseAssumeAndOvercompute(const Analyzer& analyzer) :
Parent(analyzer) {}
private:
diff --git a/src/tirx/ir/ir_mutator_with_analyzer.cc
b/src/tirx/ir/ir_mutator_with_analyzer.cc
index fb30f0f399..b1d5595520 100644
--- a/src/tirx/ir/ir_mutator_with_analyzer.cc
+++ b/src/tirx/ir/ir_mutator_with_analyzer.cc
@@ -20,7 +20,6 @@
/*!
* \file tirx/ir/ir_mutator_with_analyzer.cc
*/
-#include "ir_mutator_with_analyzer.h"
#include <tvm/backend/cuda/attr.h>
#include <tvm/ffi/cast.h>
@@ -29,14 +28,17 @@
#include <tvm/sym/iter_affine_map.h>
#include <tvm/tirx/analysis.h>
#include <tvm/tirx/op/region.h>
+#include <tvm/tirx/stmt_functor_with_analyzer.h>
#include "../../sym/constraint_helpers.h"
namespace tvm {
namespace tirx {
-void IRMutatorWithAnalyzer::InitVTable(VTable* vtable) {
StmtExprMutator::InitVTable(vtable); }
+void StmtExprMutatorWithAnalyzer::InitVTable(VTable* vtable) {
+ StmtExprMutator::InitVTable(vtable);
+}
-const IRMutatorWithAnalyzer::VTable* IRMutatorWithAnalyzer::GlobalVTable() {
+const StmtExprMutatorWithAnalyzer::VTable*
StmtExprMutatorWithAnalyzer::GlobalVTable() {
static const VTable table = [] {
VTable table;
InitVTable(&table);
@@ -50,7 +52,7 @@ using namespace tvm::prim;
using sym::detail::EnterConstraintFacts;
-void IRMutatorWithAnalyzer::MarkTensorParamShapes(const tirx::Function& func) {
+void StmtExprMutatorWithAnalyzer::MarkTensorParamShapes(const tirx::Function&
func) {
// Mark all symbolic tensor-parameter shape values as positive.
for (const tvm::Var& param : func->params) {
if (!param->ty.as<tirx::TensorTypeNode>()) {
@@ -63,7 +65,7 @@ void IRMutatorWithAnalyzer::MarkTensorParamShapes(const
tirx::Function& func) {
}
}
-ffi::Array<PrimExpr> IRMutatorWithAnalyzer::IterMapSimplifyWithContext(
+ffi::Array<PrimExpr> StmtExprMutatorWithAnalyzer::IterMapSimplifyWithContext(
const ffi::Array<PrimExpr>& indices, bool non_trivial_only) {
PrimExpr pred = IntImm::Bool(true);
for (PrimExpr val : iter_predicates_) {
@@ -83,7 +85,8 @@ ffi::Array<PrimExpr>
IRMutatorWithAnalyzer::IterMapSimplifyWithContext(
return simplified;
}
-UnchangedOr<Stmt> IRMutatorWithAnalyzer::Mutate_(const ForNode* op,
InplaceMode inplace_mode) {
+UnchangedOr<Stmt> StmtExprMutatorWithAnalyzer::Mutate_(const ForNode* op,
+ InplaceMode
inplace_mode) {
return constraint_scope_.WithNewScope([&]() -> UnchangedOr<Stmt> {
// record the loop variable as iterators
Range dom = Range::FromMinExtent(op->min, op->extent);
@@ -125,7 +128,8 @@ UnchangedOr<Stmt> IRMutatorWithAnalyzer::Mutate_(const
ForNode* op, InplaceMode
});
}
-UnchangedOr<Stmt> IRMutatorWithAnalyzer::Mutate_(const BindNode* op,
InplaceMode inplace_mode) {
+UnchangedOr<Stmt> StmtExprMutatorWithAnalyzer::Mutate_(const BindNode* op,
+ InplaceMode
inplace_mode) {
if (!op->value.as<PrimExpr>()) {
return StmtExprMutator::Mutate_(op, inplace_mode);
}
@@ -148,7 +152,7 @@ UnchangedOr<Stmt> IRMutatorWithAnalyzer::Mutate_(const
BindNode* op, InplaceMode
return Stmt(n);
}
-UnchangedOr<Stmt> IRMutatorWithAnalyzer::Mutate_(const IfNode* op, InplaceMode
inplace_mode) {
+UnchangedOr<Stmt> StmtExprMutatorWithAnalyzer::Mutate_(const IfNode* op,
InplaceMode inplace_mode) {
return constraint_scope_.WithNewScope([&]() -> UnchangedOr<Stmt> {
auto condition_result = this->Mutate(op->condition, inplace_mode);
bool condition_unchanged =
condition_result.UnchangedOrSameAs(op->condition);
@@ -211,8 +215,8 @@ UnchangedOr<Stmt> IRMutatorWithAnalyzer::Mutate_(const
IfNode* op, InplaceMode i
});
}
-UnchangedOr<Stmt> IRMutatorWithAnalyzer::Mutate_(const RegionStmtNode* op,
- InplaceMode inplace_mode) {
+UnchangedOr<Stmt> StmtExprMutatorWithAnalyzer::Mutate_(const RegionStmtNode*
op,
+ InplaceMode
inplace_mode) {
return constraint_scope_.WithNewScope([&]() -> UnchangedOr<Stmt> {
if (op->op.same_as(tirx::launch_thread_op())) {
PrimVar var = op->body_params[0].as_or_throw<PrimVar>();
@@ -241,8 +245,8 @@ UnchangedOr<Stmt> IRMutatorWithAnalyzer::Mutate_(const
RegionStmtNode* op,
});
}
-UnchangedOr<Stmt> IRMutatorWithAnalyzer::Mutate_(const AssertStmtNode* op,
- InplaceMode inplace_mode) {
+UnchangedOr<Stmt> StmtExprMutatorWithAnalyzer::Mutate_(const AssertStmtNode*
op,
+ InplaceMode
inplace_mode) {
auto condition_result = this->Mutate(op->condition, inplace_mode);
bool condition_unchanged = condition_result.UnchangedOrSameAs(op->condition);
PrimExpr condition =
std::move(condition_result).ValueOrUnchanged(op->condition);
@@ -265,7 +269,8 @@ UnchangedOr<Stmt> IRMutatorWithAnalyzer::Mutate_(const
AssertStmtNode* op,
}
}
-UnchangedOr<Expr> IRMutatorWithAnalyzer::Mutate_(const CallNode* op,
InplaceMode inplace_mode) {
+UnchangedOr<Expr> StmtExprMutatorWithAnalyzer::Mutate_(const CallNode* op,
+ InplaceMode
inplace_mode) {
// add condition context to if_then_else
static const Op if_then_else_op = Op::Get("prim.if_then_else");
if (op->op.same_as(if_then_else_op)) {
@@ -302,8 +307,8 @@ UnchangedOr<Expr> IRMutatorWithAnalyzer::Mutate_(const
CallNode* op, InplaceMode
return StmtExprMutator::Mutate_(op, inplace_mode);
}
-UnchangedOr<PrimExpr> IRMutatorWithAnalyzer::Mutate_(const prim::LetNode* op,
- InplaceMode inplace_mode)
{
+UnchangedOr<PrimExpr> StmtExprMutatorWithAnalyzer::Mutate_(const
prim::LetNode* op,
+ InplaceMode
inplace_mode) {
auto value_result = this->Mutate(op->value, inplace_mode);
bool value_unchanged = value_result.UnchangedOrSameAs(op->value);
PrimExpr value = std::move(value_result).ValueOrUnchanged(op->value);
@@ -322,8 +327,8 @@ UnchangedOr<PrimExpr> IRMutatorWithAnalyzer::Mutate_(const
prim::LetNode* op,
}
}
-UnchangedOr<PrimExpr> IRMutatorWithAnalyzer::Mutate_(const prim::SelectNode*
op,
- InplaceMode inplace_mode)
{
+UnchangedOr<PrimExpr> StmtExprMutatorWithAnalyzer::Mutate_(const
prim::SelectNode* op,
+ InplaceMode
inplace_mode) {
auto cond_result = this->Mutate(op->condition, inplace_mode);
bool cond_unchanged = cond_result.UnchangedOrSameAs(op->condition);
PrimExpr cond = std::move(cond_result).ValueOrUnchanged(op->condition);
diff --git a/src/tirx/ir/ir_visitor_with_analyzer.cc
b/src/tirx/ir/ir_visitor_with_analyzer.cc
index 0ef596ccba..4f13c2009a 100644
--- a/src/tirx/ir/ir_visitor_with_analyzer.cc
+++ b/src/tirx/ir/ir_visitor_with_analyzer.cc
@@ -20,18 +20,20 @@
/*!
* \file tirx/ir/ir_visitor_with_analyzer.cc
*/
-#include "ir_visitor_with_analyzer.h"
#include <tvm/ir/op.h>
#include <tvm/ir/prim/op.h>
#include <tvm/tirx/analysis.h>
#include <tvm/tirx/op/region.h>
+#include <tvm/tirx/stmt_functor_with_analyzer.h>
namespace tvm {
namespace tirx {
-void IRVisitorWithAnalyzer::InitVTable(VTable* vtable) {
StmtExprVisitor::InitVTable(vtable); }
+void StmtExprVisitorWithAnalyzer::InitVTable(VTable* vtable) {
+ StmtExprVisitor::InitVTable(vtable);
+}
-ffi::Optional<VisitInterrupt> IRVisitorWithAnalyzer::Visit_(const ForNode* op)
{
+ffi::Optional<VisitInterrupt> StmtExprVisitorWithAnalyzer::Visit_(const
ForNode* op) {
return constraint_scope_.WithNewScope([&]() -> ffi::Optional<VisitInterrupt>
{
analyzer_->Bind(op->loop_var, Range::FromMinExtent(op->min, op->extent));
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->Visit(op->min));
@@ -46,7 +48,7 @@ ffi::Optional<VisitInterrupt>
IRVisitorWithAnalyzer::Visit_(const ForNode* op) {
});
}
-ffi::Optional<VisitInterrupt> IRVisitorWithAnalyzer::Visit_(const BindNode*
op) {
+ffi::Optional<VisitInterrupt> StmtExprVisitorWithAnalyzer::Visit_(const
BindNode* op) {
if (!op->value.as<PrimExpr>()) return StmtExprVisitor::Visit_(op);
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->Visit(op->value));
if (ffi::Optional<PrimExpr> value = op->value.as<PrimExpr>()) {
@@ -55,7 +57,7 @@ ffi::Optional<VisitInterrupt>
IRVisitorWithAnalyzer::Visit_(const BindNode* op)
return std::nullopt;
}
-ffi::Optional<VisitInterrupt> IRVisitorWithAnalyzer::Visit_(const IfNode* op) {
+ffi::Optional<VisitInterrupt> StmtExprVisitorWithAnalyzer::Visit_(const
IfNode* op) {
return constraint_scope_.WithNewScope([&]() -> ffi::Optional<VisitInterrupt>
{
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->Visit(op->condition));
@@ -78,7 +80,7 @@ ffi::Optional<VisitInterrupt>
IRVisitorWithAnalyzer::Visit_(const IfNode* op) {
});
}
-ffi::Optional<VisitInterrupt> IRVisitorWithAnalyzer::Visit_(const
RegionStmtNode* op) {
+ffi::Optional<VisitInterrupt> StmtExprVisitorWithAnalyzer::Visit_(const
RegionStmtNode* op) {
return constraint_scope_.WithNewScope([&]() -> ffi::Optional<VisitInterrupt>
{
if (op->op.same_as(tirx::launch_thread_op())) {
PrimExpr extent = op->args[1].as_or_throw<PrimExpr>();
@@ -89,13 +91,13 @@ ffi::Optional<VisitInterrupt>
IRVisitorWithAnalyzer::Visit_(const RegionStmtNode
});
}
-ffi::Optional<VisitInterrupt> IRVisitorWithAnalyzer::Visit_(const
AssertStmtNode* op) {
+ffi::Optional<VisitInterrupt> StmtExprVisitorWithAnalyzer::Visit_(const
AssertStmtNode* op) {
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->Visit(op->condition));
constraint_scope_.Current().Emplace(analyzer_, op->condition);
return std::nullopt;
}
-ffi::Optional<VisitInterrupt> IRVisitorWithAnalyzer::Visit_(const CallNode*
op) {
+ffi::Optional<VisitInterrupt> StmtExprVisitorWithAnalyzer::Visit_(const
CallNode* op) {
// add condition context to if_then_else
static const Op if_then_else_op = Op::Get("prim.if_then_else");
if (op->op.same_as(if_then_else_op)) {
@@ -118,13 +120,13 @@ ffi::Optional<VisitInterrupt>
IRVisitorWithAnalyzer::Visit_(const CallNode* op)
return std::nullopt;
}
-ffi::Optional<VisitInterrupt> IRVisitorWithAnalyzer::Visit_(const
prim::LetNode* op) {
+ffi::Optional<VisitInterrupt> StmtExprVisitorWithAnalyzer::Visit_(const
prim::LetNode* op) {
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->Visit(op->value));
analyzer_->Bind(op->var, op->value);
return this->Visit(op->body);
}
-PrimExpr IRVisitorWithAnalyzer::ExtractRealCondition(PrimExpr condition) const
{
+PrimExpr StmtExprVisitorWithAnalyzer::ExtractRealCondition(PrimExpr condition)
const {
if (auto call = condition.as<CallNode>()) {
if (call->op.same_as(prim::likely_op())) {
return call->args[0].as_or_throw<PrimExpr>();
diff --git a/src/tirx/ir/ir_visitor_with_analyzer.h
b/src/tirx/ir/ir_visitor_with_analyzer.h
deleted file mode 100644
index 30910aeac6..0000000000
--- a/src/tirx/ir/ir_visitor_with_analyzer.h
+++ /dev/null
@@ -1,76 +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 tirx/ir/ir_visitor_with_analyzer.h
- * \brief IR visitor class with an analyzer context.
- */
-
-#ifndef TVM_TIRX_IR_IR_VISITOR_WITH_ANALYZER_H_
-#define TVM_TIRX_IR_IR_VISITOR_WITH_ANALYZER_H_
-
-#include <tvm/ir/prim/expr.h>
-#include <tvm/ir/scope_stack.h>
-#include <tvm/ir/with_context.h>
-#include <tvm/sym/analyzer.h>
-#include <tvm/tirx/stmt_functor.h>
-
-namespace tvm {
-namespace tirx {
-
-class IRVisitorWithAnalyzer : public StmtExprVisitor {
- public:
- TVM_DEFINE_OBJECT_FUNCTOR_DEFAULT_CONSTRUCTOR(IRVisitorWithAnalyzer,
StmtExprVisitor)
-
- PrimExpr Simplify(const PrimExpr& expr) { return analyzer_->Simplify(expr); }
-
- using StmtExprVisitor::Visit_;
-
- ffi::Optional<VisitInterrupt> Visit_(const ForNode* op);
- ffi::Optional<VisitInterrupt> Visit_(const BindNode* op);
- ffi::Optional<VisitInterrupt> Visit_(const IfNode* op);
- ffi::Optional<VisitInterrupt> Visit_(const RegionStmtNode* op);
- ffi::Optional<VisitInterrupt> Visit_(const AssertStmtNode* op);
- ffi::Optional<VisitInterrupt> Visit_(const CallNode* op);
- ffi::Optional<VisitInterrupt> Visit_(const prim::LetNode* op);
-
- // IRVisitorWithAnalyzer deliberately does not handle Select nodes,
- // because both sides of a Select node are visited regardless of the
- // condition.
-
- protected:
- static void InitVTable(VTable* vtable);
- explicit IRVisitorWithAnalyzer(const VTable* vtable) :
StmtExprVisitor(vtable) {}
- /*! \brief internal analyzer field. */
- sym::Analyzer analyzer_;
-
- /*! \brief Scope stack for accumulated assert constraints. */
- ScopeStack<WithGroup<sym::ConstraintContext>> constraint_scope_;
-
- /*! \brief Extract a constraint from a conditional statement
- *
- * Intended for preparing argument for use in
- * `With<ConstraintContext>`.
- */
- PrimExpr ExtractRealCondition(PrimExpr condition) const;
-};
-
-} // namespace tirx
-} // namespace tvm
-#endif // TVM_TIRX_IR_IR_VISITOR_WITH_ANALYZER_H_
diff --git a/src/tirx/transform/flatten_tensor.cc
b/src/tirx/transform/flatten_tensor.cc
index 5e22c7c719..667c78fd3c 100644
--- a/src/tirx/transform/flatten_tensor.cc
+++ b/src/tirx/transform/flatten_tensor.cc
@@ -29,11 +29,11 @@
#include <tvm/tirx/layout.h>
#include <tvm/tirx/op/memory.h>
#include <tvm/tirx/stmt_functor.h>
+#include <tvm/tirx/stmt_functor_with_analyzer.h>
#include <tvm/tirx/transform.h>
#include <unordered_set>
-#include "../ir/ir_mutator_with_analyzer.h"
#include "flattened_tensor.h"
#include "ir_utils.h"
@@ -60,10 +60,10 @@ using namespace tvm::prim;
* Every use site then only looks the pair up; a use before its definition is
* a hard error instead of a silently stale reference.
*/
-class TensorFlattener : public IRMutatorWithAnalyzer {
+class TensorFlattener : public StmtExprMutatorWithAnalyzer {
public:
- using IRMutatorWithAnalyzer::Mutate;
- using IRMutatorWithAnalyzer::Mutate_;
+ using StmtExprMutatorWithAnalyzer::Mutate;
+ using StmtExprMutatorWithAnalyzer::Mutate_;
static Function Flatten(Function func) {
if (!func->body.has_value()) return func;
sym::Analyzer ana;
@@ -103,7 +103,7 @@ class TensorFlattener : public IRMutatorWithAnalyzer {
}
public:
- explicit TensorFlattener(const sym::Analyzer& ana) :
IRMutatorWithAnalyzer(ana) {}
+ explicit TensorFlattener(const sym::Analyzer& ana) :
StmtExprMutatorWithAnalyzer(ana) {}
private:
struct FlatInfo {
@@ -188,7 +188,7 @@ class TensorFlattener : public IRMutatorWithAnalyzer {
if (call->op.same_as(alloc_tensor_op())) return MutateAllocTensor(op,
call, inplace_mode);
if (call->op.same_as(decl_tensor_op())) return MutateDeclTensor(op,
call, inplace_mode);
}
- return IRMutatorWithAnalyzer::Mutate_(op, inplace_mode);
+ return StmtExprMutatorWithAnalyzer::Mutate_(op, inplace_mode);
}
UnchangedOr<Stmt> MutateAllocTensor(const BindNode* op, const CallNode*
tensor_call,
@@ -287,7 +287,7 @@ class TensorFlattener : public IRMutatorWithAnalyzer {
}
}
}
- return IRMutatorWithAnalyzer::Mutate_(op, inplace_mode);
+ return StmtExprMutatorWithAnalyzer::Mutate_(op, inplace_mode);
}
ffi::Array<PrimExpr> FoldIndices(const FlatInfo& info, const
ffi::Array<PrimExpr>& indices) {
diff --git a/src/tirx/transform/lower_intrin.cc
b/src/tirx/transform/lower_intrin.cc
index fc3e23095f..c1f3a5f658 100644
--- a/src/tirx/transform/lower_intrin.cc
+++ b/src/tirx/transform/lower_intrin.cc
@@ -32,27 +32,27 @@
#include <tvm/target/target.h>
#include <tvm/tirx/expr.h>
#include <tvm/tirx/op/memory.h>
+#include <tvm/tirx/stmt_functor_with_analyzer.h>
#include <tvm/tirx/transform.h>
#include <limits>
#include <unordered_set>
#include "../../sym/pattern_match.h"
-#include "../ir/ir_mutator_with_analyzer.h"
namespace tvm {
namespace tirx {
using namespace tvm::prim;
-class IntrinInjecter : public IRMutatorWithAnalyzer {
+class IntrinInjecter : public StmtExprMutatorWithAnalyzer {
public:
- using IRMutatorWithAnalyzer::Mutate;
- using IRMutatorWithAnalyzer::Mutate_;
+ using StmtExprMutatorWithAnalyzer::Mutate;
+ using StmtExprMutatorWithAnalyzer::Mutate_;
using FLowerGeneral = ffi::TypedFunction<PrimExpr(PrimExpr)>;
IntrinInjecter(const sym::Analyzer& analyzer, const Target& tgt, bool
enable_fast_math)
- : IRMutatorWithAnalyzer(analyzer) {
+ : StmtExprMutatorWithAnalyzer(analyzer) {
std::string target = tgt->kind->name;
ffi::String mtriple = tgt->GetAttr<ffi::String>("mtriple").value_or("");
@@ -104,7 +104,7 @@ class IntrinInjecter : public IRMutatorWithAnalyzer {
}
}
}
- return IRMutatorWithAnalyzer::Mutate_(op, inplace_mode);
+ return StmtExprMutatorWithAnalyzer::Mutate_(op, inplace_mode);
}
UnchangedOr<PrimExpr> Mutate_(const prim::AddNode* op, InplaceMode
inplace_mode) final {
@@ -113,14 +113,14 @@ class IntrinInjecter : public IRMutatorWithAnalyzer {
} else if (const prim::MulNode* ma = op->a.as<prim::MulNode>()) {
return MakeFMA(ma->a, ma->b, op->b, op);
}
- return IRMutatorWithAnalyzer::Mutate_(op, inplace_mode);
+ return StmtExprMutatorWithAnalyzer::Mutate_(op, inplace_mode);
}
// We use floordiv for integer analysis,
// but will need to lower them to native truncdiv instructions
UnchangedOr<PrimExpr> Mutate_(const prim::FloorDivNode* op, InplaceMode
inplace_mode) final {
auto e = ffi::GetRef<PrimExpr>(op);
- PrimExpr ret = IRMutatorWithAnalyzer::Mutate_(op, inplace_mode)
+ PrimExpr ret = StmtExprMutatorWithAnalyzer::Mutate_(op, inplace_mode)
.ValueOrUnchanged(ffi::GetRef<PrimExpr>(op));
op = ret.as<prim::FloorDivNode>();
if (op == nullptr) return ret;
@@ -184,7 +184,7 @@ class IntrinInjecter : public IRMutatorWithAnalyzer {
}
UnchangedOr<PrimExpr> Mutate_(const prim::FloorModNode* op, InplaceMode
inplace_mode) final {
- PrimExpr ret = IRMutatorWithAnalyzer::Mutate_(op, inplace_mode)
+ PrimExpr ret = StmtExprMutatorWithAnalyzer::Mutate_(op, inplace_mode)
.ValueOrUnchanged(ffi::GetRef<PrimExpr>(op));
op = ret.as<prim::FloorModNode>();
if (op == nullptr) return ret;
@@ -258,7 +258,7 @@ class IntrinInjecter : public IRMutatorWithAnalyzer {
PrimExpr input = truncdiv(x, y).Eval();
return max(Mutate(input, inplace_mode).ValueOrUnchanged(input),
c.Eval());
}
- return IRMutatorWithAnalyzer::Mutate_(op, inplace_mode);
+ return StmtExprMutatorWithAnalyzer::Mutate_(op, inplace_mode);
}
UnchangedOr<PrimExpr> Mutate_(const prim::EQNode* op, InplaceMode
inplace_mode) final {
@@ -269,7 +269,7 @@ class IntrinInjecter : public IRMutatorWithAnalyzer {
PrimExpr input = (truncmod(x, y) == 0).Eval();
return Mutate(input, inplace_mode).ValueOrUnchanged(input);
}
- return IRMutatorWithAnalyzer::Mutate_(op, inplace_mode);
+ return StmtExprMutatorWithAnalyzer::Mutate_(op, inplace_mode);
}
UnchangedOr<PrimExpr> Mutate_(const prim::NENode* op, InplaceMode
inplace_mode) final {
@@ -280,7 +280,7 @@ class IntrinInjecter : public IRMutatorWithAnalyzer {
PrimExpr input = (truncmod(x, y) != 0).Eval();
return Mutate(input, inplace_mode).ValueOrUnchanged(input);
}
- return IRMutatorWithAnalyzer::Mutate_(op, inplace_mode);
+ return StmtExprMutatorWithAnalyzer::Mutate_(op, inplace_mode);
}
private:
@@ -337,7 +337,7 @@ class IntrinInjecter : public IRMutatorWithAnalyzer {
return prim::Add(mul, this->Mutate(c,
InplaceMode::kDisallow).ValueOrUnchanged(c));
}
}
- return IRMutatorWithAnalyzer::Mutate_(op, InplaceMode::kDisallow)
+ return StmtExprMutatorWithAnalyzer::Mutate_(op, InplaceMode::kDisallow)
.ValueOrUnchanged(ffi::GetRef<PrimExpr>(op));
}
diff --git a/src/tirx/transform/lower_tirx_cleanup.cc
b/src/tirx/transform/lower_tirx_cleanup.cc
index cbd8922bce..1ab309d1ee 100644
--- a/src/tirx/transform/lower_tirx_cleanup.cc
+++ b/src/tirx/transform/lower_tirx_cleanup.cc
@@ -31,6 +31,7 @@
#include <tvm/tirx/op/memory.h>
#include <tvm/tirx/stmt.h>
#include <tvm/tirx/stmt_functor.h>
+#include <tvm/tirx/stmt_functor_with_analyzer.h>
#include <tvm/tirx/transform.h>
#include <tuple>
@@ -39,17 +40,16 @@
#include <utility>
#include <vector>
-#include "../ir/ir_mutator_with_analyzer.h"
#include "flattened_tensor.h"
#include "ir_utils.h"
namespace tvm {
namespace tirx {
-class LayoutApplier : public IRMutatorWithAnalyzer {
+class LayoutApplier : public StmtExprMutatorWithAnalyzer {
public:
- using IRMutatorWithAnalyzer::Mutate;
- using IRMutatorWithAnalyzer::Mutate_;
+ using StmtExprMutatorWithAnalyzer::Mutate;
+ using StmtExprMutatorWithAnalyzer::Mutate_;
static std::pair<Stmt, ffi::Array<Var>> Flatten(const Stmt& stmt, const
ffi::Array<Var>& params,
const Target& target) {
sym::Analyzer ana;
@@ -87,7 +87,7 @@ class LayoutApplier : public IRMutatorWithAnalyzer {
public:
explicit LayoutApplier(const sym::Analyzer& analyzer, const Target& target)
- : IRMutatorWithAnalyzer(analyzer), target_(target) {}
+ : StmtExprMutatorWithAnalyzer(analyzer), target_(target) {}
protected:
ffi::Any VisitAny(const ffi::Any& any) {
@@ -120,7 +120,7 @@ class LayoutApplier : public IRMutatorWithAnalyzer {
return root.as_or_throw<TensorVar>().data();
}
}
- return IRMutatorWithAnalyzer::Mutate_(op, inplace_mode);
+ return StmtExprMutatorWithAnalyzer::Mutate_(op, inplace_mode);
}
UnchangedOr<Stmt> Mutate_(const BindNode* op, InplaceMode inplace_mode)
final {
@@ -147,7 +147,7 @@ class LayoutApplier : public IRMutatorWithAnalyzer {
.ValueOrUnchanged(binding);
}
}
- return IRMutatorWithAnalyzer::Mutate_(op, inplace_mode);
+ return StmtExprMutatorWithAnalyzer::Mutate_(op, inplace_mode);
}
UnchangedOr<Stmt> MutateAllocTensor(const BindNode* op, const CallNode*
tensor_call,
diff --git a/src/tirx/transform/remove_no_op.cc
b/src/tirx/transform/remove_no_op.cc
index 9aa90911dc..1eed6b6054 100644
--- a/src/tirx/transform/remove_no_op.cc
+++ b/src/tirx/transform/remove_no_op.cc
@@ -30,12 +30,12 @@
#include <tvm/tirx/analysis.h>
#include <tvm/tirx/stmt.h>
#include <tvm/tirx/stmt_functor.h>
+#include <tvm/tirx/stmt_functor_with_analyzer.h>
#include <tvm/tirx/transform.h>
#include <unordered_map>
#include "../../sym/const_fold.h"
-#include "../ir/ir_mutator_with_analyzer.h"
#include "ir_utils.h"
namespace tvm {
@@ -74,17 +74,17 @@ TVM_FFI_STATIC_INIT_BLOCK() {
}
// Mark the statement of each stage.
-class NoOpRemover : public IRMutatorWithAnalyzer {
+class NoOpRemover : public StmtExprMutatorWithAnalyzer {
public:
- using IRMutatorWithAnalyzer::Mutate;
- using IRMutatorWithAnalyzer::Mutate_;
+ using StmtExprMutatorWithAnalyzer::Mutate;
+ using StmtExprMutatorWithAnalyzer::Mutate_;
static Stmt Apply(Stmt stmt, const sym::Analyzer& analyzer, bool
ignore_profiler_call = false) {
auto visitor = ffi::make_object<NoOpRemover>(analyzer,
ignore_profiler_call);
return visitor->Mutate(stmt, InplaceMode::kAllow).ValueOrUnchanged(stmt);
}
private:
- using Parent = IRMutatorWithAnalyzer;
+ using Parent = StmtExprMutatorWithAnalyzer;
public:
NoOpRemover(const sym::Analyzer& analyzer, bool ignore_profiler_call = false)
diff --git a/src/tirx/transform/renormalize_split_pattern.cc
b/src/tirx/transform/renormalize_split_pattern.cc
index 1e19200b9a..9e3804751c 100644
--- a/src/tirx/transform/renormalize_split_pattern.cc
+++ b/src/tirx/transform/renormalize_split_pattern.cc
@@ -25,10 +25,10 @@
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ir/prim/op.h>
#include <tvm/tirx/stmt.h>
+#include <tvm/tirx/stmt_functor_with_analyzer.h>
#include <tvm/tirx/transform.h>
#include "../../sym/pattern_match.h"
-#include "../ir/ir_mutator_with_analyzer.h"
namespace tvm {
namespace tirx {
@@ -47,12 +47,13 @@ using namespace sym;
return RecursiveRewrite((ResExpr).Eval()); \
}
-class SplitPatternReNormalizer : public IRMutatorWithAnalyzer {
+class SplitPatternReNormalizer : public StmtExprMutatorWithAnalyzer {
public:
- using IRMutatorWithAnalyzer::Mutate;
- using IRMutatorWithAnalyzer::Mutate_;
+ using StmtExprMutatorWithAnalyzer::Mutate;
+ using StmtExprMutatorWithAnalyzer::Mutate_;
- explicit SplitPatternReNormalizer(const Analyzer& analyzer) :
IRMutatorWithAnalyzer(analyzer) {}
+ explicit SplitPatternReNormalizer(const Analyzer& analyzer)
+ : StmtExprMutatorWithAnalyzer(analyzer) {}
UnchangedOr<PrimExpr> Mutate_(const prim::FloorDivNode* op, InplaceMode
inplace_mode) final {
PrimExpr a = Mutate(op->a, inplace_mode).ValueOrUnchanged(op->a);
@@ -172,7 +173,7 @@ class SplitPatternReNormalizer : public
IRMutatorWithAnalyzer {
}
UnchangedOr<PrimExpr> Mutate_(const prim::NotNode* op, InplaceMode
inplace_mode) {
- PrimExpr ret = IRMutatorWithAnalyzer::Mutate_(op, inplace_mode)
+ PrimExpr ret = StmtExprMutatorWithAnalyzer::Mutate_(op, inplace_mode)
.ValueOrUnchanged(ffi::GetRef<PrimExpr>(op));
// Pattern var to match any expression
PVar<PrimExpr> x, y;
@@ -189,7 +190,7 @@ class SplitPatternReNormalizer : public
IRMutatorWithAnalyzer {
With<ConstraintContext> ctx1(analyzer_, op->loop_var >= op->min);
With<ConstraintContext> ctx2(analyzer_,
static_cast<PrimExpr>(op->loop_var) < op->min
+ op->extent);
- return IRMutatorWithAnalyzer::Mutate_(op, inplace_mode);
+ return StmtExprMutatorWithAnalyzer::Mutate_(op, inplace_mode);
}
// Recursive rewrite x
diff --git a/src/tirx/transform/stmt_simplify.cc
b/src/tirx/transform/stmt_simplify.cc
index 2cc1afb780..b8790de583 100644
--- a/src/tirx/transform/stmt_simplify.cc
+++ b/src/tirx/transform/stmt_simplify.cc
@@ -33,10 +33,9 @@
#include <tvm/sym/analyzer.h>
#include <tvm/tirx/analysis.h>
#include <tvm/tirx/op/memory.h>
+#include <tvm/tirx/stmt_functor_with_analyzer.h>
#include <tvm/tirx/transform.h>
-#include "../ir/ir_mutator_with_analyzer.h"
-
namespace tvm {
namespace tirx {
using namespace tvm::prim;
diff --git a/src/tirx/transform/stmt_simplify.h
b/src/tirx/transform/stmt_simplify.h
index 1e4c3e4caf..e4a5587ebc 100644
--- a/src/tirx/transform/stmt_simplify.h
+++ b/src/tirx/transform/stmt_simplify.h
@@ -26,8 +26,7 @@
#include <tvm/sym/analyzer.h>
#include <tvm/tirx/function.h>
-
-#include "../ir/ir_mutator_with_analyzer.h"
+#include <tvm/tirx/stmt_functor_with_analyzer.h>
namespace tvm {
namespace tirx {
@@ -50,18 +49,18 @@ class StmtSimplifyConfig : public ffi::ObjectRef {
StmtSimplifyConfigObj);
};
-class StmtSimplifier : public IRMutatorWithAnalyzer {
+class StmtSimplifier : public StmtExprMutatorWithAnalyzer {
public:
- using IRMutatorWithAnalyzer::Mutate;
- using IRMutatorWithAnalyzer::Mutate_;
+ using StmtExprMutatorWithAnalyzer::Mutate;
+ using StmtExprMutatorWithAnalyzer::Mutate_;
static Function Apply(Function func, const sym::Analyzer& analyzer,
ffi::Optional<StmtSimplifyConfig> config_opt =
std::nullopt);
explicit StmtSimplifier(const sym::Analyzer& analyzer, StmtSimplifyConfig
config)
- : IRMutatorWithAnalyzer(analyzer), config_(config) {}
+ : StmtExprMutatorWithAnalyzer(analyzer), config_(config) {}
protected:
- using Parent = IRMutatorWithAnalyzer;
+ using Parent = StmtExprMutatorWithAnalyzer;
StmtSimplifier(const VTable* vtable, const sym::Analyzer& analyzer,
StmtSimplifyConfig config)
: Parent(analyzer.get(), vtable), config_(config) {}
Function Run(Function func);
