fix shmem_lock

oshmem/mca/memheap/base/memheap_base_mkey.c

@@ -755,10 +755,16 @@ uint64_t mca_memheap_base_find_offset(int pe,
                                       void* rva)
 {
     map_segment_t *s;
+    int my_pe = oshmem_my_proc_id();
 
     s = __find_va(va);
 
-    return ((s && MAP_SEGMENT_IS_VALID(s)) ? ((uintptr_t)rva - (uintptr_t)(s->mkeys_cache[pe][tr_id].va_base)) : 0);
+    if (my_pe == pe) {
+        return (uintptr_t)va - (uintptr_t)s->seg_base_addr;
+    }
+    else {
+        return ((s && MAP_SEGMENT_IS_VALID(s)) ? ((uintptr_t)rva - (uintptr_t)(s->mkeys_cache[pe][tr_id].va_base)) : 0);
+    }
 }
 
 int mca_memheap_base_is_symmetric_addr(const void* va)

oshmem/shmem/c/shmem_lock.c

@@ -72,7 +72,7 @@ static int shmem_lock_pack_pe_next(void *lock,
 static void shmem_lock_increment_counter(void *lock, int lock_size);
 static int shmem_lock_decrement_counter(void *lock, int lock_size);
 
-static int shmem_lock_get_server(const void *lock);
+static int shmem_lock_get_server(void *lock);
 static int shmem_lock_is_mine(void *lock, int lock_size);
 
 static int shmem_lock_get_ticket(void *lock);
@@ -166,9 +166,10 @@ int shmem_lock_finalize()
     return OSHMEM_SUCCESS;
 }
 
-static int shmem_lock_get_server(const void *lock)
+static int shmem_lock_get_server(void *lock)
 {
-    return (int) (((uintptr_t) lock) / 8) % shmem_n_pes();
+    uint64_t offset =  MCA_MEMHEAP_CALL(find_offset(shmem_my_pe(), 0, lock, lock));
+    return (offset / 8) % shmem_n_pes();
 }
 
 static uint64_t get_lock_value(const void *lock, int lock_size)