cpu: Split the traffic generator into two classes

The traffic generator currently assumes that it is always driven from
a configuration file. Split it into a base class (BaseTrafficGen) that
handles basic packet generation and a derived class that implements
the config handling (TrafficGen).

Change-Id: I9407f04c40ad7e40a263c8d1ef29d37ff8e6f1b4
Signed-off-by: Andreas Sandberg <andreas.sandberg@arm.com>
Reviewed-by: Nikos Nikoleris <nikos.nikoleris@arm.com>
Reviewed-on: https://gem5-review.googlesource.com/11515

7 files changed, 689 insertions, 431 deletions

diff --git a/src/cpu/testers/traffic_gen/BaseTrafficGen.py b/src/cpu/testers/traffic_gen/BaseTrafficGen.py
new file mode 100644
index 000000000..2569c1ece
--- /dev/null
+++ b/src/cpu/testers/traffic_gen/BaseTrafficGen.py
@@ -0,0 +1,72 @@
+# Copyright (c) 2012, 2016, 2018 ARM Limited
+# All rights reserved.
+#
+# The license below extends only to copyright in the software and shall
+# not be construed as granting a license to any other intellectual
+# property including but not limited to intellectual property relating
+# to a hardware implementation of the functionality of the software
+# licensed hereunder.  You may use the software subject to the license
+# terms below provided that you ensure that this notice is replicated
+# unmodified and in its entirety in all distributions of the software,
+# modified or unmodified, in source code or in binary form.
+#
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are
+# met: redistributions of source code must retain the above copyright
+# notice, this list of conditions and the following disclaimer;
+# redistributions in binary form must reproduce the above copyright
+# notice, this list of conditions and the following disclaimer in the
+# documentation and/or other materials provided with the distribution;
+# neither the name of the copyright holders nor the names of its
+# contributors may be used to endorse or promote products derived from
+# this software without specific prior written permission.
+#
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
+# "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
+# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
+# A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
+# OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
+# SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
+# LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
+# DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
+# THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+#
+# Authors: Thomas Grass
+#          Andreas Hansson
+#          Sascha Bischoff
+
+from m5.params import *
+from m5.proxy import *
+from MemObject import MemObject
+
+# The traffic generator is a master module that generates stimuli for
+# the memory system, based on a collection of simple behaviours that
+# are either probabilistic or based on traces. It can be used stand
+# alone for creating test cases for interconnect and memory
+# controllers, or function as a black-box replacement for system
+# components that are not yet modelled in detail, e.g. a video engine
+# or baseband subsystem in an SoC.
+class BaseTrafficGen(MemObject):
+    type = 'BaseTrafficGen'
+    abstract = True
+    cxx_header = "cpu/testers/traffic_gen/traffic_gen.hh"
+
+    # Port used for sending requests and receiving responses
+    port = MasterPort("Master port")
+
+    # System used to determine the mode of the memory system
+    system = Param.System(Parent.any, "System this generator is part of")
+
+    # Should requests respond to back-pressure or not, if true, the
+    # rate of the traffic generator will be slowed down if requests
+    # are not immediately accepted
+    elastic_req = Param.Bool(False,
+                             "Slow down requests in case of backpressure")
+
+    # Let the user know if we have waited for a retry and not made any
+    # progress for a long period of time. The default value is
+    # somewhat arbitrary and may well have to be tuned.
+    progress_check = Param.Latency('1ms', "Time before exiting " \
+                                   "due to lack of progress")
diff --git a/src/cpu/testers/traffic_gen/SConscript b/src/cpu/testers/traffic_gen/SConscript
index 74d027970..f851b7925 100644
--- a/src/cpu/testers/traffic_gen/SConscript
+++ b/src/cpu/testers/traffic_gen/SConscript
@@ -1,6 +1,6 @@
 # -*- mode:python -*-
 
-# Copyright (c) 2012, 2017 ARM Limited
+# Copyright (c) 2012, 2017-2018 ARM Limited
 # All rights reserved.
 #
 # The license below extends only to copyright in the software and shall
@@ -39,19 +39,23 @@
 
 Import('*')
 
+
+Source('base.cc')
+Source('base_gen.cc')
+Source('dram_gen.cc')
+Source('dram_rot_gen.cc')
+Source('exit_gen.cc')
+Source('idle_gen.cc')
+Source('linear_gen.cc')
+Source('random_gen.cc')
+
+DebugFlag('TrafficGen')
+SimObject('BaseTrafficGen.py')
+
 # Only build the traffic generator if we have support for protobuf as the
 # tracing relies on it
 if env['HAVE_PROTOBUF']:
     SimObject('TrafficGen.py')
-
-    Source('base_gen.cc')
-    Source('dram_gen.cc')
-    Source('dram_rot_gen.cc')
-    Source('exit_gen.cc')
-    Source('idle_gen.cc')
-    Source('linear_gen.cc')
-    Source('random_gen.cc')
     Source('trace_gen.cc')
     Source('traffic_gen.cc')
 
-    DebugFlag('TrafficGen')
diff --git a/src/cpu/testers/traffic_gen/TrafficGen.py b/src/cpu/testers/traffic_gen/TrafficGen.py
index 1a6a6359c..f1e094821 100644
--- a/src/cpu/testers/traffic_gen/TrafficGen.py
+++ b/src/cpu/testers/traffic_gen/TrafficGen.py
@@ -1,4 +1,4 @@
-# Copyright (c) 2012, 2016 ARM Limited
+# Copyright (c) 2012, 2016, 2018 ARM Limited
 # All rights reserved.
 #
 # The license below extends only to copyright in the software and shall
@@ -38,48 +38,19 @@
 #          Sascha Bischoff
 
 from m5.params import *
