[CHERI] Fix Bounded Mem/Var Args layout on stack

veselypeta · veselypeta · commit ee9fc1571fb4 · 2025-07-04T15:07:54.000+02:00
diff --git a/llvm/lib/Target/RISCV/RISCVISelLowering.cpp b/llvm/lib/Target/RISCV/RISCVISelLowering.cpp
@@ -15234,8 +15234,10 @@ bool RISCV::CC_RISCV(const DataLayout &DL, RISCVABI::ABI ABI, unsigned ValNo,
                    ? Subtarget.typeForCapabilities()
                    : MVT();
   MVT PtrVT = DL.isFatPointer(DL.getAllocaAddrSpace()) ? CLenVT : XLenVT;
-  bool IsPureCapVarArgs = !IsFixed && RISCVABI::isCheriPureCapABI(ABI);
+  bool IsPureCap = RISCVABI::isCheriPureCapABI(ABI);
+  bool IsPureCapVarArgs = !IsFixed && IsPureCap;
   bool IsBoundedVarArgs = IsPureCapVarArgs && Subtarget.hasCheriBoundVarArg();
+  unsigned SlotSize = PtrVT.getFixedSizeInBits() / 8;
 
   // Static chain parameter must not be passed in normal argument registers,
   // so we assign t2 for it as done in GCC's __builtin_call_with_static_chain
@@ -15309,7 +15311,7 @@ bool RISCV::CC_RISCV(const DataLayout &DL, RISCVABI::ABI ABI, unsigned ValNo,
   // not apply.
   // TODO: Pure capability varargs bounds
   unsigned TwoXLenInBytes = (2 * XLen) / 8;
-  if (!IsFixed && !RISCVABI::isCheriPureCapABI(ABI) &&
+  if (!IsFixed && !IsPureCap &&
       ArgFlags.getNonZeroOrigAlign() == TwoXLenInBytes &&
       DL.getTypeAllocSize(OrigTy) == TwoXLenInBytes) {
     unsigned RegIdx = State.getFirstUnallocated(ArgGPRs);
@@ -15325,38 +15327,10 @@ bool RISCV::CC_RISCV(const DataLayout &DL, RISCVABI::ABI ABI, unsigned ValNo,
   assert(PendingLocs.size() == PendingArgFlags.size() &&
          "PendingLocs and PendingArgFlags out of sync");
 
-  // Bounded VarArgs
-  // Each bounded varargs is assigned a 2*XLen slot on the stack
-  // If the value is small enough to fit into the slot it is passed
-  // directly - otherwise a capability to the value is filled into the
-  // slot.
-  if (!IsFixed && IsBoundedVarArgs) {
-    unsigned SlotSize = CLenVT.getFixedSizeInBits() / 8;
-    // Aggregates of size 2*XLen need special handling here
-    // as LLVM with treat them as two separate XLen wide arguments
-    if(LocVT == XLenVT && OrigTy && OrigTy->isAggregateType()){
-      PendingLocs.push_back(
-          CCValAssign::getPending(ValNo, ValVT, LocVT, LocInfo));
-      PendingArgFlags.push_back(ArgFlags);
-      if(PendingLocs.size() == 2){
-        CCValAssign VA = PendingLocs[0];
-        ISD::ArgFlagsTy AF = PendingArgFlags[0];
-        PendingLocs.clear();
-        PendingArgFlags.clear();
-        return CC_RISCVAssign2XLen(XLen, State, IsPureCapVarArgs, VA, AF,
-                               ValNo, ValVT, LocVT, ArgFlags);
-      }
-      return false;
-    }
-    unsigned StackOffset = State.AllocateStack(SlotSize, Align(SlotSize));
-    State.addLoc(CCValAssign::getMem(ValNo, ValVT, StackOffset, LocVT, LocInfo));
-    return false;
-  }
-
   // Handle passing f64 on RV32D with a soft float ABI or when floating point
   // registers are exhausted. Also handle for pure capability varargs which are
   // always passed on the stack.
-  if ((UseGPRForF64 || IsPureCapVarArgs) && XLen == 32 && ValVT == MVT::f64) {
+  if (UseGPRForF64 && XLen == 32 && ValVT == MVT::f64) {
     assert(!ArgFlags.isSplit() && PendingLocs.empty() &&
            "Can't lower f64 if it is split");
     // Depending on available argument GPRS, f64 may be passed in a pair of
@@ -15366,7 +15340,9 @@ bool RISCV::CC_RISCV(const DataLayout &DL, RISCVABI::ABI ABI, unsigned ValNo,
     Register Reg = IsPureCapVarArgs ? 0 : State.AllocateReg(ArgGPRs);
     LocVT = MVT::i32;
     if (!Reg) {
-      unsigned StackOffset = State.AllocateStack(8, Align(8));
+      unsigned StackOffset =
+          IsBoundedVarArgs ? State.AllocateStack(SlotSize, Align(SlotSize))
+                           : State.AllocateStack(8, Align(8));
       State.addLoc(
           CCValAssign::getMem(ValNo, ValVT, StackOffset, LocVT, LocInfo));
       return false;
@@ -15407,8 +15383,8 @@ bool RISCV::CC_RISCV(const DataLayout &DL, RISCVABI::ABI ABI, unsigned ValNo,
     ISD::ArgFlagsTy AF = PendingArgFlags[0];
     PendingLocs.clear();
     PendingArgFlags.clear();
-    return CC_RISCVAssign2XLen(XLen, State, IsPureCapVarArgs, VA, AF,
-                               ValNo, ValVT, LocVT, ArgFlags);
+    return CC_RISCVAssign2XLen(XLen, State, IsPureCapVarArgs, VA, AF, ValNo,
+                               ValVT, LocVT, ArgFlags);
   }
 
   // Will be passed indirectly; make sure we allocate the right type of
@@ -15461,8 +15437,17 @@ bool RISCV::CC_RISCV(const DataLayout &DL, RISCVABI::ABI ABI, unsigned ValNo,
     Reg = State.AllocateReg(ArgGPRs);
   }
 
+  // Aggregate types i.e. structs/arrays which can fit into 2*XLEN
+  // Don't allocate a slot for each instead we make sure that the next element
+  // is then properly aligned.
+  bool AllocateSlot = IsBoundedVarArgs;
+  if (OrigTy && OrigTy->isAggregateType())
+    AllocateSlot = false;
   unsigned StackOffset =
-      Reg ? 0 : State.AllocateStack(StoreSizeBytes, StackAlign);
+      Reg ? 0
+          : (AllocateSlot
+                 ? State.AllocateStack(SlotSize, Align(SlotSize))
+                 : State.AllocateStack(StoreSizeBytes, StackAlign));
 
   // If we reach this point and PendingLocs is non-empty, we must be at the
   // end of a split argument that must be passed indirectly.
@@ -16339,7 +16324,7 @@ SDValue RISCVTargetLowering::LowerCall(CallLoweringInfo &CLI,
 
   // Bookkeeping for cheri varargs/memargs
   int VAArgStartOffset, VAArgEndOffset, MemArgStartOffset, MemArgEndOffset;
-  SDValue FirstAddr, FirstArgAddr;
+  SDValue FirstVAAddr, FirstArgAddr;
 
   // Copy argument values to their designated locations.
   SmallVector<std::pair<Register, SDValue>, 8> RegsToPass;
@@ -16466,34 +16451,23 @@ SDValue RISCVTargetLowering::LowerCall(CallLoweringInfo &CLI,
       SDValue Address =
           DAG.getPointerAdd(DL, StackPtr, VA.getLocMemOffset());
 
-      if (UseBoundeMemArgsCaller) {
+      if (UseBoundeMemArgsCaller && Outs[i].IsFixed) {
         if (FirstArgAddr == SDValue()) {
           FirstArgAddr = Address;
           MemArgStartOffset = VA.getLocMemOffset();
         }
         unsigned VTSize = VA.getValVT().getSizeInBits() / 8;
         MemArgEndOffset = VA.getLocMemOffset() + VTSize;
-        if (!Outs[i].IsFixed) {
-          // we need to align to 16-byte slot
-          Align OffsetAlign = Align(PtrLenBytes);
-          Type *OrigTy = CLI.getArgs()[Outs[i].OrigArgIndex].Ty;
-          if (OrigTy && OrigTy->isAggregateType())
-            OffsetAlign = Align(PtrLenBytes / 2);
-          MemArgEndOffset = alignTo(MemArgEndOffset, OffsetAlign);
-        }
       }
       if (UseBoundedVarArgs && !Outs[i].IsFixed) {
-        if (FirstAddr == SDValue()) {
-          FirstAddr = Address;
+        if (FirstVAAddr == SDValue()) {
+          FirstVAAddr = Address;
           VAArgStartOffset = VA.getLocMemOffset();
         }
         Align OffsetAlign = Align(PtrLenBytes);
-        Type *OrigTy = CLI.getArgs()[Outs[i].OrigArgIndex].Ty;
-        if (OrigTy && OrigTy->isAggregateType())
-          OffsetAlign = Align(PtrLenBytes / 2);
-
         unsigned VTSize = VA.getValVT().getSizeInBits() / 8;
         VAArgEndOffset = alignTo(VA.getLocMemOffset() + VTSize, OffsetAlign);
+        MemArgEndOffset = VAArgEndOffset;
       }
 
       // Emit the store.
@@ -16503,15 +16477,19 @@ SDValue RISCVTargetLowering::LowerCall(CallLoweringInfo &CLI,
   }
 
   if(IsVarArg && UseBoundedVarArgs && !UseBoundeMemArgsCaller) {
-    if (FirstAddr != SDValue()) {
+    if (FirstVAAddr != SDValue()) {
       SDValue VarArgs = DAG.getCSetBounds(
-          FirstAddr, DL, VAArgEndOffset - VAArgStartOffset, Align(),
+          FirstVAAddr, DL, VAArgEndOffset - VAArgStartOffset, Align(),
           "CHERI-RISCV variadic call lowering",
           cheri::SetBoundsPointerSource::Stack, "varargs call bounds setting");
       // clear write and execute permissions on varargs. Clearning other
       // permissions shouldn't be necessary since the capability is derived from
       // CSP and that shouldn't have these in the first place.
-      uint64_t PermMask = -1UL & ~(CAP_AP_X | CAP_AP_W);
+      uint64_t ExecPerm =
+          Subtarget.hasStdExtZCheriPureCap() ? (1 << 17) : (1 << 1);
+      uint64_t WritePerm =
+          Subtarget.hasStdExtZCheriPureCap() ? (1 << 1) : (1 << 3);
+      uint64_t PermMask = -1UL & ~(ExecPerm | WritePerm);
       VarArgs = DAG.getNode(
           ISD::INTRINSIC_WO_CHAIN, DL, PtrVT,
           DAG.getConstant(Intrinsic::cheri_cap_perms_and, DL, XLenVT), VarArgs,
diff --git a/llvm/test/CodeGen/RISCV/cheri/bakewell/cheri-bounded-memargs-caller.ll b/llvm/test/CodeGen/RISCV/cheri/bakewell/cheri-bounded-memargs-caller.ll
@@ -146,7 +146,8 @@ define i32 @biz() local_unnamed_addr addrspace(200) nounwind {
 ; CHECK-NEXT:    sd a0, 16(csp)
 ; CHECK-NEXT:    li a1, 32
 ; CHECK-NEXT:    scbndsr ca1, csp, a1
-; CHECK-NEXT:    li a2, -11
+; CHECK-NEXT:    lui a2, 1048544
+; CHECK-NEXT:    addiw a2, a2, -3
 ; CHECK-NEXT:    acperm ct1, ca1, a2
 ; CHECK-NEXT:    li a1, 1
 ; CHECK-NEXT:    li a2, 2
diff --git a/llvm/test/CodeGen/RISCV/cheri/purecap-bounded-varargs.ll b/llvm/test/CodeGen/RISCV/cheri/purecap-bounded-varargs.ll
@@ -6,13 +6,13 @@
 define i32 @caller_test_scalars(i32 %x, i128 %y, i64 %z, float %f, double %d) local_unnamed_addr addrspace(200) nounwind {
 ; CHECK-LABEL: caller_test_scalars:
 ; CHECK:       # %bb.0: # %entry
-; CHECK-NEXT:    cincoffset csp, csp, -192
-; CHECK-NEXT:    sc cra, 176(csp) # 16-byte Folded Spill
-; CHECK-NEXT:    sc cs0, 160(csp) # 16-byte Folded Spill
-; CHECK-NEXT:    sc cs1, 144(csp) # 16-byte Folded Spill
-; CHECK-NEXT:    sc cs2, 128(csp) # 16-byte Folded Spill
-; CHECK-NEXT:    sc cs3, 112(csp) # 16-byte Folded Spill
-; CHECK-NEXT:    sc cs4, 96(csp) # 16-byte Folded Spill
+; CHECK-NEXT:    cincoffset csp, csp, -176
+; CHECK-NEXT:    sc cra, 160(csp) # 16-byte Folded Spill
+; CHECK-NEXT:    sc cs0, 144(csp) # 16-byte Folded Spill
+; CHECK-NEXT:    sc cs1, 128(csp) # 16-byte Folded Spill
+; CHECK-NEXT:    sc cs2, 112(csp) # 16-byte Folded Spill
+; CHECK-NEXT:    sc cs3, 96(csp) # 16-byte Folded Spill
+; CHECK-NEXT:    sc cs4, 80(csp) # 16-byte Folded Spill
 ; CHECK-NEXT:    mv s0, a5
 ; CHECK-NEXT:    mv s1, a3
 ; CHECK-NEXT:    mv s2, a2
@@ -22,24 +22,24 @@ define i32 @caller_test_scalars(i32 %x, i128 %y, i64 %z, float %f, double %d) lo
 ; CHECK-NEXT:    srli a0, a0, 32
 ; CHECK-NEXT:    call __extendsfdf2
 ; CHECK-NEXT:    mv a1, a0
-; CHECK-NEXT:    sd s0, 80(csp)
-; CHECK-NEXT:    sd s1, 48(csp)
-; CHECK-NEXT:    sd s2, 32(csp)
+; CHECK-NEXT:    sd s0, 64(csp)
+; CHECK-NEXT:    sd s1, 32(csp)
+; CHECK-NEXT:    sd s2, 24(csp)
 ; CHECK-NEXT:    sd s3, 16(csp)
 ; CHECK-NEXT:    sd s4, 0(csp)
-; CHECK-NEXT:    csetbounds ca0, csp, 96
+; CHECK-NEXT:    csetbounds ca0, csp, 80
 ; CHECK-NEXT:    li a2, -11
 ; CHECK-NEXT:    candperm ct1, ca0, a2
 ; CHECK-NEXT:    li a0, 5
-; CHECK-NEXT:    sd a1, 64(csp)
+; CHECK-NEXT:    sd a1, 48(csp)
 ; CHECK-NEXT:    call callee
-; CHECK-NEXT:    lc cra, 176(csp) # 16-byte Folded Reload
-; CHECK-NEXT:    lc cs0, 160(csp) # 16-byte Folded Reload
-; CHECK-NEXT:    lc cs1, 144(csp) # 16-byte Folded Reload
-; CHECK-NEXT:    lc cs2, 128(csp) # 16-byte Folded Reload
-; CHECK-NEXT:    lc cs3, 112(csp) # 16-byte Folded Reload
-; CHECK-NEXT:    lc cs4, 96(csp) # 16-byte Folded Reload
-; CHECK-NEXT:    cincoffset csp, csp, 192
+; CHECK-NEXT:    lc cra, 160(csp) # 16-byte Folded Reload
+; CHECK-NEXT:    lc cs0, 144(csp) # 16-byte Folded Reload
+; CHECK-NEXT:    lc cs1, 128(csp) # 16-byte Folded Reload
+; CHECK-NEXT:    lc cs2, 112(csp) # 16-byte Folded Reload
+; CHECK-NEXT:    lc cs3, 96(csp) # 16-byte Folded Reload
+; CHECK-NEXT:    lc cs4, 80(csp) # 16-byte Folded Reload
+; CHECK-NEXT:    cincoffset csp, csp, 176
 ; CHECK-NEXT:    ret
 entry:
   %conv = fpext float %f to double
@@ -75,21 +75,22 @@ entry:
 define i32 @caller_test_struct(i32 %x, [2 x float] %y.coerce, i32 %z, { i8 addrspace(200)*, i64} %u.coerce) local_unnamed_addr addrspace(200) nounwind {
 ; CHECK-LABEL: caller_test_struct:
 ; CHECK:       # %bb.0: # %entry
-; CHECK-NEXT:    cincoffset csp, csp, -80
-; CHECK-NEXT:    sc cra, 64(csp) # 16-byte Folded Spill
-; CHECK-NEXT:    mv a5, a0
+; CHECK-NEXT:    cincoffset csp, csp, -96
+; CHECK-NEXT:    sc cra, 80(csp) # 16-byte Folded Spill
+; CHECK-NEXT:    mv a6, a0
+; CHECK-NEXT:    sd a5, 64(csp)
 ; CHECK-NEXT:    sc ca4, 48(csp)
 ; CHECK-NEXT:    sd a3, 32(csp)
 ; CHECK-NEXT:    sd a2, 24(csp)
 ; CHECK-NEXT:    sd a1, 16(csp)
-; CHECK-NEXT:    csetbounds ca0, csp, 64
+; CHECK-NEXT:    csetbounds ca0, csp, 80
 ; CHECK-NEXT:    li a1, -11
 ; CHECK-NEXT:    candperm ct1, ca0, a1
 ; CHECK-NEXT:    li a0, 3
-; CHECK-NEXT:    sd a5, 0(csp)
+; CHECK-NEXT:    sd a6, 0(csp)
 ; CHECK-NEXT:    call callee
-; CHECK-NEXT:    lc cra, 64(csp) # 16-byte Folded Reload
-; CHECK-NEXT:    cincoffset csp, csp, 80
+; CHECK-NEXT:    lc cra, 80(csp) # 16-byte Folded Reload
+; CHECK-NEXT:    cincoffset csp, csp, 96
 ; CHECK-NEXT:    ret
 entry:
   %call = tail call i32 (i32, ...) @callee(i32 3, i32 %x, [2 x float] %y.coerce, i32 %z, { i8 addrspace(200)*, i64} %u.coerce) nounwind
