Disable optimisation for functions that contain KLEE calls

Compilers are allowed to hoist function calls and do GVE.
This is currently done even without `--optimization` enabled.

This is unfortunate in the context of KLEE function calls that might
depend on specific code position without direct control flow
dependencies. In such cases, function calls can be hoisted.

To circumvent this, disallow to optimise functions that contain such
calls by default. This might reduce optimisation for some functions
containing such function calls but still allows it for all others.

This patch adds an additional pass, that detects all functions starting with a
prefix `klee_` and disable optimisations for functions containing such
calls.

This is enabled by default but can be disabled by
`--klee-call-optimisation=false`.

5 files changed, 86 insertions, 1 deletions

diff --git a/lib/Module/CMakeLists.txt b/lib/Module/CMakeLists.txt
index 73f003c2..341af9df 100644
--- a/lib/Module/CMakeLists.txt
+++ b/lib/Module/CMakeLists.txt
@@ -16,6 +16,7 @@ set(KLEE_MODULE_COMPONENT_SRCS
   LowerSwitch.cpp
   ModuleUtil.cpp
   Optimize.cpp
+  OptNone.cpp
   PhiCleaner.cpp
   RaiseAsm.cpp
   Scalarizer.cpp
diff --git a/lib/Module/KModule.cpp b/lib/Module/KModule.cpp
index 3ce95a9e..2a15b02f 100644
--- a/lib/Module/KModule.cpp
+++ b/lib/Module/KModule.cpp
@@ -112,6 +112,12 @@ namespace {
   DontVerify("disable-verify",
              cl::desc("Do not verify the module integrity (default=false)"),
              cl::init(false), cl::cat(klee::ModuleCat));
+
+  cl::opt<bool>
+  OptimiseKLEECall("klee-call-optimisation",
+                             cl::desc("Allow optimization of functions that "
+                                      "contain KLEE calls (default=true)"),
+                             cl::init(true), cl::cat(ModuleCat));
 }
 
 /***/
@@ -262,6 +268,14 @@ void KModule::instrument(const Interpreter::ModuleOptions &opts) {
 void KModule::optimiseAndPrepare(
     const Interpreter::ModuleOptions &opts,
     llvm::ArrayRef<const char *> preservedFunctions) {
+  // Preserve all functions containing klee-related function calls from being
+  // optimised around
+  if (!OptimiseKLEECall) {
+    LegacyLLVMPassManagerTy pm;
+    pm.add(new OptNonePass());
+    pm.run(*module);
+  }
+
   if (opts.Optimize)
     Optimize(module.get(), preservedFunctions);
 
diff --git a/lib/Module/OptNone.cpp b/lib/Module/OptNone.cpp
new file mode 100644
index 00000000..08837488
--- /dev/null
+++ b/lib/Module/OptNone.cpp
@@ -0,0 +1,59 @@
+//===-- OptNone.cpp -------------------------------------------------------===//
+//
+//                     The KLEE Symbolic Virtual Machine
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+
+#include "Passes.h"
+
+#include "klee/Config/Version.h"
+
+#include "llvm/ADT/SmallPtrSet.h"
+#include "llvm/IR/Function.h"
+#include "llvm/IR/Instruction.h"
+#include "llvm/IR/Module.h"
+
+namespace klee {
+
+char OptNonePass::ID;
+
+bool OptNonePass::runOnModule(llvm::Module &M) {
+  // Find list of functions that start with `klee_`
+  // and mark all functions that contain such call or invoke as optnone
+  llvm::SmallPtrSet<llvm::Function *,16> CallingFunctions;
+  for (auto &F : M) {
+    if (!F.hasName() || !F.getName().startswith("klee_"))
+      continue;
+#if LLVM_VERSION_CODE >= LLVM_VERSION(3, 5)
+    for (auto *U : F.users()) {
+      // skip non-calls and non-invokes
+      if (!llvm::isa<llvm::CallInst>(U) && !llvm::isa<llvm::InvokeInst>(U))
+        continue;
+      auto *Inst = llvm::cast<llvm::Instruction>(U);
+      CallingFunctions.insert(Inst->getParent()->getParent());
+    }
+#else
+    for (auto i = F.use_begin(), e = F.use_end(); i != e; ++i) {
+      if (auto Inst = llvm::dyn_cast<llvm::Instruction>(*i)) {
+        CallingFunctions.insert(Inst->getParent()->getParent());
+      }
+    }
+#endif
+  }
+
+  bool changed = false;
+  for (auto F : CallingFunctions) {
+    // Skip if already annotated
+    if (F->hasFnAttribute(llvm::Attribute::OptimizeNone))
+      continue;
+    F->addFnAttr(llvm::Attribute::OptimizeNone);
+    F->addFnAttr(llvm::Attribute::NoInline);
+    changed = true;
+  }
+
+  return changed;
+}
+} // namespace klee
diff --git a/lib/Module/Passes.h b/lib/Module/Passes.h
index 417ebe7d..6a302819 100644
--- a/lib/Module/Passes.h
+++ b/lib/Module/Passes.h
@@ -185,6 +185,14 @@ public:
   bool runOnModule(llvm::Module &M) override;
 };
 #endif
+
+/// Instruments every function that contains a KLEE function call as nonopt
+class OptNonePass : public llvm::ModulePass {
+public:
+  static char ID;
+  OptNonePass() : llvm::ModulePass(ID) {}
+  bool runOnModule(llvm::Module &M) override;
+};
 } // namespace klee
 
 #endif
diff --git a/test/Feature/srem.c b/test/Feature/srem.c
index 57bea901..1303ac63 100644
--- a/test/Feature/srem.c
+++ b/test/Feature/srem.c
@@ -1,6 +1,9 @@
 // RUN: %clang %s -emit-llvm -g %O0opt -c -o %t.bc
 // RUN: rm -rf %t.klee-out
-// RUN: %klee --output-dir=%t.klee-out -use-cex-cache=1 %t.bc 2>&1 | FileCheck %s
+// RUN: %klee --output-dir=%t.klee-out --klee-call-optimisation=false %t.bc 2>&1 | FileCheck %s
+// RUN: rm -rf %t.klee-out
+// RUN: %klee --output-dir=%t.klee-out --klee-call-optimisation=false --optimize %t.bc 2>&1 | FileCheck %s
+
 #include <stdio.h>
 #include <assert.h>