[mlir][Vector] Add a folder for vector.broadcast

[hanchung]

Summary

Fold the operation if the source is a scalar constant or splat constant.

Update transform-patterns-matmul-to-vector.mlir because the broadcast ops are folded in the conversion.

[hanchung]

Fix test

[aartbik]

mlir/lib/Dialect/Vector/VectorOps.cpp

933

is this test really required, what vector types would fail?

934

style: single statement if has no { }

939

prevOp sounds like it is a predecessor, but it is the defining op

940

these cases work for the simple tests you made, but I am not convinced this works for all the duplicate/stretch semantics? did you give that some thought?

[aartbik]

mlir/lib/Dialect/Vector/VectorOps.cpp

933

I realized this comment was a bit obscure, I meant to say that for vector types you need proper duplicate/stretch semantics and you avoid that right now for this case, but not the other. Perhaps your intention was to only deal with scalar sources for both cases?

[hanchung]

Remove the propagatin of vector.broadcast

[hanchung]

Handle splat constants.

[hanchung]

mlir/lib/Dialect/Vector/VectorOps.cpp

933

My intention is to deal with broadcasting constant scalar, so I check if this is a scalar in the first condition.

Now I fixed it to accept both scalar and splat element attrs (like vectors).

940

Actually I didn't think about this part much. I just came out of this idea when adding the folder on constant scalars. I feel that we could fold it according to the op definition -- the source operand is duplicated to match the given rank and sizes in the result vector. If we can broadcast in one step why two steps.

However, I think you are might right, there could be more complicated cases, eg, each step could have different meanings. Let's hold this off and revisit it when there are more examples?

[rriddle]

mlir/lib/Dialect/Vector/VectorOps.cpp

933

Please do not check for constants like this, operands contains the attribute values for constant operands for this operation or null if the operand isn't constant.

[aartbik]

mlir/lib/Dialect/Vector/VectorOps.cpp

940

Oh, you are probably right, I just needed some convincing and given the lack of tests I was not sure if you had considered all cases ;-) Either way works for me, but what you have now is a nice pure constant folder.

mlir/test/Dialect/Vector/canonicalize.mlir

393

can you please make these tests a bit more specific that the constant is indeed returned
something like

CHECK: [[A:.*]] = constant ....
CHECK-NOT: ...
CHECK: return %[[A]]

[rriddle]

mlir/lib/Dialect/Vector/VectorOps.cpp

933

For additional clarification as to why:

• This API is used in a wide variety of circumstances, and constants may be passed in to operands when the actual operand of this operation is not a ConstantOp(or equivalent).
• By hard coding ConstantOp you break composability with another constant operation.

[hanchung]

Add more checks on tests and handle constants properly.

[hanchung]

mlir/lib/Dialect/Vector/VectorOps.cpp

933

Thanks for pointing this out.

[aartbik]

mlir/lib/Dialect/Vector/VectorOps.cpp

935

nit: can't we use getVectorType() directly here, since this is a broadcast

I would also call it vectorType rather than shapedType (or perhaps resultType if you prefer), since thta makes it a bit more clear

[hanchung]

Use getVectorType() instead

mlir/lib/Dialect/Vector/VectorOps.cpp

935

Oh yes, it makes more sense to me too. Thanks!

[aartbik]

Good to go if River has no more feedback on this. Please make sure to run the integration tests as well prior to submitting.
And thanks for the optimization!

Also, please feel free to follow-up with the broadcast-broadcast canonicalization. It was not my intention to discourage that, just to make sure you thought and tested through all cases.

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

def Vector_BroadcastOp :
   }];
   let extraClassDeclaration = [{
     Type getSourceType() { return source().getType(); }
     VectorType getVectorType() {
       return vector().getType().cast<VectorType>();
     }
   }];
   let assemblyFormat = "$source attr-dict `:` type($source) `to` type($vector)";
+  let hasFolder = 1;
 }
 
 def Vector_ShuffleOp :
   Vector_Op<"shuffle", [NoSideEffect,
      PredOpTrait<"first operand v1 and result have same element type",
                  TCresVTEtIsSameAsOpBase<0, 0>>,
      PredOpTrait<"second operand v2 and result have same element type",
                  TCresVTEtIsSameAsOpBase<0, 1>>]>,

mlir/lib/Dialect/Vector/VectorOps.cpp

for (int64_t r = 0; r < srcRank; ++r) {
       if (srcDim != 1 && srcDim != dstDim)
         return op.emitOpError("dimension mismatch (")
                << srcDim << " vs. " << dstDim << ")";
     }
   }
   return success();
 }
 
