Enable OpenJDK VarHandle call chain inlining

nbhuiyan · nbhuiyan · commit 4aa88c0bede8 · 2024-01-24T15:25:22.000-05:00
This commit contains the necessary changes to support
refining VM INL calls in VarHandle call chains, as well
as support for inlining the call targets down to the Unsafe
operations. This is necessary for improved OpenJDK VarHandle
performance.

VarHandles must be stored in static final fields for this to
work. VarHandle final fields will undergo folding during
ILGen by default in Java 17+ as we need this folding to
happen prior to inlining. It is not possible to modify
static final fields by conventional means in JDK17+.

In general, when VarHandle object is known, we can evaluate
the steps to obtain the target MethodHandle at compile time.
For that, we obtain the result of the following calls in
MethodHandleTransformer and InterpreterEmulator:
* Invokers.directVarHandleTarget
* VarHandle.asDirect
* Invokers.checkVarHandleGenericType

Signed-off-by: Nazim Bhuiyan &lt;nubhuiyan@ibm.com&gt;
diff --git a/runtime/compiler/codegen/J9RecognizedMethodsEnum.hpp b/runtime/compiler/codegen/J9RecognizedMethodsEnum.hpp
@@ -1031,6 +1031,8 @@
    java_lang_invoke_Invokers_checkCustomized,
    java_lang_invoke_Invokers_checkExactType,
    java_lang_invoke_Invokers_getCallSiteTarget,
+   java_lang_invoke_Invokers_directVarHandleTarget,
+   java_lang_invoke_Invokers_checkVarHandleGenericType,
    java_lang_invoke_MethodHandle_doCustomizationLogic,
    java_lang_invoke_MethodHandle_asType,
    java_lang_invoke_MethodHandle_asType_instance,
@@ -1092,6 +1094,7 @@
    java_lang_invoke_FieldGetterHandle_invokeExact,
    java_lang_invoke_FieldSetterHandle_invokeExact,
    java_lang_invoke_FilterArgumentsHandle_invokeExact,
+   java_lang_invoke_VarHandle_asDirect,
    java_lang_invoke_VarHandle_get,
    java_lang_invoke_VarHandle_set,
    java_lang_invoke_VarHandle_getVolatile,
diff --git a/runtime/compiler/env/VMJ9.cpp b/runtime/compiler/env/VMJ9.cpp
@@ -5043,6 +5043,100 @@ TR_J9VMBase::getMemberNameFieldKnotIndexFromMethodHandleKnotIndex(TR::Compilatio
    return knot->getOrCreateIndex(mnObject);
    }
 
