Star1440
1504
代码
代码
Issues65
Pull Requests277
Actions
项目讨论
Wiki
分析
Star1440
1504
已合并
fix: align multinomial sampling with normalized probabilities #4500
fenglin28创建于 8月7日
fix: align multinomial sampling with normalized probabilities #4500
已合并
fenglin28创建于 8月7日
从fenglin28/ops-math_privatefix/stateless-sample-multinomial-assert合入到cann/ops-mathmaster
共 18 个文件变更+722-476
Mrandom/stateless_sample_multinomial/README.md+17-10
@@ -13,19 +13,19 @@
13 13 
14## 功能说明14## 功能说明
15 15 
16-- 算子功能：本算子是aclnnMultinomialGetWorkspaceSize接口构建计算流程时使用的内部服务算子，用于Ascend 950场景下的有放回多项分布采样路径。算子根据输入的累积分布和随机种子，从每个多项分布中抽取num_samples个样本，并将样本的类别索引存储到输出张量中。16+- 算子功能：本算子是`aclnnMultinomialGetWorkspaceSize`和`aclnnMultinomialTensorGetWorkspaceSize`接口构建计算流程时使用的内部服务算子，用于Ascend 950场景下的有放回多项分布采样路径。算子根据输入的累积分布和随机种子，从每个多项分布中抽取`num_samples`个样本，并将样本的类别索引存储到输出张量中。
17- 计算公式：17- 计算公式：
18 18 
19- 对于第d个分布的第j次抽样，生成随机数：19+ 对于第$d$个分布的第$j$次抽样，生成随机数：
20 20 
21 $$21 $$
22- u_{d,j}\sim U(0, 1]22+ u_{d,j} \sim U(0, 1]
23 $$23 $$
24 24 
25- 输出满足如下条件的最小类别索引k：25+ 输出满足如下条件的最小类别索引$k$：
26 26 
27 $$27 $$
28- x_{d,k-1}<u_{d,j}\le x_{d,k}28+ x_{d,k-1} < u_{d,j} \le x_{d,k}
29 $$29 $$
30 30 
31## 参数说明31## 参数说明
@@ -53,6 +53,13 @@
53 <td>FLOAT、FLOAT16、BFLOAT16</td>53 <td>FLOAT、FLOAT16、BFLOAT16</td>
54 <td>ND</td>54 <td>ND</td>
55 </tr>55 </tr>
56+ <tr>
57+ <td>norm_probs</td>
58+ <td>可选输入</td>
59+ <td>归一化概率张量，形状和数据类型与x相同。传入时用于采样结果的零概率类别回退。</td>
60+ <td>FLOAT、FLOAT16、BFLOAT16</td>
61+ <td>ND</td>
62+ </tr>
56 <tr>63 <tr>
57 <td>seed</td>64 <td>seed</td>
58 <td>输入</td>65 <td>输入</td>
@@ -85,14 +92,14 @@
85 92 
86## 约束说明93## 约束说明
87 94 
88-1. x仅支持1维或2维，最后一维C表示类别数，C不能超过2^24。95+1. `x`为一维或二维张量，最后一维为类别维，其长度不超过$2^{24}$。
89-2. x需要表示有效的累积分布，最后一维应单调非递减。96+2. `norm_probs`为可选输入。传入时，其形状和数据类型与`x`一致，`x`为`norm_probs`沿最后一维计算累加和的结果；缺省时，通过比较`x`的相邻元素进行零概率类别回退。
90-3. num_samples必须大于0。97+3. `num_samples`为正整数。
91-4. offset必须为4的倍数。98+4. `offset`为4的整数倍。
92-5. aclnn接口中replacement为false时，num_samples不能大于C。
93 99 
94## 调用说明100## 调用说明
95 101 
96| 调用方式 | 样例代码 | 说明 |102| 调用方式 | 样例代码 | 说明 |
97| --------- | ------------------------------------------------------------ | ------------------------------------------------------------ |103| --------- | ------------------------------------------------------------ | ------------------------------------------------------------ |
98| aclnn接口 | [test_aclnn_multinomial](./examples/test_aclnn_multinomial.cpp) | 通过[aclnnMultinomial](./docs/aclnnMultinomial.md)接口构建计算流程时，内部调用StatelessSampleMultinomial服务算子。 |104| aclnn接口 | [test_aclnn_multinomial](./examples/test_aclnn_multinomial.cpp) | 通过[aclnnMultinomial](./docs/aclnnMultinomial.md)接口构建计算流程时，内部调用StatelessSampleMultinomial服务算子。 |
105+| aclnn接口 | [test_aclnn_multinomial_tensor](./examples/test_aclnn_multinomial_tensor.cpp) | 通过[aclnnMultinomialTensor](./docs/aclnnMultinomialTensor.md)接口构建计算流程时，内部调用StatelessSampleMultinomial服务算子。 |
Mrandom/stateless_sample_multinomial/docs/aclnnMultinomial.md+5-5
@@ -349,7 +349,7 @@ int main() {
349 aclTensor* self = nullptr;349 aclTensor* self = nullptr;
350 aclTensor* out = nullptr;350 aclTensor* out = nullptr;
351 std::vector<float> selfHostData = {0, 10, 3, 0};351 std::vector<float> selfHostData = {0, 10, 3, 0};
352- std::vector<float> outHostData = {2, 0};352+ std::vector<int64_t> outHostData = {0, 0};
353 int64_t numsamples = 2;353 int64_t numsamples = 2;
354 bool replacement = false;354 bool replacement = false;
355 int64_t seed = 1234;355 int64_t seed = 1234;
@@ -382,13 +382,13 @@ int main() {
382 CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtSynchronizeStream failed. ERROR: %d\n", ret); return ret);382 CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtSynchronizeStream failed. ERROR: %d\n", ret); return ret);
383 383 
384 // 5. 获取输出的值，将device侧内存上的结果拷贝至host侧，需要根据具体API的接口定义修改384 // 5. 获取输出的值，将device侧内存上的结果拷贝至host侧，需要根据具体API的接口定义修改
385- auto size = GetShapeSize(selfShape);385+ auto size = GetShapeSize(outShape);
386- std::vector<float> resultData(size, 0);386+ std::vector<int64_t> resultData(size, 0);
387- ret = aclrtMemcpy(resultData.data(), resultData.size() * sizeof(resultData[0]), selfDeviceAddr,387+ ret = aclrtMemcpy(resultData.data(), resultData.size() * sizeof(resultData[0]), outDeviceAddr,
388 size * sizeof(resultData[0]), ACL_MEMCPY_DEVICE_TO_HOST);388 size * sizeof(resultData[0]), ACL_MEMCPY_DEVICE_TO_HOST);
389 CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("copy result from device to host failed. ERROR: %d\n", ret); return ret);389 CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("copy result from device to host failed. ERROR: %d\n", ret); return ret);
390 for (int64_t i = 0; i < size; i++) {390 for (int64_t i = 0; i < size; i++) {
391- LOG_PRINT("result[%ld] is: %f\n", i, resultData[i]);391+ LOG_PRINT("result[%ld] is: %ld\n", i, resultData[i]);
392 }392 }
393 393 
394 // 6. 释放aclTensor和aclScalar，需要根据具体API的接口定义修改394 // 6. 释放aclTensor和aclScalar，需要根据具体API的接口定义修改
Mrandom/stateless_sample_multinomial/examples/test_aclnn_multinomial.cpp+111-107
@@ -14,130 +14,134 @@
14#include "aclnnop/aclnn_multinomial.h"14#include "aclnnop/aclnn_multinomial.h"
15 15 
16#define CHECK_RET(cond, return_expr) \16#define CHECK_RET(cond, return_expr) \
17- do { \17+ do { \
18- if (!(cond)) { \18+ if (!(cond)) { \
19- return_expr; \19+ return_expr; \
20- } \20+ } \
21- } while (0)21+ } while (0)
22 22 
23-#define LOG_PRINT(message, ...) \23+#define LOG_PRINT(message, ...) \
24- do { \24+ do { \
25- printf(message, ##__VA_ARGS__); \25+ printf(message, ##__VA_ARGS__); \
26- } while (0)26+ } while (0)
27 27 
28-int64_t GetShapeSize(const std::vector<int64_t>& shape) {28+int64_t GetShapeSize(const std::vector<int64_t>& shape)
29- int64_t shapeSize = 1;29+{
30- for (auto i : shape) {30+ int64_t shapeSize = 1;
31- shapeSize *= i;31+ for (auto i : shape) {
32- }32+ shapeSize *= i;
33- return shapeSize;33+ }
34+ return shapeSize;
34}35}
35 36 
36-int Init(int32_t deviceId, aclrtStream* stream) {37+int Init(int32_t deviceId, aclrtStream* stream)
37- // 固定写法，资源初始化38+{
38- auto ret = aclInit(nullptr);39+ // 固定写法，资源初始化
39- CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclInit failed. ERROR: %d\n", ret); return ret);40+ auto ret = aclInit(nullptr);
40- ret = aclrtSetDevice(deviceId);41+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclInit failed. ERROR: %d\n", ret); return ret);
41- CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtSetDevice failed. ERROR: %d\n", ret); return ret);42+ ret = aclrtSetDevice(deviceId);
42- ret = aclrtCreateStream(stream);43+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtSetDevice failed. ERROR: %d\n", ret); return ret);
43- CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtCreateStream failed. ERROR: %d\n", ret); return ret);44+ ret = aclrtCreateStream(stream);
44- return 0;45+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtCreateStream failed. ERROR: %d\n", ret); return ret);
46+ return 0;
45}47}
46 48 
47template <typename T>49template <typename T>
48int CreateAclTensor(const std::vector<T>& hostData, const std::vector<int64_t>& shape, void** deviceAddr,50int CreateAclTensor(const std::vector<T>& hostData, const std::vector<int64_t>& shape, void** deviceAddr,
49- aclDataType dataType, aclTensor** tensor) {51+ aclDataType dataType, aclTensor** tensor)
50- auto size = GetShapeSize(shape) * sizeof(T);52+{
51- // 调用aclrtMalloc申请device侧内存53+ auto size = GetShapeSize(shape) * sizeof(T);
52- auto ret = aclrtMalloc(deviceAddr, size, ACL_MEM_MALLOC_HUGE_FIRST);54+ // 调用aclrtMalloc申请device侧内存
53- CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtMalloc failed. ERROR: %d\n", ret); return ret);55+ auto ret = aclrtMalloc(deviceAddr, size, ACL_MEM_MALLOC_HUGE_FIRST);
54- // 调用aclrtMemcpy将host侧数据拷贝到device侧内存上56+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtMalloc failed. ERROR: %d\n", ret); return ret);
55- ret = aclrtMemcpy(*deviceAddr, size, hostData.data(), size, ACL_MEMCPY_HOST_TO_DEVICE);57+ // 调用aclrtMemcpy将host侧数据拷贝到device侧内存上
56- CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtMemcpy failed. ERROR: %d\n", ret); return ret);58+ ret = aclrtMemcpy(*deviceAddr, size, hostData.data(), size, ACL_MEMCPY_HOST_TO_DEVICE);
59+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtMemcpy failed. ERROR: %d\n", ret); return ret);
57 60 
58- // 计算连续tensor的strides61+ // 计算连续tensor的strides
59- std::vector<int64_t> strides(shape.size(), 1);62+ std::vector<int64_t> strides(shape.size(), 1);
60- for (int64_t i = shape.size() - 2; i >= 0; i--) {63+ for (int64_t i = shape.size() - 2; i >= 0; i--) {
61- strides[i] = shape[i + 1] * strides[i + 1];64+ strides[i] = shape[i + 1] * strides[i + 1];
62- }65+ }
63 66 
64- // 调用aclCreateTensor接口创建aclTensor67+ // 调用aclCreateTensor接口创建aclTensor
65- *tensor = aclCreateTensor(shape.data(), shape.size(), dataType, strides.data(), 0, aclFormat::ACL_FORMAT_ND,68+ *tensor = aclCreateTensor(shape.data(), shape.size(), dataType, strides.data(), 0, aclFormat::ACL_FORMAT_ND,
66- shape.data(), shape.size(), *deviceAddr);69+ shape.data(), shape.size(), *deviceAddr);
67- return 0;70+ return 0;
68}71}
69 72 
70-int main() {73+int main()
71- // 1. （固定写法）device/stream初始化，参考acl API手册74+{
72- // 根据自己的实际device填写deviceId75+ // 1. （固定写法）device/stream初始化，参考acl API手册
73- int32_t deviceId = 0;76+ // 根据自己的实际device填写deviceId
74- aclrtStream stream;77+ int32_t deviceId = 0;
75- auto ret = Init(deviceId, &stream);78+ aclrtStream stream;
76- CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("Init acl failed. ERROR: %d\n", ret); return ret);79+ auto ret = Init(deviceId, &stream);
80+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("Init acl failed. ERROR: %d\n", ret); return ret);
77 81 
78- // 2. 构造输入与输出，需要根据API的接口自定义构造82+ // 2. 构造输入与输出，需要根据API的接口自定义构造
79- std::vector<int64_t> selfShape = {4};83+ std::vector<int64_t> selfShape = {4};
80- std::vector<int64_t> outShape = {2};84+ std::vector<int64_t> outShape = {2};
81- void* selfDeviceAddr = nullptr;85+ void* selfDeviceAddr = nullptr;
82- void* outDeviceAddr = nullptr;86+ void* outDeviceAddr = nullptr;
83- aclTensor* self = nullptr;87+ aclTensor* self = nullptr;
84- aclTensor* out = nullptr;88+ aclTensor* out = nullptr;
85- std::vector<float> selfHostData = {0, 10, 3, 0};89+ std::vector<float> selfHostData = {0, 10, 3, 0};
86- std::vector<float> outHostData = {2, 0};90+ std::vector<int64_t> outHostData = {0, 0};
87- int64_t numsamples = 2;91+ int64_t numsamples = 2;
88- bool replacement = false;92+ bool replacement = false;
89- int64_t seed = 1234;93+ int64_t seed = 1234;
90- int64_t offset = 0;94+ int64_t offset = 0;
91- // 创建self aclTensor95+ // 创建self aclTensor
92- ret = CreateAclTensor(selfHostData, selfShape, &selfDeviceAddr, aclDataType::ACL_FLOAT, &self);96+ ret = CreateAclTensor(selfHostData, selfShape, &selfDeviceAddr, aclDataType::ACL_FLOAT, &self);
93- CHECK_RET(ret == ACL_SUCCESS, return ret);97+ CHECK_RET(ret == ACL_SUCCESS, return ret);
94 98 
95- ret = CreateAclTensor(outHostData, outShape, &outDeviceAddr, aclDataType::ACL_INT64, &out);99+ ret = CreateAclTensor(outHostData, outShape, &outDeviceAddr, aclDataType::ACL_INT64, &out);
96- CHECK_RET(ret == ACL_SUCCESS, return ret);100+ CHECK_RET(ret == ACL_SUCCESS, return ret);
97 101 
98- // 3. 调用CANN算子库API，需要修改为具体的Api名称102+ // 3. 调用CANN算子库API，需要修改为具体的Api名称
99- uint64_t workspaceSize = 0;103+ uint64_t workspaceSize = 0;
100- aclOpExecutor* executor;104+ aclOpExecutor* executor;
101- // 调用aclnnMultinomial第一段接口105+ // 调用aclnnMultinomial第一段接口
102- ret = aclnnMultinomialGetWorkspaceSize(self, numsamples, replacement, seed, offset, out, &workspaceSize, &executor);106+ ret = aclnnMultinomialGetWorkspaceSize(self, numsamples, replacement, seed, offset, out, &workspaceSize, &executor);
103- CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnMultinomialGetWorkspaceSize failed. ERROR: %d\n", ret); return ret);107+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnMultinomialGetWorkspaceSize failed. ERROR: %d\n", ret); return ret);
104- // 根据第一段接口计算出的workspaceSize申请device内存108+ // 根据第一段接口计算出的workspaceSize申请device内存
105- void* workspaceAddr = nullptr;109+ void* workspaceAddr = nullptr;
106- if (workspaceSize > 0) {110+ if (workspaceSize > 0) {
107- ret = aclrtMalloc(&workspaceAddr, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST);111+ ret = aclrtMalloc(&workspaceAddr, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST);
108- CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("allocate workspace failed. ERROR: %d\n", ret); return ret);112+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("allocate workspace failed. ERROR: %d\n", ret); return ret);
109- }113+ }
110- // 调用aclnnMultinomial第二段接口114+ // 调用aclnnMultinomial第二段接口
111- ret = aclnnMultinomial(workspaceAddr, workspaceSize, executor, stream);115+ ret = aclnnMultinomial(workspaceAddr, workspaceSize, executor, stream);
112- CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnMultinomial failed. ERROR: %d\n", ret); return ret);116+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnMultinomial failed. ERROR: %d\n", ret); return ret);
113 117 
114- // 4. （固定写法）同步等待任务执行结束118+ // 4. （固定写法）同步等待任务执行结束
115- ret = aclrtSynchronizeStream(stream);119+ ret = aclrtSynchronizeStream(stream);
116- CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtSynchronizeStream failed. ERROR: %d\n", ret); return ret);120+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtSynchronizeStream failed. ERROR: %d\n", ret); return ret);
117 121 
118- // 5. 获取输出的值，将device侧内存上的结果拷贝至host侧，需要根据具体API的接口定义修改122+ // 5. 获取输出的值，将device侧内存上的结果拷贝至host侧，需要根据具体API的接口定义修改
119- auto size = GetShapeSize(selfShape);123+ auto size = GetShapeSize(outShape);
120- std::vector<float> resultData(size, 0);124+ std::vector<int64_t> resultData(size, 0);
121- ret = aclrtMemcpy(resultData.data(), resultData.size() * sizeof(resultData[0]), selfDeviceAddr,125+ ret = aclrtMemcpy(resultData.data(), resultData.size() * sizeof(resultData[0]), outDeviceAddr,
122- size * sizeof(resultData[0]), ACL_MEMCPY_DEVICE_TO_HOST);126+ size * sizeof(resultData[0]), ACL_MEMCPY_DEVICE_TO_HOST);
123- CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("copy result from device to host failed. ERROR: %d\n", ret); return ret);127+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("copy result from device to host failed. ERROR: %d\n", ret); return ret);
124- for (int64_t i = 0; i < size; i++) {128+ for (int64_t i = 0; i < size; i++) {
125- LOG_PRINT("result[%ld] is: %f\n", i, resultData[i]);129+ LOG_PRINT("result[%ld] is: %ld\n", i, resultData[i]);
126- }130+ }
127 131 
128- // 6. 释放aclTensor和aclScalar，需要根据具体API的接口定义修改132+ // 6. 释放aclTensor和aclScalar，需要根据具体API的接口定义修改
129- aclDestroyTensor(self);133+ aclDestroyTensor(self);
130- aclDestroyTensor(out);134+ aclDestroyTensor(out);
131 135 
132- // 7. 释放Device资源，需要根据具体API的接口定义修改136+ // 7. 释放Device资源，需要根据具体API的接口定义修改
133- aclrtFree(selfDeviceAddr);137+ aclrtFree(selfDeviceAddr);
134- aclrtFree(outDeviceAddr);138+ aclrtFree(outDeviceAddr);
135- if (workspaceSize > 0) {139+ if (workspaceSize > 0) {
136- aclrtFree(workspaceAddr);140+ aclrtFree(workspaceAddr);
137- }141+ }
138- aclrtDestroyStream(stream);142+ aclrtDestroyStream(stream);
139- aclrtResetDevice(deviceId);143+ aclrtResetDevice(deviceId);
140- aclFinalize();144+ aclFinalize();
141 145 
142- return 0;146+ return 0;
143-}147+}
Mrandom/stateless_sample_multinomial/op_api/aclnn_multinomial.cpp+54-46
@@ -22,24 +22,19 @@
22#include "math/cumsum/op_api/cumsum.h"22#include "math/cumsum/op_api/cumsum.h"
23#include "conversion/unsqueeze/op_host/op_api/unsqueeze.h"23#include "conversion/unsqueeze/op_host/op_api/unsqueeze.h"
24#include "random/dsa_random_uniform/op_host/op_api/dsa_random_uniform.h"24#include "random/dsa_random_uniform/op_host/op_api/dsa_random_uniform.h"
25-#include "random/stateless_random_uniform_v2/op_api/stateless_random_uniform_v2.h"
26#include "conversion/concat_d/op_api/concat_d.h"25#include "conversion/concat_d/op_api/concat_d.h"
27#include "aclnn_kernels/cast.h"26#include "aclnn_kernels/cast.h"
28#include "aclnn_kernels/reshape.h"27#include "aclnn_kernels/reshape.h"
29#include "math/greater_equal/op_api/greater_equal.h"28#include "math/greater_equal/op_api/greater_equal.h"
30#include "conversion/view_copy/op_api/view_copy.h"29#include "conversion/view_copy/op_api/view_copy.h"
31#include "aclnn_kernels/contiguous.h"30#include "aclnn_kernels/contiguous.h"
32-#include "aclnn/aclnn_base.h"
33#include "aclnn_kernels/common/op_error_check.h"31#include "aclnn_kernels/common/op_error_check.h"
34#include "op_api/aclnn_check.h"32#include "op_api/aclnn_check.h"
35#include "opdev/common_types.h"33#include "opdev/common_types.h"
36#include "opdev/shape_utils.h"34#include "opdev/shape_utils.h"
37-#include "opdev/data_type_utils.h"
38-#include "opdev/format_utils.h"
39#include "opdev/op_dfx.h"35#include "opdev/op_dfx.h"
40#include "opdev/op_executor.h"36#include "opdev/op_executor.h"
41#include "opdev/op_log.h"37#include "opdev/op_log.h"
42-#include "opdev/tensor_view_utils.h"
43#include "opdev/platform.h"38#include "opdev/platform.h"
44 39 
45using namespace op;40using namespace op;
@@ -60,29 +55,46 @@ static const std::initializer_list<op::DataType> ASCEND910_DTYPE_SUPPORT_LIST =
60static const std::initializer_list<op::DataType> ASCEND910B_DTYPE_SUPPORT_LIST = {55static const std::initializer_list<op::DataType> ASCEND910B_DTYPE_SUPPORT_LIST = {
61 op::DataType::DT_FLOAT, op::DataType::DT_FLOAT16, op::DataType::DT_DOUBLE, op::DataType::DT_BF16};56 op::DataType::DT_FLOAT, op::DataType::DT_FLOAT16, op::DataType::DT_DOUBLE, op::DataType::DT_BF16};
62 57 
63-static inline bool CheckSocVersionGe910B(void)58+static const std::initializer_list<op::DataType> ASCEND950_DTYPE_SUPPORT_LIST = {
64-{59+ op::DataType::DT_FLOAT, op::DataType::DT_FLOAT16, op::DataType::DT_DOUBLE, op::DataType::DT_BF16};
65- auto curArch = GetCurrentPlatformInfo().GetCurNpuArch();60+ 
66- return curArch == NpuArch::DAV_2201 || IsRegBase(curArch);61+static inline bool IsArch22(NpuArch arch) { return arch == NpuArch::DAV_2201; }
67-}62+ 
63+static inline bool IsArch35(NpuArch arch) { return arch == NpuArch::DAV_3510; }
68 64 
69static bool UseAicpuPath(const aclTensor* self, int64_t selfSize)65static bool UseAicpuPath(const aclTensor* self, int64_t selfSize)
70{66{
71- if (IsRegBase()) {67+ const auto curArch = GetCurrentPlatformInfo().GetCurNpuArch();
72- auto selfDtype = self->GetDataType();68+ if (IsArch35(curArch)) {
69+ const auto selfDtype = self->GetDataType();
73 return selfDtype != op::DataType::DT_FLOAT && selfDtype != op::DataType::DT_FLOAT16 &&70 return selfDtype != op::DataType::DT_FLOAT && selfDtype != op::DataType::DT_FLOAT16 &&
74 selfDtype != op::DataType::DT_BF16;71 selfDtype != op::DataType::DT_BF16;
75 }72 }
76- return !CheckSocVersionGe910B() || selfSize <= CPU_NPU_BOUNDARY;73+ if (IsArch22(curArch)) {
74+ return selfSize <= CPU_NPU_BOUNDARY;
75+ }
76+ return true;
77+}
78+ 
79+static bool ShouldCheckSeedOffsetDtype(int64_t selfSize)
80+{
81+ const auto curArch = GetCurrentPlatformInfo().GetCurNpuArch();
82+ if (IsArch22(curArch)) {
83+ return selfSize <= CPU_NPU_BOUNDARY;
84+ }
85+ return true;
77}86}
78 87 
79static inline const std::initializer_list<op::DataType>& GetDtypeSupportList()88static inline const std::initializer_list<op::DataType>& GetDtypeSupportList()
80{89{
81- if (CheckSocVersionGe910B()) {90+ const auto curArch = GetCurrentPlatformInfo().GetCurNpuArch();
82- return ASCEND910B_DTYPE_SUPPORT_LIST;91+ if (IsArch35(curArch)) {
83- } else {92+ return ASCEND950_DTYPE_SUPPORT_LIST;
84- return ASCEND910_DTYPE_SUPPORT_LIST;
85 }93 }
94+ if (IsArch22(curArch)) {
95+ return ASCEND910B_DTYPE_SUPPORT_LIST;
96+ }
97+ return ASCEND910_DTYPE_SUPPORT_LIST;
86}98}
87 99 
88static bool CheckNotNull(const aclTensor* self, const aclTensor* out)100static bool CheckNotNull(const aclTensor* self, const aclTensor* out)
@@ -128,12 +140,8 @@ static bool CheckDtypeValidTensor(const aclTensor* self, const aclTensor* seedTe
128 auto supportList = GetDtypeSupportList();140 auto supportList = GetDtypeSupportList();
129 OP_CHECK_DTYPE_NOT_SUPPORT(self, supportList, return false);141 OP_CHECK_DTYPE_NOT_SUPPORT(self, supportList, return false);
130 if (!(self->IsEmpty())) {142 if (!(self->IsEmpty())) {
131- const int64_t dimNum = static_cast<int64_t>(self->GetViewShape().GetDimNum());143+ const int64_t selfSize = self->GetViewShape().GetShapeSize();
132- int64_t selfSize = 1;144+ if (ShouldCheckSeedOffsetDtype(selfSize)) {
133- for (int64_t i = 0; i < dimNum; i++) {
134- selfSize *= self->GetViewShape().GetDim(i);
135- }
136- if (!CheckSocVersionGe910B() || selfSize <= CPU_NPU_BOUNDARY) {
137 OP_CHECK_DTYPE_NOT_SUPPORT(seedTensor, {op::DataType::DT_INT64}, return false);145 OP_CHECK_DTYPE_NOT_SUPPORT(seedTensor, {op::DataType::DT_INT64}, return false);
138 OP_CHECK_DTYPE_NOT_SUPPORT(offsetTensor, {op::DataType::DT_INT64}, return false);146 OP_CHECK_DTYPE_NOT_SUPPORT(offsetTensor, {op::DataType::DT_INT64}, return false);
139 }147 }
@@ -311,17 +319,18 @@ static const aclTensor* Run950AicoreMultinomialWithReplacement(const aclTensor*
311 auto sumSelf = l0op::ReduceSumOp(selfContiguous, dimArray, true, executor);319 auto sumSelf = l0op::ReduceSumOp(selfContiguous, dimArray, true, executor);
312 CHECK_RET(sumSelf != nullptr, nullptr);320 CHECK_RET(sumSelf != nullptr, nullptr);
313 321 
314- auto divSumSelf = l0op::RealDiv(selfContiguous, sumSelf, executor);322+ auto normProbsTensor = l0op::RealDiv(selfContiguous, sumSelf, executor);
315- CHECK_RET(divSumSelf != nullptr, nullptr);323+ CHECK_RET(normProbsTensor != nullptr, nullptr);
316 324 
317 // Build CDF via cumulative sum325 // Build CDF via cumulative sum
318 const aclTensor* dimTensor = executor->ConvertToTensor(&lastDim, 1, DataType::DT_INT32);326 const aclTensor* dimTensor = executor->ConvertToTensor(&lastDim, 1, DataType::DT_INT32);
319 CHECK_RET(dimTensor != nullptr, nullptr);327 CHECK_RET(dimTensor != nullptr, nullptr);
320- auto cdfTensor = l0op::Cumsum(divSumSelf, dimTensor, executor);328+ auto xTensor = l0op::Cumsum(normProbsTensor, dimTensor, executor);
321- CHECK_RET(cdfTensor != nullptr, nullptr);329+ CHECK_RET(xTensor != nullptr, nullptr);
322 330 
323 // Launch fused RNG + binary search kernel, output is int64 indices directly331 // Launch fused RNG + binary search kernel, output is int64 indices directly
324- auto multinomialOut = l0op::StatelessSampleMultinomial(cdfTensor, seedTensor, offsetTensor, numsamples, executor);332+ auto multinomialOut = l0op::StatelessSampleMultinomial(xTensor, normProbsTensor, seedTensor, offsetTensor,
333+ numsamples, executor);
325 CHECK_RET(multinomialOut != nullptr, nullptr);334 CHECK_RET(multinomialOut != nullptr, nullptr);
326 return multinomialOut;335 return multinomialOut;
327}336}
@@ -458,6 +467,7 @@ aclnnStatus aclnnMultinomialGetWorkspaceSize(const aclTensor* self, int64_t nums
458 for (int64_t i = 0; i < dimNum; i++) {467 for (int64_t i = 0; i < dimNum; i++) {
459 selfSize *= self->GetViewShape().GetDim(i);468 selfSize *= self->GetViewShape().GetDim(i);
460 }469 }
470+ const auto curArch = GetCurrentPlatformInfo().GetCurNpuArch();
461 471 
462 auto selfContiguous = l0op::Contiguous(self, uniqueExecutor.get());472 auto selfContiguous = l0op::Contiguous(self, uniqueExecutor.get());
463 CHECK_RET(selfContiguous != nullptr, ACLNN_ERR_PARAM_NULLPTR);473 CHECK_RET(selfContiguous != nullptr, ACLNN_ERR_PARAM_NULLPTR);
@@ -467,7 +477,7 @@ aclnnStatus aclnnMultinomialGetWorkspaceSize(const aclTensor* self, int64_t nums
467 multinomialOut = l0op::MultinomialWithReplacement(selfContiguous, numsamples, replacement, seed, offset,477 multinomialOut = l0op::MultinomialWithReplacement(selfContiguous, numsamples, replacement, seed, offset,
468 uniqueExecutor.get());478 uniqueExecutor.get());
469 } else if (!replacement || numsamples == 1) {479 } else if (!replacement || numsamples == 1) {
470- if (IsRegBase()) {480+ if (IsArch35(curArch)) {
471 // StatelessExponential takes tensor seed/offset; convert the scalar ones here.481 // StatelessExponential takes tensor seed/offset; convert the scalar ones here.
472 aclIntArray* seedList = uniqueExecutor->AllocIntArray(&seed, 1);482 aclIntArray* seedList = uniqueExecutor->AllocIntArray(&seed, 1);
473 CHECK_RET(seedList != nullptr, ACLNN_ERR_INNER_NULLPTR);483 CHECK_RET(seedList != nullptr, ACLNN_ERR_INNER_NULLPTR);
@@ -493,7 +503,7 @@ aclnnStatus aclnnMultinomialGetWorkspaceSize(const aclTensor* self, int64_t nums
493 uniqueExecutor.get());503 uniqueExecutor.get());
494 }504 }
495 } else {505 } else {
496- if (IsRegBase()) {506+ if (IsArch35(curArch)) {
497 aclIntArray* seedList = uniqueExecutor->AllocIntArray(&seed, 1);507 aclIntArray* seedList = uniqueExecutor->AllocIntArray(&seed, 1);
498 CHECK_RET(seedList != nullptr, ACLNN_ERR_INNER_NULLPTR);508 CHECK_RET(seedList != nullptr, ACLNN_ERR_INNER_NULLPTR);
499 auto seedTensor = uniqueExecutor->ConvertToTensor(seedList, op::DataType::DT_INT64);509 auto seedTensor = uniqueExecutor->ConvertToTensor(seedList, op::DataType::DT_INT64);
@@ -563,6 +573,7 @@ aclnnStatus aclnnMultinomialTensorGetWorkspaceSize(const aclTensor* self, int64_
563 for (int64_t i = 0; i < dimNum; i++) {573 for (int64_t i = 0; i < dimNum; i++) {
564 selfSize *= self->GetViewShape().GetDim(i);574 selfSize *= self->GetViewShape().GetDim(i);
565 }575 }
576+ const auto curArch = GetCurrentPlatformInfo().GetCurNpuArch();
566 577 
567 auto selfContiguous = l0op::Contiguous(self, uniqueExecutor.get());578 auto selfContiguous = l0op::Contiguous(self, uniqueExecutor.get());
568 CHECK_RET(selfContiguous != nullptr, ACLNN_ERR_PARAM_NULLPTR);579 CHECK_RET(selfContiguous != nullptr, ACLNN_ERR_PARAM_NULLPTR);
@@ -575,14 +586,7 @@ aclnnStatus aclnnMultinomialTensorGetWorkspaceSize(const aclTensor* self, int64_
575 multinomialOut = l0op::MultinomialWithReplacementTensor(selfContiguous, numsamples, replacement, seedTensor,586 multinomialOut = l0op::MultinomialWithReplacementTensor(selfContiguous, numsamples, replacement, seedTensor,
576 offsetAddOut, uniqueExecutor.get());587 offsetAddOut, uniqueExecutor.get());
577 } else if (!replacement || numsamples == 1) {588 } else if (!replacement || numsamples == 1) {
578- const aclTensor* randomUniform = nullptr;589+ if (IsArch35(curArch)) {
579- if (!IsRegBase()) {
580- randomUniform = GetRandomUniformNoReplaceMentTensor(selfContiguous, seedTensor, offsetAddOut,
581- uniqueExecutor.get());
582- CHECK_RET(randomUniform != nullptr, ACLNN_ERR_INNER_NULLPTR);
583- multinomialOut = RunMultinomialNoReplaceMent(selfContiguous, numsamples, randomUniform, out,
584- uniqueExecutor.get());
585- } else {
586 auto expInput = uniqueExecutor->AllocTensor(selfContiguous->GetViewShape(), selfContiguous->GetDataType(),590 auto expInput = uniqueExecutor->AllocTensor(selfContiguous->GetViewShape(), selfContiguous->GetDataType(),
587 selfContiguous->GetViewFormat());591 selfContiguous->GetViewFormat());
588 CHECK_RET(expInput != nullptr, ACLNN_ERR_INNER_NULLPTR);592 CHECK_RET(expInput != nullptr, ACLNN_ERR_INNER_NULLPTR);
@@ -591,19 +595,23 @@ aclnnStatus aclnnMultinomialTensorGetWorkspaceSize(const aclTensor* self, int64_
591 CHECK_RET(exp1Random != nullptr, ACLNN_ERR_PARAM_NULLPTR);595 CHECK_RET(exp1Random != nullptr, ACLNN_ERR_PARAM_NULLPTR);
592 multinomialOut = Run950AicoreMultinomialWithoutReplacement(selfContiguous, numsamples, exp1Random,596 multinomialOut = Run950AicoreMultinomialWithoutReplacement(selfContiguous, numsamples, exp1Random,
593 uniqueExecutor.get());597 uniqueExecutor.get());
598+ } else {
599+ auto randomUniform = GetRandomUniformNoReplaceMentTensor(selfContiguous, seedTensor, offsetAddOut,
600+ uniqueExecutor.get());
601+ CHECK_RET(randomUniform != nullptr, ACLNN_ERR_INNER_NULLPTR);
602+ multinomialOut = RunMultinomialNoReplaceMent(selfContiguous, numsamples, randomUniform, out,
603+ uniqueExecutor.get());
594 }604 }
595 } else {605 } else {
596- const aclTensor* randomUniform = nullptr;606+ if (IsArch35(curArch)) {
597- if (!IsRegBase()) {607+ multinomialOut = Run950AicoreMultinomialWithReplacement(selfContiguous, numsamples, seedTensor,
598- randomUniform = GetRandomUniformReplaceMentTensor(numsamples, seedTensor, offsetAddOut,608+ offsetAddOut, uniqueExecutor.get());
599- uniqueExecutor.get());609+ } else {
610+ auto randomUniform = GetRandomUniformReplaceMentTensor(numsamples, seedTensor, offsetAddOut,
611+ uniqueExecutor.get());
600 CHECK_RET(randomUniform != nullptr, ACLNN_ERR_INNER_NULLPTR);612 CHECK_RET(randomUniform != nullptr, ACLNN_ERR_INNER_NULLPTR);
601 multinomialOut = RunMultinomialReplaceMent(selfContiguous, numsamples, randomUniform, out,613 multinomialOut = RunMultinomialReplaceMent(selfContiguous, numsamples, randomUniform, out,
602 uniqueExecutor.get());614 uniqueExecutor.get());
603- } else {
604- // RegBase: use the fused 950 AICore kernel directly (offsetAddOut already carries offsetTensor + offset).
605- multinomialOut = Run950AicoreMultinomialWithReplacement(selfContiguous, numsamples, seedTensor,
606- offsetAddOut, uniqueExecutor.get());
607 }615 }
608 }616 }
609 CHECK_RET(multinomialOut != nullptr, ACLNN_ERR_PARAM_NULLPTR);617 CHECK_RET(multinomialOut != nullptr, ACLNN_ERR_PARAM_NULLPTR);
Mrandom/stateless_sample_multinomial/op_api/aclnn_multinomial.h+14-11
@@ -39,38 +39,41 @@ ACLNN_API aclnnStatus aclnnMultinomial(void* workspace, uint64_t workspaceSize,
39 * @param [in] self: npu device侧的aclTensor，shape为(N, C)或(C)，self的取值范围需要大于等于0且self与out的维度一致。39 * @param [in] self: npu device侧的aclTensor，shape为(N, C)或(C)，self的取值范围需要大于等于0且self与out的维度一致。
40 * 数据类型支持BFLOAT16、FLOAT16、FLOAT、DOUBLE。数据格式支持ND。支持非连续的Tensor。40 * 数据类型支持BFLOAT16、FLOAT16、FLOAT、DOUBLE。数据格式支持ND。支持非连续的Tensor。
41 * @param [in] numsamples: host侧的整形，从每个多项分布中抽取的样本数。41 * @param [in] numsamples: host侧的整形，从每个多项分布中抽取的样本数。
42- * numsamples为非负整数，当replacement为false时，numsamples不大于C。输入为INT64_t数据类型。42+ * numsamples为正整数，当replacement为false时，numsamples不大于C。输入为INT64_T数据类型。
43 * @param [in] replacement: host侧的布尔类型，决定了抽样时元素是否有放回。输入为BOOL数据类型。43 * @param [in] replacement: host侧的布尔类型，决定了抽样时元素是否有放回。输入为BOOL数据类型。
44- * @param [in] seedTensor: npu device侧的aclTensor，数据类型支持INT64_t。数据格式支持ND。44+ * @param [in] seedTensor: npu device侧的aclTensor，数据类型支持INT64。数据格式支持ND。
45 * 随机数生成器的种子,它影响生成的随机数序列。45 * 随机数生成器的种子,它影响生成的随机数序列。
46- * @param [in] offsetTensor: npu device侧的aclTensor，数据类型支持INT64_t。数据格式支持ND。46+ * @param [in] offsetTensor: npu device侧的aclTensor，数据类型支持INT64。数据格式支持ND。
47 * 随机数生成器的偏移量，它影响生成的随机数序列的位置。设置偏移量后，生成的随机数序列会从指定位置开始。47 * 随机数生成器的偏移量，它影响生成的随机数序列的位置。设置偏移量后，生成的随机数序列会从指定位置开始。
48 * @param [in] offset: host侧的整型，随机数生成器的偏移量，它影响生成的随机数序列的位置。输入为INT64_T数据类型。48 * @param [in] offset: host侧的整型，随机数生成器的偏移量，它影响生成的随机数序列的位置。输入为INT64_T数据类型。
49- * @param [in] out: npu device侧的aclTensor，数据类型支持INT64。shape为(N, numsamples)或(numsamples)，self与out的维度一致。49+ * @param [in] out: npu device侧的aclTensor，数据类型支持INT64。shape为(N,
50- * 支持非连续的Tensor，数据格式支持ND。50+ * numsamples)或(numsamples)，self与out的维度一致。 支持非连续的Tensor，数据格式支持ND。
51 * @param [out] workspaceSize: 返回用户需要在npu device侧申请的workspace大小。51 * @param [out] workspaceSize: 返回用户需要在npu device侧申请的workspace大小。
52 * @param [out] executor: 返回op执行器，包含算子计算流程。52 * @param [out] executor: 返回op执行器，包含算子计算流程。
53 * @return aclnnStatus: 返回状态码。53 * @return aclnnStatus: 返回状态码。
54 */54 */
55-ACLNN_API aclnnStatus aclnnMultinomialTensorGetWorkspaceSize(const aclTensor* self, int64_t numsamples, bool replacement,55+ACLNN_API aclnnStatus aclnnMultinomialTensorGetWorkspaceSize(const aclTensor* self, int64_t numsamples,
56- const aclTensor* seedTensor, const aclTensor* offsetTensor, int64_t offset,56+ bool replacement, const aclTensor* seedTensor,
57- aclTensor* out, uint64_t* workspaceSize, aclOpExecutor** executor);57+ const aclTensor* offsetTensor, int64_t offset,
58+ aclTensor* out, uint64_t* workspaceSize,
59+ aclOpExecutor** executor);
58 60 
59/**61/**
60 * @brief aclnnMultinomialTensor的第二段接口，用于执行计算。62 * @brief aclnnMultinomialTensor的第二段接口，用于执行计算。
61 *63 *
62 * 算子功能：在输入张量中根据每个对象分布的概率，抽取numsamples个样本，并将这些样本的索引存储在输出张量中。64 * 算子功能：在输入张量中根据每个对象分布的概率，抽取numsamples个样本，并将这些样本的索引存储在输出张量中。
63 * @param [in] workspace: 在npu device侧申请的workspace内存起址。65 * @param [in] workspace: 在npu device侧申请的workspace内存起址。
64- * @param [in] workspaceSize: 在npu device侧申请的workspace大小，由第一段接口aclnnMultinomialTensorGetWorkspaceSize获取。66+ * @param [in] workspaceSize: 在npu
67+ * device侧申请的workspace大小，由第一段接口aclnnMultinomialTensorGetWorkspaceSize获取。
65 * @param [in] executor: op执行器，包含了算子计算流程。68 * @param [in] executor: op执行器，包含了算子计算流程。
66 * @param [in] stream: acl stream流。69 * @param [in] stream: acl stream流。
67 * @return aclnnStatus: 返回状态码。70 * @return aclnnStatus: 返回状态码。
68 */71 */
69ACLNN_API aclnnStatus aclnnMultinomialTensor(void* workspace, uint64_t workspaceSize, aclOpExecutor* executor,72ACLNN_API aclnnStatus aclnnMultinomialTensor(void* workspace, uint64_t workspaceSize, aclOpExecutor* executor,
70- aclrtStream stream);73+ aclrtStream stream);
71 74 
72#ifdef __cplusplus75#ifdef __cplusplus
73}76}
74#endif77#endif
75 78 
76-#endif // OP_API_INC_LEVEL2_ACLNN_MULTINOMIAL_H_79+#endif // OP_API_INC_LEVEL2_ACLNN_MULTINOMIAL_H_
Mrandom/stateless_sample_multinomial/op_api/stateless_sample_multinomial.cpp+11-4
@@ -32,10 +32,11 @@ namespace l0op {
32 32 
33OP_TYPE_REGISTER(StatelessSampleMultinomial);33OP_TYPE_REGISTER(StatelessSampleMultinomial);
34 34 
35-const aclTensor* StatelessSampleMultinomial(const aclTensor* xTensor, const aclTensor* seedTensor,35+const aclTensor* StatelessSampleMultinomial(const aclTensor* xTensor, const aclTensor* normProbsTensor,
36- const aclTensor* offsetTensor, int64_t numsamples, aclOpExecutor* executor)36+ const aclTensor* seedTensor, const aclTensor* offsetTensor,
37+ int64_t numsamples, aclOpExecutor* executor)
37{38{
38- L0_DFX(StatelessSampleMultinomial, xTensor, seedTensor, offsetTensor);39+ L0_DFX(StatelessSampleMultinomial, xTensor, normProbsTensor, seedTensor, offsetTensor);
39 40 
40 auto outShape = xTensor->GetViewShape();41 auto outShape = xTensor->GetViewShape();
41 auto dimNum = outShape.GetDimNum();42 auto dimNum = outShape.GetDimNum();
@@ -45,11 +46,17 @@ const aclTensor* StatelessSampleMultinomial(const aclTensor* xTensor, const aclT
45 CHECK_RET(out != nullptr, nullptr);46 CHECK_RET(out != nullptr, nullptr);
46 47 
47 auto ret = ADD_TO_LAUNCHER_LIST_AICORE(StatelessSampleMultinomial, OP_ATTR_NAMES({"num_samples"}),48 auto ret = ADD_TO_LAUNCHER_LIST_AICORE(StatelessSampleMultinomial, OP_ATTR_NAMES({"num_samples"}),
48- OP_INPUT(xTensor, seedTensor, offsetTensor), OP_OUTPUT(out),49+ OP_INPUT(xTensor, normProbsTensor, seedTensor, offsetTensor), OP_OUTPUT(out),
49 OP_ATTR(numsamples));50 OP_ATTR(numsamples));
50 CHECK_RET(ret == ACLNN_SUCCESS, nullptr);51 CHECK_RET(ret == ACLNN_SUCCESS, nullptr);
51 52 
52 return out;53 return out;
53}54}
54 55 
56+const aclTensor* StatelessSampleMultinomial(const aclTensor* xTensor, const aclTensor* seedTensor,
57+ const aclTensor* offsetTensor, int64_t numsamples, aclOpExecutor* executor)
58+{
59+ return StatelessSampleMultinomial(xTensor, nullptr, seedTensor, offsetTensor, numsamples, executor);
60+}
61+ 
55} // namespace l0op62} // namespace l0op
Mrandom/stateless_sample_multinomial/op_api/stateless_sample_multinomial.h+12-5
@@ -20,18 +20,25 @@
20namespace l0op {20namespace l0op {
21 21 
22/**22/**
23- * @brief Generate multinomial samples with replacement using binary search on x.23+ * @brief Generate multinomial samples with replacement using binary search on a CDF.
24 * Fuses U(0,1] generation (Philox RNG) + binary search for direct index output.24 * Fuses U(0,1] generation (Philox RNG) + binary search for direct index output.
25- * Handles zero-probability categories via backward walk (PyTorch CUDA compatible).25+ * Uses normalized probabilities for zero-probability fallback when provided; otherwise falls back to
26+ * comparing adjacent CDF values.
26 * numDist and numCategories are derived from xTensor shape.27 * numDist and numCategories are derived from xTensor shape.
27 *28 *
28- * @param xTensor Input tensor (DT_FLOAT, shape [numDist, numCategories] or [numCategories])29+ * @param xTensor CDF, shape [numDist, numCategories] or [numCategories]
29- * @param seedTensor Seed tensor (INT64/UINT64, shape [1])30+ * @param normProbsTensor Optional normalized probabilities, with the same shape and dtype as xTensor
30- * @param offsetTensor Offset tensor (INT64/UINT64, shape [1])31+ * @param seedTensor Seed tensor (INT64, shape [1])
32+ * @param offsetTensor Offset tensor (INT64, shape [1])
31 * @param numsamples Number of samples per distribution33 * @param numsamples Number of samples per distribution
32 * @param executor Op executor34 * @param executor Op executor
33 * @return Output tensor with shape {numDist, numsamples}, dtype DT_INT6435 * @return Output tensor with shape {numDist, numsamples}, dtype DT_INT64
34 */36 */
37+const aclTensor* StatelessSampleMultinomial(const aclTensor* xTensor, const aclTensor* normProbsTensor,
38+ const aclTensor* seedTensor, const aclTensor* offsetTensor,
39+ int64_t numsamples, aclOpExecutor* executor);
40+ 
41+// Compatibility overload: uses adjacent CDF values when normalized probabilities are unavailable.
35const aclTensor* StatelessSampleMultinomial(const aclTensor* xTensor, const aclTensor* seedTensor,42const aclTensor* StatelessSampleMultinomial(const aclTensor* xTensor, const aclTensor* seedTensor,
36 const aclTensor* offsetTensor, int64_t numsamples, aclOpExecutor* executor);43 const aclTensor* offsetTensor, int64_t numsamples, aclOpExecutor* executor);
37 44 
Mrandom/stateless_sample_multinomial/op_host/arch35/stateless_sample_multinomial_tiling.cpp+29-2
@@ -22,8 +22,9 @@
22namespace optiling {22namespace optiling {
23 23 
24static constexpr uint16_t INPUT_IDX_X = 0;24static constexpr uint16_t INPUT_IDX_X = 0;
25-static constexpr uint16_t INPUT_IDX_SEED = 1;25+static constexpr uint16_t INPUT_IDX_NORM_PROBS = 1;
26-static constexpr uint16_t INPUT_IDX_OFFSET = 2;26+static constexpr uint16_t INPUT_IDX_SEED = 2;
27+static constexpr uint16_t INPUT_IDX_OFFSET = 3;
27static constexpr uint16_t OUTPUT_IDX_Y = 0;28static constexpr uint16_t OUTPUT_IDX_Y = 0;
28static constexpr int64_t DCACHE_SIZE = 128 * 1024;29static constexpr int64_t DCACHE_SIZE = 128 * 1024;
29static constexpr int64_t CORE_ALIGN_SIZE = 256;30static constexpr int64_t CORE_ALIGN_SIZE = 256;
@@ -36,6 +37,8 @@ OpTilingConfig StatelessSampleMultinomialTiling::BuildOpConfig()
36 config.inputCheckRules = {{INPUT_IDX_X, {{ge::DT_FLOAT, ge::DT_FLOAT16, ge::DT_BF16}, -1, {}, nullptr}},37 config.inputCheckRules = {{INPUT_IDX_X, {{ge::DT_FLOAT, ge::DT_FLOAT16, ge::DT_BF16}, -1, {}, nullptr}},
37 {INPUT_IDX_SEED, {{ge::DT_INT64}, 1, {}, nullptr}},38 {INPUT_IDX_SEED, {{ge::DT_INT64}, 1, {}, nullptr}},
38 {INPUT_IDX_OFFSET, {{ge::DT_INT64}, 1, {}, nullptr}}};39 {INPUT_IDX_OFFSET, {{ge::DT_INT64}, 1, {}, nullptr}}};
40+ config.optionalInputCheckRules = {
41+ {INPUT_IDX_NORM_PROBS, {{ge::DT_FLOAT, ge::DT_FLOAT16, ge::DT_BF16}, -1, {}, nullptr}}};
39 config.outputCheckRules = {{OUTPUT_IDX_Y, {{ge::DT_INT64}, -1, {}, nullptr}}};42 config.outputCheckRules = {{OUTPUT_IDX_Y, {{ge::DT_INT64}, -1, {}, nullptr}}};
40 43 
41 config.getOutputSize = [](gert::TilingContext* ctx, int64_t& size) {44 config.getOutputSize = [](gert::TilingContext* ctx, int64_t& size) {
@@ -65,8 +68,32 @@ OpTilingConfig StatelessSampleMultinomialTiling::BuildOpConfig()
65 return config;68 return config;
66}69}
67 70 
71+ge::graphStatus StatelessSampleMultinomialTiling::CheckXRankAndNormProbsShape()
72+{
73+ auto xShapePtr = context_->GetInputShape(INPUT_IDX_X);
74+ OP_CHECK_NULL_WITH_CONTEXT(context_, xShapePtr);
75+ const auto& xShape = xShapePtr->GetStorageShape();
76+ auto xDimNum = xShape.GetDimNum();
77+ OP_CHECK_IF(xDimNum != 1 && xDimNum != 2,
78+ OP_LOGE(context_->GetNodeName(), "x must be 1D or 2D, but got %zuD", xDimNum), return ge::GRAPH_FAILED);
79+ 
80+ auto normProbsShapePtr = context_->GetOptionalInputShape(INPUT_IDX_NORM_PROBS);
81+ if (normProbsShapePtr != nullptr) {
82+ const auto& normProbsShape = normProbsShapePtr->GetStorageShape();
83+ OP_CHECK_IF(normProbsShape != xShape,
84+ OP_LOGE(context_->GetNodeName(), "the shapes of x and norm_probs must be the same"),
85+ return ge::GRAPH_FAILED);
86+ }
87+ return ge::GRAPH_SUCCESS;
88+}
89+ 
68ge::graphStatus StatelessSampleMultinomialTiling::UniqueProcess()90ge::graphStatus StatelessSampleMultinomialTiling::UniqueProcess()
69{91{
92+ auto ret = CheckXRankAndNormProbsShape();
93+ if (ret != ge::GRAPH_SUCCESS) {
94+ return ret;
95+ }
96+ 
70 auto xShape = context_->GetInputShape(INPUT_IDX_X);97 auto xShape = context_->GetInputShape(INPUT_IDX_X);
71 if (xShape == nullptr) {98 if (xShape == nullptr) {
72 return ge::GRAPH_FAILED;99 return ge::GRAPH_FAILED;
Mrandom/stateless_sample_multinomial/op_host/arch35/stateless_sample_multinomial_tiling.h+3-1
@@ -24,13 +24,15 @@ namespace optiling {
24class StatelessSampleMultinomialTiling : public RandomTilingArch35 {24class StatelessSampleMultinomialTiling : public RandomTilingArch35 {
25public:25public:
26 explicit StatelessSampleMultinomialTiling(gert::TilingContext* context)26 explicit StatelessSampleMultinomialTiling(gert::TilingContext* context)
27- : RandomTilingArch35(context, BuildOpConfig()) {}27+ : RandomTilingArch35(context, BuildOpConfig())
28+ {}
28 29 
29protected:30protected:
30 ge::graphStatus UniqueProcess() override;31 ge::graphStatus UniqueProcess() override;
31 32 
32private:33private:
33 static OpTilingConfig BuildOpConfig();34 static OpTilingConfig BuildOpConfig();
35+ ge::graphStatus CheckXRankAndNormProbsShape();
34};36};
35 37 
36} // namespace optiling38} // namespace optiling
Drandom/stateless_sample_multinomial/op_host/config/ascend950/stateless_sample_multinomial_binary.json+0-167
@@ -1,167 +0,0 @@
1-{
2- "op_type": "StatelessSampleMultinomial",
3- "op_list": [
4- {
5- "bin_filename": "StatelessSampleMultinomial_5b71eff3d138906d9504866afac5dfd0",
6- "inputs": [
7- {
8- "name": "x",
9- "index": 0,
10- "dtype": "float32",
11- "format": "ND",
12- "paramType": "required",
13- "shape": [
14- -2
15- ]
16- },
17- {
18- "name": "seed",
19- "index": 1,
20- "dtype": "int64",
21- "format": "ND",
22- "paramType": "required",
23- "shape": [
24- -2
25- ]
26- },
27- {
28- "name": "offset",
29- "index": 2,
30- "dtype": "int64",
31- "format": "ND",
32- "paramType": "required",
33- "shape": [
34- -2
35- ]
36- }
37- ],
38- "outputs": [
39- {
40- "name": "y",
41- "index": 0,
42- "dtype": "int64",
43- "format": "ND",
44- "paramType": "required",
45- "shape": [
46- -2
47- ]
48- }
49- ],
50- "attrs": [
51- {
52- "name": "num_samples",
53- "dtype": "int",
54- "value": null
55- }
56- ]
57- },
58- {
59- "bin_filename": "StatelessSampleMultinomial_c3dbc22fef04f458489d12b4714c5c03",
60- "inputs": [
61- {
62- "name": "x",
63- "index": 0,
64- "dtype": "float16",
65- "format": "ND",
66- "paramType": "required",
67- "shape": [
68- -2
69- ]
70- },
71- {
72- "name": "seed",
73- "index": 1,
74- "dtype": "int64",
75- "format": "ND",
76- "paramType": "required",
77- "shape": [
78- -2
79- ]
80- },
81- {
82- "name": "offset",
83- "index": 2,
84- "dtype": "int64",
85- "format": "ND",
86- "paramType": "required",
87- "shape": [
88- -2
89- ]
90- }
91- ],
92- "outputs": [
93- {
94- "name": "y",
95- "index": 0,
96- "dtype": "int64",
97- "format": "ND",
98- "paramType": "required",
99- "shape": [
100- -2
101- ]
102- }
103- ],
104- "attrs": [
105- {
106- "name": "num_samples",
107- "dtype": "int",
108- "value": null
109- }
110- ]
111- },
112- {
113- "bin_filename": "StatelessSampleMultinomial_446d27c0956f8c99984a570a178bdd84",
114- "inputs": [
115- {
116- "name": "x",
117- "index": 0,
118- "dtype": "bfloat16",
119- "format": "ND",
120- "paramType": "required",
121- "shape": [
122- -2
123- ]
124- },
125- {
126- "name": "seed",
127- "index": 1,
128- "dtype": "int64",
129- "format": "ND",
130- "paramType": "required",
131- "shape": [
132- -2
133- ]
134- },
135- {
136- "name": "offset",
137- "index": 2,
138- "dtype": "int64",
139- "format": "ND",
140- "paramType": "required",
141- "shape": [
142- -2
143- ]
144- }
145- ],
146- "outputs": [
147- {
148- "name": "y",
149- "index": 0,
150- "dtype": "int64",
151- "format": "ND",
152- "paramType": "required",
153- "shape": [
154- -2
155- ]
156- }
157- ],
158- "attrs": [
159- {
160- "name": "num_samples",
161- "dtype": "int",
162- "value": null
163- }
164- ]
165- }
166- ]
167-}
Drandom/stateless_sample_multinomial/op_host/config/ascend950/stateless_sample_multinomial_simplified_key.ini+0-2
@@ -1,2 +0,0 @@
1-[StatelessSampleMultinomial]
2-default=0
Mrandom/stateless_sample_multinomial/op_host/stateless_sample_multinomial_def.cpp+5-0
@@ -25,6 +25,11 @@ public:
25 .DataType({ge::DT_FLOAT, ge::DT_FLOAT16, ge::DT_BF16})25 .DataType({ge::DT_FLOAT, ge::DT_FLOAT16, ge::DT_BF16})
26 .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})26 .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
27 .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});27 .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
28+ this->Input("norm_probs")
29+ .ParamType(OPTIONAL)
30+ .DataType({ge::DT_FLOAT, ge::DT_FLOAT16, ge::DT_BF16})
31+ .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
32+ .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
28 this->Input("seed")33 this->Input("seed")
29 .ParamType(REQUIRED)34 .ParamType(REQUIRED)
30 .DataType({ge::DT_INT64, ge::DT_INT64, ge::DT_INT64})35 .DataType({ge::DT_INT64, ge::DT_INT64, ge::DT_INT64})
Mrandom/stateless_sample_multinomial/op_kernel/arch35/stateless_sample_multinomial_impl.h+57-24
@@ -13,22 +13,22 @@
13 13 
14#include "../../random_common/arch35/random_kernel_base.h"14#include "../../random_common/arch35/random_kernel_base.h"
15#include "simt_api/asc_simt.h"15#include "simt_api/asc_simt.h"
16+#include "utils/debug/asc_assert.h"
16 17 
17namespace StatelessSampleMultinomial {18namespace StatelessSampleMultinomial {
18using namespace AscendC;19using namespace AscendC;
19using namespace RandomKernelBase;20using namespace RandomKernelBase;
20 21 
21-constexpr static uint32_t NUM_4 = 4;
22constexpr uint16_t CORE_THREAD_NUM_U32 = 256;22constexpr uint16_t CORE_THREAD_NUM_U32 = 256;
23constexpr uint16_t CORE_THREAD_NUM_U64 = 256;23constexpr uint16_t CORE_THREAD_NUM_U64 = 256;
24constexpr static int64_t SAMPLES_ALIGNMENT = 128;24constexpr static int64_t SAMPLES_ALIGNMENT = 128;
25constexpr static uint32_t UNROLL_FACTOR = 4;25constexpr static uint32_t UNROLL_FACTOR = 4;
26 26 
27-template <typename XT, typename IndexT, uint16_t THREAD_LAUNCH_BOUND>27+template <typename XT, typename NormProbsT, bool HAS_NORM_PROBS, typename IndexT, uint16_t THREAD_LAUNCH_BOUND>
28__simt_vf__ __aicore__ LAUNCH_BOUND(THREAD_LAUNCH_BOUND) inline void SimtUniformRandomBinarySearch(28__simt_vf__ __aicore__ LAUNCH_BOUND(THREAD_LAUNCH_BOUND) inline void SimtUniformRandomBinarySearch(
29- __gm__ volatile int64_t* outputGM, __gm__ volatile XT* xGM, IndexT elementNum, int64_t seed, IndexT numsamples,29+ __gm__ volatile int64_t* outputGM, __gm__ XT* xGM, __gm__ NormProbsT* normProbsGM, IndexT elementNum, int64_t seed,
30- uint32_t numCat, uint64_t nsMagic, uint64_t nsShift, IndexT samplesAligned, uint32_t baseOffsetLo,30+ IndexT numsamples, uint32_t numCat, uint64_t nsMagic, uint64_t nsShift, IndexT samplesAligned,
31- uint32_t baseOffsetHi)31+ uint32_t baseOffsetLo, uint32_t baseOffsetHi)
32{32{
33 uint32_t key[ALG_KEY_SIZE] = {0, 0};33 uint32_t key[ALG_KEY_SIZE] = {0, 0};
34 key[0] = static_cast<uint32_t>(seed);34 key[0] = static_cast<uint32_t>(seed);
@@ -45,7 +45,11 @@ __simt_vf__ __aicore__ LAUNCH_BOUND(THREAD_LAUNCH_BOUND) inline void SimtUniform
45 45 
46 IndexT d = static_cast<IndexT>(Simt::UintDiv(static_cast<uint64_t>(baseIndex), nsMagic, nsShift));46 IndexT d = static_cast<IndexT>(Simt::UintDiv(static_cast<uint64_t>(baseIndex), nsMagic, nsShift));
47 IndexT s = baseIndex - d * numsamples;47 IndexT s = baseIndex - d * numsamples;
48- __gm__ volatile XT* x = xGM + d * static_cast<IndexT>(numCat);48+ __gm__ XT* x = xGM + d * static_cast<IndexT>(numCat);
49+ __gm__ NormProbsT* normProbs = normProbsGM;
50+ if constexpr (HAS_NORM_PROBS) {
51+ normProbs += d * static_cast<IndexT>(numCat);
52+ }
49 uint64_t subsequence = static_cast<uint64_t>(d) * samplesAligned + s;53 uint64_t subsequence = static_cast<uint64_t>(d) * samplesAligned + s;
50 54 
51 for (uint32_t k = 0; k < count; k++) {55 for (uint32_t k = 0; k < count; k++) {
@@ -57,6 +61,7 @@ __simt_vf__ __aicore__ LAUNCH_BOUND(THREAD_LAUNCH_BOUND) inline void SimtUniform
57 61 
58 IndexT start = 0;62 IndexT start = 0;
59 IndexT end = numCat;63 IndexT end = numCat;
64+ assert(x[numCat - 1] > static_cast<XT>(0));
60 65 
61 while (end > start) {66 while (end > start) {
62 IndexT mid = start + ((end - start) >> 1);67 IndexT mid = start + ((end - start) >> 1);
@@ -70,8 +75,14 @@ __simt_vf__ __aicore__ LAUNCH_BOUND(THREAD_LAUNCH_BOUND) inline void SimtUniform
70 if (start >= numCat) {75 if (start >= numCat) {
71 start = numCat - 1;76 start = numCat - 1;
72 }77 }
73- while (start >= 1 && x[start] == x[start - 1]) {78+ if constexpr (HAS_NORM_PROBS) {
74- start--;79+ while (start >= 1 && normProbs[start] == static_cast<NormProbsT>(0)) {
80+ start--;
81+ }
82+ } else {
83+ while (start >= 1 && x[start] == x[start - 1]) {
84+ start--;
85+ }
75 }86 }
76 87 
77 outputGM[baseIndex + k] = static_cast<int64_t>(start);88 outputGM[baseIndex + k] = static_cast<int64_t>(start);
@@ -80,44 +91,50 @@ __simt_vf__ __aicore__ LAUNCH_BOUND(THREAD_LAUNCH_BOUND) inline void SimtUniform
80 if (++s >= numsamples) {91 if (++s >= numsamples) {
81 s = 0;92 s = 0;
82 x += numCat;93 x += numCat;
94+ if constexpr (HAS_NORM_PROBS) {
95+ normProbs += numCat;
96+ }
83 subsequence += samplesAligned - numsamples;97 subsequence += samplesAligned - numsamples;
84 }98 }
85 }99 }
86 }100 }
87}101}
88 102 
89-template <typename XT>103+template <typename XT, typename NormProbsT>
90class StatelessSampleMultinomialOp {104class StatelessSampleMultinomialOp {
91public:105public:
92 __aicore__ inline StatelessSampleMultinomialOp(){};106 __aicore__ inline StatelessSampleMultinomialOp(){};
93- __aicore__ inline void Init(GM_ADDR y, GM_ADDR x, GM_ADDR seed, GM_ADDR offset, GM_ADDR workspace,107+ __aicore__ inline void Init(GM_ADDR y, GM_ADDR x, GM_ADDR normProbs, GM_ADDR seed, GM_ADDR offset,
94- const RandomUnifiedSimtTilingDataStruct* __restrict tilingData, TPipe* pipe);108+ GM_ADDR workspace, const RandomUnifiedSimtTilingDataStruct* __restrict tilingData,
109+ TPipe* pipe);
95 __aicore__ inline void Process();110 __aicore__ inline void Process();
96 111 
97private:112private:
98 const RandomUnifiedSimtTilingDataStruct* tilingData_;113 const RandomUnifiedSimtTilingDataStruct* tilingData_;
99 GlobalTensor<int64_t> outputGM_;114 GlobalTensor<int64_t> outputGM_;
100 GM_ADDR xGM_;115 GM_ADDR xGM_;
116+ GM_ADDR normProbsGM_;
101 GM_ADDR seedGM_;117 GM_ADDR seedGM_;
102 GM_ADDR offsetGM_;118 GM_ADDR offsetGM_;
103 uint32_t blockIdx_;119 uint32_t blockIdx_;
104};120};
105 121 
106-template <typename XT>122+template <typename XT, typename NormProbsT>
107-__aicore__ inline void StatelessSampleMultinomialOp<XT>::Init(123+__aicore__ inline void StatelessSampleMultinomialOp<XT, NormProbsT>::Init(
108- GM_ADDR y, GM_ADDR x, GM_ADDR seed, GM_ADDR offset, GM_ADDR workspace,124+ GM_ADDR y, GM_ADDR x, GM_ADDR normProbs, GM_ADDR seed, GM_ADDR offset, GM_ADDR workspace,
109 const RandomUnifiedSimtTilingDataStruct* __restrict tilingData, TPipe* pipe)125 const RandomUnifiedSimtTilingDataStruct* __restrict tilingData, TPipe* pipe)
110{126{
111 tilingData_ = tilingData;127 tilingData_ = tilingData;
112 outputGM_.SetGlobalBuffer((__gm__ int64_t*)y);128 outputGM_.SetGlobalBuffer((__gm__ int64_t*)y);
113 xGM_ = x;129 xGM_ = x;
130+ normProbsGM_ = normProbs;
114 seedGM_ = seed;131 seedGM_ = seed;
115 offsetGM_ = offset;132 offsetGM_ = offset;
116 blockIdx_ = GetBlockIdx();133 blockIdx_ = GetBlockIdx();
117}134}
118 135 
119-template <typename XT>136+template <typename XT, typename NormProbsT>
120-__aicore__ inline void StatelessSampleMultinomialOp<XT>::Process()137+__aicore__ inline void StatelessSampleMultinomialOp<XT, NormProbsT>::Process()
121{138{
122 if (blockIdx_ >= tilingData_->usedCoreNum) {139 if (blockIdx_ >= tilingData_->usedCoreNum) {
123 return;140 return;
@@ -141,16 +158,32 @@ __aicore__ inline void StatelessSampleMultinomialOp<XT>::Process()
141 uint32_t baseOffsetHi = static_cast<uint32_t>(baseOffset >> 32);158 uint32_t baseOffsetHi = static_cast<uint32_t>(baseOffset >> 32);
142 159 
143 uint64_t samplesAligned = ((numsamples + SAMPLES_ALIGNMENT - 1) / SAMPLES_ALIGNMENT) * SAMPLES_ALIGNMENT;160 uint64_t samplesAligned = ((numsamples + SAMPLES_ALIGNMENT - 1) / SAMPLES_ALIGNMENT) * SAMPLES_ALIGNMENT;
144- 
145 if (useUint64Index) {161 if (useUint64Index) {
146- asc_vf_call<SimtUniformRandomBinarySearch<XT, uint64_t, CORE_THREAD_NUM_U64>>(162+ if (normProbsGM_ != nullptr) {
147- dim3(CORE_THREAD_NUM_U64), (__gm__ volatile int64_t*)(outputGM_.GetPhyAddr()), (__gm__ volatile XT*)(xGM_),163+ asc_vf_call<SimtUniformRandomBinarySearch<XT, NormProbsT, true, uint64_t, CORE_THREAD_NUM_U64>>(
148- elementNum, realSeed, numsamples, numCat, nsMagic, nsShift, samplesAligned, baseOffsetLo, baseOffsetHi);164+ dim3(CORE_THREAD_NUM_U64), (__gm__ volatile int64_t*)(outputGM_.GetPhyAddr()), (__gm__ XT*)(xGM_),
165+ (__gm__ NormProbsT*)(normProbsGM_), elementNum, realSeed, numsamples, numCat, nsMagic, nsShift,
166+ samplesAligned, baseOffsetLo, baseOffsetHi);
167+ } else {
168+ asc_vf_call<SimtUniformRandomBinarySearch<XT, NormProbsT, false, uint64_t, CORE_THREAD_NUM_U64>>(
169+ dim3(CORE_THREAD_NUM_U64), (__gm__ volatile int64_t*)(outputGM_.GetPhyAddr()), (__gm__ XT*)(xGM_),
170+ (__gm__ NormProbsT*)(normProbsGM_), elementNum, realSeed, numsamples, numCat, nsMagic, nsShift,
171+ samplesAligned, baseOffsetLo, baseOffsetHi);
172+ }
149 } else {173 } else {
150- asc_vf_call<SimtUniformRandomBinarySearch<XT, uint32_t, CORE_THREAD_NUM_U32>>(174+ if (normProbsGM_ != nullptr) {
151- dim3(CORE_THREAD_NUM_U32), (__gm__ volatile int64_t*)(outputGM_.GetPhyAddr()), (__gm__ volatile XT*)(xGM_),175+ asc_vf_call<SimtUniformRandomBinarySearch<XT, NormProbsT, true, uint32_t, CORE_THREAD_NUM_U32>>(
152- static_cast<uint32_t>(elementNum), realSeed, static_cast<uint32_t>(numsamples), numCat, nsMagic, nsShift,176+ dim3(CORE_THREAD_NUM_U32), (__gm__ volatile int64_t*)(outputGM_.GetPhyAddr()), (__gm__ XT*)(xGM_),
153- static_cast<uint32_t>(samplesAligned), baseOffsetLo, baseOffsetHi);177+ (__gm__ NormProbsT*)(normProbsGM_), static_cast<uint32_t>(elementNum), realSeed,
178+ static_cast<uint32_t>(numsamples), numCat, nsMagic, nsShift, static_cast<uint32_t>(samplesAligned),
179+ baseOffsetLo, baseOffsetHi);
180+ } else {
181+ asc_vf_call<SimtUniformRandomBinarySearch<XT, NormProbsT, false, uint32_t, CORE_THREAD_NUM_U32>>(
182+ dim3(CORE_THREAD_NUM_U32), (__gm__ volatile int64_t*)(outputGM_.GetPhyAddr()), (__gm__ XT*)(xGM_),
183+ (__gm__ NormProbsT*)(normProbsGM_), static_cast<uint32_t>(elementNum), realSeed,
184+ static_cast<uint32_t>(numsamples), numCat, nsMagic, nsShift, static_cast<uint32_t>(samplesAligned),
185+ baseOffsetLo, baseOffsetHi);
186+ }
154 }187 }
155}188}
156} // namespace StatelessSampleMultinomial189} // namespace StatelessSampleMultinomial
Mrandom/stateless_sample_multinomial/op_kernel/stateless_sample_multinomial.cpp+4-4
@@ -19,8 +19,8 @@ using namespace StatelessSampleMultinomial;
19 19 
20#define STATELESS_SAMPLE_MULTINOMIAL_DEFAULT_TILING_KEY 10020#define STATELESS_SAMPLE_MULTINOMIAL_DEFAULT_TILING_KEY 100
21 21 
22-__global__ __aicore__ void stateless_sample_multinomial(GM_ADDR x, GM_ADDR seed, GM_ADDR offset, GM_ADDR y,22+__global__ __aicore__ void stateless_sample_multinomial(GM_ADDR x, GM_ADDR normProbs, GM_ADDR seed, GM_ADDR offset,
23- GM_ADDR workspace, GM_ADDR tiling)23+ GM_ADDR y, GM_ADDR workspace, GM_ADDR tiling)
24{24{
25 REGISTER_TILING_DEFAULT(RandomUnifiedSimtTilingDataStruct);25 REGISTER_TILING_DEFAULT(RandomUnifiedSimtTilingDataStruct);
26 KERNEL_TASK_TYPE_DEFAULT(KERNEL_TYPE_MIX_AIV_1_0);26 KERNEL_TASK_TYPE_DEFAULT(KERNEL_TYPE_MIX_AIV_1_0);
@@ -28,8 +28,8 @@ __global__ __aicore__ void stateless_sample_multinomial(GM_ADDR x, GM_ADDR seed,
28 TPipe pipe;28 TPipe pipe;
29 29 
30 if (TILING_KEY_IS(STATELESS_SAMPLE_MULTINOMIAL_DEFAULT_TILING_KEY)) {30 if (TILING_KEY_IS(STATELESS_SAMPLE_MULTINOMIAL_DEFAULT_TILING_KEY)) {
31- StatelessSampleMultinomialOp<DTYPE_X> op;31+ StatelessSampleMultinomialOp<DTYPE_X, DTYPE_NORM_PROBS> op;
32- op.Init(y, x, seed, offset, workspace, &tilingData, &pipe);32+ op.Init(y, x, normProbs, seed, offset, workspace, &tilingData, &pipe);
33 op.Process();33 op.Process();
34 }34 }
35}35}
Arandom/stateless_sample_multinomial/tests/assets/golden.py+304-0
@@ -0,0 +1,304 @@
1+#!/usr/bin/env python3
2+# -*- coding: UTF-8 -*-
3+# -----------------------------------------------------------------------------
4+# Copyright (c) 2026 Huawei Technologies Co., Ltd.
5+# This program is free software, you can redistribute it and/or modify it under
6+# the terms and conditions of
7+# CANN Open Software License Agreement Version 2.0 (the "License").
8+# Please refer to the License for details. You may not use this file except in
9+# compliance with the License.
10+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
11+# EITHER EXPRESS OR IMPLIED,
12+# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR
13+# A PARTICULAR PURPOSE.
14+# See LICENSE in the root of the software repository for the full text of the
15+# License.
16+# -----------------------------------------------------------------------------
17+# NOTE: Batch tests must use --proc-no-reuse; reused workers may retain invalid
18+# ACL/torch_npu contexts between test cases.
19+ 
20+import ctypes
21+import os
22+ 
23+import numpy as np
24+import torch
25+ 
26+ 
27+__golden__ = {
28+ "aclnn": {
29+ "aclnnMultinomial": "multinomial_golden",
30+ "aclnnMultinomialTensor": "multinomial_tensor_golden",
31+ }
32+}
33+ 
34+PHILOX_W32_A = 0x9E3779B9
35+PHILOX_W32_B = 0xBB67AE85
36+PHILOX_M4X32_A = 0xD2511F53
37+PHILOX_M4X32_B = 0xCD9E8D57
38+UINT32_MASK = 0xFFFFFFFF
39+UINT64_MASK = 0xFFFFFFFFFFFFFFFF
40+ 
41+# Uniform conversion factor for the Philox u32 stream. The device kernels map a
42+# counter output u32 to a uniform float as u * RAND_2POW32_INV + RAND_2POW32_INV_HALF;
43+# golden must use the same rounding to stay bit-aligned with the operator.
44+RAND_2POW32_INV = 2.3283064e-10
45+RAND_2POW32_INV_HALF = RAND_2POW32_INV / 2.0
46+ 
47+SAMPLE_ALIGNMENT = 128
48+ 
49+_torch_npu = None
50+_acl_lib = None
51+ 
52+ 
53+def _worker_device_id():
54+ try:
55+ from ttk.core_modules.tbe_multiprocessing import get_process_context
56+ 
57+ process_context = get_process_context()
58+ if process_context is None:
59+ return None
60+ device = process_context.storage.get("device")
61+ device_id = getattr(device, "device_id", None)
62+ return None if device_id is None else int(device_id)
63+ except Exception:
64+ return None
65+ 
66+ 
67+def _acl_device_id():
68+ global _acl_lib
69+ if _acl_lib is None:
70+ try:
71+ _acl_lib = ctypes.CDLL("libascendcl.so")
72+ _acl_lib.aclrtGetDevice.restype = ctypes.c_int32
73+ _acl_lib.aclrtGetDevice.argtypes = [ctypes.POINTER(ctypes.c_int32)]
74+ except OSError:
75+ _acl_lib = False
76+ if not _acl_lib:
77+ return None
78+ 
79+ device_id = ctypes.c_int32(-1)
80+ if _acl_lib.aclrtGetDevice(ctypes.byref(device_id)) != 0:
81+ return None
82+ return device_id.value
83+ 
84+ 
85+def _load_torch_npu():
86+ global _torch_npu
87+ if _torch_npu is not None:
88+ return _torch_npu
89+ 
90+ import torch_npu
91+ 
92+ device_id = _worker_device_id()
93+ if device_id is None:
94+ device_id = _acl_device_id()
95+ override = os.getenv("MULTINOMIAL_GOLDEN_DEV")
96+ if override is not None:
97+ device_id = int(override)
98+ if device_id is not None and device_id >= 0:
99+ torch_npu.npu.set_device(device_id)
100+ _torch_npu = torch_npu
101+ return _torch_npu
102+ 
103+ 
104+def _philox_batch4_u32(seed, counter_lo, counter_hi):
105+ counter_lo = np.asarray(counter_lo, dtype=np.uint64)
106+ counter_hi = np.asarray(counter_hi, dtype=np.uint64)
107+ 
108+ c0 = (counter_lo & np.uint64(UINT32_MASK)).astype(np.uint32)
109+ c1 = (counter_lo >> np.uint64(32)).astype(np.uint32)
110+ c2 = (counter_hi & np.uint64(UINT32_MASK)).astype(np.uint32)
111+ c3 = (counter_hi >> np.uint64(32)).astype(np.uint32)
112+ key0 = np.uint32(seed & UINT32_MASK)
113+ key1 = np.uint32((seed >> 32) & UINT32_MASK)
114+ multiplier0 = np.uint64(PHILOX_M4X32_A)
115+ multiplier1 = np.uint64(PHILOX_M4X32_B)
116+ 
117+ for _ in range(10):
118+ product0 = multiplier0 * c0.astype(np.uint64)
119+ product1 = multiplier1 * c2.astype(np.uint64)
120+ lo0 = product0.astype(np.uint32)
121+ hi0 = (product0 >> np.uint64(32)).astype(np.uint32)
122+ lo1 = product1.astype(np.uint32)
123+ hi1 = (product1 >> np.uint64(32)).astype(np.uint32)
124+ c0, c1, c2, c3 = (
125+ hi1 ^ c1 ^ key0,
126+ lo1,
127+ hi0 ^ c3 ^ key1,
128+ lo0,
129+ )
130+ key0 = np.uint32((int(key0) + PHILOX_W32_A) & UINT32_MASK)
131+ key1 = np.uint32((int(key1) + PHILOX_W32_B) & UINT32_MASK)
132+ return c0, c1, c2, c3
133+ 
134+ 
135+def _uniform_u32(seed, offset, subsequence):
136+ subsequence = np.asarray(subsequence, dtype=np.uint64)
137+ counter_lo = np.full(subsequence.shape, offset, dtype=np.uint64)
138+ return _philox_batch4_u32(seed, counter_lo, subsequence)[0]
139+ 
140+ 
141+def _u32_to_float(value):
142+ inverse = np.float32(RAND_2POW32_INV)
143+ return value.astype(np.float32) * inverse + inverse / np.float32(2.0)
144+ 
145+ 
146+def _without_replacement(weights, seed, offset, numsamples):
147+ torch_npu = _load_torch_npu()
148+ weights_npu = weights.npu()
149+ generator = torch.Generator(device=weights_npu.device)
150+ generator.manual_seed(int(seed))
151+ # torch_npu Generator.set_offset requires offset to be a multiple of 4.
152+ # The device kernel aligns the offset via ceil(offset / 4), so round up to
153+ # the next multiple of 4 to start the exponential stream at the same Philox
154+ # counter the kernel will consume.
155+ aligned_offset = ((int(offset) + 3) // 4) * 4
156+ generator.set_offset(aligned_offset)
157+ exponential = torch.empty_like(weights_npu)
158+ torch_npu.npu_sim_exponential_(exponential, lambd=1.0, generator=generator)
159+ scores = torch.div(weights_npu, exponential)
160+ if numsamples == 1:
161+ result = torch.argmax(scores, dim=-1, keepdim=True)
162+ else:
163+ result = torch.topk(
164+ scores, numsamples, dim=-1, largest=True, sorted=True
165+ ).indices
166+ return result.to(torch.int64).cpu()
167+ 
168+ 
169+def _normalized_cdf(weights):
170+ _load_torch_npu()
171+ weights_npu = weights.npu()
172+ total = torch.sum(weights_npu, dim=-1, keepdim=True, dtype=weights_npu.dtype)
173+ probabilities = torch.div(weights_npu, total)
174+ cdf = torch.cumsum(probabilities, dim=-1, dtype=weights_npu.dtype)
175+ return cdf.cpu(), probabilities.cpu()
176+ 
177+ 
178+def _sample_from_cdf(cdf, probabilities, seed, offset, numsamples):
179+ squeeze_output = cdf.ndim == 1
180+ if squeeze_output:
181+ cdf = cdf.unsqueeze(0)
182+ probabilities = probabilities.unsqueeze(0)
183+ 
184+ dtype = cdf.dtype
185+ distribution_count, category_count = cdf.shape
186+ invalid_rows = torch.nonzero(~(cdf[:, -1] > 0), as_tuple=False).reshape(-1)
187+ if invalid_rows.numel():
188+ values = cdf[invalid_rows, -1].tolist()
189+ raise RuntimeError(
190+ "StatelessSampleMultinomial assertion failed: "
191+ f"cdf[:, -1] must be positive; rows={invalid_rows.tolist()}, "
192+ f"values={values}"
193+ )
194+ 
195+ aligned_samples = (
196+ (numsamples + SAMPLE_ALIGNMENT - 1) // SAMPLE_ALIGNMENT * SAMPLE_ALIGNMENT
197+ )
198+ offset_u64 = int(offset) & UINT64_MASK
199+ base_offset = ((offset_u64 + 3) & UINT64_MASK) // 4
200+ flat_index = np.arange(distribution_count * numsamples, dtype=np.uint64)
201+ distribution_index = flat_index // np.uint64(numsamples)
202+ sample_index = flat_index % np.uint64(numsamples)
203+ subsequence = distribution_index * np.uint64(aligned_samples) + sample_index
204+ random_u32 = _uniform_u32(int(seed) & UINT64_MASK, base_offset, subsequence)
205+ random_value = torch.from_numpy(_u32_to_float(random_u32)).to(dtype)
206+ random_value = random_value.reshape(distribution_count, numsamples)
207+ 
208+ start = torch.zeros((distribution_count, numsamples), dtype=torch.int64)
209+ end = torch.full_like(start, category_count)
210+ active = end > start
211+ while bool(active.any()):
212+ midpoint = start + ((end - start) >> 1)
213+ values = torch.gather(cdf, 1, midpoint.clamp_max(category_count - 1))
214+ move_right = (values < random_value) & active
215+ start = torch.where(move_right, midpoint + 1, start)
216+ end = torch.where(active & ~move_right, midpoint, end)
217+ active = end > start
218+ start.clamp_max_(category_count - 1)
219+ 
220+ category_index = torch.arange(category_count, dtype=torch.int64)
221+ category_index = category_index.unsqueeze(0).expand(
222+ distribution_count, category_count
223+ )
224+ nonzero_index = torch.where(
225+ probabilities != torch.zeros((), dtype=dtype),
226+ category_index,
227+ torch.full_like(category_index, -1),
228+ )
229+ previous_nonzero = torch.cummax(nonzero_index, dim=1).values
230+ result = torch.gather(previous_nonzero, 1, start)
231+ result = torch.where(result >= 0, result, torch.zeros_like(result))
232+ return result.reshape(numsamples) if squeeze_output else result
233+ 
234+ 
235+def _with_replacement(weights, seed, offset, numsamples):
236+ squeeze_output = weights.ndim == 1
237+ if squeeze_output:
238+ weights = weights.unsqueeze(0)
239+ cdf, probabilities = _normalized_cdf(weights)
240+ if squeeze_output:
241+ cdf = cdf.squeeze(0)
242+ probabilities = probabilities.squeeze(0)
243+ return _sample_from_cdf(cdf, probabilities, seed, offset, numsamples)
244+ 
245+ 
246+def _as_torch(value):
247+ if isinstance(value, torch.Tensor):
248+ return value
249+ if value.dtype.name == "bfloat16":
250+ return torch.frombuffer(
251+ bytearray(value.tobytes()), dtype=torch.bfloat16
252+ ).reshape(value.shape)
253+ return torch.from_numpy(np.ascontiguousarray(value))
254+ 
255+ 
256+def _compute(weights, numsamples, replacement, seed, offset):
257+ weights = _as_torch(weights).detach().cpu()
258+ numsamples = int(numsamples)
259+ if not bool(replacement) or numsamples == 1:
260+ return _without_replacement(weights, int(seed), int(offset), numsamples)
261+ return _with_replacement(weights, int(seed), int(offset), numsamples)
262+ 
263+ 
264+def multinomial_golden(self, numsamples, replacement, seed, offset, out=None, **kwargs):
265+ """
266+ Aclnn golden for aclnnMultinomial.
267+ Parameters follow @aclnnMultinomialGetWorkspaceSize without workspaceSize & executor.
268+ All the input Tensors are torch.Tensor.
269+ 
270+ replacement=False (or numsamples == 1): sample via the exponential path
271+ (weights / Exp(1.0)), then argmax (numsamples == 1) or top-k (numsamples > 1).
272+ replacement=True otherwise: normalized-CDF binary search against a Philox
273+ stream. The sample index deterministically depends on (seed, offset).
274+ """
275+ del out, kwargs
276+ return (_compute(self, numsamples, replacement, seed, offset).numpy(),)
277+ 
278+ 
279+def multinomial_tensor_golden(
280+ self,
281+ numsamples,
282+ replacement,
283+ seedTensor,
284+ offsetTensor,
285+ offset,
286+ out=None,
287+ **kwargs,
288+):
289+ """
290+ Aclnn golden for aclnnMultinomialTensor.
291+ Parameters follow @aclnnMultinomialTensorGetWorkspaceSize without workspaceSize & executor.
292+ All the input Tensors are torch.Tensor.
293+ 
294+ seedTensor / offsetTensor carry the generator seed and offset; `offset` is the
295+ stream-local intragraph offset added on top of offsetTensor. The combined offset
296+ wraps at 64 bits and is interpreted as signed for the exponential path.
297+ """
298+ del out, kwargs
299+ seed = int(_as_torch(seedTensor).reshape(-1)[0].item())
300+ tensor_offset = int(_as_torch(offsetTensor).reshape(-1)[0].item())
301+ combined_offset = (tensor_offset + int(offset)) & UINT64_MASK
302+ if combined_offset >= 1 << 63:
303+ combined_offset -= 1 << 64
304+ return (_compute(self, numsamples, replacement, seed, combined_offset).numpy(),)
Mrandom/stateless_sample_multinomial/tests/ut/op_host/arch35/test_stateless_sample_multinomial_tiling.cpp+58-64
@@ -21,15 +21,9 @@
21 21 
22class StatelessSampleMultinomialTilingTest : public testing::Test {22class StatelessSampleMultinomialTilingTest : public testing::Test {
23protected:23protected:
24- static void SetUpTestCase()24+ static void SetUpTestCase() { std::cout << "StatelessSampleMultinomialTilingTest SetUp" << std::endl; }
25- {
26- std::cout << "StatelessSampleMultinomialTilingTest SetUp" << std::endl;
27- }
28 25 
29- static void TearDownTestCase()26+ static void TearDownTestCase() { std::cout << "StatelessSampleMultinomialTilingTest TearDown" << std::endl; }
30- {
31- std::cout << "StatelessSampleMultinomialTilingTest TearDown" << std::endl;
32- }
33};27};
34 28 
35TEST_F(StatelessSampleMultinomialTilingTest, one_dim_float)29TEST_F(StatelessSampleMultinomialTilingTest, one_dim_float)
@@ -38,20 +32,20 @@ TEST_F(StatelessSampleMultinomialTilingTest, one_dim_float)
38 int64_t seedValue = 12345;32 int64_t seedValue = 12345;
39 int64_t offsetValue = 0;33 int64_t offsetValue = 0;
40 34 
41- gert::TilingContextPara tilingContextPara(35+ gert::TilingContextPara tilingContextPara("StatelessSampleMultinomial",
42- "StatelessSampleMultinomial",36+ {
43- {37+ {{{16}, {16}}, ge::DT_FLOAT, ge::FORMAT_ND},
44- {{{16}, {16}}, ge::DT_FLOAT, ge::FORMAT_ND},38+ {{{16}, {16}}, ge::DT_FLOAT, ge::FORMAT_ND},
45- {{{1}, {1}}, ge::DT_INT64, ge::FORMAT_ND, true, &seedValue},39+ {{{1}, {1}}, ge::DT_INT64, ge::FORMAT_ND, true, &seedValue},
46- {{{1}, {1}}, ge::DT_INT64, ge::FORMAT_ND, true, &offsetValue},40+ {{{1}, {1}}, ge::DT_INT64, ge::FORMAT_ND, true, &offsetValue},
47- },41+ },
48- {42+ {
49- {{{8}, {8}}, ge::DT_INT64, ge::FORMAT_ND},43+ {{{8}, {8}}, ge::DT_INT64, ge::FORMAT_ND},
50- },44+ },
51- {45+ {
52- {"num_samples", Ops::Math::AnyValue::CreateFrom<int64_t>(8)},46+ {"num_samples", Ops::Math::AnyValue::CreateFrom<int64_t>(8)},
53- },47+ },
54- &compileInfo);48+ &compileInfo);
55 49 
56 uint64_t expectTilingKey = 100;50 uint64_t expectTilingKey = 100;
57 std::vector<size_t> expectWorkspaces = {0};51 std::vector<size_t> expectWorkspaces = {0};
@@ -64,20 +58,20 @@ TEST_F(StatelessSampleMultinomialTilingTest, two_dim_float16)
64 int64_t seedValue = 7;58 int64_t seedValue = 7;
65 int64_t offsetValue = 4;59 int64_t offsetValue = 4;
66 60 
67- gert::TilingContextPara tilingContextPara(61+ gert::TilingContextPara tilingContextPara("StatelessSampleMultinomial",
68- "StatelessSampleMultinomial",62+ {
69- {63+ {{{3, 10}, {3, 10}}, ge::DT_FLOAT16, ge::FORMAT_ND},
70- {{{3, 10}, {3, 10}}, ge::DT_FLOAT16, ge::FORMAT_ND},64+ {{{3, 10}, {3, 10}}, ge::DT_FLOAT16, ge::FORMAT_ND},
71- {{{1}, {1}}, ge::DT_INT64, ge::FORMAT_ND, true, &seedValue},65+ {{{1}, {1}}, ge::DT_INT64, ge::FORMAT_ND, true, &seedValue},
72- {{{1}, {1}}, ge::DT_INT64, ge::FORMAT_ND, true, &offsetValue},66+ {{{1}, {1}}, ge::DT_INT64, ge::FORMAT_ND, true, &offsetValue},
73- },67+ },
74- {68+ {
75- {{{3, 5}, {3, 5}}, ge::DT_INT64, ge::FORMAT_ND},69+ {{{3, 5}, {3, 5}}, ge::DT_INT64, ge::FORMAT_ND},
76- },70+ },
77- {71+ {
78- {"num_samples", Ops::Math::AnyValue::CreateFrom<int64_t>(5)},72+ {"num_samples", Ops::Math::AnyValue::CreateFrom<int64_t>(5)},
79- },73+ },
80- &compileInfo);74+ &compileInfo);
81 75 
82 uint64_t expectTilingKey = 100;76 uint64_t expectTilingKey = 100;
83 std::vector<size_t> expectWorkspaces = {0};77 std::vector<size_t> expectWorkspaces = {0};
@@ -90,20 +84,20 @@ TEST_F(StatelessSampleMultinomialTilingTest, two_dim_bf16)
90 int64_t seedValue = 99;84 int64_t seedValue = 99;
91 int64_t offsetValue = 8;85 int64_t offsetValue = 8;
92 86 
93- gert::TilingContextPara tilingContextPara(87+ gert::TilingContextPara tilingContextPara("StatelessSampleMultinomial",
94- "StatelessSampleMultinomial",88+ {
95- {89+ {{{2, 64}, {2, 64}}, ge::DT_BF16, ge::FORMAT_ND},
96- {{{2, 64}, {2, 64}}, ge::DT_BF16, ge::FORMAT_ND},90+ {{{2, 64}, {2, 64}}, ge::DT_BF16, ge::FORMAT_ND},
97- {{{1}, {1}}, ge::DT_INT64, ge::FORMAT_ND, true, &seedValue},91+ {{{1}, {1}}, ge::DT_INT64, ge::FORMAT_ND, true, &seedValue},
98- {{{1}, {1}}, ge::DT_INT64, ge::FORMAT_ND, true, &offsetValue},92+ {{{1}, {1}}, ge::DT_INT64, ge::FORMAT_ND, true, &offsetValue},
99- },93+ },
100- {94+ {
101- {{{2, 16}, {2, 16}}, ge::DT_INT64, ge::FORMAT_ND},95+ {{{2, 16}, {2, 16}}, ge::DT_INT64, ge::FORMAT_ND},
102- },96+ },
103- {97+ {
104- {"num_samples", Ops::Math::AnyValue::CreateFrom<int64_t>(16)},98+ {"num_samples", Ops::Math::AnyValue::CreateFrom<int64_t>(16)},
105- },99+ },
106- &compileInfo);100+ &compileInfo);
107 101 
108 uint64_t expectTilingKey = 100;102 uint64_t expectTilingKey = 100;
109 std::vector<size_t> expectWorkspaces = {0};103 std::vector<size_t> expectWorkspaces = {0};
@@ -116,20 +110,20 @@ TEST_F(StatelessSampleMultinomialTilingTest, offset_multiple_of_four)
116 int64_t seedValue = 1;110 int64_t seedValue = 1;
117 int64_t offsetValue = 4;111 int64_t offsetValue = 4;
118 112 
119- gert::TilingContextPara tilingContextPara(113+ gert::TilingContextPara tilingContextPara("StatelessSampleMultinomial",
120- "StatelessSampleMultinomial",114+ {
121- {115+ {{{16}, {16}}, ge::DT_FLOAT, ge::FORMAT_ND},
122- {{{16}, {16}}, ge::DT_FLOAT, ge::FORMAT_ND},116+ {{{16}, {16}}, ge::DT_FLOAT, ge::FORMAT_ND},
123- {{{1}, {1}}, ge::DT_INT64, ge::FORMAT_ND, true, &seedValue},117+ {{{1}, {1}}, ge::DT_INT64, ge::FORMAT_ND, true, &seedValue},
124- {{{1}, {1}}, ge::DT_INT64, ge::FORMAT_ND, true, &offsetValue},118+ {{{1}, {1}}, ge::DT_INT64, ge::FORMAT_ND, true, &offsetValue},
125- },119+ },
126- {120+ {
127- {{{4}, {4}}, ge::DT_INT64, ge::FORMAT_ND},121+ {{{4}, {4}}, ge::DT_INT64, ge::FORMAT_ND},
128- },122+ },
129- {123+ {
130- {"num_samples", Ops::Math::AnyValue::CreateFrom<int64_t>(4)},124+ {"num_samples", Ops::Math::AnyValue::CreateFrom<int64_t>(4)},
131- },125+ },
132- &compileInfo);126+ &compileInfo);
133 127 
134 uint64_t expectTilingKey = 100;128 uint64_t expectTilingKey = 100;
135 std::vector<size_t> expectWorkspaces = {0};129 std::vector<size_t> expectWorkspaces = {0};
Mrandom/stateless_sample_multinomial/tests/ut/op_kernel/CMakeLists.txt+1-1
@@ -21,6 +21,6 @@ if(UT_TEST_ALL OR OP_KERNEL_UT)
21 AddOpTestCase(21 AddOpTestCase(
22 stateless_sample_multinomial22 stateless_sample_multinomial
23 "ascend950"23 "ascend950"
24- "-DDTYPE_X=float -DTestUtDefaultTilingStruct=RandomUnifiedSimtTilingDataStruct -I${KERNEL_STAGING_DIR}/stateless_sample_multinomial/arch35"24+ "-DDTYPE_X=float -DDTYPE_NORM_PROBS=float -DTestUtDefaultTilingStruct=RandomUnifiedSimtTilingDataStruct -I${KERNEL_STAGING_DIR}/stateless_sample_multinomial/arch35"
25 "${stateless_sample_multinomial_tiling_files}")25 "${stateless_sample_multinomial_tiling_files}")
26endif()26endif()
Mrandom/stateless_sample_multinomial/tests/ut/op_kernel/test_stateless_sample_multinomial.cpp+37-23
@@ -14,23 +14,22 @@
14#include "tikicpulib.h"14#include "tikicpulib.h"
15#include "../../../../random_common/op_kernel/arch35/random_unified_tiling_data_arch35.h"15#include "../../../../random_common/op_kernel/arch35/random_unified_tiling_data_arch35.h"
16 16 
17-extern __global__ __aicore__ void stateless_sample_multinomial(17+extern __global__ __aicore__ void stateless_sample_multinomial(GM_ADDR x, GM_ADDR normProbs, GM_ADDR seed,
18- GM_ADDR x, GM_ADDR seed, GM_ADDR offset, GM_ADDR y, GM_ADDR workspace, GM_ADDR tiling);18+ GM_ADDR offset, GM_ADDR y, GM_ADDR workspace,
19+ GM_ADDR tiling);
19 20 
20namespace {21namespace {
21using XType = DTYPE_X;22using XType = DTYPE_X;
23+using NormProbsType = DTYPE_NORM_PROBS;
22constexpr uint32_t kNumBlocks = 1;24constexpr uint32_t kNumBlocks = 1;
23constexpr uint64_t kTilingKey = 100;25constexpr uint64_t kTilingKey = 100;
24constexpr int64_t kNumDist = 2;26constexpr int64_t kNumDist = 2;
25constexpr int64_t kNumCat = 4;27constexpr int64_t kNumCat = 4;
26constexpr int64_t kNumSamples = 8;28constexpr int64_t kNumSamples = 8;
27constexpr int64_t kElementCount = kNumDist * kNumSamples;29constexpr int64_t kElementCount = kNumDist * kNumSamples;
28-constexpr int64_t kCdfElementCount = kNumDist * kNumCat;30+constexpr int64_t kInputElementCount = kNumDist * kNumCat;
29 31 
30-inline size_t Align32(size_t size)32+inline size_t Align32(size_t size) { return (size + 31U) / 32U * 32U; }
31-{
32- return (size + 31U) / 32U * 32U;
33-}
34 33 
35void FillTiling(RandomUnifiedSimtTilingDataStruct* tilingData, int64_t seed, int64_t offset)34void FillTiling(RandomUnifiedSimtTilingDataStruct* tilingData, int64_t seed, int64_t offset)
36{35{
@@ -45,32 +44,43 @@ void FillTiling(RandomUnifiedSimtTilingDataStruct* tilingData, int64_t seed, int
45 tilingData->splitBlockCount = 0;44 tilingData->splitBlockCount = 0;
46}45}
47 46 
48-void FillCdf(uint8_t* x)47+void FillX(uint8_t* xBuffer)
49{48{
50- auto* xData = reinterpret_cast<XType*>(x);49+ auto* xData = reinterpret_cast<XType*>(xBuffer);
51- const float cdf[kCdfElementCount] = {50+ const float x[kInputElementCount] = {
52- 0.10f, 0.30f, 0.60f, 1.00f,51+ 0.10f, 0.30f, 0.60f, 1.00f, 0.25f, 0.50f, 0.75f, 1.00f,
53- 0.25f, 0.50f, 0.75f, 1.00f,
54 };52 };
55- for (int64_t i = 0; i < kCdfElementCount; ++i) {53+ for (int64_t i = 0; i < kInputElementCount; ++i) {
56- xData[i] = static_cast<XType>(cdf[i]);54+ xData[i] = static_cast<XType>(x[i]);
55+ }
56+}
57+ 
58+void FillNormProbs(uint8_t* normProbsBuffer)
59+{
60+ auto* normProbsData = reinterpret_cast<NormProbsType*>(normProbsBuffer);
61+ const float normProbs[kInputElementCount] = {
62+ 0.10f, 0.20f, 0.30f, 0.40f, 0.25f, 0.25f, 0.25f, 0.25f,
63+ };
64+ for (int64_t i = 0; i < kInputElementCount; ++i) {
65+ normProbsData[i] = static_cast<NormProbsType>(normProbs[i]);
57 }66 }
58}67}
59} // namespace68} // namespace
60 69 
61-class StatelessSampleMultinomialKernelTest : public testing::Test {70+class StatelessSampleMultinomialKernelTest : public testing::Test {};
62-};
63 71 
64TEST_F(StatelessSampleMultinomialKernelTest, smoke_output_in_category_range)72TEST_F(StatelessSampleMultinomialKernelTest, smoke_output_in_category_range)
65{73{
66- auto* x = static_cast<uint8_t*>(AscendC::GmAlloc(Align32(kCdfElementCount * sizeof(XType))));74+ auto* x = static_cast<uint8_t*>(AscendC::GmAlloc(Align32(kInputElementCount * sizeof(XType))));
75+ auto* normProbs = static_cast<uint8_t*>(AscendC::GmAlloc(Align32(kInputElementCount * sizeof(NormProbsType))));
67 auto* seed = static_cast<uint8_t*>(AscendC::GmAlloc(Align32(sizeof(int64_t))));76 auto* seed = static_cast<uint8_t*>(AscendC::GmAlloc(Align32(sizeof(int64_t))));
68 auto* offset = static_cast<uint8_t*>(AscendC::GmAlloc(Align32(sizeof(int64_t))));77 auto* offset = static_cast<uint8_t*>(AscendC::GmAlloc(Align32(sizeof(int64_t))));
69 auto* y = static_cast<uint8_t*>(AscendC::GmAlloc(Align32(kElementCount * sizeof(int64_t))));78 auto* y = static_cast<uint8_t*>(AscendC::GmAlloc(Align32(kElementCount * sizeof(int64_t))));
70 auto* workspace = static_cast<uint8_t*>(AscendC::GmAlloc(Align32(16 * 1024 * 1024)));79 auto* workspace = static_cast<uint8_t*>(AscendC::GmAlloc(Align32(16 * 1024 * 1024)));
71 auto* tiling = static_cast<uint8_t*>(AscendC::GmAlloc(Align32(sizeof(RandomUnifiedSimtTilingDataStruct))));80 auto* tiling = static_cast<uint8_t*>(AscendC::GmAlloc(Align32(sizeof(RandomUnifiedSimtTilingDataStruct))));
72 81 
73- FillCdf(x);82+ FillX(x);
83+ FillNormProbs(normProbs);
74 std::memset(y, 0xFF, kElementCount * sizeof(int64_t));84 std::memset(y, 0xFF, kElementCount * sizeof(int64_t));
75 *reinterpret_cast<int64_t*>(seed) = 42;85 *reinterpret_cast<int64_t*>(seed) = 42;
76 *reinterpret_cast<int64_t*>(offset) = 0;86 *reinterpret_cast<int64_t*>(offset) = 0;
@@ -78,7 +88,7 @@ TEST_F(StatelessSampleMultinomialKernelTest, smoke_output_in_category_range)
78 88 
79 AscendC::SetKernelMode(KernelMode::AIV_MODE);89 AscendC::SetKernelMode(KernelMode::AIV_MODE);
80 ICPU_SET_TILING_KEY(kTilingKey);90 ICPU_SET_TILING_KEY(kTilingKey);
81- ICPU_RUN_KF(stateless_sample_multinomial, kNumBlocks, x, seed, offset, y, workspace, tiling);91+ ICPU_RUN_KF(stateless_sample_multinomial, kNumBlocks, x, normProbs, seed, offset, y, workspace, tiling);
82 92 
83 auto* yData = reinterpret_cast<int64_t*>(y);93 auto* yData = reinterpret_cast<int64_t*>(y);
84 for (int64_t i = 0; i < kElementCount; ++i) {94 for (int64_t i = 0; i < kElementCount; ++i) {
@@ -87,6 +97,7 @@ TEST_F(StatelessSampleMultinomialKernelTest, smoke_output_in_category_range)
87 }97 }
88 98 
89 AscendC::GmFree(x);99 AscendC::GmFree(x);
100+ AscendC::GmFree(normProbs);
90 AscendC::GmFree(seed);101 AscendC::GmFree(seed);
91 AscendC::GmFree(offset);102 AscendC::GmFree(offset);
92 AscendC::GmFree(y);103 AscendC::GmFree(y);
@@ -96,7 +107,8 @@ TEST_F(StatelessSampleMultinomialKernelTest, smoke_output_in_category_range)
96 107 
97TEST_F(StatelessSampleMultinomialKernelTest, determinism)108TEST_F(StatelessSampleMultinomialKernelTest, determinism)
98{109{
99- auto* x = static_cast<uint8_t*>(AscendC::GmAlloc(Align32(kCdfElementCount * sizeof(XType))));110+ auto* x = static_cast<uint8_t*>(AscendC::GmAlloc(Align32(kInputElementCount * sizeof(XType))));
111+ auto* normProbs = static_cast<uint8_t*>(AscendC::GmAlloc(Align32(kInputElementCount * sizeof(NormProbsType))));
100 auto* seed = static_cast<uint8_t*>(AscendC::GmAlloc(Align32(sizeof(int64_t))));112 auto* seed = static_cast<uint8_t*>(AscendC::GmAlloc(Align32(sizeof(int64_t))));
101 auto* offset = static_cast<uint8_t*>(AscendC::GmAlloc(Align32(sizeof(int64_t))));113 auto* offset = static_cast<uint8_t*>(AscendC::GmAlloc(Align32(sizeof(int64_t))));
102 auto* y1 = static_cast<uint8_t*>(AscendC::GmAlloc(Align32(kElementCount * sizeof(int64_t))));114 auto* y1 = static_cast<uint8_t*>(AscendC::GmAlloc(Align32(kElementCount * sizeof(int64_t))));
@@ -104,23 +116,25 @@ TEST_F(StatelessSampleMultinomialKernelTest, determinism)
104 auto* workspace = static_cast<uint8_t*>(AscendC::GmAlloc(Align32(16 * 1024 * 1024)));116 auto* workspace = static_cast<uint8_t*>(AscendC::GmAlloc(Align32(16 * 1024 * 1024)));
105 auto* tiling = static_cast<uint8_t*>(AscendC::GmAlloc(Align32(sizeof(RandomUnifiedSimtTilingDataStruct))));117 auto* tiling = static_cast<uint8_t*>(AscendC::GmAlloc(Align32(sizeof(RandomUnifiedSimtTilingDataStruct))));
106 118 
107- FillCdf(x);119+ FillX(x);
120+ FillNormProbs(normProbs);
108 *reinterpret_cast<int64_t*>(seed) = 12345;121 *reinterpret_cast<int64_t*>(seed) = 12345;
109 *reinterpret_cast<int64_t*>(offset) = 4;122 *reinterpret_cast<int64_t*>(offset) = 4;
110 FillTiling(reinterpret_cast<RandomUnifiedSimtTilingDataStruct*>(tiling), 12345, 4);123 FillTiling(reinterpret_cast<RandomUnifiedSimtTilingDataStruct*>(tiling), 12345, 4);
111 124 
112 AscendC::SetKernelMode(KernelMode::AIV_MODE);125 AscendC::SetKernelMode(KernelMode::AIV_MODE);
113 ICPU_SET_TILING_KEY(kTilingKey);126 ICPU_SET_TILING_KEY(kTilingKey);
114- ICPU_RUN_KF(stateless_sample_multinomial, kNumBlocks, x, seed, offset, y1, workspace, tiling);127+ ICPU_RUN_KF(stateless_sample_multinomial, kNumBlocks, x, normProbs, seed, offset, y1, workspace, tiling);
115 128 
116 std::memset(y2, 0, kElementCount * sizeof(int64_t));129 std::memset(y2, 0, kElementCount * sizeof(int64_t));
117 AscendC::SetKernelMode(KernelMode::AIV_MODE);130 AscendC::SetKernelMode(KernelMode::AIV_MODE);
118 ICPU_SET_TILING_KEY(kTilingKey);131 ICPU_SET_TILING_KEY(kTilingKey);
119- ICPU_RUN_KF(stateless_sample_multinomial, kNumBlocks, x, seed, offset, y2, workspace, tiling);132+ ICPU_RUN_KF(stateless_sample_multinomial, kNumBlocks, x, normProbs, seed, offset, y2, workspace, tiling);
120 133 
121 EXPECT_EQ(std::memcmp(y1, y2, kElementCount * sizeof(int64_t)), 0);134 EXPECT_EQ(std::memcmp(y1, y2, kElementCount * sizeof(int64_t)), 0);
122 135 
123 AscendC::GmFree(x);136 AscendC::GmFree(x);
137+ AscendC::GmFree(normProbs);
124 AscendC::GmFree(seed);138 AscendC::GmFree(seed);
125 AscendC::GmFree(offset);139 AscendC::GmFree(offset);
126 AscendC::GmFree(y1);140 AscendC::GmFree(y1);