[Mlir-commits] [mlir] [mlir] Fix type transformation in DropUnitDimFromElementwiseOps (PR #75430)

[Andrzej Warzyński]

================
@@ -1480,17 +1480,24 @@ struct DropUnitDimFromElementwiseOps final
   using OpTraitRewritePattern::OpTraitRewritePattern;
   LogicalResult matchAndRewrite(Operation *op,
                                 PatternRewriter &rewriter) const override {
-    if (op->getNumResults() != 1)
+    if (op->getNumResults() != 1 || op->getNumRegions() != 0)
       return failure();
 
-    // Check the pre-condiitions. For `Elementwise` Ops all operands
-    // are guaranteed to have identical shapes and it suffices to only check the
-    // first one.
-    auto op1 = op->getOperands()[0];
-    auto sourceVectorType = dyn_cast<VectorType>(op1.getType());
-    if (!sourceVectorType)
+    auto resultVectorType = dyn_cast<VectorType>(op->getResult(0).getType());
+    if (!resultVectorType)
       return failure();
 
+    if (llvm::any_of(op->getOperands(), [](auto operand) {
+          return !isa<VectorType>(operand.getType());
+        })) {
+      return failure();
+    }
----------------
banach-space wrote:

Is this check needed? If all operands are identical then checking `sourceVectorType` should be sufficient, right?

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