[Mosaic TPU] Allow using replicated lane offset for intermediate safe offsets in mask relayout since relayout like {0,*},(8,128) -> {0,*},(4,128) has been just supported.

PiperOrigin-RevId: 826555654

Author: yueshengys

jaxlib/mosaic/dialect/tpu/transforms/apply_vector_layout.cc

@@ -8434,7 +8434,7 @@ FailureOr<std::pair<VectorLayout, xla::Array<Value>>> changeTiling(
     retiled.Each([&](absl::Span<const int64_t> dst_idx, Value* const vreg) {
       // Recall that this is a one-to-many vregs relayout. Each destination vreg
       // will hold a different part of the source data.
-      // For example, a (8, 128) -> (4, 128) retiling with replicated 2nd minor:
+      // For example, a (8, 128) -> (4, 128) retiling with replicated minor:
       // - Destination vreg 0 comes from the first 4 sublanes of source vreg 0,
       //   with gather pattern [0, 1, 2, 3, 0, 1, 2, 3]
       // - Destination vreg 1 comes from the last 4 sublanes of source vreg 0,

jaxlib/mosaic/dialect/tpu/transforms/relayout_insertion.cc

@@ -31,6 +31,7 @@ limitations under the License.
 #include "mlir/IR/Visitors.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Support/LLVM.h"
+#include "mlir/Support/WalkResult.h"
 #include "jaxlib/mosaic/dialect/tpu/layout.h"
 #include "jaxlib/mosaic/dialect/tpu/tpu_dialect.h"
 #include "jaxlib/mosaic/dialect/tpu/util.h"
@@ -138,9 +139,8 @@ FailureOr<TypedValue<VectorType>> relayout(
     auto safe_offsets = LayoutOffsets{
         src.offsets()[0].has_value() ? *src.offsets()[0] % safe_vreg_slice[0]
                                      : LayoutOffset(),
-        src.offsets()[1].has_value()
-            ? *src.offsets()[1] % safe_vreg_slice[1]
-            : 0,  // TODO(b/452689987): change to LayoutOffset() after resolved.
+        src.offsets()[1].has_value() ? *src.offsets()[1] % safe_vreg_slice[1]
+                                     : LayoutOffset(),
     };
     auto safe_src = VectorLayout(src.bitwidth(), safe_offsets, safe_tiling,
                                  dst.implicit_dim());
@@ -281,4 +281,4 @@ std::unique_ptr<OperationPass<func::FuncOp>> createRelayoutInsertionPass(
                                                  target_shape);
 }
 
-}  // namespace mlir::tpu
+}  // namespace mlir::tpu

tests/pallas/tpu_ops_test.py

@@ -668,7 +668,8 @@ def else_0():
     self.assertEqual(output, 0)
 
   def test_retiling_with_replicated_lane(self):
-    self.skipTest("TODO(b/452689987)")
+    if not jtu.if_cloud_tpu_at_least(2025, 11, 5):
+      self.skipTest("Test requires libtpu from 2025/11/5 or later")
     shape = (32, 1)
     broadcast_shape = (32, 256)