Skip to content
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.

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

[mlir][vector] Make the in_bounds attribute mandatory #97049

Merged
merged 7 commits into from
Jul 16, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
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
45 changes: 26 additions & 19 deletions mlir/include/mlir/Dialect/Vector/IR/VectorOps.td
Original file line number Diff line number Diff line change
Expand Up @@ -1363,7 +1363,7 @@ def Vector_TransferReadOp :
AffineMapAttr:$permutation_map,
AnyType:$padding,
Optional<VectorOf<[I1]>>:$mask,
OptionalAttr<BoolArrayAttr>:$in_bounds)>,
BoolArrayAttr:$in_bounds)>,
joker-eph marked this conversation as resolved.
Show resolved Hide resolved
joker-eph marked this conversation as resolved.
Show resolved Hide resolved
Results<(outs AnyVectorOfAnyRank:$vector)> {

let summary = "Reads a supervector from memory into an SSA vector value.";
Expand Down Expand Up @@ -1401,16 +1401,19 @@ def Vector_TransferReadOp :
permutation or broadcasting. Elements whose corresponding mask element is
`0` are masked out and replaced with `padding`.

An optional boolean array attribute `in_bounds` specifies for every vector
dimension if the transfer is guaranteed to be within the source bounds. If
specified, the `in_bounds` array length has to be equal to the vector rank.
If set to "false", accesses (including the starting point) may run
For every vector dimension, the boolean array attribute `in_bounds`
specifies if the transfer is guaranteed to be within the source bounds. If
set to "false", accesses (including the starting point) may run
out-of-bounds along the respective vector dimension as the index increases.
Broadcast dimensions must always be in-bounds. In absence of the attribute,
accesses along all vector dimensions (except for broadcasts) may run
out-of-bounds. A `vector.transfer_read` can be lowered to a simple load if
all dimensions are specified to be within bounds and no `mask` was
specified. Note that non-vector dimensions *must* always be in-bounds.
Non-vector and broadcast dimensions *must* always be in-bounds. The
Copy link
Contributor

Choose a reason for hiding this comment

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

dumb question: what does Non-vector mean here?

Copy link
Contributor Author

Choose a reason for hiding this comment

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

I also don't know :) That note was introduced by @matthias-springer in https://reviews.llvm.org/D155719. Matthias, do you remember what cases that was referring to?

`in_bounds` array length has to be equal to the vector rank. This attribute
has a default value: `false` (i.e. "out-of-bounds"). When skipped in the
textual IR, the default value is assumed. Similarly, the OP printer will
omit this attribute when all dimensions are out-of-bounds (i.e. the default
value is used).

A `vector.transfer_read` can be lowered to a simple load if all dimensions
are specified to be within bounds and no `mask` was specified.

This operation is called 'read' by opposition to 'load' because the
super-vector granularity is generally not representable with a single
Expand Down Expand Up @@ -1607,7 +1610,7 @@ def Vector_TransferWriteOp :
Variadic<Index>:$indices,
AffineMapAttr:$permutation_map,
Optional<VectorOf<[I1]>>:$mask,
OptionalAttr<BoolArrayAttr>:$in_bounds)>,
BoolArrayAttr:$in_bounds)>,
Results<(outs Optional<AnyRankedTensor>:$result)> {

let summary = "The vector.transfer_write op writes a supervector to memory.";
Expand Down Expand Up @@ -1643,15 +1646,19 @@ def Vector_TransferWriteOp :
any permutation. Elements whose corresponding mask element is `0` are
masked out.

An optional boolean array attribute `in_bounds` specifies for every vector
dimension if the transfer is guaranteed to be within the source bounds. If
specified, the `in_bounds` array length has to be equal to the vector rank.
If set to "false", accesses (including the starting point) may run
For every vector dimension, the boolean array attribute `in_bounds`
specifies if the transfer is guaranteed to be within the source bounds. If
set to "false", accesses (including the starting point) may run
out-of-bounds along the respective vector dimension as the index increases.
In absence of the attribute, accesses along all vector dimensions may run
out-of-bounds. A `vector.transfer_write` can be lowered to a simple store if
all dimensions are specified to be within bounds and no `mask` was
specified. Note that non-vector dimensions *must* always be in-bounds.
Non-vector and broadcast dimensions *must* always be in-bounds. The
`in_bounds` array length has to be equal to the vector rank. This attribute
has a default value: `false` (i.e. "out-of-bounds"). When skipped in the
textual IR, the default value is assumed. Similarly, the OP printer will
omit this attribute when all dimensions are out-of-bounds (i.e. the default
value is used).

A `vector.transfer_write` can be lowered to a simple store if all
dimensions are specified to be within bounds and no `mask` was specified.

This operation is called 'write' by opposition to 'store' because the
super-vector granularity is generally not representable with a single
Expand Down
8 changes: 8 additions & 0 deletions mlir/include/mlir/IR/AffineMap.h
Original file line number Diff line number Diff line change
Expand Up @@ -146,6 +146,14 @@ class AffineMap {
/// affine map (d0, ..., dn) -> (dp, ..., dn) on the most minor dimensions.
bool isMinorIdentity() const;

/// Returns the list of broadcast dimensions (i.e. dims indicated by value 0
/// in the result).
/// Ex:
/// * (d0, d1, d2) -> (0, d1) gives [0]
/// * (d0, d1, d2) -> (d2, d1) gives []
/// * (d0, d1, d2, d4) -> (d0, 0, d1, 0) gives [1, 3]
SmallVector<unsigned> getBroadcastDims() const;
Copy link
Contributor

Choose a reason for hiding this comment

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

Have you looked at isMinorIdentityWithBroadcasting

Copy link
Contributor Author

Choose a reason for hiding this comment

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

Yes - that was my source of inspiration :) It's just below getBroadcastDims (though I am about to move it a bit).

Note, isMinorIdentityWithBroadcasting wouldn't work as maps in xfer_read/xfer_write can be permutations. I also tried re-using getBroadcastDims inside isMinorIdentityWithBroadcasting, but that didn't feel right (perhaps isMinorIdentityWithBroadcasting should be split?).


/// Returns true if this affine map is a minor identity up to broadcasted
/// dimensions which are indicated by value 0 in the result. If
/// `broadcastedDims` is not null, it will be populated with the indices of
Expand Down
6 changes: 2 additions & 4 deletions mlir/include/mlir/Interfaces/VectorInterfaces.td
Original file line number Diff line number Diff line change
Expand Up @@ -98,7 +98,7 @@ def VectorTransferOpInterface : OpInterface<"VectorTransferOpInterface"> {
dimension whether it is in-bounds or not. (Broadcast dimensions are
always in-bounds).
}],
/*retTy=*/"::std::optional<::mlir::ArrayAttr>",
/*retTy=*/"::mlir::ArrayAttr",
/*methodName=*/"getInBounds",
/*args=*/(ins)
>,
Expand Down Expand Up @@ -240,9 +240,7 @@ def VectorTransferOpInterface : OpInterface<"VectorTransferOpInterface"> {
bool isDimInBounds(unsigned dim) {
if ($_op.isBroadcastDim(dim))
return true;
if (!$_op.getInBounds())
return false;
auto inBounds = ::llvm::cast<::mlir::ArrayAttr>(*$_op.getInBounds());
auto inBounds = $_op.getInBounds();
return ::llvm::cast<::mlir::BoolAttr>(inBounds[dim]).getValue();
}

Expand Down
13 changes: 12 additions & 1 deletion mlir/lib/Dialect/Affine/Transforms/SuperVectorize.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -1223,8 +1223,19 @@ static Operation *vectorizeAffineLoad(AffineLoadOp loadOp,
LLVM_DEBUG(dbgs() << "\n[early-vect]+++++ permutationMap: ");
LLVM_DEBUG(permutationMap.print(dbgs()));

// Make sure that the in_bounds attribute corresponding to a broadcast dim
// is set to `true` - that's required by the xfer Op.
Copy link
Contributor

Choose a reason for hiding this comment

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

Nothing for this PR but this is still a constraint that seems unnecessary and makes this complicated, as we can see here.

Copy link
Contributor Author

Choose a reason for hiding this comment

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

Agreed. Expect a PR to remove this shortly :) (~next week)

// FIXME: We're not veryfying whether the corresponding access is in bounds.
// TODO: Use masking instead.
SmallVector<unsigned> broadcastedDims = permutationMap.getBroadcastDims();
SmallVector<bool> inBounds(vectorType.getRank(), false);

for (auto idx : broadcastedDims)
inBounds[idx] = true;

auto transfer = state.builder.create<vector::TransferReadOp>(
loadOp.getLoc(), vectorType, loadOp.getMemRef(), indices, permutationMap);
loadOp.getLoc(), vectorType, loadOp.getMemRef(), indices, permutationMap,
inBounds);

