Numba 功能和性能优化指南¶

本学习指南旨在帮助您全面理解 Numba 库的核心功能、编译模式、优化策略以及与外部代码的交互方式。

一、核心概念与编译模式¶

1. Numba 的基本作用是什么？¶

Numba 是一个开源的即时编译器，可以将 Python 代码（尤其是数值计算代码）转换为优化的机器码，从而显著提高执行速度。它通过 JIT (Just-in-Time) 编译技术，在代码运行时分析并编译函数。

2. Numba 的两种主要编译模式是什么？¶

NoPython 模式 (nopython mode)：这是 Numba 的推荐和默认模式（自 Numba 0.59 起）。它对 Python 代码有严格的限制，只允许使用 Numba 支持的 Python 特性和 NumPy 函数。优点是生成的代码速度最快，因为它完全绕过了 Python 解释器。
Object 模式 (object mode)：当 Numba 无法在 NoPython 模式下编译函数时，它会回退到 Object 模式（旧版本默认行为，新版本需要 forceobj=True 显式启用）。在这种模式下，Numba 会尝试编译尽可能多的代码，但会插入 Python C API 调用来处理不支持的部分。性能提升通常不如 NoPython 模式显著，甚至可能更慢。

3. `@jit` 装饰器的基本用法和编译时机¶

惰性编译 (Lazy Compilation)：当 @jit 装饰器没有传入签名时，编译会在函数第一次被调用时进行。Numba 会根据传入的参数类型推断签名并生成优化代码。
即时编译 (Eager Compilation)：当向 @jit 装饰器传入明确的函数签名时，编译会在装饰器应用时立即进行。这允许对编译器选择的类型进行细粒度控制。

二、编译选项和高级功能¶

1. `@jit` 装饰器的主要编译选项¶

选项	说明
`nopython=True`（或使用 `@njit`）	强制 Numba 在 NoPython 模式下编译。如果编译失败，Numba 会抛出错误而不是回退到 Object 模式。
`nogil=True`	当函数在 NoPython 模式下编译时，Numba 会在进入该函数时释放 Python 的全局解释器锁 (GIL)，从而允许代码与其他线程并发运行。
`cache=True`	指示 Numba 将编译结果缓存到文件中。这样可以避免在每次程序调用时重新编译函数，减少启动时间。
`parallel=True`	为函数中支持的操作启用自动并行化。此选项必须与 `nopython=True` 结合使用，可以在多核 CPU 上并行运行代码。

2. NumPy 通用函数 (ufuncs) 和广义通用函数 (gufuncs)¶

@vectorize 装饰器：用于将接受标量输入参数的 Python 函数编译成 NumPy ufuncs。这些 ufuncs 可以像 C 编写的一样快速操作 NumPy 数组，并自动获得归约、累加和广播等功能。
- 动态通用函数 (DUFuncs)：当 @vectorize 没有传入签名时，它会创建 DUFuncs。这些函数在被调用时会根据输入类型动态编译新的内核。
@guvectorize 装饰器：允许编写操作高维数组和标量，并接受和返回不同维度数组的 ufuncs，例如移动中值或卷积滤波器。与 @vectorize 不同，@guvectorize 函数通过修改数组参数来返回结果。
- 动态广义通用函数 (GUFuncs)：与 DUFuncs 类似，当 @guvectorize 没有传入类型时，它会动态编译内核。

3. C 回调 (`@cfunc`)¶

目的：使用 @cfunc 装饰器可以创建编译后的函数，这些函数可以从外部 C/C++ 代码调用，从而实现 Python 代码与原生库的交互。
用法：需要传入一个明确的 C 回调签名。编译后的 C 函数对象通过 address 属性公开其地址，也可通过 ctypes 属性获取 ctypes 回调对象。
数据处理：可以使用 numba.carray() 或 numba.farray() 函数将 C 指针和大小参数转换为 NumPy 数组视图。
C 结构体：Numba 支持通过 cffi_utils.map_type 或 numba.types.Record.make_c_struct 处理 C 结构体。

4. 自动并行化 (`@jit(parallel=True)`)¶

Numba 的自动并行化功能会尝试对函数中的某些数组操作（如逐元素运算、NumPy ufuncs、归约函数、数组创建和 dot 函数）进行并行化。
prange：可以使用 numba.prange 代替 range 来显式指定循环可以并行化。Numba 会自动推断并处理循环内的支持归约操作。
诊断报告：通过设置环境变量 NUMBA_PARALLEL_DIAGNOSTICS 或调用 .parallel_diagnostics() 方法，可以获取关于并行化转换的详细诊断信息，包括循环融合、循环序列化和循环不变代码外提等。
调度：Numba 默认使用静态调度。通过 numba.set_parallel_chunksize() 和 numba.parallel_chunksize() 上下文管理器，可以控制每个线程分配的迭代块大小，以优化负载均衡（特别是在 tbb 后端）。

三、性能优化提示¶

NoPython 模式：始终尝试在 NoPython 模式下编译代码以获得最佳性能。使用 @njit 或 @jit(nopython=True)。
循环与向量化：Numba 对 Python 循环的优化能力与 NumPy 的向量化操作相当。对于熟悉 C/Fortran 风格的用户，直接编写循环在 Numba 中同样高效。
LoopLifting (循环提升)：对于无法完全进入 NoPython 模式的函数，可以使用 @jit(forceobj=True) 尝试对其中兼容的循环应用 NoPython 编译，以提升局部性能。
Fastmath：通过设置 fastmath=True，Numba 可以放宽 IEEE 754 标准的数值严格性，从而启用额外的性能优化，例如浮点数重关联。这在对精度要求不那么高的应用中非常有用。
Parallel=True 和 prange：对于可并行化的计算，结合 parallel=True 和 prange 可以显著提高多核 CPU 上的性能。当使用 parallel=True 时，默认假定乱序执行是有效的，因此可以安全地结合 fastmath=True 以进一步提升性能。
Intel SVML：如果系统上存在 Intel 短向量数学库 (SVML)，Numba 会自动利用它来优化超越函数（如 sin, cos, sqrt），显著提升数值计算性能。fastmath 参数会影响 SVML 使用的是高精度还是低精度版本。
线性代数：Numba 在 NoPython 模式下支持大部分 numpy.linalg 操作。为了获得最佳性能，需要使用针对优化良好的 LAPACK 和 BLAS 库（例如 Intel MKL）构建的 SciPy。

四、常见限制和注意事项¶

缓存限制：缓存不是基于函数进行的，主函数调用次函数时，次函数的变化可能不会被检测到。全局变量在编译时被视为常量，缓存加载时不会重新绑定。
并发写入非线程安全容器：在 prange 并行区域中，对列表、集合和字典等容器的并发写入操作不是线程安全的。
prange 和归纳变量：prange 循环中的归纳变量在 parallel=True 时可能被类型化为 uint64，这可能导致与有符号整数操作产生浮点结果类型。
异常控制流：包含异常控制流（如断言）的 prange 循环可能无法并行化。
并行归约竞态条件：当归约到数组的切片或元素中时，如果多个并行线程同时写入同一位置，可能会发生竞态条件。
GUFuncs 广播限制：目前 GUFuncs 尚不支持广播，可能导致不正确行为。

Numba 知识测验¶

说明：请用 2-3 句话简要回答以下问题。

1. Numba 的主要目的是什么？¶

Numba 的主要目的是通过即时编译 (JIT) 技术，将 Python 代码（特别是数值计算代码）转换为优化的机器码，从而显著提高其执行速度。它针对 NumPy 数组操作和 Python 循环进行了优化。

2. 解释 NoPython 模式与 Object 模式之间的关键区别。¶

NoPython 模式强制所有代码都在没有 Python 解释器参与的情况下运行，生成最快的机器码，但对 Python 特性有限制。Object 模式则允许使用所有 Python 特性，但在遇到不支持的代码时会回退到 Python C API 调用，性能通常较低。

3. `@jit` 装饰器的惰性编译和即时编译模式有何不同？¶

惰性编译发生在函数第一次被调用时，Numba 会根据运行时传入的参数类型推断签名并生成代码。即时编译则在 @jit 装饰器应用时立即发生，需要显式指定函数签名，从而实现对类型选择的细粒度控制。

4. 什么时候会使用 `@jit(nogil=True)`？它有什么作用？¶

当 Numba 编译的函数完全在 NoPython 模式下运行时，可以使用 nogil=True 选项。它允许 Numba 在进入该函数时释放 Python 全局解释器锁 (GIL)，从而使多个线程可以并发执行该函数，实现真正的多核并行。

5. `@jit(cache=True)` 选项的主要优势和潜在限制是什么？¶

cache=True 的主要优势是避免每次程序调用时的重复编译时间，通过将编译结果保存到文件系统中来加速后续执行。然而，它的限制包括无法识别被调用函数在不同文件中的更改，以及全局变量被视为编译时常量且不会随之更新。

6. `@vectorize` 装饰器和 `@guvectorize` 装饰器分别适用于什么场景？¶

@vectorize 装饰器用于将接受标量输入并返回标量输出的 Python 函数转换为 NumPy ufuncs，实现对数组的逐元素操作。而 @guvectorize 装饰器则适用于需要操作高维数组或接受/返回不同维度数组的函数，例如移动窗口操作。

7. 在使用 `@cfunc` 创建 C 回调时，如何处理 C 结构体？¶