-from m5.proxy import *
-from MemObject import MemObject
+from BaseTrafficGen import *
 
-# The traffic generator is a master module that generates stimuli for
-# the memory system, based on a collection of simple behaviours that
-# are either probabilistic or based on traces. It can be used stand
-# alone for creating test cases for interconnect and memory
-# controllers, or function as a black-box replacement for system
-# components that are not yet modelled in detail, e.g. a video engine
-# or baseband subsystem in an SoC.
-#
-# The traffic generator has a single master port that is used to send
-# requests, independent of the specific behaviour. The behaviour of
-# the traffic generator is specified in a configuration file, and this
-# file describes a state transition graph where each state is a
-# specific generator behaviour. Examples include idling, generating
-# linear address sequences, random sequences and replay of captured
-# traces. By describing these behaviours as states, it is straight
-# forward to create very complex behaviours, simply by arranging them
-# in graphs. The graph transitions can also be annotated with
-# probabilities, effectively making it a Markov Chain.
-class TrafficGen(MemObject):
+# The behaviour of this traffic generator is specified in a
+# configuration file, and this file describes a state transition graph
+# where each state is a specific generator behaviour. Examples include
+# idling, generating linear address sequences, random sequences and
+# replay of captured traces. By describing these behaviours as states,
+# it is straight forward to create very complex behaviours, simply by
+# arranging them in graphs. The graph transitions can also be
+# annotated with probabilities, effectively making it a Markov Chain.
+class TrafficGen(BaseTrafficGen):
     type = 'TrafficGen'
     cxx_header = "cpu/testers/traffic_gen/traffic_gen.hh"
 
-    # Port used for sending requests and receiving responses
-    port = MasterPort("Master port")
-
     # Config file to parse for the state descriptions
     config_file = Param.String("Configuration file describing the behaviour")
