// SPDX-License-Identifier: MIT
/*
 * Copyright © 2025 Intel Corporation
 */

#include <drm/drm_drv.h>

#include "instructions/xe_mi_commands.h"
#include "instructions/xe_gpu_commands.h"
#include "xe_bb.h"
#include "xe_bo.h"
#include "xe_device.h"
#include "xe_exec_queue.h"
#include "xe_exec_queue_types.h"
#include "xe_gt_sriov_vf.h"
#include "xe_guc.h"
#include "xe_guc_submit.h"
#include "xe_lrc.h"
#include "xe_mem_pool.h"
#include "xe_migrate.h"
#include "xe_pm.h"
#include "xe_sriov_printk.h"
#include "xe_sriov_vf.h"
#include "xe_sriov_vf_ccs.h"
#include "xe_sriov_vf_ccs_types.h"

/**
 * DOC: VF save/restore of compression Meta Data
 *
 * VF KMD registers two special contexts/LRCAs.
 *
 * Save Context/LRCA: contain necessary cmds+page table to trigger Meta data /
 * compression control surface (Aka CCS) save in regular System memory in VM.
 *
 * Restore Context/LRCA: contain necessary cmds+page table to trigger Meta data /
 * compression control surface (Aka CCS) Restore from regular System memory in
 * VM to corresponding CCS pool.
 *
 * Below diagram explain steps needed for VF save/Restore of compression Meta Data::
 *
 *    CCS Save    CCS Restore          VF KMD                          Guc       BCS
 *     LRCA        LRCA
 *      |           |                     |                              |         |
 *      |           |                     |                              |         |
 *      |     Create Save LRCA            |                              |         |
 *     [ ]<----------------------------- [ ]                             |         |
 *      |           |                     |                              |         |
 *      |           |                     |                              |         |
 *      |           |                     |       Register save LRCA     |         |
 *      |           |                     |           with Guc           |         |
 *      |           |                    [ ]--------------------------->[ ]        |
 *      |           |                     |                              |         |
 *      |           | Create restore LRCA |                              |         |
 *      |          [ ]<------------------[ ]                             |         |
 *      |           |                     |                              |         |
 *      |           |                     |       Register restore LRCA  |         |
 *      |           |                     |           with Guc           |         |
 *      |           |                    [ ]--------------------------->[ ]        |
 *      |           |                     |                              |         |
 *      |           |                     |                              |         |
 *      |           |                    [ ]-------------------------    |         |
 *      |           |                    [ ]  Allocate main memory.  |   |         |
 *      |           |                    [ ]  Allocate CCS memory.   |   |         |
 *      |           |                    [ ]  Update Main memory &   |   |         |
 *     [ ]<------------------------------[ ]  CCS pages PPGTT + BB   |   |         |
 *      |          [ ]<------------------[ ]  cmds to save & restore.|   |         |
 *      |           |                    [ ]<------------------------    |         |
 *      |           |                     |                              |         |
 *      |           |                     |                              |         |
 *      |           |                     |                              |         |
 *      :           :                     :                              :         :
 *      ---------------------------- VF Paused -------------------------------------
 *      |           |                     |                              |         |
 *      |           |                     |                              |         |
 *      |           |                     |                              |Schedule |
 *      |           |                     |                              |CCS Save |
 *      |           |                     |                              | LRCA    |
 *      |           |                     |                             [ ]------>[ ]
 *      |           |                     |                              |         |
 *      |           |                     |                              |         |
 *      |           |                     |                              |CCS save |
 *      |           |                     |                              |completed|
 *      |           |                     |                             [ ]<------[ ]
 *      |           |                     |                              |         |
 *      :           :                     :                              :         :
 *      ---------------------------- VM Migrated -----------------------------------
 *      |           |                     |                              |         |
 *      |           |                     |                              |         |
 *      :           :                     :                              :         :
 *      ---------------------------- VF Resumed ------------------------------------
 *      |           |                     |                              |         |
 *      |           |                     |                              |         |
 *      |           |                    [ ]--------------               |         |
 *      |           |                    [ ] Fix up GGTT  |              |         |
 *      |           |                    [ ]<-------------               |         |
 *      |           |                     |                              |         |
 *      |           |                     |                              |         |
 *      |           |                     |  Notify VF_RESFIX_DONE       |         |
 *      |           |                    [ ]--------------------------->[ ]        |
 *      |           |                     |                              |         |
 *      |           |                     |                              |Schedule |
 *      |           |                     |                              |CCS      |
 *      |           |                     |                              |Restore  |
 *      |           |                     |                              |LRCA     |
 *      |           |                     |                             [ ]------>[ ]
 *      |           |                     |                              |         |
 *      |           |                     |                              |         |
 *      |           |                     |                              |CCS      |
 *      |           |                     |                              |restore  |
 *      |           |                     |                              |completed|
 *      |           |                     |                             [ ]<------[ ]
 *      |           |                     |                              |         |
 *      |           |                     |                              |         |
 *      |           |                     |  VF_RESFIX_DONE complete     |         |
 *      |           |                     |       notification           |         |
 *      |           |                    [ ]<---------------------------[ ]        |
 *      |           |                     |                              |         |
 *      |           |                     |                              |         |
 *      :           :                     :                              :         :
 *      ------------------------- Continue VM restore ------------------------------
 */

