SIRNN: A Math Library for Secure RNN Inference

Code release.

Athos/CompileSampleNetworks.py

@@ -82,8 +82,14 @@ def generate_code(params, debug=False):
         "SqueezeNetCIFAR10",
     ], "Network must be any of ResNet/DenseNet/SqueezeNetImgNet/SqueezeNetCIFAR10"
     scale = 12 if params["scale"] is None else params["scale"]
-    bitlength = 64 if params["bitlength"] is None else params["bitlength"]
     target = params["target"]
+    if params["bitlength"] is None:
+        if target == "SCI":
+            bitlength = 41
+        else:
+            bitlength = 64
+    else:
+        bitlength = params["bitlength"]
     disable_all_hlil_opts = (
         False
         if params["disable_all_hlil_opts"] is None
@@ -262,29 +268,43 @@ def generate_code(params, debug=False):
                 "Not compiling generated code. Please follow the readme and build Porthos."
             )
     elif target == "SCI":
-        sci = os.path.join(athos_dir, "..", "SCI")
-        sci_src = os.path.join(sci, "src")
-        sci_lib = os.path.join(sci, "build", "lib")
-        eigen_path = os.path.join(sci, "extern", "eigen")
-        seal_lib_path = os.path.join(sci, "extern", "SEAL", "native", "lib")
-        if os.path.exists(sci_lib):
+        sci_install = os.path.join(athos_dir, "..", "SCI", "build", "install")
+        build_dir = "build_dir"
+        os.system("rm -rf {build_dir}".format(build_dir=build_dir))
+        os.mkdir(build_dir)
+        os.chdir(build_dir)
+        cmake_file = """
+            cmake_minimum_required (VERSION 3.0)
+            project (BUILD_IT)
+            find_package(SCI REQUIRED PATHS \"{sci_install}\")
+            add_executable({prog_name} {src_file})
+            target_link_libraries({prog_name} SCI::SCI-{backend})
+        """.format(
+            sci_install=sci_install,
+            prog_name=program_name,
+            src_file=output_file,
+            backend=backend.upper(),
+        )
+        with open("CMakeLists.txt", "w") as f:
+            f.write(cmake_file)
+
+        if os.path.exists(sci_install):
+            ret = os.system("cmake --log-level=ERROR .")
+            if ret != 0:
+                sys.exit("Compilation of generated code failed. Exiting...")
+            ret = os.system("cmake --build . --parallel")
+            if ret != 0:
+                sys.exit("Compilation of generated code failed. Exiting...")
             os.system(
-                """g++ {opt_flag} -fpermissive -pthread -w -maes -msse4.1 -mavx -mavx2 -mrdseed \
-        -faligned-new -std=c++17 -fopenmp -I \"{eigen}\" -I \"{sci_src}\" \"{file}\" \
-        -L \"{sci_lib}\" -lSCI-LinearHE -L \"{seal}\" -lseal -lssl -lcrypto \
-        -o \"{output}\"""".format(
-                    eigen=eigen_path,
-                    sci_src=sci_src,
-                    file=output_file,
-                    sci_lib=sci_lib,
-                    seal=seal_lib_path,
-                    output=program_path,
-                    opt_flag=opt_flag,
+                "mv {tmp_prog} {prog_path}".format(
+                    tmp_prog=program_name, prog_path=program_path
                 )
             )
+            os.chdir("..")
+            os.system("rm -r {build_dir}".format(build_dir=build_dir))
         else:
             print(
-                "Not compiling generated code. Please follow the readme and build SCI before running this script."
+                "Not compiling generated code. Please follow the readme and build and install SCI."
             )
 
     os.chdir(cwd)

Athos/CompileTFGraph.py

@@ -101,7 +101,7 @@ def generate_code(params, role, debug=False):
     target = params["target"]
     if params["bitlength"] is None:
         if target == "SCI":
-            bitlength = 63
+            bitlength = 41
         else:
             bitlength = 64
     else:
@@ -281,31 +281,43 @@ def generate_code(params, role, debug=False):
                 "Not compiling generated code. Please follow the readme and build Porthos."
             )
     elif target == "SCI":
