diff --git a/drivers/gpu/drm/i915/i915_debugfs.c b/drivers/gpu/drm/i915/i915_debugfs.c
index 030263870ba6..4981d1f6e7a5 100644
--- a/drivers/gpu/drm/i915/i915_debugfs.c
+++ b/drivers/gpu/drm/i915/i915_debugfs.c
@@ -376,9 +376,9 @@ static void print_batch_pool_stats(struct seq_file *m,
 	print_file_stats(m, "[k]batch pool", stats);
 }
 
-static int per_file_ctx_stats(int idx, void *ptr, void *data)
+static void per_file_ctx_stats(struct i915_gem_context *ctx,
+		struct file_stats *stats)
 {
-	struct i915_gem_context *ctx = ptr;
 	struct intel_engine_cs *engine;
 	enum intel_engine_id id;
 
@@ -386,12 +386,10 @@ static int per_file_ctx_stats(int idx, void *ptr, void *data)
 		struct intel_context *ce = to_intel_context(ctx, engine);
 
 		if (ce->state)
-			per_file_stats(ce->state->obj, data);
+			per_file_stats(ce->state->obj, stats);
 		if (ce->ring)
-			per_file_stats(ce->ring->vma->obj, data);
+			per_file_stats(ce->ring->vma->obj, stats);
 	}
