Move S-TIR specific node and TensorIntrin out of TIRX

include/tvm/relax/analysis.h

@@ -29,6 +29,7 @@
 #include <tvm/relax/expr.h>
 #include <tvm/relax/op_attr_types.h>
 #include <tvm/relax/type.h>
+#include <tvm/s_tir/stmt.h>
 #include <tvm/tirx/function.h>
 #include <tvm/tirx/index_map.h>
 
@@ -646,7 +647,7 @@ TVM_DLL bool CheckWellFormed(ffi::Variant<IRModule, Function> obj, bool check_ty
  * from the object (block or buffer) to it's index map transformation.
  */
 
-TVM_DLL ffi::Map<tirx::SBlock, ffi::Map<ffi::ObjectRef, tirx::IndexMap>> SuggestLayoutTransforms(
+TVM_DLL ffi::Map<s_tir::SBlock, ffi::Map<ffi::ObjectRef, tirx::IndexMap>> SuggestLayoutTransforms(
     const Function& fn, ffi::Array<tirx::IndexMap> write_buffer_transformations);
 
 /* \brief Collect variables whose value can be computed at compile-time

include/tvm/relax/distributed/axis_group_graph.h

@@ -23,8 +23,9 @@
 #include <tvm/arith/iter_affine_map.h>
 #include <tvm/relax/distributed/type.h>
 #include <tvm/relax/expr.h>
+#include <tvm/s_tir/stmt.h>
+#include <tvm/s_tir/stmt_functor.h>
 #include <tvm/tirx/function.h>
-#include <tvm/tirx/stmt_functor.h>
 
 #include <algorithm>
 #include <limits>
@@ -65,7 +66,7 @@ Var GetShardingVarFromIndex(PrimExpr index, ffi::Map<Var, Range> var_range,
  * \brief Construct an axis group graph from a PrimFunc. Two buffer axis are connected if they
  * are accessed by the same index.
  */
-class BufferAxisGraphExtractor : public StmtExprVisitor {
+class BufferAxisGraphExtractor : public s_tir::StmtExprVisitor {
  public:
   static std::vector<std::vector<TIRVarAxis>> GetTIRVarAxisGraph(const PrimFunc& prim_func) {
     auto extractor = ffi::make_object<BufferAxisGraphExtractor>();
@@ -119,14 +120,14 @@ class BufferAxisGraphExtractor : public StmtExprVisitor {
 
  private:
   ffi::Optional<VisitInterrupt> Visit_(const BufferStoreNode* op) final {
-    TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(StmtExprVisitor::Visit_(op));
+    TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(s_tir::StmtExprVisitor::Visit_(op));
     buffer_access_indices_.push_back({op->buffer, op->indices});
 
     return std::nullopt;
   }
 
   ffi::Optional<VisitInterrupt> Visit_(const TensorLoadNode* op) final {
-    TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(StmtExprVisitor::Visit_(op));
+    TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(s_tir::StmtExprVisitor::Visit_(op));
     buffer_access_indices_.push_back({op->source.as_or_throw<tvm::tirx::BufferVar>(), op->indices});
 
     return std::nullopt;
@@ -158,12 +159,12 @@ class BufferAxisGraphExtractor : public StmtExprVisitor {
     return true;
   }
 
-  ffi::Optional<VisitInterrupt> Visit_(const SBlockNode* op) final {
+  ffi::Optional<VisitInterrupt> Visit_(const s_tir::SBlockNode* op) final {
     if (op->name_hint == "root") {
-      return StmtExprVisitor::Visit_(op);
+      return s_tir::StmtExprVisitor::Visit_(op);
     }
     buffer_access_indices_.clear();
-    TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(StmtExprVisitor::Visit_(op));
+    TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(s_tir::StmtExprVisitor::Visit_(op));
     iter_var_range_.clear();
     for (const auto& iter_var : op->iter_vars) {
       iter_var_range_.Set(iter_var->var, iter_var->dom);

include/tvm/s_tir/analysis.h

@@ -27,6 +27,7 @@
 #include <tvm/ir/module.h>
 #include <tvm/ir/prim/expr.h>
 #include <tvm/ir/transform.h>
+#include <tvm/s_tir/stmt.h>
 #include <tvm/target/target.h>
 #include <tvm/tirx/function.h>
 #include <tvm/tirx/stmt.h>
@@ -48,7 +49,7 @@ namespace tirx {
  *           - third: opaque regions
  */
 TVM_DLL ffi::Array<ffi::Array<TensorRegion>> GetSBlockAccessRegion(
-    const SBlock& block, const ffi::Map<Var, BufferVar>& buffer_var_map);
+    const s_tir::SBlock& block, const ffi::Map<Var, BufferVar>& buffer_var_map);
 
 /*!
  * \brief Auto detect the block read/write region according to its body stmt. An opaque access will
@@ -59,7 +60,7 @@ TVM_DLL ffi::Array<ffi::Array<TensorRegion>> GetSBlockAccessRegion(
  * \return An array only consisting of the read regions and write regions of the input block
  */
 TVM_DLL ffi::Array<ffi::Array<TensorRegion>> GetSBlockReadWriteRegion(
-    const SBlock& block, const ffi::Map<Var, BufferVar>& buffer_var_map);
+    const s_tir::SBlock& block, const ffi::Map<Var, BufferVar>& buffer_var_map);
 
 /*!
  * \brief Detect the lowest common ancestor(LCA) of buffer access, including both high-level
@@ -85,7 +86,7 @@ TVM_DLL ffi::Map<BufferVar, ffi::Optional<Stmt>> DetectBufferAccessLCA(const Pri
  * \param mod The input TIR module.
  * \return The anchor block if found, nullptr otherwise.
  */
-const tirx::SBlockNode* FindAnchorBlock(const IRModule& mod);
+const s_tir::SBlockNode* FindAnchorBlock(const IRModule& mod);
 
 }  // namespace tirx
 
@@ -97,6 +98,11 @@ class Analyzer;
 namespace s_tir {
 using namespace tvm::tirx;
 
+/*! \brief Verify variable/buffer definitions, load types and schedulable block boundaries. */
+TVM_DLL bool VerifyWellFormed(const tirx::PrimFunc& func, bool assert_mode = true);
+/*! \brief Verify S-TIR or mixed modules, including definitions shared across functions. */
+TVM_DLL bool VerifyWellFormed(const IRModule& mod, bool assert_mode = true);
+
 /*!
  * \brief Estimate the FLOPs of a TIR fragment.
  * \param stmt The TIR fragment to be estimated.

include/tvm/s_tir/sblock_scope.h

@@ -26,9 +26,10 @@
 #define TVM_S_TIR_SBLOCK_SCOPE_H_
 
 #include <tvm/ir/module.h>
+#include <tvm/s_tir/stmt.h>
+#include <tvm/s_tir/stmt_functor.h>
 #include <tvm/tirx/function.h>
 #include <tvm/tirx/stmt.h>
-#include <tvm/tirx/stmt_functor.h>
 
 #include <unordered_map>
 #include <utility>
@@ -41,7 +42,7 @@ namespace tirx {
  * \brief An object that refers to schedulable elements (block/for-loop) in TensorIR, aka "sref".
  *
  * Glossary
- * - SBlock sref: A StmtSRef that points to a TensorIR SBlock.
+ * - s_tir::SBlock sref: A StmtSRef that points to a TensorIR s_tir::SBlock.
  * - Loop sref: A StmtSRef that points to a TensorIR for loop.
  * - Parent sref: The parent reference of an sref is the block or loop reference to the closest
  schedulable statement. We define closest to be the nearest schedulable statement of an ancestor in
@@ -87,7 +88,7 @@ class StmtSRefNode : public ffi::Object {
    * It serves the same purpose as `ffi::ObjectRef::as`, but does not acquire strong reference to
    * `stmt`
    * \tparam StmtType The type that `this->stmt` to be downcasted to. Presumably
-   * tvm::tirx::SBlockNode or tvm::tirx::ForNode
+   * tvm::s_tir::SBlockNode or tvm::tirx::ForNode
    * \return nullptr if type check fails, otherwise the casted result for `this->stmt`
    */
   template <typename StmtType>
@@ -144,13 +145,13 @@ class StmtSRef : public ffi::ObjectRef {
   TVM_DLL static StmtSRef RootMark();
 };
 
-class SRefTreeCreator : public StmtExprVisitor {
+class SRefTreeCreator : public s_tir::StmtExprVisitor {
  public:
-  using StmtExprVisitor::Visit_;
+  using s_tir::StmtExprVisitor::Visit_;
 
   ffi::Optional<VisitInterrupt> Visit(ffi::AnyView value) override {
     if (value.as<ExprNode>()) return std::nullopt;
-    return StmtExprVisitor::Visit(value);
+    return s_tir::StmtExprVisitor::Visit(value);
   }
 
   /*!
@@ -185,7 +186,7 @@ class SRefTreeCreator : public StmtExprVisitor {
 
   ffi::Optional<VisitInterrupt> Visit_(const ForNode* loop) final;
 
-  ffi::Optional<VisitInterrupt> Visit_(const SBlockRealizeNode* realize) final;
+  ffi::Optional<VisitInterrupt> Visit_(const s_tir::SBlockRealizeNode* realize) final;
 
   ffi::Optional<VisitInterrupt> Visit_(const SeqStmtNode* seq_stmt) final;
 
@@ -251,7 +252,7 @@ class Dependency : public ffi::ObjectRef {
  * For example even leaf nodes have a scope node, even though they have no dependencies.
  *
  * Glossary:
- * - SBlock scope: A contiguous subtree of the sref tree, rooted at each SBlock sref,
+ * - s_tir::SBlock scope: A contiguous subtree of the sref tree, rooted at each s_tir::SBlock sref,
  * whose components are:
  *   - scope root: a block sref
  *   - internal srefs: loop srefs

include/tvm/s_tir/schedule/schedule.h

@@ -23,6 +23,7 @@
 #include <tvm/s_tir/random_engine.h>
 #include <tvm/s_tir/schedule/state.h>
 #include <tvm/s_tir/schedule/trace.h>
+#include <tvm/s_tir/stmt.h>
 #include <tvm/tirx/index_map.h>
 
 namespace tvm {

include/tvm/s_tir/schedule/state.h

@@ -27,6 +27,7 @@
 #include <tvm/ir/module.h>
 #include <tvm/ir/prim/expr.h>
 #include <tvm/s_tir/sblock_scope.h>
+#include <tvm/s_tir/stmt.h>
 #include <tvm/tirx/function.h>
 
 #include <unordered_map>