// SPDX-License-Identifier: MIT
/* Copyright 2015 Advanced Micro Devices, Inc. */
/* Copyright (c) 2025 Valve Corporation */

#include <linux/rbtree.h>

#include <drm/drm_print.h>
#include <drm/gpu_scheduler.h>

#include "sched_internal.h"

static __always_inline bool
drm_sched_entity_compare_before(struct rb_node *a, const struct rb_node *b)
{
        struct drm_sched_entity *ea =
                rb_entry((a), struct drm_sched_entity, rb_tree_node);
        struct drm_sched_entity *eb =
                rb_entry((b), struct drm_sched_entity, rb_tree_node);

        return ktime_before(ea->oldest_job_waiting, eb->oldest_job_waiting);
}

static void drm_sched_rq_update_prio(struct drm_sched_rq *rq)
{
        enum drm_sched_priority prio = DRM_SCHED_PRIORITY_INVALID;
        struct rb_node *rb;

        lockdep_assert_held(&rq->lock);

        rb = rb_first_cached(&rq->rb_tree_root);
        if (rb) {
                struct drm_sched_entity *entity =
                        rb_entry(rb, typeof(*entity), rb_tree_node);

                /*
                 * The normal locking order is entity then run-queue so taking
                 * the entity lock here would be a locking inversion for the
                 * case when the current head of the run-queue is different from
                 * the one we already have locked. The unlocked read is fine
                 * though, because if the priority had just changed it is no big
                 * deal for our algorithm, but just a transient reachable only
                 * by drivers with userspace dynamic priority changes API. Equal
                 * in effect to the priority change becoming visible a few
                 * instructions later.
                 */
                prio = READ_ONCE(entity->priority);
        }

        rq->head_prio = prio;
}

static void drm_sched_rq_remove_tree_locked(struct drm_sched_entity *entity,
                                            struct drm_sched_rq *rq)
{
        lockdep_assert_held(&entity->lock);
        lockdep_assert_held(&rq->lock);

        if (!RB_EMPTY_NODE(&entity->rb_tree_node)) {
                rb_erase_cached(&entity->rb_tree_node, &rq->rb_tree_root);
                RB_CLEAR_NODE(&entity->rb_tree_node);
                drm_sched_rq_update_prio(rq);
        }
}

static void drm_sched_rq_update_tree_locked(struct drm_sched_entity *entity,
                                            struct drm_sched_rq *rq,
                                            ktime_t ts)
{
        /*
         * Both locks need to be grabbed, one to protect from entity->rq change
         * for entity from within concurrent drm_sched_entity_select_rq and the
         * other to update the rb tree structure.
         */
        lockdep_assert_held(&entity->lock);
        lockdep_assert_held(&rq->lock);

        drm_sched_rq_remove_tree_locked(entity, rq);

        entity->oldest_job_waiting = ts;

        rb_add_cached(&entity->rb_tree_node, &rq->rb_tree_root,
                      drm_sched_entity_compare_before);
        drm_sched_rq_update_prio(rq);
}

/**
 * drm_sched_rq_init - initialize a given run queue struct
 * @rq: scheduler run queue
 *
 * Initializes a scheduler runqueue.
 */
void drm_sched_rq_init(struct drm_sched_rq *rq)
{
        spin_lock_init(&rq->lock);
        INIT_LIST_HEAD(&rq->entities);
        rq->rb_tree_root = RB_ROOT_CACHED;
        rq->head_prio = DRM_SCHED_PRIORITY_INVALID;
}

/*
 * Core part of the CFS-like algorithm is that the virtual runtime of lower
 * priority tasks should grow quicker than the higher priority ones, so that
 * when we then schedule entities with the aim of keeping their accumulated
 * virtual time balanced, we can approach fair distribution of GPU time.
 *
 * For converting the real GPU time into virtual we pick some multipliers with
 * the idea to achieve the following GPU time distribution:
 *
 *  - Kernel priority gets roughly 2x GPU time compared to high.
 *  - High gets ~4x relative to normal.
 *  - Normal gets ~8x relative to low.
 */
static const unsigned int vruntime_shift[] = {
        [DRM_SCHED_PRIORITY_KERNEL] = 1,
        [DRM_SCHED_PRIORITY_HIGH]   = 2,
        [DRM_SCHED_PRIORITY_NORMAL] = 4,
        [DRM_SCHED_PRIORITY_LOW]    = 7,
};

