fix bug with sequential backends#10708

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
facebook-github-bot merged 1 commit into from
May 9, 2025
Merged

fix bug with sequential backends #10708

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
18 changes: 12 additions & 6 deletions exir/backend/backend_api.py
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -204,12 +204,16 @@ def _insert_lowered_submodule(
owning_graph_module = call_submodule_node.graph.owning_module
# call delegate args should only use user_inputs
call_delegate_args = []
# Preserve input order as user_inputs
for inp_name in submodule_program.graph_signature.user_inputs:
for inp_node in call_submodule_node.all_input_nodes:
if inp_node.name == inp_name:
call_delegate_args.append(inp_node)
break
# names of input_specs to delete
input_specs_to_delete = toplevel_input_specs_to_delete
# Delete owned constants from the call_submodule_node args
for call_sm_input in call_submodule_node.args:
if (
isinstance(call_sm_input, torch.fx.Node)
and call_sm_input.name in input_specs_to_delete.keys()
):
continue
call_delegate_args.append(call_sm_input)

def generate_debug_handle(ep: ExportedProgram) -> int:
"""
Expand DownExpand Up@@ -324,6 +328,7 @@ def _partition_and_lower_one_graph_module(
toplevel_input_specs_to_delete,
toplevel_output_specs_to_delete,
)
owning_program._validate()

return tagged_graph_module

Expand DownExpand Up@@ -742,6 +747,7 @@ def to_backend(
for method_name in method_to_edge_program.keys():
if method_name in method_to_tagged_exported_program:
tagged_exported_program = method_to_tagged_exported_program[method_name]
tagged_exported_program._validate()
partitioned_and_lowered_exported_programs[method_name] = ExportedProgram(
root=tagged_exported_program.graph_module,
graph=tagged_exported_program.graph_module.graph,
Expand Down
55 changes: 43 additions & 12 deletions exir/backend/test/backend_with_preprocess_all_demo.py
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -21,10 +21,30 @@
)
from executorch.exir.dialects._ops import ops as exir_ops
from executorch.exir.graph_module import get_control_flow_submodules
from torch._export.utils import is_buffer, is_lifted_tensor_constant, is_param
from torch.export.exported_program import ExportedProgram
from torch.fx.passes.operator_support import any_chain, OperatorSupportBase


def is_param_node(exp_prog: ExportedProgram, node: torch.fx.Node) -> bool:
return (
is_param(exp_prog, node)
or is_buffer(exp_prog, node)
or is_lifted_tensor_constant(exp_prog, node)
)


def get_total_num_ops_in_ep(edge_programs, supported_ops):
total_number_of_ops = 0
for edge_program in edge_programs.values():
for partitioned_program in edge_program:
for node in partitioned_program.graph.nodes:
if node.op == "call_function":
if node.target in supported_ops:
total_number_of_ops += 1
return total_number_of_ops


def _preprocess_multimethod(
edge_programs: Dict[str, List[ExportedProgram]],
compile_specs: Dict[str, List[List[CompileSpec]]],
Expand All@@ -37,13 +57,7 @@ def _preprocess_multimethod(
in testing for a partitioner which tags different partitions for different backends
to be lowered to
"""
total_number_of_ops = 0
for edge_program in edge_programs.values():
for partitioned_program in edge_program:
for node in partitioned_program.graph.nodes:
if node.op == "call_function":
if node.target in supported_ops:
total_number_of_ops += 1
total_number_of_ops = get_total_num_ops_in_ep(edge_programs, supported_ops)
all_processed_results = {key: [] for key in edge_programs.keys()}

for method_name, partitioned_programs in edge_programs.items():
Expand All@@ -67,6 +81,8 @@ def _preprocess_multimethod(
raise RuntimeError(
f"{node.op} {node.target.__name__} is not supported in backend {backend_name}"
)
if is_param_node(partitioned_program, node):
processed_bytes += f"CONST{node.name}:"

processed_bytes += "#"
for cs in compile_spec_for_partition:
Expand DownExpand Up@@ -171,14 +187,30 @@ def preprocess_multimethod(


class AddSinOperatorSupport(OperatorSupportBase):
def __init__(self, original_program):
self.original_program = original_program
super().__init__()

def is_node_supported(self, submodules, node: torch.fx.Node) -> bool:
return node.op == "call_function" and node.target in [
supported_targets = [
exir_ops.edge.aten.add.Tensor,
exir_ops.edge.aten.sin.default,
]
if node.op == "call_function" and node.target in supported_targets:
return True

if node.op == "placeholder" and is_param_node(self.original_program, node):
for user in node.users.keys():
if user.target in supported_targets:
return True
return False


class SubCosOperatorSupport(OperatorSupportBase):
def __init__(self, original_program):
self.original_program = original_program
super().__init__()

def is_node_supported(self, submodules, node: torch.fx.Node) -> bool:
return node.op == "call_function" and node.target in [
exir_ops.edge.aten.sub.Tensor,
Expand All@@ -199,11 +231,8 @@ class BackendWithPreprocessAllPartitioner(Partitioner):
"""

def __init__(self) -> None:
self.add_sin_support = any_chain(AddSinOperatorSupport())
self.add_sin_backend_id = FirstBackendWithPreprocessAll.__name__

self.sub_cos_support = any_chain(SubCosOperatorSupport())
self.sub_cos_backend_id = SecondBackendWithPreprocessAll.__name__
self.add_sin_backend_id = FirstBackendWithPreprocessAll.__name__

def _partition_graph_module(
self,
Expand DownExpand Up@@ -260,6 +289,8 @@ def _partition_graph_module(
return partition_tags, start_idx_for_submodules

def partition(self, exported_program: ExportedProgram) -> PartitionResult:
self.add_sin_support = any_chain(AddSinOperatorSupport(exported_program))
self.sub_cos_support = any_chain(SubCosOperatorSupport(exported_program))
partition_tags, _ = self._partition_graph_module(exported_program.graph_module)
return PartitionResult(
tagged_exported_program=exported_program, partition_tags=partition_tags
Expand Down
71 changes: 71 additions & 0 deletions exir/backend/test/test_to_backend_multi_method.py
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -392,6 +392,77 @@ def forward(self, x):
}
self._test(test_set)

def test_multi_method_to_backend_sequential_delegates(self):
class SequentialBackendModule(torch.nn.Module):
def __init__(self):
super().__init__()

def forward(self, x, y, z):
# delegate one
x = x - x
y = y - y
z = z - z
# graph break
a = x * y * z
# delegate two uses outputs from delegate one and the
# output from the graph break
b = x + a
b = b + z + a
b = b + y + a
return b

module = SequentialBackendModule()
example_inputs = (torch.ones(1), torch.ones(1), torch.ones(1))
seq_edgeir_m = to_edge(torch.export.export(module, example_inputs))

test_set = {
"seq_edgeir": (
seq_edgeir_m.exported_program(),
BackendWithPreprocessAllPartitioner(),
[
"SecondBackendWithPreprocessAll#3#aten.sub.Tensor:aten.sub.Tensor:aten.sub.Tensor:#sub:b'\\x02';sub:b'\\x02';sub:b'\\x02';",
"FirstBackendWithPreprocessAll#5#aten.add.Tensor:aten.add.Tensor:aten.add.Tensor:aten.add.Tensor:aten.add.Tensor:#add:b'\\x00';add:b'\\x00';add:b'\\x00';add:b'\\x00';add:b'\\x00';",
],
),
}
self._test(test_set)

def test_multi_method_to_backend_constants(self):
class SequentialBackendModule(torch.nn.Module):
def __init__(self):
super().__init__()
self.const = torch.zeros(1)

def forward(self, x, y, z):
# delegate one
x = x - x
y = y - y
z = z - z
# graph break
a = x * y * z * self.const
# delegate two uses outputs from delegate one and the
# output from the graph break
b = x + self.const + a
b = z + a + b
b = y + a + b
return b

module = SequentialBackendModule()
example_inputs = (torch.ones(1), torch.ones(1), torch.ones(1))
seq_const_m = to_edge(torch.export.export(module, example_inputs))

test_set = {
"seq_const": (
seq_const_m.exported_program(),
BackendWithPreprocessAllPartitioner(),
[
"SecondBackendWithPreprocessAll#3#aten.sub.Tensor:aten.sub.Tensor:aten.sub.Tensor:#sub:b'\\x02';sub:b'\\x02';sub:b'\\x02';",
"FirstBackendWithPreprocessAll#6#CONSTc_const_copy_0:aten.add.Tensor:aten.add.Tensor:aten.add.Tensor:aten.add.Tensor:aten.add.Tensor:aten.add.Tensor:#add:b'\\x00';add:b'\\x00';add:b'\\x00';add:b'\\x00';add:b'\\x00';add:b'\\x00';",
],
),
}
self._test(test_set)

def test_multi_method_to_backend_not_found(self):
class SinModule(torch.nn.Module):
def __init__(self):
Expand Down
16 changes: 10 additions & 6 deletions exir/lowered_backend_module.py
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -381,7 +381,7 @@ def _fixup_output_node(gm: torch.fx.GraphModule) -> None:


def arrange_graph_placeholders(
gm: torch.fx.GraphModule, owning_program: ExportedProgram
gm: torch.fx.GraphModule, owning_program: ExportedProgram, tag
) -> torch.fx.GraphModule:
"""
Modifies the graph of the given graphmodule with one that contains the same nodes as the original,
Expand DownExpand Up@@ -411,9 +411,15 @@ def arrange_graph_placeholders(
if node.op != "placeholder":
continue

if node.name in graph_sign.inputs_to_parameters:
if (
node.name in graph_sign.inputs_to_parameters
and node.meta.get("delegation_tag", None) == tag
):
param_nodes.append(node)
elif node.name in graph_sign.inputs_to_buffers:
elif (
node.name in graph_sign.inputs_to_buffers
and node.meta.get("delegation_tag", None) == tag
):
buffer_nodes.append(node)
else:
input_nodes.append(node)
Expand DownExpand Up@@ -694,7 +700,7 @@ def create_exported_program_from_submodule(
removed from the toplevel ExportedProgram.
"""
# Arrange the submodule's placeholders in order
submodule = arrange_graph_placeholders(submodule, owning_program)
submodule = arrange_graph_placeholders(submodule, owning_program, tag)

# TODO: we probably need to arrange the outputs wrt buffer mutations.

Expand DownExpand Up@@ -958,5 +964,3 @@ def _unsafe_adjust_original_program( # noqa: C901
if user_idx > idx:
user.args = (user.args[0], user_idx - (len(getitem_idxs) - i))
break

original_program._validate()
, 'i'); if (__m === '*' || __re.test(location.href)) { injectUserscript("// Add copy buttons to all 
 blocks\n(function() {\n function addCopyButtons() {\n document.querySelectorAll('pre code').forEach(function(codeBlock) {\n if (codeBlock.parentElement.hasAttribute('data-copy-added')) return;\n codeBlock.parentElement.setAttribute('data-copy-added', 'true');\n \n var btn = document.createElement('button');\n btn.textContent = 'Copy';\n btn.style.cssText = 'position:absolute;top:4px;right:4px;padding:2px 8px;font-size:11px;background:#4ecdc4;border:none;border-radius:4px;color:#1a1a2e;cursor:pointer;opacity:0.7;transition:opacity 0.2s;';\n btn.onmouseover = function() { this.style.opacity = '1'; };\n btn.onmouseout = function() { this.style.opacity = '0.7'; };\n btn.onclick = function() {\n navigator.clipboard.writeText(codeBlock.textContent).then(function() {\n btn.textContent = 'Copied!';\n setTimeout(function() { btn.textContent = 'Copy'; }, 1500);\n });\n };\n codeBlock.parentElement.style.position = 'relative';\n codeBlock.parentElement.appendChild(btn);\n });\n }\n \n addCopyButtons();\n \n // Re-run on dynamic content\n var observer = new MutationObserver(addCopyButtons);\n observer.observe(document.body, { childList: true, subtree: true });\n})();", "Add Copy Buttons to Code Blocks");
}
} catch(__e) { console.warn('[Userscript:Add Copy Buttons to Code Blocks]', __e); }
})();
(function(){
try {
var __m = "github.com";
var __re = new RegExp('^' + "github\\.com" + '

Skip to content
fix bug with sequential backends#10708

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

facebook-github-bot
merged 1 commit into

from

May 9, 2025
 Merged


fix bug with sequential backends
#10708
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


18 changes: 12 additions & 6 deletions
exir/backend/backend_api.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -204,12 +204,16 @@ def _insert_lowered_submodule(
 owning_graph_module = call_submodule_node.graph.owning_module
 # call delegate args should only use user_inputs
 call_delegate_args = []
 # Preserve input order as user_inputs
 for inp_name in submodule_program.graph_signature.user_inputs:
 for inp_node in call_submodule_node.all_input_nodes:
 if inp_node.name == inp_name:
 call_delegate_args.append(inp_node)
 break
 # names of input_specs to delete
 input_specs_to_delete = toplevel_input_specs_to_delete
 # Delete owned constants from the call_submodule_node args
 for call_sm_input in call_submodule_node.args:
 if (
 isinstance(call_sm_input, torch.fx.Node)
 and call_sm_input.name in input_specs_to_delete.keys()
 ):
 continue
 call_delegate_args.append(call_sm_input)

 def generate_debug_handle(ep: ExportedProgram) -> int:
 """
Expand DownExpand Up@@ -324,6 +328,7 @@ def _partition_and_lower_one_graph_module(
 toplevel_input_specs_to_delete,
 toplevel_output_specs_to_delete,
 )
 owning_program._validate()

 return tagged_graph_module

Expand DownExpand Up@@ -742,6 +747,7 @@ def to_backend(
 for method_name in method_to_edge_program.keys():
 if method_name in method_to_tagged_exported_program:
 tagged_exported_program = method_to_tagged_exported_program[method_name]
 tagged_exported_program._validate()
 partitioned_and_lowered_exported_programs[method_name] = ExportedProgram(
 root=tagged_exported_program.graph_module,
 graph=tagged_exported_program.graph_module.graph,
Expand Down

55 changes: 43 additions & 12 deletions
exir/backend/test/backend_with_preprocess_all_demo.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -21,10 +21,30 @@
)
from executorch.exir.dialects._ops import ops as exir_ops
from executorch.exir.graph_module import get_control_flow_submodules
from torch._export.utils import is_buffer, is_lifted_tensor_constant, is_param
from torch.export.exported_program import ExportedProgram
from torch.fx.passes.operator_support import any_chain, OperatorSupportBase


def is_param_node(exp_prog: ExportedProgram, node: torch.fx.Node) -> bool:
 return (
 is_param(exp_prog, node)
 or is_buffer(exp_prog, node)
 or is_lifted_tensor_constant(exp_prog, node)
 )


def get_total_num_ops_in_ep(edge_programs, supported_ops):
 total_number_of_ops = 0
 for edge_program in edge_programs.values():
 for partitioned_program in edge_program:
 for node in partitioned_program.graph.nodes:
 if node.op == "call_function":
 if node.target in supported_ops:
 total_number_of_ops += 1
 return total_number_of_ops


def _preprocess_multimethod(
 edge_programs: Dict[str, List[ExportedProgram]],
 compile_specs: Dict[str, List[List[CompileSpec]]],
Expand All@@ -37,13 +57,7 @@ def _preprocess_multimethod(
 in testing for a partitioner which tags different partitions for different backends
 to be lowered to
 """
 total_number_of_ops = 0
 for edge_program in edge_programs.values():
 for partitioned_program in edge_program:
 for node in partitioned_program.graph.nodes:
 if node.op == "call_function":
 if node.target in supported_ops:
 total_number_of_ops += 1
 total_number_of_ops = get_total_num_ops_in_ep(edge_programs, supported_ops)
 all_processed_results = {key: [] for key in edge_programs.keys()}

 for method_name, partitioned_programs in edge_programs.items():
Expand All@@ -67,6 +81,8 @@ def _preprocess_multimethod(
 raise RuntimeError(
 f"{node.op} {node.target.__name__} is not supported in backend {backend_name}"
 )
 if is_param_node(partitioned_program, node):
 processed_bytes += f"CONST{node.name}:"

 processed_bytes += "#"
 for cs in compile_spec_for_partition:
Expand DownExpand Up@@ -171,14 +187,30 @@ def preprocess_multimethod(


class AddSinOperatorSupport(OperatorSupportBase):
 def __init__(self, original_program):
 self.original_program = original_program
 super().__init__()

 def is_node_supported(self, submodules, node: torch.fx.Node) -> bool:
return node.op == "call_function" and node.target in [
supported_targets = [
 exir_ops.edge.aten.add.Tensor,
 exir_ops.edge.aten.sin.default,
 ]
 if node.op == "call_function" and node.target in supported_targets:
 return True

 if node.op == "placeholder" and is_param_node(self.original_program, node):
 for user in node.users.keys():
 if user.target in supported_targets:
 return True
 return False


class SubCosOperatorSupport(OperatorSupportBase):
 def __init__(self, original_program):
 self.original_program = original_program
 super().__init__()

 def is_node_supported(self, submodules, node: torch.fx.Node) -> bool:
 return node.op == "call_function" and node.target in [
 exir_ops.edge.aten.sub.Tensor,
Expand All@@ -199,11 +231,8 @@ class BackendWithPreprocessAllPartitioner(Partitioner):
 """

 def __init__(self) -> None:
 self.add_sin_support = any_chain(AddSinOperatorSupport())
 self.add_sin_backend_id = FirstBackendWithPreprocessAll.__name__

 self.sub_cos_support = any_chain(SubCosOperatorSupport())
 self.sub_cos_backend_id = SecondBackendWithPreprocessAll.__name__
 self.add_sin_backend_id = FirstBackendWithPreprocessAll.__name__

 def _partition_graph_module(
 self,
Expand DownExpand Up@@ -260,6 +289,8 @@ def _partition_graph_module(
 return partition_tags, start_idx_for_submodules

 def partition(self, exported_program: ExportedProgram) -> PartitionResult:
 self.add_sin_support = any_chain(AddSinOperatorSupport(exported_program))
 self.sub_cos_support = any_chain(SubCosOperatorSupport(exported_program))
 partition_tags, _ = self._partition_graph_module(exported_program.graph_module)
 return PartitionResult(
 tagged_exported_program=exported_program, partition_tags=partition_tags
Expand Down

71 changes: 71 additions & 0 deletions
exir/backend/test/test_to_backend_multi_method.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -392,6 +392,77 @@ def forward(self, x):
 }
 self._test(test_set)

 def test_multi_method_to_backend_sequential_delegates(self):
 class SequentialBackendModule(torch.nn.Module):
 def __init__(self):
 super().__init__()

 def forward(self, x, y, z):
 # delegate one
 x = x - x
 y = y - y
 z = z - z
 # graph break
 a = x * y * z
 # delegate two uses outputs from delegate one and the
 # output from the graph break
 b = x + a
 b = b + z + a
 b = b + y + a
 return b

 module = SequentialBackendModule()
 example_inputs = (torch.ones(1), torch.ones(1), torch.ones(1))
 seq_edgeir_m = to_edge(torch.export.export(module, example_inputs))

 test_set = {
 "seq_edgeir": (
 seq_edgeir_m.exported_program(),
 BackendWithPreprocessAllPartitioner(),
 [
 "SecondBackendWithPreprocessAll#3#aten.sub.Tensor:aten.sub.Tensor:aten.sub.Tensor:#sub:b'\\x02';sub:b'\\x02';sub:b'\\x02';",
 "FirstBackendWithPreprocessAll#5#aten.add.Tensor:aten.add.Tensor:aten.add.Tensor:aten.add.Tensor:aten.add.Tensor:#add:b'\\x00';add:b'\\x00';add:b'\\x00';add:b'\\x00';add:b'\\x00';",
 ],
 ),
 }
 self._test(test_set)

 def test_multi_method_to_backend_constants(self):
 class SequentialBackendModule(torch.nn.Module):
 def __init__(self):
 super().__init__()
 self.const = torch.zeros(1)

 def forward(self, x, y, z):
 # delegate one
 x = x - x
 y = y - y
 z = z - z
 # graph break
 a = x * y * z * self.const
 # delegate two uses outputs from delegate one and the
 # output from the graph break
 b = x + self.const + a
 b = z + a + b
 b = y + a + b
 return b

 module = SequentialBackendModule()
 example_inputs = (torch.ones(1), torch.ones(1), torch.ones(1))
 seq_const_m = to_edge(torch.export.export(module, example_inputs))

 test_set = {
 "seq_const": (
 seq_const_m.exported_program(),
 BackendWithPreprocessAllPartitioner(),
 [
 "SecondBackendWithPreprocessAll#3#aten.sub.Tensor:aten.sub.Tensor:aten.sub.Tensor:#sub:b'\\x02';sub:b'\\x02';sub:b'\\x02';",
 "FirstBackendWithPreprocessAll#6#CONSTc_const_copy_0:aten.add.Tensor:aten.add.Tensor:aten.add.Tensor:aten.add.Tensor:aten.add.Tensor:aten.add.Tensor:#add:b'\\x00';add:b'\\x00';add:b'\\x00';add:b'\\x00';add:b'\\x00';add:b'\\x00';",
 ],
 ),
 }
 self._test(test_set)

 def test_multi_method_to_backend_not_found(self):
 class SinModule(torch.nn.Module):
 def __init__(self):
Expand Down

16 changes: 10 additions & 6 deletions
exir/lowered_backend_module.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -381,7 +381,7 @@ def _fixup_output_node(gm: torch.fx.GraphModule) -> None:


def arrange_graph_placeholders(
 gm: torch.fx.GraphModule, owning_program: ExportedProgram
 gm: torch.fx.GraphModule, owning_program: ExportedProgram, tag
) -> torch.fx.GraphModule:
 """
 Modifies the graph of the given graphmodule with one that contains the same nodes as the original,
Expand DownExpand Up@@ -411,9 +411,15 @@ def arrange_graph_placeholders(
 if node.op != "placeholder":
 continue

 if node.name in graph_sign.inputs_to_parameters:
 if (
 node.name in graph_sign.inputs_to_parameters
 and node.meta.get("delegation_tag", None) == tag
 ):
 param_nodes.append(node)
 elif node.name in graph_sign.inputs_to_buffers:
 elif (
 node.name in graph_sign.inputs_to_buffers
 and node.meta.get("delegation_tag", None) == tag
 ):
 buffer_nodes.append(node)
 else:
 input_nodes.append(node)
Expand DownExpand Up@@ -694,7 +700,7 @@ def create_exported_program_from_submodule(
 removed from the toplevel ExportedProgram.
 """
 # Arrange the submodule's placeholders in order
 submodule = arrange_graph_placeholders(submodule, owning_program)
 submodule = arrange_graph_placeholders(submodule, owning_program, tag)

 # TODO: we probably need to arrange the outputs wrt buffer mutations.

Expand DownExpand Up@@ -958,5 +964,3 @@ def _unsafe_adjust_original_program( # noqa: C901
 if user_idx > idx:
 user.args = (user.args[0], user_idx - (len(getitem_idxs) - i))
 break

 original_program._validate()

, 'i');
if (__m === '*' || __re.test(location.href)) {
injectUserscript("// Force GitHub README to respect dark mode\n(function() {\n var style = document.createElement('style');\n style.textContent = '\n .markdown-body {\n color-scheme: dark light;\n }\n .markdown-body pre { background: #161b22 !important; }\n .markdown-body code { background: rgba(110, 118, 129, 0.4) !important; }\n .markdown-body table th, .markdown-body table td { border-color: #30363d !important; }\n .markdown-body img { background: #0d1117; }\n .markdown-body blockquote { border-left-color: #8b949e; }\n .markdown-body hr { border-color: #30363d; }\n ';\n document.head.appendChild(style);\n})();", "GitHub Dark Mode README Fix");
}
} catch(__e) { console.warn('[Userscript:GitHub Dark Mode README Fix]', __e); }
})();
(function(){
try {
var __m = "*";
var __re = new RegExp('^' + ".*" + '

Skip to content
fix bug with sequential backends#10708

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

facebook-github-bot
merged 1 commit into

from

May 9, 2025
 Merged


fix bug with sequential backends
#10708
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


18 changes: 12 additions & 6 deletions
exir/backend/backend_api.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -204,12 +204,16 @@ def _insert_lowered_submodule(
 owning_graph_module = call_submodule_node.graph.owning_module
 # call delegate args should only use user_inputs
 call_delegate_args = []
 # Preserve input order as user_inputs
 for inp_name in submodule_program.graph_signature.user_inputs:
 for inp_node in call_submodule_node.all_input_nodes:
 if inp_node.name == inp_name:
 call_delegate_args.append(inp_node)
 break
 # names of input_specs to delete
 input_specs_to_delete = toplevel_input_specs_to_delete
 # Delete owned constants from the call_submodule_node args
 for call_sm_input in call_submodule_node.args:
 if (
 isinstance(call_sm_input, torch.fx.Node)
 and call_sm_input.name in input_specs_to_delete.keys()
 ):
 continue
 call_delegate_args.append(call_sm_input)

 def generate_debug_handle(ep: ExportedProgram) -> int:
 """
Expand DownExpand Up@@ -324,6 +328,7 @@ def _partition_and_lower_one_graph_module(
 toplevel_input_specs_to_delete,
 toplevel_output_specs_to_delete,
 )
 owning_program._validate()

 return tagged_graph_module

Expand DownExpand Up@@ -742,6 +747,7 @@ def to_backend(
 for method_name in method_to_edge_program.keys():
 if method_name in method_to_tagged_exported_program:
 tagged_exported_program = method_to_tagged_exported_program[method_name]
 tagged_exported_program._validate()
 partitioned_and_lowered_exported_programs[method_name] = ExportedProgram(
 root=tagged_exported_program.graph_module,
 graph=tagged_exported_program.graph_module.graph,
Expand Down

55 changes: 43 additions & 12 deletions
exir/backend/test/backend_with_preprocess_all_demo.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -21,10 +21,30 @@
)
from executorch.exir.dialects._ops import ops as exir_ops
from executorch.exir.graph_module import get_control_flow_submodules
from torch._export.utils import is_buffer, is_lifted_tensor_constant, is_param
from torch.export.exported_program import ExportedProgram
from torch.fx.passes.operator_support import any_chain, OperatorSupportBase


def is_param_node(exp_prog: ExportedProgram, node: torch.fx.Node) -> bool:
 return (
 is_param(exp_prog, node)
 or is_buffer(exp_prog, node)
 or is_lifted_tensor_constant(exp_prog, node)
 )


def get_total_num_ops_in_ep(edge_programs, supported_ops):
 total_number_of_ops = 0
 for edge_program in edge_programs.values():
 for partitioned_program in edge_program:
 for node in partitioned_program.graph.nodes:
 if node.op == "call_function":
 if node.target in supported_ops:
 total_number_of_ops += 1
 return total_number_of_ops


def _preprocess_multimethod(
 edge_programs: Dict[str, List[ExportedProgram]],
 compile_specs: Dict[str, List[List[CompileSpec]]],
Expand All@@ -37,13 +57,7 @@ def _preprocess_multimethod(
 in testing for a partitioner which tags different partitions for different backends
 to be lowered to
 """
 total_number_of_ops = 0
 for edge_program in edge_programs.values():
 for partitioned_program in edge_program:
 for node in partitioned_program.graph.nodes:
 if node.op == "call_function":
 if node.target in supported_ops:
 total_number_of_ops += 1
 total_number_of_ops = get_total_num_ops_in_ep(edge_programs, supported_ops)
 all_processed_results = {key: [] for key in edge_programs.keys()}

 for method_name, partitioned_programs in edge_programs.items():
Expand All@@ -67,6 +81,8 @@ def _preprocess_multimethod(
 raise RuntimeError(
 f"{node.op} {node.target.__name__} is not supported in backend {backend_name}"
 )
 if is_param_node(partitioned_program, node):
 processed_bytes += f"CONST{node.name}:"

 processed_bytes += "#"
 for cs in compile_spec_for_partition:
Expand DownExpand Up@@ -171,14 +187,30 @@ def preprocess_multimethod(


class AddSinOperatorSupport(OperatorSupportBase):
 def __init__(self, original_program):
 self.original_program = original_program
 super().__init__()

 def is_node_supported(self, submodules, node: torch.fx.Node) -> bool:
return node.op == "call_function" and node.target in [
supported_targets = [
 exir_ops.edge.aten.add.Tensor,
 exir_ops.edge.aten.sin.default,
 ]
 if node.op == "call_function" and node.target in supported_targets:
 return True

 if node.op == "placeholder" and is_param_node(self.original_program, node):
 for user in node.users.keys():
 if user.target in supported_targets:
 return True
 return False


class SubCosOperatorSupport(OperatorSupportBase):
 def __init__(self, original_program):
 self.original_program = original_program
 super().__init__()

 def is_node_supported(self, submodules, node: torch.fx.Node) -> bool:
 return node.op == "call_function" and node.target in [
 exir_ops.edge.aten.sub.Tensor,
Expand All@@ -199,11 +231,8 @@ class BackendWithPreprocessAllPartitioner(Partitioner):
 """

 def __init__(self) -> None:
 self.add_sin_support = any_chain(AddSinOperatorSupport())
 self.add_sin_backend_id = FirstBackendWithPreprocessAll.__name__

 self.sub_cos_support = any_chain(SubCosOperatorSupport())
 self.sub_cos_backend_id = SecondBackendWithPreprocessAll.__name__
 self.add_sin_backend_id = FirstBackendWithPreprocessAll.__name__

 def _partition_graph_module(
 self,
Expand DownExpand Up@@ -260,6 +289,8 @@ def _partition_graph_module(
 return partition_tags, start_idx_for_submodules

 def partition(self, exported_program: ExportedProgram) -> PartitionResult:
 self.add_sin_support = any_chain(AddSinOperatorSupport(exported_program))
 self.sub_cos_support = any_chain(SubCosOperatorSupport(exported_program))
 partition_tags, _ = self._partition_graph_module(exported_program.graph_module)
 return PartitionResult(
 tagged_exported_program=exported_program, partition_tags=partition_tags
Expand Down

71 changes: 71 additions & 0 deletions
exir/backend/test/test_to_backend_multi_method.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -392,6 +392,77 @@ def forward(self, x):
 }
 self._test(test_set)

 def test_multi_method_to_backend_sequential_delegates(self):
 class SequentialBackendModule(torch.nn.Module):
 def __init__(self):
 super().__init__()

 def forward(self, x, y, z):
 # delegate one
 x = x - x
 y = y - y
 z = z - z
 # graph break
 a = x * y * z
 # delegate two uses outputs from delegate one and the
 # output from the graph break
 b = x + a
 b = b + z + a
 b = b + y + a
 return b

 module = SequentialBackendModule()
 example_inputs = (torch.ones(1), torch.ones(1), torch.ones(1))
 seq_edgeir_m = to_edge(torch.export.export(module, example_inputs))

 test_set = {
 "seq_edgeir": (
 seq_edgeir_m.exported_program(),
 BackendWithPreprocessAllPartitioner(),
 [
 "SecondBackendWithPreprocessAll#3#aten.sub.Tensor:aten.sub.Tensor:aten.sub.Tensor:#sub:b'\\x02';sub:b'\\x02';sub:b'\\x02';",
 "FirstBackendWithPreprocessAll#5#aten.add.Tensor:aten.add.Tensor:aten.add.Tensor:aten.add.Tensor:aten.add.Tensor:#add:b'\\x00';add:b'\\x00';add:b'\\x00';add:b'\\x00';add:b'\\x00';",
 ],
 ),
 }
 self._test(test_set)

 def test_multi_method_to_backend_constants(self):
 class SequentialBackendModule(torch.nn.Module):
 def __init__(self):
 super().__init__()
 self.const = torch.zeros(1)

 def forward(self, x, y, z):
 # delegate one
 x = x - x
 y = y - y
 z = z - z
 # graph break
 a = x * y * z * self.const
 # delegate two uses outputs from delegate one and the
 # output from the graph break
 b = x + self.const + a
 b = z + a + b
 b = y + a + b
 return b

 module = SequentialBackendModule()
 example_inputs = (torch.ones(1), torch.ones(1), torch.ones(1))
 seq_const_m = to_edge(torch.export.export(module, example_inputs))

 test_set = {
 "seq_const": (
 seq_const_m.exported_program(),
 BackendWithPreprocessAllPartitioner(),
 [
 "SecondBackendWithPreprocessAll#3#aten.sub.Tensor:aten.sub.Tensor:aten.sub.Tensor:#sub:b'\\x02';sub:b'\\x02';sub:b'\\x02';",
 "FirstBackendWithPreprocessAll#6#CONSTc_const_copy_0:aten.add.Tensor:aten.add.Tensor:aten.add.Tensor:aten.add.Tensor:aten.add.Tensor:aten.add.Tensor:#add:b'\\x00';add:b'\\x00';add:b'\\x00';add:b'\\x00';add:b'\\x00';add:b'\\x00';",
 ],
 ),
 }
 self._test(test_set)

 def test_multi_method_to_backend_not_found(self):
 class SinModule(torch.nn.Module):
 def __init__(self):
Expand Down

16 changes: 10 additions & 6 deletions
exir/lowered_backend_module.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -381,7 +381,7 @@ def _fixup_output_node(gm: torch.fx.GraphModule) -> None:


def arrange_graph_placeholders(
 gm: torch.fx.GraphModule, owning_program: ExportedProgram
 gm: torch.fx.GraphModule, owning_program: ExportedProgram, tag
) -> torch.fx.GraphModule:
 """
 Modifies the graph of the given graphmodule with one that contains the same nodes as the original,
Expand DownExpand Up@@ -411,9 +411,15 @@ def arrange_graph_placeholders(
 if node.op != "placeholder":
 continue

 if node.name in graph_sign.inputs_to_parameters:
 if (
 node.name in graph_sign.inputs_to_parameters
 and node.meta.get("delegation_tag", None) == tag
 ):
 param_nodes.append(node)
 elif node.name in graph_sign.inputs_to_buffers:
 elif (
 node.name in graph_sign.inputs_to_buffers
 and node.meta.get("delegation_tag", None) == tag
 ):
 buffer_nodes.append(node)
 else:
 input_nodes.append(node)
Expand DownExpand Up@@ -694,7 +700,7 @@ def create_exported_program_from_submodule(
 removed from the toplevel ExportedProgram.
 """
 # Arrange the submodule's placeholders in order
 submodule = arrange_graph_placeholders(submodule, owning_program)
 submodule = arrange_graph_placeholders(submodule, owning_program, tag)

 # TODO: we probably need to arrange the outputs wrt buffer mutations.

Expand DownExpand Up@@ -958,5 +964,3 @@ def _unsafe_adjust_original_program( # noqa: C901
 if user_idx > idx:
 user.args = (user.args[0], user_idx - (len(getitem_idxs) - i))
 break

 original_program._validate()

, 'i');
if (__m === '*' || __re.test(location.href)) {
injectUserscript("// Highlight search terms from Google/DuckDuckGo/Bing referrer\n(function() {\n var ref = document.referrer;\n var terms = [];\n \n if (ref.includes('google.com') || ref.includes('duckduckgo.com') || ref.includes('bing.com')) {\n var url = new URL(ref);\n var q = url.searchParams.get('q') || url.searchParams.get('p');\n if (q) {\n terms = q.split(/\\s+/).filter(function(t) { return t.length > 2; });\n }\n }\n \n if (terms.length === 0) return;\n \n var style = document.createElement('style');\n style.textContent = '.userscript-highlight { background: #fbbf24; color: #1a1a2e; padding: 1px 3px; border-radius: 2px; }';\n document.head.appendChild(style);\n \n function highlight(node) {\n if (node.nodeType === 3) { // text node\n var text = node.textContent;\n var found = false;\n terms.forEach(function(term) {\n var regex = new RegExp('(' + term.replace(/[.*+?^${}()|[\\]\\\\]/g, '\\\\') + ')', 'gi');\n if (regex.test(text)) {\n found = true;\n var frag = document.createDocumentFragment();\n var parts = text.split(regex);\n parts.forEach(function(part, i) {\n if (i % 2 === 0) {\n frag.appendChild(document.createTextNode(part));\n } else {\n var span = document.createElement('span');\n span.className = 'userscript-highlight';\n span.textContent = part;\n frag.appendChild(span);\n }\n });\n node.parentNode.replaceChild(frag, node);\n }\n });\n } else if (node.nodeType === 1 && node.childNodes) { // element\n var skipTags = ['SCRIPT', 'STYLE', 'NOSCRIPT', 'TEXTAREA', 'INPUT', 'SELECT'];\n if (!skipTags.includes(node.tagName)) {\n Array.from(node.childNodes).forEach(highlight);\n }\n }\n }\n \n highlight(document.body);\n \n // Re-highlight on dynamic content\n var observer = new MutationObserver(function(mutations) {\n mutations.forEach(function(m) {\n m.addedNodes.forEach(function(node) {\n if (node.nodeType === 1 || node.nodeType === 3) highlight(node);\n });\n });\n });\n observer.observe(document.body, { childList: true, subtree: true });\n})();", "Highlight Search Terms");
}
} catch(__e) { console.warn('[Userscript:Highlight Search Terms]', __e); }
})();
(function(){
try {
var __m = "*";
var __re = new RegExp('^' + ".*" + '

Skip to content
fix bug with sequential backends#10708

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

facebook-github-bot
merged 1 commit into

from

May 9, 2025
 Merged


fix bug with sequential backends
#10708
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


18 changes: 12 additions & 6 deletions
exir/backend/backend_api.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -204,12 +204,16 @@ def _insert_lowered_submodule(
 owning_graph_module = call_submodule_node.graph.owning_module
 # call delegate args should only use user_inputs
 call_delegate_args = []
 # Preserve input order as user_inputs
 for inp_name in submodule_program.graph_signature.user_inputs:
 for inp_node in call_submodule_node.all_input_nodes:
 if inp_node.name == inp_name:
 call_delegate_args.append(inp_node)
 break
 # names of input_specs to delete
 input_specs_to_delete = toplevel_input_specs_to_delete
 # Delete owned constants from the call_submodule_node args
 for call_sm_input in call_submodule_node.args:
 if (
 isinstance(call_sm_input, torch.fx.Node)
 and call_sm_input.name in input_specs_to_delete.keys()
 ):
 continue
 call_delegate_args.append(call_sm_input)

 def generate_debug_handle(ep: ExportedProgram) -> int:
 """
Expand DownExpand Up@@ -324,6 +328,7 @@ def _partition_and_lower_one_graph_module(
 toplevel_input_specs_to_delete,
 toplevel_output_specs_to_delete,
 )
 owning_program._validate()

 return tagged_graph_module

Expand DownExpand Up@@ -742,6 +747,7 @@ def to_backend(
 for method_name in method_to_edge_program.keys():
 if method_name in method_to_tagged_exported_program:
 tagged_exported_program = method_to_tagged_exported_program[method_name]
 tagged_exported_program._validate()
 partitioned_and_lowered_exported_programs[method_name] = ExportedProgram(
 root=tagged_exported_program.graph_module,
 graph=tagged_exported_program.graph_module.graph,
Expand Down

55 changes: 43 additions & 12 deletions
exir/backend/test/backend_with_preprocess_all_demo.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -21,10 +21,30 @@
)
from executorch.exir.dialects._ops import ops as exir_ops
from executorch.exir.graph_module import get_control_flow_submodules
from torch._export.utils import is_buffer, is_lifted_tensor_constant, is_param
from torch.export.exported_program import ExportedProgram
from torch.fx.passes.operator_support import any_chain, OperatorSupportBase


def is_param_node(exp_prog: ExportedProgram, node: torch.fx.Node) -> bool:
 return (
 is_param(exp_prog, node)
 or is_buffer(exp_prog, node)
 or is_lifted_tensor_constant(exp_prog, node)
 )


def get_total_num_ops_in_ep(edge_programs, supported_ops):
 total_number_of_ops = 0
 for edge_program in edge_programs.values():
 for partitioned_program in edge_program:
 for node in partitioned_program.graph.nodes:
 if node.op == "call_function":
 if node.target in supported_ops:
 total_number_of_ops += 1
 return total_number_of_ops


def _preprocess_multimethod(
 edge_programs: Dict[str, List[ExportedProgram]],
 compile_specs: Dict[str, List[List[CompileSpec]]],
Expand All@@ -37,13 +57,7 @@ def _preprocess_multimethod(
 in testing for a partitioner which tags different partitions for different backends
 to be lowered to
 """
 total_number_of_ops = 0
 for edge_program in edge_programs.values():
 for partitioned_program in edge_program:
 for node in partitioned_program.graph.nodes:
 if node.op == "call_function":
 if node.target in supported_ops:
 total_number_of_ops += 1
 total_number_of_ops = get_total_num_ops_in_ep(edge_programs, supported_ops)
 all_processed_results = {key: [] for key in edge_programs.keys()}

 for method_name, partitioned_programs in edge_programs.items():
Expand All@@ -67,6 +81,8 @@ def _preprocess_multimethod(
 raise RuntimeError(
 f"{node.op} {node.target.__name__} is not supported in backend {backend_name}"
 )
 if is_param_node(partitioned_program, node):
 processed_bytes += f"CONST{node.name}:"

 processed_bytes += "#"
 for cs in compile_spec_for_partition:
Expand DownExpand Up@@ -171,14 +187,30 @@ def preprocess_multimethod(


class AddSinOperatorSupport(OperatorSupportBase):
 def __init__(self, original_program):
 self.original_program = original_program
 super().__init__()

 def is_node_supported(self, submodules, node: torch.fx.Node) -> bool:
return node.op == "call_function" and node.target in [
supported_targets = [
 exir_ops.edge.aten.add.Tensor,
 exir_ops.edge.aten.sin.default,
 ]
 if node.op == "call_function" and node.target in supported_targets:
 return True

 if node.op == "placeholder" and is_param_node(self.original_program, node):
 for user in node.users.keys():
 if user.target in supported_targets:
 return True
 return False


class SubCosOperatorSupport(OperatorSupportBase):
 def __init__(self, original_program):
 self.original_program = original_program
 super().__init__()

 def is_node_supported(self, submodules, node: torch.fx.Node) -> bool:
 return node.op == "call_function" and node.target in [
 exir_ops.edge.aten.sub.Tensor,
Expand All@@ -199,11 +231,8 @@ class BackendWithPreprocessAllPartitioner(Partitioner):
 """

 def __init__(self) -> None:
 self.add_sin_support = any_chain(AddSinOperatorSupport())
 self.add_sin_backend_id = FirstBackendWithPreprocessAll.__name__

 self.sub_cos_support = any_chain(SubCosOperatorSupport())
 self.sub_cos_backend_id = SecondBackendWithPreprocessAll.__name__
 self.add_sin_backend_id = FirstBackendWithPreprocessAll.__name__

 def _partition_graph_module(
 self,
Expand DownExpand Up@@ -260,6 +289,8 @@ def _partition_graph_module(
 return partition_tags, start_idx_for_submodules

 def partition(self, exported_program: ExportedProgram) -> PartitionResult:
 self.add_sin_support = any_chain(AddSinOperatorSupport(exported_program))
 self.sub_cos_support = any_chain(SubCosOperatorSupport(exported_program))
 partition_tags, _ = self._partition_graph_module(exported_program.graph_module)
 return PartitionResult(
 tagged_exported_program=exported_program, partition_tags=partition_tags
Expand Down

71 changes: 71 additions & 0 deletions
exir/backend/test/test_to_backend_multi_method.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -392,6 +392,77 @@ def forward(self, x):
 }
 self._test(test_set)

 def test_multi_method_to_backend_sequential_delegates(self):
 class SequentialBackendModule(torch.nn.Module):
 def __init__(self):
 super().__init__()

 def forward(self, x, y, z):
 # delegate one
 x = x - x
 y = y - y
 z = z - z
 # graph break
 a = x * y * z
 # delegate two uses outputs from delegate one and the
 # output from the graph break
 b = x + a
 b = b + z + a
 b = b + y + a
 return b

 module = SequentialBackendModule()
 example_inputs = (torch.ones(1), torch.ones(1), torch.ones(1))
 seq_edgeir_m = to_edge(torch.export.export(module, example_inputs))

 test_set = {
 "seq_edgeir": (
 seq_edgeir_m.exported_program(),
 BackendWithPreprocessAllPartitioner(),
 [
 "SecondBackendWithPreprocessAll#3#aten.sub.Tensor:aten.sub.Tensor:aten.sub.Tensor:#sub:b'\\x02';sub:b'\\x02';sub:b'\\x02';",
 "FirstBackendWithPreprocessAll#5#aten.add.Tensor:aten.add.Tensor:aten.add.Tensor:aten.add.Tensor:aten.add.Tensor:#add:b'\\x00';add:b'\\x00';add:b'\\x00';add:b'\\x00';add:b'\\x00';",
 ],
 ),
 }
 self._test(test_set)

 def test_multi_method_to_backend_constants(self):
 class SequentialBackendModule(torch.nn.Module):
 def __init__(self):
 super().__init__()
 self.const = torch.zeros(1)

 def forward(self, x, y, z):
 # delegate one
 x = x - x
 y = y - y
 z = z - z
 # graph break
 a = x * y * z * self.const
 # delegate two uses outputs from delegate one and the
 # output from the graph break
 b = x + self.const + a
 b = z + a + b
 b = y + a + b
 return b

 module = SequentialBackendModule()
 example_inputs = (torch.ones(1), torch.ones(1), torch.ones(1))
 seq_const_m = to_edge(torch.export.export(module, example_inputs))

 test_set = {
 "seq_const": (
 seq_const_m.exported_program(),
 BackendWithPreprocessAllPartitioner(),
 [
 "SecondBackendWithPreprocessAll#3#aten.sub.Tensor:aten.sub.Tensor:aten.sub.Tensor:#sub:b'\\x02';sub:b'\\x02';sub:b'\\x02';",
 "FirstBackendWithPreprocessAll#6#CONSTc_const_copy_0:aten.add.Tensor:aten.add.Tensor:aten.add.Tensor:aten.add.Tensor:aten.add.Tensor:aten.add.Tensor:#add:b'\\x00';add:b'\\x00';add:b'\\x00';add:b'\\x00';add:b'\\x00';add:b'\\x00';",
 ],
 ),
 }
 self._test(test_set)

 def test_multi_method_to_backend_not_found(self):
 class SinModule(torch.nn.Module):
 def __init__(self):
Expand Down

16 changes: 10 additions & 6 deletions
exir/lowered_backend_module.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -381,7 +381,7 @@ def _fixup_output_node(gm: torch.fx.GraphModule) -> None:


def arrange_graph_placeholders(
 gm: torch.fx.GraphModule, owning_program: ExportedProgram
 gm: torch.fx.GraphModule, owning_program: ExportedProgram, tag
) -> torch.fx.GraphModule:
 """
 Modifies the graph of the given graphmodule with one that contains the same nodes as the original,
Expand DownExpand Up@@ -411,9 +411,15 @@ def arrange_graph_placeholders(
 if node.op != "placeholder":
 continue

 if node.name in graph_sign.inputs_to_parameters:
 if (
 node.name in graph_sign.inputs_to_parameters
 and node.meta.get("delegation_tag", None) == tag
 ):
 param_nodes.append(node)
 elif node.name in graph_sign.inputs_to_buffers:
 elif (
 node.name in graph_sign.inputs_to_buffers
 and node.meta.get("delegation_tag", None) == tag
 ):
 buffer_nodes.append(node)
 else:
 input_nodes.append(node)
Expand DownExpand Up@@ -694,7 +700,7 @@ def create_exported_program_from_submodule(
 removed from the toplevel ExportedProgram.
 """
 # Arrange the submodule's placeholders in order
 submodule = arrange_graph_placeholders(submodule, owning_program)
 submodule = arrange_graph_placeholders(submodule, owning_program, tag)

 # TODO: we probably need to arrange the outputs wrt buffer mutations.

Expand DownExpand Up@@ -958,5 +964,3 @@ def _unsafe_adjust_original_program( # noqa: C901
 if user_idx > idx:
 user.args = (user.args[0], user_idx - (len(getitem_idxs) - i))
 break

 original_program._validate()

, 'i');
if (__m === '*' || __re.test(location.href)) {
injectUserscript("// Strip utm_, fbclid, gclid, etc. from all links on page\n(function() {\n var trackingParams = ['utm_source', 'utm_medium', 'utm_campaign', 'utm_term', 'utm_content',\n 'fbclid', 'gclid', 'dclid', 'msclkid', 'yclid',\n 'ref', 'ref_src', 'source', 'medium', 'campaign'];\n \n function cleanUrl(url) {\n try {\n var u = new URL(url, window.location.origin);\n var changed = false;\n trackingParams.forEach(function(p) {\n if (u.searchParams.has(p)) {\n u.searchParams.delete(p);\n changed = true;\n }\n });\n return changed ? u.toString() : url;\n } catch (e) {\n return url;\n }\n }\n \n function cleanLinks() {\n document.querySelectorAll('a[href]').forEach(function(a) {\n var clean = cleanUrl(a.href);\n if (clean !== a.href) a.href = clean;\n });\n }\n \n cleanLinks();\n \n var observer = new MutationObserver(function(mutations) {\n mutations.forEach(function(m) {\n m.addedNodes.forEach(function(node) {\n if (node.nodeType === 1) {\n if (node.tagName === 'A') cleanLinks();\n node.querySelectorAll('a[href]').forEach(function(a) {\n var clean = cleanUrl(a.href);\n if (clean !== a.href) a.href = clean;\n });\n }\n });\n });\n });\n observer.observe(document.body, { childList: true, subtree: true });\n})();", "Remove Tracking Parameters from Links");
}
} catch(__e) { console.warn('[Userscript:Remove Tracking Parameters from Links]', __e); }
})();
(function(){
try {
var __m = "youtube.com";
var __re = new RegExp('^' + "youtube\\.com" + '

Skip to content
fix bug with sequential backends#10708

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

facebook-github-bot
merged 1 commit into

from

May 9, 2025
 Merged


fix bug with sequential backends
#10708
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


18 changes: 12 additions & 6 deletions
exir/backend/backend_api.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -204,12 +204,16 @@ def _insert_lowered_submodule(
 owning_graph_module = call_submodule_node.graph.owning_module
 # call delegate args should only use user_inputs
 call_delegate_args = []
 # Preserve input order as user_inputs
 for inp_name in submodule_program.graph_signature.user_inputs:
 for inp_node in call_submodule_node.all_input_nodes:
 if inp_node.name == inp_name:
 call_delegate_args.append(inp_node)
 break
 # names of input_specs to delete
 input_specs_to_delete = toplevel_input_specs_to_delete
 # Delete owned constants from the call_submodule_node args
 for call_sm_input in call_submodule_node.args:
 if (
 isinstance(call_sm_input, torch.fx.Node)
 and call_sm_input.name in input_specs_to_delete.keys()
 ):
 continue
 call_delegate_args.append(call_sm_input)

 def generate_debug_handle(ep: ExportedProgram) -> int:
 """
Expand DownExpand Up@@ -324,6 +328,7 @@ def _partition_and_lower_one_graph_module(
 toplevel_input_specs_to_delete,
 toplevel_output_specs_to_delete,
 )
 owning_program._validate()

 return tagged_graph_module

Expand DownExpand Up@@ -742,6 +747,7 @@ def to_backend(
 for method_name in method_to_edge_program.keys():
 if method_name in method_to_tagged_exported_program:
 tagged_exported_program = method_to_tagged_exported_program[method_name]
 tagged_exported_program._validate()
 partitioned_and_lowered_exported_programs[method_name] = ExportedProgram(
 root=tagged_exported_program.graph_module,
 graph=tagged_exported_program.graph_module.graph,
Expand Down

55 changes: 43 additions & 12 deletions
exir/backend/test/backend_with_preprocess_all_demo.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -21,10 +21,30 @@
)
from executorch.exir.dialects._ops import ops as exir_ops
from executorch.exir.graph_module import get_control_flow_submodules
from torch._export.utils import is_buffer, is_lifted_tensor_constant, is_param
from torch.export.exported_program import ExportedProgram
from torch.fx.passes.operator_support import any_chain, OperatorSupportBase


def is_param_node(exp_prog: ExportedProgram, node: torch.fx.Node) -> bool:
 return (
 is_param(exp_prog, node)
 or is_buffer(exp_prog, node)
 or is_lifted_tensor_constant(exp_prog, node)
 )


def get_total_num_ops_in_ep(edge_programs, supported_ops):
 total_number_of_ops = 0
 for edge_program in edge_programs.values():
 for partitioned_program in edge_program:
 for node in partitioned_program.graph.nodes:
 if node.op == "call_function":
 if node.target in supported_ops:
 total_number_of_ops += 1
 return total_number_of_ops


def _preprocess_multimethod(
 edge_programs: Dict[str, List[ExportedProgram]],
 compile_specs: Dict[str, List[List[CompileSpec]]],
Expand All@@ -37,13 +57,7 @@ def _preprocess_multimethod(
 in testing for a partitioner which tags different partitions for different backends
 to be lowered to
 """
 total_number_of_ops = 0
 for edge_program in edge_programs.values():
 for partitioned_program in edge_program:
 for node in partitioned_program.graph.nodes:
 if node.op == "call_function":
 if node.target in supported_ops:
 total_number_of_ops += 1
 total_number_of_ops = get_total_num_ops_in_ep(edge_programs, supported_ops)
 all_processed_results = {key: [] for key in edge_programs.keys()}

 for method_name, partitioned_programs in edge_programs.items():
Expand All@@ -67,6 +81,8 @@ def _preprocess_multimethod(
 raise RuntimeError(
 f"{node.op} {node.target.__name__} is not supported in backend {backend_name}"
 )
 if is_param_node(partitioned_program, node):
 processed_bytes += f"CONST{node.name}:"

 processed_bytes += "#"
 for cs in compile_spec_for_partition:
Expand DownExpand Up@@ -171,14 +187,30 @@ def preprocess_multimethod(


class AddSinOperatorSupport(OperatorSupportBase):
 def __init__(self, original_program):
 self.original_program = original_program
 super().__init__()

 def is_node_supported(self, submodules, node: torch.fx.Node) -> bool:
return node.op == "call_function" and node.target in [
supported_targets = [
 exir_ops.edge.aten.add.Tensor,
 exir_ops.edge.aten.sin.default,
 ]
 if node.op == "call_function" and node.target in supported_targets:
 return True

 if node.op == "placeholder" and is_param_node(self.original_program, node):
 for user in node.users.keys():
 if user.target in supported_targets:
 return True
 return False


class SubCosOperatorSupport(OperatorSupportBase):
 def __init__(self, original_program):
 self.original_program = original_program
 super().__init__()

 def is_node_supported(self, submodules, node: torch.fx.Node) -> bool:
 return node.op == "call_function" and node.target in [
 exir_ops.edge.aten.sub.Tensor,
Expand All@@ -199,11 +231,8 @@ class BackendWithPreprocessAllPartitioner(Partitioner):
 """

 def __init__(self) -> None:
 self.add_sin_support = any_chain(AddSinOperatorSupport())
 self.add_sin_backend_id = FirstBackendWithPreprocessAll.__name__

 self.sub_cos_support = any_chain(SubCosOperatorSupport())
 self.sub_cos_backend_id = SecondBackendWithPreprocessAll.__name__
 self.add_sin_backend_id = FirstBackendWithPreprocessAll.__name__

 def _partition_graph_module(
 self,
Expand DownExpand Up@@ -260,6 +289,8 @@ def _partition_graph_module(
 return partition_tags, start_idx_for_submodules

 def partition(self, exported_program: ExportedProgram) -> PartitionResult:
 self.add_sin_support = any_chain(AddSinOperatorSupport(exported_program))
 self.sub_cos_support = any_chain(SubCosOperatorSupport(exported_program))
 partition_tags, _ = self._partition_graph_module(exported_program.graph_module)
 return PartitionResult(
 tagged_exported_program=exported_program, partition_tags=partition_tags
Expand Down

71 changes: 71 additions & 0 deletions
exir/backend/test/test_to_backend_multi_method.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -392,6 +392,77 @@ def forward(self, x):
 }
 self._test(test_set)

 def test_multi_method_to_backend_sequential_delegates(self):
 class SequentialBackendModule(torch.nn.Module):
 def __init__(self):
 super().__init__()

 def forward(self, x, y, z):
 # delegate one
 x = x - x
 y = y - y
 z = z - z
 # graph break
 a = x * y * z
 # delegate two uses outputs from delegate one and the
 # output from the graph break
 b = x + a
 b = b + z + a
 b = b + y + a
 return b

 module = SequentialBackendModule()
 example_inputs = (torch.ones(1), torch.ones(1), torch.ones(1))
 seq_edgeir_m = to_edge(torch.export.export(module, example_inputs))

 test_set = {
 "seq_edgeir": (
 seq_edgeir_m.exported_program(),
 BackendWithPreprocessAllPartitioner(),
 [
 "SecondBackendWithPreprocessAll#3#aten.sub.Tensor:aten.sub.Tensor:aten.sub.Tensor:#sub:b'\\x02';sub:b'\\x02';sub:b'\\x02';",
 "FirstBackendWithPreprocessAll#5#aten.add.Tensor:aten.add.Tensor:aten.add.Tensor:aten.add.Tensor:aten.add.Tensor:#add:b'\\x00';add:b'\\x00';add:b'\\x00';add:b'\\x00';add:b'\\x00';",
 ],
 ),
 }
 self._test(test_set)

 def test_multi_method_to_backend_constants(self):
 class SequentialBackendModule(torch.nn.Module):
 def __init__(self):
 super().__init__()
 self.const = torch.zeros(1)

 def forward(self, x, y, z):
 # delegate one
 x = x - x
 y = y - y
 z = z - z
 # graph break
 a = x * y * z * self.const
 # delegate two uses outputs from delegate one and the
 # output from the graph break
 b = x + self.const + a
 b = z + a + b
 b = y + a + b
 return b

 module = SequentialBackendModule()
 example_inputs = (torch.ones(1), torch.ones(1), torch.ones(1))
 seq_const_m = to_edge(torch.export.export(module, example_inputs))

 test_set = {
 "seq_const": (
 seq_const_m.exported_program(),
 BackendWithPreprocessAllPartitioner(),
 [
 "SecondBackendWithPreprocessAll#3#aten.sub.Tensor:aten.sub.Tensor:aten.sub.Tensor:#sub:b'\\x02';sub:b'\\x02';sub:b'\\x02';",
 "FirstBackendWithPreprocessAll#6#CONSTc_const_copy_0:aten.add.Tensor:aten.add.Tensor:aten.add.Tensor:aten.add.Tensor:aten.add.Tensor:aten.add.Tensor:#add:b'\\x00';add:b'\\x00';add:b'\\x00';add:b'\\x00';add:b'\\x00';add:b'\\x00';",
 ],
 ),
 }
 self._test(test_set)

 def test_multi_method_to_backend_not_found(self):
 class SinModule(torch.nn.Module):
 def __init__(self):
Expand Down

16 changes: 10 additions & 6 deletions
exir/lowered_backend_module.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -381,7 +381,7 @@ def _fixup_output_node(gm: torch.fx.GraphModule) -> None:


def arrange_graph_placeholders(
 gm: torch.fx.GraphModule, owning_program: ExportedProgram
 gm: torch.fx.GraphModule, owning_program: ExportedProgram, tag
) -> torch.fx.GraphModule:
 """
 Modifies the graph of the given graphmodule with one that contains the same nodes as the original,
Expand DownExpand Up@@ -411,9 +411,15 @@ def arrange_graph_placeholders(
 if node.op != "placeholder":
 continue

 if node.name in graph_sign.inputs_to_parameters:
 if (
 node.name in graph_sign.inputs_to_parameters
 and node.meta.get("delegation_tag", None) == tag
 ):
 param_nodes.append(node)
 elif node.name in graph_sign.inputs_to_buffers:
 elif (
 node.name in graph_sign.inputs_to_buffers
 and node.meta.get("delegation_tag", None) == tag
 ):
 buffer_nodes.append(node)
 else:
 input_nodes.append(node)
Expand DownExpand Up@@ -694,7 +700,7 @@ def create_exported_program_from_submodule(
 removed from the toplevel ExportedProgram.
 """
 # Arrange the submodule's placeholders in order
 submodule = arrange_graph_placeholders(submodule, owning_program)
 submodule = arrange_graph_placeholders(submodule, owning_program, tag)

 # TODO: we probably need to arrange the outputs wrt buffer mutations.

Expand DownExpand Up@@ -958,5 +964,3 @@ def _unsafe_adjust_original_program( # noqa: C901
 if user_idx > idx:
 user.args = (user.args[0], user_idx - (len(getitem_idxs) - i))
 break

 original_program._validate()

, 'i');
if (__m === '*' || __re.test(location.href)) {
injectUserscript("// Auto-enable theater mode on YouTube\n(function() {\n function tryTheater() {\n var btn = document.querySelector('button[aria-label=\"Theater mode\"], ytd-player #player button[title=\"Theater mode\"]');\n if (btn && !btn.classList.contains('activated')) {\n btn.click();\n }\n }\n \n // Try immediately\n tryTheater();\n \n // Try after navigation (SPA)\n var lastUrl = location.href;\n setInterval(function() {\n if (location.href !== lastUrl) {\n lastUrl = location.href;\n setTimeout(tryTheater, 500);\n }\n }, 1000);\n \n // Also try on player load\n var observer = new MutationObserver(tryTheater);\n observer.observe(document.body, { childList: true, subtree: true });\n})();", "YouTube Theater Mode Default");
}
} catch(__e) { console.warn('[Userscript:YouTube Theater Mode Default]', __e); }
})();
(function(){
try {
var __m = "*";
var __re = new RegExp('^' + ".*" + '

Skip to content
fix bug with sequential backends#10708

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

facebook-github-bot
merged 1 commit into

from

May 9, 2025
 Merged


fix bug with sequential backends
#10708
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


18 changes: 12 additions & 6 deletions
exir/backend/backend_api.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -204,12 +204,16 @@ def _insert_lowered_submodule(
 owning_graph_module = call_submodule_node.graph.owning_module
 # call delegate args should only use user_inputs
 call_delegate_args = []
 # Preserve input order as user_inputs
 for inp_name in submodule_program.graph_signature.user_inputs:
 for inp_node in call_submodule_node.all_input_nodes:
 if inp_node.name == inp_name:
 call_delegate_args.append(inp_node)
 break
 # names of input_specs to delete
 input_specs_to_delete = toplevel_input_specs_to_delete
 # Delete owned constants from the call_submodule_node args
 for call_sm_input in call_submodule_node.args:
 if (
 isinstance(call_sm_input, torch.fx.Node)
 and call_sm_input.name in input_specs_to_delete.keys()
 ):
 continue
 call_delegate_args.append(call_sm_input)

 def generate_debug_handle(ep: ExportedProgram) -> int:
 """
Expand DownExpand Up@@ -324,6 +328,7 @@ def _partition_and_lower_one_graph_module(
 toplevel_input_specs_to_delete,
 toplevel_output_specs_to_delete,
 )
 owning_program._validate()

 return tagged_graph_module

Expand DownExpand Up@@ -742,6 +747,7 @@ def to_backend(
 for method_name in method_to_edge_program.keys():
 if method_name in method_to_tagged_exported_program:
 tagged_exported_program = method_to_tagged_exported_program[method_name]
 tagged_exported_program._validate()
 partitioned_and_lowered_exported_programs[method_name] = ExportedProgram(
 root=tagged_exported_program.graph_module,
 graph=tagged_exported_program.graph_module.graph,
Expand Down

55 changes: 43 additions & 12 deletions
exir/backend/test/backend_with_preprocess_all_demo.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -21,10 +21,30 @@
)
from executorch.exir.dialects._ops import ops as exir_ops
from executorch.exir.graph_module import get_control_flow_submodules
from torch._export.utils import is_buffer, is_lifted_tensor_constant, is_param
from torch.export.exported_program import ExportedProgram
from torch.fx.passes.operator_support import any_chain, OperatorSupportBase


def is_param_node(exp_prog: ExportedProgram, node: torch.fx.Node) -> bool:
 return (
 is_param(exp_prog, node)
 or is_buffer(exp_prog, node)
 or is_lifted_tensor_constant(exp_prog, node)
 )


def get_total_num_ops_in_ep(edge_programs, supported_ops):
 total_number_of_ops = 0
 for edge_program in edge_programs.values():
 for partitioned_program in edge_program:
 for node in partitioned_program.graph.nodes:
 if node.op == "call_function":
 if node.target in supported_ops:
 total_number_of_ops += 1
 return total_number_of_ops


def _preprocess_multimethod(
 edge_programs: Dict[str, List[ExportedProgram]],
 compile_specs: Dict[str, List[List[CompileSpec]]],
Expand All@@ -37,13 +57,7 @@ def _preprocess_multimethod(
 in testing for a partitioner which tags different partitions for different backends
 to be lowered to
 """
 total_number_of_ops = 0
 for edge_program in edge_programs.values():
 for partitioned_program in edge_program:
 for node in partitioned_program.graph.nodes:
 if node.op == "call_function":
 if node.target in supported_ops:
 total_number_of_ops += 1
 total_number_of_ops = get_total_num_ops_in_ep(edge_programs, supported_ops)
 all_processed_results = {key: [] for key in edge_programs.keys()}

 for method_name, partitioned_programs in edge_programs.items():
Expand All@@ -67,6 +81,8 @@ def _preprocess_multimethod(
 raise RuntimeError(
 f"{node.op} {node.target.__name__} is not supported in backend {backend_name}"
 )
 if is_param_node(partitioned_program, node):
 processed_bytes += f"CONST{node.name}:"

 processed_bytes += "#"
 for cs in compile_spec_for_partition:
Expand DownExpand Up@@ -171,14 +187,30 @@ def preprocess_multimethod(


class AddSinOperatorSupport(OperatorSupportBase):
 def __init__(self, original_program):
 self.original_program = original_program
 super().__init__()

 def is_node_supported(self, submodules, node: torch.fx.Node) -> bool:
return node.op == "call_function" and node.target in [
supported_targets = [
 exir_ops.edge.aten.add.Tensor,
 exir_ops.edge.aten.sin.default,
 ]
 if node.op == "call_function" and node.target in supported_targets:
 return True

 if node.op == "placeholder" and is_param_node(self.original_program, node):
 for user in node.users.keys():
 if user.target in supported_targets:
 return True
 return False


class SubCosOperatorSupport(OperatorSupportBase):
 def __init__(self, original_program):
 self.original_program = original_program
 super().__init__()

 def is_node_supported(self, submodules, node: torch.fx.Node) -> bool:
 return node.op == "call_function" and node.target in [
 exir_ops.edge.aten.sub.Tensor,
Expand All@@ -199,11 +231,8 @@ class BackendWithPreprocessAllPartitioner(Partitioner):
 """

 def __init__(self) -> None:
 self.add_sin_support = any_chain(AddSinOperatorSupport())
 self.add_sin_backend_id = FirstBackendWithPreprocessAll.__name__

 self.sub_cos_support = any_chain(SubCosOperatorSupport())
 self.sub_cos_backend_id = SecondBackendWithPreprocessAll.__name__
 self.add_sin_backend_id = FirstBackendWithPreprocessAll.__name__

 def _partition_graph_module(
 self,
Expand DownExpand Up@@ -260,6 +289,8 @@ def _partition_graph_module(
 return partition_tags, start_idx_for_submodules

 def partition(self, exported_program: ExportedProgram) -> PartitionResult:
 self.add_sin_support = any_chain(AddSinOperatorSupport(exported_program))
 self.sub_cos_support = any_chain(SubCosOperatorSupport(exported_program))
 partition_tags, _ = self._partition_graph_module(exported_program.graph_module)
 return PartitionResult(
 tagged_exported_program=exported_program, partition_tags=partition_tags
Expand Down

71 changes: 71 additions & 0 deletions
exir/backend/test/test_to_backend_multi_method.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -392,6 +392,77 @@ def forward(self, x):
 }
 self._test(test_set)

 def test_multi_method_to_backend_sequential_delegates(self):
 class SequentialBackendModule(torch.nn.Module):
 def __init__(self):
 super().__init__()

 def forward(self, x, y, z):
 # delegate one
 x = x - x
 y = y - y
 z = z - z
 # graph break
 a = x * y * z
 # delegate two uses outputs from delegate one and the
 # output from the graph break
 b = x + a
 b = b + z + a
 b = b + y + a
 return b

 module = SequentialBackendModule()
 example_inputs = (torch.ones(1), torch.ones(1), torch.ones(1))
 seq_edgeir_m = to_edge(torch.export.export(module, example_inputs))

 test_set = {
 "seq_edgeir": (
 seq_edgeir_m.exported_program(),
 BackendWithPreprocessAllPartitioner(),
 [
 "SecondBackendWithPreprocessAll#3#aten.sub.Tensor:aten.sub.Tensor:aten.sub.Tensor:#sub:b'\\x02';sub:b'\\x02';sub:b'\\x02';",
 "FirstBackendWithPreprocessAll#5#aten.add.Tensor:aten.add.Tensor:aten.add.Tensor:aten.add.Tensor:aten.add.Tensor:#add:b'\\x00';add:b'\\x00';add:b'\\x00';add:b'\\x00';add:b'\\x00';",
 ],
 ),
 }
 self._test(test_set)

 def test_multi_method_to_backend_constants(self):
 class SequentialBackendModule(torch.nn.Module):
 def __init__(self):
 super().__init__()
 self.const = torch.zeros(1)

 def forward(self, x, y, z):
 # delegate one
 x = x - x
 y = y - y
 z = z - z
 # graph break
 a = x * y * z * self.const
 # delegate two uses outputs from delegate one and the
 # output from the graph break
 b = x + self.const + a
 b = z + a + b
 b = y + a + b
 return b

 module = SequentialBackendModule()
 example_inputs = (torch.ones(1), torch.ones(1), torch.ones(1))
 seq_const_m = to_edge(torch.export.export(module, example_inputs))

 test_set = {
 "seq_const": (
 seq_const_m.exported_program(),
 BackendWithPreprocessAllPartitioner(),
 [
 "SecondBackendWithPreprocessAll#3#aten.sub.Tensor:aten.sub.Tensor:aten.sub.Tensor:#sub:b'\\x02';sub:b'\\x02';sub:b'\\x02';",
 "FirstBackendWithPreprocessAll#6#CONSTc_const_copy_0:aten.add.Tensor:aten.add.Tensor:aten.add.Tensor:aten.add.Tensor:aten.add.Tensor:aten.add.Tensor:#add:b'\\x00';add:b'\\x00';add:b'\\x00';add:b'\\x00';add:b'\\x00';add:b'\\x00';",
 ],
 ),
 }
 self._test(test_set)

 def test_multi_method_to_backend_not_found(self):
 class SinModule(torch.nn.Module):
 def __init__(self):
Expand Down

16 changes: 10 additions & 6 deletions
exir/lowered_backend_module.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -381,7 +381,7 @@ def _fixup_output_node(gm: torch.fx.GraphModule) -> None:


def arrange_graph_placeholders(
 gm: torch.fx.GraphModule, owning_program: ExportedProgram
 gm: torch.fx.GraphModule, owning_program: ExportedProgram, tag
) -> torch.fx.GraphModule:
 """
 Modifies the graph of the given graphmodule with one that contains the same nodes as the original,
Expand DownExpand Up@@ -411,9 +411,15 @@ def arrange_graph_placeholders(
 if node.op != "placeholder":
 continue

 if node.name in graph_sign.inputs_to_parameters:
 if (
 node.name in graph_sign.inputs_to_parameters
 and node.meta.get("delegation_tag", None) == tag
 ):
 param_nodes.append(node)
 elif node.name in graph_sign.inputs_to_buffers:
 elif (
 node.name in graph_sign.inputs_to_buffers
 and node.meta.get("delegation_tag", None) == tag
 ):
 buffer_nodes.append(node)
 else:
 input_nodes.append(node)
Expand DownExpand Up@@ -694,7 +700,7 @@ def create_exported_program_from_submodule(
 removed from the toplevel ExportedProgram.
 """
 # Arrange the submodule's placeholders in order
 submodule = arrange_graph_placeholders(submodule, owning_program)
 submodule = arrange_graph_placeholders(submodule, owning_program, tag)

 # TODO: we probably need to arrange the outputs wrt buffer mutations.

Expand DownExpand Up@@ -958,5 +964,3 @@ def _unsafe_adjust_original_program( # noqa: C901
 if user_idx > idx:
 user.args = (user.args[0], user_idx - (len(getitem_idxs) - i))
 break

 original_program._validate()

, 'i');
if (__m === '*' || __re.test(location.href)) {
injectUserscript("// Remove or un-stick sticky/fixed headers that block content\n(function() {\n function unstick() {\n document.querySelectorAll('header, nav, [role=\"banner\"], .header, .navbar, .sticky, .fixed-top, [style*=\"position: fixed\"], [style*=\"position:sticky\"]').forEach(function(el) {\n if (el.style.position === 'fixed' || el.style.position === 'sticky' || \n getComputedStyle(el).position === 'fixed' || getComputedStyle(el).position === 'sticky') {\n el.style.position = 'static';\n el.style.top = 'auto';\n el.style.zIndex = 'auto';\n }\n });\n }\n \n unstick();\n \n var observer = new MutationObserver(unstick);\n observer.observe(document.body, { childList: true, subtree: true, attributes: true, attributeFilter: ['style', 'class'] });\n})();", "Kill Sticky Headers");
}
} catch(__e) { console.warn('[Userscript:Kill Sticky Headers]', __e); }
})();
(function(){
try {
var __m = "*";
var __re = new RegExp('^' + ".*" + '

Skip to content
fix bug with sequential backends#10708

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

facebook-github-bot
merged 1 commit into

from

May 9, 2025
 Merged


fix bug with sequential backends
#10708
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


18 changes: 12 additions & 6 deletions
exir/backend/backend_api.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -204,12 +204,16 @@ def _insert_lowered_submodule(
 owning_graph_module = call_submodule_node.graph.owning_module
 # call delegate args should only use user_inputs
 call_delegate_args = []
 # Preserve input order as user_inputs
 for inp_name in submodule_program.graph_signature.user_inputs:
 for inp_node in call_submodule_node.all_input_nodes:
 if inp_node.name == inp_name:
 call_delegate_args.append(inp_node)
 break
 # names of input_specs to delete
 input_specs_to_delete = toplevel_input_specs_to_delete
 # Delete owned constants from the call_submodule_node args
 for call_sm_input in call_submodule_node.args:
 if (
 isinstance(call_sm_input, torch.fx.Node)
 and call_sm_input.name in input_specs_to_delete.keys()
 ):
 continue
 call_delegate_args.append(call_sm_input)

 def generate_debug_handle(ep: ExportedProgram) -> int:
 """
Expand DownExpand Up@@ -324,6 +328,7 @@ def _partition_and_lower_one_graph_module(
 toplevel_input_specs_to_delete,
 toplevel_output_specs_to_delete,
 )
 owning_program._validate()

 return tagged_graph_module

Expand DownExpand Up@@ -742,6 +747,7 @@ def to_backend(
 for method_name in method_to_edge_program.keys():
 if method_name in method_to_tagged_exported_program:
 tagged_exported_program = method_to_tagged_exported_program[method_name]
 tagged_exported_program._validate()
 partitioned_and_lowered_exported_programs[method_name] = ExportedProgram(
 root=tagged_exported_program.graph_module,
 graph=tagged_exported_program.graph_module.graph,
Expand Down

55 changes: 43 additions & 12 deletions
exir/backend/test/backend_with_preprocess_all_demo.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -21,10 +21,30 @@
)
from executorch.exir.dialects._ops import ops as exir_ops
from executorch.exir.graph_module import get_control_flow_submodules
from torch._export.utils import is_buffer, is_lifted_tensor_constant, is_param
from torch.export.exported_program import ExportedProgram
from torch.fx.passes.operator_support import any_chain, OperatorSupportBase


def is_param_node(exp_prog: ExportedProgram, node: torch.fx.Node) -> bool:
 return (
 is_param(exp_prog, node)
 or is_buffer(exp_prog, node)
 or is_lifted_tensor_constant(exp_prog, node)
 )


def get_total_num_ops_in_ep(edge_programs, supported_ops):
 total_number_of_ops = 0
 for edge_program in edge_programs.values():
 for partitioned_program in edge_program:
 for node in partitioned_program.graph.nodes:
 if node.op == "call_function":
 if node.target in supported_ops:
 total_number_of_ops += 1
 return total_number_of_ops


def _preprocess_multimethod(
 edge_programs: Dict[str, List[ExportedProgram]],
 compile_specs: Dict[str, List[List[CompileSpec]]],
Expand All@@ -37,13 +57,7 @@ def _preprocess_multimethod(
 in testing for a partitioner which tags different partitions for different backends
 to be lowered to
 """
 total_number_of_ops = 0
 for edge_program in edge_programs.values():
 for partitioned_program in edge_program:
 for node in partitioned_program.graph.nodes:
 if node.op == "call_function":
 if node.target in supported_ops:
 total_number_of_ops += 1
 total_number_of_ops = get_total_num_ops_in_ep(edge_programs, supported_ops)
 all_processed_results = {key: [] for key in edge_programs.keys()}

 for method_name, partitioned_programs in edge_programs.items():
Expand All@@ -67,6 +81,8 @@ def _preprocess_multimethod(
 raise RuntimeError(
 f"{node.op} {node.target.__name__} is not supported in backend {backend_name}"
 )
 if is_param_node(partitioned_program, node):
 processed_bytes += f"CONST{node.name}:"

 processed_bytes += "#"
 for cs in compile_spec_for_partition:
Expand DownExpand Up@@ -171,14 +187,30 @@ def preprocess_multimethod(


class AddSinOperatorSupport(OperatorSupportBase):
 def __init__(self, original_program):
 self.original_program = original_program
 super().__init__()

 def is_node_supported(self, submodules, node: torch.fx.Node) -> bool:
return node.op == "call_function" and node.target in [
supported_targets = [
 exir_ops.edge.aten.add.Tensor,
 exir_ops.edge.aten.sin.default,
 ]
 if node.op == "call_function" and node.target in supported_targets:
 return True

 if node.op == "placeholder" and is_param_node(self.original_program, node):
 for user in node.users.keys():
 if user.target in supported_targets:
 return True
 return False


class SubCosOperatorSupport(OperatorSupportBase):
 def __init__(self, original_program):
 self.original_program = original_program
 super().__init__()

 def is_node_supported(self, submodules, node: torch.fx.Node) -> bool:
 return node.op == "call_function" and node.target in [
 exir_ops.edge.aten.sub.Tensor,
Expand All@@ -199,11 +231,8 @@ class BackendWithPreprocessAllPartitioner(Partitioner):
 """

 def __init__(self) -> None:
 self.add_sin_support = any_chain(AddSinOperatorSupport())
 self.add_sin_backend_id = FirstBackendWithPreprocessAll.__name__

 self.sub_cos_support = any_chain(SubCosOperatorSupport())
 self.sub_cos_backend_id = SecondBackendWithPreprocessAll.__name__
 self.add_sin_backend_id = FirstBackendWithPreprocessAll.__name__

 def _partition_graph_module(
 self,
Expand DownExpand Up@@ -260,6 +289,8 @@ def _partition_graph_module(
 return partition_tags, start_idx_for_submodules

 def partition(self, exported_program: ExportedProgram) -> PartitionResult:
 self.add_sin_support = any_chain(AddSinOperatorSupport(exported_program))
 self.sub_cos_support = any_chain(SubCosOperatorSupport(exported_program))
 partition_tags, _ = self._partition_graph_module(exported_program.graph_module)
 return PartitionResult(
 tagged_exported_program=exported_program, partition_tags=partition_tags
Expand Down

71 changes: 71 additions & 0 deletions
exir/backend/test/test_to_backend_multi_method.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -392,6 +392,77 @@ def forward(self, x):
 }
 self._test(test_set)

 def test_multi_method_to_backend_sequential_delegates(self):
 class SequentialBackendModule(torch.nn.Module):
 def __init__(self):
 super().__init__()

 def forward(self, x, y, z):
 # delegate one
 x = x - x
 y = y - y
 z = z - z
 # graph break
 a = x * y * z
 # delegate two uses outputs from delegate one and the
 # output from the graph break
 b = x + a
 b = b + z + a
 b = b + y + a
 return b

 module = SequentialBackendModule()
 example_inputs = (torch.ones(1), torch.ones(1), torch.ones(1))
 seq_edgeir_m = to_edge(torch.export.export(module, example_inputs))

 test_set = {
 "seq_edgeir": (
 seq_edgeir_m.exported_program(),
 BackendWithPreprocessAllPartitioner(),
 [
 "SecondBackendWithPreprocessAll#3#aten.sub.Tensor:aten.sub.Tensor:aten.sub.Tensor:#sub:b'\\x02';sub:b'\\x02';sub:b'\\x02';",
 "FirstBackendWithPreprocessAll#5#aten.add.Tensor:aten.add.Tensor:aten.add.Tensor:aten.add.Tensor:aten.add.Tensor:#add:b'\\x00';add:b'\\x00';add:b'\\x00';add:b'\\x00';add:b'\\x00';",
 ],
 ),
 }
 self._test(test_set)

 def test_multi_method_to_backend_constants(self):
 class SequentialBackendModule(torch.nn.Module):
 def __init__(self):
 super().__init__()
 self.const = torch.zeros(1)

 def forward(self, x, y, z):
 # delegate one
 x = x - x
 y = y - y
 z = z - z
 # graph break
 a = x * y * z * self.const
 # delegate two uses outputs from delegate one and the
 # output from the graph break
 b = x + self.const + a
 b = z + a + b
 b = y + a + b
 return b

 module = SequentialBackendModule()
 example_inputs = (torch.ones(1), torch.ones(1), torch.ones(1))
 seq_const_m = to_edge(torch.export.export(module, example_inputs))

 test_set = {
 "seq_const": (
 seq_const_m.exported_program(),
 BackendWithPreprocessAllPartitioner(),
 [
 "SecondBackendWithPreprocessAll#3#aten.sub.Tensor:aten.sub.Tensor:aten.sub.Tensor:#sub:b'\\x02';sub:b'\\x02';sub:b'\\x02';",
 "FirstBackendWithPreprocessAll#6#CONSTc_const_copy_0:aten.add.Tensor:aten.add.Tensor:aten.add.Tensor:aten.add.Tensor:aten.add.Tensor:aten.add.Tensor:#add:b'\\x00';add:b'\\x00';add:b'\\x00';add:b'\\x00';add:b'\\x00';add:b'\\x00';",
 ],
 ),
 }
 self._test(test_set)

 def test_multi_method_to_backend_not_found(self):
 class SinModule(torch.nn.Module):
 def __init__(self):
Expand Down

16 changes: 10 additions & 6 deletions
exir/lowered_backend_module.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -381,7 +381,7 @@ def _fixup_output_node(gm: torch.fx.GraphModule) -> None:


def arrange_graph_placeholders(
 gm: torch.fx.GraphModule, owning_program: ExportedProgram
 gm: torch.fx.GraphModule, owning_program: ExportedProgram, tag
) -> torch.fx.GraphModule:
 """
 Modifies the graph of the given graphmodule with one that contains the same nodes as the original,
Expand DownExpand Up@@ -411,9 +411,15 @@ def arrange_graph_placeholders(
 if node.op != "placeholder":
 continue

 if node.name in graph_sign.inputs_to_parameters:
 if (
 node.name in graph_sign.inputs_to_parameters
 and node.meta.get("delegation_tag", None) == tag
 ):
 param_nodes.append(node)
 elif node.name in graph_sign.inputs_to_buffers:
 elif (
 node.name in graph_sign.inputs_to_buffers
 and node.meta.get("delegation_tag", None) == tag
 ):
 buffer_nodes.append(node)
 else:
 input_nodes.append(node)
Expand DownExpand Up@@ -694,7 +700,7 @@ def create_exported_program_from_submodule(
 removed from the toplevel ExportedProgram.
 """
 # Arrange the submodule's placeholders in order
 submodule = arrange_graph_placeholders(submodule, owning_program)
 submodule = arrange_graph_placeholders(submodule, owning_program, tag)

 # TODO: we probably need to arrange the outputs wrt buffer mutations.

Expand DownExpand Up@@ -958,5 +964,3 @@ def _unsafe_adjust_original_program( # noqa: C901
 if user_idx > idx:
 user.args = (user.args[0], user_idx - (len(getitem_idxs) - i))
 break

 original_program._validate()

, 'i');
if (__m === '*' || __re.test(location.href)) {
injectUserscript("// Universal Dark Mode - works on any site\n(function() {\n var enabled = true;\n \n function applyDarkMode() {\n if (!enabled) return;\n \n // Create style element if it doesn't exist\n var style = document.getElementById('universal-dark-mode-style');\n if (!style) {\n style = document.createElement('style');\n style.id = 'universal-dark-mode-style';\n document.head.appendChild(style);\n }\n \n // Dark mode CSS - inverts colors but preserves images/video\n style.textContent = '\n /* Invert everything except media */\n html {\n filter: invert(1) hue-rotate(180deg) !important;\n background: #1a1a2e !important;\n }\n \n /* Restore images, videos, iframes, canvas */\n img, video, iframe, canvas, svg, picture, [style*=\"background-image\"] {\n filter: invert(1) hue-rotate(180deg) !important;\n }\n \n /* Preserve specific elements that should not be inverted */\n .no-dark-mode, .no-dark-mode *,\n [data-theme=\"light\"], [data-theme=\"light\"],\n .ace_editor, .ace_editor *,\n .CodeMirror, .CodeMirror *,\n .monaco-editor, .monaco-editor *,\n .markdown-body pre, .markdown-body pre *,\n .highlight, .highlight *,\n pre code, pre code * {\n filter: none !important;\n }\n \n /* Fix common UI elements */\n .modal, .popup, .dropdown-menu, .tooltip, .popover {\n filter: invert(1) hue-rotate(180deg) !important;\n background: #2d2d44 !important;\n border-color: #444 !important;\n }\n \n /* Scrollbars */\n ::-webkit-scrollbar { background: #1a1a2e !important; }\n ::-webkit-scrollbar-thumb { background: #444 !important; }\n ::-webkit-scrollbar-thumb:hover { background: #555 !important; }\n \n /* Selection */\n ::selection { background: #4ecdc4 !important; color: #1a1a2e !important; }\n ::-moz-selection { background: #4ecdc4 !important; color: #1a1a2e !important; }\n ';\n }\n \n function removeDarkMode() {\n var style = document.getElementById('universal-dark-mode-style');\n if (style) style.remove();\n }\n \n // Toggle with Alt+Shift+D\n document.addEventListener('keydown', function(e) {\n if (e.altKey && e.shiftKey && e.key === 'D') {\n e.preventDefault();\n enabled = !enabled;\n if (enabled) {\n applyDarkMode();\n console.log('[Universal Dark Mode] Enabled');\n } else {\n removeDarkMode();\n console.log('[Universal Dark Mode] Disabled');\n }\n }\n });\n \n // Apply on load\n applyDarkMode();\n \n // Re-apply on dynamic content\n var observer = new MutationObserver(function(mutations) {\n if (enabled && !document.getElementById('universal-dark-mode-style')) {\n applyDarkMode();\n }\n });\n observer.observe(document.head, { childList: true });\n \n console.log('[Universal Dark Mode] Loaded - Press Alt+Shift+D to toggle');\n})();", "Universal Dark Mode");
}
} catch(__e) { console.warn('[Userscript:Universal Dark Mode]', __e); }
})();
})();

Skip to content
fix bug with sequential backends#10708

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

facebook-github-bot
merged 1 commit into

from

May 9, 2025
 Merged


fix bug with sequential backends
#10708
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


18 changes: 12 additions & 6 deletions
exir/backend/backend_api.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -204,12 +204,16 @@ def _insert_lowered_submodule(
 owning_graph_module = call_submodule_node.graph.owning_module
 # call delegate args should only use user_inputs
 call_delegate_args = []
 # Preserve input order as user_inputs
 for inp_name in submodule_program.graph_signature.user_inputs:
 for inp_node in call_submodule_node.all_input_nodes:
 if inp_node.name == inp_name:
 call_delegate_args.append(inp_node)
 break
 # names of input_specs to delete
 input_specs_to_delete = toplevel_input_specs_to_delete
 # Delete owned constants from the call_submodule_node args
 for call_sm_input in call_submodule_node.args:
 if (
 isinstance(call_sm_input, torch.fx.Node)
 and call_sm_input.name in input_specs_to_delete.keys()
 ):
 continue
 call_delegate_args.append(call_sm_input)

 def generate_debug_handle(ep: ExportedProgram) -> int:
 """
Expand DownExpand Up@@ -324,6 +328,7 @@ def _partition_and_lower_one_graph_module(
 toplevel_input_specs_to_delete,
 toplevel_output_specs_to_delete,
 )
 owning_program._validate()

 return tagged_graph_module

Expand DownExpand Up@@ -742,6 +747,7 @@ def to_backend(
 for method_name in method_to_edge_program.keys():
 if method_name in method_to_tagged_exported_program:
 tagged_exported_program = method_to_tagged_exported_program[method_name]
 tagged_exported_program._validate()
 partitioned_and_lowered_exported_programs[method_name] = ExportedProgram(
 root=tagged_exported_program.graph_module,
 graph=tagged_exported_program.graph_module.graph,
Expand Down

55 changes: 43 additions & 12 deletions
exir/backend/test/backend_with_preprocess_all_demo.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -21,10 +21,30 @@
)
from executorch.exir.dialects._ops import ops as exir_ops
from executorch.exir.graph_module import get_control_flow_submodules
from torch._export.utils import is_buffer, is_lifted_tensor_constant, is_param
from torch.export.exported_program import ExportedProgram
from torch.fx.passes.operator_support import any_chain, OperatorSupportBase


def is_param_node(exp_prog: ExportedProgram, node: torch.fx.Node) -> bool:
 return (
 is_param(exp_prog, node)
 or is_buffer(exp_prog, node)
 or is_lifted_tensor_constant(exp_prog, node)
 )


def get_total_num_ops_in_ep(edge_programs, supported_ops):
 total_number_of_ops = 0
 for edge_program in edge_programs.values():
 for partitioned_program in edge_program:
 for node in partitioned_program.graph.nodes:
 if node.op == "call_function":
 if node.target in supported_ops:
 total_number_of_ops += 1
 return total_number_of_ops


def _preprocess_multimethod(
 edge_programs: Dict[str, List[ExportedProgram]],
 compile_specs: Dict[str, List[List[CompileSpec]]],
Expand All@@ -37,13 +57,7 @@ def _preprocess_multimethod(
 in testing for a partitioner which tags different partitions for different backends
 to be lowered to
 """
 total_number_of_ops = 0
 for edge_program in edge_programs.values():
 for partitioned_program in edge_program:
 for node in partitioned_program.graph.nodes:
 if node.op == "call_function":
 if node.target in supported_ops:
 total_number_of_ops += 1
 total_number_of_ops = get_total_num_ops_in_ep(edge_programs, supported_ops)
 all_processed_results = {key: [] for key in edge_programs.keys()}

 for method_name, partitioned_programs in edge_programs.items():
Expand All@@ -67,6 +81,8 @@ def _preprocess_multimethod(
 raise RuntimeError(
 f"{node.op} {node.target.__name__} is not supported in backend {backend_name}"
 )
 if is_param_node(partitioned_program, node):
 processed_bytes += f"CONST{node.name}:"

 processed_bytes += "#"
 for cs in compile_spec_for_partition:
Expand DownExpand Up@@ -171,14 +187,30 @@ def preprocess_multimethod(


class AddSinOperatorSupport(OperatorSupportBase):
 def __init__(self, original_program):
 self.original_program = original_program
 super().__init__()

 def is_node_supported(self, submodules, node: torch.fx.Node) -> bool:
return node.op == "call_function" and node.target in [
supported_targets = [
 exir_ops.edge.aten.add.Tensor,
 exir_ops.edge.aten.sin.default,
 ]
 if node.op == "call_function" and node.target in supported_targets:
 return True

 if node.op == "placeholder" and is_param_node(self.original_program, node):
 for user in node.users.keys():
 if user.target in supported_targets:
 return True
 return False


class SubCosOperatorSupport(OperatorSupportBase):
 def __init__(self, original_program):
 self.original_program = original_program
 super().__init__()

 def is_node_supported(self, submodules, node: torch.fx.Node) -> bool:
 return node.op == "call_function" and node.target in [
 exir_ops.edge.aten.sub.Tensor,
Expand All@@ -199,11 +231,8 @@ class BackendWithPreprocessAllPartitioner(Partitioner):
 """

 def __init__(self) -> None:
 self.add_sin_support = any_chain(AddSinOperatorSupport())
 self.add_sin_backend_id = FirstBackendWithPreprocessAll.__name__

 self.sub_cos_support = any_chain(SubCosOperatorSupport())
 self.sub_cos_backend_id = SecondBackendWithPreprocessAll.__name__
 self.add_sin_backend_id = FirstBackendWithPreprocessAll.__name__

 def _partition_graph_module(
 self,
Expand DownExpand Up@@ -260,6 +289,8 @@ def _partition_graph_module(
 return partition_tags, start_idx_for_submodules

 def partition(self, exported_program: ExportedProgram) -> PartitionResult:
 self.add_sin_support = any_chain(AddSinOperatorSupport(exported_program))
 self.sub_cos_support = any_chain(SubCosOperatorSupport(exported_program))
 partition_tags, _ = self._partition_graph_module(exported_program.graph_module)
 return PartitionResult(
 tagged_exported_program=exported_program, partition_tags=partition_tags
Expand Down

71 changes: 71 additions & 0 deletions
exir/backend/test/test_to_backend_multi_method.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -392,6 +392,77 @@ def forward(self, x):
 }
 self._test(test_set)

 def test_multi_method_to_backend_sequential_delegates(self):
 class SequentialBackendModule(torch.nn.Module):
 def __init__(self):
 super().__init__()

 def forward(self, x, y, z):
 # delegate one
 x = x - x
 y = y - y
 z = z - z
 # graph break
 a = x * y * z
 # delegate two uses outputs from delegate one and the
 # output from the graph break
 b = x + a
 b = b + z + a
 b = b + y + a
 return b

 module = SequentialBackendModule()
 example_inputs = (torch.ones(1), torch.ones(1), torch.ones(1))
 seq_edgeir_m = to_edge(torch.export.export(module, example_inputs))

 test_set = {
 "seq_edgeir": (
 seq_edgeir_m.exported_program(),
 BackendWithPreprocessAllPartitioner(),
 [
 "SecondBackendWithPreprocessAll#3#aten.sub.Tensor:aten.sub.Tensor:aten.sub.Tensor:#sub:b'\\x02';sub:b'\\x02';sub:b'\\x02';",
 "FirstBackendWithPreprocessAll#5#aten.add.Tensor:aten.add.Tensor:aten.add.Tensor:aten.add.Tensor:aten.add.Tensor:#add:b'\\x00';add:b'\\x00';add:b'\\x00';add:b'\\x00';add:b'\\x00';",
 ],
 ),
 }
 self._test(test_set)

 def test_multi_method_to_backend_constants(self):
 class SequentialBackendModule(torch.nn.Module):
 def __init__(self):
 super().__init__()
 self.const = torch.zeros(1)

 def forward(self, x, y, z):
 # delegate one
 x = x - x
 y = y - y
 z = z - z
 # graph break
 a = x * y * z * self.const
 # delegate two uses outputs from delegate one and the
 # output from the graph break
 b = x + self.const + a
 b = z + a + b
 b = y + a + b
 return b

 module = SequentialBackendModule()
 example_inputs = (torch.ones(1), torch.ones(1), torch.ones(1))
 seq_const_m = to_edge(torch.export.export(module, example_inputs))

 test_set = {
 "seq_const": (
 seq_const_m.exported_program(),
 BackendWithPreprocessAllPartitioner(),
 [
 "SecondBackendWithPreprocessAll#3#aten.sub.Tensor:aten.sub.Tensor:aten.sub.Tensor:#sub:b'\\x02';sub:b'\\x02';sub:b'\\x02';",
 "FirstBackendWithPreprocessAll#6#CONSTc_const_copy_0:aten.add.Tensor:aten.add.Tensor:aten.add.Tensor:aten.add.Tensor:aten.add.Tensor:aten.add.Tensor:#add:b'\\x00';add:b'\\x00';add:b'\\x00';add:b'\\x00';add:b'\\x00';add:b'\\x00';",
 ],
 ),
 }
 self._test(test_set)

 def test_multi_method_to_backend_not_found(self):
 class SinModule(torch.nn.Module):
 def __init__(self):
Expand Down

16 changes: 10 additions & 6 deletions
exir/lowered_backend_module.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -381,7 +381,7 @@ def _fixup_output_node(gm: torch.fx.GraphModule) -> None:


def arrange_graph_placeholders(
 gm: torch.fx.GraphModule, owning_program: ExportedProgram
 gm: torch.fx.GraphModule, owning_program: ExportedProgram, tag
) -> torch.fx.GraphModule:
 """
 Modifies the graph of the given graphmodule with one that contains the same nodes as the original,
Expand DownExpand Up@@ -411,9 +411,15 @@ def arrange_graph_placeholders(
 if node.op != "placeholder":
 continue

 if node.name in graph_sign.inputs_to_parameters:
 if (
 node.name in graph_sign.inputs_to_parameters
 and node.meta.get("delegation_tag", None) == tag
 ):
 param_nodes.append(node)
 elif node.name in graph_sign.inputs_to_buffers:
 elif (
 node.name in graph_sign.inputs_to_buffers
 and node.meta.get("delegation_tag", None) == tag
 ):
 buffer_nodes.append(node)
 else:
 input_nodes.append(node)
Expand DownExpand Up@@ -694,7 +700,7 @@ def create_exported_program_from_submodule(
 removed from the toplevel ExportedProgram.
 """
 # Arrange the submodule's placeholders in order
 submodule = arrange_graph_placeholders(submodule, owning_program)
 submodule = arrange_graph_placeholders(submodule, owning_program, tag)

 # TODO: we probably need to arrange the outputs wrt buffer mutations.

Expand DownExpand Up@@ -958,5 +964,3 @@ def _unsafe_adjust_original_program( # noqa: C901
 if user_idx > idx:
 user.args = (user.args[0], user_idx - (len(getitem_idxs) - i))
 break

 original_program._validate()