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

[ONNX] Fix issue with absent value in onnx.ConstantOfShape (#3713) 

Previously, if the value was absent, this conversion was creating a
dense resource of value 0 with shape equal to the result shape, then
later re-extracting a splat value. This only works if the shape is
statically known, and even when the shape is known, this is completely
unnecessary since the value's shape should be `[1]` and not the result
shape.

This patch simply sets the `splatvalue` to a `torch.constant.float 0.0`
when the onnx op's `value` attr is absent, and adds `nullptr` checks to
the subsequent conditionals to avoid them in the case where an `attr` is
not given.

Addresses <https://github.com/nod-ai/SHARK-Turbine/issues/831>.
pull/3716/head
zjgarvey 2024-09-17 14:01:01 -07:00 committed by GitHub
parent 14ef05a292
commit d2c387dd04
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 26 additions and 9 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

17
lib/Conversion/TorchOnnxToTorch/DefaultDomainAtoF.cpp
View File

@ -2802,14 +2802,14 @@ void mlir::torch::onnx_c::populateDefaultDomainAtoF(
// Get fill_value if it is present.
// Assumption : resultDType and value attr type match.
auto attr = binder.op->getAttr("torch.onnx.value");
auto resultDType = resultType.getDtype();
// Extract the fill value and dtype
// ONNX requires value attr to be a tensor
Value splatvalue;
// if no value attr is provided, default is 0.0 float value
if (!attr) {
attr =
DenseElementsAttr::get(resultType.toBuiltinTensor(),
rewriter.getFloatAttr(resultDType, 0.0));
splatvalue = rewriter.create<Torch::ConstantFloatOp>(
binder.getLoc(), rewriter.getF64FloatAttr(0.0));
}
// If its a dense resource attr we need to convert to a dense type:
@ -2830,19 +2830,18 @@ void mlir::torch::onnx_c::populateDefaultDomainAtoF(
}
Attribute splattr;
if (isa<SplatElementsAttr>(attr)) {
if (attr && isa<SplatElementsAttr>(attr)) {
auto denseAttr = cast<DenseElementsAttr>(attr);
splattr = denseAttr.getSplatValue<Attribute>();
}
if (!isa<FloatAttr, IntegerAttr>(splattr)) {
if (splattr && !isa<FloatAttr, IntegerAttr>(splattr)) {
return rewriter.notifyMatchFailure(
binder.op,
"`value` attr tensor only supports types int and float for now.");
}
Value splatvalue;
if (auto intattr = dyn_cast<IntegerAttr>(splattr)) {
if (auto intattr = dyn_cast_or_null<IntegerAttr>(splattr)) {
IntegerType intty = cast<IntegerType>(intattr.getType());
int64_t value;
if (intty.isUnsignedInteger()) {
@ -2856,7 +2855,7 @@ void mlir::torch::onnx_c::populateDefaultDomainAtoF(
rewriter.create<Torch::ConstantIntOp>(binder.getLoc(), value);
}
if (auto fpattr = dyn_cast<FloatAttr>(splattr))
if (auto fpattr = dyn_cast_or_null<FloatAttr>(splattr))
splatvalue = rewriter.create<Torch::ConstantFloatOp>(
binder.getLoc(),
rewriter.getF64FloatAttr(fpattr.getValueAsDouble()));

18
test/Conversion/TorchOnnxToTorch/simple_ops_a_to_f.mlir
View File

@ -2016,6 +2016,24 @@ func.func @test_flatten_1d_axis_1(%arg0: !torch.vtensor<[2],f32>) -> !torch.vten
// -----
// CHECK-LABEL: func.func @test_constant_of_shape_arg_input
func.func @test_constant_of_shape_arg_input(%arg0: !torch.vtensor<[2], si64>) -> !torch.vtensor<[?,?], f32> attributes {torch.onnx_meta.ir_version = 9 : si64, torch.onnx_meta.opset_version = 20 : si64, torch.onnx_meta.producer_name = "backend-test", torch.onnx_meta.producer_version = ""} {
// CHECK: %[[INT0:.*]] = torch.constant.int 0
// CHECK: %[[INT0_0:.*]] = torch.constant.int 0
// CHECK: %[[EXTRACT_0:.*]] = torch.aten.select.int %arg0, %[[INT0]], %[[INT0_0]] : !torch.vtensor<[2],si64>, !torch.int, !torch.int -> !torch.vtensor<[],si64>
// CHECK: %[[ELE_0:.*]] = torch.aten.item %[[EXTRACT_0]] : !torch.vtensor<[],si64> -> !torch.int
// CHECK: %[[INT1:.*]] = torch.constant.int 1
// CHECK: %[[EXTRACT_1:.*]] = torch.aten.select.int %arg0, %[[INT0]], %[[INT1]] : !torch.vtensor<[2],si64>, !torch.int, !torch.int -> !torch.vtensor<[],si64>
// CHECK: %[[ELE_1:.*]] = torch.aten.item %[[EXTRACT_1]] : !torch.vtensor<[],si64> -> !torch.int
// CHECK: %[[DIM_LIST:.*]] = torch.prim.ListConstruct %[[ELE_0]], %[[ELE_1]] : (!torch.int, !torch.int) -> !torch.list<int>
// CHECK: %[[NONE:.*]] = torch.constant.none
// CHECK: %[[FILL_VAL:.*]] = torch.constant.float 0.000000e+00
// CHECK: %[[ATEN_FULL:.*]] = torch.aten.full %[[DIM_LIST]], %[[FILL_VAL]], %[[NONE]], %[[NONE]], %[[NONE]], %[[NONE]] : !torch.list<int>, !torch.float, !torch.none, !torch.none, !torch.none, !torch.none -> !torch.vtensor<[?,?],f32>
%0 = "torch.operator"(%arg0) <{name = "onnx.ConstantOfShape"}> : (!torch.vtensor<[2], si64>) -> !torch.vtensor<[?,?], f32>
return %0 : !torch.vtensor<[?,?], f32>
}
// -----
// CHECK-LABEL: func.func @test_constant_of_shape_dense_float_default
func.func @test_constant_of_shape_dense_float_default() -> !torch.vtensor<[2,3,4], f32> attributes {torch.onnx_meta.ir_version = 9 : si64, torch.onnx_meta.opset_version = 20 : si64, torch.onnx_meta.producer_name = "backend-test", torch.onnx_meta.producer_version = ""} {
// CHECK: %[[SHAPE_CST:.*]] = torch.vtensor.literal(dense<[2, 3, 4]> : tensor<3xsi64>) : !torch.vtensor<[3],si64>