[AMDGPU] Add type-generic llvm.amdgcn.readfirstlane2 intrinsic

llvm/include/llvm/IR/IntrinsicsAMDGPU.td

@@ -2206,11 +2206,18 @@ class AMDGPUWaveReduce<LLVMType data_ty = llvm_anyint_ty> : Intrinsic<
 def int_amdgcn_wave_reduce_umin : AMDGPUWaveReduce;
 def int_amdgcn_wave_reduce_umax : AMDGPUWaveReduce;
 
+// i32 llvm.amdgcn.readfirstlane(i32)
 def int_amdgcn_readfirstlane :
   ClangBuiltin<"__builtin_amdgcn_readfirstlane">,
   Intrinsic<[llvm_i32_ty], [llvm_i32_ty],
             [IntrNoMem, IntrConvergent, IntrWillReturn, IntrNoCallback, IntrNoFree]>;
 
+// ty llvm.amdgcn.readfirstlane2(ty)
+// A type-generic version of readfirstlane.
+def int_amdgcn_readfirstlane2 :
+  Intrinsic<[llvm_any_ty], [LLVMMatchType<0>],
+            [IntrNoMem, IntrConvergent, IntrWillReturn, IntrNoCallback, IntrNoFree]>;
+
 // The lane argument must be uniform across the currently active threads of the
 // current wave. Otherwise, the result is undefined.
 def int_amdgcn_readlane :

llvm/lib/Target/AMDGPU/AMDGPUISelLowering.cpp

