[mlir][Vector] Support 0-D vectors in `BroadcastOp`

[michalt]

Summary

This changes the op to produce AnyVectorOfAnyRank following mostly the code for 1-D vectors.

Depends On D114598

[nicolasvasilache]

Generally looks good, thank you !
Please address the comments and this'll be good to go.

mlir/lib/Conversion/StandardToLLVM/StandardToLLVM.cpp

713 ↗

(On Diff #389552)

SplatOp requires its own independent tests that are missing here.

Can you please have a first revision that changes the semantics of SplatOp and adds a proper test?
Then you can stack this on top and keep things orthogonal.

756 ↗

(On Diff #389552)

<= 1 somehow feels more natural / in line with how we've been writing this type of condition throughout the codebase.

mlir/lib/Dialect/Vector/VectorTransforms.cpp

561

Not really in practice though: if you forward the 0-d around you will get invalid IR until all possible consumers of broadcast are updated to support 0-d.
I'd drop this micro-optimization.

[michalt]

Address review comments

[nicolasvasilache]

LGTM, thanks!

This will prob. also need to be rebased since I couldn't land the parent.

[michalt]

Rebase

[michalt]

Rebase

mlir/include/mlir/Dialect/Vector/VectorOps.td

Vector_Op<"multi_reduction", [NoSideEffect,
      PredOpTrait<"source operand and result have same element type",
                  TCresVTEtIsSameAsOpBase<0, 0>>]>,
     Arguments<(ins Vector_CombiningKindAttr:$kind,
                    AnyVector:$source,
                    I64ArrayAttr:$reduction_dims)>,
     Results<(outs AnyType:$dest)> {
   let summary = "Multi-dimensional reduction operation";
   let description = [{
-    Reduces an n-D vector into an (n-k)-D vector (or a scalar when k == n) 
+    Reduces an n-D vector into an (n-k)-D vector (or a scalar when k == n)
     using the given operation (add/mul/min/max for int/fp and and/or/xor for
     int only).
 
     Example:
 
     ```mlir
     %1 = vector.multi_reduction "add", %0 [1, 3] :
       vector<4x8x16x32xf32> into vector<4x16xf32>
def Vector_MultiDimReductionOp :
   let hasFolder = 1;
 }
 
 def Vector_BroadcastOp :
   Vector_Op<"broadcast", [NoSideEffect,
      PredOpTrait<"source operand and result have same element type",
                  TCresVTEtIsSameAsOpBase<0, 0>>]>,
     Arguments<(ins AnyType:$source)>,
-    Results<(outs AnyVector:$vector)> {
+    Results<(outs AnyVectorOfAnyRank:$vector)> {
   let summary = "broadcast operation";
   let description = [{
     Broadcasts the scalar or k-D vector value in the source operand
     to a n-D result vector such that the broadcast makes sense, i.e.,
     the source operand is duplicated to match the given rank and sizes
     in the result vector. The legality rules are:
     * the source operand must have the same element type as the result type
     * a k-D vector <s_1 x .. x s_k x type> can be broadcast to

mlir/lib/Dialect/Vector/VectorTransforms.cpp

public:
 
   LogicalResult matchAndRewrite(vector::BroadcastOp op,
                                 PatternRewriter &rewriter) const override {
     auto loc = op.getLoc();
     VectorType dstType = op.getVectorType();
     VectorType srcType = op.getSourceType().dyn_cast<VectorType>();
     Type eltType = dstType.getElementType();
 
+    // Scalar to any vector can use splat.
+    if (!srcType) {
+      rewriter.replaceOpWithNewOp<SplatOp>(op, dstType, op.source());
+      return success();
+    }
+
     // Determine rank of source and destination.
-    int64_t srcRank = srcType ? srcType.getRank() : 0;
+    int64_t srcRank = srcType.getRank();
     int64_t dstRank = dstType.getRank();
 
+    // Stretching scalar inside vector (e.g. vector<1xf32>) can use splat.
+    if (srcRank <= 1 && dstRank == 1) {
+      Value ext;
+      if (srcRank == 0)
+        ext = rewriter.create<vector::ExtractElementOp>(loc, op.source());
+      else
+        ext = rewriter.create<vector::ExtractOp>(loc, op.source(), 0);
+      rewriter.replaceOpWithNewOp<SplatOp>(op, dstType, ext);
+      return success();
+    }
+
     // Duplicate this rank.
     // For example:
     //   %x = broadcast %y  : k-D to n-D, k < n
     // becomes:
     //   %b = broadcast %y  : k-D to (n-1)-D
     //   %x = [%b,%b,%b,%b] : n-D
     // becomes:
     //   %b = [%y,%y]       : (n-1)-D
     //   %x = [%b,%b,%b,%b] : n-D
     if (srcRank < dstRank) {
-      // Scalar to any vector can use splat.
-      if (srcRank == 0) {
-        rewriter.replaceOpWithNewOp<SplatOp>(op, dstType, op.source());
-        return success();
-      }
       // Duplication.
       VectorType resType =
           VectorType::get(dstType.getShape().drop_front(), eltType);
       Value bcst =
           rewriter.create<vector::BroadcastOp>(loc, resType, op.source());
       Value result = rewriter.create<arith::ConstantOp>(
           loc, dstType, rewriter.getZeroAttr(dstType));
       for (int64_t d = 0, dim = dstType.getDimSize(0); d < dim; ++d)
for (int64_t r = 0; r < dstRank; r++)
       }
 
     // All trailing dimensions are the same. Simply pass through.
     if (m == -1) {
       rewriter.replaceOp(op, op.source());
       return success();
     }
 
-    // Stretching scalar inside vector (e.g. vector<1xf32>) can use splat.
-    if (srcRank == 1) {
-      assert(m == 0);
-      Value ext = rewriter.create<vector::ExtractOp>(loc, op.source(), 0);
-      rewriter.replaceOpWithNewOp<SplatOp>(op, dstType, ext);
-      return success();
-    }
-
     // Any non-matching dimension forces a stretch along this rank.
     // For example:
     //   %x = broadcast %y : vector<4x1x2xf32> to vector<4x2x2xf32>
     // becomes:
     //   %a = broadcast %y[0] : vector<1x2xf32> to vector<2x2xf32>
     //   %b = broadcast %y[1] : vector<1x2xf32> to vector<2x2xf32>
     //   %c = broadcast %y[2] : vector<1x2xf32> to vector<2x2xf32>
     //   %d = broadcast %y[3] : vector<1x2xf32> to vector<2x2xf32>

mlir/test/Conversion/VectorToLLVM/vector-to-llvm.mlir

 
 // CHECK-LABEL: @bitcast_index_to_i8_vector
 // CHECK-SAME:  %[[input:.*]]: vector<16xindex>
 // CHECK:       %[[T0:.*]] = builtin.unrealized_conversion_cast %[[input]] : vector<16xindex> to vector<16xi64>
 // CHECK:       llvm.bitcast %[[T0]] : vector<16xi64> to vector<128xi8>
 
 // -----
 
+func @broadcast_vec0d_from_f32(%arg0: f32) -> vector<f32> {
+  %0 = vector.broadcast %arg0 : f32 to vector<f32>
+  return %0 : vector<f32>
+}
+// CHECK-LABEL: @broadcast_vec0d_from_f32
+// CHECK-SAME:  %[[A:.*]]: f32)
+// CHECK:       %[[T0:.*]] = splat %[[A]] : vector<f32>
+// CHECK:       return %[[T0]] : vector<f32>
+
+// -----
+
+func @broadcast_vec0d_from_vec0d(%arg0: vector<f32>) -> vector<f32> {
+  %0 = vector.broadcast %arg0 : vector<f32> to vector<f32>
+  return %0 : vector<f32>
+}
+// CHECK-LABEL: @broadcast_vec0d_from_vec0d(
+// CHECK-SAME:  %[[A:.*]]: vector<f32>)
+// CHECK:       return %[[A]] : vector<f32>
+
+// -----
+
 func @broadcast_vec1d_from_f32(%arg0: f32) -> vector<2xf32> {
   %0 = vector.broadcast %arg0 : f32 to vector<2xf32>
   return %0 : vector<2xf32>
 }
 // CHECK-LABEL: @broadcast_vec1d_from_f32
 // CHECK-SAME:  %[[A:.*]]: f32)
 // CHECK:       %[[T0:.*]] = splat %[[A]] : vector<2xf32>
 // CHECK:       return %[[T0]] : vector<2xf32>
func @broadcast_vec1d_from_vec1d(%arg0: vector<2xf32>) -> vector<2xf32> {
   return %0 : vector<2xf32>
 }
 // CHECK-LABEL: @broadcast_vec1d_from_vec1d(
 // CHECK-SAME:  %[[A:.*]]: vector<2xf32>)
 // CHECK:       return %[[A]] : vector<2xf32>
 
 // -----
 
+func @broadcast_vec2d_from_vec0d(%arg0: vector<f32>) -> vector<3x2xf32> {
+  %0 = vector.broadcast %arg0 : vector<f32> to vector<3x2xf32>
+  return %0 : vector<3x2xf32>
+}
+// CHECK-LABEL: @broadcast_vec2d_from_vec0d(
+// CHECK-SAME:  %[[A:.*]]: vector<f32>)
+//       CHECK: %[[T0:.*]] = builtin.unrealized_conversion_cast %[[A]] : vector<f32> to vector<1xf32>
+//       CHECK: %[[T1:.*]] = arith.constant dense<0.000000e+00> : vector<3x2xf32>
+//       CHECK: %[[T2:.*]] = builtin.unrealized_conversion_cast %[[T1]] : vector<3x2xf32> to !llvm.array<3 x vector<2xf32>>
+//       CHECK: %[[T4:.*]] = llvm.mlir.constant(0 : index) : i64
+//       CHECK: %[[T5:.*]] = llvm.extractelement %[[T0]][%[[T4]] : i64] : vector<1xf32>
+//       CHECK: %[[T6:.*]] = splat %[[T5]] : vector<2xf32>
+//       CHECK: %[[T7:.*]] = llvm.insertvalue %[[T6]], %[[T2]][0] : !llvm.array<3 x vector<2xf32>>
+//       CHECK: %[[T8:.*]] = llvm.insertvalue %[[T6]], %[[T7]][1] : !llvm.array<3 x vector<2xf32>>
+//       CHECK: %[[T9:.*]] = llvm.insertvalue %[[T6]], %[[T8]][2] : !llvm.array<3 x vector<2xf32>>
+//       CHECK: %[[T10:.*]] = builtin.unrealized_conversion_cast %[[T9]] : !llvm.array<3 x vector<2xf32>> to vector<3x2xf32>
+//       CHECK: return %[[T10]] : vector<3x2xf32>
+
+// -----
+
 func @broadcast_vec2d_from_vec1d(%arg0: vector<2xf32>) -> vector<3x2xf32> {
   %0 = vector.broadcast %arg0 : vector<2xf32> to vector<3x2xf32>
   return %0 : vector<3x2xf32>
 }
 // CHECK-LABEL: @broadcast_vec2d_from_vec1d(
 // CHECK-SAME:  %[[A:.*]]: vector<2xf32>)
 // CHECK:       %[[T0:.*]] = arith.constant dense<0.000000e+00> : vector<3x2xf32>
 // CHECK:       %[[T1:.*]] = builtin.unrealized_conversion_cast %[[T0]] : vector<3x2xf32> to !llvm.array<3 x vector<2xf32>>

mlir/test/Integration/Dialect/Vector/CPU/test-0-d-vectors.mlir

 
 func @splat_0d(%a: f32) {
   %1 = splat %a : vector<f32>
   // CHECK: ( 42 )
   vector.print %1: vector<f32>
   return
 }
 
+func @broadcast_0d(%a: f32) {
+  %1 = vector.broadcast %a : f32 to vector<f32>
+  // CHECK: ( 42 )
+  vector.print %1: vector<f32>
+
+  %2 = vector.broadcast %1 : vector<f32> to vector<f32>
+  // CHECK: ( 42 )
+  vector.print %2: vector<f32>
+
+  %3 = vector.broadcast %1 : vector<f32> to vector<1xf32>
+  // CHECK: ( 42 )
+  vector.print %3: vector<1xf32>
+
+  %4 = vector.broadcast %1 : vector<f32> to vector<2xf32>
+  // CHECK: ( 42, 42 )
+  vector.print %4: vector<2xf32>
+
+  %5 = vector.broadcast %1 : vector<f32> to vector<2x1xf32>
+  // CHECK: ( ( 42 ), ( 42 ) )
+  vector.print %5: vector<2x1xf32>
+
+  %6 = vector.broadcast %1 : vector<f32> to vector<2x3xf32>
+  // CHECK: ( ( 42, 42, 42 ), ( 42, 42, 42 ) )
+  vector.print %6: vector<2x3xf32>
+  return
+}
+
 func @entry() {
   %0 = arith.constant 42.0 : f32
   %1 = arith.constant dense<0.0> : vector<f32>
   %2 = call  @insert_element_0d(%0, %1) : (f32, vector<f32>) -> (vector<f32>)
   call  @extract_element_0d(%2) : (vector<f32>) -> ()
 
   %3 = arith.constant dense<42.0> : vector<f32>
   call  @print_vector_0d(%3) : (vector<f32>) -> ()
 
   %4 = arith.constant 42.0 : f32
   call  @splat_0d(%4) : (f32) -> ()
+  call  @broadcast_0d(%4) : (f32) -> ()
 
   return
 }