Improvements to compiler scripts.

Athos/CompilerScripts/comparison_scripts/compare_output.sh

-# Usage: tf_output.float(floatingpt) party0_output(fixedpt) SCALING_FACTOR PRECISION(upto how many points to compare?)
+# Usage: tf_output.float(floatingpt) party0_output(fixedpt) BITLEN SCALING_FACTOR PRECISION(upto how many points to compare?)
 # This first converts unsigned fixedpt to signed
 SCRIPT_DIR="$( cd "$(dirname "$0")" >/dev/null 2>&1 ; pwd -P )"
-echo "Comparing output with tensorflow output upto $4 decimal points."
-$SCRIPT_DIR/convert_to_signed.sh $2
+echo "Comparing output with tensorflow output upto $5 decimal points."
+$SCRIPT_DIR/convert_to_signed.sh $2 $3
 #Then runs the comparison script on it.
-python3 $SCRIPT_DIR/compare_output.py $1 $2_signed $3 $4
+python3 $SCRIPT_DIR/compare_output.py $1 $2_signed $4 $5
 if [ "$?" -eq 0 ]; then
-	echo "Output matches upto ${4} decimal points"
+	echo "Output matches upto ${5} decimal points"
 else
-	echo "Output does not match upto ${4} decimal points"
+	echo "Output does not match upto ${5} decimal points"
 fi

Athos/CompilerScripts/comparison_scripts/convert_to_signed.py

 import sys
 
 if __name__ == "__main__":
-        assert(len(sys.argv) == 3)
+        assert(len(sys.argv) == 4)
         inp_fname = sys.argv[1]
         out_fname = sys.argv[2]
+        bitlen = int(sys.argv[3])
         f = open(inp_fname, 'r')
         op = [(int(line.rstrip())) for line in f]
         f.close()
         f = open(out_fname, 'w')
         for i in op:
-            f.write(str( i if (i<2**63) else i - 2**64) + '\n')
+            f.write(str( i if (i<2**(bitlen-1)) else i - 2**bitlen) + '\n')
         f.close()

Athos/CompilerScripts/comparison_scripts/convert_to_signed.sh

 SCRIPT_DIR="$( cd "$(dirname "$0")" >/dev/null 2>&1 ; pwd -P )"
 inp1=$1
+bitlen=$2
+if [ -z "$bitlen" ]
+then
+	echo "Please pass bitlen."
+
+	exit 1
+fi
+
 temp_1=${inp1}_tmp_cmp
 
 awk '$0==($0+0)' $inp1 > $temp_1
 
-python3 ${SCRIPT_DIR}/convert_to_signed.py $temp_1 ${inp1}_signed
+python3 ${SCRIPT_DIR}/convert_to_signed.py $temp_1 ${inp1}_signed $bitlen
 
 rm $temp_1

Athos/CompilerScripts/compile_tf_graph.py

@@ -30,24 +30,45 @@ Use that with this script.""")
   print("Loading processed tf graph ", model_fname)
   graph = load_pb(model_fname)
 
-  if not check_operation_exists(graph, input_t_name):
-    sys.exit(input_t_name + " input does not exist in the graph")
   if not check_operation_exists(graph, output_t_name):
     sys.exit(output_t_name + " output does not exist in the graph")
-
-  input_t = graph.get_operation_by_name(input_t_name).outputs[0]
   output_t = graph.get_operation_by_name(output_t_name).outputs[0]
+  
+  if input_t_name != "":
+    if not check_operation_exists(graph, input_t_name):
+      sys.exit(input_t_name + " input does not exist in the graph")
+
+    input_t = graph.get_operation_by_name(input_t_name).outputs[0]
 
-  # Generate random tensor as input
-  inp_shape = input_t.shape.as_list()
-  if None in inp_shape:
-    if input_shape == []:
-      sys.exit("Please supply shape for the input tensor as it is parametric (? dim) for this model. See --help.")
+    # Generate random tensor as input
+    # scalar input
+    if input_t.shape.dims == None:
+      inp_shape = []
     else:
-      inp_shape = input_shape
-  rand_inp_t = np.zeros(inp_shape)
+      inp_shape = input_t.shape.as_list()
+      if None in inp_shape:
+        if input_shape == []:
+          sys.exit("Please supply shape for the input tensor as it is parametric (? dim) for this model. See --help.")
+        else:
+          inp_shape = input_shape
+    rand_inp_t = np.zeros(inp_shape)
+
+    feed_dict = {input_t: rand_inp_t}
+  else:
+    # We can collect all placeholder nodes as inputs to the model
+    inputs = [i for i in graph.get_operations() if i.type=="Placeholder"]
+    feed_dict = {}
+    for op in inputs:
+      input_t = op.outputs[0]
+      if input_t.shape.dims == None:
+        inp_shape = []
+      else:
+        inp_shape = input_t.shape.as_list()
+        if None in inp_shape:
+          sys.exit("Please supply input names and their shapes for the input tensor as it is parametric (? dim) for this model. See --help.")
+      rand_inp_t = np.zeros(inp_shape)
+      feed_dict[input_t] = rand_inp_t
 
-  feed_dict = {input_t: rand_inp_t}
   with graph.as_default():
     with tf.Session() as sess:
       # Run initializers generated by preprocessing
@@ -65,7 +86,7 @@ Use that with this script.""")
       if save_weights:
         DumpTFMtData.updateWeightsForBN(optimized_graph_def, sess)
         weights_fname = model_name[len("mpc_processed_"):] + '_input_weights_fixedpt_scale_' + str(scaling_factor) + '.inp'
