LLVM: lib/Analysis/models/gen-regalloc-eviction-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. By construction, the mock model will always
4output the first liverange that can be evicted.
5"""
6import os
7import sys
8import tensorflow as tf
9
10POLICY_DECISION_LABEL = "index_to_evict"
11POLICY_OUTPUT_SPEC = """
12[
13 {
14 "logging_name": "index_to_evict",
15 "tensor_spec": {
16 "name": "StatefulPartitionedCall",
17 "port": 0,
18 "type": "int64_t",
19 "shape": [
20 1
21 ]
22 }
23 }
24]
25"""
26PER_REGISTER_FEATURE_LIST = ["mask"]
27NUM_REGISTERS = 33
28
29
31 """Returns (time_step_spec, action_spec) for LLVM register allocation."""
32 inputs = dict(
33 (key, tf.TensorSpec(dtype=tf.int64, shape=(NUM_REGISTERS), name=key))
34 for key in PER_REGISTER_FEATURE_LIST
35 )
36 return inputs
37
38
40 return os.path.join(path, "output_spec.json")
41
42
44 """Build and save the mock model with the given signature."""
45 module = tf.Module()
46
47
48 module.var = tf.Variable(0, dtype=tf.int64)
49
50 def action(*inputs):
51 result = (
52 tf.math.argmax(tf.cast(inputs[0]["mask"], tf.int32), axis=-1) + module.var
53 )
54 return {POLICY_DECISION_LABEL: result}
55
56 module.action = tf.function()(action)
57 action = {"action": module.action.get_concrete_function(get_input_signature())}
58 tf.saved_model.save(module, path, signatures=action)
60 with open(output_spec_path, "w") as f:
61 print(f"Writing output spec to {output_spec_path}.")
62 f.write(POLICY_OUTPUT_SPEC)
63
64
66 assert len(argv) == 2
67 model_path = argv[1]
69
70
71if __name__ == "__main__":
72 main(sys.argv)
static void print(raw_ostream &Out, object::Archive::Kind Kind, T Val)
get_output_spec_path(path)
Definition gen-regalloc-eviction-test-model.py:39
get_input_signature()
Definition gen-regalloc-eviction-test-model.py:30
main(argv)
Definition gen-regalloc-eviction-test-model.py:65
build_mock_model(path)
Definition gen-regalloc-eviction-test-model.py:43