#define __ASM_EXTABLE_H
 
 #include <linux/stringify.h>
+#include <linux/bits.h>
 #include <asm/asm-const.h>
 
-#define EX_TYPE_NONE   0
-#define EX_TYPE_FIXUP  1
-#define EX_TYPE_BPF    2
-#define EX_TYPE_UACCESS        3
+#define EX_TYPE_NONE           0
+#define EX_TYPE_FIXUP          1
+#define EX_TYPE_BPF            2
+#define EX_TYPE_UA_STORE       3
+#define EX_TYPE_UA_LOAD_MEM    4
+#define EX_TYPE_UA_LOAD_REG    5
+
+#define EX_DATA_REG_ERR_SHIFT  0
+#define EX_DATA_REG_ERR                GENMASK(3, 0)
+
+#define EX_DATA_REG_ADDR_SHIFT 4
+#define EX_DATA_REG_ADDR       GENMASK(7, 4)
+
+#define EX_DATA_LEN_SHIFT      8
+#define EX_DATA_LEN            GENMASK(11, 8)
 
 #define __EX_TABLE(_section, _fault, _target, _type)                   \
        stringify_in_c(.section _section,"a";)                          \
        stringify_in_c(.short   0;)                                     \
        stringify_in_c(.previous)
 
-#define __EX_TABLE_UA(_section, _fault, _target, _type, _reg)          \
-       stringify_in_c(.section _section,"a";)                          \
-       stringify_in_c(.align   4;)                                     \
-       stringify_in_c(.long    (_fault) - .;)                          \
-       stringify_in_c(.long    (_target) - .;)                         \
-       stringify_in_c(.short   (_type);)                               \
-       stringify_in_c(.macro extable_reg reg;)                         \
-       stringify_in_c(.set .Lfound, 0;)                                \
-       stringify_in_c(.set .Lregnr, 0;)                                \
-       stringify_in_c(.irp rs,r0,r1,r2,r3,r4,r5,r6,r7,r8,r9,r10,r11,r12,r13,r14,r15;) \
-       stringify_in_c(.ifc "\reg", "%%\rs";)                           \
-       stringify_in_c(.set .Lfound, 1;)                                \
-       stringify_in_c(.short .Lregnr;)                                 \
-       stringify_in_c(.endif;)                                         \
-       stringify_in_c(.set .Lregnr, .Lregnr+1;)                        \
-       stringify_in_c(.endr;)                                          \
-       stringify_in_c(.ifne (.Lfound != 1);)                           \
-       stringify_in_c(.error "extable_reg: bad register argument";)    \
-       stringify_in_c(.endif;)                                         \
-       stringify_in_c(.endm;)                                          \
-       stringify_in_c(extable_reg _reg;)                               \
-       stringify_in_c(.purgem extable_reg;)                            \
+#define __EX_TABLE_UA(_section, _fault, _target, _type, _regerr, _regaddr, _len)\
+       stringify_in_c(.section _section,"a";)                                  \
+       stringify_in_c(.align   4;)                                             \
+       stringify_in_c(.long    (_fault) - .;)                                  \
+       stringify_in_c(.long    (_target) - .;)                                 \
+       stringify_in_c(.short   (_type);)                                       \
+       stringify_in_c(.macro   extable_reg regerr, regaddr;)                   \
+       stringify_in_c(.set     .Lfound, 0;)                                    \
+       stringify_in_c(.set     .Lcurr, 0;)                                     \
+       stringify_in_c(.irp     rs,0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15;)      \
+       stringify_in_c(         .ifc    "\regerr", "%%r\rs";)                   \
+       stringify_in_c(                 .set    .Lfound, 1;)                    \
+       stringify_in_c(                 .set    .Lregerr, .Lcurr;)              \
+       stringify_in_c(         .endif;)                                        \
+       stringify_in_c(         .set    .Lcurr, .Lcurr+1;)                      \
+       stringify_in_c(.endr;)                                                  \
+       stringify_in_c(.ifne    (.Lfound != 1);)                                \
+       stringify_in_c(         .error  "extable_reg: bad register argument1";) \
+       stringify_in_c(.endif;)                                                 \
+       stringify_in_c(.set     .Lfound, 0;)                                    \
+       stringify_in_c(.set     .Lcurr, 0;)                                     \
+       stringify_in_c(.irp     rs,0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15;)      \
+       stringify_in_c(         .ifc    "\regaddr", "%%r\rs";)                  \
+       stringify_in_c(                 .set    .Lfound, 1;)                    \
+       stringify_in_c(                 .set    .Lregaddr, .Lcurr;)             \
+       stringify_in_c(         .endif;)                                        \
+       stringify_in_c(         .set    .Lcurr, .Lcurr+1;)                      \
+       stringify_in_c(.endr;)                                                  \
+       stringify_in_c(.ifne    (.Lfound != 1);)                                \
+       stringify_in_c(         .error  "extable_reg: bad register argument2";) \
+       stringify_in_c(.endif;)                                                 \
+       stringify_in_c(.short   .Lregerr << EX_DATA_REG_ERR_SHIFT |             \
+                               .Lregaddr << EX_DATA_REG_ADDR_SHIFT |           \
+                               _len << EX_DATA_LEN_SHIFT;)                     \
+       stringify_in_c(.endm;)                                                  \
+       stringify_in_c(extable_reg _regerr,_regaddr;)                           \
+       stringify_in_c(.purgem  extable_reg;)                                   \
        stringify_in_c(.previous)
 
 #define EX_TABLE(_fault, _target)                                      \
        __EX_TABLE(__ex_table, _fault, _target, EX_TYPE_FIXUP)
