Skip to content

Commit

Permalink
[XLA:GPU] Make VLOG explanation in SortRewriter more helpful.
Browse files Browse the repository at this point in the history
Add the dtypes for which a CUB kernel is unavailable to the log output.

PiperOrigin-RevId: 703067645
  • Loading branch information
thomasjoerg authored and tensorflower-gardener committed Dec 5, 2024
1 parent ebb9b42 commit 2e39667
Showing 1 changed file with 5 additions and 1 deletion.
6 changes: 5 additions & 1 deletion third_party/xla/xla/service/gpu/transforms/sort_rewriter.cc
Original file line number Diff line number Diff line change
Expand Up @@ -436,7 +436,11 @@ bool IsCubCompatibleSort(const HloSortInstruction* sort_op) {
return false;
}
if (!CreateRunner(*sort_analysis).ok()) {
VLOG(2) << "Unsupported operand types (no compiled CUB kernels)";
VLOG(2) << "Unsupported operand types (no compiled CUB kernels): "
<< PrimitiveType_Name(sort_analysis->key_type) << " "
<< (sort_analysis->value_type.has_value()
? PrimitiveType_Name(sort_analysis->value_type.value())
: "");
return false;
}
VLOG(2) << "Sort operation is compatible";
Expand Down

0 comments on commit 2e39667

Please sign in to comment.