[mlir] [VectorOps] Lowering of vector.extract_slices to LLVM IR

[aartbik]

This was supposed to be an update.

[merge_guards_bot]

Unit tests: pass. 61909 tests passed, 0 failed and 782 were skipped.

clang-tidy: unknown.

clang-format: pass.

Build artifacts: diff.json, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

mlir/lib/Conversion/VectorToLLVM/ConvertVectorToLLVM.cpp

 #include "mlir/IR/MLIRContext.h"
 #include "mlir/IR/Module.h"
 #include "mlir/IR/Operation.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/IR/StandardTypes.h"
 #include "mlir/IR/Types.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Pass/PassManager.h"
+#include "mlir/Support/Functional.h"
+#include "mlir/Support/MathExtras.h"
 #include "mlir/Transforms/DialectConversion.h"
 #include "mlir/Transforms/Passes.h"
 
 #include "llvm/IR/DerivedTypes.h"
 #include "llvm/IR/Module.h"
 #include "llvm/IR/Type.h"
 #include "llvm/Support/Allocator.h"
 #include "llvm/Support/ErrorHandling.h"
private:
   Operation *getPrintComma(Operation *op) const {
     return getPrint(op, lowering.getDialect(), "print_comma", {});
   }
   Operation *getPrintNewline(Operation *op) const {
     return getPrint(op, lowering.getDialect(), "print_newline", {});
   }
 };
 
+// Consume tuple elements extracted by tuple get operations.
+// Note that this is a conversion pattern, since we must also
+// rewrite tuples introduced by other rules in the lowering.
+//
+// TODO(ajcbik): express this through general folding?
+//
+class VectorTupleGetOpConversion : public LLVMOpLowering {
+public:
+  explicit VectorTupleGetOpConversion(MLIRContext *context,
+                                      LLVMTypeConverter &typeConverter)
+      : LLVMOpLowering(vector::TupleGetOp::getOperationName(), context,
+                       typeConverter) {}
+
+  PatternMatchResult
+  matchAndRewrite(Operation *op, ArrayRef<Value> operands,
+                  ConversionPatternRewriter &rewriter) const override {
+    auto tupleGetOp = cast<vector::TupleGetOp>(op);
+    // Rewrite:
+    //    %t = vector.tuple .., %e_i, ..
+    //    %x = vector.tuple_get %t, i
+    // into:
+    //    %t = vector.tuple .., %e_i, ..  // one less use
+    //    %x = %e_i
+    if (auto tupleOp = dyn_cast_or_null<TupleOp>(operands[0].getDefiningOp())) {
+      rewriter.replaceOp(tupleGetOp,
+                         tupleOp.getOperands()[tupleGetOp.getIndex()]);
+      return matchSuccess();
+    }
+    return matchFailure();
+  }
+};
+
+// The removal of a vector tuple will fail until all canonicalizations
+// occured that remove its uses (viz. all individual elements are
+// "consumed" somewhere), at which point it will apply.
+//
+// NOTE: not foolproof, since newly introduced vector tuples have
+// an empty use list; such tuples need to be legalized somehow anyway,
+// and unconditionaly erasing is the right way (even though this may
+// assert at the final rewriting if use removal rules are missing).
+//
+// TODO(ajcbik): avoid all this by simply relying on DCE
+//
+class VectorTupleOpConversion : public OpRewritePattern<TupleOp> {
+public:
+  using OpRewritePattern<TupleOp>::OpRewritePattern;
+
+  PatternMatchResult matchAndRewrite(TupleOp op,
+                                     PatternRewriter &rewriter) const override {
+    if (!op.use_empty())
+      return matchFailure();
+    rewriter.eraseOp(op);
+    return matchSuccess();
+  }
+};
+
 /// Progressive lowering of StridedSliceOp to either:
 ///   1. extractelement + insertelement for the 1-D case
 ///   2. extract + optional strided_slice + insert for the n-D case.
 class VectorStridedSliceOpConversion : public OpRewritePattern<StridedSliceOp> {
 public:
   using OpRewritePattern<StridedSliceOp>::OpRewritePattern;
 
   PatternMatchResult matchAndRewrite(StridedSliceOp op,
for (int64_t off = offset, e = offset + size * stride, idx = 0; off < e;
       }
       res = insertOne(rewriter, loc, extracted, res, idx);
     }
     rewriter.replaceOp(op, {res});
     return matchSuccess();
   }
 };
 
