如何在LLVM Pass中量化函数返回类型？稳健修改替代直接删除

在LLVM中，函数类型是不可变的，直接修改返回类型或删除旧函数容易因为未处理的引用、迭代器失效等问题引发崩溃。以下是无需直接删除旧函数的稳健转换流程：

1. 创建匹配新返回类型的函数副本 #

基于旧函数的属性，创建一个返回类型为i32的新函数，保留原函数的链接类型、调用约定、名称等关键属性：

Function *OldFunc = ...; // 通过模块遍历或getName找到目标@mulfix函数

// 构建新函数类型：i32(i32, i32)
FunctionType *NewFuncType = FunctionType::get(
    Type::getInt32Ty(OldFunc->getContext()),
    OldFunc->getArgumentTypes(),
    OldFunc->isVarArg()
);

// 创建新函数并复制旧函数的属性
Function *NewFunc = Function::Create(
    NewFuncType,
    OldFunc->getLinkage(),
    OldFunc->getName(),
    OldFunc->getParent()
);
NewFunc->copyAttributesFrom(OldFunc);
NewFunc->setDSOLocal(OldFunc->isDSOLocal());

2. 迁移旧函数的函数体并替换参数引用 #

将旧函数的所有基本块迁移到新函数中，同时替换块内对旧参数的引用为新函数的参数：

// 建立旧参数到新参数的映射表
ValueToValueMapTy ArgMap;
for (auto [OldArg, NewArg] : zip(OldFunc->args(), NewFunc->args())) {
    ArgMap[OldArg] = NewArg;
}

// 逐个迁移基本块
while (!OldFunc->empty()) {
    BasicBlock *BB = &OldFunc->back();
    BB->removeFromParent(); // 从旧函数移除，而非删除
    BB->insertInto(NewFunc);

    // 替换块内所有指令对旧参数的引用
    for (Instruction &Inst : *BB) {
        for (unsigned i = 0; i < Inst.getNumOperands(); ++i) {
            if (auto It = ArgMap.find(Inst.getOperand(i)); It != ArgMap.end()) {
                Inst.setOperand(i, It->second);
            }
        }
    }
}

3. 修改返回指令适配新类型 #

旧函数的返回值是float，需要根据你的量化逻辑将其转换为i32后返回。示例中使用FPToSI指令完成浮点到整数的转换（可根据实际量化规则调整）：

for (BasicBlock &BB : *NewFunc) {
    if (auto *RetInst = dyn_cast<ReturnInst>(BB.getTerminator())) {
        Value *OldRetVal = RetInst->getOperand(0);
        // 插入转换指令，替换原返回值
        Instruction *QuantizeConv = CastInst::Create(
            Instruction::FPToSI,
            OldRetVal,
            Type::getInt32Ty(NewFunc->getContext()),
            "quantize_ret",
            RetInst
        );
        RetInst->setOperand(0, QuantizeConv);
    }
}

4. 替换所有旧函数的调用点 #

收集并替换模块中所有对旧函数的调用，确保调用上下文适配新的返回类型：

// 提前收集所有调用指令，避免遍历过程中修改导致迭代器失效
SmallVector<CallInst *, 8> OldCalls;
for (User *U : OldFunc->users()) {
    if (auto *Call = dyn_cast<CallInst>(U)) {
        OldCalls.push_back(Call);
    }
}

// 逐个替换调用点
for (CallInst *OldCall : OldCalls) {
    // 创建新的调用指令
    CallInst *NewCall = CallInst::Create(
        NewFunc,
        OldCall->args(),
        "",
        OldCall
    );
    NewCall->copyAttributesFrom(OldCall);

    // 替换旧调用的所有使用：若原调用结果被用作float，需添加SIToFP转换（根据业务需求调整）
    OldCall->replaceAllUsesWith(NewCall);
    OldCall->eraseFromParent();
}

5. 处理旧函数的残留 #

完成上述步骤后，旧函数已无任何有效引用（调用点全部替换、函数体已迁移），LLVM的内存管理系统会在合适的时机自动清理它。无需调用eraseFromParent，避免因潜在的隐式引用（如间接函数指针、调试信息）引发崩溃。

内容的提问来源于stack exchange，提问作者Alex Xia