// SPDX-License-Identifier: GPL-2.0 WITH Linux-syscall-note OR MIT

/*
 * (C) COPYRIGHT 2024 ARM Limited or its affiliates. All rights reserved.
 */

/*
 * Part of the Mali reference arbiter
 */

#include <linux/hrtimer.h>
#include <linux/module.h>
#include <linux/platform_device.h>
#include <linux/slab.h>
#include <linux/spinlock.h>
#include <linux/workqueue.h>

#include "mali_arbiter_scheduler.h"
#include "mali_arbiter_internal.h"

/* Module parameters */
static int request_timeout = 7;
module_param(request_timeout, int, 0644);
MODULE_PARM_DESC(request_timeout, "GPU time in ms before switching to another VM");

static int yield_timeout = 7;
module_param(yield_timeout, int, 0644);
MODULE_PARM_DESC(yield_timeout,
		 "Max time in ms for VM to stop before GPU lost handling is invoked");

static int no_timeslice_yield_timeout = 5000;
module_param(no_timeslice_yield_timeout, int, 0644);
MODULE_PARM_DESC(
	no_timeslice_yield_timeout,
	"Max time in ms for a driver to stop before GPU lost handling is invoked, when it is not timesliced with other drivers");


/**
 * enum mali_arb_gsi_sched_state - gpu-subinstance scheduler states
 * @GSI_SCHED_STATE_NO_REQ: No VMs are requesting in current gpu-subinstance.
 * @GSI_SCHED_STATE_SINGLE_REQ: A single VM has requested GPU.
 * @GSI_SCHED_STATE_RUNNING: Multiple VMs are requesting GPU and arbitration
 *                     is required between them.
 * @GSI_SCHED_STATE_GPU_STOPPING: VM has been told to stop using GPU and
 *                          arbiter is waiting for a stopped event.
 *
 * Definition of the possible arbiter scheduler's states
 */
enum mali_arb_gsi_sched_state {
	GSI_SCHED_STATE_NO_REQ,
	GSI_SCHED_STATE_SINGLE_REQ,
	GSI_SCHED_STATE_RUNNING,
	GSI_SCHED_STATE_GPU_STOPPING,
};

/**
 * struct mali_arb_gpu_scheduler_rr - Represents the round robin scheduler object which schedules
 *                                 the VMs on the gpu subinstance
 * @dev: Device for arbiter (only single Arbiter device supported).
 * @gsi_ptr: Pointer to gpu-subinstance private data.
 * @gsi_idx: gpu-subinstance index within its Resource Group
 * @req_list: GPU request list
 * @sched_tmr: Timer used during RUNNING and STOPPING states for
 *           scheduler purposes
 * @sched_tmr_wq: Workqueue for safely handling timer events
 * @sched_timeout_work: Timeout work item
 * @sched_tmr_running: TRUE if timer is active, FALSE otherwise
 * @sched_tmr_expired: TRUE if timer has fired and needs to be handled
 * @sched_tmr_lock: Spinlock used to protect the #gsi_tmr_expired variable
 * @active_vm: The currently active VM, or NULL if none are active
 * @gsi_cb: Callback functions for gpu-subinstance
 * @arb_hndl: Arbiter private data handler for callback functions
 * @sched_state: gpu-subinstance scheduler state
 * @last_state_transition: Time of last dvfs state transition
 * @last_dvfs_report: Time of last utilization report
 * @gpu_active_sum: Running total of active time since last report
 * @gpu_switching_sum: Running total of switching time since last report
 * @gsi_time_sliced: TRUE if GSI is time-sliced, FALSE otherwise
 * @gsi_stopping: TRUE if the gpu-subinstance is stopping
 * @gsi_vm_stop_wait: Wait queue for blocking the calling thread whilst the
 *                 active set AW for the gpu-subinstance is stopped
 * @gsi_vm_stopped: variable to let gsi_vm_stop_wait know that
 *                  the VM it is waiting is now stopped
 * @gsi_sched_cb: gsi callback functions invoked from scheduler
 */
struct mali_arb_gpu_scheduler_rr {
	struct device *dev;
	void *gsi_ptr;
	u32 gsi_idx;
	struct list_head req_list;
	/* Scheduling: Timer */
	struct hrtimer sched_tmr;
	struct workqueue_struct *sched_tmr_wq;
	struct work_struct sched_timeout_work;
	bool sched_tmr_running;
	bool sched_tmr_expired;
	spinlock_t sched_tmr_lock;
	struct mali_arb_priv_data *active_vm;
	struct mali_arb_gsi_cb gsi_cb;
	struct mali_arb *arb_hndl;
	enum mali_arb_gsi_sched_state sched_state;
	ktime_t last_state_transition;
	ktime_t last_dvfs_report;
	u64 gpu_active_sum;
	u64 gpu_switching_sum;
	bool gsi_time_sliced;
	bool gsi_stopping;
	/* synchronous vm stop */
	wait_queue_head_t gsi_vm_stop_wait;
	bool gsi_vm_stopped;
	struct mali_arb_gsi_sched_cb gsi_sched_cb;
};

/* Forward declaration */
static void gpu_lost_complete_locked(struct mali_arb_priv_data *vm,
				     struct mali_arb_gpu_scheduler_rr *sched_ptr_rr);
static inline void reset_utilisation_counters_locked(struct mali_arb_gpu_scheduler_rr *sched_ptr_rr,
						     ktime_t t);
