73 lines · python
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 the3correct input and output parameters. By construction, the mock model will always4output the first liverange that can be evicted.5"""6import os7import sys8import tensorflow as tf9 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 121 ]22 }23 }24]25"""26PER_REGISTER_FEATURE_LIST = ["mask"]27NUM_REGISTERS = 3328 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_LIST35 )36 return inputs37 38 39def get_output_spec_path(path):40 return os.path.join(path, "output_spec.json")41 42 43def build_mock_model(path):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 correctly47 # intake it48 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.var53 )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) == 267 model_path = argv[1]68 build_mock_model(model_path)69 70 71if __name__ == "__main__":72 main(sys.argv)73