Skip to content

Commit

Permalink
Mark all outer loops as GPU block
Browse files Browse the repository at this point in the history
  • Loading branch information
antonysigma committed Nov 5, 2023
1 parent 5202a42 commit 183e240
Showing 1 changed file with 5 additions and 4 deletions.
9 changes: 5 additions & 4 deletions src/autoschedulers/mullapudi2016/AutoSchedule.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -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();
Expand All @@ -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;
}
}
}
Expand Down

0 comments on commit 183e240

Please sign in to comment.