在使用 @cfunc 处理 C 结构体时，可以通过 numba.core.typing.cffi_utils.map_type 将 cffi 类型转换为 Numba Record 类型，或者手动使用 numba.types.Record.make_c_struct 定义结构体布局。结构体通常作为 types.CPointer(my_struct) 类型的指针传递给 cfunc。

8. `numba.prange` 在自动并行化中扮演什么角色？它与 Python 的 `range` 有何不同？¶

numba.prange 替代 Python 的 range 来显式标记一个循环可以并行化。当 parallel=True 时，Numba 会尝试并行执行 prange 循环，并能自动推断支持的归约操作；如果 parallel=False，它则与 range 的行为相同。

9. `fastmath=True` 参数如何影响 Numba 编译代码的性能和数值行为？¶

fastmath=True 允许 Numba 放宽 IEEE 754 浮点标准的严格性，从而开启额外的优化，例如浮点数重关联。这通常会带来显著的性能提升，但可能会导致数值结果与严格标准计算的结果略有不同，适用于对精度要求不极高的场景。

10. 为了从 Numba 的 `numpy.linalg` 函数中获得最佳性能，需要注意什么？¶

为了从 Numba 的 numpy.linalg 函数中获得最佳性能，必须确保所使用的 SciPy 库是针对优化良好的 LAPACK 和 BLAS 库构建的。在 Anaconda 发行版中，SciPy 通常与 Intel 的 MKL 优化库绑定，这使得 Numba 能够充分利用这些高性能的线性代数例程。

建议的论文问题¶

详细比较 Numba 的 NoPython 模式和 Object 模式在性能、适用性及限制方面的差异。在何种场景下，开发者可能会被迫使用或选择 Object 模式，并应如何衡量其性能影响？
解释 Numba 的自动并行化功能 (parallel=True) 如何工作，并讨论其支持的操作类型和当前存在的限制。此外，请阐述 numba.prange 在显式并行循环中的作用，并提供一个具体例子来说明如何避免并行归约中的竞态条件。
深入探讨 @vectorize 和 @guvectorize 装饰器在创建 NumPy 通用函数方面的能力。说明这两种装饰器各自的最佳应用场景，并讨论动态通用函数（DUFuncs/GUFuncs）的机制及其潜在的优缺点。
Numba 提供了多种优化代码性能的手段，例如 nogil=True、fastmath=True 和利用 Intel SVML。请选择其中至少三个优化点，详细解释它们的工作原理、适用条件以及对代码性能和行为可能产生的影响。
讨论 Numba 与外部 C 代码交互的两种主要方式：从 Numba 调用 C 代码，以及使用 @cfunc 创建 C 回调。详细说明每种交互方式的实现步骤，包括签名规范和数据处理（例如指针和 C 结构体），并分析这些功能对于混合语言编程的重要性。

术语表¶

术语	解释
Numba	一个开源的即时 (Just-in-Time) 编译器，可以将 Python 和 NumPy 代码转换为快速的机器码。
JIT (Just-in-Time) 编译	一种在程序运行时而非编译时进行编译的技术。Numba 使用此技术来优化 Python 函数。
`@jit` 装饰器	Numba 的核心功能装饰器，用于标记 Python 函数以便 Numba 进行即时编译优化。
NoPython 模式 (nopython mode)	Numba 的一种编译模式，强制函数完全在没有 Python 解释器参与的情况下运行，生成高度优化的机器码，但对 Python 语言特性有严格限制。
Object 模式 (object mode)	Numba 的另一种编译模式，当 NoPython 模式失败时作为回退（或显式启用）。它允许使用 Python 对象，但性能提升可能不如 NoPython 模式。
惰性编译 (Lazy Compilation)	`@jit` 的默认行为，编译在函数第一次被调用时进行，Numba 会根据输入类型推断签名。
即时编译 (Eager Compilation)	`@jit` 的一种模式，通过在装饰器中提供函数签名，Numba 会立即编译函数。
全局解释器锁 (GIL)	Python 解释器的一个机制，确保在任何给定时间只有一个线程可以执行 Python 字节码，从而限制了多核 CPU 的并行性。
`nogil=True`	`@jit` 的一个选项，指示 Numba 在 NoPython 模式下编译的函数执行期间释放 GIL。
`cache=True`	`@jit` 的一个选项，指示 Numba 将编译结果缓存到磁盘，以减少后续程序运行时的编译开销。
`parallel=True`	`@jit` 的一个选项，为函数中支持的操作启用自动并行化，通常用于多核 CPU。必须与 `nopython=True` 结合使用。
`@vectorize` 装饰器	用于将接受标量输入和返回标量输出的 Python 函数编译成 NumPy 通用函数 (ufuncs)。
`@guvectorize` 装饰器	用于将操作高维数组并可接受/返回不同维度数组的 Python 函数编译成 NumPy 广义通用函数 (gufuncs)。
动态通用函数 (DUFunc)	由 `@vectorize` 在未指定签名时创建，根据运行时输入类型动态编译内核。
动态广义通用函数 (GUFunc)	由 `@guvectorize` 在未指定类型时创建，根据运行时输入类型动态编译内核。
`@cfunc` 装饰器	用于创建可从外部 C/C++ 代码调用的编译 Python 函数（C 回调）。
`numba.carray()` / `numba.farray()`	Numba 中的函数，用于从 C 指针和形状参数创建 NumPy 数组视图，分别用于 C 顺序和 Fortran 顺序的数据。
Record 类型	Numba 中表示 C 结构体的类型，可以手动创建或从 cffi 类型映射而来。
`prange`	`numba.prange` 是 `range` 的并行化版本，用于在 `parallel=True` 的 jit 函数中显式标记可并行执行的循环。
循环融合 (Loop Fusion)	一种优化技术，将具有等效边界的相邻循环组合成一个循环，以改善数据局部性。
循环序列化 (Loop Serialization)	当存在嵌套的 `prange` 循环时，只有最外层的 `prange` 循环会并行执行，内部的 `prange` 循环会被序列化为标准 `range` 循环。
循环不变代码外提 (Loop Invariant Code Motion)	一种优化，将循环体内不随循环迭代变化的计算移到循环外部，避免重复计算。
分配外提 (Allocation Hoisting)	循环不变代码外提的一种特殊情况，将循环内部的数组分配（例如 `np.zeros()`）移到循环外部，减少分配开销。
调度 (Scheduling)	在并行执行中，如何将任务分配给不同的线程。Numba 默认使用静态调度，也支持通过 `set_parallel_chunksize()` 进行动态调度。
`fastmath=True`	`@jit` 的一个选项，放宽浮点运算的 IEEE 754 标准严谨性，以启用额外的性能优化，如浮点数重关联。
Intel SVML (Short Vector Math Library)	Intel 提供的一个优化数学函数库，Numba 在可用时会自动利用它来加速超越函数计算。
LAPACK (Linear Algebra Package)	一个标准软件库，用于执行数值线性代数运算。
BLAS (Basic Linear Algebra Subprograms)	一个提供向量和矩阵基本操作的标准库，是许多高性能数值计算库的基础。
MKL (Intel Math Kernel Library)	Intel 优化过的数学函数库，包含了高性能的 BLAS 和 LAPACK 实现。

Numba：让 Python 代码飞起来的加速利器¶

Numba 是一种开源的 JIT (Just-In-Time) 编译器，它能够将 Python 代码（特别是 NumPy 和数值计算密集型代码）即时编译成快速的机器码，从而显著提高代码的执行速度。简单来说，Numba 就像一位技艺精湛的翻译官，能够将你写的 Python 代码“翻译”成更高效的低级语言，让你的程序跑得飞快。

Numba 的核心理念：¶

Numba 的核心思想是利用 类型推断 和 LLVM (Low Level Virtual Machine) 来实现性能的飞跃。

类型推断 (Type Inference): Python 是一种动态类型语言，这意味着变量的类型在运行时才能确定。而 Numba 在编译时会尝试推断出函数中变量的类型。一旦类型被确定，Numba 就可以生成更优化的机器码，避免了动态类型带来的开销。
LLVM (Low Level Virtual Machine): LLVM 是一个强大的编译器基础设施，它提供了一个中间表示（IR）和一系列优化pass，可以将代码从一种语言（如 Numba 生成的 IR）转换为多种目标机器的机器码。Numba 利用 LLVM 来生成高性能的原生代码。

Numba 的工作原理：¶

当你在 Numba 中对一个 Python 函数应用 @jit 装饰器后，Numba 会在函数第一次被调用时执行以下步骤：

解析 Python 代码: Numba 解析被装饰的 Python 函数，理解其语法和结构。
类型推断: Numba 尝试推断函数中所有变量、参数和返回值的类型。如果 Numba 无法推断出明确的类型（例如，函数中存在很多复杂的 Python 对象操作），它可能会生成一个“对象模式”的编译版本，其性能提升会比较有限，或者需要用户提供类型提示。
生成 LLVM IR: 基于推断出的类型信息，Numba 将 Python 代码转换为 LLVM 的中间表示 (IR)。
LLVM 优化: LLVM 对 IR 进行大量的优化，例如死代码消除、循环展开、内联函数等，以生成最高效的代码。
生成机器码: LLVM 将优化后的 IR 编译成特定平台的机器码。
缓存和重用: 编译后的机器码会被缓存起来。当函数再次被调用时，如果参数的类型与上次编译时相同，Numba 将直接使用缓存的机器码，无需再次编译，从而实现“即时”执行。

Numba 的主要优点：¶

