Skip to content

Commit d5d60b6

Browse files
mostafaelhoushimostafaelhoushi
authored
Merge pull request #570 from mostafaelhoushi/add-features-to-loop-opt-env
Add Observations to Loop Optimizations Environment
2 parents f8a5117 + 7d42155 commit d5d60b6

File tree

14 files changed

+300
-110
lines changed

14 files changed

+300
-110
lines changed

‎compiler_gym/service/proto/__init__.py

-1
Original file line numberDiff line numberDiff line change
@@ -109,7 +109,6 @@
109109
"GetVersionReply",
110110
"GetVersionRequest",
111111
"Int64Box",
112-
"Int64List",
113112
"Int64Range",
114113
"Int64SequenceSpace",
115114
"Int64Tensor",

‎compiler_gym/third_party/autophase/CMakeLists.txt

+1
Original file line numberDiff line numberDiff line change
@@ -51,4 +51,5 @@ cg_cc_binary(
5151
${LLVM_INCLUDE_DIRS}
5252
DEFINES
5353
${LLVM_DEFINITIONS}
54+
PUBLIC
5455
)

‎compiler_gym/third_party/programl/CMakeLists.txt

+30
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,30 @@
1+
# Copyright (c) Facebook, Inc. and its affiliates.
2+
#
3+
# This source code is licensed under the MIT license found in the
4+
# LICENSE file in the root directory of this source tree.
5+
6+
cg_add_all_subdirs()
7+
8+
llvm_map_components_to_libnames(_LLVM_LIBS analysis core irreader support)
9+
cg_cc_binary(
10+
NAME
11+
compute_programl
12+
SRCS
13+
"compute_programl.cc"
14+
COPTS
15+
"-DGOOGLE_PROTOBUF_NO_RTTI"
16+
"-fno-rtti"
17+
DEPS
18+
nlohmann_json::nlohmann_json
19+
ProGraML::graph::format::node_link_graph
20+
ProGraML::ir::llvm::llvm-10
21+
ProGraML::proto::programl_cc
22+
ABS_DEPS
23+
glog::glog
24+
${_LLVM_LIBS}
25+
INCLUDES
26+
${LLVM_INCLUDE_DIRS}
27+
DEFINES
28+
${LLVM_DEFINITIONS}
29+
PUBLIC
30+
)

‎compiler_gym/third_party/programl/compute_programl.cc

+55
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,55 @@
1+
// Copyright (c) Facebook, Inc. and its affiliates.
2+
//
3+
// This source code is licensed under the MIT license found in the
4+
// LICENSE file in the root directory of this source tree.
5+
#include <glog/logging.h>
6+
7+
#include <iostream>
8+
9+
#include "llvm/IR/Module.h"
10+
#include "llvm/IRReader/IRReader.h"
11+
#include "llvm/Support/ErrorOr.h"
12+
#include "llvm/Support/SourceMgr.h"
13+
#include "nlohmann/json.hpp"
14+
#include "programl/graph/format/node_link_graph.h"
15+
#include "programl/ir/llvm/llvm.h"
16+
17+
using nlohmann::json;
18+
19+
const programl::ProgramGraphOptions programlOptions;
20+
21+
int main(int argc, char** argv) {
22+
google::InitGoogleLogging(argv[0]);
23+
24+
CHECK(argc == 2) << "Usage: compute_programl <bitcode-path>";
25+
26+
auto buf = llvm::MemoryBuffer::getFileOrSTDIN(argv[1]);
27+
if (!buf) {
28+
LOG(FATAL) << "File not found: " << argv[1];
29+
}
30+
31+
llvm::SMDiagnostic error;
32+
llvm::LLVMContext ctx;
33+
34+
auto module = llvm::parseIRFile(argv[1], error, ctx);
35+
CHECK(module) << "Failed to parse: " << argv[1] << ": " << error.getMessage().str();
36+
37+
// Build the ProGraML graph.
38+
programl::ProgramGraph graph;
39+
auto status = programl::ir::llvm::BuildProgramGraph(*module, &graph, programlOptions);
40+
if (!status.ok()) {
41+
return -1;
42+
}
43+
44+
// Serialize the graph to a JSON node link graph.
45+
json nodeLinkGraph;
46+
status = programl::graph::format::ProgramGraphToNodeLinkGraph(graph, &nodeLinkGraph);
47+
if (!status.ok()) {
48+
return -1;
49+
}
50+
51+
nodeLinkGraph.dump();
52+
std::cout << nodeLinkGraph << "\n";
53+
54+
return 0;
55+
}

‎examples/example_compiler_gym_service/README.md

