毕昇编译器异构算子分核方案再探

摘要根据vec_cross_add接口的正确用法重新实现了前面的动态分核方案,效果进一步提升。

基于vec_cross_add接口重新实现Softmax(方案五)

方案实现

分核方案沿用方案三,经过修改,使用vec_cross_add()接口的Softmax算子实现如下。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
using data_t = float;

std::vector<data_t> ascend_softmax(std::vector<data_t> input) {
std::size_t input_sz = input.size();
std::size_t byte_count = input_sz * sizeof(data_t);

// call the host operator if input isn't enough a full block
if (byte_count < 32) {
return softmax(input);
}

// number of elements per group
const std::size_t elem_per_group = 640;
// number of repeats per group
const std::size_t repeat_per_group = elem_per_group * sizeof(data_t) / 256;
// number of elements in tail block
const std::size_t tail_elem_count = input_sz % elem_per_group;
// number of groups
// if tail block is exist, apply for one more group
const std::size_t group_num = (tail_elem_count > 0)
? ((input_sz / elem_per_group) + 1)
: (input_sz / elem_per_group);

sycl::queue Q(sycl::ascend_selector{});

// GM memory allocation
auto dev_buf = sycl::malloc_device<data_t>(group_num * elem_per_group, Q);
auto sum_res_buf = sycl::malloc_device<data_t>(group_num * (32 / sizeof(data_t)), Q);

// Host memory allocation
std::vector<data_t> sum_res(group_num * (32 / sizeof(data_t)), 0.0f);
std::vector<data_t> res(input_sz, 0.0f);

// host -> GM
Q.memcpy(dev_buf, input.data(), byte_count);

Q.launch<class Summary>(group_num, [=](sycl::group<1> group) {
bisheng::vector<data_t, elem_per_group> input_vec;
bisheng::vector<data_t, repeat_per_group> sum_vec;
std::size_t group_id = group.get_group_id();

// GM -> UB
input_vec.load(
sycl::global_ptr<data_t>(dev_buf + group_id * elem_per_group).get(),
elem_per_group);

if (tail_elem_count > 0 && group_id == group_num - 1) {
// if tail block has element and this is the last group
bisheng::vector_view<data_t> input_vec_v(input_vec.data(), tail_elem_count);

bisheng::vec_exp(input_vec_v, input_vec_v);
for (int i = 0; i < tail_elem_count; ++i)
sum_res_buf[group_id * (32 / sizeof(data_t))] += input_vec_v[i];
} else {
// full block data
bisheng::vec_exp(input_vec, input_vec);
bisheng::vec_cross_add(sum_vec.data(), input_vec);
for (int i = 0; i < repeat_per_group; ++i) {
sum_res_buf[group_id * (32 / sizeof(data_t))] += sum_vec[i];
}
}

// UB -> GM
input_vec.store(
sycl::global_ptr<data_t>(dev_buf + group_id * elem_per_group).get(),
elem_per_group);
});

// GM -> Host
Q.memcpy(sum_res.data(), sum_res_buf, group_num * (32 / sizeof(data_t)) * sizeof(data_t));
Q.wait();

data_t sum;
for (int i = 0; i < sum_res.size(); i += 32 / sizeof(data_t))
sum += sum_res[i];

Q.launch<class Softmax>(group_num, [=](sycl::group<1> group) {
// UB memory of exponent result
bisheng::vector<data_t, elem_per_group> exp_res_vec;
// UB memory of divisor
bisheng::vector<data_t, elem_per_group> divisor_vec(sum);
// UB memory of final result
bisheng::vector<data_t, elem_per_group> res_vec;
std::size_t group_id = group.get_group_id();

// GM -> UB
exp_res_vec.load(
sycl::global_ptr<data_t>(dev_buf + group_id * elem_per_group).get(),
elem_per_group);

if (tail_elem_count > 0 && group_id == group_num - 1) {
// if tail block has element and this is the last group
bisheng::vector_view<data_t> exp_res_vec_v(exp_res_vec.data(), tail_elem_count);
bisheng::vector_view<data_t> divisor_vec_v(divisor_vec.data(), tail_elem_count);
bisheng::vector_view<data_t> res_vec_v(res_vec.data(), tail_elem_count);

bisheng::vec_div(res_vec_v, exp_res_vec_v, divisor_vec_v);
} else {
// full block data
bisheng::vec_div(res_vec, exp_res_vec, divisor_vec);
}

// UB -> GM
res_vec.store(
sycl::global_ptr<data_t>(dev_buf + group_id * elem_per_group).get(),
elem_per_group);
});

// GM -> host
Q.memcpy(res.data(), dev_buf, byte_count);
Q.wait();

sycl::free(dev_buf, Q);
sycl::free(sum_res_buf, Q);

return res;
}

功能测试

功能测试全部验证正确。

性能测试

与之前效果最好的方案三对比。

