diff --git a/drivers/gpu/nvgpu/gk20a/gk20a.c b/drivers/gpu/nvgpu/gk20a/gk20a.c
index 7855493d0..9ee5f2821 100644
--- a/drivers/gpu/nvgpu/gk20a/gk20a.c
+++ b/drivers/gpu/nvgpu/gk20a/gk20a.c
@@ -424,18 +424,6 @@ done:
 	return err;
 }
 
-/*
- * Check if the device can go busy. Basically if the driver is currently
- * in the process of dying then do not let new places make the driver busy.
- */
-int gk20a_can_busy(struct gk20a *g)
-{
-	if (nvgpu_is_enabled(g, NVGPU_DRIVER_IS_DYING)) {
-		return 0;
-	}
-	return 1;
-}
-
 int gk20a_wait_for_idle(struct gk20a *g)
 {
 	int wait_length = 150; /* 3 second overall max wait. */
diff --git a/drivers/gpu/nvgpu/include/nvgpu/enabled.h b/drivers/gpu/nvgpu/include/nvgpu/enabled.h
index 074be0c3a..12acf353d 100644
--- a/drivers/gpu/nvgpu/include/nvgpu/enabled.h
+++ b/drivers/gpu/nvgpu/include/nvgpu/enabled.h
@@ -36,6 +36,7 @@ struct gk20a;
 #define NVGPU_GR_USE_DMA_FOR_FW_BOOTSTRAP 3
 #define NVGPU_FECS_TRACE_VA			4
 #define NVGPU_CAN_RAILGATE			5
+#define NVGPU_KERNEL_IS_DYING			6
 
 /*
  * ECC flags
diff --git a/drivers/gpu/nvgpu/include/nvgpu/gk20a.h b/drivers/gpu/nvgpu/include/nvgpu/gk20a.h
index a256b01f4..0424e74da 100644
--- a/drivers/gpu/nvgpu/include/nvgpu/gk20a.h
+++ b/drivers/gpu/nvgpu/include/nvgpu/gk20a.h
@@ -49,6 +49,9 @@ struct nvgpu_gpu_ctxsw_trace_filter;
 struct priv_cmd_entry;
 struct nvgpu_gpfifo_args;
 
+#ifdef __KERNEL__
+#include <linux/notifier.h>
+#endif
 #include <nvgpu/lock.h>
 #include <nvgpu/thread.h>
 
@@ -1411,6 +1414,10 @@ struct gk20a {
 	 */
 	unsigned long *enabled_flags;
 
+#ifdef __KERNEL__
+	struct notifier_block nvgpu_reboot_nb;
+#endif
+
 	nvgpu_atomic_t usage_count;
 
 	struct nvgpu_mutex ctxsw_disable_lock;
@@ -1741,7 +1748,6 @@ void gk20a_idle(struct gk20a *g);
 int __gk20a_do_idle(struct gk20a *g, bool force_reset);
 int __gk20a_do_unidle(struct gk20a *g);
 
-int gk20a_can_busy(struct gk20a *g);
 int gk20a_wait_for_idle(struct gk20a *g);
 
 #define NVGPU_GPU_ARCHITECTURE_SHIFT 4
diff --git a/drivers/gpu/nvgpu/os/linux/module.c b/drivers/gpu/nvgpu/os/linux/module.c
index 1fd7f5442..955481c88 100644
--- a/drivers/gpu/nvgpu/os/linux/module.c
+++ b/drivers/gpu/nvgpu/os/linux/module.c
@@ -24,6 +24,8 @@
 #include <linux/interrupt.h>
 #include <linux/pm_runtime.h>
 #include <linux/reset.h>
+#include <linux/reboot.h>
+#include <linux/notifier.h>
 #include <linux/platform/tegra/common.h>
 #include <linux/pci.h>
 
@@ -76,6 +78,14 @@
 #define CREATE_TRACE_POINTS
 #include <trace/events/gk20a.h>
 
+static int nvgpu_kernel_shutdown_notification(struct notifier_block *nb,
+					unsigned long event, void *unused)
+{
+	struct gk20a *g = container_of(nb, struct gk20a, nvgpu_reboot_nb);
+
+	__nvgpu_set_enabled(g, NVGPU_KERNEL_IS_DYING, true);
+	return NOTIFY_DONE;
+}
 
 struct device_node *nvgpu_get_node(struct gk20a *g)
 {
@@ -98,6 +108,22 @@ void gk20a_busy_noresume(struct gk20a *g)
 	pm_runtime_get_noresume(dev_from_gk20a(g));
 }
 
+/*
+ * Check if the device can go busy.
+ */
+static int nvgpu_can_busy(struct gk20a *g)
+{
+	/* Can't do anything if the system is rebooting/shutting down. */
+	if (nvgpu_is_enabled(g, NVGPU_KERNEL_IS_DYING))
+		return 0;
+
+	/* Can't do anything if the driver is restarting. */
+	if (nvgpu_is_enabled(g, NVGPU_DRIVER_IS_DYING))
+		return 0;
+
+	return 1;
+}
+
 int gk20a_busy(struct gk20a *g)
 {
 	struct nvgpu_os_linux *l = nvgpu_os_linux_from_gk20a(g);
@@ -111,7 +137,7 @@ int gk20a_busy(struct gk20a *g)
 
 	down_read(&l->busy_lock);
 
-	if (!gk20a_can_busy(g)) {
+	if (!nvgpu_can_busy(g)) {
 		ret = -ENODEV;
 		atomic_dec(&g->usage_count.atomic_var);
 		goto fail;
@@ -158,7 +184,7 @@ void gk20a_idle(struct gk20a *g)
 
 	dev = dev_from_gk20a(g);
 
-	if (!(dev && gk20a_can_busy(g)))
+	if (!(dev && nvgpu_can_busy(g)))
 		return;
 
 	if (pm_runtime_enabled(dev)) {
@@ -1289,6 +1315,12 @@ static int gk20a_probe(struct platform_device *dev)
 		goto return_err;
 	}
 
+	gk20a->nvgpu_reboot_nb.notifier_call =
+		nvgpu_kernel_shutdown_notification;
+	err = register_reboot_notifier(&gk20a->nvgpu_reboot_nb);
+	if (err)
+		goto return_err;
+
 	return 0;
 
 return_err:
@@ -1368,6 +1400,8 @@ static int __exit gk20a_remove(struct platform_device *pdev)
 
 	err = nvgpu_remove(dev, &nvgpu_class);
 
+	unregister_reboot_notifier(&g->nvgpu_reboot_nb);
+
 	set_gk20a(pdev, NULL);
 
 	gk20a_put(g);