oneAPI GPU 优化指南 - 使用 SYCL Joint Matrix 扩展在 Intel® XMX 上编程

本章节翻译by chenshusmail@163.com 原文：Programming Intel® XMX Using SYCL Joint Matrix Extension (intel.com)

Joint matrix 是一种新的 SYCL 扩展，用于 tensor 硬件编程。 这统一了用于CPU 的 Intel® Advanced Matrix Extensions (Intel® AMX)， 用于 GPU 的 Intel® Xe Matrix Extensions (Intel® XMX)， 以及 NVIDIA* Tensor Cores 等目标。一般来说， 像 TensorFlow 这样的框架和像 oneDNN 这样的库满足许多类型的用户和应用程序的需求。 然而，对于想要构建自己的神经网络应用程序的用户来说，这些库和框架太上层，因为用户无法进行自定义优化， 并且很臃肿，因为库的体积很大。此外，机器学习领域引入了新的操作并持续变化，对于这些操作， 框架和库没有提供及时和高效的解决方案。对于这种情况，joint matrix 比框架具有更低级别的抽象， 以提供性能、生产力和融合能力，但同时通过使用一种代码来针对不同的 tensor 硬件提供可移植性。

这个扩展的详细规范可以在 这里 找到。

Joint matrix 扩展包括一个新的类型 joint_matrix， 显式的内存操作 joint_matrix_load 和 joint_matrix_store， 用于矩阵初始化的 joint_matrix_fill， 用于实际乘法和加法操作的 joint_matrix_mad， 以及用于元素级操作的 get_wi_data()[]。

在下面的代码中，kernel 通过声明3个 joint_matrix 矩阵： tA, tB, tC 来使用 joint matrix 接口并计算操作 tC += tA * tB。 DPAS (Dot Product and Accumulate Systolic) 是在 Intel® XMX 中完成的基本操作的名称。有关使用 joint_matrix 的更多示例， 请参考 Intel llvm-test-suite repo。

为了让这个示例在 Intel® XMX 上成功运行，GPU 必须具有 Intel® XMX 硬件。 在 joint matrix 实现中，没有模拟或回退策略。

  size_t NDRangeM = M / TM;
  size_t NDRangeN = N / TN;
  sycl::buffer<bfloat16, 2> bufA(A.get_data(), sycl::range<2>(M, K));
  sycl::buffer<bfloat16, 2> bufB(B.get_data(), sycl::range<2>(K, N));
  sycl::buffer<float, 2> bufC((float *)C.get_data(), sycl::range<2>(M, N));

  sycl::queue q;
  q.submit([&](sycl::handler &cgh) {
     sycl::accessor accC(bufC, cgh, sycl::read_write, sycl::no_init);
     sycl::accessor accA(bufA, cgh, sycl::read_only, sycl::no_init);
     sycl::accessor accB(bufB, cgh, sycl::read_only, sycl::no_init);

     cgh.parallel_for(
         sycl::nd_range<2>({NDRangeM, NDRangeN * SG_SZ}, {1, 1 * SG_SZ}),
         [=](sycl::nd_item<2> spmd_item) [[intel::reqd_sub_group_size(SG_SZ)]]

         {
           // Joint matrix API 必须被 sub-group 中的所有 work-item 访问，
           // 这些函数将由 sub-group 调用一次，work-item 之间没有代码分歧。
           const auto global_idx = spmd_item.get_global_id(0);
           const auto global_idy = spmd_item.get_global_id(1);
           const auto sg_startx = global_idx - spmd_item.get_local_id(0);
           const auto sg_starty = global_idy - spmd_item.get_local_id(1);

           sycl::sub_group sg = spmd_item.get_sub_group();
           sycl::ext::oneapi::experimental::matrix::joint_matrix<
               sycl::sub_group, bfloat16, use::a, TM, TK, layout::row_major>
               sub_a;
           // 对于 B，我们假设 B 已经被转为 VNNI 指令。
           sycl::ext::oneapi::experimental::matrix::joint_matrix<
               sycl::sub_group, bfloat16, use::b, TK, TN,
               sycl::ext::intel::experimental::matrix::layout::packed>
               sub_b;
           sycl::ext::oneapi::experimental::matrix::joint_matrix<
               sycl::sub_group, float, use::accumulator, TM, TN>
               sub_c;

           joint_matrix_load(sg, sub_c,
                             accC.get_pointer() + (sg_startx * TM) * N +
                                 sg_starty / SG_SZ * TN,
                             N, layout::row_major);
           for (int k = 0; k < K / TK; k += 1) { //
             joint_matrix_load(
                 sg, sub_a, accA.get_pointer() + (sg_startx * TM) * K + k * TK,
                 K);
             joint_matrix_load(sg, sub_b,
                               accB.get_pointer() + (k * TK / 2) * (N * 2) +
                                   sg_starty / SG_SZ * TN * 2,
                               N * 2);
             sub_c = joint_matrix_mad(sg, sub_a, sub_b, sub_c);
           }
           joint_matrix_store(sg, sub_c,
                              accC.get_pointer() + (sg_startx * TM) * N +
                                  sg_starty / SG_SZ * TN,
                              N, layout::row_major);
         }); // parallel for
   }).wait();

为了在 Intel® XMX 上获得最佳性能，GEMM kernel 必须以一种使其不断向 Intel® XMX 提供数据的方式编写， 以执行最大的乘法和加法操作数/周期。其中一些调优技术如下：

• 一个 sub-group 可以执行多个 DPAS 操作。

• 对于 i， j 和 k 这三个维度，都应进行缓存局部性的阻塞。前两个维度的阻塞包含在 kernel 的全局范围内。 k 级缓存阻塞在 kernel 体内完成。 例如，通过选择 256x256x32 的块因子，全局范围结果为:

range<2> global{M / 256, N / 256 * SG_SIZE};

        然后，通过选择一个 sub-group 来执行 64x32x32 的元素，局部范围结果为:

range<2> local{256 / 64, 256 / 32 * SG_SIZE};

• 每个线程的大寄存器文件对于 GEMM kernel 来说效果最好。这可以在编译命令中明确指定，如下所示:

icpx -fsycl -fsycl-targets=spir64_gen -Xsycl-target-backend "-device pvc -internal_options -ze-opt-large-register-file" joint-matrix-bf16.cpp

上一章                                         主目录    上级目录                                                               下一章