+1-1
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,7 @@ issue](https://github.com/facebookresearch/CompilerGym/issues/new/choose).
2222
never changes. Actions never end the episode.
2323
* There are two observation spaces:
2424
* `ir` which returns the string "Hello, world!".
25-
* `features` which returns an `int64_list` of `[0, 0, 0]`.
25+
* `features` which returns an `int64_tensor` of `[0, 0, 0]`.
2626
* A single reward space `runtime` which returns 0.
2727
* It has a single dataset "benchmark://example-v0" with two programs "foo" and
2828
"bar".

‎examples/example_unrolling_service/example_without_bazel.py

+9-9
Original file line numberDiff line numberDiff line change
@@ -14,13 +14,13 @@
1414
1515
It is equivalent in behavior to the example.py script in this directory.
1616
"""
17-
import os
1817
import subprocess
1918
from pathlib import Path
2019
from typing import Iterable
2120

2221
import compiler_gym
2322
from compiler_gym.datasets import Benchmark, Dataset
23+
from compiler_gym.datasets.uri import BenchmarkUri
2424
from compiler_gym.envs.llvm.llvm_benchmark import get_system_library_flags
2525
from compiler_gym.spaces import Reward
2626
from compiler_gym.third_party import llvm
@@ -99,13 +99,13 @@ def __init__(self, *args, **kwargs):
9999
)
100100

101101
self._benchmarks = {
102-
"benchmark://unrolling-v0/offsets1": Benchmark.from_file_contents(
102+
"/offsets1": Benchmark.from_file_contents(
103103
"benchmark://unrolling-v0/offsets1",
104-
self.preprocess(os.path.join(BENCHMARKS_PATH, "offsets1.c")),
104+
self.preprocess(BENCHMARKS_PATH / "offsets1.c"),
105105
),
106-
"benchmark://unrolling-v0/conv2d": Benchmark.from_file_contents(
106+
"/conv2d": Benchmark.from_file_contents(
107107
"benchmark://unrolling-v0/conv2d",
108-
self.preprocess(os.path.join(BENCHMARKS_PATH, "conv2d.c")),
108+
self.preprocess(BENCHMARKS_PATH / "conv2d.c"),
109109
),
110110
}
111111

@@ -131,11 +131,11 @@ def preprocess(src: Path) -> bytes:
131131
)
132132

133133
def benchmark_uris(self) -> Iterable[str]:
134-
yield from self._benchmarks.keys()
134+
yield from (f"benchmark://unrolling-v0{k}" for k in self._benchmarks.keys())
135135

136-
def benchmark(self, uri: str) -> Benchmark:
137-
if uri in self._benchmarks:
138-
return self._benchmarks[uri]
136+
def benchmark_from_parsed_uri(self, uri: BenchmarkUri) -> Benchmark:
137+
if uri.path in self._benchmarks:
138+
return self._benchmarks[uri.path]
139139
else:
140140
raise LookupError("Unknown program name")
141141

‎examples/loop_optimizations_service/env_tests.py

+72-10
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,8 @@
1414
import examples.loop_optimizations_service as loop_optimizations_service
1515
from compiler_gym.envs import CompilerEnv
1616
from compiler_gym.service import SessionNotFound
17-
from compiler_gym.spaces import Box, NamedDiscrete, Scalar, Sequence
17+
from compiler_gym.spaces import Dict, NamedDiscrete, Scalar, Sequence
18+
from compiler_gym.third_party.autophase import AUTOPHASE_FEATURE_NAMES
1819
from tests.test_main import main
1920

2021

@@ -83,14 +84,41 @@ def test_action_space(env: CompilerEnv):
8384
def test_observation_spaces(env: CompilerEnv):
8485
"""Test that the environment reports the service's observation spaces."""
8586
env.reset()
86-
assert env.observation.spaces.keys() == {"ir", "features", "runtime", "size"}
87+
assert env.observation.spaces.keys() == {
88+
"ir",
89+
"Inst2vec",
90+
"Autophase",
91+
"AutophaseDict",
92+
"Programl",
93+
"runtime",
94+
"size",
95+
}
8796
assert env.observation.spaces["ir"].space == Sequence(
8897
name="ir",
8998
size_range=(0, np.iinfo(int).max),
9099
dtype=str,
91100
)
92-
assert env.observation.spaces["features"].space == Box(
93-
name="features", shape=(3,), low=0, high=1e5, dtype=int
101+
assert env.observation.spaces["Inst2vec"].space == Sequence(
102+
name="Inst2vec",
103+
size_range=(0, np.iinfo(int).max),
104+
dtype=int,
105+
)
106+
assert env.observation.spaces["Autophase"].space == Sequence(
107+
name="Autophase",
108+
size_range=(len(AUTOPHASE_FEATURE_NAMES), len(AUTOPHASE_FEATURE_NAMES)),
109+
dtype=int,
110+
)
111+
assert env.observation.spaces["AutophaseDict"].space == Dict(
112+
name="AutophaseDict",
113+
spaces={
114+
name: Scalar(name="", min=0, max=np.iinfo(np.int64).max, dtype=np.int64)
115+
for name in AUTOPHASE_FEATURE_NAMES
116+
},
117+
)
118+
assert env.observation.spaces["Programl"].space == Sequence(
119+
name="Programl",
120+
size_range=(0, np.iinfo(int).max),
121+
dtype=str,
94122
)
95123
assert env.observation.spaces["runtime"].space == Scalar(
96124
name="runtime", min=0, max=np.inf, dtype=float
@@ -160,7 +188,7 @@ def test_Step_out_of_range(env: CompilerEnv):
160188

161189

162190
def test_default_ir_observation(env: CompilerEnv):
163-
"""Test default observation space."""
191+
"""Test default IR observation space."""
164192
env.observation_space = "ir"
165193
observation = env.reset()
166194
assert len(observation) > 0
@@ -171,16 +199,48 @@ def test_default_ir_observation(env: CompilerEnv):
171199
assert reward is None
172200

173201

174-
def test_default_features_observation(env: CompilerEnv):
175-
"""Test default observation space."""
176-
env.observation_space = "features"
202+
def test_default_inst2vec_observation(env: CompilerEnv):
203+
"""Test default inst2vec observation space."""
204+
env.observation_space = "Inst2vec"
205+
observation = env.reset()
206+
assert isinstance(observation, np.ndarray)
207+
assert len(observation) >= 0
208+
assert observation.dtype == np.int64
209+
assert all(obs >= 0 for obs in observation.tolist())
210+
211+
212+
def test_default_autophase_observation(env: CompilerEnv):
213+
"""Test default autophase observation space."""
214+
env.observation_space = "Autophase"
177215
observation = env.reset()
178216
assert isinstance(observation, np.ndarray)
179-
assert observation.shape == (3,)
217+
assert observation.shape == (len(AUTOPHASE_FEATURE_NAMES),)
180218
assert observation.dtype == np.int64
181219
assert all(obs >= 0 for obs in observation.tolist())
182220

183221

222+
def test_default_autophase_dict_observation(env: CompilerEnv):
223+
"""Test default autophase dict observation space."""
224+
env.observation_space = "AutophaseDict"
225+
observation = env.reset()
226+
assert isinstance(observation, dict)
227+
assert observation.keys() == AUTOPHASE_FEATURE_NAMES
228+
assert len(observation.values()) == len(AUTOPHASE_FEATURE_NAMES)
229+
assert all(obs >= 0 for obs in observation.values())
230+
231+
232+
def test_default_programl_observation(env: CompilerEnv):
233+
"""Test default observation space."""
234+
env.observation_space = "Programl"
235+
observation = env.reset()
236+
assert len(observation) > 0
237+
238+
observation, reward, done, info = env.step(0)
239+
assert not done, info
240+
assert len(observation) > 0
241+
assert reward is None
242+
243+
184244
def test_default_reward(env: CompilerEnv):
185245
"""Test default reward space."""
186246
env.reward_space = "runtime"
@@ -195,7 +255,9 @@ def test_observations(env: CompilerEnv):
195255
"""Test observation spaces."""
196256
env.reset()
197257
assert len(env.observation["ir"]) > 0
198-
np.testing.assert_array_less([-1, -1, -1], env.observation["features"])
258+
assert all(env.observation["Inst2vec"] >= 0)
259+
assert all(env.observation["Autophase"] >= 0)
260+
assert len(env.observation["Programl"]) > 0
199261

200262

201263
def test_rewards(env: CompilerEnv):

‎examples/loop_optimizations_service/example.py

+1-1
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@
44
with compiler_gym.make(
55
"loops-opt-py-v0",
66
benchmark="loops-opt-v0/add",
7-
observation_space="ir",
7+
observation_space="AutophaseDict",
88
reward_space="runtime",
99
) as env:
1010
compiler_gym.set_debug_level(4) # TODO: check why this has no effect

‎examples/loop_optimizations_service/example_without_bazel.py

+8-7
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@
2020

2121
import compiler_gym
2222
from compiler_gym.datasets import Benchmark, Dataset
23+
from compiler_gym.datasets.uri import BenchmarkUri
2324
from compiler_gym.envs.llvm.llvm_benchmark import get_system_library_flags
2425
from compiler_gym.spaces import Reward
2526
from compiler_gym.third_party import llvm
@@ -97,15 +98,15 @@ def __init__(self, *args, **kwargs):
9798
)
9899

99100
self._benchmarks = {
100-
"benchmark://loops-opt-v0/add": Benchmark.from_file_contents(
101+
"/add": Benchmark.from_file_contents(
101102
"benchmark://loops-opt-v0/add",
102103
self.preprocess(BENCHMARKS_PATH / "add.c"),
103104
),
104-
"benchmark://loops-opt-v0/offsets1": Benchmark.from_file_contents(
105+
"/offsets1": Benchmark.from_file_contents(
105106
"benchmark://loops-opt-v0/offsets1",
106107
self.preprocess(BENCHMARKS_PATH / "offsets1.c"),
107108
),
108-
"benchmark://loops-opt-v0/conv2d": Benchmark.from_file_contents(
109+
"/conv2d": Benchmark.from_file_contents(
109110
"benchmark://loops-opt-v0/conv2d",
110111
self.preprocess(BENCHMARKS_PATH / "conv2d.c"),
111112
),
@@ -132,11 +133,11 @@ def preprocess(src: Path) -> bytes:
132133
)
133134

134135
def benchmark_uris(self) -> Iterable[str]:
135-
yield from self._benchmarks.keys()
136+
yield from (f"benchmark://loops-opt-v0{k}" for k in self._benchmarks.keys())
136137

137-
def benchmark(self, uri: str) -> Benchmark:
138-
if uri in self._benchmarks:
139-
return self._benchmarks[uri]
138+
def benchmark_from_parsed_uri(self, uri: BenchmarkUri) -> Benchmark:
139+
if uri.path in self._benchmarks:
140+
return self._benchmarks[uri.path]
140141
else:
141142
raise LookupError("Unknown program name")
142143

‎examples/loop_optimizations_service/opt_loops/CMakeLists.txt

+7
Original file line numberDiff line numberDiff line change
@@ -24,5 +24,12 @@ cg_cc_binary(
2424
${LLVM_DEFINITIONS}
2525
)
2626

27+
cg_genrule(
28+
NAME opt_loops_bin
29+
OUTS "opt_loops"
30+
SRCS "opt_loops"
31+
COMMAND "cp ${CMAKE_BINARY_DIR}/examples/loop_optimizations_service/opt_loops/opt_loops $(@D)"
32+
)
33+
2734
ADD_CUSTOM_TARGET(link_opt_loops_target ALL
2835
COMMAND ${CMAKE_COMMAND} -E create_symlink ${CMAKE_BINARY_DIR}/examples/loop_optimizations_service/opt_loops/opt_loops ${CMAKE_SOURCE_DIR}/examples/loop_optimizations_service/opt_loops/opt_loops)

‎examples/loop_optimizations_service/service_py/BUILD

+5-7
Original file line numberDiff line numberDiff line change
@@ -8,21 +8,19 @@ py_binary(
88
name = "loops-opt-service-py",
99
srcs = ["loops_opt_service.py"],
1010
data = [
11+
"//compiler_gym/third_party/autophase:compute_autophase-prelinked",
12+
"//compiler_gym/third_party/programl:compute_programl",
1113
"//examples/loop_optimizations_service/opt_loops",
1214
],
1315
main = "loops_opt_service.py",
1416
visibility = ["//visibility:public"],
1517
deps = [
16-
":utils",
1718
"//compiler_gym/service",
1819
"//compiler_gym/service/proto",
1920
"//compiler_gym/service/runtime",
21+
"//compiler_gym/spaces",
22+
"//compiler_gym/third_party/autophase",
23+
"//compiler_gym/third_party/inst2vec",
2024
"//compiler_gym/third_party/llvm",
2125
],
2226
)
23-
24-
py_library(
25-
name = "utils",
26-
srcs = ["utils.py"],
27-
visibility = ["//visibility:public"],
28-
)

‎examples/loop_optimizations_service/service_py/CMakeLists.txt

+19
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,19 @@
1+
# Copyright (c) Facebook, Inc. and its affiliates.
2+
#
3+
# This source code is licensed under the MIT license found in the
4+
# LICENSE file in the root directory of this source tree.
5+
6+
cg_add_all_subdirs()
7+
8+
cg_py_library(
9+
NAME
10+
loops_opt_service
11+
SRCS
12+
"loops_opt_service.py"
13+
DATA
14+
examples::loop_optimizations_service::opt_loops::opt_loops
15+
compiler_gym::third_party::autophase:compute_autophase
16+
compiler_gym::third_party::programl:compute_programl
17+
DEPS
18+
PUBLIC
19+
)

0 commit comments

Comments
 (0)