Fix Argmax io_stream implementation

Author: vloncar

hls4ml/templates/quartus/firmware/nnet_utils/nnet_activation_stream.h

@@ -406,22 +406,22 @@ void softmax_legacy(stream<data_T> &data, stream<res_T> &res) {
 }
 
 template<class data_T, class res_T, typename CONFIG_T>
-void softmax_argmax(stream<data_T> &data, stream<res_T> &res) {    
+void softmax_argmax(stream<data_T> &data, stream<res_T> &res) {
     #pragma ii 1
     for (int i = 0; i < CONFIG_T::n_in / res_T::size; i++) {
         data_T in_data = data.read();
         res_T out_data;
 
         #pragma unroll
-        for (int i = 0; i < CONFIG_T::n_in; i++) {
+        for (int i = 0; i < res_T::size; i++) {
             out_data[i] = (typename res_T::value_type) 0;
         }
 
         hls_register typename data_T::value_type maximum = in_data[0];
         hls_register int idx = 0; 
 
         #pragma ii 1
-        for (int i = 1; i < CONFIG_T::n_in; i++) {
+        for (int i = 1; i < res_T::size; i++) {
             if (in_data[i] > maximum) {
                 maximum = in_data[i];
                 idx = i;

hls4ml/templates/vivado/nnet_utils/nnet_activation_stream.h

@@ -315,21 +315,21 @@ void softmax_legacy(hls::stream<data_T> &data, hls::stream<res_T> &res) {
 }
 
 template<class data_T, class res_T, typename CONFIG_T>
-void softmax_argmax(hls::stream<data_T> &data, hls::stream<res_T> &res) {    
+void softmax_argmax(hls::stream<data_T> &data, hls::stream<res_T> &res) {
     for (int i = 0; i < CONFIG_T::n_in / res_T::size; i++) {
         #pragma HLS PIPELINE
         data_T in_data = data.read();
         res_T out_data;
 
-        for (int i = 0; i < CONFIG_T::n_in; i++) {
+        for (int i = 0; i < res_T::size; i++) {
             #pragma HLS UNROLL
             out_data[i] = (typename res_T::value_type) 0;
         }
 
         typename data_T::value_type maximum = in_data[0];
         int idx = 0; 
 
-        for (int i = 1; i < CONFIG_T::n_in; i++) {
+        for (int i = 1; i < res_T::size; i++) {
             #pragma HLS PIPELINE
             if (in_data[i] > maximum) {
                 maximum = in_data[i];
@@ -669,6 +669,49 @@ void prelu(hls::stream<data_T> &data, typename data_T::value_type alpha[CONFIG_T
     }
 }
 
+// *************************************************
+//       Binary TanH Activation
+// *************************************************
+template<class data_T, class res_T, typename CONFIG_T>
+void binary_tanh(hls::stream<data_T> &data, hls::stream<res_T> &res) {
+    PReLUActLoop: for (int i = 0; i < CONFIG_T::n_in / res_T::size; i++) {
+        #pragma HLS PIPELINE
+
+        data_T in_data = data.read();
+        res_T out_data;
+        #pragma HLS DATA_PACK variable=out_data
+
+        PReLUPackLoop: for (int j = 0; j < res_T::size; j++) {
+            #pragma HLS UNROLL
+            if(in_data[j] > 0) out_data[j] = (typename res_T::value_type) 1;
+            else out_data[j] = (typename res_T::value_type) -1;
+        }
+        res.write(out_data);
+    }
+}
+
+// *************************************************
+//       Ternary TanH Activation
+// *************************************************
+template<class data_T, class res_T, typename CONFIG_T>
+void ternary_tanh(hls::stream<data_T> &data, hls::stream<res_T> &res) {
+    PReLUActLoop: for (int i = 0; i < CONFIG_T::n_in / res_T::size; i++) {
+        #pragma HLS PIPELINE
+
+        data_T in_data = data.read();
+        res_T out_data;
+        #pragma HLS DATA_PACK variable=out_data
+
+        PReLUPackLoop: for (int j = 0; j < res_T::size; j++) {
+            #pragma HLS UNROLL
+            if(in_data[j] > 1) out_data[j] = (typename res_T::value_type) 1;
+            else if (in_data[j] <=-1) out_data[j] = (typename res_T::value_type) -1;
+            else out_data[j] = (typename res_T::value_type) 0;
+        }
+        res.write(out_data);
+    }
+}
+
 
 }