测试用例 640 6400 64000 640000
方案三加速比 0.224613 1.512394 13.30433 88.70575
方案五加速比 0.225836 1.330532 12.58051 101.0137

可以看到在向量长度比较大的情况下,接口的优势还是远大于for循环求和的。

毕昇编译器异构算子分核方案再探(方案六)

之前的Softmax算子实现中提到,采用动态分核方案,令每个核尽可能多的处理数据,并充分利用所以物理核带来的效果,甚至不如大量处理适量数据的逻辑核带来的效果好。当时是由于对vec_cross_add()求和接口产生了误解,所以转用for循环在逻辑核内部进行求和。

回顾一下上次的两种分核方式和性能表现。

  • 方案三:每个核处理640个数据,无论输入向量多长,都拆分为640个元素的小向量,再单独处理尾块;
  • 方案四:根据输入向量的长度和物理核心数量,动态确定逻辑核的数量,保证每个物理核都在工作,再单独处理尾块。

方案四详见基于毕昇编译器的Softmax异构算子 - 亦初 (deleter-d.github.io)

当时的性能表现是:

测试用例 640 6400 64000 640000 698880
方案三加速比 0.237432 1.478988 13.19605 87.72573 96.29706
方案四加速比 0.234373 1.436941 12.35908 80.59147 67.26953

但转念一想,得到这个结果的前提是求和均使用了核内for循环,所以动态分核策略中,输入向量长度比较长的情况下,每个逻辑核分到的向量长度就会远大于方案一的策略。可想而知,效率其实是被for循环拖慢的。

所以,利用vec_cross_add()接口再次实现动态分核方案。

方案实现

代码与之前的大同小异,只在求和的地方做了修改,改用vec_cross_add()接口。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
std::vector<data_t> ascend_softmax(std::vector<data_t> input) {
std::size_t input_sz = input.size();
std::size_t byte_count = input_sz * sizeof(data_t);

// call the host operator if input isn't enough a full block
if (byte_count < 32) {
return softmax(input);
}

// number of elements per group
std::size_t elem_per_group = 0;
if (byte_count >= PHYSICAL_CORES * UB_MAX_BYTES)
elem_per_group = UB_MAX_BYTES / sizeof(data_t);
else if (byte_count >= PHYSICAL_CORES * 51200)
elem_per_group = 51200 / sizeof(data_t);
else if (byte_count >= PHYSICAL_CORES * 25600)
elem_per_group = 25600 / sizeof(data_t);
else if (byte_count >= PHYSICAL_CORES * 12800)
elem_per_group = 12800 / sizeof(data_t);
else if (byte_count >= PHYSICAL_CORES * 5120)
elem_per_group = 5120 / sizeof(data_t);
else if (byte_count >= PHYSICAL_CORES * 2560)
elem_per_group = 2560 / sizeof(data_t);
else
elem_per_group = 1280 / sizeof(data_t);
// number of repeat per group
const std::size_t repeat_per_group = elem_per_group * sizeof(data_t) / 256;
// number of elements in tail block
const std::size_t tail_elem_count = input_sz % elem_per_group;
// number of groups
// if tail block is exist, apply for one more group
const std::size_t group_num = (tail_elem_count > 0)
? ((input_sz / elem_per_group) + 1)
: (input_sz / elem_per_group);

sycl::queue Q(sycl::ascend_selector{});

// GM memory allocation
auto dev_buf = sycl::malloc_device<data_t>(group_num * elem_per_group, Q);
auto sum_res_buf = sycl::malloc_device<data_t>(group_num * (32 / sizeof(data_t)), Q);

// Host memory allocation
std::vector<data_t> sum_res(group_num * (32 / sizeof(data_t)), 0.0f);
std::vector<data_t> res(input_sz, 0.0f);

// host -> GM
Q.memcpy(dev_buf, input.data(), byte_count);

Q.launch<class Summary>(group_num, [=](sycl::group<1> group) {
bisheng::vector<data_t, UB_MAX_BYTES / sizeof(data_t)> input_vec;
bisheng::vector<data_t, UB_MAX_BYTES / 256> sum_vec;
std::size_t group_id = group.get_group_id();

// GM -> UB
input_vec.load(
sycl::global_ptr<data_t>(dev_buf + group_id * elem_per_group).get(),
elem_per_group);

if (tail_elem_count > 0 && group_id == group_num - 1) {
// if tail block has element and this is the last group
bisheng::vector_view<data_t> input_vec_v(input_vec.data(), tail_elem_count);

bisheng::vec_exp(input_vec_v, input_vec_v);
for (int i = 0; i < tail_elem_count; ++i)
sum_res_buf[group_id * (32 / sizeof(data_t))] += input_vec_v[i];
} else {
// full block data
bisheng::vector_view<data_t> input_vec_v(input_vec.data(), elem_per_group);
bisheng::vector_view<data_t> sum_vec_v(sum_vec.data(), repeat_per_group * 8);

bisheng::vec_exp(input_vec_v, input_vec_v);
bisheng::vec_cross_add(sum_vec_v, input_vec_v);
for (int i = 0; i < repeat_per_group * 8; i += 8) {
sum_res_buf[group_id * (32 / sizeof(data_t))] += sum_vec_v[i];
}
}

// UB -> GM
input_vec.store(
sycl::global_ptr<data_t>(dev_buf + group_id * elem_per_group).get(),
elem_per_group);
});

// GM -> Host
Q.memcpy(sum_res.data(), sum_res_buf, group_num * (32 / sizeof(data_t)) * sizeof(data_t));
Q.wait();

data_t sum;
for (int i = 0; i < sum_res.size(); i += 32 / sizeof(data_t))
sum += sum_res[i];

Q.launch<class Softmax>(group_num, [=](sycl::group<1> group) {
// UB memory of exponent result
bisheng::vector<data_t, UB_MAX_BYTES / sizeof(data_t)> exp_res_vec;
// UB memory of divisor
bisheng::vector<data_t, UB_MAX_BYTES / sizeof(data_t)> divisor_vec(sum);
// UB memory of final result
bisheng::vector<data_t, UB_MAX_BYTES / sizeof(data_t)> res_vec;
std::size_t group_id = group.get_group_id();

// GM -> UB
exp_res_vec.load(
sycl::global_ptr<data_t>(dev_buf + group_id * elem_per_group).get(),
elem_per_group);

if (tail_elem_count > 0 && group_id == group_num - 1) {
// if tail block has element and this is the last group
bisheng::vector_view<data_t> exp_res_vec_v(exp_res_vec.data(), tail_elem_count);
bisheng::vector_view<data_t> divisor_vec_v(divisor_vec.data(), tail_elem_count);
bisheng::vector_view<data_t> res_vec_v(res_vec.data(), tail_elem_count);

bisheng::vec_div(res_vec_v, exp_res_vec_v, divisor_vec_v);
} else {
// full block data
bisheng::vector_view<data_t> exp_res_vec_v(exp_res_vec.data(), elem_per_group);
bisheng::vector_view<data_t> divisor_vec_v(divisor_vec.data(), elem_per_group);
bisheng::vector_view<data_t> res_vec_v(res_vec.data(), elem_per_group);
bisheng::vec_div(res_vec_v, exp_res_vec_v, divisor_vec_v);
}

// UB -> GM
res_vec.store(
sycl::global_ptr<data_t>(dev_buf + group_id * elem_per_group).get(),
elem_per_group);
});

// GM -> host
Q.memcpy(res.data(), dev_buf, byte_count);
Q.wait();

sycl::free(dev_buf, Q);
sycl::free(sum_res_buf, Q);

return res;
}

