experimental/dispatch_profiler/performance_report.py

# Copyright 2023 The IREE Authors
#
# Licensed under the Apache License v2.0 with LLVM Exceptions.
# See https://llvm.org/LICENSE.txt for license information.
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception

import csv, textwrap
import numpy as np
from collections import namedtuple
from pathlib import Path


class PerformanceResult:
  """Performance result of a single run."""

  def __init__(self, operation, configuration, verification_result, runtime):
    self.operation = operation
    self.configuration = configuration
    self.verification_result = verification_result
    self.runtime = runtime  # in milliseconds
    self.gflops = float(self.operation.flops()) / self.runtime / 1.0e6

  def print(self):
    """Prints the performance result to the console."""
    runtime = (str(self.runtime) if self.runtime != -1.0 else 'Not profiled')
    gflops = (str(float(round(self.gflops, 2)))
              if self.runtime != -1.0 else 'Not profiled')

    print('---------------------------------------------------------------- ')
    print(
        f'Dispatch      : {"_".join([self.operation.name(), self.configuration.name()])}'
    )
    print(f'Provider      : IREE Codegen')
    print(f'OpKind        : {self.operation.operation_kind}')
    print(f'Operation     : {self.operation.name()}')
    print(f'Configuration : {self.configuration.name()}')
    # Operation specific arguments.
    arg_str = ' '.join([
        f'--{key}={value}'
        for key, value in self.operation.get_argument_dict().items()
    ])
    wrapped_arg_str = textwrap.fill(arg_str,
                                    width=80,
                                    subsequent_indent='                ')
    print(f'Arguments     : {wrapped_arg_str}')
    print(f'Verification  : {self.verification_result}')
    print(f'Runtime(ms)   : {runtime}')
    print(f'GFLOPs        : {gflops}')

  def get_dict_entry(self):
    """Returns a dictionary with the performance result."""
    runtime = self.runtime if self.runtime != -1.0 else ''
    gflops = (float(round(self.gflops, 2))
              if self.runtime != -1.0 else 'Not run')
    dict_entry = {
        'Provider': 'IREE Codegen',
        'Verification': self.verification_result,
        'Runtime(ms)': runtime,
        'GFLOPs': gflops,
    }

    # Add the operation specific arguments.
    dict_entry.update(self.operation.get_dict_entry())

    # Add the configuration specific arguments.
    dict_entry.update(self.configuration.get_dict_entry())

    return dict_entry


class PerformanceReport:
  """Performance report class is used to store the performance results of multiple runs.
  The report can be written to a csv file."""

  def __init__(self, args):
    self.args = args

    # Data members extracted from the args.
    self.output_file_path = None
    if args.output != '':
      self.output_file_path = Path(args.output)

    # List of PerformanceResult.
    self.perf_result_vector = []

    # Additional tags to add to the csv report file. \
    # Useful for generating pivot tables.
    self.tags = []
    if args.tags != '':
      self.tags = args.tags.split(',')

    # Boolen to check if the header is written to the csv file.
    self.is_header_written = False

    # If the args.output set, open the file and write the header.
    self.open_mode = 'a' if self.args.append else 'w'
    if self.output_file_path:
      self.csv_file = open(self.output_file_path, self.open_mode)

  def __del__(self):
    """If the args.output set, close the file."""
    if self.output_file_path:
      print('Writing performance report to %s' % self.output_file_path)
      self.csv_file.close()

  def write_csv_header(self, operation, configuration):
    """Write the header to the csv file."""

    # Create and write the header.
    operation_specific_header = list(operation.get_dict_entry().keys())
    configuration_specific_header = list(configuration.get_dict_entry().keys())
    performance_header = ['Verification', 'Runtime(ms)', 'GFLOPs']
    csv_header = operation_specific_header + configuration_specific_header + performance_header
    csv_header = ['Provider'] + csv_header

    # If tags are present, add the tags.keys() to the begining of the csv header.
    if len(self.tags):
      tag_header = [tag.split(':')[0] for tag in self.tags]
      csv_header = tag_header + csv_header

    # Create the csv dictionary writer.
    self.csv_writer = csv.DictWriter(self.csv_file, fieldnames=csv_header)

    # Write the header if the file is being created.
    if self.open_mode == 'w':
      self.csv_writer.writeheader()

  def append_perf_result(self, performance_result):
    """Appends a performance result to the report. 
    Additionaly, if args.output set, write the csv_row entry."""
    self.perf_result_vector.append(performance_result)

    if self.output_file_path:
      # Write the header if not written.
      if not self.is_header_written:
        self.write_csv_header(performance_result.operation,
                              performance_result.configuration)
        self.is_header_written = True

      # Create the row entries for performance result.
      csv_dict_row = performance_result.get_dict_entry()

      # Create the row entries for tags.
      for tag in self.tags:
        tag_key, tag_value = tag.split(':')
        csv_dict_row[tag_key] = tag_value

      # Write the row.
      self.csv_writer.writerow(csv_dict_row)