Use perf-annotate instead of perf-record.

Author: paul-tqh-nguyen

mlir_graphblas/engine.py

@@ -541,15 +541,32 @@ def __init__(self, cli_executable=None, cli_options=None, llvmlite_engine=None):
             backing_mod = llvm.parse_assembly("")
             llvmlite_engine = llvm.create_mcjit_compiler(backing_mod, target_machine)
         self._engine = llvmlite_engine
-        self.tmp_dir = tempfile.TemporaryDirectory()
+        self.default_profile_dir = tempfile.TemporaryDirectory()
+        self.current_profile_dir: Optional[str] = None
         self.c_compiler = distutils.ccompiler.new_compiler()
         self._cli = MlirOptCli(cli_executable, cli_options)
         self.name_to_callable: Dict[str, Callable] = {}
         return
 
-    def profiled_function(self, main_callable: Callable) -> Callable:
+    @property
+    def profile_dir_name(self) -> str:
+        if self.current_profile_dir is not None:
+            # TODO consider making a context manager for setting self.current_profile_dir
+            return self.current_profile_dir
+        return self.default_profile_dir.name
+
+    def profiled_function(
+        self, main_callable: Callable, symbol_to_profile: str
+    ) -> Callable:
         """Decorator to profile a function via Linux's perf tool."""
 
+        # set this at the time that decorated_func is created and
+        # not at the time it is called as the value of
+        # self.profile_dir_name may change.
+        perf_data_file_name = os.path.join(
+            self.profile_dir_name, f"perf-{uuid.uuid4()}.data"
+        )
+
         def mp_func(queue: mp.SimpleQueue, *args):
             # send pid to decorated_func
             queue.put(os.getpid())
@@ -576,12 +593,11 @@ def decorated_func(*args) -> Any:
                 stdout=subprocess.DEVNULL,
                 stderr=subprocess.PIPE,
             )
