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

[Mosaic] apply_vector_layout C++ rewrite: Fix check for defining layout in disassemble

Browse Source 
Though `relayout` would guarantee `equivalentTo` holding true, we skip relayout when the source layout generalizes the dest layout (because it's a  no-op).

PiperOrigin-RevId: 572095024
This commit is contained in:
Tomás Longeri 2023-10-09 18:23:52 -07:00 committed by jax authors
parent 4c306381c9
commit f2cda736a9
1 changed files with 1 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
jaxlib/mosaic/dialect/tpu/transforms/apply_vector_layout.cc
View File

@ -2691,7 +2691,7 @@ FailureOr<xla::Array<Value>> disassemble(RewriteContext &ctx,
getOutLayout(*op));
const Layout def_layout = def_layouts[res_idx];
CHECK(def_layout.has_value());
CHECK(def_layout->equivalentTo(layout, std::nullopt, ctx.target_shape));
CHECK(def_layout->generalizes(layout, vty.getShape(), ctx.target_shape));
SmallVector<int64_t> layout_shape =
layout.tileArrayShape(vty.getShape(), ctx.target_shape);
if (auto roll_vectors_op = dyn_cast<RollVectorsOp>(op)) {