| """Generate a mock model for LLVM tests for Register Allocation. |
| The generated model is not a neural net - it is just a tf.function with the |
| correct input and output parameters. By construction, the mock model will always |
| output the first liverange that can be evicted. |
| """ |
| import os |
| import sys |
| import tensorflow as tf |
| |
| POLICY_DECISION_LABEL = "index_to_evict" |
| POLICY_OUTPUT_SPEC = """ |
| [ |
| { |
| "logging_name": "index_to_evict", |
| "tensor_spec": { |
| "name": "StatefulPartitionedCall", |
| "port": 0, |
| "type": "int64_t", |
| "shape": [ |
| 1 |
| ] |
| } |
| } |
| ] |
| """ |
| PER_REGISTER_FEATURE_LIST = ["mask"] |
| NUM_REGISTERS = 33 |
| |
| |
| def get_input_signature(): |
| """Returns (time_step_spec, action_spec) for LLVM register allocation.""" |
| inputs = dict( |
| (key, tf.TensorSpec(dtype=tf.int64, shape=(NUM_REGISTERS), name=key)) |
| for key in PER_REGISTER_FEATURE_LIST |
| ) |
| return inputs |
| |
| |
| def get_output_spec_path(path): |
| return os.path.join(path, "output_spec.json") |
| |
| |
| def build_mock_model(path): |
| """Build and save the mock model with the given signature.""" |
| module = tf.Module() |
| # We have to set this useless variable in order for the TF C API to correctly |
| # intake it |
| module.var = tf.Variable(0, dtype=tf.int64) |
| |
| def action(*inputs): |
| result = ( |
| tf.math.argmax(tf.cast(inputs[0]["mask"], tf.int32), axis=-1) + module.var |
| ) |
| return {POLICY_DECISION_LABEL: result} |
| |
| module.action = tf.function()(action) |
| action = {"action": module.action.get_concrete_function(get_input_signature())} |
| tf.saved_model.save(module, path, signatures=action) |
| output_spec_path = get_output_spec_path(path) |
| with open(output_spec_path, "w") as f: |
| print(f"Writing output spec to {output_spec_path}.") |
| f.write(POLICY_OUTPUT_SPEC) |
| |
| |
| def main(argv): |
| assert len(argv) == 2 |
| model_path = argv[1] |
| build_mock_model(model_path) |
| |
| |
| if __name__ == "__main__": |
| main(sys.argv) |