Add unlifting pass under private config#4

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

Closed
tugsbayasgalan wants to merge 1 commit into from
Closed

Add unlifting pass under private config #4

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
11 changes: 4 additions & 7 deletions backends/test/test_backends.py
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -618,11 +618,7 @@ def forward(self, x_raw, h, c):
).to_edge(exir.EdgeCompileConfig(_check_ir_validity=False))

program_without_delegates = (
exir.capture(
composite_m,
(input_x, input_h, input_c),
exir.CaptureConfig(pt2_mode=True),
)
exir.capture(CompositeModel(3), inputs)
.to_edge(exir.EdgeCompileConfig(_check_ir_validity=False))
.to_executorch(
config=exir.ExecutorchBackendConfig(extract_segments=extract_segments),
Expand DownExpand Up@@ -726,7 +722,7 @@ def forward(self, x_raw, h, c):

program_without_delegates = (
exir.capture(
composite_m,
CompositeModel(3),
(input_x, input_h, input_c),
exir.CaptureConfig(pt2_mode=True),
)
Expand DownExpand Up@@ -962,7 +958,8 @@ def test_quantized_with_delegate(self) -> None:
example_inputs,
exir.CaptureConfig(
pt2_mode=True,
enable_functionalization=False,
enable_aot=True,
_unlift=True,
),
).to_edge(exir.EdgeCompileConfig(_check_ir_validity=False))
FileCheck().check_count("quantize_per_tensor.default", 3).check("addmm").run(
Expand Down
160 changes: 158 additions & 2 deletions exir/__init__.py
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -5,9 +5,11 @@
from collections import namedtuple
from dataclasses import dataclass, field
from typing import Any, Callable, Dict, List, Optional, Set, Tuple, Union
from unittest.mock import patch

import sympy
import torch
import torch._export
from executorch.exir.dynamic_shape import DynamicMemoryPlanningMode
from executorch.exir.emit import emit_program, EmitterOutput
from executorch.exir.error import ExportError, ExportErrorType, InternalError
Expand All@@ -25,6 +27,7 @@
from executorch.exir.schema import Program
from executorch.exir.serialize import serialize_to_flatbuffer
from executorch.exir.tracer import (
_default_decomposition_table,
dispatch_trace,
dynamo_trace,
ExirDynamoConfig,
Expand All@@ -41,6 +44,7 @@
from torch._dynamo.eval_frame import Constraint
from torch._export import CallSpec, export, ExportGraphSignature
from torch._export.exported_program import ExportedProgram
from torch._export.passes import ReplaceViewOpsWithViewCopyOpsPass
from torch._export.passes.add_runtime_assertions_for_constraints_pass import (
InputDim,
RangeConstraint,
Expand All@@ -49,12 +53,156 @@
from torch.fx._compatibility import compatibility
from torch.fx.experimental.proxy_tensor import make_fx
from torch.fx.experimental.symbolic_shapes import ShapeEnv
from torch.fx.graph import _PyTreeCodeGen, _PyTreeInfo
from torch.utils import _pytree as pytree


Val = Any


def _unlift(gm, inp_pos_to_param_buffer_name, in_spec, out_spec, state_dict):
count = 0
# Step 1: make lifted params as get_attr
for node in gm.graph.nodes:
if node.op == "placeholder":
if count in inp_pos_to_param_buffer_name:
with gm.graph.inserting_after(node):
getattr_node = gm.graph.get_attr(
inp_pos_to_param_buffer_name[count]
)
node.replace_all_uses_with(getattr_node)
metadata = node.meta
gm.graph.erase_node(node)
getattr_node.meta = metadata
count += 1

# Step 2: Fix the input/output of the graph now that we deleted
# some args.
gm.graph.lint()
names = [f"arg_{i}" for i in range(len(in_spec.children_specs))]
gm.graph._codegen = _PyTreeCodeGen(
_PyTreeInfo(
names,
in_spec,
out_spec,
)
)
gm.recompile()

# Step 3: Find state references in HigherOrderOps and recursively
# fix them.
for node in gm.graph.nodes:
if node.op == "call_function" and node.target == torch.ops.cond:
pred, true_graph, false_graph, operands = node.args
true_gm = getattr(gm, true_graph.name)
false_gm = getattr(gm, false_graph.name)
inp_pos_to_param_buffer_name_for_submod = {}
real_operands = []
for ix, operand in enumerate(operands):
if operand.target in inp_pos_to_param_buffer_name.values():
inp_pos_to_param_buffer_name_for_submod[ix] = operand.target
true_gm.register_buffer(operand.target, state_dict[operand.target])
false_gm.register_buffer(operand.target, state_dict[operand.target])
else:
real_operands.append(operand)
node.args = (pred, true_graph, false_graph, real_operands)

_, in_spec = pytree.tree_flatten(real_operands)

_unlift(
true_gm,
inp_pos_to_param_buffer_name_for_submod,
in_spec,
None,
state_dict,
)
_unlift(
false_gm,
inp_pos_to_param_buffer_name_for_submod,
in_spec,
None,
state_dict,
)
if node.op == "call_function" and node.target.__name__ == "map_impl":
body_graph, num_mapped, *operands = node.args
body_gm = getattr(gm, body_graph.name)
inp_pos_to_buffer_name_for_submod = {}
real_operands = []
for ix, operand in enumerate(operands):
if operand.target in inp_pos_to_param_buffer_name.values():
inp_pos_to_buffer_name_for_submod[ix] = operand.target
body_gm.register_buffer(operand.target, state_dict[operand.target])
else:
real_operands.append(operand)
node.args = (body_graph, num_mapped, *real_operands)

_, in_spec = pytree.tree_flatten(real_operands)

_unlift(
body_gm, inp_pos_to_buffer_name_for_submod, in_spec, None, state_dict
)
gm.graph.lint()
gm.graph.eliminate_dead_code()
gm.recompile()
return gm


def unlift_exported_program_lifted_states(
ep: torch._export.exported_program.ExportedProgram,
):
new_gm = copy.deepcopy(ep.graph_module)

# TODO Fix the period in params/buffers names later
# maybe a pass to replace graph signature with fixed names
param_buffer_name_to_corrected_name = {}

for name, stuff in ep.state_dict.items():
if name in ep.graph_signature.buffers:
if "." in name:
new_gm.register_buffer(name.replace(".", "_"), stuff)
param_buffer_name_to_corrected_name[name] = name.replace(".", "_")
else:
new_gm.register_buffer(name, stuff)
elif name in ep.graph_signature.parameters:
if "." in name:
new_gm.register_parameter(name.replace(".", "_"), stuff)
param_buffer_name_to_corrected_name[name] = name.replace(".", "_")
else:
new_gm.register_parameter(name, stuff)
else:
raise AssertionError("encountered not registered param/buffer")

count = 0
inp_pos_to_param_buffer_name = {}
for node in new_gm.graph.nodes:
if node.op == "placeholder":
if node.name in ep.graph_signature.inputs_to_buffers:
buffer_name = ep.graph_signature.inputs_to_buffers[node.name]
if buffer_name in param_buffer_name_to_corrected_name:
inp_pos_to_param_buffer_name[
count
] = param_buffer_name_to_corrected_name[buffer_name]
else:
inp_pos_to_param_buffer_name[count] = buffer_name
if node.name in ep.graph_signature.inputs_to_parameters:
param_name = ep.graph_signature.inputs_to_parameters[node.name]
if param_name in param_buffer_name_to_corrected_name:
inp_pos_to_param_buffer_name[
count
] = param_buffer_name_to_corrected_name[param_name]
else:
inp_pos_to_param_buffer_name[count] = param_name
count += 1
new_gm = _unlift(
new_gm,
inp_pos_to_param_buffer_name,
ep.call_spec.in_spec,
ep.call_spec.out_spec,
ep.state_dict,
)
return new_gm


@compatibility(is_backward_compatible=False)
@dataclass
class CaptureConfig:
Expand All@@ -63,6 +211,7 @@ class CaptureConfig:
enable_dynamic_shape: bool = False
enable_aot: bool = False
_dynamo_config: "ExirDynamoConfig" = ExirDynamoConfig()
_unlift: bool = False


@compatibility(is_backward_compatible=False)
Expand DownExpand Up@@ -400,8 +549,15 @@ def capture(
"Functionalization is required for enable_aot.",
)

ep = export(f, args, _add_runtime_assertions=False, constraints=constraints)
return ep # pyre-ignore
# TODO remove this later
with patch("torch._export.DECOMP_TABLE", _default_decomposition_table()):
ep = export(
f, args, _add_runtime_assertions=False, constraints=constraints
)
ep = ep.transform(ReplaceViewOpsWithViewCopyOpsPass())
if not config._unlift:
return ep # pyre-ignore
graph_module = unlift_exported_program_lifted_states(ep)

elif config.enable_dynamic_shape:
if not config._dynamo_config.dynamic_shapes:
Expand Down
82 changes: 58 additions & 24 deletions exir/dialects/edge/edge.yaml
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -89,6 +89,14 @@
mat2: T0
__ret_0: T0

- func: aten::arange.start_step
namespace: edge
inherits: aten::arange.start_step
type_alias:
T0: [Byte, Char, Double, Float, Int, Long, Short]
type_constraint:
- __ret_0: T0

- func: aten::bmm
namespace: edge
inherits: aten::bmm
Expand DownExpand Up@@ -198,14 +206,43 @@
- self: T0
__ret_0: T0

- func: aten::lift_fresh_copy
- func: aten::index_select
namespace: edge
inherits: aten::lift_fresh_copy
inherits: aten::index_select
type_alias:
T0: [Bool, Byte, Char, Double, Float, Int, Long, Short]
T0: [Bool]
T1: [Byte]
T2: [Char]
T3: [Double]
T4: [Float]
T5: [Int]
T6: [Long]
T7: [Short]
type_constraint:
- self: T0
index: T6
__ret_0: T0
- self: T1
index: T6
__ret_0: T1
- self: T2
index: T6
__ret_0: T2
- self: T3
index: T6
__ret_0: T3
- self: T4
index: T6
__ret_0: T4
- self: T5
index: T6
__ret_0: T5
- self: T6
index: T6
__ret_0: T6
- self: T7
index: T6
__ret_0: T7

- func: aten::masked_fill.Scalar
namespace: edge
Expand DownExpand Up@@ -245,16 +282,6 @@
mask: T0
__ret_0: T7

- func: aten::minimum
namespace: edge
inherits: aten::minimum
type_alias:
T0: [Bool, Byte, Char, Double, Float, Int, Long, Short]
type_constraint:
- self: T0
other: T0
__ret_0: T0

- func: aten::mm
namespace: edge
inherits: aten::mm
Expand DownExpand Up@@ -324,15 +351,6 @@
- self: T0
__ret_0: T0

- func: aten::select_copy.int
namespace: edge
inherits: aten::select_copy.int
type_alias:
T0: [Bool, Byte, Char, Double, Float, Int, Long, Short]
type_constraint:
- self: T0
__ret_0: T0

- func: aten::sigmoid
namespace: edge
inherits: aten::sigmoid
Expand DownExpand Up@@ -383,9 +401,25 @@
other: T0
__ret_0: T0

- func: aten::t
- func: aten::sym_numel
namespace: edge
inherits: aten::sym_numel
type_alias:
T0: [Bool, Byte, Char, Double, Float, Int, Long, Short]
type_constraint:
- self: T0

- func: aten::sym_size.int
namespace: edge
inherits: aten::sym_size.int
type_alias:
T0: [Bool, Byte, Char, Double, Float, Int, Long, Short]
type_constraint:
- self: T0

- func: aten::t_copy
namespace: edge
inherits: aten::t
inherits: aten::t_copy
type_alias:
T0: [Bool, Byte, Char, Double, Float, Int, Long, Short]
type_constraint:
Expand Down
2 changes: 2 additions & 0 deletions exir/dialects/edge/yaml_generator.py
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -143,6 +143,8 @@ def get_test_gen_key(op_name: str) -> str:
opdb_key = opdb_key[:-5]
elif opdb_key == "sym_size":
opdb_key = "resize_"
elif opdb_key == "sym_numel":
opdb_key = "abs"
elif opdb_key == "convolution":
opdb_key = "conv_transpose2d"
elif opdb_key == "embedding":
Expand Down
Loading
, '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
Add unlifting pass under private config#4

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

 Closed

tugsbayasgalan
wants to merge
1
commit into

from

 Closed


Add unlifting pass under private config
#4
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


11 changes: 4 additions & 7 deletions
backends/test/test_backends.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -618,11 +618,7 @@ def forward(self, x_raw, h, c):
 ).to_edge(exir.EdgeCompileConfig(_check_ir_validity=False))

 program_without_delegates = (
 exir.capture(
 composite_m,
 (input_x, input_h, input_c),
 exir.CaptureConfig(pt2_mode=True),
 )
 exir.capture(CompositeModel(3), inputs)
 .to_edge(exir.EdgeCompileConfig(_check_ir_validity=False))
 .to_executorch(
 config=exir.ExecutorchBackendConfig(extract_segments=extract_segments),
Expand DownExpand Up@@ -726,7 +722,7 @@ def forward(self, x_raw, h, c):

 program_without_delegates = (
 exir.capture(
composite_m,
CompositeModel(3),
 (input_x, input_h, input_c),
 exir.CaptureConfig(pt2_mode=True),
 )
Expand DownExpand Up@@ -962,7 +958,8 @@ def test_quantized_with_delegate(self) -> None:
 example_inputs,
 exir.CaptureConfig(
 pt2_mode=True,
 enable_functionalization=False,
 enable_aot=True,
 _unlift=True,
 ),
 ).to_edge(exir.EdgeCompileConfig(_check_ir_validity=False))
 FileCheck().check_count("quantize_per_tensor.default", 3).check("addmm").run(
Expand Down

160 changes: 158 additions & 2 deletions
exir/__init__.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -5,9 +5,11 @@
from collections import namedtuple
from dataclasses import dataclass, field
from typing import Any, Callable, Dict, List, Optional, Set, Tuple, Union
from unittest.mock import patch

import sympy
import torch
import torch._export
from executorch.exir.dynamic_shape import DynamicMemoryPlanningMode
from executorch.exir.emit import emit_program, EmitterOutput
from executorch.exir.error import ExportError, ExportErrorType, InternalError
Expand All@@ -25,6 +27,7 @@
from executorch.exir.schema import Program
from executorch.exir.serialize import serialize_to_flatbuffer
from executorch.exir.tracer import (
 _default_decomposition_table,
 dispatch_trace,
 dynamo_trace,
 ExirDynamoConfig,
Expand All@@ -41,6 +44,7 @@
from torch._dynamo.eval_frame import Constraint
from torch._export import CallSpec, export, ExportGraphSignature
from torch._export.exported_program import ExportedProgram
from torch._export.passes import ReplaceViewOpsWithViewCopyOpsPass
from torch._export.passes.add_runtime_assertions_for_constraints_pass import (
 InputDim,
 RangeConstraint,
Expand All@@ -49,12 +53,156 @@
from torch.fx._compatibility import compatibility
from torch.fx.experimental.proxy_tensor import make_fx
from torch.fx.experimental.symbolic_shapes import ShapeEnv
from torch.fx.graph import _PyTreeCodeGen, _PyTreeInfo
from torch.utils import _pytree as pytree


Val = Any


def _unlift(gm, inp_pos_to_param_buffer_name, in_spec, out_spec, state_dict):
 count = 0
 # Step 1: make lifted params as get_attr
 for node in gm.graph.nodes:
 if node.op == "placeholder":
 if count in inp_pos_to_param_buffer_name:
 with gm.graph.inserting_after(node):
 getattr_node = gm.graph.get_attr(
 inp_pos_to_param_buffer_name[count]
 )
 node.replace_all_uses_with(getattr_node)
 metadata = node.meta
 gm.graph.erase_node(node)
 getattr_node.meta = metadata
 count += 1

 # Step 2: Fix the input/output of the graph now that we deleted
 # some args.
 gm.graph.lint()
 names = [f"arg_{i}" for i in range(len(in_spec.children_specs))]
 gm.graph._codegen = _PyTreeCodeGen(
 _PyTreeInfo(
 names,
 in_spec,
 out_spec,
 )
 )
 gm.recompile()

 # Step 3: Find state references in HigherOrderOps and recursively
 # fix them.
 for node in gm.graph.nodes:
 if node.op == "call_function" and node.target == torch.ops.cond:
 pred, true_graph, false_graph, operands = node.args
 true_gm = getattr(gm, true_graph.name)
 false_gm = getattr(gm, false_graph.name)
 inp_pos_to_param_buffer_name_for_submod = {}
 real_operands = []
 for ix, operand in enumerate(operands):
 if operand.target in inp_pos_to_param_buffer_name.values():
 inp_pos_to_param_buffer_name_for_submod[ix] = operand.target
 true_gm.register_buffer(operand.target, state_dict[operand.target])
 false_gm.register_buffer(operand.target, state_dict[operand.target])
 else:
 real_operands.append(operand)
 node.args = (pred, true_graph, false_graph, real_operands)

 _, in_spec = pytree.tree_flatten(real_operands)

 _unlift(
 true_gm,
 inp_pos_to_param_buffer_name_for_submod,
 in_spec,
 None,
 state_dict,
 )
 _unlift(
 false_gm,
 inp_pos_to_param_buffer_name_for_submod,
 in_spec,
 None,
 state_dict,
 )
 if node.op == "call_function" and node.target.__name__ == "map_impl":
 body_graph, num_mapped, *operands = node.args
 body_gm = getattr(gm, body_graph.name)
 inp_pos_to_buffer_name_for_submod = {}
 real_operands = []
 for ix, operand in enumerate(operands):
 if operand.target in inp_pos_to_param_buffer_name.values():
 inp_pos_to_buffer_name_for_submod[ix] = operand.target
 body_gm.register_buffer(operand.target, state_dict[operand.target])
 else:
 real_operands.append(operand)
 node.args = (body_graph, num_mapped, *real_operands)

 _, in_spec = pytree.tree_flatten(real_operands)

 _unlift(
 body_gm, inp_pos_to_buffer_name_for_submod, in_spec, None, state_dict
 )
 gm.graph.lint()
 gm.graph.eliminate_dead_code()
 gm.recompile()
 return gm


def unlift_exported_program_lifted_states(
 ep: torch._export.exported_program.ExportedProgram,
):
 new_gm = copy.deepcopy(ep.graph_module)

 # TODO Fix the period in params/buffers names later
 # maybe a pass to replace graph signature with fixed names
 param_buffer_name_to_corrected_name = {}

 for name, stuff in ep.state_dict.items():
 if name in ep.graph_signature.buffers:
 if "." in name:
 new_gm.register_buffer(name.replace(".", "_"), stuff)
 param_buffer_name_to_corrected_name[name] = name.replace(".", "_")
 else:
 new_gm.register_buffer(name, stuff)
 elif name in ep.graph_signature.parameters:
 if "." in name:
 new_gm.register_parameter(name.replace(".", "_"), stuff)
 param_buffer_name_to_corrected_name[name] = name.replace(".", "_")
 else:
 new_gm.register_parameter(name, stuff)
 else:
 raise AssertionError("encountered not registered param/buffer")

 count = 0
 inp_pos_to_param_buffer_name = {}
 for node in new_gm.graph.nodes:
 if node.op == "placeholder":
 if node.name in ep.graph_signature.inputs_to_buffers:
 buffer_name = ep.graph_signature.inputs_to_buffers[node.name]
 if buffer_name in param_buffer_name_to_corrected_name:
 inp_pos_to_param_buffer_name[
 count
 ] = param_buffer_name_to_corrected_name[buffer_name]
 else:
 inp_pos_to_param_buffer_name[count] = buffer_name
 if node.name in ep.graph_signature.inputs_to_parameters:
 param_name = ep.graph_signature.inputs_to_parameters[node.name]
 if param_name in param_buffer_name_to_corrected_name:
 inp_pos_to_param_buffer_name[
 count
 ] = param_buffer_name_to_corrected_name[param_name]
 else:
 inp_pos_to_param_buffer_name[count] = param_name
 count += 1
 new_gm = _unlift(
 new_gm,
 inp_pos_to_param_buffer_name,
 ep.call_spec.in_spec,
 ep.call_spec.out_spec,
 ep.state_dict,
 )
 return new_gm


@compatibility(is_backward_compatible=False)
@dataclass
class CaptureConfig:
Expand All@@ -63,6 +211,7 @@ class CaptureConfig:
 enable_dynamic_shape: bool = False
 enable_aot: bool = False
 _dynamo_config: "ExirDynamoConfig" = ExirDynamoConfig()
 _unlift: bool = False


@compatibility(is_backward_compatible=False)
Expand DownExpand Up@@ -400,8 +549,15 @@ def capture(
 "Functionalization is required for enable_aot.",
 )

 ep = export(f, args, _add_runtime_assertions=False, constraints=constraints)
 return ep # pyre-ignore
 # TODO remove this later
 with patch("torch._export.DECOMP_TABLE", _default_decomposition_table()):
 ep = export(
 f, args, _add_runtime_assertions=False, constraints=constraints
 )
 ep = ep.transform(ReplaceViewOpsWithViewCopyOpsPass())
 if not config._unlift:
 return ep # pyre-ignore
 graph_module = unlift_exported_program_lifted_states(ep)

 elif config.enable_dynamic_shape:
 if not config._dynamo_config.dynamic_shapes:
Expand Down

82 changes: 58 additions & 24 deletions
exir/dialects/edge/edge.yaml

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -89,6 +89,14 @@
 mat2: T0
 __ret_0: T0

- func: aten::arange.start_step
 namespace: edge
 inherits: aten::arange.start_step
 type_alias:
 T0: [Byte, Char, Double, Float, Int, Long, Short]
 type_constraint:
 - __ret_0: T0

- func: aten::bmm
 namespace: edge
 inherits: aten::bmm
Expand DownExpand Up@@ -198,14 +206,43 @@
 - self: T0
 __ret_0: T0

- func: aten::lift_fresh_copy
- func: aten::index_select
 namespace: edge
 inherits: aten::lift_fresh_copy
 inherits: aten::index_select
 type_alias:
 T0: [Bool, Byte, Char, Double, Float, Int, Long, Short]
 T0: [Bool]
 T1: [Byte]
 T2: [Char]
 T3: [Double]
 T4: [Float]
 T5: [Int]
 T6: [Long]
 T7: [Short]
 type_constraint:
 - self: T0
 index: T6
 __ret_0: T0
 - self: T1
 index: T6
 __ret_0: T1
 - self: T2
 index: T6
 __ret_0: T2
 - self: T3
 index: T6
 __ret_0: T3
 - self: T4
 index: T6
 __ret_0: T4
 - self: T5
 index: T6
 __ret_0: T5
 - self: T6
 index: T6
 __ret_0: T6
 - self: T7
 index: T6
 __ret_0: T7

- func: aten::masked_fill.Scalar
 namespace: edge
Expand DownExpand Up@@ -245,16 +282,6 @@
 mask: T0
 __ret_0: T7

- func: aten::minimum
 namespace: edge
 inherits: aten::minimum
 type_alias:
 T0: [Bool, Byte, Char, Double, Float, Int, Long, Short]
 type_constraint:
 - self: T0
 other: T0
 __ret_0: T0

- func: aten::mm
 namespace: edge
 inherits: aten::mm
Expand DownExpand Up@@ -324,15 +351,6 @@
 - self: T0
 __ret_0: T0

- func: aten::select_copy.int
 namespace: edge
 inherits: aten::select_copy.int
 type_alias:
 T0: [Bool, Byte, Char, Double, Float, Int, Long, Short]
 type_constraint:
 - self: T0
 __ret_0: T0

- func: aten::sigmoid
 namespace: edge
 inherits: aten::sigmoid
Expand DownExpand Up@@ -383,9 +401,25 @@
 other: T0
 __ret_0: T0

- func: aten::t
- func: aten::sym_numel
 namespace: edge
 inherits: aten::sym_numel
 type_alias:
 T0: [Bool, Byte, Char, Double, Float, Int, Long, Short]
 type_constraint:
 - self: T0

- func: aten::sym_size.int
 namespace: edge
 inherits: aten::sym_size.int
 type_alias:
 T0: [Bool, Byte, Char, Double, Float, Int, Long, Short]
 type_constraint:
 - self: T0

- func: aten::t_copy
 namespace: edge
 inherits: aten::t
 inherits: aten::t_copy
 type_alias:
 T0: [Bool, Byte, Char, Double, Float, Int, Long, Short]
 type_constraint:
Expand Down

2 changes: 2 additions & 0 deletions
exir/dialects/edge/yaml_generator.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -143,6 +143,8 @@ def get_test_gen_key(op_name: str) -> str:
 opdb_key = opdb_key[:-5]
 elif opdb_key == "sym_size":
 opdb_key = "resize_"
 elif opdb_key == "sym_numel":
 opdb_key = "abs"
 elif opdb_key == "convolution":
 opdb_key = "conv_transpose2d"
 elif opdb_key == "embedding":
Expand Down
Loading

, '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
Add unlifting pass under private config#4

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

 Closed

tugsbayasgalan
wants to merge
1
commit into

from

 Closed


Add unlifting pass under private config
#4
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


11 changes: 4 additions & 7 deletions
backends/test/test_backends.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -618,11 +618,7 @@ def forward(self, x_raw, h, c):
 ).to_edge(exir.EdgeCompileConfig(_check_ir_validity=False))

 program_without_delegates = (
 exir.capture(
 composite_m,
 (input_x, input_h, input_c),
 exir.CaptureConfig(pt2_mode=True),
 )
 exir.capture(CompositeModel(3), inputs)
 .to_edge(exir.EdgeCompileConfig(_check_ir_validity=False))
 .to_executorch(
 config=exir.ExecutorchBackendConfig(extract_segments=extract_segments),
Expand DownExpand Up@@ -726,7 +722,7 @@ def forward(self, x_raw, h, c):

 program_without_delegates = (
 exir.capture(
composite_m,
CompositeModel(3),
 (input_x, input_h, input_c),
 exir.CaptureConfig(pt2_mode=True),
 )
Expand DownExpand Up@@ -962,7 +958,8 @@ def test_quantized_with_delegate(self) -> None:
 example_inputs,
 exir.CaptureConfig(
 pt2_mode=True,
 enable_functionalization=False,
 enable_aot=True,
 _unlift=True,
 ),
 ).to_edge(exir.EdgeCompileConfig(_check_ir_validity=False))
 FileCheck().check_count("quantize_per_tensor.default", 3).check("addmm").run(
Expand Down

160 changes: 158 additions & 2 deletions
exir/__init__.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -5,9 +5,11 @@
from collections import namedtuple
from dataclasses import dataclass, field
from typing import Any, Callable, Dict, List, Optional, Set, Tuple, Union
from unittest.mock import patch

import sympy
import torch
import torch._export
from executorch.exir.dynamic_shape import DynamicMemoryPlanningMode
from executorch.exir.emit import emit_program, EmitterOutput
from executorch.exir.error import ExportError, ExportErrorType, InternalError
Expand All@@ -25,6 +27,7 @@
from executorch.exir.schema import Program
from executorch.exir.serialize import serialize_to_flatbuffer
from executorch.exir.tracer import (
 _default_decomposition_table,
 dispatch_trace,
 dynamo_trace,
 ExirDynamoConfig,
Expand All@@ -41,6 +44,7 @@
from torch._dynamo.eval_frame import Constraint
from torch._export import CallSpec, export, ExportGraphSignature
from torch._export.exported_program import ExportedProgram
from torch._export.passes import ReplaceViewOpsWithViewCopyOpsPass
from torch._export.passes.add_runtime_assertions_for_constraints_pass import (
 InputDim,
 RangeConstraint,
Expand All@@ -49,12 +53,156 @@
from torch.fx._compatibility import compatibility
from torch.fx.experimental.proxy_tensor import make_fx
from torch.fx.experimental.symbolic_shapes import ShapeEnv
from torch.fx.graph import _PyTreeCodeGen, _PyTreeInfo
from torch.utils import _pytree as pytree


Val = Any


def _unlift(gm, inp_pos_to_param_buffer_name, in_spec, out_spec, state_dict):
 count = 0
 # Step 1: make lifted params as get_attr
 for node in gm.graph.nodes:
 if node.op == "placeholder":
 if count in inp_pos_to_param_buffer_name:
 with gm.graph.inserting_after(node):
 getattr_node = gm.graph.get_attr(
 inp_pos_to_param_buffer_name[count]
 )
 node.replace_all_uses_with(getattr_node)
 metadata = node.meta
 gm.graph.erase_node(node)
 getattr_node.meta = metadata
 count += 1

 # Step 2: Fix the input/output of the graph now that we deleted
 # some args.
 gm.graph.lint()
 names = [f"arg_{i}" for i in range(len(in_spec.children_specs))]
 gm.graph._codegen = _PyTreeCodeGen(
 _PyTreeInfo(
 names,
 in_spec,
 out_spec,
 )
 )
 gm.recompile()

 # Step 3: Find state references in HigherOrderOps and recursively
 # fix them.
 for node in gm.graph.nodes:
 if node.op == "call_function" and node.target == torch.ops.cond:
 pred, true_graph, false_graph, operands = node.args
 true_gm = getattr(gm, true_graph.name)
 false_gm = getattr(gm, false_graph.name)
 inp_pos_to_param_buffer_name_for_submod = {}
 real_operands = []
 for ix, operand in enumerate(operands):
 if operand.target in inp_pos_to_param_buffer_name.values():
 inp_pos_to_param_buffer_name_for_submod[ix] = operand.target
 true_gm.register_buffer(operand.target, state_dict[operand.target])
 false_gm.register_buffer(operand.target, state_dict[operand.target])
 else:
 real_operands.append(operand)
 node.args = (pred, true_graph, false_graph, real_operands)

 _, in_spec = pytree.tree_flatten(real_operands)

 _unlift(
 true_gm,
 inp_pos_to_param_buffer_name_for_submod,
 in_spec,
 None,
 state_dict,
 )
 _unlift(
 false_gm,
 inp_pos_to_param_buffer_name_for_submod,
 in_spec,
 None,
 state_dict,
 )
 if node.op == "call_function" and node.target.__name__ == "map_impl":
 body_graph, num_mapped, *operands = node.args
 body_gm = getattr(gm, body_graph.name)
 inp_pos_to_buffer_name_for_submod = {}
 real_operands = []
 for ix, operand in enumerate(operands):
 if operand.target in inp_pos_to_param_buffer_name.values():
 inp_pos_to_buffer_name_for_submod[ix] = operand.target
 body_gm.register_buffer(operand.target, state_dict[operand.target])
 else:
 real_operands.append(operand)
 node.args = (body_graph, num_mapped, *real_operands)

 _, in_spec = pytree.tree_flatten(real_operands)

 _unlift(
 body_gm, inp_pos_to_buffer_name_for_submod, in_spec, None, state_dict
 )
 gm.graph.lint()
 gm.graph.eliminate_dead_code()
 gm.recompile()
 return gm


def unlift_exported_program_lifted_states(
 ep: torch._export.exported_program.ExportedProgram,
):
 new_gm = copy.deepcopy(ep.graph_module)

 # TODO Fix the period in params/buffers names later
 # maybe a pass to replace graph signature with fixed names
 param_buffer_name_to_corrected_name = {}

 for name, stuff in ep.state_dict.items():
 if name in ep.graph_signature.buffers:
 if "." in name:
 new_gm.register_buffer(name.replace(".", "_"), stuff)
 param_buffer_name_to_corrected_name[name] = name.replace(".", "_")
 else:
 new_gm.register_buffer(name, stuff)
 elif name in ep.graph_signature.parameters:
 if "." in name:
 new_gm.register_parameter(name.replace(".", "_"), stuff)
 param_buffer_name_to_corrected_name[name] = name.replace(".", "_")
 else:
 new_gm.register_parameter(name, stuff)
 else:
 raise AssertionError("encountered not registered param/buffer")

 count = 0
 inp_pos_to_param_buffer_name = {}
 for node in new_gm.graph.nodes:
 if node.op == "placeholder":
 if node.name in ep.graph_signature.inputs_to_buffers:
 buffer_name = ep.graph_signature.inputs_to_buffers[node.name]
 if buffer_name in param_buffer_name_to_corrected_name:
 inp_pos_to_param_buffer_name[
 count
 ] = param_buffer_name_to_corrected_name[buffer_name]
 else:
 inp_pos_to_param_buffer_name[count] = buffer_name
 if node.name in ep.graph_signature.inputs_to_parameters:
 param_name = ep.graph_signature.inputs_to_parameters[node.name]
 if param_name in param_buffer_name_to_corrected_name:
 inp_pos_to_param_buffer_name[
 count
 ] = param_buffer_name_to_corrected_name[param_name]
 else:
 inp_pos_to_param_buffer_name[count] = param_name
 count += 1
 new_gm = _unlift(
 new_gm,
 inp_pos_to_param_buffer_name,
 ep.call_spec.in_spec,
 ep.call_spec.out_spec,
 ep.state_dict,
 )
 return new_gm


@compatibility(is_backward_compatible=False)
@dataclass
class CaptureConfig:
Expand All@@ -63,6 +211,7 @@ class CaptureConfig:
 enable_dynamic_shape: bool = False
 enable_aot: bool = False
 _dynamo_config: "ExirDynamoConfig" = ExirDynamoConfig()
 _unlift: bool = False


@compatibility(is_backward_compatible=False)
Expand DownExpand Up@@ -400,8 +549,15 @@ def capture(
 "Functionalization is required for enable_aot.",
 )

 ep = export(f, args, _add_runtime_assertions=False, constraints=constraints)
 return ep # pyre-ignore
 # TODO remove this later
 with patch("torch._export.DECOMP_TABLE", _default_decomposition_table()):
 ep = export(
 f, args, _add_runtime_assertions=False, constraints=constraints
 )
 ep = ep.transform(ReplaceViewOpsWithViewCopyOpsPass())
 if not config._unlift:
 return ep # pyre-ignore
 graph_module = unlift_exported_program_lifted_states(ep)

 elif config.enable_dynamic_shape:
 if not config._dynamo_config.dynamic_shapes:
Expand Down

82 changes: 58 additions & 24 deletions
exir/dialects/edge/edge.yaml

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -89,6 +89,14 @@
 mat2: T0
 __ret_0: T0

- func: aten::arange.start_step
 namespace: edge
 inherits: aten::arange.start_step
 type_alias:
 T0: [Byte, Char, Double, Float, Int, Long, Short]
 type_constraint:
 - __ret_0: T0

- func: aten::bmm
 namespace: edge
 inherits: aten::bmm
Expand DownExpand Up@@ -198,14 +206,43 @@
 - self: T0
 __ret_0: T0

- func: aten::lift_fresh_copy
- func: aten::index_select
 namespace: edge
 inherits: aten::lift_fresh_copy
 inherits: aten::index_select
 type_alias:
 T0: [Bool, Byte, Char, Double, Float, Int, Long, Short]
 T0: [Bool]
 T1: [Byte]
 T2: [Char]
 T3: [Double]
 T4: [Float]
 T5: [Int]
 T6: [Long]
 T7: [Short]
 type_constraint:
 - self: T0
 index: T6
 __ret_0: T0
 - self: T1
 index: T6
 __ret_0: T1
 - self: T2
 index: T6
 __ret_0: T2
 - self: T3
 index: T6
 __ret_0: T3
 - self: T4
 index: T6
 __ret_0: T4
 - self: T5
 index: T6
 __ret_0: T5
 - self: T6
 index: T6
 __ret_0: T6
 - self: T7
 index: T6
 __ret_0: T7

- func: aten::masked_fill.Scalar
 namespace: edge
Expand DownExpand Up@@ -245,16 +282,6 @@
 mask: T0
 __ret_0: T7

- func: aten::minimum
 namespace: edge
 inherits: aten::minimum
 type_alias:
 T0: [Bool, Byte, Char, Double, Float, Int, Long, Short]
 type_constraint:
 - self: T0
 other: T0
 __ret_0: T0

- func: aten::mm
 namespace: edge
 inherits: aten::mm
Expand DownExpand Up@@ -324,15 +351,6 @@
 - self: T0
 __ret_0: T0

- func: aten::select_copy.int
 namespace: edge
 inherits: aten::select_copy.int
 type_alias:
 T0: [Bool, Byte, Char, Double, Float, Int, Long, Short]
 type_constraint:
 - self: T0
 __ret_0: T0

- func: aten::sigmoid
 namespace: edge
 inherits: aten::sigmoid
Expand DownExpand Up@@ -383,9 +401,25 @@
 other: T0
 __ret_0: T0

- func: aten::t
- func: aten::sym_numel
 namespace: edge
 inherits: aten::sym_numel
 type_alias:
 T0: [Bool, Byte, Char, Double, Float, Int, Long, Short]
 type_constraint:
 - self: T0

- func: aten::sym_size.int
 namespace: edge
 inherits: aten::sym_size.int
 type_alias:
 T0: [Bool, Byte, Char, Double, Float, Int, Long, Short]
 type_constraint:
 - self: T0

- func: aten::t_copy
 namespace: edge
 inherits: aten::t
 inherits: aten::t_copy
 type_alias:
 T0: [Bool, Byte, Char, Double, Float, Int, Long, Short]
 type_constraint:
Expand Down

2 changes: 2 additions & 0 deletions
exir/dialects/edge/yaml_generator.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -143,6 +143,8 @@ def get_test_gen_key(op_name: str) -> str:
 opdb_key = opdb_key[:-5]
 elif opdb_key == "sym_size":
 opdb_key = "resize_"
 elif opdb_key == "sym_numel":
 opdb_key = "abs"
 elif opdb_key == "convolution":
 opdb_key = "conv_transpose2d"
 elif opdb_key == "embedding":
Expand Down
Loading

, '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
Add unlifting pass under private config#4

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

 Closed

tugsbayasgalan
wants to merge
1
commit into

from

 Closed


Add unlifting pass under private config
#4
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


11 changes: 4 additions & 7 deletions
backends/test/test_backends.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -618,11 +618,7 @@ def forward(self, x_raw, h, c):
 ).to_edge(exir.EdgeCompileConfig(_check_ir_validity=False))

 program_without_delegates = (
 exir.capture(
 composite_m,
 (input_x, input_h, input_c),
 exir.CaptureConfig(pt2_mode=True),
 )
 exir.capture(CompositeModel(3), inputs)
 .to_edge(exir.EdgeCompileConfig(_check_ir_validity=False))
 .to_executorch(
 config=exir.ExecutorchBackendConfig(extract_segments=extract_segments),
Expand DownExpand Up@@ -726,7 +722,7 @@ def forward(self, x_raw, h, c):

 program_without_delegates = (
 exir.capture(
composite_m,
CompositeModel(3),
 (input_x, input_h, input_c),
 exir.CaptureConfig(pt2_mode=True),
 )
Expand DownExpand Up@@ -962,7 +958,8 @@ def test_quantized_with_delegate(self) -> None:
 example_inputs,
 exir.CaptureConfig(
 pt2_mode=True,
 enable_functionalization=False,
 enable_aot=True,
 _unlift=True,
 ),
 ).to_edge(exir.EdgeCompileConfig(_check_ir_validity=False))
 FileCheck().check_count("quantize_per_tensor.default", 3).check("addmm").run(
Expand Down

160 changes: 158 additions & 2 deletions
exir/__init__.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -5,9 +5,11 @@
from collections import namedtuple
from dataclasses import dataclass, field
from typing import Any, Callable, Dict, List, Optional, Set, Tuple, Union
from unittest.mock import patch

import sympy
import torch
import torch._export
from executorch.exir.dynamic_shape import DynamicMemoryPlanningMode
from executorch.exir.emit import emit_program, EmitterOutput
from executorch.exir.error import ExportError, ExportErrorType, InternalError
Expand All@@ -25,6 +27,7 @@
from executorch.exir.schema import Program
from executorch.exir.serialize import serialize_to_flatbuffer
from executorch.exir.tracer import (
 _default_decomposition_table,
 dispatch_trace,
 dynamo_trace,
 ExirDynamoConfig,
Expand All@@ -41,6 +44,7 @@
from torch._dynamo.eval_frame import Constraint
from torch._export import CallSpec, export, ExportGraphSignature
from torch._export.exported_program import ExportedProgram
from torch._export.passes import ReplaceViewOpsWithViewCopyOpsPass
from torch._export.passes.add_runtime_assertions_for_constraints_pass import (
 InputDim,
 RangeConstraint,
Expand All@@ -49,12 +53,156 @@
from torch.fx._compatibility import compatibility
from torch.fx.experimental.proxy_tensor import make_fx
from torch.fx.experimental.symbolic_shapes import ShapeEnv
from torch.fx.graph import _PyTreeCodeGen, _PyTreeInfo
from torch.utils import _pytree as pytree


Val = Any


def _unlift(gm, inp_pos_to_param_buffer_name, in_spec, out_spec, state_dict):
 count = 0
 # Step 1: make lifted params as get_attr
 for node in gm.graph.nodes:
 if node.op == "placeholder":
 if count in inp_pos_to_param_buffer_name:
 with gm.graph.inserting_after(node):
 getattr_node = gm.graph.get_attr(
 inp_pos_to_param_buffer_name[count]
 )
 node.replace_all_uses_with(getattr_node)
 metadata = node.meta
 gm.graph.erase_node(node)
 getattr_node.meta = metadata
 count += 1

 # Step 2: Fix the input/output of the graph now that we deleted
 # some args.
 gm.graph.lint()
 names = [f"arg_{i}" for i in range(len(in_spec.children_specs))]
 gm.graph._codegen = _PyTreeCodeGen(
 _PyTreeInfo(
 names,
 in_spec,
 out_spec,
 )
 )
 gm.recompile()

 # Step 3: Find state references in HigherOrderOps and recursively
 # fix them.
 for node in gm.graph.nodes:
 if node.op == "call_function" and node.target == torch.ops.cond:
 pred, true_graph, false_graph, operands = node.args
 true_gm = getattr(gm, true_graph.name)
 false_gm = getattr(gm, false_graph.name)
 inp_pos_to_param_buffer_name_for_submod = {}
 real_operands = []
 for ix, operand in enumerate(operands):
 if operand.target in inp_pos_to_param_buffer_name.values():
 inp_pos_to_param_buffer_name_for_submod[ix] = operand.target
 true_gm.register_buffer(operand.target, state_dict[operand.target])
 false_gm.register_buffer(operand.target, state_dict[operand.target])
 else:
 real_operands.append(operand)
 node.args = (pred, true_graph, false_graph, real_operands)

 _, in_spec = pytree.tree_flatten(real_operands)

 _unlift(
 true_gm,
 inp_pos_to_param_buffer_name_for_submod,
 in_spec,
 None,
 state_dict,
 )
 _unlift(
 false_gm,
 inp_pos_to_param_buffer_name_for_submod,
 in_spec,
 None,
 state_dict,
 )
 if node.op == "call_function" and node.target.__name__ == "map_impl":
 body_graph, num_mapped, *operands = node.args
 body_gm = getattr(gm, body_graph.name)
 inp_pos_to_buffer_name_for_submod = {}
 real_operands = []
 for ix, operand in enumerate(operands):
 if operand.target in inp_pos_to_param_buffer_name.values():
 inp_pos_to_buffer_name_for_submod[ix] = operand.target
 body_gm.register_buffer(operand.target, state_dict[operand.target])
 else:
 real_operands.append(operand)
 node.args = (body_graph, num_mapped, *real_operands)

 _, in_spec = pytree.tree_flatten(real_operands)

 _unlift(
 body_gm, inp_pos_to_buffer_name_for_submod, in_spec, None, state_dict
 )
 gm.graph.lint()
 gm.graph.eliminate_dead_code()
 gm.recompile()
 return gm


def unlift_exported_program_lifted_states(
 ep: torch._export.exported_program.ExportedProgram,
):
 new_gm = copy.deepcopy(ep.graph_module)

 # TODO Fix the period in params/buffers names later
 # maybe a pass to replace graph signature with fixed names
 param_buffer_name_to_corrected_name = {}

 for name, stuff in ep.state_dict.items():
 if name in ep.graph_signature.buffers:
 if "." in name:
 new_gm.register_buffer(name.replace(".", "_"), stuff)
 param_buffer_name_to_corrected_name[name] = name.replace(".", "_")
 else:
 new_gm.register_buffer(name, stuff)
 elif name in ep.graph_signature.parameters:
 if "." in name:
 new_gm.register_parameter(name.replace(".", "_"), stuff)
 param_buffer_name_to_corrected_name[name] = name.replace(".", "_")
 else:
 new_gm.register_parameter(name, stuff)
 else:
 raise AssertionError("encountered not registered param/buffer")

 count = 0
 inp_pos_to_param_buffer_name = {}
 for node in new_gm.graph.nodes:
 if node.op == "placeholder":
 if node.name in ep.graph_signature.inputs_to_buffers:
 buffer_name = ep.graph_signature.inputs_to_buffers[node.name]
 if buffer_name in param_buffer_name_to_corrected_name:
 inp_pos_to_param_buffer_name[
 count
 ] = param_buffer_name_to_corrected_name[buffer_name]
 else:
 inp_pos_to_param_buffer_name[count] = buffer_name
 if node.name in ep.graph_signature.inputs_to_parameters:
 param_name = ep.graph_signature.inputs_to_parameters[node.name]
 if param_name in param_buffer_name_to_corrected_name:
 inp_pos_to_param_buffer_name[
 count
 ] = param_buffer_name_to_corrected_name[param_name]
 else:
 inp_pos_to_param_buffer_name[count] = param_name
 count += 1
 new_gm = _unlift(
 new_gm,
 inp_pos_to_param_buffer_name,
 ep.call_spec.in_spec,
 ep.call_spec.out_spec,
 ep.state_dict,
 )
 return new_gm


@compatibility(is_backward_compatible=False)
@dataclass
class CaptureConfig:
Expand All@@ -63,6 +211,7 @@ class CaptureConfig:
 enable_dynamic_shape: bool = False
 enable_aot: bool = False
 _dynamo_config: "ExirDynamoConfig" = ExirDynamoConfig()
 _unlift: bool = False


@compatibility(is_backward_compatible=False)
Expand DownExpand Up@@ -400,8 +549,15 @@ def capture(
 "Functionalization is required for enable_aot.",
 )

 ep = export(f, args, _add_runtime_assertions=False, constraints=constraints)
 return ep # pyre-ignore
 # TODO remove this later
 with patch("torch._export.DECOMP_TABLE", _default_decomposition_table()):
 ep = export(
 f, args, _add_runtime_assertions=False, constraints=constraints
 )
 ep = ep.transform(ReplaceViewOpsWithViewCopyOpsPass())
 if not config._unlift:
 return ep # pyre-ignore
 graph_module = unlift_exported_program_lifted_states(ep)

 elif config.enable_dynamic_shape:
 if not config._dynamo_config.dynamic_shapes:
Expand Down

82 changes: 58 additions & 24 deletions
exir/dialects/edge/edge.yaml

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -89,6 +89,14 @@
 mat2: T0
 __ret_0: T0

- func: aten::arange.start_step
 namespace: edge
 inherits: aten::arange.start_step
 type_alias:
 T0: [Byte, Char, Double, Float, Int, Long, Short]
 type_constraint:
 - __ret_0: T0

- func: aten::bmm
 namespace: edge
 inherits: aten::bmm
Expand DownExpand Up@@ -198,14 +206,43 @@
 - self: T0
 __ret_0: T0

- func: aten::lift_fresh_copy
- func: aten::index_select
 namespace: edge
 inherits: aten::lift_fresh_copy
 inherits: aten::index_select
 type_alias:
 T0: [Bool, Byte, Char, Double, Float, Int, Long, Short]
 T0: [Bool]
 T1: [Byte]
 T2: [Char]
 T3: [Double]
 T4: [Float]
 T5: [Int]
 T6: [Long]
 T7: [Short]
 type_constraint:
 - self: T0
 index: T6
 __ret_0: T0
 - self: T1
 index: T6
 __ret_0: T1
 - self: T2
 index: T6
 __ret_0: T2
 - self: T3
 index: T6
 __ret_0: T3
 - self: T4
 index: T6
 __ret_0: T4
 - self: T5
 index: T6
 __ret_0: T5
 - self: T6
 index: T6
 __ret_0: T6
 - self: T7
 index: T6
 __ret_0: T7

- func: aten::masked_fill.Scalar
 namespace: edge
Expand DownExpand Up@@ -245,16 +282,6 @@
 mask: T0
 __ret_0: T7

- func: aten::minimum
 namespace: edge
 inherits: aten::minimum
 type_alias:
 T0: [Bool, Byte, Char, Double, Float, Int, Long, Short]
 type_constraint:
 - self: T0
 other: T0
 __ret_0: T0

- func: aten::mm
 namespace: edge
 inherits: aten::mm
Expand DownExpand Up@@ -324,15 +351,6 @@
 - self: T0
 __ret_0: T0

- func: aten::select_copy.int
 namespace: edge
 inherits: aten::select_copy.int
 type_alias:
 T0: [Bool, Byte, Char, Double, Float, Int, Long, Short]
 type_constraint:
 - self: T0
 __ret_0: T0

- func: aten::sigmoid
 namespace: edge
 inherits: aten::sigmoid
Expand DownExpand Up@@ -383,9 +401,25 @@
 other: T0
 __ret_0: T0

- func: aten::t
- func: aten::sym_numel
 namespace: edge
 inherits: aten::sym_numel
 type_alias:
 T0: [Bool, Byte, Char, Double, Float, Int, Long, Short]
 type_constraint:
 - self: T0

- func: aten::sym_size.int
 namespace: edge
 inherits: aten::sym_size.int
 type_alias:
 T0: [Bool, Byte, Char, Double, Float, Int, Long, Short]
 type_constraint:
 - self: T0

- func: aten::t_copy
 namespace: edge
 inherits: aten::t
 inherits: aten::t_copy
 type_alias:
 T0: [Bool, Byte, Char, Double, Float, Int, Long, Short]
 type_constraint:
Expand Down

2 changes: 2 additions & 0 deletions
exir/dialects/edge/yaml_generator.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -143,6 +143,8 @@ def get_test_gen_key(op_name: str) -> str:
 opdb_key = opdb_key[:-5]
 elif opdb_key == "sym_size":
 opdb_key = "resize_"
 elif opdb_key == "sym_numel":
 opdb_key = "abs"
 elif opdb_key == "convolution":
 opdb_key = "conv_transpose2d"
 elif opdb_key == "embedding":
Expand Down
Loading

, '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
Add unlifting pass under private config#4

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

 Closed

tugsbayasgalan
wants to merge
1
commit into

from

 Closed


Add unlifting pass under private config
#4
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


11 changes: 4 additions & 7 deletions
backends/test/test_backends.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -618,11 +618,7 @@ def forward(self, x_raw, h, c):
 ).to_edge(exir.EdgeCompileConfig(_check_ir_validity=False))

 program_without_delegates = (
 exir.capture(
 composite_m,
 (input_x, input_h, input_c),
 exir.CaptureConfig(pt2_mode=True),
 )
 exir.capture(CompositeModel(3), inputs)
 .to_edge(exir.EdgeCompileConfig(_check_ir_validity=False))
 .to_executorch(
 config=exir.ExecutorchBackendConfig(extract_segments=extract_segments),
Expand DownExpand Up@@ -726,7 +722,7 @@ def forward(self, x_raw, h, c):

 program_without_delegates = (
 exir.capture(
composite_m,
CompositeModel(3),
 (input_x, input_h, input_c),
 exir.CaptureConfig(pt2_mode=True),
 )
Expand DownExpand Up@@ -962,7 +958,8 @@ def test_quantized_with_delegate(self) -> None:
 example_inputs,
 exir.CaptureConfig(
 pt2_mode=True,
 enable_functionalization=False,
 enable_aot=True,
 _unlift=True,
 ),
 ).to_edge(exir.EdgeCompileConfig(_check_ir_validity=False))
 FileCheck().check_count("quantize_per_tensor.default", 3).check("addmm").run(
Expand Down

160 changes: 158 additions & 2 deletions
exir/__init__.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -5,9 +5,11 @@
from collections import namedtuple
from dataclasses import dataclass, field
from typing import Any, Callable, Dict, List, Optional, Set, Tuple, Union
from unittest.mock import patch

import sympy
import torch
import torch._export
from executorch.exir.dynamic_shape import DynamicMemoryPlanningMode
from executorch.exir.emit import emit_program, EmitterOutput
from executorch.exir.error import ExportError, ExportErrorType, InternalError
Expand All@@ -25,6 +27,7 @@
from executorch.exir.schema import Program
from executorch.exir.serialize import serialize_to_flatbuffer
from executorch.exir.tracer import (
 _default_decomposition_table,
 dispatch_trace,
 dynamo_trace,
 ExirDynamoConfig,
Expand All@@ -41,6 +44,7 @@
from torch._dynamo.eval_frame import Constraint
from torch._export import CallSpec, export, ExportGraphSignature
from torch._export.exported_program import ExportedProgram
from torch._export.passes import ReplaceViewOpsWithViewCopyOpsPass
from torch._export.passes.add_runtime_assertions_for_constraints_pass import (
 InputDim,
 RangeConstraint,
Expand All@@ -49,12 +53,156 @@
from torch.fx._compatibility import compatibility
from torch.fx.experimental.proxy_tensor import make_fx
from torch.fx.experimental.symbolic_shapes import ShapeEnv
from torch.fx.graph import _PyTreeCodeGen, _PyTreeInfo
from torch.utils import _pytree as pytree


Val = Any


def _unlift(gm, inp_pos_to_param_buffer_name, in_spec, out_spec, state_dict):
 count = 0
 # Step 1: make lifted params as get_attr
 for node in gm.graph.nodes:
 if node.op == "placeholder":
 if count in inp_pos_to_param_buffer_name:
 with gm.graph.inserting_after(node):
 getattr_node = gm.graph.get_attr(
 inp_pos_to_param_buffer_name[count]
 )
 node.replace_all_uses_with(getattr_node)
 metadata = node.meta
 gm.graph.erase_node(node)
 getattr_node.meta = metadata
 count += 1

 # Step 2: Fix the input/output of the graph now that we deleted
 # some args.
 gm.graph.lint()
 names = [f"arg_{i}" for i in range(len(in_spec.children_specs))]
 gm.graph._codegen = _PyTreeCodeGen(
 _PyTreeInfo(
 names,
 in_spec,
 out_spec,
 )
 )
 gm.recompile()

 # Step 3: Find state references in HigherOrderOps and recursively
 # fix them.
 for node in gm.graph.nodes:
 if node.op == "call_function" and node.target == torch.ops.cond:
 pred, true_graph, false_graph, operands = node.args
 true_gm = getattr(gm, true_graph.name)
 false_gm = getattr(gm, false_graph.name)
 inp_pos_to_param_buffer_name_for_submod = {}
 real_operands = []
 for ix, operand in enumerate(operands):
 if operand.target in inp_pos_to_param_buffer_name.values():
 inp_pos_to_param_buffer_name_for_submod[ix] = operand.target
 true_gm.register_buffer(operand.target, state_dict[operand.target])
 false_gm.register_buffer(operand.target, state_dict[operand.target])
 else:
 real_operands.append(operand)
 node.args = (pred, true_graph, false_graph, real_operands)

 _, in_spec = pytree.tree_flatten(real_operands)

 _unlift(
 true_gm,
 inp_pos_to_param_buffer_name_for_submod,
 in_spec,
 None,
 state_dict,
 )
 _unlift(
 false_gm,
 inp_pos_to_param_buffer_name_for_submod,
 in_spec,
 None,
 state_dict,
 )
 if node.op == "call_function" and node.target.__name__ == "map_impl":
 body_graph, num_mapped, *operands = node.args
 body_gm = getattr(gm, body_graph.name)
 inp_pos_to_buffer_name_for_submod = {}
 real_operands = []
 for ix, operand in enumerate(operands):
 if operand.target in inp_pos_to_param_buffer_name.values():
 inp_pos_to_buffer_name_for_submod[ix] = operand.target
 body_gm.register_buffer(operand.target, state_dict[operand.target])
 else:
 real_operands.append(operand)
 node.args = (body_graph, num_mapped, *real_operands)

 _, in_spec = pytree.tree_flatten(real_operands)

 _unlift(
 body_gm, inp_pos_to_buffer_name_for_submod, in_spec, None, state_dict
 )
 gm.graph.lint()
 gm.graph.eliminate_dead_code()
 gm.recompile()
 return gm


def unlift_exported_program_lifted_states(
 ep: torch._export.exported_program.ExportedProgram,
):
 new_gm = copy.deepcopy(ep.graph_module)

 # TODO Fix the period in params/buffers names later
 # maybe a pass to replace graph signature with fixed names
 param_buffer_name_to_corrected_name = {}

 for name, stuff in ep.state_dict.items():
 if name in ep.graph_signature.buffers:
 if "." in name:
 new_gm.register_buffer(name.replace(".", "_"), stuff)
 param_buffer_name_to_corrected_name[name] = name.replace(".", "_")
 else:
 new_gm.register_buffer(name, stuff)
 elif name in ep.graph_signature.parameters:
 if "." in name:
 new_gm.register_parameter(name.replace(".", "_"), stuff)
 param_buffer_name_to_corrected_name[name] = name.replace(".", "_")
 else:
 new_gm.register_parameter(name, stuff)
 else:
 raise AssertionError("encountered not registered param/buffer")

 count = 0
 inp_pos_to_param_buffer_name = {}
 for node in new_gm.graph.nodes:
 if node.op == "placeholder":
 if node.name in ep.graph_signature.inputs_to_buffers:
 buffer_name = ep.graph_signature.inputs_to_buffers[node.name]
 if buffer_name in param_buffer_name_to_corrected_name:
 inp_pos_to_param_buffer_name[
 count
 ] = param_buffer_name_to_corrected_name[buffer_name]
 else:
 inp_pos_to_param_buffer_name[count] = buffer_name
 if node.name in ep.graph_signature.inputs_to_parameters:
 param_name = ep.graph_signature.inputs_to_parameters[node.name]
 if param_name in param_buffer_name_to_corrected_name:
 inp_pos_to_param_buffer_name[
 count
 ] = param_buffer_name_to_corrected_name[param_name]
 else:
 inp_pos_to_param_buffer_name[count] = param_name
 count += 1
 new_gm = _unlift(
 new_gm,
 inp_pos_to_param_buffer_name,
 ep.call_spec.in_spec,
 ep.call_spec.out_spec,
 ep.state_dict,
 )
 return new_gm


@compatibility(is_backward_compatible=False)
@dataclass
class CaptureConfig:
Expand All@@ -63,6 +211,7 @@ class CaptureConfig:
 enable_dynamic_shape: bool = False
 enable_aot: bool = False
 _dynamo_config: "ExirDynamoConfig" = ExirDynamoConfig()
 _unlift: bool = False


@compatibility(is_backward_compatible=False)
Expand DownExpand Up@@ -400,8 +549,15 @@ def capture(
 "Functionalization is required for enable_aot.",
 )

 ep = export(f, args, _add_runtime_assertions=False, constraints=constraints)
 return ep # pyre-ignore
 # TODO remove this later
 with patch("torch._export.DECOMP_TABLE", _default_decomposition_table()):
 ep = export(
 f, args, _add_runtime_assertions=False, constraints=constraints
 )
 ep = ep.transform(ReplaceViewOpsWithViewCopyOpsPass())
 if not config._unlift:
 return ep # pyre-ignore
 graph_module = unlift_exported_program_lifted_states(ep)

 elif config.enable_dynamic_shape:
 if not config._dynamo_config.dynamic_shapes:
Expand Down

82 changes: 58 additions & 24 deletions
exir/dialects/edge/edge.yaml

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -89,6 +89,14 @@
 mat2: T0
 __ret_0: T0

- func: aten::arange.start_step
 namespace: edge
 inherits: aten::arange.start_step
 type_alias:
 T0: [Byte, Char, Double, Float, Int, Long, Short]
 type_constraint:
 - __ret_0: T0

- func: aten::bmm
 namespace: edge
 inherits: aten::bmm
Expand DownExpand Up@@ -198,14 +206,43 @@
 - self: T0
 __ret_0: T0

- func: aten::lift_fresh_copy
- func: aten::index_select
 namespace: edge
 inherits: aten::lift_fresh_copy
 inherits: aten::index_select
 type_alias:
 T0: [Bool, Byte, Char, Double, Float, Int, Long, Short]
 T0: [Bool]
 T1: [Byte]
 T2: [Char]
 T3: [Double]
 T4: [Float]
 T5: [Int]
 T6: [Long]
 T7: [Short]
 type_constraint:
 - self: T0
 index: T6
 __ret_0: T0
 - self: T1
 index: T6
 __ret_0: T1
 - self: T2
 index: T6
 __ret_0: T2
 - self: T3
 index: T6
 __ret_0: T3
 - self: T4
 index: T6
 __ret_0: T4
 - self: T5
 index: T6
 __ret_0: T5
 - self: T6
 index: T6
 __ret_0: T6
 - self: T7
 index: T6
 __ret_0: T7

- func: aten::masked_fill.Scalar
 namespace: edge
Expand DownExpand Up@@ -245,16 +282,6 @@
 mask: T0
 __ret_0: T7

- func: aten::minimum
 namespace: edge
 inherits: aten::minimum
 type_alias:
 T0: [Bool, Byte, Char, Double, Float, Int, Long, Short]
 type_constraint:
 - self: T0
 other: T0
 __ret_0: T0

- func: aten::mm
 namespace: edge
 inherits: aten::mm
Expand DownExpand Up@@ -324,15 +351,6 @@
 - self: T0
 __ret_0: T0

- func: aten::select_copy.int
 namespace: edge
 inherits: aten::select_copy.int
 type_alias:
 T0: [Bool, Byte, Char, Double, Float, Int, Long, Short]
 type_constraint:
 - self: T0
 __ret_0: T0

- func: aten::sigmoid
 namespace: edge
 inherits: aten::sigmoid
Expand DownExpand Up@@ -383,9 +401,25 @@
 other: T0
 __ret_0: T0

- func: aten::t
- func: aten::sym_numel
 namespace: edge
 inherits: aten::sym_numel
 type_alias:
 T0: [Bool, Byte, Char, Double, Float, Int, Long, Short]
 type_constraint:
 - self: T0

- func: aten::sym_size.int
 namespace: edge
 inherits: aten::sym_size.int
 type_alias:
 T0: [Bool, Byte, Char, Double, Float, Int, Long, Short]
 type_constraint:
 - self: T0

- func: aten::t_copy
 namespace: edge
 inherits: aten::t
 inherits: aten::t_copy
 type_alias:
 T0: [Bool, Byte, Char, Double, Float, Int, Long, Short]
 type_constraint:
Expand Down

2 changes: 2 additions & 0 deletions
exir/dialects/edge/yaml_generator.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -143,6 +143,8 @@ def get_test_gen_key(op_name: str) -> str:
 opdb_key = opdb_key[:-5]
 elif opdb_key == "sym_size":
 opdb_key = "resize_"
 elif opdb_key == "sym_numel":
 opdb_key = "abs"
 elif opdb_key == "convolution":
 opdb_key = "conv_transpose2d"
 elif opdb_key == "embedding":
Expand Down
Loading

, '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
Add unlifting pass under private config#4

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

 Closed

tugsbayasgalan
wants to merge
1
commit into

from

 Closed


Add unlifting pass under private config
#4
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


11 changes: 4 additions & 7 deletions
backends/test/test_backends.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -618,11 +618,7 @@ def forward(self, x_raw, h, c):
 ).to_edge(exir.EdgeCompileConfig(_check_ir_validity=False))

 program_without_delegates = (
 exir.capture(
 composite_m,
 (input_x, input_h, input_c),
 exir.CaptureConfig(pt2_mode=True),
 )
 exir.capture(CompositeModel(3), inputs)
 .to_edge(exir.EdgeCompileConfig(_check_ir_validity=False))
 .to_executorch(
 config=exir.ExecutorchBackendConfig(extract_segments=extract_segments),
Expand DownExpand Up@@ -726,7 +722,7 @@ def forward(self, x_raw, h, c):

 program_without_delegates = (
 exir.capture(
composite_m,
CompositeModel(3),
 (input_x, input_h, input_c),
 exir.CaptureConfig(pt2_mode=True),
 )
Expand DownExpand Up@@ -962,7 +958,8 @@ def test_quantized_with_delegate(self) -> None:
 example_inputs,
 exir.CaptureConfig(
 pt2_mode=True,
 enable_functionalization=False,
 enable_aot=True,
 _unlift=True,
 ),
 ).to_edge(exir.EdgeCompileConfig(_check_ir_validity=False))
 FileCheck().check_count("quantize_per_tensor.default", 3).check("addmm").run(
Expand Down

160 changes: 158 additions & 2 deletions
exir/__init__.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -5,9 +5,11 @@
from collections import namedtuple
from dataclasses import dataclass, field
from typing import Any, Callable, Dict, List, Optional, Set, Tuple, Union
from unittest.mock import patch

import sympy
import torch
import torch._export
from executorch.exir.dynamic_shape import DynamicMemoryPlanningMode
from executorch.exir.emit import emit_program, EmitterOutput
from executorch.exir.error import ExportError, ExportErrorType, InternalError
Expand All@@ -25,6 +27,7 @@
from executorch.exir.schema import Program
from executorch.exir.serialize import serialize_to_flatbuffer
from executorch.exir.tracer import (
 _default_decomposition_table,
 dispatch_trace,
 dynamo_trace,
 ExirDynamoConfig,
Expand All@@ -41,6 +44,7 @@
from torch._dynamo.eval_frame import Constraint
from torch._export import CallSpec, export, ExportGraphSignature
from torch._export.exported_program import ExportedProgram
from torch._export.passes import ReplaceViewOpsWithViewCopyOpsPass
from torch._export.passes.add_runtime_assertions_for_constraints_pass import (
 InputDim,
 RangeConstraint,
Expand All@@ -49,12 +53,156 @@
from torch.fx._compatibility import compatibility
from torch.fx.experimental.proxy_tensor import make_fx
from torch.fx.experimental.symbolic_shapes import ShapeEnv
from torch.fx.graph import _PyTreeCodeGen, _PyTreeInfo
from torch.utils import _pytree as pytree


Val = Any


def _unlift(gm, inp_pos_to_param_buffer_name, in_spec, out_spec, state_dict):
 count = 0
 # Step 1: make lifted params as get_attr
 for node in gm.graph.nodes:
 if node.op == "placeholder":
 if count in inp_pos_to_param_buffer_name:
 with gm.graph.inserting_after(node):
 getattr_node = gm.graph.get_attr(
 inp_pos_to_param_buffer_name[count]
 )
 node.replace_all_uses_with(getattr_node)
 metadata = node.meta
 gm.graph.erase_node(node)
 getattr_node.meta = metadata
 count += 1

 # Step 2: Fix the input/output of the graph now that we deleted
 # some args.
 gm.graph.lint()
 names = [f"arg_{i}" for i in range(len(in_spec.children_specs))]
 gm.graph._codegen = _PyTreeCodeGen(
 _PyTreeInfo(
 names,
 in_spec,
 out_spec,
 )
 )
 gm.recompile()

 # Step 3: Find state references in HigherOrderOps and recursively
 # fix them.
 for node in gm.graph.nodes:
 if node.op == "call_function" and node.target == torch.ops.cond:
 pred, true_graph, false_graph, operands = node.args
 true_gm = getattr(gm, true_graph.name)
 false_gm = getattr(gm, false_graph.name)
 inp_pos_to_param_buffer_name_for_submod = {}
 real_operands = []
 for ix, operand in enumerate(operands):
 if operand.target in inp_pos_to_param_buffer_name.values():
 inp_pos_to_param_buffer_name_for_submod[ix] = operand.target
 true_gm.register_buffer(operand.target, state_dict[operand.target])
 false_gm.register_buffer(operand.target, state_dict[operand.target])
 else:
 real_operands.append(operand)
 node.args = (pred, true_graph, false_graph, real_operands)

 _, in_spec = pytree.tree_flatten(real_operands)

 _unlift(
 true_gm,
 inp_pos_to_param_buffer_name_for_submod,
 in_spec,
 None,
 state_dict,
 )
 _unlift(
 false_gm,
 inp_pos_to_param_buffer_name_for_submod,
 in_spec,
 None,
 state_dict,
 )
 if node.op == "call_function" and node.target.__name__ == "map_impl":
 body_graph, num_mapped, *operands = node.args
 body_gm = getattr(gm, body_graph.name)
 inp_pos_to_buffer_name_for_submod = {}
 real_operands = []
 for ix, operand in enumerate(operands):
 if operand.target in inp_pos_to_param_buffer_name.values():
 inp_pos_to_buffer_name_for_submod[ix] = operand.target
 body_gm.register_buffer(operand.target, state_dict[operand.target])
 else:
 real_operands.append(operand)
 node.args = (body_graph, num_mapped, *real_operands)

 _, in_spec = pytree.tree_flatten(real_operands)

 _unlift(
 body_gm, inp_pos_to_buffer_name_for_submod, in_spec, None, state_dict
 )
 gm.graph.lint()
 gm.graph.eliminate_dead_code()
 gm.recompile()
 return gm


def unlift_exported_program_lifted_states(
 ep: torch._export.exported_program.ExportedProgram,
):
 new_gm = copy.deepcopy(ep.graph_module)

 # TODO Fix the period in params/buffers names later
 # maybe a pass to replace graph signature with fixed names
 param_buffer_name_to_corrected_name = {}

 for name, stuff in ep.state_dict.items():
 if name in ep.graph_signature.buffers:
 if "." in name:
 new_gm.register_buffer(name.replace(".", "_"), stuff)
 param_buffer_name_to_corrected_name[name] = name.replace(".", "_")
 else:
 new_gm.register_buffer(name, stuff)
 elif name in ep.graph_signature.parameters:
 if "." in name:
 new_gm.register_parameter(name.replace(".", "_"), stuff)
 param_buffer_name_to_corrected_name[name] = name.replace(".", "_")
 else:
 new_gm.register_parameter(name, stuff)
 else:
 raise AssertionError("encountered not registered param/buffer")

 count = 0
 inp_pos_to_param_buffer_name = {}
 for node in new_gm.graph.nodes:
 if node.op == "placeholder":
 if node.name in ep.graph_signature.inputs_to_buffers:
 buffer_name = ep.graph_signature.inputs_to_buffers[node.name]
 if buffer_name in param_buffer_name_to_corrected_name:
 inp_pos_to_param_buffer_name[
 count
 ] = param_buffer_name_to_corrected_name[buffer_name]
 else:
 inp_pos_to_param_buffer_name[count] = buffer_name
 if node.name in ep.graph_signature.inputs_to_parameters:
 param_name = ep.graph_signature.inputs_to_parameters[node.name]
 if param_name in param_buffer_name_to_corrected_name:
 inp_pos_to_param_buffer_name[
 count
 ] = param_buffer_name_to_corrected_name[param_name]
 else:
 inp_pos_to_param_buffer_name[count] = param_name
 count += 1
 new_gm = _unlift(
 new_gm,
 inp_pos_to_param_buffer_name,
 ep.call_spec.in_spec,
 ep.call_spec.out_spec,
 ep.state_dict,
 )
 return new_gm


@compatibility(is_backward_compatible=False)
@dataclass
class CaptureConfig:
Expand All@@ -63,6 +211,7 @@ class CaptureConfig:
 enable_dynamic_shape: bool = False
 enable_aot: bool = False
 _dynamo_config: "ExirDynamoConfig" = ExirDynamoConfig()
 _unlift: bool = False


@compatibility(is_backward_compatible=False)
Expand DownExpand Up@@ -400,8 +549,15 @@ def capture(
 "Functionalization is required for enable_aot.",
 )

 ep = export(f, args, _add_runtime_assertions=False, constraints=constraints)
 return ep # pyre-ignore
 # TODO remove this later
 with patch("torch._export.DECOMP_TABLE", _default_decomposition_table()):
 ep = export(
 f, args, _add_runtime_assertions=False, constraints=constraints
 )
 ep = ep.transform(ReplaceViewOpsWithViewCopyOpsPass())
 if not config._unlift:
 return ep # pyre-ignore
 graph_module = unlift_exported_program_lifted_states(ep)

 elif config.enable_dynamic_shape:
 if not config._dynamo_config.dynamic_shapes:
Expand Down

82 changes: 58 additions & 24 deletions
exir/dialects/edge/edge.yaml

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -89,6 +89,14 @@
 mat2: T0
 __ret_0: T0

- func: aten::arange.start_step
 namespace: edge
 inherits: aten::arange.start_step
 type_alias:
 T0: [Byte, Char, Double, Float, Int, Long, Short]
 type_constraint:
 - __ret_0: T0

- func: aten::bmm
 namespace: edge
 inherits: aten::bmm
Expand DownExpand Up@@ -198,14 +206,43 @@
 - self: T0
 __ret_0: T0

- func: aten::lift_fresh_copy
- func: aten::index_select
 namespace: edge
 inherits: aten::lift_fresh_copy
 inherits: aten::index_select
 type_alias:
 T0: [Bool, Byte, Char, Double, Float, Int, Long, Short]
 T0: [Bool]
 T1: [Byte]
 T2: [Char]
 T3: [Double]
 T4: [Float]
 T5: [Int]
 T6: [Long]
 T7: [Short]
 type_constraint:
 - self: T0
 index: T6
 __ret_0: T0
 - self: T1
 index: T6
 __ret_0: T1
 - self: T2
 index: T6
 __ret_0: T2
 - self: T3
 index: T6
 __ret_0: T3
 - self: T4
 index: T6
 __ret_0: T4
 - self: T5
 index: T6
 __ret_0: T5
 - self: T6
 index: T6
 __ret_0: T6
 - self: T7
 index: T6
 __ret_0: T7

- func: aten::masked_fill.Scalar
 namespace: edge
Expand DownExpand Up@@ -245,16 +282,6 @@
 mask: T0
 __ret_0: T7

- func: aten::minimum
 namespace: edge
 inherits: aten::minimum
 type_alias:
 T0: [Bool, Byte, Char, Double, Float, Int, Long, Short]
 type_constraint:
 - self: T0
 other: T0
 __ret_0: T0

- func: aten::mm
 namespace: edge
 inherits: aten::mm
Expand DownExpand Up@@ -324,15 +351,6 @@
 - self: T0
 __ret_0: T0

- func: aten::select_copy.int
 namespace: edge
 inherits: aten::select_copy.int
 type_alias:
 T0: [Bool, Byte, Char, Double, Float, Int, Long, Short]
 type_constraint:
 - self: T0
 __ret_0: T0

- func: aten::sigmoid
 namespace: edge
 inherits: aten::sigmoid
Expand DownExpand Up@@ -383,9 +401,25 @@
 other: T0
 __ret_0: T0

- func: aten::t
- func: aten::sym_numel
 namespace: edge
 inherits: aten::sym_numel
 type_alias:
 T0: [Bool, Byte, Char, Double, Float, Int, Long, Short]
 type_constraint:
 - self: T0

- func: aten::sym_size.int
 namespace: edge
 inherits: aten::sym_size.int
 type_alias:
 T0: [Bool, Byte, Char, Double, Float, Int, Long, Short]
 type_constraint:
 - self: T0

- func: aten::t_copy
 namespace: edge
 inherits: aten::t
 inherits: aten::t_copy
 type_alias:
 T0: [Bool, Byte, Char, Double, Float, Int, Long, Short]
 type_constraint:
Expand Down

2 changes: 2 additions & 0 deletions
exir/dialects/edge/yaml_generator.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -143,6 +143,8 @@ def get_test_gen_key(op_name: str) -> str:
 opdb_key = opdb_key[:-5]
 elif opdb_key == "sym_size":
 opdb_key = "resize_"
 elif opdb_key == "sym_numel":
 opdb_key = "abs"
 elif opdb_key == "convolution":
 opdb_key = "conv_transpose2d"
 elif opdb_key == "embedding":
Expand Down
Loading

, '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
Add unlifting pass under private config#4

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

 Closed

tugsbayasgalan
wants to merge
1
commit into

from

 Closed


Add unlifting pass under private config
#4
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


11 changes: 4 additions & 7 deletions
backends/test/test_backends.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -618,11 +618,7 @@ def forward(self, x_raw, h, c):
 ).to_edge(exir.EdgeCompileConfig(_check_ir_validity=False))

 program_without_delegates = (
 exir.capture(
 composite_m,
 (input_x, input_h, input_c),
 exir.CaptureConfig(pt2_mode=True),
 )
 exir.capture(CompositeModel(3), inputs)
 .to_edge(exir.EdgeCompileConfig(_check_ir_validity=False))
 .to_executorch(
 config=exir.ExecutorchBackendConfig(extract_segments=extract_segments),
Expand DownExpand Up@@ -726,7 +722,7 @@ def forward(self, x_raw, h, c):

 program_without_delegates = (
 exir.capture(
composite_m,
CompositeModel(3),
 (input_x, input_h, input_c),
 exir.CaptureConfig(pt2_mode=True),
 )
Expand DownExpand Up@@ -962,7 +958,8 @@ def test_quantized_with_delegate(self) -> None:
 example_inputs,
 exir.CaptureConfig(
 pt2_mode=True,
 enable_functionalization=False,
 enable_aot=True,
 _unlift=True,
 ),
 ).to_edge(exir.EdgeCompileConfig(_check_ir_validity=False))
 FileCheck().check_count("quantize_per_tensor.default", 3).check("addmm").run(
Expand Down

160 changes: 158 additions & 2 deletions
exir/__init__.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -5,9 +5,11 @@
from collections import namedtuple
from dataclasses import dataclass, field
from typing import Any, Callable, Dict, List, Optional, Set, Tuple, Union
from unittest.mock import patch

import sympy
import torch
import torch._export
from executorch.exir.dynamic_shape import DynamicMemoryPlanningMode
from executorch.exir.emit import emit_program, EmitterOutput
from executorch.exir.error import ExportError, ExportErrorType, InternalError
Expand All@@ -25,6 +27,7 @@
from executorch.exir.schema import Program
from executorch.exir.serialize import serialize_to_flatbuffer
from executorch.exir.tracer import (
 _default_decomposition_table,
 dispatch_trace,
 dynamo_trace,
 ExirDynamoConfig,
Expand All@@ -41,6 +44,7 @@
from torch._dynamo.eval_frame import Constraint
from torch._export import CallSpec, export, ExportGraphSignature
from torch._export.exported_program import ExportedProgram
from torch._export.passes import ReplaceViewOpsWithViewCopyOpsPass
from torch._export.passes.add_runtime_assertions_for_constraints_pass import (
 InputDim,
 RangeConstraint,
Expand All@@ -49,12 +53,156 @@
from torch.fx._compatibility import compatibility
from torch.fx.experimental.proxy_tensor import make_fx
from torch.fx.experimental.symbolic_shapes import ShapeEnv
from torch.fx.graph import _PyTreeCodeGen, _PyTreeInfo
from torch.utils import _pytree as pytree


Val = Any


def _unlift(gm, inp_pos_to_param_buffer_name, in_spec, out_spec, state_dict):
 count = 0
 # Step 1: make lifted params as get_attr
 for node in gm.graph.nodes:
 if node.op == "placeholder":
 if count in inp_pos_to_param_buffer_name:
 with gm.graph.inserting_after(node):
 getattr_node = gm.graph.get_attr(
 inp_pos_to_param_buffer_name[count]
 )
 node.replace_all_uses_with(getattr_node)
 metadata = node.meta
 gm.graph.erase_node(node)
 getattr_node.meta = metadata
 count += 1

 # Step 2: Fix the input/output of the graph now that we deleted
 # some args.
 gm.graph.lint()
 names = [f"arg_{i}" for i in range(len(in_spec.children_specs))]
 gm.graph._codegen = _PyTreeCodeGen(
 _PyTreeInfo(
 names,
 in_spec,
 out_spec,
 )
 )
 gm.recompile()

 # Step 3: Find state references in HigherOrderOps and recursively
 # fix them.
 for node in gm.graph.nodes:
 if node.op == "call_function" and node.target == torch.ops.cond:
 pred, true_graph, false_graph, operands = node.args
 true_gm = getattr(gm, true_graph.name)
 false_gm = getattr(gm, false_graph.name)
 inp_pos_to_param_buffer_name_for_submod = {}
 real_operands = []
 for ix, operand in enumerate(operands):
 if operand.target in inp_pos_to_param_buffer_name.values():
 inp_pos_to_param_buffer_name_for_submod[ix] = operand.target
 true_gm.register_buffer(operand.target, state_dict[operand.target])
 false_gm.register_buffer(operand.target, state_dict[operand.target])
 else:
 real_operands.append(operand)
 node.args = (pred, true_graph, false_graph, real_operands)

 _, in_spec = pytree.tree_flatten(real_operands)

 _unlift(
 true_gm,
 inp_pos_to_param_buffer_name_for_submod,
 in_spec,
 None,
 state_dict,
 )
 _unlift(
 false_gm,
 inp_pos_to_param_buffer_name_for_submod,
 in_spec,
 None,
 state_dict,
 )
 if node.op == "call_function" and node.target.__name__ == "map_impl":
 body_graph, num_mapped, *operands = node.args
 body_gm = getattr(gm, body_graph.name)
 inp_pos_to_buffer_name_for_submod = {}
 real_operands = []
 for ix, operand in enumerate(operands):
 if operand.target in inp_pos_to_param_buffer_name.values():
 inp_pos_to_buffer_name_for_submod[ix] = operand.target
 body_gm.register_buffer(operand.target, state_dict[operand.target])
 else:
 real_operands.append(operand)
 node.args = (body_graph, num_mapped, *real_operands)

 _, in_spec = pytree.tree_flatten(real_operands)

 _unlift(
 body_gm, inp_pos_to_buffer_name_for_submod, in_spec, None, state_dict
 )
 gm.graph.lint()
 gm.graph.eliminate_dead_code()
 gm.recompile()
 return gm


def unlift_exported_program_lifted_states(
 ep: torch._export.exported_program.ExportedProgram,
):
 new_gm = copy.deepcopy(ep.graph_module)

 # TODO Fix the period in params/buffers names later
 # maybe a pass to replace graph signature with fixed names
 param_buffer_name_to_corrected_name = {}

 for name, stuff in ep.state_dict.items():
 if name in ep.graph_signature.buffers:
 if "." in name:
 new_gm.register_buffer(name.replace(".", "_"), stuff)
 param_buffer_name_to_corrected_name[name] = name.replace(".", "_")
 else:
 new_gm.register_buffer(name, stuff)
 elif name in ep.graph_signature.parameters:
 if "." in name:
 new_gm.register_parameter(name.replace(".", "_"), stuff)
 param_buffer_name_to_corrected_name[name] = name.replace(".", "_")
 else:
 new_gm.register_parameter(name, stuff)
 else:
 raise AssertionError("encountered not registered param/buffer")

 count = 0
 inp_pos_to_param_buffer_name = {}
 for node in new_gm.graph.nodes:
 if node.op == "placeholder":
 if node.name in ep.graph_signature.inputs_to_buffers:
 buffer_name = ep.graph_signature.inputs_to_buffers[node.name]
 if buffer_name in param_buffer_name_to_corrected_name:
 inp_pos_to_param_buffer_name[
 count
 ] = param_buffer_name_to_corrected_name[buffer_name]
 else:
 inp_pos_to_param_buffer_name[count] = buffer_name
 if node.name in ep.graph_signature.inputs_to_parameters:
 param_name = ep.graph_signature.inputs_to_parameters[node.name]
 if param_name in param_buffer_name_to_corrected_name:
 inp_pos_to_param_buffer_name[
 count
 ] = param_buffer_name_to_corrected_name[param_name]
 else:
 inp_pos_to_param_buffer_name[count] = param_name
 count += 1
 new_gm = _unlift(
 new_gm,
 inp_pos_to_param_buffer_name,
 ep.call_spec.in_spec,
 ep.call_spec.out_spec,
 ep.state_dict,
 )
 return new_gm


@compatibility(is_backward_compatible=False)
@dataclass
class CaptureConfig:
Expand All@@ -63,6 +211,7 @@ class CaptureConfig:
 enable_dynamic_shape: bool = False
 enable_aot: bool = False
 _dynamo_config: "ExirDynamoConfig" = ExirDynamoConfig()
 _unlift: bool = False


@compatibility(is_backward_compatible=False)
Expand DownExpand Up@@ -400,8 +549,15 @@ def capture(
 "Functionalization is required for enable_aot.",
 )

 ep = export(f, args, _add_runtime_assertions=False, constraints=constraints)
 return ep # pyre-ignore
 # TODO remove this later
 with patch("torch._export.DECOMP_TABLE", _default_decomposition_table()):
 ep = export(
 f, args, _add_runtime_assertions=False, constraints=constraints
 )
 ep = ep.transform(ReplaceViewOpsWithViewCopyOpsPass())
 if not config._unlift:
 return ep # pyre-ignore
 graph_module = unlift_exported_program_lifted_states(ep)

 elif config.enable_dynamic_shape:
 if not config._dynamo_config.dynamic_shapes:
Expand Down

82 changes: 58 additions & 24 deletions
exir/dialects/edge/edge.yaml

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -89,6 +89,14 @@
 mat2: T0
 __ret_0: T0

- func: aten::arange.start_step
 namespace: edge
 inherits: aten::arange.start_step
 type_alias:
 T0: [Byte, Char, Double, Float, Int, Long, Short]
 type_constraint:
 - __ret_0: T0

- func: aten::bmm
 namespace: edge
 inherits: aten::bmm
Expand DownExpand Up@@ -198,14 +206,43 @@
 - self: T0
 __ret_0: T0

- func: aten::lift_fresh_copy
- func: aten::index_select
 namespace: edge
 inherits: aten::lift_fresh_copy
 inherits: aten::index_select
 type_alias:
 T0: [Bool, Byte, Char, Double, Float, Int, Long, Short]
 T0: [Bool]
 T1: [Byte]
 T2: [Char]
 T3: [Double]
 T4: [Float]
 T5: [Int]
 T6: [Long]
 T7: [Short]
 type_constraint:
 - self: T0
 index: T6
 __ret_0: T0
 - self: T1
 index: T6
 __ret_0: T1
 - self: T2
 index: T6
 __ret_0: T2
 - self: T3
 index: T6
 __ret_0: T3
 - self: T4
 index: T6
 __ret_0: T4
 - self: T5
 index: T6
 __ret_0: T5
 - self: T6
 index: T6
 __ret_0: T6
 - self: T7
 index: T6
 __ret_0: T7

- func: aten::masked_fill.Scalar
 namespace: edge
Expand DownExpand Up@@ -245,16 +282,6 @@
 mask: T0
 __ret_0: T7

- func: aten::minimum
 namespace: edge
 inherits: aten::minimum
 type_alias:
 T0: [Bool, Byte, Char, Double, Float, Int, Long, Short]
 type_constraint:
 - self: T0
 other: T0
 __ret_0: T0

- func: aten::mm
 namespace: edge
 inherits: aten::mm
Expand DownExpand Up@@ -324,15 +351,6 @@
 - self: T0
 __ret_0: T0

- func: aten::select_copy.int
 namespace: edge
 inherits: aten::select_copy.int
 type_alias:
 T0: [Bool, Byte, Char, Double, Float, Int, Long, Short]
 type_constraint:
 - self: T0
 __ret_0: T0

- func: aten::sigmoid
 namespace: edge
 inherits: aten::sigmoid
Expand DownExpand Up@@ -383,9 +401,25 @@
 other: T0
 __ret_0: T0

- func: aten::t
- func: aten::sym_numel
 namespace: edge
 inherits: aten::sym_numel
 type_alias:
 T0: [Bool, Byte, Char, Double, Float, Int, Long, Short]
 type_constraint:
 - self: T0

- func: aten::sym_size.int
 namespace: edge
 inherits: aten::sym_size.int
 type_alias:
 T0: [Bool, Byte, Char, Double, Float, Int, Long, Short]
 type_constraint:
 - self: T0

- func: aten::t_copy
 namespace: edge
 inherits: aten::t
 inherits: aten::t_copy
 type_alias:
 T0: [Bool, Byte, Char, Double, Float, Int, Long, Short]
 type_constraint:
Expand Down

2 changes: 2 additions & 0 deletions
exir/dialects/edge/yaml_generator.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -143,6 +143,8 @@ def get_test_gen_key(op_name: str) -> str:
 opdb_key = opdb_key[:-5]
 elif opdb_key == "sym_size":
 opdb_key = "resize_"
 elif opdb_key == "sym_numel":
 opdb_key = "abs"
 elif opdb_key == "convolution":
 opdb_key = "conv_transpose2d"
 elif opdb_key == "embedding":
Expand Down
Loading

, '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
Add unlifting pass under private config#4

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

 Closed

tugsbayasgalan
wants to merge
1
commit into

from

 Closed


Add unlifting pass under private config
#4
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


11 changes: 4 additions & 7 deletions
backends/test/test_backends.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -618,11 +618,7 @@ def forward(self, x_raw, h, c):
 ).to_edge(exir.EdgeCompileConfig(_check_ir_validity=False))

 program_without_delegates = (
 exir.capture(
 composite_m,
 (input_x, input_h, input_c),
 exir.CaptureConfig(pt2_mode=True),
 )
 exir.capture(CompositeModel(3), inputs)
 .to_edge(exir.EdgeCompileConfig(_check_ir_validity=False))
 .to_executorch(
 config=exir.ExecutorchBackendConfig(extract_segments=extract_segments),
Expand DownExpand Up@@ -726,7 +722,7 @@ def forward(self, x_raw, h, c):

 program_without_delegates = (
 exir.capture(
composite_m,
CompositeModel(3),
 (input_x, input_h, input_c),
 exir.CaptureConfig(pt2_mode=True),
 )
Expand DownExpand Up@@ -962,7 +958,8 @@ def test_quantized_with_delegate(self) -> None:
 example_inputs,
 exir.CaptureConfig(
 pt2_mode=True,
 enable_functionalization=False,
 enable_aot=True,
 _unlift=True,
 ),
 ).to_edge(exir.EdgeCompileConfig(_check_ir_validity=False))
 FileCheck().check_count("quantize_per_tensor.default", 3).check("addmm").run(
Expand Down

160 changes: 158 additions & 2 deletions
exir/__init__.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -5,9 +5,11 @@
from collections import namedtuple
from dataclasses import dataclass, field
from typing import Any, Callable, Dict, List, Optional, Set, Tuple, Union
from unittest.mock import patch

import sympy
import torch
import torch._export
from executorch.exir.dynamic_shape import DynamicMemoryPlanningMode
from executorch.exir.emit import emit_program, EmitterOutput
from executorch.exir.error import ExportError, ExportErrorType, InternalError
Expand All@@ -25,6 +27,7 @@
from executorch.exir.schema import Program
from executorch.exir.serialize import serialize_to_flatbuffer
from executorch.exir.tracer import (
 _default_decomposition_table,
 dispatch_trace,
 dynamo_trace,
 ExirDynamoConfig,
Expand All@@ -41,6 +44,7 @@
from torch._dynamo.eval_frame import Constraint
from torch._export import CallSpec, export, ExportGraphSignature
from torch._export.exported_program import ExportedProgram
from torch._export.passes import ReplaceViewOpsWithViewCopyOpsPass
from torch._export.passes.add_runtime_assertions_for_constraints_pass import (
 InputDim,
 RangeConstraint,
Expand All@@ -49,12 +53,156 @@
from torch.fx._compatibility import compatibility
from torch.fx.experimental.proxy_tensor import make_fx
from torch.fx.experimental.symbolic_shapes import ShapeEnv
from torch.fx.graph import _PyTreeCodeGen, _PyTreeInfo
from torch.utils import _pytree as pytree


Val = Any


def _unlift(gm, inp_pos_to_param_buffer_name, in_spec, out_spec, state_dict):
 count = 0
 # Step 1: make lifted params as get_attr
 for node in gm.graph.nodes:
 if node.op == "placeholder":
 if count in inp_pos_to_param_buffer_name:
 with gm.graph.inserting_after(node):
 getattr_node = gm.graph.get_attr(
 inp_pos_to_param_buffer_name[count]
 )
 node.replace_all_uses_with(getattr_node)
 metadata = node.meta
 gm.graph.erase_node(node)
 getattr_node.meta = metadata
 count += 1

 # Step 2: Fix the input/output of the graph now that we deleted
 # some args.
 gm.graph.lint()
 names = [f"arg_{i}" for i in range(len(in_spec.children_specs))]
 gm.graph._codegen = _PyTreeCodeGen(
 _PyTreeInfo(
 names,
 in_spec,
 out_spec,
 )
 )
 gm.recompile()

 # Step 3: Find state references in HigherOrderOps and recursively
 # fix them.
 for node in gm.graph.nodes:
 if node.op == "call_function" and node.target == torch.ops.cond:
 pred, true_graph, false_graph, operands = node.args
 true_gm = getattr(gm, true_graph.name)
 false_gm = getattr(gm, false_graph.name)
 inp_pos_to_param_buffer_name_for_submod = {}
 real_operands = []
 for ix, operand in enumerate(operands):
 if operand.target in inp_pos_to_param_buffer_name.values():
 inp_pos_to_param_buffer_name_for_submod[ix] = operand.target
 true_gm.register_buffer(operand.target, state_dict[operand.target])
 false_gm.register_buffer(operand.target, state_dict[operand.target])
 else:
 real_operands.append(operand)
 node.args = (pred, true_graph, false_graph, real_operands)

 _, in_spec = pytree.tree_flatten(real_operands)

 _unlift(
 true_gm,
 inp_pos_to_param_buffer_name_for_submod,
 in_spec,
 None,
 state_dict,
 )
 _unlift(
 false_gm,
 inp_pos_to_param_buffer_name_for_submod,
 in_spec,
 None,
 state_dict,
 )
 if node.op == "call_function" and node.target.__name__ == "map_impl":
 body_graph, num_mapped, *operands = node.args
 body_gm = getattr(gm, body_graph.name)
 inp_pos_to_buffer_name_for_submod = {}
 real_operands = []
 for ix, operand in enumerate(operands):
 if operand.target in inp_pos_to_param_buffer_name.values():
 inp_pos_to_buffer_name_for_submod[ix] = operand.target
 body_gm.register_buffer(operand.target, state_dict[operand.target])
 else:
 real_operands.append(operand)
 node.args = (body_graph, num_mapped, *real_operands)

 _, in_spec = pytree.tree_flatten(real_operands)

 _unlift(
 body_gm, inp_pos_to_buffer_name_for_submod, in_spec, None, state_dict
 )
 gm.graph.lint()
 gm.graph.eliminate_dead_code()
 gm.recompile()
 return gm


def unlift_exported_program_lifted_states(
 ep: torch._export.exported_program.ExportedProgram,
):
 new_gm = copy.deepcopy(ep.graph_module)

 # TODO Fix the period in params/buffers names later
 # maybe a pass to replace graph signature with fixed names
 param_buffer_name_to_corrected_name = {}

 for name, stuff in ep.state_dict.items():
 if name in ep.graph_signature.buffers:
 if "." in name:
 new_gm.register_buffer(name.replace(".", "_"), stuff)
 param_buffer_name_to_corrected_name[name] = name.replace(".", "_")
 else:
 new_gm.register_buffer(name, stuff)
 elif name in ep.graph_signature.parameters:
 if "." in name:
 new_gm.register_parameter(name.replace(".", "_"), stuff)
 param_buffer_name_to_corrected_name[name] = name.replace(".", "_")
 else:
 new_gm.register_parameter(name, stuff)
 else:
 raise AssertionError("encountered not registered param/buffer")

 count = 0
 inp_pos_to_param_buffer_name = {}
 for node in new_gm.graph.nodes:
 if node.op == "placeholder":
 if node.name in ep.graph_signature.inputs_to_buffers:
 buffer_name = ep.graph_signature.inputs_to_buffers[node.name]
 if buffer_name in param_buffer_name_to_corrected_name:
 inp_pos_to_param_buffer_name[
 count
 ] = param_buffer_name_to_corrected_name[buffer_name]
 else:
 inp_pos_to_param_buffer_name[count] = buffer_name
 if node.name in ep.graph_signature.inputs_to_parameters:
 param_name = ep.graph_signature.inputs_to_parameters[node.name]
 if param_name in param_buffer_name_to_corrected_name:
 inp_pos_to_param_buffer_name[
 count
 ] = param_buffer_name_to_corrected_name[param_name]
 else:
 inp_pos_to_param_buffer_name[count] = param_name
 count += 1
 new_gm = _unlift(
 new_gm,
 inp_pos_to_param_buffer_name,
 ep.call_spec.in_spec,
 ep.call_spec.out_spec,
 ep.state_dict,
 )
 return new_gm


@compatibility(is_backward_compatible=False)
@dataclass
class CaptureConfig:
Expand All@@ -63,6 +211,7 @@ class CaptureConfig:
 enable_dynamic_shape: bool = False
 enable_aot: bool = False
 _dynamo_config: "ExirDynamoConfig" = ExirDynamoConfig()
 _unlift: bool = False


@compatibility(is_backward_compatible=False)
Expand DownExpand Up@@ -400,8 +549,15 @@ def capture(
 "Functionalization is required for enable_aot.",
 )

 ep = export(f, args, _add_runtime_assertions=False, constraints=constraints)
 return ep # pyre-ignore
 # TODO remove this later
 with patch("torch._export.DECOMP_TABLE", _default_decomposition_table()):
 ep = export(
 f, args, _add_runtime_assertions=False, constraints=constraints
 )
 ep = ep.transform(ReplaceViewOpsWithViewCopyOpsPass())
 if not config._unlift:
 return ep # pyre-ignore
 graph_module = unlift_exported_program_lifted_states(ep)

 elif config.enable_dynamic_shape:
 if not config._dynamo_config.dynamic_shapes:
Expand Down

82 changes: 58 additions & 24 deletions
exir/dialects/edge/edge.yaml

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -89,6 +89,14 @@
 mat2: T0
 __ret_0: T0

- func: aten::arange.start_step
 namespace: edge
 inherits: aten::arange.start_step
 type_alias:
 T0: [Byte, Char, Double, Float, Int, Long, Short]
 type_constraint:
 - __ret_0: T0

- func: aten::bmm
 namespace: edge
 inherits: aten::bmm
Expand DownExpand Up@@ -198,14 +206,43 @@
 - self: T0
 __ret_0: T0

- func: aten::lift_fresh_copy
- func: aten::index_select
 namespace: edge
 inherits: aten::lift_fresh_copy
 inherits: aten::index_select
 type_alias:
 T0: [Bool, Byte, Char, Double, Float, Int, Long, Short]
 T0: [Bool]
 T1: [Byte]
 T2: [Char]
 T3: [Double]
 T4: [Float]
 T5: [Int]
 T6: [Long]
 T7: [Short]
 type_constraint:
 - self: T0
 index: T6
 __ret_0: T0
 - self: T1
 index: T6
 __ret_0: T1
 - self: T2
 index: T6
 __ret_0: T2
 - self: T3
 index: T6
 __ret_0: T3
 - self: T4
 index: T6
 __ret_0: T4
 - self: T5
 index: T6
 __ret_0: T5
 - self: T6
 index: T6
 __ret_0: T6
 - self: T7
 index: T6
 __ret_0: T7

- func: aten::masked_fill.Scalar
 namespace: edge
Expand DownExpand Up@@ -245,16 +282,6 @@
 mask: T0
 __ret_0: T7

- func: aten::minimum
 namespace: edge
 inherits: aten::minimum
 type_alias:
 T0: [Bool, Byte, Char, Double, Float, Int, Long, Short]
 type_constraint:
 - self: T0
 other: T0
 __ret_0: T0

- func: aten::mm
 namespace: edge
 inherits: aten::mm
Expand DownExpand Up@@ -324,15 +351,6 @@
 - self: T0
 __ret_0: T0

- func: aten::select_copy.int
 namespace: edge
 inherits: aten::select_copy.int
 type_alias:
 T0: [Bool, Byte, Char, Double, Float, Int, Long, Short]
 type_constraint:
 - self: T0
 __ret_0: T0

- func: aten::sigmoid
 namespace: edge
 inherits: aten::sigmoid
Expand DownExpand Up@@ -383,9 +401,25 @@
 other: T0
 __ret_0: T0

- func: aten::t
- func: aten::sym_numel
 namespace: edge
 inherits: aten::sym_numel
 type_alias:
 T0: [Bool, Byte, Char, Double, Float, Int, Long, Short]
 type_constraint:
 - self: T0

- func: aten::sym_size.int
 namespace: edge
 inherits: aten::sym_size.int
 type_alias:
 T0: [Bool, Byte, Char, Double, Float, Int, Long, Short]
 type_constraint:
 - self: T0

- func: aten::t_copy
 namespace: edge
 inherits: aten::t
 inherits: aten::t_copy
 type_alias:
 T0: [Bool, Byte, Char, Double, Float, Int, Long, Short]
 type_constraint:
Expand Down

2 changes: 2 additions & 0 deletions
exir/dialects/edge/yaml_generator.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -143,6 +143,8 @@ def get_test_gen_key(op_name: str) -> str:
 opdb_key = opdb_key[:-5]
 elif opdb_key == "sym_size":
 opdb_key = "resize_"
 elif opdb_key == "sym_numel":
 opdb_key = "abs"
 elif opdb_key == "convolution":
 opdb_key = "conv_transpose2d"
 elif opdb_key == "embedding":
Expand Down
Loading