Skip to content

Commit f0bf455

Browse files
banach-spacebanach-space
committed
Update the internal representation of in_bounds
This PR updates the internal representation of the `in_bounds` attribute for `xfer_read`/`xfer_write` Ops. Currently we use `ArrayAttr` - that's being updated to `DenseBoolArrayAttribute`. Note that this means that the asm format of the `xfer_{read|_write}` will change from: ```mlir vector.transfer_read %arg0[%0, %1], %cst {in_bounds = [true], permutation_map = #map3} : memref<12x16xf32>, vector<8xf32> ``` to: ```mlir vector.transfer_read %arg0[%0, %1], %cst {in_bounds = array<i1: true>, permutation_map = #map3} : memref<12x16xf32>, vector<8xf32> ```
1 parent f5c02dd commit f0bf455

File tree

82 files changed

+1149
-1169
lines changed

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

82 files changed

+1149
-1169
lines changed

‎mlir/include/mlir/Dialect/Vector/IR/VectorOps.td

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1248,7 +1248,7 @@ def Vector_TransferReadOp :
12481248
AffineMapAttr:$permutation_map,
12491249
AnyType:$padding,
12501250
Optional<VectorOf<[I1]>>:$mask,
1251-
BoolArrayAttr:$in_bounds)>,
1251+
DenseBoolArrayAttr:$in_bounds)>,
12521252
Results<(outs AnyVectorOfAnyRank:$vector)> {
12531253

12541254
let summary = "Reads a supervector from memory into an SSA vector value.";
@@ -1443,7 +1443,7 @@ def Vector_TransferReadOp :
14431443
"Value":$source,
14441444
"ValueRange":$indices,
14451445
"AffineMapAttr":$permutationMapAttr,
1446-
"ArrayAttr":$inBoundsAttr)>,
1446+
"DenseBoolArrayAttr":$inBoundsAttr)>,
14471447
/// 2. Builder that sets padding to zero and an empty mask (variant without attrs).
14481448
OpBuilder<(ins "VectorType":$vectorType,
14491449
"Value":$source,
@@ -1495,7 +1495,7 @@ def Vector_TransferWriteOp :
14951495
Variadic<Index>:$indices,
14961496
AffineMapAttr:$permutation_map,
14971497
Optional<VectorOf<[I1]>>:$mask,
1498-
BoolArrayAttr:$in_bounds)>,
1498+
DenseBoolArrayAttr:$in_bounds)>,
14991499
Results<(outs Optional<AnyRankedTensor>:$result)> {
15001500

15011501
let summary = "The vector.transfer_write op writes a supervector to memory.";
@@ -1606,13 +1606,13 @@ def Vector_TransferWriteOp :
16061606
"ValueRange":$indices,
16071607
"AffineMapAttr":$permutationMapAttr,
16081608
"Value":$mask,
1609-
"ArrayAttr":$inBoundsAttr)>,
1609+
"DenseBoolArrayAttr":$inBoundsAttr)>,
16101610
/// 2. Builder with type inference that sets an empty mask (variant with attrs).
16111611
OpBuilder<(ins "Value":$vector,
16121612
"Value":$dest,
16131613
"ValueRange":$indices,
16141614
"AffineMapAttr":$permutationMapAttr,
1615-
"ArrayAttr":$inBoundsAttr)>,
1615+
"DenseBoolArrayAttr":$inBoundsAttr)>,
16161616
/// 3. Builder with type inference that sets an empty mask (variant without attrs).
16171617
OpBuilder<(ins "Value":$vector,
16181618
"Value":$dest,

‎mlir/include/mlir/Interfaces/VectorInterfaces.td

Lines changed: 2 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -98,7 +98,7 @@ def VectorTransferOpInterface : OpInterface<"VectorTransferOpInterface"> {
9898
dimension whether it is in-bounds or not. (Broadcast dimensions are
9999
always in-bounds).
100100
}],
101-
/*retTy=*/"::mlir::ArrayAttr",
101+
/*retTy=*/"::mlir::ArrayRef<bool>",
102102
/*methodName=*/"getInBounds",
103103
/*args=*/(ins)
104104
>,
@@ -169,15 +169,6 @@ def VectorTransferOpInterface : OpInterface<"VectorTransferOpInterface"> {
169169
];
170170

