Support pltpu.roll on sublanes when not all lanes are used.

PiperOrigin-RevId: 743964736

Author: Google-ML-Automation

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

@@ -2209,16 +2209,25 @@ LogicalResult rotate_rule_impl(RewriteContext &ctx, OpTy op, Value amount,
   if (layout_in != layout) {
     return op.emitOpError("Not implemented: unsupported layout for input");
   }
-  if (layout_out != layout) {
+  // We support non-zero offsets in the output layout via lazy rotation.
+  if (layout_out.bitwidth() != layout.bitwidth() ||
+      layout_out.tiling() != layout.tiling() ||
+      layout_out.implicit_dim() != layout.implicit_dim()) {
     return op.emitOpError("Not implemented: unsupported layout for output");
   }
   auto vty = op.getResult().getType();
   if (vty.getRank() < 2) {
     return op.emitOpError("Not implemented: unsupported 1D shape");
   }
-  if (*(vty.getShape().end() - 2) % *(layout.tiling().end() - 2) != 0 ||
-      *(vty.getShape().end() - 1) % *(layout.tiling().end() - 1) != 0) {
-    return op.emitOpError("Not implemented: unsupported unaliged shape");
+  if (*(vty.getShape().end() - 2) % *(layout.tiling().end() - 2) != 0 &&
+      op.getDimension() == vty.getRank() - 2) {
+    return op.emitOpError(
+        "Not implemented: unsupported unaligned shape in sublane dimension");
+  }
+  if (*(vty.getShape().end() - 1) % *(layout.tiling().end() - 1) != 0 &&
+      op.getStride().has_value()) {
+    return op.emitOpError(
+        "Not implemented: unsupported unaligned shape in lane dimension");
   }
 
   ImplicitLocOpBuilder builder(op.getLoc(), op.getOperation());
@@ -2345,6 +2354,63 @@ LogicalResult rotate_rule_impl(RewriteContext &ctx, OpTy op, Value amount,
     return concatenate(chunks, axis);
   };
 
+  // Applies lazy rotation (see go/pltpu-roll for details).
+  auto lazyRotate = [&](const xla::Array<Value> &vregs, int64_t shift,
+                        int axis) {
+    const int tiling_dim = axis - (vregs.num_dimensions() - 2);
+    const int64_t tile_size = ctx.target_shape[tiling_dim];
+    const int64_t input_size = vty.getShape()[axis];
+    const int64_t normalized_shift = shift % input_size;
+    const int64_t start_idx = input_size - normalized_shift;
+    const int64_t start_vreg_idx = start_idx / tile_size;
+    const int64_t valid_amount = input_size % tile_size;
+
+    auto concat = concatenate({vregs, vregs}, axis);
+    auto chunks = split(concat, axis);
+    int64_t original_num_chunks = chunks.size() / 2;
+    xla::Array<Value> front_chunk_copy(chunks.front());
+
+    Value rotate_amount = mlirI32Const(valid_amount);
+    auto iota = builder.create<tpu::IotaOp>(
+        i32_vreg, builder.getI32IntegerAttr(tiling_dim));
+    auto mask = builder.create<arith::CmpIOp>(
+        arith::CmpIPredicate::sge, iota,
+        builder.create<arith::ConstantOp>(DenseElementsAttr::get(
+            i32_vreg, builder.getI32IntegerAttr(valid_amount))));
+    // overwrite padding in the last vreg with valid data from the first vreg.
+    chunks.back().Each([&](absl::Span<const int64_t> idxs, Value *v) {
+      *v = builder.create<arith::SelectOp>(
+          mask,
+          builder.create<tpu::DynamicRotateOp>(
+              res_vreg_ty, front_chunk_copy(idxs), rotate_amount, tiling_dim,
+              nullptr, nullptr),
+          *v);
+    });
+    // rotate the vregs starting from the middle vreg.
+    for (int64_t i = original_num_chunks; i < chunks.size(); ++i) {
+      chunks[i].Each([&](absl::Span<const int64_t> idxs, Value *v) {
+        *v = builder.create<tpu::DynamicRotateOp>(
+            res_vreg_ty, *v, rotate_amount, tiling_dim, nullptr, nullptr);
+      });
+    }
+    // blend the vregs to overwrite the padding.
+    for (int64_t i = original_num_chunks - 1; i < chunks.size() - 1; ++i) {
+      chunks[i].Each([&](absl::Span<const int64_t> idxs, Value *v) {
+        *v = builder.create<arith::SelectOp>(mask, chunks[i + 1](idxs), *v);
+      });
+    }
+    SmallVector<int64_t> result_dimensions =
+        layout_out.tileArrayImplicitShape(vty.getShape(), ctx.target_shape);
+    // assemble the result
+    xla::Array<Value> result(result_dimensions);
+    SmallVector<int64_t> starts(result.num_dimensions(), 0);
+    for (int64_t i = 0; i < result_dimensions[axis]; ++i) {
+      starts[axis] = i;
+      result.UpdateSlice(chunks[i + start_vreg_idx], starts);
+    }
+    return result;
+  };
+
   std::function<xla::Array<Value>(const xla::Array<Value> &, Value, int, int)>
       rotate;
   rotate = [&](const xla::Array<Value> &vregs, Value shift, int axis,
@@ -2353,7 +2419,15 @@ LogicalResult rotate_rule_impl(RewriteContext &ctx, OpTy op, Value amount,
     CHECK(axis >= 0 && axis < vregs.num_dimensions());
     int tiling_dim = axis - (vregs.num_dimensions() - 2);
     CHECK((tiling_dim != 1 && stride == 0) || (tiling_dim == 1 && stride >= 0));
+    const bool has_padding =
+        (tiling_dim == 0 || tiling_dim == 1) &&
+        vty.getShape()[axis] % ctx.target_shape[tiling_dim] != 0;
     SmallVector<xla::Array<Value>, 4> chunks;
+    // Handle rotation with static shift and padding lazily.
+    if (auto shift_cst = getIntConst(shift, /*silent=*/true);
+        succeeded(shift_cst) && has_padding) {
+      return lazyRotate(vregs, shift_cst.value(), axis);
+    }
     // Handle rotation with static shift.
     if (auto shift_cst = getIntConst(shift, /*silent=*/true);
         succeeded(shift_cst)) {
@@ -2445,7 +2519,7 @@ LogicalResult rotate_rule_impl(RewriteContext &ctx, OpTy op, Value amount,
       roll_by *= 2;
     }
     return result;
-  };
+  };  // end of rotate
 
   xla::Array<Value> out_tiles(in_tiles.dimensions());
   const auto dim = op.getDimension();

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

@@ -748,6 +748,46 @@ class VectorLayoutInferer {
     return success();
   }
 
+  // Helper function to compute the layout offsets for a dynamic rotate op.
+  LayoutOffsets compute_layout_offsets(tpu::DynamicRotateOp op) {
+    LayoutOffsets layout_offsets{0, 0};
+    const unsigned int bitwidth = op.getType().getElementTypeBitWidth();
+    const auto tiling = nativeTiling(bitwidth);
+    const int tiling_dim = op.getDimension() - (op.getType().getRank() - 2);
+    if (tiling_dim != 0 && tiling_dim != 1) {
+      return layout_offsets;
+    }
+    const int64_t tile_size = tiling[tiling_dim];
+    const int64_t dim_size = op.getType().getShape()[op.getDimension()];
+    if (dim_size % tile_size == 0) {
+      return layout_offsets;
+    }
+    auto amount = op.getAmount().getDefiningOp<arith::ConstantOp>();
+    if (!amount) {
+      return layout_offsets;
+    }
+    auto integer_attr = dyn_cast<IntegerAttr>(amount.getValue());
+    if (!integer_attr) {
+      return layout_offsets;
+    }
+    if (auto stride = op.getStride(); stride.has_value() && *stride != 0) {
+      return layout_offsets;
+    }
+    if (tiling_dim != 0 && tiling_dim != 1) {
+      return layout_offsets;
+    }
+    int64_t shift_amount = integer_attr.getValue().getSExtValue();
+    // Normalize the shift amount to the dimension size.
+    shift_amount = shift_amount % dim_size;
+    CHECK_GE(shift_amount, 0);
+    CHECK_LE(shift_amount, dim_size);
+    // Absolute offset.
+    int64_t offset = dim_size - shift_amount;
+    // Convert to relative offsets within the tile.
+    layout_offsets[tiling_dim] = offset % tile_size;
+    return layout_offsets;
+  }
+
   LogicalResult infer(tpu::DynamicRotateOp op) {
     auto bitwidth = op.getType().getElementTypeBitWidth();
     // TODO(b/347067057): Support dynamic rotate with packed dtype.
@@ -759,7 +799,9 @@ class VectorLayoutInferer {
     }
     auto layout = VectorLayout(bitwidth, {0, 0}, nativeTiling(bitwidth),
                                ImplicitDim::kNone);
-    setLayout(op, {layout, kNoLayout}, layout);
+    auto out_layout = VectorLayout(bitwidth, compute_layout_offsets(op),
+                                   nativeTiling(bitwidth), ImplicitDim::kNone);
+    setLayout(op, {layout, kNoLayout}, out_layout);
     return success();
   }
 

tests/pallas/tpu_pallas_test.py

@@ -2844,9 +2844,7 @@ def kernel(x_ref, out_ref):
     )(x)
     np.testing.assert_array_equal(out, state_utils.bitcast(x, jnp.uint32))
 
-  @only_passes_in_interpret()
   def test_roll_partial(self):
-    """b/337384645"""
     x = np.arange(8192, dtype=jnp.float32).reshape(128, 64)
 
     def kernel(x_ref, out_ref):