Simplify branch hint instrumentation

The previous branch hint instrumentation logic would introduce a scratch local to hold the condition so it could be passed into both the logging function and the original branching instruction. The de-instrumentation pass would then need to find this local and attempt to undo the data flow change. Simplify all of this by having the logging function return the condition value so it can interpose between the condition and the branch without any new locals. De-instrumentation can now just replace the call to the log function with its condition parameter.

To allow further simplification, also change the order of parameters to the logging function so the condition value is the first parameter. This ensures that we don't need to introduce a scratch local even when the condition is a `pop`, because the pop will remain the leftmost leaf expression in the catch body.

Author: tlively

scripts/fuzz_opt.py

@@ -2281,7 +2281,7 @@ def handle(self, wasm):
         for line in out.splitlines():
             if line.startswith(LOG_BRANCH_PREFIX):
                 # (1:-1 strips away the '[', ']' at the edges)
-                _, _, id_, hint, actual = line[1:-1].split(' ')
+                _, _, actual, hint, id_ = line[1:-1].split(' ')
                 all_ids.add(id_)
                 if hint != actual:
                     # This hint was misleading.

scripts/fuzz_shell.js

@@ -400,8 +400,9 @@ var baseImports = {
       });
     },
 
-    'log-branch': (id, expected, actual) => {
-      console.log(`[LoggingExternalInterface log-branch ${id} ${expected} ${actual}]`);
+    'log-branch': (actual, expected, id) => {
+      console.log(`[LoggingExternalInterface log-branch ${actual} ${expected} ${id}]`);
+      return actual;
     },
   },
   // Emscripten support.

src/passes/InstrumentBranchHints.cpp

@@ -28,9 +28,9 @@
 // into
 //
 //  @metadata.branch.hint B
-//  ;; log the ID of the condition (123), the prediction (B), and the actual
-//  ;; runtime result (temp == condition).
-//  if (temp = condition; log(123, B, temp); temp) {
+//  ;; log the actual runtime result (condition), the prediction (B), and the
+//  ;; ID (123), and return that result.
+//  if (log(condition, B, 123)) {
 //    X
 //  } else {
 //    Y
@@ -39,19 +39,20 @@
 // Concretely, we emit calls to this logging function:
 //
 //  (import "fuzzing-support" "log-branch"
-//    (func $log-branch (param i32 i32 i32)) ;; ID, prediction, actual
+//    (func $log-branch (param i32 i32 i32) (result i32))
 //  )
 //
 // This can be used to verify that branch hints are accurate, by implementing
 // the import like this for example:
 //
-//  imports['fuzzing-support']['log-branch'] = (id, prediction, actual) => {
+//  imports['fuzzing-support']['log-branch'] = (actual, prediction, id) => {
 //    // We only care about truthiness of the expected and actual values.
 //    expected = +!!expected;
 //    actual = +!!actual;
 //    // Throw if the hint said this branch would be taken, but it was not, or
 //    // vice versa.
 //    if (expected != actual) throw `Bad branch hint! (${id})`;
+//    return actual;
 //  };
 //
 // A pass to delete branch hints is also provided, which finds instrumentations
@@ -63,28 +64,28 @@
 // would do this transformation:
 //
 //  @metadata.branch.hint A
-//  if (temp = condition; log(10, A, temp); temp) { // 10 matches one of 10,20
+//  if (log(condition, A, 10)) { // 10 matches one of 10,20
 //    X
 //  }
 //  @metadata.branch.hint B
-//  if (temp = condition; log(99, B, temp); temp) { // 99 does not match
+//  if (log(condition, B, 99)) { // 99 does not match
 //    Y
 //  }
 //
 // =>
 //
 //  // Used to be a branch hint here, but it was deleted.
-//  if (temp = condition; log(10, A, temp); temp) {
+//  if (log(condition, A, 10)) {
 //    X
 //  }
 //  @metadata.branch.hint B // this one is unmodified.
