task.c

Go to the documentation of this file.

/* -*- Mode: C; c-basic-offset:4 ; indent-tabs-mode:nil ; -*- */
/*
 * See COPYRIGHT in top-level directory.
 */

#include "abti.h"

static int ABTI_task_create(ABTI_local *p_local, ABTI_pool *p_pool,
                            void (*task_func)(void *), void *arg,
                            ABTI_sched *p_sched, int refcount,
                            ABTI_task **pp_newtask);
static int ABTI_task_revive(ABTI_local *p_local, ABTI_pool *p_pool,
                            void (*task_func)(void *), void *arg,
                            ABTI_task *p_task);
static inline uint64_t ABTI_task_get_new_id(void);

int ABT_task_create(ABT_pool pool, void (*task_func)(void *), void *arg,
                    ABT_task *newtask)
{
    int abt_errno = ABT_SUCCESS;
    ABTI_local *p_local = ABTI_local_get_local();
    ABTI_task *p_newtask;
    ABTI_pool *p_pool = ABTI_pool_get_ptr(pool);
    ABTI_CHECK_NULL_POOL_PTR(p_pool);

    int refcount = (newtask != NULL) ? 1 : 0;
    abt_errno = ABTI_task_create(p_local, p_pool, task_func, arg, NULL,
                                 refcount, &p_newtask);
    ABTI_CHECK_ERROR(abt_errno);

    /* Return value */
    if (newtask) {
        *newtask = ABTI_task_get_handle(p_newtask);
    }

fn_exit:
    return abt_errno;

fn_fail:
    if (newtask)
        *newtask = ABT_TASK_NULL;
    HANDLE_ERROR_FUNC_WITH_CODE(abt_errno);
    goto fn_exit;
}

/* This routine is to create a tasklet for the scheduler. */
int ABTI_task_create_sched(ABTI_local *p_local, ABTI_pool *p_pool,
                           ABTI_sched *p_sched)
{
    int abt_errno = ABT_SUCCESS;
    ABTI_task *p_newtask;

    void *arg = (void *)ABTI_sched_get_handle(p_sched);
    /* If p_sched is reused, ABTI_task_revive() can be used. */
    if (p_sched->p_task) {
        abt_errno =
            ABTI_task_revive(p_local, p_pool, (void (*)(void *))p_sched->run,
                             arg, p_sched->p_task);
        ABTI_CHECK_ERROR(abt_errno);
        goto fn_exit;
    }

    /* Allocate a task object */
    abt_errno =
        ABTI_task_create(p_local, p_pool, (void (*)(void *))p_sched->run, arg,
                         p_sched, 1, &p_newtask);
    ABTI_CHECK_ERROR(abt_errno);

fn_exit:
    return abt_errno;

fn_fail:
    HANDLE_ERROR_FUNC_WITH_CODE(abt_errno);
    goto fn_exit;
}

int ABT_task_create_on_xstream(ABT_xstream xstream, void (*task_func)(void *),
                               void *arg, ABT_task *newtask)
{
    int abt_errno = ABT_SUCCESS;
    ABTI_local *p_local = ABTI_local_get_local();
    ABTI_task *p_newtask;

    ABTI_xstream *p_xstream = ABTI_xstream_get_ptr(xstream);
    ABTI_CHECK_NULL_XSTREAM_PTR(p_xstream);

    /* TODO: need to consider the access type of target pool */
    ABTI_pool *p_pool = ABTI_xstream_get_main_pool(p_xstream);
    int refcount = (newtask != NULL) ? 1 : 0;
    abt_errno = ABTI_task_create(p_local, p_pool, task_func, arg, NULL,
                                 refcount, &p_newtask);
    ABTI_CHECK_ERROR(abt_errno);

    /* Return value */
    if (newtask)
        *newtask = ABTI_task_get_handle(p_newtask);

fn_exit:
    return abt_errno;

fn_fail:
    if (newtask)
        *newtask = ABT_TASK_NULL;
    HANDLE_ERROR_FUNC_WITH_CODE(abt_errno);
    goto fn_exit;
}

