[FuzzMutate] New strategy `ShuffleBlockStrategy`

`ShuffleBlockStrategy` will shuffle the instructions in a basic block without breaking the dependency of instructions.
It is implemented as a topological sort, only we randomly select instructions with no dependency.

Reviewed By: arsenm

Differential Revision: https://reviews.llvm.org/D138339

llvm/include/llvm/FuzzMutate/IRMutator.h

@@ -118,6 +118,18 @@ public:
   void mutate(Instruction &Inst, RandomIRBuilder &IB) override;
 };
 
+/// Strategy to randomly select a block and shuffle the operations without
+/// affecting data dependency.
+class ShuffleBlockStrategy : public IRMutationStrategy {
+public:
+  uint64_t getWeight(size_t CurrentSize, size_t MaxSize,
+                     uint64_t CurrentWeight) override {
+    return 2;
+  }
+
+  void mutate(BasicBlock &BB, RandomIRBuilder &IB) override;
+};
+
 /// Fuzzer friendly interface for the llvm bitcode parser.
 ///
 /// \param Data Bitcode we are going to parse

llvm/lib/FuzzMutate/IRMutator.cpp

@@ -8,6 +8,8 @@
 
 #include "llvm/FuzzMutate/IRMutator.h"
 #include "llvm/ADT/Optional.h"
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/SmallSet.h"
 #include "llvm/Analysis/TargetLibraryInfo.h"
 #include "llvm/Bitcode/BitcodeReader.h"
 #include "llvm/Bitcode/BitcodeWriter.h"
@@ -297,6 +299,68 @@ void InstModificationIRStrategy::mutate(Instruction &Inst,
     RS.getSelection()();
 }
 
+void ShuffleBlockStrategy::mutate(BasicBlock &BB, RandomIRBuilder &IB) {
+
+  SmallPtrSet<Instruction *, 8> AliveInsts;
+  for (auto &I : make_early_inc_range(make_range(
+           BB.getFirstInsertionPt(), BB.getTerminator()->getIterator()))) {
+    // First gather all instructions that can be shuffled. Don't take
+    // terminator.
+    AliveInsts.insert(&I);
+    // Then remove these instructions from the block
+    I.removeFromParent();
+  }
+
+  // Shuffle these instructions using topological sort.
+  // Returns true if all current instruction's dependencies in this block have
+  // been shuffled. If so, this instruction can be shuffled too.
+  auto hasAliveParent = [&AliveInsts](Instruction *I) {
+    for (Value *O : I->operands()) {
+      Instruction *P = dyn_cast<Instruction>(O);
+      if (P && AliveInsts.count(P))
+        return true;
+    }
+    return false;
+  };
+  // Get all alive instructions that depend on the current instruction.
+  auto getAliveChildren = [&AliveInsts](Instruction *I) {
+    SmallPtrSet<Instruction *, 4> Children;
+    for (Value *U : I->users()) {
+      Instruction *P = dyn_cast<Instruction>(U);
+      if (P && AliveInsts.count(P))
+        Children.insert(P);
+    }
+    return Children;
+  };
+  SmallPtrSet<Instruction *, 8> Roots;
+  SmallVector<Instruction *, 8> Insts;
+  for (Instruction *I : AliveInsts) {
+    if (!hasAliveParent(I))
+      Roots.insert(I);
+  }
+  // Topological sort by randomly selecting a node without a parent, or root.
+  while (!Roots.empty()) {
+    auto RS = makeSampler<Instruction *>(IB.Rand);
+    for (Instruction *Root : Roots)
+      RS.sample(Root, 1);
+    Instruction *Root = RS.getSelection();
+    Roots.erase(Root);
+    AliveInsts.erase(Root);
+    Insts.push_back(Root);
+    for (Instruction *Child : getAliveChildren(Root)) {
+      if (!hasAliveParent(Child)) {
+        Roots.insert(Child);
+      }
+    }
+  }
+
+  Instruction *Terminator = BB.getTerminator();
+  // Then put instructions back.
+  for (Instruction *I : Insts) {
+    I->insertBefore(Terminator);
+  }
+}
+
 std::unique_ptr<Module> llvm::parseModule(const uint8_t *Data, size_t Size,
                                           LLVMContext &Context) {
 

llvm/unittests/FuzzMutate/StrategiesTest.cpp

@@ -6,6 +6,7 @@
 //
 //===----------------------------------------------------------------------===//
 
+#include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/AsmParser/Parser.h"
 #include "llvm/AsmParser/SlotMapping.h"
@@ -16,6 +17,7 @@
 #include "llvm/IR/Module.h"
 #include "llvm/IR/Verifier.h"
 #include "llvm/Support/SourceMgr.h"
+#include <random>
 
 #include "gtest/gtest.h"
 
@@ -307,4 +309,111 @@ TEST(InstModificationIRStrategyTest, DidntShuffleFRem) {
       }";
   VerfyDivDidntShuffle(Source);
 }
