diff --git a/drivers/gpu/nvgpu/common/linux/ioctl_ctrl.c b/drivers/gpu/nvgpu/common/linux/ioctl_ctrl.c
index cda2ce463..73a8131da 100644
--- a/drivers/gpu/nvgpu/common/linux/ioctl_ctrl.c
+++ b/drivers/gpu/nvgpu/common/linux/ioctl_ctrl.c
@@ -782,7 +782,14 @@ static int gk20a_ctrl_vsm_mapping(struct gk20a *g,
 
 	for (i = 0; i < gr->no_of_sm; i++) {
 		vsms_buf[i].gpc_index = gr->sm_to_cluster[i].gpc_index;
-		vsms_buf[i].tpc_index = gr->sm_to_cluster[i].tpc_index;
+		if (g->ops.gr.get_nonpes_aware_tpc)
+			vsms_buf[i].tpc_index =
+				g->ops.gr.get_nonpes_aware_tpc(g,
+					gr->sm_to_cluster[i].gpc_index,
+					gr->sm_to_cluster[i].tpc_index);
+		else
+			vsms_buf[i].tpc_index =
+				gr->sm_to_cluster[i].tpc_index;
 	}
 
 	err = copy_to_user((void __user *)(uintptr_t)
diff --git a/drivers/gpu/nvgpu/gk20a/gk20a.h b/drivers/gpu/nvgpu/gk20a/gk20a.h
index 49f2a34a4..be8917c77 100644
--- a/drivers/gpu/nvgpu/gk20a/gk20a.h
+++ b/drivers/gpu/nvgpu/gk20a/gk20a.h
@@ -481,6 +481,7 @@ struct gpu_ops {
 			struct channel_gk20a *c);
 		int (*commit_global_ctx_buffers)(struct gk20a *g,
 			struct channel_gk20a *c, bool patch);
+		u32 (*get_nonpes_aware_tpc)(struct gk20a *g, u32 gpc, u32 tpc);
 	} gr;
 	struct {
 		void (*init_hw)(struct gk20a *g);
diff --git a/drivers/gpu/nvgpu/gv100/hal_gv100.c b/drivers/gpu/nvgpu/gv100/hal_gv100.c
index 008404486..1d6f59b37 100644
--- a/drivers/gpu/nvgpu/gv100/hal_gv100.c
+++ b/drivers/gpu/nvgpu/gv100/hal_gv100.c
@@ -449,6 +449,7 @@ static const struct gpu_ops gv100_ops = {
 		.alloc_global_ctx_buffers = gr_gk20a_alloc_global_ctx_buffers,
 		.map_global_ctx_buffers = gr_gk20a_map_global_ctx_buffers,
 		.commit_global_ctx_buffers = gr_gk20a_commit_global_ctx_buffers,
+		.get_nonpes_aware_tpc = gr_gv11b_get_nonpes_aware_tpc,
 	},
 	.fb = {
 		.reset = gv100_fb_reset,
diff --git a/drivers/gpu/nvgpu/gv11b/gr_gv11b.c b/drivers/gpu/nvgpu/gv11b/gr_gv11b.c
index f57be9dd6..694ff8ad5 100644
--- a/drivers/gpu/nvgpu/gv11b/gr_gv11b.c
+++ b/drivers/gpu/nvgpu/gv11b/gr_gv11b.c
@@ -2661,7 +2661,7 @@ void gr_gv11b_detect_sm_arch(struct gk20a *g)
 		gr_gpc0_tpc0_sm_arch_warp_count_v(v);
 }
 
-static u32 gr_gv11b_get_nonpes_aware_tpc(struct gk20a *g, u32 gpc, u32 tpc)
+u32 gr_gv11b_get_nonpes_aware_tpc(struct gk20a *g, u32 gpc, u32 tpc)
 {
 	u32 tpc_new = 0;
 	u32 temp;
@@ -2691,7 +2691,7 @@ void gr_gv11b_program_sm_id_numbering(struct gk20a *g,
 	u32 tpc_offset = tpc_in_gpc_stride * tpc;
 	u32 global_tpc_index = g->gr.sm_to_cluster[smid].global_tpc_index;
 
-	tpc = gr_gv11b_get_nonpes_aware_tpc(g, gpc, tpc);
+	tpc = g->ops.gr.get_nonpes_aware_tpc(g, gpc, tpc);
 
 	gk20a_writel(g, gr_gpc0_tpc0_sm_cfg_r() + gpc_offset + tpc_offset,
 		gr_gpc0_tpc0_sm_cfg_tpc_id_f(global_tpc_index));
diff --git a/drivers/gpu/nvgpu/gv11b/gr_gv11b.h b/drivers/gpu/nvgpu/gv11b/gr_gv11b.h
index f6f05a3b9..fb4c8b69e 100644
--- a/drivers/gpu/nvgpu/gv11b/gr_gv11b.h
+++ b/drivers/gpu/nvgpu/gv11b/gr_gv11b.h
@@ -248,4 +248,5 @@ int gr_gv11b_create_priv_addr_table(struct gk20a *g,
 	u32 addr,
 	u32 *priv_addr_table,
 	u32 *num_registers);
+u32 gr_gv11b_get_nonpes_aware_tpc(struct gk20a *g, u32 gpc, u32 tpc);
 #endif
diff --git a/drivers/gpu/nvgpu/gv11b/hal_gv11b.c b/drivers/gpu/nvgpu/gv11b/hal_gv11b.c
index ce8f5051e..f9ac1f2a4 100644
--- a/drivers/gpu/nvgpu/gv11b/hal_gv11b.c
+++ b/drivers/gpu/nvgpu/gv11b/hal_gv11b.c
@@ -415,6 +415,7 @@ static const struct gpu_ops gv11b_ops = {
 		.alloc_global_ctx_buffers = gr_gk20a_alloc_global_ctx_buffers,
 		.map_global_ctx_buffers = gr_gk20a_map_global_ctx_buffers,
 		.commit_global_ctx_buffers = gr_gk20a_commit_global_ctx_buffers,
+		.get_nonpes_aware_tpc = gr_gv11b_get_nonpes_aware_tpc,
 	},
 	.fb = {
 		.reset = gv11b_fb_reset,