显著的性能提升: 对于数值计算、科学计算和数据处理等任务，Numba 能够带来数量级的性能提升，甚至可以媲美 C、Fortran 等编译型语言。
易于使用: Numba 的 API 设计非常简单，大多数情况下只需要在函数前加上 @jit 装饰器即可。
与 NumPy 深度集成: Numba 对 NumPy 数组操作有非常好的支持，可以高效地编译 NumPy 的各种运算。
CPU 和 GPU 加速: Numba 不仅可以加速 CPU 上的代码，还支持 GPU 加速，可以通过 @cuda.jit 装饰器将 Python 代码编译成 CUDA 核函数，充分利用 GPU 的并行计算能力。
多核并行化: Numba 提供了 @njit(parallel=True) 选项，可以轻松地将代码并行化到多核 CPU 上，进一步提升性能。
逐步优化: Numba 允许你逐步优化你的 Python 代码，你可以先用纯 Python 编写，然后选择性地使用 Numba 来加速性能瓶颈的代码，而无需重写整个项目。
兼容性: Numba 能够处理大部分 Python 的特性，并且可以很好地与 CPython 运行时集成。

Numba 的核心装饰器：¶

@jit: 这是 Numba 最基本的装饰器。它会将函数编译成机器码。
- nopython=True: 这是 Numba 最核心的优化模式。当设置为 True 时，Numba 会尝试将函数完全编译成机器码，不回退到 Python 对象模式。如果 Numba 无法在 nopython 模式下编译函数（例如，函数使用了 Numba 不支持的 Python 对象操作），它会抛出 TypingError。强烈建议在可能的情况下使用 nopython=True 来获得最佳性能。
- parallel=True: 启用自动并行化，将循环等可以并行执行的代码分发到多个 CPU 核心上。
- cache=True: 启用编译缓存，将编译后的机器码保存到磁盘，避免重复编译。
@njit: 这是 @jit(nopython=True) 的简写。当你想强制 Numba 在 nopython 模式下编译时，可以直接使用 @njit。
@vectorize: 这是一个用于创建 NumPy ufuncs (Universal Functions) 的装饰器。它允许你将一个 Python 函数编译成一个可以应用于 NumPy 数组的优化函数，而不需要显式地编写循环。
@guvectorize: 类似于 @vectorize，但它支持更通用的“广义通用函数”，可以处理形状更复杂的输入和输出。
@cuda.jit: 用于将 Python 函数编译成 CUDA 核函数，以便在 NVIDIA GPU 上执行。

Numba 的局限性：¶

不支持所有 Python 特性: Numba 仍然在不断发展，但它不支持 Python 的所有特性，例如某些复杂的动态特性、垃圾回收机制、I/O 操作、某些标准库函数等。如果你的代码大量依赖于这些特性，Numba 可能无法提供加速，或者需要回退到对象模式。
类型推断的复杂性: 对于非常复杂的函数或涉及许多动态类型操作的函数，Numba 可能难以进行有效的类型推断，导致性能提升有限。
编译开销: 第一次编译函数需要一定的时间，尤其是在 nopython 模式下。如果你的函数只被调用一两次，或者函数本身非常简单，那么编译开销可能大于执行时间，反而会减慢速度。
调试的挑战: Numba 编译后的代码是机器码，直接调试 Numba 编译的代码比调试纯 Python 代码更具挑战性。通常的策略是先在纯 Python 中进行调试，然后应用 Numba 来加速。

Numba 的使用场景：¶

科学计算和数值分析: NumPy、SciPy、Pandas 等库的性能瓶颈通常在于循环密集型的数值计算。Numba 可以显著加速这些计算。
机器学习和深度学习: 在实现自定义层、损失函数或预处理步骤时，Numba 可以提供速度上的优势。
图像处理和信号处理: 这些领域通常涉及大量的数组操作和数值计算。
高性能数据处理: 对于大规模数据的过滤、转换、聚合等操作，Numba 可以大幅提升效率。
并行计算: 利用多核 CPU 或 GPU 进行并行计算，以缩短计算时间。

如何开始使用 Numba：¶

安装:
```
pip install numba
```
如果你需要 GPU 加速，还需要安装 CUDA toolkit，并确保 Numba 能够找到它。

基本使用:

from numba import jit
import numpy as np
import time

@jit(nopython=True) # 强烈建议使用 nopython=True
def sum_array(arr):
    total = 0
    for x in arr:
        total += x
    return total

# 创建一个大的 NumPy 数组
data = np.random.rand(1_000_000)

# 第一次调用会进行编译
start_time = time.time()
result_numba = sum_array(data)
end_time = time.time()
print(f"Numba compiled and executed in {end_time - start_time:.6f} seconds")

# 第二次调用将使用缓存的机器码
start_time = time.time()
result_numba = sum_array(data)
end_time = time.time()
print(f"Numba cached execution in {end_time - start_time:.6f} seconds")

# 对比纯 Python 版本
def sum_array_python(arr):
    total = 0
    for x in arr:
        total += x
    return total

start_time = time.time()
result_python = sum_array_python(data)
end_time = time.time()
print(f"Pure Python execution in {end_time - start_time:.6f} seconds")

print(f"Numba result: {result_numba}")
print(f"Python result: {result_python}")

进阶使用：¶

类型提示 (Type Hinting): 虽然 Numba 可以自动推断类型，但对于复杂的函数，你也可以通过类型提示来帮助 Numba 进行更精确的推断，提高编译效率和生成的代码质量。
```
from numba import types
from numba.extending import typeof_impl

@typeof_impl.register(MyClass)
def typeof_myclass(val, c):
    return types.CPointer(types.void) # 这是一个简化示例
```
（更常见的是在函数签名中指定类型）

GPU 加速:

from numba import cuda
import numpy as np

@cuda.jit
def gpu_add(x, y, out):
    idx = cuda.grid(1)
    if idx < x.shape[0]:
        out[idx] = x[idx] + y[idx]

n = 1000000
x = np.arange(n).astype(np.float32)
y = np.arange(n).astype(np.float32)
out = np.empty_like(x)

# 分配 GPU 内存并拷贝数据
x_device = cuda.to_device(x)
y_device = cuda.to_device(y)
out_device = cuda.device_array_like(x)

# 设置线程块和网格的大小
threadsperblock = 128
blockspergrid = (n + (threadsperblock - 1)) // threadsperblock

gpu_add[blockspergrid, threadsperblock](x_device, y_device, out_device)

# 将结果拷贝回 CPU
result_gpu = out_device.copy_to_host()
print(result_gpu[:10])

总结：¶

Numba 是一个强大且易于使用的工具，它能够显著加速 Python 的数值计算密集型代码。通过利用 JIT 编译和 LLVM，Numba 可以让你的 Python 程序在性能上获得巨大的提升，尤其是在科学计算、数据分析和机器学习等领域。理解 Numba 的工作原理，合理使用其装饰器，并注意其局限性，将能有效地利用 Numba 来优化你的 Python 代码。

In [1]:

Copied!





from numba import jit
import numpy as np
import time
import numba
print(numba.__version__)
from numba import jit
import numpy as np
import time
import numba
print(numba.__version__)

0.61.2

In [7]:

Copied!





    @jit(nopython=True) # 强烈建议使用 nopython=True
    def sum_array(arr):
        total = 0
        for x in arr:
            total += x
        return total

    # 创建一个大的 NumPy 数组
    data = np.random.rand(10**8)

    # 第一次调用会进行编译
    start_time = time.time()
    result_numba = sum_array(data)
    end_time = time.time()
    print(f"Numba compiled and executed in {end_time - start_time:.6f} seconds")

    # 第二次调用将使用缓存的机器码
    start_time = time.time()
    result_numba = sum_array(data)
    end_time = time.time()
    print(f"Numba cached execution in {end_time - start_time:.6f} seconds")

    # 对比纯 Python 版本
    def sum_array_python(arr):
        total = 0
        for x in arr:
            total += x
        return total

    start_time = time.time()
    result_python = sum_array_python(data)
    end_time = time.time()
    print(f"Pure Python execution in {end_time - start_time:.6f} seconds")

    print(f"Numba result: {result_numba}")
    print(f"Python result: {result_python}")

    @jit(nopython=True) # 强烈建议使用 nopython=True
    def sum_array(arr):
        total = 0
        for x in arr:
            total += x
        return total

    # 创建一个大的 NumPy 数组
    data = np.random.rand(10**8)

    # 第一次调用会进行编译
    start_time = time.time()
    result_numba = sum_array(data)
    end_time = time.time()
    print(f"Numba compiled and executed in {end_time - start_time:.6f} seconds")

    # 第二次调用将使用缓存的机器码
    start_time = time.time()
    result_numba = sum_array(data)
    end_time = time.time()
    print(f"Numba cached execution in {end_time - start_time:.6f} seconds")

    # 对比纯 Python 版本
    def sum_array_python(arr):
        total = 0
        for x in arr:
            total += x
        return total

    start_time = time.time()
    result_python = sum_array_python(data)
    end_time = time.time()
    print(f"Pure Python execution in {end_time - start_time:.6f} seconds")

    print(f"Numba result: {result_numba}")
    print(f"Python result: {result_python}")

Numba compiled and executed in 0.207616 seconds
Numba cached execution in 0.139771 seconds
Pure Python execution in 12.250545 seconds
Numba result: 49996809.08289068
Python result: 49996809.08289068

语法规范¶

Numba 的语法规范主要体现在其装饰器 (Decorators) 的使用以及与 NumPy 等库的集成方式上。它本身并没有一套全新的、独立的 Python 语法，而是通过装饰器来“注解”标准的 Python 函数，指示 Numba 如何处理和编译这些函数。

