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

[onnx] Fix onnx.Shape lowering with scalar input (#3716) 

Address https://github.com/nod-ai/SHARK-Turbine/issues/826
pull/3740/head
jinchen 2024-09-27 13:30:02 -07:00 committed by GitHub
parent 9938abf25e
commit a33d1232c5
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 17 additions and 8 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

16
lib/Conversion/TorchOnnxToTorch/DefaultDomainQtoZ.cpp
View File

@ -1662,10 +1662,15 @@ void mlir::torch::onnx_c::populateDefaultDomainQtoZ(
auto shapeType = Torch::ValueTensorType::get(
binder.op->getContext(), SmallVector<int64_t>{inputRank},
resultType.getOptionalDtype());
Value shape = rewriter.create<Torch::Aten_ShapeAsTensorOp>(
binder.getLoc(), shapeType, operand);
if (inputRank == 0) {
rewriter.replaceOpWithNewOp<Torch::TensorStaticInfoCastOp>(
binder.op, resultType, shape);
return success();
}
if (start == 0 && end == -1) {
rewriter.replaceOp(binder.op, shape);
return success();
@ -1673,18 +1678,13 @@ void mlir::torch::onnx_c::populateDefaultDomainQtoZ(
Value sv = rewriter.create<Torch::ConstantIntOp>(
binder.getLoc(), rewriter.getI64IntegerAttr(start));
Value ev = rewriter.create<Torch::ConstantIntOp>(
binder.getLoc(), rewriter.getI64IntegerAttr(end));
Value step = rewriter.create<Torch::ConstantIntOp>(binder.getLoc(), 1);
Value dim = rewriter.create<Torch::ConstantIntOp>(binder.getLoc(), 0);
shape = rewriter.create<Torch::AtenSliceTensorOp>(
binder.getLoc(), resultType, shape, dim, sv, ev, step);
rewriter.replaceOp(binder.op, shape);
rewriter.replaceOpWithNewOp<Torch::AtenSliceTensorOp>(
binder.op, resultType, shape, dim, sv, ev, step);
return success();
});

9
test/Conversion/TorchOnnxToTorch/simple_ops_q_to_z.mlir
View File

@ -2833,6 +2833,15 @@ func.func @test_shape_start_1_end_negative_1(%arg0: !torch.vtensor<[3,4,5],f32>)
return %0 : !torch.vtensor<[1],si64>
}
// -----
// CHECK-LABEL: func.func @test_shape_scalar
func.func @test_shape_scalar(%arg0: !torch.vtensor<[],si64> ) -> !torch.vtensor<[?],si64> attributes {torch.onnx_meta.ir_version = 8 : si64, torch.onnx_meta.opset_version = 17 : si64, torch.onnx_meta.producer_name = "pytorch", torch.onnx_meta.producer_version = "2.1.0"} {
// CHECK: %[[SHAPE:.+]] = torch.aten._shape_as_tensor %arg0 : !torch.vtensor<[],si64> -> !torch.vtensor<[0],si64>
// CHECK: %[[CAST:.+]] = torch.tensor_static_info_cast %[[SHAPE]] : !torch.vtensor<[0],si64> to !torch.vtensor<[?],si64>
%0 = torch.operator "onnx.Shape"(%arg0) : (!torch.vtensor<[],si64>) -> !torch.vtensor<[?],si64>
return %0: !torch.vtensor<[?],si64>
}
// -----