Samsung Exynos AI LiteCore - Lower more ops and add tests to support LLM models

backends/samsung/builders/__init__.py

@@ -41,6 +41,7 @@
     op_pad,
     op_permute,
     op_pixel_shuffle,
+    op_pixel_unshuffle,
     op_placeholder,
     op_pow,
     op_prelu,
@@ -106,6 +107,7 @@
     "op_mul",
     "op_permute",
     "op_pixel_shuffle",
+    "op_pixel_unshuffle",
     "op_placeholder",
     "op_pow",
     "op_prelu",

backends/samsung/builders/op_pixel_unshuffle.py

@@ -0,0 +1,44 @@
+# Copyright (c) 2025 Samsung Electronics Co. LTD
+# All rights reserved
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+from typing import cast, Dict
+
+import torch
+from executorch.backends.samsung.builders.node_visitor import (
+    NodeVisitor,
+    register_node_visitor,
+)
+from executorch.backends.samsung.serialization.enn_graph_schema import EnnGraph
+from executorch.backends.transforms import get_shape
+
+
+@register_node_visitor
+class PixelUnshuffleVisitor(NodeVisitor):
+    target = "aten.pixel_unshuffle.default"
+
+    def __init__(self, *args) -> None:
+        super().__init__(*args)
+
+    def define_node(
+        self,
+        node: torch.fx.Node,
+        enn_graph: EnnGraph,
+        vals_to_ids: Dict[torch.Tensor, int],
+    ) -> bool:
+        if len(get_shape(node.args[0])) != 4:
+            return False
+
+        input_id = self.define_tensor(node.args[0], enn_graph, vals_to_ids)
+
+        downscale_factor = cast(int, node.args[1])
+        params = {"block_size": downscale_factor, "mode": "CRD"}
+
+        output_id = self.define_tensor(node, enn_graph, vals_to_ids)
+
+        enn_graph.define_op(
+            node.name, "SPACE_TO_DEPTH", [input_id], [output_id], params
+        )
+
+        return True

backends/samsung/builders/op_rms_norm.py

@@ -4,6 +4,7 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+import logging
 from typing import cast, Dict, List
 
 import torch
@@ -13,6 +14,7 @@
 )
 from executorch.backends.samsung.builders.utils import get_tensor
 from executorch.backends.samsung.serialization.enn_graph_schema import EnnGraph
+from executorch.backends.transforms import get_shape
 
 
 @register_node_visitor
