Merge pull request #332 from ChrisCummins/programl-dgl

[llvm] Add support for accessing ProGraML as JSON.
facebookresearch · Jul 26, 2021 · 19047db · 19047db
2 parents a534161 + 556532b
commit 19047db
Show file tree

Hide file tree

Showing 4 changed files with 38 additions and 2 deletions.
diff --git a/compiler_gym/envs/llvm/service/LlvmSession.cc b/compiler_gym/envs/llvm/service/LlvmSession.cc
@@ -331,7 +331,8 @@ Status LlvmSession::computeObservation(LlvmObservationSpace space, Observation&
       *reply.mutable_int64_list()->mutable_value() = {features.begin(), features.end()};
       break;
     }
-    case LlvmObservationSpace::PROGRAML: {
+    case LlvmObservationSpace::PROGRAML:
+    case LlvmObservationSpace::PROGRAML_JSON: {
       // Build the ProGraML graph.
       programl::ProgramGraph graph;
       auto status =

diff --git a/compiler_gym/envs/llvm/service/ObservationSpaces.cc b/compiler_gym/envs/llvm/service/ObservationSpaces.cc
@@ -102,6 +102,21 @@ std::vector<ObservationSpace> getLlvmObservationSpaceList() {
         *space.mutable_default_value()->mutable_string_value() = nodeLinkGraph.dump();
         break;
       }
+      case LlvmObservationSpace::PROGRAML_JSON: {
+        // ProGraML serializes the graph to JSON.
+        ScalarRange encodedSize;
+        encodedSize.mutable_min()->set_value(0);
+        space.set_opaque_data_format("json://");
+        *space.mutable_string_size_range() = encodedSize;
+        space.set_deterministic(true);
+        space.set_platform_dependent(false);
+        programl::ProgramGraph graph;
+        json nodeLinkGraph;
+        CHECK(programl::graph::format::ProgramGraphToNodeLinkGraph(graph, &nodeLinkGraph).ok())
+            << "Failed to serialize default ProGraML graph";
+        *space.mutable_default_value()->mutable_string_value() = nodeLinkGraph.dump();
+        break;
+      }
       case LlvmObservationSpace::CPU_INFO: {
         // Hardware info is returned as a JSON
         ScalarRange encodedSize;

diff --git a/compiler_gym/envs/llvm/service/ObservationSpaces.h b/compiler_gym/envs/llvm/service/ObservationSpaces.h
@@ -45,7 +45,7 @@ enum class LlvmObservationSpace {
    */
   AUTOPHASE,
   /**
-   * Returns the graph representation of a program.
+   * Returns the graph representation of a program as a networkx Graph.
    *
    * From:
    *
@@ -54,6 +54,16 @@ enum class LlvmObservationSpace {
    *     and Analysis. ArXiv:2003.10536. https://arxiv.org/abs/2003.10536
    */
   PROGRAML,
+  /**
+   * Returns the graph representation of a program as a JSON node-link graph.
+   *
+   * From:
+   *
+   *     Cummins, C., Fisches, Z. V., Ben-Nun, T., Hoefler, T., & Leather, H.
+   *     (2020). ProGraML: Graph-based Deep Learning for Program Optimization
+   *     and Analysis. ArXiv:2003.10536. https://arxiv.org/abs/2003.10536
+   */
+  PROGRAML_JSON,
   /** A JSON dictionary of properties describing the CPU. */
   CPU_INFO,
   /** The number of LLVM-IR instructions in the current module. */

diff --git a/tests/llvm/observation_spaces_test.py b/tests/llvm/observation_spaces_test.py
@@ -48,6 +48,7 @@ def test_observation_spaces(env: LlvmEnv):
         "Autophase",
         "AutophaseDict",
         "Programl",
+        "ProgramlJson",
         "CpuInfo",
         "Inst2vecPreprocessedText",
         "Inst2vecEmbeddingIndices",
@@ -386,6 +387,15 @@ def test_programl_observation_space(env: LlvmEnv):
     assert not space.platform_dependent
 
 
+def test_programl_json_observation_space(env: LlvmEnv):
+    env.reset("cbench-v1/crc32")
+    key = "ProgramlJson"
+    space = env.observation.spaces[key]
+    assert isinstance(space.space, Sequence)
+    graph: Dict[str, Any] = env.observation[key]
+    assert isinstance(graph, dict)
+
+
 def test_cpuinfo_observation_space(env: LlvmEnv):
     env.reset("cbench-v1/crc32")
     key = "CpuInfo"