static void run_next_vm(struct mali_arb_gpu_scheduler_rr *sched_ptr_rr);
static void set_gsi_sched_state_locked(struct mali_arb_gpu_scheduler_rr *sched_ptr_rr,
				       enum mali_arb_gsi_sched_state new_state);
static void arb_dvfs_update(struct mali_arb_gpu_scheduler_rr *sched_ptr_rr,
			    enum mali_arb_gsi_sched_state new_state);
static const char *gsi_sched_state_text(enum mali_arb_gsi_sched_state state);
static void timeout_worker(struct work_struct *data);
static void handle_gpu_lost_locked(struct mali_arb_gpu_scheduler_rr *sched_ptr_rr, bool switch_aw);
static void stop_gpu_locked(struct mali_arb_gpu_scheduler_rr *sched_ptr_rr, bool gsi_stopping);
static void stop_timer_locked(struct mali_arb_gpu_scheduler_rr *sched_ptr_rr);
static void start_timer_locked(struct mali_arb_gpu_scheduler_rr *sched_ptr_rr,
			       unsigned long timeout);
static enum hrtimer_restart timer_isr(struct hrtimer *timer);
static int remove_active_vm(struct mali_arb_gpu_scheduler_rr *sched_ptr_rr);
static bool mali_arbiter_sched_remove_vm_rr(void *sched_ptr_rr, struct mali_arb_priv_data *rem_vm,
					    bool req_again);
static void mali_arbiter_sched_resync_vm_rr(void *sched_ptr, struct mali_arb_priv_data *rem_vm);


/**
 * timer_isr() - Called when timer fires as ISR
 * @timer: timer data
 *
 * Queues a workqueue item to handle the event
 *
 * Return: enum hrtimer_restart, which tells status of the timer.
 */
static enum hrtimer_restart timer_isr(struct hrtimer *timer)
{
	struct mali_arb_gpu_scheduler_rr *sched_rr;

	sched_rr = container_of(timer, struct mali_arb_gpu_scheduler_rr, sched_tmr);

	dev_dbg(sched_rr->dev, "P%u: hrtimer_isr\n", sched_rr->gsi_idx);
	spin_lock(&sched_rr->sched_tmr_lock);
	sched_rr->sched_tmr_expired = true;
	queue_work(sched_rr->sched_tmr_wq, &sched_rr->sched_timeout_work);
	spin_unlock(&sched_rr->sched_tmr_lock);
	return HRTIMER_NORESTART;
}

/**
 * start_timer_locked() - Start the timer with a given timeout
 * @sched_ptr_rr: Pointer to round robin scheduler private data.
 * @timeout: Timeout in milliseconds.
 */
static void start_timer_locked(struct mali_arb_gpu_scheduler_rr *sched_ptr_rr,
			       unsigned long timeout)
{
	unsigned long irqflags;

	if (WARN_ON(!sched_ptr_rr))
		return;

	/* Check if lock is held */
	sched_ptr_rr->gsi_cb.cb_lock_assert_held(sched_ptr_rr->arb_hndl, sched_ptr_rr->gsi_idx);

	spin_lock_irqsave(&sched_ptr_rr->sched_tmr_lock, irqflags);
	sched_ptr_rr->sched_tmr_expired = false;
	spin_unlock_irqrestore(&sched_ptr_rr->sched_tmr_lock, irqflags);
	hrtimer_start(&sched_ptr_rr->sched_tmr, ms_to_ktime(timeout), HRTIMER_MODE_REL);
	sched_ptr_rr->sched_tmr_running = true;
}

/**
 * stop_timer_locked() - Stop the timer
 * @sched_ptr_rr: Pointer to round robin scheduler private data.
 */
static void stop_timer_locked(struct mali_arb_gpu_scheduler_rr *sched_ptr_rr)
{
	unsigned long irqflags;

	if (WARN_ON(!sched_ptr_rr))
		return;

	/* Check if lock is held */
	sched_ptr_rr->gsi_cb.cb_lock_assert_held(sched_ptr_rr->arb_hndl, sched_ptr_rr->gsi_idx);

	sched_ptr_rr->sched_tmr_running = false;
	hrtimer_cancel(&sched_ptr_rr->sched_tmr);

	spin_lock_irqsave(&sched_ptr_rr->sched_tmr_lock, irqflags);
	sched_ptr_rr->sched_tmr_expired = false;
	spin_unlock_irqrestore(&sched_ptr_rr->sched_tmr_lock, irqflags);
}

static void stop_gpu_locked(struct mali_arb_gpu_scheduler_rr *sched_ptr_rr, bool gsi_stopping)
{
	struct mali_arb_priv_data *arb_vm;
	int new_yield_timeout = yield_timeout;

	if (WARN_ON(!sched_ptr_rr))
		return;

	/* Check if lock is held */
	sched_ptr_rr->gsi_cb.cb_lock_assert_held(sched_ptr_rr->arb_hndl, sched_ptr_rr->gsi_idx);

	arb_vm = sched_ptr_rr->active_vm;
	if (!arb_vm) {
		dev_warn(sched_ptr_rr->dev,
			 "P%u: GPU_STOP cannot be issued without an active VM.\n",
			 sched_ptr_rr->gsi_idx);
		return;
	}

	set_gsi_sched_state_locked(sched_ptr_rr, GSI_SCHED_STATE_GPU_STOPPING);

	if (gsi_stopping && !sched_ptr_rr->gsi_time_sliced)
		new_yield_timeout = no_timeslice_yield_timeout;

	start_timer_locked(sched_ptr_rr, new_yield_timeout);
	sched_ptr_rr->gsi_sched_cb.gpu_stop(arb_vm);
}