int ABT_task_revive(ABT_pool pool, void (*task_func)(void *), void *arg,
                    ABT_task *task)
{
    int abt_errno = ABT_SUCCESS;
    ABTI_local *p_local = ABTI_local_get_local();

    ABTI_task *p_task = ABTI_task_get_ptr(*task);
    ABTI_CHECK_NULL_TASK_PTR(p_task);

    ABTI_pool *p_pool = ABTI_pool_get_ptr(pool);
    ABTI_CHECK_NULL_POOL_PTR(p_pool);

    abt_errno = ABTI_task_revive(p_local, p_pool, task_func, arg, p_task);
    ABTI_CHECK_ERROR(abt_errno);

fn_exit:
    return abt_errno;

fn_fail:
    HANDLE_ERROR_FUNC_WITH_CODE(abt_errno);
    goto fn_exit;
}

int ABT_task_free(ABT_task *task)
{
    int abt_errno = ABT_SUCCESS;
    ABTI_local *p_local = ABTI_local_get_local();
    ABT_task h_task = *task;
    ABTI_task *p_task = ABTI_task_get_ptr(h_task);
    ABTI_CHECK_NULL_TASK_PTR(p_task);

    /* Wait until the task terminates */
    while (ABTD_atomic_acquire_load_int(&p_task->state) !=
           ABT_TASK_STATE_TERMINATED) {
#ifndef ABT_CONFIG_DISABLE_EXT_THREAD
        if (ABTI_self_get_type(p_local) != ABT_UNIT_TYPE_THREAD) {
            ABTD_atomic_pause();
            continue;
        }
#endif
        ABTI_thread_yield(&p_local, p_local->p_thread);
    }

    /* Free the ABTI_task structure */
    ABTI_task_free(p_local, p_task);

    /* Return value */
    *task = ABT_TASK_NULL;

fn_exit:
    return abt_errno;

fn_fail:
    HANDLE_ERROR_FUNC_WITH_CODE(abt_errno);
    goto fn_exit;
}

int ABT_task_join(ABT_task task)
{
    int abt_errno = ABT_SUCCESS;
    ABTI_local *p_local = ABTI_local_get_local();

    ABTI_task *p_task = ABTI_task_get_ptr(task);
    ABTI_CHECK_NULL_TASK_PTR(p_task);

    /* TODO: better implementation */
    while (ABTD_atomic_acquire_load_int(&p_task->state) !=
           ABT_TASK_STATE_TERMINATED) {
#ifndef ABT_CONFIG_DISABLE_EXT_THREAD
        if (ABTI_self_get_type(p_local) != ABT_UNIT_TYPE_THREAD) {
            ABTD_atomic_pause();
            continue;
        }
#endif
        ABTI_thread_yield(&p_local, p_local->p_thread);
    }

fn_exit:
    return abt_errno;

fn_fail:
    HANDLE_ERROR_FUNC_WITH_CODE(abt_errno);
    goto fn_exit;
}

int ABT_task_cancel(ABT_task task)
{
#ifdef ABT_CONFIG_DISABLE_TASK_CANCEL
    return ABT_ERR_FEATURE_NA;
#else
    int abt_errno = ABT_SUCCESS;
    ABTI_task *p_task = ABTI_task_get_ptr(task);
    ABTI_CHECK_NULL_TASK_PTR(p_task);

    /* Set the cancel request */
    ABTI_task_set_request(p_task, ABTI_TASK_REQ_CANCEL);

fn_exit:
    return abt_errno;

fn_fail:
    HANDLE_ERROR_FUNC_WITH_CODE(abt_errno);
    goto fn_exit;
#endif
}

int ABT_task_self(ABT_task *task)
{
    int abt_errno = ABT_SUCCESS;
    ABTI_local *p_local = ABTI_local_get_local();

#ifndef ABT_CONFIG_DISABLE_EXT_THREAD
    /* In case that Argobots has not been initialized or this routine is called
     * by an external thread, e.g., pthread, return an error code instead of
     * making the call fail. */
    if (gp_ABTI_global == NULL) {
        abt_errno = ABT_ERR_UNINITIALIZED;
        *task = ABT_TASK_NULL;
        return abt_errno;
    }
    if (p_local == NULL) {
        abt_errno = ABT_ERR_INV_XSTREAM;
        *task = ABT_TASK_NULL;
        return abt_errno;
    }
#endif

    ABTI_task *p_task = p_local->p_task;
    if (p_task != NULL) {
        *task = ABTI_task_get_handle(p_task);
    } else {
        abt_errno = ABT_ERR_INV_TASK;
        *task = ABT_TASK_NULL;
    }

    return abt_errno;
}

