Constant-fold BATCH_MATMUL with constant f32 operands.

`BatchMatMulOp` had no folder, so a matmul of two constants survived into the
flatbuffer and was recomputed on every inference. Add `BatchMatMulOp::fold`
for f32 operands with static shapes, honoring `adj_x` / `adj_y` and
right-aligned batch broadcasting, subject to the usual `ShouldFoldOperation`
size limit.

PiperOrigin-RevId: 1003936085

tensorflow/compiler/mlir/lite/ir/tfl_ops.cc

@@ -2014,6 +2014,94 @@ LogicalResult BatchMatMulOp::verify() {
   return success();
 }
 
+// Folds a BATCH_MATMUL whose operands are both f32 constants. Left unfolded,
+// such an op survives into the flatbuffer and is recomputed on every
+// inference; delegates that require a runtime input also reject it (e.g. ML
+// Drift), splitting the graph. This shows up in DETR-style decoders, where
+// step-0 self-attention runs on constant query embeddings.
+OpFoldResult BatchMatMulOp::fold(FoldAdaptor adaptor) {
+  auto x = mlir::dyn_cast_or_null<DenseElementsAttr>(adaptor.getX());
+  auto y = mlir::dyn_cast_or_null<DenseElementsAttr>(adaptor.getY());
+  if (!x || !y) return {};
+  if (!ShouldFoldOperation(this->getOperation())) return {};
+
+  auto out_type = mlir::dyn_cast<RankedTensorType>(getType());
+  auto x_type = mlir::dyn_cast<RankedTensorType>(x.getType());
+  auto y_type = mlir::dyn_cast<RankedTensorType>(y.getType());
+  if (!out_type || !x_type || !y_type || !out_type.hasStaticShape() ||
+      !x_type.hasStaticShape() || !y_type.hasStaticShape() ||
+      !out_type.getElementType().isF32() || !x_type.getElementType().isF32() ||
+      !y_type.getElementType().isF32() || x_type.getRank() < 2 ||
+      y_type.getRank() < 2 || out_type.getRank() < 2) {
+    return {};
+  }
+
+  const ArrayRef<int64_t> x_shape = x_type.getShape();
+  const ArrayRef<int64_t> y_shape = y_type.getShape();
+  const ArrayRef<int64_t> out_shape = out_type.getShape();
+  const bool adj_x = getAdjX();
+  const bool adj_y = getAdjY();
+  const int64_t m = adj_x ? x_shape.back() : x_shape[x_shape.size() - 2];
+  const int64_t k = adj_x ? x_shape[x_shape.size() - 2] : x_shape.back();
+  const int64_t k_y = adj_y ? y_shape.back() : y_shape[y_shape.size() - 2];
+  const int64_t n = adj_y ? y_shape[y_shape.size() - 2] : y_shape.back();
+  if (k != k_y || out_shape[out_shape.size() - 2] != m ||
+      out_shape.back() != n) {
+    return {};
+  }
+
+  // Batch dims broadcast right-aligned; the verifier guarantees compatibility.
+  const ArrayRef<int64_t> out_batch = out_shape.drop_back(2);
+  const ArrayRef<int64_t> x_batch = x_shape.drop_back(2);
+  const ArrayRef<int64_t> y_batch = y_shape.drop_back(2);
+  // Computed from the batch dims directly (rather than dividing the element
+  // count by m * n) so zero-sized m or n cannot divide by zero.
+  const int64_t num_batches = ShapedType::getNumElements(out_batch);
+
+  // ShouldFoldOperation bounds tensor sizes but not compute; cap the
+  // multiply-accumulate count so canonicalization stays fast.
+  constexpr int64_t kMaxFoldMacs = int64_t{1} << 30;
+  if (num_batches * m * n * k > kMaxFoldMacs) return {};
+
+  // Maps a flat output batch index to the flat batch index of an operand.
+  auto operand_batch = [&](int64_t out_flat, ArrayRef<int64_t> batch) {
+    int64_t flat = 0;
+    int64_t out_stride = num_batches;
+    int64_t offset = out_batch.size() - batch.size();
+    for (int64_t i = 0; i < static_cast<int64_t>(out_batch.size()); ++i) {
+      out_stride /= out_batch[i];
+      const int64_t idx = (out_flat / out_stride) % out_batch[i];
+      const int64_t j = i - offset;
+      if (j >= 0) flat = flat * batch[j] + (batch[j] == 1 ? 0 : idx);
+    }
+    return flat;
+  };
+
+  const std::vector<float> x_vals(x.getValues<float>().begin(),
+                                  x.getValues<float>().end());
+  const std::vector<float> y_vals(y.getValues<float>().begin(),
+                                  y.getValues<float>().end());
+  std::vector<float> result(out_type.getNumElements(), 0.0f);
+  for (int64_t b = 0; b < num_batches; ++b) {
+    const float* xb = x_vals.data() + operand_batch(b, x_batch) * m * k;
+    const float* yb = y_vals.data() + operand_batch(b, y_batch) * k * n;
+    float* ob = result.data() + b * m * n;
+    for (int64_t i = 0; i < m; ++i) {
+      for (int64_t j = 0; j < n; ++j) {
+        // Accumulate in double to limit rounding error for large k.
+        double acc = 0.0;
+        for (int64_t p = 0; p < k; ++p) {
+          const double xv = adj_x ? xb[p * m + i] : xb[i * k + p];
+          const double yv = adj_y ? yb[j * k + p] : yb[p * n + j];
+          acc += xv * yv;
+        }
+        ob[i * n + j] = static_cast<float>(acc);
+      }
+    }
+  }
+  return DenseElementsAttr::get(out_type, llvm::ArrayRef<float>(result));
+}
+
 //===----------------------------------------------------------------------===//
 // BlockwiseQuantizeOp / BlockwiseDequantizeOp
 //===----------------------------------------------------------------------===//

