[InstSimplify] Add folds for @llvm.cheri.cap.high.get

As part of this refactor some of the CHERI get intrinsics folding code. This refactoring introduces a minor change to getoffset folding as well: we no longer fold arbitrary getoffset(setoffset(A, B)) -> B since the setoffset could have changed the bounds interpretation and therefore also the result. This optimization is unlikely to make a difference in practise and is somewhat incorrect so I believe dropping it makes sense.
CTSRD-CHERI · Feb 13, 2024 · fcc87ce · fcc87ce
1 parent 61a3130
commit fcc87ce
Show file tree

Hide file tree

Showing 7 changed files with 87 additions and 34 deletions.
diff --git a/clang/test/CodeGen/cheri/vaddr-mode.c b/clang/test/CodeGen/cheri/vaddr-mode.c
@@ -48,7 +48,9 @@ long write_uintcap(__uintcap_t *cap) {
 // OFFSET-NEXT:  entry:
 // OFFSET-NEXT:    [[TMP0:%.*]] = tail call i64 @llvm.cheri.cap.offset.get.i64(ptr addrspace(200) [[CAP:%.*]])
 // OFFSET-NEXT:    [[AND:%.*]] = and i64 [[TMP0]], 3
-// OFFSET-NEXT:    ret i64 [[AND]]
+// OFFSET-NEXT:    [[TMP1:%.*]] = tail call i8 addrspace(200)* @llvm.cheri.cap.offset.set.i64(i8 addrspace(200)* [[CAP]], i64 [[AND]])
+// OFFSET-NEXT:    [[TMP2:%.*]] = tail call i64 @llvm.cheri.cap.offset.get.i64(i8 addrspace(200)* [[TMP1]])
+// OFFSET-NEXT:    ret i64 [[TMP2]]
 //
 long get_low_bits(__uintcap_t cap) {
   return cap & 3;
@@ -126,7 +128,9 @@ __uintcap_t xor_uintcap(__uintcap_t cap, __uintcap_t cap2) {
 // OFFSET-NEXT:    [[TMP0:%.*]] = tail call i64 @llvm.cheri.cap.offset.get.i64(ptr addrspace(200) [[CAP:%.*]])
 // OFFSET-NEXT:    [[TMP1:%.*]] = tail call i64 @llvm.cheri.cap.offset.get.i64(ptr addrspace(200) [[CAP2:%.*]])
 // OFFSET-NEXT:    [[XOR:%.*]] = xor i64 [[TMP1]], [[TMP0]]
-// OFFSET-NEXT:    ret i64 [[XOR]]
+// OFFSET-NEXT:    [[TMP2:%.*]] = tail call i8 addrspace(200)* @llvm.cheri.cap.offset.set.i64(i8 addrspace(200)* [[CAP]], i64 [[XOR]])
+// OFFSET-NEXT:    [[TMP3:%.*]] = tail call i64 @llvm.cheri.cap.offset.get.i64(i8 addrspace(200)* [[TMP2]])
+// OFFSET-NEXT:    ret i64 [[TMP3]]
 //
 long xor_uintcap_return_long(__uintcap_t cap, __uintcap_t cap2) {
   return cap ^ cap2;
@@ -160,7 +164,9 @@ __uintcap_t modulo_return_uintcap(__uintcap_t cap) {
 // OFFSET-NEXT:  entry:
 // OFFSET-NEXT:    [[TMP0:%.*]] = tail call i64 @llvm.cheri.cap.offset.get.i64(ptr addrspace(200) [[CAP:%.*]])
 // OFFSET-NEXT:    [[REM:%.*]] = and i64 [[TMP0]], 31
-// OFFSET-NEXT:    ret i64 [[REM]]
+// OFFSET-NEXT:    [[TMP1:%.*]] = tail call i8 addrspace(200)* @llvm.cheri.cap.offset.set.i64(i8 addrspace(200)* [[CAP]], i64 [[REM]])
+// OFFSET-NEXT:    [[TMP2:%.*]] = tail call i64 @llvm.cheri.cap.offset.get.i64(i8 addrspace(200)* [[TMP1]])
+// OFFSET-NEXT:    ret i64 [[TMP2]]
 //
 long modulo_return_long(__uintcap_t cap) {
   return cap % 32;

diff --git a/llvm/include/llvm/IR/CheriIntrinsics.h b/llvm/include/llvm/IR/CheriIntrinsics.h
@@ -30,6 +30,8 @@ correspondingSetIntrinsic(Intrinsic::ID GetIntrin) {
     return Intrinsic::cheri_cap_address_set;
   case Intrinsic::cheri_cap_flags_get:
     return Intrinsic::cheri_cap_flags_set;
+  case Intrinsic::cheri_cap_high_get:
+    return Intrinsic::cheri_cap_high_set;
   default:
     llvm_unreachable("No matching set intrinsic");
   }
@@ -43,6 +45,8 @@ correspondingGetIntrinsic(Intrinsic::ID SetIntrin) {
     return Intrinsic::cheri_cap_address_get;
   case Intrinsic::cheri_cap_flags_set:
     return Intrinsic::cheri_cap_flags_get;
+  case Intrinsic::cheri_cap_high_set:
+    return Intrinsic::cheri_cap_high_get;
   default:
     llvm_unreachable("No matching get intrinsic");
   }

diff --git a/llvm/lib/Analysis/InstructionSimplify.cpp b/llvm/lib/Analysis/InstructionSimplify.cpp
@@ -5716,15 +5716,47 @@ stripAndAccumulateGEPsAndPointerCastsSameRepr(Value *V, const DataLayout &DL,
   return Result;
 }
 
-template <Intrinsic::ID Intrin, Intrinsic::ID SetIntrin>
+// Common helper to fold llvm.cheri.get.foo intrinsics
+template <Intrinsic::ID Intrin>
+static Value *foldCheriGetSetPair(Value *V, Type *ResultTy,
+                                  const DataLayout &DL) {
+  constexpr Intrinsic::ID SetIntrin = cheri::correspondingSetIntrinsic(Intrin);
+  // getFoo(setFoo(A, B)) -> B
+  // This is true for address and high, but not necessarily true for flags and
+  // offset since setflags truncates the value and setoffset may change the
+  // bounds representation (and thus the offset) if it is significantly OOB.
+  constexpr bool SetUpdatesAllGetBits =
+      SetIntrin != Intrinsic::cheri_cap_offset_set &&
+      SetIntrin != Intrinsic::cheri_cap_flags_set;
+  Value *IntrinArg = nullptr;
+  if (SetUpdatesAllGetBits &&
+      match(V, m_Intrinsic<SetIntrin>(m_Value(), m_Value(IntrinArg)))) {
+    return IntrinArg;
+  }
+  // All intrinsics we call this function for return zero for NULL inputs.
+  // In fact, we can extend this to any input derived from NULL where the only
+  // operations that have been address manipulations (unless we are retrieving
+  // a result depending on the address), since those cannot affect non-bounds
+  // capability metadata.
+  constexpr bool CanIgnoreAddressManipulation =
+      Intrin != Intrinsic::cheri_cap_address_get &&
+      Intrin != Intrinsic::cheri_cap_offset_get;
+  Value *Base = CanIgnoreAddressManipulation
+                    ? getBasePtrIgnoringCapabilityAddressManipulation(V, DL)
+                    : V->stripPointerCastsSameRepresentation();
+  if (isa<ConstantPointerNull>(Base)) {
+    return Constant::getNullValue(ResultTy);
+  }
+  return nullptr;
+}
+
+template <Intrinsic::ID Intrin>
 static Value *inferCapabilityOffsetOrAddr(Value *V, Type *ResultTy,
                                           const DataLayout &DL) {
+  if (Value *Ret = foldCheriGetSetPair<Intrin>(V, ResultTy, DL))
+    return Ret;
   // Try to infer the offset/address from a prior setoffset/setaddr value
   Value *IntrinArg = nullptr;
-  // getaddr(setaddr(A, B)) -> B and getoffset(setoffset(A, B)) -> B
-  if (match(V, m_Intrinsic<SetIntrin>(m_Value(), m_Value(IntrinArg)))) {
-    return IntrinArg;
-  }
   // get{addr,offset}(setaddr(NULL, B)) -> B
   if (match(V, m_Intrinsic<Intrinsic::cheri_cap_address_set>(
                    m_Zero(), m_Value(IntrinArg)))) {
@@ -5765,6 +5797,7 @@ static Value *inferCapabilityOffsetOrAddr(Value *V, Type *ResultTy,
   // We can also fold chains of constant GEPS:
   // For example: getoffset(GEP(setoffset(A, Const1), 100) -> Const1 + 100
   ConstantInt *ConstSetArg;
+  constexpr Intrinsic::ID SetIntrin = cheri::correspondingSetIntrinsic(Intrin);
   if (match(BasePtr,
             m_Intrinsic<SetIntrin>(m_Value(), m_ConstantInt(ConstSetArg)))) {
     return ConstantInt::get(ResultTy, ConstSetArg->getValue() + OffsetAPInt);
@@ -5804,7 +5837,6 @@ static Value *simplifyUnaryIntrinsic(Function *F, Value *Op0,
       return Constant::getNullValue(F->getReturnType());
     break;
   case Intrinsic::cheri_cap_sealed_get:
-  case Intrinsic::cheri_cap_flags_get:
   case Intrinsic::cheri_cap_base_get:
     // Values derived from NULL and where the only operations that have been
     // applied are address manipulations, always have the following properties:
@@ -5842,6 +5874,18 @@ static Value *simplifyUnaryIntrinsic(Function *F, Value *Op0,
             Op0, F->getReturnType(), Q.DL))
       return V;
     break;
+  case Intrinsic::cheri_cap_flags_get: {
+    if (auto *Ret = foldCheriGetSetPair<Intrinsic::cheri_cap_flags_get>(
+            Op0, F->getReturnType(), Q.DL))
+      return Ret;
+    break;
+  }
+  case Intrinsic::cheri_cap_high_get: {
+    if (auto *Ret = foldCheriGetSetPair<Intrinsic::cheri_cap_high_get>(
+            Op0, F->getReturnType(), Q.DL))
+      return Ret;
+    break;
+  }
   case Intrinsic::cheri_cap_offset_get:
     if (auto *V = inferCapabilityOffsetOrAddr<Intrinsic::cheri_cap_offset_get>(
             Op0, F->getReturnType(), Q.DL))

diff --git a/llvm/test/CodeGen/CHERI-Generic/MIPS/cheri-intrinsics-folding-broken-module-regression.ll b/llvm/test/CodeGen/CHERI-Generic/MIPS/cheri-intrinsics-folding-broken-module-regression.ll
@@ -23,20 +23,20 @@ define void @g(i64 %x, i64 %y) addrspace(200) nounwind {
 ; ASM-NEXT:    lui $1, %pcrel_hi(_CHERI_CAPABILITY_TABLE_-8)
 ; ASM-NEXT:    daddiu $1, $1, %pcrel_lo(_CHERI_CAPABILITY_TABLE_-4)
 ; ASM-NEXT:    cgetpccincoffset $c1, $1
-; ASM-NEXT:    daddu $2, $5, $4
 ; ASM-NEXT:    clcbi $c2, %captab20(d)($c1)
 ; ASM-NEXT:    clcbi $c1, %captab20(e)($c1)
 ; ASM-NEXT:    cgetoffset $1, $c2
-; ASM-NEXT:    daddu $1, $2, $1
+; ASM-NEXT:    daddu $1, $1, $4
 ; ASM-NEXT:    csetoffset $c2, $c2, $1
+; ASM-NEXT:    cincoffset $c2, $c2, $5
 ; ASM-NEXT:    cjr $c17
 ; ASM-NEXT:    csc $c2, $zero, 0($c1)
 ; CHECK-LABEL: define {{[^@]+}}@g
 ; CHECK-SAME: (i64 [[X:%.*]], i64 [[Y:%.*]]) addrspace(200) #[[ATTR0:[0-9]+]] {
 ; CHECK-NEXT:    [[TMP3:%.*]] = call i64 @llvm.cheri.cap.offset.get.i64(i8 addrspace(200)* bitcast (i64 addrspace(200)* @d to i8 addrspace(200)*))
 ; CHECK-NEXT:    [[ADD:%.*]] = add i64 [[TMP3]], [[X]]
-; CHECK-NEXT:    [[ADD1:%.*]] = add i64 [[ADD]], [[Y]]
-; CHECK-NEXT:    [[TMP11:%.*]] = call i8 addrspace(200)* @llvm.cheri.cap.offset.set.i64(i8 addrspace(200)* bitcast (i64 addrspace(200)* @d to i8 addrspace(200)*), i64 [[ADD1]])
+; CHECK-NEXT:    [[TMP5:%.*]] = call i8 addrspace(200)* @llvm.cheri.cap.offset.set.i64(i8 addrspace(200)* bitcast (i64 addrspace(200)* @d to i8 addrspace(200)*), i64 [[ADD]])
+; CHECK-NEXT:    [[TMP11:%.*]] = getelementptr i8, i8 addrspace(200)* [[TMP5]], i64 [[Y]]
 ; CHECK-NEXT:    store i8 addrspace(200)* [[TMP11]], i8 addrspace(200)* addrspace(200)* @e, align 32
 ; CHECK-NEXT:    ret void
 ;

diff --git a/llvm/test/CodeGen/CHERI-Generic/RISCV32/cheri-intrinsics-folding-broken-module-regression.ll b/llvm/test/CodeGen/CHERI-Generic/RISCV32/cheri-intrinsics-folding-broken-module-regression.ll
@@ -27,17 +27,17 @@ define void @g(i32 %x, i32 %y) addrspace(200) nounwind {
 ; ASM-NEXT:  .LBB0_2: # Label of block must be emitted
 ; ASM-NEXT:    auipcc ca4, %captab_pcrel_hi(e)
 ; ASM-NEXT:    clc ca4, %pcrel_lo(.LBB0_2)(ca4)
-; ASM-NEXT:    add a0, a1, a0
-; ASM-NEXT:    add a0, a0, a3
+; ASM-NEXT:    add a0, a3, a0
 ; ASM-NEXT:    csetoffset ca0, ca2, a0
+; ASM-NEXT:    cincoffset ca0, ca0, a1
 ; ASM-NEXT:    csc ca0, 0(ca4)
 ; ASM-NEXT:    cret
 ; CHECK-LABEL: define {{[^@]+}}@g
 ; CHECK-SAME: (i32 [[X:%.*]], i32 [[Y:%.*]]) addrspace(200) #[[ATTR0:[0-9]+]] {
 ; CHECK-NEXT:    [[TMP3:%.*]] = call i32 @llvm.cheri.cap.offset.get.i32(i8 addrspace(200)* bitcast (i32 addrspace(200)* @d to i8 addrspace(200)*))
 ; CHECK-NEXT:    [[ADD:%.*]] = add i32 [[TMP3]], [[X]]
-; CHECK-NEXT:    [[ADD1:%.*]] = add i32 [[ADD]], [[Y]]
-; CHECK-NEXT:    [[TMP11:%.*]] = call i8 addrspace(200)* @llvm.cheri.cap.offset.set.i32(i8 addrspace(200)* bitcast (i32 addrspace(200)* @d to i8 addrspace(200)*), i32 [[ADD1]])
+; CHECK-NEXT:    [[TMP5:%.*]] = call i8 addrspace(200)* @llvm.cheri.cap.offset.set.i32(i8 addrspace(200)* bitcast (i32 addrspace(200)* @d to i8 addrspace(200)*), i32 [[ADD]])
+; CHECK-NEXT:    [[TMP11:%.*]] = getelementptr i8, i8 addrspace(200)* [[TMP5]], i32 [[Y]]
 ; CHECK-NEXT:    store i8 addrspace(200)* [[TMP11]], i8 addrspace(200)* addrspace(200)* @e, align 32
 ; CHECK-NEXT:    ret void
 ;

diff --git a/llvm/test/CodeGen/CHERI-Generic/RISCV64/cheri-intrinsics-folding-broken-module-regression.ll b/llvm/test/CodeGen/CHERI-Generic/RISCV64/cheri-intrinsics-folding-broken-module-regression.ll
@@ -23,21 +23,21 @@ define void @g(i64 %x, i64 %y) addrspace(200) nounwind {
 ; ASM-NEXT:  .LBB0_1: # Label of block must be emitted
 ; ASM-NEXT:    auipcc ca2, %captab_pcrel_hi(d)
 ; ASM-NEXT:    clc ca2, %pcrel_lo(.LBB0_1)(ca2)
-; ASM-NEXT:    add a0, a1, a0
-; ASM-NEXT:  .LBB0_2: # Label of block must be emitted
-; ASM-NEXT:    auipcc ca1, %captab_pcrel_hi(e)
-; ASM-NEXT:    clc ca1, %pcrel_lo(.LBB0_2)(ca1)
 ; ASM-NEXT:    cgetoffset a3, ca2
-; ASM-NEXT:    add a0, a0, a3
+; ASM-NEXT:  .LBB0_2: # Label of block must be emitted
+; ASM-NEXT:    auipcc ca4, %captab_pcrel_hi(e)
+; ASM-NEXT:    clc ca4, %pcrel_lo(.LBB0_2)(ca4)
+; ASM-NEXT:    add a0, a3, a0
 ; ASM-NEXT:    csetoffset ca0, ca2, a0
-; ASM-NEXT:    csc ca0, 0(ca1)
+; ASM-NEXT:    cincoffset ca0, ca0, a1
+; ASM-NEXT:    csc ca0, 0(ca4)
 ; ASM-NEXT:    cret
 ; CHECK-LABEL: define {{[^@]+}}@g
 ; CHECK-SAME: (i64 [[X:%.*]], i64 [[Y:%.*]]) addrspace(200) #[[ATTR0:[0-9]+]] {
 ; CHECK-NEXT:    [[TMP3:%.*]] = call i64 @llvm.cheri.cap.offset.get.i64(i8 addrspace(200)* bitcast (i64 addrspace(200)* @d to i8 addrspace(200)*))
 ; CHECK-NEXT:    [[ADD:%.*]] = add i64 [[TMP3]], [[X]]
-; CHECK-NEXT:    [[ADD1:%.*]] = add i64 [[ADD]], [[Y]]
-; CHECK-NEXT:    [[TMP11:%.*]] = call i8 addrspace(200)* @llvm.cheri.cap.offset.set.i64(i8 addrspace(200)* bitcast (i64 addrspace(200)* @d to i8 addrspace(200)*), i64 [[ADD1]])
+; CHECK-NEXT:    [[TMP5:%.*]] = call i8 addrspace(200)* @llvm.cheri.cap.offset.set.i64(i8 addrspace(200)* bitcast (i64 addrspace(200)* @d to i8 addrspace(200)*), i64 [[ADD]])
+; CHECK-NEXT:    [[TMP11:%.*]] = getelementptr i8, i8 addrspace(200)* [[TMP5]], i64 [[Y]]
 ; CHECK-NEXT:    store i8 addrspace(200)* [[TMP11]], i8 addrspace(200)* addrspace(200)* @e, align 32
 ; CHECK-NEXT:    ret void
 ;

diff --git a/llvm/test/Transforms/InstSimplify/cheri-intrinsics-get-set.ll b/llvm/test/Transforms/InstSimplify/cheri-intrinsics-get-set.ll
@@ -12,6 +12,8 @@ declare i8 addrspace(200)* @llvm.cheri.cap.INTRINSIC.set.i64(i8 addrspace(200)*,
 
 
 ;; This is a no-op and should be folded to ret %arg
+;; NB: The exception here is @llvm.cheri.cap.high.set since it is tag-clearing
+;; and therefore could only be used on known-untagged values.
 define i8 addrspace(200)* @fold_set_of_get(i8 addrspace(200)* %arg) nounwind {
 ; ADDRESS-LABEL: define {{[^@]+}}@fold_set_of_get
 ; ADDRESS-SAME: (i8 addrspace(200)* [[ARG:%.*]]) addrspace(200) #[[ATTR0:[0-9]+]] {
@@ -85,13 +87,13 @@ define i64 @fold_get_of_set(i8 addrspace(200)* %arg, i64 %value) nounwind {
 ;
 ; HIGH-LABEL: define {{[^@]+}}@fold_get_of_set
 ; HIGH-SAME: (i8 addrspace(200)* [[ARG:%.*]], i64 [[VALUE:%.*]]) addrspace(200) #[[ATTR0]] {
-; HIGH-NEXT:    [[MODIFIED:%.*]] = tail call i8 addrspace(200)* @llvm.cheri.cap.high.set.i64(i8 addrspace(200)* [[ARG]], i64 [[VALUE]])
-; HIGH-NEXT:    [[RET:%.*]] = tail call i64 @llvm.cheri.cap.high.get.i64(i8 addrspace(200)* [[MODIFIED]])
-; HIGH-NEXT:    ret i64 [[RET]]
+; HIGH-NEXT:    ret i64 [[VALUE]]
 ;
 ; OFFSET-LABEL: define {{[^@]+}}@fold_get_of_set
 ; OFFSET-SAME: (i8 addrspace(200)* [[ARG:%.*]], i64 [[VALUE:%.*]]) addrspace(200) #[[ATTR0]] {
-; OFFSET-NEXT:    ret i64 [[VALUE]]
+; OFFSET-NEXT:    [[MODIFIED:%.*]] = tail call i8 addrspace(200)* @llvm.cheri.cap.offset.set.i64(i8 addrspace(200)* [[ARG]], i64 [[VALUE]])
+; OFFSET-NEXT:    [[RET:%.*]] = tail call i64 @llvm.cheri.cap.offset.get.i64(i8 addrspace(200)* [[MODIFIED]])
+; OFFSET-NEXT:    ret i64 [[RET]]
 ;
   %modified = tail call i8 addrspace(200)* @llvm.cheri.cap.INTRINSIC.set.i64(i8 addrspace(200)* %arg, i64 %value)
   %ret = tail call i64 @llvm.cheri.cap.INTRINSIC.get.i64(i8 addrspace(200)* %modified)
@@ -136,8 +138,7 @@ define i64 @fold_get_on_null() nounwind {
 ;
 ; HIGH-LABEL: define {{[^@]+}}@fold_get_on_null
 ; HIGH-SAME: () addrspace(200) #[[ATTR0]] {
-; HIGH-NEXT:    [[RET:%.*]] = tail call i64 @llvm.cheri.cap.high.get.i64(i8 addrspace(200)* null)
-; HIGH-NEXT:    ret i64 [[RET]]
+; HIGH-NEXT:    ret i64 0
 ;
 ; OFFSET-LABEL: define {{[^@]+}}@fold_get_on_null
 ; OFFSET-SAME: () addrspace(200) #[[ATTR0]] {
@@ -158,9 +159,7 @@ define i64 @fold_get_on_null_with_gep(i64 %value, i64 %gepoff) nounwind {
 ;
 ; HIGH-LABEL: define {{[^@]+}}@fold_get_on_null_with_gep
 ; HIGH-SAME: (i64 [[VALUE:%.*]], i64 [[GEPOFF:%.*]]) addrspace(200) #[[ATTR0]] {
-; HIGH-NEXT:    [[TMP:%.*]] = getelementptr i8, i8 addrspace(200)* null, i64 [[GEPOFF]]
-; HIGH-NEXT:    [[RET:%.*]] = tail call i64 @llvm.cheri.cap.high.get.i64(i8 addrspace(200)* [[TMP]])
-; HIGH-NEXT:    ret i64 [[RET]]
+; HIGH-NEXT:    ret i64 0
 ;
 ; OFFSET-LABEL: define {{[^@]+}}@fold_get_on_null_with_gep
 ; OFFSET-SAME: (i64 [[VALUE:%.*]], i64 [[GEPOFF:%.*]]) addrspace(200) #[[ATTR0]] {