171171
let extraSharedClassDeclaration = [{
172-
/// Return a vector of all in_bounds values as booleans (one per vector
173-
/// transfer dimension).
174-
::llvm::SmallVector<bool> getInBoundsValues() {
175-
::llvm::SmallVector<bool> inBounds;
176-
for (int64_t i = 0, e = $_op.getTransferRank(); i < e; ++i)
177-
inBounds.push_back($_op.isDimInBounds(i));
178-
return inBounds;
179-
}
180-
181172
/// Return the number of leading shaped dimensions (of the "source" operand)
182173
/// that do not participate in the permutation map.
183174
unsigned getLeadingShapedRank() {
@@ -241,7 +232,7 @@ def VectorTransferOpInterface : OpInterface<"VectorTransferOpInterface"> {
241232
if ($_op.isBroadcastDim(dim))
242233
return true;
243234
auto inBounds = $_op.getInBounds();
244-
return ::llvm::cast<::mlir::BoolAttr>(inBounds[dim]).getValue();
235+
return inBounds[dim];
245236
}
246237

247238
/// Helper function to account for the fact that `permutationMap` results

‎mlir/lib/Conversion/VectorToSCF/VectorToSCF.cpp

Lines changed: 7 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -263,13 +263,6 @@ static void generateInBoundsCheck(
263263
});
264264
}
265265

266-
/// Given an ArrayAttr, return a copy where the first element is dropped.
267-
static ArrayAttr dropFirstElem(OpBuilder &b, ArrayAttr attr) {
268-
if (!attr)
269-
return attr;
270-
return ArrayAttr::get(b.getContext(), attr.getValue().drop_front());
271-
}
272-
273266
/// Add the pass label to a vector transfer op if its rank is not the target
274267
/// rank.
275268
template <typename OpTy>
@@ -424,11 +417,11 @@ struct Strategy<TransferReadOp> {
424417
Location loc = xferOp.getLoc();
425418
auto bufferType = dyn_cast<ShapedType>(buffer.getType());
426419
auto vecType = dyn_cast<VectorType>(bufferType.getElementType());
427-
auto inBoundsAttr = dropFirstElem(b, xferOp.getInBoundsAttr());
428420
auto newXferOp = b.create<vector::TransferReadOp>(
429421
loc, vecType, xferOp.getSource(), xferIndices,
430422
AffineMapAttr::get(unpackedPermutationMap(b, xferOp)),
431-
xferOp.getPadding(), Value(), inBoundsAttr);
423+
xferOp.getPadding(), Value(),
424+
b.getDenseBoolArrayAttr(xferOp.getInBounds().drop_front()));
432425

433426
maybeApplyPassLabel(b, newXferOp, options.targetRank);
434427

@@ -511,13 +504,12 @@ struct Strategy<TransferWriteOp> {
511504

512505
Location loc = xferOp.getLoc();
513506
auto vec = b.create<memref::LoadOp>(loc, buffer, loadIndices);
514-
auto inBoundsAttr = dropFirstElem(b, xferOp.getInBoundsAttr());
515507
auto source = loopState.empty() ? xferOp.getSource() : loopState[0];
516508
Type type = isTensorOp(xferOp) ? xferOp.getShapedType() : Type();
517509
auto newXferOp = b.create<vector::TransferWriteOp>(
518510
loc, type, vec, source, xferIndices,
519511
AffineMapAttr::get(unpackedPermutationMap(b, xferOp)), Value(),
520-
inBoundsAttr);
512+
b.getDenseBoolArrayAttr(xferOp.getInBounds().drop_front()));
521513

522514
maybeApplyPassLabel(b, newXferOp, options.targetRank);
523515

@@ -1160,7 +1152,7 @@ struct ScalableTransposeTransferWriteConversion
11601152
loopIterArgs.empty() ? writeOp.getSource() : loopIterArgs.front();
11611153
auto newWriteOp = b.create<vector::TransferWriteOp>(
11621154
loc, sliceVec, dest, xferIndices,
1163-
ArrayRef<bool>(writeOp.getInBoundsValues()).drop_front());
1155+
writeOp.getInBounds().drop_front());
11641156
if (sliceMask)
11651157
newWriteOp.getMaskMutable().assign(sliceMask);
11661158

@@ -1332,11 +1324,11 @@ struct UnrollTransferReadConversion
13321324
getInsertionIndices(xferOp, insertionIndices);
13331325
insertionIndices.push_back(rewriter.getIndexAttr(i));
13341326

