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

[Jack Frankland]

================
@@ -679,6 +680,271 @@ class MatMulConverter : public OpConversionPattern<tosa::MatMulOp> {
   }
 };
 
+// Insert missing leading batch dimensions, taking their sizes from `other`.
+static Value broadcastMatMulTLeadingDimensions(PatternRewriter &rewriter,
+                                               Location loc, Value operand,
+                                               Value other) {
+  auto type = cast<RankedTensorType>(operand.getType());
+  int64_t missingDims =
+      cast<RankedTensorType>(other.getType()).getRank() - type.getRank();
+  if (missingDims <= 0)
+    return operand;
+
+  SmallVector<OpFoldResult> shape;
+  SmallVector<int64_t> dimensions;
+  for (int64_t dim = 0; dim < missingDims; ++dim) {
+    shape.push_back(tensor::getMixedSize(rewriter, loc, other, dim));
+    dimensions.push_back(dim);
+  }
+  llvm::append_range(shape, tensor::getMixedSizes(rewriter, loc, operand));
+  Value empty = tensor::EmptyOp::create(
+      rewriter, loc, shape, type.getElementType(), type.getEncoding());
+  return linalg::BroadcastOp::create(rewriter, loc, operand, empty, dimensions)
+      .getResult()
+      .front();
+}
+
+// Replace a known singleton batch dimension with `size`.
+static Value broadcastMatMulTSingletonDimension(PatternRewriter &rewriter,
----------------
FranklandJack wrote:

Good point, added :)

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