使用 oneMKL 和 OpenMP Target Offloading 求解线性系统

上一篇文章《使用 Fortran、oneMKL 和 OpenMP 加速 LU 分解*》展示了如何将 LU 分解及其后续的逆运算卸载到加速器上。文章给出了一些关于通过合并 OpenMP 卸载区域和最小化主机-设备数据传输来减少开销的技巧，但并未涉及求解线性方程组。这留给了读者作为练习，但我们现在就来完成它，以说明《oneAPI GPU 优化指南》中关于“最小化数据传输和内存分配”一章中的一些技术。

我们将展示如何使用 oneMKL 中的批量 LU 分解和求解器函数来求解一组线性方程组，并通过 OpenMP 卸载到加速器上。在实际应用中，我们很少使用逆矩阵因子，因此在本演示中我们将省略该步骤，直接从分解到求解。由于之前的示例代码是用 Fortran 编写的，而且我们也很喜欢它，所以这里将继续使用 Fortran。

让我们从几个小型线性方程组开始，以确认我们能够正确地加载和求解它们。我们在一本旧的线性代数教科书中找到了几个合适的例子（**图 1**）。矩阵是方阵，因此我们可以使用 LU 分解来求解它们。

本文的演示程序与上一篇文章中的程序类似。它首先包含 OpenMP 卸载支持的 Fortran 接口，然后决定是使用 32 位还是 64 位整数类型（**图 2**）。（有关更多信息，请参阅使用 ILP64 接口与 LP64 接口。）这个决定将在编译时做出（参见下面的编译器命令）。**图 1** 中的测试问题包含两个 3x3 的线性方程组，每个方程组都有一个右侧项 (RHS)。相应地设置了 batch_size、n、nrhs 和 stride 变量，并分配了 a、b 和 ipiv 数组来存储批量方程组。最后四个语句将两个矩阵及其右侧项分别加载到 a 和 b 中。

include "mkl_omp_offtload.t90"
program solve_batched_1inear_systems

! I Decide whether to use 32- or 64-bit integer type
#if defined(MKL_ILP64)
    use onemkl_lapack_omp_offload_ilp64 ! 64-bit
#else
    use onemkl_lapack_omp_offload_1p64  ! 32-bit
#endif

    integer, parameter         :: n = 3, batch_size = 2, nrhs = 1
    integer                    :: lda, stride_a, stride_ipiv
    integer                    :: 1db, stride_b
    real (kind=8), allocatable :: a(:,:), b(:,:)
    integer, allocatable       :: ipiv(:, 1), info(:)
    
    integer allocstat
    character (len = 132) :: allocmsg

    1da         = n
    stride_a    = n * Tda
    stride_ipiv = n
    1db         = n
    stride_b    = n * nrhs

    ! Allocate required memory
    allocate (a(stride_a, batch_size), b(n, batch_size*nrhs), &
        ipiv(stride_ipiv, batch_size),                        &
        info(batch_ size), stat = allocstat, errmsg = allocmsg)
    if (allocstat > 0) stop trim(allocmsg)

    ! Initialize the matrices. Remember that Fortran is a column-major language.
    a(:,1) = (/2.0, 1.0, -6.0, 4.0, -4.0, -9.0, -4.0, 3.0, 10.0/)
    a(:,2) = (/0.0, 2.0, 6.0, 0.0, 4.0, 5.0, 3.0, -1.0, 5.0/)
    b(:,1) = (/12.0, -21.0, -24.0/) ! x = (-4, 2, -3)
    b(:,2) = (/ 6.0, -10.0, -7.0/) ! x = (-2, -1, 2)

现在我们可以开始计算了。让我们从一个基本实现开始，该实现使用两个 OpenMP 目标区域来分派 LU 分解和线性方程组的求解（**图 3**）。我们按如下方式编译和运行示例，以验证其结果是否正确。

$ ifx -i8 -DMKL_ILP64 -qopenmp -fopenmp-targets=spir64 -fsycl -free \
> lu_solve_omp_offload_ex1_small.F90 -o lu_solve_ex1_small \
> -L${MKLROOT}/lib/intel64 -lmkl_sycl -lmkl_intel_ilp64 -lmkl_intel_thread \
> -lmkl_core -liomp5 -lpthread -ldl
$ ./lu_solve_ex1_small

=================================
 Solutions to the linear systems
=================================
   -4.0000    2.0000   -3.0000
   -2.0000   -1.0000    2.0000

! Compute the LU factorization using OpenMP offload
! On entry, "a" contains the input matrix.
! On exit, it contains the LU factorization.
!$omp target data map(tofrom:a) map(from:ipiv) map(from:info)

    !$omp dispatch
    call dgetrf_batch_strided(n, n, a, lda, stride_a, &
        ipiv, stride_ipiv, batch_size, info)

!$omp end target data

if (any(info .ne. 0)) then
    print *, 'Error: getrf_batch_strided returned with errors'
else
    ! Solving the linear systems. On exit, the solutions are stored in b.
    !$omp target data map(to:a) map(to:ipiv) map(tofrom: b) map(from:info)

        !$omp dispatch

        call dgetrs_batch_strided('N', n, nrhs, a, lda, stride_a, &
            ipiv, stride_ipiv, &
            b, 1db, stride_b, batch_size, info)

    !$omp end target data

    if (any(info .ne. 0)) then
        print *, 'Error: getrs_batch_strided returned with errors'
    else
        print *, 'Computation executed successfully'
    endif