@@ -31,6 +33,12 @@ def define_node(
 
         # input2
         normalized_shape = cast(List[int], node.args[1])
+        input_shape = get_shape(input)
+        if len(normalized_shape) != 1 or normalized_shape[0] != input_shape[-1]:
+            logging.warning(
+                "Currently, Enn backend only supports rms norm with last input dimension."
+            )
+            return False
 
         gamma_node = node.args[2]
         gamma_id = self.define_tensor(gamma_node, enn_graph, vals_to_ids)
@@ -44,6 +52,7 @@ def define_node(
         params["normalize_shape"] = normalized_shape
         params["param_num"] = 2
         params["epsilon"] = epsilon
+        params["axis"] = [len(input_shape) - 1]
 
         output_id = self.define_tensor(node, enn_graph, vals_to_ids)
 

backends/samsung/partition/enn_partitioner.py

@@ -197,6 +197,7 @@ def ops_to_not_decompose(
             torch.ops.aten.prelu.default,
             torch.ops.aten.layer_norm.default,
             torch.ops.aten.pixel_shuffle.default,
+            torch.ops.aten.pixel_unshuffle.default,
             torch.ops.aten.hardsigmoid.default,
             torch.ops.aten.silu.default,
             torch.ops.aten.pad.default,

backends/samsung/test/models/test_mobilebert_qat.py

@@ -96,10 +96,10 @@ def test_mobilebert_qat_a8w8(self):
             example_inputs[1].size(2),
             example_inputs[1].size(3),
         )
-        vector_input_ids = torch.randint(0, 256, size_input_ids).to(device)
+        vector_input_ids = torch.randint(0, 256, size_input_ids, device=device)
         vector_attention_mask = torch.zeros(
-            size_attention_mask, dtype=torch.float32
-        ).to(device)
+            size_attention_mask, dtype=torch.float32, device=device
+        )
         export_inputs = (
             vector_input_ids,
             vector_attention_mask,

backends/samsung/test/ops/test_hardsigmoid.py

@@ -0,0 +1,47 @@
+# Copyright (c) Samsung Electronics Co. LTD
+# All rights reserved
+#
+# Licensed under the BSD License (the "License"); you may not use this file
+# except in compliance with the License. See the license file in the root
+# directory of this source tree for more details.
+
+import unittest
+
+import torch
+
+from executorch.backends.samsung.serialization.compile_options import (
+    gen_samsung_backend_compile_spec,
+)
+from executorch.backends.samsung.test.tester import SamsungTester
+from executorch.backends.samsung.test.utils.utils import TestConfig
+
+
+class HardSigmoid(torch.nn.Module):
+    def __init__(self) -> None:
+        super().__init__()
+        self.module = torch.nn.Hardsigmoid()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.module(x)
+
+
+class TestHardSigmoid(unittest.TestCase):
+    def _test(self, module: torch.nn.Module, inputs):
+        tester = SamsungTester(
+            module,
+            inputs,
+            [gen_samsung_backend_compile_spec(TestConfig.chipset)],
+        )
+        (
+            tester.export()
+            .check_count({"torch.ops.aten.hardsigmoid.default": 1})
+            .to_edge_transform_and_lower()
+            .check_not(["executorch_exir_dialects_edge__ops_aten_hardsigmoid_default"])
+            .check_count({"torch.ops.higher_order.executorch_call_delegate": 1})
+            .to_executorch()
+            .run_method_and_compare_outputs(inputs=inputs)
+        )
+
+    def test_fp32_hard_sigmoid(self):
+        inputs = (torch.randn(1, 16, 32, 32),)
+        self._test(HardSigmoid(), inputs)

backends/samsung/test/ops/test_hardswish.py

@@ -0,0 +1,47 @@
+# Copyright (c) Samsung Electronics Co. LTD
+# All rights reserved
+#
+# Licensed under the BSD License (the "License"); you may not use this file
+# except in compliance with the License. See the license file in the root
+# directory of this source tree for more details.
+
+import unittest
+
+import torch
+
+from executorch.backends.samsung.serialization.compile_options import (
+    gen_samsung_backend_compile_spec,
+)
+from executorch.backends.samsung.test.tester import SamsungTester
+from executorch.backends.samsung.test.utils.utils import TestConfig
+
+
+class HardSwish(torch.nn.Module):
+    def __init__(self) -> None:
+        super().__init__()
+        self.module = torch.nn.Hardswish()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.module(x)
+
+
+class TestHardSwish(unittest.TestCase):
+    def _test(self, module: torch.nn.Module, inputs):
+        tester = SamsungTester(
+            module,
+            inputs,
+            [gen_samsung_backend_compile_spec(TestConfig.chipset)],
+        )
+        (
+            tester.export()
+            .check_count({"torch.ops.aten.hardswish.default": 1})
+            .to_edge_transform_and_lower()
+            .check_not(["executorch_exir_dialects_edge__ops_aten_hardswish_default"])
+            .check_count({"torch.ops.higher_order.executorch_call_delegate": 1})
+            .to_executorch()
+            .run_method_and_compare_outputs(inputs=inputs, atol=0.002)
+        )
+
+    def test_fp32_hard_swish(self):
+        inputs = (torch.randn(1, 16, 32, 32),)
+        self._test(HardSwish(), inputs)

backends/samsung/test/ops/test_hardtanh.py

@@ -0,0 +1,47 @@
+# Copyright (c) Samsung Electronics Co. LTD
+# All rights reserved
+#
+# Licensed under the BSD License (the "License"); you may not use this file
+# except in compliance with the License. See the license file in the root
+# directory of this source tree for more details.
+
+import unittest
+
+import torch
+
+from executorch.backends.samsung.serialization.compile_options import (
+    gen_samsung_backend_compile_spec,
+)
+from executorch.backends.samsung.test.tester import SamsungTester
+from executorch.backends.samsung.test.utils.utils import TestConfig
+
+
+class Hardtanh(torch.nn.Module):
+    def __init__(self, min_val, max_val) -> None:
+        super().__init__()
+        self.module = torch.nn.Hardtanh(min_val=min_val, max_val=max_val)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.module(x)
+
+
+class TestHardtanh(unittest.TestCase):
+    def _test(self, module: torch.nn.Module, inputs):
+        tester = SamsungTester(
+            module,
+            inputs,
+            [gen_samsung_backend_compile_spec(TestConfig.chipset)],
+        )
+        (
+            tester.export()
+            .check_count({"torch.ops.aten.hardtanh.default": 1})
+            .to_edge_transform_and_lower()
+            .check_not(["executorch_exir_dialects_edge__ops_aten_hardtanh_default"])
+            .check_count({"torch.ops.higher_order.executorch_call_delegate": 1})
+            .to_executorch()
+            .run_method_and_compare_outputs(inputs=inputs)
+        )
+
+    def test_fp32_hardtanh(self):
+        inputs = (torch.randn(1, 3, 16, 16),)
+        self._test(Hardtanh(min_val=0, max_val=1), inputs)

backends/samsung/test/ops/test_layer_norm.py

@@ -0,0 +1,47 @@
+# Copyright (c) Samsung Electronics Co. LTD
+# All rights reserved
+#
+# Licensed under the BSD License (the "License"); you may not use this file
+# except in compliance with the License. See the license file in the root
+# directory of this source tree for more details.
+
+import unittest
+
+import torch
+
+from executorch.backends.samsung.serialization.compile_options import (
+    gen_samsung_backend_compile_spec,
+)
+from executorch.backends.samsung.test.tester import SamsungTester
+from executorch.backends.samsung.test.utils.utils import TestConfig
+
+
+class LayerNorm(torch.nn.Module):
+    def __init__(self) -> None:
+        super().__init__()
+        self.module = torch.nn.LayerNorm(10)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.module(x)
+
+
+class TestLayerNorm(unittest.TestCase):
+    def _test(self, module: torch.nn.Module, inputs):
+        tester = SamsungTester(
+            module,
+            inputs,
+            [gen_samsung_backend_compile_spec(TestConfig.chipset)],
+        )
+        (
+            tester.export()
+            .check_count({"torch.ops.aten.layer_norm.default": 1})
+            .to_edge_transform_and_lower()
+            .check_not(["executorch_exir_dialects_edge__ops_aten_layer_norm_default"])
+            .check_count({"torch.ops.higher_order.executorch_call_delegate": 1})
+            .to_executorch()
+            .run_method_and_compare_outputs(inputs=inputs)
+        )
+
+    def test_fp32_layer_norm(self):
+        inputs = (torch.randn(1, 32, 10),)
+        self._test(LayerNorm(), inputs)

backends/samsung/test/ops/test_maximum.py

@@ -0,0 +1,46 @@
+# Copyright (c) Samsung Electronics Co. LTD
+# All rights reserved
+#
+# Licensed under the BSD License (the "License"); you may not use this file
+# except in compliance with the License. See the license file in the root
+# directory of this source tree for more details.
+
+import unittest
+
+import torch
+
+from executorch.backends.samsung.serialization.compile_options import (
+    gen_samsung_backend_compile_spec,
+)
+from executorch.backends.samsung.test.tester import SamsungTester
+from executorch.backends.samsung.test.utils.utils import TestConfig
+
+
+class Maximum(torch.nn.Module):
+    def __init__(self) -> None:
+        super().__init__()
+
+    def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
+        return torch.maximum(x, y)
+
+
+class TestMaximum(unittest.TestCase):
+    def _test(self, module: torch.nn.Module, inputs):
+        tester = SamsungTester(
+            module,
+            inputs,
+            [gen_samsung_backend_compile_spec(TestConfig.chipset)],
+        )
+        (
+            tester.export()
+            .check_count({"torch.ops.aten.maximum.default": 1})
+            .to_edge_transform_and_lower()
+            .check_not(["executorch_exir_dialects_edge__ops_aten_maximum_default"])
+            .check_count({"torch.ops.higher_order.executorch_call_delegate": 1})
+            .to_executorch()
+            .run_method_and_compare_outputs(inputs=inputs)
+        )
+
+    def test_fp32_maximum(self):
+        inputs = (torch.randn(1, 3, 56, 56), torch.randn(1, 3, 56, 56))
+        self._test(Maximum(), inputs)