Skip to content

Commit

Permalink
revert a change to lower_tvm_builtin.cc from #6126
Browse files Browse the repository at this point in the history
  • Loading branch information
Matthew committed Jun 17, 2021
1 parent 1fac10b commit 05f28b4
Showing 1 changed file with 10 additions and 0 deletions.
10 changes: 10 additions & 0 deletions src/tir/transforms/lower_tvm_builtin.cc
Original file line number Diff line number Diff line change
Expand Up @@ -114,6 +114,16 @@ class BuiltinLower : public StmtExprMutator {
op = stmt.as<AllocateNode>();
// Get constant allocation bound.
int64_t nbytes = GetVectorBytes(op->dtype);
if (device_type_.defined()) {
if (const auto* dev_type = device_type_.as<IntImmNode>()) {
if (dev_type->value == kDLCPU) {
int32_t constant_size = op->constant_allocation_size();
if (constant_size > 0 && constant_size * nbytes < runtime::kMaxStackAlloca) {
return stmt;
}
}
}
}
PrimExpr total_bytes = make_const(op->extents[0].dtype(), nbytes);
for (size_t i = 0; i < op->extents.size(); ++i) {
total_bytes = total_bytes * op->extents[i];
Expand Down

0 comments on commit 05f28b4

Please sign in to comment.