[Mlir-commits] [mlir] [MLIR][Linalg] Improve morphism semantics (PR #222959)

[Renato Golin]

================
@@ -90,8 +91,57 @@ static bool findIndexOfScalarOperand(GenericOp genericOp, int &index) {
   return false;
 }
 
-// Attempt to specialize unary or binary linalg.generic ops to named elementwise
-// ops or linalg.elementwise.
+// Maps a linalg.generic body operation to its corresponding elementwise kind,
+// if one exists. Handles the ops with a straightforward one-to-one mapping;
+// ops needing extra context (e.g. boolean add/mul, reciprocal, square) are
+// handled by the caller.
+static std::optional<ElementwiseKind> getElementwiseKind(Operation *op) {
+  return llvm::TypeSwitch<Operation *, std::optional<ElementwiseKind>>(op)
+      .Case<arith::AddFOp, arith::AddIOp, complex::AddOp>(
+          [](Operation *) { return ElementwiseKind::add; })
+      .Case<arith::MulIOp, arith::MulFOp, complex::MulOp>(
+          [](Operation *) { return ElementwiseKind::mul; })
+      .Case([](math::ExpOp) { return ElementwiseKind::exp; })
+      .Case([](math::AbsFOp) { return ElementwiseKind::abs; })
+      .Case([](math::CeilOp) { return ElementwiseKind::ceil; })
+      .Case([](math::FloorOp) { return ElementwiseKind::floor; })
+      .Case([](arith::NegFOp) { return ElementwiseKind::negf; })
+      .Case([](math::RoundOp) { return ElementwiseKind::round; })
+      .Case([](math::SqrtOp) { return ElementwiseKind::sqrt; })
+      .Case([](math::RsqrtOp) { return ElementwiseKind::rsqrt; })
+      .Case([](math::TanhOp) { return ElementwiseKind::tanh; })
+      .Case([](math::ErfOp) { return ElementwiseKind::erf; })
+      .Case([](math::SinOp) { return ElementwiseKind::sin; })
+      .Case([](math::CosOp) { return ElementwiseKind::cos; })
+      .Case([](math::TanOp) { return ElementwiseKind::tan; })
+      .Case([](math::AcosOp) { return ElementwiseKind::acos; })
+      .Case([](math::AcoshOp) { return ElementwiseKind::acosh; })
+      .Case([](math::AsinOp) { return ElementwiseKind::asin; })
+      .Case([](math::AsinhOp) { return ElementwiseKind::asinh; })
+      .Case([](math::AtanOp) { return ElementwiseKind::atan; })
+      .Case([](math::AtanhOp) { return ElementwiseKind::atanh; })
+      .Case([](math::LogOp) { return ElementwiseKind::log; })
+      .Case([](math::Log10Op) { return ElementwiseKind::log10; })
+      .Case([](math::Log1pOp) { return ElementwiseKind::log1p; })
+      .Case([](math::Log2Op) { return ElementwiseKind::log2; })
+      .Case<arith::SubIOp, arith::SubFOp, complex::SubOp>(
+          [](Operation *) { return ElementwiseKind::sub; })
+      .Case<arith::DivSIOp, arith::DivFOp, complex::DivOp>(
+          [](Operation *) { return ElementwiseKind::div; })
+      .Case([](arith::DivUIOp) { return ElementwiseKind::div_unsigned; })
+      .Case<arith::MaxSIOp, arith::MaximumFOp>(
+          [](Operation *) { return ElementwiseKind::max_signed; })
+      .Case<arith::MinSIOp, arith::MinimumFOp>(
+          [](Operation *) { return ElementwiseKind::min_signed; })
+      .Case([](math::PowFOp) { return ElementwiseKind::powf; })
+      .Case([](arith::MaxUIOp) { return ElementwiseKind::max_unsigned; })
+      .Case([](arith::MinUIOp) { return ElementwiseKind::min_unsigned; })
+      .Case([](arith::SelectOp) { return ElementwiseKind::select; })
----------------
rengolin wrote:

Nice catch! I'll fix this.

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