/**
 * handle_gpu_lost_locked() - Remove GPU from misbehaving VM
 * @sched_ptr_rr: Pointer to round robin scheduler private data.
 * @switch_aw: TRUE when this function is invoked during reassignment
 *
 * Forcible assign the GPU away from the currently active domain to
 * another VM in the request queue, then send a GPU lost event.
 */
static void handle_gpu_lost_locked(struct mali_arb_gpu_scheduler_rr *sched_ptr_rr, bool switch_aw)
{
	struct mali_arb_priv_data *arb_vm;
	struct mali_vm_priv_data *vm;

	if (WARN_ON(!sched_ptr_rr))
		return;

	/* Check if lock is held */
	sched_ptr_rr->gsi_cb.cb_lock_assert_held(sched_ptr_rr->arb_hndl, sched_ptr_rr->gsi_idx);

	arb_vm = sched_ptr_rr->active_vm;
	if (!arb_vm) {
		dev_warn(sched_ptr_rr->dev,
			 "P%u: GPU_LOST cannot be issued without an active VM.\n",
			 sched_ptr_rr->gsi_idx);
		return;
	}

	vm = arb_vm->vm;
	if (!vm) {
		dev_warn(sched_ptr_rr->dev, "P%u: Private VM data not available.\n",
			 sched_ptr_rr->gsi_idx);
		return;
	}

	dev_warn(sched_ptr_rr->dev, "P%u: GPU in VM%u took too long to yield.\n",
		 sched_ptr_rr->gsi_idx, arb_vm->info.id);

	list_del_init(&arb_vm->entry);
	remove_active_vm(sched_ptr_rr);

	/* Send GPU lost before reassigning to another VM. */
	arb_vm->gpu_lost = true;
#if !MALI_USE_CSF
	sched_ptr_rr->gsi_sched_cb.gpu_lost(arb_vm);
#endif

	/* Only switch to another AW if not performing a reassignment.
	 * GSI_FLAG_REASSIGNING will never set when running in Xen/PV mode.
	 */

	if (switch_aw) {
		run_next_vm(sched_ptr_rr);
	} else {
		sched_ptr_rr->gsi_stopping = false;
		sched_ptr_rr->gsi_vm_stopped = true;
		wake_up(&sched_ptr_rr->gsi_vm_stop_wait);
	}
}

/**
 * gpu_lost_complete_locked() - Cleanup after GPU lost completion
 * @arb_vm: Internal Arbiter data for VM
 * @sched_ptr_rr: Pointer to round robin scheduler private data.
 *
 * Clear GPU lost flags and power down GPU if necessary.
 */
static void gpu_lost_complete_locked(struct mali_arb_priv_data *arb_vm,
				     struct mali_arb_gpu_scheduler_rr *sched_ptr_rr)
{
	if (WARN_ON(!sched_ptr_rr))
		return;

	/* Check if lock is held */
	sched_ptr_rr->gsi_cb.cb_lock_assert_held(sched_ptr_rr->arb_hndl, sched_ptr_rr->gsi_idx);

	if (!arb_vm)
		return;

	if (arb_vm->gpu_lost) {
		arb_vm->gpu_lost = false;
		dev_dbg(sched_ptr_rr->dev, "VM%u GPU lost complete\n", arb_vm->info.id);
	}
}

/**
 * timeout_worker() - Called when timer expires
 * @data: Used to obtain arbiter data
 *
 * Handles the timeout cases in the RUNNING and GPU_STOPPING states
 */
static void timeout_worker(struct work_struct *data)
{
	struct mali_arb_gpu_scheduler_rr *sched_rr;
	struct mali_arb_priv_data *arb_vm;
	unsigned long irqflags;

	if (WARN_ON(!data))
		return;

	sched_rr = container_of(data, struct mali_arb_gpu_scheduler_rr, sched_timeout_work);

	sched_rr->gsi_cb.cb_lock(sched_rr->arb_hndl, sched_rr->gsi_idx);
	if (!sched_rr->sched_tmr_running)
		goto cleanup_mutex;

	spin_lock_irqsave(&sched_rr->sched_tmr_lock, irqflags);
	if (!sched_rr->sched_tmr_expired) {
		/* Timer was stopped/restarted before worker ran */
		spin_unlock_irqrestore(&sched_rr->sched_tmr_lock, irqflags);
		goto cleanup_mutex;
	}
	sched_rr->sched_tmr_expired = false;
	spin_unlock_irqrestore(&sched_rr->sched_tmr_lock, irqflags);

	dev_dbg(sched_rr->dev, "P%u: Timer has expired...\n", sched_rr->gsi_idx);
	sched_rr->sched_tmr_running = false;
	arb_vm = sched_rr->active_vm;
	if (!arb_vm) {
		dev_dbg(sched_rr->dev, "P%u: Timer cannot be pending without an active VM.\n",
			sched_rr->gsi_idx);
		goto cleanup_mutex;
	}
	switch (sched_rr->sched_state) {
	case GSI_SCHED_STATE_RUNNING:
		dev_dbg(sched_rr->dev, "P%u: GPU time up - suspending.\n", sched_rr->gsi_idx);
		stop_gpu_locked(sched_rr, false);
		break;
	case GSI_SCHED_STATE_GPU_STOPPING:
		/* Only forcibly remove the GPU if there is another VM waiting
		 * or the gpu-subinstance is under reassignment. Otherwise, allow the
		 * VM to take longer.
		 */

		if (sched_rr->gsi_stopping || !list_is_singular(&sched_rr->req_list))
			handle_gpu_lost_locked(sched_rr, !sched_rr->gsi_stopping);
		else
			dev_dbg(sched_rr->dev, "GPU time is up but no VMs waiting.\n");
		break;
	default:
		break;
	}

cleanup_mutex:
	sched_rr->gsi_cb.cb_unlock(sched_rr->arb_hndl, sched_rr->gsi_idx);
}

