task.hpp

Go to the documentation of this file.

/*
 *   Copyright 2025 Huawei Technologies Co., Ltd.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
 
#pragma once
 
#include <atomic>
#include <memory>
#include <queue>
#include <utility>
#include <vector>
#include <hicr/core/definitions.hpp>
#include <hicr/core/exceptions.hpp>
#include <hicr/core/executionState.hpp>
#include <hicr/core/executionUnit.hpp>
#include <hicr/core/processingUnit.hpp>
#include <hicr/core/computeManager.hpp>
#include "callbackMap.hpp"
#include "common.hpp"
 
#ifdef ENABLE_INSTRUMENTATION
  #include <tracr.hpp>
#endif
 
namespace HiCR::tasking
{
 
class Task
{
  public:
 
  enum callback_t
  {
    onTaskExecute,
 
    onTaskSuspend,
 
    onTaskFinish,
 
    onTaskSync,
  };
 
  using taskCallbackMap_t = HiCR::tasking::CallbackMap<Task *, callback_t>;
 
  Task()  = delete;
  ~Task() = default;
 
  __INLINE__ Task(std::shared_ptr<HiCR::ExecutionUnit> executionUnit, taskCallbackMap_t *callbackMap = nullptr)
    : _executionUnit(std::move(executionUnit)),
      _callbackMap(callbackMap) {};
 
  __INLINE__ void setCallbackMap(taskCallbackMap_t *callbackMap) { _callbackMap = callbackMap; }
 
  __INLINE__ taskCallbackMap_t *getCallbackMap() { return _callbackMap; }
 
  __INLINE__ void sendSyncSignal() { _callbackMap->trigger(this, HiCR::tasking::Task::callback_t::onTaskSync); }
 
  __INLINE__ const HiCR::ExecutionState::state_t getState()
  {
    // If the execution state has not been initialized then return the value expliclitly
    if (_executionState == nullptr) return HiCR::ExecutionState::state_t::uninitialized;
 
    // Otherwise just query the initial execution state
    return _executionState->getState();
  }
 
  __INLINE__ void setExecutionUnit(std::shared_ptr<HiCR::ExecutionUnit> executionUnit) { _executionUnit = std::move(executionUnit); }
 
  [[nodiscard]] __INLINE__ std::shared_ptr<HiCR::ExecutionUnit> getExecutionUnit() const { return _executionUnit; }
 
  __INLINE__ void initialize(std::unique_ptr<HiCR::ExecutionState> executionState)
  {
    if (getState() != HiCR::ExecutionState::state_t::uninitialized)
      HICR_THROW_LOGIC("Attempting to initialize a task that has already been initialized (State: %d).\n", getState());
 
    // Getting execution state as a unique pointer (to prcallback sharing the same state among different tasks)
    _executionState = std::move(executionState);
  }
 
  __INLINE__ void run()
  {
    if (getState() != HiCR::ExecutionState::state_t::initialized && getState() != HiCR::ExecutionState::state_t::suspended)
      HICR_THROW_RUNTIME("Attempting to run a task that is not in a initialized or suspended state (State: %d).\n", getState());
 
    // Triggering execution callback, if defined
    if (_callbackMap != nullptr) _callbackMap->trigger(this, callback_t::onTaskExecute);
 
// TraCR set trace of thread executing a task
#ifdef ENABLE_INSTRUMENTATION
    INSTRUMENTATION_THREAD_MARK_SET((long)0);
#endif
 
    // Now resuming the task's execution
    _executionState->resume();
 
// TraCR set trace of thread polling again
#ifdef ENABLE_INSTRUMENTATION
    INSTRUMENTATION_THREAD_MARK_SET((long)2);
#endif
 
    // Checking execution state finalization
    bool isFinished = _executionState->checkFinalization();
 
    // Getting state after execution
    const auto state = getState();
 
    if (state != HiCR::ExecutionState::state_t::suspended && state != HiCR::ExecutionState::state_t::finished)
      HICR_THROW_RUNTIME("Task has to be either in suspended or in finished state but I got State: %d. IsFinished: %b\n", state, isFinished);
 
    // If the task is suspended and callback map is defined, trigger the corresponding callback.
    if (state == HiCR::ExecutionState::state_t::suspended)
      if (_callbackMap != nullptr) _callbackMap->trigger(this, callback_t::onTaskSuspend);
 
    // If the task is still running (no suspension), then the task has fully finished executing. If so,
    // trigger the corresponding callback, if the callback map is defined. It is important that this function
    // is called from outside the context of a task to allow the upper layer to free its memory upon finishing
    if (state == HiCR::ExecutionState::state_t::finished)
      if (_callbackMap != nullptr) _callbackMap->trigger(this, callback_t::onTaskFinish);
  }
 
  __INLINE__ void suspend()
  {
    if (getState() != HiCR::ExecutionState::state_t::running) HICR_THROW_RUNTIME("Attempting to yield a task that is not in a running state (State: %d).\n", getState());
 
    // Yielding execution back to worker
    _executionState->suspend();
  }
 
  private:
 
  std::shared_ptr<HiCR::ExecutionUnit> _executionUnit;
 
  taskCallbackMap_t *_callbackMap = nullptr;
 
  std::unique_ptr<HiCR::ExecutionState> _executionState = nullptr;
 
}; // class Task
 
} // namespace HiCR::tasking