[llvm] [NVPTX] Handle comparisons of byval kernel arguments (PR #228439)

[via llvm-commits]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-backend-nvptx

Author: wenhu1024

<details>
<summary>Changes</summary>

Fixes https://github.com/llvm/llvm-project/issues/226873

This is probably not a common source pattern, but it is valid IR, and NVPTXLowerArgs should not crash on it. For example:

```llvm
define ptx_kernel void @<!-- -->kernel(ptr readonly byval(%S) align 8 %arg, ptr %out) {
  %cmp = icmp ne ptr %arg, null
  store i1 %cmp, ptr %out
  ret void
}
```


`convertToParamAS` does not handle `ICmpInst`, so a comparison involving a byval argument can currently reach `llvm_unreachable`.

This patch makes this case fall back to the existing local-copy path instead. Does this approach make sense?

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


2 Files Affected:

- (modified) llvm/lib/Target/NVPTX/NVPTXLowerArgs.cpp (+6) 
- (added) llvm/test/CodeGen/NVPTX/lower-byval-args-icmp.ll (+78) 


``````````diff
diff --git a/llvm/lib/Target/NVPTX/NVPTXLowerArgs.cpp b/llvm/lib/Target/NVPTX/NVPTXLowerArgs.cpp
index 659f72922e05e7..b1bb7cd5f1d1f6 100644
--- a/llvm/lib/Target/NVPTX/NVPTXLowerArgs.cpp
+++ b/llvm/lib/Target/NVPTX/NVPTXLowerArgs.cpp
@@ -231,6 +231,12 @@ struct ArgUseChecker : PtrUseVisitor<ArgUseChecker> {
 
   void visitPtrToIntInst(PtrToIntInst &I) { Base::visitPtrToIntInst(I); }
 
+  void visitICmpInst(ICmpInst &I) {
+    // Comparisons can expose the address of the byval copy, just like ptrtoint.
+    // They also cannot be rewritten by convertToParamAS.
+    PI.setEscaped(&I);
+  }
+
   void visitPHINodeOrSelectInst(Instruction &I) {
     assert(isa<PHINode>(I) || isa<SelectInst>(I));
     enqueueUsers(I);
diff --git a/llvm/test/CodeGen/NVPTX/lower-byval-args-icmp.ll b/llvm/test/CodeGen/NVPTX/lower-byval-args-icmp.ll
new file mode 100644
index 00000000000000..a5b69b9391d2d5
--- /dev/null
+++ b/llvm/test/CodeGen/NVPTX/lower-byval-args-icmp.ll
@@ -0,0 +1,78 @@
+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py UTC_ARGS: --version 5
+; RUN: opt < %s -mtriple nvptx64 -mcpu=sm_60 -mattr=ptx77 -nvptx-lower-args -S | FileCheck %s --check-prefixes=LOWER-ARGS
+; RUN: opt < %s -mtriple nvptx64 -mcpu=sm_70 -mattr=ptx77 -nvptx-lower-args -S | FileCheck %s --check-prefixes=LOWER-ARGS
+; RUN: opt < %s -mtriple nvptx64 -mcpu=sm_60 -mattr=ptx77 -passes=nvptx-lower-args -S | FileCheck %s --check-prefixes=LOWER-ARGS
+; RUN: opt < %s -mtriple nvptx64 -mcpu=sm_70 -mattr=ptx77 -passes=nvptx-lower-args -S | FileCheck %s --check-prefixes=LOWER-ARGS
+; RUN: llc < %s -mcpu=sm_60 -mattr=ptx77 | FileCheck %s --check-prefixes=PTX
+; RUN: llc < %s -mcpu=sm_70 -mattr=ptx77 | FileCheck %s --check-prefixes=PTX
+
+target triple = "nvptx64-nvidia-cuda"
+
+%S = type { ptr, double }
+
+; Comparing a byval pointer with null must not crash argument lowering.
+define ptx_kernel void @icmp_unused(ptr byval(%S) align 8 %arg) {
+; LOWER-ARGS-LABEL: define ptx_kernel void @icmp_unused(
+; LOWER-ARGS-SAME: ptr addrspace(101) readonly byval([[S:%.*]]) align 8 [[ARG1:%.*]]) #[[ATTR0:[0-9]+]] {
+; LOWER-ARGS-NEXT:    [[ARG:%.*]] = alloca [[S]], align 8
+; LOWER-ARGS-NEXT:    call void @llvm.memcpy.p0.p101.i64(ptr align 8 [[ARG]], ptr addrspace(101) align 8 [[ARG1]], i64 16, i1 false)
+; LOWER-ARGS-NEXT:    [[CMP:%.*]] = icmp eq ptr [[ARG]], null
+; LOWER-ARGS-NEXT:    ret void
+;
+; PTX-LABEL: icmp_unused(
+; PTX:       {
+; PTX-NEXT:    .local .align 8 .b8 __local_depot0[16];
+; PTX-NEXT:    .reg .b64 %SP;
+; PTX-NEXT:    .reg .b64 %SPL;
+; PTX-NEXT:    .reg .b64 %rd<3>;
+; PTX-EMPTY:
+; PTX-NEXT:  // %bb.0:
+; PTX-NEXT:    mov.b64 %SPL, __local_depot0;
+; PTX-NEXT:    ld.param.b64 %rd1, [icmp_unused_param_0+8];
+; PTX-NEXT:    st.local.b64 [%SPL+8], %rd1;
+; PTX-NEXT:    ld.param.b64 %rd2, [icmp_unused_param_0];
+; PTX-NEXT:    st.local.b64 [%SPL], %rd2;
+; PTX-NEXT:    ret;
+  %cmp = icmp eq ptr %arg, null
+  ret void
+}
+
+; Pointer comparisons observe the address of the private byval copy.
+define ptx_kernel void @icmp_live(ptr byval(%S) align 8 %arg, ptr %other,
+; LOWER-ARGS-LABEL: define ptx_kernel void @icmp_live(
+; LOWER-ARGS-SAME: ptr addrspace(101) readonly byval([[S:%.*]]) align 8 [[ARG1:%.*]], ptr [[OTHER:%.*]], ptr [[OUT:%.*]]) #[[ATTR0]] {
+; LOWER-ARGS-NEXT:    [[ARG:%.*]] = alloca [[S]], align 8
+; LOWER-ARGS-NEXT:    call void @llvm.memcpy.p0.p101.i64(ptr align 8 [[ARG]], ptr addrspace(101) align 8 [[ARG1]], i64 16, i1 false)
+; LOWER-ARGS-NEXT:    [[CMP:%.*]] = icmp eq ptr [[ARG]], [[OTHER]]
+; LOWER-ARGS-NEXT:    store i1 [[CMP]], ptr [[OUT]], align 1
+; LOWER-ARGS-NEXT:    ret void
+;
+; PTX-LABEL: icmp_live(
+; PTX:       {
+; PTX-NEXT:    .local .align 8 .b8 __local_depot1[16];
+; PTX-NEXT:    .reg .b64 %SP;
+; PTX-NEXT:    .reg .b64 %SPL;
+; PTX-NEXT:    .reg .pred %p<2>;
+; PTX-NEXT:    .reg .b16 %rs<2>;
+; PTX-NEXT:    .reg .b64 %rd<8>;
+; PTX-EMPTY:
+; PTX-NEXT:  // %bb.0:
+; PTX-NEXT:    mov.b64 %SPL, __local_depot1;
+; PTX-NEXT:    ld.param.b64 %rd1, [icmp_live_param_2];
+; PTX-NEXT:    cvta.to.global.u64 %rd2, %rd1;
+; PTX-NEXT:    ld.param.b64 %rd3, [icmp_live_param_1];
+; PTX-NEXT:    add.u64 %rd4, %SPL, 0;
+; PTX-NEXT:    cvta.local.u64 %rd5, %rd4;
+; PTX-NEXT:    ld.param.b64 %rd6, [icmp_live_param_0+8];
+; PTX-NEXT:    st.local.b64 [%SPL+8], %rd6;
+; PTX-NEXT:    ld.param.b64 %rd7, [icmp_live_param_0];
+; PTX-NEXT:    st.local.b64 [%SPL], %rd7;
+; PTX-NEXT:    setp.eq.b64 %p1, %rd5, %rd3;
+; PTX-NEXT:    selp.b16 %rs1, 1, 0, %p1;
+; PTX-NEXT:    st.global.b8 [%rd2], %rs1;
+; PTX-NEXT:    ret;
+  ptr %out) {
+  %cmp = icmp eq ptr %arg, %other
+  store i1 %cmp, ptr %out
+  ret void
+}

``````````

</details>


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