+TR::KnownObjectTable::Index
+TR_J9VMBase::getMethodHandleTableEntryIndex(TR::Compilation *comp, TR::KnownObjectTable::Index vhIndex, TR::KnownObjectTable::Index adIndex)
+   {
+   TR::VMAccessCriticalSection getMethodHandleTableEntryIndex(this);
+   TR::KnownObjectTable::Index result = TR::KnownObjectTable::UNKNOWN;
+   TR::KnownObjectTable *knot = comp->getKnownObjectTable();
+   if (!knot) return result;
+
+   uintptr_t varHandleObj = knot->getPointer(vhIndex);
+   uintptr_t accessDescriptorObj = knot->getPointer(adIndex);
+   uintptr_t mhTable = 0;
+   uintptr_t mtTable = 0;
+#if JAVA_SPEC_VERSION <= 17
+   uintptr_t typesAndInvokersObj = getReferenceField(varHandleObj,
+                                                      "typesAndInvokers",
+                                                      "Ljava/lang/invoke/VarHandle$TypesAndInvokers;");
+   if (!typesAndInvokersObj) return result;
+
+   mhTable = getReferenceField(typesAndInvokersObj,
+                                 "methodHandle_table",
+                                 "[Ljava/lang/invoke/MethodHandle;");
+
+   mtTable = getReferenceField(typesAndInvokersObj,
+                                    "methodType_table",
+                                    "[Ljava/lang/invoke/MethodType;");
+#else
+   mhTable = getReferenceField(varHandleObj,
+                                 "methodHandleTable",
+                                 "[Ljava/lang/invoke/MethodHandle;");
+
+   mtTable = getReferenceField(varHandleObj,
+                                 "methodTypeTable",
+                                 "[Ljava/lang/invoke/MethodType;");
+#endif // JAVA_SPEC_VERSION <= 17
+   if (!mhTable || !mtTable) return result;
+
+#if JAVA_SPEC_VERSION >= 17
+   // if the VarHandle has invokeExact behaviour, then the MethodType in
+   // AccessDescriptor.symbolicMethodTypeExact field must match the corresponding
+   // entry in the VarHandle's method type table, failing which
+   // WrongMethodTypeException should be thrown.
+   int32_t varHandleExactFieldOffset =
+         getInstanceFieldOffset(getObjectClass(varHandleObj), "exact", "Z");
+   int32_t varHandleHasInvokeExactBehaviour = getInt32FieldAt(varHandleObj, varHandleExactFieldOffset);
+   if (varHandleHasInvokeExactBehaviour)
+      {
+      int32_t mtEntryIndex = getInt32Field(accessDescriptorObj, "type");
+      uintptr_t methodTypeTableEntryObj = getReferenceElement(mtTable, mtEntryIndex);
+      if (!methodTypeTableEntryObj) return result;
+      uintptr_t symbolicMTExactObj = getReferenceField(accessDescriptorObj,
+                                                         "symbolicMethodTypeExact",
+                                                         "Ljava/lang/invoke/MethodType;");
+      if (methodTypeTableEntryObj != symbolicMTExactObj)
+         return result;
+      }
+#endif // JAVA_SPEC_VERSION >= 17
+
+   int32_t mhEntryIndex = getInt32Field(accessDescriptorObj, "mode");
+   uintptr_t methodHandleObj = getReferenceElement(mhTable, mhEntryIndex);
+
+   if (!methodHandleObj) return result;
+
+   // For the MethodHandle obtained from the VarHandle's MH table, the type must match
+   // the MethodType in AccessDescriptor.symbolicMethodTypeInvoker field, failing which
+   // the MH obtained cannot be used directly without asType conversion
+   uintptr_t methodTypeObj = getReferenceField(methodHandleObj,
+                                                "type",
+                                                "Ljava/lang/invoke/MethodType;");
+   uintptr_t symbolicMTInvokerObj = getReferenceField(accessDescriptorObj,
+                                                      "symbolicMethodTypeInvoker",
+                                                      "Ljava/lang/invoke/MethodType;");
+   if (methodTypeObj != symbolicMTInvokerObj)
+      return result;
+
+   result = knot->getOrCreateIndex(methodHandleObj);
+
+   return result;
+   }
+
+
+TR::KnownObjectTable::Index
+TR_J9VMBase::getDirectVarHandleTargetIndex(TR::Compilation* comp, TR::KnownObjectTable::Index vhIndex)
+   {
+   // Since IndirectVarHandle.asDirect can override VarHandle.asDirect, an exact type check here is necessary
+   const char * varHandleClassName = "java/lang/invoke/VarHandle";
+   TR_OpaqueClassBlock * varHandleClass =
+      getSystemClassFromClassName(varHandleClassName, strlen(varHandleClassName));
+   TR_OpaqueClassBlock * objectClass = getObjectClassFromKnownObjectIndex(comp, vhIndex);
+   if (NULL == varHandleClass || NULL == objectClass || varHandleClass != objectClass)
+      return TR::KnownObjectTable::UNKNOWN;
+
+   return vhIndex;
+   }
+
 bool
 TR_J9VMBase::isMethodHandleExpectedType(
    TR::Compilation *comp,
diff --git a/runtime/compiler/env/VMJ9.h b/runtime/compiler/env/VMJ9.h
@@ -437,6 +437,30 @@ class TR_J9VMBase : public TR_FrontEnd
       return getVolatileReferenceFieldAt(objectPointer, getInstanceFieldOffset(getObjectClass(objectPointer), fieldName, fieldSignature));
       }
 
+   /**
+    * @brief Get the known object index of a MethodHandle cached in a VarHandle's MH table
+    * corresponding to the access descriptor. When the VarHandle is known, we can evaluate
+    * the result of java/lang/invoke/Invokers.checkVarHandleGenericType at compile time and
+    * obtain the MethodHandle which can then allow us to inline the call target.
+    *
+    * @param comp the compilation
+    * @param vhIndex the VarHandle object index
+    * @param adIndex the AccessDescriptor object index
+    * @return TR::KnownObjectTable::Index the MH object index if success, TR::KnownObjectTable::UNKNOWN otherwise
+    */
+   TR::KnownObjectTable::Index getMethodHandleTableEntryIndex(TR::Compilation *comp, TR::KnownObjectTable::Index vhIndex, TR::KnownObjectTable::Index adIndex);
+
+   /**
+    * @brief Get the Direct VarHandle target object index. This function evaluates the result of
+    * java/lang/invoke/Invokers.directVarHandleTarget and java/lang/invoke/VarHandle.asDirect
+    * at compile time.
+    *
+    * @param comp the compilation
+    * @param vhIndex the VarHandle object
+    * @return TR::KnownObjectTable::Index the direct VH object index if success, TR::KnownObjectTable::UNKNOWN otherwise
+    */
+   TR::KnownObjectTable::Index getDirectVarHandleTargetIndex(TR::Compilation * comp, TR::KnownObjectTable::Index vhIndex);
+
    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/j9method.cpp b/runtime/compiler/env/j9method.cpp
