Merge pull request #20146 from nbhuiyan/mem-segment-vh-0

jdmpapin · web-flow · commit 25f0b722520f · 2024-09-23T10:10:05.000-04:00
Inline SegmentViewVarHandle operations in JDK21+
diff --git a/runtime/compiler/codegen/J9RecognizedMethodsEnum.hpp b/runtime/compiler/codegen/J9RecognizedMethodsEnum.hpp
@@ -1165,6 +1165,9 @@
    java_lang_invoke_VarHandleByteArrayAsX_ArrayHandle_method,
    java_lang_invoke_VarHandleByteArrayAsX_ByteBufferHandle_method,
 
+   jdk_internal_foreign_layout_ValueLayouts_AbstractValueLayout_accessHandle,
+   java_lang_foreign_MemorySegment_method,
+
    // Clone and Deep Copy
    java_lang_J9VMInternals_is32Bit,
    java_lang_J9VMInternals_isClassModifierPublic,
diff --git a/runtime/compiler/control/JITClientCompilationThread.cpp b/runtime/compiler/control/JITClientCompilationThread.cpp
@@ -1214,6 +1214,16 @@ handleServerMessage(JITServer::ClientStream *client, TR_J9VM *fe, JITServer::Mes
          client->write(response, mhIndex, mhObj);
          }
          break;
