[mlir][Vector] Prevent vector-to-scalar xfer patterns from triggering on sub-vectors

[dcaballe]

Summary

Patterns that convert extract(transfer_read) into a scalar load where
incorrectly triggering for cases where a sub-vector instead of a scalar
was extracted.

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

RewriteScalarExtractOfTransferReadBase(MLIRContext *context,
       : Base::OpRewritePattern(context, benefit),
         allowMultipleUses(allowMultipleUses) {}
 
   LogicalResult match(VectorExtractOp extractOp) const override {
     auto xferOp =
         extractOp.getVector().template getDefiningOp<vector::TransferReadOp>();
     if (!xferOp)
       return failure();
+    // Check that we are extracting a scalar and not a sub-vector.
+    if (isa<VectorType>(extractOp.getResult().getType()))
+      return failure();
     // 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());

 class RewriteScalarExtractElementOfTransferRead
     : public RewriteScalarExtractOfTransferReadBase<vector::ExtractElementOp> {
   using RewriteScalarExtractOfTransferReadBase::
       RewriteScalarExtractOfTransferReadBase;
 
   void rewrite(vector::ExtractElementOp extractOp,
                PatternRewriter &rewriter) const override {
     // Construct scalar load.
+    auto loc = extractOp.getLoc();
     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,
+          rewriter, loc, 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(),
+            rewriter.create<arith::ConstantIndexOp>(loc,
                                                     *getConstantIntValue(ofr));
       }
     }
     if (isa<MemRefType>(xferOp.getSource().getType())) {
       rewriter.replaceOpWithNewOp<memref::LoadOp>(extractOp, xferOp.getSource(),
                                                   newIndices);
     } else {
       rewriter.replaceOpWithNewOp<tensor::ExtractOp>(

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

 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
 }
 
+// -----
+
+// Check that patterns don't trigger for an sub-vector (not scalar) extraction.
+// CHECK-LABEL: func @subvector_extract(
+//  CHECK-SAME:   %[[m:.*]]: memref<?x?xf32>, %[[idx:.*]]: index
+//   CHECK-NOT:   memref.load
+//       CHECK:   %[[r:.*]] = vector.transfer_read %[[m]][%[[idx]], %[[idx]]]
+//       CHECK:   %[[e0:.*]] = vector.extract %[[r]][0]
+//       CHECK:   return %[[e0]]
+
+func.func @subvector_extract(%m: memref<?x?xf32>, %idx: index) -> vector<16xf32> {
+  %cst = arith.constant 0.0 : f32
+  %0 = vector.transfer_read %m[%idx, %idx], %cst {in_bounds = [true, true]} : memref<?x?xf32>, vector<8x16xf32>
+  %1 = vector.extract %0[0] : vector<8x16xf32>
+  return %1 : vector<16xf32>
+}
+