这篇文章主要介绍了在机器学习的 Python 程序和包中很流行的“JIT 装饰器”模式，通过几个例子和不同的实现策略来解释其工作原理，包括 AST 基于的 JIT、字节码基于的 JIT 和跟踪基于的 JIT。

主要观点：

• “JIT 装饰器”模式在机器学习的 Python 代码中很常见。
• 不同的 JIT 实现策略各有特点和适用场景。

关键信息和重要细节：

• JAX 示例：通过@jax.jit装饰器实现函数的 JIT 编译，如add函数，其代码在 Python 解释器中不会以正常方式执行，而是由库中的特殊编译器处理。
• Triton 语言：直接在 Python 中编写 GPU 内核，通过@triton.jit装饰器，函数的主体被解析为 Python AST，经过一系列内部 IR 最终转化为 LLVM IR 并在 GPU 上运行。

• 实现策略：

  • AST 基于的 JIT：通过astjit装饰器，获取函数的 AST 并使用_ExprCodeEmitter将其转换为 Expr，再通过llvm_jit_evaluate转换为 LLVM IR 并执行。例如 Triton 语言就是采用这种方式，其函数体经过一系列转换后在 GPU 上运行。
  • 字节码基于的 JIT：bytecodejit装饰器将函数编译到字节码，然后从字节码开始转换为 Expr IR 再到 LLVM IR 执行。Numba 编译器的工作原理类似，它使用 Python 编译器生成字节码，然后转换为自己的 IR 再到 LLVM。
  • 跟踪基于的 JIT：tracejit装饰器通过特殊的参数框来跟踪函数的执行，构建 Expr IR，再通过llvm_jit_evaluate执行。JAX 框架采用类似的跟踪方法，但其在处理数据依赖的控制流时有限制，需要使用 JAX 自己的内置函数如lax.fori_loop。
• 代码和资源：完整代码可在GitHub获取，文中还提到了一些相关的资源和注意事项，如更复杂的 Python 字节码指令支持、llvmlite 是 Numba 的子项目等。

总之，不同的 JIT 实现策略各有优缺点，在不同的场景下可以选择合适的方法来提高代码的执行效率。