1 files changed, 6 insertions, 5 deletions

diff --git a/drivers/gpu/nvgpu/gk20a/gr_gk20a.h b/drivers/gpu/nvgpu/gk20a/gr_gk20a.h
index d795a3fc..8ff2cfd4 100644
--- a/drivers/gpu/nvgpu/gk20a/gr_gk20a.h
+++ b/drivers/gpu/nvgpu/gk20a/gr_gk20a.h
@@ -28,6 +28,7 @@
 
 #include "gr_ctx_gk20a.h"
 #include "mm_gk20a.h"
+#include <nvgpu/power_features/pg.h>
 
 #include <nvgpu/comptags.h>
 #include <nvgpu/cond.h>
@@ -598,16 +599,16 @@ u32 gk20a_gr_get_sm_no_lock_down_hww_global_esr_mask(struct gk20a *g);
 #define gr_gk20a_elpg_protected_call(g, func) \
         ({ \
                 int err = 0; \
-                if ((g->support_pmu) && (g->elpg_enabled)) {\
-                        err = nvgpu_pmu_disable_elpg(g); \
+                if (g->support_pmu) {\
+                        err = nvgpu_pg_elpg_disable(g);\
                         if (err != 0) {\
-                                nvgpu_pmu_enable_elpg(g); \
+                                err = nvgpu_pg_elpg_enable(g); \
                         } \
                 } \
                 if (err == 0) { \
                         err = func; \
-                        if ((g->support_pmu) && (g->elpg_enabled)) {\
-                                nvgpu_pmu_enable_elpg(g); \
+                        if (g->support_pmu) {\
+                                (void)nvgpu_pg_elpg_enable(g); \
                         } \
                 } \
                 err; \