[TIR][VM] Revert a change to lower_tvm_builtin.cc from apache#6126 (a…

…pache#8274) * revert a change to lower_tvm_builtin.cc from apache#6126 * disable sim target on VTA tutorial fix bad refactor try again
ylc · Sep 29, 2021 · 3b9ff7a · 3b9ff7a
1 parent a6bd96a
commit 3b9ff7a
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 3 deletions.
diff --git 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];

diff --git a/vta/tutorials/frontend/deploy_classification.py b/vta/tutorials/frontend/deploy_classification.py
@@ -191,7 +191,7 @@
                 env.WGT_WIDTH,
                 start_name=pack_dict[model][0],
                 stop_name=pack_dict[model][1],
-                device_annot=(env.TARGET == "intelfocl" or env.TARGET == "sim"),
+                device_annot=(env.TARGET == "intelfocl"),
             )
     else:
         relay_prog = mod["main"]
@@ -203,7 +203,7 @@
                 relay_prog, target=target, params=params, target_host=env.target_host
             )
     else:
-        if env.TARGET == "intelfocl" or env.TARGET == "sim":
+        if env.TARGET == "intelfocl":
             # multiple targets to run both on cpu and vta
             target = {"cpu": env.target_vta_cpu, "ext_dev": target}
         with vta.build_config(opt_level=3, disabled_pass={"AlterOpLayout"}):
@@ -221,7 +221,7 @@
     remote.upload(temp.relpath("graphlib.tar"))
     lib = remote.load_module("graphlib.tar")
 
-    if env.TARGET == "intelfocl" or env.TARGET == "sim":
+    if env.TARGET == "intelfocl":
         ctxes = [remote.ext_dev(0), remote.cpu(0)]
         m = graph_executor.create(graph, lib, ctxes)
     else: