[mlir][Vector] Extend xfer_read(extract)->scalar load to support multiple uses

[dcaballe]

Summary

This patch extends the vector.extract(vector.transfer_read) -> scalar
load patterns to support vector.transfer_read with multiple uses. For
now, we check that all the uses are vector.extract operations.
Supporting multiple uses is predicated under a flag.

[hanchung]

I'm curious that is there a case we want to disable the flag? Why it is not enabled by default?

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

120

can we add a doc about what it means if allowMultipleUses is enabled and is disabled?

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

590–591

I think this can be merged into one isa?

664–665

There are same checks in other pattern. Can we refactor them to a common method?

[dcaballe]

Address comments

[dcaballe]

Great feedback! Thanks!

I'm curious that is there a case we want to disable the flag? Why it is not enabled by default?

The original comment says "xfer result must have a single use. Otherwise, it may be better to perform a vector load.". Not sure what the context is. Perhaps if there are scalar and vector uses... or if the cache line is evicted before all the scalar loads are completed (weird)...
In any case, I think it's better to preserve the original behavior for now until we can learn more about it.

[hanchung]

In D150812#4355337, @dcaballe wrote:

Great feedback! Thanks!

I'm curious that is there a case we want to disable the flag? Why it is not enabled by default?

The original comment says "xfer result must have a single use. Otherwise, it may be better to perform a vector load.". Not sure what the context is. Perhaps if there are scalar and vector uses... or if the cache line is evicted before all the scalar loads are completed (weird)...
In any case, I think it's better to preserve the original behavior for now until we can learn more about it.

Got it, thanks!

[awarzynski]

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

115–118

Sorry, it's a bit late at my end, so perhaps I'm misreading this myself :) Please double-check.

[awarzynski]

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

612

IIUC, this method will rewrite unconditionally. But match will not match unconditionally. Put differently - this comment needs updating, right?

[dcaballe]

Fix doc

[dcaballe]

Thanks! Updating doc and landing. Let me know if there are any other doc issues. I'll be happy to address them separately.

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

115–118

Thanks! It looks like it was also late at my end when I wrote that :)

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

612

Updated doc, thanks!

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

 ///
 /// [ConstantMaskOp]
 /// Progressive lowering of ConstantMaskOp to lower-D ConstantMaskOp until
 /// dim 1.
 void populateVectorMaskOpLoweringPatterns(RewritePatternSet &patterns,
                                           PatternBenefit benefit = 1);
 
 /// Collects patterns that lower scalar vector transfer ops to memref loads and
-/// stores when beneficial.
+/// stores when beneficial. If `allowMultipleUses` is set to true, the patterns
+/// are applied to vector transfer reads with any number of uses. Otherwise,
+/// only vector transfer reads with a single use will be lowered.
 void populateScalarVectorTransferLoweringPatterns(RewritePatternSet &patterns,
-                                                  PatternBenefit benefit = 1);
+                                                  PatternBenefit benefit,
+                                                  bool allowMultipleUses);
 
 /// Populate the pattern set with the following patterns:
 ///
 /// [ShapeCastOp2DDownCastRewritePattern]
 /// ShapeOp 2D -> 1D downcast serves the purpose of flattening 2-D to 1-D
 /// vectors progressively.
 ///
 /// [ShapeCastOp2DUpCastRewritePattern]

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

vector::TransferWriteOp flatWrite =
         rewriter.create<vector::TransferWriteOp>(
             loc, flatVector, collapsedSource, collapsedIndices, collapsedMap);
     flatWrite.setInBoundsAttr(rewriter.getBoolArrayAttr({true}));
     rewriter.eraseOp(transferWriteOp);
     return success();
   }
 };
 
