[FuzzMutate] Properly handle intrinsics and avoid illegal code genertion

llvm/lib/FuzzMutate/IRMutator.cpp

@@ -20,6 +20,8 @@
 #include "llvm/IR/Function.h"
 #include "llvm/IR/InstIterator.h"
 #include "llvm/IR/Instructions.h"
+#include "llvm/IR/IntrinsicInst.h"
+#include "llvm/IR/IntrinsicsAMDGPU.h"
 #include "llvm/IR/Module.h"
 #include "llvm/IR/Operator.h"
 #include "llvm/IR/PassInstrumentation.h"
@@ -115,9 +117,41 @@ InjectorIRStrategy::chooseOperation(Value *Src, RandomIRBuilder &IB) {
   return *RS;
 }
 
+static inline Instruction *getEffectiveTerminator(BasicBlock &BB) {
+  if (Instruction *I = BB.getTerminatingMustTailCall()) {
+    return I;
+  } else {
+    // Certain intrinsics, such as @llvm.amdgcn.cs.chain, must be immediately
+    // followed by an unreachable instruction..
+    if (UnreachableInst *UI = dyn_cast<UnreachableInst>(BB.getTerminator())) {
+      if (IntrinsicInst *II = dyn_cast<IntrinsicInst>(UI->getPrevNode())) {
+        return II;
+      }
+    }
+  }
+
+  return BB.getTerminator();
+}
+
+static inline BasicBlock::iterator getEndIterator(BasicBlock &BB) {
+  auto End = BB.end();
+
+  if (BB.empty()) {
+    return End;
+  }
+
+  Instruction *EffectiveTerminator = getEffectiveTerminator(BB);
+  if (EffectiveTerminator != BB.getTerminator()) {
+    // Adjust range for special cases such as tail call.
+    End = std::prev(BB.end());
+  }
+
+  return End;
+}
+
 static inline iterator_range<BasicBlock::iterator>
 getInsertionRange(BasicBlock &BB) {
-  auto End = BB.getTerminatingMustTailCall() ? std::prev(BB.end()) : BB.end();
+  auto End = getEndIterator(BB);
   return make_range(BB.getFirstInsertionPt(), End);
 }
 
@@ -409,6 +443,12 @@ static bool isUnsupportedFunction(Function *F) {
     return true;
   }
 
+  // This intrinsic has specific requirements for its parameters and the caller
+  // must adhere to certain calling conventions.
+  if (F->isIntrinsic() && F->getIntrinsicID() == Intrinsic::amdgcn_cs_chain) {
+    return true;
+  }
+
   return false;
 }
 
@@ -641,8 +681,9 @@ void ShuffleBlockStrategy::mutate(BasicBlock &BB, RandomIRBuilder &IB) {
   std::map<size_t, Instruction *> AliveInsts;
   std::map<Instruction *, size_t> AliveInstsLookup;
   size_t InsertIdx = 0;
-  for (auto &I : make_early_inc_range(make_range(
-           BB.getFirstInsertionPt(), BB.getTerminator()->getIterator()))) {
+  for (auto &I : make_early_inc_range(
+           make_range(BB.getFirstInsertionPt(),
+                      getEffectiveTerminator(BB)->getIterator()))) {
     // First gather all instructions that can be shuffled. Don't take
     // terminator.
     AliveInsts.insert({InsertIdx, &I});
@@ -702,7 +743,7 @@ void ShuffleBlockStrategy::mutate(BasicBlock &BB, RandomIRBuilder &IB) {
     }
   }
 
-  Instruction *Terminator = BB.getTerminator();
+  Instruction *Terminator = getEffectiveTerminator(BB);
   // Then put instructions back.
   for (Instruction *I : Insts) {
     I->insertBefore(Terminator->getIterator());

llvm/unittests/FuzzMutate/StrategiesTest.cpp

@@ -745,4 +745,22 @@ TEST(AllStrategies, SkipEHPad) {
   mutateAndVerifyModule<InjectorIRStrategy>(Source);
   mutateAndVerifyModule<InstModificationIRStrategy>(Source);
 }
+
+TEST(AllStrategies, SpecialTerminator) {
+  StringRef Source = "\n\
+    declare amdgpu_cs_chain void @callee(<3 x i32> inreg, { i32, ptr addrspace(5), i32, i32 })\n\
+    define amdgpu_cs_chain void @chain_to_chain(<3 x i32> inreg %sgpr, { i32, ptr addrspace(5), i32, i32 } %vgpr) {\n\
+      call void(ptr, i64, <3 x i32>, { i32, ptr addrspace(5), i32, i32 }, i32, ...) @llvm.amdgcn.cs.chain(ptr @callee, i64 -1, <3 x i32> inreg %sgpr, { i32, ptr addrspace(5), i32, i32 } %vgpr, i32 0) \n\
+      unreachable\n\
+    }\n\
+  ";
+  mutateAndVerifyModule<InjectorIRStrategy>(Source);
+  mutateAndVerifyModule<InsertCFGStrategy>(Source);
+  mutateAndVerifyModule<InsertFunctionStrategy>(Source);
+  mutateAndVerifyModule<InsertPHIStrategy>(Source);
+  mutateAndVerifyModule<InstModificationIRStrategy>(Source);
+  mutateAndVerifyModule<ShuffleBlockStrategy>(Source);
+  mutateAndVerifyModule<SinkInstructionStrategy>(Source);
+}
+
 } // namespace