+/// Progressive lowering of ExtractSlicesOp to tuple of StridedSliceOp.
+/// One:
+///   %x = vector.extract_slices
+/// is replaced by:
+///   %a = vector.strided_slice
+///   %b = vector.strided_slice
+///   ..
+///   %x = vector.tuple %a, %b, ..
+class VectorExtractSlicesOpConversion
+    : public OpRewritePattern<ExtractSlicesOp> {
+public:
+  using OpRewritePattern<ExtractSlicesOp>::OpRewritePattern;
+
+  // TODO(ajcbik): refactor slice utilities out into VectorUtils.h
+  PatternMatchResult matchAndRewrite(ExtractSlicesOp op,
+                                     PatternRewriter &rewriter) const override {
+    auto loc = op.getLoc();
+
+    VectorType vectorType = op.getSourceVectorType();
+    int64_t rank = vectorType.getRank();
+    auto shape = vectorType.getShape();
+
+    SmallVector<int64_t, 4> sizes;
+    op.getSizes(sizes);
+    SmallVector<int64_t, 4> strides;
+    op.getStrides(strides); // all-ones at the moment
+
+    // Compute the number of slices in each dimension.
+    SmallVector<int64_t, 4> dimSliceCounts(rank);
+    for (int64_t r = 0; r < rank; ++r)
+      dimSliceCounts[r] = ceilDiv(shape[r], sizes[r]);
+
+    // Compute the strides between slices in each dimension.
+    SmallVector<int64_t, 4> sliceStrides(rank);
+    sliceStrides[rank - 1] = 1;
+    for (int64_t r = rank - 2; r >= 0; --r)
+      sliceStrides[r] = sliceStrides[r + 1] * dimSliceCounts[r + 1];
+
+    // For each element in the tuple, generate the proper strided slice.
+    TupleType tupleType = op.getResultTupleType();
+    int64_t tupleSize = tupleType.size();
+    SmallVector<Value, 4> tupleValues(tupleSize);
+    for (int64_t i = 0; i < tupleSize; ++i) {
+      // Compute the vector offsets by de-linearizing the index.
+      SmallVector<int64_t, 4> vectorOffsets(rank);
+      int64_t linearIndex = i;
+      for (int64_t r = 0; r < rank; ++r) {
+        vectorOffsets[r] = linearIndex / sliceStrides[r];
+        linearIndex %= sliceStrides[r];
+      }
+      // Convert from unrolled vector-space offsets to element-space offsets.
+      auto elementOffsets = mlir::functional::zipMap(
+          [](int64_t v1, int64_t v2) { return v1 * v2; }, vectorOffsets, sizes);
+      // Compute the size of each slice.
+      SmallVector<int64_t, 4> sliceSizes(rank);
+      for (int64_t r = 0; r < rank; ++r) {
+        sliceSizes[r] = std::min(sizes[r], shape[r] - elementOffsets[r]);
+      }
+      // Insert in tuple.
+      tupleValues[i] = rewriter.create<StridedSliceOp>(
+          loc, op.vector(), elementOffsets, sliceSizes, strides);
+    }
+
+    Value tuple = rewriter.create<vector::TupleOp>(loc, tupleType, tupleValues);
+    rewriter.replaceOp(op, tuple);
+    return matchSuccess();
+  }
+};
+
 } // namespace
 
 /// Populate the given list with patterns that convert from Vector to LLVM.
 void mlir::populateVectorToLLVMConversionPatterns(
     LLVMTypeConverter &converter, OwningRewritePatternList &patterns) {
   MLIRContext *ctx = converter.getDialect()->getContext();
-  patterns.insert<VectorInsertStridedSliceOpDifferentRankRewritePattern,
+  patterns.insert<VectorExtractSlicesOpConversion,
+                  VectorInsertStridedSliceOpDifferentRankRewritePattern,
                   VectorInsertStridedSliceOpSameRankRewritePattern,
-                  VectorStridedSliceOpConversion>(ctx);
+                  VectorStridedSliceOpConversion, VectorTupleOpConversion>(ctx);
   patterns.insert<VectorBroadcastOpConversion, VectorShuffleOpConversion,
                   VectorExtractElementOpConversion, VectorExtractOpConversion,
                   VectorInsertElementOpConversion, VectorInsertOpConversion,
                   VectorOuterProductOpConversion, VectorTypeCastOpConversion,
-                  VectorPrintOpConversion>(ctx, converter);
+                  VectorPrintOpConversion, VectorTupleGetOpConversion>(
+      ctx, converter);
 }
 
 namespace {
 struct LowerVectorToLLVMPass : public ModulePass<LowerVectorToLLVMPass> {
   void runOnModule() override;
 };
 } // namespace
 

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

 //       CHECK:    llvm.call @print_comma() : () -> ()
 //       CHECK:    %[[x8:.*]] = llvm.mlir.constant(1 : index) : !llvm.i64
 //       CHECK:    %[[x9:.*]] = llvm.extractelement %[[x5]][%[[x8]] : !llvm.i64] : !llvm<"<2 x float>">
 //       CHECK:    llvm.call @print_f32(%[[x9]]) : (!llvm.float) -> ()
 //       CHECK:    llvm.call @print_close() : () -> ()
 //       CHECK:    llvm.call @print_close() : () -> ()
 //       CHECK:    llvm.call @print_newline() : () -> ()
 
