[Mlir-commits] [mlir] [mlir][tosa] Lower MATMUL_T to Linalg (PR #222661)

[Luke Hutton]

================
@@ -679,6 +680,185 @@ class MatMulConverter : public OpConversionPattern<tosa::MatMulOp> {
   }
 };
 
+// Broadcast the batch dimension of B from one to the batch dimension of A.
+// linalg.broadcast adds dimensions rather than expanding dimensions of size
+// one, so first rank-reduce B from 1xWxC to WxC.
+static Value broadcastMatMulTRhs(PatternRewriter &rewriter, Location loc,
+                                 Value a, Value b) {
+  auto bType = cast<RankedTensorType>(b.getType());
+  auto rankReducedBType =
+      RankedTensorType::get(bType.getShape().drop_front(),
+                            bType.getElementType(), bType.getEncoding());
+
+  SmallVector<OpFoldResult> offsets(3, rewriter.getIndexAttr(0));
+  SmallVector<OpFoldResult> sizes = {rewriter.getIndexAttr(1),
+                                     tensor::getMixedSize(rewriter, loc, b, 1),
+                                     tensor::getMixedSize(rewriter, loc, b, 2)};
+  SmallVector<OpFoldResult> strides(3, rewriter.getIndexAttr(1));
+  Value rankReducedB = tensor::ExtractSliceOp::create(
+      rewriter, loc, rankReducedBType, b, offsets, sizes, strides);
+
+  SmallVector<OpFoldResult> broadcastShape = {
+      tensor::getMixedSize(rewriter, loc, a, 0),
+      tensor::getMixedSize(rewriter, loc, b, 1),
+      tensor::getMixedSize(rewriter, loc, b, 2)};
+  Value broadcastInit =
+      tensor::EmptyOp::create(rewriter, loc, broadcastShape,
+                              bType.getElementType(), bType.getEncoding());
+  return linalg::BroadcastOp::create(rewriter, loc, rankReducedB, broadcastInit,
+                                     ArrayRef<int64_t>{0})
+      .getResult()
+      .front();
+}
+
+// Materialize MATMUL_T's optional B batch broadcast. When the B batch is
+// dynamic, select the broadcasted value at runtime if its size is one.
+static Value normalizeMatMulTRhsBatch(PatternRewriter &rewriter, Location loc,
+                                      Value a, Value b) {
+  auto aType = cast<RankedTensorType>(a.getType());
+  auto bType = cast<RankedTensorType>(b.getType());
+  const int64_t aBatch = aType.getDimSize(0);
+  const int64_t bBatch = bType.getDimSize(0);
+
+  if (bBatch == 1 && aBatch != 1)
+    return broadcastMatMulTRhs(rewriter, loc, a, b);
+
+  if (ShapedType::isStatic(bBatch) || aBatch == 1)
+    return b;
+
+  Value bBatchSize = tensor::DimOp::create(rewriter, loc, b, 0);
+  Value one = arith::ConstantIndexOp::create(rewriter, loc, 1);
+  Value requiresBroadcast = arith::CmpIOp::create(
----------------
lhutton1 wrote:

Thanks I think it makes sense to use the utility - we could always improve the utility if there's a more optimal lowering

https://github.com/llvm/llvm-project/pull/222661