/*
 * SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 * SPDX-License-Identifier: MIT
 *
 * Permission is hereby granted, free of charge, to any person obtaining a
 * copy of this software and associated documentation files (the "Software"),
 * to deal in the Software without restriction, including without limitation
 * the rights to use, copy, modify, merge, publish, distribute, sublicense,
 * and/or sell copies of the Software, and to permit persons to whom the
 * Software is furnished to do so, subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in
 * all copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
 * THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
 * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
 * DEALINGS IN THE SOFTWARE.
 */

#define NVOC_KERNEL_MIG_MANAGER_H_PRIVATE_ACCESS_ALLOWED

#include "kernel/gpu/mem_mgr/mem_mgr.h"
#include "kernel/gpu/mem_mgr/heap.h"
#include "gpu/mem_mgr/phys_mem_allocator/phys_mem_allocator.h"
#include "kernel/gpu/mig_mgr/kernel_mig_manager.h"
#include "kernel/gpu/fifo/kernel_fifo.h"

#include "kernel/core/locks.h"

/*!
 * @brief   Peforms checks to determine whether instancing can be enabled on
 *          this GPU, such as determining whether any partitionable engines are
 *          currently active.
 */
NV_STATUS
kmigmgrCreateGPUInstanceCheck_GB10B
(
    OBJGPU *pGpu,
    KernelMIGManager *pKernelMIGManager,
    NvBool bMemoryPartitioningNeeded
)
{
    Heap *pHeap;
    MemoryManager *pMemoryManager = GPU_GET_MEMORY_MANAGER(pGpu);
    KernelFifo *pKernelFifo = GPU_GET_KERNEL_FIFO(pGpu);
    RM_ENGINE_TYPE engines[RM_ENGINE_TYPE_LAST];
    NvU32 engineCount;
    NvU32 i;
    NvU64 largestFreeSize;
    NvU64 base;
    NvU64 unused;
    NV_RANGE freeBlock;
    NV_RANGE partitionableMemoryRange = memmgrGetMIGPartitionableMemoryRange(pGpu, pMemoryManager);

    // Ensure the engine DB is up-to-date
    NV_ASSERT_OK_OR_RETURN(gpuUpdateEngineTable(pGpu));

    // Store all engine tags of partitionable engines in the system
    engineCount = 0;
    for (i = 0; i < pGpu->engineDB.size; ++i)
    {
        if (kmigmgrIsEnginePartitionable(pGpu, pKernelMIGManager, pGpu->engineDB.pType[i]))
        {
            if (RM_ENGINE_TYPE_IS_GR(pGpu->engineDB.pType[i]) &&
                (RM_ENGINE_TYPE_GR_IDX(pGpu->engineDB.pType[i]) > 0))
            {
                //
                // This check is used during instance creation, prior to which
                // it is impossible to use GR1-7, so as an optimization, skip
                // checking for those.
                //
                continue;
            }

            engines[engineCount++] = pGpu->engineDB.pType[i];
        }
    }

    // Make sure there are no channels alive on any of these engines
    if (kfifoEngineListHasChannel(pGpu, pKernelFifo, engines, engineCount))
        return NV_ERR_STATE_IN_USE;

    pHeap = GPU_GET_HEAP(pGpu);
    if (!memmgrIsPmaInitialized(pMemoryManager))
    {
        NV_ASSERT_OK_OR_RETURN(
            heapInfo(pHeap, &unused, &unused, &unused, &base, &largestFreeSize));
    }
    else
    {
        pmaGetLargestFree(pHeap->pPmaObject, &largestFreeSize, &base, &unused);
    }

    // Make sure that no memory has been claimed from our partitionable range
    freeBlock = rangeMake(base, base + largestFreeSize - 1);
    if (!rangeContains(freeBlock, partitionableMemoryRange))
        return NV_ERR_STATE_IN_USE;

    return NV_OK;
}

/*!
 * @brief   Function to determine whether partition flag combinations are valid
 *          for this GPU
 */
