return 0;
 }
+EXPORT_SYMBOL_GPL(set_thread_tidr);
 
 #endif /* CONFIG_PPC64 */
 
 
        ctx->pid = NULL; /* Set in start work ioctl */
        mutex_init(&ctx->mapping_lock);
        ctx->mapping = NULL;
+       ctx->tidr = 0;
+       ctx->assign_tidr = false;
 
        if (cxl_is_power8()) {
                spin_lock_init(&ctx->sste_lock);
 
        struct list_head extra_irq_contexts;
 
        struct mm_struct *mm;
+
+       u16 tidr;
+       bool assign_tidr;
 };
 
 struct cxl_irq_info;
 
                 */
                attr->pid = mm->context.id;
                mmput(mm);
+               attr->tid = task->thread.tidr;
        } else {
                attr->pid = 0;
+               attr->tid = 0;
        }
-       attr->tid = 0;
        return 0;
 }
 EXPORT_SYMBOL_GPL(cxllib_get_PE_attributes);
 
         * flags are set it's invalid
         */
        if (work.reserved1 || work.reserved2 || work.reserved3 ||
-           work.reserved4 || work.reserved5 || work.reserved6 ||
+           work.reserved4 || work.reserved5 ||
            (work.flags & ~CXL_START_WORK_ALL)) {
                rc = -EINVAL;
                goto out;
                rc =  -EINVAL;
                goto out;
        }
+
        if ((rc = afu_register_irqs(ctx, work.num_interrupts)))
                goto out;
 
        if (work.flags & CXL_START_WORK_AMR)
                amr = work.amr & mfspr(SPRN_UAMOR);
 
+       if (work.flags & CXL_START_WORK_TID)
+               ctx->assign_tidr = true;
+
        ctx->mmio_err_ff = !!(work.flags & CXL_START_WORK_ERR_FF);
 
        /*
                goto out;
        }
 
-       ctx->status = STARTED;
        rc = 0;
+       if (work.flags & CXL_START_WORK_TID) {
+               work.tid = ctx->tidr;
+               if (copy_to_user(uwork, &work, sizeof(work)))
+                       rc = -EFAULT;
+       }
+
+       ctx->status = STARTED;
+
 out:
        mutex_unlock(&ctx->status_mutex);
        return rc;
 
 #include <linux/uaccess.h>
 #include <linux/delay.h>
 #include <asm/synch.h>
+#include <asm/switch_to.h>
 #include <misc/cxl-base.h>
 
 #include "cxl.h"
 static int process_element_entry_psl9(struct cxl_context *ctx, u64 wed, u64 amr)
 {
        u32 pid;
+       int rc;
 
        cxl_assign_psn_space(ctx);
 
                pid = ctx->mm->context.id;
        }
 
-       ctx->elem->common.tid = 0;
+       /* Assign a unique TIDR (thread id) for the current thread */
+       if (!(ctx->tidr) && (ctx->assign_tidr)) {
+               rc = set_thread_tidr(current);
+               if (rc)
+                       return -ENODEV;
+               ctx->tidr = current->thread.tidr;
+               pr_devel("%s: current tidr: %d\n", __func__, ctx->tidr);
+       }
+
+       ctx->elem->common.tid = cpu_to_be32(ctx->tidr);
        ctx->elem->common.pid = cpu_to_be32(pid);
 
        ctx->elem->sr = cpu_to_be64(calculate_sr(ctx));
 
        __u64 work_element_descriptor;
        __u64 amr;
        __s16 num_interrupts;
-       __s16 reserved1;
-       __s32 reserved2;
+       __u16 tid;
+       __s32 reserved1;
+       __u64 reserved2;
        __u64 reserved3;
        __u64 reserved4;
        __u64 reserved5;
-       __u64 reserved6;
 };
 
 #define CXL_START_WORK_AMR             0x0000000000000001ULL
 #define CXL_START_WORK_NUM_IRQS                0x0000000000000002ULL
 #define CXL_START_WORK_ERR_FF          0x0000000000000004ULL
+#define CXL_START_WORK_TID             0x0000000000000008ULL
 #define CXL_START_WORK_ALL             (CXL_START_WORK_AMR |\
                                         CXL_START_WORK_NUM_IRQS |\
-                                        CXL_START_WORK_ERR_FF)
+                                        CXL_START_WORK_ERR_FF |\
+                                        CXL_START_WORK_TID)
 
 
 /* Possible modes that an afu can be in */