static ktime_t
drm_sched_rq_get_min_vruntime(struct drm_sched_rq *rq)
{
        ktime_t vruntime = 0;
        struct rb_node *rb;

        lockdep_assert_held(&rq->lock);

        rb = rb_first_cached(&rq->rb_tree_root);
        if (rb) {
                struct drm_sched_entity *entity =
                        rb_entry(rb, typeof(*entity), rb_tree_node);
                struct drm_sched_entity_stats *stats = entity->stats;

                spin_lock(&stats->lock);
                vruntime = stats->vruntime;
                spin_unlock(&stats->lock);
        }

        return vruntime;
}

static void
drm_sched_entity_save_vruntime(struct drm_sched_entity *entity,
                               ktime_t min_vruntime)
{
        struct drm_sched_entity_stats *stats = entity->stats;
        ktime_t vruntime;

        spin_lock(&stats->lock);
        vruntime = stats->vruntime;
        if (min_vruntime && vruntime > min_vruntime)
                vruntime = ktime_sub(vruntime, min_vruntime);
        else
                vruntime = 0;
        stats->vruntime = vruntime;
        spin_unlock(&stats->lock);
}

static ktime_t
drm_sched_entity_restore_vruntime(struct drm_sched_entity *entity,
                                  ktime_t min_vruntime,
                                  enum drm_sched_priority rq_prio)
{
        struct drm_sched_entity_stats *stats = entity->stats;
        struct drm_gpu_scheduler *sched =
                        container_of(entity->rq, typeof(*sched), rq);
        enum drm_sched_priority prio = entity->priority;
        unsigned long avg_us, sched_avg_us;
        ktime_t vruntime;

        BUILD_BUG_ON(DRM_SCHED_PRIORITY_NORMAL < DRM_SCHED_PRIORITY_HIGH);

        spin_lock(&stats->lock);
        vruntime = stats->vruntime;
        avg_us = ewma_drm_sched_avgtime_read(&stats->avg_job_us);
        /*
         * Unlocked read of the scheduler average is fine since it is just
         * heuristics and data type is a natural word size.
         */
        sched_avg_us = ewma_drm_sched_avgtime_read(&sched->avg_job_us);

        /*
         * Special handling for entities which were picked from the top of the
         * queue and are now re-joining the top with another one already there.
         */
        if (!vruntime && rq_prio != DRM_SCHED_PRIORITY_INVALID) {
                if (prio > rq_prio) {
                        /*
                         * Lower priority should not overtake higher when re-
                         * joining at the top of the queue so push it back
                         * somewhere behind the "middle" of the run-queue,
                         * proportional to the scheduler and entity average job
                         * durations.
                         */
                        vruntime = us_to_ktime((1 + avg_us + sched_avg_us) <<
                                               vruntime_shift[prio]);
                } else if (prio < rq_prio) {
                        /*
                         * Higher priority can go first.
                         */
                        vruntime = -ns_to_ktime(rq_prio - prio);
                } else {
                        /* Favour entity with shorter jobs (interactivity). */
                        if (avg_us <= sched_avg_us)
                                vruntime = -ns_to_ktime(1);
                        else
                                vruntime = ns_to_ktime(1);
                }
        }

        /*
         * Restore saved relative position in the queue.
         */
        vruntime = ktime_add(min_vruntime, vruntime);

        stats->vruntime = vruntime;
        spin_unlock(&stats->lock);

        return vruntime;
}

static ktime_t drm_sched_entity_update_vruntime(struct drm_sched_entity *entity)
{
        struct drm_sched_entity_stats *stats = entity->stats;
        ktime_t runtime, prev;

        spin_lock(&stats->lock);
        prev = stats->prev_runtime;
        runtime = stats->runtime;
        stats->prev_runtime = runtime;
        runtime = ktime_add_ns(stats->vruntime,
                               ktime_to_ns(ktime_sub(runtime, prev)) <<
                               vruntime_shift[entity->priority]);
        stats->vruntime = runtime;
        spin_unlock(&stats->lock);

        return runtime;
}

/**
 * drm_sched_rq_add_entity - add an entity
 * @entity: scheduler entity
 *
 * Adds a scheduler entity to the run queue.
 *
 * Return: DRM scheduler selected to handle this entity or NULL if entity has
 * been stopped and cannot be submitted to.
 */