+
 #define EX_TABLE_AMODE31(_fault, _target)                              \
        __EX_TABLE(.amode31.ex_table, _fault, _target, EX_TYPE_FIXUP)
-#define EX_TABLE_UA(_fault, _target, _reg)                             \
-       __EX_TABLE_UA(__ex_table, _fault, _target, EX_TYPE_UACCESS, _reg)
+
+#define EX_TABLE_UA_STORE(_fault, _target, _regerr)                    \
+       __EX_TABLE_UA(__ex_table, _fault, _target, EX_TYPE_UA_STORE, _regerr, _regerr, 0)
+
+#define EX_TABLE_UA_LOAD_MEM(_fault, _target, _regerr, _regmem, _len)  \
+       __EX_TABLE_UA(__ex_table, _fault, _target, EX_TYPE_UA_LOAD_MEM, _regerr, _regmem, _len)
+
+#define EX_TABLE_UA_LOAD_REG(_fault, _target, _regerr, _regzero)       \
+       __EX_TABLE_UA(__ex_table, _fault, _target, EX_TYPE_UA_LOAD_REG, _regerr, _regzero, 0)
 
 #endif /* __ASM_EXTABLE_H */
 
        };
 };
 
-#define __put_get_user_asm(to, from, size, oac_spec)                   \
+#define __put_user_asm(to, from, size)                                 \
 ({                                                                     \
+       union oac __oac_spec = {                                        \
+               .oac1.as = PSW_BITS_AS_SECONDARY,                       \
+               .oac1.a = 1,                                            \
+       };                                                              \
        int __rc;                                                       \
                                                                        \
        asm volatile(                                                   \
                "0:     mvcos   %[_to],%[_from],%[_size]\n"             \
                "1:     xr      %[rc],%[rc]\n"                          \
                "2:\n"                                                  \
-               EX_TABLE_UA(0b,2b,%[rc]) EX_TABLE_UA(1b,2b,%[rc])       \
+               EX_TABLE_UA_STORE(0b, 2b, %[rc])                        \
+               EX_TABLE_UA_STORE(1b, 2b, %[rc])                        \
                : [rc] "=&d" (__rc), [_to] "+Q" (*(to))                 \
                : [_size] "d" (size), [_from] "Q" (*(from)),            \
-                 [spec] "d" (oac_spec.val)                             \
+                 [spec] "d" (__oac_spec.val)                           \
                : "cc", "0");                                           \
        __rc;                                                           \
 })
 
-#define __put_user_asm(to, from, size)                         \
-       __put_get_user_asm(to, from, size, ((union oac) {       \
-               .oac1.as = PSW_BITS_AS_SECONDARY,               \
-               .oac1.a = 1                                     \
-       }))
-
-#define __get_user_asm(to, from, size)                         \
-       __put_get_user_asm(to, from, size, ((union oac) {       \
-               .oac2.as = PSW_BITS_AS_SECONDARY,               \
-               .oac2.a = 1                                     \
-       }))                                                     \
-
 static __always_inline int __put_user_fn(void *x, void __user *ptr, unsigned long size)
 {
        int rc;
        return rc;
 }
 
