diff --git a/mlir/lib/Conversion/VectorToLLVM/ConvertVectorToLLVM.cpp b/mlir/lib/Conversion/VectorToLLVM/ConvertVectorToLLVM.cpp index b66b55ae8d57f..19cc914efae00 100644 --- a/mlir/lib/Conversion/VectorToLLVM/ConvertVectorToLLVM.cpp +++ b/mlir/lib/Conversion/VectorToLLVM/ConvertVectorToLLVM.cpp @@ -1734,6 +1734,45 @@ struct VectorSplatNdOpLowering : public ConvertOpToLLVMPattern { } }; +/// Conversion pattern for a `vector.interleave`. +/// This supports fixed-sized vectors and scalable vectors. +struct VectorInterleaveOpLowering + : public ConvertOpToLLVMPattern { + using ConvertOpToLLVMPattern::ConvertOpToLLVMPattern; + + LogicalResult + matchAndRewrite(vector::InterleaveOp interleaveOp, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + VectorType resultType = interleaveOp.getResultVectorType(); + // n-D interleaves should have been lowered already. + if (resultType.getRank() != 1) + return rewriter.notifyMatchFailure(interleaveOp, + "InterleaveOp not rank 1"); + // If the result is rank 1, then this directly maps to LLVM. + if (resultType.isScalable()) { + rewriter.replaceOpWithNewOp( + interleaveOp, typeConverter->convertType(resultType), + adaptor.getLhs(), adaptor.getRhs()); + return success(); + } + // Lower fixed-size interleaves to a shufflevector. While the + // vector.interleave2 intrinsic supports fixed and scalable vectors, the + // langref still recommends fixed-vectors use shufflevector, see: + // https://llvm.org/docs/LangRef.html#id876. + int64_t resultVectorSize = resultType.getNumElements(); + SmallVector interleaveShuffleMask; + interleaveShuffleMask.reserve(resultVectorSize); + for (int i = 0, end = resultVectorSize / 2; i < end; ++i) { + interleaveShuffleMask.push_back(i); + interleaveShuffleMask.push_back((resultVectorSize / 2) + i); + } + rewriter.replaceOpWithNewOp( + interleaveOp, adaptor.getLhs(), adaptor.getRhs(), + interleaveShuffleMask); + return success(); + } +}; + } // namespace /// Populate the given list with patterns that convert from Vector to LLVM. @@ -1758,7 +1797,8 @@ void mlir::populateVectorToLLVMConversionPatterns( VectorExpandLoadOpConversion, VectorCompressStoreOpConversion, VectorSplatOpLowering, VectorSplatNdOpLowering, VectorScalableInsertOpLowering, VectorScalableExtractOpLowering, - MaskedReductionOpConversion>(converter); + MaskedReductionOpConversion, VectorInterleaveOpLowering>( + converter); // Transfer ops with rank > 1 are handled by VectorToSCF. populateVectorTransferLoweringPatterns(patterns, /*maxTransferRank=*/1); } diff --git a/mlir/test/Conversion/VectorToLLVM/vector-to-llvm.mlir b/mlir/test/Conversion/VectorToLLVM/vector-to-llvm.mlir index 1c13b16dfd9af..a46f2e101f3c3 100644 --- a/mlir/test/Conversion/VectorToLLVM/vector-to-llvm.mlir +++ b/mlir/test/Conversion/VectorToLLVM/vector-to-llvm.mlir @@ -2460,3 +2460,40 @@ func.func @make_fixed_vector_of_scalable_vector(%f : f64) -> vector<3x[2]xf64> %res = vector.broadcast %f : f64 to vector<3x[2]xf64> return %res : vector<3x[2]xf64> } + +// ----- + +// CHECK-LABEL: @vector_interleave_0d +// CHECK-SAME: %[[LHS:.*]]: vector, %[[RHS:.*]]: vector) +func.func @vector_interleave_0d(%a: vector, %b: vector) -> vector<2xi8> { + // CHECK: %[[LHS_RANK1:.*]] = builtin.unrealized_conversion_cast %[[LHS]] : vector to vector<1xi8> + // CHECK: %[[RHS_RANK1:.*]] = builtin.unrealized_conversion_cast %[[RHS]] : vector to vector<1xi8> + // CHECK: %[[ZIP:.*]] = llvm.shufflevector %[[LHS_RANK1]], %[[RHS_RANK1]] [0, 1] : vector<1xi8> + // CHECK: return %[[ZIP]] + %0 = vector.interleave %a, %b : vector + return %0 : vector<2xi8> +} + +// ----- + +// CHECK-LABEL: @vector_interleave_1d +// CHECK-SAME: %[[LHS:.*]]: vector<8xf32>, %[[RHS:.*]]: vector<8xf32>) +func.func @vector_interleave_1d(%a: vector<8xf32>, %b: vector<8xf32>) -> vector<16xf32> +{ + // CHECK: %[[ZIP:.*]] = llvm.shufflevector %[[LHS]], %[[RHS]] [0, 8, 1, 9, 2, 10, 3, 11, 4, 12, 5, 13, 6, 14, 7, 15] : vector<8xf32> + // CHECK: return %[[ZIP]] + %0 = vector.interleave %a, %b : vector<8xf32> + return %0 : vector<16xf32> +} + +// ----- + +// CHECK-LABEL: @vector_interleave_1d_scalable +// CHECK-SAME: %[[LHS:.*]]: vector<[4]xi32>, %[[RHS:.*]]: vector<[4]xi32>) +func.func @vector_interleave_1d_scalable(%a: vector<[4]xi32>, %b: vector<[4]xi32>) -> vector<[8]xi32> +{ + // CHECK: %[[ZIP:.*]] = "llvm.intr.experimental.vector.interleave2"(%[[LHS]], %[[RHS]]) : (vector<[4]xi32>, vector<[4]xi32>) -> vector<[8]xi32> + // CHECK: return %[[ZIP]] + %0 = vector.interleave %a, %b : vector<[4]xi32> + return %0 : vector<[8]xi32> +}