以下是 Numba 语法规范的几个关键方面：

1. 装饰器的使用¶

这是 Numba 最核心的语法。通过在函数定义前加上 @ 符号，你可以将 Numba 的编译器指令应用到函数上。

@jit: 这是最基础和通用的装饰器。
- @jit: 默认情况下，@jit 会尝试使用 nopython 模式，如果失败则回退到对象模式。
- @jit(nopython=True): 强制 Numba 在 nopython 模式下编译。这是获得最佳性能的关键。如果 Numba 无法在 nopython 模式下编译（例如，使用了 Numba 不支持的 Python 对象操作），它会抛出 TypingError。
- @jit(parallel=True): 启用 自动并行化。Numba 会尝试识别函数中的可并行循环，并将其分配给多个 CPU 核心。
- @jit(cache=True): 启用 编译缓存。Numba 会将编译后的机器码保存到磁盘，避免下次运行时重复编译。
- @jit(forceobj=True): 强制 Numba 在 对象模式 下编译。这在调试或处理 Numba 暂时不支持的 Python 特性时有用，但性能提升会非常有限，甚至可能不如纯 Python。
@njit: 这是 @jit(nopython=True) 的简写。如果你希望函数总是以 nopython 模式编译，可以直接使用 @njit。
@vectorize: 用于创建 NumPy ufuncs (Universal Functions)。它接收一个返回单个值的 Python 函数，并将其编译成一个可以高效应用于 NumPy 数组的函数。
- @vectorize('return_type(arg_type1, arg_type2, ...)'): 需要指定返回类型和参数类型。例如：@vectorize('float64(float64, float32)') 表示该函数接收一个 float64 和一个 float32 参数，并返回一个 float64。
@guvectorize: 用于创建更通用的广义通用函数，可以处理任意形状的输入和输出。它的语法比 @vectorize 更复杂，需要指定输出和输入的形状签名。
@cuda.jit: 用于将 Python 函数编译成 CUDA 核函数，用于在 NVIDIA GPU 上执行。

2. 与 NumPy 的集成¶

Numba 对 NumPy 数组有非常好的支持，这是其核心优势之一。

直接操作 NumPy 数组: 你可以在 Numba 编译的函数中直接对 NumPy 数组进行切片、索引、数学运算等。Numba 会将这些操作转换为高效的机器码。
NumPy 函数的支持: Numba 支持大量的 NumPy 函数，例如 np.sum, np.mean, np.dot, np.sin, np.cos 等。
类型推断: Numba 能够很好地推断 NumPy 数组的类型（如 float64, int32）和维度，并生成针对特定类型的优化代码。

3. 变量和类型¶

动态类型 vs. 静态类型: 虽然 Python 是动态类型语言，但在 Numba 的 nopython 模式下，Numba 会尽力推断变量的类型。一旦类型被确定，它就像编译型语言一样处理。
- 例如，在一个 @jit(nopython=True) 函数中，如果你将一个整数赋值给一个变量，Numba 会将其推断为整数类型。如果你随后将一个浮点数赋值给同一个变量，Numba 可能会报错（如果类型不兼容），或者根据上下文进行隐式转换（但应避免）。
避免复杂的 Python 对象: 在 nopython 模式下，Numba 无法处理所有 Python 对象（如字典、列表、自定义类实例的复杂操作）。如果函数中存在这些操作，Numba 可能会回退到对象模式（性能下降）或报错。

4. 循环和控制流¶

Numba 能很好地编译标准的 Python 循环和控制流语句。

for 循环:

@njit
def process_list(data):
    result = 0.0
    for item in data: # Numba 可以有效处理可迭代对象
        result += item
    return result

while 循环:

@njit
def countdown(n):
    while n > 0:
        n -= 1
    return n

条件语句 (if, elif, else):

@njit
def categorize(x):
    if x < 0:
        return -1
    elif x == 0:
        return 0
    else:
        return 1

5. 函数定义和调用¶

常规函数定义: Numba 装饰器应用于标准的 Python 函数定义。
函数调用: 在 Numba 编译的函数内部，你可以调用其他 Numba 编译的函数，或者标准的 Python 函数（如果 Numba 支持）。

6. 限制和注意事项 (影响语法选择)¶

nopython=True 的重要性: 始终优先尝试使用 @njit 或 @jit(nopython=True)。这是性能提升的关键。
明确的类型: 如果 Numba 无法确定类型，或者你想要更精细的控制，可以考虑使用 numba.types 来显式声明类型（虽然不常用，通常 Numba 的推断已经足够）。
避免全局变量: 在 nopython 模式下，全局变量的读写可能会有性能开销，并且 Numba 对它们的处理可能有限制。
不支持的特性: 熟悉 Numba 支持的 Python 子集和 NumPy 函数列表非常重要。如果你遇到的函数或库不在支持范围内，Numba 可能无法加速。

总而言之，Numba 的语法规范不是一套全新的语法，而是通过装饰器来“增强”和“指示”标准的 Python 代码如何被编译。核心在于理解并正确使用这些装饰器，特别是 @jit 和 @njit 的各种选项，以及熟悉 Numba 对 NumPy 和 Python 子集的支持程度。

例如，如果你写了一个纯 NumPy 的函数，并且其中没有复杂的 Python 对象操作，那么只需要简单地加上 @njit，Numba 就会自动接管并进行优化。

import numba
import numpy as np

# 一个简单的 NumPy 操作
def numpy_operation(a, b):
    return np.sin(a) + np.cos(b) * 2

# 应用 Numba
@numba.njit
def optimized_numpy_operation(a, b):
    return np.sin(a) + np.cos(b) * 2

# 使用
arr1 = np.random.rand(1000)
arr2 = np.random.rand(1000)

# Numba 版本会更快
result = optimized_numpy_operation(arr1, arr2)

Numba支持加速的情况¶

Numba 的适用范围：让你的 Python 代码飞起来¶

Numba 的核心目标是加速 Python 的数值计算，特别是那些依赖于 NumPy 和标准 Python 语言特性的代码。它通过将 Python 代码“翻译”成高效的机器码来实现这一点。

1. Numba 能够加速的适用范围¶

Numba 主要擅长加速以下几类 Python 代码：

数值计算密集型代码: 这是 Numba 最擅长的领域。任何涉及大量数学运算、数组操作、循环计算的代码，只要符合 Numba 的要求，都能看到显著的性能提升。
NumPy 数组操作: Numba 对 NumPy 数组有非常深度的支持。几乎所有对 NumPy 数组进行的常规操作，如索引、切片、算术运算、数学函数调用、广播等，都可以被 Numba 高效地编译。
纯 Python 循环: Numba 可以将显式的 for 和 while 循环编译成高效的机器码，消除 Python 解释器在循环中的开销。
部分标准 Python 库: Numba 支持一部分常用的标准 Python 库函数，尤其是与数值计算相关的。
CPU 和 GPU 加速: Numba 不仅可以加速 CPU 上的代码，还可以通过 @cuda.jit 装饰器将代码编译到 NVIDIA GPU 上执行，充分利用 GPU 的并行计算能力。

2. Numba 对 NumPy 和 Python 的支持子集¶

Numba 的强大之处在于它对 NumPy 的广泛支持。

2.1 NumPy 的支持子集¶

Numba 对 NumPy 的支持非常全面，以下是一些 Numba 能够高效编译的 NumPy 特性：

数组 (Arrays):
- 创建、修改、索引、切片
- 多维数组操作
- 广播 (Broadcasting)
数学函数:
- np.sin, np.cos, np.tan, np.arcsin, np.arccos, np.arctan
- np.exp, np.log, np.sqrt
- np.abs, np.sign
- np.power, np.multiply, np.add, np.subtract, np.divide
- np.floor, np.ceil, np.round
- np.maximum, np.minimum
- 以及许多其他常用的数学函数。
聚合函数:
- np.sum, np.prod, np.mean, np.std, np.var
- np.min, np.max
- np.all, np.any
- np.cumsum, np.cumprod
逻辑函数:
- np.logical_and, np.logical_or, np.logical_not
排序和选择:
- np.sort, np.argsort (但可能不支持部分复杂的排序选项)
- np.argmax, np.argmin
线性代数 (部分):
- np.dot (点积)
- np.linalg.norm (范数)
- np.linalg.svd (奇异值分解) - Numba 对 linalg 模块的支持在不断完善中。
数据类型 (dtypes): Numba 支持大多数 NumPy 的基本数据类型，如 int32, int64, float32, float64, complex64, complex128, bool 等。

2.2 Python 的支持子集 (在 `nopython` 模式下)¶

在 nopython=True 模式下，Numba 会尝试将函数完全编译成机器码，因此它支持 Python 的一个受限子集。

基本数据类型:
- 整数 (int)
- 浮点数 (float)
- 布尔值 (bool)
- 复数 (complex)
控制流:
- if, elif, else
- for 循环 (支持迭代可迭代对象)
- while 循环
- break, continue
函数:
- 定义普通函数
- 调用其他 Numba 编译的函数
- 调用一部分支持的标准 Python 内置函数 (如 len, range, abs)
NumPy 数组: 如上所述，是 Numba 的核心支持对象。
元组 (tuple): Numba 支持不可变的元组，并且可以推断其类型。
字符串 (str): Numba 对字符串的支持有限，主要支持简单的字符串常量和拼接。复杂的操作（如正则表达式、大量字符串切片和格式化）可能不在支持范围内，或性能不佳。
列表 (list): Numba 对列表的支持在不断完善，但通常比 NumPy 数组更有限。在 nopython 模式下，如果列表的长度和元素类型在编译时确定，Numba 可以进行优化。否则，可能回退到对象模式。
字典 (dict): Numba 对字典的支持非常有限，在 nopython 模式下几乎不支持。