+#define __get_user_asm(to, from, size)                                 \
+({                                                                     \
+       union oac __oac_spec = {                                        \
+               .oac2.as = PSW_BITS_AS_SECONDARY,                       \
+               .oac2.a = 1,                                            \
+       };                                                              \
+       int __rc;                                                       \
+                                                                       \
+       asm volatile(                                                   \
+               "       lr      0,%[spec]\n"                            \
+               "0:     mvcos   0(%[_to]),%[_from],%[_size]\n"          \
+               "1:     xr      %[rc],%[rc]\n"                          \
+               "2:\n"                                                  \
+               EX_TABLE_UA_LOAD_MEM(0b, 2b, %[rc], %[_to], %[_ksize])  \
+               EX_TABLE_UA_LOAD_MEM(1b, 2b, %[rc], %[_to], %[_ksize])  \
+               : [rc] "=&d" (__rc), "=Q" (*(to))                       \
+               : [_size] "d" (size), [_from] "Q" (*(from)),            \
+                 [spec] "d" (__oac_spec.val), [_to] "a" (to),          \
+                 [_ksize] "K" (size)                                   \
+               : "cc", "0");                                           \
+       __rc;                                                           \
+})
+
 static __always_inline int __get_user_fn(void *x, const void __user *ptr, unsigned long size)
 {
        int rc;
        __chk_user_ptr(ptr);                                    \
        switch (sizeof(*(ptr))) {                               \
        case 1: {                                               \
-               unsigned char __x = 0;                          \
+               unsigned char __x;                              \
                __gu_err = __get_user_fn(&__x, ptr,             \
                                         sizeof(*(ptr)));       \
                (x) = *(__force __typeof__(*(ptr)) *) &__x;     \
                break;                                          \
        };                                                      \
        case 2: {                                               \
-               unsigned short __x = 0;                         \
+               unsigned short __x;                             \
                __gu_err = __get_user_fn(&__x, ptr,             \
                                         sizeof(*(ptr)));       \
                (x) = *(__force __typeof__(*(ptr)) *) &__x;     \
                break;                                          \
        };                                                      \
        case 4: {                                               \
-               unsigned int __x = 0;                           \
+               unsigned int __x;                               \
                __gu_err = __get_user_fn(&__x, ptr,             \
                                         sizeof(*(ptr)));       \
                (x) = *(__force __typeof__(*(ptr)) *) &__x;     \
                break;                                          \
        };                                                      \
        case 8: {                                               \
-               unsigned long long __x = 0;                     \
+               unsigned long __x;                              \
                __gu_err = __get_user_fn(&__x, ptr,             \
                                         sizeof(*(ptr)));       \
                (x) = *(__force __typeof__(*(ptr)) *) &__x;     \
                "0:   " insn "  %[_val],%[_to]\n"                       \
                "1:     xr      %[rc],%[rc]\n"                          \
                "2:\n"                                                  \
-               EX_TABLE_UA(0b,2b,%0) EX_TABLE_UA(1b,2b,%0)             \
+               EX_TABLE_UA_STORE(0b, 2b, %[rc])                        \
+               EX_TABLE_UA_STORE(1b, 2b, %[rc])                        \
                : [rc] "=d" (__rc), [_to] "+Q" (*(to))                  \
                : [_val] "d" (val)                                      \
                : "cc");                                                \
 
 #define __put_kernel_nofault(dst, src, type, err_label)                        \
 do {                                                                   \
-       u64 __x = (u64)(*((type *)(src)));                              \
+       unsigned long __x = (unsigned long)(*((type *)(src)));          \
        int __pk_err;                                                   \
                                                                        \
        switch (sizeof(type)) {                                         \
                "0:   " insn "  %[_val],%[_from]\n"                     \
                "1:     xr      %[rc],%[rc]\n"                          \
                "2:\n"                                                  \
-               EX_TABLE_UA(0b,2b,%0) EX_TABLE_UA(1b,2b,%0)             \
-               : [rc] "=d" (__rc), [_val] "+d" (val)                   \
+               EX_TABLE_UA_LOAD_REG(0b, 2b, %[rc], %[_val])            \
+               EX_TABLE_UA_LOAD_REG(1b, 2b, %[rc], %[_val])            \
+               : [rc] "=d" (__rc), [_val] "=d" (val)                   \
                : [_from] "Q" (*(from))                                 \
                : "cc");                                                \
        __rc;                                                           \
                                                                        \
        switch (sizeof(type)) {                                         \
        case 1: {                                                       \
-               u8 __x = 0;                                             \
+               unsigned char __x;                                      \
                                                                        \
                __gk_err = __get_kernel_asm(__x, (type *)(src), "ic");  \
                *((type *)(dst)) = (type)__x;                           \
                break;                                                  \
        };                                                              \
        case 2: {                                                       \
-               u16 __x = 0;                                            \
+               unsigned short __x;                                     \
                                                                        \
                __gk_err = __get_kernel_asm(__x, (type *)(src), "lh");  \
                *((type *)(dst)) = (type)__x;                           \
                break;                                                  \
        };                                                              \
        case 4: {                                                       \
-               u32 __x = 0;                                            \
+               unsigned int __x;                                       \
                                                                        \
                __gk_err = __get_kernel_asm(__x, (type *)(src), "l");   \
                *((type *)(dst)) = (type)__x;                           \
                break;                                                  \
        };                                                              \
        case 8: {                                                       \
-               u64 __x = 0;                                            \
+               unsigned long __x;                                      \
                                                                        \
                __gk_err = __get_kernel_asm(__x, (type *)(src), "lg");  \
                *((type *)(dst)) = (type)__x;                           \
 
 // SPDX-License-Identifier: GPL-2.0
 
+#include <linux/bitfield.h>
 #include <linux/extable.h>
+#include <linux/string.h>
 #include <linux/errno.h>
 #include <linux/panic.h>
 #include <asm/asm-extable.h>
        return true;
 }
 
-static bool ex_handler_uaccess(const struct exception_table_entry *ex, struct pt_regs *regs)
+static bool ex_handler_ua_store(const struct exception_table_entry *ex, struct pt_regs *regs)
 {
-       regs->gprs[ex->data] = -EFAULT;
+       unsigned int reg_err = FIELD_GET(EX_DATA_REG_ERR, ex->data);
+
+       regs->gprs[reg_err] = -EFAULT;
+       regs->psw.addr = extable_fixup(ex);
+       return true;
+}
+
+static bool ex_handler_ua_load_mem(const struct exception_table_entry *ex, struct pt_regs *regs)
+{
+       unsigned int reg_addr = FIELD_GET(EX_DATA_REG_ADDR, ex->data);
+       unsigned int reg_err = FIELD_GET(EX_DATA_REG_ERR, ex->data);
+       size_t len = FIELD_GET(EX_DATA_LEN, ex->data);
+
+       regs->gprs[reg_err] = -EFAULT;
+       memset((void *)regs->gprs[reg_addr], 0, len);
+       regs->psw.addr = extable_fixup(ex);
+       return true;
+}
+
+static bool ex_handler_ua_load_reg(const struct exception_table_entry *ex, struct pt_regs *regs)
+{
+       unsigned int reg_zero = FIELD_GET(EX_DATA_REG_ADDR, ex->data);
+       unsigned int reg_err = FIELD_GET(EX_DATA_REG_ERR, ex->data);
+
+       regs->gprs[reg_err] = -EFAULT;
+       regs->gprs[reg_zero] = 0;
        regs->psw.addr = extable_fixup(ex);
        return true;
 }
                return ex_handler_fixup(ex, regs);
        case EX_TYPE_BPF:
                return ex_handler_bpf(ex, regs);
-       case EX_TYPE_UACCESS:
-               return ex_handler_uaccess(ex, regs);
+       case EX_TYPE_UA_STORE:
+               return ex_handler_ua_store(ex, regs);
+       case EX_TYPE_UA_LOAD_MEM:
+               return ex_handler_ua_load_mem(ex, regs);
+       case EX_TYPE_UA_LOAD_REG:
+               return ex_handler_ua_load_reg(ex, regs);
        }
        panic("invalid exception table entry");
 }