毕昇异构算子数据搬移注意事项

数据搬移注意事项

为了更好的阅读体验，请移步笔者博客。

示例

在调用向量搬移接口load和store的时候，可能会遇到搬移出来的数据不符合预期的情况，用以下一个例子进行说明。

引入头文件，并定义两个常量作为矩阵的行和列数。

#include <iostream>
#include <sycl/sycl.hpp>
#include <bisheng/bisheng.hpp>
#include <bisheng/vector.hpp>
using namespace sycl;
using namespace bisheng;

constexpr int M = 8;
constexpr int N = 16;

定义毕昇异构算子，以实现以下功能。

代码实现如下。

void func()
{
	queue Q(ascend_selector{});

    // host数据
    float *input_ptr = new float[M * N];
    float *output_ptr = new float[M];

    // 初始化数据
    for (int i = 0; i < M; i++)
    {
        for (int j = 0; j < N; j++)
        {
            input_ptr[i * N + j] = j;
        }
    }

    // GM的buffer
    auto *input_buf = malloc_device<float>(M * N, Q);
    auto *output_buf = malloc_device<float>(M, Q);

    // 从host端拷贝到GM
    Q.memcpy(input_buf, input_ptr, M * N * sizeof(float));

    Q.launch<class Test>(M, [=](group<1> group) {
        auto group_id = group.get_id();
        // UB上的向量
        vector<float, N> vec;
        vector<float, 1> res;
        // 从GM加载到UB上
        vec.load(global_ptr<float>(input_buf + group_id * N));
        // 取正确的数放入结果向量中
        res[0] = vec[group_id];
        // 从UB存储到GM上
        res.store(global_ptr<float>(output_buf + group_id));
    });

    Q.wait();

    Q.memcpy(output_ptr, output_buf, M * sizeof(float));
    Q.memcpy(temp_ptr, temp_buf, M * N * sizeof(float));
    Q.wait();

    std::cout << "output data is:" << std::endl;
    for (int i = 0; i < M; i++)
    {
        std::cout << output_ptr[i] << " ";
    }
    std::cout << std::endl;

    free(input_buf, Q);
    free(output_buf, Q);
}

隐蔽的问题

上面的代码按照逻辑来看，似乎不存在问题，但其中有一个很隐蔽的致命问题。让我们回看核函数的代码。

Q.launch<class Test>(M, [=](group<1> group) {
    auto group_id = group.get_id();
    // UB上的向量
    vector<float, N> vec;
    vector<float, 1> res;
    // 从GM加载到UB上
    vec.load(global_ptr<float>(input_buf + group_id * N));
    // 取正确的数放入结果向量中
    res[0] = vec[group_id];
    // 从UB存储到GM上
    res.store(global_ptr<float>(output_buf + group_id));
});

不管是向量的声明，调用load()接口加载数据，还是根据group_id取得正确的数据都没有任何问题。这个隐蔽的问题隐藏在调用store接口的时候。

依照目前的写法，我们所期望的数据在内存中的搬移流程如下图所示。将Host端的数据搬移至Global Memory后，利用向量搬移接口load()搬移至Unified Buffer，每个Group读取一行数据，然后根据group_id取得正确的数据存入一个单元素向量中，再将该结果向量利用store()搬移至Global Memory中。

我们可以执行一下这段代码，输出的结果可能是这样的。

0 1 2 5.91908e-42 4 5.91908e-42 6 7

还可能是这样的。

0 1 2 5.91908e-42 4 5.91908e-42 2.8026e-45 0

还可能是。。。。

0 1 1 -1 4 5.91908e-42 2.8026e-45 7

0 1 1 3 5.91908e-42 2.8026e-45 0 2.8026e-45

0 1 2 5.91908e-42 4 5.91908e-42 2.8026e-45 7

产生的原因

这样的结果显然不符合预期。不幸的是，实际的搬移过程并不是我们想象的这样，致命的错误发生再从Unified Buffer搬移至Global Memory的过程中。

我们可以看一下load()和store()接口的源码是如何定义的。

void load(sycl::global_ptr<T> addr, size_t n = N) {
	return DMI_COPY_BLOCKS(this->data(), &addr[0], n);
}

void store(sycl::global_ptr<T> addr, size_t n = N) {
	return DMI_COPY_BLOCKS(&addr[0], this->data(), n);
}

可以观察到它们调用了一个名为DMI_COPY_BLOCKS()的接口，所以load()和store()接口显然是以块粒度搬移数据的，而根据毕昇C++的文档可以得知，一个block的大小为32B。但我们所store()的数据根本不够一个block，这就会导致在store()的过程中，实际上是搬移了一整个block的数据。除了第一个位置拥有我们取出来的数据之外，其他位置均为未定义的数据。

所以实际上的搬移情况是如下图所示的。

同时，Device端的Group之间都是异步的，甚至Group内部的一些操作都是异步的，这就导致output_buf有可能会被多个Group同时写入。显而易见地，这种情况会导致最终得到的output_buf中存在不正确的数据，如果你运气足够好，那还是有可能得到正确结果的。但计算机是科学，不是玄学。

解决方案

既然load()和store()的搬移都是块粒度的，那我们索性就避免让不同的group访问到同一个block。将上面的代码稍加修改，如下所示。

void correct()
{
    queue Q(ascend_selector{});

    float *input_ptr = new float[M * N];
    // 保证结果所在的block不重叠
    float *output_ptr = new float[M * 8];

    for (int i = 0; i < M; i++)
    {
        for (int j = 0; j < N; j++)
        {
            input_ptr[i * N + j] = j;
        }
    }

    auto *input_buf = malloc_device<float>(M * N, Q);
    // 相应的buffer也要申请更大的空间
    auto *output_buf = malloc_device<float>(M * 8, Q);

    Q.memcpy(input_buf, input_ptr, M * N * sizeof(float));

    Q.launch<class Test>(M, [=](group<1> group) {
        auto group_id = group.get_id();

        vector<float, N> vec;
        vector<float, 1> res;

        vec.load(global_ptr<float>(input_buf + group_id * N));
        
        res[0] = vec[group_id];

        // 从UB存储到GM上是，要注意位置的计算
        res.store(global_ptr<float>(output_buf + group_id * 8));
    });

    Q.wait();

    Q.memcpy(output_ptr, output_buf, M * 8 * sizeof(float));
    Q.wait();

    std::cout << "output data is:" << std::endl;
    for (int i = 0; i < M; i++)
    {
        for (int j = 0; j < 8; j++)
        {
            std::cout << output_ptr[i * 8 + j] << " ";
        }
        std::cout << std::endl;
    }

    free(input_buf, Q);
    free(output_buf, Q);
}

关键的修改在于，我们将output_buf直接申请成了一个大矩阵，上图说话。

这种方法直接避免了block的交叠，即使Group再异步，也不会产生写冲突。虽然Global Memory中的所有红色部分都是未定义的无效数据，但所带来的性能损耗几乎可以忽略不记。

需要注意该问题的情况

一定要仔细分析自己的算子逻辑，在将Unified Buffer中的数据搬移至Global Memory的过程中，理清搬移的数据大小，分析是否存在block交叠的可能性。尤其要注意输入数据很小的情况下，会不会引起该问题。如果这种情况下无法处理这个问题，可以尝试将输入数据极小的情况做单独的处理，不使用昇腾加速卡进行加速，当然也可以分析更好的解决方案。