// Register replacement for future uses in the scope.
state.registerOpVectorReplacement(loadOp, transfer);
Expand Down
25 changes: 18 additions & 7 deletions mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -1343,8 +1343,17 @@ vectorizeAsLinalgGeneric(RewriterBase &rewriter, VectorizationState &state,

SmallVector<Value> indices(linalgOp.getShape(opOperand).size(), zero);

// Make sure that the in_bounds attribute corresponding to a broadcast dim
// is `true`
SmallVector<unsigned> broadcastedDims = readMap.getBroadcastDims();
SmallVector<bool> inBounds(readType.getRank(), false);

for (auto idx : broadcastedDims)
inBounds[idx] = true;

Operation *read = rewriter.create<vector::TransferReadOp>(
loc, readType, opOperand->get(), indices, readMap);
loc, readType, opOperand->get(), indices, readMap,
ArrayRef<bool>(inBounds));
read = state.maskOperation(rewriter, read, linalgOp, maskingMap);
Value readValue = read->getResult(0);

Expand Down Expand Up @@ -2681,11 +2690,12 @@ LogicalResult LinalgCopyVTRForwardingPattern::matchAndRewrite(
// The `masked` attribute is only valid on this padded buffer.
// When forwarding to vector.transfer_read, the attribute must be reset
// conservatively.
auto vectorType = xferOp.getVectorType();
Value res = rewriter.create<vector::TransferReadOp>(
xferOp.getLoc(), xferOp.getVectorType(), in, xferOp.getIndices(),
xferOp.getLoc(), vectorType, in, xferOp.getIndices(),
xferOp.getPermutationMapAttr(), xferOp.getPadding(), xferOp.getMask(),
// in_bounds is explicitly reset
/*inBoundsAttr=*/ArrayAttr());
rewriter.getBoolArrayAttr(
SmallVector<bool>(vectorType.getRank(), false)));

