[CHERI] Align aggregate types upto stack slot

Author: veselypeta

llvm/lib/Target/RISCV/RISCVISelLowering.cpp

@@ -24,6 +24,7 @@
 #include "llvm/ADT/Statistic.h"
 #include "llvm/Analysis/MemoryLocation.h"
 #include "llvm/Analysis/VectorUtils.h"
+#include "llvm/CodeGen/Analysis.h"
 #include "llvm/CodeGen/MachineFrameInfo.h"
 #include "llvm/CodeGen/MachineFunction.h"
 #include "llvm/CodeGen/MachineInstrBuilder.h"
@@ -2240,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) &&
@@ -15361,6 +15378,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.
@@ -15440,15 +15488,9 @@ 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
-          : (AllocateSlot
+          : (IsBoundedVarArgs
                  ? State.AllocateStack(SlotSize, Align(SlotSize))
                  : State.AllocateStack(StoreSizeBytes, StackAlign));
 
@@ -16454,22 +16496,23 @@ SDValue RISCVTargetLowering::LowerCall(CallLoweringInfo &CLI,
       SDValue Address =
           DAG.getPointerAdd(DL, StackPtr, VA.getLocMemOffset());
 
-      if (UseBoundeMemArgsCaller && Outs[i].IsFixed) {
+      if (UseBoundeMemArgsCaller) {
         if (FirstArgAddr == SDValue()) {
           FirstArgAddr = Address;
           MemArgStartOffset = VA.getLocMemOffset();
         }
         unsigned VTSize = VA.getValVT().getSizeInBits() / 8;
         MemArgEndOffset = VA.getLocMemOffset() + VTSize;
       }
+
       if (UseBoundedVarArgs && !Outs[i].IsFixed) {
         if (FirstVAAddr == SDValue()) {
           FirstVAAddr = Address;
           VAArgStartOffset = VA.getLocMemOffset();
         }
-        Align OffsetAlign = Align(PtrLenBytes);
         unsigned VTSize = VA.getValVT().getSizeInBits() / 8;
-        VAArgEndOffset = alignTo(VA.getLocMemOffset() + VTSize, OffsetAlign);
+        VAArgEndOffset =
+            alignTo(VA.getLocMemOffset() + VTSize, Align(PtrLenBytes));
         MemArgEndOffset = VAArgEndOffset;
       }
 

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 ct1, 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 ct1, 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)
+