@@ -5453,6 +5453,7 @@ const char* AMDGPUTargetLowering::getTargetNodeName(unsigned Opcode) const {
   NODE_NAME_CASE(LDS)
   NODE_NAME_CASE(FPTRUNC_ROUND_UPWARD)
   NODE_NAME_CASE(FPTRUNC_ROUND_DOWNWARD)
+  NODE_NAME_CASE(READFIRSTLANE)
   NODE_NAME_CASE(DUMMY_CHAIN)
   case AMDGPUISD::FIRST_MEM_OPCODE_NUMBER: break;
   NODE_NAME_CASE(LOAD_D16_HI)

llvm/lib/Target/AMDGPU/AMDGPUISelLowering.h

@@ -541,6 +541,8 @@ enum NodeType : unsigned {
   FPTRUNC_ROUND_UPWARD,
   FPTRUNC_ROUND_DOWNWARD,
 
+  READFIRSTLANE,
+
   DUMMY_CHAIN,
   FIRST_MEM_OPCODE_NUMBER = ISD::FIRST_TARGET_MEMORY_OPCODE,
   LOAD_D16_HI,

llvm/lib/Target/AMDGPU/AMDGPUInstrInfo.td

@@ -342,6 +342,8 @@ def AMDGPUfdot2_impl : SDNode<"AMDGPUISD::FDOT2",
 
 def AMDGPUperm_impl : SDNode<"AMDGPUISD::PERM", AMDGPUDTIntTernaryOp, []>;
 
+def AMDGPUreadfirstlane_impl : SDNode<"AMDGPUISD::READFIRSTLANE", SDTIntUnaryOp>;
+
 // SI+ export
 def AMDGPUExportOp : SDTypeProfile<0, 8, [
   SDTCisInt<0>,       // i8 tgt
@@ -504,3 +506,7 @@ def AMDGPUdiv_fmas : PatFrags<(ops node:$src0, node:$src1, node:$src2, node:$vcc
 def AMDGPUperm : PatFrags<(ops node:$src0, node:$src1, node:$src2),
   [(int_amdgcn_perm node:$src0, node:$src1, node:$src2),
    (AMDGPUperm_impl node:$src0, node:$src1, node:$src2)]>;
+
+def AMDGPUreadfirstlane : PatFrags<(ops node:$src),
+  [(int_amdgcn_readfirstlane node:$src),
+   (AMDGPUreadfirstlane_impl node:$src)]>;

llvm/lib/Target/AMDGPU/AMDGPUSearchableTables.td

@@ -393,6 +393,7 @@ def UniformIntrinsics : GenericTable {
 }
 
 def : AlwaysUniform<int_amdgcn_readfirstlane>;
+def : AlwaysUniform<int_amdgcn_readfirstlane2>;
 def : AlwaysUniform<int_amdgcn_readlane>;
 def : AlwaysUniform<int_amdgcn_icmp>;
 def : AlwaysUniform<int_amdgcn_fcmp>;

llvm/lib/Target/AMDGPU/SIISelLowering.cpp

@@ -238,7 +238,6 @@ SITargetLowering::SITargetLowering(const TargetMachine &TM,
     setOperationAction(ISD::BUILD_VECTOR, MVT::v2bf16, Promote);
     AddPromotedToType(ISD::BUILD_VECTOR, MVT::v2bf16, MVT::v2i16);
   }
-
   setTruncStoreAction(MVT::v2i32, MVT::v2i16, Expand);
   setTruncStoreAction(MVT::v3i32, MVT::v3i16, Expand);
   setTruncStoreAction(MVT::v4i32, MVT::v4i16, Expand);
@@ -8452,6 +8451,26 @@ SDValue SITargetLowering::LowerINTRINSIC_WO_CHAIN(SDValue Op,
   }
   case Intrinsic::amdgcn_addrspacecast_nonnull:
     return lowerADDRSPACECAST(Op, DAG);
+  case Intrinsic::amdgcn_readfirstlane2:
+    if (VT.getSizeInBits() <= 32) {
+      MVT IntVT = MVT::getIntegerVT(VT.getSizeInBits());
+      return DAG.getBitcast(
+          VT, DAG.getAnyExtOrTrunc(
+                  DAG.getNode(AMDGPUISD::READFIRSTLANE, DL, MVT::i32,
+                              DAG.getAnyExtOrTrunc(
+                                  DAG.getBitcast(IntVT, Op.getOperand(1)), DL,
+                                  MVT::i32)),
+                  DL, IntVT));
+    }
+    if (VT.getSizeInBits() % 32 == 0) {
+      MVT VecVT = MVT::getVectorVT(MVT::i32, VT.getSizeInBits() / 32);
+      return DAG.getBitcast(
+          VT, DAG.UnrollVectorOp(
+                  DAG.getNode(AMDGPUISD::READFIRSTLANE, DL, VecVT,
+                              DAG.getBitcast(VecVT, Op.getOperand(1)))
+                      .getNode()));
+    }
+    return SDValue();
   default:
     if (const AMDGPU::ImageDimIntrinsicInfo *ImageDimIntr =
             AMDGPU::getImageDimIntrinsicInfo(IntrinsicID))

llvm/lib/Target/AMDGPU/SIInstructions.td

@@ -3405,7 +3405,7 @@ def : GCNPat<
 // FIXME: Should also do this for readlane, but tablegen crashes on
 // the ignored src1.
 def : GCNPat<
-  (int_amdgcn_readfirstlane (i32 imm:$src)),
+  (AMDGPUreadfirstlane (i32 imm:$src)),
   (S_MOV_B32 SReg_32:$src)
 >;
 

llvm/lib/Target/AMDGPU/VOP1Instructions.td

@@ -243,7 +243,7 @@ def VOP_READFIRSTLANE : VOPProfile <[i32, i32, untyped, untyped]> {
 // FIXME: Specify SchedRW for READFIRSTLANE_B32
 // TODO: There is VOP3 encoding also
 def V_READFIRSTLANE_B32 : VOP1_Pseudo <"v_readfirstlane_b32", VOP_READFIRSTLANE,
-                                       getVOP1Pat<int_amdgcn_readfirstlane,
+                                       getVOP1Pat<AMDGPUreadfirstlane,
                                                   VOP_READFIRSTLANE>.ret, 1> {
   let isConvergent = 1;
 }