+      case MessageType::VM_getLayoutVarHandle:
+         {
+         auto recv = client->getRecvData<TR::KnownObjectTable::Index>();
+         TR::KnownObjectTable::Index vhIndex = fe->getLayoutVarHandle(comp, std::get<0>(recv));
+         uintptr_t* vhObj = NULL;
+         if (vhIndex != TR::KnownObjectTable::UNKNOWN)
+            vhObj = knot->getPointerLocation(vhIndex);
+         client->write(response, vhIndex, vhObj);
+         }
+         break;
 #endif // J9VM_OPT_OPENJDK_METHODHANDLE
       case MessageType::VM_isStable:
          {
diff --git a/runtime/compiler/env/VMJ9.cpp b/runtime/compiler/env/VMJ9.cpp
@@ -5131,6 +5131,41 @@ TR_J9VMBase::getMemberNameFieldKnotIndexFromMethodHandleKnotIndex(TR::Compilatio
    return knot->getOrCreateIndex(mnObject);
    }
 
+TR::KnownObjectTable::Index
+TR_J9VMBase::getLayoutVarHandle(TR::Compilation *comp, TR::KnownObjectTable::Index layoutIndex)
+   {
+   TR::VMAccessCriticalSection getLayoutVarHandle(this);
+   TR::KnownObjectTable::Index result = TR::KnownObjectTable::UNKNOWN;
+   TR::KnownObjectTable *knot = comp->getKnownObjectTable();
+   if (!knot) return result;
+
+   const char * const layoutClassName =
+      "jdk/internal/foreign/layout/ValueLayouts$AbstractValueLayout";
+   const int layoutClassNameLen = (int)strlen(layoutClassName);
+   TR_OpaqueClassBlock *layoutClass =
+      getSystemClassFromClassName(layoutClassName, layoutClassNameLen);
+
+   TR_OpaqueClassBlock *layoutObjClass =
+      getObjectClassFromKnownObjectIndex(comp, layoutIndex);
+
+   if (layoutClass == NULL ||
+       layoutObjClass == NULL ||
+       isInstanceOf(layoutObjClass, layoutClass, true, true) != TR_yes)
+      {
+      if (comp->getOption(TR_TraceOptDetails))
+         traceMsg(comp, "getLayoutVarHandle: failed ValueLayouts$AbstractValueLayout type check.\n");
+      return result;
+      }
+
+   uintptr_t layoutObj = knot->getPointer(layoutIndex);
+   uintptr_t vhObject = getReferenceField(layoutObj,
+                                 "handle",
+                                 "Ljava/lang/invoke/VarHandle;");
+   if (!vhObject) return result;
+   result = knot->getOrCreateIndex(vhObject);
+   return result;
+   }
+
 TR::KnownObjectTable::Index
 TR_J9VMBase::getMethodHandleTableEntryIndex(TR::Compilation *comp, TR::KnownObjectTable::Index vhIndex, TR::KnownObjectTable::Index adIndex)
    {
diff --git a/runtime/compiler/env/VMJ9.h b/runtime/compiler/env/VMJ9.h
@@ -455,6 +455,18 @@ class TR_J9VMBase : public TR_FrontEnd
     */
    virtual TR::KnownObjectTable::Index getMethodHandleTableEntryIndex(TR::Compilation *comp, TR::KnownObjectTable::Index vhIndex, TR::KnownObjectTable::Index adIndex);
 
+   /**
+    * @brief Get the known object index of a cached MemorySegmentView VarHandle belonging to a value Layout object.
+    * When the Layout object is known, we can evaluate the result of
+    * jdk/internal/foreign/layout/ValueLayouts$AbstractValueLayout.accessHandle()Ljava/lang/invoke/VarHandle;
+    * as long as the layout object's handle field is not null.
+    *
+    * @param comp the compilation
+    * @param layoutIndex the ValueLayout$AbstractValueLayout object index
+    * @return TR::KnownObjectTable::Index the known object index of the MemorySegmentView VarHandle, TR::KnownObjectTable::UNKNOWN otherwise
+    */
+   virtual TR::KnownObjectTable::Index getLayoutVarHandle(TR::Compilation *comp, TR::KnownObjectTable::Index layoutIndex);
+
    virtual TR::Method * createMethod(TR_Memory *, TR_OpaqueClassBlock *, int32_t);
    virtual TR_ResolvedMethod * createResolvedMethod(TR_Memory *, TR_OpaqueMethodBlock *, TR_ResolvedMethod * = 0, TR_OpaqueClassBlock * = 0);
    virtual TR_ResolvedMethod * createResolvedMethodWithVTableSlot(TR_Memory *, uint32_t vTableSlot, TR_OpaqueMethodBlock * aMethod, TR_ResolvedMethod * owningMethod = 0, TR_OpaqueClassBlock * classForNewInstance = 0);
diff --git a/runtime/compiler/env/VMJ9Server.cpp b/runtime/compiler/env/VMJ9Server.cpp
@@ -2416,6 +2416,22 @@ TR_J9ServerVM::getMethodHandleTableEntryIndex(TR::Compilation *comp, TR::KnownOb
    return mhIndex;
    }
 
+TR::KnownObjectTable::Index
+TR_J9ServerVM::getLayoutVarHandle(TR::Compilation *comp, TR::KnownObjectTable::Index layoutIndex)
+   {
+   TR::KnownObjectTable *knot = comp->getKnownObjectTable();
+   if (!knot) return TR::KnownObjectTable::UNKNOWN;
+
+   JITServer::ServerStream *stream = _compInfoPT->getMethodBeingCompiled()->_stream;
+   stream->write(JITServer::MessageType::VM_getLayoutVarHandle, layoutIndex);
+   auto recv = stream->read<TR::KnownObjectTable::Index, uintptr_t *>();
+
+   TR::KnownObjectTable::Index vhIndex = std::get<0>(recv);
+   knot->updateKnownObjectTableAtServer(vhIndex, std::get<1>(recv));
+   return vhIndex;
+
+   }
+
 #endif /* defined(J9VM_OPT_OPENJDK_METHODHANDLE) */
 
 TR::KnownObjectTable::Index
diff --git a/runtime/compiler/env/VMJ9Server.hpp b/runtime/compiler/env/VMJ9Server.hpp
@@ -245,6 +245,7 @@ class TR_J9ServerVM: public TR_J9VM
    virtual UDATA getVMTargetOffset() override;
    virtual UDATA getVMIndexOffset() override;
    virtual TR::KnownObjectTable::Index getMethodHandleTableEntryIndex(TR::Compilation *comp, TR::KnownObjectTable::Index vhIndex, TR::KnownObjectTable::Index adIndex) override;
+   virtual TR::KnownObjectTable::Index getLayoutVarHandle(TR::Compilation *comp, TR::KnownObjectTable::Index layoutIndex) override;
 #endif
    virtual TR::KnownObjectTable::Index getMemberNameFieldKnotIndexFromMethodHandleKnotIndex(TR::Compilation *comp, TR::KnownObjectTable::Index mhIndex, const char *fieldName) override;
    virtual bool isMethodHandleExpectedType(TR::Compilation *comp, TR::KnownObjectTable::Index mhIndex, TR::KnownObjectTable::Index expectedTypeIndex) override;
diff --git a/runtime/compiler/env/j9method.cpp b/runtime/compiler/env/j9method.cpp
@@ -3678,6 +3678,12 @@ void TR_ResolvedJ9Method::construct()
       {  TR::unknownMethod}
       };
 
+   static X ValueLayoutsAbstractValueLayoutMethods[] =
+      {
+      {x(TR::jdk_internal_foreign_layout_ValueLayouts_AbstractValueLayout_accessHandle,        "accessHandle",    "()Ljava/lang/invoke/VarHandle;")},
+      {TR::unknownMethod}
+      };
+
    static X ILGenMacrosMethods[] =
       {
       {  TR::java_lang_invoke_ILGenMacros_placeholder ,      11, "placeholder",      (int16_t)-1, "*"},
@@ -4364,13 +4370,19 @@ void TR_ResolvedJ9Method::construct()
       { 0 }
       };
 
+   static Y class60[] =
+      {
+      { "jdk/internal/foreign/layout/ValueLayouts$AbstractValueLayout", ValueLayoutsAbstractValueLayoutMethods },
+      { 0 }
+      };
+
    static Y * recognizedClasses[] =
       {
       0, 0, 0, class13, class14, class15, class16, class17, class18, class19,
       class20, class21, class22, class23, class24, class25, 0, class27, class28, class29,
       class30, class31, class32, class33, class34, class35, class36, 0, class38, class39,
       class40, class41, class42, class43, class44, class45, class46, class47, class48, class49,
-      class50, 0, 0, class53, 0, class55
+      class50, 0, 0, class53, 0, class55, 0, 0, 0, 0, class60
       };
 
    const int32_t minRecognizedClassLength = 10;
@@ -4762,6 +4774,10 @@ void TR_ResolvedJ9Method::construct()
                setRecognizedMethodInfo(TR::java_lang_invoke_VarHandleByteArrayAsX_ByteBufferHandle_method);
                }
             }
