[VE] LEGALAVL and staged VVP legalization

The new LEGALAVL node annotates that the AVL refers to packs of 64bit.
We use a two-stage lowering approach with LEGALAVL:

First, standard SDNodes are translated into illegal VVP layer nodes.
Regardless of source (VP or standard), all VVP nodes have a mask and AVL
parameter. The AVL parameter refers to the element position (just as in
VP intrinsics).

Second, we legalize the AVL usage in VVP layer nodes. If the element
size is < 64bit, the EVL parameter has to be adjusted to refer to packs
of 64bits.  We wrap the legalized AVL in a LEGALAVL node to track this.

Reviewed By: kaz7

Differential Revision: https://reviews.llvm.org/D118321

Author: Simon Moll

llvm/include/llvm/CodeGen/TargetLowering.h

@@ -1071,6 +1071,11 @@ class TargetLoweringBase {
     return false;
   }
 
+  /// How to legalize this custom operation?
+  virtual LegalizeAction getCustomOperationAction(SDNode &Op) const {
+    return Legal;
+  }
+
   /// Return how this operation should be treated: either it is legal, needs to
   /// be promoted to a larger size, needs to be expanded to some other code
   /// sequence, or the target has a custom expander for it.

llvm/lib/CodeGen/SelectionDAG/LegalizeDAG.cpp