1335-
auto inBoundsAttr = dropFirstElem(b, xferOp.getInBoundsAttr());
13361327
auto newXferOp = b.create<vector::TransferReadOp>(
13371328
loc, newXferVecType, xferOp.getSource(), xferIndices,
13381329
AffineMapAttr::get(unpackedPermutationMap(b, xferOp)),
1339-
xferOp.getPadding(), Value(), inBoundsAttr);
1330+
xferOp.getPadding(), Value(),
1331+
b.getDenseBoolArrayAttr(xferOp.getInBounds().drop_front()));
13401332
maybeAssignMask(b, xferOp, newXferOp, i);
13411333
return b.create<vector::InsertOp>(loc, newXferOp, vec,
13421334
insertionIndices);
@@ -1467,7 +1459,6 @@ struct UnrollTransferWriteConversion
14671459

14681460
auto extracted =
14691461
b.create<vector::ExtractOp>(loc, vec, extractionIndices);
1470-
auto inBoundsAttr = dropFirstElem(b, xferOp.getInBoundsAttr());
14711462
Value xferVec;
14721463
if (inputVectorTy.getRank() == 1) {
14731464
// When target-rank=0, unrolling would causes the vector input
@@ -1481,7 +1472,7 @@ struct UnrollTransferWriteConversion
14811472
auto newXferOp = b.create<vector::TransferWriteOp>(
14821473
loc, sourceType, xferVec, source, xferIndices,
14831474
AffineMapAttr::get(unpackedPermutationMap(b, xferOp)), Value(),
1484-
inBoundsAttr);
1475+
b.getDenseBoolArrayAttr(xferOp.getInBounds().drop_front()));
14851476

14861477
maybeAssignMask(b, xferOp, newXferOp, i);
14871478

‎mlir/lib/Dialect/ArmSME/Transforms/VectorLegalization.cpp

Lines changed: 5 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -497,8 +497,7 @@ struct LegalizeMultiTileTransferWriteAsStoreLoop
497497
loc, slice, writeOp.getSource(), ValueRange{storeRow, storeCol},
498498
AffineMapAttr::get(writeOp.getPermutationMap().dropResult(0)),
499499
sliceMask,
500-
rewriter.getBoolArrayAttr(
501-
ArrayRef<bool>(writeOp.getInBoundsValues()).drop_front()));
500+
rewriter.getDenseBoolArrayAttr(writeOp.getInBounds().drop_front()));
502501
}
503502

504503
rewriter.eraseOp(writeOp);
@@ -691,13 +690,12 @@ struct LiftIllegalVectorTransposeToMemory
691690
transposeOp.getPermutation(), getContext());
692691
auto transposedSubview = rewriter.create<memref::TransposeOp>(
693692
loc, readSubview, AffineMapAttr::get(transposeMap));
694-
ArrayAttr inBoundsAttr = illegalRead.getInBoundsAttr();
693+
DenseBoolArrayAttr inBoundsAttr = illegalRead.getInBoundsAttr();
695694
// - The `in_bounds` attribute
696695
if (inBoundsAttr) {
697-
SmallVector<Attribute> inBoundsValues(inBoundsAttr.begin(),
698-
inBoundsAttr.end());
696+
SmallVector<bool> inBoundsValues(inBoundsAttr.asArrayRef());
699697
applyPermutationToVector(inBoundsValues, transposeOp.getPermutation());
700-
inBoundsAttr = rewriter.getArrayAttr(inBoundsValues);
698+
inBoundsAttr = rewriter.getDenseBoolArrayAttr(inBoundsValues);
701699
}
702700

703701
VectorType legalReadType = resultType.clone(readType.getElementType());
@@ -902,7 +900,7 @@ struct LowerIllegalTransposeStoreViaZA
902900
rewriter.create<arith::AddIOp>(loc, transposedCol, writeIndices[1]);
903901
auto smeWrite = rewriter.create<vector::TransferWriteOp>(
904902
loc, tile, destTensorOrMemref, ValueRange{destRow, destCol},
905-
transposeMap, subMask, writeOp.getInBounds());
903+
transposeMap, subMask, writeOp.getInBoundsAttr());
906904

907905
if (writeOp.hasPureTensorSemantics())
908906
destTensorOrMemref = smeWrite.getResult();