+         else if ((classNameLen == 31) && !strncmp(className, "java/lang/foreign/MemorySegment", 31))
+            {
+            setRecognizedMethodInfo(TR::java_lang_foreign_MemorySegment_method);
+            }
 #endif
          else if ((classNameLen >= 59 + 3 && classNameLen <= 59 + 7) && !strncmp(className, "java/lang/invoke/ArrayVarHandle$ArrayVarHandleOperations$Op", 59))
             {
@@ -6644,7 +6660,7 @@ TR_ResolvedJ9Method::getResolvedInterfaceMethod(TR::Compilation * comp, TR_Opaqu
       if (m)
          {
          c = m->classOfMethod();
-         if (c && !fej9->isInterfaceClass(c))
+         if (c)
             {
             TR::DebugCounter::incStaticDebugCounter(comp, "resources.resolvedMethods/interface");
             TR::DebugCounter::incStaticDebugCounter(comp, "resources.resolvedMethods/interface:#bytes", sizeof(TR_ResolvedJ9Method));
diff --git a/runtime/compiler/net/CommunicationStream.hpp b/runtime/compiler/net/CommunicationStream.hpp
@@ -128,7 +128,7 @@ class CommunicationStream
    // likely to lose an increment when merging/rebasing/etc.
    //
    static const uint8_t MAJOR_NUMBER = 1;
-   static const uint16_t MINOR_NUMBER = 66; // ID: x8xfrdqf9UCc8E2OiuMO
+   static const uint16_t MINOR_NUMBER = 67; // ID: eJ9ynEzY1XevZygVpTws
    static const uint8_t PATCH_NUMBER = 0;
    static uint32_t CONFIGURATION_FLAGS;
 
diff --git a/runtime/compiler/net/MessageTypes.cpp b/runtime/compiler/net/MessageTypes.cpp
@@ -191,6 +191,7 @@ const char *messageNames[] =
    "VM_inSnapshotMode",
    "VM_isInvokeCacheEntryAnArray",
    "VM_getMethodHandleTableEntryIndex",
+   "VM_getLayoutVarHandle",
    "CompInfo_isCompiled",
    "CompInfo_getPCIfCompiled",
    "CompInfo_getInvocationCount",
diff --git a/runtime/compiler/net/MessageTypes.hpp b/runtime/compiler/net/MessageTypes.hpp
@@ -200,6 +200,7 @@ enum MessageType : uint16_t
    VM_inSnapshotMode,
    VM_isInvokeCacheEntryAnArray,
    VM_getMethodHandleTableEntryIndex,
+   VM_getLayoutVarHandle,
 
    // For static TR::CompilationInfo methods
    CompInfo_isCompiled,
diff --git a/runtime/compiler/optimizer/InterpreterEmulator.cpp b/runtime/compiler/optimizer/InterpreterEmulator.cpp
@@ -869,7 +869,7 @@ void
 InterpreterEmulator::maintainStackForCall()
    {
    TR_ASSERT_FATAL(_iteratorWithState, "has to be called when the iterator has state!");
-   int32_t numOfArgs = 0;
+   int32_t numOfArgs = -1;
    TR::DataType returnType = TR::NoType;
    Operand* result = NULL;
 
@@ -908,16 +908,41 @@ InterpreterEmulator::maintainStackForCall()
             isStatic = true;
             break;
          case J9BCinvokedynamic:
-         case J9BCinvokehandle:
-            TR_ASSERT_FATAL(false, "Can't maintain stack for unresolved invokehandle");
+            {
+            // Find the signature, which corresponds to the arguments on the stack.
+            // cpIndex is really the invokedynamic call site index
+            J9ROMClass *romClass = TR::Compiler->cls.romClassOf(method()->classOfMethod());
+            J9SRP *namesAndSigs = (J9SRP*)J9ROMCLASS_CALLSITEDATA(romClass);
+            J9ROMNameAndSignature *nameAndSig = NNSRP_GET(namesAndSigs[cpIndex], J9ROMNameAndSignature*);
+            J9UTF8 *sig = J9ROMNAMEANDSIGNATURE_SIGNATURE(nameAndSig);
+
+            // Parse the signature to determine the number of arguments and
+            // whether or not there is a return value.
+            U_8 sigTypes[256]; // signatures are limited to 255 params + 1 return type
+            UDATA numParams = 0;
+            UDATA numParamSlots = 0;
+            jitParseSignature(sig, sigTypes, &numParams, &numParamSlots);
+            numOfArgs = numParams;
+
+            // returnType is only used to distinguish void return (TR::NoType)
+            // from non-void return (any other value), so it's not necessary to
+            // get the correct non-void type here.
+            if (sigTypes[numParams] == J9_NATIVE_TYPE_VOID)
+               returnType = TR::NoType;
+            else
+               returnType = TR::Int32;
             break;
+            }
 
          default:
             break;
          }
-      TR::Method * calleeMethod = comp()->fej9()->createMethod(trMemory(), _calltarget->_calleeMethod->containingClass(), cpIndex);
-      numOfArgs = calleeMethod->numberOfExplicitParameters() + (isStatic ? 0 : 1);
-      returnType = calleeMethod->returnType();
+      if (numOfArgs < 0)
+         {
+         TR::Method * calleeMethod = comp()->fej9()->createMethod(trMemory(), _calltarget->_calleeMethod->containingClass(), cpIndex);
+         numOfArgs = calleeMethod->numberOfExplicitParameters() + (isStatic ? 0 : 1);
+         returnType = calleeMethod->returnType();
+         }
       }
    maintainStackForCall(result, numOfArgs, returnType);
    }
@@ -1092,6 +1117,19 @@ InterpreterEmulator::getReturnValue(TR_ResolvedMethod *callee)
             }
          break;
          }
