Skip to content

Commit 4cf9a46

Browse files
nathanaelseefacebook-github-bot
authored andcommitted
use _to_copy or _to_dim_order_copy in I64toI32 pass (#3968)
Summary: Pull Request resolved: #3968 When converting an ExportedProgram to EdgeProgramManager via to_edge, _to_dim_order_copy or _to_copy will be used depending on the passed EdgeCompileConfig having _skip_dim_order set to true or false. For example: ``` compile_config=EdgeCompileConfig( _check_ir_validity=False, _skip_dim_order=True, ), ``` The newer _to_dim_order_copy op was added in D53747744/#1970. Adding a param to the pass init to determine if we use the newer op or not. Reviewed By: copyrightly Differential Revision: D58395616 fbshipit-source-id: 8f5289f8ed62a5c3814ec32b91fb055a10944fdb
1 parent 2d35b30 commit 4cf9a46

File tree

1 file changed

+10
-2
lines changed

1 file changed

+10
-2
lines changed

backends/transforms/i64_to_i32.py

Lines changed: 10 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,14 @@
1414

1515
class I64toI32(ExportPass):
1616

17+
def __init__(self, _skip_dim_order=False):
18+
super(I64toI32, self).__init__()
19+
self.copy_op = (
20+
exir_ops.edge.aten._to_copy.default
21+
if _skip_dim_order
22+
else exir_ops.edge.dim_order_ops._to_dim_order_copy.default
23+
)
24+
1725
def _is_i64_tensor(self, node_val):
1826
return isinstance(node_val, FakeTensor) and node_val.dtype == torch.int64
1927

@@ -34,7 +42,7 @@ def _apply_to_i32(self, graph: torch.fx.Graph):
3442
args = (node,)
3543
node_i32 = graph.create_node(
3644
"call_function",
37-
exir_ops.edge.dim_order_ops._to_dim_order_copy.default,
45+
self.copy_op,
3846
args,
3947
{"dtype": torch.int32},
4048
)
@@ -62,7 +70,7 @@ def _apply_to_i32(self, graph: torch.fx.Graph):
6270
args = (node.args[0][i],)
6371
node_i64 = graph.create_node(
6472
"call_function",
65-
exir_ops.edge.dim_order_ops._to_dim_order_copy.default,
73+
self.copy_op,
6674
args,
6775
{"dtype": torch.int64},
6876
)

0 commit comments

Comments
 (0)