/**
 * @file add_custom.cpp
 *
 * Copyright (C) 2025. Huawei Technologies Co., Ltd. All rights reserved.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
 */
#include "kernel_operator.h"

class KernelAdd {
public:
    __aicore__ inline KernelAdd() {}
    __aicore__ inline void Init(GM_ADDR x, GM_ADDR z)
    {
        xGm.SetGlobalBuffer((__gm__ float *)x , 65536 / sizeof(float));
        zGm.SetGlobalBuffer((__gm__ float *)z , 48 * 65536 / sizeof(float));

    }
    __aicore__ inline void Process()
    {
            Compute();
    }

private:

    __aicore__ inline void Compute()
    {

        AscendC::LocalTensor<float> xLocalPing(AscendC::TPosition::VECIN, xPingAddr, 65536/sizeof(float));
        AscendC::LocalTensor<float> xLocalPong(AscendC::TPosition::VECIN, xPongAddr, 65536/sizeof(float));
        AscendC::LocalTensor<float> zLocal(AscendC::TPosition::VECOUT, zAddr, 65536/sizeof(float));

        AscendC::DataCopy(xLocalPing, xGm, 65536/sizeof(float));
        AscendC::SetFlag<AscendC::HardEvent::MTE2_V>(EVENT_ID0);
        AscendC::WaitFlag<AscendC::HardEvent::MTE2_V>(EVENT_ID0);

        AscendC::DataCopy(xLocalPong, xGm, 65536/sizeof(float));
        AscendC::SetFlag<AscendC::HardEvent::MTE2_V>(EVENT_ID1);
        AscendC::WaitFlag<AscendC::HardEvent::MTE2_V>(EVENT_ID1);

        uint64_t mask = 64;
        AscendC::Mul(zLocal[0], xLocalPing[0], xLocalPong[0], mask, 128, { 1, 1, 1, 8, 8, 8 });
        // AscendC::Mul(zLocal[0], xLocalPing[0], xLocalPong[64], mask, 128, { 1, 1, 1, 8, 8, 8 });

        AscendC::DataCopy(zGm[AscendC::GetBlockIdx()* 65536 / sizeof(float)], zLocal, 65536 / sizeof(float));

    }
    static constexpr uint32_t xPingAddr = 0;
    static constexpr uint32_t xPongAddr = 65536; // 64kB
    static constexpr uint32_t zAddr = 131072;    // 128kB

    AscendC::GlobalTensor<float> xGm;
    AscendC::GlobalTensor<float> zGm;
    uint32_t blockLength;
    uint32_t tileNum;
    uint32_t tileLength;
};


extern "C" __global__ __aicore__ void add_custom(GM_ADDR x, GM_ADDR z)
{
    KernelAdd op;
    op.Init(x, z);
    op.Process();
}

AscendC::Mul(zLocal[0], xLocalPing[0], xLocalPong[0], mask, 128, { 1, 1, 1, 8, 8, 8 });

AscendC::Mul(zLocal[0], xLocalPing[0], xLocalPong[64], mask, 128, { 1, 1, 1, 8, 8, 8 });

RUN_WITH_TOOLCHAIN=1 bash run.sh -r npu -v Ascend910B2

#!/bin/bash
CURRENT_DIR=$(
    cd $(dirname ${BASH_SOURCE:-$0})
    pwd
)
cd $CURRENT_DIR

BUILD_TYPE="Debug"
INSTALL_PREFIX="${CURRENT_DIR}/out"

SHORT=r:,v:,i:,b:,p:,
LONG=run-mode:,soc-version:,install-path:,build-type:,install-prefix:,
OPTS=$(getopt -a --options $SHORT --longoptions $LONG -- "$@")
eval set -- "$OPTS"
SOC_VERSION="Ascend310P3"

while :; do
    case "$1" in
    -r | --run-mode)
        RUN_MODE="$2"
        shift 2
        ;;
    -v | --soc-version)
        SOC_VERSION="$2"
        shift 2
        ;;
    -i | --install-path)
        ASCEND_INSTALL_PATH="$2"
        shift 2
        ;;
    -b | --build-type)
        BUILD_TYPE="$2"
        shift 2
        ;;
    -p | --install-prefix)
        INSTALL_PREFIX="$2"
        shift 2
        ;;
    --)
        shift
        break
        ;;
    *)
        echo "[ERROR]: Unexpected option: $1"
        break
        ;;
    esac