+      case TR::jdk_internal_foreign_layout_ValueLayouts_AbstractValueLayout_accessHandle:
+         {
+         Operand* layoutOperand = top();
+         TR::KnownObjectTable::Index layoutIndex = layoutOperand->getKnownObjectIndex();
+         TR::KnownObjectTable *knot = comp()->getKnownObjectTable();
+         if (knot && layoutIndex != TR::KnownObjectTable::UNKNOWN && !knot->isNull(layoutIndex))
+            {
+            TR::KnownObjectTable::Index vhIndex = comp()->fej9()->getLayoutVarHandle(comp(), layoutIndex);
+            if (vhIndex != TR::KnownObjectTable::UNKNOWN)
+               result = new (trStackMemory()) KnownObjOperand(vhIndex);
+            }
+         break;
+         }
 
       default:
          break;
diff --git a/runtime/compiler/optimizer/J9EstimateCodeSize.cpp b/runtime/compiler/optimizer/J9EstimateCodeSize.cpp
@@ -159,12 +159,16 @@ class NeedsPeekingHeuristic
          {
             if (_bci.bcIndex() - _loadIndices[i] <= _distance)
             {
-               _needsPeeking = true;
                heuristicTraceIfTracerIsNotNull(_tracer, "there is a parm load at %d which is within %d of a call at %d", _loadIndices[i], _distance, _bci.bcIndex());
             }
          }
       };
 
