[Mlir-commits] [mlir] f28f09d - [mlir][Vector] Add Broadcast -> CastOp reordering to SinkVectorBroadcasting patterns. (#68257)

[llvmlistbot at llvm.org]

Author: MaheshRavishankar
Date: 2023-10-04T21:27:24-07:00
New Revision: f28f09dcf0b4a0b5b56cbce54e997d7c395ce926

URL: https://github.com/llvm/llvm-project/commit/f28f09dcf0b4a0b5b56cbce54e997d7c395ce926
DIFF: https://github.com/llvm/llvm-project/commit/f28f09dcf0b4a0b5b56cbce54e997d7c395ce926.diff

LOG: [mlir][Vector] Add Broadcast -> CastOp reordering to SinkVectorBroadcasting patterns. (#68257)

Also fix an issue with sink broadcast across elementwise where
`arith.cmpf` is elementwise, but result type is different. The result
type is not same as the operand type, creating illegal IR.
Similar issue with `vector.fma` which only accepts vector operand types,
while broadcasts can have scalar sources. Sinking broadcast across would
result in an illegal `vector.fma` (with scalar operands).

Added: 
    

Modified: 
    mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp
    mlir/test/Dialect/Vector/sink-vector-broadcast.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp b/mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp
index 77076090f354e6b..16f9d7e4d57eef7 100644
--- a/mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp
@@ -912,6 +912,15 @@ struct ReorderElementwiseOpsOnBroadcast final
       return failure();
     if (!OpTrait::hasElementwiseMappableTraits(op))
       return failure();
+    if (op->getNumOperands() == 0 ||
+        op->getResults()[0].getType() != op->getOperand(0).getType()) {
+      return failure();
+    }
+    // Avoid operations that only accept vector types, since broadcast
+    // source might be scalar types.
+    if (isa<vector::FMAOp>(op)) {
+      return failure();
+    }
 
     // Get the type of the lhs operand
     auto *lhsBcastOrSplat = op->getOperand(0).getDefiningOp();
@@ -1447,8 +1456,8 @@ void mlir::vector::
 
 void mlir::vector::populateSinkVectorBroadcastPatterns(
     RewritePatternSet &patterns, PatternBenefit benefit) {
-  patterns.add<ReorderElementwiseOpsOnBroadcast>(patterns.getContext(),
-                                                 benefit);
+  patterns.add<ReorderCastOpsOnBroadcast, ReorderElementwiseOpsOnBroadcast>(
+      patterns.getContext(), benefit);
 }
 
 //===----------------------------------------------------------------------===//

diff  --git a/mlir/test/Dialect/Vector/sink-vector-broadcast.mlir b/mlir/test/Dialect/Vector/sink-vector-broadcast.mlir
index d9d2f44e6f16c1f..e7863a9e8b7b780 100644
--- a/mlir/test/Dialect/Vector/sink-vector-broadcast.mlir
+++ b/mlir/test/Dialect/Vector/sink-vector-broadcast.mlir
@@ -105,3 +105,23 @@ func.func @broadcast_not_elementwise() -> vector<2x2xf32> {
 
   return %mm1 : vector<2x2xf32>
 }
+
+// CHECK-LABEL: func.func @dont_sink_cmp(
+//       CHECK:   %[[BROADCAST:.+]] = vector.broadcast
+//       CHECK:   %[[RETURN:.+]] = arith.cmpf uno, %[[BROADCAST]], %[[BROADCAST]]
+//       CHECK:   return %[[RETURN]]
+func.func @dont_sink_cmp(%arg0 : f32, %arg1 : vector<1xf32>) -> vector<1xi1> {
+  %0 = vector.broadcast %arg0 : f32 to vector<1xf32>
+  %1 = arith.cmpf uno, %0, %0 : vector<1xf32>
+  return %1 : vector<1xi1>
+}
+
+// CHECK-LABEL: func.func @dont_sink_fma(
+  //     CHECK:   %[[BROADCAST:.+]] = vector.broadcast
+  //     CHECK:   %[[RESULT:.+]] = vector.fma %[[BROADCAST]]
+  //     CHECK:   return %[[RESULT]]
+func.func @dont_sink_fma(%arg0 : f32) -> vector<1xf32> {
+  %0 = vector.broadcast %arg0 : f32 to vector<1xf32>
+  %1 = vector.fma %0, %0, %0 : vector<1xf32>
+  return %1 : vector<1xf32>
+}