@@ -3636,6 +3636,7 @@ void TR_ResolvedJ9Method::construct()
       {  TR::java_lang_invoke_VarHandle_getAndBitwiseXor          ,   21, "getAndBitwiseXor_impl",                 (int16_t)-1, "*"},
       {  TR::java_lang_invoke_VarHandle_getAndBitwiseXorAcquire   ,   28, "getAndBitwiseXorAcquire_impl",         (int16_t)-1, "*"},
       {  TR::java_lang_invoke_VarHandle_getAndBitwiseXorRelease   ,   28, "getAndBitwiseXorRelease_impl",         (int16_t)-1, "*"},
+      {x(TR::java_lang_invoke_VarHandle_asDirect                  ,       "asDirect",                             "()Ljava/lang/invoke/VarHandle;")},
       {  TR::unknownMethod}
       };
 
@@ -3680,10 +3681,12 @@ void TR_ResolvedJ9Method::construct()
 
    static X InvokersMethods[] =
       {
-      {TR::java_lang_invoke_Invokers_checkCustomized,            15,       "checkCustomized",             (int16_t)-1, "*"},
-      {TR::java_lang_invoke_Invokers_checkExactType,             14,       "checkExactType",              (int16_t)-1, "*"},
-      {TR::java_lang_invoke_Invokers_getCallSiteTarget,          17,       "getCallSiteTarget",           (int16_t)-1, "*"},
-      {TR::unknownMethod}
+      {  TR::java_lang_invoke_Invokers_checkCustomized,            15,     "checkCustomized",             (int16_t)-1, "*"},
+      {  TR::java_lang_invoke_Invokers_checkExactType,             14,     "checkExactType",              (int16_t)-1, "*"},
+      {  TR::java_lang_invoke_Invokers_getCallSiteTarget,          17,     "getCallSiteTarget",           (int16_t)-1, "*"},
+      {x(TR::java_lang_invoke_Invokers_directVarHandleTarget,              "directVarHandleTarget",       "(Ljava/lang/invoke/VarHandle;)Ljava/lang/invoke/VarHandle;")},
+      {x(TR::java_lang_invoke_Invokers_checkVarHandleGenericType,          "checkVarHandleGenericType",   "(Ljava/lang/invoke/VarHandle;Ljava/lang/invoke/VarHandle$AccessDescriptor;)Ljava/lang/invoke/MethodHandle;")},
+      {  TR::unknownMethod}
       };
 
    static X AsTypeHandleMethods[] =
diff --git a/runtime/compiler/optimizer/InlinerTempForJ9.cpp b/runtime/compiler/optimizer/InlinerTempForJ9.cpp
@@ -5308,6 +5308,7 @@ bool TR_J9InlinerPolicy::isJSR292SmallGetterMethod(TR_ResolvedMethod *resolvedMe
       case TR::java_lang_invoke_MethodHandleImpl_ArrayAccessor_lengthS:
       case TR::java_lang_invoke_MethodHandleImpl_ArrayAccessor_lengthC:
       case TR::java_lang_invoke_MethodHandleImpl_ArrayAccessor_lengthL:
+      case TR::java_lang_invoke_VarHandle_asDirect:
          return true;
 
       default:
diff --git a/runtime/compiler/optimizer/InterpreterEmulator.cpp b/runtime/compiler/optimizer/InterpreterEmulator.cpp
@@ -1023,6 +1023,41 @@ InterpreterEmulator::getReturnValue(TR_ResolvedMethod *callee)
             }
          break;
          }
+      case TR::java_lang_invoke_VarHandle_asDirect:
+      case TR::java_lang_invoke_Invokers_directVarHandleTarget:
+         {
+         Operand* varHandleOperand = top();
+         TR::KnownObjectTable::Index vhIndex = varHandleOperand->getKnownObjectIndex();
+         TR::KnownObjectTable *knot = comp()->getKnownObjectTable();
+         if (knot
+              && vhIndex != TR::KnownObjectTable::UNKNOWN
+              && !knot->isNull(vhIndex))
+            {
+            TR::KnownObjectTable::Index directVHIndex = comp()->fej9()->getDirectVarHandleTargetIndex(comp(), vhIndex);
+            if (directVHIndex != TR::KnownObjectTable::UNKNOWN)
+               result = new (trStackMemory()) KnownObjOperand(directVHIndex);
+            }
+         break;
+         }
+      case TR::java_lang_invoke_Invokers_checkVarHandleGenericType:
+         {
+         Operand* varHandleOperand = topn(1);
+         Operand* accessDescriptorOperand = topn(0);
+         TR::KnownObjectTable::Index vhIndex = varHandleOperand->getKnownObjectIndex();
+         TR::KnownObjectTable::Index adIndex = accessDescriptorOperand->getKnownObjectIndex();
+         TR::KnownObjectTable *knot = comp()->getKnownObjectTable();
+         if (knot
+              && vhIndex != TR::KnownObjectTable::UNKNOWN
+              && adIndex != TR::KnownObjectTable::UNKNOWN
+              && !knot->isNull(vhIndex)
+              && !knot->isNull(adIndex))
+            {
+            TR::KnownObjectTable::Index mhIndex = comp()->fej9()->getMethodHandleTableEntryIndex(comp(), vhIndex, adIndex);
+            if (mhIndex != TR::KnownObjectTable::UNKNOWN)
+               result = new (trStackMemory()) KnownObjOperand(mhIndex);
+            }
+         break;
+         }
 
       default:
          break;
