暂无图片
暂无图片
暂无图片
暂无图片
暂无图片

Meta LLM Compiler:面向编译器优化的基础模型

慢慢学 AIGC 2024-06-29
46

点击下方卡片,关注“慢慢学AIGC

摘要


大型语言模型(LLMs)在各种软件工程和编码任务中展现了卓越的能力。然而,它们在代码和编译器优化领域的应用仍未被充分探索。训练 LLMs 需要大量的 GPU 时间和海量数据收集,这可能是一个障碍。为了解决这个问题,我们推出了 Meta 大型语言模型编译器(LLM Compiler),这是一套专门为代码优化任务设计的、稳健的、公开可用的预训练模型。

LLM Compiler 建立在 Code Llama 的基础上,增强了对编译器中间表示(IRs)、汇编语言和优化技术的理解。该模型在 5460 亿个 LLVM-IR 和汇编代码 token上进行了训练,并通过指令微调来解释编译器行为。LLM Compiler 以自定义的商业许可发布,允许广泛重用,有 7B 和 13B 两种参数规模可选。

我们还展示了该模型的微调版本,证明了它在优化代码大小和将 x86_64 和 ARM 汇编反汇编回 LLVM-IR 方面的增强能力。这些模型达到了自动调优搜索优化潜力的 77%,以及 45% 的反汇编往返率(14% 的精确匹配)。此次发布旨在为学术研究人员和行业从业者在编译器优化方面的进一步研究和开发提供一个可扩展、高性价比的基础。

1. 简介


大型语言模型在软件工程任务中的应用越来越受关注,包括代码生成、代码翻译和代码测试。然而,专门针对代码优化的训练还很少。公开可用的 LLMs 可以被提示进行一些小的程序调整,甚至尝试更实质性的优化,但它们容易混淆并犯错,经常导致不正确的代码。

与之前的机器学习引导的代码优化方法相比,LLMs 可以接受完整的源程序作为输入,提供无损的表示。使用文本作为机器学习优化器的输入和输出表示具有理想的特性:文本是通用的、可移植的和易访问的接口,并且不像之前的方法那样专门针对特定任务。

然而,训练 LLMs 需要高昂的计算和数据成本。为了解决这个问题,我们发布了 LLM Compiler,这是一系列已经训练好的基础模型,它们理解编译器 IR 和汇编的语义,并能模拟编译器,允许使用最少的数据进行特定下游编译器优化任务的微调。

LLM Compiler 在两个下游编译任务上进行了微调:调整编译器标志以优化代码大小,以及将 x86_64 和ARM汇编反汇编为 LLVM-IR。与自动调优技术相比,LLM Compiler FTD 达到了 77% 的优化潜力,无需任何额外的编译。在反汇编时,LLM Compiler FTD 有 14% 的时间创建正确的反汇编。在这两项任务上,LLM Compiler FTD 模型都显著优于可比较的 LLMs Code Llama 和 GPT-4 Turbo。

我们的工作旨在为编译器优化的进一步研究和开发建立一个可扩展、高性价比的基础,既服务于学术研究人员,也服务于行业从业者。图 1 展示了我们方法的概览。

LLM Compiler 模型针对编译器优化。它们有两种模型规模:7B 和 13B 参数。LLM Compiler 模型使用相应大小的 Code Llama 模型权重进行初始化,并在额外的 5460 亿个 token 的数据上进行训练,这些数据主要包括编译器中间表示汇编代码。然后,我们进一步训练 LLM Compiler FTD 模型,使用额外的 1640 亿个 token 的数据,针对两个下游编译任务:标志调优和反汇编。在训练的所有阶段,都使用了少量来自前期阶段的代码和自然语言数据,以帮助保留基础 Code Llama 模型的能力。

2. LLM Compiler:为编译器优化定制 Code Llama


2.1 在汇编代码和编译器 IR 上的预训练

用于训练编码 LLMs 的数据通常主要由 Python 等高级源语言组成。汇编代码在这些数据集中占比很小,编译器 IR 更少。为了构建一个对这些语言有良好理解的 LLM,我们用 Code Llama 的权重初始化 LLM Compiler 模型,然后在一个以编译器为中心的数据集上训练 4010 亿个 token,这个数据集主要由汇编代码和编译器 IR 组成,如表 1 所示。

数据集:LLM Compiler 主要在由 LLVM(版本 17.0.6)生成的编译器中间表示和汇编代码上进行训练。这些数据来源于用于训练 Code Llama 的相同公开可用代码数据集。我们在表 2 中总结了这个数据集。

与 Code Llama 一样,我们也从自然语言数据集中获取少量训练批次。

2.2 编译器模拟的指令微调

为了理解代码优化的机制,我们对 LLM Compiler 模型进行指令微调以模拟编译器优化,如图 2 所示。

其思想是从有限的未优化种子程序集合中,通过对这些程序应用随机生成的编译器优化序列,生成大量示例。然后我们训练模型预测优化生成的代码。我们还训练模型预测应用优化后的代码大小

任务规范

