From 3787e932efd1e728d218a7de88419b9dae53939c Mon Sep 17 00:00:00 2001
From: Matthew <mbrookhart@octoml.ai>
Date: Thu, 17 Jun 2021 09:17:53 -0600
Subject: [PATCH] revert a change to lower_tvm_builtin.cc from #6126

---
 src/tir/transforms/lower_tvm_builtin.cc | 10 ++++++++++
 1 file changed, 10 insertions(+)

diff --git a/src/tir/transforms/lower_tvm_builtin.cc b/src/tir/transforms/lower_tvm_builtin.cc
index f5a553aa0598..99d71ebe15bd 100644
--- a/src/tir/transforms/lower_tvm_builtin.cc
+++ b/src/tir/transforms/lower_tvm_builtin.cc
@@ -113,6 +113,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];