diff --git a/drivers/gpu/drm/drm_context.c b/drivers/gpu/drm/drm_context.c
index 1febcd3..09af26c 100644
--- a/drivers/gpu/drm/drm_context.c
+++ b/drivers/gpu/drm/drm_context.c
@@ -53,6 +53,9 @@ struct drm_ctx_list {
  */
 void drm_legacy_ctxbitmap_free(struct drm_device * dev, int ctx_handle)
 {
+	if (!drm_core_check_feature(dev, DRIVER_KMS_LEGACY_CONTEXT))
+		return -EINVAL;
+
 	mutex_lock(&dev->struct_mutex);
 	idr_remove(&dev->ctx_idr, ctx_handle);
 	mutex_unlock(&dev->struct_mutex);
@@ -87,6 +90,9 @@ static int drm_legacy_ctxbitmap_next(struct drm_device * dev)
  */
 int drm_legacy_ctxbitmap_init(struct drm_device * dev)
 {
+	if (!drm_core_check_feature(dev, DRIVER_KMS_LEGACY_CONTEXT))
+		return -EINVAL;
+
 	idr_init(&dev->ctx_idr);
 	return 0;
 }
@@ -101,6 +107,9 @@ int drm_legacy_ctxbitmap_init(struct drm_device * dev)
  */
 void drm_legacy_ctxbitmap_cleanup(struct drm_device * dev)
 {
+	if (!drm_core_check_feature(dev, DRIVER_KMS_LEGACY_CONTEXT))
+		return -EINVAL;
+
 	mutex_lock(&dev->struct_mutex);
 	idr_destroy(&dev->ctx_idr);
 	mutex_unlock(&dev->struct_mutex);
@@ -119,6 +128,9 @@ void drm_legacy_ctxbitmap_flush(struct drm_device *dev, struct drm_file *file)
 {
 	struct drm_ctx_list *pos, *tmp;
 
+	if (!drm_core_check_feature(dev, DRIVER_KMS_LEGACY_CONTEXT))
+		return -EINVAL;
+
 	mutex_lock(&dev->ctxlist_mutex);
 
 	list_for_each_entry_safe(pos, tmp, &dev->ctxlist, head) {
@@ -161,6 +173,9 @@ int drm_legacy_getsareactx(struct drm_device *dev, void *data,
 	struct drm_local_map *map;
 	struct drm_map_list *_entry;
 
+	if (!drm_core_check_feature(dev, DRIVER_KMS_LEGACY_CONTEXT))
+		return -EINVAL;
+
 	mutex_lock(&dev->struct_mutex);
 
 	map = idr_find(&dev->ctx_idr, request->ctx_id);
@@ -205,6 +220,9 @@ int drm_legacy_setsareactx(struct drm_device *dev, void *data,
 	struct drm_local_map *map = NULL;
 	struct drm_map_list *r_list = NULL;
 
+	if (!drm_core_check_feature(dev, DRIVER_KMS_LEGACY_CONTEXT))
+		return -EINVAL;
+
 	mutex_lock(&dev->struct_mutex);
 	list_for_each_entry(r_list, &dev->maplist, head) {
 		if (r_list->map
@@ -311,6 +329,9 @@ int drm_legacy_resctx(struct drm_device *dev, void *data,
 	struct drm_ctx ctx;
 	int i;
 
+	if (!drm_core_check_feature(dev, DRIVER_KMS_LEGACY_CONTEXT))
+		return -EINVAL;
+
 	if (res->count >= DRM_RESERVED_CONTEXTS) {
 		memset(&ctx, 0, sizeof(ctx));
 		for (i = 0; i < DRM_RESERVED_CONTEXTS; i++) {
@@ -341,6 +362,9 @@ int drm_legacy_addctx(struct drm_device *dev, void *data,
 	struct drm_ctx_list *ctx_entry;
 	struct drm_ctx *ctx = data;
 
+	if (!drm_core_check_feature(dev, DRIVER_KMS_LEGACY_CONTEXT))
+		return -EINVAL;
+
 	ctx->handle = drm_legacy_ctxbitmap_next(dev);
 	if (_DRM_LOCKING_CONTEXT(ctx->handle) == DRM_KERNEL_CONTEXT) {
 		/* Skip kernel's context and get a new one. */
@@ -384,6 +408,9 @@ int drm_legacy_getctx(struct drm_device *dev, void *data,
 {
 	struct drm_ctx *ctx = data;
 
+	if (!drm_core_check_feature(dev, DRIVER_KMS_LEGACY_CONTEXT))
+		return -EINVAL;
+
 	/* This is 0, because we don't handle any context flags */
 	ctx->flags = 0;
 
@@ -406,6 +433,9 @@ int drm_legacy_switchctx(struct drm_device *dev, void *data,
 {
 	struct drm_ctx *ctx = data;
 
+	if (!drm_core_check_feature(dev, DRIVER_KMS_LEGACY_CONTEXT))
+		return -EINVAL;
+
 	DRM_DEBUG("%d\n", ctx->handle);
 	return drm_context_switch(dev, dev->last_context, ctx->handle);
 }
@@ -426,6 +456,9 @@ int drm_legacy_newctx(struct drm_device *dev, void *data,
 {
 	struct drm_ctx *ctx = data;
 
+	if (!drm_core_check_feature(dev, DRIVER_KMS_LEGACY_CONTEXT))
+		return -EINVAL;
+
 	DRM_DEBUG("%d\n", ctx->handle);
 	drm_context_switch_complete(dev, file_priv, ctx->handle);
 
@@ -448,6 +481,9 @@ int drm_legacy_rmctx(struct drm_device *dev, void *data,
 {
 	struct drm_ctx *ctx = data;
 
+	if (!drm_core_check_feature(dev, DRIVER_KMS_LEGACY_CONTEXT))
+		return -EINVAL;
+
 	DRM_DEBUG("%d\n", ctx->handle);
 	if (_DRM_LOCKING_CONTEXT(ctx->handle) != DRM_KERNEL_CONTEXT) {
 		if (dev->driver->context_dtor)
diff --git a/drivers/gpu/drm/drm_drv.c b/drivers/gpu/drm/drm_drv.c
index 48f7359..dc662e8 100644
--- a/drivers/gpu/drm/drm_drv.c
+++ b/drivers/gpu/drm/drm_drv.c
@@ -596,11 +596,13 @@ struct drm_device *drm_dev_alloc(struct drm_driver *driver,
 	if (drm_ht_create(&dev->map_hash, 12))
 		goto err_minors;
 
-	ret = drm_legacy_ctxbitmap_init(dev);
-	if (ret) {
-		DRM_ERROR("Cannot allocate memory for context bitmap.\n");
-		goto err_ht;
-	}
+	if (drm_core_check_feature(dev, DRIVER_KMS_LEGACY_CONTEXT))
+		ret = drm_legacy_ctxbitmap_init(dev);
+		if (ret) {
+			DRM_ERROR(
+				"Cannot allocate memory for context bitmap.\n");
+			goto err_ht;
+		}
 
 	if (drm_core_check_feature(dev, DRIVER_GEM)) {
 		ret = drm_gem_init(dev);