-/// Rewrite extractelement(transfer_read) to memref.load.
-///
-/// Rewrite only if the extractelement op is the single user of the transfer op.
-/// E.g., do not rewrite IR such as:
-/// %0 = vector.transfer_read ... : vector<1024xf32>
-/// %1 = vector.extractelement %0[%a : index] : vector<1024xf32>
-/// %2 = vector.extractelement %0[%b : index] : vector<1024xf32>
-/// Rewriting such IR (replacing one vector load with multiple scalar loads) may
-/// negatively affect performance.
-class RewriteScalarExtractElementOfTransferRead
-    : public OpRewritePattern<vector::ExtractElementOp> {
-  using OpRewritePattern::OpRewritePattern;
+/// Base class for `vector.extract/vector.extract_element(vector.transfer_read)`
+/// to `memref.load` patterns. The `match` method is shared for both
+/// `vector.extract` and `vector.extract_element`.
+template <class VectorExtractOp>
+class RewriteScalarExtractOfTransferReadBase
+    : public OpRewritePattern<VectorExtractOp> {
+  using Base = OpRewritePattern<VectorExtractOp>;
 
-  LogicalResult matchAndRewrite(vector::ExtractElementOp extractOp,
-                                PatternRewriter &rewriter) const override {
-    auto xferOp = extractOp.getVector().getDefiningOp<vector::TransferReadOp>();
+public:
+  RewriteScalarExtractOfTransferReadBase(MLIRContext *context,
+                                         PatternBenefit benefit,
+                                         bool allowMultipleUses)
+      : Base::OpRewritePattern(context, benefit),
+        allowMultipleUses(allowMultipleUses) {}
+
+  LogicalResult match(VectorExtractOp extractOp) const override {
+    auto xferOp =
+        extractOp.getVector().template getDefiningOp<vector::TransferReadOp>();
     if (!xferOp)
       return failure();
-    // xfer result must have a single use. Otherwise, it may be better to
-    // perform a vector load.
-    if (!extractOp.getVector().hasOneUse())
+    // If multiple uses are not allowed, check if xfer has a single use.
+    if (!allowMultipleUses && !xferOp.getResult().hasOneUse())
+      return failure();
+    // If multiple uses are allowed, check if all the xfer uses are extract ops.
+    if (allowMultipleUses &&
+        !llvm::all_of(xferOp->getUses(), [](OpOperand &use) {
+          return isa<vector::ExtractOp, vector::ExtractElementOp>(
+              use.getOwner());
+        }))
       return failure();
     // Mask not supported.
     if (xferOp.getMask())
       return failure();
     // Map not supported.
     if (!xferOp.getPermutationMap().isMinorIdentity())
       return failure();
-    // Cannot rewrite if the indices may be out of bounds. The starting point is
-    // always inbounds, so we don't care in case of 0d transfers.
-    if (xferOp.hasOutOfBoundsDim() && xferOp.getType().getRank() > 0)
+    // Cannot rewrite if the indices may be out of bounds.
+    if (xferOp.hasOutOfBoundsDim())
       return failure();
+    return success();
+  }
+
+private:
+  bool allowMultipleUses;
+};
+
+/// Rewrite `vector.extractelement(vector.transfer_read)` to `memref.load`.
+///
+/// All the users of the transfer op must be either `vector.extractelement` or
+/// `vector.extract` ops. If `allowMultipleUses` is set to true, rewrite
+/// transfer ops with any number of users. Otherwise, rewrite only if the
+/// extract op is the single user of the transfer op. Rewriting a single
+/// vector load with multiple scalar loads may negatively affect performance.
+class RewriteScalarExtractElementOfTransferRead
+    : public RewriteScalarExtractOfTransferReadBase<vector::ExtractElementOp> {
+  using RewriteScalarExtractOfTransferReadBase::
+      RewriteScalarExtractOfTransferReadBase;
+
+  void rewrite(vector::ExtractElementOp extractOp,
+               PatternRewriter &rewriter) const override {
     // Construct scalar load.
+    auto xferOp = extractOp.getVector().getDefiningOp<vector::TransferReadOp>();
     SmallVector<Value> newIndices(xferOp.getIndices().begin(),
                                   xferOp.getIndices().end());
     if (extractOp.getPosition()) {
       AffineExpr sym0, sym1;
       bindSymbols(extractOp.getContext(), sym0, sym1);
       OpFoldResult ofr = affine::makeComposedFoldedAffineApply(
           rewriter, extractOp.getLoc(), sym0 + sym1,
           {newIndices[newIndices.size() - 1], extractOp.getPosition()});
       if (ofr.is<Value>()) {
         newIndices[newIndices.size() - 1] = ofr.get<Value>();
       } else {
         newIndices[newIndices.size() - 1] =
             rewriter.create<arith::ConstantIndexOp>(extractOp.getLoc(),
                                                     *getConstantIntValue(ofr));
       }
     }
     if (isa<MemRefType>(xferOp.getSource().getType())) {
       rewriter.replaceOpWithNewOp<memref::LoadOp>(extractOp, xferOp.getSource(),
                                                   newIndices);
     } else {
       rewriter.replaceOpWithNewOp<tensor::ExtractOp>(
           extractOp, xferOp.getSource(), newIndices);
     }
-    return success();
   }
 };
 
-/// Rewrite extract(transfer_read) to memref.load.
+/// Rewrite `vector.extractelement(vector.transfer_read)` to `memref.load`.
+/// Rewrite `vector.extract(vector.transfer_read)` to `memref.load`.
 ///
-/// Rewrite only if the extractelement op is the single user of the transfer op.
-/// E.g., do not rewrite IR such as:
-/// %0 = vector.transfer_read ... : vector<1024xf32>
-/// %1 = vector.extract %0[0] : vector<1024xf32>
-/// %2 = vector.extract %0[5] : vector<1024xf32>
-/// Rewriting such IR (replacing one vector load with multiple scalar loads) may
-/// negatively affect performance.
+/// All the users of the transfer op must be either `vector.extractelement` or
+/// `vector.extract` ops. If `allowMultipleUses` is set to true, rewrite
+/// transfer ops with any number of users. Otherwise, rewrite only if the
+/// extract op is the single user of the transfer op. Rewriting a single
+/// vector load with multiple scalar loads may negatively affect performance.
 class RewriteScalarExtractOfTransferRead
-    : public OpRewritePattern<vector::ExtractOp> {
-  using OpRewritePattern::OpRewritePattern;
+    : public RewriteScalarExtractOfTransferReadBase<vector::ExtractOp> {
+  using RewriteScalarExtractOfTransferReadBase::
+      RewriteScalarExtractOfTransferReadBase;
 
-  LogicalResult matchAndRewrite(vector::ExtractOp extractOp,
-                                PatternRewriter &rewriter) const override {
-    // Only match scalar extracts.
-    if (isa<VectorType>(extractOp.getType()))
-      return failure();
-    auto xferOp = extractOp.getVector().getDefiningOp<vector::TransferReadOp>();
-    if (!xferOp)
-      return failure();
-    // xfer result must have a single use. Otherwise, it may be better to
-    // perform a vector load.
-    if (!extractOp.getVector().hasOneUse())
-      return failure();
-    // Mask not supported.
-    if (xferOp.getMask())
-      return failure();
-    // Map not supported.
-    if (!xferOp.getPermutationMap().isMinorIdentity())
-      return failure();
-    // Cannot rewrite if the indices may be out of bounds. The starting point is
-    // always inbounds, so we don't care in case of 0d transfers.
-    if (xferOp.hasOutOfBoundsDim() && xferOp.getType().getRank() > 0)
-      return failure();
+  void rewrite(vector::ExtractOp extractOp,
     // Construct scalar load.
+    auto xferOp = extractOp.getVector().getDefiningOp<vector::TransferReadOp>();
     SmallVector<Value> newIndices(xferOp.getIndices().begin(),
                                   xferOp.getIndices().end());
     for (const auto &it : llvm::enumerate(extractOp.getPosition())) {
       int64_t offset = cast<IntegerAttr>(it.value()).getInt();
       int64_t idx =
           newIndices.size() - extractOp.getPosition().size() + it.index();
       OpFoldResult ofr = affine::makeComposedFoldedAffineApply(
           rewriter, extractOp.getLoc(),
           rewriter.getAffineSymbolExpr(0) + offset, {newIndices[idx]});
       if (ofr.is<Value>()) {
         newIndices[idx] = ofr.get<Value>();
       } else {
         newIndices[idx] = rewriter.create<arith::ConstantIndexOp>(
             extractOp.getLoc(), *getConstantIntValue(ofr));
       }
     }
     if (isa<MemRefType>(xferOp.getSource().getType())) {
       rewriter.replaceOpWithNewOp<memref::LoadOp>(extractOp, xferOp.getSource(),
                                                   newIndices);
     } else {
       rewriter.replaceOpWithNewOp<tensor::ExtractOp>(
           extractOp, xferOp.getSource(), newIndices);
     }
-    return success();
   }
 };
 
 /// Rewrite transfer_writes of vectors of size 1 (e.g., vector<1x1xf32>)
 /// to memref.store.
 class RewriteScalarWrite : public OpRewritePattern<vector::TransferWriteOp> {
   using OpRewritePattern::OpRewritePattern;
 
void mlir::vector::transferOpflowOpt(RewriterBase &rewriter,
   rootOp->walk([&](vector::TransferWriteOp write) {
     if (isa<MemRefType>(write.getShapedType()))
       opt.deadStoreOp(write);
   });
   opt.removeDeadOp();
 }
 
 void mlir::vector::populateScalarVectorTransferLoweringPatterns(
-    RewritePatternSet &patterns, PatternBenefit benefit) {
+    RewritePatternSet &patterns, PatternBenefit benefit,
+    bool allowMultipleUses) {
   patterns.add<RewriteScalarExtractElementOfTransferRead,
-               RewriteScalarExtractOfTransferRead, RewriteScalarWrite>(
-      patterns.getContext(), benefit);
+               RewriteScalarExtractOfTransferRead>(patterns.getContext(),
+                                                   benefit, allowMultipleUses);
+  patterns.add<RewriteScalarWrite>(patterns.getContext(), benefit);
 }
 
 void mlir::vector::populateVectorTransferDropUnitDimsPatterns(
     RewritePatternSet &patterns, PatternBenefit benefit) {
   patterns
       .add<TransferReadDropUnitDimsPattern, TransferWriteDropUnitDimsPattern>(
           patterns.getContext(), benefit);
   populateShapeCastFoldingPatterns(patterns);

mlir/test/Dialect/Vector/scalar-vector-transfer-to-memref.mlir

 // RUN: mlir-opt %s -test-scalar-vector-transfer-lowering -split-input-file | FileCheck %s
+// RUN: mlir-opt %s -test-scalar-vector-transfer-lowering=allow-multiple-uses -split-input-file | FileCheck %s --check-prefix=MULTIUSE
 
 // CHECK-LABEL: func @transfer_read_0d(
 //  CHECK-SAME:     %[[m:.*]]: memref<?x?x?xf32>, %[[idx:.*]]: index
 //       CHECK:   %[[r:.*]] = memref.load %[[m]][%[[idx]], %[[idx]], %[[idx]]]
 //       CHECK:   return %[[r]]
 func.func @transfer_read_0d(%m: memref<?x?x?xf32>, %idx: index) -> f32 {
   %cst = arith.constant 0.0 : f32
   %0 = vector.transfer_read %m[%idx, %idx, %idx], %cst : memref<?x?x?xf32>, vector<f32>

 //       CHECK:   %[[cst:.*]] = arith.constant dense<5.000000e+00> : vector<1x1xf32>
 //       CHECK:   %[[extract:.*]] = vector.extract %[[cst]][0, 0] : vector<1x1xf32>
 //       CHECK:   memref.store %[[extract]], %[[m]][%[[idx]], %[[idx]], %[[idx]]]
 func.func @transfer_write_arith_constant(%m: memref<?x?x?xf32>, %idx: index) {
   %cst = arith.constant dense<5.000000e+00> : vector<1x1xf32>
   vector.transfer_write %cst, %m[%idx, %idx, %idx] : vector<1x1xf32>, memref<?x?x?xf32>
   return
 }
+
+// -----
+
+// CHECK-LABEL: func @transfer_read_multi_use(
+//  CHECK-SAME:   %[[m:.*]]: memref<?xf32>, %[[idx:.*]]: index
+//   CHECK-NOT:   memref.load
+//       CHECK:   %[[r:.*]] = vector.transfer_read %[[m]][%[[idx]]]
+//       CHECK:   %[[e0:.*]] = vector.extract %[[r]][0]
+//       CHECK:   %[[e1:.*]] = vector.extract %[[r]][1]
+//       CHECK:   return %[[e0]], %[[e1]]
+
+// MULTIUSE-LABEL: func @transfer_read_multi_use(
+//  MULTIUSE-SAME:   %[[m:.*]]: memref<?xf32>, %[[idx0:.*]]: index
+//   MULTIUSE-NOT:   vector.transfer_read
+//       MULTIUSE:   %[[r0:.*]] = memref.load %[[m]][%[[idx0]]
+//       MULTIUSE:   %[[idx1:.*]] = affine.apply
+//       MULTIUSE:   %[[r1:.*]] = memref.load %[[m]][%[[idx1]]
+//       MULTIUSE:   return %[[r0]], %[[r1]]
+
+func.func @transfer_read_multi_use(%m: memref<?xf32>, %idx: index) -> (f32, f32) {
+  %cst = arith.constant 0.0 : f32
+  %0 = vector.transfer_read %m[%idx], %cst {in_bounds = [true]} : memref<?xf32>, vector<16xf32>
+  %1 = vector.extract %0[0] : vector<16xf32>
+  %2 = vector.extract %0[1] : vector<16xf32>
+  return %1, %2 : f32, f32
+}
+

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

 };
 
 struct TestScalarVectorTransferLoweringPatterns
     : public PassWrapper<TestScalarVectorTransferLoweringPatterns,
                          OperationPass<func::FuncOp>> {
   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(
       TestScalarVectorTransferLoweringPatterns)
 
+  TestScalarVectorTransferLoweringPatterns() = default;
+  TestScalarVectorTransferLoweringPatterns(
+      const TestScalarVectorTransferLoweringPatterns &pass)
+      : PassWrapper(pass) {}
+
   StringRef getArgument() const final {
     return "test-scalar-vector-transfer-lowering";
   }
   StringRef getDescription() const final {
     return "Test lowering of scalar vector transfers to memref loads/stores.";
   }
-  TestScalarVectorTransferLoweringPatterns() = default;
 
   void getDependentDialects(DialectRegistry &registry) const override {
     registry.insert<affine::AffineDialect, memref::MemRefDialect,
                     tensor::TensorDialect, vector::VectorDialect>();
   }
 
+  Option<bool> allowMultipleUses{
+      *this, "allow-multiple-uses",
+      llvm::cl::desc("Fold transfer operations with multiple uses"),
+      llvm::cl::init(false)};
+
   void runOnOperation() override {
     MLIRContext *ctx = &getContext();
     RewritePatternSet patterns(ctx);
-    vector::populateScalarVectorTransferLoweringPatterns(patterns);
+    vector::populateScalarVectorTransferLoweringPatterns(
+        patterns, /*benefit=*/1, allowMultipleUses.getValue());
     (void)applyPatternsAndFoldGreedily(getOperation(), std::move(patterns));
   }
 };
 
 struct TestVectorTransferOpt
     : public PassWrapper<TestVectorTransferOpt, OperationPass<func::FuncOp>> {
   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestVectorTransferOpt)