3 files changed, 12 insertions, 2 deletions

diff --git a/drivers/gpu/nvgpu/common/linux/sysfs.c b/drivers/gpu/nvgpu/common/linux/sysfs.c
index b1e7d1ed..053c1a12 100644
--- a/drivers/gpu/nvgpu/common/linux/sysfs.c
+++ b/drivers/gpu/nvgpu/common/linux/sysfs.c
@@ -965,8 +965,10 @@ static ssize_t gfxp_wfi_timeout_count_store(struct device *dev,
         if (kstrtoul(buf, 10, &val) < 0)
                 return -EINVAL;
 
-        if (val >= 100*1000*1000) /* 100ms @ 1Ghz */
-                return -EINVAL;
+        if (g->ops.gr.get_max_gfxp_wfi_timeout_count) {
+                if (val >= g->ops.gr.get_max_gfxp_wfi_timeout_count(g))
+                        return -EINVAL;
+        }
 
         gr->gfxp_wfi_timeout_count = val;
 
diff --git a/drivers/gpu/nvgpu/common/linux/vgpu/gp10b/vgpu_hal_gp10b.c b/drivers/gpu/nvgpu/common/linux/vgpu/gp10b/vgpu_hal_gp10b.c
index b35cac90..acd0ad5d 100644
--- a/drivers/gpu/nvgpu/common/linux/vgpu/gp10b/vgpu_hal_gp10b.c
+++ b/drivers/gpu/nvgpu/common/linux/vgpu/gp10b/vgpu_hal_gp10b.c
@@ -209,6 +209,10 @@ static const struct gpu_ops vgpu_gp10b_ops = {
                 .set_ctxsw_preemption_mode =
                                         vgpu_gr_gp10b_set_ctxsw_preemption_mode,
                 .init_ctxsw_hdr_data = gr_gp10b_init_ctxsw_hdr_data,
+                .init_gfxp_wfi_timeout_count =
+                        gr_gp10b_init_gfxp_wfi_timeout_count,
+                .get_max_gfxp_wfi_timeout_count =
+                        gr_gp10b_get_max_gfxp_wfi_timeout_count,
         },
         .fb = {
                 .reset = fb_gk20a_reset,
diff --git a/drivers/gpu/nvgpu/common/linux/vgpu/gv11b/vgpu_hal_gv11b.c b/drivers/gpu/nvgpu/common/linux/vgpu/gv11b/vgpu_hal_gv11b.c
index a7491b20..a470377c 100644
--- a/drivers/gpu/nvgpu/common/linux/vgpu/gv11b/vgpu_hal_gv11b.c
+++ b/drivers/gpu/nvgpu/common/linux/vgpu/gv11b/vgpu_hal_gv11b.c
@@ -248,6 +248,10 @@ static const struct gpu_ops vgpu_gv11b_ops = {
                         gr_gv11b_handle_tpc_sm_ecc_exception,
                 .decode_egpc_addr = gv11b_gr_decode_egpc_addr,
                 .init_ctxsw_hdr_data = gr_gp10b_init_ctxsw_hdr_data,
+                .init_gfxp_wfi_timeout_count =
+                        gr_gv11b_init_gfxp_wfi_timeout_count,
+                .get_max_gfxp_wfi_timeout_count =
+                        gr_gv11b_get_max_gfxp_wfi_timeout_count,
         },
         .fb = {
                 .reset = gv11b_fb_reset,