解决mlc-ai/web-stable-diffusion项目中的TVMError重复符号问题

在macOS系统上部署mlc-ai/web-stable-diffusion项目时，开发者可能会遇到一个典型的TVM编译错误："IRModule contains duplicate global symbol: main"。这个问题通常出现在构建阶段，当TVM编译器检测到模块中存在重复的全局符号时触发。

问题背景

mlc-ai/web-stable-diffusion是一个基于TVM的Stable Diffusion Web部署项目。在构建过程中，TVM需要将模型的不同部分（如UNet、VAE等）编译为可执行模块。当这些模块都包含名为"main"的入口函数时，就会导致符号冲突。

错误分析

错误信息明确指出："IRModule contains duplicate global symbol: main"，这表明在IRModule中存在多个同名的全局符号"main"。在TVM的编译流程中，每个模块应该具有唯一的全局符号名称，以便运行时能够正确区分和调用不同的功能模块。

解决方案

通过修改build.py脚本，我们可以为模型的不同部分指定不同的全局符号名称。以下是关键修改点：

1. 分离数据库应用阶段：将模型的不同部分（UNet、VAE等）分别应用元调度数据库
2. 重命名全局符号：为每个主要功能模块指定唯一的符号名称
3. 处理缺失的调度：为未调度的操作添加默认GPU调度

具体实现如下：

mod_deploy = mod
print("Applying database 1 =======================")
db3 = ms.database.create(work_dir=args.db_path)
with args.target, db3, tvm.transform.PassContext(opt_level=3):
    mod_deploy = relax.transform.MetaScheduleApplyDatabase(enable_warning=True)(mod_deploy)

print("Applying database 2 =======================")
db0 = ms.database.create(work_dir=args.db_path)
with args.target, db0, tvm.transform.PassContext(opt_level=3):
    mod_deploy = relax.transform.MetaScheduleApplyDatabase(enable_warning=True)(mod_deploy)

print("Applying database 3 =======================")
db2 = ms.database.create(work_dir=args.db_path)
with args.target, db2, tvm.transform.PassContext(opt_level=3):
    mod_deploy = relax.transform.MetaScheduleApplyDatabase(enable_warning=True)(mod_deploy)

print("Generating missing schedules ==============")
with tvm.target.Target("cuda"):
    mod_deploy = tvm.tir.transform.DefaultGPUSchedule()(mod_deploy)

for gv, func in mod_deploy.functions.items():
    try:
        if func.attrs["global_symbol"] == "main" and func.attrs["num_input"] == 3: # u-net
            mod_deploy[gv] = func.with_attr("global_symbol", "unet")
        if func.attrs["global_symbol"] == "main" and func.attrs["num_input"] == 1: # vae
            mod_deploy[gv] = func.with_attr("global_symbol", "vae")
        if func.attrs["global_symbol"] == "subgraph_0":
            mod_deploy[gv] = func.with_attr("global_symbol", "clip")
    except:
        pass

环境配置建议

在macOS系统上成功运行该项目，需要注意以下环境配置：

1. Python版本：推荐使用3.11.x
2. TVM版本：建议从源码编译，确保包含必要的后端支持
3. 依赖管理：使用pip安装requirements.txt中的依赖，但可以跳过与CUDA相关的部分
4. 硬件要求：M1/M2系列芯片表现良好

验证结果

应用上述修改后，项目可以成功构建并运行。测试命令如下：

python3 deploy.py --prompt "A photo of an astronaut riding a horse on mars."

在M1 Max芯片上，生成一张512x512的图像大约需要19秒，性能表现符合预期。

总结

通过分析TVM编译过程中的符号冲突问题，我们找到了一个有效的解决方案。这种方法不仅适用于mlc-ai/web-stable-diffusion项目，也可以为其他基于TVM的项目提供参考，特别是在处理包含多个功能模块的复杂模型时。理解TVM的模块系统和符号管理机制，对于解决类似的编译问题非常有帮助。