[mlir][vector] Pattern to clean up vector.extract during distribution

[ThomasRaoux]

Summary

This prevents blocking propagation when converting between scalar and
vector<1>

mlir/lib/Dialect/Vector/Transforms/VectorDistribute.cpp

LogicalResult matchAndRewrite(WarpExecuteOnLane0Op warpOp,
     rewriter.setInsertionPointAfter(newWarpOp);
     Value broadcasted = rewriter.create<vector::BroadcastOp>(
         loc, destVecType, newWarpOp->getResult(newRetIndices[0]));
     newWarpOp->getResult(operandNumber).replaceAllUsesWith(broadcasted);
     return success();
   }
 };
 
+/// Pattern to move out vector.extract of single element vector. Those don't
+/// need to be distributed and can just be propagated outside of the region.
+struct WarpOpExtract : public OpRewritePattern<WarpExecuteOnLane0Op> {
+  using OpRewritePattern<WarpExecuteOnLane0Op>::OpRewritePattern;
+  LogicalResult matchAndRewrite(WarpExecuteOnLane0Op warpOp,
+                                PatternRewriter &rewriter) const override {
+    OpOperand *operand = getWarpResult(
+        warpOp, [](Operation *op) { return isa<vector::ExtractOp>(op); });
+    if (!operand)
+      return failure();
+    unsigned int operandNumber = operand->getOperandNumber();
+    auto extractOp = operand->get().getDefiningOp<vector::ExtractOp>();
+    if (extractOp.getVectorType().getNumElements() != 1)
+      return failure();
+    Location loc = extractOp.getLoc();
+    SmallVector<size_t> newRetIndices;
+    WarpExecuteOnLane0Op newWarpOp = moveRegionToNewWarpOpAndAppendReturns(
+        rewriter, warpOp, {extractOp.getVector()}, {extractOp.getVectorType()},
+        newRetIndices);
+    rewriter.setInsertionPointAfter(newWarpOp);
+    Value newExtract = rewriter.create<vector::ExtractOp>(
+        loc, newWarpOp->getResult(newRetIndices[0]), extractOp.getPosition());
+    newWarpOp->getResult(operandNumber).replaceAllUsesWith(newExtract);
+    return success();
+  }
+};
+
 /// Sink scf.for region out of WarpExecuteOnLane0Op. This can be done only if
 /// the scf.ForOp is the last operation in the region so that it doesn't change
 /// the order of execution. This creates a new scf.for region after the
 /// WarpExecuteOnLane0Op. The new scf.for region will contain a new
 /// WarpExecuteOnLane0Op region. Example:
 /// ```
 /// %w = vector.warp_execute_on_lane_0(%laneid) -> (vector<4xf32>) {
 ///   ...

 void mlir::vector::populateDistributeTransferWriteOpPatterns(
     RewritePatternSet &patterns, const DistributionMapFn &distributionMapFn) {
   patterns.add<WarpOpTransferWrite>(patterns.getContext(), distributionMapFn);
 }
 
 void mlir::vector::populatePropagateWarpVectorDistributionPatterns(
     RewritePatternSet &patterns) {
   patterns.add<WarpOpElementwise, WarpOpTransferRead, WarpOpDeadResult,
-               WarpOpBroadcast, WarpOpForwardOperand, WarpOpScfForOp,
-               WarpOpConstant>(patterns.getContext());
+               WarpOpBroadcast, WarpOpExtract, WarpOpForwardOperand,
+               WarpOpScfForOp, WarpOpConstant>(patterns.getContext());
 }
 
 void mlir::vector::populateDistributeReduction(
     RewritePatternSet &patterns,
     DistributedReductionFn distributedReductionFn) {
   patterns.add<WarpOpReduction>(patterns.getContext(), distributedReductionFn);
 }
 

mlir/test/Dialect/Vector/vector-warp-distribute.mlir

 //       CHECK-PROP:   return %[[C]] : vector<1xf32>
 func.func @warp_constant(%laneid: index) -> (vector<1xf32>) {
   %r = vector.warp_execute_on_lane_0(%laneid)[32] -> (vector<1xf32>) {
     %cst = arith.constant dense<2.0> : vector<32xf32>
     vector.yield %cst : vector<32xf32>
   }
   return %r : vector<1xf32>
 }
+
+// -----
+
+// CHECK-PROP-LABEL: func.func @vector_extract_simple(
+//       CHECK-PROP:   %[[R:.*]] = vector.warp_execute_on_lane_0(%{{.*}})[32] -> (vector<1xf32>) {
+//       CHECK-PROP:     %[[V:.*]] = "some_def"() : () -> vector<1xf32>
+//       CHECK-PROP:     vector.yield %[[V]] : vector<1xf32>
+//       CHECK-PROP:   }
+//       CHECK-PROP:   %[[E:.*]] = vector.extract %[[R]][0] : vector<1xf32>
+//       CHECK-PROP:   return %[[E]] : f32
+func.func @vector_extract_simple(%laneid: index) -> (f32) {
+  %r = vector.warp_execute_on_lane_0(%laneid)[32] -> (f32) {
+    %0 = "some_def"() : () -> (vector<1xf32>)
+    %1 = vector.extract %0[0] : vector<1xf32>
+    vector.yield %1 : f32
+  }
+  return %r : f32
+}