int ABT_task_self_id(uint64_t *id)
{
    int abt_errno = ABT_SUCCESS;
    ABTI_local *p_local = ABTI_local_get_local();

#ifndef ABT_CONFIG_DISABLE_EXT_THREAD
    /* In case that Argobots has not been initialized or this routine is called
     * by an external thread, e.g., pthread, return an error code instead of
     * making the call fail. */
    if (gp_ABTI_global == NULL) {
        abt_errno = ABT_ERR_UNINITIALIZED;
        return abt_errno;
    }
    if (p_local == NULL) {
        abt_errno = ABT_ERR_INV_XSTREAM;
        return abt_errno;
    }
#endif

    ABTI_task *p_task = p_local->p_task;
    if (p_task != NULL) {
        *id = ABTI_task_get_id(p_task);
    } else {
        abt_errno = ABT_ERR_INV_TASK;
    }

    return abt_errno;
}

int ABT_task_get_xstream(ABT_task task, ABT_xstream *xstream)
{
    int abt_errno = ABT_SUCCESS;
    ABTI_task *p_task = ABTI_task_get_ptr(task);
    ABTI_CHECK_NULL_TASK_PTR(p_task);

    /* Return value */
    *xstream = ABTI_xstream_get_handle(p_task->p_xstream);

fn_exit:
    return abt_errno;

fn_fail:
    HANDLE_ERROR_FUNC_WITH_CODE(abt_errno);
    goto fn_exit;
}

int ABT_task_get_state(ABT_task task, ABT_task_state *state)
{
    int abt_errno = ABT_SUCCESS;

    ABTI_task *p_task = ABTI_task_get_ptr(task);
    ABTI_CHECK_NULL_TASK_PTR(p_task);

    /* Return value */
    *state = (ABT_task_state)ABTD_atomic_acquire_load_int(&p_task->state);

fn_exit:
    return abt_errno;

fn_fail:
    HANDLE_ERROR_FUNC_WITH_CODE(abt_errno);
    goto fn_exit;
}

int ABT_task_get_last_pool(ABT_task task, ABT_pool *pool)
{
    int abt_errno = ABT_SUCCESS;

    ABTI_task *p_task = ABTI_task_get_ptr(task);
    ABTI_CHECK_NULL_TASK_PTR(p_task);

    /* Return value */
    *pool = ABTI_pool_get_handle(p_task->p_pool);

fn_exit:
    return abt_errno;

fn_fail:
    HANDLE_ERROR_FUNC_WITH_CODE(abt_errno);
    goto fn_exit;
}

int ABT_task_get_last_pool_id(ABT_task task, int *id)
{
    int abt_errno = ABT_SUCCESS;

    ABTI_task *p_task = ABTI_task_get_ptr(task);
    ABTI_CHECK_NULL_TASK_PTR(p_task);

    ABTI_ASSERT(p_task->p_pool);
    *id = (int)(p_task->p_pool->id);

fn_exit:
    return abt_errno;

fn_fail:
    HANDLE_ERROR_FUNC_WITH_CODE(abt_errno);
    goto fn_exit;
}

int ABT_task_set_migratable(ABT_task task, ABT_bool flag)
{
#ifndef ABT_CONFIG_DISABLE_MIGRATION
    int abt_errno = ABT_SUCCESS;
    ABTI_task *p_task = ABTI_task_get_ptr(task);
    ABTI_CHECK_NULL_TASK_PTR(p_task);

    p_task->migratable = flag;

fn_exit:
    return abt_errno;

fn_fail:
    HANDLE_ERROR_FUNC_WITH_CODE(abt_errno);
    goto fn_exit;
#else
    return ABT_ERR_MIGRATION_NA;
#endif
}

int ABT_task_is_migratable(ABT_task task, ABT_bool *flag)
{
#ifndef ABT_CONFIG_DISABLE_MIGRATION
    int abt_errno = ABT_SUCCESS;
    ABTI_task *p_task = ABTI_task_get_ptr(task);
    ABTI_CHECK_NULL_TASK_PTR(p_task);

    *flag = p_task->migratable;

fn_exit:
    return abt_errno;

fn_fail:
    HANDLE_ERROR_FUNC_WITH_CODE(abt_errno);
    goto fn_exit;
#else
    return ABT_ERR_MIGRATION_NA;
#endif
}

int ABT_task_equal(ABT_task task1, ABT_task task2, ABT_bool *result)
{
    ABTI_task *p_task1 = ABTI_task_get_ptr(task1);
    ABTI_task *p_task2 = ABTI_task_get_ptr(task2);
    *result = (p_task1 == p_task2) ? ABT_TRUE : ABT_FALSE;
    return ABT_SUCCESS;
}