if (maybeFillOp)
rewriter.eraseOp(maybeFillOp);
Expand Down Expand Up @@ -2739,11 +2749,12 @@ LogicalResult LinalgCopyVTWForwardingPattern::matchAndRewrite(
// The `masked` attribute is only valid on this padded buffer.
// When forwarding to vector.transfer_write, the attribute must be reset
// conservatively.
auto vector = xferOp.getVector();
rewriter.create<vector::TransferWriteOp>(
xferOp.getLoc(), xferOp.getVector(), out, xferOp.getIndices(),
xferOp.getLoc(), vector, out, xferOp.getIndices(),
xferOp.getPermutationMapAttr(), xferOp.getMask(),
// in_bounds is explicitly reset
/*inBoundsAttr=*/ArrayAttr());
rewriter.getBoolArrayAttr(
SmallVector<bool>(vector.getType().getRank(), false)));

rewriter.eraseOp(copyOp);
rewriter.eraseOp(xferOp);
Expand Down
54 changes: 34 additions & 20 deletions mlir/lib/Dialect/Vector/IR/VectorOps.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -3818,7 +3818,8 @@ void TransferReadOp::build(OpBuilder &builder, OperationState &result,
auto permutationMapAttr = AffineMapAttr::get(permutationMap);
auto inBoundsAttr = (inBounds && !inBounds.value().empty())
? builder.getBoolArrayAttr(inBounds.value())
: ArrayAttr();
: builder.getBoolArrayAttr(
SmallVector<bool>(vectorType.getRank(), false));
build(builder, result, vectorType, source, indices, permutationMapAttr,
inBoundsAttr);
}
Expand All @@ -3833,7 +3834,8 @@ void TransferReadOp::build(OpBuilder &builder, OperationState &result,
auto permutationMapAttr = AffineMapAttr::get(permutationMap);
auto inBoundsAttr = (inBounds && !inBounds.value().empty())
? builder.getBoolArrayAttr(inBounds.value())
: ArrayAttr();
: builder.getBoolArrayAttr(
SmallVector<bool>(vectorType.getRank(), false));
build(builder, result, vectorType, source, indices, permutationMapAttr,
padding,
/*mask=*/Value(), inBoundsAttr);
Expand Down Expand Up @@ -3951,17 +3953,15 @@ verifyTransferOp(VectorTransferOpInterface op, ShapedType shapedType,
<< inferredMaskType << ") and mask operand type (" << maskType
<< ") don't match";

if (inBounds) {
if (permutationMap.getNumResults() != static_cast<int64_t>(inBounds.size()))
return op->emitOpError("expects the optional in_bounds attr of same rank "
"as permutation_map results: ")
<< AffineMapAttr::get(permutationMap)
<< " vs inBounds of size: " << inBounds.size();
for (unsigned int i = 0; i < permutationMap.getNumResults(); ++i)
if (isa<AffineConstantExpr>(permutationMap.getResult(i)) &&
!llvm::cast<BoolAttr>(inBounds.getValue()[i]).getValue())
return op->emitOpError("requires broadcast dimensions to be in-bounds");
}
if (permutationMap.getNumResults() != static_cast<int64_t>(inBounds.size()))
return op->emitOpError("expects the in_bounds attr of same rank "
"as permutation_map results: ")
<< AffineMapAttr::get(permutationMap)
<< " vs inBounds of size: " << inBounds.size();
for (unsigned int i = 0, e = permutationMap.getNumResults(); i < e; ++i)
if (isa<AffineConstantExpr>(permutationMap.getResult(i)) &&
!llvm::cast<BoolAttr>(inBounds.getValue()[i]).getValue())
return op->emitOpError("requires broadcast dimensions to be in-bounds");

return success();
}
Expand Down Expand Up @@ -4037,6 +4037,13 @@ ParseResult TransferReadOp::parse(OpAsmParser &parser, OperationState &result) {
} else {
permMap = llvm::cast<AffineMapAttr>(permMapAttr).getValue();
}
auto inBoundsAttrName = TransferReadOp::getInBoundsAttrName(result.name);
Attribute inBoundsAttr = result.attributes.get(inBoundsAttrName);
if (!inBoundsAttr) {
result.addAttribute(inBoundsAttrName,
builder.getBoolArrayAttr(
SmallVector<bool>(permMap.getNumResults(), false)));
}
if (parser.resolveOperand(sourceInfo, shapedType, result.operands) ||
parser.resolveOperands(indexInfo, indexType, result.operands) ||
parser.resolveOperand(paddingInfo, shapedType.getElementType(),
Expand Down Expand Up @@ -4081,8 +4088,7 @@ LogicalResult TransferReadOp::verify() {

if (failed(verifyTransferOp(cast<VectorTransferOpInterface>(getOperation()),
shapedType, vectorType, maskType,
inferredMaskType, permutationMap,
getInBounds() ? *getInBounds() : ArrayAttr())))
inferredMaskType, permutationMap, getInBounds())))
return failure();

if (auto sourceVectorElementType =
Expand Down Expand Up @@ -4355,9 +4361,11 @@ void TransferWriteOp::build(OpBuilder &builder, OperationState &result,
AffineMap permutationMap,
std::optional<ArrayRef<bool>> inBounds) {
auto permutationMapAttr = AffineMapAttr::get(permutationMap);
auto inBoundsAttr = (inBounds && !inBounds.value().empty())
? builder.getBoolArrayAttr(inBounds.value())
: ArrayAttr();
auto inBoundsAttr =
(inBounds && !inBounds.value().empty())
? builder.getBoolArrayAttr(inBounds.value())
: builder.getBoolArrayAttr(SmallVector<bool>(
llvm::cast<VectorType>(vector.getType()).getRank(), false));
build(builder, result, vector, dest, indices, permutationMapAttr,
/*mask=*/Value(), inBoundsAttr);
}
Expand Down Expand Up @@ -4410,6 +4418,13 @@ ParseResult TransferWriteOp::parse(OpAsmParser &parser,
} else {
permMap = llvm::cast<AffineMapAttr>(permMapAttr).getValue();
}
auto inBoundsAttrName = TransferWriteOp::getInBoundsAttrName(result.name);
Attribute inBoundsAttr = result.attributes.get(inBoundsAttrName);
if (!inBoundsAttr) {
result.addAttribute(inBoundsAttrName,
builder.getBoolArrayAttr(
SmallVector<bool>(permMap.getNumResults(), false)));
}
if (parser.resolveOperand(vectorInfo, vectorType, result.operands) ||
parser.resolveOperand(sourceInfo, shapedType, result.operands) ||
parser.resolveOperands(indexInfo, indexType, result.operands))
Expand Down Expand Up @@ -4463,8 +4478,7 @@ LogicalResult TransferWriteOp::verify() {

if (failed(verifyTransferOp(cast<VectorTransferOpInterface>(getOperation()),
shapedType, vectorType, maskType,
inferredMaskType, permutationMap,
getInBounds() ? *getInBounds() : ArrayAttr())))
inferredMaskType, permutationMap, getInBounds())))
return failure();

