Minor fixes around the Arm testing framework (#5976)

Summary: Pull Request resolved: #5976

Reviewed By: mergennachin

Differential Revision: D64047755

Pulled By: digantdesai

fbshipit-source-id: 91d32e771526eff2cb9863a04b0e18a2b40173e0

Author: benkli01

backends/arm/test/runner_utils.py

@@ -457,9 +457,10 @@ def prep_data_for_save(
     data_np = np.array(data.detach(), order="C").astype(np.float32)
 
     if is_quantized:
-        assert (
-            quant_param.node_name in input_name
-        ), "These quantization params do not match the input tensor name"
+        assert quant_param.node_name in input_name, (
+            f"The quantization params name '{quant_param.node_name}' does not "
+            f"match the input tensor name '{input_name}'."
+        )
         data_np = (
             ((data_np / np.float32(quant_param.scale)) + quant_param.zp)
             .round()

backends/arm/test/tester/arm_tester.py

@@ -150,6 +150,7 @@ def __init__(
         model: torch.nn.Module,
         example_inputs: Tuple[torch.Tensor],
         compile_spec: List[CompileSpec] = None,
+        tosa_ref_model_path: str | None = None,
     ):
         """
         Args:
@@ -160,7 +161,10 @@ def __init__(
 
         # Initiate runner_util
         intermediate_path = get_intermediate_path(compile_spec)
-        self.runner_util = RunnerUtil(intermediate_path=intermediate_path)
+        self.runner_util = RunnerUtil(
+            intermediate_path=intermediate_path,
+            tosa_ref_model_path=tosa_ref_model_path,
+        )
 
         self.compile_spec = compile_spec
         super().__init__(model, example_inputs)

backends/arm/tosa_mapping.py

@@ -15,37 +15,37 @@
 import torch
 
 
+UNSUPPORTED_DTYPES = (
+    torch.float64,
+    torch.double,
+    torch.complex64,
+    torch.cfloat,
+    torch.complex128,
+    torch.cdouble,
+    torch.uint8,
+    torch.int64,
+    torch.long,
+)
+
+DTYPE_MAP = {
+    torch.float32: ts.DType.FP32,
+    torch.float: ts.DType.FP32,
+    torch.float16: ts.DType.FP16,
+    torch.half: ts.DType.FP16,
+    torch.bfloat16: ts.DType.BF16,
+    torch.int8: ts.DType.INT8,
+    torch.int16: ts.DType.INT16,
+    torch.short: ts.DType.INT16,
+    torch.int32: ts.DType.INT32,
+    torch.int: ts.DType.INT32,
+    torch.bool: ts.DType.BOOL,
+}
+
+
 def map_dtype(data_type):
-    unsupported = (
-        torch.float64,
-        torch.double,
-        torch.complex64,
-        torch.cfloat,
-        torch.complex128,
-        torch.cdouble,
-        torch.uint8,
-        torch.int64,
-        torch.long,
-    )
-
-    dmap = {
-        torch.float32: ts.DType.FP32,
-        torch.float: ts.DType.FP32,
-        torch.float16: ts.DType.FP16,
-        torch.half: ts.DType.FP16,
-        torch.bfloat16: ts.DType.BF16,
-        torch.int8: ts.DType.INT8,
-        torch.int16: ts.DType.INT16,
-        torch.short: ts.DType.INT16,
-        torch.int32: ts.DType.INT32,
-        torch.int: ts.DType.INT32,
-        torch.bool: ts.DType.BOOL,
-    }
-
-    assert unsupported.count(data_type) == 0, "Unsupported type"
-    rtype = dmap.get(data_type)
-    assert rtype is not None, "Unknown type"
-    return rtype
+    assert data_type not in UNSUPPORTED_DTYPES, f"Unsupported type: {data_type}"
+    assert data_type in DTYPE_MAP, f"Unknown type: {data_type}"
+    return DTYPE_MAP[data_type]
 
 
 # Returns the shape and type of a node