File tree Expand file tree Collapse file tree 1 file changed +5
-9
lines changed
mlir/lib/Dialect/Shape/IR Expand file tree Collapse file tree 1 file changed +5
-9
lines changed Original file line number Diff line number Diff line change @@ -47,19 +47,15 @@ bool shape::isExtentTensorType(Type type) {
47
47
LogicalResult shape::getShapeVec (Value input,
48
48
SmallVectorImpl<int64_t > &shapeValues) {
49
49
if (auto inputOp = input.getDefiningOp <ShapeOfOp>()) {
50
- auto type = inputOp.getArg ().getType ().dyn_cast <ShapedType>();
50
+ auto type = inputOp.getArg ().getType ().cast <ShapedType>();
51
51
if (!type.hasRank ())
52
52
return failure ();
53
- shapeValues = llvm::to_vector< 6 >( type.getShape ());
53
+ llvm::append_range (shapeValues, type.getShape ());
54
54
return success ();
55
55
}
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 >());
63
59
return success ();
64
60
}
65
61
return failure ();
You can’t perform that action at this time.
0 commit comments