DCMaterialization.cpp

Go to the documentation of this file.

//===- DCMaterialization.cpp - Fork/sink materialization pass ---*- C++ -*-===//
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//
//
// Contains the definitions of the Fork/sink materialization pass.
//
//===----------------------------------------------------------------------===//
 
#include "PassDetails.h"
#include "circt/Dialect/DC/DCOps.h"
#include "circt/Dialect/DC/DCPasses.h"
 
using namespace circt;
using namespace dc;
using namespace mlir;
 
static bool isDCTyped(Value v) {
  return v.getType().isa<dc::TokenType, dc::ValueType>();
}
 
static void replaceFirstUse(Operation *op, Value oldVal, Value newVal) {
  for (int i = 0, e = op->getNumOperands(); i < e; ++i)
    if (op->getOperand(i) == oldVal) {
      op->setOperand(i, newVal);
      break;
    }
}
 
// Adds a sink to the provided token or value-typed Value `v`.
static void insertSink(Value v, OpBuilder &rewriter) {
  rewriter.setInsertionPointAfterValue(v);
  if (v.getType().isa<ValueType>()) {
    // Unpack before sinking
    v = rewriter.create<UnpackOp>(v.getLoc(), v).getToken();
  }
 
  rewriter.create<SinkOp>(v.getLoc(), v);
}
 
// Adds a fork of the provided token or value-typed Value `result`.
static void insertFork(Value result, OpBuilder &rewriter) {
  rewriter.setInsertionPointAfterValue(result);
  // Get successor operations
  std::vector<Operation *> opsToProcess;
  for (auto &u : result.getUses())
    opsToProcess.push_back(u.getOwner());
 
  bool isValue = result.getType().isa<ValueType>();
  Value token = result;
  Value value;
  if (isValue) {
    auto unpack = rewriter.create<UnpackOp>(result.getLoc(), result);
    token = unpack.getToken();
    value = unpack.getOutput();
  }
 
  // Insert fork after op
  auto forkSize = opsToProcess.size();
  auto newFork = rewriter.create<ForkOp>(token.getLoc(), token, forkSize);
 
  // Modify operands of successor
  // opsToProcess may have multiple instances of same operand
  // Replace uses one by one to assign different fork outputs to them
  for (auto [op, forkOutput] : llvm::zip(opsToProcess, newFork->getResults())) {
    Value forkRes = forkOutput;
    if (isValue)
      forkRes =
          rewriter.create<PackOp>(forkRes.getLoc(), forkRes, value).getOutput();
    replaceFirstUse(op, result, forkRes);
  }
}
 
// Insert Fork Operation for every operation and function argument with more
// than one successor.
static LogicalResult addForkOps(Block &block, OpBuilder &rewriter) {
  // Materialization adds operations _after_ their definition, so we can't use
  // llvm::make_early_inc_range. Copy over all of the ops to process.
  llvm::SmallVector<Operation *> opsToProcess;
  for (auto &op : block.getOperations())
    opsToProcess.push_back(&op);
 
  for (Operation *op : opsToProcess) {
    // Ignore terminators.
    if (!op->hasTrait<OpTrait::IsTerminator>()) {
      for (auto result : op->getResults()) {
        if (!isDCTyped(result))
          continue;
        // If there is a result, it is used more than once, and it is a DC
        // type, fork it!
        if (!result.use_empty() && !result.hasOneUse())
          insertFork(result, rewriter);
      }
    }
  }
 
  for (auto barg : block.getArguments())
    if (!barg.use_empty() && !barg.hasOneUse())
      insertFork(barg, rewriter);
 
  return success();
}
 
// Create sink for every unused result
static LogicalResult addSinkOps(Block &block, OpBuilder &rewriter) {
  for (auto arg : block.getArguments()) {
    if (isDCTyped(arg) && arg.use_empty())
      insertSink(arg, rewriter);
  }
 
  // Materialization adds operations _after_ their definition, so we can't use
  // llvm::make_early_inc_range. Copy over all of the ops to process.
  llvm::SmallVector<Operation *> opsToProcess;
  for (auto &op : block.getOperations())
    opsToProcess.push_back(&op);
 
  for (Operation *op : opsToProcess) {
    if (op->getNumResults() == 0)
      continue;
 
    for (auto result : op->getResults()) {
      if (isDCTyped(result) && result.use_empty())
        insertSink(result, rewriter);
    }
  }
 
  return success();
}
 
namespace {
struct DCMaterializeForksSinksPass
    : public DCMaterializeForksSinksBase<DCMaterializeForksSinksPass> {
  void runOnOperation() override {
    auto funcOp = getOperation();
    if (funcOp.isExternal())
      return;
    OpBuilder builder(funcOp);
 
    auto walkRes = funcOp.walk([&](mlir::Block *block) {
      if (addForkOps(*block, builder).failed() ||
          addSinkOps(*block, builder).failed())
        return WalkResult::interrupt();
 
      return WalkResult::advance();
    });
 
    if (walkRes.wasInterrupted())
      return signalPassFailure();
  };
};
 
struct DCDematerializeForksSinksPass
    : public DCDematerializeForksSinksBase<DCDematerializeForksSinksPass> {
  void runOnOperation() override {
    auto funcOp = getOperation();
 
    if (funcOp.isExternal())
      return;
    funcOp.walk([&](dc::SinkOp sinkOp) { sinkOp.erase(); });
    funcOp.walk([&](dc::ForkOp forkOp) {
      for (auto res : forkOp->getResults())
        res.replaceAllUsesWith(forkOp.getOperand());
      forkOp.erase();
    });
  };
};
 
} // namespace
 
std::unique_ptr<mlir::Pass> circt::dc::createDCMaterializeForksSinksPass() {
  return std::make_unique<DCMaterializeForksSinksPass>();
}
 
std::unique_ptr<mlir::Pass> circt::dc::createDCDematerializeForksSinksPass() {
  return std::make_unique<DCDematerializeForksSinksPass>();
}