struct drm_gpu_scheduler *
drm_sched_rq_add_entity(struct drm_sched_entity *entity)
{
        struct drm_gpu_scheduler *sched;
        struct drm_sched_rq *rq;
        ktime_t ts;

        /* Add the entity to the run queue */
        spin_lock(&entity->lock);
        if (entity->stopped) {
                spin_unlock(&entity->lock);

                DRM_ERROR("Trying to push to a killed entity\n");
                return NULL;
        }

        rq = entity->rq;
        sched = container_of(rq, typeof(*sched), rq);
        spin_lock(&rq->lock);

        if (list_empty(&entity->list)) {
                atomic_inc(sched->score);
                list_add_tail(&entity->list, &rq->entities);
        }

        ts = drm_sched_rq_get_min_vruntime(rq);
        ts = drm_sched_entity_restore_vruntime(entity, ts, rq->head_prio);
        drm_sched_rq_update_tree_locked(entity, rq, ts);

        spin_unlock(&rq->lock);
        spin_unlock(&entity->lock);

        return sched;
}

/**
 * drm_sched_rq_remove_entity - remove an entity
 * @rq: scheduler run queue
 * @entity: scheduler entity
 *
 * Removes a scheduler entity from the run queue.
 */
void drm_sched_rq_remove_entity(struct drm_sched_rq *rq,
                                struct drm_sched_entity *entity)
{
        struct drm_gpu_scheduler *sched = container_of(rq, typeof(*sched), rq);

        lockdep_assert_held(&entity->lock);

        if (list_empty(&entity->list))
                return;

        spin_lock(&rq->lock);

        atomic_dec(sched->score);
        list_del_init(&entity->list);

        drm_sched_rq_remove_tree_locked(entity, rq);

        spin_unlock(&rq->lock);
}

/**
 * drm_sched_rq_pop_entity - pops an entity
 * @entity: scheduler entity
 *
 * To be called every time after a job is popped from the entity.
 */
void drm_sched_rq_pop_entity(struct drm_sched_entity *entity)
{
        struct drm_sched_job *next_job;
        struct drm_sched_rq *rq;

        /*
         * Update the entity's location in the min heap according to
         * the timestamp of the next job, if any.
         */
        spin_lock(&entity->lock);
        rq = entity->rq;
        spin_lock(&rq->lock);
        next_job = drm_sched_entity_queue_peek(entity);
        if (next_job) {
                ktime_t ts;

                ts = drm_sched_entity_update_vruntime(entity);
                drm_sched_rq_update_tree_locked(entity, rq, ts);
        } else {
                ktime_t min_vruntime;

                drm_sched_rq_remove_tree_locked(entity, rq);
                min_vruntime = drm_sched_rq_get_min_vruntime(rq);
                drm_sched_entity_save_vruntime(entity, min_vruntime);
        }
        spin_unlock(&rq->lock);
        spin_unlock(&entity->lock);
}

/**
 * drm_sched_select_entity - Select an entity which provides a job to run
 * @sched: the gpu scheduler
 *
 * Find oldest waiting ready entity.
 *
 * Return an entity if one is found; return an error-pointer (!NULL) if an
 * entity was ready, but the scheduler had insufficient credits to accommodate
 * its job; return NULL, if no ready entity was found.
 */
struct drm_sched_entity *
drm_sched_select_entity(struct drm_gpu_scheduler *sched)
{
        struct drm_sched_rq *rq = &sched->rq;
        struct rb_node *rb;

        spin_lock(&rq->lock);
        for (rb = rb_first_cached(&rq->rb_tree_root); rb; rb = rb_next(rb)) {
                struct drm_sched_entity *entity;

                entity = rb_entry(rb, struct drm_sched_entity, rb_tree_node);
                if (drm_sched_entity_is_ready(entity)) {
                        /* If we can't queue yet, preserve the current entity in
                         * terms of fairness.
                         */
                        if (!drm_sched_can_queue(sched, entity)) {
                                spin_unlock(&rq->lock);
                                return ERR_PTR(-ENOSPC);
                        }

                        reinit_completion(&entity->entity_idle);
                        break;
                }
        }
        spin_unlock(&rq->lock);

        return rb ? rb_entry(rb, struct drm_sched_entity, rb_tree_node) : NULL;
}