Skip to content

Commit

Permalink
[XLA:GPU] Allow rewrites of the dot without non-contracting dims to m…
Browse files Browse the repository at this point in the history
…ultiply+reduction if dot algorithm is F32_F32_F32

Such rewrite does not change the precision and makes this op 100x faster.

In the cl we:
1) convert the if block to the early exit version. (reduce indent)
2) relax the algorithm check by inverting the condition has_precision_config_algorithm ==> can_rewrite_dot_with_precision_config_algorithm (simplifies if-s)
3) add F32_F32_F32 algorithm to the can_rewrite_dot_with_precision_config_algorithm condition

PiperOrigin-RevId: 694491750
  • Loading branch information
loislo authored and Google-ML-Automation committed Nov 8, 2024
1 parent 47c65b1 commit 3f25d27
Show file tree
Hide file tree
Showing 2 changed files with 330 additions and 295 deletions.
Loading

0 comments on commit 3f25d27

Please sign in to comment.