[CHERIoT] Pass two-cap-sized return values and arguments in registers

Author: xdoardo

clang/lib/CodeGen/Targets/RISCV.cpp

@@ -35,6 +35,8 @@ class RISCVABIInfo : public DefaultABIInfo {
                                       llvm::Type *&Field2Ty,
                                       CharUnits &Field2Off) const;
 
+  bool shouldPassStructDirectInCapRegisters(uint64_t Size, QualType Ty) const;
+
 public:
   RISCVABIInfo(CodeGen::CodeGenTypes &CGT, unsigned XLen, unsigned FLen,
                bool EABI)
@@ -406,6 +408,20 @@ ABIArgInfo RISCVABIInfo::coerceVLSVector(QualType Ty) const {
   return ABIArgInfo::getDirect(ResType);
 }
 
+bool RISCVABIInfo::shouldPassStructDirectInCapRegisters(uint64_t Size,
+                                                        QualType Ty) const {
+  StringRef TargetABI = getTarget().getABI();
+  bool IsCheriot = TargetABI == "cheriot" || TargetABI == "cheriot-baremetal";
+
+  if (auto *RT = Ty->getAs<RecordType>()) {
+    unsigned MaxCapRegs = IsCheriot ? 2 : 1;
+    return Size == MaxCapRegs * getTarget().getCHERICapabilityWidth() &&
+           getContext().containsCapabilities(RT->getDecl());
+  }
+
+  return false;
+}
+
 ABIArgInfo RISCVABIInfo::classifyArgumentType(QualType Ty, bool IsFixed,
                                               int &ArgGPRsLeft,
                                               int &ArgFPRsLeft) const {
@@ -428,19 +444,17 @@ ABIArgInfo RISCVABIInfo::classifyArgumentType(QualType Ty, bool IsFixed,
   if (isEmptyRecord(getContext(), Ty, true) && Size == 0)
     return ABIArgInfo::getIgnore();
 
-  bool IsSingleCapRecord = false;
-  if (auto *RT = Ty->getAs<RecordType>())
-    IsSingleCapRecord = Size == getTarget().getCHERICapabilityWidth() &&
-                        getContext().containsCapabilities(RT->getDecl());
+  auto CapabilityWidth = getTarget().getCHERICapabilityWidth();
+  bool ForcePassInCapRegs = shouldPassStructDirectInCapRegisters(Size, Ty);
 
-  bool IsCapability = Ty->isCHERICapabilityType(getContext()) ||
-                      IsSingleCapRecord;
+  bool IsSingleCapability = Ty->isCHERICapabilityType(getContext()) ||
+                            (ForcePassInCapRegs && Size == CapabilityWidth);
 
   // Capabilities (including single-capability records, which are treated the
   // same as a single capability) are passed indirectly for hybrid varargs.
   // Anything larger is bigger than 2*XLEN and thus automatically passed
   // indirectly anyway.
-  if (!IsFixed && IsCapability &&
+  if (!IsFixed && IsSingleCapability &&
       !getContext().getTargetInfo().areAllPointersCapabilities()) {
     if (ArgGPRsLeft)
       ArgGPRsLeft -= 1;
@@ -495,9 +509,9 @@ ABIArgInfo RISCVABIInfo::classifyArgumentType(QualType Ty, bool IsFixed,
   // TODO: Pairs involving capabilities should be passed in registers too like
   // int/fp pairs (requires thought for fp+cap when out of FPRs).
   int NeededArgGPRs = 1;
-  if (!IsCapability && !IsFixed && NeededAlign == 2 * XLen)
+  if (!IsSingleCapability && !IsFixed && NeededAlign == 2 * XLen)
     NeededArgGPRs = 2 + (EABI && XLen == 32 ? 0 : (ArgGPRsLeft % 2));
-  else if (!IsCapability && Size > XLen && Size <= 2 * XLen)
+  else if (!IsSingleCapability && Size > XLen && Size <= 2 * XLen)
     NeededArgGPRs = 2;
 
   if (NeededArgGPRs > ArgGPRsLeft) {
@@ -528,7 +542,7 @@ ABIArgInfo RISCVABIInfo::classifyArgumentType(QualType Ty, bool IsFixed,
     return ABIArgInfo::getDirect();
   }
 
-  if (IsSingleCapRecord)
+  if (ForcePassInCapRegs)
     return ABIArgInfo::getDirect();
 
   if (const VectorType *VT = Ty->getAs<VectorType>())
@@ -593,21 +607,19 @@ RValue RISCVABIInfo::EmitVAArg(CodeGenFunction &CGF, Address VAListAddr,
   if (EABI && XLen == 32 && !IsCheriot)
     TInfo.Align = std::min(TInfo.Align, CharUnits::fromQuantity(4));
 
-  bool IsSingleCapRecord = false;
-  CharUnits CapabilityWidth =
-    CharUnits::fromQuantity(getTarget().getCHERICapabilityWidth() / 8);
-  if (const auto *RT = Ty->getAs<RecordType>())
-    IsSingleCapRecord = TInfo.Width == CapabilityWidth &&
-                        getContext().containsCapabilities(RT->getDecl());
+  uint64_t Size = TInfo.Width.getQuantity() * 8;
+  auto CapabilityWidth = getTarget().getCHERICapabilityWidth();
+  bool ForcePassInCapRegs = shouldPassStructDirectInCapRegisters(Size, Ty);
 
-  bool IsCapability = Ty->isCHERICapabilityType(getContext()) ||
-                      IsSingleCapRecord;
+  bool IsSingleCapability = Ty->isCHERICapabilityType(getContext()) ||
+                            (ForcePassInCapRegs && Size == CapabilityWidth);
 
   // Arguments bigger than 2*Xlen bytes are passed indirectly, as are
   // capabilities for the hybrid ABI.
-  bool IsIndirect = TInfo.Width > 2 * SlotSize ||
-    (!getContext().getTargetInfo().areAllPointersCapabilities() &&
-     IsCapability);
+  bool IsIndirect =
+      TInfo.Width > 2 * SlotSize ||
+      (!getContext().getTargetInfo().areAllPointersCapabilities() &&
+       IsSingleCapability);
 
   return emitVoidPtrVAArg(CGF, VAListAddr, Ty, IsIndirect, TInfo, SlotSize,
                           /*AllowHigherAlign=*/true, Slot);