tensorflow/compiler/mlir/lite/ir/tfl_ops.td

@@ -1246,6 +1246,8 @@ in the batch dimensions and broadcasting.
 
   let hasVerifier = 1;
 
+  let hasFolder = 1;
+
   let extraClassDeclaration = [{
     // DynamicRangeQuantizedOpInterface:
     bool RequireAsymmetricQuantizeInputsAttr() { return true; }

tensorflow/compiler/mlir/lite/tests/const-fold.mlir

@@ -1315,3 +1315,80 @@ func.func @slice_big_float() -> tensor<1x1x1792x256xf32> {
 
 
 
+
+// CHECK-LABEL: @ConstFoldBatchMatMul
+func.func @ConstFoldBatchMatMul() -> tensor<2x2x2xf32> {
+  %x = arith.constant dense<[[[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]], [[1.0, 0.0, 0.0], [0.0, 1.0, 0.0]]]> : tensor<2x2x3xf32>
+  %y = arith.constant dense<[[[1.0, 0.0], [0.0, 1.0], [1.0, 1.0]], [[2.0, 3.0], [4.0, 5.0], [6.0, 7.0]]]> : tensor<2x3x2xf32>
+  %0 = "tfl.batch_matmul"(%x, %y) {adj_x = false, adj_y = false} : (tensor<2x2x3xf32>, tensor<2x3x2xf32>) -> tensor<2x2x2xf32>
+  func.return %0 : tensor<2x2x2xf32>
+  // CHECK: %[[CST:.*]] = arith.constant dense<{{\[\[\[}}4.000000e+00, 5.000000e+00], [1.000000e+01, 1.100000e+01]], {{\[\[}}2.000000e+00, 3.000000e+00], [4.000000e+00, 5.000000e+00]]]> : tensor<2x2x2xf32>
+  // CHECK-NOT: tfl.batch_matmul
+  // CHECK: return %[[CST]]
+}
+
+// CHECK-LABEL: @ConstFoldBatchMatMulAdjY
+func.func @ConstFoldBatchMatMulAdjY() -> tensor<1x2x2xf32> {
+  %x = arith.constant dense<[[[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]]> : tensor<1x2x3xf32>
+  %y = arith.constant dense<[[[1.0, 0.0, 1.0], [0.0, 1.0, 1.0]]]> : tensor<1x2x3xf32>
+  %0 = "tfl.batch_matmul"(%x, %y) {adj_x = false, adj_y = true} : (tensor<1x2x3xf32>, tensor<1x2x3xf32>) -> tensor<1x2x2xf32>
+  func.return %0 : tensor<1x2x2xf32>
+  // CHECK: %[[CST:.*]] = arith.constant dense<{{\[\[\[}}4.000000e+00, 5.000000e+00], [1.000000e+01, 1.100000e+01]]]> : tensor<1x2x2xf32>
+  // CHECK-NOT: tfl.batch_matmul
+  // CHECK: return %[[CST]]
+}
+
+// CHECK-LABEL: @ConstFoldBatchMatMulAdjX
+func.func @ConstFoldBatchMatMulAdjX() -> tensor<1x2x2xf32> {
+  %x = arith.constant dense<[[[1.0, 4.0], [2.0, 5.0], [3.0, 6.0]]]> : tensor<1x3x2xf32>
+  %y = arith.constant dense<[[[1.0, 0.0], [0.0, 1.0], [1.0, 1.0]]]> : tensor<1x3x2xf32>
+  %0 = "tfl.batch_matmul"(%x, %y) {adj_x = true, adj_y = false} : (tensor<1x3x2xf32>, tensor<1x3x2xf32>) -> tensor<1x2x2xf32>
+  func.return %0 : tensor<1x2x2xf32>
+  // CHECK: %[[CST:.*]] = arith.constant dense<{{\[\[\[}}4.000000e+00, 5.000000e+00], [1.000000e+01, 1.100000e+01]]]> : tensor<1x2x2xf32>
+  // CHECK-NOT: tfl.batch_matmul
+  // CHECK: return %[[CST]]
+}
+
+// Rank-2 rhs broadcasts over the lhs batch dim.
+// CHECK-LABEL: @ConstFoldBatchMatMulBroadcastBatch
+func.func @ConstFoldBatchMatMulBroadcastBatch() -> tensor<2x2x2xf32> {
+  %x = arith.constant dense<[[[1.0, 2.0], [3.0, 4.0]], [[5.0, 6.0], [7.0, 8.0]]]> : tensor<2x2x2xf32>
+  %y = arith.constant dense<[[0.0, 1.0], [1.0, 0.0]]> : tensor<2x2xf32>
+  %0 = "tfl.batch_matmul"(%x, %y) {adj_x = false, adj_y = false} : (tensor<2x2x2xf32>, tensor<2x2xf32>) -> tensor<2x2x2xf32>
+  func.return %0 : tensor<2x2x2xf32>
+  // CHECK: %[[CST:.*]] = arith.constant dense<{{\[\[\[}}2.000000e+00, 1.000000e+00], [4.000000e+00, 3.000000e+00]], {{\[\[}}6.000000e+00, 5.000000e+00], [8.000000e+00, 7.000000e+00]]]> : tensor<2x2x2xf32>
+  // CHECK-NOT: tfl.batch_matmul
+  // CHECK: return %[[CST]]
+}
+
+// CHECK-LABEL: @NotFoldBatchMatMulRuntimeInput
+func.func @NotFoldBatchMatMulRuntimeInput(%arg0: tensor<1x2x3xf32>) -> tensor<1x2x2xf32> {
+  %y = arith.constant dense<[[[1.0, 0.0], [0.0, 1.0], [1.0, 1.0]]]> : tensor<1x3x2xf32>
+  %0 = "tfl.batch_matmul"(%arg0, %y) {adj_x = false, adj_y = false} : (tensor<1x2x3xf32>, tensor<1x3x2xf32>) -> tensor<1x2x2xf32>
+  func.return %0 : tensor<1x2x2xf32>
+  // CHECK: %[[BMM:.*]] = "tfl.batch_matmul"(%arg0
+  // CHECK: return %[[BMM]]
+}
+
+// Zero-sized m must not divide by zero.
+// CHECK-LABEL: @ConstFoldBatchMatMulZeroSized
+func.func @ConstFoldBatchMatMulZeroSized() -> tensor<1x0x2xf32> {
+  %x = arith.constant dense<> : tensor<1x0x3xf32>
+  %y = arith.constant dense<[[[1.0, 0.0], [0.0, 1.0], [1.0, 1.0]]]> : tensor<1x3x2xf32>
+  %0 = "tfl.batch_matmul"(%x, %y) {adj_x = false, adj_y = false} : (tensor<1x0x3xf32>, tensor<1x3x2xf32>) -> tensor<1x0x2xf32>
+  func.return %0 : tensor<1x0x2xf32>
+  // CHECK: %[[CST:.*]] = arith.constant dense<> : tensor<1x0x2xf32>
+  // CHECK-NOT: tfl.batch_matmul
+  // CHECK: return %[[CST]]
+}
+
+// 2 x 1024 x 1024 x 1024 = 2^31 multiply-accumulates exceeds the fold budget.
+// CHECK-LABEL: @NotFoldBatchMatMulTooExpensive
+func.func @NotFoldBatchMatMulTooExpensive() -> tensor<2x1024x1024xf32> {
+  %x = arith.constant dense<1.0> : tensor<2x1024x1024xf32>
+  %y = arith.constant dense<1.0> : tensor<1024x1024xf32>
+  %0 = "tfl.batch_matmul"(%x, %y) {adj_x = false, adj_y = false} : (tensor<2x1024x1024xf32>, tensor<1024x1024xf32>) -> tensor<2x1024x1024xf32>
+  func.return %0 : tensor<2x1024x1024xf32>
+  // CHECK: %[[BMM:.*]] = "tfl.batch_matmul"
+  // CHECK: return %[[BMM]]
+}