MLIR 源码分析：duplicate-function-elimination Pass

[xtyi]

def DuplicateFunctionEliminationPass : Pass<"duplicate-function-elimination",
    "ModuleOp"> {
  let summary = "Deduplicate functions";
  let description = [{
    Deduplicate functions that are equivalent in all aspects but their symbol
    name. The pass chooses one representative per equivalence class, erases
    the remainder, and updates function calls accordingly.
  }];
  let constructor = "mlir::func::createDuplicateFunctionEliminationPass()";
}

struct DuplicateFunctionEliminationPass
    : public impl::DuplicateFunctionEliminationPassBase<
          DuplicateFunctionEliminationPass> {

  using DuplicateFunctionEliminationPassBase<
      DuplicateFunctionEliminationPass>::DuplicateFunctionEliminationPassBase;

  void runOnOperation() override {
    auto module = getOperation();

    // Find unique representant per equivalent func ops.
    DenseSet<func::FuncOp, DuplicateFuncOpEquivalenceInfo> uniqueFuncOps;
    DenseMap<StringAttr, func::FuncOp> getRepresentant;
    DenseSet<func::FuncOp> toBeErased;
    module.walk([&](func::FuncOp f) {
      auto [repr, inserted] = uniqueFuncOps.insert(f);
      getRepresentant[f.getSymNameAttr()] = *repr;
      if (!inserted) {
        toBeErased.insert(f);
      }
    });

    // Update call ops to call unique func op representants.
    module.walk([&](func::CallOp callOp) {
      func::FuncOp callee = getRepresentant[callOp.getCalleeAttr().getAttr()];
      callOp.setCallee(callee.getSymName());
    });

    // Erase redundant func ops.
    for (auto it : toBeErased) {
      it.erase();
    }
  }
};

首先使用一个 DenseSet<func::FuncOp, DuplicateFuncOpEquivalenceInfo> uniqueFuncOps 来保存并判断 FuncOp 是否重复，由于 FuncOp 是自定义类型，DenseSet 并不知道如何计算 Hash 和判断相等，因此需要提供一个 DenseMapInfo<FuncOp>

// Define a notion of function equivalence that allows for reuse. Ignore the
// symbol name for this purpose.
struct DuplicateFuncOpEquivalenceInfo
    : public llvm::DenseMapInfo<func::FuncOp> {}

DuplicateFuncOpEquivalenceInfo 继承了 DenseMapInfo<FuncOp> 并重写了 getHashValue 和 isEqual 方法，这里主要看一下 isEqual 方法

  static bool isEqual(func::FuncOp lhs, func::FuncOp rhs) {
    if (lhs == rhs)
      return true;
    if (lhs == getTombstoneKey() || lhs == getEmptyKey() ||
        rhs == getTombstoneKey() || rhs == getEmptyKey())
      return false;
    // Check discardable attributes equivalence
    if (lhs->getDiscardableAttrDictionary() !=
        rhs->getDiscardableAttrDictionary())
      return false;

    // Check properties equivalence, ignoring the symbol name.
    // Make a copy, so that we can erase the symbol name and perform the
    // comparison.
    auto pLhs = lhs.getProperties();
    auto pRhs = rhs.getProperties();
    pLhs.sym_name = nullptr;
    pRhs.sym_name = nullptr;
    if (pLhs != pRhs)
      return false;

    // Compare inner workings.
    return OperationEquivalence::isRegionEquivalentTo(
        &lhs.getBody(), &rhs.getBody(), OperationEquivalence::IgnoreLocations);
  }

先使用 == 运算直接判断 funcOp 是否相等，调用 getDiscardableAttrDictionary 获取一个 Attribute Dictionary 判断是否相等，最后调用 OperationEquivalence::isRegionEquivalentTo 判断两个 Region 是否相等

DenseSet<func::FuncOp, DuplicateFuncOpEquivalenceInfo> uniqueFuncOps;
DenseMap<StringAttr, func::FuncOp> getRepresentant;
DenseSet<func::FuncOp> toBeErased;

回过头看 Pass，除了使用 uniqueFuncOps，还使用了一个 DenseMap<StringAttr, func::FuncOp> getRepresentant 来记录从函数名到 FuncOp 的映射，这里保存的 FuncOp 是所有等价的 FuncOp 的代表，即对于等价的函数 A，B，C，随便选出一个 FuncOp 作为代表，A，B，C 的函数名都会指向该 FuncOp

DenseSet<func::FuncOp> toBeErased 记录要被移除的 FuncOp，这里没有使用 DuplicateFuncOpEquivalenceInfo 是因为不需要重写 getHashValue 和 isEqual 方法，使用通用的方法已经够了，判断一下指针是否相等就行。

module.walk([&](func::FuncOp f) {
  auto [repr, inserted] = uniqueFuncOps.insert(f);
  getRepresentant[f.getSymNameAttr()] = *repr;
  if (!inserted) {
    toBeErased.insert(f);
  }
});

这里遍历 module 下所有 funcOp，调用 insert 插入到 uniqueFuncOps 中，这里 insert 返回一个 std::pair，first 表示元素位置的迭代器，second 是一个布尔值，表示是否成功插入，使用 [] 进行解包

没有成功插入的话，表示 Set 中已经存在该元素，所以 first 还是会指向该位置

不管插入是否成功，uniqueFuncOps 里面保存的都是准备保留下来的 funcOp，所以都设置当前函数名到 funcOp 的映射

如果插入失败，则将 funcOp 插入到 toBeErased Set 中

// Update call ops to call unique func op representants.
module.walk([&](func::CallOp callOp) {
  func::FuncOp callee = getRepresentant[callOp.getCalleeAttr().getAttr()];
  callOp.setCallee(callee.getSymName());
});

遍历 callOp，根据函数名从 getRepresentant 中找新的 Callee，然后更新成新的 Callee

// Erase redundant func ops.
  for (auto it : toBeErased) {
    it.erase();
  }

最后删除所有冗余的 funcOp