2 files changed, 17 insertions, 6 deletions

diff --git a/drivers/gpu/nvgpu/common/linux/channel.c b/drivers/gpu/nvgpu/common/linux/channel.c
index 8bfa4cfc..ea294738 100644
--- a/drivers/gpu/nvgpu/common/linux/channel.c
+++ b/drivers/gpu/nvgpu/common/linux/channel.c
@@ -753,7 +753,7 @@ int gk20a_submit_channel_gpfifo(struct channel_gk20a *c,
          */
         need_job_tracking = (flags & NVGPU_SUBMIT_GPFIFO_FLAGS_FENCE_WAIT) ||
                         (flags & NVGPU_SUBMIT_GPFIFO_FLAGS_FENCE_GET) ||
-                        c->wdt_enabled ||
+                        c->timeout.enabled ||
                         (g->can_railgate && !c->deterministic) ||
                         !skip_buffer_refcounting;
 
@@ -791,7 +791,7 @@ int gk20a_submit_channel_gpfifo(struct channel_gk20a *c,
                  */
                 need_deferred_cleanup = !c->deterministic ||
                                         need_sync_framework ||
-                                        c->wdt_enabled ||
+                                        c->timeout.enabled ||
                                         (g->can_railgate &&
                                          !c->deterministic) ||
                                         !skip_buffer_refcounting;
diff --git a/drivers/gpu/nvgpu/common/linux/ioctl_channel.c b/drivers/gpu/nvgpu/common/linux/ioctl_channel.c
index 0acaa61d..01355b78 100644
--- a/drivers/gpu/nvgpu/common/linux/ioctl_channel.c
+++ b/drivers/gpu/nvgpu/common/linux/ioctl_channel.c
@@ -319,10 +319,21 @@ static int gk20a_channel_cycle_stats_snapshot(struct channel_gk20a *ch,
 static int gk20a_channel_set_wdt_status(struct channel_gk20a *ch,
                 struct nvgpu_channel_wdt_args *args)
 {
-        if (args->wdt_status == NVGPU_IOCTL_CHANNEL_DISABLE_WDT)
-                ch->wdt_enabled = false;
-        else if (args->wdt_status == NVGPU_IOCTL_CHANNEL_ENABLE_WDT)
-                ch->wdt_enabled = true;
+        u32 status = args->wdt_status & (NVGPU_IOCTL_CHANNEL_DISABLE_WDT |
+                        NVGPU_IOCTL_CHANNEL_ENABLE_WDT);
+
+        if (status == NVGPU_IOCTL_CHANNEL_DISABLE_WDT)
+                ch->timeout.enabled = false;
+        else if (status == NVGPU_IOCTL_CHANNEL_ENABLE_WDT)
+                ch->timeout.enabled = true;
+        else
+                return -EINVAL;
+
+        if (args->wdt_status & NVGPU_IOCTL_CHANNEL_WDT_FLAG_SET_TIMEOUT)
+                ch->timeout.limit_ms = args->timeout_ms;
+
+        ch->timeout.debug_dump = (args->wdt_status &
+                        NVGPU_IOCTL_CHANNEL_WDT_FLAG_DISABLE_DUMP) == 0;
 
         return 0;
 }