mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-19 05:16:06 +00:00
Code Issues Packages Projects Releases Wiki Activity

[XLA:Mosaic] No need to assume a multiple of tile if tile dim size is 1.

Browse Source 
PiperOrigin-RevId: 642301822
This commit is contained in:
Jevin Jiang 2024-06-11 09:51:45 -07:00 committed by jax authors
parent 27140fe6de
commit 5b38549810
1 changed files with 3 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
jaxlib/mosaic/dialect/tpu/tpu_dialect.cc
View File

@ -189,6 +189,9 @@ bool isGuaranteedDivisible(Value value, int64_t divisor, int64_t fuel) {
if (fuel <= 0) {
return false;
}
if (divisor == 1) {
return true;
}
if (auto assume_op = value.getDefiningOp<tpu::AssumeMultipleOp>()) {
return assume_op.getMultiple() % divisor == 0;
}