Fix 32bit scale (#5)

* Fix functional model saturating casts

* Add more intermediate result prints

* Fix global_shift calculation

* Add tests with 32bit scale

CHANGELOG.md

@@ -8,12 +8,14 @@
 - Support for kernels without normalization and quantization for NE16
 - isort check
 - publication citation
+- support 32bit scale
 
 ### Changed
 
 - `ne16_task_init` got split into smaller parts: `ne16_task_init`, `ne16_task_set_op_to_conv`, `ne16_task_set_weight_offset`, `ne16_task_set_bits`, `ne16_task_set_norm_quant`
 - strides in `ne16_task_set_strides`, `ne16_task_set_dims`, and `ne16_task_set_ptrs` are now strides between consecutive elements in that dimension
 - `ne16_task_queue_size` is now `NE16_TASK_QUEUE_SIZE`
+- `ne16_task_set_ptrs` split into `ne16_task_set_ptrs_conv` and `ne16_task_set_ptrs_norm_quant`
 
 ### Removed
 

ne16/README.md

@@ -28,7 +28,7 @@
 - [ ] Scale type
     - [x] uint8
     - [ ] uint16
-    - [ ] uint32
+    - [x] uint32
 - [x] Bias type
     - [x] int32
 - [ ] Weight type

neureka/README.md

@@ -16,17 +16,16 @@ Github repo [link](https://github.com/siracusa-soc/ne).
     - [x] Bias (w/ and w/o)
     - [ ] Per-channel shift
     - [x] Per-layer shift
-    - [ ] Rounding
 - [x] Input type
     - [x] uint8
     - [x] int8
 - [x] Output type
     - [x] int8
     - [x] uint8 (only w/ Relu)
     - [x] int32
-- [ ] Scale type
+- [x] Scale type
     - [x] uint8
-    - [ ] uint32
+    - [x] uint32
 - [x] Bias type
     - [x] int32
 - [ ] Weight type

test/NeuralEngineFunctionalModel.py

@@ -28,24 +28,34 @@ def _norm_quant(
         bias_type: Optional[IntegerType],
         has_bias: bool,
         has_relu: bool,
+        verbose: bool,
     ) -> torch.Tensor:
         # Scale accumulators are in 48bit, so keeping the data in 64bit
         tensor = tensor * scale
         assert tensor.dtype == torch.int64
 
+        if verbose:
+            print("INTERMEDIATE RESULTS (after scale):")
+            print(tensor)
+
         if has_bias:
             assert bias is not None
             assert bias_type is not None
-            # Saturating cast to int32
+
             tensor = NeuralEngineFunctionalModel._cast(
-                tensor, bias_type, saturate=True
+                tensor, bias_type, saturate=False
             ).type(torch.int32)
 
             tensor = tensor + bias
+
             tensor = NeuralEngineFunctionalModel._cast(
-                tensor, bias_type, saturate=False
+                tensor, bias_type, saturate=True
             ).type(torch.int32)
 
+            if verbose:
+                print("INTERMEDIATE RESULTS (after bias):")
+                print(tensor)
+
         if has_relu:
             tensor = F.relu(tensor)
 
@@ -118,6 +128,7 @@ def convolution(
                 bias_type,
                 has_bias,
                 has_relu,
+                verbose,
             )
 
         return output

test/NnxTestClasses.py

@@ -254,16 +254,22 @@ def from_conf(
                 ).type(torch.int32)
             if global_shift is None:
                 global_shift = torch.Tensor([0]).type(torch.int32)
+                conv_kwargs = {
+                    **conf.__dict__,
+                    "out_type": NeuralEngineFunctionalModel.ACCUMULATOR_TYPE,
+                }
                 output = NeuralEngineFunctionalModel().convolution(
                     input,
                     weight,
                     scale,
                     bias,
                     global_shift,
-                    verbose=verbose,
-                    **conf.__dict__,
+                    verbose=False,
+                    **conv_kwargs,
+                )
+                global_shift = NnxTestGenerator._calculate_global_shift(
+                    output, conf.out_type
                 )
-                NnxTestGenerator._calculate_global_shift(output, conf.out_type)
 
         output = NeuralEngineFunctionalModel().convolution(
             input, weight, scale, bias, global_shift, verbose=verbose, **conf.__dict__

test/tests/test_116/conf.json

@@ -0,0 +1,29 @@
+{
+    "in_height": 3,
+    "in_width": 3,
+    "in_channel": 2,
+    "out_channel": 2,
+    "padding": {
+        "top": 0,
+        "bottom": 0,
+        "left": 0,
+        "right": 0
+    },
+    "kernel_shape": {
+        "height": 1,
+        "width": 1
+    },
+    "depthwise": false,
+    "stride": {
+        "height": 1,
+        "width": 1
+    },
+    "in_type": "int8",
+    "out_type": "int8",
+    "weight_type": "int8",
+    "scale_type": "uint32",
+    "bias_type": "int32",
+    "has_norm_quant": true,
+    "has_bias": true,
+    "has_relu": false
+}

test/tests/test_117/conf.json

@@ -0,0 +1,29 @@
+{
+    "in_height": 10,
+    "in_width": 10,
+    "in_channel": 10,
+    "out_channel": 10,
+    "padding": {
+        "top": 0,
+        "bottom": 0,
+        "left": 0,
+        "right": 0
+    },
+    "kernel_shape": {
+        "height": 1,
+        "width": 1
+    },
+    "depthwise": false,
+    "stride": {
+        "height": 1,
+        "width": 1
+    },
+    "in_type": "uint8",
+    "out_type": "int8",
+    "weight_type": "int8",
+    "scale_type": "uint32",
+    "bias_type": "int32",
+    "has_norm_quant": true,
+    "has_bias": true,
+    "has_relu": false
+}

test/tests/test_118/conf.json

@@ -0,0 +1,29 @@
+{
+    "in_height": 10,
+    "in_width": 10,
+    "in_channel": 128,
+    "out_channel": 128,
+    "padding": {
+        "top": 0,
+        "bottom": 0,
+        "left": 0,
+        "right": 0
+    },
+    "kernel_shape": {
+        "height": 1,
+        "width": 1
+    },
+    "depthwise": false,
+    "stride": {
+        "height": 1,
+        "width": 1
+    },
+    "in_type": "uint8",
+    "out_type": "int8",
+    "weight_type": "int8",
+    "scale_type": "uint32",
+    "bias_type": "int32",
+    "has_norm_quant": true,
+    "has_bias": true,
+    "has_relu": false
+}