[RS4GC] Copy argument attributes from call to statepoint (#68475)

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.
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

@@ -18,6 +18,7 @@
 #include "llvm/ADT/DenseSet.h"
 #include "llvm/ADT/MapVector.h"
 #include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/Sequence.h"
 #include "llvm/ADT/SetVector.h"
 #include "llvm/ADT/SmallSet.h"
 #include "llvm/ADT/SmallVector.h"
@@ -1422,14 +1423,15 @@ static constexpr Attribute::AttrKind FnAttrsToStrip[] =
   {Attribute::Memory, Attribute::NoSync, Attribute::NoFree};
 
 // Create new attribute set containing only attributes which can be transferred
-// from original call to the safepoint.
-static AttributeList legalizeCallAttributes(LLVMContext &Ctx,
-                                            AttributeList OrigAL,
+// from the original call to the safepoint.
+static AttributeList legalizeCallAttributes(CallBase *Call, bool IsMemIntrinsic,
                                             AttributeList StatepointAL) {
+  AttributeList OrigAL = Call->getAttributes();
   if (OrigAL.isEmpty())
     return StatepointAL;
 
   // Remove the readonly, readnone, and statepoint function attributes.
+  LLVMContext &Ctx = Call->getContext();
   AttrBuilder FnAttrs(Ctx, OrigAL.getFnAttrs());
   for (auto Attr : FnAttrsToStrip)
     FnAttrs.removeAttribute(Attr);
@@ -1439,8 +1441,24 @@ static AttributeList legalizeCallAttributes(LLVMContext &Ctx,
       FnAttrs.removeAttribute(A);
   }
 
-  // Just skip parameter and return attributes for now
-  return StatepointAL.addFnAttributes(Ctx, FnAttrs);
+  StatepointAL = StatepointAL.addFnAttributes(Ctx, FnAttrs);
+
+  // The memory intrinsics do not have a 1:1 correspondence of the original
+  // call arguments to the produced statepoint. Do not transfer the argument
+  // attributes to avoid putting them on incorrect arguments.
+  if (IsMemIntrinsic)
+    return StatepointAL;
+
+  // 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(Call->arg_size()))
+    StatepointAL = StatepointAL.addParamAttributes(
+        Ctx, GCStatepointInst::CallArgsBeginPos + I,
+        AttrBuilder(Ctx, OrigAL.getParamAttrs(I)));
+
+  // Return attributes are later attached to the gc.result intrinsic.
+  return StatepointAL;
 }
 
 /// Helper function to place all gc relocates necessary for the given
@@ -1630,6 +1648,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 +1689,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 +1806,10 @@ 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()));
+    SPCall->setAttributes(
+        legalizeCallAttributes(CI, IsMemIntrinsic, SPCall->getAttributes()));
 
     Token = cast<GCStatepointInst>(SPCall);
 
@@ -1812,12 +1831,10 @@ 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()));
+    SPInvoke->setAttributes(
+        legalizeCallAttributes(II, IsMemIntrinsic, 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
+}