Drop transfer_read inner most unit dimensions

[asaadaldien]

Summary

Add a pattern to take a rank-reducing subview and drop inner most
contiguous unit dim.
This is useful when lowering vector to backends with 1d vector types.

[asaadaldien]

add comment

[ThomasRaoux]

mlir/include/mlir/Dialect/Vector/VectorTransforms.h

104

I wonder if we should have a more specific name. There is a different pattern called TransferOpReduceRank that does significantly different kind of transformation.

mlir/lib/Dialect/Vector/VectorTransforms.cpp

3826

That sounds too restrictive, we should allow minor identity. (and ideally add a test for it)

3829–3830

nit: srcType.getAffineMaps().size() <= 1?

3834

nit: use getVectorType()

3843–3848

Can we just iterate reverse through the strides and break whenever it is != 1?

3846

nit: should be break in the else case?

3866–3868

Should that be outside the loop?

[asaadaldien]

Comments..

mlir/lib/Dialect/Vector/VectorTransforms.cpp

3826

Don't we have a canonization that composes this map with the memref map?

3843–3848

we need to skip the first

[asaadaldien]

mlir/include/mlir/Dialect/Vector/VectorTransforms.h

104

about about reduceTransferOpsOperandRank ?

[ThomasRaoux]

mlir/lib/Dialect/Vector/VectorTransforms.cpp

3826

minor identity transfer read is the canonical form as far as I know.

3843–3848

Can we just iterate reverse? just a suggestion though, feel free to ignore.
something like:

for(int64_t stride : llvm::reverse(scrStrides) {
  if(stride != 1)
    break;
   dimsToDrop++;
}

3870–3876

nit: can those two be merged? I realize this requires to change a bit how map is handled, just wondering if this makes the code a bit simpler, feel free to ignore.

[asaadaldien]

mlir/include/mlir/Dialect/Vector/VectorTransforms.h

104

what about reduceTransferOpsOperandRank ?

[ThomasRaoux]

mlir/include/mlir/Dialect/Vector/VectorTransforms.h

104

If this is meant to only drop inner unit dimension I would include it in the name. What do you think? Or is the idea to collapse all the contiguous inner dimension? I wonder if we can make the goal a bit more explicit in the name.

[asaadaldien]

mlir/include/mlir/Dialect/Vector/VectorTransforms.h

104

I agree, collapseInnerMostContiguousDims is the goal of these patterns. I will go with this naming. Thanks!

[asaadaldien]

Use isMinorIdentitiy

[ThomasRaoux]

mlir/test/Dialect/Vector/vector-transfer-collapse-inner-most-dims.mlir

8

Can you make the dest vector<1x8x1xf32> to show how this works when vector rank < memref rank

[asaadaldien]

Drop leading dim in test case

[asaadaldien]

update test pass name

[ThomasRaoux]

LGTM

[asaadaldien]

rebase...

mlir/include/mlir/Dialect/Vector/VectorTransforms.h

 ///                 \            |            /
 ///      <-------------------- join ------------------------->
 ///
 /// Other local patterns then kick in iteratively (including DCE) and compose
 /// to combine the ExtractStridedSlice/InsertStridedSlice.
 void populateVectorUnrollPatterns(RewritePatternSet &patterns,
                                   const UnrollVectorOptions &options);
 
+/// Collect a set of patterns to reduce the rank of the operands of vector
+/// transfer ops to operate on the largest contigious vector.
+/// These patterns are useful when lowering to dialects with 1d vector type
+/// such as llvm and it will result fewer memory reads.
+void populateVectorTransferCollapseInnerMostContiguousDimsPatterns(
+    RewritePatternSet &patterns);
+
 /// Split a vector.transfer operation into an in-bounds (i.e., no out-of-bounds
 /// masking) fastpath and a slowpath.
 /// If `ifOp` is not null and the result is `success, the `ifOp` points to the
 /// newly created conditional upon function return.
 /// To accomodate for the fact that the original vector.transfer indexing may be
 /// arbitrary and the slow path indexes @[0...0] in the temporary buffer, the
 /// scf.if op returns a view and values of type index.
 /// At this time, only vector.transfer_read case is implemented.

mlir/lib/Dialect/Vector/VectorTransforms.cpp

for (int i = 0; i < outerDim; ++i) {
       result = rewriter.create<vector::InsertElementOp>(loc, reducedValue,
                                                         result, i);
     }
     rewriter.replaceOp(multiReductionOp, result);
     return success();
   }
 };
 
