PR feedback, implement in GISel

Author: saxlungs

llvm/include/llvm/IR/IntrinsicsAMDGPU.td

@@ -2719,13 +2719,15 @@ def int_amdgcn_call_whole_wave:
             [IntrConvergent]>;
 
 // <result>
-//   llvm.amdgcn.subgroup.shuffle <value> <id>
-// value and result can be any scalar of floating-point, integer,
-// or Boolean types, but must be the same type
-def int_amdgcn_subgroup_shuffle :
-  Intrinsic<[llvm_any_ty],                                                            // return types
-            [LLVMMatchType<0>, llvm_i32_ty],                                          // arg types
-            [IntrConvergent, IntrNoMem, IntrNoFree, IntrWillReturn, IntrNoCallback]>; // flags
+//   llvm.amdgcn.wave.shuffle <value> <id>
+// value and result can be 32bit floating-point, integer,
+// or Boolean types, and must be the same type. Any index
+// value that's outside the valid range will wrap around,
+// and reading from an inactive lane will return 0.
+def int_amdgcn_wave_shuffle :
+  DefaultAttrsIntrinsic<[llvm_any_ty],                                                            // return types
+                        [LLVMMatchType<0>, llvm_i32_ty],                                          // arg types
+                        [IntrConvergent, IntrNoMem, IntrNoFree, IntrWillReturn, IntrNoCallback]>; // flags
 
 //===----------------------------------------------------------------------===//
 // CI+ Intrinsics

llvm/lib/Target/AMDGPU/AMDGPUInstructionSelector.cpp

@@ -1216,6 +1216,8 @@ bool AMDGPUInstructionSelector::selectG_INTRINSIC(MachineInstr &I) const {
   case Intrinsic::amdgcn_permlane16_swap:
   case Intrinsic::amdgcn_permlane32_swap:
     return selectPermlaneSwapIntrin(I, IntrinsicID);
+  case Intrinsic::amdgcn_wave_shuffle:
+    return selectWaveShuffleIntrin(I);
   default:
     return selectImpl(I, *CoverageInfo);
   }
@@ -3852,6 +3854,129 @@ bool AMDGPUInstructionSelector::selectWaveAddress(MachineInstr &MI) const {
   return true;
 }
 
