text: # -*- coding: utf-8 -*-
  """STREAM memory bandwidth benchmarking.
  
  This module handles the downloading, building, and running GCC compiled STREAM
  OMP.
  """
  
  import logging
  import os
  import stat
  import statistics
  import time
  
  from spet.lib.utilities import download
  from spet.lib.utilities import execute
  from spet.lib.utilities import file
  from spet.lib.utilities import grep
  from spet.lib.utilities import optimize
  from spet.lib.utilities import prettify
  
  
  class STREAM:
      """GCC compiled STREAM memory bandwidth benchmarking.
  
      Notes:
          * Requires `gcc` to be installed.
  
      Args:
          version (str): Version number for STREAM.
          root_dir (str): The main directory for SPET.
          results_dir (str): The SPET run's result directory.
  
      Attributes:
          version (str): Version number for STREAM.
          src_dir (str): The source directory for installing packages.
          stream_dir (str): The source directory for the STREAM.
          results_dir (str): The results directory for the STREAM results.
      """
  
      def __init__(self, version, root_dir, results_dir):
          self.version = version
          self.src_dir = root_dir + "/src"
          self.stream_dir = self.src_dir + "/stream"
          self.results_dir = results_dir + "/stream"
          self.commands = []
  
      def download(self):
          """Download STREAM source.
  
          Returns:
              Boolean: True if download was successful otherwise False.
          """
          url = "https://www.cs.virginia.edu/stream/FTP/Code/stream.c"
          target_file = self.stream_dir + "/stream.c"
  
          if os.path.isfile(target_file):
              return True
  
          logging.info("Downloading STREAM.")
  
          os.makedirs(self.stream_dir, exist_ok=True)
          download.file(url, target_file)
  
          if os.path.isfile(target_file):
              return True
          return False
  
      @staticmethod
      def __set_array_size(cache, sockets):
          """Determines how large the STREAM array size should be.
  
          Args:
              cache (int): The cache size in bytes (usually L3).
              sockets (int): The number of sockets on the system.
  
          Returns:
              Integer: STREAM array size
          """
          multiplier = 4
          min_array_size = 10000000
  
          stream_array_size = multiplier * (cache * sockets) / 8
  
          if stream_array_size < min_array_size:
              stream_array_size = min_array_size
  
          return int(stream_array_size)
  
      def build(self, cache, sockets, cflags=None, stream_array_size=None):
          """Compiles STREAM with GCC.
  
          Args:
              cache (int): The cache size in bytes (usually L3).
              sockets (int): The number of sockets on the system.
              cflags (str, optional): The CFLAGS for GCC.
              stream_array_size (int, optional): The array size for STREAM.
  
          Returns:
              Boolean: True if compilation was successful otherwise False.
          """
          if cflags is None:
              cflags = "-march=native -mtune=native"
          if "-O" not in cflags:
              cflags += " -O3 "
  
          cflags += " -fopenmp "
  
          stream_file = self.stream_dir + "/stream.c"
          stream_exe = self.stream_dir + "/stream"
          shell_env = os.environ.copy()
          mpi_root = self.src_dir + "/openmpi/build"
          mpi_path = mpi_root + "/bin"
          mpi_lib = mpi_root + "/lib"
          shell_env["PATH"] += ":" + mpi_path
          if "LD_LIBRARY_PATH" in shell_env:
              shell_env["LD_LIBRARY_PATH"] += mpi_lib
          else:
              shell_env["LD_LIBRARY_PATH"] = mpi_lib
  
          mpicc = mpi_path + "/mpicc"
  
          if os.path.isfile(stream_exe):
              return True
  
          if not os.path.isfile(stream_file):
              prettify.error_message('Cannot compile STREAM because "{}" could '
                                     "not be found.".format(stream_file))
              return False
  
          if stream_array_size is None:
              stream_array_size = self.__set_array_size(cache, sockets)
  
          if stream_array_size > 4000000000:
              cflags += " -mcmodel=medium "
  
          cflags += " -D_OPENMP "
          cflags += " -DSTREAM_ARRAY_SIZE={} ".format(stream_array_size)
          cflags += " -DNTIMES=1000 "
  
          build_cmd = "{} {} stream.c -o stream".format(mpicc, cflags)
  
          logging.info(
              'Compiling STREAM with %d array size, and "%s" CFLAGS.',
              stream_array_size,
              cflags,
          )
  
          self.commands.append("Build: " + build_cmd)
  
          output = execute.output(build_cmd,
                                  working_dir=self.stream_dir,
                                  environment=shell_env)
  
          logging.debug(output)
  
          if os.path.isfile(stream_exe):
              status = os.stat(stream_exe)
              os.chmod(stream_exe, status.st_mode | stat.S_IEXEC)
  
          if os.path.isfile(stream_exe):
              return True
          return False
  
      def run(self, threads):
          """Run GCC compiled STREAM three times.
  
          Args:
              threads (int): The total number of threads on the system.
  
          Returns:
              If success, a dict containing (unit, run1, run2, run3, average,
              median).
  
                  unit (str): Score units.
                  run1 (float): Score for the first run.
                  run2 (float): Score for the second run.
                  run3 (float): Score for the third run.
                  average (float): Average of run1, run2, and run3.
                  median (float): Median of run1, run2, and run3.
  
              Else, a dict containing (error).
  
                  error (str): Error message.
          """
          stream_bin = self.stream_dir + "/stream"
          shell_env = os.environ.copy()
          shell_env["OMP_NUM_THREADS"] = str(threads)
          mpi_root = self.src_dir + "/openmpi/build"
          mpi_path = mpi_root + "/bin"
          mpi_lib = mpi_root + "/lib"
          shell_env["PATH"] += ":" + mpi_path
          if "LD_LIBRARY_PATH" in shell_env:
              shell_env["LD_LIBRARY_PATH"] += mpi_lib
          else:
              shell_env["LD_LIBRARY_PATH"] = mpi_lib
          results = {"unit": "MB/s"}
  
          shell_env["OMP_PROC_BIND"] = "true"
  
          if not os.path.isfile(stream_bin):
              text = 'Cannot run STREAM because "{}" could not be found.'.format(
                  stream_bin)
              prettify.error_message(text)
              return {"error": text}
  
          logging.info("Running STREAM with %d OMP threads.", threads)
  
          os.makedirs(self.results_dir, exist_ok=True)
  
          tmp_results = []
  
          cmd = "./stream"
  
          self.commands.append("Run: OMP_NUM_THREADS = " + str(threads))
          self.commands.append("Run: OMP_PROC_BIND = true")
          self.commands.append("Run: " + cmd)
  
          for count in range(1, 4):
              run_num = "run" + str(count)
              result_file = "{}/stream_{}.txt".format(self.results_dir, run_num)
  
              optimize.prerun()
              time.sleep(10)
  
              output = execute.output(cmd,
                                      working_dir=self.stream_dir,
                                      environment=shell_env)
  
              file.write(result_file, output)
  
              result = grep.text(output, "Triad")
              result = result[0].split()[1]  # 2nd word
              result = float(result)
              results[run_num] = result
              tmp_results.append(result)
  
          results["average"] = statistics.mean(tmp_results)
          results["median"] = statistics.median(tmp_results)
          results["variance"] = statistics.variance(tmp_results)
          sorted_results = sorted(tmp_results)
          results["range"] = sorted_results[-1] - sorted_results[0]
  
          logging.info("STREAM results: %s", str(results))
  
          return results
  