Loop opt service

examples/example_unrolling_service/__init__.py

@@ -2,7 +2,7 @@
 #
 # This source code is licensed under the MIT license found in the
 # LICENSE file in the root directory of this source tree.
-"""This module demonstrates how to """
+"""This module defines and registers the example gym environments."""
 import subprocess
 from pathlib import Path
 from typing import Iterable

examples/example_unrolling_service/example.py

@@ -10,32 +10,34 @@
 import compiler_gym
 import examples.example_unrolling_service as unrolling_service  # noqa Register environments.
 
-env = compiler_gym.make(
+with compiler_gym.make(
     "unrolling-py-v0",
     benchmark="unrolling-v0/offsets1",
     observation_space="features",
     reward_space="runtime",
-)
-compiler_gym.set_debug_level(4)  # TODO: check why this has no effect
+) as env:
+    compiler_gym.set_debug_level(4)  # TODO: check why this has no effect
 
-observation = env.reset()
-print("observation: ", observation)
+    observation = env.reset()
+    print("observation: ", observation)
 
-print()
+    print()
 
-observation, reward, done, info = env.step(env.action_space.sample())
-print("observation: ", observation)
-print("reward: ", reward)
-print("done: ", done)
-print("info: ", info)
+    observation, reward, done, info = env.step(env.action_space.sample())
+    print("observation: ", observation)
+    print("reward: ", reward)
+    print("done: ", done)
+    print("info: ", info)
 
-print()
+    print()
 
-observation, reward, done, info = env.step(env.action_space.sample())
-print("observation: ", observation)
-print("reward: ", reward)
-print("done: ", done)
-print("info: ", info)
+    observation, reward, done, info = env.step(env.action_space.sample())
+    print("observation: ", observation)
+    print("reward: ", reward)
+    print("done: ", done)
+    print("info: ", info)
 
-# TODO: implement write_bitcode(..) or write_ir(..)
-# env.write_bitcode("/tmp/output.bc")
+    env.reset()
+
+    # TODO: implement write_bitcode(..) or write_ir(..)
+    # env.write_bitcode("/tmp/output.bc")

examples/example_unrolling_service/example_without_bazel.py

@@ -154,32 +154,32 @@ def benchmark(self, uri: str) -> Benchmark:
     },
 )
 
-env = compiler_gym.make(
+with compiler_gym.make(
     "unrolling-py-v0",
     benchmark="unrolling-v0/offsets1",
     observation_space="features",
     reward_space="runtime",
-)
-compiler_gym.set_debug_level(4)  # TODO: check why this has no effect
+) as env:
+    compiler_gym.set_debug_level(4)  # TODO: check why this has no effect
 
-observation = env.reset()
-print("observation: ", observation)
+    observation = env.reset()
+    print("observation: ", observation)
 
-print()
+    print()
 
-observation, reward, done, info = env.step(env.action_space.sample())
-print("observation: ", observation)
-print("reward: ", reward)
-print("done: ", done)
-print("info: ", info)
+    observation, reward, done, info = env.step(env.action_space.sample())
+    print("observation: ", observation)
+    print("reward: ", reward)
+    print("done: ", done)
+    print("info: ", info)
 
-print()
+    print()
 
-observation, reward, done, info = env.step(env.action_space.sample())
-print("observation: ", observation)
-print("reward: ", reward)
-print("done: ", done)
-print("info: ", info)
+    observation, reward, done, info = env.step(env.action_space.sample())
+    print("observation: ", observation)
+    print("reward: ", reward)
+    print("done: ", done)
+    print("info: ", info)
 
-# TODO: implement write_bitcode(..) or write_ir(..)
-# env.write_bitcode("/tmp/output.bc")
+    # TODO: implement write_bitcode(..) or write_ir(..)
+    # env.write_bitcode("/tmp/output.bc")

examples/example_unrolling_service/loop_unroller/CMakeLists.txt

@@ -22,5 +22,5 @@ cg_cc_binary(
     ${LLVM_DEFINITIONS}
 )
 
