[examples] Add implementations of LLVM autotuners.

Author: ChrisCummins

compiler_gym/envs/llvm/compute_observation.py

@@ -11,6 +11,7 @@
 import google.protobuf.text_format
 
 from compiler_gym.service.proto import Observation
+from compiler_gym.util.gym_type_hints import ObservationType
 from compiler_gym.util.runfiles_path import runfiles_path
 from compiler_gym.util.shell_format import plural
 from compiler_gym.views.observation_space_spec import ObservationSpaceSpec
@@ -35,7 +36,7 @@ def pascal_case_to_enum(pascal_case: str) -> str:
 
 def compute_observation(
     observation_space: ObservationSpaceSpec, bitcode: Path, timeout: float = 300
-):
+) -> ObservationType:
     """Compute an LLVM observation.
 
     This is a utility function that uses a standalone C++ binary to compute an

examples/llvm_autotuning/autotuners/__init__.py

@@ -0,0 +1,96 @@
+# 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 modules defines a class for describing LLVM autotuners."""
+import tempfile
+from pathlib import Path
+from typing import Any, Dict
+
+from llvm_autotuning.autotuners.greedy import greedy  # noqa autotuner
+from llvm_autotuning.autotuners.nevergrad_ import nevergrad  # noqa autotuner
+from llvm_autotuning.autotuners.opentuner_ import opentuner_ga  # noqa autotuner
+from llvm_autotuning.autotuners.random_ import random  # noqa autotuner
+from llvm_autotuning.optimization_target import OptimizationTarget
+from pydantic import BaseModel, validator
+
+from compiler_gym.compiler_env_state import CompilerEnvState
+from compiler_gym.envs import CompilerEnv
+from compiler_gym.util.capture_output import capture_output
+from compiler_gym.util.runfiles_path import transient_cache_path
+from compiler_gym.util.temporary_working_directory import temporary_working_directory
+from compiler_gym.util.timer import Timer
+
+
+class Autotuner(BaseModel):
+
+    algorithm: str
+    """The name of the autotuner algorithm."""
+
+    optimization_target: OptimizationTarget
+    """The target that the autotuner is optimizing for."""
+
+    search_time_seconds: int
+    """The search budget of the autotuner."""
+
+    algorithm_config: Dict[str, Any] = {}
+    """An optional dictionary of keyword arguments for the autotuner function."""
+
+    @property
+    def autotune(self):
+        """Return the autotuner function for this algorithm.
+
+        An autotuner function takes a single CompilerEnv argument and optional
+        keyword configuration arguments (determined by algorithm_config) and
+        tunes the environment, returning nothing.
+        """
+        try:
+            return globals()[self.algorithm]
+        except KeyError as e:
+            raise ValueError(
+                f"Unknown autotuner: {self.algorithm}.\n"
+                f"Make sure the {self.algorithm}() function definition is available "
+                "in the global namespace of {__file__}."
+            ) from e
+
+    @property
+    def autotune_kwargs(self) -> Dict[str, Any]:
+        """Get the keyword arguments dictionary for the autotuner."""
+        kwargs = {
+            "optimization_target": self.optimization_target,
+            "search_time_seconds": self.search_time_seconds,
+        }
+        kwargs.update(self.algorithm_config)
+        return kwargs
+
+    def __call__(self, env: CompilerEnv, seed: int = 0xCC) -> CompilerEnvState:
+        """Autotune the given environment.
+
+        :param env: The environment to autotune.
+
+        :param seed: The random seed for the autotuner.
+
+        :returns: A CompilerEnvState tuple describing the autotuning result.
+        """
+        # Run the autotuner in a temporary working directory and capture the
+        # stdout/stderr.
+        with tempfile.TemporaryDirectory(
+            dir=transient_cache_path("."), prefix="autotune-"
+        ) as tmpdir:
+            with temporary_working_directory(Path(tmpdir)):
+                with capture_output():
+                    with Timer() as timer:
+                        self.autotune(env, seed=seed, **self.autotune_kwargs)
+
+        return CompilerEnvState(
+            benchmark=env.benchmark.uri,
+            commandline=env.commandline(),
+            walltime=timer.time,
+            reward=self.optimization_target.final_reward(env),
+        )
+
+    # === Start of implementation details. ===
+
+    @validator("algorithm_config", pre=True)
+    def validate_algorithm_config(cls, value) -> Dict[str, Any]:
+        return value or {}

examples/llvm_autotuning/autotuners/greedy.py