NvBool
kmigmgrIsGPUInstanceCombinationValid_GB10B
(
    OBJGPU *pGpu,
    KernelMIGManager *pKernelMIGManager,
    NvU32 gpuInstanceFlag
)
{
    NvU32 computeSizeFlag = DRF_VAL(2080_CTRL_GPU, _PARTITION_FLAG, _COMPUTE_SIZE, gpuInstanceFlag);
    NvU32 smallestComputeSizeFlag;

    if (!kmigmgrIsGPUInstanceFlagValid_HAL(pGpu, pKernelMIGManager, gpuInstanceFlag))
    {
        return NV_FALSE;
    }

    smallestComputeSizeFlag = kmigmgrSmallestComputeProfileSize(pGpu, pKernelMIGManager);
    NV_CHECK_OR_RETURN(LEVEL_ERROR, smallestComputeSizeFlag != KMIGMGR_COMPUTE_SIZE_INVALID, NV_FALSE);

    // JPG_OFA profile is only available on the smallest available partition
    if (FLD_TEST_REF(NV2080_CTRL_GPU_PARTITION_FLAG_REQ_DEC_JPG_OFA, _ENABLE, gpuInstanceFlag))
    {
        if (computeSizeFlag != smallestComputeSizeFlag)
        {
            return NV_FALSE;
        }
    }

    //
    // MEMORY=_FULL/COMPUTE=_FULL implicitly includes OFA=_ENABLE so need to
    // report this profile separately.
    //
    if (FLD_TEST_REF(NV2080_CTRL_GPU_PARTITION_FLAG_REQ_DEC_JPG_OFA, _ENABLE, gpuInstanceFlag) &&
        FLD_TEST_REF(NV2080_CTRL_GPU_PARTITION_FLAG_MEMORY_SIZE, _FULL, gpuInstanceFlag) &&
        FLD_TEST_REF(NV2080_CTRL_GPU_PARTITION_FLAG_COMPUTE_SIZE, _FULL, gpuInstanceFlag))
    {
        return NV_FALSE;
    }

    return NV_TRUE;
}

/*!
 * @brief   Function to determine whether gpu instance flags are valid
 *          for this GPU
 *
 * @param[IN]   pGpu
 * @param[IN]   pKernelMIGManager
 * @param[IN]   gpuInstanceFlag       NV2080_CTRL_GPU_PARTITION_FLAG_*
 *
 * @return  Returns true if flags are valid
 */
NvBool
kmigmgrIsGPUInstanceFlagValid_GB10B
(
    OBJGPU *pGpu,
    KernelMIGManager *pKernelMIGManager,
    NvU32 gpuInstanceFlag
)
{
    NvU32 memSizeFlag = DRF_VAL(2080_CTRL_GPU, _PARTITION_FLAG,
                                _MEMORY_SIZE, gpuInstanceFlag);
    NvU32 computeSizeFlag = DRF_VAL(2080_CTRL_GPU, _PARTITION_FLAG,
                                    _COMPUTE_SIZE, gpuInstanceFlag);
    NvU32 gfxSizeFlag = DRF_VAL(2080_CTRL_GPU, _PARTITION_FLAG,
                                    _GFX_SIZE, gpuInstanceFlag);

    NV_CHECK_OR_RETURN(LEVEL_ERROR,
        kmigmgrIsGPUInstanceFlagLegal(pGpu, pKernelMIGManager, gpuInstanceFlag),
        NV_FALSE);

    if (!FLD_TEST_REF(NV2080_CTRL_GPU_PARTITION_FLAG_REQ_ALL_MEDIA, _DEFAULT, gpuInstanceFlag))
    {
        return NV_FALSE;
    }

    switch (memSizeFlag)
    {
        case NV2080_CTRL_GPU_PARTITION_FLAG_MEMORY_SIZE_FULL:
            break;
        default:
            return NV_FALSE;
    }

    switch (computeSizeFlag)
    {
        case NV2080_CTRL_GPU_PARTITION_FLAG_COMPUTE_SIZE_FULL:
        case NV2080_CTRL_GPU_PARTITION_FLAG_COMPUTE_SIZE_HALF:
        case NV2080_CTRL_GPU_PARTITION_FLAG_COMPUTE_SIZE_MINI_HALF:
            break;
        default:
            return NV_FALSE;
    }

    switch (gfxSizeFlag)
    {
        case NV2080_CTRL_GPU_PARTITION_FLAG_GFX_SIZE_FULL:
            if (!pKernelMIGManager->bIsSmgEnabled)
            {
                return NV_FALSE;
            }
            break;
        case NV2080_CTRL_GPU_PARTITION_FLAG_GFX_SIZE_NONE:
            break;
        default:
            return NV_FALSE;
    }

    return NV_TRUE;
}

