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

[Mosaic TPU] Propagate the memory space change for memref bitcast and reshape.

Browse Source 
PiperOrigin-RevId: 674067380
This commit is contained in:
Jevin Jiang 2024-09-12 17:13:46 -07:00 committed by jax authors
parent 178fb03050
commit 8d93e101b9
1 changed files with 8 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
jaxlib/mosaic/dialect/tpu/transforms/memory_space_specialization.cc
View File

@ -78,6 +78,14 @@ LogicalResult specializeMemorySpace(TypedValue<MemRefType> value,
updateResultFrom(op, op.getInput().getType());
continue;
}
if (auto op = dyn_cast<tpu::MemRefBitcastOp>(some_op)) {
updateResultFrom(op, op.getInput().getType());
continue;
}
if (auto op = dyn_cast<tpu::MemRefReshapeOp>(some_op)) {
updateResultFrom(op, op.getInput().getType());
continue;
}
if (auto op = dyn_cast<tpu::EraseLayoutOp>(some_op)) {
updateResultFrom(op, op.getOperand().getType());
continue;