diff --git a/test/modules/op/prelu.py b/test/modules/op/prelu.py index 87b20b45..250333bb 100644 --- a/test/modules/op/prelu.py +++ b/test/modules/op/prelu.py @@ -28,3 +28,17 @@ def forward(self, x): def get_example_inputs(self): return (torch.randn(1, 2, 3, 3),), {} + + +class ChannelWisePReLU(TestModuleBase): + """Exercise PReLU with one learned slope per NCHW channel.""" + + def __init__(self): + super().__init__() + self.prelu = torch.nn.PReLU(num_parameters=2) + + def forward(self, x): + return self.prelu(x) + + def get_example_inputs(self): + return (torch.randn(1, 2, 3, 3),), {} diff --git a/test/unit_test/passes/test_legalize_predefined_layout_operators.py b/test/unit_test/passes/test_legalize_predefined_layout_operators.py new file mode 100644 index 00000000..30c29f67 --- /dev/null +++ b/test/unit_test/passes/test_legalize_predefined_layout_operators.py @@ -0,0 +1,130 @@ +# Copyright (c) 2026 Samsung Electronics Co., Ltd. All Rights Reserved +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import torch + +from tico.passes.legalize_predefined_layout_operators import ( + LegalizePreDefinedLayoutOperators, +) + +from test.support.helper import num_of_ops +from test.support.pass_value_test import SinglePassValueTest + + +class ChannelWisePReLUNet(torch.nn.Module): + """Apply one PReLU slope per channel to an NCHW input.""" + + def __init__(self) -> None: + super().__init__() + self.prelu = torch.nn.PReLU(num_parameters=4) + + def forward(self, input_: torch.Tensor) -> torch.Tensor: + """Return the channel-wise PReLU output.""" + return self.prelu(input_) + + def get_example_inputs(self): + """Return representative inputs for export and value comparison.""" + return (torch.randn(1, 4, 5, 7),), {} + + +class SharedPReLUNet(torch.nn.Module): + """Apply one PReLU slope shared by every input element.""" + + def __init__(self) -> None: + super().__init__() + self.prelu = torch.nn.PReLU(num_parameters=1) + + def forward(self, input_: torch.Tensor) -> torch.Tensor: + """Return the shared-slope PReLU output.""" + return self.prelu(input_) + + def get_example_inputs(self): + """Return representative inputs for export and value comparison.""" + return (torch.randn(1, 4, 5, 7),), {} + + +class ConvChannelWisePReLUNet(torch.nn.Module): + """Apply channel-wise PReLU directly after an NCHW convolution.""" + + def __init__(self) -> None: + super().__init__() + self.conv = torch.nn.Conv2d(3, 4, kernel_size=3, padding=1) + self.prelu = torch.nn.PReLU(num_parameters=4) + + def forward(self, input_: torch.Tensor) -> torch.Tensor: + """Return the convolution followed by channel-wise PReLU.""" + return self.prelu(self.conv(input_)) + + def get_example_inputs(self): + """Return representative inputs for export and value comparison.""" + return (torch.randn(1, 3, 5, 7),), {} + + +class LegalizeChannelWisePReLUTest(SinglePassValueTest): + """Verify channel-wise PReLU layout legalization and numerical parity.""" + + def test_pass(self): + """Replace aten.prelu with the channel-last Circle custom operator.""" + self.setup(ChannelWisePReLUNet()) + self.assertEqual( + num_of_ops(self.exported_program(), [torch.ops.aten.prelu.default]), 1 + ) + + self.run_value_test(LegalizePreDefinedLayoutOperators()) + + self.assertEqual( + num_of_ops(self.exported_program(), [torch.ops.aten.prelu.default]), 0 + ) + self.assertEqual( + num_of_ops(self.exported_program(), [torch.ops.circle_custom.prelu]), 1 + ) + self.assertEqual( + num_of_ops(self.exported_program(), [torch.ops.aten.permute.default]), 2 + ) + + def test_shared_slope_is_not_rewritten(self): + """Keep a broadcast-safe shared PReLU slope in native aten form.""" + self.setup(SharedPReLUNet()) + self.run_value_test(LegalizePreDefinedLayoutOperators()) + + self.assertEqual( + num_of_ops(self.exported_program(), [torch.ops.aten.prelu.default]), 1 + ) + self.assertEqual( + num_of_ops(self.exported_program(), [torch.ops.circle_custom.prelu]), 0 + ) + + def test_reuses_channel_last_convolution_output(self): + """Avoid an inverse permutation pair between Conv2D and PReLU.""" + self.setup(ConvChannelWisePReLUNet()) + self.run_value_test(LegalizePreDefinedLayoutOperators()) + + self.assertEqual( + num_of_ops(self.exported_program(), [torch.ops.circle_custom.conv2d]), 1 + ) + self.assertEqual( + num_of_ops(self.exported_program(), [torch.ops.circle_custom.prelu]), 1 + ) + + graph_nodes = list(self.exported_program().graph.nodes) + circle_conv = next( + node + for node in graph_nodes + if node.target == torch.ops.circle_custom.conv2d + ) + circle_prelu = next( + node for node in graph_nodes if node.target == torch.ops.circle_custom.prelu + ) + + self.assertIs(circle_prelu.args[0], circle_conv) diff --git a/test/unit_test/utils/test_register_custom_op.py b/test/unit_test/utils/test_register_custom_op.py index af7c98f8..fee5bb5c 100644 --- a/test/unit_test/utils/test_register_custom_op.py +++ b/test/unit_test/utils/test_register_custom_op.py @@ -314,6 +314,39 @@ def test_circle_instance_norm_with_custom_params(self): # Check output shape self.assertEqual(list(result.shape), list(input_tensor.shape)) + def test_circle_prelu_channelwise(self): + """Test CirclePReLU with one slope per NHWC channel.""" + input_tensor = torch.randn(2, 4, 5, 3) + weight_tensor = torch.randn(3) + + result = torch.ops.circle_custom.prelu(input_tensor, weight_tensor) + expected = torch.where( + input_tensor >= 0, + input_tensor, + input_tensor * weight_tensor.reshape(1, 1, 1, 3), + ) + + torch.testing.assert_close(result, expected) + + def test_circle_prelu_shared_slope(self): + """Test CirclePReLU with one slope shared by all channels.""" + input_tensor = torch.randn(2, 4, 5, 3) + weight_tensor = torch.randn(1) + + result = torch.ops.circle_custom.prelu(input_tensor, weight_tensor) + expected = torch.where( + input_tensor >= 0, + input_tensor, + input_tensor * weight_tensor, + ) + + torch.testing.assert_close(result, expected) + + def test_circle_prelu_rejects_channel_mismatch(self): + """Test that CirclePReLU rejects an incompatible channel count.""" + with self.assertRaisesRegex(RuntimeError, "match the last input dimension"): + torch.ops.circle_custom.prelu(torch.randn(1, 4, 5, 3), torch.randn(2)) + def test_circle_mx_fake_quantize_basic(self): """Test CircleMXFakeQuantize basic functionality.""" input_tensor = torch.randn(2, 32, 32, 3) diff --git a/tico/passes/legalize_predefined_layout_operators.py b/tico/passes/legalize_predefined_layout_operators.py index 5e9bd690..c579f4ba 100644 --- a/tico/passes/legalize_predefined_layout_operators.py +++ b/tico/passes/legalize_predefined_layout_operators.py @@ -35,6 +35,8 @@ DequantizePerTensorArgs, InstanceNormArgs, MaxPool2dWithIndicesArgs, + PermuteArgs, + PReLUArgs, ) @@ -94,11 +96,112 @@ class LegalizePreDefinedLayoutOperators(PassBase): Input[NCHW] ---- aten.permute(NCHW_to_NHWC) ---- circle_cumstom.depthwise_conv2d[NHWC] ---- aten.permute(NHWC_to_NCHW) ---- OUTPUT[NCHW] Weight[CNHW] - (aten.dequantize) - aten.permute(CNHW_to_NHWC) ---/ Bias ----------(aten.dequantize) -------------------------------/ + + [3] aten.prelu with one slope per PyTorch channel + + [BEFORE PASS] + Input[N,C,...] ---- aten.prelu[N,C,...] ---- OUTPUT[N,C,...] + Weight[C] --------/ + + [AFTER PASS] + Input[N,C,...] ---- aten.permute(channel_first_to_last) ---- circle_custom.prelu[N,...,C] ---- aten.permute(channel_last_to_first) ---- OUTPUT[N,C,...] + Weight[C] --------------------------------------------------/ """ def __init__(self): super().__init__() + def legalize_prelu( + self, exported_program: ExportedProgram, node: torch.fx.Node + ) -> bool: + """Lower channel-wise PyTorch PReLU to Circle PReLU. + + PyTorch applies a rank-1 PReLU weight along input dimension 1. Circle + PReLU follows regular right-aligned broadcasting, so this pass moves + dimension 1 to the last position before lowering the operation. The + output is permuted back to preserve the PyTorch-visible layout. + + A shared one-element slope does not require legalization. Rank-2 input + also needs no transformation because its channel dimension is already + the last dimension. + """ + logger = logging.getLogger(__name__) + + graph_module = exported_program.graph_module + graph = graph_module.graph + + args = PReLUArgs(*node.args, **node.kwargs) # type: ignore[arg-type] + input_ = args.input + weight = args.weight + + input_shape = extract_shape(input_) + weight_shape = extract_shape(weight) + + if len(weight_shape) != 1: + raise NotYetSupportedError( + "Only rank-1 PReLU weights are supported: " + f"weight shape={weight_shape}." + ) + + num_parameters = weight_shape[0] + if num_parameters == 1 or len(input_shape) <= 2: + return False + + input_channels = input_shape[1] + if isinstance(input_channels, int) and isinstance(num_parameters, int): + if input_channels != num_parameters: + raise NotYetSupportedError( + "PReLU weight size must match input dimension 1: " + f"input shape={input_shape}, weight shape={weight_shape}." + ) + + rank = len(input_shape) + channel_first_to_last = [0, *range(2, rank), 1] + channel_last_to_first = [0, rank - 1, *range(1, rank - 1)] + + # Reuse an already channel-last producer when the input is the output + # permutation inserted for another predefined-layout operator. + prelu_input = input_ + if is_target_node(input_, [torch.ops.aten.permute.default]): + input_permute_args = PermuteArgs( + *input_.args, **input_.kwargs # type: ignore[arg-type] + ) + if list(input_permute_args.dims) == channel_last_to_first: + prelu_input = input_permute_args.input + else: + with graph.inserting_after(input_): + prelu_input = create_node( + graph, + torch.ops.aten.permute.default, + args=(input_, channel_first_to_last), + origin=input_, + ) + else: + with graph.inserting_after(input_): + prelu_input = create_node( + graph, + torch.ops.aten.permute.default, + args=(input_, channel_first_to_last), + origin=input_, + ) + + with graph.inserting_before(node): + circle_prelu = create_node( + graph, + torch.ops.circle_custom.prelu, + args=(prelu_input, weight), + origin=node, + ) + prelu_out_permute = create_node( + graph, + torch.ops.aten.permute.default, + args=(circle_prelu, channel_last_to_first), + ) + + node.replace_all_uses_with(prelu_out_permute, propagate_meta=True) + logger.debug(f"{node.name} is replaced with {circle_prelu.name}") + return True + def legalize_conv2d(self, exported_program, node) -> bool: logger = logging.getLogger(__name__) modified = False @@ -436,6 +539,7 @@ def legalize_avg_pool2d(self, exported_program, node) -> bool: def call(self, exported_program: ExportedProgram) -> PassResult: target_to_legalize_func = { + torch.ops.aten.prelu.default: self.legalize_prelu, torch.ops.aten.conv2d.default: self.legalize_conv2d, torch.ops.aten.conv2d.padding: self.legalize_conv2d, torch.ops.aten.conv_transpose2d.input: self.legalize_conv_transpose2d, diff --git a/tico/serialize/operators/op_prelu.py b/tico/serialize/operators/op_prelu.py index 7f89f0be..79d1d277 100644 --- a/tico/serialize/operators/op_prelu.py +++ b/tico/serialize/operators/op_prelu.py @@ -29,7 +29,12 @@ @register_node_visitor class PReLUVisitor(NodeVisitor): - target: List[torch._ops.OpOverload] = [torch.ops.aten.prelu.default] + """Serialize native or layout-legalized PReLU nodes as Circle PRELU.""" + + target: List[torch._ops.OpOverload] = [ + torch.ops.aten.prelu.default, + torch.ops.circle_custom.prelu, + ] def __init__(self, op_codes: Dict[OpCode, int], graph: CircleSubgraph): super().__init__(op_codes, graph) diff --git a/tico/utils/convert.py b/tico/utils/convert.py index 9e7afcce..d8522d6f 100644 --- a/tico/utils/convert.py +++ b/tico/utils/convert.py @@ -113,6 +113,7 @@ def run_decompositions_v25(ep: ExportedProgram): torch.ops.aten.instance_norm.default, torch.ops.aten._safe_softmax.default, torch.ops.aten.relu6.default, # Do not decompose to hardtanh + torch.ops.aten.prelu.default, torch.ops.aten.linear.default, torch.ops.aten.upsample_nearest2d.vec, torch.ops.aten.rms_norm.default, diff --git a/tico/utils/register_custom_op.py b/tico/utils/register_custom_op.py index fb1f46ee..f3b34d2d 100644 --- a/tico/utils/register_custom_op.py +++ b/tico/utils/register_custom_op.py @@ -46,6 +46,66 @@ def _(input_: torch.Tensor, size: List[int]): return result +def CirclePReLU(): + """Register a channel-last PReLU operator for Circle lowering. + + PyTorch PReLU treats dimension 1 as the channel dimension, while Circle + PReLU applies regular right-aligned broadcasting. This internal operator + makes the channel-last contract explicit after layout legalization. + """ + + @custom_op("circle_custom::prelu", mutates_args=()) + def prelu(input_: torch.Tensor, weight: torch.Tensor) -> torch.Tensor: + """Apply PReLU with channel-last broadcasting semantics.""" + if input_.dim() == 0: + raise RuntimeError( + "CirclePReLU requires an input tensor with rank greater than zero." + ) + if weight.dim() != 1: + raise RuntimeError( + "CirclePReLU requires a rank-1 weight tensor, " + f"but received rank {weight.dim()}." + ) + + channels = input_.size(-1) + num_parameters = weight.numel() + if num_parameters not in (1, channels): + raise RuntimeError( + "CirclePReLU weight must contain one value or match the last " + "input dimension: " + f"weight={num_parameters}, channels={channels}." + ) + + broadcast_shape = [1] * (input_.dim() - 1) + [num_parameters] + alpha = weight.reshape(broadcast_shape) + return torch.where(input_ >= 0, input_, input_ * alpha) + + @register_fake("circle_custom::prelu") + def _(input_: torch.Tensor, weight: torch.Tensor) -> torch.Tensor: + """Infer metadata for the internal channel-last PReLU operator.""" + if input_.dim() == 0: + raise RuntimeError( + "CirclePReLU requires an input tensor with rank greater than zero." + ) + if weight.dim() != 1: + raise RuntimeError( + "CirclePReLU requires a rank-1 weight tensor, " + f"but received rank {weight.dim()}." + ) + + channels = input_.shape[-1] + num_parameters = weight.shape[0] + if isinstance(channels, int) and isinstance(num_parameters, int): + if num_parameters not in (1, channels): + raise RuntimeError( + "CirclePReLU weight must contain one value or match the last " + "input dimension: " + f"weight={num_parameters}, channels={channels}." + ) + + return input_.new_empty(input_.size()) + + def CircleConv2d(): """ Note that this op follows the input spec of `aten.conv2d.default` whose number @@ -902,6 +962,7 @@ def _(params: torch.Tensor, indices: torch.Tensor) -> torch.Tensor: # Add custom ops to the torch namespace def RegisterOps(): CircleResizeNearestNeighbor() + CirclePReLU() CircleDepthwiseConv2d() CircleDepthwiseConv2dPadding() CircleConv2d()