Sign
功能说明
按元素执行Sign操作,Sign是指返回输入数据的符号,如果为0则返回0,如果为正数则返回1,如果为负数则返回-1。
函数原型
- 通过sharedTmpBuffer入参传入临时空间
- 源操作数Tensor全部/部分参与计算
1 2
template <typename T, bool isReuseSource = false> __aicore__ inline void Sign(const LocalTensor<T> &dstTensor, const LocalTensor<T> &srcTensor, const LocalTensor<uint8_t> &sharedTmpBuffer, const uint32_t calCount)
- 源操作数Tensor全部参与计算
1 2
template <typename T, bool isReuseSource = false> __aicore__ inline void Sign(const LocalTensor<T>& dstTensor, const LocalTensor<T>& srcTensor, const LocalTensor<uint8_t>& sharedTmpBuffer)
- 源操作数Tensor全部/部分参与计算
- 接口框架申请临时空间
- 源操作数Tensor全部/部分参与计算
1 2
template <typename T, bool isReuseSource = false> __aicore__ inline void Sign(const LocalTensor<T> &dstTensor, const LocalTensor<T> &srcTensor, const uint32_t calCount)
- 源操作数Tensor全部参与计算
1 2
template <typename T, bool isReuseSource = false> __aicore__ inline void Sign(const LocalTensor<T> &dstTensor, const LocalTensor<T> &srcTensor)
- 源操作数Tensor全部/部分参与计算
由于该接口的内部实现中涉及复杂的数学计算,需要额外的临时空间来存储计算过程中的中间变量。临时空间支持开发者通过sharedTmpBuffer入参传入和接口框架申请两种方式。
- 通过sharedTmpBuffer入参传入,使用该tensor作为临时空间进行处理,接口框架不再申请。该方式开发者可以自行管理sharedTmpBuffer内存空间,并在接口调用完成后,复用该部分内存,内存不会反复申请释放,灵活性较高,内存利用率也较高。
- 接口框架申请临时空间,开发者无需申请,但是需要预留临时空间的大小。
通过sharedTmpBuffer传入的情况,开发者需要为tensor申请空间;接口框架申请的方式,开发者需要预留临时空间。临时空间大小BufferSize的获取方式如下:通过GetSignMaxMinTmpSize中提供的接口获取需要预留空间范围的大小。
参数说明
参数名 |
描述 |
---|---|
T |
操作数的数据类型。 |
isReuseSource |
是否允许修改源操作数。该参数预留,传入默认值false即可。 |
参数名 |
输入/输出 |
描述 |
---|---|---|
dstTensor |
输出 |
目的操作数。 类型为LocalTensor,支持的TPosition为VECIN/VECCALC/VECOUT。 Atlas A2训练系列产品/Atlas 800I A2推理产品,支持的数据类型为:half/float Atlas推理系列产品(Ascend 310P处理器)AI Core,支持的数据类型为:half/float |
srcTensor |
输入 |
源操作数。 类型为LocalTensor,支持的TPosition为VECIN/VECCALC/VECOUT。 源操作数的数据类型需要与目的操作数保持一致。 Atlas A2训练系列产品/Atlas 800I A2推理产品,支持的数据类型为:half/float Atlas推理系列产品(Ascend 310P处理器)AI Core,支持的数据类型为:half/float |
sharedTmpBuffer |
输入 |
临时缓存。 类型为LocalTensor,支持的TPosition为VECIN/VECCALC/VECOUT。 用于Sign内部复杂计算时存储中间变量,由开发者提供。 临时空间大小BufferSize的获取方式请参考GetSignMaxMinTmpSize。 Atlas A2训练系列产品/Atlas 800I A2推理产品,支持的数据类型为:uint8_t Atlas推理系列产品(Ascend 310P处理器)AI Core,支持的数据类型为: uint8_t |
calCount |
输入 |
实际计算元素个数,calCount∈[0, srcTensor.GetSize()]。 |
返回值
无
支持的型号
Atlas A2训练系列产品/Atlas 800I A2推理产品
Atlas推理系列产品(Ascend 310P处理器)AI Core
约束说明
- 不支持源操作数与目的操作数地址重叠。
- 当前仅支持ND格式的输入,不支持其他格式。
- calCount需要保证小于或等于srcTensor和dstTensor存储的元素范围。
- 不支持sharedTmpBuffer与源操作数和目的操作数地址重叠。
- 操作数地址偏移对齐要求请参见通用约束。
调用示例
kernel侧sign_custom.cpp
#include "kernel_operator.h" constexpr int32_t BUFFER_NUM = 1; class KernelSign { public: __aicore__ inline KernelSign() {} __aicore__ inline void Init(GM_ADDR x, GM_ADDR y, uint32_t totalLength, uint32_t tilenum, uint32_t tmpSize, uint32_t mcount) { this->totalLength = totalLength; this->blockLength = totalLength / AscendC::GetBlockNum(); this->tilenum = tilenum; this->tmpSize = tmpSize; this->mcount = mcount; this->tileLength = this->blockLength / tilenum / BUFFER_NUM; xGm.SetGlobalBuffer((__gm__ half *)x + this->blockLength * AscendC::GetBlockIdx(), this->blockLength); yGm.SetGlobalBuffer((__gm__ half *)y + this->blockLength * AscendC::GetBlockIdx(), this->blockLength); if (this->tmpSize != 0) { pipe.InitBuffer(tmpQueue, BUFFER_NUM, this->tmpSize); } pipe.InitBuffer(inQueueX, BUFFER_NUM, this->tileLength * sizeof(half)); pipe.InitBuffer(outQueueY, BUFFER_NUM, this->tileLength * sizeof(half)); } __aicore__ inline void Process() { int32_t loopCount = this->tilenum * BUFFER_NUM; for (int32_t i = 0; i < loopCount; i++) { CopyIn(i); Compute(i); CopyOut(i); } } private: __aicore__ inline void CopyIn(int32_t progress) { AscendC::LocalTensor<half> xLocal = inQueueX.AllocTensor<half>(); AscendC::DataCopy(xLocal, xGm[progress * this->tileLength], this->tileLength); inQueueX.EnQue(xLocal); } __aicore__ inline void Compute(int32_t progress) { AscendC::LocalTensor<half> xLocal = inQueueX.DeQue<half>(); AscendC::LocalTensor<half> yLocal = outQueueY.AllocTensor<half>(); if (this->tmpSize != 0) { // 传入sharedTmpBuffer AscendC::LocalTensor<uint8_t> tmpLocal = tmpQueue.AllocTensor<uint8_t>(); if (this->mcount != this->totalLength) { // 是否传入calCount AscendC::Sign(yLocal, xLocal, tmpLocal, this->mcount); } else { AscendC::Sign(yLocal, xLocal, tmpLocal); } tmpQueue.FreeTensor(tmpLocal); } else { // 不传入sharedTmpBuffer if (this->mcount != this->totalLength) { AscendC::Sign(yLocal, xLocal, this->mcount); } else { AscendC::Sign(yLocal, xLocal); } } outQueueY.EnQue<half>(yLocal); inQueueX.FreeTensor(xLocal); } __aicore__ inline void CopyOut(int32_t progress) { AscendC::LocalTensor<half> yLocal = outQueueY.DeQue<half>(); AscendC::DataCopy(yGm[progress * this->tileLength], yLocal, this->tileLength); outQueueY.FreeTensor(yLocal); } private: AscendC::TPipe pipe; AscendC::TQue<AscendC::QuePosition::VECIN, BUFFER_NUM> inQueueX; AscendC::TQue<AscendC::QuePosition::VECIN, BUFFER_NUM> tmpQueue; AscendC::TQue<AscendC::QuePosition::VECOUT, BUFFER_NUM> outQueueY; AscendC::GlobalTensor<half> xGm; AscendC::GlobalTensor<half> yGm; uint32_t blockLength; uint32_t tilenum; uint32_t tileLength; uint32_t tmpSize; uint32_t mcount; uint32_t totalLength; }; extern "C" __global__ __aicore__ void sign_custom(GM_ADDR x, GM_ADDR y, GM_ADDR workspace, GM_ADDR tiling) { GET_TILING_DATA(tilingData, tiling); KernelSign op; op.Init(x, y, tilingData.totalLength, tilingData.tilenum, tilingData.tmpSize, tilingData.mcount); if (TILING_KEY_IS(1)) { op.Process(); } }
host侧sign_custom_tiling.h
#include "register/tilingdata_base.h" namespace optiling { BEGIN_TILING_DATA_DEF(SignCustomTilingData) TILING_DATA_FIELD_DEF(uint32_t, totalLength); TILING_DATA_FIELD_DEF(uint32_t, tmpSize); TILING_DATA_FIELD_DEF(uint32_t, tilenum); TILING_DATA_FIELD_DEF(uint32_t, mcount); END_TILING_DATA_DEF; REGISTER_TILING_DATA_CLASS(SignCustom, SignCustomTilingData) }
host侧sign_custom.cpp
#include "sign_custom_tiling.h" #include "register/op_def_registry.h" #include "tiling/tiling_api.h" namespace optiling { static ge::graphStatus TilingFunc(gert::TilingContext* context) { SignCustomTilingData tiling; const gert::RuntimeAttrs* cosattrs = context->GetAttrs(); const uint32_t tilenum = *(cosattrs->GetAttrPointer<uint32_t>(0)); const uint32_t blockdim = *(cosattrs->GetAttrPointer<uint32_t>(1)); const uint32_t sizeflag = *(cosattrs->GetAttrPointer<uint32_t>(2)); const uint32_t countflag = *(cosattrs->GetAttrPointer<uint32_t>(3)); uint32_t totalLength = context->GetInputTensor(0)->GetShapeSize(); auto dt = context->GetInputTensor(0)->GetDataType(); context->SetBlockDim(blockdim); tiling.set_totalLength(totalLength); tiling.set_tilenum(tilenum); if (countflag == 1) { tiling.set_mcount(totalLength); } std::vector<int64_t> shape_vec = {totalLength}; ge::Shape srcShape(shape_vec); uint32_t maxValue = 0; uint32_t minValue = 0; uint32_t dtypesize; if (dt == ge::DT_FLOAT16) { dtypesize = 2; } else { dtypesize = 4; } bool isReuseSource = false; AscendC::GetSignMaxMinTmpSize(srcShape, dtypesize, isReuseSource, maxValue, minValue); if (sizeflag == 0) { // sizeflag为0传入最小size的sharedTmpBuffer;为1传入最大size的sharedTmpBuffer;为2则相当于不传sharedTmpBuffer tiling.set_tmpSize(minValue); } else if (sizeflag == 1) { tiling.set_tmpSize(maxValue); } else if (sizeflag == 2) { tiling.set_tmpSize(0); } tiling.SaveToBuffer(context->GetRawTilingData()->GetData(), context->GetRawTilingData()->GetCapacity()); context->GetRawTilingData()->SetDataSize(tiling.GetDataSize()); context->SetTilingKey(1); size_t *currentWorkspace = context->GetWorkspaceSizes(1); currentWorkspace[0] = 0; return ge::GRAPH_SUCCESS; } } namespace ge { static ge::graphStatus InferShape(gert::InferShapeContext* context) { const gert::Shape* x1_shape = context->GetInputShape(0); gert::Shape* y_shape = context->GetOutputShape(0); *y_shape = *x1_shape; return GRAPH_SUCCESS; } } namespace ops { class SignCustom : public OpDef { public: explicit SignCustom(const char* name) : OpDef(name) { this->Input("x") .ParamType(REQUIRED) .DataType({ge::DT_FLOAT16}) .Format({ge::FORMAT_ND}); this->Output("y") .ParamType(REQUIRED) .DataType({ge::DT_FLOAT16}) .Format({ge::FORMAT_ND}); this->SetInferShape(ge::InferShape); this->Attr("tilenum") .AttrType(REQUIRED) .Int(0); this->Attr("blockdim") .AttrType(REQUIRED) .Int(0); this->Attr("sizeflag") .AttrType(REQUIRED) .Int(0); this->Attr("countflag") .AttrType(REQUIRED) .Int(0); this->AICore() .SetTiling(optiling::TilingFunc); this->AICore().AddConfig("ascendxxx"); // 这里的xxx根据不同的芯片填写不同信息 } };
输入输出的数据类型为float,一维向量包含8个数字; 输入数据(srcLocal): [-np.inf, -2.0, -0.0, 0.0, np.nan, -np.nan, 2.0, np.inf] 输出数据(dstLocal): [-1, -1, 0, 0, 0, 0, 1, 1]