+bool AMDGPUInstructionSelector::selectWaveShuffleIntrin(
+    MachineInstr &MI) const {
+  assert(MI.getNumOperands() == 4);
+  MachineBasicBlock *MBB = MI.getParent();
+  const DebugLoc &DL = MI.getDebugLoc();
+
+  Register DstReg = MI.getOperand(0).getReg();
+  Register ValReg = MI.getOperand(2).getReg();
+  Register IdxReg = MI.getOperand(3).getReg();
+
+  const LLT DstTy = MRI->getType(DstReg);
+  unsigned DstSize = DstTy.getSizeInBits();
+  const RegisterBank *DstRB = RBI.getRegBank(DstReg, *MRI, TRI);
+  const TargetRegisterClass *DstRC =
+      TRI.getRegClassForSizeOnBank(DstSize, *DstRB);
+
+  assert(DstTy == LLT::scalar(32));
+
+  // If we can bpermute across the whole wave, then just do that
+  if (Subtarget->supportsWaveWideBPermute()) {
+    Register ShiftIdxReg = MRI->createVirtualRegister(DstRC);
+    BuildMI(*MBB, MI, DL, TII.get(AMDGPU::V_LSHLREV_B32_e64), ShiftIdxReg)
+        .addImm(2)
+        .addReg(IdxReg);
+
+    BuildMI(*MBB, MI, DL, TII.get(AMDGPU::DS_BPERMUTE_B32), DstReg)
+        .addReg(ShiftIdxReg)
+        .addReg(ValReg)
+        .addImm(0);
+  } else {
+    // Otherwise, we need to make use of whole wave mode
+    assert(Subtarget->isWave64());
+
+    // Set inactive lanes to poison
+    Register UndefValReg =
+        MRI->createVirtualRegister(TRI.getRegClass(AMDGPU::SReg_32RegClassID));
+    BuildMI(*MBB, MI, DL, TII.get(AMDGPU::IMPLICIT_DEF), UndefValReg);
+
+    Register UndefExecReg = MRI->createVirtualRegister(
+        TRI.getRegClass(AMDGPU::SReg_64_XEXECRegClassID));
+    BuildMI(*MBB, MI, DL, TII.get(AMDGPU::IMPLICIT_DEF), UndefExecReg);
+
+    Register PoisonValReg = MRI->createVirtualRegister(DstRC);
+    BuildMI(*MBB, MI, DL, TII.get(AMDGPU::V_SET_INACTIVE_B32), PoisonValReg)
+        .addImm(0)
+        .addReg(ValReg)
+        .addImm(0)
+        .addReg(UndefValReg)
+        .addReg(UndefExecReg);
+
+    // ds_bpermute requires index to be multiplied by 4
+    Register ShiftIdxReg = MRI->createVirtualRegister(DstRC);
+    BuildMI(*MBB, MI, DL, TII.get(AMDGPU::V_LSHLREV_B32_e64), ShiftIdxReg)
+        .addImm(2)
+        .addReg(IdxReg);
+
+    Register PoisonIdxReg = MRI->createVirtualRegister(DstRC);
+    BuildMI(*MBB, MI, DL, TII.get(AMDGPU::V_SET_INACTIVE_B32), PoisonIdxReg)
+        .addImm(0)
+        .addReg(ShiftIdxReg)
+        .addImm(0)
+        .addReg(UndefValReg)
+        .addReg(UndefExecReg);
+
+    // Get permutation of each half, then we'll select which one to use
+    Register SameSidePermReg = MRI->createVirtualRegister(DstRC);
+    BuildMI(*MBB, MI, DL, TII.get(AMDGPU::DS_BPERMUTE_B32), SameSidePermReg)
+        .addReg(PoisonIdxReg)
+        .addReg(PoisonValReg)
+        .addImm(0);
+
+    Register SwappedValReg = MRI->createVirtualRegister(DstRC);
+    BuildMI(*MBB, MI, DL, TII.get(AMDGPU::V_PERMLANE64_B32), SwappedValReg)
+        .addReg(PoisonValReg);
+
+    Register OppSidePermReg = MRI->createVirtualRegister(DstRC);
+    BuildMI(*MBB, MI, DL, TII.get(AMDGPU::DS_BPERMUTE_B32), OppSidePermReg)
+        .addReg(PoisonIdxReg)
+        .addReg(SwappedValReg)
+        .addImm(0);
+
+    Register WWMSwapPermReg = MRI->createVirtualRegister(DstRC);
+    BuildMI(*MBB, MI, DL, TII.get(AMDGPU::STRICT_WWM), WWMSwapPermReg)
+        .addReg(OppSidePermReg);
+
+    // Select which side to take the permute from
+    // We can get away with only using mbcnt_lo here since we're only
+    // trying to detect which side of 32 each lane is on, and mbcnt_lo
+    // returns 32 for lanes 32-63.
+    Register ThreadIDReg = MRI->createVirtualRegister(DstRC);
+    BuildMI(*MBB, MI, DL, TII.get(AMDGPU::V_MBCNT_LO_U32_B32_e64), ThreadIDReg)
+        .addImm(-1)
+        .addImm(0);
+
+    Register XORReg = MRI->createVirtualRegister(DstRC);
+    BuildMI(*MBB, MI, DL, TII.get(AMDGPU::V_XOR_B32_e64), XORReg)
+        .addReg(ThreadIDReg)
+        .addReg(PoisonIdxReg);
+
+    Register ANDReg = MRI->createVirtualRegister(DstRC);
+    BuildMI(*MBB, MI, DL, TII.get(AMDGPU::V_AND_B32_e64), ANDReg)
+        .addReg(XORReg)
+        .addImm(32);
+
+    Register CompareReg = MRI->createVirtualRegister(
+        TRI.getRegClass(AMDGPU::SReg_64_XEXECRegClassID));
+    BuildMI(*MBB, MI, DL, TII.get(AMDGPU::V_CMP_EQ_U32_e64), CompareReg)
+        .addReg(ANDReg)
+        .addImm(0);
+
+    // Finally do the selection
+    BuildMI(*MBB, MI, DL, TII.get(AMDGPU::V_CNDMASK_B32_e64), DstReg)
+        .addImm(0)
+        .addReg(WWMSwapPermReg)
+        .addImm(0)
+        .addReg(SameSidePermReg)
+        .addReg(CompareReg);
+  }
+
+  MI.eraseFromParent();
+  return true;
+}
+
 // Match BITOP3 operation and return a number of matched instructions plus
 // truth table.
 static std::pair<unsigned, uint8_t> BitOp3_Op(Register R,

llvm/lib/Target/AMDGPU/AMDGPUInstructionSelector.h

@@ -156,6 +156,7 @@ class AMDGPUInstructionSelector final : public InstructionSelector {
   bool selectSBarrierSignalIsfirst(MachineInstr &I, Intrinsic::ID IID) const;
   bool selectSGetBarrierState(MachineInstr &I, Intrinsic::ID IID) const;
   bool selectSBarrierLeave(MachineInstr &I) const;
+  bool selectWaveShuffleIntrin(MachineInstr &I) const;
 
   std::pair<Register, unsigned> selectVOP3ModsImpl(Register Src,
                                                    bool IsCanonicalizing = true,

llvm/lib/Target/AMDGPU/AMDGPURegisterBankInfo.cpp

@@ -5230,11 +5230,20 @@ AMDGPURegisterBankInfo::getInstrMapping(const MachineInstr &MI) const {
       OpdsMapping[2] = AMDGPU::getValueMapping(regBankID, OpSize);
       break;
     }
-    case Intrinsic::amdgcn_s_bitreplicate:
+    case Intrinsic::amdgcn_s_bitreplicate: {
       Register MaskReg = MI.getOperand(2).getReg();
       unsigned MaskBank = getRegBankID(MaskReg, MRI, AMDGPU::SGPRRegBankID);
       OpdsMapping[0] = AMDGPU::getValueMapping(AMDGPU::SGPRRegBankID, 64);
       OpdsMapping[2] = AMDGPU::getValueMapping(MaskBank, 32);
+      break;
+    }
+    case Intrinsic::amdgcn_wave_shuffle: {
+      unsigned OpSize = MRI.getType(MI.getOperand(0).getReg()).getSizeInBits();
+      OpdsMapping[0] = AMDGPU::getValueMapping(AMDGPU::VGPRRegBankID, OpSize);
+      OpdsMapping[2] = AMDGPU::getValueMapping(AMDGPU::VGPRRegBankID, OpSize);
+      OpdsMapping[3] = AMDGPU::getValueMapping(AMDGPU::VGPRRegBankID, OpSize);
+      break;
+    }
     }
     break;
   }

llvm/lib/Target/AMDGPU/GCNSubtarget.h

@@ -1896,7 +1896,7 @@ class GCNSubtarget final : public AMDGPUGenSubtargetInfo,
   }
 
   bool supportsWaveWideBPermute() const {
-    return ((getGeneration() == AMDGPUSubtarget::GFX12) || isWave32());
+    return getGeneration() == AMDGPUSubtarget::GFX12 || isWave32();
   }
 };
 