-
-    # System used to determine the mode of the memory system
-    system = Param.System(Parent.any, "System this generator is part of")
-
-    # Should requests respond to back-pressure or not, if true, the
-    # rate of the traffic generator will be slowed down if requests
-    # are not immediately accepted
-    elastic_req = Param.Bool(False,
-                             "Slow down requests in case of backpressure")
-
-    # Let the user know if we have waited for a retry and not made any
-    # progress for a long period of time. The default value is
-    # somewhat arbitrary and may well have to be tuned.
-    progress_check = Param.Latency('1ms', "Time before exiting " \
-                                   "due to lack of progress")
diff --git a/src/cpu/testers/traffic_gen/base.cc b/src/cpu/testers/traffic_gen/base.cc
new file mode 100644
index 000000000..9ea6c9398
--- /dev/null
+++ b/src/cpu/testers/traffic_gen/base.cc
@@ -0,0 +1,306 @@
+/*
+ * Copyright (c) 2012-2013, 2016-2018 ARM Limited
+ * All rights reserved
+ *
+ * The license below extends only to copyright in the software and shall
+ * not be construed as granting a license to any other intellectual
+ * property including but not limited to intellectual property relating
+ * to a hardware implementation of the functionality of the software
+ * licensed hereunder.  You may use the software subject to the license
+ * terms below provided that you ensure that this notice is replicated
+ * unmodified and in its entirety in all distributions of the software,
+ * modified or unmodified, in source code or in binary form.
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are
+ * met: redistributions of source code must retain the above copyright
+ * notice, this list of conditions and the following disclaimer;
+ * redistributions in binary form must reproduce the above copyright
+ * notice, this list of conditions and the following disclaimer in the
+ * documentation and/or other materials provided with the distribution;
+ * neither the name of the copyright holders nor the names of its
+ * contributors may be used to endorse or promote products derived from
+ * this software without specific prior written permission.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
+ * "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
+ * LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
+ * A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
+ * OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
+ * SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
+ * LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
+ * DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
+ * THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+ * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ * Authors: Thomas Grass
+ *          Andreas Hansson
+ *          Sascha Bischoff
+ */
+#include "cpu/testers/traffic_gen/base.hh"
+
+#include <sstream>
+
+#include "base/intmath.hh"
+#include "base/random.hh"
+#include "cpu/testers/traffic_gen/base_gen.hh"
+#include "debug/Checkpoint.hh"
+#include "debug/TrafficGen.hh"
+#include "params/BaseTrafficGen.hh"
+#include "sim/sim_exit.hh"
+#include "sim/stats.hh"
+#include "sim/system.hh"
+
+using namespace std;
+
+BaseTrafficGen::BaseTrafficGen(const BaseTrafficGenParams* p)
+    : MemObject(p),
+      system(p->system),
+      elasticReq(p->elastic_req),
+      progressCheck(p->progress_check),
+      noProgressEvent([this]{ noProgress(); }, name()),
+      nextTransitionTick(0),
+      nextPacketTick(0),
+      port(name() + ".port", *this),
+      retryPkt(NULL),
+      retryPktTick(0),
+      updateEvent([this]{ update(); }, name()),
+      numSuppressed(0),
+      masterID(system->getMasterId(this))
+{
+}
+
+BaseMasterPort&
+BaseTrafficGen::getMasterPort(const string& if_name, PortID idx)
+{
+    if (if_name == "port") {
+        return port;
+    } else {
+        return MemObject::getMasterPort(if_name, idx);
+    }
+}
+
+void
+BaseTrafficGen::init()
+{
+    MemObject::init();
+
+    if (!port.isConnected())
+        fatal("The port of %s is not connected!\n", name());
+}
+
+DrainState
+BaseTrafficGen::drain()
+{
+    if (!updateEvent.scheduled()) {
+        // no event has been scheduled yet (e.g. switched from atomic mode)
+        return DrainState::Drained;
+    }
+
+    if (retryPkt == NULL) {
+        // shut things down
+        nextPacketTick = MaxTick;
+        nextTransitionTick = MaxTick;
+        deschedule(updateEvent);
+        return DrainState::Drained;
+    } else {
+        return DrainState::Draining;
+    }
+}
+
+void
+BaseTrafficGen::serialize(CheckpointOut &cp) const
+{
+    DPRINTF(Checkpoint, "Serializing BaseTrafficGen\n");
+
+    // save ticks of the graph event if it is scheduled
+    Tick nextEvent = updateEvent.scheduled() ? updateEvent.when() : 0;
+
+    DPRINTF(TrafficGen, "Saving nextEvent=%llu\n", nextEvent);
+
+    SERIALIZE_SCALAR(nextEvent);
+
+    SERIALIZE_SCALAR(nextTransitionTick);
+
+    SERIALIZE_SCALAR(nextPacketTick);
+}
+
+void
+BaseTrafficGen::unserialize(CheckpointIn &cp)
+{
+    // restore scheduled events
+    Tick nextEvent;
+    UNSERIALIZE_SCALAR(nextEvent);
+    if (nextEvent != 0)
+        schedule(updateEvent, nextEvent);
+
+    UNSERIALIZE_SCALAR(nextTransitionTick);
+
+    UNSERIALIZE_SCALAR(nextPacketTick);
+}
+
+void
+BaseTrafficGen::update()
+{
+    // shift our progress-tracking event forward
+    reschedule(noProgressEvent, curTick() + progressCheck, true);
+
+    // if we have reached the time for the next state transition, then
+    // perform the transition
+    if (curTick() >= nextTransitionTick) {
+        transition();
+    } else {
+        assert(curTick() >= nextPacketTick);
+        // get the next packet and try to send it
+        PacketPtr pkt = activeGenerator->getNextPacket();
+
+        // suppress packets that are not destined for a memory, such as
+        // device accesses that could be part of a trace
+        if (system->isMemAddr(pkt->getAddr())) {
+            numPackets++;
+            if (!port.sendTimingReq(pkt)) {
+                retryPkt = pkt;
+                retryPktTick = curTick();
+            }
+        } else {
+            DPRINTF(TrafficGen, "Suppressed packet %s 0x%x\n",
+                    pkt->cmdString(), pkt->getAddr());
+
+            ++numSuppressed;
+            if (numSuppressed % 10000)
+                warn("%s suppressed %d packets with non-memory addresses\n",
+                     name(), numSuppressed);
+
+            delete pkt;
+            pkt = nullptr;
+        }
+    }
+
+    // if we are waiting for a retry, do not schedule any further
+    // events, in the case of a transition or a successful send, go
+    // ahead and determine when the next update should take place
+    if (retryPkt == NULL) {
+        nextPacketTick = activeGenerator->nextPacketTick(elasticReq, 0);
+        scheduleUpdate();
+    }
+}
+
+void
+BaseTrafficGen::transition()
+{
+    if (activeGenerator)
+        activeGenerator->exit();
+
+    activeGenerator = nextGenerator();
+
+    if (activeGenerator) {
+        const Tick duration = activeGenerator->duration;
+        if (duration != MaxTick && duration != 0) {
+            // we could have been delayed and not transitioned on the
+            // exact tick when we were supposed to (due to back
+            // pressure when sending a packet)
+            nextTransitionTick = curTick() + duration;
+        } else {
+            nextTransitionTick = MaxTick;
+        }
+
+        activeGenerator->enter();
+        nextPacketTick = activeGenerator->nextPacketTick(elasticReq, 0);
+    } else {
+        nextPacketTick = MaxTick;
+        nextTransitionTick = MaxTick;
+        assert(!updateEvent.scheduled());
+    }
+}
+
+void
+BaseTrafficGen::scheduleUpdate()
+{
+    // schedule next update event based on either the next execute
+    // tick or the next transition, which ever comes first
+    const Tick nextEventTick = std::min(nextPacketTick, nextTransitionTick);
+    if (nextEventTick == MaxTick)
+        return;
+
+    DPRINTF(TrafficGen, "Next event scheduled at %lld\n", nextEventTick);
+
+    // The next transition tick may be in the past if there was a
+    // retry, so ensure that we don't schedule anything in the past.
+    schedule(updateEvent, std::max(curTick(), nextEventTick));
+}
+
+void
+BaseTrafficGen::start()
+{
+    transition();
+    scheduleUpdate();
+}
+
+void
+BaseTrafficGen::recvReqRetry()
+{
+    assert(retryPkt != NULL);
+
+    DPRINTF(TrafficGen, "Received retry\n");
+    numRetries++;
+    // attempt to send the packet, and if we are successful start up
+    // the machinery again
+    if (port.sendTimingReq(retryPkt)) {
+        retryPkt = NULL;
+        // remember how much delay was incurred due to back-pressure
+        // when sending the request, we also use this to derive
+        // the tick for the next packet
+        Tick delay = curTick() - retryPktTick;
+        retryPktTick = 0;
+        retryTicks += delay;
+
+        if (drainState() != DrainState::Draining) {
+            // packet is sent, so find out when the next one is due
+            nextPacketTick = activeGenerator->nextPacketTick(elasticReq,
+                                                             delay);
+            scheduleUpdate();
+        } else {
+            // shut things down
+            nextPacketTick = MaxTick;
+            nextTransitionTick = MaxTick;
+            signalDrainDone();
+        }
+    }
+}
+
+void
+BaseTrafficGen::noProgress()
+{
+    fatal("BaseTrafficGen %s spent %llu ticks without making progress",
+          name(), progressCheck);
+}
+
+void
+BaseTrafficGen::regStats()
+{
+    ClockedObject::regStats();
+
+    // Initialise all the stats
+    using namespace Stats;
+
+    numPackets
+        .name(name() + ".numPackets")
+        .desc("Number of packets generated");
+
+    numRetries
+        .name(name() + ".numRetries")
+        .desc("Number of retries");
+
+    retryTicks
+        .name(name() + ".retryTicks")
+        .desc("Time spent waiting due to back-pressure (ticks)");
+}
+
+bool
+BaseTrafficGen::TrafficGenPort::recvTimingResp(PacketPtr pkt)
+{
+    delete pkt;
+
+    return true;
+}
diff --git a/src/cpu/testers/traffic_gen/base.hh b/src/cpu/testers/traffic_gen/base.hh
new file mode 100644
index 000000000..a2cad8707
--- /dev/null
+++ b/src/cpu/testers/traffic_gen/base.hh
@@ -0,0 +1,208 @@
+/*
+ * Copyright (c) 2012-2013, 2016-2018 ARM Limited
+ * All rights reserved
+ *
+ * The license below extends only to copyright in the software and shall
+ * not be construed as granting a license to any other intellectual
+ * property including but not limited to intellectual property relating
+ * to a hardware implementation of the functionality of the software
+ * licensed hereunder.  You may use the software subject to the license
+ * terms below provided that you ensure that this notice is replicated
+ * unmodified and in its entirety in all distributions of the software,
+ * modified or unmodified, in source code or in binary form.
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are
+ * met: redistributions of source code must retain the above copyright
+ * notice, this list of conditions and the following disclaimer;
+ * redistributions in binary form must reproduce the above copyright
+ * notice, this list of conditions and the following disclaimer in the
+ * documentation and/or other materials provided with the distribution;
+ * neither the name of the copyright holders nor the names of its
+ * contributors may be used to endorse or promote products derived from
+ * this software without specific prior written permission.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
+ * "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
+ * LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
+ * A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
+ * OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
+ * SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
+ * LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
+ * DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
+ * THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+ * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ * Authors: Thomas Grass
+ *          Andreas Hansson
+ *          Sascha Bischoff
+ */
+
+#ifndef __CPU_TRAFFIC_GEN_BASE_HH__
+#define __CPU_TRAFFIC_GEN_BASE_HH__
+
+#include <memory>
+#include <tuple>
+
+#include "base/statistics.hh"
+#include "mem/mem_object.hh"
+#include "mem/qport.hh"
+
+class BaseGen;
+class System;
+struct BaseTrafficGenParams;
+
+/**
+ * The traffic generator is a master module that generates stimuli for
+ * the memory system, based on a collection of simple generator
+ * behaviours that are either probabilistic or based on traces. It can
+ * be used stand alone for creating test cases for interconnect and
+ * memory controllers, or function as a black box replacement for
+ * system components that are not yet modelled in detail, e.g. a video
+ * engine or baseband subsystem.
+ */
+class BaseTrafficGen : public MemObject
+{
+  protected: // Params
+    /**
+     * The system used to determine which mode we are currently operating
+     * in.
+     */
+    System *const system;
+
+    /**
+     * Determine whether to add elasticity in the request injection,
+     * thus responding to backpressure by slowing things down.
+     */
+    const bool elasticReq;
+
+    /**
+     * Time to tolerate waiting for retries (not making progress),
+     * until we declare things broken.
+     */
+    const Tick progressCheck;
+
+  private:
+    /**
+     * Receive a retry from the neighbouring port and attempt to
+     * resend the waiting packet.
+     */
+    void recvReqRetry();
+
+    /** Transition to the next generator */
+    void transition();
+
+    /**
+     * Schedule the update event based on nextPacketTick and
+     * nextTransitionTick.
+     */
+    void scheduleUpdate();
+
+    /**
+     * Method to inform the user we have made no progress.
+     */
+    void noProgress();
+
+    /**
+     * Event to keep track of our progress, or lack thereof.
+     */
+    EventFunctionWrapper noProgressEvent;
+
+    /** Time of next transition */
+    Tick nextTransitionTick;
+
+    /** Time of the next packet. */
+    Tick nextPacketTick;
+
+
+    /** Master port specialisation for the traffic generator */
+    class TrafficGenPort : public MasterPort
+    {
+      public:
+
+        TrafficGenPort(const std::string& name, BaseTrafficGen& traffic_gen)
+            : MasterPort(name, &traffic_gen), trafficGen(traffic_gen)
+        { }
+
+      protected:
+
+        void recvReqRetry() { trafficGen.recvReqRetry(); }
+
+        bool recvTimingResp(PacketPtr pkt);
+
+        void recvTimingSnoopReq(PacketPtr pkt) { }
+
+        void recvFunctionalSnoop(PacketPtr pkt) { }
+
+        Tick recvAtomicSnoop(PacketPtr pkt) { return 0; }
+
+      private:
+
+        BaseTrafficGen& trafficGen;
+
+    };
+
+    /**
+     * Schedules event for next update and generates a new packet or
+     * requests a new generatoir depending on the current time.
+     */
+    void update();
+
+    /** The instance of master port used by the traffic generator. */
+    TrafficGenPort port;
+
+    /** Packet waiting to be sent. */
+    PacketPtr retryPkt;
+
+    /** Tick when the stalled packet was meant to be sent. */
+    Tick retryPktTick;
+
+    /** Event for scheduling updates */
+    EventFunctionWrapper updateEvent;
+
+    uint64_t numSuppressed;
+
+  private: // Stats
+    /** Count the number of generated packets. */
+    Stats::Scalar numPackets;
+
+    /** Count the number of retries. */
+    Stats::Scalar numRetries;
+
+    /** Count the time incurred from back-pressure. */
+    Stats::Scalar retryTicks;
+
+  public:
+    BaseTrafficGen(const BaseTrafficGenParams* p);
+
+    ~BaseTrafficGen() {}
+
+    BaseMasterPort& getMasterPort(const std::string &if_name,
+                                  PortID idx = InvalidPortID) override;
+
+    void init() override;
+
+    DrainState drain() override;
+
+    void serialize(CheckpointOut &cp) const override;
+    void unserialize(CheckpointIn &cp) override;
+
+    /** Register statistics */
+    void regStats() override;
+
+  protected:
+    void start();
+
+    virtual std::shared_ptr<BaseGen> nextGenerator() = 0;
+
+    /**
+     * MasterID used in generated requests.
+     */
+    const MasterID masterID;
+
+    /** Currently active generator */
+    std::shared_ptr<BaseGen> activeGenerator;
+};
+
+#endif //__CPU_TRAFFIC_GEN_BASE_HH__
diff --git a/src/cpu/testers/traffic_gen/traffic_gen.cc b/src/cpu/testers/traffic_gen/traffic_gen.cc
index d262fd9f7..b3f73e2e5 100644
--- a/src/cpu/testers/traffic_gen/traffic_gen.cc
+++ b/src/cpu/testers/traffic_gen/traffic_gen.cc
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2012-2013, 2016-2017 ARM Limited
+ * Copyright (c) 2012-2013, 2016-2018 ARM Limited
  * All rights reserved
  *
  * The license below extends only to copyright in the software and shall