@@ -1212,7 +1212,7 @@ void SelectionDAGLegalize::LegalizeOp(SDNode *Node) {
     break;
   default:
     if (Node->getOpcode() >= ISD::BUILTIN_OP_END) {
-      Action = TargetLowering::Legal;
+      Action = TLI.getCustomOperationAction(*Node);
     } else {
       Action = TLI.getOperationAction(Node->getOpcode(), Node->getValueType(0));
     }

llvm/lib/Target/VE/CMakeLists.txt

@@ -26,6 +26,7 @@ add_llvm_target(VECodeGen
   VERegisterInfo.cpp
   VESubtarget.cpp
   VETargetMachine.cpp
+  VVPISelLowering.cpp
 
   LINK_COMPONENTS
   Analysis

llvm/lib/Target/VE/VECustomDAG.cpp

@@ -42,6 +42,32 @@ Optional<unsigned> getVVPOpcode(unsigned Opcode) {
   return None;
 }
 
+bool maySafelyIgnoreMask(SDValue Op) {
+  auto VVPOpc = getVVPOpcode(Op->getOpcode());
+  auto Opc = VVPOpc.getValueOr(Op->getOpcode());
+
+  switch (Opc) {
+  case VEISD::VVP_SDIV:
+  case VEISD::VVP_UDIV:
+  case VEISD::VVP_FDIV:
+  case VEISD::VVP_SELECT:
+    return false;
+
+  default:
+    return true;
+  }
+}
+
+bool isVVPOrVEC(unsigned Opcode) {
+  switch (Opcode) {
+  case VEISD::VEC_BROADCAST:
+#define ADD_VVP_OP(VVPNAME, ...) case VEISD::VVPNAME:
+#include "VVPNodes.def"
+    return true;
+  }
+  return false;
+}
+
 bool isVVPBinaryOp(unsigned VVPOpcode) {
   switch (VVPOpcode) {
 #define ADD_BINARY_VVP_OP(VVPNAME, ...)                                        \
@@ -52,6 +78,44 @@ bool isVVPBinaryOp(unsigned VVPOpcode) {
   return false;
 }
 
+// Return the AVL operand position for this VVP or VEC Op.
+Optional<int> getAVLPos(unsigned Opc) {
+  // This is only available for VP SDNodes
+  auto PosOpt = ISD::getVPExplicitVectorLengthIdx(Opc);
+  if (PosOpt)
+    return *PosOpt;
+
+  // VVP Opcodes.
+  if (isVVPBinaryOp(Opc))
+    return 3;
+
+  // VM Opcodes.
+  switch (Opc) {
+  case VEISD::VEC_BROADCAST:
+    return 1;
+  case VEISD::VVP_SELECT:
+    return 3;
+  }
+
+  return None;
+}
+
+bool isLegalAVL(SDValue AVL) { return AVL->getOpcode() == VEISD::LEGALAVL; }
+
+SDValue getNodeAVL(SDValue Op) {
+  auto PosOpt = getAVLPos(Op->getOpcode());
+  return PosOpt ? Op->getOperand(*PosOpt) : SDValue();
+}
+
+std::pair<SDValue, bool> getAnnotatedNodeAVL(SDValue Op) {
+  SDValue AVL = getNodeAVL(Op);
+  if (!AVL)
+    return {SDValue(), true};
+  if (isLegalAVL(AVL))
+    return {AVL->getOperand(0), true};
+  return {AVL, false};
+}
+
 SDValue VECustomDAG::getConstant(uint64_t Val, EVT VT, bool IsTarget,
                                  bool IsOpaque) const {
   return DAG.getConstant(Val, DL, VT, IsTarget, IsOpaque);
@@ -78,4 +142,10 @@ SDValue VECustomDAG::getBroadcast(EVT ResultVT, SDValue Scalar,
   return getNode(VEISD::VEC_BROADCAST, ResultVT, {Scalar, AVL});
 }
 
+SDValue VECustomDAG::annotateLegalAVL(SDValue AVL) const {
+  if (isLegalAVL(AVL))
+    return AVL;
+  return getNode(VEISD::LEGALAVL, AVL.getValueType(), AVL);
+}
+
 } // namespace llvm

llvm/lib/Target/VE/VECustomDAG.h

@@ -27,6 +27,52 @@ bool isVVPBinaryOp(unsigned Opcode);
 
 bool isPackedVectorType(EVT SomeVT);
 
+bool isVVPOrVEC(unsigned);
+
+bool maySafelyIgnoreMask(SDValue Op);
+
+/// The VE backend uses a two-staged process to lower and legalize vector
+/// instructions:
+//
+/// 1. VP and standard vector SDNodes are lowered to SDNodes of the VVP_* layer.
+//
+//     All VVP nodes have a mask and an Active Vector Length (AVL) parameter.
+//     The AVL parameters refers to the element position in the vector the VVP
+//     node operates on.
+//
+//
+//  2. The VVP SDNodes are legalized. The AVL in a legal VVP node refers to
+//     chunks of 64bit. We track this by wrapping the AVL in a LEGALAVL node.
+//
+//     The AVL mechanism in the VE architecture always refers to chunks of
+//     64bit, regardless of the actual element type vector instructions are
+//     operating on. For vector types v256.32 or v256.64 nothing needs to be
+//     legalized since each element occupies a 64bit chunk - there is no
+//     difference between counting 64bit chunks or element positions. However,
+//     all vector types with > 256 elements store more than one logical element
+//     per 64bit chunk and need to be transformed.
+//     However legalization is performed, the resulting legal VVP SDNodes will
+//     have a LEGALAVL node as their AVL operand. The LEGALAVL nodes wraps
+//     around an AVL that refers to 64 bit chunks just as the architecture
+//     demands - that is, the wrapped AVL is the correct setting for the VL
+//     register for this VVP operation to get the desired behavior.
+//
+/// AVL Functions {
+// The AVL operand position of this node.
+Optional<int> getAVLPos(unsigned);
+
+// Whether this is a LEGALAVL node.
+bool isLegalAVL(SDValue AVL);
+
+// The AVL operand of this node.
+SDValue getNodeAVL(SDValue);
+
+// Return the AVL operand of this node. If it is a LEGALAVL node, unwrap it.
+// Return with the boolean whether unwrapping happened.
+std::pair<SDValue, bool> getAnnotatedNodeAVL(SDValue);
+
+/// } AVL Functions
+
 class VECustomDAG {
   SelectionDAG &DAG;
   SDLoc DL;
@@ -72,6 +118,9 @@ class VECustomDAG {
                       bool IsOpaque = false) const;
 
   SDValue getBroadcast(EVT ResultVT, SDValue Scalar, SDValue AVL) const;
+
+  // Wrap AVL in a LEGALAVL node (unless it is one already).
+  SDValue annotateLegalAVL(SDValue AVL) const;
 };
 
 } // namespace llvm

llvm/lib/Target/VE/VEISelDAGToDAG.cpp

@@ -335,6 +335,12 @@ void VEDAGToDAGISel::Select(SDNode *N) {
   }
 
   switch (N->getOpcode()) {
+
+  // Late eliminate the LEGALAVL wrapper
+  case VEISD::LEGALAVL:
+    ReplaceNode(N, N->getOperand(0).getNode());
+    return;
+
   case VEISD::GLOBAL_BASE_REG:
     ReplaceNode(N, getGlobalBaseReg());
     return;

llvm/lib/Target/VE/VEISelLowering.cpp

@@ -902,6 +902,8 @@ const char *VETargetLowering::getTargetNodeName(unsigned Opcode) const {
     TARGET_NODE_CASE(REPL_I32)
     TARGET_NODE_CASE(REPL_F32)
 
+    TARGET_NODE_CASE(LEGALAVL)
+
     // Register the VVP_* SDNodes.
 #define ADD_VVP_OP(VVP_NAME, ...) TARGET_NODE_CASE(VVP_NAME)
 #include "VVPNodes.def"
@@ -1658,18 +1660,24 @@ SDValue VETargetLowering::lowerBUILD_VECTOR(SDValue Op,
   // Else emit a broadcast.
   if (SDValue ScalarV = getSplatValue(Op.getNode())) {
     unsigned NumEls = ResultVT.getVectorNumElements();
-    // TODO: Legalize packed-mode AVL.
-    //       For now, cap the AVL at 256.
-    auto CappedLength = std::min<unsigned>(256, NumEls);
-    auto AVL = CDAG.getConstant(CappedLength, MVT::i32);
+    auto AVL = CDAG.getConstant(NumEls, MVT::i32);
     return CDAG.getBroadcast(ResultVT, Op.getOperand(0), AVL);
   }
 
   // Expand
   return SDValue();
 }
 
+TargetLowering::LegalizeAction
+VETargetLowering::getCustomOperationAction(SDNode &Op) const {
+  // Custom lower to legalize AVL for packed mode.
+  if (isVVPOrVEC(Op.getOpcode()))
+    return Custom;
+  return Legal;
+}
+
 SDValue VETargetLowering::LowerOperation(SDValue Op, SelectionDAG &DAG) const {
+  LLVM_DEBUG(dbgs() << "::LowerOperation"; Op->print(dbgs()););
   unsigned Opcode = Op.getOpcode();
   if (ISD::isVPOpcode(Opcode))
     return lowerToVVP(Op, DAG);
@@ -1721,6 +1729,16 @@ SDValue VETargetLowering::LowerOperation(SDValue Op, SelectionDAG &DAG) const {
   case ISD::EXTRACT_VECTOR_ELT:
     return lowerEXTRACT_VECTOR_ELT(Op, DAG);
 
+  // Legalize the AVL of this internal node.
+  case VEISD::VEC_BROADCAST:
+#define ADD_VVP_OP(VVP_NAME, ...) case VEISD::VVP_NAME:
+#include "VVPNodes.def"
+    // AVL already legalized.
+    if (getAnnotatedNodeAVL(Op).second)
+      return Op;
+    return legalizeInternalVectorOp(Op, DAG);
+
+    // Translate into a VEC_*/VVP_* layer operation.
 #define ADD_VVP_OP(VVP_NAME, ISD_NAME) case ISD::ISD_NAME:
 #include "VVPNodes.def"
     return lowerToVVP(Op, DAG);

llvm/lib/Target/VE/VEISelLowering.h

@@ -43,12 +43,19 @@ enum NodeType : unsigned {
   REPL_I32,
   REPL_F32, // Replicate subregister to other half.
 
+  // Annotation as a wrapper. LEGALAVL(VL) means that VL refers to 64bit of
+  // data, whereas the raw EVL coming in from VP nodes always refers to number
+  // of elements, regardless of their size.
+  LEGALAVL,
+
 // VVP_* nodes.
 #define ADD_VVP_OP(VVP_NAME, ...) VVP_NAME,
 #include "VVPNodes.def"
 };
 }
 
+class VECustomDAG;
+
 class VETargetLowering : public TargetLowering {
   const VESubtarget *Subtarget;
 
@@ -105,6 +112,9 @@ class VETargetLowering : public TargetLowering {
   }
 
   /// Custom Lower {
+  TargetLoweringBase::LegalizeAction
+  getCustomOperationAction(SDNode &) const override;
+
   SDValue LowerOperation(SDValue Op, SelectionDAG &DAG) const override;
   unsigned getJumpTableEncoding() const override;
   const MCExpr *LowerCustomJumpTableEntry(const MachineJumpTableInfo *MJTI,
@@ -170,6 +180,8 @@ class VETargetLowering : public TargetLowering {
 
   /// VVP Lowering {
   SDValue lowerToVVP(SDValue Op, SelectionDAG &DAG) const;
+  SDValue legalizeInternalVectorOp(SDValue Op, SelectionDAG &DAG) const;
+  SDValue legalizePackedAVL(SDValue Op, VECustomDAG &CDAG) const;
   /// } VVPLowering
 
   /// Custom DAGCombine {

llvm/lib/Target/VE/VVPISelLowering.cpp

@@ -0,0 +1,74 @@
+//===-- VVPISelLowering.cpp - VE DAG Lowering Implementation --------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This file implements the lowering and legalization of vector instructions to
+// VVP_*layer SDNodes.
+//
+//===----------------------------------------------------------------------===//
+
+#include "VECustomDAG.h"
+#include "VEISelLowering.h"
+
+using namespace llvm;
+
+#define DEBUG_TYPE "ve-lower"
+
+SDValue VETargetLowering::legalizeInternalVectorOp(SDValue Op,
+                                                   SelectionDAG &DAG) const {
+  VECustomDAG CDAG(DAG, Op);
+  // TODO: Implement odd/even splitting.
+  return legalizePackedAVL(Op, CDAG);
+}
+
+SDValue VETargetLowering::legalizePackedAVL(SDValue Op,
+                                            VECustomDAG &CDAG) const {
+  LLVM_DEBUG(dbgs() << "::legalizePackedAVL\n";);
+  // Only required for VEC and VVP ops.
+  if (!isVVPOrVEC(Op->getOpcode()))
+    return Op;
+
+  // Operation already has a legal AVL.
+  auto AVL = getNodeAVL(Op);
+  if (isLegalAVL(AVL))
+    return Op;
+
+  // Half and round up EVL for 32bit element types.
+  SDValue LegalAVL = AVL;
+  if (isPackedVectorType(Op.getValueType())) {
+    assert(maySafelyIgnoreMask(Op) &&
+           "TODO Shift predication from EVL into Mask");
+
+    if (auto *ConstAVL = dyn_cast<ConstantSDNode>(AVL)) {
+      LegalAVL = CDAG.getConstant((ConstAVL->getZExtValue() + 1) / 2, MVT::i32);
+    } else {
+      auto ConstOne = CDAG.getConstant(1, MVT::i32);
+      auto PlusOne = CDAG.getNode(ISD::ADD, MVT::i32, {AVL, ConstOne});
+      LegalAVL = CDAG.getNode(ISD::SRL, MVT::i32, {PlusOne, ConstOne});
+    }
+  }
+
+  SDValue AnnotatedLegalAVL = CDAG.annotateLegalAVL(LegalAVL);
+
+  // Copy the operand list.
+  int NumOp = Op->getNumOperands();
+  auto AVLPos = getAVLPos(Op->getOpcode());
+  std::vector<SDValue> FixedOperands;
+  for (int i = 0; i < NumOp; ++i) {
+    if (AVLPos && (i == *AVLPos)) {
+      FixedOperands.push_back(AnnotatedLegalAVL);
+      continue;
+    }
+    FixedOperands.push_back(Op->getOperand(i));
+  }
+
+  // Clone the operation with fixed operands.
+  auto Flags = Op->getFlags();
+  SDValue NewN =
+      CDAG.getNode(Op->getOpcode(), Op->getVTList(), FixedOperands, Flags);
+  return NewN;
+}