From 2657050cbb66ed09c45f3c51452e9619ea8e3a19 Mon Sep 17 00:00:00 2001
From: Rahil Shah <rahil@ca.ibm.com>
Date: Mon, 20 Jan 2025 15:45:33 -0500
Subject: [PATCH] Add codegen query for supporting LXA instructions

Load indexed address register instructions are used to generate address to
access element in the array. This commit adds a code-gen query for using LXA
instruction to used at various places while doing pattern matching.

Signed-off-by: Rahil Shah <rahil@ca.ibm.com>
---
 compiler/z/codegen/OMRCodeGenerator.cpp | 16 ++++++++++------
 compiler/z/codegen/OMRCodeGenerator.hpp |  7 +++++--
 compiler/z/codegen/OMRTreeEvaluator.cpp |  7 +------
 3 files changed, 16 insertions(+), 14 deletions(-)

diff --git a/compiler/z/codegen/OMRCodeGenerator.cpp b/compiler/z/codegen/OMRCodeGenerator.cpp
index 6b531a05c6a..aed657b0195 100644
--- a/compiler/z/codegen/OMRCodeGenerator.cpp
+++ b/compiler/z/codegen/OMRCodeGenerator.cpp
@@ -408,17 +408,12 @@ OMR::Z::CodeGenerator::lowerTreeIfNeeded(
          }
       }
 
-   static const bool canEmulateLXA = TR::InstOpCode(TR::InstOpCode::LXAB).canEmulate() &&
-                                     TR::InstOpCode(TR::InstOpCode::LXAH).canEmulate() &&
-                                     TR::InstOpCode(TR::InstOpCode::LXAF).canEmulate() &&
-                                     TR::InstOpCode(TR::InstOpCode::LXAG).canEmulate() &&
-                                     TR::InstOpCode(TR::InstOpCode::LXAQ).canEmulate();
    static bool disableLXAUncommoning = feGetEnv("TR_disableLXAUncommoning") != NULL;
 
    if (!disableLXAUncommoning &&
        (node->getOpCodeValue() == TR::aiadd || node->getOpCodeValue() == TR::aladd) &&
        !parent->getOpCode().isLoad() &&
-       (canEmulateLXA || self()->comp()->target().cpu.isAtLeast(OMR_PROCESSOR_S390_ZNEXT)))
+       self()->getUseLXAInstructions())
       {
       // To enable generating more LXAs, perform uncommoning on trees that look like this:
       // axadd
@@ -650,6 +645,15 @@ OMR::Z::CodeGenerator::initialize()
       comp->setOption(TR_DisableVectorBCD);
       }
 
+   static bool canEmulateLXA = TR::InstOpCode(TR::InstOpCode::LXAB).canEmulate() &&
+                                     TR::InstOpCode(TR::InstOpCode::LXAH).canEmulate() &&
+                                     TR::InstOpCode(TR::InstOpCode::LXAF).canEmulate() &&
+                                     TR::InstOpCode(TR::InstOpCode::LXAG).canEmulate() &&
+                                     TR::InstOpCode(TR::InstOpCode::LXAQ).canEmulate();
+   if (canEmulateLXA || comp->target().cpu.isAtLeast(OMR_PROCESSOR_S390_ZNEXT))
+      {
+      cg->setUseLXAInstructions(true);
+      }
    // Be pessimistic until we can prove we don't exit after doing code-generation
    cg->setExitPointsInMethod(true);
 
diff --git a/compiler/z/codegen/OMRCodeGenerator.hpp b/compiler/z/codegen/OMRCodeGenerator.hpp
index fc9210b00cd..3c6f8ac66f0 100644
--- a/compiler/z/codegen/OMRCodeGenerator.hpp
+++ b/compiler/z/codegen/OMRCodeGenerator.hpp
@@ -505,6 +505,9 @@ class OMR_EXTENSIBLE CodeGenerator : public OMR::CodeGenerator
    bool getCondCodeShouldBePreserved() { return _cgFlags.testAny(S390CG_condCodeShouldBePreserved); }
    void setCondCodeShouldBePreserved(bool b) { _cgFlags.set(S390CG_condCodeShouldBePreserved, b); }
 
+   bool getUseLXAInstructions() { return _cgFlags.testAny(S390CG_useLXAInstructions); }
+   void setUseLXAInstructions(bool b) { _cgFlags.set(S390CG_useLXAInstructions, b); }
+
    uint8_t getFCondMoveBranchOpCond() { return fCondMoveBranchOpCond; }
    void setFCondMoveBranchOpCond(TR::InstOpCode::S390BranchCondition b) { fCondMoveBranchOpCond = (getMaskForBranchCondition(b)) & 0xF; }
 
@@ -727,7 +730,7 @@ class OMR_EXTENSIBLE CodeGenerator : public OMR::CodeGenerator
     */
    bool yankIndexScalingOp()
       {
-      return false;
+      return !getUseLXAInstructions();
       }
 
    bool excludeInvariantsFromGRAEnabled();
@@ -836,7 +839,7 @@ class OMR_EXTENSIBLE CodeGenerator : public OMR::CodeGenerator
       S390CG_condCodeShouldBePreserved   = 0x00004000,
       S390CG_enableBranchPreload         = 0x00008000,
       S390CG_globalStaticBaseRegisterOn  = 0x00010000,
-      // Available                       = 0x00020000,
+      S390CG_useLXAInstructions          = 0x00020000,
       S390CG_canExceptByTrap             = 0x00040000,
       S390CG_enableTLHPrefetching        = 0x00080000,
       S390CG_enableBranchPreloadForCalls = 0x00100000,
diff --git a/compiler/z/codegen/OMRTreeEvaluator.cpp b/compiler/z/codegen/OMRTreeEvaluator.cpp
index 7b8a2493f1b..f1b43ab8107 100644
--- a/compiler/z/codegen/OMRTreeEvaluator.cpp
+++ b/compiler/z/codegen/OMRTreeEvaluator.cpp
@@ -7877,13 +7877,8 @@ OMR::Z::TreeEvaluator::axaddEvaluator(TR::Node * node, TR::CodeGenerator * cg)
    TR::MemoryReference * axaddMR = generateS390MemoryReference(cg);
    TR::InstOpCode::Mnemonic loadOp;
    static const bool disableLXAaxaddZNext = feGetEnv("TR_disableLXAaxaddZNext") != NULL;
-   static const bool canEmulateLXA = TR::InstOpCode(TR::InstOpCode::LXAB).canEmulate() &&
-                                     TR::InstOpCode(TR::InstOpCode::LXAH).canEmulate() &&
-                                     TR::InstOpCode(TR::InstOpCode::LXAF).canEmulate() &&
-                                     TR::InstOpCode(TR::InstOpCode::LXAG).canEmulate() &&
-                                     TR::InstOpCode(TR::InstOpCode::LXAQ).canEmulate();
 
-   axaddMR->populateAddTree(node, cg, &loadOp, !disableLXAaxaddZNext && (cg->comp()->target().cpu.isAtLeast(OMR_PROCESSOR_S390_ZNEXT) || canEmulateLXA));
+   axaddMR->populateAddTree(node, cg, &loadOp, !disableLXAaxaddZNext && cg->getUseLXAInstructions());
    axaddMR->eliminateNegativeDisplacement(node, cg);
    axaddMR->enforceDisplacementLimit(node, cg, NULL);