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())
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 {
70 Addr.ChangeToES(SymName: Sym.getSymbolName(), TargetFlags: Sym.getTargetFlags());
71 }
72
73 if (MRI.use_empty(RegNo: Mov->getOperand(i: 0).getReg()))
74 Mov->eraseFromParent();
75
76 return true;
77}
78
79static bool foldAddresses(MachineFunction &MF) {
80 MachineRegisterInfo &MRI = MF.getRegInfo();
81
82 bool Changed = false;
83 for (MachineBasicBlock &MBB : MF)
84 for (MachineInstr &MI : make_early_inc_range(Range&: MBB))
85 if (MI.mayLoadOrStore()) {
86 const int AddrIdx =
87 NVPTX::getNamedOperandIdx(Opcode: MI.getOpcode(), Name: NVPTX::OpName::addr);
88 if (AddrIdx >= 0 && MI.getOperand(i: AddrIdx).isReg())
89 Changed |= foldAddress(MI, Addr&: MI.getOperand(i: AddrIdx), MRI);
90 }
91
92 return Changed;
93}
94
95/// ----------------------------------------------------------------------------
96/// Pass (Manager) Boilerplate
97/// ----------------------------------------------------------------------------
98
99namespace {
100struct NVPTXAddressFolderPass : public MachineFunctionPass {
101 static char ID;
102 NVPTXAddressFolderPass() : MachineFunctionPass(ID) {}
103
104 bool runOnMachineFunction(MachineFunction &MF) override;
105
106 void getAnalysisUsage(AnalysisUsage &AU) const override {
107 MachineFunctionPass::getAnalysisUsage(AU);
108 }
109};
110} // namespace
111
112char NVPTXAddressFolderPass::ID = 0;
113
114INITIALIZE_PASS(NVPTXAddressFolderPass, "nvptx-address-folder",
115 "NVPTX Address Folder", false, false)
116
117bool NVPTXAddressFolderPass::runOnMachineFunction(MachineFunction &MF) {
118 return foldAddresses(MF);
119}
120
121MachineFunctionPass *llvm::createNVPTXAddressFolderPass() {
122 return new NVPTXAddressFolderPass();
123}
124