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
|
"""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)
|
