[AMDGPU] selecting v_sat_pk instruction, version 2

llvm/lib/Target/AMDGPU/AMDGPUISelLowering.cpp

@@ -5498,6 +5498,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 handle v_sat_pk to avoid v2i8
+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

@@ -816,6 +816,35 @@ 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
+      setOperationAction(ISD::TRUNCATE_SSAT_U,
+                         {
+                             MVT::v2i16,
+                             MVT::v4i16,
+                             MVT::v8i16,
+                             MVT::v16i16,
+                             MVT::v32i16,
+                             MVT::v64i16,
+                             MVT::v128i16,
+                             MVT::v256i16,
+                             MVT::v512i16,
+                             MVT::v2i8,
+                             MVT::v4i8,
+                             MVT::v8i8,
+                             MVT::v16i8,
+                             MVT::v32i8,
+                             MVT::v64i8,
+                             MVT::v128i8,
+                             MVT::v256i8,
+                             MVT::v512i8,
+                         },
+                         Custom);
+    }
   }
 
   setOperationAction({ISD::FNEG, ISD::FABS}, MVT::v4f16, Custom);
@@ -1975,6 +2004,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);
 }
 
@@ -6606,6 +6641,43 @@ 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());
+
+    unsigned EleNo = SrcVT.getVectorNumElements();
+    assert(EleNo == DstVT.getVectorNumElements());
+
+    if (EleNo == 2) {
+      SDValue Op =
+          DAG.getNode(AMDGPUISD::SAT_PK_CAST, SL, MVT::i16, N->getOperand(0));
+      Op = DAG.getNode(ISD::BITCAST, SL, N->getValueType(0), Op);
+      Results.push_back(Op);
+    } else {
+      // Must be even number
+      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 + 1 < EleNo; i = 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

@@ -3324,6 +3324,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 <