diff --git a/src/arm-ubuntu/files/after_boot.sh b/src/arm-ubuntu/files/after_boot.sh
index c503c0f1f..70c370993 100755
--- a/src/arm-ubuntu/files/after_boot.sh
+++ b/src/arm-ubuntu/files/after_boot.sh
@@ -40,7 +40,7 @@ else
             # If we can't read the script exit the simulation. If we cannot exit the
             # simulation, this probably means that we are running in QEMU. So, ignore
             # future calls to gem5-bridge.
-            if ! gem5-bridge exit; then
+            if ! echo 0 > /dev/gem5/exit; then
                 # Useful for booting the disk image in (e.g.,) qemu for debugging
                 printf "gem5-bridge exit failed, dropping to shell.\n"
                 IGNORE_M5=1 /bin/bash
diff --git a/src/arm-ubuntu/files/flash0.img b/src/arm-ubuntu/files/flash0.img
new file mode 100644
index 000000000..4c1f81384
Binary files /dev/null and b/src/arm-ubuntu/files/flash0.img differ
diff --git a/src/arm-ubuntu/files/gem5_init.sh b/src/arm-ubuntu/files/gem5_init.sh
index 5ed2017a2..45a7c20e9 100755
--- a/src/arm-ubuntu/files/gem5_init.sh
+++ b/src/arm-ubuntu/files/gem5_init.sh
@@ -7,9 +7,25 @@ mount -t sysfs /sys /sys
 cmdline=$(cat /proc/cmdline)
 no_systemd=false
 
+# Load gem5_bridge driver
+## Default parameters (ARM64)
+gem5_bridge_baseaddr=0x10010000
+gem5_bridge_rangesize=0x10000
+## Try to read overloads from kernel arguments
+if [[ $cmdline =~ gem5_bridge_baseaddr=([[:alnum:]]+) ]]; then
+    gem5_bridge_baseaddr=${BASH_REMATCH[1]}
+fi
+if [[ $cmdline =~ gem5_bridge_rangesize=([[:alnum:]]+) ]]; then
+    gem5_bridge_rangesize=${BASH_REMATCH[1]}
+fi
+## Insert driver
+modprobe gem5_bridge \
+    gem5_bridge_baseaddr=$gem5_bridge_baseaddr \
+    gem5_bridge_rangesize=$gem5_bridge_rangesize
+
 # gem5-bridge exit signifying that kernel is booted
 printf "Kernel booted, starting gem5 init...\n"
-gem5-bridge exit
+echo 0 > /dev/gem5/exit # TODO: Make this a specialized event.
 
 if [[ $cmdline == *"no_systemd"* ]]; then
     no_systemd=true
@@ -17,9 +33,9 @@ fi
 
 # Run systemd via exec if not disabled
 if [[ $no_systemd == false ]]; then
-    # gem5-bridge exit signifying that systemd will be booted
     printf "Starting systemd...\n"
     exec /lib/systemd/systemd
 else
+    printf "Dropping to shell as gem5 user...\n"
     exec su - gem5
 fi
diff --git a/src/arm-ubuntu/scripts/post-installation.sh b/src/arm-ubuntu/scripts/post-installation.sh
index 8bd0bbbed..f5970ccda 100755
--- a/src/arm-ubuntu/scripts/post-installation.sh
+++ b/src/arm-ubuntu/scripts/post-installation.sh
@@ -15,6 +15,13 @@ apt-get install -y build-essential
 echo "Installing serial service for autologin after systemd"
 mv /home/gem5/serial-getty@.service /lib/systemd/system/
 
+# Make sure the headers are installed to extract the kernel that DKMS
+# packages will be built against.
+sudo apt -y install "linux-headers-$(uname -r)" "linux-modules-extra-$(uname -r)"
+
+echo "Extracting linux kernel $(uname -r) to /home/gem5/vmlinux-arm-ubuntu"
+sudo bash -c "/usr/src/linux-headers-$(uname -r)/scripts/extract-vmlinux /boot/vmlinuz-$(uname -r) > /home/gem5/vmlinux-arm-ubuntu"
+
 echo "Installing the gem5 init script in /sbin"
 mv /home/gem5/gem5_init.sh /sbin
 mv /sbin/init /sbin/init.old
