[TIR][Runtime] Enforce host-evaluable assumptions at runtime

src/transform/common/attr.h

@@ -33,6 +33,10 @@ 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 *kAssumeRequiresRuntimeCheck =
+    "tl.assume_requires_runtime_check";
 
 // Attributes to implement SourceCodeBlock
 constexpr const char *kCodeBlockSource = "code_block_source";

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::kAssumeRequiresRuntimeCheck, 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";
       }

src/transform/make_packed_api.cc

@@ -179,6 +179,75 @@ class SubroutineCallRewriter : public StmtExprMutator {
   bool made_change_{false};
 };
 
+struct AssumeRuntimeCheck {
+  PrimExpr condition;
+  Array<StringImm> message_parts;
+  Span span;
+};
+
+class AssumeRuntimeCheckExtractor : public StmtMutator {
+public:
+  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 != tl::attr::kAssumeRequiresRuntimeCheck) {
+      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()));
+    }
+
+    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
 
 inline Stmt MakeAssertEQ(PrimExpr lhs, PrimExpr rhs, std::string msg) {
@@ -250,6 +319,9 @@ PrimFunc MakePackedAPI(PrimFunc func) {
   }
 
   auto *func_ptr = func.CopyOnWrite();
+  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());
@@ -545,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

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_.
+// 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.
+//    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,
@@ -92,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);
   }
@@ -113,6 +116,124 @@ class HostDeviceSplitter : public tirx::StmtMutator {
   bool found_device_region() const { return found_device_region_; }
 
 private:
+  struct RuntimeCheckInfo {
+    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 DeviceRuntimeCheckCollector : public tirx::StmtVisitor {
+  public:
+    static std::vector<RuntimeCheckInfo>
+    Collect(const Stmt &body, const Array<tirx::Var> &params) {
+      DeviceRuntimeCheckCollector collector(params);
+      collector(body);
+      return collector.runtime_checks_;
+    }
+
+  private:
+    explicit DeviceRuntimeCheckCollector(const Array<tirx::Var> &params)
+        : params_(params) {}
+
+    void VisitStmt_(const tirx::AttrStmtNode *op) final {
+      if (op->attr_key == tl::attr::kAssumeRequiresRuntimeCheck &&
+          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());
+          }();
+          runtime_checks_.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<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::kAssumeRequiresRuntimeCheck) {
+        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_;
@@ -266,6 +387,17 @@ class HostDeviceSplitter : public tirx::StmtMutator {
     return body;
   }
 
+  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::kAssumeRequiresRuntimeCheck,
+                      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 +425,19 @@ class HostDeviceSplitter : public tirx::StmtMutator {
               use_def.undefined_buffers_};
     }();
 
+    // 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
     // incorrectly renaming variables when it processes multiple functions.
@@ -397,30 +542,29 @@ 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);
       tirx::Stmt assert_success = tirx::AssertStmt(
           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 WrapBodyWithRuntimeChecks(std::move(host_call), host_runtime_checks);
   }
 
   // target ir module
   IRModule *device_mod_;
   // 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,