-        sci = os.path.join(athos_dir, "..", "SCI")
-        sci_src = os.path.join(sci, "src")
-        sci_lib = os.path.join(sci, "build", "lib")
-        eigen_path = os.path.join(sci, "extern", "eigen")
-        seal_lib_path = os.path.join(sci, "extern", "SEAL", "native", "lib")
-        seal_inc_path = os.path.join(sci, "extern", "SEAL", "native", "src")
-        if os.path.exists(sci_lib):
+        sci_install = os.path.join(athos_dir, "..", "SCI", "build", "install")
+        build_dir = "build_dir"
+        os.system("rm -r {build_dir}".format(build_dir=build_dir))
+        os.mkdir(build_dir)
+        os.chdir(build_dir)
+        cmake_file = """
+            cmake_minimum_required (VERSION 3.0)
+            project (BUILD_IT)
+            find_package(SCI REQUIRED PATHS \"{sci_install}\")
+            add_executable({prog_name} {src_file})
+            target_link_libraries({prog_name} SCI::SCI-{backend})
+        """.format(
+            sci_install=sci_install,
+            prog_name=program_name,
+            src_file=output_file,
+            backend=backend.upper(),
+        )
+        with open("CMakeLists.txt", "w") as f:
+            f.write(cmake_file)
+
+        if os.path.exists(sci_install):
+            ret = os.system("cmake --log-level=ERROR .")
+            if ret != 0:
+                sys.exit("Compilation of generated code failed. Exiting...")
+            ret = os.system("cmake --build . --parallel")
+            if ret != 0:
+                sys.exit("Compilation of generated code failed. Exiting...")
             os.system(
-                """g++ {opt_flag} -fpermissive -pthread -w -maes -msse4.1 -mavx -mavx2 -mrdseed \
-        -faligned-new -std=c++17 -fopenmp -I \"{eigen}\" -I \"{seal_inc_path}\" -I \"{sci_src}\" \"{file}\" \
-        -L \"{sci_lib}\" -lSCI-LinearHE -L \"{seal}\" -lseal -lssl -lcrypto \
-        -o \"{output}\"""".format(
-                    eigen=eigen_path,
-                    sci_src=sci_src,
-                    file=output_file,
-                    sci_lib=sci_lib,
-                    seal_inc_path=seal_inc_path,
-                    seal=seal_lib_path,
-                    output=program_path,
-                    opt_flag=opt_flag,
+                "mv {tmp_prog} {prog_path}".format(
+                    tmp_prog=program_name, prog_path=program_path
                 )
             )
+            os.chdir("..")
+            os.system("rm -r {build_dir}".format(build_dir=build_dir))
         else:
             print(
-                "Not compiling generated code. Please follow the readme and build SCI."
+                "Not compiling generated code. Please follow the readme and build and install SCI."
             )
 
     os.chdir(cwd)

Athos/CompilerScripts/compile_tf.py

@@ -121,14 +121,34 @@ def get_unsupported_ops(graph):
     return unsupported_ops
 
 
+def get_op_names_from_tensors(tensor_names):
+    op_names = []
+    for name in tensor_names:
+        if ":" in name:
+            try:
+                op_name, out_n = name.split(":")
+                out_n = int(out_n)
+            except:
+                raise ValueError(
+                    "The tensor name {} looks like a tensor name but is not a valid one".format(
+                        name
+                    )
+                )
+            op_names.append(op_name)
+        else:
+            op_names.append(name)
+    return op_names
+
+
 # Generates the computation graph and tensor size metadata and saves them in
 # the model directory.
 # Optionaly dumps model weights as fixedpt in specified scaling factor
 def compile(model_fname, input_t_info, output_t_names, scaling_factor, save_weights):
     model_name = os.path.basename(model_fname)[:-3]
     print("Loading tf graph ", model_fname)
     graph = tf_graph_io.load_pb(model_fname)
-    assert tensors_exist(graph, output_t_names)
+    output_op_names = get_op_names_from_tensors(output_t_names)
+    assert tensors_exist(graph, output_op_names)
 
     unsupported_ops = get_unsupported_ops(graph)
     if len(unsupported_ops) != 0:
@@ -145,7 +165,7 @@ def compile(model_fname, input_t_info, output_t_names, scaling_factor, save_weig
         tensors_exist(graph, list(input_t_info.keys()))
         graph = set_input_shapes(graph, input_t_info)
     input_t_names = list(input_t_info.keys())
-    graph_def = grappler.optimize(graph, input_t_names, output_t_names)
+    graph_def = grappler.optimize(graph, input_t_names, output_op_names)
     graph_def = grappler.convert_consts_to_var(graph_def)
     graph = get_graph_from(graph_def)
 
@@ -170,7 +190,7 @@ def compile(model_fname, input_t_info, output_t_names, scaling_factor, save_weig
             # these constants. We strip them away in a new graph def which is amenable
             # to codegen but leave them in the graph.
             optimized_graph_def = DumpTFMtData.strip_variable_init_constants(
-                graph_def, input_t_names, output_t_names
+                graph_def, input_t_names, output_op_names
             )
 
             tf_graph_io.dump_graph_def_pb(

Athos/CompilerScripts/sample_networks/run_demo_2pc.sh

@@ -46,8 +46,8 @@ TIME_CMD="/usr/bin/time --format \"%M %U %S %e\""
 PARTY0_IP="127.0.0.1"
 PORT=12345
 
-PARTY0_RUN_CMD="${MODEL_BINARY_PATH} r=1 p=${PORT} < ${MODEL_WEIGHT_PATH}"
-PARTY1_RUN_CMD="${MODEL_BINARY_PATH} r=2 ip=${PARTY0_IP} p=${PORT} < ${MODEL_INPUT_PATH}"
+PARTY0_RUN_CMD="${MODEL_BINARY_PATH} r=1 port=${PORT} < ${MODEL_WEIGHT_PATH}"
+PARTY1_RUN_CMD="${MODEL_BINARY_PATH} r=2 ip=${PARTY0_IP} port=${PORT} < ${MODEL_INPUT_PATH}"
 
 PARTY0_DUMP_CMD="> ${MODEL_DIR}/party0_mpc_output.out 2> ${MODEL_DIR}/party0_stats"
 PARTY1_DUMP_CMD="> ${MODEL_DIR}/party1_mpc_output.out 2> ${MODEL_DIR}/party1_stats"

Athos/Networks/ChestXRay/setup_and_run.sh

@@ -49,4 +49,4 @@ echo "This will take some time"
 echo "--------------------------------------------------------------------------------"
 echo -e "\n\n"
 python3 ChestXRay_tf_main.py --runPrediction True --scalingFac $scale --saveImgAndWtData True
-echo -e "\n\n"
+echo -e "\n\n"

Athos/Networks/DenseNet/setup_and_run.sh

@@ -40,12 +40,12 @@ if [[  ! -f "PreTrainedModel/tf-densenet121.ckpt.data-00000-of-00001"  ||  ! -f
   tar -xvzf ${filename}
   cd -
 fi
-exit
+#exit
 echo -e "\n\n"
 echo "--------------------------------------------------------------------------------"
 echo "Running DenseNet network and dumping computation graph, inputs and model weights"
 echo "This will take some time"
 echo "--------------------------------------------------------------------------------"
 echo -e "\n\n"
 python3 DenseNet_main.py --runPrediction True --scalingFac $scale --saveImgAndWtData True
-echo -e "\n\n"
+echo -e "\n\n"

Athos/Networks/ResNet/setup_and_run.sh

@@ -37,12 +37,12 @@ if [[  ! -d "PreTrainedModel/resnet_v2_fp32_savedmodel_NHWC" ]]; then
   tar -xvzf ${filename}
   cd -
 fi
-exit 
+#exit
 echo -e "\n\n"
 echo "--------------------------------------------------------------------------------"
 echo "Running ResNet network and dumping computation graph, inputs and model weights"
 echo "This will take some time"
 echo "--------------------------------------------------------------------------------"
 echo -e "\n\n"
 python3 ResNet_main.py --runPrediction True --scalingFac $scale --saveImgAndWtData True
-echo -e "\n\n"
+echo -e "\n\n"

Athos/SeeDot/Compiler.py

@@ -131,6 +131,9 @@ def fixOuputScale(self, res: (IR.Prog, IR.Expr), compiler: IRBuilderCSF):
                     argsDict[IR.Int(scale_down, 32)] = "consSF"
                     funcCall = IR.FuncCall(funcName, argsDict)
                     scaledown_cmd_list.append(funcCall)
+                # ArgMax sets scale to -1
+                if output_scale == -1:
+                    continue
                 if output_scale < Util.Config.consSF:
                     assert (
                         False