static u64 get_ccs_bb_pool_size(struct xe_device *xe)
{
        u64 sys_mem_size, ccs_mem_size, ptes, bb_pool_size;
        struct sysinfo si;

        si_meminfo(&si);
        sys_mem_size = si.totalram * si.mem_unit;
        ccs_mem_size = div64_u64(sys_mem_size, NUM_BYTES_PER_CCS_BYTE(xe));
        ptes = DIV_ROUND_UP_ULL(sys_mem_size + ccs_mem_size, XE_PAGE_SIZE);

        /**
         * We need below BB size to hold PTE mappings and some DWs for copy
         * command. In reality, we need space for many copy commands. So, let
         * us allocate double the calculated size which is enough to holds GPU
         * instructions for the whole region.
         */
        bb_pool_size = ptes * sizeof(u32);

        return round_up(bb_pool_size * 2, SZ_1M);
}

static int alloc_bb_pool(struct xe_tile *tile, struct xe_sriov_vf_ccs_ctx *ctx)
{
        struct xe_mem_pool *pool;
        struct xe_device *xe = tile_to_xe(tile);
        u32 *pool_cpu_addr, *last_dw_addr;
        u64 bb_pool_size;
        int err;

        bb_pool_size = get_ccs_bb_pool_size(xe);
        xe_sriov_info(xe, "Allocating %s CCS BB pool size = %lldMB\n",
                      ctx->ctx_id ? "Restore" : "Save", bb_pool_size / SZ_1M);

        pool = xe_mem_pool_init(tile, bb_pool_size, sizeof(u32),
                                XE_MEM_POOL_BO_FLAG_INIT_SHADOW_COPY);
        if (IS_ERR(pool)) {
                xe_sriov_err(xe, "xe_mem_pool_init failed with error: %pe\n",
                             pool);
                err = PTR_ERR(pool);
                return err;
        }

        pool_cpu_addr = xe_mem_pool_cpu_addr(pool);
        memset(pool_cpu_addr, 0, bb_pool_size);

        last_dw_addr = pool_cpu_addr + (bb_pool_size / sizeof(u32)) - 1;
        *last_dw_addr = MI_BATCH_BUFFER_END;

        /**
         * Sync the main copy and shadow copy so that the shadow copy is
         * replica of main copy. We sync only BBs after init part. So, we
         * need to make sure the main pool and shadow copy are in sync after
         * this point. This is needed as GuC may read the BB commands from
         * shadow copy.
         */
        xe_mem_pool_sync(pool);

        ctx->mem.ccs_bb_pool = pool;
        return 0;
}