+      void setNeedsPeekingToTrue()
+      {
+      _needsPeeking = true;
+      }
+
       void processByteCode()
          {
             if (!_hasArgumentsInfo)
@@ -837,6 +841,20 @@ TR_J9EstimateCodeSize::processBytecodeAndGenerateCFG(TR_CallTarget *calltarget,
             break;
          case J9BCinvokeinterface:
             cpIndex = bci.next2Bytes();
+#if JAVA_SPEC_VERSION >= 21
+            {
+            TR::Method *meth = comp()->fej9()->createMethod(comp()->trMemory(), calltarget->_calleeMethod->containingClass(), cpIndex);
+            if (meth)
+               {
+               const char * sig = meth->signature(comp()->trMemory());
+               if (sig && !strncmp(sig, "java/lang/foreign/MemorySegment", 31))
+                  {
+                  nph.setNeedsPeekingToTrue();
+                  heuristicTrace(tracer(), "Depth %d: invokeinterface call at bc index %d has Signature %s, enabled peeking for caller to fold layout field load necessary for VarHandle operation inlining.", _recursionDepth, i, sig);
+                  }
+               }
+            }
+#endif // JAVA_SPEC_VERSION >= 21
             flags[i].set(InterpreterEmulator::BytecodePropertyFlag::isUnsanitizeable);
             break;
          case J9BCgetfield:
@@ -1326,6 +1344,13 @@ TR_J9EstimateCodeSize::realEstimateCodeSize(TR_CallTarget *calltarget, TR_CallSt
    bool inlineLambdaFormGeneratedMethod = comp()->fej9()->isLambdaFormGeneratedMethod(calltarget->_calleeMethod) &&
                                    (!disableMethodHandleInliningAfterFirstPass || _inliner->firstPass());
 
+   TR::Block * * blocks =
+         (TR::Block * *) comp()->trMemory()->allocateStackMemory(maxIndex
+               * sizeof(TR::Block *));
+   memset(blocks, 0, maxIndex * sizeof(TR::Block *));
+
+   TR::CFG &cfg = processBytecodeAndGenerateCFG(calltarget, cfgRegion, bci, nph, blocks, flags);
+
    // No need to peek LF methods, as we'll always interprete the method with state in order to propagate object info
    // through bytecodes to find call targets
    if (!inlineLambdaFormGeneratedMethod &&
@@ -1350,12 +1375,6 @@ TR_J9EstimateCodeSize::realEstimateCodeSize(TR_CallTarget *calltarget, TR_CallSt
       methodSymbol->getResolvedMethod()->genMethodILForPeekingEvenUnderMethodRedefinition(methodSymbol, comp(), false, NULL);
       }
 
-   TR::Block * * blocks =
-         (TR::Block * *) comp()->trMemory()->allocateStackMemory(maxIndex
-               * sizeof(TR::Block *));
-   memset(blocks, 0, maxIndex * sizeof(TR::Block *));
-
-   TR::CFG &cfg = processBytecodeAndGenerateCFG(calltarget, cfgRegion, bci, nph, blocks, flags);
    int size = calltarget->_fullSize;
 
    // Adjust call frequency for unknown or direct calls, for which we don't get profiling information
@@ -1467,6 +1486,16 @@ TR_J9EstimateCodeSize::realEstimateCodeSize(TR_CallTarget *calltarget, TR_CallSt
       bci.prepareToFindAndCreateCallsites(blocks, flags, callSites, &cfg, &newBCInfo, _recursionDepth, &callStack);
       bool iteratorWithState = (inlineArchetypeSpecimen && !mhInlineWithPeeking) || inlineLambdaFormGeneratedMethod;
 
+// in JDK21+, iterating bytecodes of MemorySegment methods with state is necesssary to obtain the segment view
+// VarHandle object through folding ValueLayouts$AbstractValueLayout.accessHandle() if the VarHandle has been created
+#if JAVA_SPEC_VERSION >= 21
+      if (calltarget->_calleeMethod->getRecognizedMethod() == TR::java_lang_foreign_MemorySegment_method)
+         {
+         heuristicTrace(tracer(), "Callee method is a MemorySegment method. Setting iteratorWithState to true.\n");
+         iteratorWithState = true;
+         }
+#endif // JAVA_SPEC_VERSION >= 21
+
       if (!bci.findAndCreateCallsitesFromBytecodes(wasPeekingSuccessfull, iteratorWithState))
          {
          heuristicTrace(tracer(), "*** Depth %d: ECS end for target %p signature %s. bci.findAndCreateCallsitesFromBytecode failed", _recursionDepth, calltarget, callerName);
diff --git a/runtime/compiler/optimizer/J9TransformUtil.cpp b/runtime/compiler/optimizer/J9TransformUtil.cpp
@@ -678,6 +678,8 @@ bool J9::TransformUtil::foldFinalFieldsIn(TR_OpaqueClassBlock *clazz, const char
       return true; // We can ONLY do this opt to fields that are never victimized by setAccessible
    else if (classNameLength >= 18 && !strncmp(className, "java/lang/reflect/", 18))
       return true;
+   else if (classNameLength >= 18 && !strncmp(className, "java/lang/foreign/", 18))
+      return true;
    else if (classNameLength >= 30 && !strncmp(className, "java/lang/String$UnsafeHelpers", 30))
       return true;
 
