/* Only used by new UAPI. */
        struct xarray mappings;
+       struct host1x_memory_context *memory_context;
 };
 
 struct tegra_drm_client_ops {
        int (*submit)(struct tegra_drm_context *context,
                      struct drm_tegra_submit *args, struct drm_device *drm,
                      struct drm_file *file);
+       int (*get_streamid_offset)(struct tegra_drm_client *client, u32 *offset);
+       int (*can_use_memory_ctx)(struct tegra_drm_client *client, bool *supported);
 };
 
 int tegra_drm_submit(struct tegra_drm_context *context,
 
        struct tegra_drm_submit_data *job_data = job->user_data;
        u32 i;
 
+       if (job->memory_context)
+               host1x_memory_context_put(job->memory_context);
+
        for (i = 0; i < job_data->num_used_mappings; i++)
                tegra_drm_mapping_put(job_data->used_mappings[i].mapping);
 
                goto put_job;
        }
 
+       if (context->client->ops->get_streamid_offset) {
+               err = context->client->ops->get_streamid_offset(
+                       context->client, &job->engine_streamid_offset);
+               if (err) {
+                       SUBMIT_ERR(context, "failed to get streamid offset: %d", err);
+                       goto unpin_job;
+               }
+       }
+
+       if (context->memory_context && context->client->ops->can_use_memory_ctx) {
+               bool supported;
+
+               err = context->client->ops->can_use_memory_ctx(context->client, &supported);
+               if (err) {
+                       SUBMIT_ERR(context, "failed to detect if engine can use memory context: %d", err);
+                       goto unpin_job;
+               }
+
+               if (supported) {
+                       job->memory_context = context->memory_context;
+                       host1x_memory_context_get(job->memory_context);
+               }
+       } else if (context->client->ops->get_streamid_offset) {
+#ifdef CONFIG_IOMMU_API
+               struct iommu_fwspec *spec;
+
+               /*
+                * Job submission will need to temporarily change stream ID,
+                * so need to tell it what to change it back to.
+                */
+               spec = dev_iommu_fwspec_get(context->client->base.dev);
+               if (spec && spec->num_ids > 0)
+                       job->engine_fallback_streamid = spec->ids[0] & 0xffff;
+               else
+                       job->engine_fallback_streamid = 0x7f;
+#else
+               job->engine_fallback_streamid = 0x7f;
+#endif
+       }
+
        /* Boot engine. */
        err = pm_runtime_resume_and_get(context->client->base.dev);
        if (err < 0) {
                SUBMIT_ERR(context, "could not power up engine: %d", err);
-               goto unpin_job;
+               goto put_memory_context;
        }
 
        job->user_data = job_data;
 
        goto put_job;
 
+put_memory_context:
+       if (job->memory_context)
+               host1x_memory_context_put(job->memory_context);
 unpin_job:
        host1x_job_unpin(job);
 put_job:
 
        struct tegra_drm_mapping *mapping;
        unsigned long id;
 
+       if (context->memory_context)
+               host1x_memory_context_put(context->memory_context);
+
        xa_for_each(&context->mappings, id, mapping)
                tegra_drm_mapping_put(mapping);
 
 
 int tegra_drm_ioctl_channel_open(struct drm_device *drm, void *data, struct drm_file *file)
 {
+       struct host1x *host = tegra_drm_to_host1x(drm->dev_private);
        struct tegra_drm_file *fpriv = file->driver_priv;
        struct tegra_drm *tegra = drm->dev_private;
        struct drm_tegra_channel_open *args = data;
                }
        }
 
+       /* Only allocate context if the engine supports context isolation. */
+       if (device_iommu_mapped(client->base.dev) && client->ops->can_use_memory_ctx) {
+               bool supported;
+
+               err = client->ops->can_use_memory_ctx(client, &supported);
+               if (err)
+                       goto put_channel;
+
+               if (supported)
+                       context->memory_context = host1x_memory_context_alloc(
+                               host, get_task_pid(current, PIDTYPE_TGID));
+
+               if (IS_ERR(context->memory_context)) {
+                       if (PTR_ERR(context->memory_context) != -EOPNOTSUPP) {
+                               err = PTR_ERR(context->memory_context);
+                               goto put_channel;
+                       } else {
+                               /*
+                                * OK, HW does not support contexts or contexts
+                                * are disabled.
+                                */
+                               context->memory_context = NULL;
+                       }
+               }
+       }
+
        err = xa_alloc(&fpriv->contexts, &args->context, context, XA_LIMIT(1, U32_MAX),
                       GFP_KERNEL);
        if (err < 0)
-               goto put_channel;
+               goto put_memctx;
 
        context->client = client;
        xa_init_flags(&context->mappings, XA_FLAGS_ALLOC1);
 
        return 0;
 
+put_memctx:
+       if (context->memory_context)
+               host1x_memory_context_put(context->memory_context);
 put_channel:
        host1x_channel_put(context->channel);
 free:
        struct tegra_drm_mapping *mapping;
        struct tegra_drm_context *context;
        enum dma_data_direction direction;
+       struct device *mapping_dev;
        int err = 0;
 
        if (args->flags & ~DRM_TEGRA_CHANNEL_MAP_READ_WRITE)
 
        kref_init(&mapping->ref);
 
+       if (context->memory_context)
+               mapping_dev = &context->memory_context->dev;
+       else
+               mapping_dev = context->client->base.dev;
+
        mapping->bo = tegra_gem_lookup(file, args->handle);
        if (!mapping->bo) {
                err = -EINVAL;
                goto put_gem;
        }
 
-       mapping->map = host1x_bo_pin(context->client->base.dev, mapping->bo, direction, NULL);
+       mapping->map = host1x_bo_pin(mapping_dev, mapping->bo, direction, NULL);
        if (IS_ERR(mapping->map)) {
                err = PTR_ERR(mapping->map);
                goto put_gem;