static void ccs_rw_update_ring(struct xe_sriov_vf_ccs_ctx *ctx)
{
        u64 addr = xe_mem_pool_gpu_addr(ctx->mem.ccs_bb_pool);
        struct xe_lrc *lrc = xe_exec_queue_lrc(ctx->mig_q);
        u32 dw[10], i = 0;

        /*
         * XXX: Save/restore fixes — for some reason, the GuC only accepts the
         * save/restore context if the LRC head pointer is zero. This is evident
         * from repeated VF migrations failing when the LRC head pointer is
         * non-zero.
         */
        lrc->ring.tail = 0;
        xe_lrc_set_ring_head(lrc, 0);

        dw[i++] = MI_ARB_ON_OFF | MI_ARB_ENABLE;
        dw[i++] = MI_BATCH_BUFFER_START | XE_INSTR_NUM_DW(3);
        dw[i++] = lower_32_bits(addr);
        dw[i++] = upper_32_bits(addr);
        dw[i++] = MI_NOOP;
        dw[i++] = MI_NOOP;

        xe_lrc_write_ring(lrc, dw, i * sizeof(u32));
        xe_lrc_set_ring_tail(lrc, lrc->ring.tail);
}

/**
 * xe_sriov_vf_ccs_rebase - Rebase GGTT addresses for CCS save / restore
 * @xe: the &xe_device.
 */
void xe_sriov_vf_ccs_rebase(struct xe_device *xe)
{
        enum xe_sriov_vf_ccs_rw_ctxs ctx_id;

        if (!IS_VF_CCS_READY(xe))
                return;

        for_each_ccs_rw_ctx(ctx_id) {
                struct xe_sriov_vf_ccs_ctx *ctx =
                        &xe->sriov.vf.ccs.contexts[ctx_id];

                ccs_rw_update_ring(ctx);
        }
}

static int register_save_restore_context(struct xe_sriov_vf_ccs_ctx *ctx)
{
        int ctx_type;

        switch (ctx->ctx_id) {
        case XE_SRIOV_VF_CCS_READ_CTX:
                ctx_type = GUC_CONTEXT_COMPRESSION_SAVE;
                break;
        case XE_SRIOV_VF_CCS_WRITE_CTX:
                ctx_type = GUC_CONTEXT_COMPRESSION_RESTORE;
                break;
        default:
                return -EINVAL;
        }

        xe_guc_register_vf_exec_queue(ctx->mig_q, ctx_type);
        return 0;
}

/**
 * xe_sriov_vf_ccs_register_context - Register read/write contexts with guc.
 * @xe: the &xe_device to register contexts on.
 *
 * This function registers read and write contexts with Guc. Re-registration
 * is needed whenever resuming from pm runtime suspend.
 *
 * Return: 0 on success. Negative error code on failure.
 */
int xe_sriov_vf_ccs_register_context(struct xe_device *xe)
{
        enum xe_sriov_vf_ccs_rw_ctxs ctx_id;
        struct xe_sriov_vf_ccs_ctx *ctx;
        int err;

        xe_assert(xe, IS_VF_CCS_READY(xe));

        for_each_ccs_rw_ctx(ctx_id) {
                ctx = &xe->sriov.vf.ccs.contexts[ctx_id];
                err = register_save_restore_context(ctx);
                if (err)
                        return err;
        }

        return err;
}

/*
 * Whether GuC requires CCS copy BBs for VF migration.
 * @xe: the &xe_device instance.
 *
 * Only selected platforms require VF KMD to maintain CCS copy BBs and linked LRCAs.
 *
 * Return: true if VF driver must participate in the CCS migration, false otherwise.
 */
static bool vf_migration_ccs_bb_needed(struct xe_device *xe)
{
        xe_assert(xe, IS_SRIOV_VF(xe));

        return !IS_DGFX(xe) && xe_device_has_flat_ccs(xe);
}

/*
 * Check for disable migration due to no CCS BBs support in GuC FW.
 * @xe: the &xe_device instance.
 *
 * Performs late disable of VF migration feature in case GuC FW cannot support it.
 *
 * Returns: True if VF migration with CCS BBs is supported, false otherwise.
 */
