2 files changed, 4 insertions, 2 deletions

diff --git a/drivers/gpu/nvgpu/common/pmu/pmu.c b/drivers/gpu/nvgpu/common/pmu/pmu.c
index 58108722..63597d10 100644
--- a/drivers/gpu/nvgpu/common/pmu/pmu.c
+++ b/drivers/gpu/nvgpu/common/pmu/pmu.c
@@ -16,6 +16,7 @@
 #include <nvgpu/log.h>
 #include <nvgpu/pmuif/nvgpu_gpmu_cmdif.h>
 #include <nvgpu/enabled.h>
+#include <nvgpu/barrier.h>
 
 #include "gk20a/gk20a.h"
 
@@ -394,7 +395,7 @@ void nvgpu_pmu_state_change(struct gk20a *g, u32 pmu_state,
         }
 
         /* make status visible */
-        smp_mb();
+        nvgpu_smp_mb();
 }
 
 static int nvgpu_pg_init_task(void *arg)
diff --git a/drivers/gpu/nvgpu/common/pmu/pmu_pg.c b/drivers/gpu/nvgpu/common/pmu/pmu_pg.c
index 935ae95a..b435f4a7 100644
--- a/drivers/gpu/nvgpu/common/pmu/pmu_pg.c
+++ b/drivers/gpu/nvgpu/common/pmu/pmu_pg.c
@@ -14,6 +14,7 @@
 #include <nvgpu/pmu.h>
 #include <nvgpu/log.h>
 #include <nvgpu/pmuif/nvgpu_gpmu_cmdif.h>
+#include <nvgpu/barrier.h>
 
 #include "gk20a/gk20a.h"
 
@@ -84,7 +85,7 @@ static void pmu_handle_pg_elpg_msg(struct gk20a *g, struct pmu_msg *msg,
                                         true);
                                 WRITE_ONCE(pmu->mscg_stat, PMU_MSCG_DISABLED);
                                 /* make status visible */
-                                smp_mb();
+                                nvgpu_smp_mb();
                         } else
                                 nvgpu_pmu_state_change(g, PMU_STATE_ELPG_BOOTED,
                                         true);