-//  if (temp = condition; log(99, B, temp); temp) {
+//  if (log(condition, B, 99)) {
 //    Y
 //  }
 //
 // A pass to undo the instrumentation is also provided, which does
 //
-//  if (temp = condition; log(123, A, temp); temp) {
+//  if (log(condition, A, 123)) {
 //    X
 //  }
 //
@@ -95,14 +96,8 @@
 //  }
 //
 
-#include "ir/drop.h"
 #include "ir/effects.h"
-#include "ir/eh-utils.h"
-#include "ir/find_all.h"
-#include "ir/local-graph.h"
 #include "ir/names.h"
-#include "ir/parents.h"
-#include "ir/properties.h"
 #include "ir/utils.h"
 #include "pass.h"
 #include "support/string.h"
@@ -133,8 +128,6 @@ int branchId = 1;
 struct InstrumentBranchHints
   : public WalkerPass<PostWalker<InstrumentBranchHints>> {
 
-  using Super = WalkerPass<PostWalker<InstrumentBranchHints>>;
-
   // The internal name of our import.
   Name logBranch;
 
@@ -148,8 +141,6 @@ struct InstrumentBranchHints
 
   // TODO: BrOn, but the condition there is not an i32
 
-  bool addedInstrumentation = false;
-
   template<typename T> void processCondition(T* curr) {
     if (curr->condition->type == Type::unreachable) {
       // This branch is not even reached.
@@ -167,41 +158,32 @@ struct InstrumentBranchHints
     int id = branchId++;
 
     // Instrument the condition.
-    auto tempLocal = builder.addVar(getFunction(), Type::i32);
-    auto* set = builder.makeLocalSet(tempLocal, curr->condition);
     auto* idConst = builder.makeConst(Literal(int32_t(id)));
     auto* guess = builder.makeConst(Literal(int32_t(*likely)));
-    auto* get1 = builder.makeLocalGet(tempLocal, Type::i32);
-    auto* log = builder.makeCall(logBranch, {idConst, guess, get1}, Type::none);
-    auto* get2 = builder.makeLocalGet(tempLocal, Type::i32);
-    curr->condition = builder.makeBlock({set, log, get2});
-    addedInstrumentation = true;
-  }
-
-  void doWalkFunction(Function* func) {
-    Super::doWalkFunction(func);
 
-    // Our added blocks may have caused nested pops.
-    if (addedInstrumentation) {
-      EHUtils::handleBlockNestedPops(func, *getModule());
-      addedInstrumentation = false;
-    }
+    curr->condition =
+      builder.makeCall(logBranch, {curr->condition, guess, idConst}, Type::i32);
   }
 
   void doWalkModule(Module* module) {
     if (auto existing = getLogBranchImport(module)) {
       // This file already has our import. We nop it out, as whatever the
       // current code does may be dangerous (it may log incorrect hints).
       auto* func = module->getFunction(existing);
-      func->body = Builder(*module).makeNop();
+      Builder builder(*module);
+      if (func->getSig().results == Type::none) {
+        func->body = builder.makeNop();
+      } else {
+        func->body = builder.makeUnreachable();
+      }
       func->module = func->base = Name();
       func->type = func->type.with(Exact);
     }
 
     // Add our import.
     auto* func = module->addFunction(Builder::makeFunction(
       Names::getValidFunctionName(*module, BASE),
-      Type(Signature({Type::i32, Type::i32, Type::i32}, Type::none),
+      Type(Signature({Type::i32, Type::i32, Type::i32}, Type::i32),
            NonNullable,
            Inexact),
       {}));
@@ -210,7 +192,7 @@ struct InstrumentBranchHints
     logBranch = func->name;
 
     // Walk normally, using logBranch as we go.
-    Super::doWalkModule(module);
+    PostWalker<InstrumentBranchHints>::doWalkModule(module);
 
     // Update ref.func type changes.
     ReFinalize().run(getPassRunner(), module);
@@ -228,12 +210,6 @@ struct InstrumentationProcessor : public WalkerPass<PostWalker<Sub>> {
   // The internal name of our import.
   Name logBranch;
 
-  // A LocalGraph, so we can identify the pattern.
-  std::unique_ptr<LocalGraph> localGraph;
-
-  // A map of expressions to their parents, so we can identify the pattern.
-  std::unique_ptr<Parents> parents;
-
   Sub* self() { return static_cast<Sub*>(this); }
 
   void visitIf(If* curr) { self()->processCondition(curr); }
@@ -246,15 +222,6 @@ struct InstrumentationProcessor : public WalkerPass<PostWalker<Sub>> {
 
   // TODO: BrOn, but the condition there is not an i32
 
-  void doWalkFunction(Function* func) {
-    localGraph = std::make_unique<LocalGraph>(func, this->getModule());
-    localGraph->computeSetInfluences();
-
-    parents = std::make_unique<Parents>(func->body);
-
-    Super::doWalkFunction(func);
-  }
-
   void doWalkModule(Module* module) {
     logBranch = getLogBranchImport(module);
     if (!logBranch) {
@@ -267,73 +234,14 @@ struct InstrumentationProcessor : public WalkerPass<PostWalker<Sub>> {
 
   // Helpers
 
-  // Instrumentation info for a chunk of code that is the result of the
-  // instrumentation pass.
-  struct Instrumentation {
-    // The condition before the instrumentation (a pointer to it, so we can
-    // replace it).
-    Expression** originalCondition;
-    // The local that the original condition is stored in temporarily.
-    Index tempLocal;
-    // The call to the logging that the instrumentation added.
-    Call* call;
-  };
-
-  // Check if an expression's condition is an instrumentation, and return the
-  // info if so.
-  std::optional<Instrumentation> getInstrumentation(Expression* condition) {
-    // We must identify this pattern:
-    //
-    //  (br_if
-    //    (block
-    //      (local.set $temp (condition))
-    //      (call $log (id, prediction, (local.get $temp)))
-    //      (local.get $temp)
-    //    )
-    //
-    // The block may vanish during roundtrip though, so we just follow back from
-    // the last local.get, which appears in the condition:
-    //
-    //  (local.set $temp (condition))
-    //  (call $log (id, prediction, (local.get $temp)))
-    //  (br_if
-    //    (local.get $temp)
-    //
-    auto* fallthrough = Properties::getFallthrough(
-      condition, this->getPassOptions(), *this->getModule());
-    auto* get = fallthrough->template dynCast<LocalGet>();
-    if (!get) {
-      return {};
-    }
-    auto& sets = localGraph->getSets(get);
-    if (sets.size() != 1) {
-      return {};
-    }
-    auto* set = *sets.begin();
-    if (!set) {
-      return {};
-    }
-    auto& gets = localGraph->getSetInfluences(set);
-    if (gets.size() != 2) {
-      return {};
-    }
-    // The set has two gets: the get in the condition we began at, and
-    // another.
-    LocalGet* otherGet = nullptr;
-    for (auto* get2 : gets) {
-      if (get2 != get) {
-        otherGet = get2;
-      }
-    }
-    assert(otherGet);
-    // See if that other get is used in a logging. The parent should be a
-    // logging call.
-    auto* call = parents->getParent(otherGet)->template dynCast<Call>();
+  // Check if an expression's condition is instrumented, and return the
+  // instrumentation call if so. Otherwise return null.
+  Call* getInstrumentation(Expression* condition) {
+    auto* call = condition->dynCast<Call>();
     if (!call || call->target != logBranch) {
-      return {};
+      return nullptr;
     }
-    // Great, this is indeed a prior instrumentation.
-    return Instrumentation{&set->value, set->index, call};
+    return call;
   }
 };
 
@@ -344,8 +252,8 @@ struct DeleteBranchHints : public InstrumentationProcessor<DeleteBranchHints> {
   std::unordered_set<Index> idsToDelete;
 
   template<typename T> void processCondition(T* curr) {
-    if (auto info = getInstrumentation(curr->condition)) {
-      if (auto* c = info->call->operands[0]->template dynCast<Const>()) {
+    if (auto* call = getInstrumentation(curr->condition)) {
+      if (auto* c = call->operands[2]->template dynCast<Const>()) {
         auto id = c->value.geti32();
         if (idsToDelete.contains(id)) {
           // Remove the branch hint.
@@ -368,78 +276,31 @@ struct DeleteBranchHints : public InstrumentationProcessor<DeleteBranchHints> {
 };
 
 struct DeInstrumentBranchHints
-  : public InstrumentationProcessor<DeInstrumentBranchHints> {
+  : public WalkerPass<PostWalker<DeInstrumentBranchHints>> {
 
-  template<typename T> void processCondition(T* curr) {
-    if (auto info = getInstrumentation(curr->condition)) {
-      // Replace the instrumented condition with the original one (swap so that
-      // the IR remains valid: we cannot use the same expression twice in our
-      // IR, and the original condition is still used in another place, until
-      // we remove the logging calls; since we will remove the calls anyhow, we
-      // just need some valid IR there).
-      //
-      // Check for dangerous effects in the condition we are about to replace,
-      // to avoid a situation where the condition looks like this:
-      //
-      //  (set $temp (original condition))
-      //   ..effects..
-      //  (local.get $temp)
-      //
-      // We cannot replace all this with the original condition, as it would
-      // remove the effects. (Even in that case we will remove the actual call
-      // to log the branch hint, below, so this just prevents some cleanup that
-      // is normally safe - the cleanup is mainly useful to allow inspection of
-      // testcases for debugging.)
-      EffectAnalyzer effects(getPassOptions(), *getModule(), curr->condition);
-      // The only condition we allow is a write to the temp local from the
-      // instrumentation, which getInstrumentation() verified has no other uses
-      // than us.
-      effects.localsWritten.erase(info->tempLocal);
-      if (!effects.hasUnremovableSideEffects()) {
-        std::swap(curr->condition, *info->originalCondition);
-      }
-    }
-  }
+  // The internal name of our import.
+  Name logBranch;
 
-  void visitFunction(Function* func) {
-    if (func->imported()) {
-      return;
-    }
-    // At the very end, remove all logging calls (we use them during the main
-    // walk to identify instrumentation).
-    for (auto** callp : FindAllPointers<Call>(func->body).list) {
-      auto* call = (*callp)->cast<Call>();
-      if (call->target == logBranch) {
-        Builder builder(*getModule());
-        Expression* last;
-        if (call->type == Type::none) {
-          last = builder.makeNop();
-        } else {
-          last = builder.makeUnreachable();
-        }
-        *callp = getDroppedChildrenAndAppend(call,
-                                             *getModule(),
-                                             getPassOptions(),
-                                             last,
-                                             // We know the call is removable.
-                                             DropMode::IgnoreParentEffects);
-      }
+  void visitCall(Call* curr) {
+    if (curr->target == logBranch) {
+      // Replace the call with its first operand (the original condition).
+      replaceCurrent(curr->operands[0]);
     }
   }
 
   void doWalkModule(Module* module) {
-    auto logBranchImport = getLogBranchImport(module);
-    if (!logBranchImport) {
+    logBranch = getLogBranchImport(module);
+    if (!logBranch) {
       Fatal()
         << "No branch hint logging import found. Was this code instrumented?";
     }
 
     // Mark the log-branch import as having no side effects - we are removing it
     // entirely here, and its effect should not stop us when we compute effects.
-    module->getFunction(logBranchImport)->effects =
+    module->getFunction(logBranch)->effects =
       std::make_shared<EffectAnalyzer>(getPassOptions(), *module);
 
-    InstrumentationProcessor<DeInstrumentBranchHints>::doWalkModule(module);
+    WalkerPass<PostWalker<DeInstrumentBranchHints>>::doWalkModule(module);
   }
 };