[Mlir-commits] [mlir] [mlir][Vector] Move insert/extractelement distribution patterns to insert/extract (PR #116425)

[Diego Caballero]

================
@@ -1413,31 +1397,60 @@ struct WarpOpExtractElement : public OpRewritePattern<WarpExecuteOnLane0Op> {
   WarpShuffleFromIdxFn warpShuffleFromIdxFn;
 };
 
-struct WarpOpInsertElement : public OpRewritePattern<WarpExecuteOnLane0Op> {
+/// Pattern to convert vector.extractelement to vector.extract.
+struct WarpOpExtractElement : public OpRewritePattern<WarpExecuteOnLane0Op> {
+  WarpOpExtractElement(MLIRContext *ctx, PatternBenefit b = 1)
+      : OpRewritePattern<WarpExecuteOnLane0Op>(ctx, b) {}
+  LogicalResult matchAndRewrite(WarpExecuteOnLane0Op warpOp,
+                                PatternRewriter &rewriter) const override {
+    OpOperand *operand =
+        getWarpResult(warpOp, llvm::IsaPred<vector::ExtractElementOp>);
+    if (!operand)
+      return failure();
+    auto extractOp = operand->get().getDefiningOp<vector::ExtractElementOp>();
+    rewriter.setInsertionPoint(extractOp);
+    if (auto pos = extractOp.getPosition()) {
+      rewriter.replaceOpWithNewOp<vector::ExtractOp>(
+          extractOp, extractOp.getVector(), pos);
+    } else {
+      rewriter.replaceOpWithNewOp<vector::ExtractOp>(
+          extractOp, extractOp.getVector(), ArrayRef<int64_t>{});
----------------
dcaballe wrote:

some of these if-else statements might be cleaner if you move the op creation after the if-else and only conditionally add `pos` to an empty position list?

https://github.com/llvm/llvm-project/pull/116425