[Mlir-commits] [mlir] [TOSA] Convert tosa.elementwise to linalg.elementwise (PR #224344)

[Federico Bruzzone]

================
@@ -603,6 +749,84 @@ class DepthwiseConvConverter
   }
 };
 
+// Lower TOSA elementwise ops that have a direct linalg.elementwise named op
+// counterpart. Unlike linalg.generic, linalg.elementwise dispatches on a single
+// function kind and infers the body, so the conversion is a handful of checks
+// (kind, ranks, static shapes) followed by construction of the indexing maps
+// and the named op.
+template <typename TosaOp>
+class ElementwiseConverter : public OpConversionPattern<TosaOp> {
+public:
+  using OpConversionPattern<TosaOp>::OpConversionPattern;
+  using typename OpConversionPattern<TosaOp>::OpAdaptor;
+
+  LogicalResult
+  matchAndRewrite(TosaOp op, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const final {
+    Operation *operation = op;
+    if (!isConvertibleToLinalgElementwise(operation))
+      return rewriter.notifyMatchFailure(
+          op, "op has no linalg.elementwise counterpart");
+
+    Location loc = op.getLoc();
+    auto resultTy = cast<RankedTensorType>(op.getType());
+    const int64_t rank = resultTy.getRank();
+
+    std::optional<linalg::ElementwiseKind> kind =
+        getElementwiseKind(operation, resultTy.getElementType());
+
+    // Trailing metadata operands (e.g. MulOp's shift, NegateOp's zero points)
+    // are not forwarded to linalg.elementwise.
+    ValueRange inputs = adaptor.getOperands().take_front(getArity(operation));
+
+    // Build one indexing map per input, broadcasting static size-1 dimensions,
+    // followed by the identity map for the output.
+    SmallVector<AffineMap> indexingMaps;
+    for (Value operand : inputs) {
+      ArrayRef<int64_t> shape = cast<ShapedType>(operand.getType()).getShape();
+      SmallVector<AffineExpr> exprs;
+      for (const auto &dim : llvm::enumerate(shape)) {
+        bool requiresBroadcast =
+            dim.value() == 1 && resultTy.getDimSize(dim.index()) != 1;
+        exprs.push_back(requiresBroadcast
+                            ? rewriter.getAffineConstantExpr(0)
+                            : rewriter.getAffineDimExpr(dim.index()));
+      }
----------------
FedericoBruzzone wrote:

I'm still learning and want to know what happened at the TOSA level from a high-level POV.

Given that, I think the following could be useful for figuring out what to do:
https://github.com/llvm/llvm-project/blob/be40d49f89d7873a2060e702e9a5b3350aa614e2/mlir/include/mlir/Dialect/Tosa/IR/TosaOpBase.td#L463-L470

- `SameOperandsAndResultRank` is used to guarantee the same rank.
- `ResultsBroadcastableShape` is used to broadcast on size-1 dims only.

So, can we reuse `computeTargetSize`?

P.S. I think that merging functionality could be a follow-up, given the nature of this PR.

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