gemini-code-assist[bot] commented on code in PR #19984:
URL: https://github.com/apache/tvm/pull/19984#discussion_r3564553468
##########
python/tvm/s_tir/dlight/gpu/reduction.py:
##########
@@ -268,10 +270,36 @@ def _sch_inner_spatial(
sch.decompose_reduction(rf, r)
# Schedule the write back block
sch.reverse_compute_at(block, bx, preserve_unit_loops=True)
- _, r, *s = sch.get_loops(block)
+
+ all_loops = sch.get_loops(block)[1:]
+ all_iter_vars = sch.get(block).iter_vars[1:]
+
+ if len(all_loops) != len(all_iter_vars):
+ return None
+
+ s_loops = []
+ r_loops = []
+ for loop_rv, iter_var in zip(all_loops, all_iter_vars):
+ if iter_var.iter_type == tirx.IterVar.DataPar:
+ s_loops.append(loop_rv)
+ elif iter_var.iter_type == tirx.IterVar.CommReduce:
+ r_loops.append(loop_rv)
+
+ if not s_loops or not r_loops:
+ return None
+
+ r = r_loops[0] # Use first (usually only) reduction loop
+
if unroll_spatial_factor:
- assert len(s) == len(loop_order)
- new_order_s = [s[loop_order[i]] for i in range(len(s))]
+ num_original_spatial = len([i for i in block_info.iters if i.kind
== "S"])
+ if len(s_loops) != num_original_spatial:
+ return None
+ # Safe to apply loop_order reordering (assumes write-back block
spatial loops)
+ # correspond to original block spatial loops)
Review Comment:
There is a parenthesis typo in the comment. The opening parenthesis on line
297 is closed prematurely, and line 298 has an extra closing parenthesis.
```suggestion
# Safe to apply loop_order reordering (assumes write-back block
spatial loops
# correspond to original block spatial loops)
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