int ABT_task_get_id(ABT_task task, uint64_t *task_id)
{
    int abt_errno = ABT_SUCCESS;

    ABTI_task *p_task = ABTI_task_get_ptr(task);
    ABTI_CHECK_NULL_TASK_PTR(p_task);

    *task_id = ABTI_task_get_id(p_task);

fn_exit:
    return abt_errno;

fn_fail:
    HANDLE_ERROR_FUNC_WITH_CODE(abt_errno);
    goto fn_exit;
}

int ABT_task_get_arg(ABT_task task, void **arg)
{
    int abt_errno = ABT_SUCCESS;

    ABTI_task *p_task = ABTI_task_get_ptr(task);
    ABTI_CHECK_NULL_TASK_PTR(p_task);

    *arg = p_task->p_arg;

fn_exit:
    return abt_errno;

fn_fail:
    HANDLE_ERROR_FUNC_WITH_CODE(abt_errno);
    goto fn_exit;
}

/*****************************************************************************/
/* Private APIs                                                              */
/*****************************************************************************/

static int ABTI_task_create(ABTI_local *p_local, ABTI_pool *p_pool,
                            void (*task_func)(void *), void *arg,
                            ABTI_sched *p_sched, int refcount,
                            ABTI_task **pp_newtask)
{
    int abt_errno = ABT_SUCCESS;
    ABTI_task *p_newtask;
    ABT_task h_newtask;
    ABTI_CHECK_NULL_POOL_PTR(p_pool);

    /* Allocate a task object */
    p_newtask = ABTI_mem_alloc_task(p_local);

    p_newtask->p_xstream = NULL;
    ABTD_atomic_relaxed_store_int(&p_newtask->state, ABT_TASK_STATE_READY);
    ABTD_atomic_relaxed_store_uint32(&p_newtask->request, 0);
    p_newtask->f_task = task_func;
    p_newtask->p_arg = arg;
#ifndef ABT_CONFIG_DISABLE_STACKABLE_SCHED
    p_newtask->is_sched = p_sched;
#endif
    p_newtask->p_pool = p_pool;
    p_newtask->refcount = refcount;
    p_newtask->p_keytable = NULL;
#ifndef ABT_CONFIG_DISABLE_MIGRATION
    p_newtask->migratable = ABT_TRUE;
#endif
    p_newtask->id = ABTI_TASK_INIT_ID;

    /* Create a wrapper work unit */
    h_newtask = ABTI_task_get_handle(p_newtask);
    p_newtask->unit = p_pool->u_create_from_task(h_newtask);

    LOG_EVENT("[T%" PRIu64 "] created\n", ABTI_task_get_id(p_newtask));

    /* Add this task to the scheduler's pool */
#ifdef ABT_CONFIG_DISABLE_POOL_PRODUCER_CHECK
    ABTI_pool_push(p_pool, p_newtask->unit);
#else
    abt_errno = ABTI_pool_push(p_pool, p_newtask->unit,
                               ABTI_self_get_native_thread_id(p_local));
    if (abt_errno != ABT_SUCCESS) {
        ABTI_task_free(p_local, p_newtask);
        goto fn_fail;
    }
#endif

    /* Return value */
    *pp_newtask = p_newtask;

fn_exit:
    return abt_errno;

fn_fail:
    HANDLE_ERROR_FUNC_WITH_CODE(abt_errno);
    goto fn_exit;
}

static int ABTI_task_revive(ABTI_local *p_local, ABTI_pool *p_pool,
                            void (*task_func)(void *), void *arg,
                            ABTI_task *p_task)
{
    int abt_errno = ABT_SUCCESS;

    ABTI_CHECK_TRUE(ABTD_atomic_relaxed_load_int(&p_task->state) ==
                        ABT_TASK_STATE_TERMINATED,
                    ABT_ERR_INV_TASK);

    p_task->p_xstream = NULL;
    ABTD_atomic_relaxed_store_int(&p_task->state, ABT_TASK_STATE_READY);
    ABTD_atomic_relaxed_store_uint32(&p_task->request, 0);
    p_task->f_task = task_func;
    p_task->p_arg = arg;
    p_task->refcount = 1;
    p_task->p_keytable = NULL;

    if (p_task->p_pool != p_pool) {
        /* Free the unit for the old pool */
        p_task->p_pool->u_free(&p_task->unit);

        /* Set the new pool */
        p_task->p_pool = p_pool;

        /* Create a wrapper work unit */
        ABT_task task = ABTI_task_get_handle(p_task);
        p_task->unit = p_pool->u_create_from_task(task);
    }

    LOG_EVENT("[T%" PRIu64 "] revived\n", ABTI_task_get_id(p_task));

    /* Add this task to the scheduler's pool */
#ifdef ABT_CONFIG_DISABLE_POOL_PRODUCER_CHECK
    ABTI_pool_push(p_pool, p_task->unit);
#else
    abt_errno = ABTI_pool_push(p_pool, p_task->unit,
                               ABTI_self_get_native_thread_id(p_local));
    ABTI_CHECK_ERROR(abt_errno);
#endif

fn_exit:
    return abt_errno;

fn_fail:
    HANDLE_ERROR_FUNC_WITH_CODE(abt_errno);
    goto fn_exit;
}