@@ -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.
+from time import time
+
+
+def greedy(env, search_time_seconds: int, **kwargs) -> None:
+    """A greedy search policy.
+
+    At each step, the policy evaluates all possible actions and selects the
+    action with the highest reward. The search stops when no action produces a
+    positive reward.
+
+    :param env: The environment to optimize.
+    """
+
+    def eval_action(env, action: int):
+        with env.fork() as fkd:
+            return (fkd.step(action)[1], action)
+
+    end_time = time() + search_time_seconds
+    while time() < end_time:
+        best = max(eval_action(env, action) for action in range(env.action_space.n))
+        if best[0] <= 0 or env.step(best[1])[2]:
+            return

examples/llvm_autotuning/autotuners/greedy_test.py

@@ -0,0 +1,30 @@
+# 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.
+"""Integration tests for the LLVM autotuners."""
+import pytest
+from llvm_autotuning.autotuners import Autotuner
+
+import compiler_gym
+
+
+@pytest.mark.skip(reason="greedy takes a long time")
+def test_autotune():
+    with compiler_gym.make("llvm-v0", reward_space="IrInstructionCount") as env:
+        env.reset(benchmark="benchmark://cbench-v1/crc32")
+
+        autotuner = Autotuner(
+            algorithm="greedy",
+            optimization_target="codesize",
+            search_time_seconds=3,
+        )
+
+        result = autotuner(env)
+        print(result)
+        assert result.benchmark == "benchmark://cbench-v1/crc32"
+        assert result.walltime >= 3
+        assert result.commandline == env.commandline()
+        assert env.episode_reward
+        assert env.benchmark == "benchmark://cbench-v1/crc32"
+        assert env.reward_space == "IrInstructionCount"

examples/llvm_autotuning/autotuners/nevergrad_.py

@@ -0,0 +1,64 @@
+# 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.
+from functools import lru_cache
+from time import time
+from typing import Tuple
+
+import nevergrad as ng
+from llvm_autotuning.optimization_target import OptimizationTarget
+
+from compiler_gym.envs import CompilerEnv
+
+
+def nevergrad(
+    env: CompilerEnv,
+    optimization_target: OptimizationTarget,
+    search_time_seconds: int,
+    seed: int,
+    episode_length: int = 100,
+    optimizer: str = "DiscreteLenglerOnePlusOne",
+    **kwargs
+) -> None:
+    """Optimize an environment using nevergrad.
+
+    Nevergrad is a gradient-free optimization platform that provides
+    implementations of various black box optimizations techniques:
+
+        https://facebookresearch.github.io/nevergrad/
+    """
+    if optimization_target == OptimizationTarget.RUNTIME:
+
+        def calculate_negative_reward(actions: Tuple[int]) -> float:
+            env.reset()
+            env.step(actions)
+            return -env.episode_reward
+
+    else:
+        # Only cache the deterministic non-runtime rewards.
+        @lru_cache(maxsize=int(1e4))
+        def calculate_negative_reward(actions: Tuple[int]) -> float:
+            env.reset()
+            env.step(actions)
+            return -env.episode_reward
+
+    params = ng.p.Choice(
+        choices=range(env.action_space.n),
+        repetitions=episode_length,
+        deterministic=True,
+    )
+    params.random_state.seed(seed)
+
+    optimizer_class = getattr(ng.optimizers, optimizer)
+    optimizer = optimizer_class(parametrization=params, budget=1, num_workers=1)
+
+    end_time = time() + search_time_seconds
+    while time() < end_time:
+        x = optimizer.ask()
+        optimizer.tell(x, calculate_negative_reward(x.value))
+
+    # Get best solution and replay it.
+    recommendation = optimizer.provide_recommendation()
+    env.reset()
+    env.step(recommendation.value)

examples/llvm_autotuning/autotuners/nevergrad_test.py

@@ -0,0 +1,29 @@
+# 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.
+"""Integration tests for the LLVM autotuners."""
+from llvm_autotuning.autotuners import Autotuner
+
+import compiler_gym
+
+
+def test_autotune():
+    with compiler_gym.make("llvm-v0", reward_space="IrInstructionCount") as env:
+        env.reset(benchmark="benchmark://cbench-v1/crc32")
+        env.reward_space = "IrInstructionCount"
+
+        autotuner = Autotuner(
+            algorithm="nevergrad",
+            optimization_target="codesize",
+            search_time_seconds=3,
+        )
+
+        result = autotuner(env)
+        print(result)
+        assert result.benchmark == "benchmark://cbench-v1/crc32"
+        assert result.walltime >= 3
+        assert result.commandline == env.commandline()
+        assert env.episode_reward >= 0
+        assert env.benchmark == "benchmark://cbench-v1/crc32"
+        assert env.reward_space == "IrInstructionCount"