[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);
+
+  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(
+      rewriter, loc, arith::CmpIPredicate::eq, bBatchSize, one);
+
+  auto thenBuilder = [&](OpBuilder &builder, Location nestedLoc) {
+    Value broadcastedB = broadcastMatMulTRhs(rewriter, nestedLoc, a, b);
+    Value castB =
+        builder.createOrFold<tensor::CastOp>(nestedLoc, bType, broadcastedB);
+    scf::YieldOp::create(builder, nestedLoc, castB);
+  };
+  auto elseBuilder = [&](OpBuilder &builder, Location nestedLoc) {
+    scf::YieldOp::create(builder, nestedLoc, b);
+  };
+  return scf::IfOp::create(rewriter, loc, requiresBroadcast, thenBuilder,
+                           elseBuilder)
+      .getResult(0);
+}
+
+class MatMulTConverter : public OpConversionPattern<tosa::MatMulTOp> {
+public:
+  using OpConversionPattern<tosa::MatMulTOp>::OpConversionPattern;
+  LogicalResult
+  matchAndRewrite(tosa::MatMulTOp op, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const final {
+    Location loc = op.getLoc();
+    auto aType = dyn_cast<RankedTensorType>(adaptor.getA().getType());
+    auto bType = dyn_cast<RankedTensorType>(adaptor.getB().getType());
+    auto outputType = dyn_cast<RankedTensorType>(op.getType());
+    if (!aType || !bType || !outputType)
+      return rewriter.notifyMatchFailure(op, "only supports ranked tensors");
+
+    if (aType.getRank() != 3 || bType.getRank() != 3)
+      return rewriter.notifyMatchFailure(
+          op, "only supports rank-3 A and B tensors");
+
+    const int64_t aBatch = aType.getDimSize(0);
+    const int64_t bBatch = bType.getDimSize(0);
+    if (aBatch == 1 && ShapedType::isStatic(bBatch) && bBatch != 1)
+      return rewriter.notifyMatchFailure(
+          op, "does not support broadcasting the A matrix");
+
+    FailureOr<int64_t> maybeAZp = op.getAZeroPoint();
+    FailureOr<int64_t> maybeBZp = op.getBZeroPoint();
+    if (failed(maybeAZp))
+      return rewriter.notifyMatchFailure(
+          op, "input a zero point cannot be statically determined");
----------------
sahas3 wrote:

```suggestion
          op, "input A zero point cannot be statically determined");
```
and similar changes elsewhere 

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