Line data    Source code

       1              : /**
       2              :  * Copyright (c) 2025 Huawei Technologies Co., Ltd.
       3              :  * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
       4              :  * CANN Open Software License Agreement Version 2.0 (the "License").
       5              :  * Please refer to the License for details. You may not use this file except in compliance with the License.
       6              :  * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
       7              :  * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
       8              :  * See LICENSE in the root of the software repository for the full text of the License.
       9              :  */
      10              : 
      11              : #include "operator_kernel_mark_step.h"
      12              : 
      13              : #include "aicpusd_status.h"
      14              : #include "aicpusd_util.h"
      15              : #include "aicpusd_model_execute.h"
      16              : 
      17              : namespace AicpuSchedule {
      18              : namespace {
      19              : const std::string KERNEL_MARK_STEP = "markStep";
      20              : } // namespace
      21              : 
      22            6 : int32_t OperatorKernelMarkStep::Compute(const AicpuTaskInfo& kernelTaskInfo, const RunContext& taskContext)
      23              : {
      24            6 :     aicpusd_info("Begin to execute MarkStep modelId[%u].", taskContext.modelId);
      25              :     MarkStepInfo* const bufInfo =
      26            6 :         PtrToPtr<void, MarkStepInfo>(ValueToPtr(static_cast<uintptr_t>(kernelTaskInfo.paraBase)));
      27            6 :     if (CheckMarkStepPara(bufInfo) != AICPU_SCHEDULE_OK) {
      28            1 :         aicpusd_err(
      29              :             "MarkStep para check failed, modelId[%u], streamId[%u], taskId[%u]", taskContext.modelId,
      30              :             taskContext.streamId, kernelTaskInfo.taskID);
      31            1 :         return AICPU_SCHEDULE_ERROR_PARAMETER_NOT_VALID;
      32              :     }
      33              : 
      34              :     // bufInfo parse
      35            5 :     AicpuModel* const model = AicpuModelManager::GetInstance().GetModel(taskContext.modelId);
      36            5 :     if (model == nullptr) {
      37            1 :         aicpusd_err("cannot get model by modelId:[%u]!", taskContext.modelId);
      38            1 :         return AICPU_SCHEDULE_ERROR_PARAMETER_NOT_VALID;
      39              :     }
      40              : 
      41            4 :     uint64_t* const stepIdAddr = PtrToPtr<void, uint64_t>(ValueToPtr(static_cast<uintptr_t>(bufInfo->stepIdAddr)));
      42            4 :     if (static_cast<uint32_t>(bufInfo->headFlag) == 1U) {
      43            1 :         aicpusd_debug("Not is head node, modelId=%u", taskContext.modelId);
      44            1 :         model->SetHeadNodeFlag(false);
      45            1 :         model->SetStepIdInfo(std::move(StepIdInfo(stepIdAddr, 0U)));
      46            1 :         return AICPU_SCHEDULE_OK;
      47              :     }
      48              : 
      49            3 :     const uint64_t iteratorCount = model->GetIteratorId();
      50            3 :     if (AicpuUtil::IsUint64MulOverflow(iteratorCount, bufInfo->groupTotalCount)) {
      51            1 :         aicpusd_err(
      52              :             "modelId[%u], iteratorCount:[%lu], groupIndex:%u, totalCnt:%u!", taskContext.modelId, iteratorCount,
      53              :             bufInfo->groupIndex, bufInfo->groupTotalCount);
      54            1 :         return AICPU_SCHEDULE_ERROR_OVERFLOW;
      55              :     }
      56            2 :     uint64_t tempCnt = iteratorCount * bufInfo->groupTotalCount;
      57            2 :     if ((std::numeric_limits<uint64_t>::max() - tempCnt) <= bufInfo->groupIndex) {
      58            1 :         aicpusd_err(
      59              :             "modelId[%u], iteratorCount:[%lu], groupIndex:%u, totalCnt:%u!", taskContext.modelId, iteratorCount,
      60              :             bufInfo->groupIndex, bufInfo->groupTotalCount);
      61            1 :         return AICPU_SCHEDULE_ERROR_OVERFLOW;
      62              :     }
      63            1 :     tempCnt += bufInfo->groupIndex;
      64            1 :     if ((std::numeric_limits<uint64_t>::max() - tempCnt) <= static_cast<uint64_t>(bufInfo->groupTotalCount - 1U)) {
      65            0 :         aicpusd_err(
      66              :             "modelId[%u], iteratorCount:[%lu], groupIndex:%u, totalCnt:%u!", taskContext.modelId, iteratorCount,
      67              :             bufInfo->groupIndex, bufInfo->groupTotalCount);
      68            0 :         return AICPU_SCHEDULE_ERROR_OVERFLOW;
      69              :     }
      70              : 
      71              :     // because transId starts from 1, the following operations is needed
      72            1 :     const uint64_t stepId =
      73            1 :         (bufInfo->groupIndex == 0U) ? tempCnt + static_cast<uint64_t>(bufInfo->groupTotalCount - 1U) : tempCnt - 1UL;
      74            1 :     aicpusd_info(
      75              :         "[MarkStep] headFlag[true], iteratorCount[%lu], totalCnt[%u], groupIndex[%u], stepId:[%lu].", iteratorCount,
      76              :         bufInfo->groupTotalCount, bufInfo->groupIndex, stepId);
      77            1 :     *stepIdAddr = stepId;
      78              : 
      79            1 :     model->SetHeadNodeFlag(true);
      80            1 :     model->SetStepIdInfo(std::move(StepIdInfo(stepIdAddr, stepId)));
      81            1 :     return AICPU_SCHEDULE_OK;
      82              : }
      83              : 
      84            6 : int32_t OperatorKernelMarkStep::CheckMarkStepPara(const MarkStepInfo* const markStepInfo) const
      85              : {
      86            6 :     if ((markStepInfo == nullptr) || (markStepInfo->stepIdAddr == 0UL) || (markStepInfo->groupTotalCount == 0U)) {
      87            1 :         return AICPU_SCHEDULE_ERROR_PARAMETER_NOT_VALID;
      88              :     }
      89              : 
      90            5 :     return AICPU_SCHEDULE_OK;
      91              : }
      92              : 
      93            6 : REGISTER_OPERATOR_KERNEL(KERNEL_MARK_STEP, OperatorKernelMarkStep);
      94              : } // namespace AicpuSchedule