给定未优化的 LLVM-IR(由 clang 前端生成),一个优化 pass 列表,和一个初始代码大小,生成应用这些优化后的结果代码和结果代码大小。这个任务有两种形式:在第一种中,模型需要输出编译器 IR在第二种中,模型需要输出汇编代码两种形式的输入 IR、优化 pass 和代码大小是相同的。提示决定了所需的输出格式。每种提示的示例在附录清单 2 和 3 中提供。

    B.1 Compiler emulation prompts
    Prompt:
    [ INST ] Give the LLVM - IR for the following code when optimized using opt -p ’ module ( default <Oz >) ’:
    <code >; ModuleID = ’< stdin >’
    source_filename = " -"
    target datalayout = "e -m:e - p270 :32:32 - p271 :32:32 - p272 :64:64 - i64 :64 - f80 :128 - n8 :16:32:64 - S128 "
    target triple = " x86_64 - unknown - linux - gnu "
    ; Function Attrs : minsize nounwind optsize uwtable
    define dso_local i32 @add_two ( i32 noundef %0 , i32 noundef %1) #0 {
    %3 = alloca i32 , align 4
    %4 = alloca i32 , align 4
    store i32 %0 , ptr %3 , align 4 , ! tbaa !5
    store i32 %1 , ptr %4 , align 4 , ! tbaa !5
    %5 = load i32 , ptr %3 , align 4 , ! tbaa !5
    %6 = load i32 , ptr %4 , align 4, ! tbaa !5
    %7 = add nsw i32 %5 , %6
    ret i32 %7
    }
    attributes #0 = { minsize nounwind optsize uwtable " min - legal - vector - width "="0" "no - trapping - math "=" true "
    ,
    " stack - protector - buffer - size "="8" " target - cpu "=" x86 -64" " target - features "="+ cmov ,+ cx8 ,+ fxsr ,+ mmx ,+ sse
    ,
    →,+ sse2 ,+ x87 " " tune - cpu "=" generic " }
    ! llvm . module . flags = !{!0 , !1 , !2 , !3}
    ! llvm . ident = !{!4}
    !0 = !{ i32 1, !" wchar_size ", i32 4}
    !1 = !{ i32 8, !" PIC Level ", i32 2}
    !2 = !{ i32 7, !" PIE Level ", i32 2}
    !3 = !{ i32 7, !" uwtable " , i32 2}
    !4 = !{!" clang version 17.0.6 ( git@github . com : fairinternal CodeGen . git
    ,
    →b05db9bbf7a92019267416c1bb9996fe6134e3f1 ) "}
    !5 = !{!6 , !6 , i64 0}
    !6 = !{!" int ", !7 , i64 0}
    !7 = !{!" omnipotent char ", !8 , i64 0}
    !8 = !{!" Simple C/C ++ TBAA "}
    </ code >
    The input code has instruction count 8 and binary size 65 bytes . [/ INST ]
    Answer:
    The LLVM - IR will have instruction count 2 and binary sise 53 bytes :
    <code >; ModuleID = ’< stdin >’
    source_filename = " -"
    target datalayout = "e -m:e - p270 :32:32 - p271 :32:32 - p272 :64:64 - i64 :64 - f80 :128 - n8 :16:32:64 - S128 "
    target triple = " x86_64 - unknown - linux - gnu "
    ; Function Attrs : minsize mustprogress nofree norecurse nosync nounwind optsize willreturn memory ( none )
    ,
    →uwtable
    define dso_local i32 @add_two ( i32 noundef %0 , i32 noundef %1) local_unnamed_addr #0 {
    %3 = add nsw i32 %1 , %0
    ret i32 %3
    }
    attributes #0 = { minsize mustprogress nofree norecurse nosync nounwind optsize willreturn memory ( none )
    ,
    →uwtable " min - legal - vector - width "="0" "no - trapping - math "=" true " " stack - protector - buffer - size "="8" "
    ,
    →target - cpu "=" x86 -64" " target - features "="+ cmov ,+ cx8 ,+ fxsr ,+ mmx ,+ sse ,+ sse2 ,+ x87 " " tune - cpu "=" generic " }
    ! llvm . module . flags = !{!0 , !1 , !2 , !3}
    ! llvm . ident = !{!4}
    !0 = !{ i32 1, !" wchar_size ", i32 4}
    !1 = !{ i32 8, !" PIC Level ", i32 2}
    !2 = !{ i32 7, !" PIE Level ", i32 2}
    !3 = !{ i32 7, !" uwtable " , i32 2}
    !4 = !{!" clang version 17.0.6 ( git@github . com : fairinternal CodeGen . git
    ,
    →b05db9bbf7a92019267416c1bb9996fe6134e3f1 ) "}
    </ code >


    Prompt:
    [ INST ] Give the assembly for the following code when optimized using opt -p ’ module ( default <Oz >) ’:
    <code >; ModuleID = ’< stdin >’
    source_filename = " -"
    target datalayout = "e -m:e - p270 :32:32 - p271 :32:32 - p272 :64:64 - i64 :64 - f80 :128 - n8 :16:32:64 - S128 "
    target triple = " x86_64 - unknown - linux - gnu "
    ; Function Attrs : minsize nounwind optsize uwtable
    define dso_local i32 @add_two ( i32 noundef %0 , i32 noundef %1) #0 {
    %3 = alloca i32 , align 4
    %4 = alloca i32 , align 4
    store i32 %0 , ptr %3 , align 4 , ! tbaa !5
    store i32 %1 , ptr %4 , align 4 , ! tbaa !5
    %5 = load i32 , ptr %3 , align 4 , ! tbaa !5
    %6 = load i32 , ptr %4 , align 4, ! tbaa !5
    %7 = add nsw i32 %5 , %6
    ret i32 %7
    }
    attributes #0 = { minsize nounwind optsize uwtable " min - legal - vector - width "="0" "no - trapping - math "=" true "
    ,
    → " stack - protector - buffer - size "="8" " target - cpu "=" x86 -64" " target - features "="+ cmov ,+ cx8 ,+ fxsr ,+ mmx ,+ sse
    ,
    →,+ sse2 ,+ x87 " " tune - cpu "=" generic " }
    ! llvm . module . flags = !{!0 , !1 , !2 , !3}
    ! llvm . ident = !{!4}
    !0 = !{ i32 1, !" wchar_size ", i32 4}
    !1 = !{ i32 8, !" PIC Level ", i32 2}
    !2 = !{ i32 7, !" PIE Level ", i32 2}
    !3 = !{ i32 7, !" uwtable " , i32 2}
    !4 = !{!" clang version 17.0.6 ( git@github . com : fairinternal CodeGen . git
    ,
    →b05db9bbf7a92019267416c1bb9996fe6134e3f1 ) "}
    !5 = !{!6 , !6 , i64 0}
    !6 = !{!" int ", !7 , i64 0}
    !7 = !{!" omnipotent char ", !8 , i64 0}
    !8 = !{!" Simple C/C ++ TBAA "}
    </ code >
    The input code has instruction count 8 and binary size 65 bytes . [/ INST ]
    Answer:
    The assembly will have instruction count 2 and binary sise 53 bytes :
    <code > . text
    . file " -"
    . globl add_two # -- Begin function add_two
    . type add_two , @function
    add_two : # @add_two
    . cfi_startproc
    # % bb .0:
    movl % esi , % eax
    addl % edi , % eax
    retq
    . Lfunc_end0 :
    . size add_two , . Lfunc_end0 - add_two
    . cfi_endproc
    # -- End function
    . ident " clang version 17.0.6 ( git@github . com : fairinternal CodeGen . git
    ,
    →b05db9bbf7a92019267416c1bb9996fe6134e3f1 )"
    . section ". note . GNU - stack " ,"" , @progbits
    . addrsig
    </ code >

    代码大小

    我们使用两个指标来衡量代码大小:IR 指令数和二进制大小。二进制大小通过将 IR 或汇编降级为目标文件后,.TEXT 和 .DATA 段大小的总和计算得出。我们排除 .BSS 段,因为它不影响磁盘上的大小。

    优化 pass

    在这项工作中,我们针对 LLVM 17.0.6,并使用新 pass 管理器(PM, 2021),它为不同级别(如模块、函数、循环等)以及转换和分析 pass 分类 pass。转换 pass 改变给定的输入 IR,而分析 pass 生成影响后续转换的信息。

    在 opt 的 346 个可能的 pass 参数中,我们选择了 167 个使用。这包括每个默认优化流水线(例如 module(default<Oz>)),单独的优化转换 pass (例如module(constmerge)),但排除了非优化实用程序 pass (例如module(dot-callgraph))和不保留语义的转换 pass (例如module(internalize))。我们排除了分析 pass,因为它们没有副作用,我们依赖 pass 管理器根据需要注入依赖的分析 pass。对于接受参数的 pass,我们使用默认值(例如module(licm<allowspeculation>))。表 9 包含了所有使用的 pass 列表。我们使用 LLVM 的 opt 工具应用 pass 列表,并使用 clang 将结果 IR 降级为目标文件。清单 1 显示了使用的命令。

      opt input.bc -o output.bc -p 'module(default<Oz>),module(iroutliner)'
      clang output.bc -o output.o
      size output.o

      数据集

      我们通过对表 2 中总结的未优化程序应用 1 到 50 个随机优化 pass 列表生成编译器模拟数据集。每个 pass 列表的长度是均匀随机选择的。pass 列表是通过从上述 167 个 pass 集合中均匀采样生成的。导致编译器崩溃或在 120 秒后超时的 pass 列表被排除。

      3. LLM Compiler FTD:扩展到下游编译器任务


      3.1 优化标志调优的指令微调

      操作编译器标志众所周知对运行时性能和代码大小都有相当大的影响(Fursin et al., 2005)。我们训练 LLM Compiler FTD 模型执行下游任务,即为 LLVM 的 IR 优化工具 opt 选择标志,以生成最小的代码大小。标志调优的机器学习方法以前已经显示出良好的结果,但在不同程序之间的泛化方面存在困难(Cummins et al., 2022)。以前的工作通常需要编译新程序数十或数百次,以尝试不同的配置并找出性能最佳的选项。我们通过预测标志来最小化未见程序的代码大小,在这个任务的零样本版本上训练和评估 LLM Compiler FTD 模型。我们的方法不依赖于所选择的编译器和优化指标,我们打算在未来针对运行时性能。目前,优化代码大小简化了训练数据的收集。

      任务规范

      我们向 LLM Compiler FTD 模型呈现一个未优化的 LLVM-IR(由 clang 前端生成),并要求它生成应该应用的 opt 标志列表,这些优化应用前后的二进制大小,以及输出代码。如果无法对输入代码进行改进,则生成一个只包含未优化二进制大小的简短输出消息。清单 4 和 5 提供了此任务的提示和输出模板。

        Prompt:
        [ INST ] Tell me how to optimize this LLVM - IR for object file size :
        <code >; ModuleID = ’< stdin >’
        source_filename = " -"
        target datalayout = "e -m:e - p270 :32:32 - p271 :32:32 - p272 :64:64 - i64 :64 - f80 :128 - n8 :16:32:64 - S128 "
        target triple = " x86_64 - unknown - linux - gnu "
        ; Function Attrs : minsize nounwind optsize uwtable
        define dso_local i32 @add_two ( i32 noundef %0 , i32 noundef %1) #0 {
        %3 = alloca i32 , align 4
        %4 = alloca i32 , align 4
        store i32 %0 , ptr %3 , align 4 , ! tbaa !5
        store i32 %1 , ptr %4 , align 4 , ! tbaa !5
        %5 = load i32 , ptr %3 , align 4 , ! tbaa !5
        %6 = load i32 , ptr %4 , align 4, ! tbaa !5
        %7 = add nsw i32 %5 , %6
        ret i32 %7
        }
        attributes #0 = { minsize nounwind optsize uwtable " min - legal - vector - width "="0" "no - trapping - math "=" true "
        ,
        " stack - protector - buffer - size "="8" " target - cpu "=" x86 -64" " target - features "="+ cmov ,+ cx8 ,+ fxsr ,+ mmx ,+ sse
        ,
        →,+ sse2 ,+ x87 " " tune - cpu "=" generic " }
        ! llvm . module . flags = !{!0 , !1 , !2 , !3}
        ! llvm . ident = !{!4}
        !0 = !{ i32 1, !" wchar_size ", i32 4}
        !1 = !{ i32 8, !" PIC Level ", i32 2}
        !2 = !{ i32 7, !" PIE Level ", i32 2}
        !3 = !{ i32 7, !" uwtable " , i32 2}
        !4 = !{!" clang version 17.0.6 ( git@github . com : fairinternal CodeGen . git
        ,
        →b05db9bbf7a92019267416c1bb9996fe6134e3f1 ) "}
        !5 = !{!6 , !6 , i64 0}
        !6 = !{!" int ", !7 , i64 0}
        !7 = !{!" omnipotent char ", !8 , i64 0}
        !8 = !{!" Simple C/C ++ TBAA "} </ code > [/ INST ]
        Answer:
        The code has object file size of 65 which can be reduced to 53 by running :
        ‘ opt -p ’ module ( default <Oz >) ’‘
        This will produce code :
        <code >; ModuleID = ’< stdin >’
        source_filename = " -"
        target datalayout = "e -m:e - p270 :32:32 - p271 :32:32 - p272 :64:64 - i64 :64 - f80 :128 - n8 :16:32:64 - S128 "
        target triple = " x86_64 - unknown - linux - gnu "
        ; Function Attrs : minsize mustprogress nofree norecurse nosync nounwind optsize willreturn memory ( none )
        ,
        →uwtable
        define dso_local i32 @add_two ( i32 noundef %0 , i32 noundef %1) local_unnamed_addr #0 {
        %3 = add nsw i32 %1 , %0
        ret i32 %3
        }
        attributes #0 = { minsize mustprogress nofree norecurse nosync nounwind optsize willreturn memory ( none )
        ,
        →uwtable " min - legal - vector - width "="0" "no - trapping - math "=" true " " stack - protector - buffer - size "="8" "
        ,
        →target - cpu "=" x86 -64" " target - features "="+ cmov ,+ cx8 ,+ fxsr ,+ mmx ,+ sse ,+ sse2 ,+ x87 " " tune - cpu "=" generic " }
        ! llvm . module . flags = !{!0 , !1 , !2 , !3}
        ! llvm . ident = !{!4}
        !0 = !{ i32 1, !" wchar_size ", i32 4}
        !1 = !{ i32 8, !" PIC Level ", i32 2}
        !2 = !{ i32 7, !" PIE Level ", i32 2}
        !3 = !{ i32 7, !" uwtable " , i32 2}
        !4 = !{!" clang version 17.0.6 ( git@github . com : fairinternal / CodeGen . git
        ,
        →b05db9bbf7a92019267416c1bb9996fe6134e3f1 ) "} </ code >


        Prompt:
        [ INST ] Tell me how to optimize this LLVM - IR for object file size :
        <code >; ModuleID = ’< stdin >’
        source_filename = " -"
        target datalayout = "e -m:e - p270 :32:32 - p271 :32:32 - p272 :64:64 - i64 :64 - f80 :128 - n8 :16:32:64 - S128 "
        target triple = " x86_64 - unknown - linux - gnu "
        ; Function Attrs : minsize mustprogress nofree norecurse nosync nounwind optsize willreturn memory ( none )
        ,
        →uwtable
        define dso_local i32 @add_two ( i32 noundef %0 , i32 noundef %1) local_unnamed_addr #0 {
        %3 = add nsw i32 %1 , %0
        ret i32 %3
        }
        attributes #0 = { minsize mustprogress nofree norecurse nosync nounwind optsize willreturn memory ( none )
        ,
        →uwtable " min - legal - vector - width "="0" "no - trapping - math "=" true " " stack - protector - buffer - size "="8" "
        ,
        →target - cpu "=" x86 -64" " target - features "="+ cmov ,+ cx8 ,+ fxsr ,+ mmx ,+ sse ,+ sse2 ,+ x87 " " tune - cpu "=" generic " }
        ! llvm . module . flags = !{!0 , !1 , !2 , !3}
        ! llvm . ident = !{!4}
        !0 = !{ i32 1, !" wchar_size ", i32 4}
        !1 = !{ i32 8, !" PIC Level ", i32 2}
        !2 = !{ i32 7, !" PIE Level ", i32 2}
        !3 = !{ i32 7, !" uwtable " , i32 2}
        !4 = !{!" clang version 17.0.6 ( git@github . com : fairinternal / CodeGen . git
        ,
        →b05db9bbf7a92019267416c1bb9996fe6134e3f1 ) "} </ code > [/ INST ]
        Answer:
        The code has object file size of 53 which cannot be reduced further

        我们使用了与编译器模拟任务相同的受限优化 pass 集,并以相同的方式计算二进制大小。

        图 3 说明了用于生成训练数据的过程(如下所述)以及如何在推理时使用模型。

        在评估时只需要生成的 pass 列表。我们从模型输出中提取 pass 列表,并使用给定的参数运行 opt。然后我们可以评估模型预测的二进制大小和优化输出代码的准确性,但这些是辅助学习任务,不是使用所必需的。

        正确性

        LLVM 的优化器并非没有错误,以意外或未经测试的顺序运行优化 pass 可能会暴露出微妙的正确性错误,从而降低模型的实用性。为了缓解这种风险,我们开发了 PassListEval,这是一个工具,用于帮助自动识别破坏程序语义或导致编译器崩溃的 pass 列表。图 4 显示了该工具的概览。

        PassListEval 接受候选 pass 列表作为输入,并在 164 个自测试 C++ 程序套件上评估它,这些程序取自 HumanEval-X(Zheng et al., 2023)。每个程序都包含一个编程挑战的参考解决方案,例如"检查给定数字向量中是否有两个数字之间的距离小于给定阈值",以及验证正确性的单元测试套件。我们将候选 pass 列表应用于参考解决方案,然后将它们与测试套件链接以生成二进制文件。执行时,如果任何测试失败,二进制文件将崩溃。如果任何二进制崩溃,或者任何编译器调用失败,我们就拒绝该候选 pass 列表。

        数据集

        我们在源自 450 万个未优化 IR 的标志调优示例数据集上训练了 LLM Compiler FTD 模型。这些 IR 用于预训练。为生成每个程序的最佳 pass 列表示例,我们进行了广泛的迭代编译过程,如图 3 所示。

        1. 我们使用大规模随机搜索为程序生成初始候选最佳 pass 列表。对每个程序,我们独立生成最多 50 个 pass 的随机列表,从之前描述的 167 个可搜索 pass 集合中均匀采样。每次我们评估一个程序的 pass 列表时,都记录生成的二进制大小。然后选择产生最小二进制大小的每个程序 pass 列表。我们运行了 220 亿次独立编译,平均每个程序 4,877 次。

        2. 随机搜索生成的 pass 列表可能包含冗余 pass,这些 pass 对最终结果没有影响。此外,一些 pass 顺序是可交换的,重新排序不会影响最终结果。由于这些会在我们的训练数据中引入噪声,我们开发了一个最小化过程,并将其应用于每个 pass 列表。最小化包括三个步骤:冗余 pass 消除、冒泡排序和插入搜索。在冗余 pass 消除中,我们通过迭代删除单个 pass 来最小化最佳 pass 列表,看它们是否对二进制大小有贡献。如果没有,就丢弃它们。重复此过程,直到不能再丢弃 pass。然后冒泡排序尝试为 pass 子序列提供统一排序,根据关键字对 pass 进行排序。最后,插入排序通过遍历 pass 列表中的每个 pass 并尝试在其之前插入 167 个搜索 pass 中的每一个来执行局部搜索。如果这样做改善了二进制大小,就保留这个新的 pass 列表。整个最小化管道循环直到达到固定点。最小化后的 pass 列表长度分布如图 9 所示。平均 pass 列表长度为 3.84。

        3. 我们将之前描述的 PassListEval 应用于候选最佳 pass 列表。通过这个过程,我们确定了 1,704,443 个独立 pass 列表中的 167,971 个(9.85%)会导致编译时或运行时错误。

        4. 我们将 100 个最常见的最优 pass 列表广播到所有程序,如果发现改进就更新每个程序的最佳 pass 列表。之后,唯一最佳 pass 列表的总数从 1,536,472 减少到 581,076。

        上述自动调优管道相比 -Oz 产生了 7.1% 的几何平均二进制大小减少。图 10 显示了单个 pass 的频率。对我们来说,这种自动调优作为每个程序优化的黄金标准。虽然发现的二进制大小节省很显著,但这需要 280 亿次额外编译,计算成本超过 21,000 个 CPU 天。对 LLM Compiler FTD 进行指令微调以执行标志调优任务的目标是在不需要运行编译器数千次的情况下达到自动调优器性能的一部分。

        3.2 反汇编的指令微调

        将代码从汇编提升回更高级结构的能力使得可以对直接与应用程序代码集成的库代码运行额外的优化,或将遗留代码移植到新架构。反编译领域在应用机器学习技术生成可读和准确的代码方面取得了进展。几项研究探索了机器学习在反编译任务中的应用,如将二进制文件提升到中间表示以针对合成 C 程序进行评估(Cao et al., 2022),利用遗传算法等进化方法进行程序分析(Schulte et al., 2018),以及提出 XLIR 等方法匹配不同编程语言的二进制代码(Gui et al., 2022)。Armengol-Estapé et al.(2024)训练了一个语言模型将 x86 汇编反编译为高级 C 代码。在本研究中,我们通过对 LLM Compiler FTD 进行反汇编微调,展示了它如何学习汇编代码和编译器 IR 之间的关系。任务是学习 clang -xir - -o - -S 的逆转换,如图 5 所示。

        往返转换

        使用 LLM 进行反汇编会导致正确性问题。提升的代码必须由等价性检查器验证(并非总是可行),或手动验证正确性,或进行足够的测试用例以提供信心。然而,可以通过往返转换找到正确性的下限。也就是说,通过将提升的 IR 编译回汇编,如果汇编是相同的,那么 IR 是正确的。这为使用 LLM 的结果提供了一种简单的方法,也为衡量反汇编模型的实用性提供了一种简单的方法。

        任务规范

        我们向模型提供汇编代码,并训练它生成相应的反汇编 IR。清单 7 显示了提示格式。此任务的上下文长度设置为输入汇编代码 8k tokens 和输出 IR 8k tokens。

          Prompt:
          [ INST ] Disassemble this code to LLVM - IR :
          <code > . text
          . file " -"
          . globl add_two # -- Begin function add_two
          . type add_two , @function
          add_two : # @add_two
          . cfi_startproc
          # % bb .0:
          movl % edi , -4(% rsp )
          movl % esi , -8(% rsp )
          movl -4(% rsp ) , % eax
          addl -8(% rsp ) , % eax
          retq
          . Lfunc_end0 :
          . size add_two , . Lfunc_end0 - add_two
          . cfi_endproc
          # -- End function
          . ident " clang version 17.0.6 ( git@github . com : fairinternal / CodeGen . git
          ,
          →b05db9bbf7a92019267416c1bb9996fe6134e3f1 )"
          . section ". note . GNU - stack " ,"" , @progbits
          . addrsig
          </ code > [/ INST ]
          Answer:
          <code >; ModuleID = ’< stdin >’
          source_filename = " -"
          target datalayout = "e -m:e - p270 :32:32 - p271 :32:32 - p272 :64:64 - i64 :64 - f80 :128 - n8 :16:32:64 - S128 "
          target triple = " x86_64 - unknown - linux - gnu "
          ; Function Attrs : minsize nounwind optsize uwtable
          define dso_local i32 @add_two ( i32 noundef %0 , i32 noundef %1) #0 {
          %3 = alloca i32 , align 4
          %4 = alloca i32 , align 4
          store i32 %0 , ptr %3 , align 4 , ! tbaa !5
          store i32 %1 , ptr %4 , align 4 , ! tbaa !5
          %5 = load i32 , ptr %3 , align 4 , ! tbaa !5
          %6 = load i32 , ptr %4 , align 4, ! tbaa !5
          %7 = add nsw i32 %5 , %6
          ret i32 %7
          }
          attributes #0 = { minsize nounwind optsize uwtable " min - legal - vector - width "="0" "no - trapping - math "=" true "
          ,
          " stack - protector - buffer - size "="8" " target - cpu "=" x86 -64" " target - features "="+ cmov ,+ cx8 ,+ fxsr ,+ mmx ,+ sse
          ,
          →,+ sse2 ,+ x87 " " tune - cpu "=" generic " }
          ! llvm . module . flags = !{!0 , !1 , !2 , !3}
          ! llvm . ident = !{!4}
          !0 = !{ i32 1, !" wchar_size ", i32 4}
          !1 = !{ i32 8, !" PIC Level ", i32 2}
          !2 = !{ i32 7, !" PIE Level ", i32 2}
          !3 = !{ i32 7, !" uwtable " , i32 2}
          !4 = !{!" clang version 17.0.6 ( git@github . com : fairinternal / CodeGen . git
          ,
          →b05db9bbf7a92019267416c1bb9996fe6134e3f1 ) "}
          !5 = !{!6 , !6 , i64 0}
          !6 = !{!" int ", !7 , i64 0}
          !7 = !{!" omnipotent char ", !8 , i64 0}
          !8 = !{!" Simple C/C ++ TBAA "}
          </ code >

          数据集

          我们从与之前任务相同的数据集中派生汇编代码和 IR 对。我们的微调数据集包含 470 万个样本。输入 IR 在降级为 x86 汇编之前已用 -Oz 优化。

          4. 训练参数


          数据通过字节对编码 BPE (Gage, 1994) 进行标记化,使用与 Code Llama、Llama (Touvron et al., 2023a)和 Llama 2 (Touvron et al., 2023b)相同的分词器。

          我们对所有四个训练阶段使用相同的训练参数。我们使用的大多数训练参数与 Code Llama 基础模型相同。我们使用 AdamW (Loshchilov & Hutter, 2017)优化器,β1 和 β2 值为 0.9 和 0.95。我们使用余弦学习率,预热步骤为 1000,并将最终学习率设置为峰值学习率的 1/30。与 Code Llama 基础模型相比,我们将单个序列的上下文长度从 4,096 增加到 16,384,但保持批量大小恒定为 4M tokens。为了适应更长的上下文,我们将学习率设置为 2e-5,并修改了 RoPE 位置嵌入的参数(Su et al., 2024),我们重置频率,基值 θ=10^6。这些设置与 Code Llama 基础模型的长上下文训练一致。

          5. 评估


          在本节中,我们评估 LLM Compiler 模型在标志调优和反汇编任务、编译器模拟、下一个 token 预测以及软件工程任务上的表现。

          5.1 标志调优任务

          方法: 我们评估 LLM Compiler FTD 在未见程序的优化标志调优任务上的表现,并与 GPT-4 Turbo 和 Code Llama - Instruct 进行比较。我们对每个模型运行推理,从模型输出中提取优化 pass 列表。然后我们使用这个 pass 列表来优化特定程序并记录二进制大小。基线是使用 -Oz 优化时程序的二进制大小。

          对于 GPT-4 Turbo 和 Code Llama - Instruct,我们在提示后附加一个后缀,提供额外上下文以进一步描述问题和预期输出格式。经过一些实验,我们发现清单 6 中显示的提示后缀提供了最佳性能。

            Provide your answer as a list of command line options to opt version 17.0.6 , using the format :
            "$ opt -p ’< passes > ’"
            Where < passes > is a list of passes for the new pass manager , e. g. " function ( dce ) , module ( default <Oz >) ,
            ,
            →function ( load - store - vectorizer ) ".
            Only include the passes list . Do not include file paths or other flags such as -o . Terminate the opt
            ,
            →command line options with a newline .
            Then report the optimized code that will be produced , delimited by <code > and </ code > tags .
            Finally , report the binary size the code before and after optimization using the template :
            " Before optimization : X bytes . After optimization : Y bytes ."
            Where X and Y are placeholders for integer binary sizes in bytes . Binary size is the summation of the .
            ,
            →text and . data segment sizes of the object file generated by ‘ clang -17 output . bc -c ‘ , as reported by
            ,
            →the ‘ size ‘ tool .
            Include no other text in your response .

            所有模型生成的 pass 列表都使用 PassListEval 进行验证,如果验证失败则使用 -Oz 作为替代。为进一步验证模型生成的 pass 列表的正确性,我们链接最终的程序二进制文件,并将其输出与使用保守的 -O2 优化管道优化的基准输出进行差分测试。

            数据集

            我们使用从 MiBench 基准套件(Guthaus et al., 2001)提取的 2,398 个测试提示进行评估。为生成这些提示,我们取构成 24 个 MiBench 基准的所有 713 个翻译单元,并从每个单元生成未优化的 IR。然后我们按照清单 4 的格式将它们格式化为提示。如果生成的提示超过 15k tokens,我们使用 llvm-extract 将代表该翻译单元的 LLVM 模块分割成更小的模块,每个函数一个。这导致 1,985 个提示适合 15k token 上下文窗口,剩下 443 个翻译单元不适合。在计算性能分数时,我们对 443 个被排除的翻译单元使用 -Oz。表 10 总结了基准。

            结果

            表 3 显示了所有模型在标志调优任务上的零样本性能。只有 LLM Compiler FTD 模型比 -Oz 有所改进,13B 参数模型略优于较小的模型,在 61% 的情况下生成比 -Oz 更小的目标文件。

            在某些情况下,模型生成的 pass 列表导致比 -Oz 更大的目标文件大小。例如,LLM Compiler FTD 13B 在 12% 的情况下有退化。这些退化可以通过简单地编译程序两次来避免:一次使用模型生成的 pass 列表,一次使用 -Oz,然后选择产生最佳结果的 pass 列表。通过消除相对于 -Oz 的退化,这些 -Oz 备份分数将 LLM Compiler FTD 13B 相对于 -Oz 的总体改进提高到 5.26%,并使 Code Llama - Instruct 和 GPT-4 Turbo 相对于 -Oz 有适度的改进。图 6 显示了每个模型在各个基准上的性能细分。

            二进制大小准确性

            虽然模型生成的二进制大小预测对实际编译没有影响,但我们可以评估模型在预测优化前后的二进制大小方面的性能,以了解每个模型对优化的理解程度。图 7 显示了结果。

            LLM Compiler FTD 的二进制大小预测与实际情况相关性良好,7B 参数模型对未优化和优化的二进制大小分别达到了 0.083 和 0.225 的 MAPE 值。13B 参数模型的 MAPE 值相似,分别为 0.082 和 0.225。Code Llama - Instruct 和 GPT-4 Turbo 的二进制大小预测与实际情况几乎没有相关性。我们注意到,LLM Compiler FTD 对优化代码的错误略高于未优化代码。特别是 LLM Compiler FTD 偶尔有高估优化效果的趋势,导致预测的二进制大小低于实际情况。

            消融研究

            表 4 对模型在 500 个提示的小型保留验证集上的性能进行了消融研究,这些提示来自与我们训练数据相同的分布(但未在训练中使用)。我们在图 1 所示训练管道的每个阶段进行标志调优训练,以比较性能。如图所示,反汇编训练导致性能从平均 5.15% 略微下降到 5.12%(相对于 -Oz 的改进)。我们还展示了用于生成第 2 节所述训练数据的自动调优器的性能。LLM Compiler FTD 达到了自动调优器 77% 的性能。

            5.2 反汇编任务

            方法

            我们评估 LLM 生成的代码在将汇编代码反汇编到 LLVM-IR 时的功能正确性。与 5.1 节一样,我们评估 LLM Compiler FTD 并与 Code Llama - Instruct 和 GPT-4 Turbo 进行比较,发现需要额外的提示后缀(如清单 8 所示)才能从这些模型中提取最佳性能。

              Use LLVM version 17.0.6. Provide the IR enclosed by <code > and </ code > tags .
              Include no other text .

              后缀提供了关于任务和预期输出格式的额外上下文。为评估模型的性能,我们将模型生成的反汇编 IR 往返降级回汇编。这使我们能够通过比较原始汇编与往返结果的 BLEU 分数(Papineni et al., 2002)来评估反汇编的准确性。从汇编到 IR 的无损完美反汇编将有 1.0 的往返 BLEU 分数(精确匹配)。

              数据集

              我们使用从 MiBench 基准套件提取的 2,015 个测试提示进行评估。我们取用于上述标志调优评估的 2,398 个翻译单元,生成反汇编提示。然后我们根据最大 8k token 长度过滤提示,允许 8k tokens 用于模型输出,剩下 2,015 个。表 11 总结了基准。

              结果

              表 5 显示了模型在反汇编任务上的性能。

              LLM Compiler FTD 7B 的往返成功率略高于 LLM Compiler FTD 13B,但 LLM Compiler FTD 13B 具有最高的往返汇编准确性(往返 BLEU)和最频繁产生完美反汇编(往返精确匹配)。Code Llama - Instruct 和 GPT-4 Turbo 在生成语法正确的 LLVM-IR 方面存在困难。图 8 显示了所有模型的往返 BLEU 分数分布。

              消融研究

              表 6 对模型在 500 个提示的小型保留验证集上的性能进行了消融研究,这些提示取自之前使用的 MiBench 数据集。

              我们在图 1 所示训练管道的每个阶段进行反汇编训练,以比较性能。往返率在通过整个训练数据堆栈时最高,并随每个训练阶段持续下降,尽管往返 BLEU 在每个阶段变化不大。

              5.3 基础模型任务

              方法

              我们在下一个 token 预测和编译器模拟两个基础模型任务上对 LLM Compiler 模型进行消融研究。我们在训练管道的每个阶段进行这种评估,以了解为每个连续任务训练如何影响性能。对于下一个 token 预测,我们在所有优化级别的 LLVM-IR 和汇编代码的小样本上计算困惑度。我们使用两个指标评估编译器模拟:生成的 IR 或汇编代码是否编译,以及生成的 IR 或汇编代码是否与编译器产生的完全匹配。

              数据集

              对于下一个 token 预测,我们使用从与我们训练数据相同分布但未用于训练的小型保留验证数据集。我们使用混合的优化级别,包括未优化代码、用 -Oz 优化的代码和随机生成的 pass 列表。对于编译器模拟,我们使用从 MiBench 生成的 500 个提示进行评估,这些提示使用第 2.2 节描述的方式随机生成的 pass 列表。

              结果

              表 7 显示了 LLM Compiler FTD 在所有训练阶段在两个基础模型训练任务(下一个 token 预测和编译器模拟)上的性能。下一个 token 预测性能在 Code Llama 之后急剧上升,后者几乎没有见过 IR 和汇编,并在随后的每个微调阶段略有下降。

              对于编译器模拟,Code Llama 基础模型和预训练模型表现不佳,因为它们没有在这个任务上训练过。在编译器模拟训练之后直接达到最高性能,其中 LLM Compiler FTD 13B 生成的 95.6% 的 IR 和汇编可以编译,20% 与编译器完全匹配。在进行标志调优和反汇编微调后,性能下降。

              5.4 软件工程任务

              方法

              虽然 LLM Compiler FTD 的目的是为代码优化提供基础模型,但它建立在为软件工程任务训练的基础 Code Llama 模型之上。为评估 LLM Compiler FTD 的额外训练如何影响代码生成的性能,我们使用与 Code Llama 相同的基准套件,评估 LLM 从自然语言提示生成 Python 代码的能力,如"编写一个函数,找出可以从给定的对集合形成的最长链。"

              数据集

              我们使用 HumanEval (Chen et al., 2021)和 MBPP (Austin et al., 2021)基准,与 Code Llama 相同。

              结果

              表 8 显示了从 Code Llama 基础模型开始的所有模型训练阶段和模型大小的贪婪解码性能(pass@1)。它还显示了模型在 pass@10 和 pass@100 上的分数,这些分数是用 p=0.95 和 temperature=0.6 生成的。每个以编译器为中心的训练阶段都导致 Python 编程能力略有退化。在 HumanEval 和 MBPP 上,LLM Compiler 的 pass@1 性能最多下降 18% 和 5%,LLM Compiler FTD 在额外的标志调优和反汇编微调后最多下降 29% 和 22%。所有模型在这两个任务上仍然优于 Llama 2。

              6. 相关工作


              代码语言模型

              对源代码推理和生成的 LLM 越来越受关注(Jiang et al., 2024; Hou et al., 2023)。这一领域进展的主要推动力是预训练的基础模型,供其他人在此基础上构建,包括Code Llama (Rozière et al., 2023)、StarCoder (Lozhkov et al., 2024)、Magicoder (Wei et al., 2024)、DeepSeek-Coder (Guo et al., 2024)、GPT-4 (OpenAI, 2023)等(Wang et al.,2023年;Allal等人,2023年;冯等人,2020年)。一些现有的模型是开源的(Rozière等人,2023年;Lozhkov等人,2024年;魏等人,2024年;Allal等人,2023年),而其他一些则是闭源的(陈等人,2021年;OpenAI,2023年;李等人,2022年;Gunasekar等人,2023年)。我们通过一系列专门针对中间代码表示进行训练的模型扩展了代码基础模型的集合,这些模型的许可证允许广泛重用。

              语言模型已经被应用于执行程序模糊测试(夏等人,2023年a;邓等人,2023年)、测试生成(Schäfer等人,2023年)、自动程序修复(夏等人,2023年b)和源级算法优化(Madaan等人,2023年)。引入填充中间部分的能力对软件工程用例(如代码补全)特别有用,并已在最近的代码模型中变得普遍,如InCoder(Fried等人,2023年)、SantaCoder(Allal等人,2023年)、StarCoder(Lozhkov等人,2024年)和Code Llama(Rozière等人,2023年)。虽然已经探索了大量有用的 LLM 应用,但只有很少直接专注于编译任务。

              IR 上的语言模型

              尽管 LLM 在编码任务中得到了广泛应用,但很少有模型在编译器层面上运作。Gallagher等人(2022年)在 LLVM-IR 上训练了一个 RoBERTA 架构,用于代码弱点识别,而 Transcoder-IR(Szafraniec等人,2022年)使用 LLVM-IR 作为源到源转换的枢纽。很少有 LLM 在训练中包含编译器 IR,即使包含,与其他编程语言相比,IR 也只占数据的很小一部分。StarCoder 2(Lozhkov等人,2024年)和 DeepSeek-Coder(郭等人,2024年)在其训练数据中分别包含了 7.7 GB(0.4%)和 0.91 GB(0.1%)的 LLVM-IR。LLM Compiler 预训练了 422 GB 的 LLVM-IR,在微调期间还包含了额外的 LLVM-IR 和汇编代码,这些至少占总训练数据的 85%。

              Paul 等人(2024年)创建了 SLTrans,这是一个包含 26B tokens 的数据集,将高级源代码与相应的 LLVM-IR 配对。像我们的数据集一样,他们的 IR 包括不同的源语言和优化级别,但他们的优化仅限于 -Oz 和 -O3。他们在 SLTrans 的 800 M个标记上训练 IRCoder,并展示了它如何提高底层基础模型的代码推理能力。IRCoder 和 StarCoder 2 向他们的模型呈现 LLVM-IR。我们既包括 LLVM-IR,也包括来自多个源语言和多个架构目标的原生汇编代码。

              随着人们对 IR 越来越感兴趣以提高代码生成模型的性能,新的数据集正在出现。例如 ComPile(Grossman等人,2024年)是一个 2.4 TB 的未优化 LLVM-IR 数据集。

              编译器中的机器学习

              许多工作已经在编译器中应用了机器学习(Leather & Cummins,2020年;Ashouri等人,2022年;Cummins等人,2017年;Phothilimthana等人,2021年;Seeker等人,2024年)。编译器通道排序已被利用了几十年。多年来,使用机器学习的方法有几种(梁等人,2023年;Agakov等人,2006年;Ogilvie等人,2017年;Jayatilaka等人,2021年;Queiroz Jr & da Silva,2023年;Grubisic等人,2024年a)。神经机器翻译是一个新兴领域,使用语言模型将代码从一种语言转换为另一种语言。先前的例子包括将 C 编译为汇编(Armengol-Estapé & O'Boyle,2021年)、汇编到 C(Armengol-Estapé等人,2024年;Hosseini & Dolan-Gavitt,2022年)和源到源(Lachaux等人,2020年)。

              7. 讨论


              在本文中,我们介绍了 LLM Compiler,这是一个专门设计用于解决代码和编译器优化挑战的新型大型语言模型系列。通过扩展基础 Code Llama 模型的能力,LLM Compiler 提供了一个强大的、预训练的平台,显著增强了对编译器中间表示和汇编语言的理解和操作。

              我们根据定制的商业许可发布 LLM Compiler,以促进广泛访问和协作,使学术研究人员和行业从业者都能根据其特定需求探索、修改和扩展该模型。

              模型 HuggingFace 地址:

              https://huggingface.co/collections/facebook/llm-compiler-667c5b05557fe99a9edd25cb

              7.1 局限性

              我们已经展示了 LLM Compiler 在编译器优化任务上表现良好,并且相比先前的工作,对编译器表示和汇编代码的理解有所改进,但仍存在一些局限性。主要限制是输入的有限序列长度(上下文窗口)。LLM Compiler 支持 16k tokens 的上下文窗口,但程序代码可能远远超过这个长度。例如,当格式化为标志调优提示时,67% 的 MiBench 翻译单元超过了这个上下文窗口,如表 10 所示。

              为了缓解这一问题,我们将较大的翻译单元拆分为单独的函数,尽管这限制了可以执行的优化范围,而且仍有 18% 的拆分翻译单元对模型来说太大,无法作为输入接受。研究人员正在采用不断增加的上下文窗口(丁等人,2023年),但有限的上下文窗口仍然是 LLM 的一个普遍问题。

              第二个限制,也是所有 LLM 的共同问题,是模型输出的准确性。建议 LLM Compiler 的用户使用特定于编译器的评估基准来评估他们的模型。鉴于编译器并非无 bug,任何建议的编译器优化都必须经过严格测试。当模型反编译汇编代码时,其准确性应通过往返、人工检查或单元测试来确认。对于某些应用,LLM 生成可以被限制在正则表达式内(Grubisic等人,2024年b),或与自动验证相结合以确保正确性(Taneja等人,2024年)。

              (全文完)

              文章转载自慢慢学 AIGC,如果涉嫌侵权,请发送邮件至:contact@modb.pro进行举报,并提供相关证据,一经查实,墨天轮将立刻删除相关内容。

              评论