void ABTI_task_free(ABTI_local *p_local, ABTI_task *p_task)
{
    LOG_EVENT("[T%" PRIu64 "] freed\n", ABTI_task_get_id(p_task));

    /* Free the unit */
    p_task->p_pool->u_free(&p_task->unit);

    /* Free the key-value table */
    if (p_task->p_keytable) {
        ABTI_ktable_free(p_task->p_keytable);
    }

    ABTI_mem_free_task(p_local, p_task);
}

void ABTI_task_print(ABTI_task *p_task, FILE *p_os, int indent)
{
    char *prefix = ABTU_get_indent_str(indent);

    if (p_task == NULL) {
        fprintf(p_os, "%s== NULL TASKLET ==\n", prefix);
        goto fn_exit;
    }

    ABTI_xstream *p_xstream = p_task->p_xstream;
    int xstream_rank = p_xstream ? p_xstream->rank : 0;
    char *state;
    switch (ABTD_atomic_acquire_load_int(&p_task->state)) {
        case ABT_TASK_STATE_READY:
            state = "READY";
            break;
        case ABT_TASK_STATE_RUNNING:
            state = "RUNNING";
            break;
        case ABT_TASK_STATE_TERMINATED:
            state = "TERMINATED";
            break;
        default:
            state = "UNKNOWN";
    }

    fprintf(p_os,
            "%s== TASKLET (%p) ==\n"
            "%sid        : %" PRIu64 "\n"
            "%sstate     : %s\n"
            "%sES        : %p (%d)\n"
#ifndef ABT_CONFIG_DISABLE_STACKABLE_SCHED
            "%sis_sched  : %p\n"
#endif
            "%spool      : %p\n"
#ifndef ABT_CONFIG_DISABLE_MIGRATION
            "%smigratable: %s\n"
#endif
            "%srefcount  : %u\n"
            "%srequest   : 0x%x\n"
            "%sp_arg     : %p\n"
            "%skeytable  : %p\n",
            prefix, (void *)p_task, prefix, ABTI_task_get_id(p_task), prefix,
            state, prefix, (void *)p_task->p_xstream, xstream_rank,
#ifndef ABT_CONFIG_DISABLE_STACKABLE_SCHED
            prefix, (void *)p_task->is_sched,
#endif
            prefix, (void *)p_task->p_pool,
#ifndef ABT_CONFIG_DISABLE_MIGRATION
            prefix, (p_task->migratable == ABT_TRUE) ? "TRUE" : "FALSE",
#endif
            prefix, p_task->refcount, prefix,
            ABTD_atomic_acquire_load_uint32(&p_task->request), prefix,
            p_task->p_arg, prefix, (void *)p_task->p_keytable);

fn_exit:
    fflush(p_os);
    ABTU_free(prefix);
}

static ABTD_atomic_uint64 g_task_id = ABTD_ATOMIC_UINT64_STATIC_INITIALIZER(0);
void ABTI_task_reset_id(void)
{
    ABTD_atomic_release_store_uint64(&g_task_id, 0);
}

uint64_t ABTI_task_get_id(ABTI_task *p_task)
{
    if (p_task->id == ABTI_TASK_INIT_ID) {
        p_task->id = ABTI_task_get_new_id();
    }
    return p_task->id;
}

/*****************************************************************************/
/* Internal static functions                                                 */
/*****************************************************************************/

static inline uint64_t ABTI_task_get_new_id(void)
{
    return ABTD_atomic_fetch_add_uint64(&g_task_id, 1);
}