Skip to content

Commit c45c53b

Browse files
committed
[Shape] Simplify getShapeVec a bit. NFCI.
1 parent 634da7a commit c45c53b

File tree

1 file changed

+5
-9
lines changed

1 file changed

+5
-9
lines changed

mlir/lib/Dialect/Shape/IR/Shape.cpp

Lines changed: 5 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -47,19 +47,15 @@ bool shape::isExtentTensorType(Type type) {
4747
LogicalResult shape::getShapeVec(Value input,
4848
SmallVectorImpl<int64_t> &shapeValues) {
4949
if (auto inputOp = input.getDefiningOp<ShapeOfOp>()) {
50-
auto type = inputOp.getArg().getType().dyn_cast<ShapedType>();
50+
auto type = inputOp.getArg().getType().cast<ShapedType>();
5151
if (!type.hasRank())
5252
return failure();
53-
shapeValues = llvm::to_vector<6>(type.getShape());
53+
llvm::append_range(shapeValues, type.getShape());
5454
return success();
5555
}
56-
if (auto inputOp = input.getDefiningOp<ConstShapeOp>()) {
57-
shapeValues = llvm::to_vector<6>(inputOp.getShape().getValues<int64_t>());
58-
return success();
59-
}
60-
if (auto inputOp = input.getDefiningOp<arith::ConstantOp>()) {
61-
shapeValues = llvm::to_vector<6>(
62-
inputOp.getValue().cast<DenseIntElementsAttr>().getValues<int64_t>());
56+
DenseIntElementsAttr attr;
57+
if (matchPattern(input, m_Constant(&attr))) {
58+
llvm::append_range(shapeValues, attr.getValues<int64_t>());
6359
return success();
6460
}
6561
return failure();

0 commit comments

Comments
 (0)