[llvm] [NVPTX] Forward grid constant kernel params used outside their block (PR #202379)

[via llvm-commits]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-backend-nvptx

Author: Tim Besard (maleadt)

<details>
<summary>Changes</summary>

`NVPTXForwardParams` turns a byval parameter that was materialized into a register (the "move param" form) back into a direct `ld.param`, which is cheaper. Today it only does this for device-function parameters: the `MOV{32,64}_PARAM` produced by `NVPTXISD::MoveParam`.

Kernel grid constant parameters are lowered to a bare parameter symbol instead. SelectionDAG can fold that symbol into `ld.param [sym+off]` only when the load is in the same basic block. Once a load is in another block, the symbol is carried across the block boundary in a vreg, ISel materializes it with the generic `MOV_B{32,64}_sym`, and the loads become register-relative. So a kernel that reads a byval field behind any control flow (a bounds check, a loop body) ends up with:

```
mov.b64       %rd1, kernel_param_0;
ld.param.b64  %rd2, [%rd1];
ld.param.b64  %rd3, [%rd1+8];
```

instead of:

```
ld.param.b64  %rd2, [kernel_param_0];
ld.param.b64  %rd3, [kernel_param_0+8];
```

The parameter stays in parameter space in both cases, with no local copy and no spill. The only difference is the redundant address arithmetic.

This patch extends the pass to cover that case:

- It recognizes the symbol mov (`MOV_B{32,64}_sym`) when the operand names a parameter of the function. That opcode is also used for global-address materialization, so the operand is matched against the function's parameter symbols.
- It scans the whole function rather than just the entry block, because MachineSink usually sinks the mov down to the block that uses it.
- It selects the address space from the function kind: `.param::entry` for kernels, `.param::func` for device functions (previously a fixed constant).

The device-function path is unchanged. A parameter whose address escapes (passed to a call, stored, used in inline asm, or carried through a phi) keeps its mov, since not all of its uses are loads.

This turned up from a CUDA.jl performance regression, where the front-end stopped pre-lowering simple-load byval kernel arguments and left them for the back-end to read from parameter space, which exposed the missing fold.

---
Full diff: https://github.com/llvm/llvm-project/pull/202379.diff


2 Files Affected:

- (modified) llvm/lib/Target/NVPTX/NVPTXForwardParams.cpp (+62-6) 
- (modified) llvm/test/CodeGen/NVPTX/forward-ld-param.ll (+34) 


``````````diff
diff --git a/llvm/lib/Target/NVPTX/NVPTXForwardParams.cpp b/llvm/lib/Target/NVPTX/NVPTXForwardParams.cpp
index 2b59286295d2b..e60c4d3828ed1 100644
--- a/llvm/lib/Target/NVPTX/NVPTXForwardParams.cpp
+++ b/llvm/lib/Target/NVPTX/NVPTXForwardParams.cpp
@@ -36,15 +36,26 @@
 // We do this by simply traversing uses of the param "mov" instructions an
 // trivially checking if they are all loads.
 //
+// The same "mov" also appears for grid constant kernel parameters: those are
+// lowered to a bare parameter symbol that SelectionDAG can fold directly into a
+// `ld.param` only within a single basic block. When such a parameter is used
+// outside its defining block its address is instead materialized with a `mov`,
+// and this pass forwards it back into the loads just like the device parameter
+// case above.
+//
 //===----------------------------------------------------------------------===//
 
 #include "NVPTX.h"
+#include "NVPTXSubtarget.h"
+#include "NVVMProperties.h"
 #include "llvm/ADT/SmallVector.h"
+#include "llvm/ADT/StringSet.h"
 #include "llvm/CodeGen/MachineFunctionPass.h"
 #include "llvm/CodeGen/MachineInstr.h"
 #include "llvm/CodeGen/MachineOperand.h"
 #include "llvm/CodeGen/MachineRegisterInfo.h"
 #include "llvm/CodeGen/TargetRegisterInfo.h"
+#include "llvm/IR/Function.h"
 #include "llvm/Support/ErrorHandling.h"
 
 using namespace llvm;
@@ -82,7 +93,8 @@ static bool traverseMoveUse(MachineInstr &U, const MachineRegisterInfo &MRI,
 }
 
 static bool eliminateMove(MachineInstr &Mov, const MachineRegisterInfo &MRI,
-                          SmallVectorImpl<MachineInstr *> &RemoveList) {
+                          SmallVectorImpl<MachineInstr *> &RemoveList,
+                          unsigned ParamAddrSpace) {
   SmallVector<MachineInstr *, 16> MaybeRemoveList;
   SmallVector<MachineInstr *, 16> LoadInsts;
 
@@ -102,20 +114,64 @@ static bool eliminateMove(MachineInstr &Mov, const MachineRegisterInfo &MRI,
     (LI->uses().begin() + LDInstBasePtrOpIdx)
         ->ChangeToES(ParamSymbol->getSymbolName());
     (LI->uses().begin() + LDInstAddrSpaceOpIdx)
-        ->ChangeToImmediate(NVPTX::AddressSpace::DeviceParam);
+        ->ChangeToImmediate(ParamAddrSpace);
   }
   return true;
 }
 
+// Return true if \p MI is a `mov` that materializes the address of a byval
+// parameter and is a candidate for forwarding. Device function parameters are
+// lowered through NVPTXISD::MoveParam (MOV{32,64}_PARAM), while grid constant
+// kernel parameters are lowered to a bare parameter symbol that ISel
+// materializes with the generic symbol `mov` (MOV_B{32,64}_sym) when the
+// parameter is used outside its defining block. The latter opcode is shared
+// with global address materialization, so check that the operand actually
+// names a parameter of this function.
+static bool isForwardableParamMove(const MachineInstr &MI,
+                                   const StringSet<> &ParamSymbols) {
+  switch (MI.getOpcode()) {
+  case NVPTX::MOV32_PARAM:
+  case NVPTX::MOV64_PARAM:
+    return true;
+  case NVPTX::MOV_B32_sym:
+  case NVPTX::MOV_B64_sym: {
+    const MachineOperand &Src = MI.getOperand(1);
+    return Src.isSymbol() && ParamSymbols.contains(Src.getSymbolName());
+  }
+  default:
+    return false;
+  }
+}
+
 static bool forwardDeviceParams(MachineFunction &MF) {
   const auto &MRI = MF.getRegInfo();
+  const Function &F = MF.getFunction();
+
+  // Kernel parameters live in the read-only ".param::entry" space, device
+  // function parameters in ".param::func"; both are read with `ld.param`.
+  const bool IsKernel = isKernelFunction(F);
+  const unsigned ParamAddrSpace = IsKernel ? NVPTX::AddressSpace::EntryParam
+                                           : NVPTX::AddressSpace::DeviceParam;
+
+  // Grid constant kernel parameters are materialized with the generic symbol
+  // `mov`, which is also used for global addresses. Collect the parameter
+  // symbol names so they can be told apart (see isForwardableParamMove).
+  StringSet<> ParamSymbols;
+  if (IsKernel) {
+    const NVPTXTargetLowering *TLI =
+        MF.getSubtarget<NVPTXSubtarget>().getTargetLowering();
+    for (const Argument &Arg : F.args())
+      ParamSymbols.insert(TLI->getParamName(&F, Arg.getArgNo()));
+  }
 
+  // The `mov` may have been sunk out of the entry block (e.g. to the single
+  // block that uses it), so scan the whole function rather than just the entry.
   bool Changed = false;
   SmallVector<MachineInstr *, 16> RemoveList;
-  for (auto &MI : make_early_inc_range(*MF.begin()))
-    if (MI.getOpcode() == NVPTX::MOV32_PARAM ||
-        MI.getOpcode() == NVPTX::MOV64_PARAM)
-      Changed |= eliminateMove(MI, MRI, RemoveList);
+  for (MachineBasicBlock &MBB : MF)
+    for (auto &MI : make_early_inc_range(MBB))
+      if (isForwardableParamMove(MI, ParamSymbols))
+        Changed |= eliminateMove(MI, MRI, RemoveList, ParamAddrSpace);
 
   for (auto *MI : RemoveList)
     MI->eraseFromParent();
diff --git a/llvm/test/CodeGen/NVPTX/forward-ld-param.ll b/llvm/test/CodeGen/NVPTX/forward-ld-param.ll
index 4f1454d3788a4..aff28be4a37f1 100644
--- a/llvm/test/CodeGen/NVPTX/forward-ld-param.ll
+++ b/llvm/test/CodeGen/NVPTX/forward-ld-param.ll
@@ -125,3 +125,37 @@ end:
   %v = phi i32 [ %v2, %if ], [ %v3, %else ]
   ret i32 %v
 }
+
+; A grid constant kernel parameter loaded at a constant offset from a block
+; other than the one that defines its address must still be read directly from
+; parameter space (ld.param [param+off]), not via a materialized base register.
+define ptx_kernel void @test_kernel_cross_block(ptr byval([2 x i64]) align 8 %r, i64 %n, ptr addrspace(1) %out) {
+; CHECK-LABEL: test_kernel_cross_block(
+; CHECK:       {
+; CHECK-NEXT:    .reg .pred %p<2>;
+; CHECK-NEXT:    .reg .b64 %rd<6>;
+; CHECK-EMPTY:
+; CHECK-NEXT:  // %bb.0:
+; CHECK-NEXT:    ld.param.b64 %rd2, [test_kernel_cross_block_param_1];
+; CHECK-NEXT:    setp.lt.s64 %p1, %rd2, 1;
+; CHECK-NEXT:    @%p1 bra $L__BB6_2;
+; CHECK-NEXT:  // %bb.1: // %body
+; CHECK-NEXT:    ld.param.b64 %rd1, [test_kernel_cross_block_param_2];
+; CHECK-NEXT:    ld.param.b64 %rd3, [test_kernel_cross_block_param_0];
+; CHECK-NEXT:    ld.param.b64 %rd4, [test_kernel_cross_block_param_0+8];
+; CHECK-NEXT:    mul.lo.s64 %rd5, %rd3, %rd4;
+; CHECK-NEXT:    st.global.b64 [%rd1], %rd5;
+; CHECK-NEXT:  $L__BB6_2: // %done
+; CHECK-NEXT:    ret;
+  %c = icmp sgt i64 %n, 0
+  br i1 %c, label %body, label %done
+body:
+  %v0 = load i64, ptr %r, align 8
+  %p1 = getelementptr inbounds i8, ptr %r, i64 8
+  %v1 = load i64, ptr %p1, align 8
+  %m = mul i64 %v0, %v1
+  store i64 %m, ptr addrspace(1) %out, align 8
+  br label %done
+done:
+  ret void
+}

``````````

</details>


https://github.com/llvm/llvm-project/pull/202379