Skip to content

[Enhancement] Add thread count validation for ReduceOp fragment layout inference #1225

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
LeiWang1999 merged 2 commits into from
Nov 11, 2025
Merged

[Enhancement] Add thread count validation for ReduceOp fragment layout inference #1225

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
5b11be1
[Enhancement] Add thread count validation for ReduceOp fragment layou…
LeiWang1999 Nov 11, 2025
2938111
lint fix
LeiWang1999 Nov 11, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
29 changes: 29 additions & 0 deletions src/op/reduce.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -389,6 +389,35 @@ LayoutMap ReduceOpNode::InferLayout(const LayoutInferArgs &T,
}
auto thd = src_layout->ForwardThread(
fwd, FloorDiv(ReplicationPlaceholder(), indice_rep_extent));

// Ensure the thread count is divisible by the replicate extent.
// Otherwise, we cannot infer a valid fragment<->fragment layout.
{
arith::Analyzer analyzer;
PrimExpr num_threads = T.thread_bounds->extent;
// Though the dest_buffer_rep_extent will be compressed at
// CondenseReplicateVar, we need to check the divisibility here to avoid
// the issue that the thread count is not divisible by the replicate
// extent.
if (!analyzer.CanProve(FloorMod(num_threads, dest_buffer_rep_extent) ==
0) &&
!analyzer.CanProve(FloorMod(dest_buffer_rep_extent, num_threads) ==
0)) {
ICHECK(false) << "ReduceOp fragment layout inference failed: "
"num_threads % replicate_extent != 0. "
<< "This mapping requires the block's thread count to be "
"divisible by the "
<< "replicate extent. "
<< "Try one of: (1) choose a thread block size divisible "
"by replicate_extent; "
<< "(2) pick a different reduce dimension or adjust the "
"source fragment layout; "
<< "Details: num_threads=" << num_threads
<< ", replicate_extent=" << indice_rep_extent
<< ", src=" << src << ", dst=" << dst;
}
Comment on lines +406 to +418
Copy link
Contributor

@coderabbitai coderabbitai bot Nov 11, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟡 Minor

Error message should print the actual value being checked.

The error message on line 416 prints indice_rep_extent, but the divisibility check validates dest_buffer_rep_extent (which equals indice_rep_extent * src_rep_extent from line 378). This discrepancy could confuse users debugging the error.

Apply this diff to improve error message accuracy:

                      << "Details: num_threads=" << num_threads
-                     << ", replicate_extent=" << indice_rep_extent
+                     << ", replicate_extent=" << dest_buffer_rep_extent
                      << ", src=" << src << ", dst=" << dst;
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
ICHECK(false) << "ReduceOp fragment layout inference failed: "
"num_threads % replicate_extent != 0. "
<< "This mapping requires the block's thread count to be "
"divisible by the "
<< "replicate extent. "
<< "Try one of: (1) choose a thread block size divisible "
"by replicate_extent; "
<< "(2) pick a different reduce dimension or adjust the "
"source fragment layout; "
<< "Details: num_threads=" << num_threads
<< ", replicate_extent=" << indice_rep_extent
<< ", src=" << src << ", dst=" << dst;
}
ICHECK(false) << "ReduceOp fragment layout inference failed: "
"num_threads % replicate_extent != 0. "
<< "This mapping requires the block's thread count to be "
"divisible by the "
<< "replicate extent. "
<< "Try one of: (1) choose a thread block size divisible "
"by replicate_extent; "
<< "(2) pick a different reduce dimension or adjust the "
"source fragment layout; "
<< "Details: num_threads=" << num_threads
<< ", replicate_extent=" << dest_buffer_rep_extent
<< ", src=" << src << ", dst=" << dst;
}
🤖 Prompt for AI Agents 
In src/op/reduce.cc around lines 406 to 418, the error message prints
indice_rep_extent but the divisibility check is actually on
dest_buffer_rep_extent (which equals indice_rep_extent * src_rep_extent); update
the ICHECK message to print dest_buffer_rep_extent (or include both
dest_buffer_rep_extent and indice_rep_extent) so the logged values match the
condition being checked and give accurate debugging info.

}

Fragment dst_layout =
Fragment(dst->shape, {}, thd, dest_buffer_rep_extent, std::nullopt)
->CondenseReplicateVar()
Expand Down
2 changes: 0 additions & 2 deletions testing/python/language/test_tilelang_language_reduce.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -116,7 +116,6 @@ def test_reduce_sum():

def test_reduce_sum_shared():
run_reduce_sum(64, 64, mode="ss")
run_reduce_sum(32, 96, mode="ss")


def test_reduce_max():
Expand All @@ -127,7 +126,6 @@ def test_reduce_max():

def test_reduce_max_shared():
run_shared_reduce(reduce_max_ss, lambda A: A.max(dim=1).values, 64, 64, "float32")
run_shared_reduce(reduce_max_ss, lambda A: A.max(dim=1).values, 96, 48, "float32")


def test_reduce_min_shared():
Expand Down
Loading