Merge pull request #16171 from gita-omr/mask_reduction_coerced

knn-k · web-flow · commit c23ec23b1087 · 2022-11-10T15:40:52.000+09:00
Handle mask reduction intrinsic in VectorAPIExpansion
diff --git a/runtime/compiler/codegen/J9RecognizedMethodsEnum.hpp b/runtime/compiler/codegen/J9RecognizedMethodsEnum.hpp
@@ -454,6 +454,7 @@
    jdk_internal_vm_vector_VectorSupport_blend,
    jdk_internal_vm_vector_VectorSupport_compare,
    jdk_internal_vm_vector_VectorSupport_fromBitsCoerced,
+   jdk_internal_vm_vector_VectorSupport_maskReductionCoerced,
    jdk_internal_vm_vector_VectorSupport_reductionCoerced,
    jdk_internal_vm_vector_VectorSupport_ternaryOp,
    jdk_internal_vm_vector_VectorSupport_test,
diff --git a/runtime/compiler/env/j9method.cpp b/runtime/compiler/env/j9method.cpp
@@ -3003,6 +3003,7 @@ void TR_ResolvedJ9Method::construct()
       {x(TR::jdk_internal_vm_vector_VectorSupport_blend, "blend", "(Ljava/lang/Class;Ljava/lang/Class;Ljava/lang/Class;ILjdk/internal/vm/vector/VectorSupport$Vector;Ljdk/internal/vm/vector/VectorSupport$Vector;Ljdk/internal/vm/vector/VectorSupport$VectorMask;Ljdk/internal/vm/vector/VectorSupport$VectorBlendOp;)Ljdk/internal/vm/vector/VectorSupport$Vector;")},
       {x(TR::jdk_internal_vm_vector_VectorSupport_compare, "compare", "(ILjava/lang/Class;Ljava/lang/Class;Ljava/lang/Class;ILjdk/internal/vm/vector/VectorSupport$Vector;Ljdk/internal/vm/vector/VectorSupport$Vector;Ljdk/internal/vm/vector/VectorSupport$VectorMask;Ljdk/internal/vm/vector/VectorSupport$VectorCompareOp;)Ljdk/internal/vm/vector/VectorSupport$VectorMask;")},
       {x(TR::jdk_internal_vm_vector_VectorSupport_fromBitsCoerced, "fromBitsCoerced", "(Ljava/lang/Class;Ljava/lang/Class;IJILjdk/internal/vm/vector/VectorSupport$VectorSpecies;Ljdk/internal/vm/vector/VectorSupport$FromBitsCoercedOperation;)Ljdk/internal/vm/vector/VectorSupport$VectorPayload;")},
+      {x(TR::jdk_internal_vm_vector_VectorSupport_maskReductionCoerced, "maskReductionCoerced", "(ILjava/lang/Class;Ljava/lang/Class;ILjdk/internal/vm/vector/VectorSupport$VectorMask;Ljdk/internal/vm/vector/VectorSupport$VectorMaskOp;)J")},
       {x(TR::jdk_internal_vm_vector_VectorSupport_reductionCoerced, "reductionCoerced", "(ILjava/lang/Class;Ljava/lang/Class;Ljava/lang/Class;ILjdk/internal/vm/vector/VectorSupport$Vector;Ljdk/internal/vm/vector/VectorSupport$VectorMask;Ljdk/internal/vm/vector/VectorSupport$ReductionOperation;)J")},
       {x(TR::jdk_internal_vm_vector_VectorSupport_ternaryOp, "ternaryOp", "(ILjava/lang/Class;Ljava/lang/Class;Ljava/lang/Class;ILjdk/internal/vm/vector/VectorSupport$Vector;Ljdk/internal/vm/vector/VectorSupport$Vector;Ljdk/internal/vm/vector/VectorSupport$Vector;Ljdk/internal/vm/vector/VectorSupport$VectorMask;Ljdk/internal/vm/vector/VectorSupport$TernaryOperation;)Ljdk/internal/vm/vector/VectorSupport$Vector;")},
       {x(TR::jdk_internal_vm_vector_VectorSupport_test, "test", "(ILjava/lang/Class;Ljava/lang/Class;ILjdk/internal/vm/vector/VectorSupport$VectorMask;Ljdk/internal/vm/vector/VectorSupport$VectorMask;Ljava/util/function/BiFunction;)Z")},
