[mlir][VectorOps] Add vector.interleave operation (1/4) #80965
Conversation
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.
|
@llvm/pr-subscribers-mlir Author: Benjamin Maxwell (MacDue) ChangesThe 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 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: %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. Full diff: https://github.com/llvm/llvm-project/pull/80965.diff 2 Files Affected:
diff --git a/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td b/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td
index bc08f8d07fb0de..6d50b0654bc575 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 2f8530e7c171aa..79a80be4f8b202 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>
+}
|
|
@llvm/pr-subscribers-mlir-vector Author: Benjamin Maxwell (MacDue) ChangesThe 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 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: %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. Full diff: https://github.com/llvm/llvm-project/pull/80965.diff 2 Files Affected:
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>
+}
|
MacDue
changed the title
| 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 |
There was a problem hiding this comment.
Thanks for proposing this, 2 questions come to mind:
- what is the state of the 1-D support for vector.shuffle / shufflevector in MLIR / LLVM for scalable types? Is this generally supported for various permutations or is interleaving a desirable outlier (because e.g. the HW can really do well on this but not on more general permutations)?
- would this be necessary if 1fe6568 had introduced shuffles to be on the innermost dim? I can't remember the rationale but shuffling the outer dims is more or less a no-op in many cases since these dimensions get unrolled when lowering, on most HW...
I am trying to gauge whether we would be better off having a vector.shuffle_innermost_dim that takes a general shuffle mask, for which vector.interleave is a special case that is easy to detect.
In the absence of compelling HW reasons to only have interleave, I'd favor a "better shuffle" because it is very easy to detect that [0, n, 1, n+1 ... n-1, 2n-1] is an interleave.
There was a problem hiding this comment.
In LLVM for scalable vectors shufflevector can only perform a splat, there's no support for any other shuffle masks. An interleave for scalable vectors has to use the special interleave2 intrinsic: https://llvm.org/docs/LangRef.html#llvm-experimental-vector-interleave2-intrinsic.
This is because the fixed-size shuffle mask does not make sense in the context of scalable vectors. The length of the mask may be less than the total number of elements in the two vectors.
There was a problem hiding this comment.
See also the shufflevector docs: https://llvm.org/docs/LangRef.html#id189
For scalable vectors, the only valid mask values at present are zeroinitializer, undef and poison, since we cannot write all indices as literals for a vector with a length unknown at compile time.
There was a problem hiding this comment.
So, in short:
- There's basically no support for
shufflevector,vector.shuffle, etc for scalable vectors - I believe so
For 2., while you can hack it and claim an arbitrary pattern means interleave for scalable vectors, there's no pattern that'd actually make sense with the current semantics of vector.shuffle. So I think that's just adding tricky edge cases in the meaning of the op.
There was a problem hiding this comment.
cool, thanks for digging in, makes sense !
There was a problem hiding this comment.
Yes, plus this is also similar to having a vector.broadcast vs just using a shuffle for that. Interleaving/deinterleaving are very common patterns and having something with a bit higher level is helpful.
| 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 |
There was a problem hiding this comment.
Yes, plus this is also similar to having a vector.broadcast vs just using a shuffle for that. Interleaving/deinterleaving are very common patterns and having something with a bit higher level is helpful.
I'm not sure a meaningful negative test can be written given all the types are inferred from a single input |
+1, since there are no custom messages, I believe everything is covered by the default verifiers |
|
Yep, that's correct the errors are just " |
|
If there's no objection, I'll land this patch in the next few days :) |
…ges (#102125) - Use vector.interleave rather than the LLVM intrinsic - Remove dependency on LLVM dialect - Remove manual outerproduct erases (these are now trivially dead) - Remove comment explaining issues with previous tile allocator - Update pipeline in `multi-tile-matmul-mixed-types.mlir` Recent changes: #90448, #80965
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:
Note: This change alone does not add any lowerings.