[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

118

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

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

595–596

I think this can be merged into one isa?

667–675

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–116

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

617

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–116

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

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

617

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.
 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

 /// %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;
 
+public:
+  RewriteScalarExtractElementOfTransferRead(MLIRContext *context,
+                                            PatternBenefit benefit,
+                                            bool allowMultipleUses)
+      : OpRewritePattern(context, benefit),
+        allowMultipleUses(allowMultipleUses) {}
+
   LogicalResult matchAndRewrite(vector::ExtractElementOp extractOp,
                                 PatternRewriter &rewriter) const override {
     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())
+    // 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>(use.getOwner()) ||
+                 isa<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();
     // Construct scalar load.
     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 (xferOp.getSource().getType().isa<MemRefType>()) {
       rewriter.replaceOpWithNewOp<memref::LoadOp>(extractOp, xferOp.getSource(),
                                                   newIndices);
     } else {
       rewriter.replaceOpWithNewOp<tensor::ExtractOp>(
           extractOp, xferOp.getSource(), newIndices);
     }
     return success();
   }
+
+private:
+  bool allowMultipleUses;
 };
 
 /// Rewrite extract(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.
 class RewriteScalarExtractOfTransferRead
     : public OpRewritePattern<vector::ExtractOp> {
   using OpRewritePattern::OpRewritePattern;
 
+public:
+  RewriteScalarExtractOfTransferRead(MLIRContext *context,
+                                     PatternBenefit benefit,
+                                     bool allowMultipleUses)
+      : OpRewritePattern(context, benefit),
+        allowMultipleUses(allowMultipleUses) {}
+
   LogicalResult matchAndRewrite(vector::ExtractOp extractOp,
                                 PatternRewriter &rewriter) const override {
     // Only match scalar extracts.
     if (extractOp.getType().isa<VectorType>())
       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())
+    // 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>(use.getOwner()) ||
+                 isa<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();
     // Construct scalar load.
     SmallVector<Value> newIndices(xferOp.getIndices().begin(),
                                   xferOp.getIndices().end());
     for (const auto &it : llvm::enumerate(extractOp.getPosition())) {
       int64_t offset = it.value().cast<IntegerAttr>().getInt();
       int64_t idx =
           newIndices.size() - extractOp.getPosition().size() + it.index();
if (xferOp.getSource().getType().isa<MemRefType>()) {
       rewriter.replaceOpWithNewOp<memref::LoadOp>(extractOp, xferOp.getSource(),
                                                   newIndices);
     } else {
       rewriter.replaceOpWithNewOp<tensor::ExtractOp>(
           extractOp, xferOp.getSource(), newIndices);
     }
     return success();
   }
+
+private:
+  bool allowMultipleUses;
 };
 
 /// Rewrite transfer_writes of vectors of size 1 (e.g., vector<1x1xf32>)
 /// to memref.store.
 class RewriteScalarWrite : public OpRewritePattern<vector::TransferWriteOp> {
   using OpRewritePattern::OpRewritePattern;
 
   LogicalResult matchAndRewrite(vector::TransferWriteOp xferOp,
void mlir::vector::transferOpflowOpt(RewriterBase &rewriter,
   rootOp->walk([&](vector::TransferWriteOp write) {
     if (write.getShapedType().isa<MemRefType>())
       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)