[FxImporter] Replace local_scalar_dense in fx_importer (#3180)

pull/3134/head
penguin_wwy 2024-04-17 22:45:47 +08:00 committed by GitHub
parent d2ba956e69
commit 3aa81f78d8
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 11 additions and 11 deletions

View File

@ -387,7 +387,6 @@ TORCHDYNAMO_CRASHING_SET = {
}
FX_IMPORT_XFAIL_SET = {
'AddIntModule_basic',
'AllBoolFalseModule_basic',
'AllBoolTrueModule_basic',
'AnyBoolFalseModule_basic',
@ -399,10 +398,7 @@ FX_IMPORT_XFAIL_SET = {
'AtenIntBoolOpConstFalseModule_basic',
'AtenIntBoolOpConstTrueModule_basic',
'AtenIntBoolOpModule_basic',
'AtenIntTensorByteDtypeModule_basic',
'AtenIntTensorCharDtypeModule_basic',
'AtenItemFpOpModule_basic',
'AtenItemIntOpModule_basic',
'AtenMatmulQMixedSigni8Transpose_basic',
'AtenMatmulQMixedSigni8_basic',
'AtenMatmulQint8MV_basic',
@ -465,7 +461,6 @@ FX_IMPORT_XFAIL_SET = {
'MaxPool3dStaticCeilModeTrueModule_basic',
'MaxPool3dStaticModule_basic',
'MulFloatModule_basic',
'MulIntModule_basic',
'NativeGroupNormBackwardModule_basic',
'NeFloatIntModule_basic',
'NeIntModule_basic',
@ -496,9 +491,6 @@ FX_IMPORT_XFAIL_SET = {
'RsubInt0d_NumToTensor_Module_basic',
'ScalarConstantTupleModule_basic',
'ScalarImplicitFloatModule_basic',
'ScalarImplicitIntModule_basic',
'ScatterValueFloatModule_basic',
'ScatterValueIntModule_basic',
'SortIntListReverse_basic',
'SortIntList_basic',
'SplitDimDynamicModule_basic',
@ -506,14 +498,11 @@ FX_IMPORT_XFAIL_SET = {
'SqrtIntConstantModule_basic',
'SqrtIntModule_basic',
'SubFloatModule_basic',
'SubIntModule_basic',
'TModuleRank0_basic',
'TensorToBoolZeroRank_basic',
'TensorToBool_basic',
'TensorToFloatZeroRank_basic',
'TensorToFloat_basic',
'TensorToIntZeroRank_basic',
'TensorToInt_basic',
'TestMultipleTensorAndPrimitiveTypesReturn_basic',
'ThresholdBackward2dMixedModule_basic',
'TorchPrimLoopForLikeModule_basic',

View File

@ -1397,6 +1397,7 @@ class GraphNodeImporter:
def _import_torch_op_overload(
self, loc: Location, node: torch_fx.Node, target: TorchOpOverload
):
# TODO: Convert this cascade of ifs to a table-driven
# replace lift_fresh_copy with clone op
if target == torch.ops.aten.lift_fresh_copy.default:
node.target = target = torch.ops.aten.clone.default
@ -1420,6 +1421,16 @@ class GraphNodeImporter:
for key_node in node.users:
if key_node.target == torch.ops.aten.baddbmm.default:
node.target = target = torch.ops.aten.zeros.default
elif target == torch.ops.aten._local_scalar_dense.default:
input_type = node.args[0].meta["tensor_meta"].dtype
if input_type.is_floating_point:
node.target = target = torch.ops.aten.Float.Tensor
else:
node.target = target = torch.ops.aten.Int.Tensor
node.args = (node.args[0],)
elif target == torch.ops.aten._assert_async.msg:
# TODO: A more suitable op to replace it?
return
schema = target._schema
assert isinstance(schema, FunctionSchema)