/**
 * gsi_sched_state_text() - Helper to convert state to string (for debug)
 * @state: State to convert.
 *
 * Return: A null terminated string
 */
static const char *gsi_sched_state_text(enum mali_arb_gsi_sched_state state)
{
	switch (state) {
	case GSI_SCHED_STATE_NO_REQ:
		return "GSI_SCHED_STATE_NO_REQ";
	case GSI_SCHED_STATE_SINGLE_REQ:
		return "GSI_SCHED_STATE_SINGLE_REQ";
	case GSI_SCHED_STATE_RUNNING:
		return "GSI_SCHED_STATE_RUNNING";
	case GSI_SCHED_STATE_GPU_STOPPING:
		return "GSI_SCHED_STATE_GPU_STOPPING";
	default:
		return "UNKNOWN";
	}
}

/**
 * arb_dvfs_update() - update dvfs data.
 * @sched_ptr_rr: Pointer to round robin scheduler private data.
 * @new_state: state being transition to.
 *
 * this function updates the running totals for active and switching times
 * and/or restarts the dvfs timestamp as necessary before we transition state.
 */
static void arb_dvfs_update(struct mali_arb_gpu_scheduler_rr *sched_ptr_rr,
			    enum mali_arb_gsi_sched_state new_state)
{
	ktime_t now = ktime_get();
	u64 since_last;

	if (WARN_ON(!sched_ptr_rr))
		return;

	/* Check if lock is held */
	sched_ptr_rr->gsi_cb.cb_lock_assert_held(sched_ptr_rr->arb_hndl, sched_ptr_rr->gsi_idx);

	since_last = ktime_to_ns(ktime_sub(now, sched_ptr_rr->last_state_transition));

	switch (sched_ptr_rr->sched_state) {
	case GSI_SCHED_STATE_NO_REQ:
		switch (new_state) {
		/* idle --> active */
		case GSI_SCHED_STATE_SINGLE_REQ:
			dev_dbg(sched_ptr_rr->dev, "P%u: dvfs: idle --> active\n",
				sched_ptr_rr->gsi_idx);
			sched_ptr_rr->last_state_transition = now;
			break;
		default:
			break;
		}
		break;
	case GSI_SCHED_STATE_SINGLE_REQ:
	case GSI_SCHED_STATE_RUNNING:
		switch (new_state) {
		/* active --> switching */
		case GSI_SCHED_STATE_GPU_STOPPING:
			dev_dbg(sched_ptr_rr->dev, "P%u dvfs: active --> switching\n",
				sched_ptr_rr->gsi_idx);
			sched_ptr_rr->gpu_active_sum += since_last;
			sched_ptr_rr->last_state_transition = now;
			break;
		default:
			break;
		}
		break;
	case GSI_SCHED_STATE_GPU_STOPPING:
		switch (new_state) {
		/* switching --> active */
		case GSI_SCHED_STATE_RUNNING:
		case GSI_SCHED_STATE_SINGLE_REQ:
			dev_dbg(sched_ptr_rr->dev, "P%u dvfs: switching --> active\n",
				sched_ptr_rr->gsi_idx);
			sched_ptr_rr->gpu_switching_sum += since_last;
			sched_ptr_rr->last_state_transition = now;
			break;
		/* switching --> idle */
		case GSI_SCHED_STATE_NO_REQ:
			sched_ptr_rr->gpu_switching_sum += since_last;
			dev_dbg(sched_ptr_rr->dev, "P%u dvfs: switching --> idle\n",
				sched_ptr_rr->gsi_idx);
			break;
		default:
			break;
		}
		break;
	default:
		break;
	}
}

/**
 * set_gsi_sched_state_locked() - Used to modify the gpu-subinstance state
 * @sched_ptr_rr: Pointer to round robin scheduler private data.
 * @new_state: New gpu-subinstance state.
 *
 * Mutex must be held when calling this function.
 */
static void set_gsi_sched_state_locked(struct mali_arb_gpu_scheduler_rr *sched_ptr_rr,
				       enum mali_arb_gsi_sched_state new_state)
{
	if (WARN_ON(!sched_ptr_rr))
		return;

	/* Check if lock is held */
	sched_ptr_rr->gsi_cb.cb_lock_assert_held(sched_ptr_rr->arb_hndl, sched_ptr_rr->gsi_idx);

	if (sched_ptr_rr->sched_state != new_state) {
		dev_dbg(sched_ptr_rr->dev, "%s -> %s\n",
			gsi_sched_state_text(sched_ptr_rr->sched_state),
			gsi_sched_state_text(new_state));
		arb_dvfs_update(sched_ptr_rr, new_state);
		sched_ptr_rr->sched_state = new_state;
	}
}

/**
 * run_next_vm() - Run the next VM
 * @sched_ptr_rr: Pointer to round robin scheduler private data.
 *
 * Makes the necessary state transitions based on the current state and VM list.
 * The mutex lock must already be held when calling this function.
 */
