generalize pooling

fastmachinelearning · Dec 21, 2023 · f41951c · f41951c
1 parent 07c5bb6
commit f41951c
Showing 1 changed file with 27 additions and 3 deletions.
diff --git a/hls4ml/templates/vivado/nnet_utils/nnet_pooling.h b/hls4ml/templates/vivado/nnet_utils/nnet_pooling.h
@@ -27,15 +27,39 @@ template <int W, int N> ap_int<W> avg(ap_int<W> (&x)[N]) {
     return tmp;
 }
 
-template <int W, int I, int N> ap_fixed<W, I> avg(ap_fixed<W, I> (&x)[N]) {
+template <int W, int N> ap_int<W> avg(ap_uint<W> (&x)[N]) {
     // Use a wider accumulator than the input to avoid overflow
-    ap_fixed<W + ceillog2(N), I + ceillog2(N)> tmp = 0;
+    ap_uint<W + ceillog2(N)> tmp = 0;
     for (int i = 0; i < N; i++) {
         tmp += x[i];
     }
     tmp /= N;
     // Now cast back to original type
-    ap_fixed<W, I> y = tmp;
+    ap_uint<W> y = tmp;
+    return tmp;
+}
+
+template <int W, int I, int N, ap_q_mode Q, ap_o_mode O> ap_fixed<W, I, Q, O> avg(ap_fixed<W, I, Q, O> (&x)[N]) {
+    // Use a wider accumulator than the input to avoid overflow
+    ap_fixed<W + ceillog2(N), I + ceillog2(N), Q, O> tmp = 0;
+    for (int i = 0; i < N; i++) {
+        tmp += x[i];
+    }
+    tmp /= N;
+    // Now cast back to original type
+    ap_fixed<W, I, Q, O> y = tmp;
+    return y;
+}
+
+template <int W, int I, int N, ap_q_mode Q, ap_o_mode O> ap_ufixed<W, I, Q, O> avg(ap_ufixed<W, I, Q, O> (&x)[N]) {
+    // Use a wider accumulator than the input to avoid overflow
+    ap_ufixed<W + ceillog2(N), I + ceillog2(N), Q, O> tmp = 0;
+    for (int i = 0; i < N; i++) {
+        tmp += x[i];
+    }
+    tmp /= N;
+    // Now cast back to original type
+    ap_ufixed<W, I, Q, O> y = tmp;
     return y;
 }