From a3a54d38d509d963bec1b756f14020f32a3a8afc Mon Sep 17 00:00:00 2001
From: LeiWang1999 <leiwang1999@outlook.com>
Date: Mon, 13 Jul 2026 18:44:05 +0800
Subject: [PATCH 1/3] Add runtime checks for host-evaluable assumptions

---
 src/transform/make_packed_api.cc              |  31 ++++
 src/transform/split_host_device.cc            | 139 +++++++++++++++++-
 .../language/test_tilelang_language_assume.py |  26 ++++
 ...test_tilelang_transform_make_packed_api.py |  34 +++++
 ...st_tilelang_transform_split_host_device.py |  70 +++++++++
 5 files changed, 292 insertions(+), 8 deletions(-)

diff --git a/src/transform/make_packed_api.cc b/src/transform/make_packed_api.cc
index e8a5376140..38de7aec09 100644
--- a/src/transform/make_packed_api.cc
+++ b/src/transform/make_packed_api.cc
@@ -179,6 +179,36 @@ class SubroutineCallRewriter : public StmtExprMutator {
   bool made_change_{false};
 };
 
+class AssumeToAssertRewriter : public StmtMutator {
+public:
+  static Stmt Rewrite(Stmt body) {
+    AssumeToAssertRewriter rewriter;
+    return rewriter(std::move(body));
+  }
+
+private:
+  Stmt VisitStmt_(const AttrStmtNode *op) final {
+    if (op->attr_key != tirx::attr::tilelang_assume) {
+      return StmtMutator::VisitStmt_(op);
+    }
+
+    PrimExpr condition = Downcast<PrimExpr>(op->node);
+    Array<StringImm> message_parts;
+    if (const auto *message = op->value.as<StringImmNode>()) {
+      message_parts.push_back(GetRef<StringImm>(message));
+    } else {
+      std::ostringstream os;
+      os << "Assume: " << condition;
+      message_parts.push_back(StringImm(os.str()));
+    }
+
+    Stmt assertion = AssertStmt(condition, StringImm("RuntimeError"),
+                                std::move(message_parts), op->span);
+    Stmt body = VisitStmt(op->body);
+    return SeqStmt::Flatten(std::move(assertion), std::move(body));
+  }
+};
+
 } // namespace
 
 inline Stmt MakeAssertEQ(PrimExpr lhs, PrimExpr rhs, std::string msg) {
@@ -250,6 +280,7 @@ PrimFunc MakePackedAPI(PrimFunc func) {
   }
 
   auto *func_ptr = func.CopyOnWrite();
+  func_ptr->body = AssumeToAssertRewriter::Rewrite(func_ptr->body);
   // set the global symbol to the packed function name
   const Stmt nop = Evaluate(0);
   int num_args = static_cast<int>(func_ptr->params.size());
diff --git a/src/transform/split_host_device.cc b/src/transform/split_host_device.cc
index 9097e35b76..15e7dc7af5 100644
--- a/src/transform/split_host_device.cc
+++ b/src/transform/split_host_device.cc
@@ -34,7 +34,9 @@
 #include <tvm/tirx/stmt_functor.h>
 #include <tvm/tirx/transform.h>
 
+#include <sstream>
 #include <unordered_set>
+#include <vector>
 
 #include "../op/builtin.h"
 #include "common/assume.h"
@@ -47,13 +49,14 @@ namespace tl {
 using namespace ffi;
 using namespace tirx;
 
-// This pass traverses the AST, split the target function into host part and
-// device part and copies all assume attribute statements to the device side.
+// This pass traverses the AST, splits the target function into host and device
+// parts, and preserves assumptions on every side that can evaluate them.
 
 // 1. Traverse AST and collect all assume statements into host_assumes_.
 // 2. Until the first AttrStmtNode with tvm::attr::kTarget.
 // 3. Call SplitDeviceFunc, which will create a new device function and replace
-//    the original body with a call to that function.
+//    the original body with a call to that function. Host-evaluable assumptions
+//    found in the device body are attached to that host-side call.
 class HostDeviceSplitter : public tirx::StmtMutator {
 public:
   explicit HostDeviceSplitter(IRModule *device_mod,
@@ -113,6 +116,108 @@ class HostDeviceSplitter : public tirx::StmtMutator {
   bool found_device_region() const { return found_device_region_; }
 
 private:
+  struct AssumeInfo {
+    PrimExpr condition;
+    StringImm message;
+    Span span;
+  };
+
+  class HostEvaluableConditionChecker : public tirx::StmtExprVisitor {
+  public:
+    explicit HostEvaluableConditionChecker(const Array<tirx::Var> &params) {
+      for (const tirx::Var &param : params) {
+        if (!param->dtype.is_handle()) {
+          params_.insert(param);
+        }
+      }
+    }
+
+    bool Check(const PrimExpr &condition) {
+      if (!condition->dtype.is_bool() || condition->dtype.lanes() != 1) {
+        return false;
+      }
+      VisitExpr(condition);
+      return is_host_evaluable_;
+    }
+
+  private:
+    void VisitExpr_(const tirx::VarNode *op) final {
+      if (!params_.count(GetRef<tirx::Var>(op))) {
+        is_host_evaluable_ = false;
+      }
+    }
+
+    void VisitExpr_(const tirx::BufferLoadNode *op) final {
+      is_host_evaluable_ = false;
+    }
+
+    void VisitExpr_(const tirx::CallNode *op) final {
+      // Device intrinsics and memory-dependent calls cannot be evaluated by
+      // the host launcher. Keep this deliberately conservative; ordinary
+      // shape constraints use arithmetic expression nodes rather than calls.
+      is_host_evaluable_ = false;
+    }
+
+    std::unordered_set<tirx::Var, ObjectPtrHash, ObjectPtrEqual> params_;
+    bool is_host_evaluable_{true};
+  };
+
+  class DeviceAssumeCollector : public tirx::StmtVisitor {
+  public:
+    static std::vector<AssumeInfo> Collect(const Stmt &body,
+                                           const Array<tirx::Var> &params) {
+      DeviceAssumeCollector collector(params);
+      collector(body);
+      return collector.assumes_;
+    }
+
+  private:
+    explicit DeviceAssumeCollector(const Array<tirx::Var> &params)
+        : params_(params) {}
+
+    void VisitStmt_(const tirx::AttrStmtNode *op) final {
+      if (op->attr_key == tirx::attr::tilelang_assume &&
+          conditional_scope_depth_ == 0) {
+        PrimExpr condition = Downcast<PrimExpr>(op->node);
+        HostEvaluableConditionChecker checker(params_);
+        if (checker.Check(condition)) {
+          StringImm message = [&]() {
+            if (const auto *string_imm = op->value.as<StringImmNode>()) {
+              return GetRef<StringImm>(string_imm);
+            }
+            std::ostringstream os;
+            os << "Assume: " << condition;
+            return StringImm(os.str());
+          }();
+          assumes_.push_back({condition, message, op->span});
+        }
+      }
+      tirx::StmtVisitor::VisitStmt_(op);
+    }
+
+    void VisitStmt_(const tirx::IfThenElseNode *op) final {
+      ++conditional_scope_depth_;
+      tirx::StmtVisitor::VisitStmt_(op);
+      --conditional_scope_depth_;
+    }
+
+    void VisitStmt_(const tirx::ForNode *op) final {
+      ++conditional_scope_depth_;
+      tirx::StmtVisitor::VisitStmt_(op);
+      --conditional_scope_depth_;
+    }
+
+    void VisitStmt_(const tirx::WhileNode *op) final {
+      ++conditional_scope_depth_;
+      tirx::StmtVisitor::VisitStmt_(op);
+      --conditional_scope_depth_;
+    }
+
+    Array<tirx::Var> params_;
+    std::vector<AssumeInfo> assumes_;
+    int conditional_scope_depth_{0};
+  };
+
   bool found_device_region_{false};
   Map<tirx::Var, tirx::Buffer> host_buffer_map_;
   Array<tirx::Var> non_restrict_params_;
@@ -266,6 +371,15 @@ class HostDeviceSplitter : public tirx::StmtMutator {
     return body;
   }
 
+  static Stmt
+  WrapBodyWithDeviceAssumes(Stmt body, const std::vector<AssumeInfo> &assumes) {
+    for (auto it = assumes.rbegin(); it != assumes.rend(); ++it) {
+      body = AttrStmt(it->condition, tirx::attr::tilelang_assume, it->message,
+                      std::move(body), it->span);
+    }
+    return body;
+  }
+
   tirx::Stmt SplitDeviceFunc(tirx::Stmt body, tvm::Target device_target) {
     code_block_source_ = std::nullopt;
     code_block_entry_name_ = std::nullopt;
@@ -293,6 +407,15 @@ class HostDeviceSplitter : public tirx::StmtMutator {
               use_def.undefined_buffers_};
     }();
 
+    // Assumes written inside a device region are otherwise removed from the
+    // host function when the region is replaced by a kernel launch. Preserve
+    // conditions that the host can evaluate from scalar kernel parameters so
+    // MakePackedAPI can turn them into runtime assertions. Assumes involving
+    // thread/block variables, local bindings, buffer loads, or calls remain
+    // device-only optimization facts.
+    std::vector<AssumeInfo> host_evaluable_device_assumes =
+        DeviceAssumeCollector::Collect(body, old_params);
+
     // Create new parameter variables for the device function to avoid sharing
     // Var objects with the host function. This prevents ConvertSSA from
     // incorrectly renaming variables when it processes multiple functions.
@@ -397,6 +520,7 @@ class HostDeviceSplitter : public tirx::StmtMutator {
     Array<PrimExpr> args =
         old_params.Map([](const tirx::Var &var) -> PrimExpr { return var; });
 
+    Stmt host_call;
     if (can_propagate_errors) {
       tirx::Var kernel_error_code("kernel_error_code", success->dtype);
       tirx::Call kernel_call(success->dtype, kernel_symbol_global, args);
@@ -404,15 +528,14 @@ class HostDeviceSplitter : public tirx::StmtMutator {
           kernel_error_code == success, tirx::StringImm("RuntimeError"),
           Array<tirx::StringImm>(
               {tirx::StringImm("Error executing compute kernel")}));
-      tirx::Stmt let_check = tirx::SeqStmt(
+      host_call = tirx::SeqStmt(
           {tirx::Bind(kernel_error_code, kernel_call), assert_success});
-
-      return let_check;
-
     } else {
-      return tirx::Evaluate(
+      host_call = tirx::Evaluate(
           tirx::Call(DataType::Void(), kernel_symbol_global, args));
     }
+    return WrapBodyWithDeviceAssumes(std::move(host_call),
+                                     host_evaluable_device_assumes);
   }
 
   // target ir module
diff --git a/testing/python/language/test_tilelang_language_assume.py b/testing/python/language/test_tilelang_language_assume.py
index 691a7f32a2..35f765819c 100644
--- a/testing/python/language/test_tilelang_language_assume.py
+++ b/testing/python/language/test_tilelang_language_assume.py
@@ -1,3 +1,6 @@
+import pytest
+import torch
+
 import tilelang
 import tilelang.language as T
 import tilelang.testing
@@ -105,5 +108,28 @@ def main(A: T.Tensor((N,), T.float32), out: T.Tensor((1,), T.float32), l: T.int3
     assert "if (" not in source
 
 
+@tilelang.testing.requires_cuda
+def test_host_evaluable_assume_is_checked_at_runtime():
+    @tilelang.jit
+    def kernel_with_runtime_check():
+        N = T.dynamic("N")
+
+        @T.prim_func
+        def main(A: T.Tensor((N,), T.float32)):
+            with T.Kernel(1, threads=32):
+                T.assume(N % 4 == 0)
+                tx = T.get_thread_binding()
+                if tx < N:
+                    A[tx] = 1.0
+
+        return main
+
+    jit_kernel = kernel_with_runtime_check()
+    jit_kernel(torch.empty(4, device="cuda"))
+
+    with pytest.raises(RuntimeError, match="Assume: N % 4 == 0"):
+        jit_kernel(torch.empty(2, device="cuda"))
+
+
 if __name__ == "__main__":
     tilelang.testing.main()
diff --git a/testing/python/transform/test_tilelang_transform_make_packed_api.py b/testing/python/transform/test_tilelang_transform_make_packed_api.py
index 58d187a5e9..51811ef479 100644
--- a/testing/python/transform/test_tilelang_transform_make_packed_api.py
+++ b/testing/python/transform/test_tilelang_transform_make_packed_api.py
@@ -51,6 +51,17 @@ def _visitor(stmt):
     return result
 
 
+def _collect_nodes(stmt, node_type):
+    nodes = []
+
+    def _visitor(node):
+        if isinstance(node, node_type):
+            nodes.append(node)
+
+    tvm.tirx.stmt_functor.post_order_visit(stmt, _visitor)
+    return nodes
+
+
 @pytest.mark.parametrize("use_global_symbol", [False])
 def test_no_op_when_global_symbol_is_absent(use_global_symbol):
     func_attr = {"target": tvm.target.Target("llvm", host="llvm")}
@@ -125,6 +136,29 @@ def subroutine(A_data: T.handle("float32")):
     )
 
 
+def test_assume_is_lowered_to_runtime_assert():
+    n = tirx.Var("n", "int32")
+    condition = n % 4 == 0
+    message = "n must be divisible by 4"
+    body = tirx.AttrStmt(
+        condition,
+        "tl.assume",
+        tirx.StringImm(message),
+        tirx.Evaluate(0),
+    )
+    before = tirx.PrimFunc([n], body).with_attr("global_symbol", "main")
+    before = before.with_attr("target", tvm.target.Target("cuda", host="llvm"))
+
+    after = tilelang.transform.MakePackedAPI()(tvm.IRModule.from_expr(before))["main"]
+    assert not [stmt for stmt in _collect_nodes(after.body, tirx.AttrStmt) if stmt.attr_key == "tl.assume"]
+
+    matching_asserts = [
+        stmt for stmt in _collect_nodes(after.body, tirx.AssertStmt) if any(part.value == message for part in stmt.message_parts)
+    ]
+    assert len(matching_asserts) == 1
+    tvm.ir.assert_structural_equal(matching_asserts[0].condition, condition, map_free_vars=True)
+
+
 def test_subroutine_call_to_externally_visible_subroutine():
     """Externally-visible subroutines should use the PackedFunc API
 
diff --git a/testing/python/transform/test_tilelang_transform_split_host_device.py b/testing/python/transform/test_tilelang_transform_split_host_device.py
index 89779c639d..e2bdc604d9 100644
--- a/testing/python/transform/test_tilelang_transform_split_host_device.py
+++ b/testing/python/transform/test_tilelang_transform_split_host_device.py
@@ -59,6 +59,18 @@ def collect_vars_from_expr(expr):
     return assume_vars
 
 
+def collect_assume_conditions(func: tvm.tirx.PrimFunc):
+    """Collect conditions stored in tl.assume attributes."""
+    conditions = []
+
+    def collect_assumes(stmt):
+        if isinstance(stmt, tirx.AttrStmt) and stmt.attr_key == "tl.assume":
+            conditions.append(stmt.node)
+
+    tirx.stmt_functor.post_order_visit(func.body, collect_assumes)
+    return conditions
+
+
 def get_var_name(var):
     """Get the name of a Var, handling different TVM versions."""
     if hasattr(var, "name_hint"):
@@ -250,5 +262,63 @@ def main(a: T.Tensor[(n,), T.int32]):
                 )
 
 
+@tilelang.testing.requires_cuda
+def test_split_host_device_lifts_host_evaluable_device_assume():
+    @T.prim_func
+    def main(a: T.Tensor[(1,), T.int32], n: T.int32):
+        with T.Kernel(1, threads=128):
+            T.assume(n % 4 == 0)
+            a[0] = n
+
+    mod = run_split_host_device_passes(main)
+    host_func = get_host_func(mod)
+    device_func = get_device_func(mod)
+    assert host_func is not None
+    assert device_func is not None
+
+    expected = main.params[1] % 4 == 0
+    assert any(tvm.ir.structural_equal(condition, expected) for condition in collect_assume_conditions(host_func))
+    assert any(tvm.ir.structural_equal(condition, expected, map_free_vars=True) for condition in collect_assume_conditions(device_func))
+
+
+@tilelang.testing.requires_cuda
+def test_split_host_device_keeps_device_local_assume_off_host():
+    @T.prim_func
+    def main(a: T.Tensor[(1,), T.int32]):
+        with T.Kernel(1, threads=128):
+            tx = T.get_thread_binding()
+            T.assume(tx < 128)
+            if tx == 0:
+                a[0] = tx
+
+    mod = run_split_host_device_passes(main)
+    host_func = get_host_func(mod)
+    device_func = get_device_func(mod)
+    assert host_func is not None
+    assert device_func is not None
+
+    assert not collect_assume_conditions(host_func)
+    assert collect_assume_conditions(device_func)
+
+
+@tilelang.testing.requires_cuda
+def test_split_host_device_does_not_lift_conditional_device_assume():
+    @T.prim_func
+    def main(a: T.Tensor[(1,), T.int32], n: T.int32):
+        with T.Kernel(1, threads=128):
+            if n % 4 == 0:
+                T.assume(n > 0)
+                a[0] = n
+
+    mod = run_split_host_device_passes(main)
+    host_func = get_host_func(mod)
+    device_func = get_device_func(mod)
+    assert host_func is not None
+    assert device_func is not None
+
+    assert not collect_assume_conditions(host_func)
+    assert collect_assume_conditions(device_func)
+
+
 if __name__ == "__main__":
     tilelang.testing.main()

From f09faa357c3e7839a0b91f727247e78646c9e4e4 Mon Sep 17 00:00:00 2001
From: LeiWang1999 <leiwang1999@outlook.com>
Date: Mon, 13 Jul 2026 19:56:45 +0800
Subject: [PATCH 2/3] Preserve nullable buffer semantics for assume checks

---
 src/transform/common/attr.h                   |  3 +
 src/transform/inject_assumes.cc               | 25 +++---
 src/transform/make_packed_api.cc              | 67 +++++++++++++---
 src/transform/split_host_device.cc            | 77 ++++++++++++-------
 ...test_tilelang_transform_make_packed_api.py | 44 ++++++++++-
 ...st_tilelang_transform_split_host_device.py | 30 ++++++--
 6 files changed, 188 insertions(+), 58 deletions(-)

diff --git a/src/transform/common/attr.h b/src/transform/common/attr.h
index a8c25455b3..fa55c0349e 100644
--- a/src/transform/common/attr.h
+++ b/src/transform/common/attr.h
@@ -33,6 +33,9 @@ constexpr const char *kHasGridSync = "has_cuda_pdl_sync";
 constexpr const char *volatile_scope = "volatile_scope";
 constexpr const char *coproc_scope = "coproc_scope";
 constexpr const char *pipeline_exec_scope = "pipeline_exec_scope";
+// Marks user-authored assumptions that require a host runtime check. The
+// corresponding tl.assume remains in the IR as an optimizer fact.
+constexpr const char *kAssumeRuntimeCheck = "tl.assume_runtime_check";
 
 // Attributes to implement SourceCodeBlock
 constexpr const char *kCodeBlockSource = "code_block_source";
diff --git a/src/transform/inject_assumes.cc b/src/transform/inject_assumes.cc
index 75aa02204b..a98f7babaa 100644
--- a/src/transform/inject_assumes.cc
+++ b/src/transform/inject_assumes.cc
@@ -1,7 +1,7 @@
 /*!
  * \file inject_assumes.cc
  * \brief Inject assumes on buffer's shape boundary check. Also convert
- * existing assumes to AttrNodes.
+ * user-authored assumes to optimizer and runtime-check AttrNodes.
  */
 
 #include "common/assume.h"
@@ -173,14 +173,12 @@ class AssumeInjector : public tvm::tirx::StmtExprMutator {
       //    Stmt1
       //    Stmt2
       //    T.assume(cond2)
-      // This SeqStmt will be converted to:
-      //    With(attr::tilelang_assume, cond1) {
-      //      Stmt1
-      //      Stmt2
-      //    }
-      //    With(attr::tilelang_assume, cond2) {
-      //      ...
-      //    }
+      // Each user-authored assume is converted into two nested attributes:
+      // the outer runtime-check marker is consumed by MakePackedAPI, while the
+      // inner tilelang_assume remains available to optimization passes. Buffer
+      // shape/stride assumptions created by AssumeCreator intentionally do not
+      // receive the runtime-check marker because nullable buffers bind missing
+      // symbolic shapes to zero.
       if (auto e = GetAssumeExprInEvaluateForm(stmt)) {
         groups.push_back(AssumeGroup{*e, {}});
       } else {
@@ -193,9 +191,12 @@ class AssumeInjector : public tvm::tirx::StmtExprMutator {
         Stmt body = g.stmts.size() == 1 ? g.stmts[0] : SeqStmt(g.stmts);
         std::stringstream ss;
         ss << "Assume: " << *(g.e);
-        AttrStmt attr = AttrStmt(*g.e, tirx::attr::tilelang_assume,
-                                 StringImm(ss.str()), body);
-        groups[i - 1].stmts.push_back(attr);
+        StringImm message(ss.str());
+        body = AttrStmt(*g.e, tirx::attr::tilelang_assume, message,
+                        std::move(body));
+        body = AttrStmt(*g.e, tl::attr::kAssumeRuntimeCheck, message,
+                        std::move(body));
+        groups[i - 1].stmts.push_back(std::move(body));
       } else {
         ICHECK(i == 0) << "only the first group can have no assume";
       }
diff --git a/src/transform/make_packed_api.cc b/src/transform/make_packed_api.cc
index 38de7aec09..00e8f23532 100644
--- a/src/transform/make_packed_api.cc
+++ b/src/transform/make_packed_api.cc
@@ -179,16 +179,27 @@ class SubroutineCallRewriter : public StmtExprMutator {
   bool made_change_{false};
 };
 
-class AssumeToAssertRewriter : public StmtMutator {
+struct AssumeRuntimeCheck {
+  PrimExpr condition;
+  Array<StringImm> message_parts;
+  Span span;
+};
+
+class AssumeRuntimeCheckExtractor : public StmtMutator {
 public:
-  static Stmt Rewrite(Stmt body) {
-    AssumeToAssertRewriter rewriter;
-    return rewriter(std::move(body));
+  static Stmt Extract(Stmt body,
+                      std::vector<AssumeRuntimeCheck> *runtime_checks) {
+    AssumeRuntimeCheckExtractor extractor(runtime_checks);
+    return extractor(std::move(body));
   }
 
 private:
+  explicit AssumeRuntimeCheckExtractor(
+      std::vector<AssumeRuntimeCheck> *runtime_checks)
+      : runtime_checks_(runtime_checks) {}
+
   Stmt VisitStmt_(const AttrStmtNode *op) final {
-    if (op->attr_key != tirx::attr::tilelang_assume) {
+    if (op->attr_key != tl::attr::kAssumeRuntimeCheck) {
       return StmtMutator::VisitStmt_(op);
     }
 
@@ -202,11 +213,39 @@ class AssumeToAssertRewriter : public StmtMutator {
       message_parts.push_back(StringImm(os.str()));
     }
 
-    Stmt assertion = AssertStmt(condition, StringImm("RuntimeError"),
-                                std::move(message_parts), op->span);
-    Stmt body = VisitStmt(op->body);
-    return SeqStmt::Flatten(std::move(assertion), std::move(body));
+    if (conditional_scope_depth_ != 0) {
+      // Conditional checks cannot be moved to the packed-function entry
+      // without changing their execution semantics. Keep only the paired
+      // tl.assume, matching SplitHostDevice's conservative lifting policy.
+      return VisitStmt(op->body);
+    }
+    runtime_checks_->push_back({condition, std::move(message_parts), op->span});
+    return VisitStmt(op->body);
+  }
+
+  Stmt VisitStmt_(const IfThenElseNode *op) final {
+    ++conditional_scope_depth_;
+    Stmt result = StmtMutator::VisitStmt_(op);
+    --conditional_scope_depth_;
+    return result;
   }
+
+  Stmt VisitStmt_(const ForNode *op) final {
+    ++conditional_scope_depth_;
+    Stmt result = StmtMutator::VisitStmt_(op);
+    --conditional_scope_depth_;
+    return result;
+  }
+
+  Stmt VisitStmt_(const WhileNode *op) final {
+    ++conditional_scope_depth_;
+    Stmt result = StmtMutator::VisitStmt_(op);
+    --conditional_scope_depth_;
+    return result;
+  }
+
+  std::vector<AssumeRuntimeCheck> *runtime_checks_;
+  int conditional_scope_depth_{0};
 };
 
 } // namespace
@@ -280,7 +319,9 @@ PrimFunc MakePackedAPI(PrimFunc func) {
   }
 
   auto *func_ptr = func.CopyOnWrite();
-  func_ptr->body = AssumeToAssertRewriter::Rewrite(func_ptr->body);
+  std::vector<AssumeRuntimeCheck> assume_runtime_checks;
+  func_ptr->body = AssumeRuntimeCheckExtractor::Extract(func_ptr->body,
+                                                        &assume_runtime_checks);
   // set the global symbol to the packed function name
   const Stmt nop = Evaluate(0);
   int num_args = static_cast<int>(func_ptr->params.size());
@@ -576,6 +617,12 @@ PrimFunc MakePackedAPI(PrimFunc func) {
         symbol::tvm_ffi_symbol_prefix + global_symbol.value()}});
 
   Stmt body = ReturnRewriter(v_result)(func_ptr->body);
+  for (auto it = assume_runtime_checks.rbegin();
+       it != assume_runtime_checks.rend(); ++it) {
+    Stmt assertion = AssertStmt(it->condition, StringImm("RuntimeError"),
+                                it->message_parts, it->span);
+    body = SeqStmt::Flatten(std::move(assertion), std::move(body));
+  }
   body = AttrStmt(make_zero(DataType::Int(32)), tirx::attr::compute_scope,
                   StringImm(name_hint + "_compute_"), body);
   // Set device context
diff --git a/src/transform/split_host_device.cc b/src/transform/split_host_device.cc
index 15e7dc7af5..5a3898ade6 100644
--- a/src/transform/split_host_device.cc
+++ b/src/transform/split_host_device.cc
@@ -52,11 +52,11 @@ using namespace tirx;
 // This pass traverses the AST, splits the target function into host and device
 // parts, and preserves assumptions on every side that can evaluate them.
 
-// 1. Traverse AST and collect all assume statements into host_assumes_.
+// 1. Traverse AST and collect all optimizer assumptions into host_assumes_.
 // 2. Until the first AttrStmtNode with tvm::attr::kTarget.
 // 3. Call SplitDeviceFunc, which will create a new device function and replace
-//    the original body with a call to that function. Host-evaluable assumptions
-//    found in the device body are attached to that host-side call.
+//    the original body with a call to that function. Host-evaluable runtime
+//    checks found in the device body are attached to that host-side call.
 class HostDeviceSplitter : public tirx::StmtMutator {
 public:
   explicit HostDeviceSplitter(IRModule *device_mod,
@@ -95,7 +95,7 @@ class HostDeviceSplitter : public tirx::StmtMutator {
       // We first push back the outside assume, then visit the child.
       // So when moving assumes to device side, we need to do the building
       // process in a reverse order.
-      host_assumes_.push_back(op);
+      host_assumes_.push_back(GetRef<tirx::AttrStmt>(op));
     }
     return tirx::StmtMutator::VisitStmt_(op);
   }
@@ -116,7 +116,7 @@ class HostDeviceSplitter : public tirx::StmtMutator {
   bool found_device_region() const { return found_device_region_; }
 
 private:
-  struct AssumeInfo {
+  struct RuntimeCheckInfo {
     PrimExpr condition;
     StringImm message;
     Span span;
@@ -162,21 +162,21 @@ class HostDeviceSplitter : public tirx::StmtMutator {
     bool is_host_evaluable_{true};
   };
 
-  class DeviceAssumeCollector : public tirx::StmtVisitor {
+  class DeviceRuntimeCheckCollector : public tirx::StmtVisitor {
   public:
-    static std::vector<AssumeInfo> Collect(const Stmt &body,
-                                           const Array<tirx::Var> &params) {
-      DeviceAssumeCollector collector(params);
+    static std::vector<RuntimeCheckInfo>
+    Collect(const Stmt &body, const Array<tirx::Var> &params) {
+      DeviceRuntimeCheckCollector collector(params);
       collector(body);
-      return collector.assumes_;
+      return collector.runtime_checks_;
     }
 
   private:
-    explicit DeviceAssumeCollector(const Array<tirx::Var> &params)
+    explicit DeviceRuntimeCheckCollector(const Array<tirx::Var> &params)
         : params_(params) {}
 
     void VisitStmt_(const tirx::AttrStmtNode *op) final {
-      if (op->attr_key == tirx::attr::tilelang_assume &&
+      if (op->attr_key == tl::attr::kAssumeRuntimeCheck &&
           conditional_scope_depth_ == 0) {
         PrimExpr condition = Downcast<PrimExpr>(op->node);
         HostEvaluableConditionChecker checker(params_);
@@ -189,7 +189,7 @@ class HostDeviceSplitter : public tirx::StmtMutator {
             os << "Assume: " << condition;
             return StringImm(os.str());
           }();
-          assumes_.push_back({condition, message, op->span});
+          runtime_checks_.push_back({condition, message, op->span});
         }
       }
       tirx::StmtVisitor::VisitStmt_(op);
@@ -214,10 +214,26 @@ class HostDeviceSplitter : public tirx::StmtMutator {
     }
 
     Array<tirx::Var> params_;
-    std::vector<AssumeInfo> assumes_;
+    std::vector<RuntimeCheckInfo> runtime_checks_;
     int conditional_scope_depth_{0};
   };
 
+  class RuntimeCheckMarkerRemover : public tirx::StmtMutator {
+  public:
+    static Stmt Remove(Stmt body) {
+      RuntimeCheckMarkerRemover remover;
+      return remover(std::move(body));
+    }
+
+  private:
+    Stmt VisitStmt_(const tirx::AttrStmtNode *op) final {
+      if (op->attr_key == tl::attr::kAssumeRuntimeCheck) {
+        return VisitStmt(op->body);
+      }
+      return tirx::StmtMutator::VisitStmt_(op);
+    }
+  };
+
   bool found_device_region_{false};
   Map<tirx::Var, tirx::Buffer> host_buffer_map_;
   Array<tirx::Var> non_restrict_params_;
@@ -371,11 +387,13 @@ class HostDeviceSplitter : public tirx::StmtMutator {
     return body;
   }
 
-  static Stmt
-  WrapBodyWithDeviceAssumes(Stmt body, const std::vector<AssumeInfo> &assumes) {
-    for (auto it = assumes.rbegin(); it != assumes.rend(); ++it) {
+  static Stmt WrapBodyWithRuntimeChecks(
+      Stmt body, const std::vector<RuntimeCheckInfo> &runtime_checks) {
+    for (auto it = runtime_checks.rbegin(); it != runtime_checks.rend(); ++it) {
       body = AttrStmt(it->condition, tirx::attr::tilelang_assume, it->message,
                       std::move(body), it->span);
+      body = AttrStmt(it->condition, tl::attr::kAssumeRuntimeCheck, it->message,
+                      std::move(body), it->span);
     }
     return body;
   }
@@ -407,14 +425,18 @@ class HostDeviceSplitter : public tirx::StmtMutator {
               use_def.undefined_buffers_};
     }();
 
-    // Assumes written inside a device region are otherwise removed from the
-    // host function when the region is replaced by a kernel launch. Preserve
-    // conditions that the host can evaluate from scalar kernel parameters so
-    // MakePackedAPI can turn them into runtime assertions. Assumes involving
-    // thread/block variables, local bindings, buffer loads, or calls remain
-    // device-only optimization facts.
-    std::vector<AssumeInfo> host_evaluable_device_assumes =
-        DeviceAssumeCollector::Collect(body, old_params);
+    // Runtime-check markers written inside a device region are otherwise
+    // removed from the host function when the region is replaced by a kernel
+    // launch. Preserve checks that the host can evaluate from scalar kernel
+    // parameters so MakePackedAPI can turn them into runtime assertions.
+    // Checks involving thread/block variables, local bindings, buffer loads,
+    // calls, or conditional scopes remain device-only optimizer assumptions.
+    std::vector<RuntimeCheckInfo> host_runtime_checks =
+        DeviceRuntimeCheckCollector::Collect(body, old_params);
+
+    // Runtime checks execute in the host launcher. Their paired tl.assume
+    // attributes remain in the device body as optimizer facts.
+    body = RuntimeCheckMarkerRemover::Remove(std::move(body));
 
     // Create new parameter variables for the device function to avoid sharing
     // Var objects with the host function. This prevents ConvertSSA from
@@ -534,8 +556,7 @@ class HostDeviceSplitter : public tirx::StmtMutator {
       host_call = tirx::Evaluate(
           tirx::Call(DataType::Void(), kernel_symbol_global, args));
     }
-    return WrapBodyWithDeviceAssumes(std::move(host_call),
-                                     host_evaluable_device_assumes);
+    return WrapBodyWithRuntimeChecks(std::move(host_call), host_runtime_checks);
   }
 
   // target ir module
@@ -543,7 +564,7 @@ class HostDeviceSplitter : public tirx::StmtMutator {
   // Generate new GlobalVar for the kernel
   std::function<GlobalVar()> var_supply_;
   // Collect assumes in host side
-  Array<const tirx::AttrStmtNode *> host_assumes_;
+  Array<tirx::AttrStmt> host_assumes_;
 };
 
 tirx::PrimFunc SplitHostDevice(tirx::PrimFunc func, IRModule *device_mod,
diff --git a/testing/python/transform/test_tilelang_transform_make_packed_api.py b/testing/python/transform/test_tilelang_transform_make_packed_api.py
index 51811ef479..4ce3f16e7d 100644
--- a/testing/python/transform/test_tilelang_transform_make_packed_api.py
+++ b/testing/python/transform/test_tilelang_transform_make_packed_api.py
@@ -136,7 +136,7 @@ def subroutine(A_data: T.handle("float32")):
     )
 
 
-def test_assume_is_lowered_to_runtime_assert():
+def test_assume_runtime_check_is_lowered_to_assert():
     n = tirx.Var("n", "int32")
     condition = n % 4 == 0
     message = "n must be divisible by 4"
@@ -146,11 +146,25 @@ def test_assume_is_lowered_to_runtime_assert():
         tirx.StringImm(message),
         tirx.Evaluate(0),
     )
+    body = tirx.AttrStmt(
+        condition,
+        "tl.assume_runtime_check",
+        tirx.StringImm(message),
+        body,
+    )
+    body = tirx.AttrStmt(
+        condition,
+        "tl.assume",
+        tirx.StringImm("compiler-generated assumption"),
+        body,
+    )
     before = tirx.PrimFunc([n], body).with_attr("global_symbol", "main")
     before = before.with_attr("target", tvm.target.Target("cuda", host="llvm"))
 
-    after = tilelang.transform.MakePackedAPI()(tvm.IRModule.from_expr(before))["main"]
-    assert not [stmt for stmt in _collect_nodes(after.body, tirx.AttrStmt) if stmt.attr_key == "tl.assume"]
+    after_mod = tilelang.transform.MakePackedAPI()(tvm.IRModule.from_expr(before))
+    after = after_mod["main"]
+    assert not [stmt for stmt in _collect_nodes(after.body, tirx.AttrStmt) if stmt.attr_key == "tl.assume_runtime_check"]
+    assert [stmt for stmt in _collect_nodes(after.body, tirx.AttrStmt) if stmt.attr_key == "tl.assume"]
 
     matching_asserts = [
         stmt for stmt in _collect_nodes(after.body, tirx.AssertStmt) if any(part.value == message for part in stmt.message_parts)
@@ -158,6 +172,30 @@ def test_assume_is_lowered_to_runtime_assert():
     assert len(matching_asserts) == 1
     tvm.ir.assert_structural_equal(matching_asserts[0].condition, condition, map_free_vars=True)
 
+    simplified = tilelang.transform.Simplify()(after_mod)["main"]
+    matching_asserts = [
+        stmt for stmt in _collect_nodes(simplified.body, tirx.AssertStmt) if any(part.value == message for part in stmt.message_parts)
+    ]
+    assert len(matching_asserts) == 1
+    tvm.ir.assert_structural_equal(matching_asserts[0].condition, condition, map_free_vars=True)
+
+
+def test_optimizer_only_assume_is_not_lowered_to_assert():
+    n = tirx.Var("n", "int32")
+    message = "compiler-generated shape assumption"
+    body = tirx.AttrStmt(
+        n > 0,
+        "tl.assume",
+        tirx.StringImm(message),
+        tirx.Evaluate(0),
+    )
+    before = tirx.PrimFunc([n], body).with_attr("global_symbol", "main")
+    before = before.with_attr("target", tvm.target.Target("cuda", host="llvm"))
+
+    after = tilelang.transform.MakePackedAPI()(tvm.IRModule.from_expr(before))["main"]
+    assert [stmt for stmt in _collect_nodes(after.body, tirx.AttrStmt) if stmt.attr_key == "tl.assume"]
+    assert not [stmt for stmt in _collect_nodes(after.body, tirx.AssertStmt) if any(part.value == message for part in stmt.message_parts)]
+
 
 def test_subroutine_call_to_externally_visible_subroutine():
     """Externally-visible subroutines should use the PackedFunc API
diff --git a/testing/python/transform/test_tilelang_transform_split_host_device.py b/testing/python/transform/test_tilelang_transform_split_host_device.py
index e2bdc604d9..039392e807 100644
--- a/testing/python/transform/test_tilelang_transform_split_host_device.py
+++ b/testing/python/transform/test_tilelang_transform_split_host_device.py
@@ -59,18 +59,26 @@ def collect_vars_from_expr(expr):
     return assume_vars
 
 
-def collect_assume_conditions(func: tvm.tirx.PrimFunc):
-    """Collect conditions stored in tl.assume attributes."""
+def collect_attr_conditions(func: tvm.tirx.PrimFunc, attr_key: str):
+    """Collect conditions stored in attributes with the requested key."""
     conditions = []
 
-    def collect_assumes(stmt):
-        if isinstance(stmt, tirx.AttrStmt) and stmt.attr_key == "tl.assume":
+    def collect_attrs(stmt):
+        if isinstance(stmt, tirx.AttrStmt) and stmt.attr_key == attr_key:
             conditions.append(stmt.node)
 
-    tirx.stmt_functor.post_order_visit(func.body, collect_assumes)
+    tirx.stmt_functor.post_order_visit(func.body, collect_attrs)
     return conditions
 
 
+def collect_assume_conditions(func: tvm.tirx.PrimFunc):
+    return collect_attr_conditions(func, "tl.assume")
+
+
+def collect_runtime_check_conditions(func: tvm.tirx.PrimFunc):
+    return collect_attr_conditions(func, "tl.assume_runtime_check")
+
+
 def get_var_name(var):
     """Get the name of a Var, handling different TVM versions."""
     if hasattr(var, "name_hint"):
@@ -119,6 +127,8 @@ def main(a: T.Tensor[(n,), T.int32]):
 
     assert device_func is not None, "Device function not found"
     assert host_func is not None, "Host function not found"
+    assert collect_runtime_check_conditions(host_func)
+    assert not collect_runtime_check_conditions(device_func)
 
     # Check that device function has assume statements
     device_str = str(device_func)
@@ -156,8 +166,12 @@ def main(a: T.Tensor[(n, m), T.float32]):
 
     mod = run_split_host_device_passes(main)
 
+    host_func = get_host_func(mod)
     device_func = get_device_func(mod)
+    assert host_func is not None
     assert device_func is not None
+    assert not collect_runtime_check_conditions(host_func)
+    assert not collect_runtime_check_conditions(device_func)
 
     # Check that assumes exist
     device_str = str(device_func)
@@ -279,6 +293,8 @@ def main(a: T.Tensor[(1,), T.int32], n: T.int32):
     expected = main.params[1] % 4 == 0
     assert any(tvm.ir.structural_equal(condition, expected) for condition in collect_assume_conditions(host_func))
     assert any(tvm.ir.structural_equal(condition, expected, map_free_vars=True) for condition in collect_assume_conditions(device_func))
+    assert any(tvm.ir.structural_equal(condition, expected) for condition in collect_runtime_check_conditions(host_func))
+    assert not collect_runtime_check_conditions(device_func)
 
 
 @tilelang.testing.requires_cuda
@@ -299,6 +315,8 @@ def main(a: T.Tensor[(1,), T.int32]):
 
     assert not collect_assume_conditions(host_func)
     assert collect_assume_conditions(device_func)
+    assert not collect_runtime_check_conditions(host_func)
+    assert not collect_runtime_check_conditions(device_func)
 
 
 @tilelang.testing.requires_cuda
@@ -318,6 +336,8 @@ def main(a: T.Tensor[(1,), T.int32], n: T.int32):
 
     assert not collect_assume_conditions(host_func)
     assert collect_assume_conditions(device_func)
+    assert not collect_runtime_check_conditions(host_func)
+    assert not collect_runtime_check_conditions(device_func)
 
 
 if __name__ == "__main__":

From ae7315ae969e906a041b759dbf9eac65f9be8409 Mon Sep 17 00:00:00 2001
From: LeiWang1999 <leiwang1999@outlook.com>
Date: Tue, 14 Jul 2026 11:34:38 +0800
Subject: [PATCH 3/3] Rename assume runtime-check marker

---
 src/transform/common/attr.h                               | 3 ++-
 src/transform/inject_assumes.cc                           | 2 +-
 src/transform/make_packed_api.cc                          | 2 +-
 src/transform/split_host_device.cc                        | 8 ++++----
 .../transform/test_tilelang_transform_make_packed_api.py  | 6 +++---
 .../test_tilelang_transform_split_host_device.py          | 2 +-
 6 files changed, 12 insertions(+), 11 deletions(-)

diff --git a/src/transform/common/attr.h b/src/transform/common/attr.h
index fa55c0349e..a02e4a1940 100644
--- a/src/transform/common/attr.h
+++ b/src/transform/common/attr.h
@@ -35,7 +35,8 @@ constexpr const char *coproc_scope = "coproc_scope";
 constexpr const char *pipeline_exec_scope = "pipeline_exec_scope";
 // Marks user-authored assumptions that require a host runtime check. The
 // corresponding tl.assume remains in the IR as an optimizer fact.
-constexpr const char *kAssumeRuntimeCheck = "tl.assume_runtime_check";
+constexpr const char *kAssumeRequiresRuntimeCheck =
+    "tl.assume_requires_runtime_check";
 
 // Attributes to implement SourceCodeBlock
 constexpr const char *kCodeBlockSource = "code_block_source";
diff --git a/src/transform/inject_assumes.cc b/src/transform/inject_assumes.cc
index a98f7babaa..9fb6eeeaef 100644
--- a/src/transform/inject_assumes.cc
+++ b/src/transform/inject_assumes.cc
@@ -194,7 +194,7 @@ class AssumeInjector : public tvm::tirx::StmtExprMutator {
         StringImm message(ss.str());
         body = AttrStmt(*g.e, tirx::attr::tilelang_assume, message,
                         std::move(body));
-        body = AttrStmt(*g.e, tl::attr::kAssumeRuntimeCheck, message,
+        body = AttrStmt(*g.e, tl::attr::kAssumeRequiresRuntimeCheck, message,
                         std::move(body));
         groups[i - 1].stmts.push_back(std::move(body));
       } else {
diff --git a/src/transform/make_packed_api.cc b/src/transform/make_packed_api.cc
index 00e8f23532..9593403aa5 100644
--- a/src/transform/make_packed_api.cc
+++ b/src/transform/make_packed_api.cc
@@ -199,7 +199,7 @@ class AssumeRuntimeCheckExtractor : public StmtMutator {
       : runtime_checks_(runtime_checks) {}
 
   Stmt VisitStmt_(const AttrStmtNode *op) final {
-    if (op->attr_key != tl::attr::kAssumeRuntimeCheck) {
+    if (op->attr_key != tl::attr::kAssumeRequiresRuntimeCheck) {
       return StmtMutator::VisitStmt_(op);
     }
 
diff --git a/src/transform/split_host_device.cc b/src/transform/split_host_device.cc
index 5a3898ade6..d24e4cf520 100644
--- a/src/transform/split_host_device.cc
+++ b/src/transform/split_host_device.cc
@@ -176,7 +176,7 @@ class HostDeviceSplitter : public tirx::StmtMutator {
         : params_(params) {}
 
     void VisitStmt_(const tirx::AttrStmtNode *op) final {
-      if (op->attr_key == tl::attr::kAssumeRuntimeCheck &&
+      if (op->attr_key == tl::attr::kAssumeRequiresRuntimeCheck &&
           conditional_scope_depth_ == 0) {
         PrimExpr condition = Downcast<PrimExpr>(op->node);
         HostEvaluableConditionChecker checker(params_);
@@ -227,7 +227,7 @@ class HostDeviceSplitter : public tirx::StmtMutator {
 
   private:
     Stmt VisitStmt_(const tirx::AttrStmtNode *op) final {
-      if (op->attr_key == tl::attr::kAssumeRuntimeCheck) {
+      if (op->attr_key == tl::attr::kAssumeRequiresRuntimeCheck) {
         return VisitStmt(op->body);
       }
       return tirx::StmtMutator::VisitStmt_(op);
@@ -392,8 +392,8 @@ class HostDeviceSplitter : public tirx::StmtMutator {
     for (auto it = runtime_checks.rbegin(); it != runtime_checks.rend(); ++it) {
       body = AttrStmt(it->condition, tirx::attr::tilelang_assume, it->message,
                       std::move(body), it->span);
-      body = AttrStmt(it->condition, tl::attr::kAssumeRuntimeCheck, it->message,
-                      std::move(body), it->span);
+      body = AttrStmt(it->condition, tl::attr::kAssumeRequiresRuntimeCheck,
+                      it->message, std::move(body), it->span);
     }
     return body;
   }
diff --git a/testing/python/transform/test_tilelang_transform_make_packed_api.py b/testing/python/transform/test_tilelang_transform_make_packed_api.py
index 4ce3f16e7d..abe1f43cc5 100644
--- a/testing/python/transform/test_tilelang_transform_make_packed_api.py
+++ b/testing/python/transform/test_tilelang_transform_make_packed_api.py
@@ -136,7 +136,7 @@ def subroutine(A_data: T.handle("float32")):
     )
 
 
-def test_assume_runtime_check_is_lowered_to_assert():
+def test_assume_requires_runtime_check_is_lowered_to_assert():
     n = tirx.Var("n", "int32")
     condition = n % 4 == 0
     message = "n must be divisible by 4"
@@ -148,7 +148,7 @@ def test_assume_runtime_check_is_lowered_to_assert():
     )
     body = tirx.AttrStmt(
         condition,
-        "tl.assume_runtime_check",
+        "tl.assume_requires_runtime_check",
         tirx.StringImm(message),
         body,
     )
@@ -163,7 +163,7 @@ def test_assume_runtime_check_is_lowered_to_assert():
 
     after_mod = tilelang.transform.MakePackedAPI()(tvm.IRModule.from_expr(before))
     after = after_mod["main"]
-    assert not [stmt for stmt in _collect_nodes(after.body, tirx.AttrStmt) if stmt.attr_key == "tl.assume_runtime_check"]
+    assert not [stmt for stmt in _collect_nodes(after.body, tirx.AttrStmt) if stmt.attr_key == "tl.assume_requires_runtime_check"]
     assert [stmt for stmt in _collect_nodes(after.body, tirx.AttrStmt) if stmt.attr_key == "tl.assume"]
 
     matching_asserts = [
diff --git a/testing/python/transform/test_tilelang_transform_split_host_device.py b/testing/python/transform/test_tilelang_transform_split_host_device.py
index 039392e807..68ca4ad6fd 100644
--- a/testing/python/transform/test_tilelang_transform_split_host_device.py
+++ b/testing/python/transform/test_tilelang_transform_split_host_device.py
@@ -76,7 +76,7 @@ def collect_assume_conditions(func: tvm.tirx.PrimFunc):
 
 
 def collect_runtime_check_conditions(func: tvm.tirx.PrimFunc):
-    return collect_attr_conditions(func, "tl.assume_runtime_check")
+    return collect_attr_conditions(func, "tl.assume_requires_runtime_check")
 
 
 def get_var_name(var):