[Mlir-commits] [mlir] [mlir][VectorOps] Add conversion of 1-D vector.interleave ops to LLVM (PR #80966)
Benjamin Maxwell
llvmlistbot at llvm.org
Wed Feb 7 02:14:42 PST 2024
https://github.com/MacDue created https://github.com/llvm/llvm-project/pull/80966
The 1-D case directly maps to LLVM intrinsics. The n-D case will be
handled by unrolling to 1-D first (in a later patch).
Depends on: #80965
>From 4b442eeb377cfcb60a0d05cf6f1ec6ba735c3152 Mon Sep 17 00:00:00 2001
From: Benjamin Maxwell <benjamin.maxwell at arm.com>
Date: Wed, 7 Feb 2024 09:53:22 +0000
Subject: [PATCH 1/2] [mlir][VectorOps] Add vector.interleave operation
The interleave operation constructs a new vector by interleaving the
elements from the trailing (or final) dimension of two input vectors,
returning a new vector where the trailing dimension is twice the size.
Note that for the n-D case this differs from the interleaving possible
with `vector.shuffle`, which would only operate on the leading
dimension.
Another key difference is this operation supports scalable vectors,
though currently a general LLVM lowering is limited to the case where
only the trailing dimension is scalable.
Example:
```mlir
%0 = vector.interleave %a, %b
: vector<[4]xi32> ; yields vector<[8]xi32>
%1 = vector.interleave %c, %d
: vector<8xi8> ; yields vector<16xi8>
%2 = vector.interleave %e, %f
: vector<f16> ; yields vector<2xf16>
%3 = vector.interleave %g, %h
: vector<2x4x[2]xf64> ; yields vector<2x4x[4]xf64>
%4 = vector.interleave %i, %j
: vector<6x3xf32> ; yields vector<6x6xf32>
```
Note: This change alone does not add any lowerings.
---
.../mlir/Dialect/Vector/IR/VectorOps.td | 63 +++++++++++++++++++
mlir/test/Dialect/Vector/ops.mlir | 35 +++++++++++
2 files changed, 98 insertions(+)
diff --git a/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td b/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td
index bc08f8d07fb0d..6d50b0654bc57 100644
--- a/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td
+++ b/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td
@@ -478,6 +478,69 @@ def Vector_ShuffleOp :
let hasCanonicalizer = 1;
}
+def Vector_InterleaveOp :
+ Vector_Op<"interleave", [Pure,
+ AllTypesMatch<["lhs", "rhs"]>,
+ TypesMatchWith<
+ "type of 'result' is double the width of the inputs",
+ "lhs", "result",
+ [{
+ [&]() -> ::mlir::VectorType {
+ auto vectorType = ::llvm::cast<mlir::VectorType>($_self);
+ ::mlir::VectorType::Builder builder(vectorType);
+ if (vectorType.getRank() == 0) {
+ static constexpr int64_t v2xty_shape[] = { 2 };
+ return builder.setShape(v2xty_shape);
+ }
+ auto lastDim = vectorType.getRank() - 1;
+ return builder.setDim(lastDim, vectorType.getDimSize(lastDim) * 2);
+ }()
+ }]>]> {
+ let summary = "constructs a vector by interleaving two input vectors";
+ let description = [{
+ The interleave operation constructs a new vector by interleaving the
+ elements from the trailing (or final) dimension of two input vectors,
+ returning a new vector where the trailing dimension is twice the size.
+
+ Note that for the n-D case this differs from the interleaving possible with
+ `vector.shuffle`, which would only operate on the leading dimension.
+
+ Another key difference is this operation supports scalable vectors, though
+ currently a general LLVM lowering is limited to the case where only the
+ trailing dimension is scalable.
+
+ Example:
+ ```mlir
+ %0 = vector.interleave %a, %b
+ : vector<[4]xi32> ; yields vector<[8]xi32>
+ %1 = vector.interleave %c, %d
+ : vector<8xi8> ; yields vector<16xi8>
+ %2 = vector.interleave %e, %f
+ : vector<f16> ; yields vector<2xf16>
+ %3 = vector.interleave %g, %h
+ : vector<2x4x[2]xf64> ; yields vector<2x4x[4]xf64>
+ %4 = vector.interleave %i, %j
+ : vector<6x3xf32> ; yields vector<6x6xf32>
+ ```
+ }];
+
+ let arguments = (ins AnyVectorOfAnyRank:$lhs, AnyVectorOfAnyRank:$rhs);
+ let results = (outs AnyVector:$result);
+
+ let assemblyFormat = [{
+ $lhs `,` $rhs attr-dict `:` type($lhs)
+ }];
+
+ let extraClassDeclaration = [{
+ VectorType getSourceVectorType() {
+ return ::llvm::cast<VectorType>(getLhs().getType());
+ }
+ VectorType getResultVectorType() {
+ return ::llvm::cast<VectorType>(getResult().getType());
+ }
+ }];
+}
+
def Vector_ExtractElementOp :
Vector_Op<"extractelement", [Pure,
TypesMatchWith<"result type matches element type of vector operand",
diff --git a/mlir/test/Dialect/Vector/ops.mlir b/mlir/test/Dialect/Vector/ops.mlir
index 2f8530e7c171a..79a80be4f8b20 100644
--- a/mlir/test/Dialect/Vector/ops.mlir
+++ b/mlir/test/Dialect/Vector/ops.mlir
@@ -1081,3 +1081,38 @@ func.func @fastmath(%x: vector<42xf32>) -> f32 {
%min = vector.reduction <minnumf>, %x fastmath<reassoc,nnan,ninf> : vector<42xf32> into f32
return %min: f32
}
+
+// CHECK-LABEL: @interleave_0d
+func.func @interleave_0d(%a: vector<f32>, %b: vector<f32>) -> vector<2xf32> {
+ // CHECK: vector.interleave %{{.*}}, %{{.*}} : vector<f32>
+ %0 = vector.interleave %a, %b : vector<f32>
+ return %0 : vector<2xf32>
+}
+
+// CHECK-LABEL: @interleave_1d
+func.func @interleave_1d(%a: vector<4xf32>, %b: vector<4xf32>) -> vector<8xf32> {
+ // CHECK: vector.interleave %{{.*}}, %{{.*}} : vector<4xf32>
+ %0 = vector.interleave %a, %b : vector<4xf32>
+ return %0 : vector<8xf32>
+}
+
+// CHECK-LABEL: @interleave_1d_scalable
+func.func @interleave_1d_scalable(%a: vector<[8]xi16>, %b: vector<[8]xi16>) -> vector<[16]xi16> {
+ // CHECK: vector.interleave %{{.*}}, %{{.*}} : vector<[8]xi16>
+ %0 = vector.interleave %a, %b : vector<[8]xi16>
+ return %0 : vector<[16]xi16>
+}
+
+// CHECK-LABEL: @interleave_2d
+func.func @interleave_2d(%a: vector<2x8xf32>, %b: vector<2x8xf32>) -> vector<2x16xf32> {
+ // CHECK: vector.interleave %{{.*}}, %{{.*}} : vector<2x8xf32>
+ %0 = vector.interleave %a, %b : vector<2x8xf32>
+ return %0 : vector<2x16xf32>
+}
+
+// CHECK-LABEL: @interleave_2d_scalable
+func.func @interleave_2d_scalable(%a: vector<2x[2]xf64>, %b: vector<2x[2]xf64>) -> vector<2x[4]xf64> {
+ // CHECK: vector.interleave %{{.*}}, %{{.*}} : vector<2x[2]xf64>
+ %0 = vector.interleave %a, %b : vector<2x[2]xf64>
+ return %0 : vector<2x[4]xf64>
+}
>From c2047adb5672ec8eaa05e16f34dbbb794a0eba6c Mon Sep 17 00:00:00 2001
From: Benjamin Maxwell <benjamin.maxwell at arm.com>
Date: Wed, 7 Feb 2024 09:56:42 +0000
Subject: [PATCH 2/2] [mlir][VectorOps] Add conversion of 1-D vector.interleave
ops to LLVM
The 1-D case directly maps to LLVM intrinsics. The n-D case will be
handled by unrolling to 1-D first (in a later patch).
---
.../VectorToLLVM/ConvertVectorToLLVM.cpp | 41 ++++++++++++++++++-
.../VectorToLLVM/vector-to-llvm.mlir | 37 +++++++++++++++++
2 files changed, 77 insertions(+), 1 deletion(-)
diff --git a/mlir/lib/Conversion/VectorToLLVM/ConvertVectorToLLVM.cpp b/mlir/lib/Conversion/VectorToLLVM/ConvertVectorToLLVM.cpp
index b66b55ae8d57f..0d9a451d11ca8 100644
--- a/mlir/lib/Conversion/VectorToLLVM/ConvertVectorToLLVM.cpp
+++ b/mlir/lib/Conversion/VectorToLLVM/ConvertVectorToLLVM.cpp
@@ -1734,6 +1734,44 @@ struct VectorSplatNdOpLowering : public ConvertOpToLLVMPattern<SplatOp> {
}
};
+/// Conversion pattern for a `vector.interleave`.
+/// This supports fixed-sized vectors and scalable vectors.
+struct VectorInterleaveOpLowering
+ : public ConvertOpToLLVMPattern<vector::InterleaveOp> {
+ 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 failure();
+ // If the result is rank 1, then this directly maps to LLVM.
+ if (resultType.isScalable()) {
+ rewriter.replaceOpWithNewOp<LLVM::experimental_vector_interleave2>(
+ 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<int32_t> 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<LLVM::ShuffleVectorOp>(
+ interleaveOp, adaptor.getLhs(), adaptor.getRhs(),
+ interleaveShuffleMask);
+ return success();
+ }
+};
+
} // namespace
/// Populate the given list with patterns that convert from Vector to LLVM.
@@ -1758,7 +1796,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<i8>, %[[RHS:.*]]: vector<i8>)
+func.func @vector_interleave_0d(%a: vector<i8>, %b: vector<i8>) -> vector<2xi8> {
+ // CHECK: %[[LHS_RANK1:.*]] = builtin.unrealized_conversion_cast %[[LHS]] : vector<i8> to vector<1xi8>
+ // CHECK: %[[RHS_RANK1:.*]] = builtin.unrealized_conversion_cast %[[RHS]] : vector<i8> to vector<1xi8>
+ // CHECK: %[[ZIP:.*]] = llvm.shufflevector %[[LHS_RANK1]], %[[RHS_RANK1]] [0, 1] : vector<1xi8>
+ // CHECK: return %[[ZIP]]
+ %0 = vector.interleave %a, %b : vector<i8>
+ 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>
+}
More information about the Mlir-commits
mailing list