Replace module_paths with corpus_description.json (#65)

* Replace module_paths with metadata.json

- This adds 'global_command_override' functionality for lld-thinlto corpora, exposed in metadata.json
- Setting the above in the metadata.json file will default the runners to using that as the base command, rather than looking for a .cmd file

* Rename to corpus_description.json

* Fix pytype error

* Make UNSPECIFIED a module constant

* Move UNSPECIFIED to constant.py. Resolved comments

Author: Northbadge

compiler_opt/rl/constant.py

@@ -23,6 +23,9 @@
 # Delta to add when computing reward.
 DELTA = 0.01
 
+# Default of global_command_override in corpus_description.json
+UNSPECIFIED_OVERRIDE = ['<UNSPECIFIED>']
+
 
 @gin.constants_from_enum
 class AgentName(enum.Enum):

compiler_opt/rl/corpus.py

@@ -14,11 +14,14 @@
 # limitations under the License.
 """ModuleSpec definition and utility command line parsing functions."""
 
+from absl import logging
 from dataclasses import dataclass
-from typing import Iterable, Optional, Tuple, List
+from typing import List, Dict, Tuple, Any
 
+import json
 import os
-import tensorflow as tf
+
+from compiler_opt.rl import constant
 
 
 @dataclass(frozen=True)
@@ -34,77 +37,91 @@ def build_modulespecs_from_datapath(
     additional_flags: Tuple[str, ...] = (),
     delete_flags: Tuple[str, ...] = ()
 ) -> List[ModuleSpec]:
-  module_paths: List[str] = _load_module_paths(data_path)
-
-  has_thinlto: bool = _has_thinlto_index(module_paths)
+  # TODO: (b/233935329) Per-corpus *fdo profile paths can be read into
+  # {additional|delete}_flags here
+  with open(
+      os.path.join(data_path, 'corpus_description.json'), 'r',
+      encoding='utf-8') as f:
+    corpus_description: Dict[str, Any] = json.load(f)
+
+  module_paths = corpus_description['modules']
+  if len(module_paths) == 0:
+    raise ValueError(f'{data_path}\'s corpus_description contains no modules.')
+
+  has_thinlto: bool = corpus_description['has_thinlto']
+
+  cmd_override = ()
+  if 'global_command_override' in corpus_description:
+    if corpus_description[
+        'global_command_override'] == constant.UNSPECIFIED_OVERRIDE:
+      raise ValueError(
+          'global_command_override in corpus_description.json not filled.')
+    cmd_override = tuple(corpus_description['global_command_override'])
+    if len(additional_flags) > 0:
+      logging.warning('Additional flags are specified together with override.')
+    if len(delete_flags) > 0:
+      logging.warning('Delete flags are specified together with override.')
 
   module_specs: List[ModuleSpec] = []
 
   # This takes ~7s for 30k modules
   for module_path in module_paths:
     exec_cmd = _load_and_parse_command(
-        ir_file=module_path + '.bc',
-        cmd_file=(module_path + '.cmd'),
-        thinlto_file=(module_path + '.thinlto.bc') if has_thinlto else None,
+        module_path=os.path.join(data_path, module_path),
+        has_thinlto=has_thinlto,
         additional_flags=additional_flags,
-        delete_flags=delete_flags)
+        delete_flags=delete_flags,
+        cmd_override=cmd_override)
     module_specs.append(ModuleSpec(name=module_path, exec_cmd=tuple(exec_cmd)))
 
   return module_specs
 
 
-def _has_thinlto_index(module_paths: Iterable[str]) -> bool:
-  return tf.io.gfile.exists(next(iter(module_paths)) + '.thinlto.bc')
-
-
-def _load_module_paths(data_path) -> List[str]:
-  module_paths_path = os.path.join(data_path, 'module_paths')
-  with open(module_paths_path, 'r', encoding='utf-8') as f:
-    ret = [os.path.join(data_path, name.rstrip('\n')) for name in f]
-    if len(ret) == 0:
-      raise ValueError(f'{module_paths_path} is empty.')
-    return ret
-
-
 def _load_and_parse_command(
-    ir_file: str,
-    cmd_file: str,
-    thinlto_file: Optional[str] = None,
+    module_path: str,
+    has_thinlto: bool,
     additional_flags: Tuple[str, ...] = (),
-    delete_flags: Tuple[str, ...] = ()
+    delete_flags: Tuple[str, ...] = (),
+    cmd_override: Tuple[str, ...] = ()
 ) -> List[str]:
   """Cleans up base command line.
 
   Remove certain unnecessary flags, and add the .bc file to compile and, if
   given, the thinlto index.
 
   Args:
-    cmd_file: Path to a .cmd file (from corpus).
-    ir_file: The path to the ir file to compile.
-    thinlto_file: The path to the thinlto index, or None.
+    module_path: Absolute path to the module without extension (from corpus).
+    has_thinlto: Whether to add thinlto flags.
     additional_flags: Tuple of clang flags to add.
     delete_flags: Tuple of clang flags to remove.
+    cmd_override: Tuple of strings to use as the base command line.
 
   Returns:
     The argument list to pass to the compiler process.
   """
   cmdline = []
 
-  with open(cmd_file, encoding='utf-8') as f:
-    option_iterator = iter(f.read().split('\0'))
+  if cmd_override:
+    option_iterator = iter(cmd_override)
+  else:
+    with open(module_path + '.cmd', encoding='utf-8') as f:
+      option_iterator = iter(f.read().split('\0'))
+  option = next(option_iterator, None)
+
+  while option is not None:
+    if any(option.startswith(flag) for flag in delete_flags):
+      if '=' not in option:
+        next(option_iterator, None)
+    else:
+      cmdline.append(option)
     option = next(option_iterator, None)
-    while option is not None:
-      if any(option.startswith(flag) for flag in delete_flags):
-        if '=' not in option:
-          next(option_iterator, None)
-      else:
-        cmdline.append(option)
-      option = next(option_iterator, None)
-  cmdline.extend(['-x', 'ir', ir_file])
-
-  if thinlto_file:
-    cmdline.extend(
-        [f'-fthinlto-index={thinlto_file}', '-mllvm', '-thinlto-assume-merged'])
+  cmdline.extend(['-x', 'ir', module_path + '.bc'])
+
+  if has_thinlto:
+    cmdline.extend([
+        f'-fthinlto-index={module_path}.thinlto.bc', '-mllvm',
+        '-thinlto-assume-merged'
+    ])
 
   cmdline.extend(additional_flags)