-
-	return 0;
 }
 
 static void print_context_stats(struct seq_file *m,
@@ -405,11 +403,15 @@ static void print_context_stats(struct seq_file *m,
 
 	mutex_lock(&dev->struct_mutex);
 	if (dev_priv->kernel_context)
-		per_file_ctx_stats(0, dev_priv->kernel_context, &stats);
+		per_file_ctx_stats(dev_priv->kernel_context, &stats);
 
 	list_for_each_entry(file, &dev->filelist, lhead) {
 		struct drm_i915_file_private *fpriv = file->driver_priv;
-		idr_for_each(&fpriv->context_idr, per_file_ctx_stats, &stats);
+		struct i915_gem_context *ctx;
+		unsigned long index;
+
+		xa_for_each(&fpriv->contexts, index, ctx)
+			per_file_ctx_stats(ctx, &stats);
 	}
 	mutex_unlock(&dev->struct_mutex);
 
@@ -2078,16 +2080,14 @@ static int i915_swizzle_info(struct seq_file *m, void *data)
 	return 0;
 }
 
-static int per_file_ctx(int id, void *ptr, void *data)
+static void per_file_ctx(struct i915_gem_context *ctx, struct seq_file *m)
 {
-	struct i915_gem_context *ctx = ptr;
-	struct seq_file *m = data;
 	struct i915_hw_ppgtt *ppgtt = ctx->ppgtt;
 
 	if (!ppgtt) {
 		seq_printf(m, "  no ppgtt for context %d\n",
 			   ctx->user_handle);
-		return 0;
+		return;
 	}
 
 	if (i915_gem_context_is_default(ctx))
@@ -2095,8 +2095,6 @@ static int per_file_ctx(int id, void *ptr, void *data)
 	else
 		seq_printf(m, "  context %d:\n", ctx->user_handle);
 	ppgtt->debug_dump(ppgtt, m);
-
-	return 0;
 }
 
 static void gen8_ppgtt_info(struct seq_file *m,
@@ -2176,6 +2174,8 @@ static int i915_ppgtt_info(struct seq_file *m, void *data)
 	list_for_each_entry_reverse(file, &dev->filelist, lhead) {
 		struct drm_i915_file_private *file_priv = file->driver_priv;
 		struct task_struct *task;
+		struct i915_gem_context *ctx;
+		unsigned long index;
 
 		task = get_pid_task(file->pid, PIDTYPE_PID);
 		if (!task) {
@@ -2184,8 +2184,8 @@ static int i915_ppgtt_info(struct seq_file *m, void *data)
 		}
 		seq_printf(m, "\nproc: %s\n", task->comm);
 		put_task_struct(task);
-		idr_for_each(&file_priv->context_idr, per_file_ctx,
-			     (void *)(unsigned long)m);
+		xa_for_each(&file_priv->contexts, index, ctx)
+			per_file_ctx(ctx, m);
 	}
 
 out_rpm:
diff --git a/drivers/gpu/drm/i915/i915_drv.h b/drivers/gpu/drm/i915/i915_drv.h
index 4b5ce517cbcf..1b7663258f42 100644
--- a/drivers/gpu/drm/i915/i915_drv.h
+++ b/drivers/gpu/drm/i915/i915_drv.h
@@ -213,7 +213,7 @@ struct drm_i915_file_private {
  */
 #define DRM_I915_THROTTLE_JIFFIES msecs_to_jiffies(20)
 	} mm;
-	struct idr context_idr;
+	struct xarray contexts;
 
 	struct intel_rps_client {
 		atomic_t boosts;
@@ -3118,7 +3118,7 @@ void i915_gem_object_save_bit_17_swizzle(struct drm_i915_gem_object *obj,
 static inline struct i915_gem_context *
 __i915_gem_context_lookup_rcu(struct drm_i915_file_private *file_priv, u32 id)
 {
-	return idr_find(&file_priv->context_idr, id);
+	return xa_load(&file_priv->contexts, id);
 }
 
 static inline struct i915_gem_context *
diff --git a/drivers/gpu/drm/i915/i915_gem_context.c b/drivers/gpu/drm/i915/i915_gem_context.c
index 9db04b2e65cf..033a6fd72e40 100644
--- a/drivers/gpu/drm/i915/i915_gem_context.c
+++ b/drivers/gpu/drm/i915/i915_gem_context.c
@@ -350,12 +350,12 @@ __create_hw_context(struct drm_i915_private *dev_priv,
 	/* Default context will never have a file_priv */
 	ret = DEFAULT_CONTEXT_HANDLE;
 	if (file_priv) {
-		ret = idr_alloc(&file_priv->context_idr, ctx,
-				DEFAULT_CONTEXT_HANDLE, 0, GFP_KERNEL);
+		ret = xa_alloc(&file_priv->contexts, &ctx->user_handle, ctx,
+				XA_LIMIT(DEFAULT_CONTEXT_HANDLE, UINT_MAX),
+				GFP_KERNEL);
 		if (ret < 0)
 			goto err_lut;
 	}
-	ctx->user_handle = ret;
 
 	ctx->file_priv = file_priv;
 	if (file_priv) {
@@ -384,7 +384,7 @@ __create_hw_context(struct drm_i915_private *dev_priv,
 
 err_pid:
 	put_pid(ctx->pid);
-	idr_remove(&file_priv->context_idr, ctx->user_handle);
+	xa_erase(&file_priv->contexts, ctx->user_handle);
 err_lut:
 	context_close(ctx);
 	return ERR_PTR(ret);
@@ -393,7 +393,7 @@ __create_hw_context(struct drm_i915_private *dev_priv,
 static void __destroy_hw_context(struct i915_gem_context *ctx,
 				 struct drm_i915_file_private *file_priv)
 {
-	idr_remove(&file_priv->context_idr, ctx->user_handle);
+	xa_erase(&file_priv->contexts, ctx->user_handle);
 	context_close(ctx);
 }
 
@@ -597,29 +597,19 @@ void i915_gem_contexts_fini(struct drm_i915_private *i915)
 	ida_destroy(&i915->contexts.hw_ida);
 }
 
-static int context_idr_cleanup(int id, void *p, void *data)
-{
-	struct i915_gem_context *ctx = p;
-
-	context_close(ctx);
-	return 0;
-}
-
 int i915_gem_context_open(struct drm_i915_private *i915,
 			  struct drm_file *file)
 {
 	struct drm_i915_file_private *file_priv = file->driver_priv;
 	struct i915_gem_context *ctx;
 
-	idr_init(&file_priv->context_idr);
+	xa_init_flags(&file_priv->contexts, XA_FLAGS_ALLOC);
 
 	mutex_lock(&i915->drm.struct_mutex);
 	ctx = i915_gem_create_context(i915, file_priv);
 	mutex_unlock(&i915->drm.struct_mutex);
-	if (IS_ERR(ctx)) {
-		idr_destroy(&file_priv->context_idr);
+	if (IS_ERR(ctx))
 		return PTR_ERR(ctx);
-	}
 
 	GEM_BUG_ON(i915_gem_context_is_kernel(ctx));
 
@@ -629,11 +619,14 @@ int i915_gem_context_open(struct drm_i915_private *i915,
 void i915_gem_context_close(struct drm_file *file)
 {
 	struct drm_i915_file_private *file_priv = file->driver_priv;
+	struct i915_gem_context *ctx;
+	unsigned long index;
 
 	lockdep_assert_held(&file_priv->dev_priv->drm.struct_mutex);
 
-	idr_for_each(&file_priv->context_idr, context_idr_cleanup, NULL);
-	idr_destroy(&file_priv->context_idr);
+	xa_for_each(&file_priv->contexts, index, ctx)
+		context_close(ctx);
+	xa_destroy(&file_priv->contexts);
 }
 
 static struct i915_request *