-            record_process.stdin.write(
-                f"""
-            perf record -p {execution_process_id}
+            record_command = f"""
+            perf record -p {execution_process_id} --output={perf_data_file_name}
             exit
-            """.encode()
-            )
+            """
+            record_process.stdin.write(record_command.encode())
             record_process.stdin.flush()
 
             # wait for profiling to initialize
@@ -600,14 +616,14 @@ def decorated_func(*args) -> Any:
                 pass
 
             # gather profiling results
-            report_command = f"perf report --pid={execution_process_id} | cat"
-            report_process = subprocess.Popen(
-                report_command,
+            annotate_command = f"perf annotate --stdio --symbol {symbol_to_profile} -l --input={perf_data_file_name}"
+            annotate_process = subprocess.Popen(
+                annotate_command,
                 shell=True,
                 stdout=subprocess.PIPE,
                 stderr=subprocess.STDOUT,
             )  # combines stdout and stderr
-            stdout_string, _ = report_process.communicate()
+            stdout_string, _ = annotate_process.communicate()
             stdout_string = stdout_string.decode()
 
             # print results
@@ -645,14 +661,14 @@ def _add_llvm_module(
             assert len(o_file_bytes_list) in (1, 2)
             o_file_bytes = o_file_bytes_list[-1]
             o_file_name = f"mod-{uuid.uuid4()}.o"
-            o_file_name = os.path.join(self.tmp_dir.name, o_file_name)
+            o_file_name = os.path.join(self.profile_dir_name, o_file_name)
             with open(o_file_name, "wb") as f:
                 f.write(o_file_bytes_list[-1])
             files_to_link.append(o_file_name)
             if isinstance(profile, ctypes.CDLL):
                 files_to_link.append(profile._name)
             so_file_name = f"shared-{uuid.uuid4()}.so"
-            so_file_name = os.path.join(self.tmp_dir.name, so_file_name)
+            so_file_name = os.path.join(self.profile_dir_name, so_file_name)
             self.c_compiler.link_shared_object(files_to_link, so_file_name)
             ctypes.cdll.LoadLibrary(so_file_name)
             shared_lib = ctypes.CDLL(so_file_name)
@@ -753,7 +769,7 @@ def python_callable(mlir_function, encoders, c_callable, decoder, *args):
                 python_callable, mlir_function, encoders, c_callable, decoder
             )
             if shared_lib is not None:
-                bound_func = self.profiled_function(bound_func)
+                bound_func = self.profiled_function(bound_func, name)
             name_to_callable[name] = bound_func
 
         return name_to_callable
@@ -791,7 +807,7 @@ def _lower_types_to_strings(
 
     def _generate_mlir_string_for_multivalued_functions(
         self, mlir_functions: Iterable[mlir.astnodes.Function], passes: List[str]
-    ) -> Tuple[str, str]:
+    ) -> Tuple[str, List[str], List[str]]:
 
         result_type_name_to_lowered_result_type_name = self._lower_types_to_strings(
             sum((mlir_function.result_types for mlir_function in mlir_functions), []),
@@ -800,9 +816,8 @@ def _generate_mlir_string_for_multivalued_functions(
 
         # Generate conglomerate MLIR string for all wrappers
         mlir_wrapper_texts: List[str] = []
-        wrapper_names = [
-            mlir_function.name.value + "wrapper" for mlir_function in mlir_functions
-        ]
+        names = [mlir_function.name.value for mlir_function in mlir_functions]
+        wrapper_names = [name + "wrapper" for name in names]
         for mlir_function, wrapper_name in zip(mlir_functions, wrapper_names):
             lowered_result_type_names = [
                 result_type_name_to_lowered_result_type_name[result_type.dump()]
@@ -860,7 +875,7 @@ def _generate_mlir_string_for_multivalued_functions(
             mlir_wrapper_texts.append(mlir_wrapper_text)
 
         mlir_text = "\n".join(mlir_wrapper_texts)
-        return mlir_text, wrapper_names
+        return mlir_text, names, wrapper_names
 
     def _generate_multivalued_functions(
         self,
@@ -870,9 +885,11 @@ def _generate_multivalued_functions(
     ) -> Dict[str, Callable]:
         name_to_callable: Dict[str, Callable] = {}
 
-        mlir_text, wrapper_names = self._generate_mlir_string_for_multivalued_functions(
-            mlir_functions, passes
-        )
+        (
+            mlir_text,
+            names,
+            wrapper_names,
+        ) = self._generate_mlir_string_for_multivalued_functions(mlir_functions, passes)
 
         # this is guaranteed to not raise exceptions since the user-provided
         # code was already added (failures would occur then)
@@ -882,7 +899,9 @@ def _generate_multivalued_functions(
         assert bool(wrapper_shared_lib) == bool(internal_shared_lib)
 
         # Generate callables
-        for mlir_function, wrapper_name in zip(mlir_functions, wrapper_names):
+        for mlir_function, name, wrapper_name in zip(
+            mlir_functions, names, wrapper_names
+        ):
             ctypes_input_types, input_encoders = mlir_function_input_encoders(
                 mlir_function
             )
@@ -954,7 +973,7 @@ def python_callable(
                 decoders,
             )
             if wrapper_shared_lib is not None:
-                bound_func = self.profiled_function(bound_func)
+                bound_func = self.profiled_function(bound_func, name)
             name_to_callable[mlir_function.name.value] = bound_func
 
         return name_to_callable
@@ -975,10 +994,18 @@ def add(
         mlir_text: Union[str, bytes],
         passes: Tuple[str],
         *,
-        debug=False,
-        profile=False,
+        debug: bool = False,
+        profile: bool = False,
+        profile_result_directory: Optional[str] = None,
     ) -> Union[List[str], DebugResult]:
         """List of new function names added."""
+        if profile_result_directory is not None:
+            if not profile:
+                raise ValueError(
+                    "Cannot specify a profile result directory without also enabling profiling."
+                )
+            self.current_profile_dir = profile_result_directory
+
         if profile:
             if not sys.platform.startswith("linux"):
                 raise NotImplementedError("Profiling only supported on linux.")
@@ -1054,6 +1081,9 @@ def add(
 
             self.name_to_callable[name] = python_callable
 
+        if profile_result_directory is not None:
+            self.current_profile_dir = None
+
         return function_names
 
     def __getitem__(self, func_name: str) -> Callable:

mlir_graphblas/tests/test_profile.py

@@ -1,5 +1,8 @@
 import sys
+import os
 import pytest
+import uuid
+import tempfile
 import numpy as np
 
 from io import StringIO
@@ -35,18 +38,20 @@ def captured_stdout() -> Generator[List[str], None, None]:
 def test_profile_single_return(
     engine: MlirJitEngine,
 ):
-    mlir_text = """
-func @single_value_slow_mul(%input_val: i32) -> i32 {
-  %ci0 = constant 0 : i32
+    num_iterations = 12345678
+    function_name = f"single_value_slow_mul_{int(uuid.uuid4())}"
+    mlir_text = f"""
+func @{function_name}(%input_val: i64) -> i64 {{
+  %ci0 = constant 0 : i64
   %c0 = constant 0 : index
   %c1 = constant 1 : index
-  %num_iterations = constant 99 : index
-  %answer = scf.for %i = %c0 to %num_iterations step %c1 iter_args(%current_sum=%ci0) -> (i32) {
-    %new_sum = addi %current_sum, %input_val : i32
-    scf.yield %new_sum : i32
-  }
-  return %answer : i32
-}
+  %num_iterations = constant {num_iterations} : index
+  %answer = scf.for %i = %c0 to %num_iterations step %c1 iter_args(%current_sum=%ci0) -> (i64) {{
+    %new_sum = addi %current_sum, %input_val : i64
+    scf.yield %new_sum : i64
+  }}
+  return %answer : i64
+}}
 """
     try:
         engine.add(mlir_text, STANDARD_PASSES, profile=True)
@@ -62,42 +67,47 @@ def test_profile_single_return(
             pytest.skip(f"Profiling not permitted.")
         else:
             raise e
-    compiled_func = getattr(engine, "single_value_slow_mul")
+    compiled_func = getattr(engine, function_name)
     input_val = 12
 
     with captured_stdout() as stdout_string_container:
         result = compiled_func(input_val)
 
-    assert result == input_val * 99
+    assert result == input_val * num_iterations
 
     (stdout_string,) = stdout_string_container
-    assert "Overhead" in stdout_string
-    assert "Command" in stdout_string
-    assert "Shared Object" in stdout_string
-    assert "Symbol" in stdout_string
-    assert "python3" in stdout_string
-
+    assert "Source code & Disassembly of" in stdout_string
+    assert "Sorted summary for file" in stdout_string
+    assert "Percent" in stdout_string
     return
 
 
 def test_profile_multiple_returns(
     engine: MlirJitEngine,
 ):
-    mlir_text = """
-func @multiple_value_slow_mul(%input_val: i32) -> (i32, i32) {
-  %ci0 = constant 0 : i32
+    result_dir = tempfile.TemporaryDirectory()
+    num_iterations = 12345678
+    function_name = f"multivalue_slow_mul_{int(uuid.uuid4())}"
+    mlir_text = f"""
+func @{function_name}(%input_val: i64) -> (i64, i64) {{
+  %ci0 = constant 0 : i64
   %c0 = constant 0 : index
   %c1 = constant 1 : index
-  %num_iterations = constant 99 : index
-  %answer = scf.for %i = %c0 to %num_iterations step %c1 iter_args(%current_sum=%ci0) -> (i32) {
-    %new_sum = addi %current_sum, %input_val : i32
-    scf.yield %new_sum : i32
-  }
-  return %input_val, %answer : i32, i32
-}
+  %num_iterations = constant {num_iterations} : index
+  %answer = scf.for %i = %c0 to %num_iterations step %c1 iter_args(%current_sum=%ci0) -> (i64) {{
+    %new_sum = addi %current_sum, %input_val : i64
+    scf.yield %new_sum : i64
+  }}
+  return %input_val, %answer : i64, i64
+}}
 """
     try:
-        engine.add(mlir_text, STANDARD_PASSES, profile=True)
+        engine.add(
+            mlir_text,
+            STANDARD_PASSES,
+            profile=True,
+            profile_result_directory=result_dir.name,
+        )
     except NotImplementedError as e:
         if not sys.platform.startswith("linux"):
             pytest.skip(f"Profiling not supported on {sys.platform}.")
@@ -110,21 +120,30 @@ def test_profile_multiple_returns(
             pytest.skip(f"Profiling not permitted.")
         else:
             raise e
-    compiled_func = getattr(engine, "multiple_value_slow_mul")
+    compiled_func = getattr(engine, function_name)
     input_val = 12
 
     with captured_stdout() as stdout_string_container:
         redundant_input_val, result = compiled_func(input_val)
 
     assert redundant_input_val == input_val
-    assert result == input_val * 99
+    assert result == input_val * num_iterations
 
     (stdout_string,) = stdout_string_container
-    assert "Overhead" in stdout_string
-    assert "Command" in stdout_string
-    assert "Shared Object" in stdout_string
-    assert "Symbol" in stdout_string
-    assert "python3" in stdout_string
+    assert "Source code & Disassembly of" in stdout_string
+    assert "Sorted summary for file" in stdout_string
+    assert "Percent" in stdout_string
+
+    assert (
+        len(
+            [
+                file_name
+                for file_name in os.listdir(result_dir.name)
+                if file_name.startswith("perf-") and file_name.endswith(".data")
+            ]
+        )
+        == 1
+    )
     return