-ADD_CUSTOM_TARGET(link_target ALL
+ADD_CUSTOM_TARGET(link_loop_unroller_target ALL
                   COMMAND ${CMAKE_COMMAND} -E create_symlink ${CMAKE_BINARY_DIR}/examples/example_unrolling_service/loop_unroller/loop_unroller ${CMAKE_SOURCE_DIR}/examples/example_unrolling_service/loop_unroller/loop_unroller)

examples/loop_optimizations_service/BUILD

@@ -0,0 +1,38 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+load("@rules_python//python:defs.bzl", "py_library", "py_test")
+
+py_library(
+    name = "loop_optimizations_service",
+    srcs = ["__init__.py"],
+    data = [
+        "//examples/loop_optimizations_service/benchmarks",
+        "//examples/loop_optimizations_service/service_py:loops-opt-service-py",
+    ],
+    visibility = ["//visibility:public"],
+    deps = [
+        "//compiler_gym/envs/llvm",
+        "//compiler_gym/util",
+    ],
+)
+
+py_test(
+    name = "env_tests",
+    srcs = ["env_tests.py"],
+    deps = [
+        ":loop_optimizations_service",
+        "//compiler_gym",
+        "//tests:test_main",
+    ],
+)
+
+py_binary(
+    name = "example",
+    srcs = ["example.py"],
+    deps = [
+        ":loop_optimizations_service",
+        "//compiler_gym",
+    ],
+)

examples/loop_optimizations_service/CMakeLists.txt

@@ -0,0 +1,6 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+cg_add_all_subdirs()

examples/loop_optimizations_service/README.md

@@ -0,0 +1,30 @@
+# Loop Optimizations CompilerGym Service
+
+A CompilerGym environment dedicated to loop optimizations. All paths listed below are relative to the path of this README file.
+
+* Actions: this environment currently focuses on unrolling and vectorization optimizations. The plan is to extend that to other loop optimizations. The actions are the different vectorization factors and unrolling factors.
+    - The actions are listed in `action_spaces` struct in `service_py/loops_opt_service.py`
+    - The actions are implemented in `apply_action(...)` function in `service_py/loops_opt_service.py`
+* Observations: the observations are: textual form of the LLVM IR, statistical features of different types of IR instructions, runtime execution, or code size
+    - The observations are listed in `observation_spaces` struct in `service_py/loops_opt_service.py`.
+    - The observations are implemented in `get_observation(...)` function in `service_py/loops_opt_service.py`
+* Rewards: the rewards could be runtime or code size.
+    - The rewards are implemented in `__init__.py` and they reuse the runtime and code size observations mentioned above
+* Benchmarks: this environment expects your benchmarks to follow the templates from the [Neruovectorizer repo](https://github.com/intel/neuro-vectorizer/tree/master/training_data) repo, that was in turn adapted from the [LLVM loop test suite](https://github.com/llvm/llvm-test-suite/blob/main/SingleSource/UnitTests/Vectorizer/gcc-loops.cpp).
+    - To implement your benchmark, you need to: include the `header.h` file, implement your benchmark in a custom function, then invoke it using `BENCH` macro inside the `main()` function.
+    - Following this template is necessary in order for the benchmark to measure the execution runtime and write it to stdout, which is in turn parsed by this environment to measure the runtime reward.
+    - You can view and add examples of benchmarks in `benchmarks` directory
+    - Also, when adding your own benchmark, you need to add it to the `LoopsDataset` class in `__init__.py`
+
+## Usage
+
+Run `example.py` example:
+```sh
+$ bazel run //examples/loop_optimizations_service:example
+```
+
+Run `env_tests.py` unit tests:
+
+```sh
+$ bazel test //examples/loop_optimizations_service:env_tests
+```

examples/loop_optimizations_service/__init__.py

@@ -0,0 +1,146 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+"""This module registers the Loop Optimizations CompilerGym environment """
+import subprocess
+from pathlib import Path
+from typing import Iterable
+
+from compiler_gym.datasets import Benchmark, Dataset
+from compiler_gym.envs.llvm.llvm_benchmark import get_system_includes
+from compiler_gym.spaces import Reward
+from compiler_gym.third_party import llvm
+from compiler_gym.util.registration import register
+from compiler_gym.util.runfiles_path import runfiles_path
+
+LOOPS_OPT_PY_SERVICE_BINARY: Path = runfiles_path(
+    "examples/loop_optimizations_service/service_py/loops-opt-service-py"
+)
+
+BENCHMARKS_PATH: Path = runfiles_path("examples/loop_optimizations_service/benchmarks")
+
+NEURO_VECTORIZER_HEADER: Path = runfiles_path(
+    "compiler_gym/third_party/neuro-vectorizer/header.h"
+)
+
+
+class RuntimeReward(Reward):
+    """An example reward that uses changes in the "runtime" observation value
+    to compute incremental reward.
+    """
+
+    def __init__(self):
+        super().__init__(
+            id="runtime",
+            observation_spaces=["runtime"],
+            default_value=0,
+            default_negates_returns=True,
+            deterministic=False,
+            platform_dependent=True,
+        )
+        self.baseline_runtime = 0
+
+    def reset(self, benchmark: str, observation_view):
+        del benchmark  # unused
+        self.baseline_runtime = observation_view["runtime"]
+
+    def update(self, action, observations, observation_view):
+        del action  # unused
+        del observation_view  # unused
+        return float(self.baseline_runtime - observations[0]) / self.baseline_runtime
+
+
+class SizeReward(Reward):
+    """An example reward that uses changes in the "size" observation value
+    to compute incremental reward.
+    """
+
+    def __init__(self):
+        super().__init__(
+            id="size",
+            observation_spaces=["size"],
+            default_value=0,
+            default_negates_returns=True,
+            deterministic=False,
+            platform_dependent=True,
+        )
+        self.baseline_size = 0
+
+    def reset(self, benchmark: str, observation_view):
+        del benchmark  # unused
+        self.baseline_runtime = observation_view["size"]
+
+    def update(self, action, observations, observation_view):
+        del action  # unused
+        del observation_view  # unused
+        return float(self.baseline_size - observations[0]) / self.baseline_size
+
+
+class LoopsDataset(Dataset):
+    def __init__(self, *args, **kwargs):
+        super().__init__(
+            name="benchmark://loops-opt-v0",
+            license="MIT",
+            description="Loops optimization dataset",
+        )
+
+        self._benchmarks = {
+            "benchmark://loops-opt-v0/add": Benchmark.from_file_contents(
+                "benchmark://loops-opt-v0/add",
+                self.preprocess(BENCHMARKS_PATH / "add.c"),
+            ),
+            "benchmark://loops-opt-v0/offsets1": Benchmark.from_file_contents(
+                "benchmark://loops-opt-v0/offsets1",
+                self.preprocess(BENCHMARKS_PATH / "offsets1.c"),
+            ),
+            "benchmark://loops-opt-v0/conv2d": Benchmark.from_file_contents(
+                "benchmark://loops-opt-v0/conv2d",
+                self.preprocess(BENCHMARKS_PATH / "conv2d.c"),
+            ),
+        }
+
+    @staticmethod
+    def preprocess(src: Path) -> bytes:
+        """Front a C source through the compiler frontend."""
+        # TODO(github.com/facebookresearch/CompilerGym/issues/325): We can skip
+        # this pre-processing, or do it on the service side, once support for
+        # multi-file benchmarks lands.
+        cmd = [
+            str(llvm.clang_path()),
+            "-E",
+            "-o",
+            "-",
+            "-I",
+            str(NEURO_VECTORIZER_HEADER.parent),
+            src,
+        ]
+        for directory in get_system_includes():
+            cmd += ["-isystem", str(directory)]
+        return subprocess.check_output(
+            cmd,
+            timeout=300,
+        )
+
+    def benchmark_uris(self) -> Iterable[str]:
+        yield from self._benchmarks.keys()
+
+    def benchmark(self, uri: str) -> Benchmark:
+        if uri in self._benchmarks:
+            return self._benchmarks[uri]
+        else:
+            raise LookupError("Unknown program name")
+
+
+# Register the unrolling example service on module import. After importing this module,
+# the loops-opt-py-v0 environment will be available to gym.make(...).
+
+register(
+    id="loops-opt-py-v0",
+    entry_point="compiler_gym.envs:CompilerEnv",
+    kwargs={
+        "service": LOOPS_OPT_PY_SERVICE_BINARY,
+        "rewards": [RuntimeReward(), SizeReward()],
+        "datasets": [LoopsDataset()],
+    },
+)

examples/loop_optimizations_service/benchmarks/BUILD

@@ -0,0 +1,15 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+filegroup(
+    name = "benchmarks",
+    srcs = [
+        "add.c",
+        "conv2d.c",
+        "offsets1.c",
+        "//compiler_gym/third_party/neuro-vectorizer:header",
+    ],
+    visibility = ["//visibility:public"],
+)

examples/loop_optimizations_service/benchmarks/add.c

@@ -0,0 +1,26 @@
+// Copyright (c) Facebook, Inc. and its affiliates.
+//
+// This source code is licensed under the MIT license found in the
+// LICENSE file in the root directory of this source tree.
+#include "header.h"
+
+#ifndef N
+#define N 1000000
+#endif
+
+int A[N];
+int B[N];
+
+__attribute__((noinline)) void add(int* ret) {
+  for (int i = 0; i < N; i++) A[i] = A[i] + B[i];
+
+  *ret = A[N - 1];
+}
+
+__attribute__((optnone)) int main(int argc, char* argv[]) {
+  int dummy = 0;
+  // TODO: initialize tensors
+  BENCH("add", add(&dummy), 100, dummy);
+
+  return 0;
+}