From 4e89f97969fd9a3067fda7f6c64749a7e8ac4930 Mon Sep 17 00:00:00 2001
From: Ehsan Kiani Far <ehsan.kianifar@gmail.com>
Date: Fri, 22 Nov 2024 16:46:00 -0500
Subject: [PATCH] Z: Recognize Thread.onSpinWait()

Thread.onSpinWait() is a simple `nop` instruction on z.
Added instructions to recognize that method and inline it in jit code.
Enabled by default. Disable by setting the `TR_noPauseOnSpinWait`
environment variable.

Signed-off-by: Ehsan Kiani Far <ehsan.kianifar@gmail.com>
---
 runtime/compiler/z/codegen/J9CodeGenerator.cpp |  8 ++++++++
 runtime/compiler/z/codegen/J9TreeEvaluator.cpp | 15 +++++++++++++++
 runtime/compiler/z/codegen/J9TreeEvaluator.hpp | 14 ++++++++++++++
 3 files changed, 37 insertions(+)

diff --git a/runtime/compiler/z/codegen/J9CodeGenerator.cpp b/runtime/compiler/z/codegen/J9CodeGenerator.cpp
index ca878950dbd..4d2b60a5a51 100644
--- a/runtime/compiler/z/codegen/J9CodeGenerator.cpp
+++ b/runtime/compiler/z/codegen/J9CodeGenerator.cpp
@@ -4043,6 +4043,14 @@ J9::Z::CodeGenerator::inlineDirectCall(
             return resultReg != NULL;
             }
          break;
+      case TR::java_lang_Thread_onSpinWait:
+         static const bool enableOSW = feGetEnv("TR_noPauseOnSpinWait") == NULL;
+         if (enableOSW)
+            {
+            resultReg = TR::TreeEvaluator::inlineOnSpinWait(node, cg);
+            return true;
+            }
+         break;
 
       default:
          break;
diff --git a/runtime/compiler/z/codegen/J9TreeEvaluator.cpp b/runtime/compiler/z/codegen/J9TreeEvaluator.cpp
index 3976af7e72c..21a1d097d54 100644
--- a/runtime/compiler/z/codegen/J9TreeEvaluator.cpp
+++ b/runtime/compiler/z/codegen/J9TreeEvaluator.cpp
@@ -14665,3 +14665,18 @@ J9::Z::TreeEvaluator::inlineIntegerStringSize(TR::Node* node, TR::CodeGenerator*
 
    return node->setRegister(lengthReg);
    }
+
+TR::Register*
+J9::Z::TreeEvaluator::inlineOnSpinWait(TR::Node *node, TR::CodeGenerator *cg)
+   {
+   TR::Compilation *comp = cg->comp();
+   if (comp->getOption(TR_TraceCG))
+      {
+      traceMsg(comp, "Inlining Thread.onSpinWait call on node %p to NOP.\n", node);
+      }
+
+   // onSpinWait() method calls VM_AtomicSupport::yieldCPU() which is a simple NOP instruction on Z.
+   // Check omr/include_core/AtomicSupport.hpp for the JNI implementation.
+   TR::Instruction* cursor = new (cg->trHeapMemory()) TR::S390NOPInstruction(TR::InstOpCode::NOP, 2, node, cg);
+   return NULL;
+   }
\ No newline at end of file
diff --git a/runtime/compiler/z/codegen/J9TreeEvaluator.hpp b/runtime/compiler/z/codegen/J9TreeEvaluator.hpp
index 7744037efb0..81e6ad16e32 100644
--- a/runtime/compiler/z/codegen/J9TreeEvaluator.hpp
+++ b/runtime/compiler/z/codegen/J9TreeEvaluator.hpp
@@ -141,6 +141,20 @@ class OMR_EXTENSIBLE TreeEvaluator: public J9::TreeEvaluator
    static TR::Register* inlineUTF16BEEncode    (TR::Node *node, TR::CodeGenerator *cg);
    static TR::Register *inlineCRC32CUpdateBytes(TR::Node *node, TR::CodeGenerator *cg, bool isDirectBuffer);
 
+   /**
+   * \brief
+   * Accelerate inlining onSpinWait() method.
+   *
+   * \details
+   * onSpinWait() method calls VM_AtomicSupport::yieldCPU() which is a simple NOP instruction on Z.
+   *
+   * \param node the method call node.
+   * \param cg the code generator.
+   *
+   * \return NULL
+   */
+   static TR::Register *inlineOnSpinWait(TR::Node *node, TR::CodeGenerator *cg);
+
    static TR::Register *zdloadEvaluator(TR::Node *node, TR::CodeGenerator *cg);
    static TR::Register *zdloadiEvaluator(TR::Node *node, TR::CodeGenerator *cg);
    static TR::Register *zdstoreEvaluator(TR::Node *node, TR::CodeGenerator *cg);