From 7ff28d4f56c951cf47f86dd553de0b157dc24944 Mon Sep 17 00:00:00 2001
From: Mostafa Elhoushi <melhoushi@fb.com>
Date: Sat, 26 Feb 2022 00:01:42 -0500
Subject: [PATCH] add AutoPaseDict observation

---
 .../loop_optimizations_service/example.py     |  2 +-
 .../service_py/loops_opt_service.py           | 36 +++++++++++++++++--
 2 files changed, 34 insertions(+), 4 deletions(-)

diff --git a/examples/loop_optimizations_service/example.py b/examples/loop_optimizations_service/example.py
index efc0a86646..e10a322444 100644
--- a/examples/loop_optimizations_service/example.py
+++ b/examples/loop_optimizations_service/example.py
@@ -4,7 +4,7 @@
 with compiler_gym.make(
     "loops-opt-py-v0",
     benchmark="loops-opt-v0/add",
-    observation_space="Autophase",
+    observation_space="AutophaseDict",
     reward_space="runtime",
 ) as env:
     compiler_gym.set_debug_level(4)  # TODO: check why this has no effect
diff --git a/examples/loop_optimizations_service/service_py/loops_opt_service.py b/examples/loop_optimizations_service/service_py/loops_opt_service.py
index 3d173f1157..248a24f0aa 100755
--- a/examples/loop_optimizations_service/service_py/loops_opt_service.py
+++ b/examples/loop_optimizations_service/service_py/loops_opt_service.py
@@ -19,6 +19,7 @@
 from compiler_gym.service.proto import (
     ActionSpace,
     Benchmark,
+    DictEvent,
     DictSpace,
     DoubleRange,
     Event,
@@ -85,6 +86,19 @@ class LoopsOptCompilationSession(CompilationSession):
         ),
         ObservationSpace(
             name="Autophase",
+            space=Space(
+                int64_sequence=Int64SequenceSpace(
+                    length_range=Int64Range(
+                        min=len(AUTOPHASE_FEATURE_NAMES),
+                        max=len(AUTOPHASE_FEATURE_NAMES),
+                    )
+                ),
+            ),
+            deterministic=True,
+            platform_dependent=False,
+        ),
+        ObservationSpace(
+            name="AutophaseDict",
             space=Space(
                 space_dict=DictSpace(
                     space={
@@ -256,7 +270,7 @@ def get_observation(self, observation_space: ObservationSpace) -> Event:
                 int64_tensor=Int64Tensor(shape=[len(Inst2vec_ids)], value=Inst2vec_ids)
             )
         elif observation_space.name == "Autophase":
-            Autophase_ids_str = run_command(
+            Autophase_str = run_command(
                 [
                     runfiles_path(
                         "compiler_gym/third_party/autophase/compute_autophase-prelinked"
@@ -265,12 +279,28 @@ def get_observation(self, observation_space: ObservationSpace) -> Event:
                 ],
                 timeout=30,
             )
-            Autophase_ids_int = list(map(int, list(Autophase_ids_str.split(" "))))
+            Autophase_list = list(map(int, list(Autophase_str.split(" "))))
             return Event(
                 int64_tensor=Int64Tensor(
-                    shape=[len(Autophase_ids_int)], value=Autophase_ids_int
+                    shape=[len(Autophase_list)], value=Autophase_list
                 )
             )
+        elif observation_space.name == "AutophaseDict":
+            Autophase_str = run_command(
+                [
+                    runfiles_path(
+                        "compiler_gym/third_party/autophase/compute_autophase-prelinked"
+                    ),
+                    self._llvm_path,
+                ],
+                timeout=30,
+            )
+            Autophase_list = list(map(int, list(Autophase_str.split(" "))))
+            Autophase_dict = {
+                name: Event(int64_value=val)
+                for name, val in zip(AUTOPHASE_FEATURE_NAMES, Autophase_list)
+            }
+            return Event(event_dict=DictEvent(event=Autophase_dict))
         elif observation_space.name == "Programl":
             Programl_str = run_command(
                 [