mirror of https://github.com/Jittor/Jittor
59 lines
2.3 KiB
C++
59 lines
2.3 KiB
C++
#include <iostream>
|
||
#include <vector>
|
||
#include "aclnn.h"
|
||
|
||
int64_t GetShapeSize(const std::vector<int64_t>& shape) {
|
||
int64_t shapeSize = 1;
|
||
for (auto i : shape) {
|
||
shapeSize *= i;
|
||
}
|
||
return shapeSize;
|
||
}
|
||
|
||
void PrintOutResult(std::vector<int64_t> &shape, void** deviceAddr) {
|
||
auto size = GetShapeSize(shape);
|
||
std::vector<int> resultData(size, 0);
|
||
auto ret = aclrtMemcpy(resultData.data(), resultData.size() * sizeof(resultData[0]),
|
||
*deviceAddr, size * sizeof(resultData[0]), ACL_MEMCPY_DEVICE_TO_HOST);
|
||
CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("copy result from device to host failed. ERROR: %d\n", ret); return);
|
||
for (int64_t i = 0; i < size; i++) {
|
||
LOG_PRINT("mean result[%ld] is: %d\n", i, resultData[i]);
|
||
}
|
||
}
|
||
|
||
/*int Init(int32_t deviceId) {
|
||
// 固定写法，AscendCL初始化
|
||
auto ret = aclInit(nullptr);
|
||
CHECK_RET(ret == ACL_SUCCESS or ret == 100002, LOG_PRINT("aclInit failed. ERROR: %d\n", ret); return ret);
|
||
ret = aclrtSetDevice(deviceId);
|
||
CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtSetDevice failed. ERROR: %d\n", ret); return ret);
|
||
//ret = aclrtCreateStream(stream);
|
||
//CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtCreateStream failed. ERROR: %d\n", ret); return ret);
|
||
return 0;
|
||
}*/
|
||
|
||
/*
|
||
template <typename T>
|
||
int CreateAclTensor(const std::vector<T>& hostData, const std::vector<int64_t>& shape, void** deviceAddr,
|
||
aclDataType dataType, aclTensor** tensor) {
|
||
auto size = GetShapeSize(shape) * sizeof(T);
|
||
// 调用aclrtMalloc申请device侧内存
|
||
auto ret = aclrtMalloc(deviceAddr, size, ACL_MEM_MALLOC_HUGE_FIRST);
|
||
CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtMalloc failed. ERROR: %d\n", ret); return ret);
|
||
// 调用aclrtMemcpy将host侧数据拷贝到device侧内存上
|
||
ret = aclrtMemcpy(*deviceAddr, size, hostData.data(), size, ACL_MEMCPY_HOST_TO_DEVICE);
|
||
CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtMemcpy failed. ERROR: %d\n", ret); return ret);
|
||
|
||
// 计算连续tensor的strides
|
||
std::vector<int64_t> strides(shape.size(), 1);
|
||
for (int64_t i = shape.size() - 2; i >= 0; i--) {
|
||
strides[i] = shape[i + 1] * strides[i + 1];
|
||
}
|
||
|
||
// 调用aclCreateTensor接口创建aclTensor
|
||
*tensor = aclCreateTensor(shape.data(), shape.size(), dataType, strides.data(), 0, aclFormat::ACL_FORMAT_ND,
|
||
shape.data(), shape.size(), *deviceAddr);
|
||
return 0;
|
||
}*/
|
||
|