return 0;
 }
 
+static void proto_context_close(struct i915_gem_proto_context *pc)
+{
+       if (pc->vm)
+               i915_vm_put(pc->vm);
+       kfree(pc);
+}
+
+static struct i915_gem_proto_context *
+proto_context_create(struct drm_i915_private *i915, unsigned int flags)
+{
+       struct i915_gem_proto_context *pc, *err;
+
+       pc = kzalloc(sizeof(*pc), GFP_KERNEL);
+       if (!pc)
+               return ERR_PTR(-ENOMEM);
+
+       pc->user_flags = BIT(UCONTEXT_BANNABLE) |
+                        BIT(UCONTEXT_RECOVERABLE);
+       if (i915->params.enable_hangcheck)
+               pc->user_flags |= BIT(UCONTEXT_PERSISTENCE);
+       pc->sched.priority = I915_PRIORITY_NORMAL;
+
+       if (flags & I915_CONTEXT_CREATE_FLAGS_SINGLE_TIMELINE) {
+               if (!HAS_EXECLISTS(i915)) {
+                       err = ERR_PTR(-EINVAL);
+                       goto proto_close;
+               }
+               pc->single_timeline = true;
+       }
+
+       return pc;
+
+proto_close:
+       proto_context_close(pc);
+       return err;
+}
+
 static struct i915_address_space *
 context_get_vm_rcu(struct i915_gem_context *ctx)
 {
 }
 
 static struct i915_gem_context *
-__create_context(struct drm_i915_private *i915)
+__create_context(struct drm_i915_private *i915,
+                const struct i915_gem_proto_context *pc)
 {
        struct i915_gem_context *ctx;
        struct i915_gem_engines *e;
 
        kref_init(&ctx->ref);
        ctx->i915 = i915;
-       ctx->sched.priority = I915_PRIORITY_NORMAL;
+       ctx->sched = pc->sched;
        mutex_init(&ctx->mutex);
        INIT_LIST_HEAD(&ctx->link);
 
         * is no remap info, it will be a NOP. */
        ctx->remap_slice = ALL_L3_SLICES(i915);
 
-       i915_gem_context_set_bannable(ctx);
-       i915_gem_context_set_recoverable(ctx);
-       __context_set_persistence(ctx, true /* cgroup hook? */);
+       ctx->user_flags = pc->user_flags;
 
        for (i = 0; i < ARRAY_SIZE(ctx->hang_timestamp); i++)
                ctx->hang_timestamp[i] = jiffies - CONTEXT_FAST_HANG_JIFFIES;
 }
 
 static struct i915_gem_context *
-i915_gem_create_context(struct drm_i915_private *i915, unsigned int flags)
+i915_gem_create_context(struct drm_i915_private *i915,
+                       const struct i915_gem_proto_context *pc)
 {
        struct i915_gem_context *ctx;
        int ret;
 
-       if (flags & I915_CONTEXT_CREATE_FLAGS_SINGLE_TIMELINE &&
-           !HAS_EXECLISTS(i915))
-               return ERR_PTR(-EINVAL);
-
-       ctx = __create_context(i915);
+       ctx = __create_context(i915, pc);
        if (IS_ERR(ctx))
                return ctx;
 
-       if (HAS_FULL_PPGTT(i915)) {
+       if (pc->vm) {
+               /* __assign_ppgtt() requires this mutex to be held */
+               mutex_lock(&ctx->mutex);
+               __assign_ppgtt(ctx, pc->vm);
+               mutex_unlock(&ctx->mutex);
+       } else if (HAS_FULL_PPGTT(i915)) {
                struct i915_ppgtt *ppgtt;
 
                ppgtt = i915_ppgtt_create(&i915->gt);
                        return ERR_CAST(ppgtt);
                }
 
+               /* __assign_ppgtt() requires this mutex to be held */
                mutex_lock(&ctx->mutex);
                __assign_ppgtt(ctx, &ppgtt->vm);
                mutex_unlock(&ctx->mutex);
 
+               /* __assign_ppgtt() takes another reference for us */
                i915_vm_put(&ppgtt->vm);
        }
 
-       if (flags & I915_CONTEXT_CREATE_FLAGS_SINGLE_TIMELINE) {
+       if (pc->single_timeline) {
                ret = drm_syncobj_create(&ctx->syncobj,
                                         DRM_SYNCOBJ_CREATE_SIGNALED,
                                         NULL);
                          struct drm_file *file)
 {
        struct drm_i915_file_private *file_priv = file->driver_priv;
+       struct i915_gem_proto_context *pc;
        struct i915_gem_context *ctx;
        int err;
        u32 id;
        /* 0 reserved for invalid/unassigned ppgtt */
        xa_init_flags(&file_priv->vm_xa, XA_FLAGS_ALLOC1);
 
-       ctx = i915_gem_create_context(i915, 0);
+       pc = proto_context_create(i915, 0);
+       if (IS_ERR(pc)) {
+               err = PTR_ERR(pc);
+               goto err;
+       }
+
+       ctx = i915_gem_create_context(i915, pc);
+       proto_context_close(pc);
        if (IS_ERR(ctx)) {
                err = PTR_ERR(ctx);
                goto err;
 {
        struct drm_i915_private *i915 = to_i915(dev);
        struct drm_i915_gem_context_create_ext *args = data;
+       struct i915_gem_proto_context *pc;
        struct create_ext ext_data;
        int ret;
        u32 id;
                return -EIO;
        }
 
-       ext_data.ctx = i915_gem_create_context(i915, args->flags);
+       pc = proto_context_create(i915, args->flags);
+       if (IS_ERR(pc))
+               return PTR_ERR(pc);
+
+       ext_data.ctx = i915_gem_create_context(i915, pc);
+       proto_context_close(pc);
        if (IS_ERR(ext_data.ctx))
                return PTR_ERR(ext_data.ctx);
 
 
 struct i915_gem_context *
 live_context(struct drm_i915_private *i915, struct file *file)
 {
+       struct i915_gem_proto_context *pc;
        struct i915_gem_context *ctx;
        int err;
        u32 id;
 
-       ctx = i915_gem_create_context(i915, 0);
+       pc = proto_context_create(i915, 0);
+       if (IS_ERR(pc))
+               return ERR_CAST(pc);
+
+       ctx = i915_gem_create_context(i915, pc);
+       proto_context_close(pc);
        if (IS_ERR(ctx))
                return ctx;
 
 kernel_context(struct drm_i915_private *i915)
 {
        struct i915_gem_context *ctx;
+       struct i915_gem_proto_context *pc;
+
+       pc = proto_context_create(i915, 0);
+       if (IS_ERR(pc))
+               return ERR_CAST(pc);
 
-       ctx = i915_gem_create_context(i915, 0);
+       ctx = i915_gem_create_context(i915, pc);
+       proto_context_close(pc);
        if (IS_ERR(ctx))
                return ctx;