需要注意的 Python 特性，Numba 在 nopython 模式下不支持或支持有限：

类 (Classes) 和对象: Numba 无法直接编译涉及自定义类实例的复杂操作（如方法调用、属性访问）。
动态属性访问: getattr, setattr
异常处理: try, except, finally (部分支持)
高阶函数: map, filter, reduce (虽然可以通过 numba.vectorize 等方式实现类似功能)
生成器 (Generators): yield 语句
装饰器 (Decorators): Numba 装饰器本身不是被装饰函数的一部分。
I/O 操作: print (部分支持), open, read, write 等文件操作。
很多标准库模块: 尤其是那些不直接涉及数值计算的模块。

3. `@njit` 能够加速的情况有哪些？¶

@njit 相当于 @jit(nopython=True)。它指示 Numba 强制以 nopython 模式 来编译函数。这意味着 Numba 会尽力将整个函数编译成独立的机器码，不依赖于 Python 解释器。

@njit 能够加速的情况主要包括：

显式的 for 和 while 循环:
- 场景: 对 NumPy 数组进行逐元素的计算、累加、过滤等。
- 例子:

In [6]:

Copied!





        import numba
        import numpy as np
        import time

        @numba.njit
        def sum_of_squares(arr):
            result = 0.0
            for x in arr:
                result += x * x
            return result

        # --- 性能对比 ---
        data = np.random.rand(10_000_000)

        start = time.time()
        numba_result = sum_of_squares(data)
        numba_time = time.time() - start
        print(f"Numba (njit) time: {numba_time:.6f} seconds")

        # 纯 Python 对比
        def python_sum_of_squares(arr):
            result = 0.0
            for x in arr:
                result += x * x
            return result

        start = time.time()
        python_result = python_sum_of_squares(data)
        python_time = time.time() - start
        print(f"Pure Python time: {python_time:.6f} seconds")
        
        

        import numba
        import numpy as np
        import time

        @numba.njit
        def sum_of_squares(arr):
            result = 0.0
            for x in arr:
                result += x * x
            return result

        # --- 性能对比 ---
        data = np.random.rand(10_000_000)

        start = time.time()
        numba_result = sum_of_squares(data)
        numba_time = time.time() - start
        print(f"Numba (njit) time: {numba_time:.6f} seconds")

        # 纯 Python 对比
        def python_sum_of_squares(arr):
            result = 0.0
            for x in arr:
                result += x * x
            return result

        start = time.time()
        python_result = python_sum_of_squares(data)
        python_time = time.time() - start
        print(f"Pure Python time: {python_time:.6f} seconds")
        
        

Numba (njit) time: 0.332334 seconds
Pure Python time: 1.994897 seconds

在这个例子中，@njit 移除了 Python 解释器在循环每次迭代中的开销，并生成了高效的机器码来执行计算。

NumPy 数组的高级操作:
- 场景: 对 NumPy 数组进行复杂的数学运算，包括多个函数调用、条件判断等。
- 例子:

In [7]:

Copied!





    
        import numba
        import numpy as np
        import time

        @numba.njit
        def complex_array_ops(arr):
            result = np.zeros_like(arr)
            for i in range(arr.shape[0]):
                if arr[i] > 0.5:
                    result[i] = np.sin(arr[i]) * 2.0
                else:
                    result[i] = np.cos(arr[i]) / 1.5
            return result

        # --- 性能对比 ---
        data = np.random.rand(10_000_000)

        start = time.time()
        numba_result = complex_array_ops(data)
        numba_time = time.time() - start
        print(f"Numba (njit) time: {numba_time:.6f} seconds")

        # 纯 Python 对比 (注意：这里尽量模拟 Numba 的编译行为，但纯 Python 循环的性能会差很多)
        def python_complex_array_ops(arr):
            result = np.zeros_like(arr)
            for i in range(arr.shape[0]):
                if arr[i] > 0.5:
                    result[i] = np.sin(arr[i]) * 2.0
                else:
                    result[i] = np.cos(arr[i]) / 1.5
            return result

        start = time.time()
        python_result = python_complex_array_ops(data)
        python_time = time.time() - start
        print(f"Pure Python time: {python_time:.6f} seconds")

        
    
        import numba
        import numpy as np
        import time

        @numba.njit
        def complex_array_ops(arr):
            result = np.zeros_like(arr)
            for i in range(arr.shape[0]):
                if arr[i] > 0.5:
                    result[i] = np.sin(arr[i]) * 2.0
                else:
                    result[i] = np.cos(arr[i]) / 1.5
            return result

        # --- 性能对比 ---
        data = np.random.rand(10_000_000)

        start = time.time()
        numba_result = complex_array_ops(data)
        numba_time = time.time() - start
        print(f"Numba (njit) time: {numba_time:.6f} seconds")

        # 纯 Python 对比 (注意：这里尽量模拟 Numba 的编译行为，但纯 Python 循环的性能会差很多)
        def python_complex_array_ops(arr):
            result = np.zeros_like(arr)
            for i in range(arr.shape[0]):
                if arr[i] > 0.5:
                    result[i] = np.sin(arr[i]) * 2.0
                else:
                    result[i] = np.cos(arr[i]) / 1.5
            return result

        start = time.time()
        python_result = python_complex_array_ops(data)
        python_time = time.time() - start
        print(f"Pure Python time: {python_time:.6f} seconds")

        

Numba (njit) time: 0.567412 seconds
Pure Python time: 15.903088 seconds

@njit 能够将 np.sin, np.cos 这些 NumPy 函数调用以及 if/else 逻辑，都高效地编译成机器码，并且避免了 Python 解释器在循环中处理这些操作的开销。

函数调用链:
- 场景: 当一个 Numba 编译的函数调用另一个 Numba 编译的函数时。
- 例子:

In [8]:

Copied!





        import numba
        import numpy as np

        @numba.njit
        def helper_func(x):
            return x * x + 1

        @numba.njit
        def main_func(arr):
            result = np.zeros_like(arr)
            for i in range(arr.shape[0]):
                result[i] = helper_func(arr[i]) * 2
            return result

        data = np.arange(1000000)
        print(main_func(data)[:10])
  
        

        import numba
        import numpy as np

        @numba.njit
        def helper_func(x):
            return x * x + 1

        @numba.njit
        def main_func(arr):
            result = np.zeros_like(arr)
            for i in range(arr.shape[0]):
                result[i] = helper_func(arr[i]) * 2
            return result

        data = np.arange(1000000)
        print(main_func(data)[:10])
  
        

[  2   4  10  20  34  52  74 100 130 164]

    Numba 会将 `helper_func` 的调用也内联或优化，使得整个调用链非常高效。

并行计算 (配合 parallel=True):
- 场景: 函数中存在独立的、可以并行执行的循环。
- 例子:

In [18]:

Copied!





import numba
from numba import prange
import numpy as np
import time

@numba.jit(nopython=True, parallel=True)
def parallel_sum_squares(arr):
    local_sum = 0.0
    for i in prange(len(arr)):
        local_sum += arr[i] * arr[i]
    return local_sum



        # --- 性能对比 ---
data = np.random.rand(50_000_0000) # 数据量更大以突出并行优势

start = time.time()
numba_result = parallel_sum_squares(data)
numba_time = time.time() - start
print(f"Numba (njit, parallel=True) time: {numba_time:.6f} seconds")

        # 不并行化 Numba 版本
@numba.njit
def non_parallel_sum_squares(arr):
    result = 0.0
    for x in arr:
        result += x * x
    return result

start = time.time()
numba_non_parallel_result = non_parallel_sum_squares(data)
numba_non_parallel_time = time.time() - start
print(f"Numba (njit) time: {numba_non_parallel_time:.6f} seconds")
  
        

import numba
from numba import prange
import numpy as np
import time

@numba.jit(nopython=True, parallel=True)
def parallel_sum_squares(arr):
    local_sum = 0.0
    for i in prange(len(arr)):
        local_sum += arr[i] * arr[i]
    return local_sum



        # --- 性能对比 ---
data = np.random.rand(50_000_0000) # 数据量更大以突出并行优势

start = time.time()
numba_result = parallel_sum_squares(data)
numba_time = time.time() - start
print(f"Numba (njit, parallel=True) time: {numba_time:.6f} seconds")

        # 不并行化 Numba 版本
@numba.njit
def non_parallel_sum_squares(arr):
    result = 0.0
    for x in arr:
        result += x * x
    return result

start = time.time()
numba_non_parallel_result = non_parallel_sum_squares(data)
numba_non_parallel_time = time.time() - start
print(f"Numba (njit) time: {numba_non_parallel_time:.6f} seconds")
  
        

Numba (njit, parallel=True) time: 9.821023 seconds
Numba (njit) time: 3.274496 seconds

在这个例子中，parallel=True 允许 Numba 将 for x in arr: 循环分割到多个 CPU 核心上，显著缩短了计算时间。

总结来说，@njit 能够加速的根本原因在于它强制 Numba 摆脱 Python 解释器，直接生成底层的、高度优化的机器码。这对于涉及大量计算、数组操作和循环的函数来说，可以带来数量级的性能提升。

4. `@vectorize` 的例子¶

