Adding BatchNorm translation in HPVM frontend translator

llvm/projects/keras/frontend/hpvm_dfg_translator.py

@@ -187,7 +187,31 @@ class HPVMTranslator:
     return node_footer_str
 
 
+  # NOTE: genHpvmNodeEdges is replaced by genHpvmEdges 
+  def genHpvmNodeEdges2(self, hpvm_node_id, input_vars):
+
+    hpvm_edge_str = "\n  void* " + hpvm_node_id + " = "
+    hpvm_edge_str += "__visc__createNodeND(0, " + hpvm_node_id + "_node); \n\n"
+
+    it = 0
+    for input_var_name in input_vars:
+      if input_var_name in self.filter_names:
+        input_index = self.filter_names[input_var_name]
+        index1 = input_index * 2
+        index2 = index1 + 1      
+        hpvm_edge_str += "  __visc__bindIn(" + hpvm_node_id + ", " + str(index1) + ", " + str(it*2) + ", 0); \n"
+        hpvm_edge_str += "  __visc__bindIn(" + hpvm_node_id + ", " + str(index2) + ", " + str(it*2+1) + ", 0); \n"
+
+      elif input_var_name in self.hpvm_node_names:
+        hpvm_edge_str += "  __visc__edge(" + input_var_name + ", " + hpvm_node_id + ", 1, 0, " + str(it*2) + ", 0); \n"
+        hpvm_edge_str += "  __visc__edge(" + input_var_name + ", " + hpvm_node_id + ", 1, 1, " + str(it*2+1) + ", 0); \n"
+
+      it += 1
+      
+    return hpvm_edge_str
 
+  
+  # Fix: replace deprecated  genHpvmNodeEdges with  genHpvmEdges
   def genHpvmNodeEdges(self, out_var_name, input_var_name, input_var_name2):
 
     print ("input_var_name2 = ", input_var_name2)
@@ -331,6 +355,40 @@ class HPVMTranslator:
     
 
 
+  def genBatchNormNode(self, cur_node):
+    #input_var_name = self.getSingleInputName(cur_node)
+    out_var_name = self.getVariableName(cur_node)
+    
+    header_str = self.genNodeHeader(out_var_name, 5)
+    inst_str = header_str 
+
+    inst_str += "  void *r = __visc__tensor_batchnorm(t1, t2, t3, t4, t5, "
+    inst_str += str(cur_node.epsilon) + "); \n"
+
+    footer_str = self.genNodeFooter(5)
+    inst_str += footer_str
+        
+    self.node_str += inst_str
+
+    layer_name = cur_node.layer_name
+    input_id = self.getSingleInputName(cur_node)
+    gamma_id = layer_name + "_gamma"
+    beta_id = layer_name + "_beta"
+    mean_id = layer_name + "_mean"
+    variance_id = layer_name + "_variance"
+    
+    input_vars = []
+    input_vars.append(input_id)
+    input_vars.append(gamma_id)
+    input_vars.append(beta_id)
+    input_vars.append(mean_id)
+    input_vars.append(variance_id)
+
+    self.root_str +=  self.genHpvmNodeEdges2(out_var_name, input_vars)
+
+
+    
+
   def genBiasNode(self, cur_node):
     input_var_name = self.output_map[cur_node.layer_name]
     out_var_name = self.getVariableName(cur_node)    
@@ -431,7 +489,10 @@ class HPVMTranslator:
 
     if layer_type == "DepthwiseConv2D":
       self.genDepthwiseConvNode(cur_node)
-      
+
+    if layer_type == "BatchNormalization":
+      self.genBatchNormNode(cur_node)
+
     if layer_type == "Dense":
       self.genDenseNode(cur_node)