Add an optimization pattern that tranform const<[a, 1]> @ <[1, b]> to <[1, b]> * const<[a, 1]>.

sirakiin · tensorflower-gardener · commit 2e79b0980310 · 2025-07-22T21:51:47.000-07:00
PiperOrigin-RevId: 786124861
diff --git a/tensorflow/compiler/mlir/lite/transforms/optimize_patterns.td b/tensorflow/compiler/mlir/lite/transforms/optimize_patterns.td
@@ -2220,3 +2220,20 @@ def DegenerateFCtoMul : Pat<
    (IsLastDimensionEqualOne $input),
    (SameElementType $input, $filter),
    (IsNoneType $bias)]>;
+
+// Replace matmul where inputs & weights have inner dimension of 1 with an
+// elementwise multiplication that broadcasts, i.e. replace:
+//   const [a, 1] @ [1, b] => [a, b]
+// with:
+//   [1, b] * const [a, 1] => [a, b]
+def BMMAsBroadCastMultoFc : Pat<
+  (TFL_BatchMatMulOp
+    (Arith_ConstantOp:$const_input $constInputVal),
+    $filter, ConstBoolAttrFalse, ConstBoolAttrFalse, ConstBoolAttrFalse),
+  (TFL_MulOp
+    $filter, $const_input,
+    TFL_AF_None),
+  [(HasRank<2> $const_input),
+   (HasRank<2> $filter),
+   (IsLastDimensionEqualOne $const_input),
+   (SameElementType $const_input, $filter)]>;
