[mlir][vector] vector.fma is not ElementwiseMappable

[Hardcode84]

ElementwiseMappable implies Scalarizable and Tensorizable but vector.fma only supports vector inputs.

[llvmbot]

@llvm/pr-subscribers-mlir-vector

Author: Ivan Butygin (Hardcode84)

Changes

ElementwiseMappable implies Scalarizable and Tensorizable but vector.fma only supports vector inputs.

---

Full diff: https://github.com/llvm/llvm-project/pull/132611.diff

2 Files Affected:

• (modified) mlir/include/mlir/Dialect/Vector/IR/VectorOps.td (+3-2)
• (modified) mlir/lib/Dialect/Vector/Transforms/VectorUnroll.cpp (+2-1)

diff --git a/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td b/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td
index fbbf817ecff98..c895e32839acd 100644
--- a/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td
+++ b/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td
@@ -759,8 +759,9 @@ def Vector_ExtractOp :
 def Vector_FMAOp :
   Op<Vector_Dialect, "fma", [
        Pure, AllTypesMatch<["lhs", "rhs", "acc", "result"]>,
-       DeclareOpInterfaceMethods<VectorUnrollOpInterface, ["getShapeForUnroll"]>
-     ] # ElementwiseMappable.traits>,
+       DeclareOpInterfaceMethods<VectorUnrollOpInterface, ["getShapeForUnroll"]>,
+       Elementwise, Vectorizable
+     ] >,
     Arguments<(ins VectorOfAnyRankOf<[AnyFloat]>:$lhs,
                    VectorOfAnyRankOf<[AnyFloat]>:$rhs,
                    VectorOfAnyRankOf<[AnyFloat]>:$acc)>,
diff --git a/mlir/lib/Dialect/Vector/Transforms/VectorUnroll.cpp b/mlir/lib/Dialect/Vector/Transforms/VectorUnroll.cpp
index 04c38f9f7b2e3..2d3febffddc74 100644
--- a/mlir/lib/Dialect/Vector/Transforms/VectorUnroll.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/VectorUnroll.cpp
@@ -434,7 +434,8 @@ struct UnrollElementwisePattern : public RewritePattern {
 
   LogicalResult matchAndRewrite(Operation *op,
                                 PatternRewriter &rewriter) const override {
-    if (!OpTrait::hasElementwiseMappableTraits(op) || op->getNumResults() != 1)
+    if (!op->hasTrait<OpTrait::Elementwise>() ||
+        !op->hasTrait<OpTrait::Vectorizable>() || op->getNumResults() != 1)
       return failure();
     auto targetShape = getTargetShape(options, op);
     if (!targetShape)

[llvmbot]

@llvm/pr-subscribers-mlir

Author: Ivan Butygin (Hardcode84)

Changes

ElementwiseMappable implies Scalarizable and Tensorizable but vector.fma only supports vector inputs.

---

Full diff: https://github.com/llvm/llvm-project/pull/132611.diff

2 Files Affected:

• (modified) mlir/include/mlir/Dialect/Vector/IR/VectorOps.td (+3-2)
• (modified) mlir/lib/Dialect/Vector/Transforms/VectorUnroll.cpp (+2-1)

diff --git a/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td b/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td
index fbbf817ecff98..c895e32839acd 100644
--- a/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td
+++ b/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td
@@ -759,8 +759,9 @@ def Vector_ExtractOp :
 def Vector_FMAOp :
   Op<Vector_Dialect, "fma", [
        Pure, AllTypesMatch<["lhs", "rhs", "acc", "result"]>,
-       DeclareOpInterfaceMethods<VectorUnrollOpInterface, ["getShapeForUnroll"]>
-     ] # ElementwiseMappable.traits>,
+       DeclareOpInterfaceMethods<VectorUnrollOpInterface, ["getShapeForUnroll"]>,
+       Elementwise, Vectorizable
+     ] >,
     Arguments<(ins VectorOfAnyRankOf<[AnyFloat]>:$lhs,
                    VectorOfAnyRankOf<[AnyFloat]>:$rhs,
                    VectorOfAnyRankOf<[AnyFloat]>:$acc)>,
diff --git a/mlir/lib/Dialect/Vector/Transforms/VectorUnroll.cpp b/mlir/lib/Dialect/Vector/Transforms/VectorUnroll.cpp
index 04c38f9f7b2e3..2d3febffddc74 100644
--- a/mlir/lib/Dialect/Vector/Transforms/VectorUnroll.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/VectorUnroll.cpp
@@ -434,7 +434,8 @@ struct UnrollElementwisePattern : public RewritePattern {
 
   LogicalResult matchAndRewrite(Operation *op,
                                 PatternRewriter &rewriter) const override {
-    if (!OpTrait::hasElementwiseMappableTraits(op) || op->getNumResults() != 1)
+    if (!op->hasTrait<OpTrait::Elementwise>() ||
+        !op->hasTrait<OpTrait::Vectorizable>() || op->getNumResults() != 1)
       return failure();
     auto targetShape = getTargetShape(options, op);
     if (!targetShape)

[Hardcode84]

And as a side question, why do we even need vector.fma? Looks like it's 1:1 equivalent of math.fma on vectors, except math.fma actually does support scalars and tensors.

[kuhar]

And as a side question, why do we even need vector.fma? Looks like it's 1:1 equivalent of math.fma on vectors, except math.fma actually does support scalars and tensors.

+1. I'd be in favor of deprecating it and eventually dropping.

[banach-space]

And as a side question, why do we even need vector.fma? Looks like it's 1:1 equivalent of math.fma on vectors, except math.fma actually does support scalars and tensors.

+1. I'd be in favor of deprecating it and eventually dropping.

I don't have an opinion just yet.

It would be good to identify the original rationale for introducing vector.fma - why didn't people use math.fma? Also, removing vector.fma would mean introducing a dependency of Vector on Math - perhaps that's the reason vector.fma exists?

So, in principle, no objections (removing redundancy is good), but lets make sure it's an informed decision.

[dcaballe]

IIRC, math.fma represents a library implementation of FMA, which would result in a library call, whereas vector.fma is expected to be lowered to an fma LLVM or hardware intrinsic...