88 lines
3.3 KiB
Python
88 lines
3.3 KiB
Python
"""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)
|