[RS4GC] Copy argument attributes from call to statepoint

The current implementation completely ignores argument attributes on calls, discarding them completely when creating a statepoint from a call instruction.
This is problematic in some scenarios as the argument attributes affect the ABI of the call, leading to undefined behavior if called with the wrong ABI attributes.
Note that this cannot be solved either by just having the function declaration annotated with the right parameter attributes as the call might be indirect, therefore requiring them to be present on the arguments.

This PR simply copies all parameter attributes over from the original call to the created statepoint.
declare i32 @personality_function()
declare i32 @personality_function()
Note that some argument attributes become invalid after the lowering as they imply memory effects that no longer hold with the statepoints. These do not need to be explicitly handled in this PR as they are removed by the `stripNonValidDataFromBody`.

Author: zero9178

llvm/lib/Transforms/Scalar/RewriteStatepointsForGC.cpp

@@ -1423,7 +1423,7 @@ static constexpr Attribute::AttrKind FnAttrsToStrip[] =
 
 // Create new attribute set containing only attributes which can be transferred
 // from original call to the safepoint.
-static AttributeList legalizeCallAttributes(LLVMContext &Ctx,
+static AttributeList legalizeCallAttributes(LLVMContext &Ctx, unsigned NumArgs,
                                             AttributeList OrigAL,
                                             AttributeList StatepointAL) {
   if (OrigAL.isEmpty())
@@ -1439,7 +1439,15 @@ static AttributeList legalizeCallAttributes(LLVMContext &Ctx,
       FnAttrs.removeAttribute(A);
   }
 
-  // Just skip parameter and return attributes for now
+  // Attach the argument attributes from the original call at the corresponding
+  // arguments in the statepoint. Note that any argument attributes that are
+  // invalid after lowering are stripped in stripNonValidDataFromBody.
+  for (unsigned I : llvm::seq(NumArgs))
+    StatepointAL = StatepointAL.addParamAttributes(
+        Ctx, GCStatepointInst::CallArgsBeginPos + I,
+        AttrBuilder(Ctx, OrigAL.getParamAttrs(I)));
+
+  // Return attributes are later attached to the gc.result intrinsic.
   return StatepointAL.addFnAttributes(Ctx, FnAttrs);
 }
 
@@ -1630,6 +1638,7 @@ makeStatepointExplicitImpl(CallBase *Call, /* to replace */
   // with a return value, we lower then as never returning calls to
   // __llvm_deoptimize that are followed by unreachable to get better codegen.
   bool IsDeoptimize = false;
+  bool IsMemIntrinsic = false;
 
   StatepointDirectives SD =
       parseStatepointDirectivesFromAttrs(Call->getAttributes());
@@ -1670,6 +1679,8 @@ makeStatepointExplicitImpl(CallBase *Call, /* to replace */
       IsDeoptimize = true;
     } else if (IID == Intrinsic::memcpy_element_unordered_atomic ||
                IID == Intrinsic::memmove_element_unordered_atomic) {
+      IsMemIntrinsic = true;
+
       // Unordered atomic memcpy and memmove intrinsics which are not explicitly
       // marked as "gc-leaf-function" should be lowered in a GC parseable way.
       // Specifically, these calls should be lowered to the
@@ -1785,12 +1796,19 @@ makeStatepointExplicitImpl(CallBase *Call, /* to replace */
     SPCall->setTailCallKind(CI->getTailCallKind());
     SPCall->setCallingConv(CI->getCallingConv());
 
-    // Currently we will fail on parameter attributes and on certain
-    // function attributes.  In case if we can handle this set of attributes -
-    // set up function attrs directly on statepoint and return attrs later for
+    // Set up function attrs directly on statepoint and return attrs later for
     // gc_result intrinsic.
-    SPCall->setAttributes(legalizeCallAttributes(
-        CI->getContext(), CI->getAttributes(), SPCall->getAttributes()));
+    AttributeList AttributeList = CI->getAttributes();
+    // The memory intrinsics do not have a 1:1 correspondence of the original
+    // call arguments to the produced statepoint. Remove the parameter
+    // attributes to avoid putting them on incorrect arguments.
+    if (IsMemIntrinsic)
+      AttributeList =
+          AttributeList::get(CI->getContext(), AttributeList.getFnAttrs(),
+                             AttributeList.getRetAttrs(), {});
+    SPCall->setAttributes(legalizeCallAttributes(CI->getContext(),
+                                                 CI->arg_size(), AttributeList,
+                                                 SPCall->getAttributes()));
 
     Token = cast<GCStatepointInst>(SPCall);
 
@@ -1812,12 +1830,19 @@ makeStatepointExplicitImpl(CallBase *Call, /* to replace */
 
     SPInvoke->setCallingConv(II->getCallingConv());
 
-    // Currently we will fail on parameter attributes and on certain
-    // function attributes.  In case if we can handle this set of attributes -
-    // set up function attrs directly on statepoint and return attrs later for
+    // Set up function attrs directly on statepoint and return attrs later for
     // gc_result intrinsic.
-    SPInvoke->setAttributes(legalizeCallAttributes(
-        II->getContext(), II->getAttributes(), SPInvoke->getAttributes()));
+    AttributeList AttributeList = II->getAttributes();
+    // The memory intrinsics do not have a 1:1 correspondence of the original
+    // call arguments to the produced statepoint. Remove the parameter
+    // attributes to avoid putting them on incorrect arguments.
+    if (IsMemIntrinsic)
+      AttributeList =
+          AttributeList::get(II->getContext(), AttributeList.getFnAttrs(),
+                             AttributeList.getRetAttrs(), {});
+    SPInvoke->setAttributes(
+        legalizeCallAttributes(II->getContext(), II->arg_size(),
+                               II->getAttributes(), SPInvoke->getAttributes()));
 
     Token = cast<GCStatepointInst>(SPInvoke);
 

llvm/test/Transforms/RewriteStatepointsForGC/call-argument-attributes.ll

@@ -0,0 +1,42 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 3
+; RUN: opt < %s -passes=rewrite-statepoints-for-gc -S | FileCheck %s
+
+declare i8 @callee(ptr, i8, float, ptr)
+
+define i8 @test(ptr %arg) gc "statepoint-example" {
+; CHECK-LABEL: define i8 @test(
+; CHECK-SAME: ptr [[ARG:%.*]]) gc "statepoint-example" {
+; CHECK-NEXT:    [[STATEPOINT_TOKEN:%.*]] = call token (i64, i32, ptr, i32, i32, ...) @llvm.experimental.gc.statepoint.p0(i64 2882400000, i32 0, ptr elementtype(i8 (ptr, i8, float, ptr)) @callee, i32 4, i32 0, ptr nocapture sret({ i64, i64 }) align 8 null, i8 signext 8, float inreg 1.000000e+00, ptr [[ARG]], i32 0, i32 0)
+; CHECK-NEXT:    [[R1:%.*]] = call zeroext i8 @llvm.experimental.gc.result.i8(token [[STATEPOINT_TOKEN]])
+; CHECK-NEXT:    ret i8 [[R1]]
+;
+  %r = call zeroext i8 @callee(ptr sret({i64, i64}) noalias align 8 nocapture null, i8 signext 8, float inreg 1.0, ptr writeonly %arg)
+  ret i8 %r
+}
+
+declare i32 @personality_function()
+
+define i8 @test_invoke(ptr %arg) gc "statepoint-example" personality ptr @personality_function {
+; CHECK-LABEL: define i8 @test_invoke(
+; CHECK-SAME: ptr [[ARG:%.*]]) gc "statepoint-example" personality ptr @personality_function {
+; CHECK-NEXT:    [[STATEPOINT_TOKEN:%.*]] = invoke token (i64, i32, ptr, i32, i32, ...) @llvm.experimental.gc.statepoint.p0(i64 2882400000, i32 0, ptr elementtype(i8 (ptr, i8, float, ptr)) @callee, i32 4, i32 0, ptr nocapture sret({ i64, i64 }) align 8 null, i8 signext 8, float inreg 1.000000e+00, ptr [[ARG]], i32 0, i32 0)
+; CHECK-NEXT:    to label [[NORMAL_RETURN:%.*]] unwind label [[EXCEPTIONAL_RETURN:%.*]]
+; CHECK:       normal_return:
+; CHECK-NEXT:    [[R1:%.*]] = call zeroext i8 @llvm.experimental.gc.result.i8(token [[STATEPOINT_TOKEN]])
+; CHECK-NEXT:    ret i8 [[R1]]
+; CHECK:       exceptional_return:
+; CHECK-NEXT:    [[LANDING_PAD4:%.*]] = landingpad token
+; CHECK-NEXT:    cleanup
+; CHECK-NEXT:    ret i8 0
+;
+  %r = invoke zeroext i8 @callee(ptr sret({i64, i64}) noalias align 8 nocapture null, i8 signext 8, float inreg 1.0, ptr writeonly %arg)
+  to label %normal_return unwind label %exceptional_return
+
+normal_return:
+  ret i8 %r
+
+exceptional_return:
+  %landing_pad4 = landingpad token
+  cleanup
+  ret i8 0
+}