2 files changed, 6 insertions, 2 deletions

diff --git a/drivers/gpu/nvgpu/gp10b/gr_gp10b.c b/drivers/gpu/nvgpu/gp10b/gr_gp10b.c
index 03462d5f..33a52db5 100644
--- a/drivers/gpu/nvgpu/gp10b/gr_gp10b.c
+++ b/drivers/gpu/nvgpu/gp10b/gr_gp10b.c
@@ -475,7 +475,7 @@ static int gr_gp10b_alloc_gr_ctx(struct gk20a *g,
         if (err)
                 return err;
 
-        if (flags == NVGPU_GR_PREEMPTION_MODE_GFXP) {
+        if (flags & NVGPU_ALLOC_OBJ_FLAGS_GFXP) {
                 u32 spill_size =
                         gr_gpc0_swdx_rm_spill_buffer_size_256b_default_v();
                 u32 pagepool_size = g->ops.gr.pagepool_default_size(g) *
@@ -528,7 +528,7 @@ static int gr_gp10b_alloc_gr_ctx(struct gk20a *g,
         }
 
         if (class == PASCAL_COMPUTE_A) {
-                if (flags == NVGPU_GR_PREEMPTION_MODE_CILP)
+                if (flags & NVGPU_ALLOC_OBJ_FLAGS_CILP)
                         (*gr_ctx)->preempt_mode = NVGPU_GR_PREEMPTION_MODE_CILP;
                 else
                         (*gr_ctx)->preempt_mode = NVGPU_GR_PREEMPTION_MODE_CTA;
diff --git a/include/uapi/linux/nvgpu-t18x.h b/include/uapi/linux/nvgpu-t18x.h
index 76814415..b2a75143 100644
--- a/include/uapi/linux/nvgpu-t18x.h
+++ b/include/uapi/linux/nvgpu-t18x.h
@@ -39,6 +39,10 @@
  */
 #define NVGPU_ALLOC_GPFIFO_FLAGS_REPLAYABLE_FAULTS_ENABLE   (1 << 2)
 
+/* Flags in nvgpu_alloc_obj_ctx_args.flags */
+#define NVGPU_ALLOC_OBJ_FLAGS_GFXP              (1 << 1)
+#define NVGPU_ALLOC_OBJ_FLAGS_CILP              (1 << 2)
+
 #endif /* _UAPI__LINUX_NVGPU_T18X_IOCTL_H_ */