Skip to content

Commit c866abc

Browse files
mshr-hmshr-h
authored
[Relax] Fix wrong memory planning when only lower bound was provided (#18663)
This PR fixes an issue in StaticPlanBlockMemory where dynamic shapes were incorrectly planned as static memory when only a lower bound was provided for TIR variables. Repro: <details> <summary>repro_dynamic_memory_plan.py</summary> ```python import tvm from tvm import relax, testing from tvm.relax.frontend.torch import from_exported_program from torch.export import Dim, export import torch class SimpleConv(torch.nn.Module): def __init__(self): super().__init__() self.conv = torch.nn.Conv2d(3, 64, kernel_size=3, padding=1) def forward(self, x): return self.conv(x) def main(): model = SimpleConv().eval() example = torch.randn(2, 3, 32, 32) batch = Dim("batch") # No max= specified, so upper bound is unknown exported = export(model, (example,), dynamic_shapes={"x": {0: batch}}) mod = from_exported_program(exported) mod = relax.transform.DecomposeOpsForInference()(mod) target = tvm.target.Target("llvm") exe = tvm.compile(mod, target=target) vm = relax.VirtualMachine(exe, tvm.cpu()) inp = tvm.runtime.from_dlpack(example) out = vm["main"](inp) expected = model(example).detach().numpy() actual = out[0].numpy() testing.assert_allclose(actual, expected, rtol=1e-4, atol=1e-4) if __name__ == "__main__": main() ``` </details> This will fail with the following error. <details> <summary>output</summary> ``` $ uv run python repro_dynamic_memory_plan.py /home/ubuntu/data/project/tvm-example/.venv/lib/python3.12/site-packages/torch/cuda/__init__.py:182: UserWarning: CUDA initialization: Unexpected error from cudaGetDeviceCount(). Did you run some cuda functions before calling NumCudaDevices() that might have already set an error? Error 804: forward compatibility was attempted on non supported HW (Triggered internally at /pytorch/c10/cuda/CUDAFunctions.cpp:119.) return torch._C._cuda_getDeviceCount() > 0 Traceback (most recent call last): File "/home/ubuntu/data/project/tvm-example/frontend/repro_dynamic_memory_plan.py", line 40, in <module> main() File "/home/ubuntu/data/project/tvm-example/frontend/repro_dynamic_memory_plan.py", line 32, in main out = vm["main"](inp) ^^^^^^^^^^^^^^^ File "python/tvm_ffi/cython/function.pxi", line 923, in tvm_ffi.core.Function.__call__ File "/home/ubuntu/data/project/tvm-example/tvm/src/runtime/vm/vm.cc", line 549, in tvm::runtime::vm::VirtualMachineImpl::InvokeClosurePacked(tvm::ffi::ObjectRef const&, tvm::ffi::PackedArgs, tvm::ffi::Any*) clo->impl.CallPacked(ffi::PackedArgs(packed_args.data(), packed_args.size()), rv); File "/home/ubuntu/data/project/tvm-example/tvm/src/runtime/vm/vm.cc", line 622, in operator() *rv = static_cast<VirtualMachineImpl*>(ctx_ptr)->InvokeBytecode(gf_idx, inputs); File "/home/ubuntu/data/project/tvm-example/tvm/src/runtime/vm/vm.cc", line 693, in tvm::runtime::vm::VirtualMachineImpl::InvokeBytecode(long, std::vector<tvm::ffi::Any, std::allocator<tvm::ffi::Any> > const&) RunLoop(); File "/home/ubuntu/data/project/tvm-example/tvm/src/runtime/vm/vm.cc", line 816, in tvm::runtime::vm::VirtualMachineImpl::RunLoop() this->RunInstrCall(curr_frame, instr); File "/home/ubuntu/data/project/tvm-example/tvm/src/runtime/vm/vm.cc", line 767, in tvm::runtime::vm::VirtualMachineImpl::RunInstrCall(tvm::runtime::vm::VMFrame*, tvm::runtime::vm::Instruction) this->InvokeClosurePacked(func_pool_[instr.func_idx].cast<ObjectRef>(), args, &ret); File "/home/ubuntu/data/project/tvm-example/tvm/src/runtime/vm/builtin.cc", line 405, in operator() *rv = sobj->AllocTensor(offset, shape, dtype); File "/home/ubuntu/data/project/tvm-example/tvm/src/runtime/memory/memory_manager.cc", line 98, in tvm::runtime::memory::StorageObj::AllocTensor(long, tvm::ffi::Shape, DLDataType) ICHECK(offset + needed_size <= this->buffer.size) File "/home/ubuntu/data/project/tvm-example/tvm/include/tvm/runtime/logging.h", line 321, in tvm::runtime::detail::LogFatal::~LogFatal() GetEntry().Finalize(); File "/home/ubuntu/data/project/tvm-example/tvm/include/tvm/runtime/logging.h", line 337, in tvm::runtime::detail::LogFatal::Entry::Finalize() InternalError error(file_, lineno_, stream_.str()); tvm.error.InternalError: Check failed: (offset + needed_size <= this->buffer.size) is false: storage allocation failure, attempted to allocate 524288 at offset 0 in region that is 262144bytes ``` </details>
1 parent 2b4b58d commit c866abc

2 files changed

Lines changed: 246 additions & 28 deletions

File tree

‎src/relax/transform/static_plan_block_memory.cc‎

Lines changed: 7 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -437,15 +437,18 @@ void SetTIRVarRangeConstraints(Function func, arith::Analyzer* ana,
437437
auto it_upper = var_upper_bound_attr.find(tir_var->name_hint);
438438
auto it_lower = var_lower_bound_attr.find(tir_var->name_hint);
439439

440-
if (it_upper != var_upper_bound_attr.end() || it_lower != var_lower_bound_attr.end()) {
440+
// Only bind the variable to a range if an upper bound is explicitly provided.
441+
// Without an upper bound, memory planning cannot determine the required storage size,
442+
// so we skip binding and let the variable remain unbounded.
443+
if (it_upper != var_upper_bound_attr.end()) {
441444
int64_t lower = (it_lower != var_lower_bound_attr.end()) ? it_lower->second->value : 0;
442-
int64_t upper = (it_upper != var_upper_bound_attr.end())
443-
? it_upper->second->value
444-
: std::numeric_limits<int64_t>::max();
445+
int64_t upper = it_upper->second->value;
445446
tvm::Range range = tvm::Range::FromMinExtent(
446447
tvm::IntImm(DataType::Int(64), lower), tvm::IntImm(DataType::Int(64), upper - lower + 1));
447448
ana->Bind(tir_var, range);
448449
dom_map->Set(tir_var, arith::IntSet::FromRange(range));
450+
} else if (it_lower != var_lower_bound_attr.end() && it_lower->second->value >= 0) {
451+
ana->MarkGlobalNonNegValue(tir_var);
449452
} else if (non_negative_var_attr.count(tir_var->name_hint)) {
450453
ana->MarkGlobalNonNegValue(tir_var);
451454
}

‎tests/python/relax/test_transform_static_plan_block_memory.py‎

Lines changed: 239 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -1018,6 +1018,245 @@ def main(x: R.Tensor((2, "n"), dtype="float32")) -> R.Tensor(("2 * n + 2",), dty
10181018
tvm.ir.assert_structural_equal(mod, Expected)
10191019

10201020

1021+
def test_lower_bound_only():
1022+
# fmt: off
1023+
@tvm.script.ir_module
1024+
class Module:
1025+
@T.prim_func
1026+
def add(rxplaceholder: T.handle, rxplaceholder_1: T.handle, T_add: T.handle):
1027+
T.evaluate(0)
1028+
1029+
@T.prim_func
1030+
def reshape(rxplaceholder: T.handle, T_reshape: T.handle):
1031+
T.evaluate(0)
1032+
1033+
@T.prim_func
1034+
def relu(rxplaceholder: T.handle, compute: T.handle):
1035+
T.evaluate(0)
1036+
1037+
@T.prim_func
1038+
def log(rxplaceholder: T.handle, compute: T.handle):
1039+
T.evaluate(0)
1040+
1041+
@T.prim_func
1042+
def exp(rxplaceholder: T.handle, compute: T.handle):
1043+
T.evaluate(0)
1044+
1045+
@T.prim_func
1046+
def pad(rxplaceholder: T.handle, PadInput: T.handle):
1047+
T.evaluate(0)
1048+
1049+
@R.function
1050+
def main(x: R.Tensor((2, "n"), dtype="float32")) -> R.Tensor(("2 * n + 2",), dtype="float32"):
1051+
R.func_attr({"tir_var_lower_bound": {"n": 2}, "relax.force_pure": True})
1052+
n = T.int64()
1053+
cls = Module
1054+
alloc: R.Tensor((2, n), dtype="float32") = R.builtin.alloc_tensor(R.shape([2, n]), dtype="float32", runtime_device_index=0)
1055+
_: R.Tuple() = cls.exp(x, alloc)
1056+
lv: R.Tensor((2, n), dtype="float32") = alloc
1057+
lv1: R.Tensor((2 * n,), dtype="float32") = R.reshape(lv, (2 * n,))
1058+
alloc1: R.Tensor((2 * n,), dtype="float32") = R.builtin.alloc_tensor(R.shape([2 * n]), dtype="float32", runtime_device_index=0)
1059+
_1: R.Tuple() = cls.relu(lv1, alloc1)
1060+
lv2: R.Tensor((2 * n,), dtype="float32") = alloc1
1061+
alloc2: R.Tensor((2 * n,), dtype="float32") = R.builtin.alloc_tensor(R.shape([2 * n]), dtype="float32", runtime_device_index=0)
1062+
_2: R.Tuple() = cls.add(lv2, R.const(1, "float32"), alloc2)
1063+
lv3: R.Tensor((2 * n,), dtype="float32") = alloc2
1064+
alloc3: R.Tensor((2 * n + 2,), dtype="float32") = R.builtin.alloc_tensor(R.shape([2 * n + 2]), dtype="float32", runtime_device_index=0)
1065+
_3: R.Tuple() = cls.pad(lv3, alloc3)
1066+
lv4: R.Tensor((2 * n + 2,), dtype="float32") = alloc3
1067+
alloc4: R.Tensor((2 * n + 2,), dtype="float32") = R.builtin.alloc_tensor(R.shape([10]), dtype="float32", runtime_device_index=0)
1068+
_4: R.Tuple() = cls.log(lv4, alloc4)
1069+
gv: R.Tensor((2 * n + 2,), dtype="float32") = alloc4
1070+
return gv
1071+
1072+
@I.ir_module
1073+
class Expected:
1074+
@T.prim_func
1075+
def add(rxplaceholder: T.handle, rxplaceholder_1: T.handle, T_add: T.handle):
1076+
T.evaluate(0)
1077+
1078+
@T.prim_func
1079+
def exp(rxplaceholder: T.handle, compute: T.handle):
1080+
T.evaluate(0)
1081+
1082+
@T.prim_func
1083+
def log(rxplaceholder: T.handle, compute: T.handle):
1084+
T.evaluate(0)
1085+
1086+
@T.prim_func
1087+
def pad(rxplaceholder: T.handle, PadInput: T.handle):
1088+
T.evaluate(0)
1089+
1090+
@T.prim_func
1091+
def relu(rxplaceholder: T.handle, compute: T.handle):
1092+
T.evaluate(0)
1093+
1094+
@T.prim_func
1095+
def reshape(rxplaceholder: T.handle, T_reshape: T.handle):
1096+
T.evaluate(0)
1097+
1098+
@R.function
1099+
def main(x: R.Tensor((2, "n"), dtype="float32")) -> R.Tensor(("2 * n + 2",), dtype="float32"):
1100+
n = T.int64()
1101+
R.func_attr({"tir_var_lower_bound": {"n": 2}, "relax.force_pure": True})
1102+
cls = Expected
1103+
storage: R.Object = R.memory.alloc_storage(R.shape([8 * n]), R.prim_value(0), R.str("global"), R.dtype("float32"))
1104+
alloc: R.Tensor((2, n), dtype="float32") = R.memory.alloc_tensor(storage, R.prim_value(0), R.shape([2, n]), R.dtype("float32"), R.prim_value(0))
1105+
_: R.Tuple = cls.exp(x, alloc)
1106+
lv: R.Tensor((2, n), dtype="float32") = alloc
1107+
lv1: R.Tensor((2 * n,), dtype="float32") = R.reshape(lv, R.shape([2 * n]))
1108+
storage1: R.Object = R.memory.alloc_storage(R.shape([4 * (2 * n)]), R.prim_value(0), R.str("global"), R.dtype("float32"))
1109+
alloc1: R.Tensor((2 * n,), dtype="float32") = R.memory.alloc_tensor(storage1, R.prim_value(0), R.shape([2 * n]), R.dtype("float32"))
1110+
_1: R.Tuple = cls.relu(lv1, alloc1)
1111+
lv2: R.Tensor((2 * n,), dtype="float32") = alloc1
1112+
alloc2: R.Tensor((2 * n,), dtype="float32") = R.memory.alloc_tensor(storage, R.prim_value(0), R.shape([2 * n]), R.dtype("float32"))
1113+
_2: R.Tuple = cls.add(lv2, R.const(1, "float32"), alloc2)
1114+
lv3: R.Tensor((2 * n,), dtype="float32") = alloc2
1115+
storage2: R.Object = R.memory.alloc_storage(R.shape([4 * (2 * n + 2)]), R.prim_value(0), R.str("global"), R.dtype("float32"))
1116+
alloc3: R.Tensor((2 * n + 2,), dtype="float32") = R.memory.alloc_tensor(storage2, R.prim_value(0), R.shape([2 * n + 2]), R.dtype("float32"), R.prim_value(0))
1117+
_3: R.Tuple = cls.pad(lv3, alloc3)
1118+
lv4: R.Tensor((2 * n + 2,), dtype="float32") = alloc3
1119+
alloc4: R.Tensor((2 * n + 2,), dtype="float32") = R.builtin.alloc_tensor(R.shape([10]), R.dtype("float32"), R.prim_value(0))
1120+
_4: R.Tuple = cls.log(lv4, alloc4)
1121+
gv: R.Tensor((2 * n + 2,), dtype="float32") = alloc4
1122+
return gv
1123+
# fmt: on
1124+
1125+
mod = relax.transform.StaticPlanBlockMemory()(Module)
1126+
tvm.ir.assert_structural_equal(mod, Expected)
1127+
1128+
1129+
def test_upper_and_lower_bounds():
1130+
# fmt: off
1131+
@tvm.script.ir_module
1132+
class Module:
1133+
@T.prim_func
1134+
def add(rxplaceholder: T.handle, rxplaceholder_1: T.handle, T_add: T.handle):
1135+
T.evaluate(0)
1136+
1137+
@T.prim_func
1138+
def reshape(rxplaceholder: T.handle, T_reshape: T.handle):
1139+
T.evaluate(0)
1140+
1141+
@T.prim_func
1142+
def relu(rxplaceholder: T.handle, compute: T.handle):
1143+
T.evaluate(0)
1144+
1145+
@T.prim_func
1146+
def log(rxplaceholder: T.handle, compute: T.handle):
1147+
T.evaluate(0)
1148+
1149+
@T.prim_func
1150+
def exp(rxplaceholder: T.handle, compute: T.handle):
1151+
T.evaluate(0)
1152+
1153+
@T.prim_func
1154+
def pad(rxplaceholder: T.handle, PadInput: T.handle):
1155+
T.evaluate(0)
1156+
1157+
@R.function
1158+
def main(x: R.Tensor((2, "n"), dtype="float32")) -> R.Tensor(("2 * n + 2",), dtype="float32"):
1159+
R.func_attr({"tir_var_upper_bound": {"n": 4}, "tir_var_lower_bound": {"n": 2}, "relax.force_pure": True})
1160+
n = T.int64()
1161+
cls = Module
1162+
alloc: R.Tensor((2, n), dtype="float32") = R.builtin.alloc_tensor(R.shape([2, n]), dtype="float32", runtime_device_index=0)
1163+
_: R.Tuple() = cls.exp(x, alloc)
1164+
lv: R.Tensor((2, n), dtype="float32") = alloc
1165+
lv1: R.Tensor((2 * n,), dtype="float32") = R.reshape(lv, (2 * n,))
1166+
alloc1: R.Tensor((2 * n,), dtype="float32") = R.builtin.alloc_tensor(R.shape([2 * n]), dtype="float32", runtime_device_index=0)
1167+
_1: R.Tuple() = cls.relu(lv1, alloc1)
1168+
lv2: R.Tensor((2 * n,), dtype="float32") = alloc1
1169+
alloc2: R.Tensor((2 * n,), dtype="float32") = R.builtin.alloc_tensor(R.shape([2 * n]), dtype="float32", runtime_device_index=0)
1170+
_2: R.Tuple() = cls.add(lv2, R.const(1, "float32"), alloc2)
1171+
lv3: R.Tensor((2 * n,), dtype="float32") = alloc2
1172+
alloc3: R.Tensor((2 * n + 2,), dtype="float32") = R.builtin.alloc_tensor(R.shape([2 * n + 2]), dtype="float32", runtime_device_index=0)
1173+
_3: R.Tuple() = cls.pad(lv3, alloc3)
1174+
lv4: R.Tensor((2 * n + 2,), dtype="float32") = alloc3
1175+
alloc4: R.Tensor((2 * n + 2,), dtype="float32") = R.builtin.alloc_tensor(R.shape([10]), dtype="float32", runtime_device_index=0)
1176+
_4: R.Tuple() = cls.log(lv4, alloc4)
1177+
gv: R.Tensor((2 * n + 2,), dtype="float32") = alloc4
1178+
return gv
1179+
1180+
@I.ir_module
1181+
class Expected:
1182+
@T.prim_func
1183+
def add(rxplaceholder: T.handle, rxplaceholder_1: T.handle, T_add: T.handle):
1184+
T.evaluate(0)
1185+
1186+
@T.prim_func
1187+
def exp(rxplaceholder: T.handle, compute: T.handle):
1188+
T.evaluate(0)
1189+
1190+
@T.prim_func
1191+
def log(rxplaceholder: T.handle, compute: T.handle):
1192+
T.evaluate(0)
1193+
1194+
@T.prim_func
1195+
def pad(rxplaceholder: T.handle, PadInput: T.handle):
1196+
T.evaluate(0)
1197+
1198+
@T.prim_func
1199+
def relu(rxplaceholder: T.handle, compute: T.handle):
1200+
T.evaluate(0)
1201+
1202+
@T.prim_func
1203+
def reshape(rxplaceholder: T.handle, T_reshape: T.handle):
1204+
T.evaluate(0)
1205+
1206+
@R.function
1207+
def main(x: R.Tensor((2, "n"), dtype="float32")) -> R.Tensor(("2 * n + 2",), dtype="float32"):
1208+
n = T.int64()
1209+
R.func_attr({"tir_var_upper_bound": {"n": 4}, "tir_var_lower_bound": {"n": 2}, "relax.force_pure": True})
1210+
cls = Expected
1211+
storage: R.Object = R.memory.alloc_storage(R.shape([32]), R.prim_value(0), R.str("global"), R.dtype("float32"))
1212+
alloc: R.Tensor((2, n), dtype="float32") = R.memory.alloc_tensor(storage, R.prim_value(0), R.shape([2, n]), R.dtype("float32"))
1213+
_: R.Tuple = cls.exp(x, alloc)
1214+
lv: R.Tensor((2, n), dtype="float32") = alloc
1215+
lv1: R.Tensor((2 * n,), dtype="float32") = R.reshape(lv, R.shape([2 * n]))
1216+
storage1: R.Object = R.memory.alloc_storage(R.shape([40]), R.prim_value(0), R.str("global"), R.dtype("float32"))
1217+
alloc1: R.Tensor((2 * n,), dtype="float32") = R.memory.alloc_tensor(storage1, R.prim_value(0), R.shape([2 * n]), R.dtype("float32"))
1218+
_1: R.Tuple = cls.relu(lv1, alloc1)
1219+
lv2: R.Tensor((2 * n,), dtype="float32") = alloc1
1220+
alloc2: R.Tensor((2 * n,), dtype="float32") = R.memory.alloc_tensor(storage, R.prim_value(0), R.shape([2 * n]), R.dtype("float32"))
1221+
_2: R.Tuple = cls.add(lv2, R.const(1, "float32"), alloc2)
1222+
lv3: R.Tensor((2 * n,), dtype="float32") = alloc2
1223+
alloc3: R.Tensor((2 * n + 2,), dtype="float32") = R.memory.alloc_tensor(storage1, R.prim_value(0), R.shape([2 * n + 2]), R.dtype("float32"))
1224+
_3: R.Tuple = cls.pad(lv3, alloc3)
1225+
lv4: R.Tensor((2 * n + 2,), dtype="float32") = alloc3
1226+
alloc4: R.Tensor((2 * n + 2,), dtype="float32") = R.builtin.alloc_tensor(R.shape([10]), R.dtype("float32"), R.prim_value(0))
1227+
_4: R.Tuple = cls.log(lv4, alloc4)
1228+
gv: R.Tensor((2 * n + 2,), dtype="float32") = alloc4
1229+
return gv
1230+
# fmt: on
1231+
1232+
mod = relax.transform.StaticPlanBlockMemory()(Module)
1233+
tvm.ir.assert_structural_equal(mod, Expected)
1234+
1235+
1236+
def test_invalid_tir_var_upper_bound():
1237+
@tvm.script.ir_module
1238+
class Module:
1239+
@R.function
1240+
def main(x: R.Tensor((2, "n"), dtype="float32")):
1241+
R.func_attr({"tir_var_upper_bound": {"n": [4]}, "relax.force_pure": True})
1242+
return x
1243+
1244+
with pytest.raises((TVMError, TypeError)):
1245+
relax.transform.StaticPlanBlockMemory()(Module)
1246+
1247+
1248+
def test_invalid_tir_var_lower_bound():
1249+
@tvm.script.ir_module
1250+
class Module:
1251+
@R.function
1252+
def main(x: R.Tensor((2, "n"), dtype="float32")):
1253+
R.func_attr({"tir_var_lower_bound": {"n": [4]}, "relax.force_pure": True})
1254+
return x
1255+
1256+
with pytest.raises((TVMError, TypeError)):
1257+
relax.transform.StaticPlanBlockMemory()(Module)
1258+
1259+
10211260
def test_tir_var_decreasing_monotone():
10221261
# fmt: off
10231262
@I.ir_module
@@ -1335,30 +1574,6 @@ def func2(x: R.Tensor((10,), dtype="float32")) -> R.Tensor((10,), dtype="float32
13351574
tvm.ir.assert_structural_equal(mod, Expected)
13361575

13371576

1338-
def test_invalid_tir_var_upper_bound():
1339-
@tvm.script.ir_module
1340-
class Module:
1341-
@R.function
1342-
def main(x: R.Tensor((2, "n"), dtype="float32")):
1343-
R.func_attr({"tir_var_upper_bound": {"n": [4]}, "relax.force_pure": True})
1344-
return x
1345-
1346-
with pytest.raises((TVMError, TypeError)):
1347-
relax.transform.StaticPlanBlockMemory()(Module)
1348-
1349-
1350-
def test_invalid_tir_var_lower_bound():
1351-
@tvm.script.ir_module
1352-
class Module:
1353-
@R.function
1354-
def main(x: R.Tensor((2, "n"), dtype="float32")):
1355-
R.func_attr({"tir_var_lower_bound": {"n": [4]}, "relax.force_pure": True})
1356-
return x
1357-
1358-
with pytest.raises((TVMError, TypeError)):
1359-
relax.transform.StaticPlanBlockMemory()(Module)
1360-
1361-
13621577
def test_add():
13631578
@I.ir_module
13641579
class Module:

0 commit comments

Comments
 (0)