| """Code generator for Code Completion Model Inference. |
| |
| Tool runs on the Decision Forest model defined in {model} directory. |
| It generates two files: {output_dir}/{filename}.h and {output_dir}/{filename}.cpp |
| The generated files defines the Example class named {cpp_class} having all the features as class members. |
| The generated runtime provides an `Evaluate` function which can be used to score a code completion candidate. |
| """ |
| |
| import argparse |
| import json |
| import struct |
| |
| |
| class CppClass: |
| """Holds class name and names of the enclosing namespaces.""" |
| |
| def __init__(self, cpp_class): |
| ns_and_class = cpp_class.split("::") |
| self.ns = [ns for ns in ns_and_class[0:-1] if len(ns) > 0] |
| self.name = ns_and_class[-1] |
| if len(self.name) == 0: |
| raise ValueError("Empty class name.") |
| |
| def ns_begin(self): |
| """Returns snippet for opening namespace declarations.""" |
| open_ns = ["namespace %s {" % ns for ns in self.ns] |
| return "\n".join(open_ns) |
| |
| def ns_end(self): |
| """Returns snippet for closing namespace declarations.""" |
| close_ns = [ |
| "} // namespace %s" % ns for ns in reversed(self.ns)] |
| return "\n".join(close_ns) |
| |
| |
| def header_guard(filename): |
| '''Returns the header guard for the generated header.''' |
| return "GENERATED_DECISION_FOREST_MODEL_%s_H" % filename.upper() |
| |
| |
| def boost_node(n, label, next_label): |
| """Returns code snippet for a leaf/boost node.""" |
| return "%s: return %sf;" % (label, n['score']) |
| |
| |
| def if_greater_node(n, label, next_label): |
| """Returns code snippet for a if_greater node. |
| Jumps to true_label if the Example feature (NUMBER) is greater than the threshold. |
| Comparing integers is much faster than comparing floats. Assuming floating points |
| are represented as IEEE 754, it order-encodes the floats to integers before comparing them. |
| Control falls through if condition is evaluated to false.""" |
| threshold = n["threshold"] |
| return "%s: if (E.get%s() >= %s /*%s*/) goto %s;" % ( |
| label, n['feature'], order_encode(threshold), threshold, next_label) |
| |
| |
| def if_member_node(n, label, next_label): |
| """Returns code snippet for a if_member node. |
| Jumps to true_label if the Example feature (ENUM) is present in the set of enum values |
| described in the node. |
| Control falls through if condition is evaluated to false.""" |
| members = '|'.join([ |
| "BIT(%s_type::%s)" % (n['feature'], member) |
| for member in n["set"] |
| ]) |
| return "%s: if (E.get%s() & (%s)) goto %s;" % ( |
| label, n['feature'], members, next_label) |
| |
| |
| def node(n, label, next_label): |
| """Returns code snippet for the node.""" |
| return { |
| 'boost': boost_node, |
| 'if_greater': if_greater_node, |
| 'if_member': if_member_node, |
| }[n['operation']](n, label, next_label) |
| |
| |
| def tree(t, tree_num, node_num): |
| """Returns code for inferencing a Decision Tree. |
| Also returns the size of the decision tree. |
| |
| A tree starts with its label `t{tree#}`. |
| A node of the tree starts with label `t{tree#}_n{node#}`. |
| |
| The tree contains two types of node: Conditional node and Leaf node. |
| - Conditional node evaluates a condition. If true, it jumps to the true node/child. |
| Code is generated using pre-order traversal of the tree considering |
| false node as the first child. Therefore the false node is always the |
| immediately next label. |
| - Leaf node adds the value to the score and jumps to the next tree. |
| """ |
| label = "t%d_n%d" % (tree_num, node_num) |
| code = [] |
| |
| if t["operation"] == "boost": |
| code.append(node(t, label=label, next_label="t%d" % (tree_num + 1))) |
| return code, 1 |
| |
| false_code, false_size = tree( |
| t['else'], tree_num=tree_num, node_num=node_num+1) |
| |
| true_node_num = node_num+false_size+1 |
| true_label = "t%d_n%d" % (tree_num, true_node_num) |
| |
| true_code, true_size = tree( |
| t['then'], tree_num=tree_num, node_num=true_node_num) |
| |
| code.append(node(t, label=label, next_label=true_label)) |
| |
| return code+false_code+true_code, 1+false_size+true_size |
| |
| |
| def gen_header_code(features_json, cpp_class, filename): |
| """Returns code for header declaring the inference runtime. |
| |
| Declares the Example class named {cpp_class} inside relevant namespaces. |
| The Example class contains all the features as class members. This |
| class can be used to represent a code completion candidate. |
| Provides `float Evaluate()` function which can be used to score the Example. |
| """ |
| setters = [] |
| getters = [] |
| for f in features_json: |
| feature = f["name"] |
| |
| if f["kind"] == "NUMBER": |
| # Floats are order-encoded to integers for faster comparison. |
| setters.append( |
| "void set%s(float V) { %s = OrderEncode(V); }" % ( |
| feature, feature)) |
| elif f["kind"] == "ENUM": |
| setters.append( |
| "void set%s(unsigned V) { %s = 1LL << V; }" % (feature, feature)) |
| else: |
| raise ValueError("Unhandled feature type.", f["kind"]) |
| |
| # Class members represent all the features of the Example. |
| class_members = [ |
| "uint%d_t %s = 0;" |
| % (64 if f["kind"] == "ENUM" else 32, f['name']) |
| for f in features_json |
| ] |
| getters = [ |
| "LLVM_ATTRIBUTE_ALWAYS_INLINE uint%d_t get%s() const { return %s; }" |
| % (64 if f["kind"] == "ENUM" else 32, f['name'], f['name']) |
| for f in features_json |
| ] |
| nline = "\n " |
| guard = header_guard(filename) |
| return """#ifndef %s |
| #define %s |
| #include <cstdint> |
| #include "llvm/Support/Compiler.h" |
| |
| %s |
| class %s { |
| public: |
| // Setters. |
| %s |
| |
| // Getters. |
| %s |
| |
| private: |
| %s |
| |
| // Produces an integer that sorts in the same order as F. |
| // That is: a < b <==> orderEncode(a) < orderEncode(b). |
| static uint32_t OrderEncode(float F); |
| }; |
| |
| float Evaluate(const %s&); |
| %s |
| #endif // %s |
| """ % (guard, guard, cpp_class.ns_begin(), cpp_class.name, |
| nline.join(setters), |
| nline.join(getters), |
| nline.join(class_members), |
| cpp_class.name, cpp_class.ns_end(), guard) |
| |
| |
| def order_encode(v): |
| i = struct.unpack('<I', struct.pack('<f', v))[0] |
| TopBit = 1 << 31 |
| # IEEE 754 floats compare like sign-magnitude integers. |
| if (i & TopBit): # Negative float |
| return (1 << 32) - i # low half of integers, order reversed. |
| return TopBit + i # top half of integers |
| |
| |
| def evaluate_func(forest_json, cpp_class): |
| """Generates evaluation functions for each tree and combines them in |
| `float Evaluate(const {Example}&)` function. This function can be |
| used to score an Example.""" |
| |
| code = "" |
| |
| # Generate evaluation function of each tree. |
| code += "namespace {\n" |
| tree_num = 0 |
| for tree_json in forest_json: |
| code += "LLVM_ATTRIBUTE_NOINLINE float EvaluateTree%d(const %s& E) {\n" % (tree_num, cpp_class.name) |
| code += " " + \ |
| "\n ".join( |
| tree(tree_json, tree_num=tree_num, node_num=0)[0]) + "\n" |
| code += "}\n\n" |
| tree_num += 1 |
| code += "} // namespace\n\n" |
| |
| # Combine the scores of all trees in the final function. |
| # MSAN will timeout if these functions are inlined. |
| code += "float Evaluate(const %s& E) {\n" % cpp_class.name |
| code += " float Score = 0;\n" |
| for tree_num in range(len(forest_json)): |
| code += " Score += EvaluateTree%d(E);\n" % tree_num |
| code += " return Score;\n" |
| code += "}\n" |
| |
| return code |
| |
| |
| def gen_cpp_code(forest_json, features_json, filename, cpp_class): |
| """Generates code for the .cpp file.""" |
| # Headers |
| # Required by OrderEncode(float F). |
| angled_include = [ |
| '#include <%s>' % h |
| for h in ["cstring", "limits"] |
| ] |
| |
| # Include generated header. |
| qouted_headers = {filename + '.h', 'llvm/ADT/bit.h'} |
| # Headers required by ENUM features used by the model. |
| qouted_headers |= {f["header"] |
| for f in features_json if f["kind"] == "ENUM"} |
| quoted_include = ['#include "%s"' % h for h in sorted(qouted_headers)] |
| |
| # using-decl for ENUM features. |
| using_decls = "\n".join("using %s_type = %s;" % ( |
| feature['name'], feature['type']) |
| for feature in features_json |
| if feature["kind"] == "ENUM") |
| nl = "\n" |
| return """%s |
| |
| %s |
| |
| #define BIT(X) (1LL << X) |
| |
| %s |
| |
| %s |
| |
| uint32_t %s::OrderEncode(float F) { |
| static_assert(std::numeric_limits<float>::is_iec559, ""); |
| constexpr uint32_t TopBit = ~(~uint32_t{0} >> 1); |
| |
| // Get the bits of the float. Endianness is the same as for integers. |
| uint32_t U = llvm::bit_cast<uint32_t>(F); |
| std::memcpy(&U, &F, sizeof(U)); |
| // IEEE 754 floats compare like sign-magnitude integers. |
| if (U & TopBit) // Negative float. |
| return 0 - U; // Map onto the low half of integers, order reversed. |
| return U + TopBit; // Positive floats map onto the high half of integers. |
| } |
| |
| %s |
| %s |
| """ % (nl.join(angled_include), nl.join(quoted_include), cpp_class.ns_begin(), |
| using_decls, cpp_class.name, evaluate_func(forest_json, cpp_class), |
| cpp_class.ns_end()) |
| |
| |
| def main(): |
| parser = argparse.ArgumentParser('DecisionForestCodegen') |
| parser.add_argument('--filename', help='output file name.') |
| parser.add_argument('--output_dir', help='output directory.') |
| parser.add_argument('--model', help='path to model directory.') |
| parser.add_argument( |
| '--cpp_class', |
| help='The name of the class (which may be a namespace-qualified) created in generated header.' |
| ) |
| ns = parser.parse_args() |
| |
| output_dir = ns.output_dir |
| filename = ns.filename |
| header_file = "%s/%s.h" % (output_dir, filename) |
| cpp_file = "%s/%s.cpp" % (output_dir, filename) |
| cpp_class = CppClass(cpp_class=ns.cpp_class) |
| |
| model_file = "%s/forest.json" % ns.model |
| features_file = "%s/features.json" % ns.model |
| |
| with open(features_file) as f: |
| features_json = json.load(f) |
| |
| with open(model_file) as m: |
| forest_json = json.load(m) |
| |
| with open(cpp_file, 'w+t') as output_cc: |
| output_cc.write( |
| gen_cpp_code(forest_json=forest_json, |
| features_json=features_json, |
| filename=filename, |
| cpp_class=cpp_class)) |
| |
| with open(header_file, 'w+t') as output_h: |
| output_h.write(gen_header_code( |
| features_json=features_json, |
| cpp_class=cpp_class, |
| filename=filename)) |
| |
| |
| if __name__ == '__main__': |
| main() |
