[mlir] [VectorOps] Improve lowering of extract_strided_slice (and friends like shape_cast)

[aartbik]

Summary

Using a shuffle for the last recursive step in progressive lowering not only
results in much more compact IR, but also more efficient code (since the
backend is no longer confused on subvector aliasing for longer vectors).

E.g. the following

%f = vector.shape_cast %v0: vector<1024xf32> to vector<32x32xf32>

yields much better x86-64 code that runs 3x faster than the original.

[bkramer]

lg

[aartbik]

trigger tests

[nicolasvasilache]

Great, thanks for fixing this perf bug !

mlir/lib/Conversion/VectorToLLVM/ConvertVectorToLLVM.cpp

Operation *getPrintComma(Operation *op) const {
     return getPrint(op, typeConverter.getDialect(), "print_comma", {});
   }
   Operation *getPrintNewline(Operation *op) const {
     return getPrint(op, typeConverter.getDialect(), "print_newline", {});
   }
 };
 
 /// Progressive lowering of ExtractStridedSliceOp to either:
-///   1. extractelement + insertelement for the 1-D case
-///   2. extract + optional strided_slice + insert for the n-D case.
-class VectorStridedSliceOpConversion
+///   1. express single offset extract as a direct shuffle.
+///   2. extract + lower rank strided_slice + insert for the n-D case.
+class VectorExtractStridedSliceOpConversion
     : public OpRewritePattern<ExtractStridedSliceOp> {
 public:
   using OpRewritePattern<ExtractStridedSliceOp>::OpRewritePattern;
 
   LogicalResult matchAndRewrite(ExtractStridedSliceOp op,
                                 PatternRewriter &rewriter) const override {
     auto dstType = op.getResult().getType().cast<VectorType>();
 
     assert(!op.offsets().getValue().empty() && "Unexpected empty offsets");
 
     int64_t offset =
         op.offsets().getValue().front().cast<IntegerAttr>().getInt();
     int64_t size = op.sizes().getValue().front().cast<IntegerAttr>().getInt();
     int64_t stride =
         op.strides().getValue().front().cast<IntegerAttr>().getInt();
 
     auto loc = op.getLoc();
     auto elemType = dstType.getElementType();
     assert(elemType.isSignlessIntOrIndexOrFloat());
+
+    // Single offset can be more efficiently shuffled.
+    if (op.offsets().getValue().size() == 1) {
+      SmallVector<int64_t, 4> offsets;
+      offsets.reserve(size);
+      for (int64_t off = offset, e = offset + size * stride; off < e;
+           off += stride)
+        offsets.push_back(off);
+      rewriter.replaceOpWithNewOp<ShuffleOp>(op, dstType, op.vector(),
+                                             op.vector(),
+                                             rewriter.getI64ArrayAttr(offsets));
+      return success();
+    }
+
+    // Extract/insert on a lower ranked extract strided slice op.
     Value zero = rewriter.create<ConstantOp>(loc, elemType,
                                              rewriter.getZeroAttr(elemType));
     Value res = rewriter.create<SplatOp>(loc, dstType, zero);
     for (int64_t off = offset, e = offset + size * stride, idx = 0; off < e;
          off += stride, ++idx) {
-      Value extracted = extractOne(rewriter, loc, op.vector(), off);
-      if (op.offsets().getValue().size() > 1) {
-        extracted = rewriter.create<ExtractStridedSliceOp>(
-            loc, extracted, getI64SubArray(op.offsets(), /* dropFront=*/1),
-            getI64SubArray(op.sizes(), /* dropFront=*/1),
-            getI64SubArray(op.strides(), /* dropFront=*/1));
-      }
+      Value one = extractOne(rewriter, loc, op.vector(), off);
+      Value extracted = rewriter.create<ExtractStridedSliceOp>(
+          loc, one, getI64SubArray(op.offsets(), /* dropFront=*/1),
       res = insertOne(rewriter, loc, extracted, res, idx);
     }
-    rewriter.replaceOp(op, {res});
+    rewriter.replaceOp(op, res);
     return success();
   }
   /// This pattern creates recursive ExtractStridedSliceOp, but the recursion is
   /// bounded as the rank is strictly decreasing.
   bool hasBoundedRewriteRecursion() const final { return true; }
 };
 
 } // namespace
 
 /// Populate the given list with patterns that convert from Vector to LLVM.
 void mlir::populateVectorToLLVMConversionPatterns(
     LLVMTypeConverter &converter, OwningRewritePatternList &patterns,
     bool reassociateFPReductions) {
   MLIRContext *ctx = converter.getDialect()->getContext();
   // clang-format off
   patterns.insert<VectorFMAOpNDRewritePattern,
                   VectorInsertStridedSliceOpDifferentRankRewritePattern,
                   VectorInsertStridedSliceOpSameRankRewritePattern,
-                  VectorStridedSliceOpConversion>(ctx);
+                  VectorExtractStridedSliceOpConversion>(ctx);
   patterns.insert<VectorReductionOpConversion>(
       ctx, converter, reassociateFPReductions);
   patterns
       .insert<VectorShuffleOpConversion,
               VectorExtractElementOpConversion,
               VectorExtractOpConversion,
               VectorFMAOp1DConversion,
               VectorInsertElementOpConversion,

mlir/test/Conversion/VectorToLLVM/vector-to-llvm.mlir

 //       CHECK:    llvm.call @print_close() : () -> ()
 //       CHECK:    llvm.call @print_close() : () -> ()
 //       CHECK:    llvm.call @print_newline() : () -> ()
 
 func @extract_strided_slice1(%arg0: vector<4xf32>) -> vector<2xf32> {
   %0 = vector.extract_strided_slice %arg0 {offsets = [2], sizes = [2], strides = [1]} : vector<4xf32> to vector<2xf32>
   return %0 : vector<2xf32>
 }
-// CHECK-LABEL: llvm.func @extract_strided_slice1
-//       CHECK:    llvm.mlir.constant(0.000000e+00 : f32) : !llvm.float
-//       CHECK:    llvm.mlir.constant(dense<0.000000e+00> : vector<2xf32>) : !llvm.vec<2 x float>
-//       CHECK:    llvm.mlir.constant(2 : index) : !llvm.i64
-//       CHECK:    llvm.extractelement %{{.*}}[%{{.*}} : !llvm.i64] : !llvm.vec<4 x float>
-//       CHECK:    llvm.mlir.constant(0 : index) : !llvm.i64
-//       CHECK:    llvm.insertelement %{{.*}}, %{{.*}}[%{{.*}} : !llvm.i64] : !llvm.vec<2 x float>
-//       CHECK:    llvm.mlir.constant(3 : index) : !llvm.i64
-//       CHECK:    llvm.extractelement %{{.*}}[%{{.*}} : !llvm.i64] : !llvm.vec<4 x float>
-//       CHECK:    llvm.mlir.constant(1 : index) : !llvm.i64
-//       CHECK:    llvm.insertelement %{{.*}}, %{{.*}}[%{{.*}} : !llvm.i64] : !llvm.vec<2 x float>
+// CHECK-LABEL: llvm.func @extract_strided_slice1(
+//  CHECK-SAME:    %[[A:.*]]: !llvm.vec<4 x float>)
+//       CHECK:    %[[T0:.*]] = llvm.shufflevector %[[A]], %[[A]] [2, 3] : !llvm.vec<4 x float>, !llvm.vec<4 x float>
+//       CHECK:    llvm.return %[[T0]] : !llvm.vec<2 x float>
 
 func @extract_strided_slice2(%arg0: vector<4x8xf32>) -> vector<2x8xf32> {
   %0 = vector.extract_strided_slice %arg0 {offsets = [2], sizes = [2], strides = [1]} : vector<4x8xf32> to vector<2x8xf32>
   return %0 : vector<2x8xf32>
 }
-// CHECK-LABEL: llvm.func @extract_strided_slice2
-//       CHECK:    llvm.mlir.constant(0.000000e+00 : f32) : !llvm.float
-//       CHECK:    llvm.mlir.constant(dense<0.000000e+00> : vector<2x8xf32>) : !llvm.array<2 x vec<8 x float>>
-//       CHECK:    llvm.extractvalue %{{.*}}[2] : !llvm.array<4 x vec<8 x float>>
-//       CHECK:    llvm.insertvalue %{{.*}}, %{{.*}}[0] : !llvm.array<2 x vec<8 x float>>
-//       CHECK:    llvm.extractvalue %{{.*}}[3] : !llvm.array<4 x vec<8 x float>>
-//       CHECK:    llvm.insertvalue %{{.*}}, %{{.*}}[1] : !llvm.array<2 x vec<8 x float>>
+// CHECK-LABEL: llvm.func @extract_strided_slice2(
+//  CHECK-SAME:    %[[A:.*]]: !llvm.array<4 x vec<8 x float>>)
+//       CHECK:    %[[T0:.*]] = llvm.mlir.undef : !llvm.array<2 x vec<8 x float>>
+//       CHECK:    %[[T1:.*]] = llvm.extractvalue %[[A]][2] : !llvm.array<4 x vec<8 x float>>
+//       CHECK:    %[[T2:.*]] = llvm.insertvalue %[[T1]], %[[T0]][0] : !llvm.array<2 x vec<8 x float>>
+//       CHECK:    %[[T3:.*]] = llvm.extractvalue %[[A]][3] : !llvm.array<4 x vec<8 x float>>
+//       CHECK:    %[[T4:.*]] = llvm.insertvalue %[[T3]], %[[T2]][1] : !llvm.array<2 x vec<8 x float>>
+//       CHECK:    llvm.return %[[T4]] : !llvm.array<2 x vec<8 x float>>
 
 func @extract_strided_slice3(%arg0: vector<4x8xf32>) -> vector<2x2xf32> {
   %0 = vector.extract_strided_slice %arg0 {offsets = [2, 2], sizes = [2, 2], strides = [1, 1]} : vector<4x8xf32> to vector<2x2xf32>
   return %0 : vector<2x2xf32>
 }
-// CHECK-LABEL: llvm.func @extract_strided_slice3
-//       CHECK:    llvm.mlir.constant(0.000000e+00 : f32) : !llvm.float
-//       CHECK:    llvm.mlir.constant(dense<0.000000e+00> : vector<2x2xf32>) : !llvm.array<2 x vec<2 x float>>
-//
-// Subvector vector<8xf32> @2
-//       CHECK:    llvm.extractvalue {{.*}}[2] : !llvm.array<4 x vec<8 x float>>
-//       CHECK:    llvm.mlir.constant(0.000000e+00 : f32) : !llvm.float
-//       CHECK:    llvm.mlir.constant(dense<0.000000e+00> : vector<2xf32>) : !llvm.vec<2 x float>
-//       CHECK:    llvm.mlir.constant(2 : index) : !llvm.i64
-//       CHECK:    llvm.extractelement {{.*}}[{{.*}} : !llvm.i64] : !llvm.vec<8 x float>
-//       CHECK:    llvm.mlir.constant(0 : index) : !llvm.i64
-//       CHECK:    llvm.insertelement {{.*}}, {{.*}}[{{.*}} : !llvm.i64] : !llvm.vec<2 x float>
-//       CHECK:    llvm.mlir.constant(3 : index) : !llvm.i64
-//       CHECK:    llvm.extractelement {{.*}}[{{.*}} : !llvm.i64] : !llvm.vec<8 x float>
-//       CHECK:    llvm.mlir.constant(1 : index) : !llvm.i64
-//       CHECK:    llvm.insertelement {{.*}}, {{.*}}[{{.*}} : !llvm.i64] : !llvm.vec<2 x float>
-//       CHECK:    llvm.insertvalue {{.*}}, {{.*}}[0] : !llvm.array<2 x vec<2 x float>>
-//
-// Subvector vector<8xf32> @3
-//       CHECK:    llvm.extractvalue {{.*}}[3] : !llvm.array<4 x vec<8 x float>>
-//       CHECK:    llvm.mlir.constant(0.000000e+00 : f32) : !llvm.float
-//       CHECK:    llvm.mlir.constant(dense<0.000000e+00> : vector<2xf32>) : !llvm.vec<2 x float>
-//       CHECK:    llvm.mlir.constant(2 : index) : !llvm.i64
-//       CHECK:    llvm.extractelement {{.*}}[{{.*}} : !llvm.i64] : !llvm.vec<8 x float>
-//       CHECK:    llvm.mlir.constant(0 : index) : !llvm.i64
-//       CHECK:    llvm.insertelement {{.*}}, {{.*}}[{{.*}} : !llvm.i64] : !llvm.vec<2 x float>
-//       CHECK:    llvm.mlir.constant(3 : index) : !llvm.i64
-//       CHECK:    llvm.extractelement {{.*}}[{{.*}} : !llvm.i64] : !llvm.vec<8 x float>
-//       CHECK:    llvm.mlir.constant(1 : index) : !llvm.i64
-//       CHECK:    llvm.insertelement {{.*}}, {{.*}}[{{.*}} : !llvm.i64] : !llvm.vec<2 x float>
-//       CHECK:    llvm.insertvalue {{.*}}, {{.*}}[1] : !llvm.array<2 x vec<2 x float>>
+// CHECK-LABEL: llvm.func @extract_strided_slice3(
+//  CHECK-SAME:    %[[A:.*]]: !llvm.array<4 x vec<8 x float>>)
+//       CHECK:    %[[T1:.*]] = llvm.mlir.constant(dense<0.000000e+00> : vector<2x2xf32>) : !llvm.array<2 x vec<2 x float>>
+//       CHECK:    %[[T2:.*]] = llvm.extractvalue %[[A]][2] : !llvm.array<4 x vec<8 x float>>
+//       CHECK:    %[[T3:.*]] = llvm.shufflevector %[[T2]], %[[T2]] [2, 3] : !llvm.vec<8 x float>, !llvm.vec<8 x float>
+//       CHECK:    %[[T4:.*]] = llvm.insertvalue %[[T3]], %[[T1]][0] : !llvm.array<2 x vec<2 x float>>
+//       CHECK:    %[[T5:.*]] = llvm.extractvalue %[[A]][3] : !llvm.array<4 x vec<8 x float>>
+//       CHECK:    %[[T6:.*]] = llvm.shufflevector %[[T5]], %[[T5]] [2, 3] : !llvm.vec<8 x float>, !llvm.vec<8 x float>
+//       CHECK:    %[[T7:.*]] = llvm.insertvalue %[[T6]], %[[T4]][1] : !llvm.array<2 x vec<2 x float>>
+//       CHECK:    llvm.return %[[T7]] : !llvm.array<2 x vec<2 x float>>
 
 func @insert_strided_slice1(%b: vector<4x4xf32>, %c: vector<4x4x4xf32>) -> vector<4x4x4xf32> {
   %0 = vector.insert_strided_slice %b, %c {offsets = [2, 0, 0], strides = [1, 1]} : vector<4x4xf32> into vector<4x4x4xf32>
   return %0 : vector<4x4x4xf32>
 }
 // CHECK-LABEL: llvm.func @insert_strided_slice1
 //       CHECK:    llvm.extractvalue {{.*}}[2] : !llvm.array<4 x array<4 x vec<4 x float>>>
 //  CHECK-NEXT:    llvm.insertvalue {{.*}}, {{.*}}[2] : !llvm.array<4 x array<4 x vec<4 x float>>>

 func @extract_strides(%arg0: vector<3x3xf32>) -> vector<1x1xf32> {
   %0 = vector.extract_slices %arg0, [2, 2], [1, 1]
     : vector<3x3xf32> into tuple<vector<2x2xf32>, vector<2x1xf32>, vector<1x2xf32>, vector<1x1xf32>>
   %1 = vector.tuple_get %0, 3 : tuple<vector<2x2xf32>, vector<2x1xf32>, vector<1x2xf32>, vector<1x1xf32>>
   return %1 : vector<1x1xf32>
 }
 // CHECK-LABEL: llvm.func @extract_strides(
 // CHECK-SAME: %[[A:.*]]: !llvm.array<3 x vec<3 x float>>)
-//      CHECK: %[[s0:.*]] = llvm.mlir.constant(dense<0.000000e+00> : vector<1x1xf32>) : !llvm.array<1 x vec<1 x float>>
-//      CHECK: %[[s1:.*]] = llvm.extractvalue %[[A]][2] : !llvm.array<3 x vec<3 x float>>
-//      CHECK: %[[s3:.*]] = llvm.mlir.constant(dense<0.000000e+00> : vector<1xf32>) : !llvm.vec<1 x float>
-//      CHECK: %[[s4:.*]] = llvm.mlir.constant(2 : index) : !llvm.i64
-//      CHECK: %[[s5:.*]] = llvm.extractelement %[[s1]][%[[s4]] : !llvm.i64] : !llvm.vec<3 x float>
-//      CHECK: %[[s6:.*]] = llvm.mlir.constant(0 : index) : !llvm.i64
-//      CHECK: %[[s7:.*]] = llvm.insertelement %[[s5]], %[[s3]][%[[s6]] : !llvm.i64] : !llvm.vec<1 x float>
-//      CHECK: %[[s8:.*]] = llvm.insertvalue %[[s7]], %[[s0]][0] : !llvm.array<1 x vec<1 x float>>
-//      CHECK: llvm.return %[[s8]] : !llvm.array<1 x vec<1 x float>>
+//      CHECK: %[[T1:.*]] = llvm.mlir.constant(dense<0.000000e+00> : vector<1x1xf32>) : !llvm.array<1 x vec<1 x float>>
+//      CHECK: %[[T2:.*]] = llvm.extractvalue %[[A]][2] : !llvm.array<3 x vec<3 x float>>
+//      CHECK: %[[T3:.*]] = llvm.shufflevector %[[T2]], %[[T2]] [2] : !llvm.vec<3 x float>, !llvm.vec<3 x float>
+//      CHECK: %[[T4:.*]] = llvm.insertvalue %[[T3]], %[[T1]][0] : !llvm.array<1 x vec<1 x float>>
+//      CHECK: llvm.return %[[T4]] : !llvm.array<1 x vec<1 x float>>
 
 // CHECK-LABEL: llvm.func @vector_fma(
 //  CHECK-SAME: %[[A:.*]]: !llvm.vec<8 x float>, %[[B:.*]]: !llvm.array<2 x vec<4 x float>>)
 //  CHECK-SAME: -> !llvm.struct<(vec<8 x float>, array<2 x vec<4 x float>>)> {
 func @vector_fma(%a: vector<8xf32>, %b: vector<2x4xf32>) -> (vector<8xf32>, vector<2x4xf32>) {
   //         CHECK: "llvm.intr.fmuladd"(%[[A]], %[[A]], %[[A]]) :
   //    CHECK-SAME:   (!llvm.vec<8 x float>, !llvm.vec<8 x float>, !llvm.vec<8 x float>) -> !llvm.vec<8 x float>
   %0 = vector.fma %a, %a, %a : vector<8xf32>