@@ -47,29 +47,24 @@
 
 #include "base/intmath.hh"
 #include "base/random.hh"
-#include "debug/Checkpoint.hh"
+#include "cpu/testers/traffic_gen/dram_gen.hh"
+#include "cpu/testers/traffic_gen/dram_rot_gen.hh"
+#include "cpu/testers/traffic_gen/exit_gen.hh"
+#include "cpu/testers/traffic_gen/idle_gen.hh"
+#include "cpu/testers/traffic_gen/linear_gen.hh"
+#include "cpu/testers/traffic_gen/random_gen.hh"
+#include "cpu/testers/traffic_gen/trace_gen.hh"
 #include "debug/TrafficGen.hh"
+#include "params/TrafficGen.hh"
 #include "sim/stats.hh"
 #include "sim/system.hh"
 
 using namespace std;
 
 TrafficGen::TrafficGen(const TrafficGenParams* p)
-    : MemObject(p),
-      system(p->system),
-      masterID(system->getMasterId(this)),
+    : BaseTrafficGen(p),
       configFile(p->config_file),
-      elasticReq(p->elastic_req),
-      progressCheck(p->progress_check),
-      noProgressEvent([this]{ noProgress(); }, name()),
-      nextTransitionTick(0),
-      nextPacketTick(0),
-      currState(0),
-      port(name() + ".port", *this),
-      retryPkt(NULL),
-      retryPktTick(0),
-      updateEvent([this]{ update(); }, name()),
-      numSuppressed(0)
+      currState(0)
 {
 }
 
