1//===- NVPTXAddressFolder.cpp - Fold symbol addresses into memory ops -----===//
2//
3// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4// See https://llvm.org/LICENSE.txt for license information.
5// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6//
7//===----------------------------------------------------------------------===//
8//
9// SelectionDAG folds a symbol address (a kernel parameter, global variable, or
10// external symbol) directly into the address operand of a memory access, but
11// only within a single basic block. When the address is carried across a block
12// boundary in a register it is materialized with a generic `mov`
13// (MOV_B{32,64}_sym) and the accesses become register-relative:
14//
15// mov.b64 %rd1, kernel_param_0;
16// ld.param.b64 %rd2, [%rd1];
17// ld.param.b64 %rd3, [%rd1+8];
18//
19// This pass folds the symbol back into those address operands, eliminating the
20// redundant address arithmetic (and the `mov` itself once no use remains):
21//
22// ld.param.b64 %rd2, [kernel_param_0];
23// ld.param.b64 %rd3, [kernel_param_0+8];
24//
25// Shared-memory accesses are left alone: `mov`s of shared symbols are
26// deliberately kept CSE-able rather than duplicated into their uses, as
27// rematerializing them has caused performance regressions before (see
28// MovSymInst in NVPTXInstrInfo.td).
29//
30//===----------------------------------------------------------------------===//
31
32#include "NVPTX.h"
33#include "llvm/CodeGen/MachineFunctionPass.h"
34#include "llvm/CodeGen/MachineInstr.h"
35#include "llvm/CodeGen/MachineOperand.h"
36#include "llvm/CodeGen/MachineRegisterInfo.h"
37
38using namespace llvm;
39
40// Try to fold the definition of \p Addr into \p MI's address operand. The
41// defining instruction is erased once it has no uses left; it is kept for any
42// remaining use, e.g. because the address also feeds arithmetic or escapes.
43static bool foldAddress(MachineInstr &MI, MachineOperand &Addr,
44 MachineRegisterInfo &MRI) {
45 assert(Addr.isReg() && "Expected an address register");
46 assert(Addr.getReg().isVirtual() && "Expected a virtual address register");
47
48 MachineInstr *Mov = MRI.getVRegDef(Reg: Addr.getReg());
49 if (!Mov || (Mov->getOpcode() != NVPTX::MOV_B32_sym &&
50 Mov->getOpcode() != NVPTX::MOV_B64_sym))
51 return false;
52
53 const MachineOperand &Sym = Mov->getOperand(i: 1);
54 if (!Sym.isGlobal() && !Sym.isSymbol() && !Sym.isMCSymbol())
55 return false;
56
57 // The accessed address space must be known and must not be shared.
58 const int AddrSpaceIdx =
59 NVPTX::getNamedOperandIdx(Opcode: MI.getOpcode(), Name: NVPTX::OpName::addsp);
60 if (AddrSpaceIdx < 0)
61 return false;
62 const auto AddrSpace = MI.getOperand(i: AddrSpaceIdx).getImm();
63 if (AddrSpace == NVPTX::AddressSpace::Shared ||
64 AddrSpace == NVPTX::AddressSpace::SharedCluster)
65 return false;
66
67 if (Sym.isGlobal())
68 Addr.ChangeToGA(GV: Sym.getGlobal(), Offset: Sym.getOffset(), TargetFlags: Sym.getTargetFlags());
69 else if (Sym.isSymbol())
70 Addr.ChangeToES(SymName: Sym.getSymbolName(), TargetFlags: Sym.getTargetFlags());
71 else
72 Addr.ChangeToMCSymbol(Sym: Sym.getMCSymbol(), TargetFlags: Sym.getTargetFlags());
73
74 if (MRI.use_empty(RegNo: Mov->getOperand(i: 0).getReg()))
75 Mov->eraseFromParent();
76
77 return true;
78}
79
80static bool foldAddresses(MachineFunction &MF) {
81 MachineRegisterInfo &MRI = MF.getRegInfo();
82
83 bool Changed = false;
84 for (MachineBasicBlock &MBB : MF)
85 for (MachineInstr &MI : make_early_inc_range(Range&: MBB))
86 if (MI.mayLoadOrStore()) {
87 const int AddrIdx =
88 NVPTX::getNamedOperandIdx(Opcode: MI.getOpcode(), Name: NVPTX::OpName::addr);
89 if (AddrIdx >= 0 && MI.getOperand(i: AddrIdx).isReg())
90 Changed |= foldAddress(MI, Addr&: MI.getOperand(i: AddrIdx), MRI);
91 }
92
93 return Changed;
94}
95
96/// ----------------------------------------------------------------------------
97/// Pass (Manager) Boilerplate
98/// ----------------------------------------------------------------------------
99
100namespace {
101struct NVPTXAddressFolderLegacyPass : public MachineFunctionPass {
102 static char ID;
103 NVPTXAddressFolderLegacyPass() : MachineFunctionPass(ID) {}
104
105 bool runOnMachineFunction(MachineFunction &MF) override {
106 if (skipFunction(F: MF.getFunction()))
107 return false;
108 return foldAddresses(MF);
109 }
110
111 void getAnalysisUsage(AnalysisUsage &AU) const override {
112 AU.setPreservesCFG();
113 MachineFunctionPass::getAnalysisUsage(AU);
114 }
115};
116} // namespace
117
118char NVPTXAddressFolderLegacyPass::ID = 0;
119
120INITIALIZE_PASS(NVPTXAddressFolderLegacyPass, "nvptx-address-folder",
121 "NVPTX Address Folder", false, false)
122
123MachineFunctionPass *llvm::createNVPTXAddressFolderLegacyPass() {
124 return new NVPTXAddressFolderLegacyPass();
125}
126
127PreservedAnalyses
128NVPTXAddressFolderPass::run(MachineFunction &MF,
129 MachineFunctionAnalysisManager &MFAM) {
130 if (!foldAddresses(MF))
131 return PreservedAnalyses::all();
132 return getMachineFunctionPassPreservedAnalyses().preserveSet<CFGAnalyses>();
133}
134