Skip to content

Add RemoveAssert pass to remove _assert_tensor_metadata nodes #7277

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Dec 16, 2024
Merged

Add RemoveAssert pass to remove _assert_tensor_metadata nodes #7277

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
23 changes: 23 additions & 0 deletions exir/passes/remove_graph_asserts_pass.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,7 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult:
torch.ops.aten._assert_scalar.default,
torch.ops.aten.sym_constrain_range_for_size.default,
torch.ops.aten.sym_constrain_range.default,
torch.ops.aten._assert_tensor_metadata.default,
)
):
module.graph.erase_node(node)
Expand All @@ -37,3 +38,25 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult:
module.graph.eliminate_dead_code()

return PassResult(graph_module, True)


class RemoveNonCoreAtenOpGraphAssertsPass(PassBase):
"""
Remove assert ops from the graph that're not Aten Canonical.
"""

def call(self, graph_module: torch.fx.GraphModule) -> PassResult:
for module in graph_module.modules():
if not isinstance(module, torch.fx.GraphModule):
continue

for node in module.graph.nodes:
if node.op == "call_function" and (
node.target in (torch.ops.aten._assert_tensor_metadata.default,)
):
module.graph.erase_node(node)

module.recompile()
module.graph.eliminate_dead_code()

return PassResult(graph_module, True)
15 changes: 12 additions & 3 deletions exir/program/_program.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -40,7 +40,10 @@
from executorch.exir.passes.normalize_view_copy_base_pass import (
NormalizeViewCopyBasePass,
)
from executorch.exir.passes.remove_graph_asserts_pass import RemoveGraphAssertsPass
from executorch.exir.passes.remove_graph_asserts_pass import (
RemoveGraphAssertsPass,
RemoveNonCoreAtenOpGraphAssertsPass,
)
from executorch.exir.passes.remove_mixed_type_operators import RemoveMixedTypeOperators
from executorch.exir.passes.replace_aten_with_edge_pass import aten_to_edge
from executorch.exir.passes.replace_view_copy_with_view_pass import (
Expand Down Expand Up @@ -722,13 +725,20 @@ def _generate_edge_program(
program: ExportedProgram,
ops_set_to_not_decompose: Optional[List[torch._ops.OpOverload]] = None,
) -> ExportedProgram:

# Remove invalid assert ops, such as _assert_tensor_metadata
gm = program.graph_module
gm_res = RemoveNonCoreAtenOpGraphAssertsPass()(gm)
assert gm_res is not None
gm = gm_res.graph_module

if config._check_ir_validity:
try:
EXIRATenDialectVerifier(
edge_compile_config=config,
class_only=False,
exception_list=ops_set_to_not_decompose,
)(program.graph_module)
)(gm)
except ExportError as e:
logging.info(f"Input program {name} is not in ATen dialect.")
raise e
Expand All @@ -745,7 +755,6 @@ def _generate_edge_program(
if not config._skip_dim_order:
passes.append(MemoryFormatOpsPass())

gm = program.graph_module
for p in passes:
gm_res = p(gm)
assert gm_res is not None
Expand Down