[mlir][vector] add pattern to cast away leading unit dim for elementwise op

[mravishankar]

Just one minor typo (unrelated to this change, but good to fix)

mlir/lib/Dialect/Vector/VectorTransforms.cpp

3840

Just noticed typo here CastAwayBrodcastLeadingOneDim -> CastAwayBroadcastLeadingOneDim

[ThomasRaoux]

Fix typo

mlir/lib/Dialect/Vector/VectorTransforms.cpp

3840

Thanks for catching this.

[stella.stamenova]

mlir/lib/Dialect/Vector/VectorTransforms.cpp

3791

The default statement that you added in your previous change is causing the following warning:

/mnt/vss/_work/5/s/mlir/lib/Dialect/Vector/VectorTransforms.cpp:3790:7: error: default label in switch which covers all enumeration values [-Werror,-Wcovered-switch-default]
      default:

I am not sure what warning was there before, but there's definitely one now. Can you have a look?

[ThomasRaoux]

mlir/lib/Dialect/Vector/VectorTransforms.cpp

3791

Thanks for pointing that out. Somehow this warning doesn't show up for me locally. This commit should fix it:
https://github.com/llvm/llvm-project/commit/6aaf06f92988c6e2b91f90ab7ed3a6d21981480a

mlir/lib/Dialect/Vector/VectorTransforms.cpp

auto newVector = rewriter.create<vector::ShapeCastOp>(
         write.getLoc(), newType, write.vector());
     rewriter.replaceOpWithNewOp<vector::TransferWriteOp>(
         write, newVector, write.source(), write.indices(), newMap, inBounds);
 
     return success();
   }
 };
 
-struct CastAwayBrodcastLeadingOneDim
+struct CastAwayBroadcastLeadingOneDim
     : public OpRewritePattern<vector::BroadcastOp> {
   using OpRewritePattern::OpRewritePattern;
 
   LogicalResult matchAndRewrite(vector::BroadcastOp broadcastOp,
                                 PatternRewriter &rewriter) const override {
     VectorType newDstType = trimLeadingOneDims(broadcastOp.getVectorType());
     if (newDstType == broadcastOp.getVectorType())
       return failure();
     Location loc = broadcastOp.getLoc();
     VectorType srcVecType = broadcastOp.getSourceType().dyn_cast<VectorType>();
     if (srcVecType)
       srcVecType = trimLeadingOneDims(srcVecType);
     Value source = broadcastOp.source();
     if (srcVecType && srcVecType != broadcastOp.getSourceType()) {
       source = rewriter.create<vector::ShapeCastOp>(loc, srcVecType, source);
     }
     Value newBroadcastOp =
         rewriter.create<vector::BroadcastOp>(loc, newDstType, source);
     rewriter.replaceOpWithNewOp<vector::ShapeCastOp>(
         broadcastOp, broadcastOp.getVectorType(), newBroadcastOp);
     return success();
   }
 };
 
+class CastAwayElementwiseLeadingOneDim : public RewritePattern {
+public:
+  CastAwayElementwiseLeadingOneDim(MLIRContext *context)
+      : RewritePattern(MatchAnyOpTypeTag(), /*benefit=*/1, context) {}
+
+  LogicalResult matchAndRewrite(Operation *op,
+                                PatternRewriter &rewriter) const override {
+    if (!OpTrait::hasElementwiseMappableTraits(op) || op->getNumResults() != 1)
+      return failure();
+    auto vecType = op->getResultTypes()[0].dyn_cast<VectorType>();
+    if (!vecType)
+      return failure();
+    VectorType newVecType = trimLeadingOneDims(vecType);
+    if (newVecType == vecType)
+      return failure();
+
+    SmallVector<Value, 4> newOperands;
+    for (Value operand : op->getOperands()) {
+      if (auto opVecType = operand.getType().dyn_cast<VectorType>()) {
+        auto newType =
+            VectorType::get(newVecType.getShape(), opVecType.getElementType());
+        newOperands.push_back(rewriter.create<vector::ShapeCastOp>(
+            op->getLoc(), newType, operand));
+      } else {
+        newOperands.push_back(operand);
+      }
+    }
+    OperationState state(op->getLoc(), op->getName());
+    state.addAttributes(op->getAttrs());
+    state.addOperands(newOperands);
+    state.addTypes(newVecType);
+    Operation *newOp = rewriter.createOperation(state);
+    rewriter.replaceOpWithNewOp<vector::ShapeCastOp>(op, vecType,
+                                                     newOp->getResult(0));
+    return success();
+  }
+};
+
 // Returns the values in `arrayAttr` as an integer vector.
 static SmallVector<int64_t, 4> getIntValueVector(ArrayAttr arrayAttr) {
   return llvm::to_vector<4>(
       llvm::map_range(arrayAttr.getAsRange<IntegerAttr>(),
                       [](IntegerAttr attr) { return attr.getInt(); }));
 }
 
 // Shuffles vector.bitcast op after vector.extract op.
