[serialize] Support prelu

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),), {}

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)

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)

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,

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)

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,