+
+static void VerifyBlockShuffle(StringRef Source) {
+  LLVMContext Ctx;
+  auto Mutator = createMutator<ShuffleBlockStrategy>();
+  ASSERT_TRUE(Mutator);
+
+  std::unique_ptr<Module> M = parseAssembly(Source.data(), Ctx);
+  Function *F = &*M->begin();
+  DenseMap<BasicBlock *, int> PreShuffleInstCnt;
+  for (BasicBlock &BB : F->getBasicBlockList()) {
+    PreShuffleInstCnt.insert({&BB, BB.size()});
+  }
+  std::mt19937 mt(Seed);
+  std::uniform_int_distribution<int> RandInt(INT_MIN, INT_MAX);
+  for (int i = 0; i < 100; i++) {
+    Mutator->mutateModule(*M, RandInt(mt), Source.size(), Source.size() + 1024);
+    for (BasicBlock &BB : F->getBasicBlockList()) {
+      int PostShuffleIntCnt = BB.size();
+      EXPECT_EQ(PostShuffleIntCnt, PreShuffleInstCnt[&BB]);
+    }
+    EXPECT_FALSE(verifyModule(*M, &errs()));
+  }
+}
+
+TEST(ShuffleBlockStrategy, ShuffleBlocks) {
+  StringRef Source = "\n\
+        define i64 @test(i1 %0, i1 %1, i1 %2, i32 %3, i32 %4) { \n\
+        Entry:  \n\
+          %A = alloca i32, i32 8, align 4 \n\
+          %E.1 = and i32 %3, %4 \n\
+          %E.2 = add i32 %4 , 1 \n\
+          %A.GEP.1 = getelementptr i32, ptr %A, i32 0 \n\
+          %A.GEP.2 = getelementptr i32, ptr %A.GEP.1, i32 1 \n\
+          %L.2 = load i32, ptr %A.GEP.2 \n\
+          %L.1 = load i32, ptr %A.GEP.1 \n\
+          %E.3 = sub i32 %E.2, %L.1 \n\
+          %Cond.1 = icmp eq i32 %E.3, %E.2 \n\
+          %Cond.2 = and i1 %0, %1 \n\
+          %Cond = or i1 %Cond.1, %Cond.2 \n\
+          br i1 %Cond, label %BB0, label %BB1  \n\
+        BB0:  \n\
+          %Add = add i32 %L.1, %L.2 \n\
+          %Sub = sub i32 %L.1, %L.2 \n\
+          %Sub.1 = sub i32 %Sub, 12 \n\
+          %Cast.1 = bitcast i32 %4 to float \n\
+          %Add.2 = add i32 %3, 1 \n\
+          %Cast.2 = bitcast i32 %Add.2 to float \n\
+          %FAdd = fadd float %Cast.1, %Cast.2 \n\
+          %Add.3 = add i32 %L.2, %L.1 \n\
+          %Cast.3 = bitcast float %FAdd to i32 \n\
+          %Sub.2 = sub i32 %Cast.3, %Sub.1 \n\
+          %SExt = sext i32 %Cast.3 to i64 \n\
+          %A.GEP.3 = getelementptr i64, ptr %A, i32 1 \n\
+          store i64 %SExt, ptr %A.GEP.3 \n\
+          br label %Exit  \n\
+        BB1:  \n\
+          %PHI.1 = phi i32 [0, %Entry] \n\
+          %SExt.1 = sext i1 %Cond.2 to i32 \n\
+          %SExt.2 = sext i1 %Cond.1 to i32 \n\
+          %E.164 = zext i32 %E.1 to i64 \n\
+          %E.264 = zext i32 %E.2 to i64 \n\
+          %E.1264 = mul i64 %E.164, %E.264 \n\
+          %E.12 = trunc i64 %E.1264 to i32 \n\
+          %A.GEP.4 = getelementptr i32, ptr %A, i32 2 \n\
+          %A.GEP.5 = getelementptr i32, ptr %A.GEP.4, i32 2 \n\
+          store i32 %E.12, ptr %A.GEP.5 \n\
+          br label %Exit  \n\
+        Exit:  \n\
+          %PHI.2 = phi i32 [%Add, %BB0], [%E.3, %BB1] \n\
+          %PHI.3 = phi i64 [%SExt, %BB0], [%E.1264, %BB1] \n\
+          %ZExt = zext i32 %PHI.2 to i64 \n\
+          %Add.5 = add i64 %PHI.3, 3 \n\
+          ret i64 %Add.5  \n\
+      }";
+  VerifyBlockShuffle(Source);
+}
+
+TEST(ShuffleBlockStrategy, ShuffleLoop) {
+  StringRef Source = "\n\
+    define i32 @foo(i32 %Left, i32 %Right) { \n\
+    Entry: \n\
+      %LPtr = alloca i32, align 4 \n\
+      %RPtr = alloca i32, align 4 \n\
+      %RetValPtr = alloca i32, align 4 \n\
+      store i32 %Left, ptr %LPtr, align 4 \n\
+      store i32 %Right, ptr %RPtr, align 4 \n\
+      store i32 0, ptr %RetValPtr, align 4 \n\
+      br label %LoopHead \n\
+    LoopHead: \n\
+      %L = load i32, ptr %LPtr, align 4 \n\
+      %R = load i32, ptr %RPtr, align 4 \n\
+      %C = icmp slt i32 %L, %R \n\
+      br i1 %C, label %LoopBody, label %Exit \n\
+    LoopBody: \n\
+      %OldL = load i32, ptr %LPtr, align 4 \n\
+      %NewL = add nsw i32 %OldL, 1 \n\
+      store i32 %NewL, ptr %LPtr, align 4 \n\
+      %OldRetVal = load i32, ptr %RetValPtr, align 4 \n\
+      %NewRetVal = add nsw i32 %OldRetVal, 1 \n\
+      store i32 %NewRetVal, ptr %RetValPtr, align 4 \n\
+      br label %LoopHead \n\
+    Exit: \n\
+      %RetVal = load i32, ptr %RetValPtr, align 4 \n\
+      ret i32 %RetVal \n\
+    }";
+  VerifyBlockShuffle(Source);
+}
 } // namespace