[AMDGPU] Improved Lowering of abs(i8/i16) and -abs(i8/i16)

llvm/include/llvm/CodeGen/TargetLowering.h

@@ -5576,8 +5576,8 @@ class LLVM_ABI TargetLowering : public TargetLoweringBase {
   /// \param N Node to expand
   /// \param IsNegative indicate negated abs
   /// \returns The expansion result or SDValue() if it fails.
-  SDValue expandABS(SDNode *N, SelectionDAG &DAG,
-                    bool IsNegative = false) const;
+  virtual SDValue expandABS(SDNode *N, SelectionDAG &DAG,
+                            bool IsNegative = false) const;
 
   /// Expand ABDS/ABDU nodes. Expands vector/scalar ABDS/ABDU nodes.
   /// \param N Node to expand

llvm/lib/Target/AMDGPU/AMDGPUISelLowering.cpp

@@ -5287,6 +5287,27 @@ SDValue AMDGPUTargetLowering::performRcpCombine(SDNode *N,
   return DCI.DAG.getConstantFP(One / Val, SDLoc(N), N->getValueType(0));
 }
 
+SDValue AMDGPUTargetLowering::expandABS(SDNode *N, SelectionDAG &DAG,
+                                        bool IsNegative) const {
+  if (N->isDivergent() ||
+      (N->getValueType(0) != MVT::i8 && N->getValueType(0) != MVT::i16))
+    return TargetLowering::expandABS(N, DAG, IsNegative);
+
+  //(abs i8/i16 (i8/i16 op1)) -> (trunc i8/i16 (abs i32 (sext i32 (i8/i16
+  // op1))))
+  SDValue Src = N->getOperand(0);
+  SDLoc DL(Src);
+  SDValue SExtSrc = DAG.getSExtOrTrunc(Src, DL, MVT::i32);
+  SDValue ExtAbs = DAG.getNode(ISD::ABS, DL, MVT::i32, SExtSrc);
+  SDValue TruncResult =
+      DAG.getNode(ISD::TRUNCATE, DL, N->getValueType(0), ExtAbs);
+
+  if (!IsNegative)
+    return TruncResult;
+
+  return DAG.getNegative(TruncResult, DL, N->getValueType(0));
+}
+
 SDValue AMDGPUTargetLowering::PerformDAGCombine(SDNode *N,
                                                 DAGCombinerInfo &DCI) const {
   SelectionDAG &DAG = DCI.DAG;

llvm/lib/Target/AMDGPU/AMDGPUISelLowering.h

@@ -45,6 +45,9 @@ class AMDGPUTargetLowering : public TargetLowering {
   /// original size will not change the value.
   static unsigned numBitsSigned(SDValue Op, SelectionDAG &DAG);
 
+  virtual SDValue expandABS(SDNode *N, SelectionDAG &DAG,
+                            bool IsNegative = false) const override;
+
 protected:
   SDValue LowerEXTRACT_SUBVECTOR(SDValue Op, SelectionDAG &DAG) const;
   SDValue LowerCONCAT_VECTORS(SDValue Op, SelectionDAG &DAG) const;