From 183e2404087ac74a984bb93b9482b35f3306091c Mon Sep 17 00:00:00 2001 From: Antony Chan Date: Sat, 4 Nov 2023 12:43:23 -0700 Subject: [PATCH] Mark all outer loops as GPU block --- src/autoschedulers/mullapudi2016/AutoSchedule.cpp | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) 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; } } }