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

[Sayan Saha]

================
@@ -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);
----------------
sahas3 wrote:

I think you can use the util `tensor::createCanonicalRankReducingExtractSliceOp` here. 

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