-        print("Dumping model weights in ", weights_fname, ". These are to be used as input for party which owns the model")
+        print("\nDumping model weights in ", weights_fname, ". These are to be used as input for party which owns the model\n")
         DumpTFMtData.dumpTrainedWeightsInt(sess, trainVars, weights_fname, scaling_factor, 'w')
 
 def boolean_string(s):
@@ -76,7 +97,7 @@ def boolean_string(s):
 def parse_args():
   parser = argparse.ArgumentParser()
   parser.add_argument("--modelName", required=True, type=str, help="Name of processed tensorflow model (mpc_processed*.pb)")
-  parser.add_argument("--inputTensorName", required=True, type=str, help="Name of the input tensor for the model. (Op name, dont add '/:0' suffix)")
+  parser.add_argument("--inputTensorName", type=str, default='', help="Name of the input tensor for the model. (Op name, dont add '/:0' suffix)")
   parser.add_argument("--outputTensorName", required=True, type=str, help="Name of the input tensor for the model. (Op name, dont add '/:0' suffix)")
   parser.add_argument("--sf", default=12, type=int, help="scaling factor (int)")
   parser.add_argument("--saveWeights", type=boolean_string, default=False, help="Dump model weights in fixedpt {True/False}")

Athos/CompilerScripts/preprocess_frozen_tf_graph.py

@@ -3,6 +3,10 @@ from tf_graph_trans import *
 import sys
 import time
 import os 
+
+import argparse
+import os.path
+
 # Transpose nodes require perm as compile time constants for parametric codegen
 # So we don't eliminate the constants we need dring compile time
 def get_const_names(graph):
@@ -11,39 +15,66 @@ def get_const_names(graph):
   slice_begin_ops = set(i.inputs[1].op.name for i in graph.get_operations() if i.type == 'Slice')
   slice_size_ops = set(i.inputs[2].op.name for i in graph.get_operations() if i.type == 'Slice')
   mean_axes_ops = set(i.inputs[1].op.name for i in graph.get_operations() if i.type == 'Mean')
-  white_list = transp_perm_ops | padding_ops | slice_begin_ops | slice_size_ops | mean_axes_ops
+  split_dim_ops = set(i.inputs[0].op.name for i in graph.get_operations() if i.type == 'Split')
+  white_list = transp_perm_ops | padding_ops | slice_begin_ops | slice_size_ops | mean_axes_ops | split_dim_ops
   all_const_ops = set(i.name for i in graph.get_operations() if i.type == 'Const')
   return list(all_const_ops - white_list)
 
-if __name__ == "__main__":
-  if len(sys.argv) != 2:
-    print("Usage: python preprocess_frozen_tf_graph.py tf_model_name.pb")
-    sys.exit()
-  else:
-    input_fname = sys.argv[1]
+def check_operation_exists(graph, tensor_name):
+  op_list = [i.name for i in graph.get_operations()]
+  return tensor_name in op_list
 
+def optimize(input_fname, output_t_name):
+  if not input_fname.endswith('.pb'):
+    sys.exit("Please supply a valid tensorflow protobuf model (.pb extension)")
+  
   actual_fname = os.path.basename(input_fname) 
   dirname = os.path.dirname(input_fname) 
   output_fname = os.path.join(dirname, "mpc_processed_" + actual_fname)
   print("Loading ", input_fname, "for processing.")