-
 func @strided_slice(%arg0: vector<4xf32>, %arg1: vector<4x8xf32>, %arg2: vector<4x8x16xf32>) {
 // CHECK-LABEL: llvm.func @strided_slice(
   %0 = vector.strided_slice %arg0 {offsets = [2], sizes = [2], strides = [1]} : vector<4xf32> to vector<2xf32>
 //       CHECK:    llvm.mlir.constant(0.000000e+00 : f32) : !llvm.float
 //       CHECK:    llvm.mlir.constant(dense<0.000000e+00> : vector<2xf32>) : !llvm<"<2 x float>">
 //       CHECK:    llvm.mlir.constant(2 : index) : !llvm.i64
 //       CHECK:    llvm.extractelement %{{.*}}[%{{.*}} : !llvm.i64] : !llvm<"<4 x float>">
 //       CHECK:    llvm.mlir.constant(0 : index) : !llvm.i64

 //  CHECK-NEXT:    llvm.extractelement {{.*}}[{{.*}} : !llvm.i64] : !llvm<"<2 x float>">
 //  CHECK-NEXT:    llvm.mlir.constant(3 : index) : !llvm.i64
 //  CHECK-NEXT:    llvm.insertelement {{.*}}, {{.*}}[{{.*}} : !llvm.i64] : !llvm<"<4 x float>">
 //  CHECK-NEXT:    llvm.insertvalue {{.*}}, {{.*}}[3] : !llvm<"[4 x <4 x float>]">
 
   return
 }
 
+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: extract_strides(%arg0: !llvm<"[3 x <3 x float>]">)
+//      CHECK: %[[s0:.*]] = llvm.mlir.constant(dense<0.000000e+00> : vector<1x1xf32>) : !llvm<"[1 x <1 x float>]">
+//      CHECK: %[[s1:.*]] = llvm.extractvalue %arg0[2] : !llvm<"[3 x <3 x float>]">
+//      CHECK: %[[s3:.*]] = llvm.mlir.constant(dense<0.000000e+00> : vector<1xf32>) : !llvm<"<1 x float>">
+//      CHECK: %[[s4:.*]] = llvm.mlir.constant(2 : index) : !llvm.i64
+//      CHECK: %[[s5:.*]] = llvm.extractelement %[[s1]][%[[s4]] : !llvm.i64] : !llvm<"<3 x float>">
+//      CHECK: %[[s6:.*]] = llvm.mlir.constant(0 : index) : !llvm.i64
+//      CHECK: %[[s7:.*]] = llvm.insertelement %[[s5]], %[[s3]][%[[s6]] : !llvm.i64] : !llvm<"<1 x float>">
+//      CHECK: %[[s8:.*]] = llvm.insertvalue %[[s7]], %[[s0]][0] : !llvm<"[1 x <1 x float>]">
+//      CHECK: llvm.return %[[s8]] : !llvm<"[1 x <1 x float>]">