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