endif

target 结构将控制权转移到加速器设备。第一个区域将输入矩阵 [map(tofrom:a)] 传输到设备内存；将就地 LU 分解 (dgetrf_batch_strided) 分派到设备；并从设备内存中检索 LU 分解后的矩阵（现在存储在 a 中）、枢轴索引 [map(from:ipiv)] 和状态 [map(from:info)]。如果分解成功，程序将继续执行下一个 OpenMP 区域。LU 分解后的矩阵和枢轴索引被传输回设备内存 [map(to:a) 和 map(to:ipiv)]，并在设备上求解线性方程组 (dgetrs_batch_strided)。右侧项被覆盖为解向量，并从设备内存 [map(tofrom:b)] 中检索，同时检索计算的状态 [map(from:info)]。主机-设备数据传输在**图 4** 中以示意图的形式展示。

然而，我们不必依赖对主机-设备数据传输的感性理解。我们可以通过请求 OpenMP 运行时库的调试信息来直接检查数据移动，如最小化数据传输和内存分配中所述。测试问题太小，不足以进行加速，因此我们将问题规模增加到更有趣的程度。让我们（以双精度）求解一批八个 8000 x 8000 的线性方程组，每个方程组有一个右侧项。

我们高亮显示了使用 OpenMP map 构造显式传输的数组，并添加了一个空格来分隔两个 OpenMP 目标区域。未高亮显示的数据移动是正在映射到目标设备之数组的数据描述符。我们可以忽略它们，因为它们通常很小。

在第一个目标区域中，数组 a (4,096,000,000 字节) 从主机传输到目标（hst→tgt）[map(tofrom:a)]，然后再从目标传输回主机（tgt→hst）。在第二个目标区域中，它被传输回目标设备 [map(to:a)]。枢轴（ipiv，512,000 字节）在第一个目标区域在设备上计算，被检索 [map(from:ipiv), tgt→hst]，然后在第二个目标区域传输回设备 [map(to:ipiv), hst→tgt] 以计算解。存储在数组 b (512,000 字节) 中的右侧项被传输到设备，被解向量覆盖，然后传输回主机 [map(tofrom:b)]。状态数组 info (64 字节) 在两个目标区域结束时从设备检索 [map(from:info), tgt→hst]。

在不连续的内存之间移动数据需要时间和精力，因此我们需要密切关注高亮显示的传输（总计 12,290,048,128 字节）。考虑到这一点，让我们看看如何改进我们的初始实现。

上一篇关于 LU 分解的文章讨论了在可以容纳一个目标区域的情况下使用两个 OpenMP 目标区域的缺点。首先，将控制流转移到目标设备会产生开销。在初始实现（**图 3**）中，这会发生两次。其次，需要冗余的主机-设备数据传输（**图 4**）。在求解线性方程组时，我们只需要将输入数组和右侧项复制到设备，并从设备检索解。我们还需要检索状态数组，但这些相对较小。枢轴仅由设备使用，因此我们将直接在设备内存中分配 ipiv 数组 [map(alloc:ipiv(1:stride_ipiv, 1:batch_size))]。将两个 OpenMP 目标区域合并后，代码会更清晰、更简洁（**图 5**），并且需要更少的数据移动（**图 6**）。

!$omp target data map(to:a) map(tofrom: Db)      &
    map(alloc:ipiv(1:stride_ipiv, l:batch_size)) &
    map(from:info_rf, info_rs)

    !$omp dispatch
    call dgetrf_batch_strided(n, n, a, lda, stride_a, &
        ipiv, stride_ipiv, batch_size, info_rf)

    !$omp dispatch
    call dgetrs_batch_strided('N', n, nrhs, a, lda, stride_a, &
        ipiv, stride_ipiv, &
        b, 1db, stride_b, batch_size, info_rs)
!$omp end target data

if (any(info_rf .ne. 0)) then
    print *, 'Error: getrf_batch_strided returned with errors.'
elseif (any(info_rs .ne. 0)) then
    print *, 'Error: getrs_batch_strided returned with errors.'
else
    print *, 'Computation executed successfully'
endif

如果我们忽略数据描述符，**图 5** 中的单个 OpenMP 目标区域比初始实现中的两个目标区域（12,290,048,128 字节）需要显著更少的主机-设备数据传输（4,097,024,128 字节）（**图 3**）。

在异构并行计算中，最小化主机-设备数据传输是首要考虑的问题。第二个示例（**图 5**）中的代码性能明显优于原始代码（**图 3**），尤其是在线性方程组规模增大时（**表 1**）。

本文使用的示例程序可在以下存储库中找到。您可以在免费的Intel® Developer Cloud上体验 OpenMP 加速器卸载，该云平台拥有最新的 Intel 硬件和软件。

OpenMP target 构造有许多额外的选项，可让程序员精细控制加速器卸载。我们鼓励您查阅OpenMP 规范和示例代码，并查看使用 OpenMP API 进行 GPU 编程入门和最小化数据传输和内存分配教程。