[mlir][Vector] Add a Broadcast::createBroadcastOp helper

[nicolasvasilache]

Summary

This helper handles non trivial cases of broadcast + optional transpose creation
that should not leak to the outside world.

[springerm]

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

459

I think this should be a top-level method in VectorOps.h or VectorUtils.h because it it is more than just a BroadcastOp builder.

462

const llvm::SetVector<int64_t> &

mlir/lib/Dialect/Vector/IR/VectorOps.cpp

1718–1719

I would move this comment to computeBroadcastedUnitDims.

1745–1747

Should this be an assert?

1785

nit: nextSrcShapeDim

mlir/test/lib/Dialect/Vector/TestVectorTransforms.cpp

846–863

I would walk the IR for unregistered ops like "test.create_broadcast"(%a) {broadcasted_dims = [2]} : (vector<3x4xf32>) -> (vector<3x4x5xf32>) instead of hardcoding these here.

[nicolasvasilache]

Address comments.

[nicolasvasilache]

Address review.

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

459

true .. OTOH is is not a builder but a separate createOrFold method that is clearly documented and I really want people to find it rather than reinvent a broken wheel ..
If I move, I'd like it to be in a common place where we have clear guidelines that people must go an check it out first before building or manipulating op internals (i.e. things that should be methods of the op .. and we're back to square 1)

mlir/lib/Dialect/Vector/IR/VectorOps.cpp

1745–1747

yes, I reworked a bunch of things around asserts to guarantee this creates or fold something.

mlir/test/lib/Dialect/Vector/TestVectorTransforms.cpp

846–863

fair enough, thanks!

[nicolasvasilache]

Drop spurious includes.

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

let extraClassDeclaration = [{
     Type getSourceType() { return getSource().getType(); }
     VectorType getVectorType() {
       return getVector().getType().cast<VectorType>();
     }
 
     /// Return the dimensions of the result vector that were formerly ones in the
     /// source tensor and thus correspond to "dim-1" broadcasting.
     llvm::SetVector<int64_t> computeBroadcastedUnitDims();
+
+    /// Broadcast `value` to a vector of `dstShape`, knowing that exactly the 
+    /// `broadcastedDims` dimensions in the dstShape are broadcasted.
+    /// This requires (and asserts) that the broadcast is free of dim-1 
+    /// broadcasting.
+    /// Since vector.broadcast only allows expanding leading dimensions, an extra
+    /// vector.transpose may be inserted to make the broadcast possible.
+    static Value createBroadcastOp(
+      OpBuilder &b, Value value,
+      ArrayRef<int64_t> dstShape,
+      llvm::SetVector<int64_t> &broadcastedDims);
   }];
   let assemblyFormat = "$source attr-dict `:` type($source) `to` type($vector)";
   let hasFolder = 1;
   let hasCanonicalizer = 1;
   let hasVerifier = 1;
 }
 
 def Vector_ShuffleOp :

mlir/lib/Dialect/Vector/IR/VectorOps.cpp

 Optional<SmallVector<int64_t, 4>> FMAOp::getShapeForUnroll() {
   return llvm::to_vector<4>(getVectorType().getShape());
 }
 
 //===----------------------------------------------------------------------===//
 // BroadcastOp
 //===----------------------------------------------------------------------===//
 
