Skip to content
Snippets Groups Projects
hpvm_codegen.py 4.54 KiB
class HpvmCodeGen:
    def __init__(self, DFG, weights_dir, test_data=None, test_labels=None):
        self.program_str = ""
        self.graph = DFG.graph
        self.tensors = DFG.tensors
        self.nodes = DFG.nodes
        self.var_cnt = 0
        self.weights_dir = weights_dir
        self.test_data = test_data
        self.test_labels = test_labels
        self.filter_names = {}  # essentially tensors

    def emit_header(self):
        headers = "\n#include <stdio.h> \n"
        headers += "#include <stdlib.h> \n"
        headers += "#include <unistd.h> \n"
        headers += "#include <fcntl.h> \n"
        headers += "#include <sys/stat.h> \n"
        headers += "#include <cstring> \n"
        headers += "#include <visc.h> \n"
        headers += "#include <tensorTypes.h> \n"
        headers += "#include <tensorUtils.h> \n\n"
        self.program_str += headers

    def emit_root(self):
        def emitRootNodeHeader():
            root_signature = "void root("
            index = 0
            for f_name in self.filter_names:
                if index > 0:
                    root_signature += "\t  "
                self.filter_names[f_name] = index
                root_signature += "void* " + f_name + ", "
                root_signature += "size_t " + f_name + "_bytes"
                if index < len(self.filter_names) - 1:
                    root_signature += ", \n"
                index += 1
            root_signature += "){ \n\n"
            root_signature += "\n  __visc__hint(visc::CPU_TARGET); \n"
            root_signature += "  __visc__attributes(" + \
                str(len(self.filter_names)) + ", "
            index = 0
            for f_name in self.filter_names:
                root_signature += f_name
                if index < len(self.filter_names) - 1:
                    root_signature += ", "
                    index += 1
            root_signature += ", 0); \n\n"
            return root_signature

        def emitRootNodeFooter(self):
            last_node = self.dfg.last_node
            output_var = self.output_map[last_node.layer_name]
            # Binding output of last DFG node to the Root Node output
            root_footer_str = "\n  __visc__bindOut(" + \
                output_var + ", 0, 0, 0); \n"
            root_footer_str += "  __visc__bindOut(" + \
                output_var + ", 1, 1, 0); \n"
            root_footer_str += "\n}\n\n"
            return root_footer_str

        def emitRootStructure(self):
            root_struct = ""
            root_struct += "struct ret_t {\n"
            root_struct += "  void* tensor; \n"
            root_struct += "  size_t bytes; \n"
            root_struct += "}; \n\n"
            root_struct += "typedef struct __attribute__((__packed__)) {\n"
            for f_name in self.filter_names:
                root_struct += "  void* " + f_name + "; \n"
                root_struct += "  size_t " + f_name + "_bytes; \n"
            root_struct += "\n  struct ret_t r; \n"
            root_struct += "}\nRootIn;\n\n"
            return root_struct

        self.program_str += emitRootNodeHeader()
        self.program_str += emitRootStructure()
        # self.codegen(self.dfg)
        self.program_str += emitRootNodeFooter()

    def emit_main(self, test_data):
        main_func_str = "int main(){ \n\n"
        #main_func_str += self.weight_str
        #main_func_str += self.input_str
        main_func_str += "\n__visc__init(); \n"
        main_func_str += "RootIn* args = static_cast<RootIn*>(malloc(sizeof(RootIn))); \n\n"
        for f_name in self.filter_names:
            main_func_str += "args->" + f_name + " = " + f_name + "; \n"
            main_func_str += "args->" + f_name + "_bytes = 0; \n"
        main_func_str += "\nvoid* dfg = __visc__launch(0, root, (void*) args); \n\n"
        main_func_str += "__visc__wait(dfg); \n\n"
        main_func_str += "void *result = static_cast<RootIn*>(args)->input; \n"
        main_func_str += "hpvm_request_tensor(result, 0); \n\n"
        main_func_str += "__visc__cleanup(); \n "
        main_func_str += "computeAccuracy3(labels, result); \n"
        main_func_str += "return 0; \n\n"
        main_func_str += "} \n"
        self.program_str += main_func_str

    def emit_source(self, dir_prefix):
        print(self.program_str)
        f = open(dir_prefix + "/approxhpvm_src.cc", "w+")
        f.write(self.program_str)
        f.close()

    def compile(self):
        self.emit_header()
        # self.emitRoot()
        self.emit_main(self.test_data)
        # dump generated program string to source file
        self.emit_source(self.weights_dir)