static void run_next_vm(struct mali_arb_gpu_scheduler_rr *sched_ptr_rr)
{
	struct mali_arb_priv_data *prev_vm;

	if (WARN_ON(!sched_ptr_rr))
		return;

	/* Check if lock is held */
	sched_ptr_rr->gsi_cb.cb_lock_assert_held(sched_ptr_rr->arb_hndl, sched_ptr_rr->gsi_idx);

	if (sched_ptr_rr->sched_tmr_running)
		stop_timer_locked(sched_ptr_rr);

	if (list_empty(&sched_ptr_rr->req_list)) {
		if (sched_ptr_rr->sched_state != GSI_SCHED_STATE_NO_REQ) {
			remove_active_vm(sched_ptr_rr);
			sched_ptr_rr->gsi_cb.cb_idle(sched_ptr_rr->arb_hndl, sched_ptr_rr->gsi_idx);
		}
		set_gsi_sched_state_locked(sched_ptr_rr, GSI_SCHED_STATE_NO_REQ);
	} else if (list_is_singular(&sched_ptr_rr->req_list))
		set_gsi_sched_state_locked(sched_ptr_rr, GSI_SCHED_STATE_SINGLE_REQ);
	else
		set_gsi_sched_state_locked(sched_ptr_rr, GSI_SCHED_STATE_RUNNING);

	if (sched_ptr_rr->sched_state != GSI_SCHED_STATE_NO_REQ) {
		int ret = 0;

		prev_vm = sched_ptr_rr->active_vm;
		sched_ptr_rr->active_vm =
			list_first_entry(&sched_ptr_rr->req_list, struct mali_arb_priv_data, entry);

		if (prev_vm == sched_ptr_rr->active_vm)
			goto start_timer;

		ret = sched_ptr_rr->gsi_sched_cb.grant_gpu(sched_ptr_rr->gsi_ptr,
							   sched_ptr_rr->active_vm);

		if (ret) {
			sched_ptr_rr->active_vm = prev_vm;
			return;
		}


start_timer:
		if (sched_ptr_rr->sched_state == GSI_SCHED_STATE_RUNNING)
			start_timer_locked(sched_ptr_rr, request_timeout);
	}
}

/**
 * reset_utilisation_counters_locked() - Reset utilization counters
 * @sched_ptr_rr: Pointer to round robin scheduler private data.
 * @t: Timestamp to be used.
 *
 * Called from get_utilisation function to reset the counters.
 */
static inline void reset_utilisation_counters_locked(struct mali_arb_gpu_scheduler_rr *sched_ptr_rr,
						     ktime_t t)
{
	sched_ptr_rr->gpu_active_sum = 0;
	sched_ptr_rr->gpu_switching_sum = 0;
	sched_ptr_rr->last_state_transition = t;
	sched_ptr_rr->last_dvfs_report = t;
}

/* Shift used for dvfs.time busy/idle units of (1 << 8) ns
 * This gives a maximum period between utilization reports of 2^(32+8) ns
 * Exceeding this will cause overflow
 */
#define ARBITER_TIME_SHIFT 8

/** Get gpu-subinstance utilization info
 */
static void mali_arbiter_sched_get_utilisation_rr(void *sched_ptr, u32 *gsi_busytime,
						  u32 *gsi_totaltime)
{
	ktime_t now;
	u64 time_since_last_report, time_since_last_transition;
	struct mali_arb_gpu_scheduler_rr *sched_ptr_rr =
		(struct mali_arb_gpu_scheduler_rr *)sched_ptr;

	if (WARN_ON(!sched_ptr_rr))
		return;

	/* Check if lock is held */
	sched_ptr_rr->gsi_cb.cb_lock_assert_held(sched_ptr_rr->arb_hndl, sched_ptr_rr->gsi_idx);

	now = ktime_get();
	time_since_last_report = ktime_to_ns(ktime_sub(now, sched_ptr_rr->last_dvfs_report));
	time_since_last_transition =
		ktime_to_ns(ktime_sub(now, sched_ptr_rr->last_state_transition));

	switch (sched_ptr_rr->sched_state) {
	case GSI_SCHED_STATE_RUNNING:
	case GSI_SCHED_STATE_SINGLE_REQ: /* active */
		sched_ptr_rr->gpu_active_sum += time_since_last_transition;
		break;
	case GSI_SCHED_STATE_GPU_STOPPING: /* switching */
		sched_ptr_rr->gpu_switching_sum += time_since_last_transition;
		break;
	default:
		break;
	}

	/* get GPU time and log utilization info */
	*gsi_busytime = sched_ptr_rr->gpu_active_sum >> ARBITER_TIME_SHIFT;
	*gsi_totaltime = (time_since_last_report - sched_ptr_rr->gpu_switching_sum) >>
			 ARBITER_TIME_SHIFT;

	/* Reset counters for next time */
	reset_utilisation_counters_locked(sched_ptr_rr, now);
}

/**
 * remove_active_vm() - Remove active VM from the gpu-subinstance
 * @sched_ptr_rr: Pointer to round robin scheduler private data.
 *
 * Removes the active domain on the arbiter and attempts to unassign and reset
 * the gpu-subinstance or a whole GPU.
 *
 * Return: 0 on success, or error code
 */
static int remove_active_vm(struct mali_arb_gpu_scheduler_rr *sched_ptr_rr)
{
	int ret;

	if (WARN_ON(!sched_ptr_rr))
		return -EINVAL;

	/* Check if lock is held */
	sched_ptr_rr->gsi_cb.cb_lock_assert_held(sched_ptr_rr->arb_hndl, sched_ptr_rr->gsi_idx);
	sched_ptr_rr->active_vm = NULL;

	ret = sched_ptr_rr->gsi_sched_cb.unassign_vm(sched_ptr_rr->gsi_ptr);
	return ret;
}