diff --git a/runtime/compiler/optimizer/VectorAPIExpansion.cpp b/runtime/compiler/optimizer/VectorAPIExpansion.cpp
@@ -682,8 +682,8 @@ TR_VectorAPIExpansion::validateSymRef(int32_t id, int32_t i, vec_sz_t &classLeng
                methodType != classType)
          {
          if (_trace)
-            traceMsg(comp(), "%s invalidating6 class #%d due to symref #%d method type %d, seen type %d\n",
-                     OPT_DETAILS_VECTOR, id, i, (int)methodType, (int)classType);
+            traceMsg(comp(), "%s invalidating6 class #%d due to symref #%d method type %s, seen type %s\n",
+                     OPT_DETAILS_VECTOR, id, i, TR::DataType::getName(methodType), TR::DataType::getName(classType));
          return false;
          }
       }
@@ -1602,6 +1602,14 @@ TR::Node *TR_VectorAPIExpansion::binaryIntrinsicHandler(TR_VectorAPIExpansion *o
    return naryIntrinsicHandler(opt, treeTop, node, elementType, vectorLength, numLanes, mode, 2, Other);
    }
 
+TR::Node *TR_VectorAPIExpansion::maskReductionCoercedIntrinsicHandler(TR_VectorAPIExpansion *opt, TR::TreeTop *treeTop, TR::Node *node,
+                                                                  TR::DataType elementType, TR::VectorLength vectorLength, int32_t numLanes,
+                                                                  handlerMode mode)
+   {
+   return naryIntrinsicHandler(opt, treeTop, node, elementType, vectorLength, numLanes, mode, 1, MaskReduction);
+   }
+
+
 TR::Node *TR_VectorAPIExpansion::reductionCoercedIntrinsicHandler(TR_VectorAPIExpansion *opt, TR::TreeTop *treeTop, TR::Node *node,
                                                                   TR::DataType elementType, TR::VectorLength vectorLength, int32_t numLanes,
                                                                   handlerMode mode)
