[TPU][Mosaic] Add support for a no-op reshape where sublane_tiling = …

…1 and the res_tiled and src_tiled shapes both fill a full vreg (1024)

PiperOrigin-RevId: 690796348

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

@@ -4112,6 +4112,13 @@ LogicalResult vector_shape_cast_rule(RewriteContext &ctx, Operation &op,
              dst_tiled_dims[1] % dst_vreg_slice[1] == 0) {
     // Shapecast (..., 128) -> (..., m * 128 * packing).
     no_op = true;
+  } else if (layout_in.offsets() == LayoutOffsets{0, 0} &&
+             layout_out.offsets() == LayoutOffsets{0, 0} &&
+             layout_in.tiling()[0] == 1 && layout_out.tiling()[0] == 1 &&
+             src_vreg_slice[1] == dst_vreg_slice[1] &&
+             src_tiled_dims[1] % src_vreg_slice[1] == 0 &&
+             dst_tiled_dims[1] % dst_vreg_slice[1] == 0) {
+    no_op = true;
   }
   FAILUREOR_ASSIGN_OR_RETURN(
       xla::Array<Value> src_vregs,

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

@@ -1448,7 +1448,13 @@ class VectorLayoutInferer {
       // 1D tilings that use 1 in the sublane dimension.
       int64_t sublane_tiling = vreg_slice[0];
       do {
-        if (src_tiled_ishape[1] == res_tiled_ishape[1] &&
+        auto src_res_tiled_equal = src_tiled_ishape[1] == res_tiled_ishape[1];
+        auto vreg_num_elements = target_shape_[0] * target_shape_[1];
+        auto single_subline_mod_1024 =
+            (sublane_tiling == 1 &&
+             src_tiled_ishape[1] % vreg_num_elements == 0 &&
+             res_tiled_ishape[1] % vreg_num_elements == 0);
+        if ((src_res_tiled_equal || single_subline_mod_1024) &&
             src_tiled_ishape[0] % sublane_tiling == 0 &&
             res_tiled_ishape[0] % sublane_tiling == 0) {
           std::array<int64_t, 2> tiling = {sublane_tiling, target_shape_[1]};