@vectorize 允许你创建一个可以应用于 NumPy 数组的自定义函数，而不需要显式地编写循环。

例子：自定义激活函数

In [19]:

Copied!





    import numba
    import numpy as np
    import time

    # 定义一个 sigmoid 函数
    @numba.vectorize('float64(float64)')
    def sigmoid(x):
        return 1.0 / (1.0 + np.exp(-x))

    # --- 性能对比 ---
    data = np.random.rand(10_000_000) * 10 - 5 # 范围在 -5 到 5 之间

    # 使用 numba 编译的 sigmoid
    start = time.time()
    numba_result = sigmoid(data)
    numba_time = time.time() - start
    print(f"Numba (vectorize) sigmoid time: {numba_time:.6f} seconds")

    # 纯 Python 实现 (非常慢)
    def python_sigmoid(x):
        return 1.0 / (1.0 + np.exp(-x))

    start = time.time()
    # 需要使用 np.vectorize 来获得一些优化，但不如 Numba
    python_vectorized_sigmoid = np.vectorize(python_sigmoid)
    python_result = python_vectorized_sigmoid(data)
    python_time = time.time() - start
    print(f"Pure Python (np.vectorize) sigmoid time: {python_time:.6f} seconds")

    # NumPy 内置的 sigmoid (如果存在，通常会更快，因为是 C 实现)
    # 这里我们只是为了展示 Numba 的能力

    

    import numba
    import numpy as np
    import time

    # 定义一个 sigmoid 函数
    @numba.vectorize('float64(float64)')
    def sigmoid(x):
        return 1.0 / (1.0 + np.exp(-x))

    # --- 性能对比 ---
    data = np.random.rand(10_000_000) * 10 - 5 # 范围在 -5 到 5 之间

    # 使用 numba 编译的 sigmoid
    start = time.time()
    numba_result = sigmoid(data)
    numba_time = time.time() - start
    print(f"Numba (vectorize) sigmoid time: {numba_time:.6f} seconds")

    # 纯 Python 实现 (非常慢)
    def python_sigmoid(x):
        return 1.0 / (1.0 + np.exp(-x))

    start = time.time()
    # 需要使用 np.vectorize 来获得一些优化，但不如 Numba
    python_vectorized_sigmoid = np.vectorize(python_sigmoid)
    python_result = python_vectorized_sigmoid(data)
    python_time = time.time() - start
    print(f"Pure Python (np.vectorize) sigmoid time: {python_time:.6f} seconds")

    # NumPy 内置的 sigmoid (如果存在，通常会更快，因为是 C 实现)
    # 这里我们只是为了展示 Numba 的能力

    

Numba (vectorize) sigmoid time: 0.266099 seconds
Pure Python (np.vectorize) sigmoid time: 13.478786 seconds

在这个例子中，@numba.vectorize('float64(float64)') 创建了一个能够直接应用于 NumPy 数组的 sigmoid 函数。Numba 在底层为这个函数生成了高效的、针对数组操作的机器码。

5. `@cuda.jit` 的例子 (GPU 加速)¶

例子：GPU 上的数组加法

In [20]:

Copied!





  
    import numba
    import numpy as np
    import time

    @numba.cuda.jit
    def gpu_add_arrays(x, y, out):
        idx = numba.cuda.grid(1) # 获取当前线程的全局索引
        if idx < x.shape[0]:
            out[idx] = x[idx] + y[idx]

    # --- GPU 加速 ---
    n = 10_000_000
    # 在 CPU 上创建数据
    a_cpu = np.arange(n, dtype=np.float32)
    b_cpu = np.arange(n, dtype=np.float32)
    result_gpu = np.empty_like(a_cpu)

    # 将数据拷贝到 GPU
    a_gpu = numba.cuda.to_device(a_cpu)
    b_gpu = numba.cuda.to_device(b_cpu)
    result_gpu_device = numba.cuda.device_array_like(a_gpu)

    # 配置 GPU 执行参数 (线程块和网格大小)
    threadsperblock = 128
    blockspergrid = (n + (threadsperblock - 1)) // threadsperblock

    # 执行 GPU 核函数
    start_time = time.time()
    gpu_add_arrays[blockspergrid, threadsperblock](a_gpu, b_gpu, result_gpu_device)
    # 等待 GPU 计算完成 (隐式或显式)
    numba.cuda.current_context().synchronize()
    gpu_time = time.time() - start_time
    print(f"Numba (cuda.jit) GPU addition time: {gpu_time:.6f} seconds")

    # 将结果拷贝回 CPU
    result_gpu = result_gpu_device.copy_to_host()

    # CPU 版本对比 (纯 NumPy 操作通常已经非常快，但 Numba CUDA 针对大规模并行有优势)
    start_time = time.time()
    result_cpu = a_cpu + b_cpu
    cpu_time = time.time() - start_time
    print(f"NumPy CPU addition time: {cpu_time:.6f} seconds")

    # 验证结果
    assert np.allclose(result_gpu, result_cpu)
    
    
  
    import numba
    import numpy as np
    import time

    @numba.cuda.jit
    def gpu_add_arrays(x, y, out):
        idx = numba.cuda.grid(1) # 获取当前线程的全局索引
        if idx < x.shape[0]:
            out[idx] = x[idx] + y[idx]

    # --- GPU 加速 ---
    n = 10_000_000
    # 在 CPU 上创建数据
    a_cpu = np.arange(n, dtype=np.float32)
    b_cpu = np.arange(n, dtype=np.float32)
    result_gpu = np.empty_like(a_cpu)

    # 将数据拷贝到 GPU
    a_gpu = numba.cuda.to_device(a_cpu)
    b_gpu = numba.cuda.to_device(b_cpu)
    result_gpu_device = numba.cuda.device_array_like(a_gpu)

    # 配置 GPU 执行参数 (线程块和网格大小)
    threadsperblock = 128
    blockspergrid = (n + (threadsperblock - 1)) // threadsperblock

    # 执行 GPU 核函数
    start_time = time.time()
    gpu_add_arrays[blockspergrid, threadsperblock](a_gpu, b_gpu, result_gpu_device)
    # 等待 GPU 计算完成 (隐式或显式)
    numba.cuda.current_context().synchronize()
    gpu_time = time.time() - start_time
    print(f"Numba (cuda.jit) GPU addition time: {gpu_time:.6f} seconds")

    # 将结果拷贝回 CPU
    result_gpu = result_gpu_device.copy_to_host()

    # CPU 版本对比 (纯 NumPy 操作通常已经非常快，但 Numba CUDA 针对大规模并行有优势)
    start_time = time.time()
    result_cpu = a_cpu + b_cpu
    cpu_time = time.time() - start_time
    print(f"NumPy CPU addition time: {cpu_time:.6f} seconds")

    # 验证结果
    assert np.allclose(result_gpu, result_cpu)
    
    

---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
Cell In[20], line 5
      2 import numpy as np
      3 import time
----> 5 @numba.cuda.jit
      6 def gpu_add_arrays(x, y, out):
      7     idx = numba.cuda.grid(1) # 获取当前线程的全局索引
      8     if idx < x.shape[0]:

AttributeError: module 'numba' has no attribute 'cuda'

@cuda.jit 将 Python 代码编译成 GPU 上的 CUDA 核函数，通过 numba.cuda.grid(1) 等函数利用 GPU 的大量并行处理能力。

总而言之，Numba 的加速能力主要体现在数值计算、NumPy 操作和显式循环上。通过 @njit 强制进入 nopython 模式，可以最大限度地榨取性能。同时，@vectorize 和 @cuda.jit 提供了更高级的抽象，分别用于创建自定义 NumPy ufuncs 和进行 GPU 加速。

In [4]:

Copied!





import numba
import numpy as np
import time

# Numba 编译的函数，使用Numba优化的矩阵乘法
@numba.jit(nopython=True, cache=True)
def matrix_multiply_numba(A, B):
    # A: (m, k), B: (k, n) -> Result: (m, n)
    return np.dot(A, B)

# 生成测试数据
m, k, n = 500, 600, 700
A = np.random.rand(m, k)
B = np.random.rand(k, n)

# --- 性能测试 ---

# 第一次调用，Numba 会编译
start_time = time.time()
result_numba = matrix_multiply_numba(A, B)
compile_and_run_time = time.time() - start_time
print(f"Numba (njit) compile and run time: {compile_and_run_time:.6f} seconds")

# 第二次调用，使用缓存的机器码
start_time = time.time()
result_numba_cached = matrix_multiply_numba(A, B)
run_time = time.time() - start_time
print(f"Numba (njit) run time (cached): {run_time:.6f} seconds")

# 对比 NumPy 原生实现
start_time = time.time()
result_numpy = np.einsum('ik,kj->ij', A, B)
numpy_time = time.time() - start_time
print(f"NumPy np.einsum time: {numpy_time:.6f} seconds")

# 对比 NumPy dot 实现
start_time = time.time()
result_numpy_dot = np.dot(A, B)
numpy_dot_time = time.time() - start_time
print(f"NumPy np.dot time: {numpy_dot_time:.6f} seconds")

# 验证结果
print(f"Numba and NumPy einsum results match: {np.allclose(result_numba, result_numpy)}")
print(f"Numba and NumPy dot results match: {np.allclose(result_numba, result_numpy_dot)}")
import numba
import numpy as np
import time

# Numba 编译的函数，使用Numba优化的矩阵乘法
@numba.jit(nopython=True, cache=True)
def matrix_multiply_numba(A, B):
    # A: (m, k), B: (k, n) -> Result: (m, n)
    return np.dot(A, B)