return verifyPermutationMap(permutationMap,
Expand Down
4 changes: 2 additions & 2 deletions mlir/lib/Dialect/Vector/Transforms/LowerVectorMask.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -224,7 +224,7 @@ struct MaskedTransferReadOpPattern
rewriter.replaceOpWithNewOp<TransferReadOp>(
maskingOp.getOperation(), readOp.getVectorType(), readOp.getSource(),
readOp.getIndices(), readOp.getPermutationMap(), readOp.getPadding(),
maskingOp.getMask(), readOp.getInBounds().value_or(ArrayAttr()));
maskingOp.getMask(), readOp.getInBounds());
return success();
}
};
Expand All @@ -246,7 +246,7 @@ struct MaskedTransferWriteOpPattern
rewriter.replaceOpWithNewOp<TransferWriteOp>(
maskingOp.getOperation(), resultType, writeOp.getVector(),
writeOp.getSource(), writeOp.getIndices(), writeOp.getPermutationMap(),
maskingOp.getMask(), writeOp.getInBounds().value_or(ArrayAttr()));
maskingOp.getMask(), writeOp.getInBounds());
return success();
}
};
Expand Down
8 changes: 2 additions & 6 deletions mlir/lib/Dialect/Vector/Transforms/LowerVectorTransfer.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -133,9 +133,7 @@ struct TransferReadPermutationLowering

