Fix use after move to get exegesis annotator working in pipeline

gematria/datasets/bhive_to_exegesis.cc

@@ -47,7 +47,7 @@ namespace gematria {
 
 BHiveToExegesis::BHiveToExegesis(
     LlvmArchitectureSupport& ArchitectureSupport,
-    llvm::exegesis::LLVMState&& LLVMExegesisState,
+    std::unique_ptr<llvm::exegesis::LLVMState>&& LLVMExegesisState,
     std::unique_ptr<ExegesisAnnotator>&& LLVMExegesisAnnotator)
     : LLVMAnnotator(std::move(LLVMExegesisAnnotator)),
       ExegesisState(std::move(LLVMExegesisState)),
@@ -60,13 +60,15 @@ Expected<std::unique_ptr<BHiveToExegesis>> BHiveToExegesis::create(
     LlvmArchitectureSupport& ArchitectureSupport) {
   Expected<LLVMState> LLVMStateOrErr = LLVMState::Create("", "native");
   if (!LLVMStateOrErr) return LLVMStateOrErr.takeError();
+  std::unique_ptr<LLVMState> LLVMStateOwner(
+      new LLVMState(std::move(*LLVMStateOrErr)));
 
   Expected<std::unique_ptr<ExegesisAnnotator>> AnnotatorOrErr =
-      ExegesisAnnotator::create(*LLVMStateOrErr);
+      ExegesisAnnotator::create(*LLVMStateOwner);
   if (!AnnotatorOrErr) return AnnotatorOrErr.takeError();
 
   return std::unique_ptr<BHiveToExegesis>(
-      new BHiveToExegesis(ArchitectureSupport, std::move(*LLVMStateOrErr),
+      new BHiveToExegesis(ArchitectureSupport, std::move(LLVMStateOwner),
                           std::move(*AnnotatorOrErr)));
 }
 

gematria/datasets/bhive_to_exegesis.h

@@ -49,15 +49,15 @@ class BHiveToExegesis {
 
  private:
   std::unique_ptr<ExegesisAnnotator> LLVMAnnotator;
-  llvm::exegesis::LLVMState ExegesisState;
+  std::unique_ptr<llvm::exegesis::LLVMState> ExegesisState;
   gematria::X86Canonicalizer Canonicalizer;
   gematria::BHiveImporter BHiveImporter;
   LlvmArchitectureSupport &ArchSupport;
   std::unique_ptr<llvm::MCInstPrinter> InstPrinter;
 
   BHiveToExegesis(
       LlvmArchitectureSupport &ArchitectureSupport,
-      llvm::exegesis::LLVMState &&LLVMExegesisState,
+      std::unique_ptr<llvm::exegesis::LLVMState> &&LLVMExegesisState,
       std::unique_ptr<gematria::ExegesisAnnotator> &&LLVMExegesisAnnotator);
 
   absl::StatusOr<ExecutionAnnotations> getAccessedAddrs(

gematria/datasets/pipelines/BUILD.bazel

@@ -40,6 +40,9 @@ gematria_py_test(
     env = {
         "PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION": "python",
     },
+    tags = [
+        "perf_counters",
+    ],
     deps = [
         ":compile_modules_lib",
         "//gematria/testing/python:ir_utils",

gematria/datasets/pipelines/compile_modules.py

@@ -47,7 +47,10 @@
 )
 
 _ANNOTATOR_TYPE = flags.DEFINE_enum(
-    'annotator_type', 'fast', ['fast'], 'The type of annotator to use.'
+    'annotator_type',
+    'fast',
+    ['fast', 'exegesis'],
+    'The type of annotator to use.',
 )
 
 _MAX_ANNOTATION_ATTEMPTS = flags.DEFINE_integer(
@@ -56,9 +59,10 @@
     'The maximum number of times to try annotating a block before giving up',
 )
 
-# TODO(boomanaiden154): Currently, only the fast annotator works. Eventually
-# this should be fixed so we can use the exegesis annotator too.
-ANNOTATOR_MAPPING = {'fast': bhive_to_exegesis.AnnotatorType.fast}
+ANNOTATOR_MAPPING = {
+    'fast': bhive_to_exegesis.AnnotatorType.fast,
+    'exegesis': bhive_to_exegesis.AnnotatorType.exegesis,
+}
 
 
 def main(argv) -> None:

gematria/datasets/pipelines/compile_modules_lib_test.py

@@ -16,6 +16,7 @@
 import textwrap
 
 from absl.testing import absltest
+from absl.testing import parameterized
 import apache_beam as beam
 from apache_beam.testing import test_pipeline
 from apache_beam.testing import util as beam_test
@@ -27,7 +28,7 @@
 from gematria.io.python import tfrecord
 
 
-class CompileModulesTests(absltest.TestCase):
+class CompileModulesTests(parameterized.TestCase):
 
   def test_optimize_modules(self):
     ir_string = textwrap.dedent("""\
@@ -88,10 +89,12 @@ def test_deduplicate_values(self):
       output = input | compile_modules_lib.DeduplicateValues()
       beam_test.assert_that(output, beam_test.equal_to(['aa', 'ab', 'bc']))
 
-  def test_annotate_bbs(self):
-    annotator = compile_modules_lib.AnnotateBBs(
-        bhive_to_exegesis.AnnotatorType.fast, 50
-    )
+  @parameterized.parameters([
+      bhive_to_exegesis.AnnotatorType.fast,
+      bhive_to_exegesis.AnnotatorType.exegesis,
+  ])
+  def test_annotate_bbs(self, annotator_type):
+    annotator = compile_modules_lib.AnnotateBBs(annotator_type, 50)
     annotator.setup()
 
     annotated_blocks = list(
@@ -150,7 +153,11 @@ def test_process_and_filter_bbs(self):
     filtered_bbs = list(process_and_filter_transform.process(bb_hex))
     self.assertLen(filtered_bbs, 0)
 
-  def test_get_bbs(self):
+  @parameterized.parameters([
+      bhive_to_exegesis.AnnotatorType.fast,
+      bhive_to_exegesis.AnnotatorType.exegesis,
+  ])
+  def test_get_bbs(self, annotator_type):
     ir_string1 = textwrap.dedent("""\
       define i32 @a() {
         ret i32 1
@@ -175,7 +182,7 @@ def test_get_bbs(self):
         test_parquet_file.full_path,
         output_file_pattern,
         False,
-        bhive_to_exegesis.AnnotatorType.fast,
+        annotator_type,
         50,
         vocab_output_file_pattern,
     )