diff --git a/mlir/lib/Conversion/ShapeToSCF/ShapeToSCF.cpp b/mlir/lib/Conversion/ShapeToSCF/ShapeToSCF.cpp --- a/mlir/lib/Conversion/ShapeToSCF/ShapeToSCF.cpp +++ b/mlir/lib/Conversion/ShapeToSCF/ShapeToSCF.cpp @@ -185,7 +185,7 @@ auto rankVal = rewriter.create(loc, tensorVal); auto i64Ty = rewriter.getI64Type(); auto memTy = MemRefType::get({ShapedType::kDynamicSize}, i64Ty); - auto memVal = rewriter.create(loc, memTy, ValueRange({rankVal})); + auto memVal = rewriter.create(loc, memTy, ValueRange{rankVal}); // Copy shape extents to stack-allocated memory. auto zeroVal = rewriter.create(loc, 0); @@ -200,11 +200,9 @@ }); // Load extents to tensor value. - auto shapeIntVal = rewriter.create(loc, memVal); - auto indexTy = rewriter.getIndexType(); - auto shapeTy = RankedTensorType::get({ShapedType::kDynamicSize}, indexTy); - rewriter.replaceOpWithNewOp(op.getOperation(), shapeIntVal, - shapeTy); + auto extentTensorInt = rewriter.create(loc, memVal); + rewriter.replaceOpWithNewOp(op.getOperation(), extentTensorInt, + op.getType()); return success(); }