LLVM  16.0.0git
gen-inline-oz-test-model.py
Go to the documentation of this file.
1 """Generate a mock model for LLVM tests.
2 
3 The generated model is not a neural net - it is just a tf.function with the
4 correct input and output parameters. By construction, the mock model will always
5 output 1.
6 """
7 
8 import os
9 import importlib.util
10 import sys
11 
12 import tensorflow as tf
13 
14 POLICY_DECISION_LABEL = 'inlining_decision'
15 POLICY_OUTPUT_SPEC = """
16 [
17  {
18  "logging_name": "inlining_decision",
19  "tensor_spec": {
20  "name": "StatefulPartitionedCall",
21  "port": 0,
22  "type": "int64_t",
23  "shape": [
24  1
25  ]
26  }
27  }
28 ]
29 """
30 
31 
32 # pylint: disable=g-complex-comprehension
34  """Returns the list of features for LLVM inlining."""
35  # int64 features
36  inputs = [
37  tf.TensorSpec(dtype=tf.int64, shape=(), name=key) for key in [
38  'caller_basic_block_count',
39  'caller_conditionally_executed_blocks',
40  'caller_users',
41  'callee_basic_block_count',
42  'callee_conditionally_executed_blocks',
43  'callee_users',
44  'nr_ctant_params',
45  'node_count',
46  'edge_count',
47  'callsite_height',
48  'cost_estimate',
49  'inlining_default',
50  'sroa_savings',
51  'sroa_losses',
52  'load_elimination',
53  'call_penalty',
54  'call_argument_setup',
55  'load_relative_intrinsic',
56  'lowered_call_arg_setup',
57  'indirect_call_penalty',
58  'jump_table_penalty',
59  'case_cluster_penalty',
60  'switch_penalty',
61  'unsimplified_common_instructions',
62  'num_loops',
63  'dead_blocks',
64  'simplified_instructions',
65  'constant_args',
66  'constant_offset_ptr_args',
67  'callsite_cost',
68  'cold_cc_penalty',
69  'last_call_to_static_bonus',
70  'is_multiple_blocks',
71  'nested_inlines',
72  'nested_inline_cost_estimate',
73  'threshold',
74  ]
75  ]
76 
77  # float32 features
78  inputs.extend([
79  tf.TensorSpec(dtype=tf.float32, shape=(), name=key)
80  for key in ['discount', 'reward']
81  ])
82 
83  # int32 features
84  inputs.extend([
85  tf.TensorSpec(dtype=tf.int32, shape=(), name=key)
86  for key in ['step_type']
87  ])
88  return inputs
89 
90 
92  return POLICY_DECISION_LABEL
93 
94 
96  return POLICY_OUTPUT_SPEC
97 
99  return os.path.join(path, 'output_spec.json')
100 
101 
102 def build_mock_model(path, signature):
103  """Build and save the mock model with the given signature"""
104  module = tf.Module()
105 
106  # We have to set this useless variable in order for the TF C API to correctly
107  # intake it
108  module.var = tf.Variable(0.)
109 
110  def action(*inputs):
111  s = tf.reduce_sum([tf.cast(x, tf.float32) for x in tf.nest.flatten(inputs)])
112  return {signature['output']: float('inf') + s + module.var}
113 
114  module.action = tf.function()(action)
115  action = {'action': module.action.get_concrete_function(signature['inputs'])}
116  tf.saved_model.save(module, path, signatures=action)
117 
118  output_spec_path = get_output_spec_path(path)
119  with open(output_spec_path, 'w') as f:
120  print(f'Writing output spec to {output_spec_path}.')
121  f.write(signature['output_spec'])
122 
123 
125  return {
126  'inputs': get_input_signature(),
127  'output': get_output_signature(),
128  'output_spec': get_output_spec()
129  }
130 
131 
132 def main(argv):
133  assert len(argv) == 2
134  model_path = argv[1]
135 
136  print(f'Output model to: [{argv[1]}]')
137  signature = get_signature()
138  build_mock_model(model_path, signature)
139 
140 
141 if __name__ == '__main__':
142  main(sys.argv)
gen-inline-oz-test-model.get_signature
def get_signature()
Definition: gen-inline-oz-test-model.py:124
print
static void print(raw_ostream &Out, object::Archive::Kind Kind, T Val)
Definition: ArchiveWriter.cpp:189
gen-inline-oz-test-model.get_input_signature
def get_input_signature()
Definition: gen-inline-oz-test-model.py:33
gen-inline-oz-test-model.get_output_spec
def get_output_spec()
Definition: gen-inline-oz-test-model.py:95
gen-inline-oz-test-model.get_output_spec_path
def get_output_spec_path(path)
Definition: gen-inline-oz-test-model.py:98
gen-inline-oz-test-model.get_output_signature
def get_output_signature()
Definition: gen-inline-oz-test-model.py:91
gen-inline-oz-test-model.main
def main(argv)
Definition: gen-inline-oz-test-model.py:132
gen-inline-oz-test-model.build_mock_model
def build_mock_model(path, signature)
Definition: gen-inline-oz-test-model.py:102