[MLGO] Add tests for scripts (#78878)

This patch adds integration tests for the script entry points. The tests
don't exercise all functionality, as that case is better covered by the
unit testing already checked in. This ensures that things like flag
parsing work and that the scripts are syntactically valid.

llvm/utils/mlgo-utils/CMakeLists.txt

@@ -5,7 +5,7 @@ configure_lit_site_cfg(
 
 add_lit_testsuite(check-mlgo-utils "Running mlgo-utils tests"
   ${CMAKE_CURRENT_BINARY_DIR}
-  DEPENDS "FileCheck" "not" "count"
+  DEPENDS "FileCheck" "not" "count" "split-file" "yaml2obj" "llvm-objcopy"
 )
 
 set_target_properties(check-mlgo-utils PROPERTIES FOLDER "Tests")

llvm/utils/mlgo-utils/tests/corpus/combine_training_corpus_script.test

@@ -0,0 +1,31 @@
+# REQUIRES: python-38, absl, system-linux
+
+## Testing that the combine_trainig_corpus script works as expected when
+## invoked.
+
+# RUN: rm -rf %t.dir && mkdir %t.dir
+# RUN: split-file %s %t.dir
+# RUN: %python %scripts_dir/corpus/combine_training_corpus.py --root_dir=%t.dir
+# RUN: cat %t.dir/corpus_description.json | FileCheck %s
+
+## Check that we end up with the same properties as the original corpora
+# CHECK: "has_thinlto": false
+
+## Check that the modules end up in the combined corpus. Order does not matter.
+# CHECK-DAG: "subcorpus1/test1.o"
+# CHECK-DAG: "subcorpus2/test2.o"
+
+#--- subcorpus1/corpus_description.json
+{
+  "has_thinlto": false,
+  "modules": [
+    "test1.o"
+  ]
+}
+#--- subcorpus2/corpus_description.json
+{
+  "has_thinlto": false,
+  "modules": [
+    "test2.o"
+  ]
+}

llvm/utils/mlgo-utils/tests/corpus/extract_ir_script.test

@@ -0,0 +1,46 @@
+# REQUIRES: python-38, absl, system-linux
+
+## Test that invoking the extract_ir script work as expected.
+
+# RUN: rm -rf %t.dir && mkdir %t.dir
+# RUN: yaml2obj %s -o %t.dir/test1.o
+# RUN: yaml2obj %s -o %t.dir/test2.o
+# RUN: rm -rf %t.dir.out && mkdir %t.dir.out
+
+# RUN: %python %scripts_dir/corpus/extract_ir.py --input=%t.dir --input_type=directory --output_dir=%t.dir.out --llvm_objcopy_path=llvm-objcopy
+# RUN: cat %t.dir.out/corpus_description.json | FileCheck %s
+
+## Check that this is not a thinLTO build
+# CHECK: "has_thinlto": false
+## Check that the expected modules end up in the corpus description
+# CHECK-DAG: "test1.o"
+# CHECK-DAG: "test2.o"
+
+# RUN: ls %t.dir.out | FileCheck %s --check-prefix CHECK-DIR
+
+# CHECK-DIR: test1.o.bc
+# CHECK-DIR: test1.o.cmd
+# CHECK-DIR: test2.o.bc
+# CHECK-DIR: test2.o.cmd
+
+--- !ELF
+FileHeader:
+  Class:           ELFCLASS64
+  Data:            ELFDATA2LSB
+  Type:            ET_REL
+  Machine:         EM_X86_64
+  SectionHeaderStringTable: .strtab
+Sections:
+  - Name:            .text
+    Type:            SHT_PROGBITS
+    Flags:           [ SHF_ALLOC, SHF_EXECINSTR ]
+    AddressAlign:    0x10
+    Content:         55
+  - Name:            .llvmbc
+    Type:            SHT_PROGBITS
+    AddressAlign:    0x1
+    Content:         55
+  - Name:            .llvmcmd
+    Type:            SHT_PROGBITS
+    AddressAlign:    0x1
+    Content:         ff

llvm/utils/mlgo-utils/tests/corpus/make_corpus_script.test

@@ -0,0 +1,24 @@
+# REQUIRES: python-38, absl, system-linux
+
+## Testing that the make_corpus script works as expected when invoked.
+
+# RUN: rm -rf %t.dir && mkdir %t.dir
+# RUN: touch %t.dir/test1.bc
+# RUN: touch %t.dir/test2.bc
+# RUN: rm -rf %t.out.dir && mkdir %t.out.dir
+
+# RUN: %python %scripts_dir/corpus/make_corpus.py --input_dir=%t.dir --output_dir=%t.out.dir --default_args="-test"
+
+# RUN: cat %t.out.dir/corpus_description.json | FileCheck %s
+
+## Check that we get the expected command in the global command override
+# CHECK: "-test"
+# CHECK: "has_thinlto": false
+## Check that the modules are in the corpus description
+# CHECK: "test1"
+# CHECK: "test2"
+
+# RUN: ls %t.out.dir | FileCheck %s --check-prefix CHECK-DIR
+
+# CHECK-DIR: test1.bc
+# CHECK-DIR: test2.bc

llvm/utils/mlgo-utils/tests/lit.cfg

@@ -1,15 +1,23 @@
+import os
+
 import lit.formats
 
 from lit.llvm import llvm_config
 
 config.name = "mlgo-utils"
 config.test_format = lit.formats.ShTest(execute_external=False)
 
-config.suffixes = [".py"]
+config.suffixes = [".py", ".test"]
 
 config.test_source_root = os.path.dirname(__file__)
 config.test_exec_root = config.obj_root
 
 config.environment["PYTHONPATH"] = os.path.join(config.src_root, "utils", "mlgo-utils")
 
 llvm_config.use_default_substitutions()
+config.substitutions.append(("split-file", llvm_config.use_llvm_tool("split-file")))
+config.substitutions.append(("yaml2obj", llvm_config.use_llvm_tool("yaml2obj")))
+config.substitutions.append(("llvm-objcopy", llvm_config.use_llvm_tool("llvm-objcopy")))
+
+scripts_dir = os.path.join(config.src_root, "utils/mlgo-utils/mlgo")
+config.substitutions.append(("%scripts_dir", scripts_dir))