@@ -79,156 +74,46 @@ TrafficGenParams::create()
     return new TrafficGen(this);
 }
 
-BaseMasterPort&
-TrafficGen::getMasterPort(const string& if_name, PortID idx)
-{
-    if (if_name == "port") {
-        return port;
-    } else {
-        return MemObject::getMasterPort(if_name, idx);
-    }
-}
-
 void
 TrafficGen::init()
 {
-    if (!port.isConnected())
-        fatal("The port of %s is not connected!\n", name());
-
-    // if the system is in timing mode active the request generator
-    if (system->isTimingMode()) {
-        DPRINTF(TrafficGen, "Timing mode, activating request generator\n");
-
-        parseConfig();
+    BaseTrafficGen::init();
 
-        // enter initial state
-        enterState(currState);
-    } else {
-        DPRINTF(TrafficGen,
-                "Traffic generator is only active in timing mode\n");
-    }
+    parseConfig();
 }
 
 void
 TrafficGen::initState()
 {
+    BaseTrafficGen::initState();
+
     // when not restoring from a checkpoint, make sure we kick things off
     if (system->isTimingMode()) {
-        // call nextPacketTick on the state to advance it
-        nextPacketTick = states[currState]->nextPacketTick(elasticReq, 0);
-        schedule(updateEvent, std::min(nextPacketTick, nextTransitionTick));
+        DPRINTF(TrafficGen, "Timing mode, activating request generator\n");
+        start();
     } else {
         DPRINTF(TrafficGen,
                 "Traffic generator is only active in timing mode\n");
     }
 }
 
-DrainState
-TrafficGen::drain()
-{
-    if (!updateEvent.scheduled()) {
-        // no event has been scheduled yet (e.g. switched from atomic mode)
-        return DrainState::Drained;
-    }
-
-    if (retryPkt == NULL) {
-        // shut things down
-        nextPacketTick = MaxTick;
-        nextTransitionTick = MaxTick;
-        deschedule(updateEvent);
-        return DrainState::Drained;
-    } else {
-        return DrainState::Draining;
-    }
-}
-
 void
 TrafficGen::serialize(CheckpointOut &cp) const
 {
-    DPRINTF(Checkpoint, "Serializing TrafficGen\n");
-
-    // save ticks of the graph event if it is scheduled
-    Tick nextEvent = updateEvent.scheduled() ? updateEvent.when() : 0;
-
-    DPRINTF(TrafficGen, "Saving nextEvent=%llu\n", nextEvent);
-
-    SERIALIZE_SCALAR(nextEvent);
-
-    SERIALIZE_SCALAR(nextTransitionTick);
-
-    SERIALIZE_SCALAR(nextPacketTick);
-
     SERIALIZE_SCALAR(currState);
+
+    BaseTrafficGen::serialize(cp);
 }
 
 void
 TrafficGen::unserialize(CheckpointIn &cp)
 {
-    // restore scheduled events
-    Tick nextEvent;
-    UNSERIALIZE_SCALAR(nextEvent);
-    if (nextEvent != 0) {
-        schedule(updateEvent, nextEvent);
-    }
-
-    UNSERIALIZE_SCALAR(nextTransitionTick);
-
-    UNSERIALIZE_SCALAR(nextPacketTick);
-
     // @todo In the case of a stateful generator state such as the
     // trace player we would also have to restore the position in the
     // trace playback and the tick offset
     UNSERIALIZE_SCALAR(currState);
-}
 
