Repository: llvm/llvm-project
File: mlir/lib/Dialect/Vector/IR/VectorOps.cpp (C++)

Code before:
/// AFTER:
///    %2 = vector.from_elements %c1, %c2, %c3 : vector<1x3xf32>
class FoldShapeCastOfFromElements final : public OpRewritePattern<ShapeCastOp> {
public:
  using Base::Base;

  LogicalResult matchAndRewrite(ShapeCastOp shapeCastOp,
                                PatternRewriter &rewriter) const override {
    auto fromElements = shapeCastOp.getSource().getDefiningOp<FromElementsOp>();
    if (!fromElements)
      return failure();

    rewriter.replaceOpWithNewOp<FromElementsOp>(
        shapeCastOp, shapeCastOp.getResultVectorType(),
        fromElements.getElements());
    return success();
  }
};

} // namespace

void ShapeCastOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                              MLIRContext *context) {
  results.add<ShapeCastCreateMaskFolderTrailingOneDim, ShapeCastBroadcastFolder,
              FoldShapeCastOfFromElements>(context);
}

//===----------------------------------------------------------------------===//
// VectorBitCastOp
//===----------------------------------------------------------------------===//

LogicalResult BitCastOp::verify() {
  auto sourceVectorType = getSourceVectorType();
  auto resultVectorType = getResultVectorType();

  for (int64_t i = 0, e = sourceVectorType.getRank() - 1; i < e; i++) {
    if (sourceVectorType.getDimSize(i) != resultVectorType.getDimSize(i))
      return emitOpError("dimension size mismatch at: ") << i;
  }

  DataLayout dataLayout = DataLayout::closest(*this);
  auto sourceElementBits =
      dataLayout.getTypeSizeInBits(sourceVectorType.getElementType());
  auto resultElementBits =
      dataLayout.getTypeSizeInBits(resultVectorType.getElementType());

  if (sourceVectorType.getRank() == 0) {
    if (sourceElementBits != resultElementBits)
      return emitOpError("source/result bitwidth of the 0-D vector element "
                         "types must be equal");
  } else if (sourceElementBits * sourceVectorType.getShape().back() !=

Review comment (suggestion):
Instead of adding a new pattern, you should be able to simply extend this: https://github.com/llvm/llvm-project/blob/8f378ea7e6fa31179266c69368be56a866b631e1/mlir/lib/Dialect/Vector/IR/VectorOps.cpp?plain=1#L6481-L6521 

IIUC, that's what Mehdi and Kunwar had in mind.

Code after:
/// AFTER:
///    %2 = vector.from_elements %c1, %c2, %c3 : vector<1x3xf32>
///
/// Note: this transformation is implemented as an OpRewritePattern, not as a
/// fold, because we have to create new op FromElementsOp with updated result
/// type. This cannot be done with a fold, because fold cannot create new ops
/// and the existing FromElementsOp result type differs from the ShapeCastOp
/// result type. Mutating the FromElementsOp (not root op) would violate the
/// fold contract and break other users.
class FoldShapeCastOfFromElements final : public OpRewritePattern<ShapeCastOp> {
public:
  using Base::Base;

  LogicalResult matchAndRewrite(ShapeCastOp shapeCastOp,
                                PatternRewriter &rewriter) const override {
    auto fromElements = shapeCastOp.getSource().getDefiningOp<FromElementsOp>();
    if (!fromElements)
      return failure();

    rewriter.replaceOpWithNewOp<FromElementsOp>(
        shapeCastOp, shapeCastOp.getResultVectorType(),
        fromElements.getElements());
    return success();
  }
};

} // namespace

void ShapeCastOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                              MLIRContext *context) {
  results.add<ShapeCastCreateMaskFolderTrailingOneDim, ShapeCastBroadcastFolder,
              FoldShapeCastOfFromElements>(context);
}

//===----------------------------------------------------------------------===//
// VectorBitCastOp
//===----------------------------------------------------------------------===//

LogicalResult BitCastOp::verify() {
  auto sourceVectorType = getSourceVectorType();
  auto resultVectorType = getResultVectorType();

  for (int64_t i = 0, e = sourceVectorType.getRank() - 1; i < e; i++) {
    if (sourceVectorType.getDimSize(i) != resultVectorType.getDimSize(i))
      return emitOpError("dimension size mismatch at: ") << i;
  }

  DataLayout dataLayout = DataLayout::closest(*this);
  auto sourceElementBits =
      dataLayout.getTypeSizeInBits(sourceVectorType.getElementType());
  auto resultElementBits =