/*!
 * @brief This function is introduced because GB10B has uneven Static GPCs for
 * HALF partitions(1 GPC(MIG 1):2 GPC(MIG 2)). With the current logic of the function
 * kmigmgrGetComputeProfileFromGpcCount_IMPL, a request for GPC count 2 will return a
 * FULL compute profile based on maxGPC calculation(3/2 == 1). For GB10B, If GPC count
 * is 3  a FULL profile is returned, for a GPC count of 2, HALF profile is returned and
 * for GPC count 1, a MINI_HALF profile is returned.
 *
 * @param[IN]   pGpu
 * @param[IN]   pKernelMIGManager
 * @param[IN]   gpcCount                GPC Count to look up the associated compute profile
 * @param[OUT]  pProfile                Pointer to  NV2080_CTRL_INTERNAL_MIGMGR_COMPUTE_PROFILE struct filled with
 *                                      a copy of the compute profile info associated with the GPC count
 */
NV_STATUS
kmigmgrGetComputeProfileFromGpcCount_GB10B
(
    OBJGPU *pGpu,
    KernelMIGManager *pKernelMIGManager,
    NvU32 gpcCount,
    NV2080_CTRL_INTERNAL_MIGMGR_COMPUTE_PROFILE *pProfile
)
{
    KernelGraphicsManager *pKernelGraphicsManager = GPU_GET_KERNEL_GRAPHICS_MANAGER(pGpu);
    const KERNEL_MIG_MANAGER_STATIC_INFO *pStaticInfo = kmigmgrGetStaticInfo(pGpu, pKernelMIGManager);
    NvU32 compSize = NV2080_CTRL_GPU_PARTITION_FLAG_COMPUTE_SIZE__SIZE;
    NvU32 maxGpc;
    NvU32 i;
    NvU32 actualMaxGpc = 0;

    NV_ASSERT_OR_RETURN(pProfile != NULL, NV_ERR_INVALID_ARGUMENT);
    NV_CHECK_OR_RETURN(LEVEL_ERROR, pStaticInfo != NULL, NV_ERR_OBJECT_NOT_FOUND);
    NV_CHECK_OR_RETURN(LEVEL_WARNING, pStaticInfo->pCIProfiles != NULL, NV_ERR_OBJECT_NOT_FOUND);

    maxGpc = kgrmgrGetLegacyKGraphicsStaticInfo(pGpu, pKernelGraphicsManager)->pGrInfo->infoList[NV2080_CTRL_GR_INFO_INDEX_MAX_PARTITIONABLE_GPCS].data;
    actualMaxGpc = maxGpc;

    for (i = 0; i < pStaticInfo->pCIProfiles->profileCount; i++)
    {
        if (pStaticInfo->pCIProfiles->profiles[i].computeSize == NV2080_CTRL_GPU_PARTITION_FLAG_COMPUTE_SIZE_FULL)
        {
            actualMaxGpc = pStaticInfo->pCIProfiles->profiles[i].gpcCount;
        }
    }

    if (gpcCount == 1)
    {
        compSize = NV2080_CTRL_GPU_PARTITION_FLAG_COMPUTE_SIZE_MINI_HALF;
    }
    else if ((gpcCount == ((maxGpc / 2) + 1)) && (actualMaxGpc == maxGpc))
    {
        compSize = NV2080_CTRL_GPU_PARTITION_FLAG_COMPUTE_SIZE_HALF;
    }
    else if (gpcCount == actualMaxGpc)
    {
        compSize = NV2080_CTRL_GPU_PARTITION_FLAG_COMPUTE_SIZE_FULL;
    }

    for (i = 0; i < pStaticInfo->pCIProfiles->profileCount; i++)
    {
        if (pStaticInfo->pCIProfiles->profiles[i].computeSize == compSize)
        {
            NV_PRINTF(LEVEL_INFO, "Found matching Compute Profile:%d for gpcCount=%d\n", compSize, gpcCount);
            portMemCopy(pProfile, sizeof(*pProfile), &pStaticInfo->pCIProfiles->profiles[i], sizeof(pStaticInfo->pCIProfiles->profiles[i]));
            return NV_OK;
        }
    }

    NV_PRINTF(LEVEL_INFO, "Found no Compute Profile for gpcCount=%d\n", gpcCount);
    return NV_ERR_OBJECT_NOT_FOUND;
}

