[llvm] [AMDGPU] Add waterfall intrinsics (PR #192409)
Nicolai Hähnle via llvm-commits
llvm-commits at lists.llvm.org
Mon Jun 8 12:02:15 PDT 2026
================
@@ -0,0 +1,798 @@
+//===- AMDGPUInsertWaterfall.cpp - insert waterfall loops around markers -====//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+/// \file
+/// Replace intrinsics used to mark waterfall regions with actual waterfall
+/// loops. This is done at MachineIR level rather than LLVM-IR due to the use of
+/// exec mask in this operation.
+///
+//===----------------------------------------------------------------------===//
+
+#include "AMDGPUInsertWaterfall.h"
+#include "AMDGPU.h"
+#include "GCNSubtarget.h"
+#include "MCTargetDesc/AMDGPUMCTargetDesc.h"
+#include "SIInstrInfo.h"
+#include "llvm/CodeGen/MachineFunctionPass.h"
+
+using namespace llvm;
+
+#define DEBUG_TYPE "amdgpu-insert-waterfall"
+
+namespace {
+
+static unsigned getWFBeginSize(const unsigned Opcode) {
+ switch (Opcode) {
+ case AMDGPU::SI_WATERFALL_BEGIN_V1:
+ return 1;
+ case AMDGPU::SI_WATERFALL_BEGIN_V2:
+ return 2;
+ case AMDGPU::SI_WATERFALL_BEGIN_V4:
+ return 4;
+ case AMDGPU::SI_WATERFALL_BEGIN_V8:
+ return 8;
+ default:
+ break;
+ }
+
+ return 0; // Not SI_WATERFALL_BEGIN_*
+}
+
+static unsigned getWFRFLSize(const unsigned Opcode) {
+ switch (Opcode) {
+ case AMDGPU::SI_WATERFALL_READFIRSTLANE_V1:
+ return 1;
+ case AMDGPU::SI_WATERFALL_READFIRSTLANE_V2:
+ return 2;
+ case AMDGPU::SI_WATERFALL_READFIRSTLANE_V4:
+ return 4;
+ case AMDGPU::SI_WATERFALL_READFIRSTLANE_V8:
+ return 8;
+ default:
+ break;
+ }
+
+ return 0; // Not SI_WATERFALL_READFIRSTLANE_*
+}
+
+static unsigned getWFEndSize(const unsigned Opcode) {
+ switch (Opcode) {
+ case AMDGPU::SI_WATERFALL_END_V1:
+ return 1;
+ case AMDGPU::SI_WATERFALL_END_V2:
+ return 2;
+ case AMDGPU::SI_WATERFALL_END_V4:
+ return 4;
+ case AMDGPU::SI_WATERFALL_END_V8:
+ return 8;
+ default:
+ break;
+ }
+
+ return 0; // Not SI_WATERFALL_END_*
+}
+
+static void readFirstLaneReg(MachineBasicBlock &MBB, MachineRegisterInfo *MRI,
+ const SIRegisterInfo *RI, const SIInstrInfo *TII,
+ MachineBasicBlock::iterator &I, const DebugLoc &DL,
+ Register RFLReg, const MachineOperand &RFLSrcOp) {
+ Register RFLSrcReg = RFLSrcOp.getReg();
+ const TargetRegisterClass *RFLRegRC = MRI->getRegClass(RFLReg);
+ uint32_t RegSize = RI->getRegSizeInBits(*RFLRegRC) / 32;
+ assert(RI->hasVGPRs(MRI->getRegClass(RFLSrcReg)) &&
+ "unexpected uniform operand for readfirstlane");
+
+ if (RegSize == 1) {
+ MRI->constrainRegClass(RFLReg, &AMDGPU::SReg_32_XM0RegClass);
+ BuildMI(MBB, I, DL, TII->get(AMDGPU::V_READFIRSTLANE_B32), RFLReg)
+ .addReg(RFLSrcReg, getUndefRegState(RFLSrcOp.isUndef()),
+ RFLSrcOp.getSubReg());
+ } else {
+ SmallVector<Register, 8> TRegs;
+ for (unsigned i = 0; i < RegSize; ++i) {
+ Register TReg = MRI->createVirtualRegister(&AMDGPU::SReg_32_XM0RegClass);
+ BuildMI(MBB, I, DL, TII->get(AMDGPU::V_READFIRSTLANE_B32), TReg)
+ .addReg(RFLSrcReg, {}, RI->getSubRegFromChannel(i));
+ TRegs.push_back(TReg);
+ }
+ MachineInstrBuilder MIB =
+ BuildMI(MBB, I, DL, TII->get(AMDGPU::REG_SEQUENCE), RFLReg);
+ for (unsigned i = 0; i < RegSize; ++i) {
+ MIB.addReg(TRegs[i]);
+ MIB.addImm(RI->getSubRegFromChannel(i));
+ }
+ }
+}
+
+// Check if operand is uniform by checking:
+// 1. Trivially detectable as operand in SGPR
+// 2. Direct def is from an SGPR->VGPR copy (which may happen if assumed
+// non-uniform value turns out to be uniform)
+static Register getUniformOperandReplacementReg(MachineRegisterInfo *MRI,
+ const SIRegisterInfo *RI,
+ Register Reg) {
+ auto RegRC = MRI->getRegClass(Reg);
+ if (!RI->hasVGPRs(RegRC)) {
+ return Reg;
+ }
+
+ // Check for operand def being a copy from SGPR
+ MachineInstr *DefMI = MRI->getVRegDef(Reg);
+ if (DefMI->isFullCopy()) {
+ auto const &DefSrcOp = DefMI->getOperand(1);
+ if (DefSrcOp.isReg() && DefSrcOp.getReg().isVirtual()) {
+ Register ReplaceReg = DefSrcOp.getReg();
+ if (!RI->hasVGPRs(MRI->getRegClass(ReplaceReg)))
+ return ReplaceReg;
+ }
+ }
+ return AMDGPU::NoRegister;
+}
+
+static void compareIdxUsingCmpx(
+ MachineBasicBlock &MBB, MachineRegisterInfo *MRI, const SIRegisterInfo *RI,
+ const SIInstrInfo *TII, MachineBasicBlock::iterator &I, const DebugLoc &DL,
+ Register CurrentIdxReg, const MachineOperand &IndexOp, bool IsWave32) {
+
+ unsigned CmpxEqU32Opc = AMDGPU::V_CMPX_EQ_U32_nosdst_e32;
+
+ Register IndexReg = IndexOp.getReg();
+ const TargetRegisterClass *IndexRC =
+ RI->getRegClassForOperandReg(*MRI, IndexOp);
+
+ // Iterate over the index in dword chunks
+ uint32_t RegSize = RI->getRegSizeInBits(*IndexRC) / 32;
+ if (RegSize == 1) {
+ BuildMI(MBB, I, DL, TII->get(CmpxEqU32Opc))
+ .addReg(CurrentIdxReg)
+ .addReg(IndexReg, {}, IndexOp.getSubReg());
+
+ } else {
+ for (unsigned Idx = 0; Idx < RegSize; ++Idx) {
+ BuildMI(MBB, I, DL, TII->get(CmpxEqU32Opc))
+ .addReg(CurrentIdxReg, {}, RI->getSubRegFromChannel(Idx))
+ .addReg(IndexReg, {}, RI->getSubRegFromChannel(Idx));
+ }
+ }
+}
+
+static Register compareIdx(MachineBasicBlock &MBB, MachineRegisterInfo *MRI,
+ const SIRegisterInfo *RI, const SIInstrInfo *TII,
+ MachineBasicBlock::iterator &I, const DebugLoc &DL,
+ Register CurrentIdxReg,
+ const MachineOperand &IndexOp, Register CondReg,
+ bool IsWave32) {
+ // Iterate over the index in dword chunks and'ing the result with the
+ // CondReg
+ // Optionally CondReg is passed in from a previous compareIdx call
+ Register IndexReg = IndexOp.getReg();
+ auto IndexRC = RI->getRegClassForOperandReg(*MRI, IndexOp);
+ unsigned AndOpc = IsWave32 ? AMDGPU::S_AND_B32 : AMDGPU::S_AND_B64;
+ const auto *BoolXExecRC = TII->getRegisterInfo().getWaveMaskRegClass();
+
+ uint32_t RegSize = RI->getRegSizeInBits(*IndexRC) / 32;
+
+ if (RegSize == 1) {
+ Register TReg = MRI->createVirtualRegister(BoolXExecRC);
+ BuildMI(MBB, I, DL, TII->get(AMDGPU::V_CMP_EQ_U32_e64), TReg)
+ .addReg(CurrentIdxReg)
+ .addReg(IndexReg, {}, IndexOp.getSubReg());
+
+ if (CondReg != AMDGPU::NoRegister) {
+ Register TReg2 = MRI->createVirtualRegister(BoolXExecRC);
+ BuildMI(MBB, I, DL, TII->get(AndOpc), TReg2).addReg(CondReg).addReg(TReg);
+ CondReg = TReg2;
+ } else {
+ CondReg = TReg;
+ }
+ } else {
+ unsigned StartCount;
+ Register TReg;
+ if (CondReg != AMDGPU::NoRegister) {
+ TReg = CondReg;
+ StartCount = 0;
+ } else {
+ TReg = MRI->createVirtualRegister(BoolXExecRC);
+ BuildMI(MBB, I, DL, TII->get(AMDGPU::V_CMP_EQ_U32_e64), TReg)
+ .addReg(CurrentIdxReg, {}, AMDGPU::sub0)
+ .addReg(IndexReg, {}, AMDGPU::sub0);
+ StartCount = 1;
+ }
+
+ for (unsigned i = StartCount; i < RegSize; ++i) {
+ Register TReg2 = MRI->createVirtualRegister(BoolXExecRC);
+ BuildMI(MBB, I, DL, TII->get(AMDGPU::V_CMP_EQ_U32_e64), TReg2)
+ .addReg(CurrentIdxReg, {}, RI->getSubRegFromChannel(i))
+ .addReg(IndexReg, {}, RI->getSubRegFromChannel(i));
+ Register TReg3 = MRI->createVirtualRegister(BoolXExecRC);
+ BuildMI(MBB, I, DL, TII->get(AndOpc), TReg3).addReg(TReg).addReg(TReg2);
+ TReg = TReg3;
+ }
+ CondReg = TReg;
+ }
+ return CondReg;
+}
+
+// Replace all registers From with To.
+// Also handles From and To MachineOperands having sub registers.
+// Note: MRI->replaceRegWith doesn't handle sub registers since it is
+// register based and subreg is carried on the operand.
+static void replaceRegIncSubReg(const MachineRegisterInfo *MRI,
+ const TargetRegisterInfo *TRI,
+ const MachineOperand *From,
+ const MachineOperand *To) {
+ for (auto &O : make_early_inc_range(MRI->reg_operands(From->getReg())))
+ O.substVirtReg(To->getReg(), To->getSubReg(), *TRI);
+}
+
+class AMDGPUInsertWaterfall : public MachineFunctionPass {
+private:
+ struct WaterfallWorkitem {
+ const SIInstrInfo *TII;
+ const MachineRegisterInfo *MRI;
+ Register TokReg; // This is always the token from the last begin intrinsic
+ MachineInstr *Final;
+
+ SmallVector<MachineInstr *, 4> BeginList;
+ SmallVector<MachineInstr *, 4> RFLList;
+ SmallVector<MachineInstr *, 4> EndList;
+
+ // List of corresponding init, newdst and phi registers used in loop for
+ // end pseudos
+ std::vector<std::pair<MachineOperand *, MachineOperand *>> EndRegs;
+ std::vector<Register> RFLRegs;
+
+ WaterfallWorkitem() = default;
+ WaterfallWorkitem(MachineInstr *_Begin, const SIInstrInfo *_TII,
+ MachineRegisterInfo *_MRI)
+ : TII(_TII), MRI(_MRI), Final(nullptr) {
+
+ const MachineOperand *TokMO =
+ TII->getNamedOperand(*_Begin, AMDGPU::OpName::tok_ret);
+
+ assert(tokIsStart(TII->getNamedOperand(*_Begin, AMDGPU::OpName::tok)) &&
+ "first begin does not have an undefined input token as expected");
+
+ assert(TokMO &&
+ "Unable to extract tok operand from SI_WATERFALL_BEGIN pseudo op");
+
+ BeginList.push_back(_Begin);
+ TokReg = TokMO->getReg();
+ }
+
+ WaterfallWorkitem(const SIInstrInfo *_TII, MachineRegisterInfo *_MRI)
+ : TII(_TII), MRI(_MRI), TokReg(AMDGPU::NoRegister), Final(nullptr) {}
+
+ MachineInstr *getDefInstr(const MachineOperand *MO) const {
+ if (MO->isReg() && MRI->hasOneDef(MO->getReg())) {
+ return MRI->def_begin(MO->getReg())->getParent();
+ }
+ return nullptr;
+ }
+
+ bool tokIsStart(const MachineOperand *MO) const {
+ MachineInstr *defInstr = getDefInstr(MO);
+ if (!defInstr)
+ return false;
+ if (defInstr->getOpcode() == AMDGPU::IMPLICIT_DEF)
+ return true;
+ if (defInstr->isCopy())
+ return tokIsStart(&defInstr->getOperand(1));
+ // S_MOV_B32 0 for backwards compatibility with hand-written MIR.
+ if (defInstr->getOpcode() == AMDGPU::S_MOV_B32) {
+ auto CopySrcOp = TII->getNamedOperand(*defInstr, AMDGPU::OpName::src0);
+ if (CopySrcOp && CopySrcOp->isImm()) {
+ if (CopySrcOp->getImm() == 0)
+ return true;
+ }
+ }
+ return false;
+ }
+
+ bool addCandidate(MachineInstr *Cand) {
+ unsigned Opcode = Cand->getOpcode();
+
+ assert((getWFBeginSize(Opcode) || getWFRFLSize(Opcode) ||
+ getWFEndSize(Opcode)) &&
+ "expected a waterfall instruction in addCandidate");
+
+ auto CandTokMO = TII->getNamedOperand(*Cand, AMDGPU::OpName::tok);
+ // There are a couple of scenarios at this point:
+ // 1. Standard - there's already been a begin that's been processed and
+ // set up the WaterfallWorkItem. In which case the token is
+ // valid and needs to be checked to ensure well- formed
+ // waterfall groups.
+ // 2. Begins removed - begins were uniform - and all of them have been
+ // removed. Need to process the rest of the instructions
+ // in the group, and verify that they have a undefined
+ // token
+ if (TokReg == AMDGPU::NoRegister) {
+ // All begins have been removed - continue to process the rest of the
+ // grouping ready for them to be removed in the next stage
+ assert(!getWFBeginSize(Opcode) &&
+ "unexpected begin instruction for addCandidate");
+ assert(tokIsStart(CandTokMO) &&
+ "waterfall group with no begin doesn't have undef tok input");
+
+ TokReg = CandTokMO->getReg();
+ }
+ if (CandTokMO->getReg() == TokReg) {
+ if (getWFBeginSize(Opcode)) {
+ auto TokRetMO = TII->getNamedOperand(*Cand, AMDGPU::OpName::tok_ret);
+ assert(TokRetMO && "Unable to extract tok_ret operand from "
+ "SI_WATERFALL_BEGIN pseudo op");
+ BeginList.push_back(Cand);
+ TokReg = TokRetMO->getReg();
+ return true;
+ } else if (getWFRFLSize(Opcode)) {
+ RFLList.push_back(Cand);
+ return true;
+ } else if (getWFEndSize(Opcode)) {
+ EndList.push_back(Cand);
+ Final = Cand;
+ return true;
+ } else {
+ report_fatal_error("Unknown opcode, expected waterfall intrinsic");
+ }
+ }
+ LLVM_DEBUG(dbgs() << "malformed waterfall instruction group");
+ return false;
+ }
+
+ void eraseFromParent() {
+ for (auto BeginMI : BeginList)
+ BeginMI->eraseFromParent();
+ for (auto RFLMI : RFLList)
+ RFLMI->eraseFromParent();
+ for (auto EndMI : EndList)
+ EndMI->eraseFromParent();
+ }
+ };
+
+ std::vector<WaterfallWorkitem> Worklist;
+
+ const GCNSubtarget *ST;
+ const SIInstrInfo *TII;
+ MachineRegisterInfo *MRI;
+ const SIRegisterInfo *RI;
+
+public:
+ static char ID;
+
+ AMDGPUInsertWaterfall() : MachineFunctionPass(ID) {
+ initializeAMDGPUInsertWaterfallPass(*PassRegistry::getPassRegistry());
+ }
+
+ void getAnalysisUsage(AnalysisUsage &AU) const override {
+ MachineFunctionPass::getAnalysisUsage(AU);
+ }
+
+ bool removeRedundantWaterfall(WaterfallWorkitem &Item);
+ bool processWaterfall(MachineBasicBlock &MBB);
+
+ Register getToken(MachineInstr *MI);
+
+ bool runOnMachineFunction(MachineFunction &MF) override;
+};
+
+} // End anonymous namespace.
+
+INITIALIZE_PASS(AMDGPUInsertWaterfall, DEBUG_TYPE, "AMDGPU Insert waterfalls",
+ false, false)
+
+char AMDGPUInsertWaterfall::ID = 0;
+
+char &llvm::AMDGPUInsertWaterfallID = AMDGPUInsertWaterfall::ID;
+
+FunctionPass *llvm::createAMDGPUInsertWaterfallPass() {
+ return new AMDGPUInsertWaterfall;
+}
+
+PreservedAnalyses
+AMDGPUInsertWaterfallPass::run(MachineFunction &MF,
+ MachineFunctionAnalysisManager &MFAM) {
+ AMDGPUInsertWaterfall Impl;
+ if (!Impl.runOnMachineFunction(MF))
+ return PreservedAnalyses::all();
+ return PreservedAnalyses::none();
+}
+
+bool AMDGPUInsertWaterfall::removeRedundantWaterfall(WaterfallWorkitem &Item) {
+ // In some cases, the waterfall is actually redundant
+ // If all the readfirstlane intrinsics are actually for uniform values and
+ // the token used in the begin/end isn't used in anything else the waterfall
+ // can be removed.
+ // Alternatively, prior passes may have removed the readfirstlane intrinsics
+ // altogether, in this case the begin/end intrinsics are now redundant and can
+ // also be removed.
+ // The readfirstlane intrinsics are replaced with the uniform source value,
+ // the loop is removed and the defs in the end intrinsics are just replaced
+ // with the input operands
+ // We can also have cases where the begins are all removed (all the indices
+ // were actually uniform).
----------------
nhaehnle wrote:
It would be better to achieve the same thing using AMDGPUUniformIntrinsicCombine.
https://github.com/llvm/llvm-project/pull/192409
More information about the llvm-commits
mailing list