-
Notifications
You must be signed in to change notification settings - Fork 15.2k
[mlir][VectorOps] Extend vector.constant_mask to support 'all true' scalable dims #66638
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
Changes from all commits
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -105,7 +105,6 @@ class ConstantMaskOpLowering : public OpRewritePattern<vector::ConstantMaskOp> { | |
PatternRewriter &rewriter) const override { | ||
auto loc = op.getLoc(); | ||
auto dstType = op.getType(); | ||
auto eltType = dstType.getElementType(); | ||
auto dimSizes = op.getMaskDimSizes(); | ||
int64_t rank = dstType.getRank(); | ||
|
||
|
@@ -115,43 +114,43 @@ class ConstantMaskOpLowering : public OpRewritePattern<vector::ConstantMaskOp> { | |
bool value = cast<IntegerAttr>(dimSizes[0]).getInt() == 1; | ||
rewriter.replaceOpWithNewOp<arith::ConstantOp>( | ||
op, dstType, | ||
DenseIntElementsAttr::get( | ||
VectorType::get(ArrayRef<int64_t>{}, rewriter.getI1Type()), | ||
ArrayRef<bool>{value})); | ||
DenseIntElementsAttr::get(VectorType::get({}, rewriter.getI1Type()), | ||
value)); | ||
return success(); | ||
} | ||
|
||
// Scalable constant masks can only be lowered for the "none set" case. | ||
if (cast<VectorType>(dstType).isScalable()) { | ||
rewriter.replaceOpWithNewOp<arith::ConstantOp>( | ||
op, DenseElementsAttr::get(dstType, false)); | ||
return success(); | ||
} | ||
|
||
int64_t trueDim = std::min(dstType.getDimSize(0), | ||
cast<IntegerAttr>(dimSizes[0]).getInt()); | ||
int64_t trueDimSize = cast<IntegerAttr>(dimSizes[0]).getInt(); | ||
|
||
if (rank == 1) { | ||
// Express constant 1-D case in explicit vector form: | ||
// [T,..,T,F,..,F]. | ||
SmallVector<bool> values(dstType.getDimSize(0)); | ||
for (int64_t d = 0; d < trueDim; d++) | ||
values[d] = true; | ||
rewriter.replaceOpWithNewOp<arith::ConstantOp>( | ||
op, dstType, rewriter.getBoolVectorAttr(values)); | ||
if (trueDimSize == 0 || trueDimSize == dstType.getDimSize(0)) { | ||
// Use constant splat for 'all set' or 'none set' dims. | ||
// This produces correct code for scalable dimensions (it will lower to | ||
// a constant splat). | ||
rewriter.replaceOpWithNewOp<arith::ConstantOp>( | ||
op, DenseElementsAttr::get(dstType, trueDimSize != 0)); | ||
} else { | ||
// Express constant 1-D case in explicit vector form: | ||
// [T,..,T,F,..,F]. | ||
// Note: The verifier would reject this case for scalable vectors. | ||
SmallVector<bool> values(dstType.getDimSize(0), false); | ||
for (int64_t d = 0; d < trueDimSize; d++) | ||
values[d] = true; | ||
rewriter.replaceOpWithNewOp<arith::ConstantOp>( | ||
op, dstType, rewriter.getBoolVectorAttr(values)); | ||
} | ||
return success(); | ||
} | ||
|
||
VectorType lowType = | ||
VectorType::get(dstType.getShape().drop_front(), eltType); | ||
SmallVector<int64_t> newDimSizes; | ||
for (int64_t r = 1; r < rank; r++) | ||
newDimSizes.push_back(cast<IntegerAttr>(dimSizes[r]).getInt()); | ||
if (dstType.getScalableDims().front()) | ||
return rewriter.notifyMatchFailure( | ||
op, "Cannot unroll leading scalable dim in dstType"); | ||
|
||
|
||
VectorType lowType = VectorType::Builder(dstType).dropDim(0); | ||
Value trueVal = rewriter.create<vector::ConstantMaskOp>( | ||
loc, lowType, rewriter.getI64ArrayAttr(newDimSizes)); | ||
loc, lowType, rewriter.getArrayAttr(dimSizes.getValue().drop_front())); | ||
Value result = rewriter.create<arith::ConstantOp>( | ||
loc, dstType, rewriter.getZeroAttr(dstType)); | ||
for (int64_t d = 0; d < trueDim; d++) | ||
for (int64_t d = 0; d < trueDimSize; d++) | ||
result = | ||
rewriter.create<vector::InsertOp>(loc, dstType, trueVal, result, d); | ||
rewriter.replaceOp(op, result); | ||
|
Uh oh!
There was an error while loading. Please reload this page.