llvm/lib/Target/AMDGPU/SIISelLowering.cpp

@@ -7280,24 +7280,25 @@ static SDValue lowerLaneOp(const SITargetLowering &TLI, SDNode *N,
   return DAG.getBitcast(VT, UnrolledLaneOp);
 }
 
-static SDValue lowerSubgroupShuffle(const SITargetLowering &TLI, SDNode *N,
+static SDValue lowerWaveShuffle(const SITargetLowering &TLI, SDNode *N,
                                     SelectionDAG &DAG) {
   EVT VT = N->getValueType(0);
   unsigned ValSize = VT.getSizeInBits();
+  assert(ValSize == 32);
   SDLoc SL(N);
 
   SDValue Value = N->getOperand(1);
   SDValue Index = N->getOperand(2);
 
   // ds_bpermute requires index to be multiplied by 4
-  SDValue ShiftAmount = DAG.getTargetConstant(2, SL, MVT::i32);
+  SDValue ShiftAmount = DAG.getShiftAmountConstant(2, MVT::i32, SL);
   SDValue ShiftedIndex = DAG.getNode(ISD::SHL, SL, Index.getValueType(), Index,
                                    ShiftAmount);
 
   // Intrinsics will require i32 to operate on
-  SDValue Value32 = Value;
-  if ((ValSize != 32) || (VT.isFloatingPoint()))
-    Value32 = DAG.getBitcast(MVT::i32, Value);
+  SDValue ValueI32 = Value;
+  if (VT.isFloatingPoint())
+    ValueI32 = DAG.getBitcast(MVT::i32, Value);
 
   auto MakeIntrinsic = [&DAG, &SL](unsigned IID, MVT RetVT,
                                    SmallVector<SDValue> IntrinArgs) -> SDValue {
@@ -7307,54 +7308,55 @@ static SDValue lowerSubgroupShuffle(const SITargetLowering &TLI, SDNode *N,
     return DAG.getNode(ISD::INTRINSIC_WO_CHAIN, SL, RetVT, Operands);
   };
 
+  // If we can bpermute across the whole wave, then just do that
   if (TLI.getSubtarget()->supportsWaveWideBPermute()) {
-    // If we can bpermute across the whole wave, then just do that
     SDValue BPermute = MakeIntrinsic(Intrinsic::amdgcn_ds_bpermute, MVT::i32,
-                                     {ShiftedIndex, Value32});
+                                     {ShiftedIndex, ValueI32});
     return DAG.getBitcast(VT, BPermute);
-  } else {
-    assert(TLI.getSubtarget()->isWave64());
-
-    // Otherwise, we need to make use of whole wave mode
-    SDValue PoisonVal = DAG.getPOISON(Value32->getValueType(0));
-    SDValue PoisonIndex = DAG.getPOISON(ShiftedIndex->getValueType(0));
-
-    // Set inactive lanes to poison
-    SDValue WWMValue = MakeIntrinsic(Intrinsic::amdgcn_set_inactive, MVT::i32,
-                                     {Value32, PoisonVal});
-    SDValue WWMIndex = MakeIntrinsic(Intrinsic::amdgcn_set_inactive, MVT::i32,
-                                     {ShiftedIndex, PoisonIndex});
-
-    SDValue Swapped =
-        MakeIntrinsic(Intrinsic::amdgcn_permlane64, MVT::i32, {WWMValue});
-
-    // Get permutation of each half, then we'll select which one to use
-    SDValue BPermSameHalf = MakeIntrinsic(Intrinsic::amdgcn_ds_bpermute,
-                                          MVT::i32, {WWMIndex, WWMValue});
-    SDValue BPermOtherHalf = MakeIntrinsic(Intrinsic::amdgcn_ds_bpermute,
-                                           MVT::i32, {WWMIndex, Swapped});
-    SDValue BPermOtherHalfWWM =
-        MakeIntrinsic(Intrinsic::amdgcn_wwm, MVT::i32, {BPermOtherHalf});
-
-    // Select which side to take the permute from
-    SDValue ThreadIDMask = DAG.getTargetConstant(UINT32_MAX, SL, MVT::i32);
-    SDValue ThreadIDLo =
-        MakeIntrinsic(Intrinsic::amdgcn_mbcnt_lo, MVT::i32,
-                      {ThreadIDMask, DAG.getTargetConstant(0, SL, MVT::i32)});
-    SDValue ThreadID = MakeIntrinsic(Intrinsic::amdgcn_mbcnt_hi, MVT::i32,
-                                     {ThreadIDMask, ThreadIDLo});
-
-    SDValue SameOrOtherHalf =
-        DAG.getNode(ISD::AND, SL, MVT::i32,
-                    DAG.getNode(ISD::XOR, SL, MVT::i32, ThreadID, Index),
-                    DAG.getTargetConstant(32, SL, MVT::i32));
-    SDValue UseSameHalf =
-        DAG.getSetCC(SL, MVT::i1, SameOrOtherHalf,
-                     DAG.getConstant(0, SL, MVT::i32), ISD::SETEQ);
-    SDValue Result = DAG.getSelect(SL, MVT::i32, UseSameHalf, BPermSameHalf,
-                                   BPermOtherHalfWWM);
-    return DAG.getBitcast(VT, Result);
   }
+
+  assert(TLI.getSubtarget()->isWave64());
+
+  // Otherwise, we need to make use of whole wave mode
+  SDValue PoisonVal = DAG.getPOISON(ValueI32->getValueType(0));
+  SDValue PoisonIndex = DAG.getPOISON(ShiftedIndex->getValueType(0));
+
+  // Set inactive lanes to poison
+  SDValue WWMValue = MakeIntrinsic(Intrinsic::amdgcn_set_inactive, MVT::i32,
+                                   {ValueI32, PoisonVal});
+  SDValue WWMIndex = MakeIntrinsic(Intrinsic::amdgcn_set_inactive, MVT::i32,
+                                   {ShiftedIndex, PoisonIndex});
+
+  SDValue Swapped =
+      MakeIntrinsic(Intrinsic::amdgcn_permlane64, MVT::i32, {WWMValue});
+
+  // Get permutation of each half, then we'll select which one to use
+  SDValue BPermSameHalf = MakeIntrinsic(Intrinsic::amdgcn_ds_bpermute,
+                                        MVT::i32, {WWMIndex, WWMValue});
+  SDValue BPermOtherHalf = MakeIntrinsic(Intrinsic::amdgcn_ds_bpermute,
+                                         MVT::i32, {WWMIndex, Swapped});
+  SDValue BPermOtherHalfWWM =
+      MakeIntrinsic(Intrinsic::amdgcn_wwm, MVT::i32, {BPermOtherHalf});
+
+  // Select which side to take the permute from
+  SDValue ThreadIDMask = DAG.getAllOnesConstant(SL, MVT::i32);
+  // We can get away with only using mbcnt_lo here since we're only
+  // trying to detect which side of 32 each lane is on, and mbcnt_lo
+  // returns 32 for lanes 32-63.
+  SDValue ThreadID =
+      MakeIntrinsic(Intrinsic::amdgcn_mbcnt_lo, MVT::i32,
+                    {ThreadIDMask, DAG.getTargetConstant(0, SL, MVT::i32)});
+
+  SDValue SameOrOtherHalf =
+      DAG.getNode(ISD::AND, SL, MVT::i32,
+                  DAG.getNode(ISD::XOR, SL, MVT::i32, ThreadID, Index),
+                  DAG.getTargetConstant(32, SL, MVT::i32));
+  SDValue UseSameHalf =
+      DAG.getSetCC(SL, MVT::i1, SameOrOtherHalf,
+                   DAG.getConstant(0, SL, MVT::i32), ISD::SETEQ);
+  SDValue Result = DAG.getSelect(SL, MVT::i32, UseSameHalf, BPermSameHalf,
+                                 BPermOtherHalfWWM);
+  return DAG.getBitcast(VT, Result);
 }
 
 void SITargetLowering::ReplaceNodeResults(SDNode *N,
@@ -10264,8 +10266,8 @@ SDValue SITargetLowering::LowerINTRINSIC_WO_CHAIN(SDValue Op,
       Poisons.push_back(DAG.getPOISON(ValTy));
     return DAG.getMergeValues(Poisons, SDLoc(Op));
   }
-  case Intrinsic::amdgcn_subgroup_shuffle:
-    return lowerSubgroupShuffle(*this, Op.getNode(), DAG);
+  case Intrinsic::amdgcn_wave_shuffle:
+    return lowerWaveShuffle(*this, Op.getNode(), DAG);
   default:
     if (const AMDGPU::ImageDimIntrinsicInfo *ImageDimIntr =
             AMDGPU::getImageDimIntrinsicInfo(IntrinsicID))