[Mlir-commits] [mlir] bc946f5 - [mlir][vector] Add 1D vector.deinterleave lowering (#93042)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Thu May 30 01:42:40 PDT 2024
Author: Mubashar Ahmad
Date: 2024-05-30T09:42:35+01:00
New Revision: bc946f52870a51b49a0c7c3c508cef1905a11bc0
URL: https://github.com/llvm/llvm-project/commit/bc946f52870a51b49a0c7c3c508cef1905a11bc0
DIFF: https://github.com/llvm/llvm-project/commit/bc946f52870a51b49a0c7c3c508cef1905a11bc0.diff
LOG: [mlir][vector] Add 1D vector.deinterleave lowering (#93042)
This patch implements the lowering of vector.deinterleave
for 1D vectors.
For fixed vector types, the operation is lowered to two
llvm shufflevector operations. One for even indexed
elements and the other for odd indexed elements. A poison
operation is used to satisfy the parameters of the
shufflevector parameters.
For scalable vectors, the llvm vector.deinterleave2
intrinsic is used for lowering. As such the results
found by extraction and used to form the result
struct for the intrinsic.
Added:
mlir/test/Integration/Dialect/Vector/CPU/ArmSVE/Emulated/test-scalable-deinterleave.mlir
mlir/test/Integration/Dialect/Vector/CPU/test-deinterleave.mlir
Modified:
mlir/lib/Conversion/VectorToLLVM/ConvertVectorToLLVM.cpp
mlir/test/Conversion/VectorToLLVM/vector-to-llvm.mlir
Removed:
################################################################################
diff --git a/mlir/lib/Conversion/VectorToLLVM/ConvertVectorToLLVM.cpp b/mlir/lib/Conversion/VectorToLLVM/ConvertVectorToLLVM.cpp
index fe6bcc1c8b667..8671c8412f4ec 100644
--- a/mlir/lib/Conversion/VectorToLLVM/ConvertVectorToLLVM.cpp
+++ b/mlir/lib/Conversion/VectorToLLVM/ConvertVectorToLLVM.cpp
@@ -1761,6 +1761,70 @@ struct VectorInterleaveOpLowering
}
};
+/// Conversion pattern for a `vector.deinterleave`.
+/// This supports fixed-sized vectors and scalable vectors.
+struct VectorDeinterleaveOpLowering
+ : public ConvertOpToLLVMPattern<vector::DeinterleaveOp> {
+ using ConvertOpToLLVMPattern::ConvertOpToLLVMPattern;
+
+ LogicalResult
+ matchAndRewrite(vector::DeinterleaveOp deinterleaveOp, OpAdaptor adaptor,
+ ConversionPatternRewriter &rewriter) const override {
+ VectorType resultType = deinterleaveOp.getResultVectorType();
+ VectorType sourceType = deinterleaveOp.getSourceVectorType();
+ auto loc = deinterleaveOp.getLoc();
+
+ // Note: n-D deinterleave operations should be lowered to the 1-D before
+ // converting to LLVM.
+ if (resultType.getRank() != 1)
+ return rewriter.notifyMatchFailure(deinterleaveOp,
+ "DeinterleaveOp not rank 1");
+
+ if (resultType.isScalable()) {
+ auto llvmTypeConverter = this->getTypeConverter();
+ auto deinterleaveResults = deinterleaveOp.getResultTypes();
+ auto packedOpResults =
+ llvmTypeConverter->packOperationResults(deinterleaveResults);
+ auto intrinsic = rewriter.create<LLVM::vector_deinterleave2>(
+ loc, packedOpResults, adaptor.getSource());
+
+ auto evenResult = rewriter.create<LLVM::ExtractValueOp>(
+ loc, intrinsic->getResult(0), 0);
+ auto oddResult = rewriter.create<LLVM::ExtractValueOp>(
+ loc, intrinsic->getResult(0), 1);
+
+ rewriter.replaceOp(deinterleaveOp, ValueRange{evenResult, oddResult});
+ return success();
+ }
+ // Lower fixed-size deinterleave to two shufflevectors. While the
+ // vector.deinterleave2 intrinsic supports fixed and scalable vectors, the
+ // langref still recommends fixed-vectors use shufflevector, see:
+ // https://llvm.org/docs/LangRef.html#id889.
+ int64_t resultVectorSize = resultType.getNumElements();
+ SmallVector<int32_t> evenShuffleMask;
+ SmallVector<int32_t> oddShuffleMask;
+
+ evenShuffleMask.reserve(resultVectorSize);
+ oddShuffleMask.reserve(resultVectorSize);
+
+ for (int i = 0; i < sourceType.getNumElements(); ++i) {
+ if (i % 2 == 0)
+ evenShuffleMask.push_back(i);
+ else
+ oddShuffleMask.push_back(i);
+ }
+
+ auto poison = rewriter.create<LLVM::PoisonOp>(loc, sourceType);
+ auto evenShuffle = rewriter.create<LLVM::ShuffleVectorOp>(
+ loc, adaptor.getSource(), poison, evenShuffleMask);
+ auto oddShuffle = rewriter.create<LLVM::ShuffleVectorOp>(
+ loc, adaptor.getSource(), poison, oddShuffleMask);
+
+ rewriter.replaceOp(deinterleaveOp, ValueRange{evenShuffle, oddShuffle});
+ return success();
+ }
+};
+
} // namespace
/// Populate the given list with patterns that convert from Vector to LLVM.
@@ -1785,8 +1849,8 @@ void mlir::populateVectorToLLVMConversionPatterns(
VectorExpandLoadOpConversion, VectorCompressStoreOpConversion,
VectorSplatOpLowering, VectorSplatNdOpLowering,
VectorScalableInsertOpLowering, VectorScalableExtractOpLowering,
- MaskedReductionOpConversion, VectorInterleaveOpLowering>(
- converter);
+ MaskedReductionOpConversion, VectorInterleaveOpLowering,
+ VectorDeinterleaveOpLowering>(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 a7a0ca3d43b01..245edb6789d30 100644
--- a/mlir/test/Conversion/VectorToLLVM/vector-to-llvm.mlir
+++ b/mlir/test/Conversion/VectorToLLVM/vector-to-llvm.mlir
@@ -2542,3 +2542,25 @@ func.func @vector_interleave_2d_scalable(%a: vector<2x[8]xi16>, %b: vector<2x[8]
%0 = vector.interleave %a, %b : vector<2x[8]xi16> -> vector<2x[16]xi16>
return %0 : vector<2x[16]xi16>
}
+
+// -----
+
+// CHECK-LABEL: @vector_deinterleave_1d
+// CHECK-SAME: (%[[SRC:.*]]: vector<4xi32>) -> (vector<2xi32>, vector<2xi32>)
+func.func @vector_deinterleave_1d(%a: vector<4xi32>) -> (vector<2xi32>, vector<2xi32>) {
+ // CHECK: %[[POISON:.*]] = llvm.mlir.poison : vector<4xi32>
+ // CHECK: llvm.shufflevector %[[SRC]], %[[POISON]] [0, 2] : vector<4xi32>
+ // CHECK: llvm.shufflevector %[[SRC]], %[[POISON]] [1, 3] : vector<4xi32>
+ %0, %1 = vector.deinterleave %a : vector<4xi32> -> vector<2xi32>
+ return %0, %1 : vector<2xi32>, vector<2xi32>
+}
+
+// CHECK-LABEL: @vector_deinterleave_1d_scalable
+// CHECK-SAME: %[[SRC:.*]]: vector<[4]xi32>) -> (vector<[2]xi32>, vector<[2]xi32>)
+func.func @vector_deinterleave_1d_scalable(%a: vector<[4]xi32>) -> (vector<[2]xi32>, vector<[2]xi32>) {
+ // CHECK: %[[RES:.*]] = "llvm.intr.vector.deinterleave2"(%[[SRC]]) : (vector<[4]xi32>) -> !llvm.struct<(vector<[2]xi32>, vector<[2]xi32>)>
+ // CHECK: llvm.extractvalue %[[RES]][0] : !llvm.struct<(vector<[2]xi32>, vector<[2]xi32>)>
+ // CHECK: llvm.extractvalue %[[RES]][1] : !llvm.struct<(vector<[2]xi32>, vector<[2]xi32>)>
+ %0, %1 = vector.deinterleave %a : vector<[4]xi32> -> vector<[2]xi32>
+ return %0, %1 : vector<[2]xi32>, vector<[2]xi32>
+}
diff --git a/mlir/test/Integration/Dialect/Vector/CPU/ArmSVE/Emulated/test-scalable-deinterleave.mlir b/mlir/test/Integration/Dialect/Vector/CPU/ArmSVE/Emulated/test-scalable-deinterleave.mlir
new file mode 100644
index 0000000000000..e6c561437132f
--- /dev/null
+++ b/mlir/test/Integration/Dialect/Vector/CPU/ArmSVE/Emulated/test-scalable-deinterleave.mlir
@@ -0,0 +1,30 @@
+// DEFINE: %{entry_point} = entry
+// DEFINE: %{compile} = mlir-opt %s -test-lower-to-llvm
+// DEFINE: %{run} = %mcr_aarch64_cmd -march=aarch64 -mattr=+sve \
+// DEFINE: -e %{entry_point} -entry-point-result=void \
+// DEFINE: -shared-libs=%mlir_c_runner_utils,%mlir_arm_runner_utils
+
+// RUN: %{compile} | %{run} | FileCheck %s
+
+func.func @entry() {
+ // Set the vector length to 256-bit (equivalent to vscale=2).
+ // This allows the checks (below) to look at an entire vector.
+ %c256 = arith.constant 256 : i32
+ func.call @setArmVLBits(%c256) : (i32) -> ()
+ func.call @test_deinterleave() : () -> ()
+ return
+}
+
+func.func @test_deinterleave() {
+ %step_vector = llvm.intr.experimental.stepvector : vector<[4]xi8>
+ vector.print %step_vector : vector<[4]xi8>
+ // CHECK: ( 0, 1, 2, 3, 4, 5, 6, 7 )
+ %v1, %v2 = vector.deinterleave %step_vector : vector<[4]xi8> -> vector<[2]xi8>
+ vector.print %v1 : vector<[2]xi8>
+ vector.print %v2 : vector<[2]xi8>
+ // CHECK: ( 0, 2, 4, 6 )
+ // CHECK: ( 1, 3, 5, 7 )
+ return
+}
+
+func.func private @setArmVLBits(%bits : i32)
diff --git a/mlir/test/Integration/Dialect/Vector/CPU/test-deinterleave.mlir b/mlir/test/Integration/Dialect/Vector/CPU/test-deinterleave.mlir
new file mode 100644
index 0000000000000..4915a3cde124d
--- /dev/null
+++ b/mlir/test/Integration/Dialect/Vector/CPU/test-deinterleave.mlir
@@ -0,0 +1,18 @@
+// RUN: mlir-opt %s -test-lower-to-llvm | \
+// RUN: mlir-cpu-runner -e entry -entry-point-result=void \
+// RUN: -shared-libs=%mlir_c_runner_utils | \
+// RUN: FileCheck %s
+
+func.func @entry() {
+ %v0 = arith.constant dense<[1, 2, 3, 4]> : vector<4xi8>
+ vector.print %v0 : vector<4xi8>
+ // CHECK: ( 1, 2, 3, 4 )
+
+ %v1, %v2 = vector.deinterleave %v0 : vector<4xi8> -> vector<2xi8>
+ vector.print %v1 : vector<2xi8>
+ vector.print %v2 : vector<2xi8>
+ // CHECK: ( 1, 3 )
+ // CHECK: ( 2, 4 )
+
+ return
+}
More information about the Mlir-commits
mailing list