static bool vf_migration_ccs_bb_support_check(struct xe_device *xe)
{
        struct xe_gt *gt = xe_root_mmio_gt(xe);
        struct xe_uc_fw_version guc_version;

        xe_gt_sriov_vf_guc_versions(gt, NULL, &guc_version);
        if (MAKE_GUC_VER_STRUCT(guc_version) < MAKE_GUC_VER(1, 23, 0)) {
                xe_sriov_vf_migration_disable(xe,
                                              "CCS migration requires GuC ABI >= 1.23 but only %u.%u found",
                                              guc_version.major, guc_version.minor);
                return false;
        }

        return true;
}

static void xe_sriov_vf_ccs_fini(void *arg)
{
        struct xe_sriov_vf_ccs_ctx *ctx = arg;
        struct xe_lrc *lrc = xe_exec_queue_lrc(ctx->mig_q);

        /*
         * Make TAIL = HEAD in the ring so that no issues are seen if Guc
         * submits this context to HW on VF pause after unbinding device.
         */
        xe_lrc_set_ring_tail(lrc, xe_lrc_ring_head(lrc));
        xe_exec_queue_put(ctx->mig_q);
}

/**
 * xe_sriov_vf_ccs_init - Setup LRCA for save & restore.
 * @xe: the &xe_device to start recovery on
 *
 * This function shall be called only by VF. It initializes
 * LRCA and suballocator needed for CCS save & restore.
 *
 * Return: 0 on success. Negative error code on failure.
 */
int xe_sriov_vf_ccs_init(struct xe_device *xe)
{
        struct xe_tile *tile = xe_device_get_root_tile(xe);
        enum xe_sriov_vf_ccs_rw_ctxs ctx_id;
        struct xe_sriov_vf_ccs_ctx *ctx;
        struct xe_exec_queue *q;
        u32 flags;
        int err;

        xe_assert(xe, IS_SRIOV_VF(xe));

        if (!xe_sriov_vf_migration_supported(xe) ||
            !vf_migration_ccs_bb_needed(xe) ||
            !vf_migration_ccs_bb_support_check(xe))
                return 0;

        for_each_ccs_rw_ctx(ctx_id) {
                ctx = &xe->sriov.vf.ccs.contexts[ctx_id];
                ctx->ctx_id = ctx_id;

                flags = EXEC_QUEUE_FLAG_KERNEL |
                        EXEC_QUEUE_FLAG_PERMANENT |
                        EXEC_QUEUE_FLAG_MIGRATE;
                q = xe_exec_queue_create_bind(xe, tile, NULL, flags, 0);
                if (IS_ERR(q)) {
                        err = PTR_ERR(q);
                        goto err_ret;
                }
                ctx->mig_q = q;

                err = alloc_bb_pool(tile, ctx);
                if (err)
                        goto err_free_queue;

                ccs_rw_update_ring(ctx);

                err = register_save_restore_context(ctx);
                if (err)
                        goto err_free_queue;

                err = devm_add_action_or_reset(xe->drm.dev,
                                               xe_sriov_vf_ccs_fini,
                                               ctx);
                if (err)
                        goto err_ret;
        }

        xe->sriov.vf.ccs.initialized = 1;

        return 0;

err_free_queue:
        xe_exec_queue_put(q);

err_ret:
        return err;
}

#define XE_SRIOV_VF_CCS_RW_BB_ADDR_OFFSET       (2 * sizeof(u32))
void xe_sriov_vf_ccs_rw_update_bb_addr(struct xe_sriov_vf_ccs_ctx *ctx)
{
        u64 addr = xe_mem_pool_gpu_addr(ctx->mem.ccs_bb_pool);
        struct xe_lrc *lrc = xe_exec_queue_lrc(ctx->mig_q);
        struct xe_device *xe = gt_to_xe(ctx->mig_q->gt);

        xe_device_wmb(xe);
        xe_map_wr(xe, &lrc->bo->vmap, XE_SRIOV_VF_CCS_RW_BB_ADDR_OFFSET, u32, addr);
        xe_device_wmb(xe);
}

