Added stacked traces class and GPU combine algorithms

4 files changed, 85 insertions, 47 deletions

diff --git a/pyecsca/sca/__init__.py b/pyecsca/sca/__init__.py
index 1aae9d9..5b359b8 100644
--- a/pyecsca/sca/__init__.py
+++ b/pyecsca/sca/__init__.py
@@ -5,3 +5,4 @@ from .scope import *
 from .target import *
 from .trace import *
 from .trace_set import *
+from .stacked_traces import *
diff --git a/pyecsca/sca/stacked_trace/__init__.py b/pyecsca/sca/stacked_trace/__init__.py
deleted file mode 100644
index 0e77c18..0000000
--- a/pyecsca/sca/stacked_trace/__init__.py
+++ /dev/null
@@ -1,3 +0,0 @@
-
-
-from .stacked_trace import *
\ No newline at end of file
diff --git a/pyecsca/sca/stacked_traces/__init__.py b/pyecsca/sca/stacked_traces/__init__.py
new file mode 100644
index 0000000..090c00c
--- /dev/null
+++ b/pyecsca/sca/stacked_traces/__init__.py
@@ -0,0 +1 @@
+from .stacked_traces import *
\ No newline at end of file
diff --git a/pyecsca/sca/stacked_trace/stacked_trace.py b/pyecsca/sca/stacked_traces/stacked_traces.py
index 9ab74b1..7af96e3 100644
--- a/pyecsca/sca/stacked_trace/stacked_trace.py
+++ b/pyecsca/sca/stacked_traces/stacked_traces.py
@@ -1,4 +1,3 @@
-from audioop import avg
 from numba import cuda, float32
 import numpy as np
 from public import public
@@ -13,15 +12,15 @@ class StackedTraces:
     """Samples of multiple traces and metadata"""
 
     meta: Mapping[str, Any]
-    traces: np.ndarray
+    samples: np.ndarray
 
     def __init__(
-                 self, traces: np.ndarray,
+                 self, samples: np.ndarray,
                  meta: Mapping[str, Any] = None) -> None:
         if meta is None:
             meta = dict()
         self.meta = meta