+// Drop inner most contiguous unit dimensions from transfer_read operand.
+class DropInnerMostUnitDims : public OpRewritePattern<vector::TransferReadOp> {
+  using OpRewritePattern<vector::TransferReadOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(vector::TransferReadOp readOp,
+                                PatternRewriter &rewriter) const override {
+    auto srcType = readOp.source().getType().cast<MemRefType>();
+    if (!srcType || !srcType.hasStaticShape())
+      return failure();
+
+    if (!readOp.permutation_map().isMinorIdentity())
+      return failure();
+
+    if (!srcType.getAffineMaps().empty() &&
+        srcType.getAffineMaps().size() != 1) {
+      return failure();
+    }
+
+    auto targetType = readOp.getVectorType();
+    if (targetType.getRank() <= 1)
+      return failure();
+
+    SmallVector<int64_t> srcStrides;
+    int64_t srcOffset;
+    if (failed(getStridesAndOffset(srcType, srcStrides, srcOffset)))
+      return failure();
+
+    size_t dimsToDrop = 0;
+    for (size_t i = 1; i < srcStrides.size(); ++i) {
+      int dim = srcType.getRank() - i - 1;
+      if (srcStrides[dim] == 1) {
+        dimsToDrop++;
+      } else {
+        break;
+      }
+    }
+    if (dimsToDrop == 0)
+      return failure();
+
+    auto resultTargetVecType =
+        VectorType::get(targetType.getShape().drop_back(dimsToDrop),
+                        targetType.getElementType());
+
+    MemRefType resultMemrefType;
+    if (!srcType.getAffineMaps().empty()) {
+      AffineMap map = srcType.getAffineMaps()[0];
+      int numResultDims = map.getNumDims() - dimsToDrop;
+      int numSymbols = map.getNumSymbols();
+      for (size_t i = 0; i < dimsToDrop; ++i) {
+        int dim = srcType.getRank() - i - 1;
+        map = map.replace(rewriter.getAffineDimExpr(dim),
+                          rewriter.getAffineConstantExpr(0), numResultDims,
+                          numSymbols);
+      }
+      resultMemrefType = MemRefType::get(
+          srcType.getShape().drop_back(dimsToDrop), srcType.getElementType(),
+          map, srcType.getMemorySpaceAsInt());
+    } else {
+      resultMemrefType = MemRefType::get(
+          srcType.getShape().drop_back(dimsToDrop), srcType.getElementType(),
+          {}, srcType.getMemorySpaceAsInt());
+    }
+
+    auto loc = readOp.getLoc();
+    SmallVector<int64_t> offsets(srcType.getRank(), 0);
+    SmallVector<int64_t> strides(srcType.getRank(), 1);
+    Value rankedReducedView = rewriter.create<memref::SubViewOp>(
+        loc, resultMemrefType, readOp.source(), offsets, srcType.getShape(),
+        strides);
+    Value result = rewriter.create<vector::TransferReadOp>(
+        loc, resultTargetVecType, rankedReducedView,
+        readOp.indices().drop_back(dimsToDrop));
+    rewriter.replaceOpWithNewOp<vector::ShapeCastOp>(readOp, targetType,
+                                                     result);
+
+    return success();
+  }
+};
+
 void mlir::vector::populateVectorMaskMaterializationPatterns(
     RewritePatternSet &patterns, bool enableIndexOptimizations) {
   patterns.add<VectorCreateMaskOpConversion,
                MaterializeTransferMask<vector::TransferReadOp>,
                MaterializeTransferMask<vector::TransferWriteOp>>(
       patterns.getContext(), enableIndexOptimizations);
 }
 

 }
 
 void mlir::vector::populateVectorUnrollPatterns(
     RewritePatternSet &patterns, const UnrollVectorOptions &options) {
   patterns.add<UnrollTransferReadPattern, UnrollTransferWritePattern,
                UnrollContractionPattern, UnrollElementwisePattern>(
       patterns.getContext(), options);
 }
+
+void mlir::vector::
+    populateVectorTransferCollapseInnerMostContiguousDimsPatterns(
+        RewritePatternSet &patterns) {
+  patterns.add<DropInnerMostUnitDims>(patterns.getContext());
+}

mlir/test/Dialect/Vector/vector-transfer-collapse-inner-most-dims.mlir

