(tvm) branch main updated: [FIX][TIRx] Preserve renewal domains and allocation-bound lifetimes (#20502)

[tqchen]

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 0587e59a69 [FIX][TIRx] Preserve renewal domains and allocation-bound 
lifetimes (#20502)
0587e59a69 is described below

commit 0587e59a696cb435ad4d328726a07daa8d966f3c
Author: Tianqi Chen <tqc...@users.noreply.github.com>
AuthorDate: Wed Sep 30 13:09:58 2026 -0400

    [FIX][TIRx] Preserve renewal domains and allocation-bound lifetimes (#20502)
    
    Preserve iteration-domain extents when renewing S-TIR definitions with
    nonzero minima by reconstructing ranges from their minimum and extent.
    
    Keep integer-limit expressions alive while Metal and WebGPU code
    generation check symbolic allocation bounds, so the comparisons retain
    valid backing nodes.
---
 src/backend/metal/codegen/codegen_metal.cc   | 3 ++-
 src/backend/webgpu/codegen/codegen_webgpu.cc | 3 ++-
 src/s_tir/transform/renew_defs.cc            | 5 +++--
 3 files changed, 7 insertions(+), 4 deletions(-)

diff --git a/src/backend/metal/codegen/codegen_metal.cc 
b/src/backend/metal/codegen/codegen_metal.cc
index caa5ae1bd5..565290d04a 100644
--- a/src/backend/metal/codegen/codegen_metal.cc
+++ b/src/backend/metal/codegen/codegen_metal.cc
@@ -385,7 +385,8 @@ void CodeGenMetal::DispatchAllocBuffer(const BindNode* op, 
const CallNode* buffe
       // An integer dtype's intrinsic maximum is not a program-derived 
allocation bound.
       TVM_FFI_ICHECK(dim_size != sym::ConstIntBound::kPosInf)
           << "Metal allocation extent requires a finite compile-time upper 
bound, but got " << dim;
-      if (const auto* dtype_max = 
max_value(dim->ty.as_or_throw<PrimType>()).as<IntImmNode>()) {
+      PrimExpr dtype_max_value = max_value(dim->ty.as_or_throw<PrimType>());
+      if (const auto* dtype_max = dtype_max_value.as<IntImmNode>()) {
         TVM_FFI_ICHECK_LT(dim_size, dtype_max->value)
             << "Metal allocation extent requires a finite compile-time upper 
bound, but got "
             << dim;
diff --git a/src/backend/webgpu/codegen/codegen_webgpu.cc 
b/src/backend/webgpu/codegen/codegen_webgpu.cc
index d313454099..1a505f966f 100644
--- a/src/backend/webgpu/codegen/codegen_webgpu.cc
+++ b/src/backend/webgpu/codegen/codegen_webgpu.cc
@@ -752,7 +752,8 @@ void CodeGenWebGPU::DispatchAllocBuffer(const BindNode* op, 
const CallNode* buff
     int64_t dim_size = dim_imm ? static_cast<int64_t>(dim_imm->value)
                                : 
analyzer_->const_int_bound(dim.as_or_throw<PrimExpr>())->max_value;
     if (dim_imm == nullptr) {
-      const auto* dtype_max = 
max_value(dim->ty.as_or_throw<PrimType>()).as<IntImmNode>();
+      PrimExpr dtype_max_value = max_value(dim->ty.as_or_throw<PrimType>());
+      const auto* dtype_max = dtype_max_value.as<IntImmNode>();
       // An integer dtype's intrinsic maximum is not a program-derived 
allocation bound.
       TVM_FFI_ICHECK(dtype_max && dim_size < dtype_max->value)
           << "WebGPU allocation extent requires a finite compile-time upper 
bound, but got " << dim;
diff --git a/src/s_tir/transform/renew_defs.cc 
b/src/s_tir/transform/renew_defs.cc
index a8250b8b88..81826778d6 100644
--- a/src/s_tir/transform/renew_defs.cc
+++ b/src/s_tir/transform/renew_defs.cc
@@ -144,8 +144,9 @@ class RenewDefMutator : public StmtExprMutator {
         Mutate(iter_var->dom->min, 
InplaceMode::kDisallow).ValueOrUnchanged(iter_var->dom->min);
     PrimExpr extent = Mutate(iter_var->dom->extent, InplaceMode::kDisallow)
                           .ValueOrUnchanged(iter_var->dom->extent);
-    IterVar new_iter_var(Range(min, extent), 
ReDefineVar(iter_var->var).as_or_throw<PrimVar>(),
-                         iter_var->iter_type, iter_var->thread_tag);
+    IterVar new_iter_var(Range::FromMinExtent(min, extent),
+                         ReDefineVar(iter_var->var).as_or_throw<PrimVar>(), 
iter_var->iter_type,
+                         iter_var->thread_tag);
     VarRemapSet(iter_var, new_iter_var);
     return new_iter_var;
   }