-/// Return the dimensions of the result vector that were formerly ones in the
-/// source tensor and thus correspond to "dim-1" broadcasting.
-llvm::SetVector<int64_t> BroadcastOp::computeBroadcastedUnitDims() {
-  VectorType srcVectorType = getSourceType().dyn_cast<VectorType>();
-  // Scalar broadcast is without any unit dim broadcast.
-  if (!srcVectorType)
-    return {};
-  ArrayRef<int64_t> srcShape = srcVectorType.getShape();
-  ArrayRef<int64_t> dstShape = getVectorType().getShape();
+static llvm::SetVector<int64_t>
+computeBroadcastedUnitDims(ArrayRef<int64_t> srcShape,
+                           ArrayRef<int64_t> dstShape) {
   int64_t rankDiff = dstShape.size() - srcShape.size();
   int64_t dstDim = rankDiff;
   llvm::SetVector<int64_t> res;
   for (auto [s1, s2] : llvm::zip(srcShape, dstShape.drop_front(rankDiff))) {
     if (s1 != s2) {
       assert(s1 == 1 && "expected dim-1 broadcasting");
       res.insert(dstDim);
     }
     ++dstDim;
   }
   return res;
 }
 
+/// Return the dimensions of the result vector that were formerly ones in the
+/// source tensor and thus correspond to "dim-1" broadcasting.
+llvm::SetVector<int64_t> BroadcastOp::computeBroadcastedUnitDims() {
+  // Scalar broadcast is without any unit dim broadcast.
+  auto srcVectorType = getSourceType().dyn_cast<VectorType>();
+  if (!srcVectorType)
+    return {};
+  return ::computeBroadcastedUnitDims(srcVectorType.getShape(),
+                                      getVectorType().getShape());
+}
+
+static bool allBitsSet(llvm::SmallBitVector &bv, int64_t lb, int64_t ub) {
+  for (int64_t i = lb; i < ub; ++i)
+    if (!bv.test(i))
+      return false;
+  return true;
+}
+
+/// Broadcast `value` to a vector of `dstShape`, knowing that exactly the
+/// `broadcastedDims` dimensions in the dstShape are broadcasted.
+/// This requires (and asserts) that the broadcast is free of dim-1
+/// broadcasting.
+/// Since `vector.broadcast` only allows expanding leading dimensions, an extra
+/// `vector.transpose` may be inserted to make the broadcast possible.
+Value BroadcastOp::createBroadcastOp(
+    OpBuilder &b, Value value, ArrayRef<int64_t> dstShape,
+    llvm::SetVector<int64_t> &broadcastedDims) {
+  // Step 1. If no dstShape to broadcast to, just return `value`.
+  if (dstShape.empty())
+    return value;
+
+  Location loc = value.getLoc();
+  Type elementType = getElementTypeOrSelf(value.getType());
+  VectorType srcVectorType = value.getType().dyn_cast<VectorType>();
+  VectorType dstVectorType = VectorType::get(dstShape, elementType);
+
+  // Step 2. If scalar -> dstShape broadcast, just do it.
+  if (!srcVectorType) {
+    // Well-formedness check, all broadcastedDims must be set.
+    for (int64_t i = 0, e = broadcastedDims.size(); i < e; ++i)
+      assert(broadcastedDims.contains(i) && "broadcastedDims is ill-formed");
+    return b.createOrFold<vector::BroadcastOp>(loc, dstVectorType, value);
+  }
+
+  // Step 3. Since vector.broadcast only allows creating leading dims,
+  //   vector -> dstShape broadcast may require a transpose.
+  // Traverse the dims in order and construct:
+  //   1. The leading entries of the broadcastShape that is guaranteed to be
+  //      achievable by a simple broadcast.
+  //   2. The induced permutation for the subsequent vector.transpose that will
+  //      bring us from `broadcastShape` back to he desired `dstShape`.
+  // If the induced permutation is not the identity, create a vector.transpose.
+  SmallVector<int64_t> broadcastShape, permutation(dstShape.size(), -1);
+  broadcastShape.reserve(dstShape.size());
+  // Consider the example:
+  //   srcShape     = 2x4
+  //   dstShape     = 1x2x3x4x5
+  //   broadcastedDims = [0, 2, 4]
+  //
+  // We want to build:
+  //   broadcastShape  = 1x3x5x2x4
+  //   permutation     = [0, 2, 4,                 1, 3]
+  //                      ---V---           -----V-----
+  //            leading broadcast part      src shape part
+  //
+  // Note that the trailing dims of broadcastShape are exactly the srcShape
+  // by construction.
+  // nextsrcShapeDim is used to keep track of where in the permutation the
+  // "src shape part" occurs.
+  int64_t nextsrcShapeDim = broadcastedDims.size();
+  for (int64_t i = 0, e = dstShape.size(); i < e; ++i) {
+    if (broadcastedDims.contains(i)) {
+      // 3.a. For each dim in the dst shape, if it is a broadcasted dim,
+      // bring it to the head of the broadcastShape.
+      // It will need to be permuted back from `broadcastShape.size() - 1` into
+      // position `i`.
+      broadcastShape.push_back(dstShape[i]);
+      permutation[i] = broadcastShape.size() - 1;
+    } else {
+      // 3.b. Otherwise, the dim is not broadcasted, it comes from the src
+      // shape and needs to be permuted into position `i`.
+      // Don't touch `broadcastShape` here, the whole srcShape will be
+      // appended after.
+      permutation[i] = nextsrcShapeDim++;
+    }
+  }
+  // 3.c. Append the srcShape.
+  llvm::append_range(broadcastShape, srcVectorType.getShape());
+
+  // Ensure there are no dim-1 broadcasts.
+  assert(::computeBroadcastedUnitDims(srcVectorType.getShape(), broadcastShape)
+             .empty() &&
+         "unexpected dim-1 broadcast");
+
+  VectorType broadcastType = VectorType::get(broadcastShape, elementType);
+  assert(vector::isBroadcastableTo(value.getType(), broadcastType) ==
+             vector::BroadcastableToResult::Success &&
+         "must be broadcastable");
+  Value res = b.createOrFold<vector::BroadcastOp>(loc, broadcastType, value);
+  // Step 4. If we find any dimension that indeed needs to be permuted,
+  // immediately return a new vector.transpose.
+  for (int64_t i = 0, e = permutation.size(); i < e; ++i)
+    if (permutation[i] != i)
+      return b.createOrFold<vector::TransposeOp>(loc, res, permutation);
+  // Otherwise return res.
+  return res;
+}
+
 BroadcastableToResult
 mlir::vector::isBroadcastableTo(Type srcType, VectorType dstVectorType,
                                 std::pair<int, int> *mismatchingDims) {
   // Broadcast scalar to vector of the same element type.
   if (srcType.isIntOrIndexOrFloat() && dstVectorType &&
       getElementTypeOrSelf(srcType) == getElementTypeOrSelf(dstVectorType))
     return BroadcastableToResult::Success;
   // From now on, only vectors broadcast.

mlir/test/Dialect/Vector/test-create-broadcast.mlir

+// RUN: mlir-opt %s -test-create-vector-broadcast --split-input-file | FileCheck %s
+
+// The test hardcodes scalar is broadcasted to targetShape 1x2 with
+// broadcastedDims = [0, 1].
+func.func @foo(%a : f32) {
+  // CHECK: vector.broadcast {{.*}} : f32 to vector<1x2xf32>
+  // CHECK-NOT: vector.transpose
+  return
+}
+
+// -----
+
+// The test hardcodes 2x2 is broadcasted to targetShape 2x2x3xf32 with
+// broadcastedDims = [2].
+func.func @foo(%a : vector<2x2xf32>) {
+  // CHECK: vector.broadcast {{.*}} : vector<2x2xf32> to vector<3x2x2xf32>
+  // CHECK: vector.transpose {{.*}}, [1, 2, 0] : vector<3x2x2xf32> to vector<2x2x3xf32>
+  return
+}
+
+// -----
+
+// The test hardcodes 3x3 is broadcasted to targetShape 4x3x3xf32 with
+// broadcastedDims = [0].
+func.func @foo(%a : vector<3x3xf32>) {
+  // CHECK: vector.broadcast {{.*}} : vector<3x3xf32> to vector<4x3x3xf32>
+  // CHECK-NOT: vector.transpose
+  return
+}
+
+// -----
+
+// The test hardcodes 2x4 is broadcasted to targetShape 1x2x3x4x5 with
+// broadcastedDims = [0, 2, 4].
+func.func @foo(%a : vector<2x4xf32>) {
+  // CHECK: vector.broadcast {{.*}} : vector<2x4xf32> to vector<1x3x5x2x4xf32>
+  // CHECK: vector.transpose {{.*}}, [0, 3, 1, 4, 2] : vector<1x3x5x2x4xf32> to vector<1x2x3x4x5xf32>
+  return
+}

mlir/test/lib/Dialect/Vector/TestVectorTransforms.cpp

 #include "mlir/Dialect/Linalg/Passes.h"
 #include "mlir/Dialect/Linalg/Transforms/Transforms.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/SCF/IR/SCF.h"
 #include "mlir/Dialect/Vector/IR/VectorOps.h"
 #include "mlir/Dialect/Vector/Transforms/VectorDistribution.h"
 #include "mlir/Dialect/Vector/Transforms/VectorRewritePatterns.h"
 #include "mlir/Dialect/Vector/Transforms/VectorTransforms.h"
+#include "mlir/IR/Value.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Pass/PassManager.h"
 #include "mlir/Support/LLVM.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+#include "llvm/ADT/SmallBitVector.h"
+#include "llvm/ADT/SmallVector.h"
 
 using namespace mlir;
 using namespace mlir::linalg;
 using namespace mlir::vector;
 
 namespace {
 
 struct TestVectorToVectorLowering
struct TestVectorExtractStridedSliceLowering
   }
   void runOnOperation() override {
     RewritePatternSet patterns(&getContext());
     populateVectorExtractStridedSliceToExtractInsertChainPatterns(patterns);
     (void)applyPatternsAndFoldGreedily(getOperation(), std::move(patterns));
   }
 };
 
+struct TestCreateVectorBroadcast
+    : public PassWrapper<TestCreateVectorBroadcast,
+                         OperationPass<func::FuncOp>> {
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestCreateVectorBroadcast)
+
+  StringRef getArgument() const final { return "test-create-vector-broadcast"; }
+  StringRef getDescription() const final {
+    return "Test optimization transformations for transfer ops";
+  }
+  void getDependentDialects(DialectRegistry &registry) const override {
+    registry.insert<vector::VectorDialect>();
+  }
+
+  void runOnOperation() override {
+    OpBuilder b(getOperation());
+    b.setInsertionPointToStart(&getOperation().getBody().front());
+    for (BlockArgument bbArg : getOperation().getArguments()) {
+      auto vectorType = bbArg.getType().dyn_cast<VectorType>();
+      llvm::SetVector<int64_t> broadcastedDims;
+      SmallVector<int64_t> targetShape;
+      if (!vectorType) {
+        targetShape = {1, 2};
+        ArrayRef<int64_t> dims{0, 1};
+        broadcastedDims.insert(dims.begin(), dims.end());
+      } else if (vectorType.getShape().equals({2, 2})) {
+        targetShape = {2, 2, 3};
+        ArrayRef<int64_t> dims{2};
+        broadcastedDims.insert(dims.begin(), dims.end());
+      } else if (vectorType.getShape().equals({3, 3})) {
+        targetShape = {4, 3, 3};
+        ArrayRef<int64_t> dims{0};
+        broadcastedDims.insert(dims.begin(), dims.end());
+      } else if (vectorType.getShape().equals({2, 4})) {
+        targetShape = {1, 2, 3, 4, 5};
+        ArrayRef<int64_t> dims{0, 2, 4};
+        broadcastedDims.insert(dims.begin(), dims.end());
+      } else {
+        continue;
+      }
+      vector::BroadcastOp::createBroadcastOp(b, bbArg, targetShape,
+                                             broadcastedDims);
+    }
+  }
+};
+
 } // namespace
 
 namespace mlir {
 namespace test {
 void registerTestVectorLowerings() {
   PassRegistration<TestVectorToVectorLowering>();
 
   PassRegistration<TestVectorContractionLowering>();
void registerTestVectorLowerings() {
 
   PassRegistration<TestFlattenVectorTransferPatterns>();
 
   PassRegistration<TestVectorScanLowering>();
 
   PassRegistration<TestVectorDistribution>();
 
   PassRegistration<TestVectorExtractStridedSliceLowering>();
+
+  PassRegistration<TestCreateVectorBroadcast>();
 }
 } // namespace test
 } // namespace mlir