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