Apply review comment

- split test into PowScalar and PowTensor
- raise RuntimeError when expected nodes aren't found.
- fix perf mode on aot_compiler.py.
-  added -> None to define_node.

Signed-off-by: jiseong.oh <jiseong.oh@samsung.com>

Author: Jiseong-oh

backends/samsung/_passes/compose_rms_norm.py

@@ -19,18 +19,38 @@ class RecomposeRmsNorm(ExportPass):
     def __init__(self):
         super().__init__()
 
+    _ADD_TARGETS = (
+        exir_ops.edge.aten.add.Tensor,
+        exir_ops.edge.aten.add.Scalar,
+        torch.ops.aten.add.Tensor,
+        torch.ops.aten.add.Scalar,
+    )
+
     def _get_eps_node(self, nodes):
         # eps: one of inputs of add node
-        add_node = [n for n in nodes if hasattr(n, "name") and "add" in n.name][0]
+        add_nodes = [
+            n
+            for n in nodes
+            if isinstance(n, torch.fx.Node)
+            and n.op == "call_function"
+            and n.target in self._ADD_TARGETS
+        ]
+        if not add_nodes:
+            raise RuntimeError("Failed to locate add node in RMSNorm partition")
+        add_node = add_nodes[0]
         for a in add_node.args:
-            if isinstance(a, float) or a.op != "call_function":
+            if isinstance(a, float) or (
+                isinstance(a, torch.fx.Node) and a.op != "call_function"
+            ):
                 return a
+        raise RuntimeError("Failed to locate eps argument in RMSNorm add node")
 
     def _get_gamma_node(self, output_node):
         # gamma: one of inputs of output node
         for a in output_node.args:
-            if a.op != "call_function":
+            if isinstance(a, torch.fx.Node) and a.op != "call_function":
                 return a
+        raise RuntimeError("Failed to locate gamma argument in RMSNorm output node")
 
     def call(self, graph_module: torch.fx.GraphModule):
         graph = graph_module.graph

backends/samsung/builders/op_split_with_sizes_copy.py

@@ -23,7 +23,7 @@ def define_node(
         node: torch.fx.Node,
         enn_graph: EnnGraph,
         vals_to_ids: Dict[torch.Tensor, int],
-    ):
+    ) -> None:
         input = node.args[0]
         input_id = self.define_tensor(input, enn_graph, vals_to_ids)
 

backends/samsung/builders/op_topk.py

@@ -67,8 +67,8 @@ def define_node(
                 raise AssertionError("Not supported largest = False.")
 
         if len(node.args) > 4:
-            sorted = cast(bool, node.args[4])
-            if not sorted:
+            is_sorted = cast(bool, node.args[4])
+            if not is_sorted:
                 raise AssertionError("Not supported sorted = False.")
 
         enn_graph.define_op(node.name, "TopK", [input_id], all_output_tensors, params)

backends/samsung/serialization/compile_options.py

@@ -12,6 +12,7 @@
 from enum import IntEnum, unique
 
 from importlib.resources import files
+from typing import Optional
 
 from executorch.exir._serialize._dataclass import _DataclassEncoder
 from executorch.exir._serialize._flatbuffer import _flatc_compile
@@ -73,7 +74,7 @@ def gen_samsung_backend_compile_spec_core(options: EnnExecuTorchOptions) -> Comp
 
 def gen_samsung_backend_compile_spec(
     chipset: str,
-    perf_mode: PerformanceMode = None,
+    perf_mode: Optional[PerformanceMode] = None,
 ):
     """
     A function to generate an ExecuTorch binary for Samsung Backend.

backends/samsung/test/ops/test_pow.py

@@ -16,31 +16,43 @@
 from executorch.backends.samsung.test.utils.utils import TestConfig
 
 
-class Pow(torch.nn.Module):
+class PowScalar(torch.nn.Module):
     def __init__(self) -> None:
         super().__init__()
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         return x**2.0
 
 
+class PowTensor(torch.nn.Module):
+    def __init__(self) -> None:
+        super().__init__()
+
+    def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
+        return torch.pow(x, y)
+
+
 class TestPow(unittest.TestCase):
-    def _test(self, module: torch.nn.Module, inputs):
+    def _test(self, module: torch.nn.Module, inputs, expected_op: str):
         tester = SamsungTester(
             module,
             inputs,
             [gen_samsung_backend_compile_spec(TestConfig.chipset)],
         )
         (
             tester.export()
-            .check_count({"torch.ops.aten.pow.Tensor_Scalar": 1})
+            .check_count({expected_op: 1})
             .to_edge_transform_and_lower()
             .check_not(["executorch_exir_dialects_edge__ops_aten_pow_default"])
             .check_count({"torch.ops.higher_order.executorch_call_delegate": 1})
             .to_executorch()
             .run_method_and_compare_outputs()
         )
 
-    def test_fp32_pow(self):
+    def test_fp32_pow_scalar(self):
         inputs = (torch.randn(1, 1, 16, 8),)
-        self._test(Pow(), inputs)
+        self._test(PowScalar(), inputs, "torch.ops.aten.pow.Tensor_Scalar")
+
+    def test_fp32_pow_tensor(self):
+        inputs = (torch.randn(1, 1, 16, 8), torch.randn(1, 1, 16, 8))
+        self._test(PowTensor(), inputs, "torch.ops.aten.pow.Tensor_Tensor")

examples/samsung/aot_compiler.py

@@ -55,6 +55,15 @@
         help=f"Model name. Valid ones: {SUPPORT_MODEL_NAMES}",
     )
     parser.add_argument("-o", "--output_dir", default=".", help="output directory")
+    parser.add_argument(
+        "--perf_mode",
+        default=PerformanceMode.DEFAULT.name,
+        choices=[m.name for m in PerformanceMode],
+        help=(
+            "Performance mode. HIGH_PERFORMANCE is experimental and must be "
+            "verified on the exynos device farm before deploying on a phone."
+        ),
+    )
 
     args = parser.parse_args()
 
@@ -78,7 +87,7 @@
     compile_specs = [
         gen_samsung_backend_compile_spec(
             args.chipset,
-            PerformanceMode.HIGH_PERFORMANCE,
+            PerformanceMode[args.perf_mode],
         )
     ]
     edge = to_edge_transform_and_lower_to_enn(