// Transpose in_bounds attribute.
ArrayAttr newInBoundsAttr =
op.getInBounds() ? inverseTransposeInBoundsAttr(
rewriter, op.getInBounds().value(), permutation)
: ArrayAttr();
inverseTransposeInBoundsAttr(rewriter, op.getInBounds(), permutation);

// Generate new transfer_read operation.
VectorType newReadType = VectorType::get(
Expand Down Expand Up @@ -208,9 +206,7 @@ struct TransferWritePermutationLowering

// Transpose in_bounds attribute.
ArrayAttr newInBoundsAttr =
op.getInBounds() ? inverseTransposeInBoundsAttr(
rewriter, op.getInBounds().value(), permutation)
: ArrayAttr();
inverseTransposeInBoundsAttr(rewriter, op.getInBounds(), permutation);

// Generate new transfer_write operation.
Value newVec = rewriter.create<vector::TransposeOp>(
Expand Down
13 changes: 13 additions & 0 deletions mlir/lib/IR/AffineMap.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -158,6 +158,19 @@ bool AffineMap::isMinorIdentity() const {
getMinorIdentityMap(getNumDims(), getNumResults(), getContext());
}

SmallVector<unsigned> AffineMap::getBroadcastDims() const {
SmallVector<unsigned> broadcastedDims;
for (const auto &[resIdx, expr] : llvm::enumerate(getResults())) {
if (auto constExpr = dyn_cast<AffineConstantExpr>(expr)) {
if (constExpr.getValue() != 0)
continue;
broadcastedDims.push_back(resIdx);
}
}

return broadcastedDims;
}

/// Returns true if this affine map is a minor identity up to broadcasted
/// dimensions which are indicated by value 0 in the result.
bool AffineMap::isMinorIdentityWithBroadcasting(
Expand Down
6 changes: 3 additions & 3 deletions mlir/test/Conversion/VectorToSCF/vector-to-scf.mlir
Original file line number Diff line number Diff line change
Expand Up @@ -133,7 +133,7 @@ func.func @materialize_read(%M: index, %N: index, %O: index, %P: index) {
affine.for %i1 = 0 to %N {
affine.for %i2 = 0 to %O {
affine.for %i3 = 0 to %P step 5 {
%f = vector.transfer_read %A[%i0, %i1, %i2, %i3], %f0 {permutation_map = affine_map<(d0, d1, d2, d3) -> (d3, 0, d0)>} : memref<?x?x?x?xf32>, vector<5x4x3xf32>
%f = vector.transfer_read %A[%i0, %i1, %i2, %i3], %f0 {in_bounds = [false, true, false], permutation_map = affine_map<(d0, d1, d2, d3) -> (d3, 0, d0)>} : memref<?x?x?x?xf32>, vector<5x4x3xf32>
// Add a dummy use to prevent dead code elimination from removing
// transfer read ops.
"dummy_use"(%f) : (vector<5x4x3xf32>) -> ()
Expand Down Expand Up @@ -507,7 +507,7 @@ func.func @transfer_read_with_tensor(%arg: tensor<f32>) -> vector<1xf32> {
// CHECK-NEXT: %[[RESULT:.*]] = vector.broadcast %[[EXTRACTED]] : f32 to vector<1xf32>
// CHECK-NEXT: return %[[RESULT]] : vector<1xf32>
%f0 = arith.constant 0.0 : f32
%0 = vector.transfer_read %arg[], %f0 {permutation_map = affine_map<()->(0)>} :
%0 = vector.transfer_read %arg[], %f0 {in_bounds = [true], permutation_map = affine_map<()->(0)>} :
tensor<f32>, vector<1xf32>
return %0: vector<1xf32>
}
Expand Down Expand Up @@ -746,7 +746,7 @@ func.func @cannot_lower_transfer_read_with_leading_scalable(%arg0: memref<?x4xf3
func.func @does_not_crash_on_unpack_one_dim(%subview: memref<1x1x1x1xi32>, %mask: vector<1x1xi1>) -> vector<1x1x1x1xi32> {
%c0 = arith.constant 0 : index
%c0_i32 = arith.constant 0 : i32
%3 = vector.transfer_read %subview[%c0, %c0, %c0, %c0], %c0_i32, %mask {permutation_map = #map1}
%3 = vector.transfer_read %subview[%c0, %c0, %c0, %c0], %c0_i32, %mask {in_bounds = [false, true, true, false], permutation_map = #map1}
: memref<1x1x1x1xi32>, vector<1x1x1x1xi32>
return %3 : vector<1x1x1x1xi32>
}
Expand Down
Loading
Loading