功能测试

功能测试全部正确。

性能测试

把几个方案放在一起对比,其中:

  • 方案三和方案五:每个核心固定640个元素;
  • 方案四和方案六:动态分核。
测试用例 640 6400 64000 640000 698880
方案三加速比 0.237432 1.478988 13.19605 87.72573 96.29706
方案四加速比 0.234373 1.436941 12.35908 80.59147 67.26953
方案五加速比 0.225836 1.330532 12.58051 101.0137 94.52815
方案六加速比 0.205704 1.29488 12.22252 120.7111 123.515

在输入向量较小的情况下,由于动态分核方案可能出现某些物理核不干活儿的情况,但在大向量的情况下,这种策略是占绝对优势的。加速比的一些波动是可接受的。同时由于使用了接口进行求和,大大减小了for循环的压力,整体效率进一步提升了。

目前还能够想到的优化手段可能就不是靠编译器能够自动完成的了,进一步的优化可能需要借助手动同步等手段追求极致性能了。

与mindspore对比

所用的mindspore版本为2.1.0,加速卡为昇腾910,测试代码如下:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
import numpy as np
import mindspore as ms
import mindspore.ops as ops
import time

INPUT = 640
input = ms.Tensor(np.random.random(INPUT), ms.float32)

start = time.perf_counter_ns()

ops.softmax(input)

end = time.perf_counter_ns()
runTime = end-start
print(f'Total time: {runTime}')

针对不同的测试用例只需修改INPUT即可,与上面效果最好的方案六做对比,各测试用例的加速比如下。

测试用例 640 6400 64000 640000 698880
加速比 1.636923067 1.72635719 1.672068458 1.71248809 1.636225576

加速比基本在1.7左右,虽然python解释器需要一定时间来解析,但通过一些手段测试后,这个时间的影响可以忽略。

毕昇编译器异构算子分核方案再探

https://deleter-d.github.io/posts/63148/

作者

亦初

发布于

2023-10-16

更新于

2024-06-19

许可协议

评论

:D 一言句子获取中...

加载中,最新评论有1分钟缓存...