@@ -31,10 +38,11 @@ rm /etc/update-motd.d/*
 echo "Building and installing gem5-bridge (m5) and libm5"
 
 # Just get the files we need
-git clone https://github.com/gem5/gem5.git --depth=1 --filter=blob:none --no-checkout --sparse --single-branch --branch=stable
+git clone https://github.com/nkrim/gem5.git --depth=1 --filter=blob:none --no-checkout --sparse --single-branch --branch=gem5-bridge
 pushd gem5
 # Checkout just the files we need
 git sparse-checkout add util/m5
+git sparse-checkout add util/gem5_bridge
 git sparse-checkout add include
 git checkout
 # Install the headers globally so that other benchmarks can use them
@@ -46,6 +54,12 @@ scons build/arm64/out/m5
 cp build/arm64/out/m5 /usr/local/bin/
 cp build/arm64/out/libm5.a /usr/local/lib/
 popd
+
+# Build and insert the gem5-bridge driver
+pushd util/gem5_bridge
+make build install
+depmod --quick
+popd
 popd
 
 # rename the m5 binary to gem5-bridge
@@ -58,7 +72,7 @@ chmod u+s /usr/local/bin/gem5-bridge
 ln -s /usr/local/bin/gem5-bridge /usr/local/bin/m5
 
 # delete the git repo for gem5
-rm -rf gem5
+#rm -rf gem5
 echo "Done building and installing gem5-bridge (m5) and libm5"
 
 # You can extend this script to install your own packages here.
diff --git a/src/x86-ubuntu/files/after_boot.sh b/src/x86-ubuntu/files/after_boot.sh
index c503c0f1f..3cc11d28a 100755
--- a/src/x86-ubuntu/files/after_boot.sh
+++ b/src/x86-ubuntu/files/after_boot.sh
@@ -28,7 +28,7 @@ fi
 printf "Interactive mode: $interactive\n"
 
 if [[ $interactive == true ]]; then
-    printf "Interactive mode enabled, dropping to shell."
+    printf "Interactive mode enabled, dropping to shell.\n"
     /bin/bash
 else
     # Try to read the file from the host when running with gem5
diff --git a/src/x86-ubuntu/files/exit.sh b/src/x86-ubuntu/files/exit.sh
index 4b147e01f..355a4fbb0 100755
--- a/src/x86-ubuntu/files/exit.sh
+++ b/src/x86-ubuntu/files/exit.sh
@@ -3,4 +3,4 @@
 # Copyright (c) 2020 The Regents of the University of California.
 # SPDX-License-Identifier: BSD 3-Clause
 
-m5 exit
+gem5-bridge exit
diff --git a/src/x86-ubuntu/files/gem5_init.sh b/src/x86-ubuntu/files/gem5_init.sh
index 553531732..fd3ae7bdc 100755
--- a/src/x86-ubuntu/files/gem5_init.sh
+++ b/src/x86-ubuntu/files/gem5_init.sh
@@ -15,6 +15,22 @@ mount -t sysfs /sys /sys
 cmdline=$(cat /proc/cmdline)
 no_systemd=false
 
+# Load gem5_bridge driver
+## Default parameters (x86_64)
+gem5_bridge_baseaddr=0xffff0000
+gem5_bridge_rangesize=0x10000
+## Try to read overloads from kernel arguments
+if [[ $cmdline =~ gem5_bridge_baseaddr=([[:alnum:]]+) ]]; then
+    gem5_bridge_baseaddr=${BASH_REMATCH[1]}
+fi
+if [[ $cmdline =~ gem5_bridge_rangesize=([[:alnum:]]+) ]]; then
+    gem5_bridge_rangesize=${BASH_REMATCH[1]}
+fi
+## Insert driver
+modprobe gem5_bridge \
+    gem5_bridge_baseaddr=$gem5_bridge_baseaddr \
+    gem5_bridge_rangesize=$gem5_bridge_rangesize
+
 # gem5-bridge exit signifying that kernel is booted
 # This will cause the simulation to exit. Note that this will
 # cause qemu to fail.
diff --git a/src/x86-ubuntu/scripts/post-installation.sh b/src/x86-ubuntu/scripts/post-installation.sh
index 0b3df372f..c8172f678 100755
--- a/src/x86-ubuntu/scripts/post-installation.sh
+++ b/src/x86-ubuntu/scripts/post-installation.sh
@@ -41,10 +41,11 @@ rm /etc/update-motd.d/*
 echo "Building and installing gem5-bridge (m5) and libm5"
 
 # Just get the files we need
-git clone https://github.com/gem5/gem5.git --depth=1 --filter=blob:none --no-checkout --sparse --single-branch --branch=stable
+git clone https://github.com/nkrim/gem5.git --depth=1 --filter=blob:none --no-checkout --sparse --single-branch --branch=gem5-bridge
 pushd gem5
 # Checkout just the files we need
 git sparse-checkout add util/m5
+git sparse-checkout add util/gem5_bridge
 git sparse-checkout add include
 git checkout
 # Install the headers globally so that other benchmarks can use them
@@ -56,6 +57,12 @@ scons build/x86/out/m5
 cp build/x86/out/m5 /usr/local/bin/
 cp build/x86/out/libm5.a /usr/local/lib/
 popd
+
+# Build and insert the gem5-bridge driver
+pushd util/gem5_bridge
+make build install
+depmod --quick
+popd
 popd
 
 # rename the m5 binary to gem5-bridge
@@ -68,7 +75,7 @@ chmod u+s /usr/local/bin/gem5-bridge
 ln -s /usr/local/bin/gem5-bridge /usr/local/bin/m5
 
 # delete the git repo for gem5
-rm -rf gem5
+# rm -rf gem5
 echo "Done building and installing gem5-bridge (m5) and libm5"
 
 # You can extend this script to install your own packages here or by modifying the `x86-ubuntu.pkr.hcl` file.
@@ -76,14 +83,14 @@ echo "Done building and installing gem5-bridge (m5) and libm5"
 # Disable network by default
 echo "Disabling network by default"
 echo "See README.md for instructions on how to enable network"
-if [ -f /etc/netplan/50-cloud-init.yaml ]; then
-    mv /etc/netplan/50-cloud-init.yaml /etc/netplan/50-cloud-init.yaml.bak
-elif [ -f /etc/netplan/00-installer-config.yaml ]; then
-    mv /etc/netplan/00-installer-config.yaml /etc/netplan/00-installer-config.yaml.bak
-    netplan apply
-fi
+# if [ -f /etc/netplan/50-cloud-init.yaml ]; then
+#     mv /etc/netplan/50-cloud-init.yaml /etc/netplan/50-cloud-init.yaml.bak
+# elif [ -f /etc/netplan/00-installer-config.yaml ]; then
+#     mv /etc/netplan/00-installer-config.yaml /etc/netplan/00-installer-config.yaml.bak
+#     netplan apply
+# fi
 # Disable systemd service that waits for network to be online
-systemctl disable systemd-networkd-wait-online.service
-systemctl mask systemd-networkd-wait-online.service
+# systemctl disable systemd-networkd-wait-online.service
+# systemctl mask systemd-networkd-wait-online.service
 
 echo "Post Installation Done"
diff --git a/src/x86-ubuntu/x86-ubuntu.pkr.hcl b/src/x86-ubuntu/x86-ubuntu.pkr.hcl
index 9e913cc7d..10c2b368e 100644
--- a/src/x86-ubuntu/x86-ubuntu.pkr.hcl
+++ b/src/x86-ubuntu/x86-ubuntu.pkr.hcl
@@ -78,8 +78,8 @@ build {
   }
 
   provisioner "file" {
-  source      = "/home/gem5/vmlinux-x86-ubuntu"
-  destination = "./disk-image/vmlinux-x86-ubuntu"
-  direction   = "download"
+    source      = "/home/gem5/vmlinux-x86-ubuntu"
+    destination = "./disk-image/vmlinux-x86-ubuntu"
+    direction   = "download"
   }
 }