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

[Federico Bruzzone]

================
@@ -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; })
----------------
FedericoBruzzone wrote:

Please correct me if I'm missing something here :D

The `SelectOp` was not present in the previous implementation. 
The documentation of `specializeLinalgElementwise` specifies:
https://github.com/llvm/llvm-project/blob/e1a4305b37f20140819a1c768cd873f93add61b5/mlir/lib/Dialect/Linalg/Transforms/Specialize.cpp#L143-L144
But with this PR, we will end up using it also with the `SelectOp`, which is neither unary nor binary: 
https://github.com/llvm/llvm-project/blob/e1a4305b37f20140819a1c768cd873f93add61b5/mlir/lib/Dialect/Linalg/Transforms/Specialize.cpp#L239-L246
My concern is about the swap within the `replaceOp` closure:
https://github.com/llvm/llvm-project/blob/e1a4305b37f20140819a1c768cd873f93add61b5/mlir/lib/Dialect/Linalg/Transforms/Specialize.cpp#L192-L195
Because `hasSwappedOperands` does not considers ternary ops:
https://github.com/llvm/llvm-project/blob/e1a4305b37f20140819a1c768cd873f93add61b5/mlir/lib/Dialect/Linalg/Transforms/Specialize.cpp#L180

Should ternary ops be reordered here too?
For instance, consider `select %c, %f, %t` over `ins(%C, %T, %F)` becomes `elementwise <select>(C, T, F)`, which is fine. But what if the condition is the second input?




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