Fix the behavior of decompose_reduction to keep thread_binding inform…

…ation (apache#59)
yzh119 · Oct 26, 2022 · a2f4447 · a2f4447
1 parent f045903
commit a2f4447
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/src/tir/schedule/primitive/reduction.cc b/src/tir/schedule/primitive/reduction.cc
@@ -296,7 +296,9 @@ StmtSRef DecomposeReduction(ScheduleState self, const StmtSRef& block_sref,
                /*min=*/old_loop->min,
                /*extent=*/old_loop->extent,
                /*kind=*/old_loop->kind,
-               /*body=*/body);
+               /*body=*/body,
+               /*thread_binding=*/old_loop->thread_binding,
+               /*annotations=*/old_loop->annotations);
   }
   body = Substitute(body, loop_var_map);
   // Step 6. Mutate IR