[Pallas/Mosaic TPU] Allow non-leading and non-matching batch dimensions in dot_general.

The constraints on `lhs_batch_dims` and `rhs_batch_dims` for `dot_general` in Pallas/Mosaic on TPU are now relaxed. Batch dimensions do not have to be at the front of the shape, and the dimension indices used for batching on the LHS and RHS can be different (which requires to update the semantics of `output_dim_order` in `tpu.dot_dimension_numbers`).

The remaining gap compared to JAX is the lack of support for multiple batch dimensions.

PiperOrigin-RevId: 831015967

Author: yueshengys

jax/_src/pallas/mosaic/lowering.py

@@ -2104,7 +2104,9 @@ def _proxy_fun(val, *, shape, broadcast_dimensions):
   return vector.broadcast(out_type, val)
 
 
-def jax_dot_dims_to_tpu_dot_dot_dims(dimension_numbers, lhs_shape, rhs_shape):
+def jax_dot_dims_to_tpu_dot_dot_dims(
+    ctx, dimension_numbers, lhs_shape, rhs_shape
+):
   """Converts a jax dot dimension numbers to a tpu dot dimension numbers.
 
   Jax dot dimension numbers are given as a tuple of tuples of sequences of ints
@@ -2142,6 +2144,11 @@ def jax_dot_dims_to_tpu_dot_dot_dims(dimension_numbers, lhs_shape, rhs_shape):
     output_dim_order.append(0)
     output_dim_order.append(lhs_dim_map[dim])
 
+  if not (ctx.is_cloud_tpu_older_than(2025, 11, 18) and ctx.forward_compatible):
+    for dim in rhs_batch_dims:
+      output_dim_order.append(1)
+      output_dim_order.append(rhs_dim_map[dim])
+
   for dim in lhs_non_contracting_dims:
     output_dim_order.append(0)
     output_dim_order.append(lhs_dim_map[dim])
@@ -2267,7 +2274,7 @@ def _dot_general_lowering_rule(
     return vector.shape_cast(out_type, red)
 
   tpu_dot_dims = jax_dot_dims_to_tpu_dot_dot_dims(
-      dimension_numbers, lhs_aval.shape, rhs_aval.shape
+      ctx, dimension_numbers, lhs_aval.shape, rhs_aval.shape
   )
 
   if precision is not None:

jaxlib/mosaic/dialect/tpu/tpu_ops.cc

@@ -1114,15 +1114,8 @@ LogicalResult MatmulOp::verify() {
     const std::optional<int64_t> batch_dim_rhs =
         rhs_batch_dims.empty() ? std::nullopt
                                : std::optional<int64_t>(rhs_batch_dims[0]);
-    if (batch_dim_lhs != batch_dim_rhs) {
-      emitOpError("Not Implemented: batch dims must be equal");
-      return failure();
-    }
-    if (batch_dim_lhs.has_value() && (batch_dim_lhs.value() != 0)) {
-      emitOpError("Not Implemented: batch dims pos must be 0");
-      return failure();
-    }
-    // Invariant above enforces only 1 batch dim atm, and that both are eq
+
+    // Invariant above enforces only 1 batch dim atm.
     std::optional<int64_t> batch_size = std::nullopt;
     if (batch_dim_lhs.has_value()) {
       batch_size = lhs_ty.getShape()[batch_dim_lhs.value()];
@@ -1142,30 +1135,26 @@ LogicalResult MatmulOp::verify() {
           "Illegal: output dim order must have an even number of elements.");
       return failure();
     }
-    if (batch_size.has_value()) {
-      if (output_dim_order[0] != 0 || output_dim_order[1] != 0) {
-        emitOpError(
-            "Not implemented: Output with batch size must be the lhs 0 idx for "
-            "now.");
-        return failure();
-      }
-    }
 
-    // Invariants above enforce a single batch idx for now, and that it is in
-    // position 0. Future extensions to this will be to:
-    // 1. Support multiple batch dims
-    // 2. Support batch dims in any position in the output dim order
+    // Invariants above enforce a single batch idx for now. Future extension to
+    // this will be to support multiple batch dims.
 
-    // Verify that the output dim order is always in the form of [0, batch_dims,
-    // 0, lhs_non_contracting_dims, 1, rhs_non_contracting_dims].
+    // Verify that the output dim order is always in the form of [0,
+    // lhs_batch_dims, 1, rhs_batch_dims, 0, lhs_non_contracting_dims, 1,
+    // rhs_non_contracting_dims].
     llvm::SmallVector<int64_t> expected_output_dim_order;
-    expected_output_dim_order.reserve(2 * (lhs_batch_dims.size() +
-                                           lhs_non_contracting_dims.size() +
-                                           rhs_non_contracting_dims.size()));
+    expected_output_dim_order.reserve(
+        2 *
+        (lhs_batch_dims.size() + rhs_batch_dims.size() +
+         lhs_non_contracting_dims.size() + rhs_non_contracting_dims.size()));
     for (int64_t dim : lhs_batch_dims) {
       expected_output_dim_order.push_back(0);
       expected_output_dim_order.push_back(dim);
     }
+    for (int64_t dim : rhs_batch_dims) {
+      expected_output_dim_order.push_back(1);
+      expected_output_dim_order.push_back(dim);
+    }
     for (int64_t dim : lhs_non_contracting_dims) {
       expected_output_dim_order.push_back(0);
       expected_output_dim_order.push_back(dim);
@@ -1177,7 +1166,7 @@ LogicalResult MatmulOp::verify() {
     if (!absl::c_equal(output_dim_order, expected_output_dim_order)) {
       emitOpError(
           "Illegal: output dim order must be in the form of [0, "
-          "batch_dims, 0, lhs_non_contracting_dims, 1, "
+          "lhs_batch_dims, 1, rhs_batch_dims, 0, lhs_non_contracting_dims, 1, "
           "rhs_non_contracting_dims]");
       return failure();
     }

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

@@ -132,24 +132,24 @@ class CanonicalBuilder : public ImplicitLocOpBuilder {
   Operation *op_;
 };
 
-// Ensures both lhs and rhs have contiguous non-contracting and contracting
-// dimensions by inserting transposes if needed. Returns lhs, rhs, and new
-// dimension numbers if a transpose was inserted, otherwise returns
-// std::nullopt.
+// Ensures both lhs and rhs are in form of [batch_dims, non_contracting_dims,
+// contracting_dims] or [batch_dims, contracting_dims, non_contracting_dims] by
+// inserting transposes if needed. Returns lhs, rhs, and new dimension numbers
+// if a transpose was inserted, otherwise returns std::nullopt.
 std::optional<std::tuple<TypedValue<VectorType>, TypedValue<VectorType>,
                          DotDimensionNumbersAttr>>
 ensure_matmul_contiguous_dims(
     CanonicalBuilder& builder, TypedValue<VectorType> lhs,
     TypedValue<VectorType> rhs,
     const DotDimensionNumbersAttr& dimension_numbers) {
-  // Returns a tuple of [new_operand, new_non_contracting_dims,
+  // Returns a tuple of [new_operand, new_batch_dims, new_non_contracting_dims,
   // new_contracting_dims]. new_operand is nullptr if no transpose is inserted.
   auto maybe_insert_transpose =
       [&](TypedValue<VectorType> operand, ArrayRef<int64_t> batch_dims,
           ArrayRef<int64_t> non_contracting_dims,
           ArrayRef<int64_t> contracting_dims, bool is_lhs)
       -> std::tuple<TypedValue<VectorType>, SmallVector<int64_t>,
-                    SmallVector<int64_t>> {
+                    SmallVector<int64_t>, SmallVector<int64_t>> {
     VectorType vty = operand.getType();
     auto shape = vty.getShape();
     auto rank = shape.size();
@@ -170,7 +170,8 @@ ensure_matmul_contiguous_dims(
                     contracting_dims.end());
     // Already in [B..., NC..., C...].
     if (is_identity(perm_BNC)) {
-      return {nullptr, llvm::to_vector(non_contracting_dims),
+      return {nullptr, llvm::to_vector(batch_dims),
+              llvm::to_vector(non_contracting_dims),
               llvm::to_vector(contracting_dims)};
     }
 
@@ -183,7 +184,8 @@ ensure_matmul_contiguous_dims(
                     non_contracting_dims.end());
     // Already in [B..., C..., NC...].
     if (is_identity(perm_BCN)) {
-      return {nullptr, llvm::to_vector(non_contracting_dims),
+      return {nullptr, llvm::to_vector(batch_dims),
+              llvm::to_vector(non_contracting_dims),
               llvm::to_vector(contracting_dims)};
     }
 
@@ -246,18 +248,21 @@ ensure_matmul_contiguous_dims(
     };
 
     // Map the dimension indices to the new dimension order.
+    SmallVector<int64_t> new_b = map_dims(batch_dims);
     SmallVector<int64_t> new_c = map_dims(contracting_dims);
     SmallVector<int64_t> new_nc = map_dims(non_contracting_dims);
 
-    return {new_operand, new_nc, new_c};
+    return {new_operand, new_b, new_nc, new_c};
   };
 
-  auto [new_lhs, new_lhs_non_contracting_dims, new_lhs_contracting_dims] =
+  auto [new_lhs, new_lhs_batch_dims, new_lhs_non_contracting_dims,
+        new_lhs_contracting_dims] =
       maybe_insert_transpose(lhs, dimension_numbers.getLhsBatchDims(),
                              dimension_numbers.getLhsNonContractingDims(),
                              dimension_numbers.getLhsContractingDims(),
                              /*is_lhs=*/true);
-  auto [new_rhs, new_rhs_non_contracting_dims, new_rhs_contracting_dims] =
+  auto [new_rhs, new_rhs_batch_dims, new_rhs_non_contracting_dims,
+        new_rhs_contracting_dims] =
       maybe_insert_transpose(rhs, dimension_numbers.getRhsBatchDims(),
                              dimension_numbers.getRhsNonContractingDims(),
                              dimension_numbers.getRhsContractingDims(),
@@ -267,13 +272,18 @@ ensure_matmul_contiguous_dims(
   }
 
   SmallVector<int64_t> new_output_dim_order;
-  new_output_dim_order.reserve(2 * (dimension_numbers.getLhsBatchDims().size() +
+  new_output_dim_order.reserve(2 * (new_lhs_batch_dims.size() +
+                                    new_rhs_batch_dims.size() +
                                     new_lhs_non_contracting_dims.size() +
                                     new_rhs_non_contracting_dims.size()));
-  for (int64_t batch_dim : dimension_numbers.getLhsBatchDims()) {
+  for (int64_t batch_dim : new_lhs_batch_dims) {
     new_output_dim_order.push_back(0);
     new_output_dim_order.push_back(batch_dim);
   }
+  for (int64_t batch_dim : new_rhs_batch_dims) {
+    new_output_dim_order.push_back(1);
+    new_output_dim_order.push_back(batch_dim);
+  }
   for (int64_t non_contracting_dim : new_lhs_non_contracting_dims) {
     new_output_dim_order.push_back(0);
     new_output_dim_order.push_back(non_contracting_dim);
@@ -286,8 +296,7 @@ ensure_matmul_contiguous_dims(
   DotDimensionNumbersAttr new_dimension_numbers = DotDimensionNumbersAttr::get(
       builder.getContext(), new_lhs_contracting_dims, new_rhs_contracting_dims,
       new_lhs_non_contracting_dims, new_rhs_non_contracting_dims,
-      new_output_dim_order, dimension_numbers.getLhsBatchDims(),
-      dimension_numbers.getRhsBatchDims());
+      new_output_dim_order, new_lhs_batch_dims, new_rhs_batch_dims);
 
   return std::make_tuple(new_lhs ? new_lhs : lhs, new_rhs ? new_rhs : rhs,
                          new_dimension_numbers);
@@ -419,6 +428,10 @@ collapse_matmul_non_contracting_dims(
     new_output_dim_order.push_back(batch_dim);
     new_acc_shape.push_back(lhs.getType().getDimSize(batch_dim));
   }
+  for (int64_t batch_dim : dimension_numbers.getRhsBatchDims()) {
+    new_output_dim_order.push_back(1);
+    new_output_dim_order.push_back(batch_dim);
+  }
   for (int64_t non_contracting_dim : new_lhs_non_contracting_dims) {
     new_output_dim_order.push_back(0);
     new_output_dim_order.push_back(non_contracting_dim);

tests/pallas/ops_test.py

@@ -1498,6 +1498,45 @@ def kernel(x_ref, y_ref, out_ref):
         expected,
     )
 
+  @parameterized.product(
+      shapes_and_dims_numbers=(
+          ((3, 4, 128), (4, 2, 128), (((2,), (2,)), ((1,), (0,)))),
+          ((3, 4, 128), (2, 4, 128), (((2,), (2,)), ((1,), (1,)))),
+          ((3, 4, 256), (2, 3, 256), (((2,), (2,)), ((0,), (1,)))),
+          ((4, 3, 2, 32), (2, 128, 32, 2), (((3,), (2,)), ((2,), (3,)))),
+      ),
+  )
+  def test_dot_general_non_front_batch_dims(self, shapes_and_dims_numbers):
+    if jtu.test_device_matches(["gpu"]):
+      self.skipTest("TPU only test")
+
+    if jtu.test_device_matches(["tpu"]) and not jtu.if_cloud_tpu_at_least(
+        2025, 11, 18
+    ):
+      self.skipTest("Requires libtpu built after 2025-11-18")
+
+    x_shape, y_shape, dims_numbers = shapes_and_dims_numbers
+
+    k1, k2 = random.split(jax.random.key(0))
+    x = jax.random.normal(k1, x_shape, dtype=jnp.float32)
+    y = jax.random.normal(k2, y_shape, dtype=jnp.float32)
+
+    # Just infer shape from jax.
+    expected = jax.lax.dot_general(x, y, dimension_numbers=dims_numbers)
+
+    @functools.partial(
+        self.pallas_call,
+        out_shape=jax.ShapeDtypeStruct(expected.shape, jnp.float32),
+    )
+    def kernel(x_ref, y_ref, out_ref):
+      out_ref[...] = jax.lax.dot_general(
+          x_ref[...],
+          y_ref[...],
+          dimension_numbers=dims_numbers,
+      )
+
+    np.testing.assert_allclose(kernel(x, y), expected, atol=1e-5, rtol=1e-5)
+
   @parameterized.product(
       batch_size=(None, 1, 2),
       # dims_numbers is without batch dims