From 93ba4a135b0e87fecf19b979c5148a6fc809e4b5 Mon Sep 17 00:00:00 2001 From: "A. Unique TensorFlower" Date: Wed, 21 Oct 2026 07:48:57 -0700 Subject: [PATCH] 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 | 88 +++++++++++++++++++ tensorflow/compiler/mlir/lite/ir/tfl_ops.td | 2 + .../compiler/mlir/lite/tests/const-fold.mlir | 77 ++++++++++++++++ 3 files changed, 167 insertions(+) diff --git a/tensorflow/compiler/mlir/lite/ir/tfl_ops.cc b/tensorflow/compiler/mlir/lite/ir/tfl_ops.cc index 4cd59932c74..f6c569f48c6 100644 --- a/tensorflow/compiler/mlir/lite/ir/tfl_ops.cc +++ b/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(adaptor.getX()); + auto y = mlir::dyn_cast_or_null(adaptor.getY()); + if (!x || !y) return {}; + if (!ShouldFoldOperation(this->getOperation())) return {}; + + auto out_type = mlir::dyn_cast(getType()); + auto x_type = mlir::dyn_cast(x.getType()); + auto y_type = mlir::dyn_cast(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 x_shape = x_type.getShape(); + const ArrayRef y_shape = y_type.getShape(); + const ArrayRef 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 out_batch = out_shape.drop_back(2); + const ArrayRef x_batch = x_shape.drop_back(2); + const ArrayRef 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 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(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 x_vals(x.getValues().begin(), + x.getValues().end()); + const std::vector y_vals(y.getValues().begin(), + y.getValues().end()); + std::vector 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(acc); + } + } + } + return DenseElementsAttr::get(out_type, llvm::ArrayRef(result)); +} + //===----------------------------------------------------------------------===// // BlockwiseQuantizeOp / BlockwiseDequantizeOp //===----------------------------------------------------------------------===// diff --git a/tensorflow/compiler/mlir/lite/ir/tfl_ops.td b/tensorflow/compiler/mlir/lite/ir/tfl_ops.td index 7a2778dc978..aeb8703956a 100644 --- a/tensorflow/compiler/mlir/lite/ir/tfl_ops.td +++ b/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; } diff --git a/tensorflow/compiler/mlir/lite/tests/const-fold.mlir b/tensorflow/compiler/mlir/lite/tests/const-fold.mlir index dcdda38f84b..d67b8cf7153 100644 --- a/tensorflow/compiler/mlir/lite/tests/const-fold.mlir +++ b/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]] +}