论文
面向可扩展等式饱和的LLM引导策略合成
LLM-Guided Strategy Synthesis for Scalable Equality Saturation
摘要
等式饱和(EqSat)是一种强大的优化范式,它在e-graph中紧凑地表示许多等价程序,并延迟提交,直到抽取阶段选出成本最低的程序。因此,让EqSat有效不仅需要领域特定的重写规则,还需要领域特定的策略。如今,这类策略设计大多仍靠手工完成,成为自动化基于e-graph的编译器的一大障碍。近期的规则合成框架能够从语义规约自动推断庞大的重写词汇,但这也扩大了重写空间并进一步加剧e-graph爆炸。尽管大语言模型(LLM)使自动化策略合成成为可能,但直接演化后端代码在实践中仍然低效。这种搜索缺乏可复用的策略抽象与可执行的反馈,且极易触发e-graph爆炸或收敛到糟糕的设计。我们提出EggMind,一个LLM引导的端到端框架,用于合成可复用的EqSat策略。其核心是EggMind引入领域特定语言EqSatL,将EqSat策略表示为显式且可检查的产物。随后它提出一个LLM引导的智能体工作流,配备证明派生重写母题缓存与可处理性引导等新技术,以高效搜索高质量策略,同时保持合成在e-graph增长下的稳定。评估表明,EggMind显著改善了向量化基准上的资源-质量权衡,相对完整EqSat将最终成本降低45.1%、峰值内存降低69.1%。我们进一步表明,同一方法可有效地迁移到基于XLA的张量编译器,并在一个具有增强重写空间的逻辑综合案例研究中展示了其实用潜力。