Skip to content

[mlir][vector] Remove unneeded mask restriction #113742

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
Oct 26, 2024
Merged

[mlir][vector] Remove unneeded mask restriction #113742

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
52 changes: 29 additions & 23 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 @@ -1819,17 +1819,17 @@ def Vector_MaskedLoadOp :
Vector_Op<"maskedload">,
Arguments<(ins Arg<AnyMemRef, "", [MemRead]>:$base,
Variadic<Index>:$indices,
VectorOfRankAndType<[1], [I1]>:$mask,
VectorOfRank<[1]>:$pass_thru)>,
Results<(outs VectorOfRank<[1]>:$result)> {
VectorOf<[I1]>:$mask,
AnyVector:$pass_thru)>,
Results<(outs AnyVector:$result)> {
Copy link
Contributor

@dcaballe dcaballe Oct 25, 2024

Choose a reason for hiding this comment

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

Probably good to add constraints to ensure that pass_thru and result's types match and that the number of elements in mask and one of the other vectors match as well.

Copy link
Member Author

@jpienaar jpienaar Oct 26, 2024

Choose a reason for hiding this comment

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

Ha this let to finding a bug in the generated inference with variadics (well known limitations there)


let summary = "loads elements from memory into a vector as defined by a mask vector";

let description = [{
The masked load reads elements from memory into a 1-D vector as defined
by a base with indices and a 1-D mask vector. When the mask is set, the
The masked load reads elements from memory into a vector as defined
by a base with indices and a mask vector. When the mask is set, the
element is read from memory. Otherwise, the corresponding element is taken
from a 1-D pass-through vector. Informally the semantics are:
from a pass-through vector. Informally the semantics are:
```
result[0] := if mask[0] then base[i + 0] else pass_thru[0]
result[1] := if mask[1] then base[i + 1] else pass_thru[1]
Expand Down Expand Up @@ -1882,14 +1882,14 @@ def Vector_MaskedStoreOp :
Vector_Op<"maskedstore">,
Arguments<(ins Arg<AnyMemRef, "", [MemWrite]>:$base,
Variadic<Index>:$indices,
VectorOfRankAndType<[1], [I1]>:$mask,
VectorOfRank<[1]>:$valueToStore)> {
VectorOf<[I1]>:$mask,
AnyVector:$valueToStore)> {

let summary = "stores elements from a vector into memory as defined by a mask vector";

let description = [{
The masked store operation writes elements from a 1-D vector into memory
as defined by a base with indices and a 1-D mask vector. When the mask is
The masked store operation writes elements from a vector into memory
as defined by a base with indices and a mask vector. When the mask is
set, the corresponding element from the vector is written to memory. Otherwise,
no action is taken for the element. Informally the semantics are:
```
Expand Down Expand Up @@ -2076,23 +2076,26 @@ def Vector_ExpandLoadOp :
Vector_Op<"expandload">,
Arguments<(ins Arg<AnyMemRef, "", [MemRead]>:$base,
Variadic<Index>:$indices,
VectorOfRankAndType<[1], [I1]>:$mask,
VectorOfRank<[1]>:$pass_thru)>,
Results<(outs VectorOfRank<[1]>:$result)> {
VectorOf<[I1]>:$mask,
AnyVector:$pass_thru)>,
Results<(outs AnyVector:$result)> {

let summary = "reads elements from memory and spreads them into a vector as defined by a mask";

let description = [{
The expand load reads elements from memory into a 1-D vector as defined
by a base with indices and a 1-D mask vector. When the mask is set, the
next element is read from memory. Otherwise, the corresponding element
is taken from a 1-D pass-through vector. Informally the semantics are:
The expand load reads elements from memory into a vector as defined by a
base with indices and a mask vector. Expansion only applies to the innermost
dimension. When the mask is set, the next element is read from memory.
Otherwise, the corresponding element is taken from a pass-through vector.
Informally the semantics are:

```
index = i
result[0] := if mask[0] then base[index++] else pass_thru[0]
result[1] := if mask[1] then base[index++] else pass_thru[1]
etc.
```

Note that the index increment is done conditionally.

If a mask bit is set and the corresponding index is out-of-bounds for the
Expand Down Expand Up @@ -2140,22 +2143,25 @@ def Vector_CompressStoreOp :
Vector_Op<"compressstore">,
Arguments<(ins Arg<AnyMemRef, "", [MemWrite]>:$base,
Variadic<Index>:$indices,
VectorOfRankAndType<[1], [I1]>:$mask,
VectorOfRank<[1]>:$valueToStore)> {
VectorOf<[I1]>:$mask,
AnyVector:$valueToStore)> {
Copy link
Contributor

@dcaballe dcaballe Oct 25, 2024

Choose a reason for hiding this comment

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

For compress and expand, I would state in the doc that the compression/expansion applies only along the innermost dimension.

Copy link
Member Author

@jpienaar jpienaar Oct 26, 2024

Choose a reason for hiding this comment

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

SG


let summary = "writes elements selectively from a vector as defined by a mask";

let description = [{
The compress store operation writes elements from a 1-D vector into memory
as defined by a base with indices and a 1-D mask vector. When the mask is
set, the corresponding element from the vector is written next to memory.
Otherwise, no action is taken for the element. Informally the semantics are:
The compress store operation writes elements from a vector into memory as
defined by a base with indices and a mask vector. Compression only applies
to the innermost dimension. When the mask is set, the corresponding element
from the vector is written next to memory. Otherwise, no action is taken
for the element. Informally the semantics are:

```
index = i
if (mask[0]) base[index++] = value[0]
if (mask[1]) base[index++] = value[1]
etc.
```

Note that the index increment is done conditionally.

If a mask bit is set and the corresponding index is out-of-bounds for the
Expand Down
8 changes: 4 additions & 4 deletions mlir/lib/Dialect/Vector/IR/VectorOps.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4977,8 +4977,8 @@ LogicalResult MaskedLoadOp::verify() {
return emitOpError("base and result element type should match");
if (llvm::size(getIndices()) != memType.getRank())
return emitOpError("requires ") << memType.getRank() << " indices";
if (resVType.getDimSize(0) != maskVType.getDimSize(0))
return emitOpError("expected result dim to match mask dim");
if (resVType.getShape() != maskVType.getShape())
return emitOpError("expected result shape to match mask shape");
if (resVType != passVType)
return emitOpError("expected pass_thru of same type as result type");
return success();
Expand Down Expand Up @@ -5030,8 +5030,8 @@ LogicalResult MaskedStoreOp::verify() {
return emitOpError("base and valueToStore element type should match");
if (llvm::size(getIndices()) != memType.getRank())
return emitOpError("requires ") << memType.getRank() << " indices";
if (valueVType.getDimSize(0) != maskVType.getDimSize(0))
return emitOpError("expected valueToStore dim to match mask dim");
if (valueVType.getShape() != maskVType.getShape())
return emitOpError("expected valueToStore shape to match mask shape");
return success();
}

Expand Down
4 changes: 2 additions & 2 deletions mlir/test/Dialect/Vector/invalid.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1356,7 +1356,7 @@ func.func @maskedload_base_type_mismatch(%base: memref<?xf64>, %mask: vector<16x

func.func @maskedload_dim_mask_mismatch(%base: memref<?xf32>, %mask: vector<15xi1>, %pass: vector<16xf32>) {
%c0 = arith.constant 0 : index
// expected-error@+1 {{'vector.maskedload' op expected result dim to match mask dim}}
// expected-error@+1 {{'vector.maskedload' op expected result shape to match mask shape}}
%0 = vector.maskedload %base[%c0], %mask, %pass : memref<?xf32>, vector<15xi1>, vector<16xf32> into vector<16xf32>
}

Expand Down Expand Up @@ -1387,7 +1387,7 @@ func.func @maskedstore_base_type_mismatch(%base: memref<?xf64>, %mask: vector<16

func.func @maskedstore_dim_mask_mismatch(%base: memref<?xf32>, %mask: vector<15xi1>, %value: vector<16xf32>) {
%c0 = arith.constant 0 : index
// expected-error@+1 {{'vector.maskedstore' op expected valueToStore dim to match mask dim}}
// expected-error@+1 {{'vector.maskedstore' op expected valueToStore shape to match mask shape}}
vector.maskedstore %base[%c0], %mask, %value : memref<?xf32>, vector<15xi1>, vector<16xf32>
}

Expand Down
Loading