LLVM 20.0.0git
gen-regalloc-eviction-test-model.py
Go to the documentation of this file.
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
30def get_input_signature():
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 # We have to set this useless variable in order for the TF C API to correctly
47 # intake it
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)
59 output_spec_path = get_output_spec_path(path)
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
65def main(argv):
66 assert len(argv) == 2
67 model_path = argv[1]
68 build_mock_model(model_path)
69
70
71if __name__ == "__main__":
72 main(sys.argv)
static void print(raw_ostream &Out, object::Archive::Kind Kind, T Val)