# 生成测试数据
m, k, n = 500, 600, 700
A = np.random.rand(m, k)
B = np.random.rand(k, n)

# --- 性能测试 ---

# 第一次调用，Numba 会编译
start_time = time.time()
result_numba = matrix_multiply_numba(A, B)
compile_and_run_time = time.time() - start_time
print(f"Numba (njit) compile and run time: {compile_and_run_time:.6f} seconds")

# 第二次调用，使用缓存的机器码
start_time = time.time()
result_numba_cached = matrix_multiply_numba(A, B)
run_time = time.time() - start_time
print(f"Numba (njit) run time (cached): {run_time:.6f} seconds")

# 对比 NumPy 原生实现
start_time = time.time()
result_numpy = np.einsum('ik,kj->ij', A, B)
numpy_time = time.time() - start_time
print(f"NumPy np.einsum time: {numpy_time:.6f} seconds")

# 对比 NumPy dot 实现
start_time = time.time()
result_numpy_dot = np.dot(A, B)
numpy_dot_time = time.time() - start_time
print(f"NumPy np.dot time: {numpy_dot_time:.6f} seconds")

# 验证结果
print(f"Numba and NumPy einsum results match: {np.allclose(result_numba, result_numpy)}")
print(f"Numba and NumPy dot results match: {np.allclose(result_numba, result_numpy_dot)}")

Numba (njit) compile and run time: 0.187086 seconds
Numba (njit) run time (cached): 0.004119 seconds
NumPy np.einsum time: 0.161889 seconds
NumPy np.dot time: 0.003919 seconds
Numba and NumPy einsum results match: True
Numba and NumPy dot results match: True

In [5]:

Copied!





import numba
import numpy as np
import time

# Numba 编译的函数
@numba.jit(nopython=True, cache=True)
def matrix_multiply_numba(A, B):
    return A @ B

# 测试不同矩阵大小
sizes = [
    (100, 100, 100),   # 小矩阵
    (500, 600, 700),   # 中等矩阵
    (1000, 1200, 1400), # 大矩阵
    (2000, 2400, 2800)  # 很大矩阵
]

for m, k, n in sizes:
    print(f"\nMatrix size: {m}x{k} x {k}x{n}")
    
    # 生成测试数据
    A = np.random.rand(m, k).astype(np.float64)
    B = np.random.rand(k, n).astype(np.float64)
    
    # Numba 第一次调用（编译+运行）
    start = time.time()
    result_numba = matrix_multiply_numba(A, B)
    compile_and_run_time = time.time() - start
    
    # Numba 第二次调用（缓存）
    start = time.time()
    result_numba_cached = matrix_multiply_numba(A, B)
    run_time_numba = time.time() - start
    
    # NumPy dot
    start = time.time()
    result_numpy_dot = np.dot(A, B)
    numpy_dot_time = time.time() - start
    
    # NumPy einsum
    start = time.time()
    result_numpy_einsum = np.einsum('ik,kj->ij', A, B)
    numpy_einsum_time = time.time() - start
    
    # 验证结果
    assert np.allclose(result_numba, result_numpy_dot)
    assert np.allclose(result_numba, result_numpy_einsum)
    
    # 打印结果
    print(f"Numba (first): {compile_and_run_time:.6f} seconds")
    print(f"Numba (cached): {run_time_numba:.6f} seconds")
    print(f"NumPy dot: {numpy_dot_time:.6f} seconds")
    print(f"NumPy einsum: {numpy_einsum_time:.6f} seconds")
    print(f"Numba speedup over dot (cached): {numpy_dot_time/run_time_numba:.2f}x")
import numba
import numpy as np
import time

# Numba 编译的函数
@numba.jit(nopython=True, cache=True)
def matrix_multiply_numba(A, B):
    return A @ B

# 测试不同矩阵大小
sizes = [
    (100, 100, 100),   # 小矩阵
    (500, 600, 700),   # 中等矩阵
    (1000, 1200, 1400), # 大矩阵
    (2000, 2400, 2800)  # 很大矩阵
]

for m, k, n in sizes:
    print(f"\nMatrix size: {m}x{k} x {k}x{n}")
    
    # 生成测试数据
    A = np.random.rand(m, k).astype(np.float64)
    B = np.random.rand(k, n).astype(np.float64)
    
    # Numba 第一次调用（编译+运行）
    start = time.time()
    result_numba = matrix_multiply_numba(A, B)
    compile_and_run_time = time.time() - start
    
    # Numba 第二次调用（缓存）
    start = time.time()
    result_numba_cached = matrix_multiply_numba(A, B)
    run_time_numba = time.time() - start
    
    # NumPy dot
    start = time.time()
    result_numpy_dot = np.dot(A, B)
    numpy_dot_time = time.time() - start
    
    # NumPy einsum
    start = time.time()
    result_numpy_einsum = np.einsum('ik,kj->ij', A, B)
    numpy_einsum_time = time.time() - start
    
    # 验证结果
    assert np.allclose(result_numba, result_numpy_dot)
    assert np.allclose(result_numba, result_numpy_einsum)
    
    # 打印结果
    print(f"Numba (first): {compile_and_run_time:.6f} seconds")
    print(f"Numba (cached): {run_time_numba:.6f} seconds")
    print(f"NumPy dot: {numpy_dot_time:.6f} seconds")
    print(f"NumPy einsum: {numpy_einsum_time:.6f} seconds")
    print(f"Numba speedup over dot (cached): {numpy_dot_time/run_time_numba:.2f}x")

Matrix size: 100x100 x 100x100
Numba (first): 0.381446 seconds
Numba (cached): 0.000290 seconds
NumPy dot: 0.000293 seconds
NumPy einsum: 0.000331 seconds
Numba speedup over dot (cached): 1.01x

Matrix size: 500x600 x 600x700
Numba (first): 0.005252 seconds
Numba (cached): 0.007056 seconds
NumPy dot: 0.004691 seconds
NumPy einsum: 0.108998 seconds
Numba speedup over dot (cached): 0.66x

Matrix size: 1000x1200 x 1200x1400
Numba (first): 0.027031 seconds
Numba (cached): 0.025676 seconds
NumPy dot: 0.038295 seconds
NumPy einsum: 1.238479 seconds
Numba speedup over dot (cached): 1.49x

Matrix size: 2000x2400 x 2400x2800
Numba (first): 0.283099 seconds
Numba (cached): 0.254683 seconds
NumPy dot: 0.216929 seconds
NumPy einsum: 9.073445 seconds
Numba speedup over dot (cached): 0.85x

In [21]:

Copied!





import numpy as np
import time
from numba import njit, prange

# 纯Python实现的热传导模拟器
def heat_equation_python(u0, alpha, dx, dt, steps):
    """
    使用纯Python实现二维热传导方程的数值解
    
    参数:
    u0: 初始温度分布 (2D numpy数组)
    alpha: 热扩散系数
    dx: 空间步长
    dt: 时间步长
    steps: 模拟步数
    
    返回:
    最终温度分布 (2D numpy数组)
    """
    nx, ny = u0.shape
    u = u0.copy()
    u_new = np.zeros_like(u)
    
    # 稳定性条件检查
    stability = alpha * dt / (dx * dx)
    if stability > 0.25:
        raise ValueError(f"稳定性条件不满足: {stability} > 0.25")
    
    for _ in range(steps):
        # 内部点更新
        for i in range(1, nx-1):
            for j in range(1, ny-1):
                u_new[i, j] = u[i, j] + alpha * dt / (dx * dx) * (
                    u[i+1, j] + u[i-1, j] + u[i, j+1] + u[i, j-1] - 4 * u[i, j]
                )
        
        # 边界条件 (绝热边界)
        u_new[0, :] = u_new[1, :]
        u_new[-1, :] = u_new[-2, :]
        u_new[:, 0] = u_new[:, 1]
        u_new[:, -1] = u_new[:, -2]
        
        # 交换数组
        u, u_new = u_new, u
    
    return u

# 使用Numba优化的热传导模拟器
@njit
def heat_equation_numba(u0, alpha, dx, dt, steps):
    """
    使用Numba优化的二维热传导方程数值解
    
    参数和返回值与纯Python版本相同
    """
    nx, ny = u0.shape
    u = u0.copy()
    u_new = np.zeros_like(u)
    
    # 稳定性条件检查
    stability = alpha * dt / (dx * dx)
    if stability > 0.25:
        raise ValueError(f"稳定性条件不满足: {stability} > 0.25")
    
    for _ in range(steps):
        # 内部点更新
        for i in range(1, nx-1):
            for j in range(1, ny-1):
                u_new[i, j] = u[i, j] + alpha * dt / (dx * dx) * (
                    u[i+1, j] + u[i-1, j] + u[i, j+1] + u[i, j-1] - 4 * u[i, j]
                )
        
        # 边界条件 (绝热边界)
        u_new[0, :] = u_new[1, :]
        u_new[-1, :] = u_new[-2, :]
        u_new[:, 0] = u_new[:, 1]
        u_new[:, -1] = u_new[:, -2]
        
        # 交换数组
        u, u_new = u_new, u
    
    return u

