Sign

#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 is passed.
            AscendC::LocalTensor<uint8_t> tmpLocal = tmpQueue.AllocTensor<uint8_t>();
            if (this->mcount != this->totalLength)
            { // Whether to pass calCount.
                AscendC::Sign(yLocal, xLocal, tmpLocal, this->mcount);
            }
            else
            {
                AscendC::Sign(yLocal, xLocal, tmpLocal);
            }
            tmpQueue.FreeTensor(tmpLocal);
        }
        else
        { // sharedTmpBuffer is not passed.
            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();
    }
}

#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)
}

#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) {  // If sizeflag is set to 0, sharedTmpBuffer of the minimum size is passed. If it set to 1, sharedTmpBuffer of the maximum size is passed. If it set to 2, sharedTmpBuffer is not passed.
      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");  // Replace xxx with the actual processor information.
    }
};
}