Merge pull request #105 from ChrisCummins/fix/ProGraML/57

ml4pl/llvm2graph: Add support for struct inlining

Author: ChrisCummins

deeplearning/ml4pl/graphs/llvm2graph/BUILD

@@ -146,6 +146,7 @@ py_binary(
         "//labm8/py:app",
         "//labm8/py:bazelutil",
         "//labm8/py:decorators",
+        "//labm8/py:fs",
         "//third_party/py/numpy",
     ],
 )

deeplearning/ml4pl/graphs/llvm2graph/node_encoder.py

@@ -17,7 +17,9 @@
 
 When executed as a binary, this program reads a single program graph from
 stdin, encodes it, and writes a graph to stdout. Use --stdin_fmt and
---stdout_fmt to convert between different graph types.
+--stdout_fmt to convert between different graph types, and --ir to read the
+IR file that the graph was constructed from, required for resolving struct
+definitions.
 
 Example usage:
 
@@ -26,10 +28,12 @@
     $ bazel run //deeplearning/ml4pl/graphs/llvm2graph:node_encoder -- \
         --stdin_fmt=pb \
         --stdout_fmt=pbtxt \
+        --ir=/tmp/source.ll \
         < /tmp/proto.pb > /tmp/proto.pbtxt
 """
 import pickle
 from typing import List
+from typing import Optional
 
 import networkx as nx
 import numpy as np
@@ -40,9 +44,17 @@
 from labm8.py import app
 from labm8.py import bazelutil
 from labm8.py import decorators
+from labm8.py import fs
 
 
 FLAGS = app.FLAGS
+app.DEFINE_output_path(
+  "ir",
+  None,
+  "The path of the IR file that was used to construct the graph. This is "
+  "required to inline struct definitions. This argument may be omitted when "
+  "struct definitions do not need to be inlined.",
+)
 
 DICTIONARY = bazelutil.DataPath(
   "phd/deeplearning/ml4pl/graphs/llvm2graph/node_embeddings/inst2vec_augmented_dictionary.pickle"
@@ -69,20 +81,35 @@ def __init__(self):
     with open(str(AUGMENTED_INST2VEC_EMBEDDINGS), "rb") as f:
       self.node_text_embeddings = pickle.load(f)
 
-  def EncodeNodes(self, g: nx.DiGraph) -> None:
+  def EncodeNodes(self, g: nx.DiGraph, ir: Optional[str] = None) -> None:
     """Pre-process the node text and set the text embedding index.
 
     For each node, this sets the 'preprocessed_text', 'x', and 'y' attributes.
 
     Args:
       g: The graph to encode the nodes of.
+      ir: The LLVM IR that was used to construct the graph. This is required for
+        struct inlining. If struct inlining is not required, this may be
+        omitted.
     """
     # Pre-process the statements of the graph in a single pass.
     lines = [
       [data["text"]]
       for _, data in g.nodes(data=True)
       if data["type"] == programl_pb2.Node.STATEMENT
     ]
+
+    if ir:
+      # NOTE(github.com/ChrisCummins/ProGraML/issues/57): Extract the struct
+      # definitions from the IR and inline their definitions in place of the
+      # struct names. This is brittle string substitutions, in the future we
+      # should do this inlining in llvm2graph where we have a parsed
+      # llvm::Module.
+      structs = inst2vec_preprocess.GetStructTypes(ir)
+      for line in lines:
+        for struct, definition in structs.items():
+          line[0] = line[0].replace(struct, definition)
+
     preprocessed_lines, _ = inst2vec_preprocess.preprocess(lines)
     preprocessed_texts = [
       inst2vec_preprocess.PreprocessStatement(x[0] if len(x) else "")
@@ -122,7 +149,8 @@ def Main():
   proto = programl.ReadStdin()
   g = programl.ProgramGraphToNetworkX(proto)
   encoder = GraphNodeEncoder()
-  encoder.EncodeNodes(g)
+  ir = fs.Read(FLAGS.ir) if FLAGS.ir else None
+  encoder.EncodeNodes(g, ir=ir)
   programl.WriteStdout(programl.NetworkXToProgramGraph(g))
 
 

deeplearning/ncc/inst2vec/inst2vec_preprocess.py

@@ -25,6 +25,7 @@
 import os
 import pickle
 import re
+from typing import Dict
 
 import networkx as nx
 
@@ -2959,6 +2960,21 @@ def inline_struct_types_in_file(G, dic, specific_struct_name_pattern):
   return G
 
 
+def GetStructTypes(ir: str) -> Dict[str, str]:
+  """Extract a dictionary of struct definitions from the given IR.
+
+  Args:
+    ir: A string of LLVM IR.
+
+  Returns:
+    A dictionary of <name, def> entries, where <name> is the name of a struct
+    definition (e.g. "%struct.foo"), and <def> is the definition of the member
+    types, e.g. "{ i32 }".
+  """
+  _, dict_temp = construct_struct_types_dictionary_for_file(ir.split("\n"))
+  return dict_temp
+
+
 def inline_struct_types(
   G, data_with_struct_def, file_name, specific_struct_name_pattern
 ):

tools/BUILD

@@ -2,6 +2,10 @@
 # Some of these scripts are meant for execution in-tree, and cannot be run
 # using bazel, please see the comments.
 
+exports_files([
+    "bazel",
+])
+
 sh_binary(
     name = "whoami",
     srcs = ["whoami.sh"],

tools/bazel

@@ -21,6 +21,52 @@
 # Your mileage may vary.
 set -eu
 
+build_path() {
+    local _dir=""
+    local _path=""
+
+    for _dir in "$@"
+    do
+        if [ -d $_dir ]; then
+            _path=$_path:$_dir
+        fi
+    done
+
+    _path=${_path:1}
+    echo $_path
+}
+
+# Accepts an array of directories and returns a colon separated path
+# of all of the directories that exist, in order.  Example usage:
+#
+#    dirs=("/usr/local/bin" /usr/bin "/not a real path")
+#    unset FOO
+#    FOO=$(build_path "${dirs[@]}")
+#    echo $FOO
+#    # Outputs: /usr/local/bin:/usr/bin
+path_dirs=( \
+    /usr/local/opt/llvm/bin \
+    /usr/local/opt/gnu-sed/libexec/gnubin \
+    /usr/bin \
+    /usr/local/bin \
+    /bin \
+)
+
+if [[ -f "/usr/local/opt/llvm/bin/clang" ]]; then
+  CC=/usr/local/opt/llvm/bin/clang
+  CXX=/usr/local/opt/llvm/bin/clang++
+else
+  CC="$(which gcc)"
+  CXX="$(which g++)"
+fi
+
 # PULLET_TIMEOUT to increase the timeout on docker image pulls from the default
 # 600s. See: https://github.com/bazelbuild/rules_docker
-env -i TERM="$TERM" PATH=/usr/local/opt/gnu-sed/libexec/gnubin:/usr/bin:/usr/local/bin:/bin PULLER_TIMEOUT=3600 "$BAZEL_REAL" "$@"
+set +u
+env -i \
+  TERM="$TERM" \
+  PATH="$(build_path ${path_dirs[@]})" \
+  CC=$CC \
+  CXX=$CXX \
+  PULLER_TIMEOUT=3600 \
+  "$BAZEL_REAL" "$@"