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

[Chuanqi Xu]

================
@@ -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()));
+      }
+      indexingMaps.push_back(AffineMap::get(rank, /*symbolCount=*/0, exprs,
+                                            rewriter.getContext()));
+    }
+    indexingMaps.push_back(rewriter.getMultiDimIdentityMap(rank));
+
+    // Compute the output sizes, taking each dynamic result dimension from the
+    // first input that is not broadcast (statically 1) along that dimension.
+    SmallVector<OpFoldResult> sizes;
+    for (int64_t i = 0; i < rank; ++i) {
+      if (!resultTy.isDynamicDim(i)) {
+        sizes.push_back(rewriter.getIndexAttr(resultTy.getDimSize(i)));
+        continue;
+      }
+      Value dimSource = inputs.front();
+      for (Value operand : inputs) {
+        auto operandTy = cast<ShapedType>(operand.getType());
+        if (operandTy.getRank() == rank && operandTy.getDimSize(i) != 1) {
----------------
ChuanqiXu9 wrote:

What if the dim size is dynamic ?

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