llvm/lib/Analysis/models/inlining/config.py

   1 """Inlining Training config."""
   2
   3 import tensorflow as tf
   4
   5 POLICY_DECISION_LABEL = 'inlining_decision'
   6 POLICY_OUTPUT_SPEC = """
   7 [
   8     {
   9         "logging_name": "inlining_decision",
  10         "tensor_spec": {
  11             "name": "StatefulPartitionedCall",
  12             "port": 0,
  13             "type": "int64_t",
  14             "shape": [
  15                 1
  16             ]
  17         }
  18     }
  19 ]
  20 """
  21
  22
  23 # pylint: disable=g-complex-comprehension
  24 def get_input_signature():
  25   """Returns the list of features for LLVM inlining."""
  26   # int64 features
  27   inputs = [
  28       tf.TensorSpec(dtype=tf.int64, shape=(), name=key) for key in [
  29           'caller_basic_block_count',
  30           'caller_conditionally_executed_blocks',
  31           'caller_users',
  32           'callee_basic_block_count',
  33           'callee_conditionally_executed_blocks',
  34           'callee_users',
  35           'nr_ctant_params',
  36           'node_count',
  37           'edge_count',
  38           'callsite_height',
  39           'cost_estimate',
  40           'inlining_default',
  41           'sroa_savings',
  42           'sroa_losses',
  43           'load_elimination',
  44           'call_penalty',
  45           'call_argument_setup',
  46           'load_relative_intrinsic',
  47           'lowered_call_arg_setup',
  48           'indirect_call_penalty',
  49           'jump_table_penalty',
  50           'case_cluster_penalty',
  51           'switch_penalty',
  52           'unsimplified_common_instructions',
  53           'num_loops',
  54           'dead_blocks',
  55           'simplified_instructions',
  56           'constant_args',
  57           'constant_offset_ptr_args',
  58           'callsite_cost',
  59           'cold_cc_penalty',
  60           'last_call_to_static_bonus',
  61           'is_multiple_blocks',
  62           'nested_inlines',
  63           'nested_inline_cost_estimate',
  64           'threshold',
  65       ]
  66   ]
  67
  68   # float32 features
  69   inputs.extend([
  70       tf.TensorSpec(dtype=tf.float32, shape=(), name=key)
  71       for key in ['discount', 'reward']
  72   ])
  73
  74   # int32 features
  75   inputs.extend([
  76       tf.TensorSpec(dtype=tf.int32, shape=(), name=key)
  77       for key in ['step_type']
  78   ])
  79   return inputs
  80
  81
  82 def get_output_signature():
  83   return POLICY_DECISION_LABEL
  84
  85
  86 def get_output_spec():
  87   return POLICY_OUTPUT_SPEC