‎mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -659,7 +659,7 @@ static Value buildVectorWrite(RewriterBase &rewriter, Value value,
659659
if (auto maskOp = dyn_cast<vector::MaskingOpInterface>(write)) {
660660
auto maskedWriteOp = cast<vector::TransferWriteOp>(maskOp.getMaskableOp());
661661
SmallVector<bool> inBounds(maskedWriteOp.getVectorType().getRank(), true);
662-
maskedWriteOp.setInBoundsAttr(rewriter.getBoolArrayAttr(inBounds));
662+
maskedWriteOp.setInBoundsAttr(rewriter.getDenseBoolArrayAttr(inBounds));
663663
}
664664

665665
LDBG("vectorized op: " << *write << "\n");
@@ -1399,7 +1399,7 @@ vectorizeAsLinalgGeneric(RewriterBase &rewriter, VectorizationState &state,
13991399
if (auto maskOp = dyn_cast<vector::MaskingOpInterface>(read)) {
14001400
SmallVector<bool> inBounds(readType.getRank(), true);
14011401
cast<vector::TransferReadOp>(maskOp.getMaskableOp())
1402-
.setInBoundsAttr(rewriter.getBoolArrayAttr(inBounds));
1402+
.setInBoundsAttr(rewriter.getDenseBoolArrayAttr(inBounds));
14031403
}
14041404

14051405
// 3.c. Not all ops support 0-d vectors, extract the scalar for now.
@@ -2432,7 +2432,7 @@ struct PadOpVectorizationWithTransferReadPattern
24322432
rewriter.modifyOpInPlace(xferOp, [&]() {
24332433
SmallVector<bool> inBounds(xferOp.getVectorType().getRank(), false);
24342434
xferOp->setAttr(xferOp.getInBoundsAttrName(),
2435-
rewriter.getBoolArrayAttr(inBounds));
2435+
rewriter.getDenseBoolArrayAttr(inBounds));
24362436
xferOp.getSourceMutable().assign(padOp.getSource());
24372437
xferOp.getPaddingMutable().assign(padValue);
24382438
});
@@ -2511,7 +2511,7 @@ struct PadOpVectorizationWithTransferWritePattern
25112511
auto newXferOp = rewriter.replaceOpWithNewOp<vector::TransferWriteOp>(
25122512
xferOp, padOp.getSource().getType(), xferOp.getVector(),
25132513
padOp.getSource(), xferOp.getIndices(), xferOp.getPermutationMapAttr(),
2514-
xferOp.getMask(), rewriter.getBoolArrayAttr(inBounds));
2514+
xferOp.getMask(), rewriter.getDenseBoolArrayAttr(inBounds));
25152515
rewriter.replaceOp(trimPadding, newXferOp->getResult(0));
25162516

25172517
return success();
@@ -2815,7 +2815,7 @@ LogicalResult LinalgCopyVTRForwardingPattern::matchAndRewrite(
28152815
Value res = rewriter.create<vector::TransferReadOp>(
28162816
xferOp.getLoc(), vectorType, in, xferOp.getIndices(),
28172817
xferOp.getPermutationMapAttr(), xferOp.getPadding(), xferOp.getMask(),
2818-
rewriter.getBoolArrayAttr(
2818+
rewriter.getDenseBoolArrayAttr(
28192819
SmallVector<bool>(vectorType.getRank(), false)));
28202820

28212821
if (maybeFillOp)
@@ -2874,7 +2874,7 @@ LogicalResult LinalgCopyVTWForwardingPattern::matchAndRewrite(
28742874
rewriter.create<vector::TransferWriteOp>(
28752875
xferOp.getLoc(), vector, out, xferOp.getIndices(),
28762876
xferOp.getPermutationMapAttr(), xferOp.getMask(),
2877-
rewriter.getBoolArrayAttr(
2877+
rewriter.getDenseBoolArrayAttr(
28782878
SmallVector<bool>(vector.getType().getRank(), false)));
28792879

28802880
rewriter.eraseOp(copyOp);
@@ -3381,7 +3381,7 @@ struct Conv1DGenerator
33813381
SmallVector<bool> inBounds(maskShape.size(), true);
33823382
auto xferOp = cast<VectorTransferOpInterface>(opToMask);
33833383
xferOp->setAttr(xferOp.getInBoundsAttrName(),
3384-
rewriter.getBoolArrayAttr(inBounds));
3384+
rewriter.getDenseBoolArrayAttr(inBounds));
33853385

33863386
SmallVector<OpFoldResult> mixedDims = vector::getMixedSizesXfer(
33873387
cast<LinalgOp>(op).hasPureTensorSemantics(), opToMask, rewriter);

0 commit comments

Comments
 (0)