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

llvm/lib/Target/RISCV/MCTargetDesc/RISCVBaseInfo.cpp

@@ -131,7 +131,7 @@ MCRegister getSCSPReg(ABI TargetABI) {
 }
 
 // Returns the register used for bounded mem/var args
-MCRegister getCheriBoundedArgReg() { return RISCV::C6; }
+MCRegister getCheriBoundedArgReg() { return RISCV::C31; }
 
 } // namespace RISCVABI
 

llvm/lib/Target/RISCV/RISCVISelLowering.cpp

@@ -24,7 +24,7 @@
 #include "llvm/ADT/Statistic.h"
 #include "llvm/Analysis/MemoryLocation.h"
 #include "llvm/Analysis/VectorUtils.h"
-#include "llvm/CHERI/cheri-compressed-cap/cheri_compressed_cap.h"
+#include "llvm/CodeGen/Analysis.h"
 #include "llvm/CodeGen/MachineFrameInfo.h"
 #include "llvm/CodeGen/MachineFunction.h"
 #include "llvm/CodeGen/MachineInstrBuilder.h"
@@ -2241,6 +2241,22 @@ unsigned RISCVTargetLowering::combineRepeatedFPDivisors() const {
   return NumRepeatedDivisors;
 }
 
+bool RISCVTargetLowering::functionArgumentNeedsConsecutiveRegisters(
+    Type *Ty, CallingConv::ID CallConv, bool isVarArg,
+    const DataLayout &DL) const {
+  const bool IsPureCapABI =
+      RISCVABI::isCheriPureCapABI(Subtarget.getTargetABI());
+  const bool HasBoundedVarArgs =
+      IsPureCapABI && Subtarget.hasCheriBoundVarArg();
+  if (!Ty->isArrayTy())
+    return false;
+
+  // All non aggregate members of the type must have the same type
+  SmallVector<EVT> ValueVTs;
+  ComputeValueVTs(*this, DL, Ty, ValueVTs);
+  return all_equal(ValueVTs) && isVarArg && HasBoundedVarArgs;
+}
+
 static SDValue getVLOperand(SDValue Op) {
   assert((Op.getOpcode() == ISD::INTRINSIC_WO_CHAIN ||
           Op.getOpcode() == ISD::INTRINSIC_W_CHAIN) &&
@@ -6502,7 +6518,8 @@ SDValue RISCVTargetLowering::lowerVASTARTCap(SDValue Op, SelectionDAG &DAG) cons
   unsigned PtrSize = MF.getDataLayout().getPointerSize(AllocaAS);
   bool IsPurecap = RISCVABI::isCheriPureCapABI(Subtarget.getTargetABI());
   bool UseBoundedMemArgsCallee =
-      IsPurecap && Subtarget.hasCheriBoundMemArgCallee();
+      IsPurecap && Subtarget.hasCheriBoundMemArgCallee() &&
+      MF.getFunction().getCallingConv() != CallingConv::Fast;
 
   int Index = FuncInfo->getPureCapVarArgsIndex();
   SDValue FI = DAG.getFrameIndex(Index, PtrVT);
@@ -6511,7 +6528,11 @@ SDValue RISCVTargetLowering::lowerVASTARTCap(SDValue Op, SelectionDAG &DAG) cons
                   MachinePointerInfo::getStack(MF, 0), Align(PtrSize));
   SDValue Chain = VarPtr.getOperand(0);
   if (UseBoundedMemArgsCallee) {
-    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 << 0) : (1 << 3);
+    uint64_t PermMask = -1UL & ~(ExecPerm | WritePerm);
     VarPtr = DAG.getNode(ISD::INTRINSIC_WO_CHAIN, DL, PtrVT,
                          DAG.getConstant(Intrinsic::cheri_cap_perms_and, DL,
                                          Subtarget.getXLenVT()),
@@ -15234,8 +15255,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 +15332,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 +15348,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 +15361,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;
@@ -15382,6 +15379,37 @@ bool RISCV::CC_RISCV(const DataLayout &DL, RISCVABI::ABI ABI, unsigned ValNo,
   if (ValVT.isFixedLengthVector())
     LocVT = TLI.getContainerForFixedLengthVector(LocVT);
 
+  // For purecap bounded varargs - aggregate types which can fit into a stack
+  // slot are passed to CC_RISCV as separate arguments. We need to align the
+  // first argument to a CLEN alignment.
+  if (IsBoundedVarArgs && ArgFlags.isInConsecutiveRegs()) {
+    PendingLocs.push_back(
+        CCValAssign::getPending(ValNo, ValVT, LocVT, LocInfo));
+    PendingArgFlags.push_back(ArgFlags);
+    if (!ArgFlags.isInConsecutiveRegsLast())
+      return false;
+  }
+
+  if (IsBoundedVarArgs && ArgFlags.isInConsecutiveRegsLast()) {
+    for (size_t I = 0, E = PendingLocs.size(); I < E; I++){
+      CCValAssign VA = PendingLocs[I];
+      unsigned Size =
+          VA.getValVT() == CLenVT ? DL.getPointerSize(200) : XLen / 8;
+      Align Alignment(Size);
+      // For consecutive types the first item needs to be aligned.
+      if (I == 0)
+        Alignment = Align(SlotSize);
+
+      unsigned StackOffset = State.AllocateStack(Size, Alignment);
+      State.addLoc(CCValAssign::getMem(VA.getValNo(), VA.getValVT(),
+                                       StackOffset, VA.getLocVT(),
+                                       VA.getLocInfo()));
+    }
+    PendingLocs.clear();
+    PendingArgFlags.clear();
+    return false;
+  }
+
   // Split arguments might be passed indirectly, so keep track of the pending
   // values. Split vectors are passed via a mix of registers and indirectly, so
   // treat them as we would any other argument.
@@ -15407,8 +15435,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
@@ -15462,7 +15490,10 @@ bool RISCV::CC_RISCV(const DataLayout &DL, RISCVABI::ABI ABI, unsigned ValNo,
   }
 
   unsigned StackOffset =
-      Reg ? 0 : State.AllocateStack(StoreSizeBytes, StackAlign);
+      Reg ? 0
+          : (IsBoundedVarArgs
+                 ? 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.
@@ -15669,7 +15700,8 @@ static SDValue unpackFromMemLoc(SelectionDAG &DAG, SDValue Chain,
   const RISCVSubtarget &STI = MF.getSubtarget<RISCVSubtarget>();
   const bool IsPureCapABI = RISCVABI::isCheriPureCapABI(STI.getTargetABI());
   const bool UseBoundedMemArgsCallee =
-      IsPureCapABI && STI.hasCheriBoundMemArgCallee();
+      IsPureCapABI && STI.hasCheriBoundMemArgCallee() &&
+      MF.getFunction().getCallingConv() != CallingConv::Fast;
 
   EVT LocVT = VA.getLocVT();
   EVT ValVT = VA.getValVT();
@@ -16011,11 +16043,13 @@ SDValue RISCVTargetLowering::LowerFormalArguments(
   bool HasMemArgs = false;
   const bool IsCheriPureCapABI =
       RISCVABI::isCheriPureCapABI(Subtarget.getTargetABI());
+  const bool IsFastCC = CallConv == CallingConv::Fast;
   const bool UseBoundedMemArgsCallee =
-      IsCheriPureCapABI && Subtarget.hasCheriBoundMemArgCallee();
+      IsCheriPureCapABI && Subtarget.hasCheriBoundMemArgCallee() && !IsFastCC;
   const bool UseBoundedMemArgsCaller =
-      IsCheriPureCapABI && Subtarget.hasCheriBoundMemArgCaller();
-  const bool UseBoundedVarArgs = IsCheriPureCapABI && Subtarget.hasCheriBoundVarArg();
+      IsCheriPureCapABI && Subtarget.hasCheriBoundMemArgCaller() && !IsFastCC;
+  const bool UseBoundedVarArgs =
+      IsCheriPureCapABI && Subtarget.hasCheriBoundVarArg();
   if (UseBoundedMemArgsCallee) {
     for (size_t I = 0; I < Ins.size(); I++) {
       CCValAssign &VA = ArgLocs[I];
@@ -16257,11 +16291,13 @@ SDValue RISCVTargetLowering::LowerCall(CallLoweringInfo &CLI,
 
   MachineFunction &MF = DAG.getMachineFunction();
   bool PureCapABI = RISCVABI::isCheriPureCapABI(Subtarget.getTargetABI());
-  bool UseBoundedVarArgs = PureCapABI && Subtarget.hasCheriBoundVarArg();
+  bool IsFastCC = CallConv == CallingConv::Fast;
+  bool UseBoundedVarArgs =
+      PureCapABI && Subtarget.hasCheriBoundVarArg() && !IsFastCC;
   bool UseBoundeMemArgsCaller =
-      PureCapABI && Subtarget.hasCheriBoundMemArgCaller();
+      PureCapABI && Subtarget.hasCheriBoundMemArgCaller() && !IsFastCC;
   bool UseBoundeMemArgsCallee =
-      PureCapABI && Subtarget.hasCheriBoundMemArgCallee();
+      PureCapABI && Subtarget.hasCheriBoundMemArgCallee() && !IsFastCC;
 
   // Analyze the operands of the call, assigning locations to each operand.
   SmallVector<CCValAssign, 16> ArgLocs;
@@ -16339,7 +16375,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;
@@ -16473,27 +16509,17 @@ SDValue RISCVTargetLowering::LowerCall(CallLoweringInfo &CLI,
         }
         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);
+        VAArgEndOffset =
+            alignTo(VA.getLocMemOffset() + VTSize, Align(PtrLenBytes));
+        MemArgEndOffset = VAArgEndOffset;
       }
 
       // Emit the store.
@@ -16503,15 +16529,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 << 0) : (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,
@@ -16535,7 +16565,7 @@ SDValue RISCVTargetLowering::LowerCall(CallLoweringInfo &CLI,
           std::make_pair(RISCVABI::getCheriBoundedArgReg(), MemArgs));
     } else {
       bool ShouldClearArgReg = IsVarArg;
-      if (!ShouldClearArgReg && UseBoundeMemArgsCallee) {
+      if (!ShouldClearArgReg && !UseBoundeMemArgsCallee) {
         auto *G = dyn_cast<GlobalAddressSDNode>(Callee);
         ShouldClearArgReg = !G || !G->getGlobal()->hasInternalLinkage();
       }

llvm/lib/Target/RISCV/RISCVISelLowering.h

@@ -977,6 +977,10 @@ class RISCVTargetLowering : public TargetLowering {
   /// For available scheduling models FDIV + two independent FMULs are much
   /// faster than two FDIVs.
   unsigned combineRepeatedFPDivisors() const override;
+
+  bool functionArgumentNeedsConsecutiveRegisters(
+      Type *Ty, CallingConv::ID CallConv, bool isVarArg,
+      const DataLayout &DL) const override;
 };
 
 namespace RISCV {

llvm/test/CodeGen/RISCV/cheri/bakewell/cheri-bound-varargs-aggregates.ll

@@ -0,0 +1,70 @@
+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py UTC_ARGS: --version 2
+; RUN: llc -mtriple riscv64 -mattr=+zcheripurecap,+zcherihybrid,+cap-mode,+cheri-bounded-vararg,+cheri-bounded-memarg-callee,+cheri-bounded-memarg-caller -target-abi l64pc128d %s -o - | FileCheck %s
+
+define dso_local i32 @g() addrspace(200) {
+; CHECK-LABEL: g:
+; CHECK:       # %bb.0: # %entry
+; CHECK-NEXT:    caddi csp, csp, -48
+; CHECK-NEXT:    .cfi_def_cfa_offset 48
+; CHECK-NEXT:    sc cra, 32(csp) # 16-byte Folded Spill
+; CHECK-NEXT:    .cfi_offset ra, -16
+; CHECK-NEXT:    li a0, 200
+; CHECK-NEXT:    sd a0, 24(csp)
+; CHECK-NEXT:    li a0, 100
+; CHECK-NEXT:    sd a0, 16(csp)
+; CHECK-NEXT:    li a0, 8
+; CHECK-NEXT:    li a1, 32
+; CHECK-NEXT:    scbndsr ct6, csp, a1
+; CHECK-NEXT:    li a1, 1
+; CHECK-NEXT:    li a2, 2
+; CHECK-NEXT:    li a3, 3
+; CHECK-NEXT:    li a4, 4
+; CHECK-NEXT:    li a5, 5
+; CHECK-NEXT:    li a6, 6
+; CHECK-NEXT:    li a7, 7
+; CHECK-NEXT:    sd a0, 0(csp)
+; CHECK-NEXT:    li a0, 0
+; CHECK-NEXT:    call f
+; CHECK-NEXT:    lc cra, 32(csp) # 16-byte Folded Reload
+; CHECK-NEXT:    caddi csp, csp, 48
+; CHECK-NEXT:    ret
+entry:
+  %call = tail call signext i32 (i32, i32, i32, i32, i32, i32, i32, i32, i32, ...) @f(i32 0, i32 1, i32 2, i32 3, i32 4, i32 5, i32 6, i32 7, i32 8, [2 x i64] [i64 100, i64 200])
+  ret i32 %call
+}
+
+declare i32 @f(i32, i32, i32, i32, i32, i32, i32, i32, i32, ...) addrspace(200)
+
+define dso_local i128 @foo() addrspace(200) {
+; CHECK-LABEL: foo:
+; CHECK:       # %bb.0: # %entry
+; CHECK-NEXT:    caddi csp, csp, -48
+; CHECK-NEXT:    .cfi_def_cfa_offset 48
+; CHECK-NEXT:    sc cra, 32(csp) # 16-byte Folded Spill
+; CHECK-NEXT:    .cfi_offset ra, -16
+; CHECK-NEXT:    sd zero, 24(csp)
+; CHECK-NEXT:    li a0, 789
+; CHECK-NEXT:    sd a0, 16(csp)
+; CHECK-NEXT:    sd zero, 8(csp)
+; CHECK-NEXT:    li a1, 456
+; CHECK-NEXT:    li a0, 32
+; CHECK-NEXT:    scbndsr ct6, csp, a0
+; CHECK-NEXT:    caddi ca0, cnull, 1
+; CHECK-NEXT:    li a5, 123
+; CHECK-NEXT:    sd a1, 0(csp)
+; CHECK-NEXT:    li a1, 0
+; CHECK-NEXT:    li a2, 0
+; CHECK-NEXT:    li a3, 0
+; CHECK-NEXT:    li a4, 0
+; CHECK-NEXT:    li a6, 0
+; CHECK-NEXT:    call bar
+; CHECK-NEXT:    lc cra, 32(csp) # 16-byte Folded Reload
+; CHECK-NEXT:    caddi csp, csp, 48
+; CHECK-NEXT:    ret
+entry:
+  %call = call i128 (ptr addrspace(200), i128, i128, i128, ...) @bar(ptr addrspace(200) getelementptr (i8, ptr addrspace(200) null, i64 1), i128 0, i128 0, i128 123, i128 456, i128 789)
+  ret i128 %call
+}
+
+declare i128 @bar(ptr addrspace(200), i128, i128, i128, ...) addrspace(200)
+