| """Utility for testing InteractiveModelRunner. |
| |
| Use it from pass-specific tests by providing a main .py which calls this library's |
| `run_interactive` with an appropriate callback to provide advice. |
| |
| From .ll tests, just call the above-mentioned main as a prefix to the opt/llc |
| invocation (with the appropriate flags enabling the interactive mode) |
| |
| Examples: |
| test/Transforms/Inline/ML/interactive-mode.ll |
| test/CodeGen/MLRegalloc/interactive-mode.ll |
| """ |
| |
| import ctypes |
| import log_reader |
| import io |
| import math |
| import os |
| import subprocess |
| from typing import Callable, List, Union |
| |
| |
| def send(f: io.BufferedWriter, value: Union[int, float], spec: log_reader.TensorSpec): |
| """Send the `value` - currently just a scalar - formatted as per `spec`.""" |
| |
| # just int64 for now |
| assert spec.element_type == ctypes.c_int64 |
| to_send = ctypes.c_int64(int(value)) |
| assert f.write(bytes(to_send)) == ctypes.sizeof(spec.element_type) * math.prod( |
| spec.shape |
| ) |
| f.flush() |
| |
| |
| def run_interactive( |
| temp_rootname: str, |
| make_response: Callable[[List[log_reader.TensorValue]], Union[int, float]], |
| process_and_args: List[str], |
| ): |
| """Host the compiler. |
| Args: |
| temp_rootname: the base file name from which to construct the 2 pipes for |
| communicating with the compiler. |
| make_response: a function that, given the current tensor values, provides a |
| response. |
| process_and_args: the full commandline for the compiler. It it assumed it |
| contains a flag poiting to `temp_rootname` so that the InteractiveModeRunner |
| would attempt communication on the same pair as this function opens. |
| |
| This function sets up the communication with the compiler - via 2 files named |
| `temp_rootname`.in and `temp_rootname`.out - prints out the received features, |
| and sends back to the compiler an advice (which it gets from `make_response`). |
| It's used for testing, and also to showcase how to set up communication in an |
| interactive ML ("gym") environment. |
| """ |
| to_compiler = temp_rootname + ".in" |
| from_compiler = temp_rootname + ".out" |
| try: |
| os.mkfifo(to_compiler, 0o666) |
| os.mkfifo(from_compiler, 0o666) |
| compiler_proc = subprocess.Popen( |
| process_and_args, stderr=subprocess.PIPE, stdout=subprocess.DEVNULL |
| ) |
| with io.BufferedWriter(io.FileIO(to_compiler, "wb")) as tc: |
| with io.BufferedReader(io.FileIO(from_compiler, "rb")) as fc: |
| tensor_specs, _, advice_spec = log_reader.read_header(fc) |
| context = None |
| while compiler_proc.poll() is None: |
| next_event = fc.readline() |
| if not next_event: |
| break |
| ( |
| last_context, |
| observation_id, |
| features, |
| _, |
| ) = log_reader.read_one_observation( |
| context, next_event, fc, tensor_specs, None |
| ) |
| if last_context != context: |
| print(f"context: {last_context}") |
| context = last_context |
| print(f"observation: {observation_id}") |
| tensor_values = [] |
| for fv in features: |
| log_reader.pretty_print_tensor_value(fv) |
| tensor_values.append(fv) |
| send(tc, make_response(tensor_values), advice_spec) |
| _, err = compiler_proc.communicate() |
| print(err.decode("utf-8")) |
| compiler_proc.wait() |
| |
| finally: |
| os.unlink(to_compiler) |
| os.unlink(from_compiler) |