+OpFoldResult BroadcastOp::fold(ArrayRef<Attribute> operands) {
+  if (source().getType().isIntOrIndexOrFloat()) {
+    if (auto cst = source().getDefiningOp<ConstantOp>()) {
+      return DenseElementsAttr::get(getResult().getType().cast<ShapedType>(),
+                                    cst.value());
+    }
+  }
+  if (auto prevOp = source().getDefiningOp<vector::BroadcastOp>()) {
+    setOperand(prevOp.source());
+    return getResult();
+  }
+  return {};
+}
+
 //===----------------------------------------------------------------------===//
 // ShuffleOp
 //===----------------------------------------------------------------------===//
 
 void ShuffleOp::build(OpBuilder &builder, OperationState &result, Value v1,
                       Value v2, ArrayRef<int64_t> mask) {
   result.addOperands({v1, v2});
   auto maskAttr = getVectorSubscriptAttr(builder, mask);

mlir/test/Dialect/Linalg/transform-patterns-matmul-to-vector.mlir

 // RUN: mlir-opt %s -test-linalg-transform-patterns=test-matmul-to-vector-patterns-tile-1d | FileCheck %s
 // RUN: mlir-opt %s -test-linalg-transform-patterns=test-matmul-to-vector-patterns-tile-2d | FileCheck %s
 // RUN: mlir-opt %s -test-linalg-transform-patterns=test-contraction-to-vector-patterns | FileCheck %s --check-prefix=VECTOR-CONTRACTION
 
 func @matmul(%A: memref<1584x1584xf32, offset: 0, strides: [1584, 1]>,
                   %B: memref<1584x1584xf32, offset: 0, strides: [1584, 1]>,
                   %C: memref<1584x1584xf32, offset: 0, strides: [1584, 1]>) {
   linalg.matmul %A, %B, %C {__internal_linalg_transform__ = "START"} :
     (memref<1584x1584xf32, offset: 0, strides: [1584, 1]>,
      memref<1584x1584xf32, offset: 0, strides: [1584, 1]>,
      memref<1584x1584xf32, offset: 0, strides: [1584, 1]>)
   return
 }
 
 // CHECK-LABEL:func @matmul
-//      CHECK: vector.broadcast {{.*}} : f32 to vector<8x16xf32>
 //      CHECK: store {{.*}}[] : memref<vector<8x16xf32>>
-//
-//      CHECK: vector.broadcast {{.*}} : f32 to vector<16x12xf32>
 //      CHECK: store {{.*}}[] : memref<vector<16x12xf32>>
-//
-//      CHECK: vector.broadcast {{.*}} : f32 to vector<8x12xf32>
 //      CHECK: store {{.*}}[] : memref<vector<8x12xf32>>
 //
 //      CHECK: linalg.copy
 //      CHECK: linalg.copy
 //      CHECK: linalg.copy
 //
 //      CHECK: vector.contract
 // CHECK-SAME:   iterator_types = ["parallel", "parallel", "reduction"]

mlir/test/Dialect/Vector/canonicalize.mlir

 //  CHECK-SAME:   %[[B:.*]]: vector<2xi32>
 //  CHECK:        return %[[A]], %[[B]] : vector<4x8xf32>, vector<2xi32>
 func @bitcast_folding(%I1: vector<4x8xf32>, %I2: vector<2xi32>) -> (vector<4x8xf32>, vector<2xi32>) {
   %0 = vector.bitcast %I1 : vector<4x8xf32> to vector<4x8xf32>
   %1 = vector.bitcast %I2 : vector<2xi32> to vector<4xi16>
   %2 = vector.bitcast %1 : vector<4xi16> to vector<2xi32>
   return %0, %2 : vector<4x8xf32>, vector<2xi32>
 }
+
+// -----
+
+// CHECK-LABEL: broadcast_folding1
+//       CHECK: constant
+//   CHECK-NOT: vector.broadcast
+func @broadcast_folding1() -> vector<4xf32> {
+  %0 = constant 0.0 : f32
+  %1 = vector.broadcast %0 : f32 to vector<4xf32>
+  return %1 : vector<4xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @broadcast_folding2
+//  CHECK-SAME: %[[ARG0:.*]]: f32
+//       CHECK: %[[RES:.*]] = vector.broadcast %[[ARG0]] : f32 to vector<4x16xf32>
+//       CHECK: return %[[RES]]
+func @broadcast_folding2(%arg0: f32) -> vector<4x16xf32> {
+  %0 = vector.broadcast %arg0 : f32 to vector<16xf32>
+  %1 = vector.broadcast %0 : vector<16xf32> to vector<4x16xf32>
+  return %1 : vector<4x16xf32>
+}