[mlir][vector] Add patterns to propagate vector distribution

[ThomasRaoux]

Summary

Add patterns to propagate vector distribution and remove dead
arguments. This handles propagation for several vector operations.

co-authored with @springerm

[csigg]

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

223

How do the two vector types above imply this map?

228

This comment doesn't seem to match the code.

410

I'm starting to think I'm not understanding the semantic of this op.

What is the type of %arg0?

How does the vector<32xf32> operand type from vector.yield and the vector<1xf32> return type of vector.warp_execute_on_lane_0 relate?

[ThomasRaoux]

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

223

The map indicates which dimensions are distributed. Here based on the vector shape we know that dimension 0 and 2 are different in the yield and result meaning they must be distributed on threads.

410

I have tried to describe it in the op but probably could use some feedback on better wording. %arg0 is the threadId, just an ID identifying the different threads we want to distribute on.
The Value being yield is distributed on multiple lanes/threads. Therefore the vector<32xf32> value outside of the region is distributed so that the 32 lanes each old 1 element of the vector hence the vector<1xf32>. This conversion implicitly goes through memory to transfer the vector<32xf32> from lane 0 to a Value distributed on all the threads.

[ThomasRaoux]

Update comment

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

228

Updated the comment.

[springerm]

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

186

and is not dead.

202

operation

204

RewriterBase &rewriter

210

rewriter

468

I think we need a loop iterating over all uses, then wrap each update them one-by-one. Each of them wrapped in a rewriter.updateRootInplace(use.getOwner(), ...)

528

Same as above, we need a loop with updateRootInPlace.

557

Same as above, we need a loop with updateRootInPlace. At this point it may be worth adding a helper function.

I suggested deleting the original op during moveRegionToNewWarpOpAndReplaceReturns in the previous change. In that case, you can't iterate over warpOp.getResults here.

601

Use rewriter.

625

This is probably safe, but I would use the rewriter for consistency.

682

updateRootInplace

697–710

I think this can be simplified with rewriter.mergeBlocks. The replacements for the bbArgs can be specified as the third argument. The bvm should not be necessary then.

This is also more efficient because the ops are not cloned. The terminator is also moved, so just remove that one afterwards.

719–720

Use rewriter

[ThomasRaoux]

Address review comments.

[ThomasRaoux]

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

468

I don't see user being notified when the uses are replaced in replaceOp or other cases. I'd rather avoid having a special case here. If it is needed we most likely need to update Rewriter class to support such thing. I think at this point it is safe so I'd prefer keeping it like that.

528

same comment.

697–710

Good point, changed it.

mlir/include/mlir/Dialect/Vector/Transforms/VectorDistribution.h

 /// vector.transfer_write %v, %A[%id] : vector<1xf32>, memref<128xf32>
 void populateDistributeTransferWriteOpPatterns(
     RewritePatternSet &patterns, DistributionMapFn distributionMapFn);
 
 /// Move scalar operations with no dependency on the warp op outside of the
 /// region.
 void moveScalarUniformCode(WarpExecuteOnLane0Op op);
 
+/// Collect patterns to propagate warp distribution.
+void populatePropagateWarpVectorDistributionPatterns(
+    RewritePatternSet &pattern);
+
 } // namespace vector
 } // namespace mlir
 #endif // MLIR_DIALECT_VECTOR_TRANSFORMS_VECTORDISTRIBUTION_H_

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

 //===- VectorDistribute.cpp - patterns to do vector distribution ----------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/SCF/SCF.h"
 #include "mlir/Dialect/Vector/Transforms/VectorDistribution.h"
