def test_not_equal():
named_discrete = NamedDiscrete(["a", "b"], name="test_named_discrete")
assert named_discrete != NamedDiscrete(["a", "bb"],
name="test_named_discrete")
assert named_discrete != NamedDiscrete(["a", "b", "c"],
name="test_named_discrete")
assert named_discrete != NamedDiscrete(["a", "b"],
name="test_named_discrete_2")
assert named_discrete != "not_a_named_discrete"
def test_action_space(env: CompilerEnv):
assert env.action_spaces == [
NamedDiscrete(
name="default",
items=["a", "b", "c"],
)
]
def test_action_space(env: CompilerEnv):
"""Test that the environment reports the service's action spaces."""
assert env.action_spaces == [
NamedDiscrete(
name="default",
items=["a", "b", "c"],
)
]
def test_mlir_rl_wrapper_env_action_space(env: MlirEnv):
wrapper_env = make_mlir_rl_wrapper_env(env)
action_space = wrapper_env.action_space
tile_size = NamedDiscrete(
name=None,
items=[
"1", "2", "4", "8", "16", "32", "64", "128", "256", "512", "1024"
],
)
expected_subspace = Dict(
name=None,
spaces={
"tile_options":
Dict(
name=None,
spaces={
"interchange_vector":
Discrete(name=None, n=6),
"tile_sizes":
TupleSpace(name=None,
spaces=[tile_size, tile_size, tile_size]),
"promote":
NamedDiscrete(name=None, items=["False", "True"]),
"promote_full_tile":
NamedDiscrete(name=None, items=["False", "True"]),
"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":
NamedDiscrete(name=None, items=["False", "True"]),
},
),
},
)
assert action_space[0] == expected_subspace
for i in range(1, 4):
assert action_space[i]["is_present"] == NamedDiscrete(
name=None, items=["False", "True"])
assert action_space[i]["space"] == expected_subspace
def _make_action_space(self, name: str, entries: List[str]) -> Space:
"""Create an action space from the given values.
Subclasses may override this method to produce specialized action
spaces.
:param name: The name of the action space.
:param entries: The entries in the action space.
:return: A :code:`gym.Space` instance.
"""
return NamedDiscrete(entries, name)
def test_action_space(env: CompilerEnv):
"""Test that the environment reports the service's action spaces."""
assert env.action_spaces == [
NamedDiscrete(
name="unrolling",
items=[
"-loop-unroll -unroll-count=2",
"-loop-unroll -unroll-count=4",
"-loop-unroll -unroll-count=8",
],
)
]
def test_action_space(env: CompilerEnv):
"""Test that the environment reports the service's action spaces."""
assert env.action_spaces == [
NamedDiscrete(
name="loop-opt",
items=[
"--loop-unroll --unroll-count=2",
"--loop-unroll --unroll-count=4",
"--loop-unroll --unroll-count=8",
"--loop-unroll --unroll-count=16",
"--loop-unroll --unroll-count=32",
"--loop-vectorize -force-vector-width=2",
"--loop-vectorize -force-vector-width=4",
"--loop-vectorize -force-vector-width=8",
"--loop-vectorize -force-vector-width=16",
"--loop-vectorize -force-vector-width=32",
],
)
]
def test_to_string():
space = NamedDiscrete(["foo", "bar"])
assert space.to_string(0) == "foo"
assert space.to_string([0]) == "foo"
assert space.to_string([0, 0, 1]) == "foo foo bar"
def test_repr():
space = NamedDiscrete(["foo", "bar"])
assert str(space) == "NamedDiscrete([foo, bar])"
def test_name_lookup():
space = NamedDiscrete(["a", "b", "c"])
assert space["a"] == 0
assert space["b"] == 1
assert space["c"] == 2
def test_space_size():
space = NamedDiscrete(["a", "b", "c"])
assert space.n == 3
def test_invalid_name_lookup():
space = NamedDiscrete(["foo"])
with pytest.raises(ValueError):
_ = space["bar"]
def test_empty_space():
space = NamedDiscrete([])
assert space.n == 0
assert space.names == []
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
def test_equal():
assert NamedDiscrete(["a", "b"],
name="test_named_discrete") == NamedDiscrete(
["a", "b"], name="test_named_discrete")
def test_convert_to_named_discrete_space_message():
space = NamedDiscrete(name=None, items=["a", "b"])
converted_space = py_converters.convert_to_named_discrete_space_message(space)
assert isinstance(converted_space, NamedDiscreteSpace)
assert np.array_equal(space.names, converted_space.name)