auto getKindStr = [](vector::CombiningKind kind) {
         return "max";
       case vector::CombiningKind::AND:
         return "and";
       case vector::CombiningKind::OR:
         return "or";
       case vector::CombiningKind::XOR:
         return "xor";
       default:
         llvm_unreachable("unknown combining kind");
       }
     };
 
     for (int i = 0; i < outerDim; ++i) {
       auto v = rewriter.create<vector::ExtractOp>(
           loc, multiReductionOp.source(), ArrayRef<int64_t>{i});
       auto reducedValue = rewriter.create<vector::ReductionOp>(
           loc, multiReductionOp.getDestVectorType().getElementType(),
void mlir::vector::populateSplitVectorTransferPatterns(
     RewritePatternSet &patterns,
     std::function<bool(Operation *)> ignoreFilter) {
   patterns.add<SplitTransferReadOp, SplitTransferWriteOp>(patterns.getContext(),
                                                           ignoreFilter);
 }
 
 void mlir::vector::populateCastAwayVectorLeadingOneDimPatterns(
     RewritePatternSet &patterns) {
-  patterns.add<CastAwayExtractStridedSliceLeadingOneDim,
-               CastAwayInsertStridedSliceLeadingOneDim,
-               CastAwayTransferReadLeadingOneDim,
-               CastAwayTransferWriteLeadingOneDim,
-               CastAwayBrodcastLeadingOneDim, ShapeCastOpFolder>(
+  patterns
+      .add<CastAwayExtractStridedSliceLeadingOneDim,
+           CastAwayTransferWriteLeadingOneDim, CastAwayBroadcastLeadingOneDim,
+           CastAwayElementwiseLeadingOneDim, ShapeCastOpFolder>(
           patterns.getContext());
 }
 
 void mlir::vector::populateBubbleVectorBitCastOpPatterns(
     RewritePatternSet &patterns) {
   patterns.add<BubbleDownVectorBitCastForExtract,
                BubbleDownBitCastForStridedSliceExtract,
                BubbleUpBitCastForStridedSliceInsert>(patterns.getContext());
 }

mlir/test/Dialect/Vector/vector-transforms.mlir

func @cast_away_broadcast_leading_one_dims(
   %1 = vector.broadcast %arg1 : f32 to vector<1x1x4xf32>
   // CHECK:  vector.shape_cast %{{.*}} : vector<1x4xf32> to vector<4xf32>
   // CHECK:  vector.broadcast %{{.*}} : vector<4xf32> to vector<3x4xf32>
   // CHECK:  vector.shape_cast %{{.*}} : vector<3x4xf32> to vector<1x3x4xf32>
   %2 = vector.broadcast %arg2 : vector<1x4xf32> to vector<1x3x4xf32>
   return %0, %1, %2: vector<1x1x8xf32>, vector<1x1x4xf32>, vector<1x3x4xf32>
 }
 
+// CHECK-LABEL: func @cast_away_elementwise_leading_one_dims
+func @cast_away_elementwise_leading_one_dims(
+  %arg0: vector<1x1x8xf32>, %arg1: f32, %arg2: vector<1x4xf32>,
+  %arg3: vector<1x4xf32>, %arg4: i1) ->
+  (vector<1x1x8xf32>, vector<1x4xi1>, vector<1x4xf32>, vector<1x4xf32>) {
+  // CHECK:  vector.shape_cast %{{.*}} : vector<1x1x8xf32> to vector<8xf32>
+  // CHECK:  vector.shape_cast %{{.*}} : vector<1x1x8xf32> to vector<8xf32>
+  // CHECK:  addf %{{.*}}, %{{.*}} : vector<8xf32>
+  // CHECK:  vector.shape_cast %{{.*}} : vector<8xf32> to vector<1x1x8xf32>
+  %0 = addf %arg0, %arg0 : vector<1x1x8xf32>
+  // CHECK:  vector.shape_cast %{{.*}} : vector<1x4xf32> to vector<4xf32>
+  // CHECK:  vector.shape_cast %{{.*}} : vector<1x4xf32> to vector<4xf32>
+  // CHECK:  cmpf ogt, %{{.*}}, %{{.*}} : vector<4xf32>
+  // CHECK:  vector.shape_cast %{{.*}} : vector<4xi1> to vector<1x4xi1>
+  %1 = cmpf ogt, %arg2, %arg3 : vector<1x4xf32>
+  // CHECK:  vector.shape_cast %{{.*}} : vector<1x4xf32> to vector<4xf32>
+  // CHECK:  vector.shape_cast %{{.*}} : vector<1x4xf32> to vector<4xf32>
+  // CHECK:  select %{{.*}}, %{{.*}}, %{{.*}} : vector<4xi1>, vector<4xf32>
+  // CHECK:  vector.shape_cast %{{.*}} : vector<4xf32> to vector<1x4xf32>
+  %2 = select %1, %arg3, %arg2 : vector<1x4xi1>, vector<1x4xf32>
+  // CHECK:  vector.shape_cast %{{.*}} : vector<1x4xf32> to vector<4xf32>
+  // CHECK:  vector.shape_cast %{{.*}} : vector<1x4xf32> to vector<4xf32>
+  // CHECK:  select %arg4, %12, %{{.*}} : vector<4xf32>
+  // CHECK:  vector.shape_cast %{{.*}} : vector<4xf32> to vector<1x4xf32>
+  %3 = select %arg4, %arg3, %arg2 : vector<1x4xf32>
+  return %0, %1, %2, %3: vector<1x1x8xf32>, vector<1x4xi1>, vector<1x4xf32>, vector<1x4xf32>
+}
+
 // CHECK-LABEL: func @bubble_down_bitcast_in_extract
 //  CHECK-SAME: %[[SRC:.+]]: vector<4xf32>
 func @bubble_down_bitcast_in_extract(%src: vector<4xf32>) -> (f16, f16) {
   %0 = vector.bitcast %src : vector<4xf32> to vector<8xf16>
   // CHECK: %[[EXTRACT1:.+]] = vector.extract %[[SRC]][1] : vector<4xf32>
   // CHECK:    %[[CAST1:.+]] = vector.bitcast %[[EXTRACT1]] : vector<1xf32> to vector<2xf16>
   // CHECK: %[[EXTRACT2:.+]] = vector.extract %[[CAST1]][1] : vector<2xf16>
   %1 = vector.extract %0[3] : vector<8xf16>