done

RUN_MODE_LIST="cpu sim npu"
if [[ " $RUN_MODE_LIST " != *" $RUN_MODE "* ]]; then
    echo "[ERROR]: RUN_MODE error, This sample only support specify cpu, sim or npu!"
    exit -1
fi

VERSION_LIST="Ascend910A Ascend910B Ascend310B1 Ascend310B2 Ascend310B3 Ascend310B4 Ascend310P1 Ascend310P3 Ascend910B1 Ascend910B2 Ascend910B3 Ascend910B4"
if [[ " $VERSION_LIST " != *" $SOC_VERSION "* ]]; then
    echo "[ERROR]: SOC_VERSION should be in [$VERSION_LIST]"
    exit -1
fi

if [ -n "$ASCEND_INSTALL_PATH" ]; then
    _ASCEND_INSTALL_PATH=$ASCEND_INSTALL_PATH
elif [ -n "$ASCEND_HOME_PATH" ]; then
    _ASCEND_INSTALL_PATH=$ASCEND_HOME_PATH
else
    if [ -d "$HOME/Ascend/ascend-toolkit/latest" ]; then
        _ASCEND_INSTALL_PATH=$HOME/Ascend/ascend-toolkit/latest
    else
        _ASCEND_INSTALL_PATH=/usr/local/Ascend/ascend-toolkit/latest
    fi
fi

export ASCEND_TOOLKIT_HOME=${_ASCEND_INSTALL_PATH}
export ASCEND_HOME_PATH=${_ASCEND_INSTALL_PATH}
echo "[INFO]: Current compile soc version is ${SOC_VERSION}"
source ${_ASCEND_INSTALL_PATH}/bin/setenv.bash
if [ "${RUN_MODE}" = "sim" ]; then
    # in case of running op in simulator, use stub .so instead
    export LD_LIBRARY_PATH=${_ASCEND_INSTALL_PATH}/tools/simulator/${SOC_VERSION}/lib:$LD_LIBRARY_PATH
elif [ "${RUN_MODE}" = "cpu" ]; then
    export LD_LIBRARY_PATH=${_ASCEND_INSTALL_PATH}/tools/tikicpulib/lib:${_ASCEND_INSTALL_PATH}/tools/tikicpulib/lib/${SOC_VERSION}:${_ASCEND_INSTALL_PATH}/tools/simulator/${SOC_VERSION}/lib:$LD_LIBRARY_PATH
fi

set -e
rm -rf build out
mkdir -p build
cmake -B build \
    -DRUN_MODE=${RUN_MODE} \
    -DSOC_VERSION=${SOC_VERSION} \
    -DCMAKE_BUILD_TYPE=${BUILD_TYPE} \
    -DCMAKE_INSTALL_PREFIX=${INSTALL_PREFIX} \
    -DASCEND_CANN_PACKAGE_PATH=${_ASCEND_INSTALL_PATH}
cmake --build build -j
cmake --install build

rm -f ascendc_kernels_bbit
cp ./out/bin/ascendc_kernels_bbit ./
rm -rf input output
mkdir -p input output
(
    export LD_LIBRARY_PATH=$(pwd)/out/lib:$(pwd)/out/lib64:${_ASCEND_INSTALL_PATH}/lib64:$LD_LIBRARY_PATH
    if [[ "$RUN_WITH_TOOLCHAIN" -eq 1 ]]; then
        if [ "${RUN_MODE}" = "npu" ]; then
            msprof op --aic-metrics=Roofline --application=./ascendc_kernels_bbit
        elif [ "${RUN_MODE}" = "sim" ]; then
            msprof op simulator --soc-version=Ascend910B2 --application=./ascendc_kernels_bbit
        elif [ "${RUN_MODE}" = "cpu" ]; then
            ./ascendc_kernels_bbit
        fi
    else
        ./ascendc_kernels_bbit
    fi
)
# tidy folder by delete log files
if [ "${RUN_MODE}" = "sim" ]; then
    rm -f *.log *.dump *.vcd *.toml *_log
fi
md5sum output/*.bin