-  exec_graph = load_pb(input_fname)
-  
+  graph = load_pb(input_fname)
+
+  if not check_operation_exists(graph, output_t_name):
+    sys.exit(output_t_name + " output does not exist in the graph")
+  input_names = [i.name for i in graph.get_operations() if i.type=="Placeholder"]
+
+  #graph = remove_dead_nodes(graph, input_names, [output_t_name]) 
+
   print("\n\nThis process will take some time to run as we execute portions of the graph.\n\n")
-  time.sleep(5)
+  time.sleep(1)
   # Fold away all static computations
+
+  print("Running fold splits")
+  graph = fold_splits(graph)
+  print(graph.get_operations(),end="\n\n")
   print("Running constant folding")
-  exec_graph = fold_splits(exec_graph)
-  exec_graph = fold_constants(exec_graph)
+  graph = fold_constants(graph)
+
   # Convert constants to variables so as to separate the data and the generated code
   # Otherwise huge arrays will show up as constants in the generated code, thereby
   # increasing binary size.
   print("Convert frozen constants to variables")
-  exec_graph = convert_consts_to_var(exec_graph, get_const_names(exec_graph))
+  graph = convert_consts_to_var(graph, get_const_names(graph))
+
   
+  input_names = [i.name for i in graph.get_operations() if i.type=="Placeholder"]
+  #graph = remove_dead_nodes(graph, input_names, [output_t_name]) 
+
   # At this stage the graph still has constants embedded in it
   # in the assign nodes for variables. We cannot execute the graph without
   # these constants. However after inferring the size, we can call remove_dead_nodes
   # to optimize away the constants and assign nodes and make the graph amenable
   # for codegen
-  dump_pb(exec_graph, output_fname)
+  dump_pb(graph, output_fname)
   print("The processed graph is dumped in ", output_fname)
+
+def parse_args():
+  parser = argparse.ArgumentParser()
+  parser.add_argument("--modelName", required=True, type=str, help="Name of tensorflow model (*.pb)")
+  parser.add_argument("--outputTensorName", required=True, type=str, help="Name of the output tensor for the model. (Op name, dont add '/:0' suffix)")
+  args = parser.parse_args()
+  return args
+
+if __name__ == '__main__':
+  args = parse_args()
+  optimize(args.modelName, args.outputTensorName)

Athos/CompilerScripts/tf_graph_trans.py

@@ -16,26 +16,26 @@ def delete_nodes(graph, ops):
     tf.import_graph_def(new_gd, name="")
   return new_graph
 
-def remove_dead_nodes(graph, input_tensors, output_tensors):
+def remove_dead_nodes(graph, in_list, out_list):
   transforms = ['remove_nodes(op=Identity)', 'strip_unused_nodes']
-  in_list = [i.name for i in input_tensors]
-  out_list = [i.name for i in output_tensors]
   optimized_graph_def = TransformGraph(graph.as_graph_def(), in_list, out_list, transforms)
   with tf.Graph().as_default() as opt_graph:
     tf.import_graph_def(optimized_graph_def, name="")
     return opt_graph
 
 def convert_consts_to_var(graph, const_names_list):
+
   const_var_names_pairs = []
   ops_to_delete = []
   with graph.as_default():
+    preexisting_vars = [tf.get_variable(i.name, i.outputs[0].shape) for i in graph.get_operations() if i.type=="VariableV2" or i.type=="Variable"]
+
     var_list = []
     for name in const_names_list:
-      #tensor = graph.get_tensor_by_name('{}:0'.format(name))
       tensor = graph.get_operation_by_name(name).outputs[0]
       with tf.Session() as sess:
         t_value = sess.run(tensor)
-      t_name = '{}_const_var'.format(name)
+      t_name = '{}_mpc_const_var'.format(name)
       var = tf.Variable(t_value, name=t_name)
       const_var_names_pairs.append((name, t_name))
       var_list.append(var)
@@ -45,7 +45,8 @@ def convert_consts_to_var(graph, const_names_list):
       var_op = graph.get_operation_by_name('{}/read'.format(var_name))
       ge.swap_outputs(ge.sgv(const_op), ge.sgv(var_op))
       ops_to_delete.append(const_op)
-    tf.compat.v1.variables_initializer(var_list, 'init_constvars')
+
+    tf.compat.v1.variables_initializer(var_list + preexisting_vars, 'init_constvars')
   return delete_nodes(graph, ops_to_delete)
 
 def get_inputs(op):