案例4：全局值编号实现——消除冗余计算

1. 任务目标

本案例要实现一个基于 MLIR 的全局值编号（Global Value Numbering, GVN）Pass，用于消除中间表示中的冗余计算。你将完成一个完整的 Pass 编写、注册、加载与测试闭环，最终通过 mlir-opt 工具验证冗余的算术运算被正确消除。

2. 核心原理与实现策略

全局值编号的核心思想是：为每个计算操作分配一个唯一的编号（Value Number），该编号由操作的算子类型和操作数决定。如果两个操作具有相同的编号，则它们计算的是同一个值，后出现的操作可以被替换为先前的结果，从而消除冗余。

本案例采用 MLIR 的 DominanceInfo 来保证替换的安全性——只有当前操作被冗余操作的支配时才可替换。实现分为三个步骤：

• 哈希键构造：将操作的算子名与操作数组合成哈希键
• 冗余检测：在支配关系约束下查找相同哈希键的历史操作
• 重写替换：将冗余操作的输出替换为已存在值，并删除冗余操作

3. 完整代码实现

创建文件 GVNPass.cpp，代码如下：

#include "mlir/IR/BuiltinOps.h"
#include "mlir/IR/Dominance.h"
#include "mlir/IR/PatternMatch.h"
#include "mlir/Pass/Pass.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
#include "llvm/ADT/DenseMap.h"
#include "llvm/ADT/Hashing.h"
#include "llvm/ADT/STLExtras.h"
#include <set>

using namespace mlir;

namespace {

// 哈希键：算子名 + 操作数列表
struct ValueNumberKey {
  StringRef opName;
  SmallVector<Value, 4> operands;

  bool operator==(const ValueNumberKey &other) const {
    return opName == other.opName && operands == other.operands;
  }
};

struct ValueNumberKeyHash {
  size_t operator()(const ValueNumberKey &key) const {
    return llvm::hash_combine(
        key.opName,
        llvm::hash_combine_range(key.operands.begin(), key.operands.end()));
  }
};

struct GVNPass : public PassWrapper<GVNPass, OperationPass<func::FuncOp>> {
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(GVNPass)

  StringRef getArgument() const final { return "my-gvn"; }
  StringRef getDescription() const final {
    return "Global Value Numbering: eliminate redundant computations";
  }

  void runOnOperation() override {
    func::FuncOp func = getOperation();
    DominanceInfo &domInfo = getAnalysis<DominanceInfo>();

    // 存储每个值编号对应的代表操作
    llvm::DenseMap<ValueNumberKey, Value, ValueNumberKeyHash> valueNumberToValue;
    // 记录已删除的操作，避免重复处理
    std::set<Operation *> erasedOps;

    // 收集所有可编号的算术操作
    SmallVector<Operation *, 16> candidates;
    func.walk([&](Operation *op) {
      // 只处理纯计算操作（无副作用），排除常量、函数调用等
      if (op->getNumResults() == 1 &&
          op->hasTrait<OpTrait::ConstantLike>() == false &&
          isMemoryEffectFree(op)) {
        candidates.push_back(op);
      }
    });

    // 按支配序处理（先处理靠前的操作）
    llvm::sort(candidates, [&](Operation *a, Operation *b) {
      return domInfo.dominates(a, b);
    });

    for (Operation *op : candidates) {
      if (erasedOps.count(op))
        continue;

      // 构造哈希键
      ValueNumberKey key;
      key.opName = op->getName().getStringRef();
      for (Value operand : op->getOperands()) {
        // 跳过块参数（函数参数等），只对可追踪的SSA值编号
        if (auto *defOp = operand.getDefiningOp()) {
          if (erasedOps.count(defOp))
            continue;
        }
        key.operands.push_back(operand);
      }

      // 查找是否已有等价操作
      auto it = valueNumberToValue.find(key);
      if (it != valueNumberToValue.end()) {
        // 检查支配关系：只有代表操作支配当前操作时才可替换
        Operation *domOp = it->second.getDefiningOp();
        if (domOp && domInfo.dominates(domOp, op)) {
          // 替换所有使用当前操作结果的用户
          op->replaceAllUsesWith(it->second);
          erasedOps.insert(op);
          op->erase();
          continue;
        }
      }

      // 未找到冗余，记录当前操作
      valueNumberToValue[key] = op->getResult(0);
    }
  }

private:
  // 检查操作是否无内存副作用
  static bool isMemoryEffectFree(Operation *op) {
    if (auto effectInterface = dyn_cast<MemoryEffectOpInterface>(op)) {
      return effectInterface.hasNoEffect();
    }
    // 默认认为无接口的操作可能有副作用，保守处理
    return false;
  }
};

} // namespace

namespace mlir {
void registerGVNPass() { PassRegistration<GVNPass>(); }
} // namespace mlir

4. 注册与构建配置

在 CMakeLists.txt 中添加构建规则：

add_mlir_pass_library(GVNPass
  GVNPass.cpp
  DEPENDS
  mlirPass
  mlirIR
  mlirTransforms
  mlirAnalysis
  )

在 Pass 注册文件中（如 Passes.h）声明：

void registerGVNPass();

在 mlir-opt 的 main 函数中调用注册：

registerGVNPass();

5. 测试输入与运行方式

创建测试文件 test-gvn.mlir：

func.func @test_gvn(%arg0: i32, %arg1: i32) -> i32 {
  %0 = arith.addi %arg0, %arg1 : i32
  %1 = arith.addi %arg0, %arg1 : i32   // 冗余：与%0相同
  %2 = arith.muli %0, %1 : i32
  %3 = arith.addi %arg0, %arg1 : i32   // 冗余：与%0相同
  %4 = arith.subi %2, %3 : i32
  return %4 : i32
}

运行命令：

./bin/mlir-opt --my-gvn test-gvn.mlir

6. 预期输出与验证

正确执行后，输出应如下（冗余操作已被消除）：

func.func @test_gvn(%arg0: i32, %arg1: i32) -> i32 {
  %0 = arith.addi %arg0, %arg1 : i32
  %1 = arith.muli %0, %0 : i32
  %2 = arith.subi %1, %0 : i32
  return %2 : i32
}

验证要点：

• 三个 arith.addi %arg0, %arg1 中只有第一个被保留，后两个被替换为 %0
• arith.muli 的两个操作数都变为 %0（因为第二个操作数原本是冗余的 %1，已被替换）
• 最终 arith.subi 的第二个操作数从 %3 变为 %0
• 函数语义保持不变：(a+b)*(a+b) - (a+b)

7. 边界情况与扩展思考

本实现处理了以下关键边界情况：

• 支配关系检查：确保不会用后续定义的值替换前面的操作
• 副作用过滤：仅处理无内存副作用的纯计算操作
• 常量处理：跳过常量操作，避免与 MLIR 内置的常量折叠冲突

若要扩展到更复杂的场景，可考虑：

• 支持交换律操作（如 addi 的 %a+%b 与 %b+%a 等价）
• 结合 ScalarEvolution 处理循环不变表达式
• 与 MLIR 的 CSE（公共子表达式消除）Pass 对比，理解两者差异

8. 案例验证总结

通过本案例，你完成了：

• 一个完整的 MLIR Pass 编写流程（继承、注册、构建）
• 基于哈希的全局值编号核心算法实现
• 利用支配关系保证替换安全性
• 通过 mlir-opt 工具验证了冗余消除的正确性

该 Pass 可直接集成到优化流水线中，作为 --my-gvn 使用，为后续更复杂的优化（如循环不变量外提、死代码消除）奠定基础。