1
2# Copyright (C) 2003  CAMP
3# Please see the accompanying LICENSE file for further information.
4
5import sys
6import time
7import math
8
9import numpy as np
10from ase.utils.timing import Timer
11
12import gpaw.mpi as mpi
13
14
15class NullTimer:
16    """Compatible with Timer and StepTimer interfaces.  Does nothing."""
17    def __init__(self):
18        pass
19
20    def print_info(self, calc):
21        pass
22
23    def start(self, name):
24        pass
25
26    def stop(self, name=None):
27        pass
28
29    def get_time(self, name):
30        return 0.0
31
32    def write(self, out=sys.stdout):
33        pass
34
35    def write_now(self, mark=''):
36        pass
37
38    def add(self, timer):
39        pass
40
41    def __call__(self, name):
42        return self
43
44    def __enter__(self):
45        pass
46
47    def __exit__(self, *args):
48        pass
49
50
51nulltimer = NullTimer()
52
53
54class DebugTimer(Timer):
55    def __init__(self, print_levels=1000, comm=mpi.world, txt=sys.stdout):
56        Timer.__init__(self, print_levels)
57        ndigits = 1 + int(math.log10(comm.size))
58        self.srank = '%0*d' % (ndigits, comm.rank)
59        self.txt = txt
60
61    def start(self, name):
62        Timer.start(self, name)
63        abstime = time.time()
64        t = self.timers[tuple(self.running)] + abstime
65        self.txt.write('T%s >> %15.8f %s (%7.5fs) started\n'
66                       % (self.srank, abstime, name, t))
67
68    def stop(self, name=None):
69        if name is None:
70            name = self.running[-1]
71        abstime = time.time()
72        t = self.timers[tuple(self.running)] + abstime
73        self.txt.write('T%s << %15.8f %s (%7.5fs) stopped\n'
74                       % (self.srank, abstime, name, t))
75        Timer.stop(self, name)
76
77
78class ParallelTimer(DebugTimer):
79    """Like DebugTimer but writes timings from all ranks.
80
81    Each rank writes to timings.<rank>.txt.  Also timings.metadata.txt
82    will contain information about the parallelization layout.  The idea
83    is that the output from this timer can be used for plots and to
84    determine bottlenecks in the parallelization.
85
86    See the tool gpaw-plot-parallel-timings."""
87    def __init__(self, prefix='timings', flush=False):
88        ndigits = len(str(mpi.world.size - 1))
89        ranktxt = '%0*d' % (ndigits, mpi.world.rank)
90        fname = '%s.%s.txt' % (prefix, ranktxt)
91        txt = open(fname, 'w', buffering=1 if flush else -1)
92        DebugTimer.__init__(self, comm=mpi.world, txt=txt)
93        self.prefix = prefix
94
95    def print_info(self, calc):
96        """Print information about parallelization into a file."""
97        fd = open('%s.metadata.txt' % self.prefix, 'w')
98        DebugTimer.print_info(self, calc)
99        wfs = calc.wfs
100
101        # We won't have to type a lot if everyone just sends all their numbers.
102        myranks = np.array([wfs.world.rank, wfs.kd.comm.rank,
103                            wfs.bd.comm.rank, wfs.gd.comm.rank])
104        allranks = None
105        if wfs.world.rank == 0:
106            allranks = np.empty(wfs.world.size * 4, dtype=int)
107        wfs.world.gather(myranks, 0, allranks)
108        if wfs.world.rank == 0:
109            for itsranks in allranks.reshape(-1, 4):
110                fd.write('r=%d k=%d b=%d d=%d\n' % tuple(itsranks))
111        fd.close()
112
113
114class StepTimer(Timer):
115    """Step timer to print out timing used in computation steps.
116
117    Use it like this::
118
119      from gpaw.utilities.timing import StepTimer
120      st = StepTimer()
121      ...
122      st.write_now('step 1')
123      ...
124      st.write_now('step 2')
125
126    The parameter write_as_master_only can be used to force the timer to
127    print from processess that are not the mpi master process.
128    """
129
130    def __init__(self, out=sys.stdout, name=None, write_as_master_only=True):
131        Timer.__init__(self)
132        if name is None:
133            name = '<%s>' % sys._getframe(1).f_code.co_name
134        self.name = name
135        self.out = out
136        self.alwaysprint = not write_as_master_only
137        self.now = 'temporary now'
138        self.start(self.now)
139
140    def write_now(self, mark=''):
141        self.stop(self.now)
142        if self.alwaysprint or mpi.rank == 0:
143            print(self.name, mark, self.get_time(self.now), file=self.out)
144        self.out.flush()
145        del self.timers[self.now]
146        self.start(self.now)
147
148
149class HPMTimer(Timer):
150    """HPMTimer requires installation of the IBM BlueGene/P HPM
151    middleware interface to the low-level UPC library. This will
152    most likely only work at ANL's BlueGene/P. Must compile
153    with GPAW_HPM macro in customize.py. Note that HPM_Init
154    and HPM_Finalize are called in _gpaw.c and not in the Python
155    interface. Timer must be called on all ranks in node, otherwise
156    HPM will hang. Hence, we only call HPM_start/stop on a list
157    subset of timers."""
158
159    top_level = 'GPAW.calculator'  # HPM needs top level timer
160    compatible = ['Initialization', 'SCF-cycle']
161
162    def __init__(self):
163        Timer.__init__(self)
164        from _gpaw import hpm_start, hpm_stop
165        self.hpm_start = hpm_start
166        self.hpm_stop = hpm_stop
167        hpm_start(self.top_level)
168
169    def start(self, name):
170        Timer.start(self, name)
171        if name in self.compatible:
172            self.hpm_start(name)
173
174    def stop(self, name=None):
175        Timer.stop(self, name)
176        if name in self.compatible:
177            self.hpm_stop(name)
178
179    def write(self, out=sys.stdout):
180        Timer.write(self, out)
181        self.hpm_stop(self.top_level)
182
183
184class CrayPAT_timer(Timer):
185    """Interface to CrayPAT API. In addition to regular timers,
186    the corresponding regions are profiled by CrayPAT. The gpaw-python has
187    to be compiled under CrayPAT.
188    """
189
190    def __init__(self, print_levels=4):
191        Timer.__init__(self, print_levels)
192        from _gpaw import craypat_region_begin, craypat_region_end
193        self.craypat_region_begin = craypat_region_begin
194        self.craypat_region_end = craypat_region_end
195        self.regions = {}
196        self.region_id = 5  # leave room for regions in C
197
198    def start(self, name):
199        Timer.start(self, name)
200        if name in self.regions:
201            id = self.regions[name]
202        else:
203            id = self.region_id
204            self.regions[name] = id
205            self.region_id += 1
206        self.craypat_region_begin(id, name)
207
208    def stop(self, name=None):
209        Timer.stop(self, name)
210        id = self.regions[name]
211        self.craypat_region_end(id)
212