Skip to content

Lower Conv1d atomically through TOSA Conv2d (#22282) #22282

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

Open
mcremon-meta wants to merge 1 commit into
base: main
Choose a base branch
from
+289 −8
Open

Lower Conv1d atomically through TOSA Conv2d (#22282) #22282

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
19 changes: 18 additions & 1 deletion backends/arm/_passes/conv1d_unsqueeze_pass.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,8 @@

from typing import Set, Type

import torch

from executorch.backends.arm._passes import ArmOpTargetedPass
from executorch.backends.arm._passes.convert_squeezes_to_view import (
ConvertSqueezesToViewPass,
Expand All @@ -21,11 +23,26 @@


class Conv1dUnsqueezePass(ConvertConv1dToConv2dPass, ArmOpTargetedPass):
"""Arm wrapper for the shared Conv1d-to-Conv2d transform."""
"""Convert ConvTranspose1d into ConvTranspose2d.

Forward Conv1d is lowered atomically by ``RewriteConvPass`` so its layout
transforms can be emitted in the rank-3 domain. TOSA has no native
transpose-convolution 1D operator, so this pass retains the rank-4
expansion for ConvTranspose1d.

"""

_passes_required_after: Set[Type[ExportPass]] = {
ConvertSqueezesToViewPass,
RewriteConvPass,
SizeAdjustInputPass,
}
target_ops = (exir_ops.edge.aten.convolution.default,)

def _conv1d_weight_node(self, node: torch.fx.Node) -> torch.fx.Node | None:
weight_node = super()._conv1d_weight_node(node)
# A non-None weight node means the base pass matched the full nine-arg
# convolution signature, so args[6] (transposed) is always present.
if weight_node is None or not node.args[6]:
return None
return weight_node
122 changes: 115 additions & 7 deletions backends/arm/_passes/rewrite_conv_pass.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -121,14 +121,14 @@ def _adjust_pad_if_needed(

return pad - mod_remainder

def _is_depthwise_conv2d(self, node: torch.fx.Node) -> bool:
def _is_depthwise_conv(self, node: torch.fx.Node) -> bool:
if (
node.op != "call_function"
or node.target != exir_ops.edge.aten.convolution.default
):
return False
input_tensor = get_first_fake_tensor(node.all_input_nodes[0])
if len(input_tensor.shape) != 4:
if len(input_tensor.shape) not in (3, 4):
return False
groups = node.args[-1]
in_channels = input_tensor.shape[1]
Expand Down Expand Up @@ -558,8 +558,26 @@ def _insert_layout_permute(
input_node: torch.fx.Node,
input_fake_tensor: FakeTensor,
dims: tuple[int, ...],
squeeze_spatial: bool = False,
) -> tuple[torch.fx.Node, FakeTensor]:
"""Insert the mandatory TOSA-to-Edge output layout permutation."""
"""Insert the mandatory TOSA-to-Edge output layout permutation.

``squeeze_spatial`` drops the singleton spatial dimension a rank-three
convolution carries through TOSA Conv2d, so the layout permutation is
emitted in the rank-three domain the exported graph expects.

"""
if squeeze_spatial:
squeezed_fake_tensor = cast(FakeTensor, input_fake_tensor.squeeze(1))
with graph_module.graph.inserting_after(input_node):
input_node = create_node(
graph=graph_module.graph,
op_target=exir_ops.edge.aten.view_copy.default,
args=(input_node, list(squeezed_fake_tensor.shape)),
from_node=source_node,
)
input_node.meta["val"] = squeezed_fake_tensor
input_fake_tensor = squeezed_fake_tensor
with graph_module.graph.inserting_after(input_node):
output = create_node(
graph=graph_module.graph,
Expand All @@ -579,6 +597,7 @@ def _insert_a16w8_output_branches(
tosa_node_fake_tensor: torch.Tensor,
default_rescale: torch.fx.Node,
post_permute_dims: tuple[int, ...],
squeeze_spatial: bool = False,
) -> None:
"""Route A16W8 convolution users through the required output types.

Expand Down Expand Up @@ -652,6 +671,7 @@ def _insert_a16w8_output_branches(
int32_user,
int32_fake_tensor,
post_permute_dims,
squeeze_spatial,
)
for user in previous_users:
user.replace_input_with(int32_user, int32_output)
Expand Down Expand Up @@ -702,6 +722,7 @@ def _insert_a16w8_output_branches(
widened_rescale,
widened_fake_tensor,
post_permute_dims,
squeeze_spatial,
)

# Clone the source permutation for the widened branch. The
Expand Down Expand Up @@ -911,7 +932,70 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult: # noqa: C901
dilation = tuple(dilation_list)
pad = pad_attr

if self._is_conv3d(len(input_shape), group):
if spatial_rank == 1:
target_op = (
exir_ops.backend.tosa.DEPTHWISE_CONV2D.default
if self._is_depthwise_conv(node)
else exir_ops.backend.tosa.CONV2D.default
)
pre_permute_dims = (0, 2, 1)
post_permute_dims = (0, 2, 1)
with graph_module.graph.inserting_before(node):
x = create_node(
graph=graph_module.graph,
op_target=exir_ops.edge.aten.permute_copy.default,
args=(x, list(pre_permute_dims)),
from_node=node,
)
permuted_input_fake = permute_fake_tensor_metadata(
input_fake_tensor, pre_permute_dims
)
x.meta["val"] = permuted_input_fake
input_tensor_for_tosa_fake = permuted_input_fake.unsqueeze(1)
x = create_node(
graph=graph_module.graph,
op_target=exir_ops.edge.aten.view_copy.default,
args=(x, list(input_tensor_for_tosa_fake.shape)),
from_node=node,
)
x.meta["val"] = input_tensor_for_tosa_fake

kernel_width = weight_shape[2]
if target_op == exir_ops.backend.tosa.DEPTHWISE_CONV2D.default:
in_channels = input_fake_tensor.shape[1]
channel_multiplier = weight_shape[0] // in_channels
weight = self._rewrite_weight(
graph_module,
weight,
node,
permute_dims=(1, 2, 0),
name_suffix="hwicm",
reshape_dims=(
1,
kernel_width,
in_channels,
channel_multiplier,
),
)
else:
weight = self._rewrite_weight(
graph_module,
weight,
node,
permute_dims=(0, 2, 1),
name_suffix="ohwi",
reshape_dims=(
weight_shape[0],
1,
kernel_width,
weight_shape[1],
),
)
weight_fake_tensor = get_first_fake_tensor(weight)
stride = (1, stride[0])
dilation = (1, dilation[0])
pad = [0, 0, pad[0], pad[1]]
elif self._is_conv3d(len(input_shape), group):
target_op = exir_ops.backend.tosa.CONV3D.default
pre_permute_dims = ODHWI_ORDER
post_permute_dims = ODHWI_INVERSE_ORDER
Expand All @@ -934,7 +1018,7 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult: # noqa: C901
name_suffix="odhwi",
)
weight_fake_tensor = get_first_fake_tensor(weight)
elif self._is_depthwise_conv2d(node):
elif self._is_depthwise_conv(node):
target_op = exir_ops.backend.tosa.DEPTHWISE_CONV2D.default
pre_permute_dims = NHWC_ORDER
post_permute_dims = NHWC_INVERSE_ORDER
Expand Down Expand Up @@ -1039,7 +1123,28 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult: # noqa: C901

if post_permute_dims is None:
raise RuntimeError("Expected post permute dims for explicit layout")
output_conversion_node = node_replacement
post_permute_input = node_replacement
squeeze_view: torch.fx.Node | None = None
if spatial_rank == 1:
squeezed_output_fake = cast(
FakeTensor, node_replacement_fake_tensor.squeeze(1)
)
special_dtype = node_replacement.meta.get(TosaSpecialDtype.meta_key())
with graph_module.graph.inserting_after(node_replacement):
node_replacement = create_node(
graph=graph_module.graph,
op_target=exir_ops.edge.aten.view_copy.default,
args=(node_replacement, list(squeezed_output_fake.shape)),
from_node=node,
)
node_replacement.meta["val"] = squeezed_output_fake
if special_dtype:
node_replacement.meta[TosaSpecialDtype.meta_key()] = special_dtype
squeeze_view = node_replacement
post_permute_input = node_replacement
node_replacement_fake_tensor = squeezed_output_fake

with graph_module.graph.inserting_after(node_replacement):
node_replacement = create_node(
graph=graph_module.graph,
Expand Down Expand Up @@ -1068,8 +1173,9 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult: # noqa: C901
node,
tosa_op,
tosa_node_fake_tensor,
post_permute_input,
output_conversion_node,
post_permute_dims,
spatial_rank == 1,
)
# Only users not moved to widened branches remain on the
# source node. Route those through the declared INT16 output,
Expand All @@ -1078,7 +1184,9 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult: # noqa: C901
node.replace_all_uses_with(node_replacement)
else:
graph_module.graph.erase_node(node_replacement)
graph_module.graph.erase_node(post_permute_input)
if squeeze_view is not None:
graph_module.graph.erase_node(squeeze_view)
graph_module.graph.erase_node(output_conversion_node)
else:
node.replace_all_uses_with(node_replacement)

Expand Down
36 changes: 36 additions & 0 deletions backends/arm/test/passes/test_remove_permutes_around_elementwise_tosa_ops.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,8 @@
RESCALE_TARGET = exir_ops.backend.tosa.RESCALE.default
MUL_TARGET = exir_ops.edge.aten.mul.Tensor
ADD_TARGET = exir_ops.edge.aten.add.Tensor
SUB_TARGET = exir_ops.edge.aten.sub.Tensor
VIEW_TARGET = exir_ops.edge.aten.view_copy.default
ERF_TARGET = exir_ops.edge.aten.erf.default


Expand Down Expand Up @@ -150,6 +152,40 @@ def test_remove_permutes_around_rescale_tosa_INT() -> None:
assert _count_nodes(result.graph_module, RESCALE_TARGET) == 1


def test_sink_view_preserves_layout_through_rescale_to_broadcast_tosa_INT() -> None:
graph = torch.fx.Graph()
x = graph.placeholder("x")
x.meta["val"] = torch.randn(1, 4, 1, 1)
direct = graph.placeholder("direct")
direct.meta["val"] = torch.randn(1, 8, 4)

permute = graph.create_node("call_function", PERMUTE_TARGET, args=(x, [0, 2, 3, 1]))
permute.meta["val"] = torch.randn(1, 1, 1, 4)
mul = graph.create_node("call_function", MUL_TARGET, args=(permute, permute))
mul.meta["val"] = torch.randn(1, 1, 1, 4)
sink = graph.create_node("call_function", VIEW_TARGET, args=(mul, [1, 1, 4]))
sink.meta["val"] = torch.randn(1, 1, 4)
rescale = graph.create_node(
"call_function",
RESCALE_TARGET,
args=(sink, torch.int8, [1.0], 0, 0),
)
rescale.meta["val"] = torch.randn(1, 1, 4)
sub = graph.create_node("call_function", SUB_TARGET, args=(direct, rescale))
sub.meta["val"] = torch.randn(1, 8, 4)
graph.output(sub)

graph_module = torch.fx.GraphModule({}, graph)
with TosaLoweringContext(TOSA_INT_SPEC):
result = RemovePermutesAroundElementwiseTosaOps(_fake_exported_program()).call(
graph_module
)

assert not result.modified
assert _count_nodes(result.graph_module, PERMUTE_TARGET) == 1
assert sub.args == (direct, rescale)


def test_remove_permutes_around_gelu_with_folded_scalar_constants_tosa_FP() -> None:
graph = torch.fx.Graph()
x = graph.placeholder("x")
Expand Down
Loading
Loading