/*!
 * @brief   Returns the span covered by the swizzId
 */
NV_RANGE
kmigmgrSwizzIdToSpan_GB10B
(
    OBJGPU *pGpu,
    KernelMIGManager *pKernelMIGManager,
    NvU32 swizzId
)
{
    KernelGraphicsManager *pKernelGraphicsManager = GPU_GET_KERNEL_GRAPHICS_MANAGER(pGpu);
    NV_RANGE ret;
    NvU8 spanLen;

    NV_ASSERT_OR_RETURN(kgrmgrGetLegacyKGraphicsStaticInfo(pGpu, pKernelGraphicsManager)->bInitialized, NV_RANGE_EMPTY);
    NV_ASSERT_OR_RETURN(kgrmgrGetLegacyKGraphicsStaticInfo(pGpu, pKernelGraphicsManager)->pGrInfo != NULL, NV_RANGE_EMPTY);

    spanLen = kgrmgrGetLegacyKGraphicsStaticInfo(pGpu, pKernelGraphicsManager)->pGrInfo->infoList[NV2080_CTRL_GR_INFO_INDEX_MAX_PARTITIONABLE_GPCS].data;

    switch (swizzId)
    {
        case 0:
            ret = rangeMake(0, spanLen - 1);
            break;
        case 1:
            ret = rangeMake(0, ((spanLen - 1)/2));
            break;
        case 2:
            ret = rangeMake((spanLen - 1), (spanLen - 1));
            break;
        default:
            NV_PRINTF(LEVEL_ERROR, "Unsupported swizzid 0x%x\n", swizzId);
            ret = NV_RANGE_EMPTY;
            break;
    }

    return ret;
}

/*!
 * @brief   Returns the GR span covered by the swizzId
 */
NV_RANGE
kmigmgrSwizzIdToGrSpan_GB10B
(
    OBJGPU *pGpu,
    KernelMIGManager *pKernelMIGManager,
    NvU32 swizzId
)
{
    KernelGraphicsManager *pKernelGraphicsManager = GPU_GET_KERNEL_GRAPHICS_MANAGER(pGpu);
    NV_RANGE ret;
    NvU8 spanLen;

    NV_ASSERT_OR_RETURN(kgrmgrGetLegacyKGraphicsStaticInfo(pGpu, pKernelGraphicsManager)->bInitialized, NV_RANGE_EMPTY);
    NV_ASSERT_OR_RETURN(kgrmgrGetLegacyKGraphicsStaticInfo(pGpu, pKernelGraphicsManager)->pGrInfo != NULL, NV_RANGE_EMPTY);

    spanLen = kgrmgrGetLegacyKGraphicsStaticInfo(pGpu, pKernelGraphicsManager)->pGrInfo->infoList[NV2080_CTRL_GR_INFO_INDEX_MAX_MIG_ENGINES].data;

    switch (swizzId)
    {
        case 0:
            ret = rangeMake(0, spanLen - 1);
            break;
        case 1:
            ret = rangeMake(0, (spanLen/2) - 1);
            break;
        case 2:
            ret = rangeMake(spanLen/2, spanLen - 1);
            break;
        default:
            NV_PRINTF(LEVEL_ERROR, "Unsupported swizzid 0x%x\n", swizzId);
            ret = NV_RANGE_EMPTY;
            break;
    }

    return ret;
}

