diff --git a/src/autoschedulers/mullapudi2016/AutoSchedule.cpp b/src/autoschedulers/mullapudi2016/AutoSchedule.cpp index 04b841c4c441..e67d7d7be2d9 100644 --- a/src/autoschedulers/mullapudi2016/AutoSchedule.cpp +++ b/src/autoschedulers/mullapudi2016/AutoSchedule.cpp @@ -1105,6 +1105,7 @@ class GPUTilingDedup { } void mark_gpu_threads(AutoSchedule &sched) const { + bool is_gpu_block_marked = false; for (const auto &v : ordering) { const auto &v_name = v.name(); @@ -1115,17 +1116,17 @@ class GPUTilingDedup { continue; } + // Skip all gpu_blocks if the current Stage is "compute_at" another + // stage, in which the gpu_blocks are already specified. if (is_compute_at) { continue; } - // Skip all gpu_blocks if the current Stage is "compute_at" another - // stage, in which the gpu_blocks are already specified. - if (is_outer(v_name)) { + if (is_outer(v_name) || is_gpu_block_marked) { // Mark as gpu blocks; f.gpu_blocks(v); sched.push_schedule(f.name(), stage_num, "gpu_blocks(" + v_name + ")", {v_name}); - continue; + is_gpu_block_marked = true; } } }