1# Copyright 2015 The Chromium OS Authors. All rights reserved.
2# Use of this source code is governed by a BSD-style license
3# that can be found in the LICENSE file.
4
5import argparse
6import logging
7import mmap
8import os
9import signal
10import struct
11import sys
12import threading
13import time
14
15# some magic numbers: see http://goo.gl/ecAgke for Intel docs
16PCI_IMC_BAR_OFFSET = 0x48
17IMC_DRAM_GT_REQUESTS = 0x5040 # GPU
18IMC_DRAM_IA_REQUESTS = 0x5044 # CPU
19IMC_DRAM_IO_REQUESTS = 0x5048 # PCIe, Display Engine, USB, etc.
20IMC_DRAM_DATA_READS = 0x5050  # read traffic
21IMC_DRAM_DATA_WRITES = 0x5054 # write traffic
22IMC_MMAP_SIZE = 0x6000
23
24CACHE_LINE = 64.0
25MEGABYTE = 1048576.0
26
27RATE_FIELD_FORMAT = '%s: %5d MB/s'
28RAW_FIELD_FORMAT = '%s: %d'
29
30class IMCCounter:
31    """Small struct-like class to keep track of the
32    location and attributes for each counter.
33
34    Parameters:
35      name: short, unique identifying token for this
36        counter type
37      idx: offset into the IMC memory where we can find
38        this counter
39      total: True if we should count this in the number
40        for total bandwidth
41    """
42    def __init__(self, name, idx, total):
43        self.name = name
44        self.idx = idx
45        self.total = total
46
47
48counters = [
49#              name          idx           total
50    IMCCounter("GT", IMC_DRAM_GT_REQUESTS, False),
51    IMCCounter("IA", IMC_DRAM_IA_REQUESTS, False),
52    IMCCounter("IO", IMC_DRAM_IO_REQUESTS, False),
53    IMCCounter("RD", IMC_DRAM_DATA_READS,  True),
54    IMCCounter("WR", IMC_DRAM_DATA_WRITES, True),
55]
56
57
58class MappedFile:
59    """Helper class to wrap mmap calls in a context
60    manager so they are always cleaned up, and to
61    help extract values from the bytes.
62
63    Parameters:
64      filename: name of file to mmap
65      offset: offset from beginning of file to mmap
66        from
67      size: amount of the file to mmap
68    """
69    def __init__(self, filename, offset, size):
70        self._filename = filename
71        self._offset = offset
72        self._size = size
73
74
75    def __enter__(self):
76        self._f = open(self._filename, 'rb')
77        try:
78            self._mm = mmap.mmap(self._f.fileno(),
79                                 self._size,
80                                 mmap.MAP_SHARED,
81                                 mmap.PROT_READ,
82                                 offset=self._offset)
83        except mmap.error:
84            self._f.close()
85            raise
86        return self
87
88
89    def __exit__(self, exc_type, exc_val, exc_tb):
90        self._mm.close()
91        self._f.close()
92
93
94    def bytes_to_python(self, offset, fmt):
95        """Grab a portion of an mmapped file and return the bytes
96        as a python object.
97
98        Parameters:
99          offset: offset into the mmapped file to start at
100          fmt: string containing the struct type to extract from the
101            file
102        Returns: a Struct containing the bytes starting at offset
103          into the mmapped file, reified as python values
104        """
105        s = struct.Struct(fmt)
106        return s.unpack(self._mm[offset:offset+s.size])
107
108
109def file_bytes_to_python(f, offset, fmt):
110    """Grab a portion of a regular file and return the bytes
111    as a python object.
112
113    Parameters:
114      f: file-like object to extract from
115      offset: offset into the mmapped file to start at
116      fmt: string containing the struct type to extract from the
117        file
118    Returns: a Struct containing the bytes starting at offset into
119      f, reified as python values
120    """
121    s = struct.Struct(fmt)
122    f.seek(0)
123    bs = f.read()
124    if len(bs) >= offset + s.size:
125        return s.unpack(bs[offset:offset+s.size])
126    else:
127        raise IOError('Invalid seek in file')
128
129
130def uint32_diff(l, r):
131    """Compute the difference of two 32-bit numbers as
132    another 32-bit number.
133
134    Since the counters are monotonically increasing, we
135    always want the unsigned difference.
136    """
137    return l - r if l >= r else l - r + 0x100000000
138
139
140class MemoryBandwidthLogger(threading.Thread):
141    """Class for gathering memory usage in MB/s on x86 systems.
142    raw: dump raw counter values
143    seconds_period: time period between reads
144
145    If you are using non-raw mode and your seconds_period is
146    too high, your results might be nonsense because the counters
147    might have wrapped around.
148
149    Parameters:
150      raw: True if you want to dump raw counters. These will simply
151        tell you the number of cache-line-size transactions that
152        have occurred so far.
153      seconds_period: Duration to wait before dumping counters again.
154        Defaults to 2 seconds.
155      """
156    def __init__(self, raw, seconds_period=2):
157        super(MemoryBandwidthLogger, self).__init__()
158        self._raw = raw
159        self._seconds_period = seconds_period
160        self._running = True
161
162
163    def run(self):
164        # get base address register and align to 4k
165        try:
166            bar_addr = self._get_pci_imc_bar()
167        except IOError:
168            logging.error('Cannot read base address register')
169            return
170        bar_addr = (bar_addr // 4096) * 4096
171
172        # set up the output formatting. raw counters don't have any
173        # particular meaning in MB/s since they count how many cache
174        # lines have been read from or written to up to that point,
175        # and so don't represent a rate.
176        # TOTAL is always given as a rate, though.
177        rate_factor = CACHE_LINE / (self._seconds_period * MEGABYTE)
178        if self._raw:
179            field_format = RAW_FIELD_FORMAT
180        else:
181            field_format = RATE_FIELD_FORMAT
182
183        # get /dev/mem and mmap it
184        with MappedFile('/dev/mem', bar_addr, IMC_MMAP_SIZE) as mm:
185            # take initial samples, then take samples every seconds_period
186            last_values = self._take_samples(mm)
187            while self._running:
188                time.sleep(self._seconds_period)
189                values = self._take_samples(mm)
190                # we need to calculate the MB differences no matter what
191                # because the "total" field uses it even when we are in
192                # raw mode
193                mb_diff = { c.name:
194                    uint32_diff(values[c.name], last_values[c.name])
195                        * rate_factor for c in counters }
196                output_dict = values if self._raw else mb_diff
197                output = list((c.name, output_dict[c.name]) for c in counters)
198
199                total_rate = sum(mb_diff[c.name] for c in counters if c.total)
200                output_str = \
201                    ' '.join(field_format % (k, v) for k, v in output) + \
202                    ' ' + (RATE_FIELD_FORMAT % ('TOTAL', total_rate))
203
204                logging.debug(output_str)
205                last_values = values
206
207
208    def stop(self):
209        self._running = False
210
211
212    def _get_pci_imc_bar(self):
213        """Get the base address register for the IMC (integrated
214        memory controller). This is later used to extract counter
215        values.
216
217        Returns: physical address for the IMC.
218        """
219        with open('/proc/bus/pci/00/00.0', 'rb') as pci:
220            return file_bytes_to_python(pci, PCI_IMC_BAR_OFFSET, '=Q')[0]
221
222
223    def _take_samples(self, mm):
224        """Get samples for each type of memory transaction.
225
226        Parameters:
227          mm: MappedFile representing physical memory
228        Returns: dictionary mapping counter type to counter value
229        """
230        return { c.name: mm.bytes_to_python(c.idx, '=I')[0]
231            for c in counters }
232