/*!
 * @brief  Patch the CI capacity based on the GR count
 */
NV_STATUS
kmigmgrPatchCICapacity_GB10B
(
    OBJGPU *pGpu,
    KernelMIGManager *pKernelMIGManager,
    const NV2080_CTRL_INTERNAL_MIGMGR_PROFILE_INFO *pProfile,
    NVC637_CTRL_EXEC_PARTITIONS_GET_PROFILE_CAPACITY_PARAMS *pParams,
    NvU32 veidSlotCount
)
{
    NV_ASSERT_OR_RETURN(pProfile != NULL, NV_ERR_INVALID_STATE);
    NV_ASSERT_OR_RETURN(pParams != NULL, NV_ERR_INVALID_STATE);

    /*
     * Max syspipe for any GI is 1. CI capacity is calculated based on
     * VEID mask which could result in the following condition passing
     * for a FULL GI with HALF/MINI_HALF CI combination.
     * If a CI is already created in the GI, veidSlotCount will have a
     * value less than the totalSpansCount which indicates that the only
     * available syspipe is already occupied. Clear the profileCount and
     * availableSpansCount in that case.
     */
    if (pParams->totalSpansCount > pProfile->grCount)
    {
        NV_PRINTF(LEVEL_INFO, "Total Span count:0x%x more than GR Count:0x%x\n",
                  pParams->totalSpansCount, pProfile->grCount);
        if (veidSlotCount < pParams->totalSpansCount)
        {
            NV_PRINTF(LEVEL_INFO, "Clear profile and available spans count\n");
            pParams->profileCount = 0;
            pParams->availableSpansCount = 0;
        }
    }
    return NV_OK;
}

/*
 * @brief   Function to configure PM settings for current MIG mode
 *
 * This workaround handles configuring dynamic PM for the current MIG
 * mode.  When MIG is enabled, then dynamic PM is disabled and vice versa.
 *
 * @param[IN]   pGpu
 * @param[IN]   pMIGManager
 */
NV_STATUS
kmigmgrConfigurePMSettings_GB10B
(
    OBJGPU *pGpu,
    KernelMIGManager *pKernelMIGManager
)
{
    NvBool bDisablePM = NV_FALSE;
    NvBool bEnablePM = NV_FALSE;

    NV_ASSERT_OR_RETURN(rmapiLockIsWriteOwner() && rmGpuLockIsOwner(),
                        NV_ERR_INVALID_LOCK_STATE);

    //
    // Nothing prevents redundant calls to enable/disable MIG from
    // userspace.  To avoid repeat (and unnecessary) updates to the
    // dynamic PM refcnt we only do them when the MIG mode actually
    // changes.
    //
    bDisablePM = (pKernelMIGManager->bMIGEnabled &&
                  !pKernelMIGManager->bIsDynamicPMDisabled);
    bEnablePM = (!pKernelMIGManager->bMIGEnabled &&
                  pKernelMIGManager->bIsDynamicPMDisabled);

    if (bDisablePM)
    {
        pKernelMIGManager->bIsDynamicPMDisabled = NV_TRUE;
    }
    else if (bEnablePM)
    {
        pKernelMIGManager->bIsDynamicPMDisabled = NV_FALSE;
    }

    //
    // We need to drop (and reacquire) the GPU locks around the dynamic PM
    // refcnt calls.
    //
    rmGpuLocksRelease(GPUS_LOCK_FLAGS_NONE, NULL);

    if (bDisablePM)
    {
        osRefGpuAccessNeeded(pGpu->pOsGpuInfo);
    }
    else if (bEnablePM)
    {
        osUnrefGpuAccessNeeded(pGpu->pOsGpuInfo);
    }

    NV_ASSERT_OK_OR_RETURN(rmGpuLocksAcquire(GPUS_LOCK_FLAGS_NONE,
                                             RM_LOCK_MODULES_NONE));

    return NV_OK;
}
