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

[Renato Golin]

================
@@ -18,15 +18,161 @@
 #include "mlir/Dialect/Tosa/IR/TosaOps.h"
 #include "mlir/Dialect/Tosa/Utils/ConversionUtils.h"
 #include "mlir/Dialect/Utils/ReshapeOpsUtils.h"
+#include "mlir/IR/Matchers.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/Transforms/DialectConversion.h"
 #include "llvm/ADT/SmallVectorExtras.h"
+#include "llvm/ADT/TypeSwitch.h"
 
+#include <optional>
 #include <type_traits>
 
 using namespace mlir;
 using namespace mlir::tosa;
 
+// Number of leading data operands the op forwards to linalg.elementwise,
+// matching its unary/binary/ternary arity group. Trailing metadata operands
+// (e.g. MulOp's shift, NegateOp's zero points) are excluded.
+static unsigned getArity(Operation *op) {
+  return llvm::TypeSwitch<Operation *, unsigned>(op)
+      .Case<tosa::AbsOp, tosa::CeilOp, tosa::FloorOp, tosa::ExpOp, tosa::LogOp,
+            tosa::RsqrtOp, tosa::SinOp, tosa::CosOp, tosa::TanhOp, tosa::ErfOp,
+            tosa::ReciprocalOp, tosa::NegateOp>([](Operation *) { return 1u; })
+      .Case<tosa::AddOp, tosa::SubOp, tosa::IntDivOp, tosa::MulOp, tosa::PowOp,
+            tosa::MaximumOp, tosa::MinimumOp>([](Operation *) { return 2u; })
+      .Case<tosa::SelectOp>([](Operation *) { return 3u; })
----------------
rengolin wrote:

Even better, we don't need the helper at all. Turns out I always use them _after_ getting the linalg elementwise kind, so can use the linalg helper to figure that out. Code changed.

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