llvm
/
torch-mlir
mirror of https://github.com/llvm/torch-mlir
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Fix TorchToIREE lowering. 

We needed to resize the list, not just reserve capacity.
pull/300/head
Sean Silva 2021-08-31 22:01:50 +00:00
parent 600cc6b9c7
commit ed2afe43e7
2 changed files with 4 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
lib/Conversion/TorchToIREE/TorchToIREE.cpp
View File

@ -35,10 +35,11 @@ public:
matchAndRewrite(PrimListConstructOp op, ArrayRef<Value> operands,
ConversionPatternRewriter &rewriter) const override {
auto type = getTypeConverter()->convertType(op.getType());
auto capacity =
auto size =
rewriter.create<ConstantIndexOp>(op.getLoc(), op->getNumOperands());
auto ireeList =
rewriter.replaceOpWithNewOp<iree::ListCreateOp>(op, type, capacity);
rewriter.replaceOpWithNewOp<iree::ListCreateOp>(op, type, size);
rewriter.create<iree::ListResizeOp>(op.getLoc(), ireeList, size);
for (int i = 0, e = operands.size(); i != e; ++i) {
auto index = rewriter.create<ConstantIndexOp>(op.getLoc(), i);
rewriter.create<iree::ListSetOp>(op.getLoc(), ireeList, index,

1
test/Conversion/TorchToIREE/basic.mlir
View File

@ -7,6 +7,7 @@
// CHECK: %[[ALSO_ARG:.*]] = torch_c.to_f64 %[[ARG_TORCH]]
// CHECK: %[[C2:.*]] = constant 2 : index
// CHECK: %[[LIST:.*]] = iree.list.create %[[C2]] : !iree.list<f64>
// CHECK: iree.list.resize %[[LIST]], %[[C2]] : !iree.list<f64>
// CHECK: %[[C0:.*]] = constant 0 : index
// CHECK: iree.list.set %[[LIST]][%[[C0]]], %[[ARG]] : !iree.list<f64>, f64
// CHECK: %[[C1:.*]] = constant 1 : index