[Wisp] Support new fast locking from backport of 8291555

Summary: Wisp should support the new fast locking from backport of 8291555

Test Plan: CICD

Reviewed-by: kuaiwei, yulei, ddh

Issue: #697

src/hotspot/cpu/aarch64/c2_CodeStubs_aarch64.cpp

@@ -35,6 +35,7 @@ int C2HandleAnonOMOwnerStub::max_size() const {
   // Max size of stub has been determined by testing with 0, in which case
   // C2CodeStubList::emit() will throw an assertion and report the actual size that
   // is needed.
+  if (UseWispMonitor) return 36;
   return 24;
 }
 
@@ -43,10 +44,15 @@ void C2HandleAnonOMOwnerStub::emit(C2_MacroAssembler& masm) {
   Register mon = monitor();
   Register t = tmp();
   assert(t != noreg, "need tmp register");
-
+  if (UseWispMonitor) {
+    __ ldr(rthread, Address(rthread, JavaThread::current_coroutine_offset()));
+    __ ldr(rthread, Address(rthread, Coroutine::wisp_thread_offset()));
+  }
   // Fix owner to be the current thread.
   __ str(rthread, Address(mon, ObjectMonitor::owner_offset_in_bytes()));
-
+  if (UseWispMonitor) {
+    __ ldr(rthread, Address(rthread, WispThread::thread_offset()));
+  }
   // Pop owner object from lock-stack.
   __ ldrw(t, Address(rthread, JavaThread::lock_stack_top_offset()));
   __ subw(t, t, oopSize);

src/hotspot/cpu/aarch64/sharedRuntime_aarch64.cpp

@@ -3156,17 +3156,18 @@ void create_switchTo_contents(MacroAssembler *masm, int start, OopMapSet* oop_ma
   __ ldr(target_coroutine, Address(target_coroutine, java_dyn_CoroutineBase::get_data_offset()));
 
   Register temp = r4;
-  Register temp2 = r5;
+  Register temp2 = r9;
+  Register temp3 = r10;
+  Register temp4 = r11;
   {
     //////////////////////////////////////////////////////////////////////////
     // store information into the old coroutine's object
     //
-    // valid registers: rsi = old Coroutine, rdx = target Coroutine
+    // valid registers: r1 = old Coroutine, r2 = target Coroutine
 
     Register old_coroutine_obj = j_rarg0;
     Register old_coroutine = r5;
     Register old_stack = r6;
-
     // check that we're dealing with sane objects...
     __ ldr(old_coroutine, Address(old_coroutine_obj, java_dyn_CoroutineBase::get_data_offset()));
     __ ldr(old_stack, Address(old_coroutine, Coroutine::stack_offset()));
@@ -3203,6 +3204,43 @@ void create_switchTo_contents(MacroAssembler *masm, int start, OopMapSet* oop_ma
 
     __ mov(temp, sp);
     __ str(temp, Address(old_stack, CoroutineStack::last_sp_offset())); // str cannot use sp as an argument
+
+    if (UseWispMonitor && UseAltFastLocking) {
+#ifdef ASSERT
+      assert(WispThread::lock_stack_top_offset() == JavaThread::lock_stack_top_offset(), "Should be");
+      assert(WispThread::lock_stack_base_offset() == JavaThread::lock_stack_base_offset(), "Should be");
+#endif
+      __ ldr(temp2, Address(old_coroutine, Coroutine::wisp_thread_offset()));
+      __ ldrw(temp3, Address(thread, JavaThread::lock_stack_top_offset()));
+      __ strw(temp3, Address(temp2, WispThread::lock_stack_top_offset()));
+      __ add(temp3, temp3, thread);
+
+      __ lea(temp, Address(thread, JavaThread::lock_stack_base_offset()));
+      __ lea(temp2, Address(temp2, WispThread::lock_stack_base_offset()));
+
+      Label loop, test;
+      __ br(Assembler::AL, test);
+      __ bind(loop);
+      __ ldr(temp4, Address(temp, 0));
+      __ str(temp4, Address(temp2, 0));
+      __ add(temp, temp, oopSize);
+      __ add(temp2, temp2, oopSize);
+      __ bind(test);
+      __ cmp(temp, temp3);
+      __ br(Assembler::LO, loop);
+#ifdef ASSERT
+      __ ldr(temp3, Address(old_coroutine, Coroutine::wisp_thread_offset()));
+      __ add(temp3, temp3, static_cast<int32_t>(LockStack::end_offset()));
+      Label loop1, test1;
+      __ br(Assembler::AL, test1);
+      __ bind(loop1);
+      __ str(zr, Address(temp2, 0));
+      __ add(temp2, temp2, oopSize);
+      __ bind(test1);
+      __ cmp(temp2, temp3);
+      __ br(Assembler::LO, loop1);
+#endif
+    }
   }
   Register target_stack = rheapbase;
   __ ldr(target_stack, Address(target_coroutine, Coroutine::stack_offset()));
@@ -3211,8 +3249,7 @@ void create_switchTo_contents(MacroAssembler *masm, int start, OopMapSet* oop_ma
     //////////////////////////////////////////////////////////////////////////
     // perform the switch to the new stack
     //
-    // valid registers: rdx = target Coroutine
-
+    // valid registers: r2 = target Coroutine
     __ movw(temp, Coroutine::_current);
     __ strw(temp, Address(target_coroutine, Coroutine::state_offset()));
     {
@@ -3244,6 +3281,39 @@ void create_switchTo_contents(MacroAssembler *masm, int start, OopMapSet* oop_ma
       __ str(temp, Address(thread, JavaThread::monitor_chunks_offset()));
       __ ldrb(temp, Address(target_coroutine, Coroutine::do_not_unlock_if_synchronized_offset()));
       __ strb(temp, Address(thread, JavaThread::do_not_unlock_if_synchronized_offset()));
+
+      if (UseWispMonitor && UseAltFastLocking) {
+        __ ldr(temp2, Address(target_coroutine, Coroutine::wisp_thread_offset()));
+        __ ldrw(temp, Address(temp2, WispThread::lock_stack_top_offset()));
+        __ strw(temp, Address(thread, JavaThread::lock_stack_top_offset()));
+        __ mov(temp3, temp);
+        __ add(temp3, temp3, temp2);
+
+        __ lea(temp, Address(temp2, WispThread::lock_stack_base_offset()));
+        __ lea(temp2, Address(thread, JavaThread::lock_stack_base_offset()));
+
+        Label loop, test;
+        __ br(Assembler::AL, test);
+        __ bind(loop);
+        __ ldr(temp4, Address(temp, 0));
+        __ str(temp4, Address(temp2, 0));
+        __ add(temp, temp, oopSize);
+        __ add(temp2, temp2, oopSize);
+        __ bind(test);
+        __ cmp(temp, temp3);
+        __ br(Assembler::LO, loop);
+#ifdef ASSERT
+        __ lea(temp3, Address(thread, LockStack::end_offset()));
+        Label loop1, test1;
+        __ br(Assembler::AL, test1);
+        __ bind(loop1);
+        __ str(zr, Address(temp2, 0));
+        __ add(temp2, temp2, oopSize);
+        __ bind(test1);
+        __ cmp(temp2, temp3);
+        __ br(Assembler::LO, loop1);
+#endif
+      }
 #ifdef ASSERT
       __ str(zr, Address(target_coroutine, Coroutine::handle_area_offset()));
       __ str(zr, Address(target_coroutine, Coroutine::resource_area_offset()));

src/hotspot/cpu/x86/c2_CodeStubs_x86.cpp

@@ -36,14 +36,24 @@ int C2HandleAnonOMOwnerStub::max_size() const {
   // Max size of stub has been determined by testing with 0, in which case
   // C2CodeStubList::emit() will throw an assertion and report the actual size that
   // is needed.
+  if (UseWispMonitor) {
+    return DEBUG_ONLY(50) NOT_DEBUG(35);
+  }
   return DEBUG_ONLY(36) NOT_DEBUG(21);
 }
 
 void C2HandleAnonOMOwnerStub::emit(C2_MacroAssembler& masm) {
   __ bind(entry());
   Register mon = monitor();
   Register t = tmp();
+  if (UseWispMonitor) {
+    __ movptr(r15_thread, Address(r15_thread, JavaThread::current_coroutine_offset()));
+    __ movptr(r15_thread, Address(r15_thread, Coroutine::wisp_thread_offset()));
+  }
   __ movptr(Address(mon, OM_OFFSET_NO_MONITOR_VALUE_TAG(owner)), r15_thread);
+  if (UseWispMonitor) {
+    __ movptr(r15_thread, Address(r15_thread, WispThread::thread_offset()));
+  }
   __ subl(Address(r15_thread, JavaThread::lock_stack_top_offset()), oopSize);
 #ifdef ASSERT
   __ movl(t, Address(r15_thread, JavaThread::lock_stack_top_offset()));

src/hotspot/cpu/x86/sharedRuntime_x86_64.cpp

@@ -4154,6 +4154,9 @@ void create_switchTo_contents(MacroAssembler *masm, int start, OopMapSet* oop_ma
     Register old_coroutine = r9;
     Register old_stack = r10;
     Register temp = r8;
+    Register temp2 = r11;
+    Register temp3 = rax;
+    Register temp4 = rcx;
 
     // check that we're dealing with sane objects...
     DEBUG_ONLY(stop_if_null(masm, old_coroutine_obj, "null old_coroutine"));
@@ -4198,6 +4201,45 @@ void create_switchTo_contents(MacroAssembler *masm, int start, OopMapSet* oop_ma
     __ movbool(Address(old_coroutine, Coroutine::do_not_unlock_if_synchronized_offset()), temp);
 
     __ movptr(Address(old_stack, CoroutineStack::last_sp_offset()), rsp);
+
+    if (UseWispMonitor && UseAltFastLocking) {
+#ifdef ASSERT
+      assert(WispThread::lock_stack_top_offset() == JavaThread::lock_stack_top_offset(), "Should be");
+      assert(WispThread::lock_stack_base_offset() == JavaThread::lock_stack_base_offset(), "Should be");
+#endif
+      // temp2: WispThread*
+      __ movptr(temp2, Address(old_coroutine, Coroutine::wisp_thread_offset()));
+      __ movl(temp3, Address(thread, JavaThread::lock_stack_top_offset()));
+      __ movl(Address(temp2, WispThread::lock_stack_top_offset()), temp3);
+      __ addptr(temp3, thread);
+
+      __ lea(temp, Address(thread, JavaThread::lock_stack_base_offset()));
+      __ lea(temp2, Address(temp2, WispThread::lock_stack_base_offset()));
+
+      Label loop, test;
+      __ jmp(test);
+      __ bind(loop);
+      __ movptr(temp4, Address(temp, 0));
+      __ movptr(Address(temp2, 0), temp4);
+      __ addptr(temp, oopSize);
+      __ addptr(temp2, oopSize);
+      __ bind(test);
+      __ cmpptr(temp, temp3);
+      __ jcc(Assembler::below, loop);
+
+#ifdef ASSERT
+      __ movptr(temp3, Address(old_coroutine, Coroutine::wisp_thread_offset()));
+      __ addptr(temp3, static_cast<int32_t>(LockStack::end_offset()));
+      Label loop1, test1;
+      __ jmp(test1);
+      __ bind(loop1);
+      __ movptr(Address(temp2, 0), (intptr_t)NULL_WORD);
+      __ addptr(temp2, oopSize);
+      __ bind(test1);
+      __ cmpptr(temp2, temp3);
+      __ jcc(Assembler::below, loop1);
+#endif
+    }
   }
   Register target_stack = r12;
   __ movptr(target_stack, Address(target_coroutine, Coroutine::stack_offset()));
@@ -4212,6 +4254,8 @@ void create_switchTo_contents(MacroAssembler *masm, int start, OopMapSet* oop_ma
 
     Register temp = r8;
     Register temp2 = r9;
+    Register temp3 = rax;
+    Register temp4 = r11;
     {
       Register thread = r15;
       __ movptr(Address(thread, JavaThread::current_coroutine_offset()), target_coroutine);
@@ -4241,6 +4285,41 @@ void create_switchTo_contents(MacroAssembler *masm, int start, OopMapSet* oop_ma
       __ movptr(Address(thread, JavaThread::monitor_chunks_offset()), temp);
       __ movbool(temp, Address(target_coroutine, Coroutine::do_not_unlock_if_synchronized_offset()));
       __ movbool(Address(thread, JavaThread::do_not_unlock_if_synchronized_offset()), temp);
+
+      if (UseWispMonitor && UseAltFastLocking) {
+        // temp2: WispThread*
+        __ movptr(temp2, Address(target_coroutine, Coroutine::wisp_thread_offset()));
+        __ movl(temp, Address(temp2, WispThread::lock_stack_top_offset()));
+        __ movl(Address(thread, JavaThread::lock_stack_top_offset()), temp);
+        __ movl(temp3, temp);
+        __ addptr(temp3, temp2);
+
+        __ lea(temp, Address(temp2, WispThread::lock_stack_base_offset()));
+        __ lea(temp2, Address(thread, JavaThread::lock_stack_base_offset()));
+
+        Label loop, test;
+        __ jmp(test);
+        __ bind(loop);
+        __ movptr(temp4, Address(temp, 0));
+        __ movptr(Address(temp2, 0), temp4);
+        __ addptr(temp, oopSize);
+        __ addptr(temp2, oopSize);
+        __ bind(test);
+        __ cmpptr(temp, temp3);
+        __ jcc(Assembler::below, loop);
+#ifdef ASSERT
+        __ lea(temp3, Address(thread, LockStack::end_offset()));
+        Label loop1, test1;
+        __ jmp(test1);
+        __ bind(loop1);
+        __ movptr(Address(temp2, 0), (intptr_t)NULL_WORD);
+        __ addptr(temp2, oopSize);
+        __ bind(test1);
+        __ cmpptr(temp2, temp3);
+        __ jcc(Assembler::below, loop1);
+#endif
+      }
+
 #ifdef ASSERT
       __ movptr(Address(target_coroutine, Coroutine::handle_area_offset()), (intptr_t)NULL_WORD);
       __ movptr(Address(target_coroutine, Coroutine::resource_area_offset()), (intptr_t)NULL_WORD);

src/hotspot/share/runtime/coroutine.cpp

@@ -648,18 +648,18 @@ void WispThread::set_wisp_booted(Thread* thread) {
 /*
  * Avoid coroutine switch in the following scenarios:
  *
- * - _wisp_booted: 
- *   We guarantee the classes referenced by WispTask.park(called in WispThread::park in native) 
+ * - _wisp_booted:
+ *   We guarantee the classes referenced by WispTask.park(called in WispThread::park in native)
  *   are already loaded after _wisp_booted is set(as true). Otherwise it might result in loading class during execution of WispTask.park.
  *   Coroutine switch caused by object monitors in class loading might lead to recursive deadlock.
  *
  * - !com_alibaba_wisp_engine_WispCarrier::is_critical(_coroutine->wisp_engine()):
- *   If the program is already running in kernel code of wisp engine(marked by WispEngine.isInCritical at Java level),  we don't expect 
+ *   If the program is already running in kernel code of wisp engine(marked by WispEngine.isInCritical at Java level),  we don't expect
  *   the switch while coroutine running into 'synchronized' block which is heavily used by Java NIO library.
  *   Otherwise, it might lead to potential recursive deadlock.
- * 
+ *
  * - monitor->object() != java_lang_ref_Reference::pending_list_lock():
- *   pending_list_lock(PLL) is special ObjectMonitor used in GC vm operation. 
+ *   pending_list_lock(PLL) is special ObjectMonitor used in GC vm operation.
  *   if we treated it as normal monitor(T10965418)
  *   - 'dead lock' in jni_critical case:
  *     Given coroutine A, B running in the same thread,
@@ -1002,6 +1002,9 @@ const char* WispThread::print_blocking_status(int status) {
 
 void WispThread::oops_do(OopClosure *f, CodeBlobClosure *cf) {
   f->do_oop((oop*) &_threadObj);
+  if (UseWispMonitor && UseAltFastLocking) {
+    lock_stack().oops_do(f);
+  }
 }
 
 void Coroutine::after_safepoint(JavaThread* thread) {
@@ -1048,7 +1051,7 @@ void Coroutine::after_safepoint(JavaThread* thread) {
   }
 
   coroutine->_is_yielding = true;
-  // "yield" will immediately switch context to execute other coroutines. 
+  // "yield" will immediately switch context to execute other coroutines.
   // After all the runnable coroutines has been executed, we'll switch back.
   //
   // - The preempt mechanism should be disabled when current coroutine is calling "yield"

src/hotspot/share/runtime/coroutine.hpp

@@ -27,6 +27,7 @@
 
 #include "runtime/jniHandles.hpp"
 #include "runtime/handles.hpp"
+#include "runtime/lockStack.hpp"
 #include "memory/allocation.hpp"
 #include "memory/resourceArea.hpp"
 #include "runtime/javaFrameAnchor.hpp"
@@ -476,7 +477,7 @@ class WispThread: public JavaThread {
 
   static WispThread* current(Thread* thread) {
     assert(thread->is_Java_thread(), "invariant") ;
-    return thread->is_Wisp_thread() ? (WispThread*) thread : 
+    return thread->is_Wisp_thread() ? (WispThread*) thread :
       ((JavaThread*) thread)->current_coroutine()->wisp_thread();
   }
 

src/hotspot/share/runtime/objectMonitor.cpp

@@ -980,6 +980,7 @@ void ObjectMonitor::exit(bool not_suspended, TRAPS) {
                 ((WispThread*)Self)->coroutine()->stack()->stack_base(),
                 ((WispThread*)Self)->coroutine()->stack()->stack_base() - ((WispThread*)Self)->coroutine()->stack()->stack_size(),
                 _owner, ((WispThread*)_owner)->coroutine());
+        printf("DEBUG: _owner: %p, WispThread::current(_owner): %p, THREAD: %p\n", _owner, WispThread::current((JavaThread*)_owner), THREAD);
       }
       assert(false, "Non-balanced monitor enter/exit! Likely JNI locking");
       return;

src/hotspot/share/runtime/objectMonitor.inline.hpp

@@ -34,7 +34,12 @@ inline intptr_t ObjectMonitor::is_entered(TRAPS) const {
   if (UseAltFastLocking) {
     if (is_owner_anonymous()) {
       assert(THREAD->is_Java_thread(), "sanity");
-      JavaThread* jt = (JavaThread*)THREAD;
+      JavaThread* jt;
+      if (UseWispMonitor) {
+        jt = ((WispThread*) THREAD)->thread();
+      } else {
+        jt = (JavaThread*)THREAD;
+      }
       return jt->lock_stack().contains((oop)object()) ? 1 : 0;
     } else {
       return THREAD == _owner ? 1 : 0;