bool test_and_clear_guest_dirty(struct mm_struct *mm, unsigned long address);
 int set_guest_storage_key(struct mm_struct *mm, unsigned long addr,
                          unsigned char key, bool nq);
-unsigned long get_guest_storage_key(struct mm_struct *mm, unsigned long addr);
+int get_guest_storage_key(struct mm_struct *mm, unsigned long addr,
+                         unsigned char *key);
 
 /*
  * Certain architectures need to do special things when PTEs
 
 {
        uint8_t *keys;
        uint64_t hva;
-       unsigned long curkey;
        int i, r = 0;
 
        if (args->flags != 0)
                        break;
                }
 
-               curkey = get_guest_storage_key(current->mm, hva);
-               if (IS_ERR_VALUE(curkey)) {
-                       r = curkey;
+               r = get_guest_storage_key(current->mm, hva, &keys[i]);
+               if (r)
                        break;
-               }
-               keys[i] = curkey;
        }
        up_read(¤t->mm->mmap_sem);
 
 
 }
 EXPORT_SYMBOL(set_guest_storage_key);
 
-unsigned long get_guest_storage_key(struct mm_struct *mm, unsigned long addr)
+int get_guest_storage_key(struct mm_struct *mm, unsigned long addr,
+                         unsigned char *key)
 {
-       unsigned char key;
        spinlock_t *ptl;
        pgste_t pgste;
        pte_t *ptep;
                return -EFAULT;
 
        pgste = pgste_get_lock(ptep);
-       key = (pgste_val(pgste) & (PGSTE_ACC_BITS | PGSTE_FP_BIT)) >> 56;
+       *key = (pgste_val(pgste) & (PGSTE_ACC_BITS | PGSTE_FP_BIT)) >> 56;
        if (!(pte_val(*ptep) & _PAGE_INVALID))
-               key = page_get_storage_key(pte_val(*ptep) & PAGE_MASK);
+               *key = page_get_storage_key(pte_val(*ptep) & PAGE_MASK);
        /* Reflect guest's logical view, not physical */
-       key |= (pgste_val(pgste) & (PGSTE_GR_BIT | PGSTE_GC_BIT)) >> 48;
+       *key |= (pgste_val(pgste) & (PGSTE_GR_BIT | PGSTE_GC_BIT)) >> 48;
        pgste_set_unlock(ptep, pgste);
        pte_unmap_unlock(ptep, ptl);
-       return key;
+       return 0;
 }
 EXPORT_SYMBOL(get_guest_storage_key);
 #endif