Skip to content

[mlir][vector] Add folders for full constant transfer masks #71676

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
Nov 8, 2023
Merged

[mlir][vector] Add folders for full constant transfer masks #71676

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
10 changes: 10 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 @@ -2276,6 +2276,16 @@ def Vector_ConstantMaskOp :
```
}];

let extraClassDeclaration = [{
/// Return the result type of this op.
VectorType getVectorType() {
return cast<VectorType>(getOperation()->getResultTypes()[0]);
}

/// Return whether the mask is a uniform vector of `1`s.
bool isFullMask();
}];

let assemblyFormat = "$mask_dim_sizes attr-dict `:` type(results)";
let hasVerifier = 1;
}
Expand Down
37 changes: 37 additions & 0 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 @@ -3937,6 +3937,23 @@ static LogicalResult foldTransferInBoundsAttribute(TransferOp op) {
return success();
}

template <typename TransferOp>
static LogicalResult foldTransferFullMask(TransferOp op) {
auto mask = op.getMask();
if (!mask)
return failure();

auto constantMask = mask.template getDefiningOp<vector::ConstantMaskOp>();
if (!constantMask)
return failure();

if (!constantMask.isFullMask())
return failure();

op.getMaskMutable().clear();
return success();
}

/// ```
/// %w0 = vector.transfer_write %v0, %arg0[%c1, %c0] {in_bounds = [true, true]}
/// : vector<1x4xf32>, tensor<4x4xf32>
Expand Down Expand Up @@ -3969,6 +3986,8 @@ OpFoldResult TransferReadOp::fold(FoldAdaptor) {
/// transfer_read(memrefcast) -> transfer_read
if (succeeded(foldTransferInBoundsAttribute(*this)))
return getResult();
if (succeeded(foldTransferFullMask(*this)))
return getResult();
if (succeeded(memref::foldMemRefCast(*this)))
return getResult();
if (succeeded(tensor::foldTensorCast(*this)))
Expand Down Expand Up @@ -4334,6 +4353,8 @@ LogicalResult TransferWriteOp::fold(FoldAdaptor adaptor,
return success();
if (succeeded(foldTransferInBoundsAttribute(*this)))
return success();
if (succeeded(foldTransferFullMask(*this)))
return success();
return memref::foldMemRefCast(*this);
}

Expand Down Expand Up @@ -5601,6 +5622,22 @@ LogicalResult ConstantMaskOp::verify() {
return success();
}

bool ConstantMaskOp::isFullMask() {
auto resultType = getVectorType();
Copy link
Contributor

Choose a reason for hiding this comment

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

I think @aartbik had already created some utilities to check if a mask is all ones or zeros.

Copy link
Contributor Author

Choose a reason for hiding this comment

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

Do you happen to know where those utilities are? I poked around and couldn't find anything.

Copy link
Contributor

Choose a reason for hiding this comment

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

I thought Aart would know by heart but I found it :) https://github.com/llvm/llvm-project/blob/main/mlir/lib/Dialect/Vector/IR/VectorOps.cpp#L56-L120

Copy link
Contributor Author

Choose a reason for hiding this comment

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

Oh nice, although I don't think this would be something we should run as a folder, given that it might iterate over the elements of an arith.constant op. What about extracting out the constant mask (and maybe create_mask) handling alone and using it for the folder? Or we could add a canonicalization pattern instead.

Copy link
Contributor Author

@qedawkins qedawkins Nov 11, 2023 •
edited
Loading

Choose a reason for hiding this comment

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

A canonicalization would be nice because we could turn an all-false transfer_read into a vector.splat, and for transfer_write just drop the op altogether. (Just catching up to this work, maybe some of it already exists)

Copy link
Contributor

Choose a reason for hiding this comment

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

Yes, I would expect this to be a canonicalization.

// Check the corner case of 0-D vectors first.
if (resultType.getRank() == 0) {
assert(getMaskDimSizes().size() == 1 && "invalid sizes for zero rank mask");
return llvm::cast<IntegerAttr>(getMaskDimSizes()[0]).getInt() == 1;
}
for (const auto [resultSize, intAttr] :
llvm::zip_equal(resultType.getShape(), getMaskDimSizes())) {
int64_t maskDimSize = llvm::cast<IntegerAttr>(intAttr).getInt();
if (maskDimSize < resultSize)
return false;
}
return true;
}

//===----------------------------------------------------------------------===//
// CreateMaskOp
//===----------------------------------------------------------------------===//
Expand Down
23 changes: 23 additions & 0 deletions mlir/test/Dialect/Vector/canonicalize.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -854,6 +854,29 @@ func.func @canonicalize_broadcast_shapecast_to_shapecast(%arg0: vector<3x4xf32>)

// -----

// CHECK-LABEL: fold_vector_transfer_masks
func.func @fold_vector_transfer_masks(%A: memref<?x?xf32>) -> (vector<4x8xf32>) {
// CHECK: %[[C0:.+]] = arith.constant 0 : index
%c0 = arith.constant 0 : index
// CHECK: %[[F0:.+]] = arith.constant 0.000000e+00 : f32
%f0 = arith.constant 0.0 : f32

%mask = vector.constant_mask [8, 4] : vector<8x4xi1>

// CHECK: vector.transfer_read %{{.*}}, %[[F0]] {permutation_map
%1 = vector.transfer_read %A[%c0, %c0], %f0, %mask
{permutation_map = affine_map<(d0, d1) -> (d1, d0)>} : memref<?x?xf32>, vector<4x8xf32>

// CHECK: vector.transfer_write {{.*}}[%[[C0]], %[[C0]]] {permutation_map
vector.transfer_write %1, %A[%c0, %c0], %mask
{permutation_map = affine_map<(d0, d1) -> (d1, d0)>} : vector<4x8xf32>, memref<?x?xf32>

// CHECK: return
return %1 : vector<4x8xf32>
}

// -----

// CHECK-LABEL: fold_vector_transfers
func.func @fold_vector_transfers(%A: memref<?x8xf32>) -> (vector<4x8xf32>, vector<4x9xf32>) {
%c0 = arith.constant 0 : index
Expand Down