/**
 * xe_sriov_vf_ccs_attach_bo - Insert CCS read write commands in the BO.
 * @bo: the &buffer object to which batch buffer commands will be added.
 * @new_mem: the (not yet committed) destination resource @bo is being moved
 *          into; bo->ttm.resource is still the old resource at this point.
 *
 * This function shall be called only by VF. It inserts the PTEs and copy
 * command instructions in the BO by calling xe_migrate_ccs_rw_copy()
 * function.
 *
 * Returns: 0 if successful, negative error code on failure.
 */
int xe_sriov_vf_ccs_attach_bo(struct xe_bo *bo, struct ttm_resource *new_mem)
{
        struct xe_device *xe = xe_bo_device(bo);
        enum xe_sriov_vf_ccs_rw_ctxs ctx_id;
        struct xe_sriov_vf_ccs_ctx *ctx;
        struct xe_mem_pool_node *bb;
        struct xe_tile *tile;
        int err = 0;

        xe_assert(xe, IS_VF_CCS_READY(xe));

        tile = xe_device_get_root_tile(xe);

        for_each_ccs_rw_ctx(ctx_id) {
                bb = bo->bb_ccs[ctx_id];
                /* bb should be NULL here. Assert if not NULL */
                xe_assert(xe, !bb);

                ctx = &xe->sriov.vf.ccs.contexts[ctx_id];
                err = xe_migrate_ccs_rw_copy(tile, ctx->mig_q, bo, new_mem, ctx_id);
                if (err)
                        goto err_unwind;
        }
        return 0;

err_unwind:
        /*
         * Clean up any contexts already attached. Can't reuse
         * xe_sriov_vf_ccs_detach_bo() here as it requires both contexts
         * attached before cleaning up either one.
         */
        for_each_ccs_rw_ctx(ctx_id) {
                if (bo->bb_ccs[ctx_id])
                        xe_migrate_ccs_rw_copy_clear(bo, ctx_id, true);
        }
        return err;
}

/**
 * xe_sriov_vf_ccs_detach_bo - Remove CCS read write commands from the BO.
 * @bo: the &buffer object from which batch buffer commands will be removed.
 *
 * This function shall be called only by VF. It removes the PTEs and copy
 * command instructions from the BO. Make sure to update the BB with MI_NOOP
 * before freeing.
 *
 * Returns: 0 if successful.
 */
int xe_sriov_vf_ccs_detach_bo(struct xe_bo *bo)
{
        struct xe_device *xe = xe_bo_device(bo);
        enum xe_sriov_vf_ccs_rw_ctxs ctx_id;
        struct xe_mem_pool_node *bb;
        bool bound;
        int idx;

        xe_assert(xe, IS_VF_CCS_READY(xe));

        if (!xe_bo_has_valid_ccs_bb(bo))
                return 0;

        bound = drm_dev_enter(&xe->drm, &idx);

        for_each_ccs_rw_ctx(ctx_id) {
                bb = bo->bb_ccs[ctx_id];
                if (!bb)
                        continue;

                xe_migrate_ccs_rw_copy_clear(bo, ctx_id, bound);
        }

        if (bound)
                drm_dev_exit(idx);

        return 0;
}

/**
 * xe_sriov_vf_ccs_print - Print VF CCS details.
 * @xe: the &xe_device
 * @p: the &drm_printer
 *
 * This function is for VF use only.
 */
void xe_sriov_vf_ccs_print(struct xe_device *xe, struct drm_printer *p)
{
        enum xe_sriov_vf_ccs_rw_ctxs ctx_id;
        struct xe_mem_pool *bb_pool;

        if (!IS_VF_CCS_READY(xe))
                return;

        guard(xe_pm_runtime)(xe);
        for_each_ccs_rw_ctx(ctx_id) {
                bb_pool = xe->sriov.vf.ccs.contexts[ctx_id].mem.ccs_bb_pool;
                if (!bb_pool)
                        break;

                drm_printf(p, "ccs %s bb suballoc info\n", ctx_id ? "write" : "read");
                drm_printf(p, "-------------------------\n");
                xe_mem_pool_dump(bb_pool, p);
                drm_puts(p, "\n");
        }
}