2 files changed, 2 insertions, 3 deletions

diff --git a/drivers/gpu/nvgpu/common/linux/kmem.c b/drivers/gpu/nvgpu/common/linux/kmem.c
index 87aeb184..a2e49e7a 100644
--- a/drivers/gpu/nvgpu/common/linux/kmem.c
+++ b/drivers/gpu/nvgpu/common/linux/kmem.c
@@ -16,7 +16,6 @@
 
 #include <linux/mutex.h>
 #include <linux/slab.h>
-#include <linux/atomic.h>
 #include <linux/rbtree.h>
 #include <linux/debugfs.h>
 #include <linux/spinlock.h>
@@ -25,6 +24,7 @@
 #include <linux/stacktrace.h>
 
 #include <nvgpu/kmem.h>
+#include <nvgpu/atomic.h>
 
 #include "gk20a/gk20a.h"
 
diff --git a/drivers/gpu/nvgpu/common/mm/lockless_allocator.c b/drivers/gpu/nvgpu/common/mm/lockless_allocator.c
index 93c870b5..dc72d8bf 100644
--- a/drivers/gpu/nvgpu/common/mm/lockless_allocator.c
+++ b/drivers/gpu/nvgpu/common/mm/lockless_allocator.c
@@ -14,8 +14,7 @@
  * along with this program.  If not, see <http://www.gnu.org/licenses/>.
  */
 
-#include <linux/atomic.h>
-
+#include <nvgpu/atomic.h>
 #include <nvgpu/allocator.h>
 #include <nvgpu/kmem.h>