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_lock_table.h"
      12              : 
      13              : #include "aicpusd_status.h"
      14              : #include "aicpusd_model_execute.h"
      15              : #include "aicpusd_resource_manager.h"
      16              : 
      17              : namespace AicpuSchedule {
      18              : namespace {
      19              : const std::string KERNEL_LOCK_TABLE = "lockTable";
      20              : } // namespace
      21              : 
      22            7 : int32_t OperatorKernelLockTable::Compute(const AicpuTaskInfo& kernelTaskInfo, const RunContext& taskContext)
      23              : {
      24            7 :     aicpusd_info(
      25              :         "Start ModelLockTable. modelId=%u, streamId=%u, taskId=%u.", taskContext.modelId, kernelTaskInfo.streamID,
      26              :         kernelTaskInfo.taskID);
      27            7 :     if (kernelTaskInfo.paraBase == 0UL) {
      28            1 :         aicpusd_err("kernelTaskInfo.paraBase is null");
      29            1 :         return AICPU_SCHEDULE_ERROR_PARAMETER_NOT_VALID;
      30              :     }
      31              : 
      32            6 :     const auto model = AicpuModelManager::GetInstance().GetModel(taskContext.modelId);
      33            6 :     if (model == nullptr) {
      34            1 :         aicpusd_err(
      35              :             "Cannot get model by modelId:[%u], streamId[%u], taskId[%u].", taskContext.modelId, taskContext.streamId,
      36              :             kernelTaskInfo.taskID);
      37            1 :         return AICPU_SCHEDULE_ERROR_INNER_ERROR;
      38              :     }
      39              : 
      40            5 :     const LockTableTaskParam* const lockParam = PtrToPtr<void, LockTableTaskParam>(ValueToPtr(kernelTaskInfo.paraBase));
      41            5 :     const int32_t lockType = lockParam->lockType;
      42            5 :     const uint32_t tableId = lockParam->tableId;
      43              : 
      44            5 :     const auto triedTable = model->GetTableTryLock();
      45            5 :     if ((triedTable != INVALID_TABLE_ID) && (triedTable != static_cast<int64_t>(tableId))) {
      46            1 :         aicpusd_err(
      47              :             "model[%u] was tring to lock table[%d], cannot try to lock table[%u]", taskContext.modelId, triedTable,
      48              :             tableId);
      49            1 :         return AICPU_SCHEDULE_ERROR_PARAMETER_NOT_VALID;
      50              :     }
      51            4 :     model->SetTableTryLock(static_cast<int64_t>(tableId));
      52              : 
      53            4 :     EventWaitManager::TableUnlockWaitManager().ResetEventState(static_cast<size_t>(taskContext.modelId));
      54              :     do {
      55            4 :         bool lockRet = false;
      56            4 :         if (lockType == 0) {
      57            2 :             lockRet = TableLockManager::GetInstance().RdLockTable(tableId);
      58            2 :         } else if (lockType == 1) {
      59            1 :             lockRet = TableLockManager::GetInstance().WrLockTable(tableId);
      60              :         } else {
      61            1 :             aicpusd_err("Invalid lockType[%d].", lockType);
      62            1 :             return AICPU_SCHEDULE_ERROR_PARAMETER_NOT_VALID;
      63              :         }
      64              : 
      65            3 :         if (lockRet) {
      66            2 :             model->RecordLockedTable(tableId);
      67            2 :             aicpusd_info("model[%u] lock table[%u], type[%d] success.", taskContext.modelId, tableId, lockType);
      68            2 :             model->SetTableTryLock(INVALID_TABLE_ID);
      69            2 :             break;
      70              :         }
      71              : 
      72            1 :         bool needWait = false;
      73            1 :         EventWaitManager::TableUnlockWaitManager().WaitEvent(
      74            1 :             static_cast<size_t>(taskContext.modelId), taskContext.streamId, needWait);
      75            1 :         if (needWait) {
      76              :             // pending
      77            1 :             bool* const pending = const_cast<bool*>(&taskContext.pending);
      78            1 :             *pending = true;
      79            1 :             break;
      80              :         }
      81            0 :     } while (true);
      82              : 
      83            3 :     return AICPU_SCHEDULE_OK;
      84              : }
      85              : 
      86            6 : REGISTER_OPERATOR_KERNEL(KERNEL_LOCK_TABLE, OperatorKernelLockTable);
      87              : } // namespace AicpuSchedule