static struct mali_arb_priv_data *mali_arbiter_sched_get_active_vm_rr(void *sched_ptr)
{
	struct mali_arb_gpu_scheduler_rr *sched_ptr_rr =
		(struct mali_arb_gpu_scheduler_rr *)sched_ptr;

	return sched_ptr_rr->active_vm;
}


static void mali_arbiter_sched_stop_idle_vm_rr(void *sched_ptr, struct mali_arb_priv_data *arb_vm)
{
	struct mali_arb_gpu_scheduler_rr *sched_ptr_rr =
		(struct mali_arb_gpu_scheduler_rr *)sched_ptr;

	if (WARN_ON(!arb_vm) || WARN_ON(!sched_ptr_rr))
		return;

	if (arb_vm == sched_ptr_rr->active_vm &&
	    (sched_ptr_rr->sched_state == GSI_SCHED_STATE_RUNNING ||
	     sched_ptr_rr->sched_state == GSI_SCHED_STATE_SINGLE_REQ)) {
		dev_dbg(sched_ptr_rr->dev, "P%u: GPU %u is idle - giving up...\n",
			sched_ptr_rr->gsi_idx, arb_vm->info.id);
		if (sched_ptr_rr->sched_state == GSI_SCHED_STATE_RUNNING)
			stop_timer_locked(sched_ptr_rr);

		stop_gpu_locked(sched_ptr_rr, false);
	} else {
		dev_dbg(sched_ptr_rr->dev, "P%u: GPU %u is idle\n", sched_ptr_rr->gsi_idx,
			arb_vm->info.id);
	}
}

static void mali_arbiter_sched_stop_rr(void *sched_ptr)
{
	int32_t aw_id = 0;
	bool needs_blocking, is_running = false;
	int ret;
	struct mali_arb_gpu_scheduler_rr *sched_ptr_rr =
		(struct mali_arb_gpu_scheduler_rr *)sched_ptr;

	if (WARN_ON(!sched_ptr_rr))
		return;

	/* Check if lock is held */
	sched_ptr_rr->gsi_cb.cb_lock_assert_held(sched_ptr_rr->arb_hndl, sched_ptr_rr->gsi_idx);

	sched_ptr_rr->gsi_stopping = true;

	if (sched_ptr_rr->active_vm)
		is_running = true;
	needs_blocking = is_running || sched_ptr_rr->sched_state == GSI_SCHED_STATE_GPU_STOPPING;

	if (needs_blocking) {
		sched_ptr_rr->gsi_vm_stopped = false;
		if (is_running) {
			/* Issue a stop command */
			if (sched_ptr_rr->sched_tmr_running)
				stop_timer_locked(sched_ptr_rr);

			stop_gpu_locked(sched_ptr_rr, true);
		}

		/* Wait for the gsi_ptr to be relinquished */
		sched_ptr_rr->gsi_cb.cb_unlock(sched_ptr_rr->arb_hndl, sched_ptr_rr->gsi_idx);
		wait_event(sched_ptr_rr->gsi_vm_stop_wait, sched_ptr_rr->gsi_vm_stopped);
		sched_ptr_rr->gsi_cb.cb_lock(sched_ptr_rr->arb_hndl, sched_ptr_rr->gsi_idx);
	}

	/* De-assign VM if it hasn't been done during gpu_stopped(). */
	ret = sched_ptr_rr->gsi_sched_cb.get_assigned_vm(sched_ptr_rr->gsi_ptr, &aw_id);
	if (ret) {
		dev_err(sched_ptr_rr->dev, "P%u: Failed to get assigned AW. (err = %d)\n",
			sched_ptr_rr->gsi_idx, ret);
		return;
	}

	if (aw_id != VM_UNASSIGNED) {
		ret = remove_active_vm(sched_ptr_rr);
		if (ret) {
			dev_err(sched_ptr_rr->dev,
				"P%u: gpu-subinstance reset failed. (err = %d)\n",
				sched_ptr_rr->gsi_idx, ret);
			return;
		}
	}
}

static void mali_arbiter_sched_start_rr(void *sched_ptr)
{
	struct mali_arb_gpu_scheduler_rr *sched_ptr_rr =
		(struct mali_arb_gpu_scheduler_rr *)sched_ptr;

	sched_ptr_rr->gsi_stopping = false;
	run_next_vm(sched_ptr_rr);
}

