[NVPTX] Assign correct memory attrs to args in NVPTXLowerArgs (#207461)
diff --git a/llvm/lib/Target/NVPTX/NVPTXLowerArgs.cpp b/llvm/lib/Target/NVPTX/NVPTXLowerArgs.cpp
index 3e8eb66..659f729 100644
--- a/llvm/lib/Target/NVPTX/NVPTXLowerArgs.cpp
+++ b/llvm/lib/Target/NVPTX/NVPTXLowerArgs.cpp
@@ -268,6 +268,13 @@
Arg.getParamAlign().value_or(DL.getPrefTypeAlign(ByValType)));
Arg.replaceAllUsesWith(AllocA);
+ // If the parameter is never read (writeonly or readnone), there is nothing to
+ // copy in; the alloca above already provides the writable local storage the
+ // body needs, and reading the param here would contradict the attribute.
+ if (Arg.hasAttribute(Attribute::ReadNone) ||
+ Arg.hasAttribute(Attribute::WriteOnly))
+ return;
+
// Be sure to propagate alignment to this copy; LLVM doesn't know that NVPTX
// addrspacecast preserves alignment. Since params are constant, this copy
// is definitely not volatile.
@@ -343,6 +350,20 @@
copyByValParam(F, OldArg, NewParamArg);
}
+// Mark a param-space byval argument as non-writable.
+static void markArgNonWritable(Argument &Arg) {
+ if (Arg.onlyReadsMemory())
+ return;
+
+ if (Arg.hasAttribute(Attribute::WriteOnly)) {
+ Arg.removeAttr(Attribute::WriteOnly);
+ Arg.addAttr(Attribute::ReadNone);
+ return;
+ }
+
+ Arg.addAttr(Attribute::ReadOnly);
+}
+
// Rewrite a kernel's signature so that each byval argument is declared directly
// as a pointer in the param address space, then lower the body to match. This
// creates a new function, moves the body across, and erases \p F.
@@ -366,11 +387,11 @@
// ISel reads the param symbol directly for kernel byval arguments; this is
// valid because the signature rewrite above puts them in the param address
- // space. Mark them readonly: any mutation is redirected to a local copy
+ // space. Mark them non-writable: any mutation is redirected to a local copy
// below, so the param itself is never written.
for (Argument &NewArg : NF->args())
if (NewArg.hasByValAttr())
- NewArg.addAttr(Attribute::ReadOnly);
+ markArgNonWritable(NewArg);
// Take over F's name and uses (e.g. @llvm.used, nvvm.annotations metadata),
// then move the body across.
diff --git a/llvm/test/CodeGen/NVPTX/lower-byval-args-mem-attrs.ll b/llvm/test/CodeGen/NVPTX/lower-byval-args-mem-attrs.ll
new file mode 100644
index 0000000..5bfdca8
--- /dev/null
+++ b/llvm/test/CodeGen/NVPTX/lower-byval-args-mem-attrs.ll
@@ -0,0 +1,84 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 5
+; RUN: opt < %s -mcpu=sm_70 -mattr=ptx77 -passes=nvptx-lower-args -S | FileCheck %s --check-prefixes=CHECK
+; RUN: opt < %s -mcpu=sm_70 -mattr=ptx77 -passes=nvptx-copy-byval-args -S | FileCheck %s --check-prefixes=COPY
+
+target triple = "nvptx64-nvidia-cuda"
+
+%struct.S = type { i32, i32 }
+
+declare void @escape(ptr)
+
+define ptx_kernel void @byval_write_only(ptr writeonly byval(%struct.S) align 4 %s) {
+; CHECK-LABEL: define ptx_kernel void @byval_write_only(
+; CHECK-SAME: ptr addrspace(101) readnone byval([[STRUCT_S:%.*]]) align 4 [[S1:%.*]]) #[[ATTR0:[0-9]+]] {
+; CHECK-NEXT: [[S:%.*]] = alloca [[STRUCT_S]], align 4
+; CHECK-NEXT: store i32 0, ptr [[S]], align 4
+; CHECK-NEXT: ret void
+;
+; COPY-LABEL: define ptx_kernel void @byval_write_only(
+; COPY-SAME: ptr writeonly byval([[STRUCT_S:%.*]]) align 4 [[S:%.*]]) #[[ATTR0:[0-9]+]] {
+; COPY-NEXT: [[S1:%.*]] = alloca [[STRUCT_S]], align 4
+; COPY-NEXT: store i32 0, ptr [[S1]], align 4
+; COPY-NEXT: ret void
+;
+ store i32 0, ptr %s, align 4
+ ret void
+}
+
+define ptx_kernel void @byval_read_none(ptr readnone byval(%struct.S) align 4 %s, ptr %out) {
+; CHECK-LABEL: define ptx_kernel void @byval_read_none(
+; CHECK-SAME: ptr addrspace(101) readnone byval([[STRUCT_S:%.*]]) align 4 [[S1:%.*]], ptr [[OUT:%.*]]) #[[ATTR0]] {
+; CHECK-NEXT: [[S:%.*]] = alloca [[STRUCT_S]], align 4
+; CHECK-NEXT: [[I:%.*]] = ptrtoint ptr [[S]] to i64
+; CHECK-NEXT: store i64 [[I]], ptr [[OUT]], align 8
+; CHECK-NEXT: ret void
+;
+; COPY-LABEL: define ptx_kernel void @byval_read_none(
+; COPY-SAME: ptr readnone byval([[STRUCT_S:%.*]]) align 4 [[S:%.*]], ptr [[OUT:%.*]]) #[[ATTR0]] {
+; COPY-NEXT: [[S1:%.*]] = alloca [[STRUCT_S]], align 4
+; COPY-NEXT: [[I:%.*]] = ptrtoint ptr [[S1]] to i64
+; COPY-NEXT: store i64 [[I]], ptr [[OUT]], align 8
+; COPY-NEXT: ret void
+;
+ %i = ptrtoint ptr %s to i64
+ store i64 %i, ptr %out, align 8
+ ret void
+}
+
+define ptx_kernel void @byval_plain(ptr byval(%struct.S) align 4 %s) {
+; CHECK-LABEL: define ptx_kernel void @byval_plain(
+; CHECK-SAME: ptr addrspace(101) readonly byval([[STRUCT_S:%.*]]) align 4 [[S1:%.*]]) #[[ATTR0]] {
+; CHECK-NEXT: [[S:%.*]] = alloca [[STRUCT_S]], align 4
+; CHECK-NEXT: call void @llvm.memcpy.p0.p101.i64(ptr align 4 [[S]], ptr addrspace(101) align 4 [[S1]], i64 8, i1 false)
+; CHECK-NEXT: call void @escape(ptr [[S]])
+; CHECK-NEXT: ret void
+;
+; COPY-LABEL: define ptx_kernel void @byval_plain(
+; COPY-SAME: ptr byval([[STRUCT_S:%.*]]) align 4 [[S:%.*]]) #[[ATTR0]] {
+; COPY-NEXT: [[S1:%.*]] = alloca [[STRUCT_S]], align 4
+; COPY-NEXT: call void @llvm.memcpy.p0.p0.i64(ptr align 4 [[S1]], ptr align 4 [[S]], i64 8, i1 false)
+; COPY-NEXT: call void @escape(ptr [[S1]])
+; COPY-NEXT: ret void
+;
+ call void @escape(ptr %s)
+ ret void
+}
+
+define ptx_kernel void @byval_read_only(ptr readonly byval(%struct.S) align 4 %s) {
+; CHECK-LABEL: define ptx_kernel void @byval_read_only(
+; CHECK-SAME: ptr addrspace(101) readonly byval([[STRUCT_S:%.*]]) align 4 [[S1:%.*]]) #[[ATTR0]] {
+; CHECK-NEXT: [[S:%.*]] = alloca [[STRUCT_S]], align 4
+; CHECK-NEXT: call void @llvm.memcpy.p0.p101.i64(ptr align 4 [[S]], ptr addrspace(101) align 4 [[S1]], i64 8, i1 false)
+; CHECK-NEXT: call void @escape(ptr [[S]])
+; CHECK-NEXT: ret void
+;
+; COPY-LABEL: define ptx_kernel void @byval_read_only(
+; COPY-SAME: ptr readonly byval([[STRUCT_S:%.*]]) align 4 [[S:%.*]]) #[[ATTR0]] {
+; COPY-NEXT: [[S1:%.*]] = alloca [[STRUCT_S]], align 4
+; COPY-NEXT: call void @llvm.memcpy.p0.p0.i64(ptr align 4 [[S1]], ptr align 4 [[S]], i64 8, i1 false)
+; COPY-NEXT: call void @escape(ptr [[S1]])
+; COPY-NEXT: ret void
+;
+ call void @escape(ptr %s)
+ ret void
+}
