获取用户使用的workspace指针。
__aicore__ inline GM_ADDR GetUserWorkspace(GM_ADDR workspace)
参数名 |
输入/输出 |
描述 |
---|---|---|
workspace |
输入 |
传入workspace的指针,包括系统workspace和用户使用的workspace。 |
Atlas 训练系列产品
Atlas推理系列产品AI Core
Atlas A2训练系列产品
无
用户使用workspace指针。
// 用户自定义的tiling函数 static ge::graphStatus TilingFunc(gert::TilingContext* context) { AddApiTiling tiling; ... size_t usrSize = 256; // 设置用户需要使用的workspace大小。 // 如需要使用workspace需要默认设置系统workspace的大小为16MB。 size_t sysWorkspaceSize = 16 * 1024 * 1024; size_t *currentWorkspace = context->GetWorkspaceSizes(1); // 通过框架获取workspace的指针,GetWorkspaceSizes入参为所需workspace的块数。当前限制使用一块。 currentWorkspace[0] = usrSize + sysWorkspaceSize; // 设置总的workspace的数值大小,总的workspace空间由框架来申请并管理。 ... }
// 用户写的Kernel函数,核函数必须包括GM_ADDR workspace入参,位置需要放在tiling之前 extern "C" __global__ __aicore__ void add_custom(GM_ADDR x, GM_ADDR y, GM_ADDR z, GM_ADDR workspace, GM_ADDR tiling) { ... GM_ADDR usrWorkspace = AscendC::GetUserWorkspace(workspace); // 获取用户workspace指针。 ... }