static void mali_arbiter_sched_add_vm_rr(void *sched_ptr, struct mali_arb_priv_data *add_vm
)
{
	struct mali_arb_priv_data *cur_vm;
	struct mali_arb_gpu_scheduler_rr *sched_ptr_rr =
		(struct mali_arb_gpu_scheduler_rr *)sched_ptr;

	if (WARN_ON(!sched_ptr_rr || !add_vm))
		return;


	/* While reassigning, just add the VM to req_list and return */
	if (sched_ptr_rr->gsi_stopping) {
		dev_info(sched_ptr_rr->dev, "Reassigning VM-%d to GSI-%d", add_vm->info.id,
			 sched_ptr_rr->gsi_idx);

		/* Enable power if needed. */
		if (list_empty(&sched_ptr_rr->req_list))
			sched_ptr_rr->gsi_cb.cb_active(sched_ptr_rr->arb_hndl,
						       sched_ptr_rr->gsi_idx);

		/* Add VM to req_list. */
		list_add_tail(&add_vm->entry, &sched_ptr_rr->req_list);
		return;
	}

	/* gpu_lost operation is complete for the VM */
	gpu_lost_complete_locked(add_vm, sched_ptr_rr);

	/* enable slice first when queue is empty */
	if (list_empty(&sched_ptr_rr->req_list))
		sched_ptr_rr->gsi_cb.cb_active(sched_ptr_rr->arb_hndl, sched_ptr_rr->gsi_idx);

	/* If GPU_REQUEST arrives after sending GPU_STOP
	 * then it is likely that a GPU_STOPPED message was missed.
	 * Adjust the state to stopped and add this VM to the
	 * stopped-requested list
	 */
	list_for_each_entry(cur_vm, &sched_ptr_rr->req_list, entry) {
		if (cur_vm == add_vm) {
			dev_dbg(sched_ptr_rr->dev,
				"P%u: VM %u requested GPU when already requested\n",
				sched_ptr_rr->gsi_idx, add_vm->info.id);
			mali_arbiter_sched_remove_vm_rr(sched_ptr_rr, add_vm, true);
			return;
		}
	}

	/* handle req_list */
	list_add_tail(&add_vm->entry, &sched_ptr_rr->req_list);

	if (!list_is_singular(&sched_ptr_rr->req_list))
		sched_ptr_rr->gsi_time_sliced = true;

	/* It might be the case that the active VM was requested to stop
	 * but did not do so in time. If there were no other VMs waiting
	 * to run then GPU lost will not have been done. Now there is a VM
	 * requesting the GPU we should forcibly remove it.
	 */
	if (sched_ptr_rr->sched_state == GSI_SCHED_STATE_GPU_STOPPING &&
	    !sched_ptr_rr->sched_tmr_running) {
		dev_dbg(sched_ptr_rr->dev, "New VM request, GPU lost now required\n");
		handle_gpu_lost_locked(sched_ptr_rr, true);
	}

	if (sched_ptr_rr->sched_state == GSI_SCHED_STATE_NO_REQ ||
	    sched_ptr_rr->sched_state == GSI_SCHED_STATE_SINGLE_REQ)
		run_next_vm(sched_ptr_rr);
}

static void mali_arbiter_sched_resync_vm_rr(void *sched_ptr, struct mali_arb_priv_data *rem_vm)
{
	struct mali_arb_gpu_scheduler_rr *sched_ptr_rr =
		(struct mali_arb_gpu_scheduler_rr *)sched_ptr;

	if (WARN_ON(!sched_ptr_rr || !rem_vm))
		return;

	if (sched_ptr_rr->active_vm == NULL) {
		int32_t aw_id = 0;
		int err;
		/* If active_vm is NULL, this can mean arbiter is restarted.
		 * The VM to be removed might be running on current gpu-subinstance.
		 * In sush case, check vm_id of assigned VM with assign interface
		 * and do soft-stop.
		 */
		err = sched_ptr_rr->gsi_sched_cb.get_assigned_vm(sched_ptr_rr->gsi_ptr, &aw_id);
		if (unlikely(err)) {
			dev_err(sched_ptr_rr->dev, "P%u: get assigned VM failed. (err = %d)\n",
				sched_ptr_rr->gsi_idx, err);
			return;
		}

		dev_dbg(sched_ptr_rr->dev, "gsi-remove-vm: Assigned VM-id = %d", aw_id);
		/* Stop VM when aw_id is valid and matches to vm_id of rem_vm */
		if (aw_id != VM_UNASSIGNED && rem_vm->info.id == (uint32_t)aw_id) {
			dev_dbg(sched_ptr_rr->dev,
				"gsi-remove-vm: VM-#%d is an active VM, stopping it from GSI-#%d",
				rem_vm->info.id, sched_ptr_rr->gsi_idx);
			sched_ptr_rr->active_vm = rem_vm;
			stop_gpu_locked(sched_ptr_rr, false);
		}
	}
}