-void
-TrafficGen::update()
-{
-    // shift our progress-tracking event forward
-    reschedule(noProgressEvent, curTick() + progressCheck, true);
-
-    // if we have reached the time for the next state transition, then
-    // perform the transition
-    if (curTick() >= nextTransitionTick) {
-        transition();
-    } else {
-        assert(curTick() >= nextPacketTick);
-        // get the next packet and try to send it
-        PacketPtr pkt = states[currState]->getNextPacket();
-
-        // suppress packets that are not destined for a memory, such as
-        // device accesses that could be part of a trace
-        if (system->isMemAddr(pkt->getAddr())) {
-            numPackets++;
-            if (!port.sendTimingReq(pkt)) {
-                retryPkt = pkt;
-                retryPktTick = curTick();
-            }
-        } else {
-            DPRINTF(TrafficGen, "Suppressed packet %s 0x%x\n",
-                    pkt->cmdString(), pkt->getAddr());
-
-            ++numSuppressed;
-            if (numSuppressed % 10000)
-                warn("%s suppressed %d packets with non-memory addresses\n",
-                     name(), numSuppressed);
-
-            delete pkt;
-            pkt = nullptr;
-        }
-    }
-
-    // if we are waiting for a retry, do not schedule any further
-    // events, in the case of a transition or a successful send, go
-    // ahead and determine when the next update should take place
-    if (retryPkt == NULL) {
-        // schedule next update event based on either the next execute
-        // tick or the next transition, which ever comes first
-        nextPacketTick = states[currState]->nextPacketTick(elasticReq, 0);
-        Tick nextEventTick = std::min(nextPacketTick, nextTransitionTick);
-        DPRINTF(TrafficGen, "Next event scheduled at %lld\n", nextEventTick);
-        schedule(updateEvent, nextEventTick);
-    }
+    BaseTrafficGen::unserialize(cp);
 }
 
 std::string
@@ -298,14 +183,14 @@ TrafficGen::parseConfig()
                     is >> traceFile >> addrOffset;
                     traceFile = resolveFile(traceFile);
 
-                    states[id] = new TraceGen(name(), masterID, duration,
-                                              traceFile, addrOffset);
+                    states[id].reset(new TraceGen(name(), masterID, duration,
+                                                  traceFile, addrOffset));
                     DPRINTF(TrafficGen, "State: %d TraceGen\n", id);
                 } else if (mode == "IDLE") {
-                    states[id] = new IdleGen(name(), masterID, duration);
+                    states[id].reset(new IdleGen(name(), masterID, duration));
                     DPRINTF(TrafficGen, "State: %d IdleGen\n", id);
                 } else if (mode == "EXIT") {
-                    states[id] = new ExitGen(name(), masterID, duration);
+                    states[id].reset(new ExitGen(name(), masterID, duration));
                     DPRINTF(TrafficGen, "State: %d ExitGen\n", id);
                 } else if (mode == "LINEAR" || mode == "RANDOM" ||
                            mode == "DRAM"   || mode == "DRAM_ROTATE") {
@@ -338,18 +223,18 @@ TrafficGen::parseConfig()
                         fatal("%s cannot have min_period > max_period", name());
 
                     if (mode == "LINEAR") {
-                        states[id] = new LinearGen(name(), masterID,
+                        states[id].reset(new LinearGen(name(), masterID,
                                                    duration, start_addr,
                                                    end_addr, blocksize,
                                                    min_period, max_period,
-                                                   read_percent, data_limit);
+                                                   read_percent, data_limit));
                         DPRINTF(TrafficGen, "State: %d LinearGen\n", id);
                     } else if (mode == "RANDOM") {
-                        states[id] = new RandomGen(name(), masterID,
+                        states[id].reset(new RandomGen(name(), masterID,
                                                    duration, start_addr,
                                                    end_addr, blocksize,
                                                    min_period, max_period,
-                                                   read_percent, data_limit);
+                                                   read_percent, data_limit));
                         DPRINTF(TrafficGen, "State: %d RandomGen\n", id);
                     } else if (mode == "DRAM" || mode == "DRAM_ROTATE") {
                         // stride size (bytes) of the request for achieving
@@ -387,7 +272,7 @@ TrafficGen::parseConfig()
                         }
 
                         if (mode == "DRAM") {
-                            states[id] = new DramGen(name(), masterID,
+                            states[id].reset(new DramGen(name(), masterID,
                                                      duration, start_addr,
                                                      end_addr, blocksize,
                                                      min_period, max_period,
@@ -396,7 +281,7 @@ TrafficGen::parseConfig()
                                                      nbr_of_banks_DRAM,
                                                      nbr_of_banks_util,
                                                      addr_mapping,
-                                                     nbr_of_ranks);
+                                                     nbr_of_ranks));
                             DPRINTF(TrafficGen, "State: %d DramGen\n", id);
                         } else {
                             // Will rotate to the next rank after rotating
@@ -407,7 +292,7 @@ TrafficGen::parseConfig()
                                 (read_percent == 50) ? nbr_of_banks_util * 2
                                                      : nbr_of_banks_util;
 
-                            states[id] = new DramRotGen(name(), masterID,
+                            states[id].reset(new DramRotGen(name(), masterID,
                                                      duration, start_addr,
                                                      end_addr, blocksize,
                                                      min_period, max_period,
@@ -417,7 +302,7 @@ TrafficGen::parseConfig()
                                                      nbr_of_banks_util,
                                                      addr_mapping,
                                                      nbr_of_ranks,
-                                                     max_seq_count_per_rank);
+                                                     max_seq_count_per_rank));
                             DPRINTF(TrafficGen, "State: %d DramRotGen\n", id);
                         }
                     }
