[Mosaic] Fix handling of i1 splat constants

PiperOrigin-RevId: 694248723
jax-ml · Nov 7, 2024 · 04a6652 · 04a6652
1 parent 3b2e4a1
commit 04a6652
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 2 deletions.
diff --git a/jaxlib/mosaic/dialect/tpu/transforms/apply_vector_layout.cc b/jaxlib/mosaic/dialect/tpu/transforms/apply_vector_layout.cc
@@ -3190,8 +3190,8 @@ LogicalResult arith_constant_rule(RewriteContext &ctx, Operation &op,
     }
     const VectorLayout &layout_out = *layouts_out.front();
     DenseElementsAttr value = cast<DenseElementsAttr>(constant_op.getValue());
-    const VectorType target_vty =
-        getNativeVregType(vty.getElementType(), ctx.target_shape);
+    const VectorType target_vty = getNativeVregOrVmaskType(
+        vty.getElementType(), layout_out.bitwidth(), ctx.target_shape);
     if (value.isSplat()) {
       if (layout_out.offsets() != LayoutOffsets{std::nullopt, std::nullopt}) {
         return op.emitOpError(

diff --git a/jaxlib/mosaic/dialect/tpu/transforms/infer_vector_layout.cc b/jaxlib/mosaic/dialect/tpu/transforms/infer_vector_layout.cc
@@ -365,6 +365,11 @@ class VectorLayoutInferer {
       TPU_CHECK_OP(ty.getRank() > 0, "rank 0 vectors unsupported");
       TPU_CHECK_OP(elems, "expected vector constants to use DenseElementsAttr");
       auto bitwidth = ty.getElementTypeBitWidth();
+      if (bitwidth == 1) {
+        // i1 is a special case where the layout bitwidth can be different from
+        // the element bitwidth, see comment in VectorLayout class
+        bitwidth = kNativeBitwidth;
+      }
       if (elems.isSplat()) {
         if (ty.getRank() == 1) {
           // Here, we choose to lay out along lanes arbitrarily. It would be