static bool mali_arbiter_sched_remove_vm_rr(void *sched_ptr, struct mali_arb_priv_data *rem_vm,
					    bool req_again)
{
	struct mali_arb_priv_data *cur_vm, *tmp_vm;
	bool requested = false;
	struct mali_arb_gpu_scheduler_rr *sched_ptr_rr =
		(struct mali_arb_gpu_scheduler_rr *)sched_ptr;

	if (WARN_ON(!sched_ptr_rr || !rem_vm))
		return requested;

	if (!list_empty(&rem_vm->entry)) {
		/* Handle if it's an active VM */
		if (sched_ptr_rr->active_vm == rem_vm) {
			if (req_again)
				/* ensure GRANTED is sent again */
				list_move_tail(&rem_vm->entry, &sched_ptr_rr->req_list);
			else
				list_del_init(&rem_vm->entry);
			dev_dbg(sched_ptr_rr->dev,
				"gsi-remove-vm: VM-#%d is an active VM, removing it from GSI-#%d",
				rem_vm->info.id, sched_ptr_rr->gsi_idx);

			/* If active_vm is the one, do hard-stop and run next VM. */
			remove_active_vm(sched_ptr_rr);
			if (!sched_ptr_rr->gsi_stopping)
				run_next_vm(sched_ptr_rr);
		}
	} else if (req_again) {
		/* The VM was previously forcibly removed because it didn't
		 * release the GPU in time. The VM is no longer in the arbiter
		 * request list so it needs to be added again if the req_again
		 * parameter is true, otherwise the VM will not be able to
		 * access the GPU.
		 */
		dev_dbg(sched_ptr_rr->dev, "P%u - vm->entry is empty. requesting gpu for VM %u\n",
			sched_ptr_rr->gsi_idx, rem_vm->info.id);

		/* enable slice first when queue is empty */
		if (list_empty(&sched_ptr_rr->req_list))
			sched_ptr_rr->gsi_cb.cb_active(sched_ptr_rr->arb_hndl,
						       sched_ptr_rr->gsi_idx);

		list_add_tail(&rem_vm->entry, &sched_ptr_rr->req_list);

		if (sched_ptr_rr->sched_state == GSI_SCHED_STATE_NO_REQ ||
		    sched_ptr_rr->sched_state == GSI_SCHED_STATE_SINGLE_REQ) {
			WARN_ON(sched_ptr_rr->sched_tmr_running);

			if (!sched_ptr_rr->gsi_stopping)
				run_next_vm(sched_ptr_rr);
		}
	}

	/* Handle GPU lost if it was being done. */
	gpu_lost_complete_locked(rem_vm, sched_ptr_rr);
	if (!req_again) {
		/* Remove from req_list if it is requested. */
		list_for_each_entry_safe(cur_vm, tmp_vm, &sched_ptr_rr->req_list, entry) {
			if (cur_vm->info.id == rem_vm->info.id) {
				dev_dbg(sched_ptr_rr->dev,
					"gsi-remove-vm: Removing VM-#%d from req_list of GSI-#%d",
					rem_vm->info.id, sched_ptr_rr->gsi_idx);
				list_del_init(&cur_vm->entry);
				requested = true;
				break;
			}
		}
	}

	/* Wake up the slice assignment thread */
	if (sched_ptr_rr->gsi_stopping) {
		sched_ptr_rr->gsi_stopping = false;
		sched_ptr_rr->gsi_vm_stopped = true;
		wake_up(&sched_ptr_rr->gsi_vm_stop_wait);
	}

	return requested;
}

/** Creates a scheduler
 */
void *mali_arbiter_sched_create_rr(struct device *dev, void *gsi_ptr, u32 gsi_idx,
				   struct mali_arb_gsi_cb gsi_cb, struct mali_arb *arb_hndl,
				   struct mali_arb_gsi_sched_cb gsi_sched_cb,
				   struct mali_arb_gsi_sched_ops *sched_ops)
{
	struct mali_arb_gpu_scheduler_rr *sched_ptr_rr;

	if (!gsi_ptr)
		return NULL;

	sched_ptr_rr = kzalloc(sizeof(struct mali_arb_gpu_scheduler_rr), GFP_KERNEL);

	if (!sched_ptr_rr)
		return NULL;

	/* Module param check */
	if (no_timeslice_yield_timeout < yield_timeout)
		dev_warn_once(dev, "no_timeslice_yield_timeout is shorter than yield_timeout");

	sched_ptr_rr->gsi_time_sliced = false;
	sched_ptr_rr->gsi_stopping = false;
	spin_lock_init(&sched_ptr_rr->sched_tmr_lock);
	hrtimer_init(&sched_ptr_rr->sched_tmr, CLOCK_MONOTONIC, HRTIMER_MODE_REL);
	sched_ptr_rr->sched_tmr.function = timer_isr;
	INIT_LIST_HEAD(&sched_ptr_rr->req_list);
	sched_ptr_rr->sched_tmr_wq =
		alloc_workqueue("mali_arb_D%uP%u", WQ_UNBOUND, 0, dev->id, gsi_idx);
	INIT_WORK(&sched_ptr_rr->sched_timeout_work, timeout_worker);

	sched_ptr_rr->dev = dev;
	sched_ptr_rr->gsi_ptr = gsi_ptr;
	sched_ptr_rr->gsi_idx = gsi_idx;
	sched_ptr_rr->gsi_cb = gsi_cb;
	sched_ptr_rr->arb_hndl = arb_hndl;
	sched_ptr_rr->gsi_sched_cb = gsi_sched_cb;

	sched_ptr_rr->last_dvfs_report = ktime_get();
	sched_ptr_rr->last_state_transition = sched_ptr_rr->last_dvfs_report;

	init_waitqueue_head(&sched_ptr_rr->gsi_vm_stop_wait);
	sched_ptr_rr->gsi_vm_stopped = false;

	sched_ops->sched_add_vm = mali_arbiter_sched_add_vm_rr;
	sched_ops->sched_stop_idle_vm = mali_arbiter_sched_stop_idle_vm_rr;
	sched_ops->sched_get_utilisation = mali_arbiter_sched_get_utilisation_rr;
	sched_ops->sched_remove_vm = mali_arbiter_sched_remove_vm_rr;
	sched_ops->sched_resync_vm = mali_arbiter_sched_resync_vm_rr;
	sched_ops->sched_start = mali_arbiter_sched_start_rr;
	sched_ops->sched_stop = mali_arbiter_sched_stop_rr;
	sched_ops->sched_get_active_vm = mali_arbiter_sched_get_active_vm_rr;


	return sched_ptr_rr;
}

/** Destroys a scheduler
 */
void mali_arbiter_sched_destroy_rr(void *sched_ptr)
{
	struct mali_arb_gpu_scheduler_rr *sched_ptr_rr =
		(struct mali_arb_gpu_scheduler_rr *)sched_ptr;

	if (!sched_ptr_rr)
		return;


	if (sched_ptr_rr->sched_tmr_running)
		hrtimer_cancel(&sched_ptr_rr->sched_tmr);

	if (sched_ptr_rr->sched_tmr_wq)
		destroy_workqueue(sched_ptr_rr->sched_tmr_wq);

	kfree(sched_ptr_rr);
}