@@ -478,13 +363,9 @@ TrafficGen::parseConfig()
     infile.close();
 }
 
-void
-TrafficGen::transition()
+size_t
+TrafficGen::nextState()
 {
-    // exit the current state
-    states[currState]->exit();
-
-    // determine next state
     double p = random_mt.random<double>();
     assert(currState < transitionMatrix.size());
     double cumulative = 0.0;
@@ -494,87 +375,17 @@ TrafficGen::transition()
         ++i;
     } while (cumulative < p && i < transitionMatrix[currState].size());
 
-    enterState(i - 1);
-}
-
-void
-TrafficGen::enterState(uint32_t newState)
-{
-    DPRINTF(TrafficGen, "Transition to state %d\n", newState);
-
-    currState = newState;
-    // we could have been delayed and not transitioned on the exact
-    // tick when we were supposed to (due to back pressure when
-    // sending a packet)
-    nextTransitionTick = curTick() + states[currState]->duration;
-    states[currState]->enter();
-}
-
-void
-TrafficGen::recvReqRetry()
-{
-    assert(retryPkt != NULL);
-
-    DPRINTF(TrafficGen, "Received retry\n");
-    numRetries++;
-    // attempt to send the packet, and if we are successful start up
-    // the machinery again
-    if (port.sendTimingReq(retryPkt)) {
-        retryPkt = NULL;
-        // remember how much delay was incurred due to back-pressure
-        // when sending the request, we also use this to derive
-        // the tick for the next packet
-        Tick delay = curTick() - retryPktTick;
-        retryPktTick = 0;
-        retryTicks += delay;
-
-        if (drainState() != DrainState::Draining) {
-            // packet is sent, so find out when the next one is due
-            nextPacketTick = states[currState]->nextPacketTick(elasticReq,
-                                                               delay);
-            Tick nextEventTick = std::min(nextPacketTick, nextTransitionTick);
-            schedule(updateEvent, std::max(curTick(), nextEventTick));
-        } else {
-            // shut things down
-            nextPacketTick = MaxTick;
-            nextTransitionTick = MaxTick;
-            signalDrainDone();
-        }
-    }
-}
-
-void
-TrafficGen::noProgress()
-{
-    fatal("TrafficGen %s spent %llu ticks without making progress",
-          name(), progressCheck);
-}
-
-void
-TrafficGen::regStats()
-{
-    ClockedObject::regStats();
-
-    // Initialise all the stats
-    using namespace Stats;
-
-    numPackets
-        .name(name() + ".numPackets")
-        .desc("Number of packets generated");
-
-    numRetries
-        .name(name() + ".numRetries")
-        .desc("Number of retries");
-
-    retryTicks
-        .name(name() + ".retryTicks")
-        .desc("Time spent waiting due to back-pressure (ticks)");
+    return i - 1;
 }
 
-bool
-TrafficGen::TrafficGenPort::recvTimingResp(PacketPtr pkt)
+std::shared_ptr<BaseGen>
+TrafficGen::nextGenerator()
 {
-    delete pkt;
+    // Return the initial state if there isn't an active generator,
+    // otherwise perform a state transition.
+    if (activeGenerator)
+        currState = nextState();
 
-    return true;
+    DPRINTF(TrafficGen, "Transition to state %d\n", currState);
+    return states[currState];
 }
diff --git a/src/cpu/testers/traffic_gen/traffic_gen.hh b/src/cpu/testers/traffic_gen/traffic_gen.hh
index 6c2c16382..5514e3be6 100644
--- a/src/cpu/testers/traffic_gen/traffic_gen.hh
+++ b/src/cpu/testers/traffic_gen/traffic_gen.hh
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2012-2013, 2016-2017 ARM Limited
+ * Copyright (c) 2012-2013, 2016-2018 ARM Limited
  * All rights reserved
  *
  * The license below extends only to copyright in the software and shall
@@ -44,45 +44,39 @@
 
 #include <unordered_map>
 
-#include "base/statistics.hh"
-#include "cpu/testers/traffic_gen/base_gen.hh"
-#include "cpu/testers/traffic_gen/dram_gen.hh"
-#include "cpu/testers/traffic_gen/dram_rot_gen.hh"
-#include "cpu/testers/traffic_gen/exit_gen.hh"
-#include "cpu/testers/traffic_gen/idle_gen.hh"
-#include "cpu/testers/traffic_gen/linear_gen.hh"
-#include "cpu/testers/traffic_gen/random_gen.hh"
-#include "cpu/testers/traffic_gen/trace_gen.hh"
-#include "mem/mem_object.hh"
-#include "mem/qport.hh"
-#include "params/TrafficGen.hh"
+#include "cpu/testers/traffic_gen/base.hh"
+
+struct TrafficGenParams;
 
 /**
  * The traffic generator is a master module that generates stimuli for
- * the memory system, based on a collection of simple generator
- * behaviours that are either probabilistic or based on traces. It can
- * be used stand alone for creating test cases for interconnect and
- * memory controllers, or function as a black box replacement for
- * system components that are not yet modelled in detail, e.g. a video
- * engine or baseband subsystem.
+ * the memory system, based on a collection of simple behaviours that
+ * are either probabilistic or based on traces. It can be used stand
+ * alone for creating test cases for interconnect and memory
+ * controllers, or function as a black-box replacement for system
+ * components that are not yet modelled in detail, e.g. a video engine
+ * or baseband subsystem in an SoC.
+ *
+ * The traffic generator has a single master port that is used to send
+ * requests, independent of the specific behaviour. The behaviour of
+ * the traffic generator is specified in a configuration file, and this
+ * file describes a state transition graph where each state is a
+ * specific generator behaviour. Examples include idling, generating
+ * linear address sequences, random sequences and replay of captured
+ * traces. By describing these behaviours as states, it is straight
+ * forward to create very complex behaviours, simply by arranging them
+ * in graphs. The graph transitions can also be annotated with
+ * probabilities, effectively making it a Markov Chain.
  */