# 使用Numba并行优化的热传导模拟器
@njit(parallel=True)
def heat_equation_numba_parallel(u0, alpha, dx, dt, steps):
    """
    使用Numba并行优化的二维热传导方程数值解
    
    参数和返回值与纯Python版本相同
    """
    nx, ny = u0.shape
    u = u0.copy()
    u_new = np.zeros_like(u)
    
    # 稳定性条件检查
    stability = alpha * dt / (dx * dx)
    if stability > 0.25:
        raise ValueError(f"稳定性条件不满足: {stability} > 0.25")
    
    for _ in range(steps):
        # 内部点更新 - 使用并行循环
        for i in prange(1, nx-1):
            for j in range(1, ny-1):
                u_new[i, j] = u[i, j] + alpha * dt / (dx * dx) * (
                    u[i+1, j] + u[i-1, j] + u[i, j+1] + u[i, j-1] - 4 * u[i, j]
                )
        
        # 边界条件 (绝热边界)
        u_new[0, :] = u_new[1, :]
        u_new[-1, :] = u_new[-2, :]
        u_new[:, 0] = u_new[:, 1]
        u_new[:, -1] = u_new[:, -2]
        
        # 交换数组
        u, u_new = u_new, u
    
    return u

# 设置模拟参数
nx, ny = 200, 200  # 网格大小
dx = 1.0           # 空间步长
alpha = 1.0        # 热扩散系数
dt = 0.1           # 时间步长
steps = 1000       # 模拟步数

# 创建初始温度分布 - 中心热点
u0 = np.zeros((nx, ny))
cx, cy = nx // 2, ny // 2
radius = min(nx, ny) // 10
for i in range(nx):
    for j in range(ny):
        if (i - cx)**2 + (j - cy)**2 < radius**2:
            u0[i, j] = 100.0  # 中心热点温度为100

# 测试纯Python版本
print("运行纯Python版本...")
start_time = time.time()
u_python = heat_equation_python(u0, alpha, dx, dt, steps)
python_time = time.time() - start_time
print(f"纯Python版本耗时: {python_time:.4f} 秒")

# 测试Numba优化版本
print("运行Numba优化版本...")
start_time = time.time()
u_numba = heat_equation_numba(u0, alpha, dx, dt, steps)
numba_time = time.time() - start_time
print(f"Numba优化版本耗时: {numba_time:.4f} 秒")

# 测试Numba并行优化版本
print("运行Numba并行优化版本...")
start_time = time.time()
u_numba_parallel = heat_equation_numba_parallel(u0, alpha, dx, dt, steps)
numba_parallel_time = time.time() - start_time
print(f"Numba并行优化版本耗时: {numba_parallel_time:.4f} 秒")

# 计算加速比
numba_speedup = python_time / numba_time
numba_parallel_speedup = python_time / numba_parallel_time

print(f"\nNumba优化版本相对于纯Python版本的加速比: {numba_speedup:.2f}x")
print(f"Numba并行优化版本相对于纯Python版本的加速比: {numba_parallel_speedup:.2f}x")

# 验证结果是否一致
max_diff_numba = np.max(np.abs(u_python - u_numba))
max_diff_parallel = np.max(np.abs(u_python - u_numba_parallel))
print(f"\n纯Python和Numba版本的最大差异: {max_diff_numba:.10f}")
print(f"纯Python和Numba并行版本的最大差异: {max_diff_parallel:.10f}")
import numpy as np
import time
from numba import njit, prange

# 纯Python实现的热传导模拟器
def heat_equation_python(u0, alpha, dx, dt, steps):
    """
    使用纯Python实现二维热传导方程的数值解
    
    参数:
    u0: 初始温度分布 (2D numpy数组)
    alpha: 热扩散系数
    dx: 空间步长
    dt: 时间步长
    steps: 模拟步数
    
    返回:
    最终温度分布 (2D numpy数组)
    """
    nx, ny = u0.shape
    u = u0.copy()
    u_new = np.zeros_like(u)
    
    # 稳定性条件检查
    stability = alpha * dt / (dx * dx)
    if stability > 0.25:
        raise ValueError(f"稳定性条件不满足: {stability} > 0.25")
    
    for _ in range(steps):
        # 内部点更新
        for i in range(1, nx-1):
            for j in range(1, ny-1):
                u_new[i, j] = u[i, j] + alpha * dt / (dx * dx) * (
                    u[i+1, j] + u[i-1, j] + u[i, j+1] + u[i, j-1] - 4 * u[i, j]
                )
        
        # 边界条件 (绝热边界)
        u_new[0, :] = u_new[1, :]
        u_new[-1, :] = u_new[-2, :]
        u_new[:, 0] = u_new[:, 1]
        u_new[:, -1] = u_new[:, -2]
        
        # 交换数组
        u, u_new = u_new, u
    
    return u

# 使用Numba优化的热传导模拟器
@njit
def heat_equation_numba(u0, alpha, dx, dt, steps):
    """
    使用Numba优化的二维热传导方程数值解
    
    参数和返回值与纯Python版本相同
    """
    nx, ny = u0.shape
    u = u0.copy()
    u_new = np.zeros_like(u)
    
    # 稳定性条件检查
    stability = alpha * dt / (dx * dx)
    if stability > 0.25:
        raise ValueError(f"稳定性条件不满足: {stability} > 0.25")
    
    for _ in range(steps):
        # 内部点更新
        for i in range(1, nx-1):
            for j in range(1, ny-1):
                u_new[i, j] = u[i, j] + alpha * dt / (dx * dx) * (
                    u[i+1, j] + u[i-1, j] + u[i, j+1] + u[i, j-1] - 4 * u[i, j]
                )
        
        # 边界条件 (绝热边界)
        u_new[0, :] = u_new[1, :]
        u_new[-1, :] = u_new[-2, :]
        u_new[:, 0] = u_new[:, 1]
        u_new[:, -1] = u_new[:, -2]
        
        # 交换数组
        u, u_new = u_new, u
    
    return u

# 使用Numba并行优化的热传导模拟器
@njit(parallel=True)
def heat_equation_numba_parallel(u0, alpha, dx, dt, steps):
    """
    使用Numba并行优化的二维热传导方程数值解
    
    参数和返回值与纯Python版本相同
    """
    nx, ny = u0.shape
    u = u0.copy()
    u_new = np.zeros_like(u)
    
    # 稳定性条件检查
    stability = alpha * dt / (dx * dx)
    if stability > 0.25:
        raise ValueError(f"稳定性条件不满足: {stability} > 0.25")
    
    for _ in range(steps):
        # 内部点更新 - 使用并行循环
        for i in prange(1, nx-1):
            for j in range(1, ny-1):
                u_new[i, j] = u[i, j] + alpha * dt / (dx * dx) * (
                    u[i+1, j] + u[i-1, j] + u[i, j+1] + u[i, j-1] - 4 * u[i, j]
                )
        
        # 边界条件 (绝热边界)
        u_new[0, :] = u_new[1, :]
        u_new[-1, :] = u_new[-2, :]
        u_new[:, 0] = u_new[:, 1]
        u_new[:, -1] = u_new[:, -2]
        
        # 交换数组
        u, u_new = u_new, u
    
    return u

# 设置模拟参数
nx, ny = 200, 200  # 网格大小
dx = 1.0           # 空间步长
alpha = 1.0        # 热扩散系数
dt = 0.1           # 时间步长
steps = 1000       # 模拟步数

# 创建初始温度分布 - 中心热点
u0 = np.zeros((nx, ny))
cx, cy = nx // 2, ny // 2
radius = min(nx, ny) // 10
for i in range(nx):
    for j in range(ny):
        if (i - cx)**2 + (j - cy)**2 < radius**2:
            u0[i, j] = 100.0  # 中心热点温度为100

# 测试纯Python版本
print("运行纯Python版本...")
start_time = time.time()
u_python = heat_equation_python(u0, alpha, dx, dt, steps)
python_time = time.time() - start_time
print(f"纯Python版本耗时: {python_time:.4f} 秒")

# 测试Numba优化版本
print("运行Numba优化版本...")
start_time = time.time()
u_numba = heat_equation_numba(u0, alpha, dx, dt, steps)
numba_time = time.time() - start_time
print(f"Numba优化版本耗时: {numba_time:.4f} 秒")

# 测试Numba并行优化版本
print("运行Numba并行优化版本...")
start_time = time.time()
u_numba_parallel = heat_equation_numba_parallel(u0, alpha, dx, dt, steps)
numba_parallel_time = time.time() - start_time
print(f"Numba并行优化版本耗时: {numba_parallel_time:.4f} 秒")

# 计算加速比
numba_speedup = python_time / numba_time
numba_parallel_speedup = python_time / numba_parallel_time

print(f"\nNumba优化版本相对于纯Python版本的加速比: {numba_speedup:.2f}x")
print(f"Numba并行优化版本相对于纯Python版本的加速比: {numba_parallel_speedup:.2f}x")

# 验证结果是否一致
max_diff_numba = np.max(np.abs(u_python - u_numba))
max_diff_parallel = np.max(np.abs(u_python - u_numba_parallel))
print(f"\n纯Python和Numba版本的最大差异: {max_diff_numba:.10f}")
print(f"纯Python和Numba并行版本的最大差异: {max_diff_parallel:.10f}")

运行纯Python版本...
纯Python版本耗时: 64.3132 秒
运行Numba优化版本...
Numba优化版本耗时: 5.2572 秒
运行Numba并行优化版本...
Numba并行优化版本耗时: 2.9786 秒

Numba优化版本相对于纯Python版本的加速比: 12.23x
Numba并行优化版本相对于纯Python版本的加速比: 21.59x

纯Python和Numba版本的最大差异: 0.0000000000
纯Python和Numba并行版本的最大差异: 0.0000000000

In [ ]: