1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
|
"""Generate a mock model for LLVM tests for Register Allocation.
The generated model is not a neural net - it is just a tf.function with the
correct input and output parameters.
"""
## By construction, the mock model will always output the first liverange that can be evicted.
import os
import sys
import tensorflow as tf
POLICY_DECISION_LABEL = "priority"
POLICY_OUTPUT_SPEC = """
[
{
"logging_name": "priority",
"tensor_spec": {
"name": "StatefulPartitionedCall",
"port": 0,
"type": "float",
"shape": [
1
]
}
}
]
"""
PER_LIVEINTERVAL_INT64_FEATURE_LIST = ["li_size", "stage"]
PER_LIVEINTERVAL_FLOAT32_FEATURE_LIST = ["weight"]
PER_LIVEINTERVAL_FEATURE_LIST = (
PER_LIVEINTERVAL_FLOAT32_FEATURE_LIST + PER_LIVEINTERVAL_INT64_FEATURE_LIST
)
CONTEXT_FEATURE_LIST = ("discount", "reward", "step_type")
def get_input_signature():
"""Returns (time_step_spec, action_spec) for LLVM register allocation."""
inputs = dict(
(key, tf.TensorSpec(dtype=tf.int64, shape=(), name=key))
for key in PER_LIVEINTERVAL_INT64_FEATURE_LIST
)
inputs.update(
dict(
(key, tf.TensorSpec(dtype=tf.float32, shape=(), name=key))
for key in PER_LIVEINTERVAL_FLOAT32_FEATURE_LIST
)
)
inputs.update(
dict(
(key, tf.TensorSpec(dtype=tf.float32, shape=(), name=key))
for key in ["discount", "reward"]
)
)
inputs.update(
dict(
(key, tf.TensorSpec(dtype=tf.int32, shape=(), name=key))
for key in ["step_type"]
)
)
return inputs
def get_output_spec_path(path):
return os.path.join(path, "output_spec.json")
def build_mock_model(path):
"""Build and save the mock model with the given signature."""
module = tf.Module()
# We have to set this useless variable in order for the TF C API to correctly
# intake it
module.var = tf.Variable(0, dtype=tf.float32)
def action(*inputs):
s1 = tf.reduce_sum(
[
tf.cast(inputs[0][key], tf.float32)
for key in PER_LIVEINTERVAL_FEATURE_LIST
],
axis=0,
)
s2 = tf.reduce_sum(
[tf.cast(inputs[0][key], tf.float32) for key in CONTEXT_FEATURE_LIST]
)
# Add a large number so s won't be 0.
s = s1 + s2
result = s + module.var
return {POLICY_DECISION_LABEL: result}
module.action = tf.function()(action)
action = {"action": module.action.get_concrete_function(get_input_signature())}
tf.saved_model.save(module, path, signatures=action)
output_spec_path = get_output_spec_path(path)
with open(output_spec_path, "w") as f:
print(f"Writing output spec to {output_spec_path}.")
f.write(POLICY_OUTPUT_SPEC)
def main(argv):
assert len(argv) == 2
model_path = argv[1]
build_mock_model(model_path)
if __name__ == "__main__":
main(sys.argv)
|
