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

Code before:
    return success();
  }
};

/// Pattern to rewrite Y = ToElements(ShapeCast(X)) as Y = ToElements(X)
///
/// BEFORE:
///    %1 = vector.shape_cast %0 : vector<6xf32> to vector<2x3xf32>
///    %2:6 = vector.to_elements %1 : vector<2x3xf32>
/// AFTER:
///    %2:6 = vector.to_elements %0 : vector<6xf32>
struct FoldToElementsOfShapeCast final : public OpRewritePattern<ToElementsOp> {
  using Base::Base;

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

    rewriter.replaceOpWithNewOp<ToElementsOp>(toElementsOp,
                                              shapeCast.getSource());
    return success();
  }
};

void ToElementsOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                               MLIRContext *context) {
  results.add<ToElementsOfBroadcast, FoldToElementsOfShapeCast>(context);
}

//===----------------------------------------------------------------------===//
// FromElementsOp
//===----------------------------------------------------------------------===//

/// Folds vector.from_elements(vector.to_elements(%vector)) into %vector.
///
/// Case #1: Input and output vectors are the same.
///
///   %0:3 = vector.to_elements %a : vector<3xf32>
///   %1 = vector.from_elements %0#0, %0#1, %0#2 : vector<3xf32>
///   user_op %1
///
/// becomes:
///
///   user_op %a
///
static OpFoldResult foldFromElementsToElements(FromElementsOp fromElementsOp) {
  OperandRange fromElemsOperands = fromElementsOp.getElements();
  if (fromElemsOperands.empty())
    return {};

Review comment (suggestion):
This should be a fold. Have a look at ToElementsOp::fold . You can read more about the difference here: https://mlir.llvm.org/docs/Canonicalization/

Code after:
        srcIdx[k] = (srcShape[k] == 1) ? 0 : dstIdx[dstRank - srcRank + k];
      // Convert per-dimension source indices back to a linear index.
      int64_t srcLin = linearize(srcIdx, srcStrides);
      replacements.push_back(srcElems.getResult(srcLin));
    }

    rewriter.replaceOp(toElementsOp, replacements);
    return success();
  }
};

void ToElementsOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                               MLIRContext *context) {
  results.add<ToElementsOfBroadcast>(context);
}

//===----------------------------------------------------------------------===//
// FromElementsOp
//===----------------------------------------------------------------------===//

/// Folds vector.from_elements(vector.to_elements(%vector)) into %vector.
///
/// Case #1: Input and output vectors are the same.
///
///   %0:3 = vector.to_elements %a : vector<3xf32>
///   %1 = vector.from_elements %0#0, %0#1, %0#2 : vector<3xf32>
///   user_op %1
///
/// becomes:
///
///   user_op %a
///
static OpFoldResult foldFromElementsToElements(FromElementsOp fromElementsOp) {
  OperandRange fromElemsOperands = fromElementsOp.getElements();
  if (fromElemsOperands.empty())
    return {};

  auto toElementsOp = fromElemsOperands[0].getDefiningOp<ToElementsOp>();
  if (!toElementsOp)
    return {};

  if (!haveSameDefiningOp(fromElemsOperands, toElementsOp))
    return {};

  // Case #1: Input and output vectors are the same. Forward the input vector.
  Value toElementsInput = toElementsOp.getSource();
  if (fromElementsOp.getType() == toElementsInput.getType() &&
      llvm::equal(fromElemsOperands, toElementsOp.getResults())) {
    return toElementsInput;
  }