-class TrafficGen : public MemObject
+class TrafficGen : public BaseTrafficGen
 {
-
-  private:
-
+  private: // Params
     /**
-     * Determine next state and perform the transition.
-     */
-    void transition();
-
-    /**
-     * Enter a new state.
-     *
-     * @param newState identifier of state to enter
+     * The config file to parse.
      */
-    void enterState(uint32_t newState);
+    const std::string configFile;
 
+  private:
     /**
      * Resolve a file path in the configuration file.
      *
@@ -98,29 +92,16 @@ class TrafficGen : public MemObject
      */
     std::string resolveFile(const std::string &name);
 
-    /**
-     * Parse the config file and build the state map and
-     * transition matrix.
-     */
+     /**
+      * Parse the config file and build the state map and
+      * transition matrix.
+      */
     void parseConfig();
 
     /**
-     * Schedules event for next update and executes an update on the
-     * state graph, either performing a state transition or executing
-     * the current state, depending on the current time.
-     */
-    void update();
-
-    /**
-     * Receive a retry from the neighbouring port and attempt to
-     * resend the waiting packet.
+     * Use the transition matrix to find the next state index.
      */
-    void recvReqRetry();
-
-    /**
-     * Method to inform the user we have made no progress.
-     */
-    void noProgress();
+    size_t nextState();
 
     /** Struct to represent a probabilistic transition during parsing. */
     struct Transition {
@@ -129,45 +110,6 @@ class TrafficGen : public MemObject
         double p;
     };
 
-    /**
-     * The system used to determine which mode we are currently operating
-     * in.
-     */
-    System* system;
-
-    /**
-     * MasterID used in generated requests.
-     */
-    MasterID masterID;
-
-    /**
-     * The config file to parse.
-     */
-    const std::string configFile;
-
-    /**
-     * Determine whether to add elasticity in the request injection,
-     * thus responding to backpressure by slowing things down.
-     */
-    const bool elasticReq;
-
-    /**
-     * Time to tolerate waiting for retries (not making progress),
-     * until we declare things broken.
-     */
-    const Tick progressCheck;
-
-    /**
-     * Event to keep track of our progress, or lack thereof.
-     */
-    EventFunctionWrapper noProgressEvent;
-
-    /** Time of next transition */
-    Tick nextTransitionTick;
-
-    /** Time of the next packet. */
-    Tick nextPacketTick;
-
     /** State transition matrix */
     std::vector<std::vector<double> > transitionMatrix;
 
@@ -175,57 +117,10 @@ class TrafficGen : public MemObject
     uint32_t currState;
 
     /** Map of generator states */
-    std::unordered_map<uint32_t, BaseGen*> states;
-
-    /** Master port specialisation for the traffic generator */
-    class TrafficGenPort : public MasterPort
-    {
-      public:
-
-        TrafficGenPort(const std::string& name, TrafficGen& traffic_gen)
-            : MasterPort(name, &traffic_gen), trafficGen(traffic_gen)
-        { }
-
-      protected:
-
-        void recvReqRetry() { trafficGen.recvReqRetry(); }
-
-        bool recvTimingResp(PacketPtr pkt);
-
-        void recvTimingSnoopReq(PacketPtr pkt) { }
-
-        void recvFunctionalSnoop(PacketPtr pkt) { }
+    std::unordered_map<uint32_t, std::shared_ptr<BaseGen>> states;
 
-        Tick recvAtomicSnoop(PacketPtr pkt) { return 0; }
-
-      private:
-
-        TrafficGen& trafficGen;
-
-    };
-
-    /** The instance of master port used by the traffic generator. */
-    TrafficGenPort port;
-
-    /** Packet waiting to be sent. */
-    PacketPtr retryPkt;
-
-    /** Tick when the stalled packet was meant to be sent. */
-    Tick retryPktTick;
-
-    /** Event for scheduling updates */
-    EventFunctionWrapper updateEvent;
-
-    uint64_t numSuppressed;
-
-    /** Count the number of generated packets. */
-    Stats::Scalar numPackets;
-
-    /** Count the number of retries. */
-    Stats::Scalar numRetries;
-
-    /** Count the time incurred from back-pressure. */
-    Stats::Scalar retryTicks;
+  protected: // BaseTrafficGen
+    std::shared_ptr<BaseGen> nextGenerator() override;
 
   public:
 
@@ -233,21 +128,12 @@ class TrafficGen : public MemObject
 
     ~TrafficGen() {}
 
-    BaseMasterPort& getMasterPort(const std::string &if_name,
-                                  PortID idx = InvalidPortID) override;
-
     void init() override;
-
     void initState() override;
 
-    DrainState drain() override;
-
     void serialize(CheckpointOut &cp) const override;
     void unserialize(CheckpointIn &cp) override;
 
-    /** Register statistics */
-    void regStats() override;
-
 };
 
 #endif //__CPU_TRAFFIC_GEN_TRAFFIC_GEN_HH__