From f11f6b4e4362686c8714651ed00ef7ff2164a039 Mon Sep 17 00:00:00 2001
From: Tianqi Chen <tqchen@users.noreply.github.com>
Date: Fri, 9 Mar 2018 11:39:58 -0800
Subject: [PATCH] prevent aggressive unrolling in vthread (#983)
---
src/pass/inject_virtual_thread.cc | 3 ++-
1 file changed, 2 insertions(+), 1 deletion(-)
diff --git a/src/pass/inject_virtual_thread.cc b/src/pass/inject_virtual_thread.cc
index c0e7b4b3f..833513756 100644
--- a/src/pass/inject_virtual_thread.cc
+++ b/src/pass/inject_virtual_thread.cc
@@ -418,7 +418,8 @@ class VTInjector : public IRMutator {
// reset the flags after processing.
vt_loop_injected_ = false;
visit_touched_var_ = false;
- if (max_loop_depth_ == 0) {
+ // only unroll if number of vthreads are small
+ if (max_loop_depth_ == 0 && num_threads_ < 16) {
// do unrolling if it is inside innermost content.
Stmt blk = Substitute(stmt, {{var_, make_zero(var_.type())}});
for (int i = 1; i < num_threads_; ++i) {
--
GitLab