-        self.traces = traces
+        self.samples = samples
     
     @classmethod
     def fromarray(cls, traces: MutableSequence[np.ndarray],
@@ -47,10 +46,11 @@ class StackedTraces:
         yield from self.traces
 
 
+@public
 class GPUTraceManager:
     @staticmethod
     def average(traces: StackedTraces) -> np.ndarray:
-        samples = traces.traces
+        samples = traces.samples
         samples_global = cuda.to_device(samples)
         device_result = cuda.device_array(samples.shape[1])
 
@@ -61,11 +61,13 @@ class GPUTraceManager:
         res = device_result.copy_to_host()
         return res
     
+    @staticmethod
     def conditional_average(traces: StackedTraces) -> np.ndarray:
         raise NotImplementedError
     
+    @staticmethod
     def standard_deviation(traces: StackedTraces) -> np.ndarray:
-        samples = traces.traces
+        samples = traces.samples
         samples_global = cuda.to_device(samples)
         device_result = cuda.device_array(samples.shape[1])
 
@@ -77,12 +79,12 @@ class GPUTraceManager:
         return res
 
 
-@cuda.jit
-def gpu_average(samples: np.ndarray, result: np.ndarray):
-    col = cuda.grid(1)
+@cuda.jit(device=True)
+def _gpu_average(col: int, samples: np.ndarray, result: np.ndarray):
+    # col = cuda.grid(1)
 
-    if col >= samples.shape[1]:
-        return
+    # if col >= samples.shape[1]:
+    #     return
     
     acc = 0.
     for row in range(samples.shape[0]):
@@ -90,45 +92,80 @@ def gpu_average(samples: np.ndarray, result: np.ndarray):
     result[col] = acc / samples.shape[0]
 
 
-@cuda.jit()
-def gpu_std_dev(samples: np.ndarray, result: np.ndarray):
+@cuda.jit
+def gpu_average(samples: np.ndarray, result: np.ndarray):
     col = cuda.grid(1)
 
     if col >= samples.shape[1]:
         return
 
-    avg = 0.
-    for row in range(samples.shape[0]):
-        avg += samples[row, col]
-    avg /= samples.shape[0]
+    _gpu_average(col, samples, result)
 
+
+@cuda.jit(device=True)
+def _gpu_var_from_avg(col: int, samples: np.ndarray, averages: np.ndarray, result: np.ndarray):
     var = 0.
     for row in range(samples.shape[0]):
-        current = samples[row, col] - avg
+        current = samples[row, col] - averages[col]
         var += current * current
-    result[col] = sqrt(var / samples.shape[0])
+    result[col] = var / samples.shape[0]
 
 
-@cuda.jit()
-def gpu_variance(samples: np.ndarray, result: np.ndarray):
-    col = cuda.grid(1)
+@cuda.jit(device=True)
+def _gpu_variance(col: int, samples: np.ndarray, result: np.ndarray):
+    # col = cuda.grid(1)
     
+    # if col >= samples.shape[1]:
+    #     return
+
+    # avg = 0.
+    # for row in range(samples.shape[0]):
+    #     avg += samples[row, col]
+    # avg /= samples.shape[0]
+
+    _gpu_average(col, samples, result)
+    _gpu_var_from_avg(col, samples, result, result)
+    # var = 0.
+    # for row in range(samples.shape[0]):
+    #     current = samples[row, col] - result[col]
+    #     var += current * current
+    # result[col] = var / samples.shape[0]
+
+
+@cuda.jit
+def gpu_std_dev(samples: np.ndarray, result: np.ndarray):
+    col = cuda.grid(1)
+
     if col >= samples.shape[1]:
         return
 
-    avg = 0.
-    for row in range(samples.shape[0]):
-        avg += samples[row, col]
-    avg /= samples.shape[0]
+    # avg = 0.
+    # for row in range(samples.shape[0]):
+    #     avg += samples[row, col]
+    # avg /= samples.shape[0]
 
-    var = 0.
-    for row in range(samples.shape[0]):
-        current = samples[row, col] - avg
-        var += current * current
-    result[col] = var / samples.shape[0]
+    # var = 0.
+    # for row in range(samples.shape[0]):
+    #     current = samples[row, col] - result[col]
+    #     var += current * current
+    # result[col] = sqrt(var / samples.shape[0])
+
+    _gpu_variance(col, samples, result)
+
+    result[col] = sqrt(result[col])
+
+
+@cuda.jit
+def gpu_variance(samples: np.ndarray, result: np.ndarray):
+    col = cuda.grid(1)
 
+    if col >= samples.shape[1]:
+        return
+
+    _gpu_variance(col, samples, result)
 
-@cuda.jit()
+
+@cuda.jit
 def gpu_avg_var(samples: np.ndarray, result_avg: np.ndarray,
                 result_var: np.ndarray):
     col = cuda.grid(1)
@@ -136,20 +173,22 @@ def gpu_avg_var(samples: np.ndarray, result_avg: np.ndarray,
     if col >= samples.shape[1]:
         return
     
-    avg = 0.
-    for row in range(samples.shape[0]):
-        avg += samples[row, col]
-    avg /= samples.shape[0]
+    # avg = 0.
+    # for row in range(samples.shape[0]):
+    #     avg += samples[row, col]
+    # avg /= samples.shape[0]
 
-    var = 0.
-    for row in range(samples.shape[0]):
-        current = samples[row, col] - avg
-        var += current * current
-    result_avg[col] = avg
-    result_var[col] = var
+    _gpu_average(samples, result_avg)
+    _gpu_var_from_avg(col, samples, result_avg, result_var)
 
+    # var = 0.
+    # for row in range(samples.shape[0]):
+    #     current = samples[row, col] - result_avg[col]
+    #     var += current * current
+    # result_var[col] = var / samples.shape[0]
 
-@cuda.jit()
+
+@cuda.jit
 def gpu_add(samples: np.ndarray, result: np.ndarray):
     col = cuda.grid(1)
 
@@ -162,7 +201,7 @@ def gpu_add(samples: np.ndarray, result: np.ndarray):
     result[col] = res
 
 
-@cuda.jit()
+@cuda.jit
 def gpu_subtract(samples_one: np.ndarray, samples_other: np.ndarray,
                  result: np.ndarray):
     col = cuda.grid(1)