def test_not_equal():
box = Box(low=0, high=1, name="test_box", shape=[1, 2], dtype=int)
assert box != Box(
low=0, high=1, name="test_box_2", shape=[1, 2], dtype=int)
assert box != Box(low=-1, high=1, name="test_box", shape=[1, 2], dtype=int)
assert box != Box(low=0, high=2, name="test_box", shape=[1, 2], dtype=int)
assert box != Box(low=0, high=1, name="test_box", shape=[1, 3], dtype=int)
assert box != "not_a_box"
def test_equal():
assert Box(low=0, high=1, name="test_box", shape=[1, 2],
dtype=int) == Box(low=0,
high=1,
name="test_box",
shape=[1, 2],
dtype=int)
assert Box(low=0, high=1, name="test_box", shape=[1, 2],
dtype=int) == Box(low=0,
high=1,
name="test_box",
shape=[1, 2],
dtype=float)
def test_observation_spaces(env: CompilerEnv):
"""Test that the environment reports the service's observation spaces."""
env.reset()
assert env.observation.spaces.keys() == {
"ir", "features", "runtime", "size"
}
assert env.observation.spaces["ir"].space == Sequence(
name="ir",
size_range=(0, np.iinfo(np.int64).max),
dtype=str,
opaque_data_format=None,
)
assert env.observation.spaces["features"].space == Box(name="features",
shape=(3, ),
low=0,
high=100000,
dtype=int)
assert env.observation.spaces["runtime"].space == Scalar(name="runtime",
min=0,
max=np.inf,
dtype=float)
assert env.observation.spaces["size"].space == Scalar(name="size",
min=0,
max=np.inf,
dtype=float)
def test_add_derived_space(env: LlvmEnv):
env.reset()
with pytest.deprecated_call(
match=
"Use the derived_observation_spaces argument to CompilerEnv constructor."
):
env.observation.add_derived_space(
id="IrLen",
base_id="Ir",
space=Box(name="IrLen",
low=0,
high=float("inf"),
shape=(1, ),
dtype=int),
translate=lambda base: [15],
)
value = env.observation["IrLen"]
assert isinstance(value, list)
assert value == [15]
# Repeat the above test using the generated bound method.
value = env.observation.IrLen()
assert isinstance(value, list)
assert value == [15]
def test_mlir_rl_wrapper_env_observation_space(env: MlirEnv):
wrapper_env = make_mlir_rl_wrapper_env(env)
observation_space = wrapper_env.observation_space
assert observation_space == Box(name="Runtime",
shape=[1],
low=0,
high=np.inf,
dtype=float)
def test_convert_to_int64_box_message():
box = Box(
low=np.array([[1], [2]]), high=np.array([[3], [4]]), name=None, dtype=np.int64
)
converted_box = py_converters.convert_to_box_message(box)
assert isinstance(converted_box, Int64Box)
assert isinstance(converted_box.low, Int64Tensor)
assert np.array_equal(converted_box.low.shape, box.shape)
assert np.array_equal(converted_box.low.value, box.low.flatten())
assert isinstance(converted_box.high, Int64Tensor)
assert np.array_equal(converted_box.high.shape, box.shape)
assert np.array_equal(converted_box.high.value, box.high.flatten())
def test_convert_to_double_box_message():
box = Box(
low=np.array([[1.0], [2.0]]),
high=np.array([[3.0], [4.0]]),
name=None,
dtype=np.float64,
)
converted_box = py_converters.convert_to_box_message(box)
assert isinstance(converted_box, DoubleBox)
assert isinstance(converted_box.low, DoubleTensor)
assert np.array_equal(converted_box.low.shape, box.shape)
assert np.array_equal(converted_box.low.value, box.low.flatten())
assert isinstance(converted_box.high, DoubleTensor)
assert np.array_equal(converted_box.high.shape, box.shape)
assert np.array_equal(converted_box.high.value, box.high.flatten())
def test_convert_to_float_box_message():
box = Box(
low=np.array([[1], [2]], dtype=np.float32),
high=np.array([[3], [4]], dtype=np.float32),
name=None,
dtype=np.float32,
)
converted_box = py_converters.convert_to_box_message(box)
assert isinstance(converted_box, FloatBox)
assert isinstance(converted_box.low, FloatTensor)
assert np.array_equal(converted_box.low.shape, box.shape)
assert np.array_equal(converted_box.low.value, box.low.flatten())
assert isinstance(converted_box.high, FloatTensor)
assert np.array_equal(converted_box.high.shape, box.shape)
assert np.array_equal(converted_box.high.value, box.high.flatten())
def test_convert_to_boolean_box_message():
box = Box(
low=np.array([[False], [True]]),
high=np.array([[False], [True]]),
name=None,
dtype=bool,
)
converted_box = py_converters.convert_to_box_message(box)
assert isinstance(converted_box, BooleanBox)
assert isinstance(converted_box.low, BooleanTensor)
assert np.array_equal(converted_box.low.shape, box.shape)
assert np.array_equal(converted_box.low.value, box.low.flatten())
assert isinstance(converted_box.high, BooleanTensor)
assert np.array_equal(converted_box.high.shape, box.shape)
assert np.array_equal(converted_box.high.value, box.high.flatten())
def test_to_space_message_default_converter():
space = Tuple(
name=None,
spaces=[
Dict(
name=None,
spaces={"key": Box(name=None, low=0, high=1, shape=[1, 2])},
)
],
)
converted_space = py_converters.to_space_message_default_converter()(space)
assert isinstance(converted_space, Space)
assert isinstance(
converted_space.space_list.space[0].space_dict.space["key"].float_box,
FloatBox,
)
def test_convert_to_byte_box_message():
box = Box(
low=np.array([[1], [2]]), high=np.array([[3], [4]]), name=None, dtype=np.int8
)
converted_box = py_converters.convert_to_box_message(box)
assert isinstance(converted_box, ByteBox)
assert isinstance(converted_box.low, ByteTensor)
assert np.array_equal(converted_box.low.shape, box.shape)
assert np.array_equal(
np.frombuffer(converted_box.low.value, dtype=np.int8), box.low.flatten()
)
assert isinstance(converted_box.high, ByteTensor)
assert np.array_equal(converted_box.high.shape, box.shape)
assert np.array_equal(
np.frombuffer(converted_box.high.value, dtype=np.int8), box.high.flatten()
)
def __init__(
self,
*args,
benchmark: Optional[Union[str, Benchmark]] = None,
datasets_site_path: Optional[Path] = None,
**kwargs,
):
# First perform a one-time download of LLVM binaries that are needed by
# the LLVM service and are not included by the pip-installed package.
download_llvm_files()
self.inst2vec = _INST2VEC_ENCODER
super().__init__(
*args,
**kwargs,
# Set a default benchmark for use.
benchmark=benchmark or "cbench-v1/qsort",
datasets=_get_llvm_datasets(site_data_base=datasets_site_path),
rewards=[
CostFunctionReward(
name="IrInstructionCount",
cost_function="IrInstructionCount",
init_cost_function="IrInstructionCountO0",
default_negates_returns=True,
deterministic=True,
platform_dependent=False,
),
NormalizedReward(
name="IrInstructionCountNorm",
cost_function="IrInstructionCount",
init_cost_function="IrInstructionCountO0",
max=1,
default_negates_returns=True,
deterministic=True,
platform_dependent=False,
),
BaselineImprovementNormalizedReward(
name="IrInstructionCountO3",
cost_function="IrInstructionCount",
baseline_cost_function="IrInstructionCountO3",
init_cost_function="IrInstructionCountO0",
success_threshold=1,
default_negates_returns=True,
deterministic=True,
platform_dependent=False,
),
BaselineImprovementNormalizedReward(
name="IrInstructionCountOz",
cost_function="IrInstructionCount",
baseline_cost_function="IrInstructionCountOz",
init_cost_function="IrInstructionCountO0",
success_threshold=1,
default_negates_returns=True,
deterministic=True,
platform_dependent=False,
),
CostFunctionReward(
name="ObjectTextSizeBytes",
cost_function="ObjectTextSizeBytes",
init_cost_function="ObjectTextSizeO0",
default_negates_returns=True,
deterministic=True,
platform_dependent=True,
),
NormalizedReward(
name="ObjectTextSizeNorm",
cost_function="ObjectTextSizeBytes",
init_cost_function="ObjectTextSizeO0",
max=1,
default_negates_returns=True,
deterministic=True,
platform_dependent=True,
),
BaselineImprovementNormalizedReward(
name="ObjectTextSizeO3",
cost_function="ObjectTextSizeBytes",
init_cost_function="ObjectTextSizeO0",
baseline_cost_function="ObjectTextSizeO3",
success_threshold=1,
default_negates_returns=True,
deterministic=True,
platform_dependent=True,
),
BaselineImprovementNormalizedReward(
name="ObjectTextSizeOz",
cost_function="ObjectTextSizeBytes",
init_cost_function="ObjectTextSizeO0",
baseline_cost_function="ObjectTextSizeOz",
success_threshold=1,
default_negates_returns=True,
deterministic=True,
platform_dependent=True,
),
CostFunctionReward(
name="TextSizeBytes",
cost_function="TextSizeBytes",
init_cost_function="TextSizeO0",
default_negates_returns=True,
deterministic=True,
platform_dependent=True,
),
NormalizedReward(
name="TextSizeNorm",
cost_function="TextSizeBytes",
init_cost_function="TextSizeO0",
max=1,
default_negates_returns=True,
deterministic=True,
platform_dependent=True,
),
BaselineImprovementNormalizedReward(
name="TextSizeO3",
cost_function="TextSizeBytes",
init_cost_function="TextSizeO0",
baseline_cost_function="TextSizeO3",
success_threshold=1,
default_negates_returns=True,
deterministic=True,
platform_dependent=True,
),
BaselineImprovementNormalizedReward(
name="TextSizeOz",
cost_function="TextSizeBytes",
init_cost_function="TextSizeO0",
baseline_cost_function="TextSizeOz",
success_threshold=1,
default_negates_returns=True,
deterministic=True,
platform_dependent=True,
),
],
derived_observation_spaces=[
{
"id":
"Inst2vecPreprocessedText",
"base_id":
"Ir",
"space":
Sequence(name="Inst2vecPreprocessedText",
size_range=(0, None),
dtype=str),
"translate":
self.inst2vec.preprocess,
"default_value":
"",
},
{
"id":
"Inst2vecEmbeddingIndices",
"base_id":
"Ir",
"space":
Sequence(
name="Inst2vecEmbeddingIndices",
size_range=(0, None),
dtype=np.int32,
),
"translate":
lambda base_observation: self.inst2vec.encode(
self.inst2vec.preprocess(base_observation)),
"default_value":
np.array([self.inst2vec.vocab["!UNK"]]),
},
{
"id":
"Inst2vec",
"base_id":
"Ir",
"space":
Sequence(name="Inst2vec",
size_range=(0, None),
dtype=np.ndarray),
"translate":
lambda base_observation: self.inst2vec.embed(
self.inst2vec.encode(
self.inst2vec.preprocess(base_observation))),
"default_value":
np.vstack([
self.inst2vec.embeddings[self.inst2vec.vocab["!UNK"]]
]),
},
{
"id":
"InstCountDict",
"base_id":
"InstCount",
"space":
DictSpace(
{
f"{name}Count": Scalar(name=f"{name}Count",
min=0,
max=None,
dtype=int)
for name in INST_COUNT_FEATURE_NAMES
},
name="InstCountDict",
),
"translate":
lambda base_observation: {
f"{name}Count": val
for name, val in zip(INST_COUNT_FEATURE_NAMES,
base_observation)
},
},
{
"id":
"InstCountNorm",
"base_id":
"InstCount",
"space":
Box(
name="InstCountNorm",
low=0,
high=1,
shape=(len(INST_COUNT_FEATURE_NAMES) - 1, ),
dtype=np.float32,
),
"translate":
lambda base_observation: (base_observation[1:] / max(
base_observation[0], 1)).astype(np.float32),
},
{
"id":
"InstCountNormDict",
"base_id":
"InstCountNorm",
"space":
DictSpace(
{
f"{name}Density": Scalar(name=f"{name}Density",
min=0,
max=None,
dtype=int)
for name in INST_COUNT_FEATURE_NAMES[1:]
},
name="InstCountNormDict",
),
"translate":
lambda base_observation: {
f"{name}Density": val
for name, val in zip(INST_COUNT_FEATURE_NAMES[1:],
base_observation)
},
},
{
"id":
"AutophaseDict",
"base_id":
"Autophase",
"space":
DictSpace(
{
name: Scalar(name=name, min=0, max=None, dtype=int)
for name in AUTOPHASE_FEATURE_NAMES
},
name="AutophaseDict",
),
"translate":
lambda base_observation: {
name: val
for name, val in zip(AUTOPHASE_FEATURE_NAMES,
base_observation)
},
},
],
)
# Mutable runtime configuration options that must be set on every call
# to reset.
self._runtimes_per_observation_count: Optional[int] = None
self._runtimes_warmup_per_observation_count: Optional[int] = None
cpu_info_spaces = [
Sequence(name="name", size_range=(0, None), dtype=str),
Scalar(name="cores_count", min=None, max=None, dtype=int),
Scalar(name="l1i_cache_size", min=None, max=None, dtype=int),
Scalar(name="l1i_cache_count", min=None, max=None, dtype=int),
Scalar(name="l1d_cache_size", min=None, max=None, dtype=int),
Scalar(name="l1d_cache_count", min=None, max=None, dtype=int),
Scalar(name="l2_cache_size", min=None, max=None, dtype=int),
Scalar(name="l2_cache_count", min=None, max=None, dtype=int),
Scalar(name="l3_cache_size", min=None, max=None, dtype=int),
Scalar(name="l3_cache_count", min=None, max=None, dtype=int),
Scalar(name="l4_cache_size", min=None, max=None, dtype=int),
Scalar(name="l4_cache_count", min=None, max=None, dtype=int),
]
self.observation.spaces["CpuInfo"].space = DictSpace(
{space.name: space
for space in cpu_info_spaces},
name="CpuInfo",
)
def test_action_space(env: MlirEnv):
expected_action_space = SpaceSequence(
name="MatrixMultiplication",
size_range=[1, 4],
space=Dict(
name=None,
spaces={
"tile_options":
Dict(
name=None,
spaces={
"interchange_vector":
Permutation(
name=None,
scalar_range=Scalar(name=None,
min=0,
max=2,
dtype=int),
),
"tile_sizes":
Box(
name=None,
low=np.array([1] * 3, dtype=int),
high=np.array([2**32] * 3, dtype=int),
dtype=np.int64,
),
"promote":
Scalar(name=None, min=False, max=True, dtype=bool),
"promote_full_tile":
Scalar(name=None, min=False, max=True, dtype=bool),
"loop_type":
NamedDiscrete(
name=None,
items=["loops", "affine_loops"],
),
},
),
"vectorize_options":
Dict(
name=None,
spaces={
"vectorize_to":
NamedDiscrete(
name=None,
items=["dot", "matmul", "outer_product"],
),
"vector_transfer_split":
NamedDiscrete(
name=None,
items=["none", "linalg_copy", "vector_transfer"],
),
"unroll_vector_transfers":
Scalar(
name=None,
min=False,
max=True,
dtype=bool,
),
},
),
},
),
)
assert expected_action_space == env.action_space