+// RUN: mlir-opt %s -test-vector-transfer-collapse-inner-most-dims -split-input-file | FileCheck %s
+
+#map1 = affine_map<(d0, d1, d2, d3)[s0] -> (d0 * 3072 + s0 + d1 * 8 + d2 + d3)>
+func @contiguous_inner_most_view(%in: memref<1x1x8x1xf32, #map1>) -> vector<1x8x1xf32>{
+  %c0 = constant 0 : index
+  %cst = constant 0.0 : f32
+  %0 = vector.transfer_read %in[%c0, %c0, %c0, %c0], %cst {in_bounds = [true, true, true]} : memref<1x1x8x1xf32, #map1>, vector<1x8x1xf32>
+  return %0 : vector<1x8x1xf32>
+}
+//  CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1, d2, d3)[s0] -> (d0 * 3072 + s0 + d1 * 8 + d2 + d3)>
+//  CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2)[s0] -> (d0 * 3072 + s0 + d1 * 8 + d2)>
+//      CHECK: func @contiguous_inner_most_view(%[[SRC:.+]]: memref<1x1x8x1xf32, #[[MAP0]]>
+//      CHECK:   %[[SRC_0:.+]] = memref.subview %[[SRC]]
+// CHECK-SAME:    memref<1x1x8x1xf32, #[[MAP0]]> to memref<1x1x8xf32, #[[MAP1]]>
+//      CHECK:   %[[VEC:.+]] = vector.transfer_read %[[SRC_0]]
+// CHECK-SAME:    memref<1x1x8xf32, #[[MAP1]]>, vector<1x8xf32>
+//      CHECK:   %[[RESULT:.+]] = vector.shape_cast %[[VEC]]
+//      CHECK:   return %[[RESULT]]
+// -----
+
+func @contiguous_inner_most_dim(%A: memref<16x1xf32>, %i:index, %j:index) -> (vector<8x1xf32>) {
+  %c0 = constant 0 : index
+  %f0 = constant 0.0 : f32
+  %1 = vector.transfer_read %A[%i, %j], %f0 : memref<16x1xf32>, vector<8x1xf32>
+  return %1 : vector<8x1xf32>
+}
+//      CHECK: func @contiguous_inner_most_dim(%[[SRC:.+]]: memref<16x1xf32>, %[[I:.+]]: index, %[[J:.+]]: index) -> vector<8x1xf32>
+//      CHECK:   %[[SRC_0:.+]] = memref.subview %[[SRC]]
+// CHECK-SAME:     memref<16x1xf32> to memref<16xf32>
+//      CHECK:   %[[V:.+]] = vector.transfer_read %[[SRC_0]]
+//      CHECK:   %[[RESULT]] = vector.shape_cast %[[V]] : vector<8xf32> to vector<8x1xf32>
+//      CHECK:   return %[[RESULT]]

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

Option<bool> useOuterReductions{
       llvm::cl::init(false)};
   void runOnFunction() override {
     RewritePatternSet patterns(&getContext());
     populateVectorMultiReductionLoweringPatterns(patterns, !useOuterReductions);
     (void)applyPatternsAndFoldGreedily(getFunction(), std::move(patterns));
   }
 };
 
+struct TestVectorTransferRankReductionPatterns
+    : public PassWrapper<TestVectorTransferRankReductionPatterns,
+                         FunctionPass> {
+  TestVectorTransferRankReductionPatterns() = default;
+  TestVectorTransferRankReductionPatterns(
+      const TestVectorTransferRankReductionPatterns &pass) {}
+
+  void getDependentDialects(DialectRegistry &registry) const override {
+    registry.insert<memref::MemRefDialect, AffineDialect>();
+  }
+
+  StringRef getArgument() const final {
+    return "test-vector-transfer-collapse-inner-most-dims";
+  }
+
+  StringRef getDescription() const final {
+    return "Test conversion patterns that reducedes the rank of the vector "
+           "transfer memory and vector operands.";
+  }
+
+  void runOnFunction() override {
+    RewritePatternSet patterns(&getContext());
+    populateVectorTransferCollapseInnerMostContiguousDimsPatterns(patterns);
+    (void)applyPatternsAndFoldGreedily(getFunction(), std::move(patterns));
+  }
+};
+
 } // end anonymous namespace
 
 namespace mlir {
 namespace test {
 void registerTestVectorConversions() {
   PassRegistration<TestVectorToVectorConversion>();
 
   PassRegistration<TestVectorContractionConversion>();
 
   PassRegistration<TestVectorUnrollingPatterns>();
 
   PassRegistration<TestVectorTransferUnrollingPatterns>();
 
   PassRegistration<TestVectorTransferFullPartialSplitPatterns>();
 
   PassRegistration<TestVectorDistributePatterns>();
 
   PassRegistration<TestVectorToLoopPatterns>();
 
   PassRegistration<TestVectorTransferOpt>();
 
   PassRegistration<TestVectorTransferLoweringPatterns>();
 
   PassRegistration<TestVectorMultiReductionLoweringPatterns>();
+
+  PassRegistration<TestVectorTransferRankReductionPatterns>();
 }
 } // namespace test
 } // namespace mlir