昇腾AI原生创新算子挑战赛（S1赛季）复盘-ScatterSub算子
发表于2024-05-25 02:57:45
0 查看
- 正文前感谢昇腾各位工作人员，没有你们的辛勤就没有我们的进步
- 本文立意交流大赛ScatterSub算子编译过程
- 这道题难点在于
- 如何划分tiling，是以哪个输入变量shape作为划分的标准，从而获得相对应tileLength
- 如何设计copy_in来满足对var数据的sub操作，是否需要对var数据进行读入
- 参考了0xcccccccc大神的代码，只能用一个字来形容，妙呀，短小精悍
- 对原子操作SetAtomicAdd有独到见解
- 对copy_in数据流设计不需要引入var，并修改var，通过datacopy将updates数据施加原子加法操作，导入Indices指定位置
- 参考了blankcat代码
- 比较喜欢DataCopyPadCustom_GM2UB，DataCopyPadCustom_UB2GM 两个API构建
- gen_data.py的设计
- 开发板并不兼容tensorflow.自行编译步骤繁琐，而且真的太耗费编译时间。个人尝试编译1天1夜也无功而返
- 使用电脑进行同代码生成数据时，发现导出数据错误，input_var 与golden出现数据偏差
- 所以只能自行改写gen_data.py，摒弃tensorflow操作，使用pytorch进行数据的生成与导出
- 通过DataCopyPadCustom_GM2UB进行数据搬迁，注意单包情况
- 实际需要传输的是dimLength,其余都置0
点赞
收藏
回复
分享
举报
import torch import numpy as np def is_integer_dto(dtype): return dtype in [torch.int8, torch.int16, torch.int32, torch.int64] def is_float_dto(dtype): return dtype in [torch.float16, torch.float32, torch.float64] def gen_golden_data_simple2(): dtype = torch.float32 if(is_float_dto(dtype)): var = torch.randn(1024* 1024).reshape(1024, 1024).to(dtype)*20-10 indices = torch.randint(0, 1023, (1024,)).to(torch.int64) updates = torch.randn(1024* 1024).reshape(1024, 1024).to(dtype)*20-10 else: var = torch.randint(-10, 10, (3, 41, 52, 63)).to(dtype) indices = torch.randint(0, 3, (3,)).to(torch.int64) updates = torch.randint(-10, 10, (3, 41, 52, 63)).to(dtype) golden = var.clone() for i in range(indices.size(0)): index = indices[i] golden[index] = golden[index] - updates[i] var.to(dtype).detach().cpu().numpy().tofile("./input/input_var.bin") indices.to(torch.int32).detach().cpu().numpy().tofile("./input/input_indices.bin") updates.to(dtype).detach().cpu().numpy().tofile("./input/input_updates.bin") golden.to(dtype).detach().cpu().numpy().tofile("./output/golden.bin") if __name__ == "__main__": gen_golden_data_simple2()__aicore__ inline void CopyOut(int32_t indices,int32_t j,int32_t progress) { LocalTensor<DTYPE_VAR> varLocal = outQueueVar.DeQue<DTYPE_VAR>(); SetAtomicAdd<DTYPE_VAR>(); if (progress == this->tileNum - 1) { if (progress == 0) { //如果只有一包，则搬运的起始地址为0，tileLength为实际分块的数据量 DataCopy(varGm[indices * this->lastdim ], varLocal, this->tileLength); } else { //将最后一个分块的起始地址向前移动tileLength-lasttileLength DataCopy(varGm[indices * this->lastdim + (progress - 1) * this->tileLength + this->lasttileLength], varLocal, this->tileLength); } } else { DataCopy(varGm[indices * this->lastdim + progress * this->tileLength], varLocal, this->tileLength); } SetAtomicNone(); outQueueVar.FreeTensor(varLocal); }__aicore__ inline void CopyIn(int32_t indices, int32_t j,int32_t progress) { LocalTensor<DTYPE_UPDATES> updatesLocal = inQueueUpdates.AllocTensor<DTYPE_UPDATES>(); if (progress == this->tileNum - 1) { if (progress == 0) { DataCopyPadCustom_GM2UB(updatesLocal,updatesGm[j * this->lastdim],this->dimLength); for(int i = this->dimLength; i < this->tileLength; i++) { updatesLocal(i) = static_cast<DTYPE_UPDATES>(0); } } else { DataCopy(updatesLocal[0], updatesGm[j * this->lastdim + (progress - 1) * this->tileLength + this->lasttileLength], this->tileLength); } } else { DataCopy(updatesLocal[0], updatesGm[j * this->lastdim + progress * this->tileLength], this->tileLength); } inQueueUpdates.EnQue(updatesLocal); }所谓的this->dimLength = lastdim/GetBlockNum();
lastdim为所有非0轴的其他轴shape，也就是乘积，以此为tiling分割
最终贴出测试结果