[AMDGPU] selecting v_sat_pk instruction, version 2

llvm/lib/Target/AMDGPU/AMDGPUISelLowering.cpp

@@ -5518,6 +5518,7 @@ const char* AMDGPUTargetLowering::getTargetNodeName(unsigned Opcode) const {
   NODE_NAME_CASE(UMIN3)
   NODE_NAME_CASE(FMED3)
   NODE_NAME_CASE(SMED3)
+  NODE_NAME_CASE(SAT_PK_CAST)
   NODE_NAME_CASE(UMED3)
   NODE_NAME_CASE(FMAXIMUM3)
   NODE_NAME_CASE(FMINIMUM3)

llvm/lib/Target/AMDGPU/AMDGPUISelLowering.h

@@ -461,6 +461,7 @@ enum NodeType : unsigned {
   FMED3,
   SMED3,
   UMED3,
+  SAT_PK_CAST,
   FMAXIMUM3,
   FMINIMUM3,
   FDOT2,

llvm/lib/Target/AMDGPU/AMDGPUInstrInfo.td

@@ -332,6 +332,9 @@ def AMDGPUumed3 : SDNode<"AMDGPUISD::UMED3", AMDGPUDTIntTernaryOp,
   []
 >;
 
+// Special node to pack v2i8 into i16 for v_sat_pk lowering.
+def AMDGPUsat_pk_cast : SDNode<"AMDGPUISD::SAT_PK_CAST", SDTUnaryOp, []>;
+
 def AMDGPUfmed3_impl : SDNode<"AMDGPUISD::FMED3", SDTFPTernaryOp, []>;
 
 def AMDGPUfdot2_impl : SDNode<"AMDGPUISD::FDOT2",

llvm/lib/Target/AMDGPU/SIISelLowering.cpp

@@ -824,6 +824,25 @@ SITargetLowering::SITargetLowering(const TargetMachine &TM,
                          {MVT::v4f32, MVT::v8f32, MVT::v16f32, MVT::v32f32},
                          Custom);
     }
+
+    // True 16 instruction is current not supported
+    // FIXME: Add support for true 16 when supported
+    if (!Subtarget->hasTrue16BitInsts() || !Subtarget->useRealTrue16Insts()) {
+      // MVT::vNi16 for src type check in foldToSaturated
+      // MVT::vNi8 for dst type check in CustomLowerNode
+      // FIXME: Handle N = 2, 4, 8 first, should change verification logic from
+      //        LLVM side, like break bigger vector into legal small vectors
+      setOperationAction(ISD::TRUNCATE_SSAT_U,
+                         {
+                             MVT::v2i16,
+                             MVT::v4i16,
+                             MVT::v8i16,
+                             MVT::v2i8,
+                             MVT::v4i8,
+                             MVT::v8i8,
+                         },
+                         Custom);
+    }
   }
 
   setOperationAction({ISD::FNEG, ISD::FABS}, MVT::v4f16, Custom);
@@ -1983,6 +2002,12 @@ bool SITargetLowering::isTypeDesirableForOp(unsigned Op, EVT VT) const {
   if (VT == MVT::i1 && Op == ISD::SETCC)
     return false;
 
+  // Special case for vNi8 handling where N is even
+  if (Op == ISD::TRUNCATE_SSAT_U && VT.isVector() &&
+      VT.getVectorElementType() == MVT::i8 &&
+      ((VT.getVectorNumElements() & 1) == 0))
+    return true;
+
   return TargetLowering::isTypeDesirableForOp(Op, VT);
 }
 
@@ -6615,6 +6640,45 @@ void SITargetLowering::ReplaceNodeResults(SDNode *N,
     Results.push_back(lowerFSQRTF16(SDValue(N, 0), DAG));
     break;
   }
+  case ISD::TRUNCATE_SSAT_U: {
+    SDLoc SL(N);
+    SDValue Src = N->getOperand(0);
+    EVT SrcVT = Src.getValueType();
+    EVT DstVT = N->getValueType(0);
+
+    assert(SrcVT.isVector() && DstVT.isVector());
+    assert(DstVT.getVectorElementType() == MVT::i8);
+    assert(SrcVT.getVectorElementType() == MVT::i16);
+
+    unsigned EleNo = SrcVT.getVectorNumElements();
+    assert(EleNo == DstVT.getVectorNumElements());
+
+    if (EleNo == 2) {
+      SDValue Op = DAG.getNode(AMDGPUISD::SAT_PK_CAST, SL, MVT::i16, Src);
+      Op = DAG.getNode(ISD::BITCAST, SL, N->getValueType(0), Op);
+      Results.push_back(Op);
+      break;
+    }
+
+    // Vector case, number of element must be even
+    assert((EleNo & 1) == 0);
+    SmallVector<SDValue> DstPairs;
+    EVT SrcEleVT = SrcVT.getVectorElementType();
+    EVT DstEleVT = DstVT.getVectorElementType();
+    EVT SrcPairVT = EVT::getVectorVT(*DAG.getContext(), SrcEleVT, 2);
+    EVT DstPairVT = EVT::getVectorVT(*DAG.getContext(), DstEleVT, 2);
+    for (unsigned i = 0; i != EleNo; i += 2) {
+      SDValue SrcPair = DAG.getNode(ISD::EXTRACT_SUBVECTOR, SL, SrcPairVT, Src,
+                                    DAG.getConstant(i, SL, MVT::i32));
+      SDValue SatPk =
+          DAG.getNode(AMDGPUISD::SAT_PK_CAST, SL, MVT::i16, SrcPair);
+      SDValue DstPair = DAG.getNode(ISD::BITCAST, SL, DstPairVT, SatPk);
+      DstPairs.push_back(DstPair);
+    }
+    SDValue Op = DAG.getNode(ISD::CONCAT_VECTORS, SL, DstVT, DstPairs);
+    Results.push_back(Op);
+    break;
+  }
   default:
     AMDGPUTargetLowering::ReplaceNodeResults(N, Results, DAG);
     break;

llvm/lib/Target/AMDGPU/SIInstructions.td

@@ -3382,6 +3382,21 @@ def : GCNPat <
   (v2i16 (V_LSHL_OR_B32_e64 $src1, (i32 16), (i32 (V_AND_B32_e64 (i32 (V_MOV_B32_e32 (i32 0xffff))), $src0))))
 >;
 
+multiclass V_SAT_PK_Pat<Instruction inst> {
+  def : GCNPat<
+    (i16 (AMDGPUsat_pk_cast v2i16:$src)),
+    (inst VRegSrc_32:$src)
+  >;
+}
+
+let OtherPredicates = [NotHasTrue16BitInsts] in {
+  defm : V_SAT_PK_Pat<V_SAT_PK_U8_I16_e64>;
+} // End OtherPredicates = [NotHasTrue16BitInsts]
+
+let True16Predicate = UseFakeTrue16Insts in {
+  defm : V_SAT_PK_Pat<V_SAT_PK_U8_I16_fake16_e64>;
+} // End True16Predicate = UseFakeTrue16Insts
+
 // With multiple uses of the shift, this will duplicate the shift and
 // increase register pressure.
 def : GCNPat <