diff --git a/runtime/compiler/optimizer/J9TransformUtil.cpp b/runtime/compiler/optimizer/J9TransformUtil.cpp
@@ -2537,6 +2537,15 @@ J9::TransformUtil::knownObjectFromFinalStatic(
       return TR::KnownObjectTable::UNKNOWN;
 
    static const char *foldVarHandle = feGetEnv("TR_FoldVarHandleWithoutFear");
+
+   bool canFoldVarHandleWithoutFear = false;
+
+   // in Java 17 and above, it is not possible to remove the final attribute of a field
+   // to make it writable via reflection.
+#if JAVA_SPEC_VERSION >= 17
+   canFoldVarHandleWithoutFear = true;
+#endif
+
    int32_t clazzNameLength = 0;
    char *clazzName = fej9->getClassNameChars(declaringClass, clazzNameLength);
    bool createKnownObject = false;
@@ -2545,7 +2554,7 @@ J9::TransformUtil::knownObjectFromFinalStatic(
       {
       createKnownObject = true;
       }
-   else if (foldVarHandle
+   else if ((foldVarHandle || canFoldVarHandleWithoutFear)
             && (clazzNameLength != 16 || strncmp(clazzName, "java/lang/System", 16)))
       {
       TR_OpaqueClassBlock *varHandleClass =  fej9->getSystemClassFromClassName("java/lang/invoke/VarHandle", 26);
diff --git a/runtime/compiler/optimizer/MethodHandleTransformer.cpp b/runtime/compiler/optimizer/MethodHandleTransformer.cpp
@@ -404,6 +404,44 @@ TR_MethodHandleTransformer::getObjectInfoOfNode(TR::Node* node)
                comp(), dmhIndex, trace());
             }
 #endif
+         case TR::java_lang_invoke_Invokers_directVarHandleTarget:
+         case TR::java_lang_invoke_VarHandle_asDirect:
+            {
+            auto vhIndex = getObjectInfoOfNode(node->getLastChild());
+            if (knot && isKnownObject(vhIndex) && !knot->isNull(vhIndex))
+               {
+               auto directVHIndex = comp()->fej9()->getDirectVarHandleTargetIndex(comp(), vhIndex);
+               if (directVHIndex == TR::KnownObjectTable::UNKNOWN)
+                  break;
+               if (trace())
+                  {
+                  if (rm == TR::java_lang_invoke_Invokers_directVarHandleTarget)
+                     traceMsg(comp(), "Invokers_directVarHandleTarget with known VarHandle object %d, updating node n%dn with known object info\n", directVHIndex, node->getGlobalIndex());
+                  else traceMsg(comp(), "VarHandle_asDirect with known VarHandle object %d, updating node n%dn with known object info\n", directVHIndex, node->getGlobalIndex());
+                  }
+               node->setKnownObjectIndex(directVHIndex);
+               return directVHIndex;
+               }
+            break;
+            }
+         case TR::java_lang_invoke_Invokers_checkVarHandleGenericType:
+            {
+            auto vhIndex = getObjectInfoOfNode(node->getFirstArgument());
+            auto adIndex = getObjectInfoOfNode(node->getLastChild());
+            if (knot
+               && isKnownObject(adIndex)
+               && isKnownObject(vhIndex)
+               && !knot->isNull(vhIndex)
+               && !knot->isNull(adIndex))
+               {
+               auto mhIndex = comp()->fej9()->getMethodHandleTableEntryIndex(comp(), vhIndex, adIndex);
+               if (trace())
+                 traceMsg(comp(), "Invokers_checkVarHandleGenericType with known VarHandle object %d, updating node n%dn with known MH object %d from MH table\n", vhIndex, node->getGlobalIndex(), mhIndex);
+               node->setKnownObjectIndex(mhIndex);
+               return mhIndex;
+               }
+            break;
+            }
 
          default:
             break;
