Skip to content

[mlir][VectorOps] Add vector.interleave operation (1/4) #80965

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Feb 12, 2024
Merged

[mlir][VectorOps] Add vector.interleave operation (1/4) #80965

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
63 changes: 63 additions & 0 deletions mlir/include/mlir/Dialect/Vector/IR/VectorOps.td
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Copy link
Contributor

@nicolasvasilache nicolasvasilache Feb 7, 2024 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for proposing this, 2 questions come to mind:

  1. 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)?
  2. 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.

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Member Author

@MacDue MacDue Feb 7, 2024 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

So, in short:

  1. There's basically no support for shufflevector, vector.shuffle, etc for scalable vectors
  2. 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.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

cool, thanks for digging in, makes sense !

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

`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",
Expand Down
35 changes: 35 additions & 0 deletions mlir/test/Dialect/Vector/ops.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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>
}