[RS4GC] Copy argument attributes from call to statepoint

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
+}