[mlir][vector] Add pattern to distribute splat constant

[ThomasRaoux]

Summary

Distribute splat constant out of WarpExecuteOnLane0Op region.

[springerm]

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

560

should be wrapped in loop + updateRootInPlace.

[ThomasRaoux]

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

560

The pattern is all over this file, let me do a separate clean up patch for this.

[nicolasvasilache]

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

527

Sink really refers to increasing the nesting level.

How about:
Hoist out ... after the warp_execute_on_lane_0 op?

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

LogicalResult matchAndRewrite(WarpExecuteOnLane0Op warpOp,
     Operation *newOp = cloneOpWithOperandsAndTypes(
         rewriter, loc, elementWise, newOperands,
         {newWarpOp.getResult(operandIndex).getType()});
     newWarpOp.getResult(operandIndex).replaceAllUsesWith(newOp->getResult(0));
     return success();
   }
 };
 
+/// Sink out splat constant op feeding into a warp op yield.
+/// ```
+/// %0 = vector.warp_execute_on_lane_0(%arg0) -> (vector<1xf32>) {
+///   ...
+///   %cst = arith.constant dense<2.0> : vector<32xf32>
+///   vector.yield %cst : vector<32xf32>
+/// }
+/// ```
+/// To
+/// ```
+/// vector.warp_execute_on_lane_0(%arg0 {
+///   ...
+/// }
+/// %0 = arith.constant dense<2.0> : vector<1xf32>
+struct WarpOpConstant : public OpRewritePattern<WarpExecuteOnLane0Op> {
+  using OpRewritePattern<WarpExecuteOnLane0Op>::OpRewritePattern;
+  LogicalResult matchAndRewrite(WarpExecuteOnLane0Op warpOp,
+                                PatternRewriter &rewriter) const override {
+    OpOperand *yieldOperand = getWarpResult(
+        warpOp, [](Operation *op) { return isa<arith::ConstantOp>(op); });
+    if (!yieldOperand)
+      return failure();
+    auto constantOp = yieldOperand->get().getDefiningOp<arith::ConstantOp>();
+    auto dense = constantOp.getValue().dyn_cast<SplatElementsAttr>();
+    if (!dense)
+      return failure();
+    unsigned operandIndex = yieldOperand->getOperandNumber();
+    Attribute scalarAttr = dense.getSplatValue<Attribute>();
+    Attribute newAttr = DenseElementsAttr::get(
+        warpOp.getResult(operandIndex).getType(), scalarAttr);
+    Location loc = warpOp.getLoc();
+    rewriter.setInsertionPointAfter(warpOp);
+    Value distConstant = rewriter.create<arith::ConstantOp>(loc, newAttr);
+    warpOp.getResult(operandIndex).replaceAllUsesWith(distConstant);
+    return success();
+  }
+};
+
 /// Sink out transfer_read op feeding into a warp op yield.
 /// ```
 /// %0 = vector.warp_execute_on_lane_0(%arg0) -> (vector<1xf32>) {
 ///   ...
 //    %2 = vector.transfer_read %src[%c0], %cst : memref<1024xf32>,
 //    vector<32xf32>
 ///   vector.yield %2 : vector<32xf32>
 /// }

 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>(
-      patterns.getContext());
+               WarpOpBroadcast, 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

%r:2 = vector.warp_execute_on_lane_0(%laneid)[32] -> (vector<1xf32>, vector<1xf32>) {
     %4 = arith.addf %2, %3 : vector<32xf32>
     %5 = arith.addf %2, %2 : vector<32xf32>
 // CHECK-PROP-NOT:   arith.addf
 //     CHECK-PROP:   vector.yield %{{.*}}, %{{.*}} : vector<32xf32>, vector<32xf32>
     vector.yield %4, %5 : vector<32xf32>, vector<32xf32>
   }
   return %r#0, %r#1 : vector<1xf32>, vector<1xf32>
 }
+
+// -----
+
+// CHECK-PROP-LABEL: func @warp_constant(
+//       CHECK-PROP:   %[[C:.*]] = arith.constant dense<2.000000e+00> : vector<1xf32>
+//       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>
+}