[LLHD] Add memory to block argument promotion pass (#52)

* [LLHD] Add memory to block argument promotion pass

* Promote memory locations allocated with llhd.var to block arguments to allow further optimizations.
* This pass is also required to lower behavioral LLHD to structural LLHD, because all control-flow has to be eliminated and structural LLHD has no memory model
* For this version pointers may not be used in operations apart from llhd.load and llhd.store, it may also not be created by a function call.
* It is limited to llhd.proc for now as functions have to be completely inlined for the lowering anyways
* There may not be any operations with nested regions in llhd.proc which are not isolated from above. This only affects the llhd.for operation which has to be completely unrolled for the lowering anyways
* Some of those restrictions can be lifted in a future commit

* Fix typos and better data structures

include/circt/Dialect/LLHD/Transforms/Passes.h

@@ -24,6 +24,8 @@ std::unique_ptr<OperationPass<ModuleOp>> createProcessLoweringPass();
 
 std::unique_ptr<OperationPass<ModuleOp>> createFunctionEliminationPass();
 
+std::unique_ptr<OperationPass<ProcOp>> createMemoryToBlockArgumentPass();
+
 std::unique_ptr<OperationPass<ProcOp>> createEarlyCodeMotionPass();
 
 /// Register the LLHD Transformation passes.

include/circt/Dialect/LLHD/Transforms/Passes.td

@@ -18,6 +18,57 @@ def ProcessLowering : Pass<"llhd-process-lowering", "ModuleOp"> {
   let constructor = "mlir::llhd::createProcessLoweringPass()";
 }
 
+def MemoryToBlockArgument : Pass<"llhd-memory-to-block-argument",
+                                 "llhd::ProcOp"> {
+  let summary = "Promote memory to block arguments.";
+  let description = [{
+    Promotes memory locations allocated with `llhd.var` to block arguments. This
+    enables other optimizations and is required to be able to lower behavioral
+    LLHD to structural LLHD. This is because there are no memory model and
+    control flow in structural LLHD. After executing this pass, the
+    "-llhd-block-argument-to-mux" pass can be used to convert the block
+    arguments to multiplexers to enable more control-flow elimination.
+
+    Example:
+
+    ```mlir
+    llhd.proc @check_simple(%condsig : !llhd.sig<i1>) -> () {
+      %c5 = llhd.const 5 : i32
+      %cond = llhd.prb %condsig : !llhd.sig<i1>
+      %ptr = llhd.var %c5 : i32
+      cond_br %cond, ^bb1, ^bb2
+    ^bb1:
+      %c6 = llhd.const 6 : i32
+      llhd.store %ptr, %c6 : !llhd.ptr<i32>
+      br ^bb2
+    ^bb2:
+      %ld = llhd.load %ptr : !llhd.ptr<i32>
+      %res = llhd.not %ld : i32
+      llhd.halt
+    }
+    ```
+
+    is transformed to
+
+    ```mlir
+    llhd.proc @check_simple(%condsig : !llhd.sig<i1>) -> () {
+      %c5 = llhd.const 5 : i32
+      %cond = llhd.prb %condsig : !llhd.sig<i1>
+      cond_br %cond, ^bb1, ^bb2(%c5 : i32)
+    ^bb1:
+      %c6 = llhd.const 6 : i32
+      br ^bb2(%c6 : i32)
+    ^bb2(%arg : i32):
+      %res = llhd.not %arg : i32
+      llhd.halt
+    }
+    ```
+
+  }];
+
+  let constructor = "mlir::llhd::createMemoryToBlockArgumentPass()";
+}
+
 def FunctionElimination : Pass<"llhd-function-elimination", "ModuleOp"> {
   let summary = "Deletes all functions.";
   let description = [{

lib/Dialect/LLHD/Transforms/CMakeLists.txt

@@ -3,6 +3,7 @@ add_mlir_dialect_library(MLIRLLHDTransforms
   PassRegistration.cpp
   ProcessLoweringPass.cpp
   FunctionEliminationPass.cpp
+  MemoryToBlockArgumentPass.cpp
   EarlyCodeMotionPass.cpp
 
   DEPENDS

lib/Dialect/LLHD/Transforms/MemoryToBlockArgumentPass.cpp

@@ -0,0 +1,237 @@
+//===- VarToBlockArgumentPass.cpp - Implement Var to Block Argument Pass --===//
+//
+// Implement pass to promote memory to block arguments.
+//
+//===----------------------------------------------------------------------===//
+
+#include "PassDetails.h"
+#include "circt/Dialect/LLHD/Transforms/Passes.h"
+#include "mlir/Dialect/StandardOps/IR/Ops.h"
+#include "mlir/IR/Dominance.h"
+#include <set>
+
+using namespace mlir;
+
+namespace {
+
+struct MemoryToBlockArgumentPass
+    : public llhd::MemoryToBlockArgumentBase<MemoryToBlockArgumentPass> {
+  void runOnOperation() override;
+};
+
+} // anonymous namespace
+
+/// Add the dominance fontier blocks of 'frontierOf' to the 'df' set
+static void getDominanceFrontier(Block *frontierOf, Operation *op,
+                                 std::set<Block *> &df) {
+  DominanceInfo dom(op);
+  for (Block &block : op->getRegion(0).getBlocks()) {
+    for (Block *pred : block.getPredecessors()) {
+      if (dom.dominates(frontierOf, pred) &&
+          !dom.properlyDominates(frontierOf, &block) &&
+          block.getOps<llhd::VarOp>().empty()) {
+        df.insert(&block);
+      }
+    }
+  }
+}
+
+/// Add the blocks in the closure of the dominance fontier relation of all the
+/// block in 'initialSet' to 'closure'
+static void getDFClosure(SmallVectorImpl<Block *> &initialSet, Operation *op,
+                         std::set<Block *> &closure) {
+  unsigned numElements;
+  for (Block *block : initialSet) {
+    getDominanceFrontier(block, op, closure);
+  }
+  do {
+    numElements = closure.size();
+    for (Block *block : closure) {
+      getDominanceFrontier(block, op, closure);
+    }
+  } while (numElements < closure.size());
+}
+
+/// Add a block argument to a given terminator. Only 'std.br', 'std.cond_br' and
+/// 'llhd.wait' are supported. The successor block has to be provided for the
+/// 'std.cond_br' terminator which has two possible successors.
+static void addBlockOperandToTerminator(Operation *terminator,
+                                        Block *successsor, Value toAppend) {
+  if (auto wait = dyn_cast<llhd::WaitOp>(terminator)) {
+    wait.destOpsMutable().append(toAppend);
+  } else if (auto br = dyn_cast<BranchOp>(terminator)) {
+    br.destOperandsMutable().append(toAppend);
+  } else if (auto condBr = dyn_cast<CondBranchOp>(terminator)) {
+    if (condBr.falseDest() == successsor) {
+      condBr.falseDestOperandsMutable().append(toAppend);
+    } else {
+      condBr.trueDestOperandsMutable().append(toAppend);
+    }
+  } else {
+    llvm_unreachable("unsupported terminator op");
+  }
+}
+
+void MemoryToBlockArgumentPass::runOnOperation() {
+  Operation *operation = getOperation();
+  OpBuilder builder(operation);
+
+  // No operations that have their own region and are not isolated from above
+  // are allowed for now.
+  WalkResult result = operation->walk([](Operation *op) -> WalkResult {
+    if (op->getNumRegions() > 0 && !op->isKnownIsolatedFromAbove())
+      return WalkResult::interrupt();
+    return WalkResult::advance();
+  });
+  if (result.wasInterrupted())
+    return;
+
+  // Get all variables defined in the body of this operation
+  // Note that variables that are passed as a function argument are not
+  // considered.
+  SmallVector<Value, 16> vars;
+  for (llhd::VarOp var : operation->getRegion(0).getOps<llhd::VarOp>()) {
+    vars.push_back(var.result());
+  }
+
+  // Don't consider variables that are used in other operations than load and
+  // store (e.g. as an argument to a call)
+  for (auto var = vars.begin(); var != vars.end(); ++var) {
+    for (Operation *user : var->getUsers()) {
+      if (!isa<llhd::LoadOp>(user) && !isa<llhd::StoreOp>(user)) {
+        vars.erase(var--);
+        break;
+      }
+    }
+  }
+
+  // For each variable find the blocks where a value is stored to it
+  for (Value var : vars) {
+    SmallVector<Block *, 16> defBlocks;
+    defBlocks.push_back(
+        var.getDefiningOp<llhd::VarOp>().getOperation()->getBlock());
+    operation->walk([&](llhd::StoreOp op) {
+      if (op.pointer() == var)
+        defBlocks.push_back(op.getOperation()->getBlock());
+    });
+    // Remove duplicates from the list
+    std::sort(defBlocks.begin(), defBlocks.end());
+    defBlocks.erase(std::unique(defBlocks.begin(), defBlocks.end()),
+                    defBlocks.end());
+
+    // Calculate initial set of join points
+    std::set<Block *> joinPoints;
+    getDFClosure(defBlocks, operation, joinPoints);
+
+    for (Block *jp : joinPoints) {
+      // Add a block argument for the variable at each join point
+      BlockArgument phi = jp->addArgument(
+          var.getType().cast<llhd::PtrType>().getUnderlyingType());
+
+      // Add a load at the end of every predecessor and pass the loaded value as
+      // the block argument
+      for (Block *pred : jp->getPredecessors()) {
+        // Set insertion point before terminator to insert the load operation
+        builder.setInsertionPoint(pred->getTerminator());
+        Value load = builder.create<llhd::LoadOp>(
+            pred->getTerminator()->getLoc(),
+            var.getType().cast<llhd::PtrType>().getUnderlyingType(), var);
+        // Add the loaded value as additional block argument
+        addBlockOperandToTerminator(pred->getTerminator(), jp, load);
+      }
+      // Insert a store at the beginning of the join point to make removal of
+      // all the memory operations easier later on
+      builder.setInsertionPointToStart(jp);
+      builder.create<llhd::StoreOp>(jp->front().getLoc(), var, phi);
+    }
+
+    // Basically reaching definitions analysis and replacing the loaded values
+    // by the values stored
+    DenseMap<Block *, Value> outputMap;
+    SmallPtrSet<Block *, 32> workQueue;
+    SmallPtrSet<Block *, 32> workDone;
+
+    workQueue.insert(&operation->getRegion(0).front());
+
+    while (!workQueue.empty()) {
+      // Pop block to process at the front
+      Block *block = *workQueue.begin();
+      workQueue.erase(block);
+
+      // Remember the currently stored value
+      Value currStoredValue;
+
+      // Update the value stored for the current variable at the start of this
+      // block
+      for (Block *pred : block->getPredecessors()) {
+        // Get the value at the end of a predecessor block and set it to the
+        // currently stored value at the start of this block
+        // NOTE: because we added block arguments and a store at the beginning
+        // of each join point we can assume that all predecessors have the same
+        // value stored at their end here because if that is not the case the
+        // first instruction in this block is a store instruction and will
+        // update the currently stored value to a correct one
+        if (!currStoredValue && outputMap.count(pred)) {
+          currStoredValue = outputMap[pred];
+          break;
+        }
+      }
+
+      // Iterate through all operations of the current block
+      for (auto op = block->begin(); op != block->end(); ++op) {
+        // Update currStoredValue at every store operation that stores to the
+        // variable we are currently considering
+        if (auto store = dyn_cast<llhd::StoreOp>(op)) {
+          if (store.pointer() == var) {
+            currStoredValue = store.value();
+            op = std::prev(op);
+            store.getOperation()->dropAllReferences();
+            store.erase();
+          }
+          // Set currStoredValue to the initializer value of the variable
+          // operation that created the variable we are currently considering,
+          // note that before that currStoredValue is uninitialized
+        } else if (auto varOp = dyn_cast<llhd::VarOp>(op)) {
+          if (varOp.result() == var)
+            currStoredValue = varOp.init();
+          // Replace the value returned by a load from the variable we are
+          // currently considering with the currStoredValue and delete the load
+          // operation
+        } else if (auto load = dyn_cast<llhd::LoadOp>(op)) {
+          if (load.pointer() == var && currStoredValue) {
+            op = std::prev(op);
+            load.result().replaceAllUsesWith(currStoredValue);
+            load.getOperation()->dropAllReferences();
+            load.erase();
+          }
+        }
+      }
+
+      if (currStoredValue)
+        outputMap.insert(std::make_pair(block, currStoredValue));
+
+      workDone.insert(block);
+
+      // Add all successors of this block to the work queue if they are not
+      // already processed
+      for (Block *succ : block->getSuccessors()) {
+        if (!workDone.count(succ))
+          workQueue.insert(succ);
+      }
+    }
+  }
+
+  // Remove all variable declarations for which we already removed all loads and
+  // stores
+  for (Value var : vars) {
+    Operation *op = var.getDefiningOp();
+    op->dropAllDefinedValueUses();
+    op->dropAllReferences();
+    op->erase();
+  }
+}
+
+std::unique_ptr<OperationPass<llhd::ProcOp>>
+mlir::llhd::createMemoryToBlockArgumentPass() {
+  return std::make_unique<MemoryToBlockArgumentPass>();
+}