+#include "mlir/IR/BlockAndValueMapping.h"
 #include "mlir/Transforms/SideEffectUtils.h"
 
 using namespace mlir;
 using namespace mlir::vector;
 
 static LogicalResult
 rewriteWarpOpToScfFor(RewriterBase &rewriter, WarpExecuteOnLane0Op warpOp,
                       const WarpExecuteOnLane0LoweringOptions &options) {

 
 /// Helper to know if an op can be hoisted out of the region.
 static bool canBeHoisted(Operation *op,
                          function_ref<bool(Value)> definedOutside) {
   return llvm::all_of(op->getOperands(), definedOutside) &&
          isSideEffectFree(op) && op->getNumRegions() == 0;
 }
 
+/// Return a value yielded by `warpOp` which statifies the filter lamdba
+/// condition and is not dead.
+static OpOperand *getWarpResult(WarpExecuteOnLane0Op warpOp,
+                                std::function<bool(Operation *)> fn) {
+  auto yield = cast<vector::YieldOp>(
+      warpOp.getBodyRegion().getBlocks().begin()->getTerminator());
+  for (OpOperand &yieldOperand : yield->getOpOperands()) {
+    Value yieldValues = yieldOperand.get();
+    Operation *definedOp = yieldValues.getDefiningOp();
+    if (definedOp && fn(definedOp)) {
+      if (!warpOp.getResult(yieldOperand.getOperandNumber()).use_empty())
+        return &yieldOperand;
+    }
+  }
+  return {};
+}
+
+// Clones `op` into a new operation that takes `operands` and returns
+// `resultTypes`.
+static Operation *cloneOpWithOperandsAndTypes(RewriterBase &rewriter,
+                                              Location loc, Operation *op,
+                                              ArrayRef<Value> operands,
+                                              ArrayRef<Type> resultTypes) {
+  OperationState res(loc, op->getName().getStringRef(), operands, resultTypes,
+                     op->getAttrs());
+  return rewriter.create(res);
+}
+
+/// Currently the distribution map is implicit based on the vector shape. In the
+/// future it will be part of the op.
+/// Example:
+/// ```
+/// %0 = vector.warp_execute_on_lane_0(%arg0) -> (vector<1x16x2xf32>) {
+///   ...
+///   vector.yield %3 : vector<32x16x64xf32>
+/// }
+/// ```
+/// Would have an implicit map of:
+/// `(d0, d1, d2) -> (d0, d2)`
+static AffineMap calculateImplicitMap(Value yield, Value ret) {
+  auto srcType = yield.getType().cast<VectorType>();
+  auto dstType = ret.getType().cast<VectorType>();
+  SmallVector<AffineExpr> perm;
+  // Check which dimensions of the yield value are different than the dimensions
+  // of the result to know the distributed dimensions. Then associate each
+  // distributed dimension to an ID in order.
+  for (unsigned i = 0, e = srcType.getRank(); i < e; i++) {
+    if (srcType.getDimSize(i) != dstType.getDimSize(i))
+      perm.push_back(getAffineDimExpr(i, yield.getContext()));
+  }
+  auto map = AffineMap::get(srcType.getRank(), 0, perm, yield.getContext());
+  return map;
+}
+
 namespace {
 
 struct WarpOpToScfForPattern : public OpRewritePattern<WarpExecuteOnLane0Op> {
   WarpOpToScfForPattern(MLIRContext *context,
                         const WarpExecuteOnLane0LoweringOptions &options,
                         PatternBenefit benefit = 1)
       : OpRewritePattern<WarpExecuteOnLane0Op>(context, benefit),
         options(options) {}
LogicalResult matchAndRewrite(vector::TransferWriteOp writeOp,
 
     return failure();
   }
 
 private:
   DistributionMapFn distributionMapFn;
 };
 
+/// Sink out elementwise op feeding into a warp op yield.
+/// ```
+/// %0 = vector.warp_execute_on_lane_0(%arg0) -> (vector<1xf32>) {
+///   ...
+///   %3 = arith.addf %1, %2 : vector<32xf32>
+///   vector.yield %3 : vector<32xf32>
+/// }
+/// ```
+/// To
+/// ```
+/// %r:3 = vector.warp_execute_on_lane_0(%arg0) -> (vector<1xf32>,
+/// vector<1xf32>, vector<1xf32>) {
+///   ...
+///   %4 = arith.addf %2, %3 : vector<32xf32>
+///   vector.yield %4, %2, %3 : vector<32xf32>, vector<32xf32>,
+///   vector<32xf32>
+/// }
+/// %0 = arith.addf %r#1, %r#2 : vector<1xf32>
+struct WarpOpElementwise : public OpRewritePattern<WarpExecuteOnLane0Op> {
+  using OpRewritePattern<WarpExecuteOnLane0Op>::OpRewritePattern;
+  LogicalResult matchAndRewrite(WarpExecuteOnLane0Op warpOp,
+                                PatternRewriter &rewriter) const override {
+    OpOperand *yieldOperand = getWarpResult(warpOp, [](Operation *op) {
+      return OpTrait::hasElementwiseMappableTraits(op);
+    });
+    if (!yieldOperand)
+      return failure();
+    Operation *elementWise = yieldOperand->get().getDefiningOp();
+    unsigned operandIndex = yieldOperand->getOperandNumber();
+    Value distributedVal = warpOp.getResult(operandIndex);
+    SmallVector<Value> yieldValues;
+    SmallVector<Type> retTypes;
+    Location loc = warpOp.getLoc();
+    for (OpOperand &operand : elementWise->getOpOperands()) {
+      Type targetType;
+      if (auto vecType = distributedVal.getType().dyn_cast<VectorType>()) {
+        // If the result type is a vector, the operands must also be vectors.
+        auto operandType = operand.get().getType().cast<VectorType>();
+        targetType =
+            VectorType::get(vecType.getShape(), operandType.getElementType());
+      } else {
+        auto operandType = operand.get().getType();
+        assert(!operandType.isa<VectorType>() &&
+               "unexpected yield of vector from op with scalar result type");
+        targetType = operandType;
+      }
+      retTypes.push_back(targetType);
+      yieldValues.push_back(operand.get());
+    }
+    WarpExecuteOnLane0Op newWarpOp = moveRegionToNewWarpOpAndAppendReturns(
+        rewriter, warpOp, yieldValues, retTypes);
+    rewriter.setInsertionPointAfter(newWarpOp);
+    SmallVector<Value> newOperands(elementWise->getOperands().begin(),
+                                   elementWise->getOperands().end());
+    for (unsigned i : llvm::seq(unsigned(0), elementWise->getNumOperands())) {
+      newOperands[i] = newWarpOp.getResult(i + warpOp.getNumResults());
+    }
+    OpBuilder::InsertionGuard g(rewriter);
+    rewriter.setInsertionPointAfter(newWarpOp);
+    Operation *newOp = cloneOpWithOperandsAndTypes(
+        rewriter, loc, elementWise, newOperands,
+        {newWarpOp.getResult(operandIndex).getType()});
+    newWarpOp.getResult(operandIndex).replaceAllUsesWith(newOp->getResult(0));
+    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>
+/// }
+/// ```
+/// To
+/// ```
+/// %dead = vector.warp_execute_on_lane_0(%arg0) -> (vector<1xf32>,
+/// vector<1xf32>, vector<1xf32>) {
+///   ...
+///   %2 = vector.transfer_read %src[%c0], %cst : memref<1024xf32>,
+///   vector<32xf32> vector.yield %2 : vector<32xf32>
+/// }
+/// %0 = vector.transfer_read %src[%c0], %cst : memref<1024xf32>, vector<1xf32>
+struct WarpOpTransferRead : public OpRewritePattern<WarpExecuteOnLane0Op> {
+  using OpRewritePattern<WarpExecuteOnLane0Op>::OpRewritePattern;
+  LogicalResult matchAndRewrite(WarpExecuteOnLane0Op warpOp,
+                                PatternRewriter &rewriter) const override {
+    OpOperand *operand = getWarpResult(
+        warpOp, [](Operation *op) { return isa<vector::TransferReadOp>(op); });
+    if (!operand)
+      return failure();
+    auto read = operand->get().getDefiningOp<vector::TransferReadOp>();
+    unsigned operandIndex = operand->getOperandNumber();
+    Value distributedVal = warpOp.getResult(operandIndex);
+
+    SmallVector<Value, 4> indices(read.getIndices().begin(),
+                                  read.getIndices().end());
+    AffineMap map = calculateImplicitMap(read.getResult(), distributedVal);
+    AffineMap indexMap = map.compose(read.getPermutationMap());
+    OpBuilder::InsertionGuard g(rewriter);
+    rewriter.setInsertionPointAfter(warpOp);
+    for (auto it : llvm::zip(indexMap.getResults(), map.getResults())) {
+      AffineExpr d0, d1;
+      bindDims(read.getContext(), d0, d1);
+      auto indexExpr = std::get<0>(it).dyn_cast<AffineDimExpr>();
+      if (!indexExpr)
+        continue;
+      unsigned indexPos = indexExpr.getPosition();
+      unsigned vectorPos = std::get<1>(it).cast<AffineDimExpr>().getPosition();
+      int64_t scale =
+          distributedVal.getType().cast<VectorType>().getDimSize(vectorPos);
+      indices[indexPos] =
+          makeComposedAffineApply(rewriter, read.getLoc(), d0 + scale * d1,
+                                  {indices[indexPos], warpOp.getLaneid()});
+    }
+    Value newRead = rewriter.create<vector::TransferReadOp>(
+        read.getLoc(), distributedVal.getType(), read.getSource(), indices,
+        read.getPermutationMapAttr(), read.getPadding(), read.getMask(),
+        read.getInBoundsAttr());
+    distributedVal.replaceAllUsesWith(newRead);
+    return success();
+  }
+};
+
+/// Remove any result that has no use along with the matching yieldOp operand.
+// TODO: Move this in WarpExecuteOnLane0Op canonicalization.
+struct WarpOpDeadResult : public OpRewritePattern<WarpExecuteOnLane0Op> {
+  using OpRewritePattern<WarpExecuteOnLane0Op>::OpRewritePattern;
+  LogicalResult matchAndRewrite(WarpExecuteOnLane0Op warpOp,
+                                PatternRewriter &rewriter) const override {
+    SmallVector<Type> resultTypes;
+    SmallVector<Value> yieldValues;
+    auto yield = cast<vector::YieldOp>(
+        warpOp.getBodyRegion().getBlocks().begin()->getTerminator());
+    for (OpResult result : warpOp.getResults()) {
+      if (result.use_empty())
+        continue;
+      resultTypes.push_back(result.getType());
+      yieldValues.push_back(yield.getOperand(result.getResultNumber()));
+    }
+    if (yield.getNumOperands() == yieldValues.size())
+      return failure();
+    WarpExecuteOnLane0Op newWarpOp = moveRegionToNewWarpOpAndReplaceReturns(
+        rewriter, warpOp, yieldValues, resultTypes);
+    unsigned resultIndex = 0;
+    for (OpResult result : warpOp.getResults()) {
+      if (result.use_empty())
+        continue;
+      result.replaceAllUsesWith(newWarpOp.getResult(resultIndex++));
+    }
+    rewriter.eraseOp(warpOp);
+    return success();
+  }
+};
+
+// If an operand is directly yielded out of the region we can forward it
+// directly and it doesn't need to go through the region.
+struct WarpOpForwardOperand : public OpRewritePattern<WarpExecuteOnLane0Op> {
+  using OpRewritePattern<WarpExecuteOnLane0Op>::OpRewritePattern;
+  LogicalResult matchAndRewrite(WarpExecuteOnLane0Op warpOp,
+                                PatternRewriter &rewriter) const override {
+    SmallVector<Type> resultTypes;
+    SmallVector<Value> yieldValues;
+    auto yield = cast<vector::YieldOp>(
+        warpOp.getBodyRegion().getBlocks().begin()->getTerminator());
+    Value valForwarded;
+    unsigned resultIndex;
+    for (OpOperand &operand : yield->getOpOperands()) {
+      Value result = warpOp.getResult(operand.getOperandNumber());
+      if (result.use_empty())
+        continue;
+
+      // Assume all the values coming from above are uniform.
+      if (!warpOp.getBodyRegion().isAncestor(operand.get().getParentRegion())) {
+        if (result.getType() != operand.get().getType())
+          continue;
+        valForwarded = operand.get();
+        resultIndex = operand.getOperandNumber();
+        break;
+      }
+      auto arg = operand.get().dyn_cast<BlockArgument>();
+      if (!arg || arg.getOwner()->getParentOp() != warpOp.getOperation())
+        continue;
+      Value warpOperand = warpOp.getArgs()[arg.getArgNumber()];
+      if (result.getType() != warpOperand.getType())
+        continue;
+      valForwarded = warpOperand;
+      resultIndex = operand.getOperandNumber();
+      break;
+    }
+    if (!valForwarded)
+      return failure();
+    warpOp.getResult(resultIndex).replaceAllUsesWith(valForwarded);
+    return success();
+  }
+};
+
+struct WarpOpBroadcast : public OpRewritePattern<WarpExecuteOnLane0Op> {
+  using OpRewritePattern<WarpExecuteOnLane0Op>::OpRewritePattern;
+  LogicalResult matchAndRewrite(WarpExecuteOnLane0Op warpOp,
+                                PatternRewriter &rewriter) const override {
+    OpOperand *operand = getWarpResult(
+        warpOp, [](Operation *op) { return isa<vector::BroadcastOp>(op); });
+    if (!operand)
+      return failure();
+    unsigned int operandNumber = operand->getOperandNumber();
+    auto broadcastOp = operand->get().getDefiningOp<vector::BroadcastOp>();
+    Location loc = broadcastOp.getLoc();
+    auto destVecType =
+        warpOp->getResultTypes()[operandNumber].cast<VectorType>();
+    WarpExecuteOnLane0Op newWarpOp = moveRegionToNewWarpOpAndAppendReturns(
+        rewriter, warpOp, {broadcastOp.getSource()},
+        {broadcastOp.getSource().getType()});
+    rewriter.setInsertionPointAfter(newWarpOp);
+    Value broadcasted = rewriter.create<vector::BroadcastOp>(
+        loc, destVecType, newWarpOp->getResults().back());
+    newWarpOp->getResult(operandNumber).replaceAllUsesWith(broadcasted);
+
+    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>) {
+///   ...
+///   %v1 = scf.for %arg3 = %c0 to %c128 step %c1 iter_args(%arg4 = %v)
+///   -> (vector<128xf32>) {
+///     ...
+///     scf.yield %r : vector<128xf32>
+///   }
+///   vector.yield %v1 : vector<128xf32>
+/// }
+/// ```
+/// To:
+/// %w0 = vector.warp_execute_on_lane_0(%arg0) -> (vector<4xf32>) {
+///   ...
+///   vector.yield %v : vector<128xf32>
+/// }
+/// %w = scf.for %arg3 = %c0 to %c128 step %c1 iter_args(%varg = %q0)
+///   -> (vector<4xf32>) {
+///     %iw = vector.warp_execute_on_lane_0(%laneid)
+///     args(%varg : vector<4xf32>) -> (vector<4xf32>) {
+///     ^bb0(%arg: vector<128xf32>):
+///       ...
+///       vector.yield %ir : vector<128xf32>
+///     }
+///     scf.yield %iw : vector<4xf32>
+///  }
+/// ```
+struct WarpOpScfForOp : public OpRewritePattern<WarpExecuteOnLane0Op> {
+  using OpRewritePattern<WarpExecuteOnLane0Op>::OpRewritePattern;
+  LogicalResult matchAndRewrite(WarpExecuteOnLane0Op warpOp,
+                                PatternRewriter &rewriter) const override {
+    auto yield = cast<vector::YieldOp>(
+        warpOp.getBodyRegion().getBlocks().begin()->getTerminator());
+    // Only pick up forOp if it is the last op in the region.
+    Operation *lastNode = yield->getPrevNode();
+    auto forOp = dyn_cast_or_null<scf::ForOp>(lastNode);
+    if (!forOp)
+      return failure();
+    SmallVector<Value> newOperands;
+    SmallVector<unsigned> resultIdx;
+    // Collect all the outputs coming from the forOp.
+    for (OpOperand &yieldOperand : yield->getOpOperands()) {
+      if (yieldOperand.get().getDefiningOp() != forOp.getOperation())
+        continue;
+      auto forResult = yieldOperand.get().cast<OpResult>();
+      newOperands.push_back(warpOp.getResult(yieldOperand.getOperandNumber()));
+      yieldOperand.set(forOp.getIterOperands()[forResult.getResultNumber()]);
+      resultIdx.push_back(yieldOperand.getOperandNumber());
+    }
+    OpBuilder::InsertionGuard g(rewriter);
+    rewriter.setInsertionPointAfter(warpOp);
+    // Create a new for op outside the region with a WarpExecuteOnLane0Op region
+    // inside.
+    auto newForOp = rewriter.create<scf::ForOp>(
+        forOp.getLoc(), forOp.getLowerBound(), forOp.getUpperBound(),
+        forOp.getStep(), newOperands);
+    rewriter.setInsertionPoint(newForOp.getBody(), newForOp.getBody()->begin());
+    auto innerWarp = rewriter.create<WarpExecuteOnLane0Op>(
+        warpOp.getLoc(), newForOp.getResultTypes(), warpOp.getLaneid(),
+        warpOp.getWarpSize(), newForOp.getRegionIterArgs(),
+        forOp.getResultTypes());
+
+    SmallVector<Value> argMapping;
+    argMapping.push_back(newForOp.getInductionVar());
+    for (Value args : innerWarp.getBody()->getArguments()) {
+      argMapping.push_back(args);
+    }
+    SmallVector<Value> yieldOperands;
+    for (Value operand : forOp.getBody()->getTerminator()->getOperands())
+      yieldOperands.push_back(operand);
+    rewriter.eraseOp(forOp.getBody()->getTerminator());
+    rewriter.mergeBlocks(forOp.getBody(), innerWarp.getBody(), argMapping);
+    rewriter.setInsertionPoint(innerWarp.getBody(), innerWarp.getBody()->end());
+    rewriter.create<vector::YieldOp>(innerWarp.getLoc(), yieldOperands);
+    rewriter.setInsertionPointAfter(innerWarp);
+    rewriter.create<scf::YieldOp>(forOp.getLoc(), innerWarp.getResults());
+    rewriter.eraseOp(forOp);
+    // Replace the warpOp result coming from the original ForOp.
+    for (const auto &res : llvm::enumerate(resultIdx)) {
+      warpOp.getResult(res.value())
+          .replaceAllUsesWith(newForOp.getResult(res.index()));
+      newForOp->setOperand(res.index() + 3, warpOp.getResult(res.value()));
+    }
+    return success();
+  }
+};
+
 } // namespace
 
 void mlir::vector::populateWarpExecuteOnLane0OpToScfForPattern(
     RewritePatternSet &patterns,
     const WarpExecuteOnLane0LoweringOptions &options) {
   patterns.add<WarpOpToScfForPattern>(patterns.getContext(), options);
 }
 
 void mlir::vector::populateDistributeTransferWriteOpPatterns(
     RewritePatternSet &patterns, DistributionMapFn distributionMapFn) {
   patterns.add<WarpOpTransferWrite>(patterns.getContext(), distributionMapFn);
 }
 
+void mlir::vector::populatePropagateWarpVectorDistributionPatterns(
+    RewritePatternSet &patterns) {
+  patterns.add<WarpOpElementwise, WarpOpTransferRead, WarpOpDeadResult,
+               WarpOpBroadcast, WarpOpForwardOperand, WarpOpScfForOp>(
+      patterns.getContext());
+}
+
 void mlir::vector::moveScalarUniformCode(WarpExecuteOnLane0Op warpOp) {
   Block *body = warpOp.getBody();
 
   // Keep track of the ops we want to hoist.
   llvm::SmallSetVector<Operation *, 8> opsToMove;
 
   // Helper to check if a value is or will be defined outside of the region.
   auto isDefinedOutsideOfBody = [&](Value value) {

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

 // RUN: mlir-opt %s -allow-unregistered-dialect -split-input-file -test-vector-warp-distribute=rewrite-warp-ops-to-scf-if | FileCheck %s --check-prefix=CHECK-SCF-IF
 // RUN: mlir-opt %s -allow-unregistered-dialect -split-input-file -test-vector-warp-distribute="hoist-uniform" | FileCheck --check-prefixes=CHECK-HOIST %s
 // RUN: mlir-opt %s -allow-unregistered-dialect -split-input-file -test-vector-warp-distribute="hoist-uniform distribute-transfer-write" | FileCheck --check-prefixes=CHECK-D %s
+// RUN: mlir-opt %s -allow-unregistered-dialect -split-input-file -test-vector-warp-distribute=propagate-distribution -canonicalize | FileCheck --check-prefixes=CHECK-PROP %s
 
 // CHECK-SCF-IF-DAG: memref.global "private" @__shared_32xf32 : memref<32xf32, 3>
 // CHECK-SCF-IF-DAG: memref.global "private" @__shared_64xf32 : memref<64xf32, 3>
 // CHECK-SCF-IF-DAG: memref.global "private" @__shared_128xf32 : memref<128xf32, 3>
 // CHECK-SCF-IF-DAG: memref.global "private" @__shared_256xf32 : memref<256xf32, 3>
 
 // CHECK-SCF-IF-LABEL: func @rewrite_warp_op_to_scf_if(
 //  CHECK-SCF-IF-SAME:     %[[laneid:.*]]: index,

 func.func @warp_extract(%laneid: index, %arg1: memref<1024xf32>, %gid : index) {
   vector.warp_execute_on_lane_0(%laneid)[32] {
     %sa = memref.subview %arg1[%gid] [128] [1] : memref<1024xf32> to memref<128xf32, #map2>
     %c0 = arith.constant 0 : index
     %v = "test.dummy_op"() : () -> (vector<1xf32>)
     vector.transfer_write %v, %sa[%c0] : vector<1xf32>, memref<128xf32, #map2>
   }
   return
-}
+}
+
+// -----
+
+// CHECK-PROP-LABEL:   func @warp_dead_result(
+func.func @warp_dead_result(%laneid: index) -> (vector<1xf32>) {
+  // CHECK-PROP: %[[R:.*]] = vector.warp_execute_on_lane_0(%{{.*}})[32] -> (vector<1xf32>)
+  %r:3 = vector.warp_execute_on_lane_0(%laneid)[32] ->
+    (vector<1xf32>, vector<1xf32>, vector<1xf32>) {
+    %2 = "some_def"() : () -> (vector<32xf32>)
+    %3 = "some_def"() : () -> (vector<32xf32>)
+    %4 = "some_def"() : () -> (vector<32xf32>)
+  // CHECK-PROP:   vector.yield %{{.*}} : vector<32xf32>
+    vector.yield %2, %3, %4 : vector<32xf32>, vector<32xf32>, vector<32xf32>
+  }
+  // CHECK-PROP: return %[[R]] : vector<1xf32>
+  return %r#1 : vector<1xf32>
+}
+
+// -----
+
+// CHECK-PROP-LABEL:   func @warp_propagate_operand(
+//  CHECK-PROP-SAME:   %[[ID:.*]]: index, %[[V:.*]]: vector<4xf32>)
+func.func @warp_propagate_operand(%laneid: index, %v0: vector<4xf32>)
+  -> (vector<4xf32>) {
+  %r = vector.warp_execute_on_lane_0(%laneid)[32]
+     args(%v0 : vector<4xf32>) -> (vector<4xf32>) {
+     ^bb0(%arg0 : vector<128xf32>) :
+    vector.yield %arg0 : vector<128xf32>
+  }
+  // CHECK-PROP: return %[[V]] : vector<4xf32>
+  return %r : vector<4xf32>
+}
+
+// -----
+
+#map0 = affine_map<()[s0] -> (s0 * 2)>
+
+// CHECK-PROP-LABEL:   func @warp_propagate_elementwise(
+func.func @warp_propagate_elementwise(%laneid: index, %dest: memref<1024xf32>) {
+  %c0 = arith.constant 0 : index
+  %c32 = arith.constant 0 : index
+  %cst = arith.constant 0.000000e+00 : f32
+  // CHECK-PROP: %[[R:.*]]:4 = vector.warp_execute_on_lane_0(%{{.*}})[32] -> (vector<1xf32>, vector<1xf32>, vector<2xf32>, vector<2xf32>)
+  %r:2 = vector.warp_execute_on_lane_0(%laneid)[32] ->
+    (vector<1xf32>, vector<2xf32>) {
+    // CHECK-PROP: %[[V0:.*]] = "some_def"() : () -> vector<32xf32>
+    // CHECK-PROP: %[[V1:.*]] = "some_def"() : () -> vector<32xf32>
+    // CHECK-PROP: %[[V2:.*]] = "some_def"() : () -> vector<64xf32>
+    // CHECK-PROP: %[[V3:.*]] = "some_def"() : () -> vector<64xf32>
+    // CHECK-PROP: vector.yield %[[V0]], %[[V1]], %[[V2]], %[[V3]] : vector<32xf32>, vector<32xf32>, vector<64xf32>, vector<64xf32>
+    %2 = "some_def"() : () -> (vector<32xf32>)
+    %3 = "some_def"() : () -> (vector<32xf32>)
+    %4 = "some_def"() : () -> (vector<64xf32>)
+    %5 = "some_def"() : () -> (vector<64xf32>)
+    %6 = arith.addf %2, %3 : vector<32xf32>
+    %7 = arith.addf %4, %5 : vector<64xf32>
+    vector.yield %6, %7 : vector<32xf32>, vector<64xf32>
+  }
+  // CHECK-PROP: %[[A0:.*]] = arith.addf %[[R]]#2, %[[R]]#3 : vector<2xf32>
+  // CHECK-PROP: %[[A1:.*]] = arith.addf %[[R]]#0, %[[R]]#1 : vector<1xf32>
+  %id2 = affine.apply #map0()[%laneid]
+  // CHECK-PROP: vector.transfer_write %[[A1]], {{.*}} : vector<1xf32>, memref<1024xf32>
+  // CHECK-PROP: vector.transfer_write %[[A0]], {{.*}} : vector<2xf32>, memref<1024xf32>
+  vector.transfer_write %r#0, %dest[%laneid] : vector<1xf32>, memref<1024xf32>
+  vector.transfer_write %r#1, %dest[%id2] : vector<2xf32>, memref<1024xf32>
+  return
+}
+
+// -----
+
+// CHECK-PROP-LABEL: func @warp_propagate_scalar_arith(
+//       CHECK-PROP:   %[[r:.*]]:2 = vector.warp_execute_on_lane_0{{.*}} {
+//       CHECK-PROP:     %[[some_def0:.*]] = "some_def"
+//       CHECK-PROP:     %[[some_def1:.*]] = "some_def"
+//       CHECK-PROP:     vector.yield %[[some_def0]], %[[some_def1]]
+//       CHECK-PROP:   }
+//       CHECK-PROP:   arith.addf %[[r]]#0, %[[r]]#1 : f32
+func.func @warp_propagate_scalar_arith(%laneid: index) {
+  %r = vector.warp_execute_on_lane_0(%laneid)[32] -> (f32) {
+    %0 = "some_def"() : () -> (f32)
+    %1 = "some_def"() : () -> (f32)
+    %2 = arith.addf %0, %1 : f32
+    vector.yield %2 : f32
+  }
+  vector.print %r : f32
+  return
+}
+
+// -----
+
+// CHECK-PROP-LABEL: func @warp_propagate_cast(
+//   CHECK-PROP-NOT:   vector.warp_execute_on_lane_0
+//       CHECK-PROP:   %[[result:.*]] = arith.sitofp %{{.*}} : i32 to f32
+//       CHECK-PROP:   return %[[result]]
+func.func @warp_propagate_cast(%laneid : index, %i : i32) -> (f32) {
+  %r = vector.warp_execute_on_lane_0(%laneid)[32] -> (f32) {
+    %casted = arith.sitofp %i : i32 to f32
+    vector.yield %casted : f32
+  }
+  return %r : f32
+}
+
+// -----
+
+#map0 = affine_map<()[s0] -> (s0 * 2)>
+
+//  CHECK-PROP-DAG: #[[MAP0:.*]] = affine_map<()[s0] -> (s0 * 2)>
+
+// CHECK-PROP:   func @warp_propagate_read
+//  CHECK-PROP-SAME:     (%[[ID:.*]]: index
+func.func @warp_propagate_read(%laneid: index, %src: memref<1024xf32>, %dest: memref<1024xf32>) {
+// CHECK-PROP-NOT: warp_execute_on_lane_0
+// CHECK-PROP-DAG: %[[R0:.*]] = vector.transfer_read %arg1[%[[ID]]], %{{.*}} : memref<1024xf32>, vector<1xf32>
+// CHECK-PROP-DAG: %[[ID2:.*]] = affine.apply #[[MAP0]]()[%[[ID]]]
+// CHECK-PROP-DAG: %[[R1:.*]] = vector.transfer_read %arg1[%[[ID2]]], %{{.*}} : memref<1024xf32>, vector<2xf32>
+// CHECK-PROP: vector.transfer_write %[[R0]], {{.*}} : vector<1xf32>, memref<1024xf32>
+// CHECK-PROP: vector.transfer_write %[[R1]], {{.*}} : vector<2xf32>, memref<1024xf32>
+  %c0 = arith.constant 0 : index
+  %c32 = arith.constant 0 : index
+  %cst = arith.constant 0.000000e+00 : f32
+  %r:2 = vector.warp_execute_on_lane_0(%laneid)[32] ->(vector<1xf32>, vector<2xf32>) {
+    %2 = vector.transfer_read %src[%c0], %cst : memref<1024xf32>, vector<32xf32>
+    %3 = vector.transfer_read %src[%c32], %cst : memref<1024xf32>, vector<64xf32>
+    vector.yield %2, %3 : vector<32xf32>, vector<64xf32>
+  }
+  %id2 = affine.apply #map0()[%laneid]
+  vector.transfer_write %r#0, %dest[%laneid] : vector<1xf32>, memref<1024xf32>
+  vector.transfer_write %r#1, %dest[%id2] : vector<2xf32>, memref<1024xf32>
+  return
+}
+
+// -----
+
+// CHECK-PROP-LABEL: func @fold_vector_broadcast(
+//       CHECK-PROP:   %[[r:.*]] = vector.warp_execute_on_lane_0{{.*}} -> (vector<1xf32>)
+//       CHECK-PROP:     %[[some_def:.*]] = "some_def"
+//       CHECK-PROP:     vector.yield %[[some_def]] : vector<1xf32>
+//       CHECK-PROP:   vector.print %[[r]] : vector<1xf32>
+func.func @fold_vector_broadcast(%laneid: index) {
+  %r = vector.warp_execute_on_lane_0(%laneid)[32] -> (vector<1xf32>) {
+    %0 = "some_def"() : () -> (vector<1xf32>)
+    %1 = vector.broadcast %0 : vector<1xf32> to vector<32xf32>
+    vector.yield %1 : vector<32xf32>
+  }
+  vector.print %r : vector<1xf32>
+  return
+}
+
+// -----
+
+// CHECK-PROP-LABEL: func @extract_vector_broadcast(
+//       CHECK-PROP:   %[[r:.*]] = vector.warp_execute_on_lane_0{{.*}} -> (vector<1xf32>)
+//       CHECK-PROP:     %[[some_def:.*]] = "some_def"
+//       CHECK-PROP:     vector.yield %[[some_def]] : vector<1xf32>
+//       CHECK-PROP:   %[[broadcasted:.*]] = vector.broadcast %[[r]] : vector<1xf32> to vector<2xf32>
+//       CHECK-PROP:   vector.print %[[broadcasted]] : vector<2xf32>
+func.func @extract_vector_broadcast(%laneid: index) {
+  %r = vector.warp_execute_on_lane_0(%laneid)[32] -> (vector<2xf32>) {
+    %0 = "some_def"() : () -> (vector<1xf32>)
+    %1 = vector.broadcast %0 : vector<1xf32> to vector<64xf32>
+    vector.yield %1 : vector<64xf32>
+  }
+  vector.print %r : vector<2xf32>
+  return
+}
+
+// -----
+
+// CHECK-PROP-LABEL: func @extract_scalar_vector_broadcast(
+//       CHECK-PROP:   %[[r:.*]] = vector.warp_execute_on_lane_0{{.*}} -> (f32)
+//       CHECK-PROP:     %[[some_def:.*]] = "some_def"
+//       CHECK-PROP:     vector.yield %[[some_def]] : f32
+//       CHECK-PROP:   %[[broadcasted:.*]] = vector.broadcast %[[r]] : f32 to vector<2xf32>
+//       CHECK-PROP:   vector.print %[[broadcasted]] : vector<2xf32>
+func.func @extract_scalar_vector_broadcast(%laneid: index) {
+  %r = vector.warp_execute_on_lane_0(%laneid)[32] -> (vector<2xf32>) {
+    %0 = "some_def"() : () -> (f32)
+    %1 = vector.broadcast %0 : f32 to vector<64xf32>
+    vector.yield %1 : vector<64xf32>
+  }
+  vector.print %r : vector<2xf32>
+  return
+}
+
+// -----
+
+// CHECK-PROP-LABEL:   func @warp_scf_for(
+// CHECK-PROP: %[[INI:.*]] = vector.warp_execute_on_lane_0(%{{.*}})[32] -> (vector<4xf32>) {
+// CHECK-PROP:   %[[INI1:.*]] = "some_def"() : () -> vector<128xf32>
+// CHECK-PROP:   vector.yield %[[INI1]] : vector<128xf32>
+// CHECK-PROP: }
+// CHECK-PROP: %[[F:.*]] = scf.for %{{.*}} = %{{.*}} to %{{.*}} step %{{.*}} iter_args(%[[FARG:.*]] = %[[INI]]) -> (vector<4xf32>) {
+// CHECK-PROP:   %[[W:.*]] = vector.warp_execute_on_lane_0(%{{.*}})[32] args(%[[FARG]] : vector<4xf32>) -> (vector<4xf32>) {
+// CHECK-PROP:    ^bb0(%[[ARG:.*]]: vector<128xf32>):
+// CHECK-PROP:      %[[ACC:.*]] = "some_def"(%[[ARG]]) : (vector<128xf32>) -> vector<128xf32>
+// CHECK-PROP:      vector.yield %[[ACC]] : vector<128xf32>
+// CHECK-PROP:   }
+// CHECK-PROP:   scf.yield %[[W]] : vector<4xf32>
+// CHECK-PROP: }
+// CHECK-PROP: "some_use"(%[[F]]) : (vector<4xf32>) -> ()
+func.func @warp_scf_for(%arg0: index) {
+  %c128 = arith.constant 128 : index
+  %c1 = arith.constant 1 : index
+  %c0 = arith.constant 0 : index
+  %0 = vector.warp_execute_on_lane_0(%arg0)[32] -> (vector<4xf32>) {
+    %ini = "some_def"() : () -> (vector<128xf32>)
+    %3 = scf.for %arg3 = %c0 to %c128 step %c1 iter_args(%arg4 = %ini) -> (vector<128xf32>) {
+      %acc = "some_def"(%arg4) : (vector<128xf32>) -> (vector<128xf32>)
+      scf.yield %acc : vector<128xf32>
+    }
+    vector.yield %3 : vector<128xf32>
+  }
+  "some_use"(%0) : (vector<4xf32>) -> ()
+  return
+}
+
+// -----
+
+// CHECK-PROP-LABEL:   func @warp_scf_for_swap(
+// CHECK-PROP: %[[INI:.*]]:2 = vector.warp_execute_on_lane_0(%{{.*}})[32] -> (vector<4xf32>, vector<4xf32>) {
+// CHECK-PROP:   %[[INI1:.*]] = "some_def"() : () -> vector<128xf32>
+// CHECK-PROP:   %[[INI2:.*]] = "some_def"() : () -> vector<128xf32>
+// CHECK-PROP:   vector.yield %[[INI1]], %[[INI2]] : vector<128xf32>, vector<128xf32>
+// CHECK-PROP: }
+// CHECK-PROP: %[[F:.*]]:2 = scf.for %{{.*}} = %{{.*}} to %{{.*}} step %{{.*}} iter_args(%[[FARG1:.*]] = %[[INI]]#0, %[[FARG2:.*]] = %[[INI]]#1) -> (vector<4xf32>, vector<4xf32>) {
+// CHECK-PROP:   %[[W:.*]]:2 = vector.warp_execute_on_lane_0(%{{.*}})[32] args(%[[FARG1]], %[[FARG2]] : vector<4xf32>, vector<4xf32>) -> (vector<4xf32>, vector<4xf32>) {
+// CHECK-PROP:    ^bb0(%[[ARG1:.*]]: vector<128xf32>, %[[ARG2:.*]]: vector<128xf32>):
+// CHECK-PROP:      %[[ACC1:.*]] = "some_def"(%[[ARG1]]) : (vector<128xf32>) -> vector<128xf32>
+// CHECK-PROP:      %[[ACC2:.*]] = "some_def"(%[[ARG2]]) : (vector<128xf32>) -> vector<128xf32>
+// CHECK-PROP:      vector.yield %[[ACC2]], %[[ACC1]] : vector<128xf32>, vector<128xf32>
+// CHECK-PROP:   }
+// CHECK-PROP:   scf.yield %[[W]]#0, %[[W]]#1 : vector<4xf32>, vector<4xf32>
+// CHECK-PROP: }
+// CHECK-PROP: "some_use"(%[[F]]#0) : (vector<4xf32>) -> ()
+// CHECK-PROP: "some_use"(%[[F]]#1) : (vector<4xf32>) -> ()
+func.func @warp_scf_for_swap(%arg0: index) {
+  %c128 = arith.constant 128 : index
+  %c1 = arith.constant 1 : index
+  %c0 = arith.constant 0 : index
+  %0:2 = vector.warp_execute_on_lane_0(%arg0)[32] -> (vector<4xf32>, vector<4xf32>) {
+    %ini1 = "some_def"() : () -> (vector<128xf32>)
+    %ini2 = "some_def"() : () -> (vector<128xf32>)
+    %3:2 = scf.for %arg3 = %c0 to %c128 step %c1 iter_args(%arg4 = %ini1, %arg5 = %ini2) -> (vector<128xf32>, vector<128xf32>) {
+      %acc1 = "some_def"(%arg4) : (vector<128xf32>) -> (vector<128xf32>)
+      %acc2 = "some_def"(%arg5) : (vector<128xf32>) -> (vector<128xf32>)
+      scf.yield %acc2, %acc1 : vector<128xf32>, vector<128xf32>
+    }
+    vector.yield %3#0, %3#1 : vector<128xf32>, vector<128xf32>
+  }
+  "some_use"(%0#0) : (vector<4xf32>) -> ()
+  "some_use"(%0#1) : (vector<4xf32>) -> ()
+  return
+}
+
+// -----
+
+#map = affine_map<()[s0] -> (s0 * 4)>
+#map1 = affine_map<()[s0] -> (s0 * 128 + 128)>
+#map2 = affine_map<()[s0] -> (s0 * 4 + 128)>
+
+// CHECK-PROP-LABEL:   func @warp_scf_for_multiple_yield(
+//       CHECK-PROP:   vector.warp_execute_on_lane_0(%{{.*}})[32] -> (vector<1xf32>) {
+//  CHECK-PROP-NEXT:     "some_def"() : () -> vector<32xf32>
+//  CHECK-PROP-NEXT:     vector.yield %{{.*}} : vector<32xf32>
+//  CHECK-PROP-NEXT:   }
+//   CHECK-PROP-NOT:   vector.warp_execute_on_lane_0
+//       CHECK-PROP:   vector.transfer_read {{.*}} : memref<?xf32>, vector<4xf32>
+//       CHECK-PROP:   vector.transfer_read {{.*}} : memref<?xf32>, vector<4xf32>
+//       CHECK-PROP:   %{{.*}}:2 = scf.for {{.*}} -> (vector<4xf32>, vector<4xf32>) {
+//   CHECK-PROP-NOT:     vector.warp_execute_on_lane_0
+//       CHECK-PROP:     vector.transfer_read {{.*}} : memref<?xf32>, vector<4xf32>
+//       CHECK-PROP:     vector.transfer_read {{.*}} : memref<?xf32>, vector<4xf32>
+//       CHECK-PROP:     arith.addf {{.*}} : vector<4xf32>
+//       CHECK-PROP:     arith.addf {{.*}} : vector<4xf32>
+//       CHECK-PROP:     scf.yield {{.*}} : vector<4xf32>, vector<4xf32>
+//       CHECK-PROP:   }
+func.func @warp_scf_for_multiple_yield(%arg0: index, %arg1: memref<?xf32>, %arg2: memref<?xf32>) {
+  %c256 = arith.constant 256 : index
+  %c128 = arith.constant 128 : index
+  %c1 = arith.constant 1 : index
+  %c0 = arith.constant 0 : index
+  %cst = arith.constant 0.000000e+00 : f32
+  %0:3 = vector.warp_execute_on_lane_0(%arg0)[32] ->
+  (vector<1xf32>, vector<4xf32>, vector<4xf32>) {
+    %def = "some_def"() : () -> (vector<32xf32>)
+    %r1 = vector.transfer_read %arg2[%c0], %cst {in_bounds = [true]} : memref<?xf32>, vector<128xf32>
+    %r2 = vector.transfer_read %arg2[%c128], %cst {in_bounds = [true]} : memref<?xf32>, vector<128xf32>
+    %3:2 = scf.for %arg3 = %c0 to %c128 step %c1 iter_args(%arg4 = %r1, %arg5 = %r2)
+    -> (vector<128xf32>, vector<128xf32>) {
+      %o1 = affine.apply #map1()[%arg3]
+      %o2 = affine.apply #map2()[%arg3]
+      %4 = vector.transfer_read %arg1[%o1], %cst {in_bounds = [true]} : memref<?xf32>, vector<128xf32>
+      %5 = vector.transfer_read %arg1[%o2], %cst {in_bounds = [true]} : memref<?xf32>, vector<128xf32>
+      %6 = arith.addf %4, %arg4 : vector<128xf32>
+      %7 = arith.addf %5, %arg5 : vector<128xf32>
+      scf.yield %6, %7 : vector<128xf32>, vector<128xf32>
+    }
+    vector.yield %def, %3#0, %3#1 :  vector<32xf32>, vector<128xf32>, vector<128xf32>
+  }
+  %1 = affine.apply #map()[%arg0]
+  vector.transfer_write %0#1, %arg2[%1] {in_bounds = [true]} : vector<4xf32>, memref<?xf32>
+  %2 = affine.apply #map2()[%arg0]
+  vector.transfer_write %0#2, %arg2[%2] {in_bounds = [true]} : vector<4xf32>, memref<?xf32>
+  "some_use"(%0#0) : (vector<1xf32>) -> ()
+  return
+}

mlir/test/Integration/Dialect/Vector/GPU/CUDA/test-warp-distribute.mlir

 // Run the test cases without distributing ops to test default lowering. Run
 // everything on the same thread.
 // RUN: mlir-opt %s -test-vector-warp-distribute=rewrite-warp-ops-to-scf-if -canonicalize | \
 // RUN: mlir-opt -convert-scf-to-cf -convert-cf-to-llvm -convert-vector-to-llvm -convert-arith-to-llvm \
 // RUN:  -gpu-kernel-outlining \
 // RUN:  -pass-pipeline='gpu.module(strip-debuginfo,convert-gpu-to-nvvm,reconcile-unrealized-casts,gpu-to-cubin)' \
 // RUN:  -gpu-to-llvm -reconcile-unrealized-casts |\
 // RUN: mlir-cpu-runner -e main -entry-point-result=void \
 // RUN:   -shared-libs=%mlir_runner_utils_dir/libmlir_cuda_runtime%shlibext \
 // RUN:   -shared-libs=%mlir_runner_utils_dir/libmlir_c_runner_utils%shlibext \
 // RUN:   -shared-libs=%mlir_runner_utils_dir/libmlir_runner_utils%shlibext | \
 // RUN: FileCheck %s
 
+// Run the same test cases with distribution and propagation.
+// RUN: mlir-opt %s  -test-vector-warp-distribute="hoist-uniform distribute-transfer-write" \
+// RUN:   -test-vector-warp-distribute=rewrite-warp-ops-to-scf-if -canonicalize | \
+// RUN: mlir-opt -convert-scf-to-cf -convert-cf-to-llvm -convert-vector-to-llvm -convert-arith-to-llvm \
+// RUN:  -gpu-kernel-outlining \
+// RUN:  -pass-pipeline='gpu.module(strip-debuginfo,convert-gpu-to-nvvm,reconcile-unrealized-casts,gpu-to-cubin)' \
+// RUN:  -gpu-to-llvm -reconcile-unrealized-casts |\
+// RUN: mlir-cpu-runner -e main -entry-point-result=void \
+// RUN:   -shared-libs=%mlir_runner_utils_dir/libmlir_cuda_runtime%shlibext \
+// RUN:   -shared-libs=%mlir_runner_utils_dir/libmlir_c_runner_utils%shlibext \
+// RUN:   -shared-libs=%mlir_runner_utils_dir/libmlir_runner_utils%shlibext | \
+// RUN: FileCheck %s
+
+// RUN: mlir-opt %s  -test-vector-warp-distribute="hoist-uniform distribute-transfer-write propagate-distribution" \
+// RUN:   -test-vector-warp-distribute=rewrite-warp-ops-to-scf-if -canonicalize | \
+// RUN: mlir-opt -convert-scf-to-cf -convert-cf-to-llvm -convert-vector-to-llvm -convert-arith-to-llvm \
+// RUN:  -gpu-kernel-outlining \
+// RUN:  -pass-pipeline='gpu.module(strip-debuginfo,convert-gpu-to-nvvm,reconcile-unrealized-casts,gpu-to-cubin)' \
+// RUN:  -gpu-to-llvm -reconcile-unrealized-casts |\
+// RUN: mlir-cpu-runner -e main -entry-point-result=void \
+// RUN:   -shared-libs=%mlir_runner_utils_dir/libmlir_cuda_runtime%shlibext \
+// RUN:   -shared-libs=%mlir_runner_utils_dir/libmlir_c_runner_utils%shlibext \
+// RUN:   -shared-libs=%mlir_runner_utils_dir/libmlir_runner_utils%shlibext | \
+// RUN: FileCheck %s
+
 func.func @gpu_func(%arg1: memref<32xf32>, %arg2: memref<32xf32>) {
   %c1 = arith.constant 1 : index
   %c0 = arith.constant 0 : index
   %c32 = arith.constant 32 : index
   %cst = arith.constant 0.000000e+00 : f32
   gpu.launch blocks(%arg3, %arg4, %arg5)
   in (%arg9 = %c1, %arg10 = %c1, %arg11 = %c1)
   threads(%arg6, %arg7, %arg8) in (%arg12 = %c32, %arg13 = %c1, %arg14 = %c1) {

mlir/test/lib/Dialect/Vector/TestVectorTransforms.cpp

Option<bool> distributeTransferWriteOps{
       *this, "distribute-transfer-write",
       llvm::cl::desc("Test distribution of transfer write"),
       llvm::cl::init(false)};
 
   Option<bool> hoistUniform{*this, "hoist-uniform",
                             llvm::cl::desc("Test hoist uniform"),
                             llvm::cl::init(false)};
 
+  Option<bool> propagateDistribution{
+      *this, "propagate-distribution",
+      llvm::cl::desc("Test distribution propgation"), llvm::cl::init(false)};
+
   void runOnOperation() override {
     RewritePatternSet patterns(&getContext());
 
     getOperation().walk([&](Operation *op) {
       if (auto warpOp = dyn_cast<WarpExecuteOnLane0Op>(op)) {
         if (hoistUniform) {
           moveScalarUniformCode(warpOp);
         }
if (distributeTransferWriteOps) {
         auto map =
             AffineMap::get(vecRank, 0, builder.getAffineDimExpr(vecRank - 1));
         return map;
       };
       RewritePatternSet patterns(ctx);
       populateDistributeTransferWriteOpPatterns(patterns, distributionFn);
       (void)applyPatternsAndFoldGreedily(getOperation(), std::move(patterns));
     }
-
+    if (propagateDistribution) {
+      RewritePatternSet patterns(ctx);
+      vector::populatePropagateWarpVectorDistributionPatterns(patterns);
+      (void)applyPatternsAndFoldGreedily(getOperation(), std::move(patterns));
+    }
     WarpExecuteOnLane0LoweringOptions options;
     options.warpAllocationFn = allocateGlobalSharedMemory;
     options.warpSyncronizationFn = [](Location loc, OpBuilder &builder,
                                       WarpExecuteOnLane0Op warpOp) {
       builder.create<gpu::BarrierOp>(loc);
     };
     // Test on one pattern in isolation.
     if (warpOpToSCF) {