@@ -1633,11 +1641,16 @@ TR::Node *TR_VectorAPIExpansion::naryIntrinsicHandler(TR_VectorAPIExpansion *opt
    TR::Node *opcodeNode = node->getFirstChild();
    int firstOperand = 5;
 
-   if (opCodeType == Test)
+   if (opCodeType == Test || opCodeType == MaskReduction)
       firstOperand = 4;
 
-   TR::Node *maskNode = node->getChild(firstOperand + numChildren);  // each intrinsic has a mask argument
-   bool withMask = !maskNode->isConstZeroValue();
+   bool withMask = false;
+
+   if (opCodeType != MaskReduction)
+      {
+      TR::Node *maskNode = node->getChild(firstOperand + numChildren);  // each intrinsic has a mask argument
+      withMask = !maskNode->isConstZeroValue();
+      }
 
    if (withMask) numChildren++;
 
@@ -1758,12 +1771,27 @@ TR::Node *TR_VectorAPIExpansion::fromBitsCoercedIntrinsicHandler(TR_VectorAPIExp
    {
    TR::Compilation *comp = opt->comp();
 
+   TR::Node *broadcastTypeNode = node->getChild(4);
+
+   if (!broadcastTypeNode->getOpCode().isLoadConst())
+      {
+      if (opt->_trace) traceMsg(comp, "Unknown broadcast type in node %p\n", node);
+      return NULL;
+      }
+
+   int32_t broadcastType = broadcastTypeNode->get32bitIntegralValue();
+
+   TR_ASSERT_FATAL(broadcastType == MODE_BROADCAST || broadcastType == MODE_BITS_COERCED_LONG_TO_MASK,
+                   "Unexpected broadcast type in node %p\n", node);
+
+   bool mask = (broadcastType == MODE_BITS_COERCED_LONG_TO_MASK);
+
    if (mode == checkScalarization)
-      return node;
+      return mask ? NULL : node;
 
    if (mode == checkVectorization)
       {
-      TR::ILOpCodes splatsOpCode = TR::ILOpCode::createVectorOpCode(TR::vsplats,
+      TR::ILOpCodes splatsOpCode = TR::ILOpCode::createVectorOpCode(mask ? TR::mLongBitsToMask : TR::vsplats,
                                                                     TR::DataType::createVectorType(elementType, vectorLength));
 
       if (!comp->cg()->getSupportsOpCodeForAutoSIMD(splatsOpCode))
@@ -1829,7 +1857,7 @@ TR::Node *TR_VectorAPIExpansion::fromBitsCoercedIntrinsicHandler(TR_VectorAPIExp
       {
       node->setAndIncChild(0, newNode);
       node->setNumChildren(1);
-      TR::ILOpCodes splatsOpCode = TR::ILOpCode::createVectorOpCode(TR::vsplats,
+      TR::ILOpCodes splatsOpCode = TR::ILOpCode::createVectorOpCode(mask ? TR::mLongBitsToMask : TR::vsplats,
                                                                     TR::DataType::createVectorType(elementType, vectorLength));
 
       TR::Node::recreate(node, splatsOpCode);
@@ -1926,7 +1954,7 @@ TR::ILOpCodes TR_VectorAPIExpansion::ILOpcodeFromVectorAPIOpcode(int32_t vectorA
          case VECTOR_OP_MIN: return scalar ? TR::BadILOp : TR::ILOpCode::createVectorOpCode(TR::vreductionMin, vectorType);
          case VECTOR_OP_MAX: return scalar ? TR::BadILOp : TR::ILOpCode::createVectorOpCode(TR::vreductionMax, vectorType);
          case VECTOR_OP_AND: return scalar ? TR::BadILOp : TR::ILOpCode::createVectorOpCode(TR::vreductionAnd, vectorType);
-         case VECTOR_OP_OR:  return scalar ? TR::BadILOp : TR::ILOpCode::createVectorOpCode(TR::vreductionOr, vectorType);
+         case VECTOR_OP_OR:  return scalar ? TR::BadILOp : TR::ILOpCode::createVectorOpCode(TR::vreductionOr,  vectorType);
          case VECTOR_OP_XOR: return scalar ? TR::BadILOp : TR::ILOpCode::createVectorOpCode(TR::vreductionXor, vectorType);
             // These don't seem to be generated by the library:
             // vreductionOrUnchecked
@@ -1935,6 +1963,18 @@ TR::ILOpCodes TR_VectorAPIExpansion::ILOpcodeFromVectorAPIOpcode(int32_t vectorA
             return TR::BadILOp;
          }
       }
+   else if (opCodeType == MaskReduction)
+      {
+      switch (vectorAPIOpCode)
+         {
+         case VECTOR_OP_MASK_TRUECOUNT: return scalar ? TR::BadILOp : TR::ILOpCode::createVectorOpCode(TR::mTrueCount, vectorType);
+         case VECTOR_OP_MASK_FIRSTTRUE: return scalar ? TR::BadILOp : TR::ILOpCode::createVectorOpCode(TR::mFirstTrue, vectorType);
+         case VECTOR_OP_MASK_LASTTRUE:  return scalar ? TR::BadILOp : TR::ILOpCode::createVectorOpCode(TR::mLastTrue, vectorType);
+         case VECTOR_OP_MASK_TOLONG:    return scalar ? TR::BadILOp : TR::ILOpCode::createVectorOpCode(TR::mToLongBits, vectorType);
+         default:
+            return TR::BadILOp;
+         }
+      }
    else if (withMask)
       {
       switch (vectorAPIOpCode)
@@ -2113,16 +2153,17 @@ TR::Node *TR_VectorAPIExpansion::transformNary(TR_VectorAPIExpansion *opt, TR::T
 TR_VectorAPIExpansion::methodTableEntry
 TR_VectorAPIExpansion::methodTable[] =
    {
-   {loadIntrinsicHandler,             Unknown, {Unknown, ElementType, NumLanes}},                                           // jdk_internal_vm_vector_VectorSupport_load
-   {storeIntrinsicHandler,            Unknown, {Unknown, ElementType, NumLanes, Unknown, Unknown, Vector}},                 // jdk_internal_vm_vector_VectorSupport_store
-   {binaryIntrinsicHandler,           Vector,  {Unknown, Unknown, Unknown, ElementType, NumLanes, Vector, Vector, Mask}},   // jdk_internal_vm_vector_VectorSupport_binaryOp
-   {blendIntrinsicHandler,            Vector,  {Unknown, Unknown, ElementType, NumLanes, Vector, Vector, Vector, Unknown}}, // jdk_internal_vm_vector_VectorSupport_blend
-   {compareIntrinsicHandler,          Mask,    {Unknown, Unknown, Unknown, ElementType, NumLanes, Vector, Vector, Mask}},   // jdk_internal_vm_vector_VectorSupport_compare
-   {fromBitsCoercedIntrinsicHandler,  Vector,  {Unknown, ElementType, NumLanes, Unknown, Unknown, Unknown}},                // jdk_internal_vm_vector_VectorSupport_fromBitsCoerced
-   {reductionCoercedIntrinsicHandler, Scalar,  {Unknown, Unknown, Unknown, ElementType, NumLanes, Vector, Mask}},           // jdk_internal_vm_vector_VectorSupport_reductionCoerced
-   {ternaryIntrinsicHandler,          Vector,  {Unknown, Unknown, Unknown, ElementType, NumLanes, Vector, Vector, Vector, Mask}},  // jdk_internal_vm_vector_VectorSupport_ternaryOp
-   {testIntrinsicHandler,             Scalar,  {Unknown, Unknown, ElementType, NumLanes, Mask, Mask, Unknown}},             // jdk_internal_vm_vector_VectorSupport_test
-   {unaryIntrinsicHandler,            Vector,  {Unknown, Unknown, Unknown, ElementType, NumLanes, Vector, Mask}},           // jdk_internal_vm_vector_VectorSupport_unaryOp
+   {loadIntrinsicHandler,                 Unknown, {Unknown, ElementType, NumLanes}},                                           // jdk_internal_vm_vector_VectorSupport_load
+   {storeIntrinsicHandler,                Unknown, {Unknown, ElementType, NumLanes, Unknown, Unknown, Vector}},                 // jdk_internal_vm_vector_VectorSupport_store
+   {binaryIntrinsicHandler,               Vector,  {Unknown, Unknown, Unknown, ElementType, NumLanes, Vector, Vector, Mask}},   // jdk_internal_vm_vector_VectorSupport_binaryOp
+   {blendIntrinsicHandler,                Vector,  {Unknown, Unknown, ElementType, NumLanes, Vector, Vector, Vector, Unknown}}, // jdk_internal_vm_vector_VectorSupport_blend
+   {compareIntrinsicHandler,              Mask,    {Unknown, Unknown, Unknown, ElementType, NumLanes, Vector, Vector, Mask}},   // jdk_internal_vm_vector_VectorSupport_compare
+   {fromBitsCoercedIntrinsicHandler,      Vector,  {Unknown, ElementType, NumLanes, Unknown, Unknown, Unknown}},                // jdk_internal_vm_vector_VectorSupport_fromBitsCoerced
+   {maskReductionCoercedIntrinsicHandler, Scalar,  {Unknown, Unknown, ElementType, NumLanes, Mask}},                            // jdk_internal_vm_vector_VectorSupport_maskReductionCoerced
+   {reductionCoercedIntrinsicHandler,     Scalar,  {Unknown, Unknown, Unknown, ElementType, NumLanes, Vector, Mask}},           // jdk_internal_vm_vector_VectorSupport_reductionCoerced
+   {ternaryIntrinsicHandler,              Vector,  {Unknown, Unknown, Unknown, ElementType, NumLanes, Vector, Vector, Vector, Mask}},  // jdk_internal_vm_vector_VectorSupport_ternaryOp
+   {testIntrinsicHandler,                 Scalar,  {Unknown, Unknown, ElementType, NumLanes, Mask, Mask, Unknown}},             // jdk_internal_vm_vector_VectorSupport_test
+   {unaryIntrinsicHandler,                Vector,  {Unknown, Unknown, Unknown, ElementType, NumLanes, Vector, Mask}},           // jdk_internal_vm_vector_VectorSupport_unaryOp
    };
 
 
diff --git a/runtime/compiler/optimizer/VectorAPIExpansion.hpp b/runtime/compiler/optimizer/VectorAPIExpansion.hpp
@@ -120,7 +120,14 @@ class TR_VectorAPIExpansion : public TR::Optimization
    static int32_t const VECTOR_OP_URSHIFT = 16;
 
    static int32_t const VECTOR_OP_CAST        = 17;
-   static int32_t const VECTOR_OP_REINTERPRET = 18;
+   static int32_t const VECTOR_OP_UCAST       = 18;
+   static int32_t const VECTOR_OP_REINTERPRET = 19;
+
+   // Mask manipulation operations
+   static int32_t const VECTOR_OP_MASK_TRUECOUNT = 20;
+   static int32_t const VECTOR_OP_MASK_FIRSTTRUE = 21;
+   static int32_t const VECTOR_OP_MASK_LASTTRUE  = 22;
+   static int32_t const VECTOR_OP_MASK_TOLONG    = 23;
 
    // Compare
    static int32_t const BT_eq = 0;
@@ -137,6 +144,10 @@ class TR_VectorAPIExpansion : public TR::Optimization
    static int32_t const BT_ult = BT_lt | BT_unsigned_compare;
    static int32_t const BT_ugt = BT_gt | BT_unsigned_compare;
 
+   // Various broadcasting modes.
+   static int32_t const MODE_BROADCAST = 0;
+   static int32_t const MODE_BITS_COERCED_LONG_TO_MASK = 1;
+
   /** \brief
    *  Is passed to methods handlers during analysis and transforamtion phases
    *
@@ -175,6 +186,7 @@ class TR_VectorAPIExpansion : public TR::Optimization
    enum vapiOpCodeType
       {
       Compare,
+      MaskReduction,
       Reduction,
       Test,
       Other
@@ -800,6 +812,38 @@ class TR_VectorAPIExpansion : public TR::Optimization
    */
    static TR::Node *binaryIntrinsicHandler(TR_VectorAPIExpansion *opt, TR::TreeTop *treeTop, TR::Node *node, TR::DataType elementType, TR::VectorLength vectorLength, int32_t numLanes, handlerMode mode);
 
+  /** \brief
+   *    Scalarizes or vectorizes a node that is a call to \c VectorSupport.maskReductionCoerced() intrinsic.
+   *    In both cases, the node is modified in place.
+   *    In the case of scalarization, extra nodes are created(number of lanes minus one)
+   *
+   *   \param opt
+   *      This optimization object
+   *
+   *   \param treeTop
+   *      Tree top of the \c node
+   *
+   *   \param node
+   *      Node to transform
+   *
+   *   \param elementType
+   *      Element type
+   *
+   *   \param vectorLength
+   *      Vector length
+   *
+   *   \param numLanes
+   *      Number of elements
+   *
+   *   \param mode
+   *      Handler mode
+   *
+   *   \return
+   *      Transformed node
+   *
+   */
+   static TR::Node *maskReductionCoercedIntrinsicHandler(TR_VectorAPIExpansion *opt, TR::TreeTop *treeTop, TR::Node *node, TR::DataType elementType, TR::VectorLength vectorLength, int32_t numLanes, handlerMode mode);
+
   /** \brief
    *    Scalarizes or vectorizes a node that is a call to \c VectorSupport.reductionCoerced() intrinsic.
    *    In both cases, the node is modified in place.
