/*
 * SPDX-FileCopyrightText: Copyright (c) 2012-2025 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.
 */

#include "ctrl/ctrla081.h"
#include "class/clc637.h" // for AMPERE_SMC_PARTITION_REF
#include "virtualization/kernel_hostvgpudeviceapi.h"

#include "core/system.h"
#include "os/os.h"
#include "rmapi/control.h"
#include "virtualization/hypervisor/hypervisor.h"
#include "virtualization/kernel_vgpu_mgr.h"
#include "gpu/device/device.h"
#include "vgpu/vgpu_version.h"
#include "vgpu/vgpu_events.h"
#include "rmapi/rmapi_utils.h"
#include "rmapi/rs_utils.h"
#include "gpu/gpu.h"
#include "gpu/mem_sys/kern_mem_sys.h"
#include "gpu/mem_mgr/mem_mgr.h"
#include "gpu/mem_mgr/heap.h"
#include "gpu/mem_mgr/phys_mem_allocator/phys_mem_allocator.h"
#include "gpu/fifo/kernel_fifo.h"
#include "containers/eheap_old.h"
#include "kernel/gpu/mig_mgr/kernel_mig_manager.h"
#include "nvdevid.h"
#include "nvmisc.h"

#define __VGPU_ALIAS_PGPU_LIST__
#include "g_vgpu_resman_specific.h" // _get_chip_id_for_alias_pgpu
#undef __VGPU_ALIAS_PGPU_LIST__

#define MAX_NVU32_TO_CONVERTED_STR_LEN   0x8

#define DEVID_ENCODED_VALUE_SIZE          0x4
#define SUBDEVID_ENCODED_VALUE_SIZE       0x4
#define PGPU_CAP_ENCODED_VALUE_SIZE       0x8

#define PGPU_CAP_ECC_ON                  (1 << 0)
#define PGPU_CAP_SRIOV_ON                (1 << 1)

#define VGPU_LOWER_FB_PROFILE_SIZE_1_GB (1024 * 1024 * 1024)

// Mapping of vGPU MIG profiles types to corresponding RM partiotionFlag
#define GENERATE_vgpuSmcTypeIdMappings
#include "g_vgpu_resman_specific.h"
#undef GENERATE_vgpuSmcTypeIdMappings
//
// ODB functions
//

NV_STATUS
kvgpumgrConstruct_IMPL(KernelVgpuMgr *pKernelVgpuMgr)
{
    NV_PRINTF(LEVEL_INFO, "Enter function\n");

    pKernelVgpuMgr->pHeap = NULL;

    /* Default vGPU version is set to the host supported vGPU version range */
    pKernelVgpuMgr->user_min_supported_version
        = GRIDSW_VERSION_EXTERNAL(NV_VGPU_MIN_SUPPORTED_GRIDSW_VERSION_EXTERNAL_MAJOR,
                                  NV_VGPU_MIN_SUPPORTED_GRIDSW_VERSION_EXTERNAL_MINOR);
    pKernelVgpuMgr->user_max_supported_version
        = GRIDSW_VERSION_EXTERNAL(NV_VGPU_MAX_SUPPORTED_GRIDSW_VERSION_EXTERNAL_MAJOR,
                                  NV_VGPU_MAX_SUPPORTED_GRIDSW_VERSION_EXTERNAL_MINOR);

    listInit(&(pKernelVgpuMgr->listVgpuTypeHead), portMemAllocatorGetGlobalNonPaged());
    listInit(&(pKernelVgpuMgr->listVgpuGuestHead), portMemAllocatorGetGlobalNonPaged());
    listInit(&(pKernelVgpuMgr->listRequestVgpuHead), portMemAllocatorGetGlobalNonPaged());

    return NV_OK;
}

void
kvgpumgrDestruct_IMPL(KernelVgpuMgr *pKernelVgpuMgr)
{
    if (pKernelVgpuMgr->pHeap != NULL)
    {
        pKernelVgpuMgr->pHeap->eheapDestruct(pKernelVgpuMgr->pHeap);
        portMemFree(pKernelVgpuMgr->pHeap);
    }

    listDestroy(&(pKernelVgpuMgr->listRequestVgpuHead));
    listDestroy(&(pKernelVgpuMgr->listVgpuTypeHead));
    listDestroy(&(pKernelVgpuMgr->listVgpuGuestHead));
}

//
// Get max instance for a vgpu profile.
//
NV_STATUS
kvgpumgrGetMaxInstanceOfVgpu(NvU32 vgpuTypeId, NvU32 *maxInstanceVgpu)
{
    VGPU_TYPE  *vgpuTypeInfo;
    NV_STATUS   status;

    if (maxInstanceVgpu == NULL)
    {
        return NV_ERR_INVALID_ARGUMENT;
    }

    status = kvgpumgrGetVgpuTypeInfo(vgpuTypeId, &vgpuTypeInfo);
    if (status == NV_OK)
    {
        *maxInstanceVgpu = vgpuTypeInfo->maxInstance;
    }

    return status;
}

//
// vGPU Manager functions
//


static NV_STATUS
_kvgpumgrUpdateHeterogeneousCreatablePlacementIds(OBJGPU *pGpu, NvU16 placementId,
                                                  NvU32 vgpuTypeId, NvBool isAlloc,
                                                  NvU32 swizzId);

static NV_STATUS
_kvgpumgrUpdateHomogeneousCreatablePlacementIds(OBJGPU *pGpu, NvU16 placementId,
                                                NvU32 vgpuTypeId, NvBool isAlloc,
                                                NvU32 swizzId);

static NV_STATUS _kvgpuMgrSetHeterogeneousResourcesPerGI(OBJGPU *pGpu, NvU32 swizzId, KERNEL_PHYS_GPU_INFO *pPgpuInfo);

static NV_STATUS _kvgpumgrSetHomogeneousResourcesPerGI(OBJGPU *pGpu, NvU32 swizzId, KERNEL_PHYS_GPU_INFO *pPgpuInfo);

NV_STATUS
kvgpumgrGetPgpuIndex(KernelVgpuMgr *pKernelVgpuMgr, NvU32 gpuPciId, NvU32* index)
{
    NvU32 i;
    for (i = 0; i < NV_ARRAY_ELEMENTS(pKernelVgpuMgr->pgpuInfo); i++)
    {
        if ((pKernelVgpuMgr->pgpuInfo[i].gpuPciId == gpuPciId) &&
            (pKernelVgpuMgr->pgpuInfo[i].isAttached == NV_TRUE))
        {
            *index = i;
            return NV_OK;
        }
    }
    return NV_ERR_OBJECT_NOT_FOUND;
}

NvBool
kvgpumgrIsHeterogeneousVgpuTypeSupported(void)
{
    /* This support is currently limited to VMware, Device VM and KVM */
    return ((osIsVgpuVfioPresent() == NV_OK) ||
            (osIsVgpuDeviceVmPresent() == NV_OK));
}

NV_STATUS
_kvgpumgrCheckPartitionFlag(OBJGPU *pGpu, NvU32 swizzId, NvU32 vgpuTypeId, NvBool *bMatch)
{
    if (IS_MIG_IN_USE(pGpu) && swizzId != KMIGMGR_SWIZZID_INVALID)
    {
        KernelMIGManager *pKernelMIGManager = GPU_GET_KERNEL_MIG_MANAGER(pGpu);
        KERNEL_MIG_GPU_INSTANCE *pKernelMIGGpuInstance;
        NvU32 partitionFlag;

        NV_ASSERT_OK_OR_RETURN(kmigmgrGetGPUInstanceInfo(pGpu, pKernelMIGManager, swizzId, &pKernelMIGGpuInstance));

        kvgpumgrGetPartitionFlag(vgpuTypeId, &partitionFlag);

        *bMatch = (partitionFlag == pKernelMIGGpuInstance->partitionFlag) ? NV_TRUE : NV_FALSE;

        return NV_OK;
    }
    else
    {
        return NV_ERR_NOT_SUPPORTED;
    }
}

/*
 * This function should be called to check at runtime whether homogeneous
 * vGPU placement mode is supported / enabled or not. Homogeneous vGPU placement
 * mode is enabled in mutual exclusion with heterogeneous mode on the Ada+
 * vGPU profiles.
 */
NV_STATUS
kvgpumgrCheckHomogeneousPlacementSupported(OBJGPU *pGpu, NvU32 swizzId)
{
    OBJSYS *pSys = SYS_GET_INSTANCE();
    KernelVgpuMgr *pKernelVgpuMgr = SYS_GET_KERNEL_VGPUMGR(pSys);
    KERNEL_PHYS_GPU_INFO *pPgpuInfo;
    NvU32 index;
    NvBool bHeterogeneousModeEnabled;

    NV_CHECK_OR_RETURN(LEVEL_ERROR, pGpu != NULL, NV_ERR_INVALID_ARGUMENT);

    NV_CHECK_OK_OR_RETURN(LEVEL_ERROR,
            kvgpumgrGetPgpuIndex(pKernelVgpuMgr, pGpu->gpuId, &index));

    pPgpuInfo = &(pKernelVgpuMgr->pgpuInfo[index]);

    NV_CHECK_OK_OR_RETURN(LEVEL_ERROR, kvgpuMgrGetHeterogeneousMode(pGpu, swizzId, &bHeterogeneousModeEnabled));

    if (!bHeterogeneousModeEnabled &&
        (pPgpuInfo->homogeneousPlacementSupported))
    {
        return NV_OK;
    }

    return NV_ERR_NOT_SUPPORTED;
}

NvBool
kvgpumgrIsVgpuWarmUpdateSupported(void)
{
    /*This support is currently limited to VMware and KVM*/
    return ((osIsVgpuVfioPresent() == NV_OK) ||
            (osIsVgpuDeviceVmPresent() == NV_OK));
}

NV_STATUS
kvgpumgrSetVgpuType(OBJGPU *pGpu, KERNEL_PHYS_GPU_INFO *pPhysGpuInfo, NvU32 vgpuTypeId, NvU32 swizzId)
{
    NvU32 i;

    /*
     * If heterogeneous timesliced mode is enabled, supportedTypeIds is already
     * prepared when creatable placementIds array is updated in
     * _kvgpumgrUpdateHeterogeneousCreatablePlacementIds() function
     */
    if (pGpu != NULL)
    {
        NvBool bHeterogeneousModeEnabled;

        NV_CHECK_OK_OR_RETURN(LEVEL_ERROR, kvgpuMgrGetHeterogeneousMode(pGpu, swizzId, &bHeterogeneousModeEnabled));

        if (bHeterogeneousModeEnabled)
            return NV_OK;
    }

    for (i = 0; i < MAX_VGPU_TYPES_PER_PGPU; i++)
        pPhysGpuInfo->supportedTypeIds[i] = NVA081_CTRL_VGPU_CONFIG_INVALID_TYPE;

    if (vgpuTypeId == NVA081_CTRL_VGPU_CONFIG_INVALID_TYPE)
    {
        for (i = 0; i < pPhysGpuInfo->numVgpuTypes; i++)
            pPhysGpuInfo->supportedTypeIds[i] = pPhysGpuInfo->vgpuTypes[i]->vgpuTypeId;

        return NV_OK;
    }

    if (kvgpumgrIsHeterogeneousVgpuTypeSupported())
    {
        VGPU_TYPE *vgpuTypeInfo = NULL;
        NV_STATUS  status       = kvgpumgrGetVgpuTypeInfo(vgpuTypeId, &vgpuTypeInfo);
        NvU32      numSupported = 0;

        if (status != NV_OK)
            return status;

        NvU64 totalFbAllocation = vgpuTypeInfo->profileSize;

        for (i = 0; i < pPhysGpuInfo->numVgpuTypes; i++)
        {
            if (pPhysGpuInfo->vgpuTypes[i]->profileSize == totalFbAllocation)
            {
                pPhysGpuInfo->supportedTypeIds[numSupported] = pPhysGpuInfo->vgpuTypes[i]->vgpuTypeId;
                numSupported++;
            }
        }
    }
    else
    {
        pPhysGpuInfo->supportedTypeIds[0] = vgpuTypeId;
    }

    return NV_OK;
}

// non-MIG vgpu types
static NV_STATUS
_kvgpumgrCheckVgpuTypeSupported(OBJGPU *pGpu, KERNEL_PHYS_GPU_INFO *pPhysGpuInfo,
                                VGPU_TYPE *vgpuTypeInfo, NvU32 existingVgpus)
{
    NvU32 i;

    if (existingVgpus == 0)
        return NV_OK;

    for (i = 0; i < pPhysGpuInfo->numVgpuTypes; i++)
    {
        if (pPhysGpuInfo->supportedTypeIds[i] == vgpuTypeInfo->vgpuTypeId)
            break;
        if (pPhysGpuInfo->supportedTypeIds[i] == NVA081_CTRL_VGPU_CONFIG_INVALID_TYPE)
            return NV_ERR_NOT_SUPPORTED;
    }

    /*
     * In heterogeneous vGPU mode, creatable vGPU type is not dependent on a vGPU types
     * maxInstance. If the type is present in supportedTypeIds (populated based on placementIds),
     * it should suffice
     */
    if (!pGpu->getProperty(pGpu, PDB_PROP_GPU_IS_VGPU_HETEROGENEOUS_MODE))
    {
        if (existingVgpus >= vgpuTypeInfo->maxInstance)
            return NV_ERR_INSUFFICIENT_RESOURCES;
    }

    return NV_OK;
}

NV_STATUS
// non-MIG vgpu types
kvgpumgrCheckVgpuTypeCreatable(OBJGPU *pGpu, KERNEL_PHYS_GPU_INFO *pPhysGpuInfo, VGPU_TYPE *vgpuTypeInfo)
{
    if (osIsVgpuVfioPresent() == NV_OK)
        return _kvgpumgrCheckVgpuTypeSupported(pGpu, pPhysGpuInfo, vgpuTypeInfo, pPhysGpuInfo->numCreatedVgpu);
    else
        return _kvgpumgrCheckVgpuTypeSupported(pGpu, pPhysGpuInfo, vgpuTypeInfo, pPhysGpuInfo->numActiveVgpu);
}

static NV_STATUS
// non-MIG vgpu types
_kvgpumgrCheckVgpuTypeActivatable(OBJGPU *pGpu, KERNEL_PHYS_GPU_INFO *pPhysGpuInfo, VGPU_TYPE *vgpuTypeInfo)
{
    return _kvgpumgrCheckVgpuTypeSupported(pGpu, pPhysGpuInfo, vgpuTypeInfo, pPhysGpuInfo->numActiveVgpu);
}

static NV_STATUS
_addVgpuTypeToCreatableType(NvU32 vgpuTypeId,
                             NvU32 *numVgpuTypes,
                             NvU32 *vgpuTypes,
                             KERNEL_PHYS_GPU_INFO *pgpuInfo,
                             KERNEL_MIG_GPU_INSTANCE *pKernelMIGGpuInstance )
{
    NvU32 j = 0;
    NV_STATUS rmStatus = NV_OK;
    // Check if we have already included this vGPU type ID.
    // If not, then add vgpuTypeID to out vgpuTypes array.
    if (vgpuTypeId != NVA081_CTRL_VGPU_CONFIG_INVALID_TYPE)
    {
        for(j = 0; j < *numVgpuTypes; j++)
        {
            if (vgpuTypes[j] == vgpuTypeId)
            {
                break;
            }
        }

        // vgpuTypeId not included in vgputypes array yet, include it.
        if (j == *numVgpuTypes)
        {
            vgpuTypes[*numVgpuTypes] = vgpuTypeId;
            (*numVgpuTypes)++;
        }
    }
    return rmStatus;
}

static NV_STATUS
_kvgpumgrGetCreatedVgpuOnGpuInstance(OBJGPU *pGpu, NvU32 swizzId, NvU32 *vgpuType)
{
    OBJSYS *pSys = SYS_GET_INSTANCE();
    KernelVgpuMgr *pKernelVgpuMgr = SYS_GET_KERNEL_VGPUMGR(pSys);
    NvU32 pgpuIndex;
    KERNEL_PHYS_GPU_INFO *pPgpuInfo;
    REQUEST_VGPU_INFO_NODE *pRequestVgpu = NULL;

    NV_CHECK_OK_OR_RETURN(LEVEL_ERROR, kvgpumgrGetPgpuIndex(pKernelVgpuMgr, pGpu->gpuId, &pgpuIndex));
    pPgpuInfo = &pKernelVgpuMgr->pgpuInfo[pgpuIndex];

    if (osIsVgpuVfioPresent() == NV_OK)
    {
        for (pRequestVgpu = listHead(&pKernelVgpuMgr->listRequestVgpuHead);
             pRequestVgpu != NULL;
             pRequestVgpu = listNext(&pKernelVgpuMgr->listRequestVgpuHead, pRequestVgpu))
        {
            /* Check for an existing Vgpu on the same physical gpu and same GI */
            if ((pRequestVgpu->gpuPciId == pGpu->gpuId) && (pRequestVgpu->swizzId == swizzId))
            {
                *vgpuType = pRequestVgpu->vgpuTypeId;
                return NV_OK;
            }
        }
    } else {
        KERNEL_HOST_VGPU_DEVICE *pKernelHostVgpuDevice;

        for (pKernelHostVgpuDevice = listHead(&(pPgpuInfo->listHostVgpuDeviceHead));
             pKernelHostVgpuDevice != NULL;
             pKernelHostVgpuDevice = listNext(&(pPgpuInfo->listHostVgpuDeviceHead), pKernelHostVgpuDevice))
        {
            if (pKernelHostVgpuDevice->swizzId == swizzId)
            {
                *vgpuType = pKernelHostVgpuDevice->vgpuType;
                return NV_OK;
            }
        }
    }
    return NV_ERR_OBJECT_NOT_FOUND;
}

NV_STATUS
kvgpumgrGetCreatableVgpuTypes(OBJGPU *pGpu, KernelVgpuMgr *pKernelVgpuMgr, NvU32 pgpuIndex,
                              NvU32 gpuInstanceId, NvU32* numVgpuTypes, NvU32* vgpuTypes)
{
    NvU32 i                        = 0;
    KERNEL_PHYS_GPU_INFO *pgpuInfo = &pKernelVgpuMgr->pgpuInfo[pgpuIndex];

    *numVgpuTypes = 0;

    if (IS_MIG_ENABLED(pGpu))
    {
        KernelMIGManager *pKernelMIGManager = GPU_GET_KERNEL_MIG_MANAGER(pGpu);
        NvU64 swizzIdInUseMask = 0;
        NvU32 id;

        // No GPU instances created, none of the profiles are creatable. 
        if (!IS_MIG_IN_USE(pGpu))
            return NV_OK;

        swizzIdInUseMask = kmigmgrGetSwizzIdInUseMask(pGpu, pKernelMIGManager);

        for (id = 0; id < KMIGMGR_MAX_GPU_SWIZZID; id++)
        {
            KERNEL_MIG_GPU_INSTANCE *pKernelMIGGpuInstance;
            NvU32 j;
            VGPU_TYPE *pVgpuTypeInfo;
            VGPU_TYPE *existingVgpuTypeInfo = NULL;
            NvBool bHeterogeneousModeEnabled;
            NvU32 *pSupportedTypeId;
            NV_STATUS rmStatus = NV_OK;

            if (!(NVBIT64(id) & swizzIdInUseMask))
                 continue;

            //If specified GI is present, ignore other GIs
            if ((gpuInstanceId != KMIGMGR_SWIZZID_INVALID) && (gpuInstanceId != id))
                 continue;

            NV_ASSERT_OK_OR_RETURN(
                kmigmgrGetGPUInstanceInfo(pGpu, pKernelMIGManager, id, &pKernelMIGGpuInstance));
            rmStatus = kvgpuMgrGetHeterogeneousMode(pGpu,
                                                    id,
                                                    &bHeterogeneousModeEnabled);
            // Skip out-of-range GI and return success with creatable vGPU count as 0
            if (rmStatus == NV_ERR_OUT_OF_RANGE)
                continue;

            if (bHeterogeneousModeEnabled)
            {
                pSupportedTypeId = kvgpuMgrGetVgpuCreatableTypeIdFromSwizzId(pGpu, pgpuInfo, id);

                for (i = 0; i < pgpuInfo->numVgpuTypes; i++)
                {
                    if (pSupportedTypeId[i] == NVA081_CTRL_VGPU_CONFIG_INVALID_TYPE)
                        break;
                    _addVgpuTypeToCreatableType(pSupportedTypeId[i],
                                                numVgpuTypes, vgpuTypes, pgpuInfo, pKernelMIGGpuInstance);
                }
                continue;
            }
            for (j = 0; j < NV_ARRAY_ELEMENTS(vgpuSmcTypeIdMappings); j++)
            {

                //Check for Ids with same partition flag on same DevID (GPU) 
                if ((vgpuSmcTypeIdMappings[j].devId == pGpu->idInfo.PCIDeviceID) &&
                    (vgpuSmcTypeIdMappings[j].partitionFlag == pKernelMIGGpuInstance->partitionFlag))
                {
                    if (pgpuInfo->assignedSwizzIdVgpuCount[id] == 0)
                    {
                        // check to make sure the profile is supported
                        for (i = 0; i < pgpuInfo->numVgpuTypes; i++)
                        {
                            if (pgpuInfo->vgpuTypes[i] == NULL)
                                continue;

                            if (!kvgpumgrIsMigTimeslicingModeEnabled(pGpu) && (pgpuInfo->vgpuTypes[i]->maxInstancePerGI != 1))
                                continue;

                            if (pgpuInfo->vgpuTypes[i]->vgpuTypeId == vgpuSmcTypeIdMappings[j].vgpuTypeId)
                            {
                                _addVgpuTypeToCreatableType(vgpuSmcTypeIdMappings[j].vgpuTypeId,
                                                    numVgpuTypes, vgpuTypes, pgpuInfo, pKernelMIGGpuInstance);
                                break;
                            }
                        }

                    }
                    else
                    {
                        if (pgpuInfo->migTimeslicingModeEnabled)
                        {
                            NvU32 vgpuTypeId;
                            pVgpuTypeInfo = NULL;
                            // check to make sure the profile is supported
                            for (i = 0; i < pgpuInfo->numVgpuTypes; i++)
                            {
                                if (pgpuInfo->vgpuTypes[i] == NULL)
                                    continue;
                                if (pgpuInfo->vgpuTypes[i]->vgpuTypeId == vgpuSmcTypeIdMappings[j].vgpuTypeId)
                                {
                                    pVgpuTypeInfo = pgpuInfo->vgpuTypes[i];
                                }
                            }
                            // If the vgpu type id is not supported, skip and search the next one
                            if (pVgpuTypeInfo == NULL)
                                continue;

                            if (_kvgpumgrGetCreatedVgpuOnGpuInstance(pGpu, id, &vgpuTypeId) != NV_OK)
                                continue;

                            NV_ASSERT_OK_OR_RETURN(kvgpumgrGetVgpuTypeInfo(vgpuTypeId,
                                                    &existingVgpuTypeInfo));

                            // Check if maximum allowed number of vGPUs for this type has already been created
                            if (pgpuInfo->assignedSwizzIdVgpuCount[id] >= existingVgpuTypeInfo->maxInstancePerGI)
                                continue;
                            
                            if(kvgpumgrIsHeterogeneousVgpuTypeSupported() == NV_TRUE)
                            {
                                if (existingVgpuTypeInfo->profileSize == pVgpuTypeInfo->profileSize)
                                    _addVgpuTypeToCreatableType(vgpuSmcTypeIdMappings[j].vgpuTypeId, numVgpuTypes, 
                                                                vgpuTypes, pgpuInfo, pKernelMIGGpuInstance);
                            }
                            else
                            {
                                if (existingVgpuTypeInfo->vgpuTypeId == pVgpuTypeInfo->vgpuTypeId)
                                     _addVgpuTypeToCreatableType(vgpuSmcTypeIdMappings[j].vgpuTypeId, numVgpuTypes, 
                                                                 vgpuTypes, pgpuInfo, pKernelMIGGpuInstance);
                            }
                        }
                    }
                }
            }
        }
    }
    else
    {
        *numVgpuTypes = 0;
        for (i = 0; i < pgpuInfo->numVgpuTypes; i++)
        {
            // Consider only non-MIG (gpuInstanceSize == 0) profiles
            if (!pgpuInfo->vgpuTypes[i]->gpuInstanceSize &&
                (kvgpumgrCheckVgpuTypeCreatable(pGpu, pgpuInfo, pgpuInfo->vgpuTypes[i]) == NV_OK))
            {
                vgpuTypes[*numVgpuTypes] = pgpuInfo->vgpuTypes[i]->vgpuTypeId;
                (*numVgpuTypes)++;
            }
        }
    }

    return NV_OK;
}

static NV_STATUS
_kvgpumgrCreateVgpuType(NVA081_CTRL_VGPU_INFO *pVgpuInfo, VGPU_TYPE **ppVgpuType)
{
    OBJSYS *pSys = SYS_GET_INSTANCE();
    KernelVgpuMgr *pKernelVgpuMgr = SYS_GET_KERNEL_VGPUMGR(pSys);
    VGPU_TYPE *pVgpuTypeNode;

    pVgpuTypeNode = listAppendNew(&(pKernelVgpuMgr->listVgpuTypeHead));
    if (pVgpuTypeNode == NULL)
        return NV_ERR_NO_MEMORY;

    vgpuMgrFillVgpuType(pVgpuInfo, pVgpuTypeNode);

    *ppVgpuType = pVgpuTypeNode;

    return NV_OK;
}

NV_STATUS
kvgpumgrGetVgpuTypeInfo(NvU32 vgpuTypeId, VGPU_TYPE **vgpuType)
{
    OBJSYS *pSys = SYS_GET_INSTANCE();
    KernelVgpuMgr *pKernelVgpuMgr = SYS_GET_KERNEL_VGPUMGR(pSys);
    VGPU_TYPE *pVgpuTypeNode;

    for (pVgpuTypeNode = listHead(&(pKernelVgpuMgr->listVgpuTypeHead));
         pVgpuTypeNode != NULL;
         pVgpuTypeNode = listNext(&(pKernelVgpuMgr->listVgpuTypeHead), pVgpuTypeNode))
    {
        if (vgpuTypeId == pVgpuTypeNode->vgpuTypeId)
        {
            *vgpuType = pVgpuTypeNode;
            return NV_OK;
        }
    }
    return NV_ERR_OBJECT_NOT_FOUND;
}

static void
kvgpumgrCopyFromVgpuTypeToVgpuInfo(NVA081_CTRL_VGPU_INFO *pVgpuInfo, VGPU_TYPE *pVgpuType)
{
    pVgpuInfo->vgpuType            = pVgpuType->vgpuTypeId;
    pVgpuInfo->maxInstance         = pVgpuType->maxInstance;
    pVgpuInfo->numHeads            = pVgpuType->numHeads;
    pVgpuInfo->maxResolutionX      = pVgpuType->maxResolutionX;
    pVgpuInfo->maxResolutionY      = pVgpuType->maxResolutionY;
    pVgpuInfo->maxPixels           = pVgpuType->maxPixels;
    pVgpuInfo->frlConfig           = pVgpuType->frlConfig;
    pVgpuInfo->cudaEnabled         = pVgpuType->cudaEnabled;
    pVgpuInfo->eccSupported        = pVgpuType->eccSupported;
    pVgpuInfo->gpuInstanceSize     = pVgpuType->gpuInstanceSize;
    pVgpuInfo->multiVgpuSupported  = pVgpuType->multiVgpuSupported;
    pVgpuInfo->vdevId              = pVgpuType->vdevId;
    pVgpuInfo->pdevId              = pVgpuType->pdevId;
    pVgpuInfo->profileSize         = pVgpuType->profileSize;
    pVgpuInfo->fbLength            = pVgpuType->fbLength;
    pVgpuInfo->gspHeapSize         = pVgpuType->gspHeapSize;
    pVgpuInfo->fbReservation       = pVgpuType->fbReservation;
    pVgpuInfo->mappableVideoSize   = pVgpuType->mappableVideoSize;
    pVgpuInfo->encoderCapacity     = pVgpuType->encoderCapacity;
    pVgpuInfo->bar1Length          = pVgpuType->bar1Length;
    pVgpuInfo->gpuDirectSupported  = pVgpuType->gpuDirectSupported;
    pVgpuInfo->nvlinkP2PSupported  = pVgpuType->nvlinkP2PSupported;
    pVgpuInfo->maxInstancePerGI    = pVgpuType->maxInstancePerGI;
    pVgpuInfo->multiVgpuExclusive  = pVgpuType->multiVgpuExclusive;
    pVgpuInfo->frlEnable           = pVgpuType->frlEnable;

    portStringCopy(
        (char *) pVgpuInfo->vgpuName, sizeof(pVgpuInfo->vgpuName),
        (char *) pVgpuType->vgpuName, sizeof(pVgpuType->vgpuName));
    portStringCopy(
        (char *) pVgpuInfo->vgpuClass, sizeof(pVgpuInfo->vgpuClass),
        (char *) pVgpuType->vgpuClass, sizeof(pVgpuType->vgpuClass));
    portStringCopy(
        (char *) pVgpuInfo->license, sizeof(pVgpuInfo->license),
        (char *) pVgpuType->license, sizeof(pVgpuType->license));
    portStringCopy(
        (char *) pVgpuInfo->vgpuExtraParams, sizeof(pVgpuInfo->vgpuExtraParams),
        (char *) pVgpuType->vgpuExtraParams, sizeof(pVgpuType->vgpuExtraParams));
    portMemCopy(
        (char *) pVgpuInfo->vgpuSignature, sizeof(pVgpuInfo->vgpuSignature),
        (char *) pVgpuType->vgpuSignature, sizeof(pVgpuType->vgpuSignature));
}

NV_STATUS
kvgpumgrSendAllVgpuTypesToGsp(OBJGPU *pGpu)
{
    NV_STATUS             rmStatus       = NV_OK;
    OBJSYS               *pSys           = SYS_GET_INSTANCE();
    KernelVgpuMgr        *pKernelVgpuMgr = SYS_GET_KERNEL_VGPUMGR(pSys);
    RM_API               *pRmApi         = GPU_GET_PHYSICAL_RMAPI(pGpu);
    KERNEL_PHYS_GPU_INFO *pPgpuInfo;
    NV2080_CTRL_VGPU_MGR_INTERNAL_PGPU_ADD_VGPU_TYPE_PARAMS *pAddVgpuTypeParams;
    NvU32                 index, i;

    NV_CHECK_OK_OR_RETURN(LEVEL_ERROR,
            kvgpumgrGetPgpuIndex(pKernelVgpuMgr, pGpu->gpuId, &index));

    pPgpuInfo = &(pKernelVgpuMgr->pgpuInfo[index]);

    // Early exit if no vgpu types found
    if (pPgpuInfo->numVgpuTypes == 0)
    {
        return NV_OK;
    }

    pAddVgpuTypeParams = portMemAllocNonPaged(sizeof(*pAddVgpuTypeParams));
    if (pAddVgpuTypeParams == NULL)
    {
        NV_PRINTF(LEVEL_ERROR, "Failed to alloc memory for add vGPU type params\n");
        return NV_ERR_NO_MEMORY;
    }

    portMemSet(pAddVgpuTypeParams, 0, sizeof(*pAddVgpuTypeParams));

    for (i = 0; i < pPgpuInfo->numVgpuTypes; i++)
    {
        kvgpumgrCopyFromVgpuTypeToVgpuInfo(&pAddVgpuTypeParams->vgpuInfo[i],
                                           pPgpuInfo->vgpuTypes[i]);
    }

    pAddVgpuTypeParams->discardVgpuTypes = NV_TRUE;
    pAddVgpuTypeParams->vgpuInfoCount = pPgpuInfo->numVgpuTypes;

    NV_CHECK_OK_OR_GOTO(rmStatus, LEVEL_ERROR,
            pRmApi->Control(pRmApi, pGpu->hInternalClient, pGpu->hInternalSubdevice,
                            NV2080_CTRL_CMD_VGPU_MGR_INTERNAL_PGPU_ADD_VGPU_TYPE,
                            pAddVgpuTypeParams, sizeof(*pAddVgpuTypeParams)),
            kvgpumgrSendAllVgpuTypesToGsp_exit);

kvgpumgrSendAllVgpuTypesToGsp_exit:
    portMemFree(pAddVgpuTypeParams);
    return rmStatus;
}

NV_STATUS
kvgpumgrPgpuAddVgpuType
(
    OBJGPU *pGpu,
    NvBool discardVgpuTypes,
    NVA081_CTRL_VGPU_INFO *pVgpuInfo
)
{
    NV_STATUS rmStatus = NV_OK;
    NvU32 index;
    OBJSYS *pSys = SYS_GET_INSTANCE();
    KernelVgpuMgr *pKernelVgpuMgr = SYS_GET_KERNEL_VGPUMGR(pSys);
    RM_API *pRmApi = GPU_GET_PHYSICAL_RMAPI(pGpu);
    VGPU_TYPE *pVgpuTypeInfo;
    KERNEL_PHYS_GPU_INFO *pPgpuInfo;
    NV2080_CTRL_VGPU_MGR_INTERNAL_PGPU_ADD_VGPU_TYPE_PARAMS *pAddVgpuTypeParams;

    // TODO: Eliminate the heap allocation here
    // NV2080_CTRL_VGPU_MGR_INTERNAL_PGPU_ADD_VGPU_TYPE_PARAMS has massive size
    // because it has an embedded struct VGPU_TYPE instead of a pointer.
    // FINN won't be able to figure out serialization for a pointer at this moment.
    pAddVgpuTypeParams = portMemAllocNonPaged(sizeof(*pAddVgpuTypeParams));
    if (pAddVgpuTypeParams == NULL)
    {
        NV_PRINTF(LEVEL_ERROR, "Failed to alloc add vGPU type ctrl params.\n");
        return NV_ERR_NO_MEMORY;
    }
    portMemSet(pAddVgpuTypeParams, 0, sizeof(*pAddVgpuTypeParams));

    pAddVgpuTypeParams->discardVgpuTypes = discardVgpuTypes;
    pAddVgpuTypeParams->vgpuInfoCount = 1;
    pAddVgpuTypeParams->vgpuInfo[0] = *pVgpuInfo;

    rmStatus = pRmApi->Control(pRmApi, pGpu->hInternalClient, pGpu->hInternalSubdevice,
                               NV2080_CTRL_CMD_VGPU_MGR_INTERNAL_PGPU_ADD_VGPU_TYPE,
                               pAddVgpuTypeParams, sizeof(*pAddVgpuTypeParams));

    portMemFree(pAddVgpuTypeParams);

    if (rmStatus != NV_OK)
    {
        NV_PRINTF(LEVEL_ERROR, "Failed to create vGPU type on GSP.\n");
        return rmStatus;
    }

    if (kvgpumgrGetPgpuIndex(pKernelVgpuMgr, pGpu->gpuId, &index) != NV_OK)
    {
        return NV_ERR_OBJECT_NOT_FOUND;
    }

    pPgpuInfo = &(pKernelVgpuMgr->pgpuInfo[index]);

    if (pPgpuInfo->numVgpuTypes == NVA081_MAX_VGPU_TYPES_PER_PGPU && !discardVgpuTypes)
    {
        return NV_ERR_NOT_SUPPORTED;
    }

    // TODO: This is a WAR. Need to do this change in VGPUD.
     if (pPgpuInfo->fractionalMultiVgpu == 0 && pVgpuInfo->maxInstance != 1)
    {
        pVgpuInfo->multiVgpuSupported = 0;
    }

    rmStatus = kvgpumgrGetVgpuTypeInfo(pVgpuInfo->vgpuType, &pVgpuTypeInfo);
    if (rmStatus == NV_ERR_OBJECT_NOT_FOUND)
    {
        rmStatus = _kvgpumgrCreateVgpuType(pVgpuInfo, &pVgpuTypeInfo);
        if (rmStatus != NV_OK)
        {
            NV_PRINTF(LEVEL_ERROR, "Failed to create vGPU type.\n");
            return rmStatus;
        }
    }
    else if (rmStatus != NV_OK)
    {
        NV_PRINTF(LEVEL_ERROR, "Failed to get vGPU type.\n");
        return rmStatus;
    }

    if (discardVgpuTypes)
    {
        pPgpuInfo->numVgpuTypes = 0;
    }

    if (pPgpuInfo->numVgpuTypes == 0)
    {
        pPgpuInfo->sriovEnabled = gpuIsSriovEnabled(pGpu);
    }

    pPgpuInfo->vgpuTypes[pPgpuInfo->numVgpuTypes] = pVgpuTypeInfo;
    pPgpuInfo->numVgpuTypes++;

    vgpuMgrSetNominalFbSize(&pPgpuInfo->vgpuNominalFbSize, pVgpuTypeInfo);

    return NV_OK;
}

NV_STATUS
kvgpumgrAttachGpu(NvU32 gpuPciId)
{
    NvU32 index;
    OBJSYS *pSys = SYS_GET_INSTANCE();
    KernelVgpuMgr *pKernelVgpuMgr = SYS_GET_KERNEL_VGPUMGR(pSys);
    KERNEL_PHYS_GPU_INFO *pPhysGpuInfo;

    NV_PRINTF(LEVEL_INFO, "Enter function\n");

    NV_ASSERT_OR_RETURN((pKernelVgpuMgr != NULL), NV_ERR_NOT_SUPPORTED);

    index = 0;
    while ((index < pKernelVgpuMgr->pgpuCount) &&
           ((pKernelVgpuMgr->pgpuInfo[index].isAttached == NV_TRUE) ||
           (pKernelVgpuMgr->pgpuInfo[index].gpuPciId != gpuPciId)))
    {
        index++;
    }
    NV_ASSERT_OR_RETURN((index != NV_MAX_DEVICES), NV_ERR_INSUFFICIENT_RESOURCES);

    pPhysGpuInfo = &(pKernelVgpuMgr->pgpuInfo[index]);

    /* Probe call, RmInit is not done yet, so send pGpu as NULL */
    kvgpumgrSetVgpuType(NULL, pPhysGpuInfo, NVA081_CTRL_VGPU_CONFIG_INVALID_TYPE, KMIGMGR_SWIZZID_INVALID);
    pPhysGpuInfo->numActiveVgpu    = 0;
    pPhysGpuInfo->isAttached       = NV_TRUE;
    pPhysGpuInfo->numCreatedVgpu   = 0;

    listInit(&(pKernelVgpuMgr->pgpuInfo[index].listVgpuConfigEventsHead), portMemAllocatorGetGlobalNonPaged());
    listInit(&(pKernelVgpuMgr->pgpuInfo[index].listHostVgpuDeviceHead), portMemAllocatorGetGlobalNonPaged());

    if (index == pKernelVgpuMgr->pgpuCount)
    {
        pPhysGpuInfo->gpuPciId     = gpuPciId;
        pPhysGpuInfo->numVgpuTypes = 0;
        pPhysGpuInfo->vgpuConfigState = NVA081_CTRL_VGPU_CONFIG_STATE_UNINITIALIZED;
        portMemSet(&pPhysGpuInfo->vgpuNominalFbSize, 0, sizeof(VGPU_NOMINAL_FB_SIZE));
        pKernelVgpuMgr->pgpuCount++;
    }
    // Clear placement info for GI
    portMemSet(pPhysGpuInfo->kernelVgpuTypePlacementInfoGi, 0,
               sizeof(pPhysGpuInfo->kernelVgpuTypePlacementInfoGi));

    return NV_OK;
}

NV_STATUS
kvgpumgrMigTimeslicingModeEnabled(OBJGPU *pGpu)
{
    OBJSYS *pSys = SYS_GET_INSTANCE();
    KernelVgpuMgr *pKernelVgpuMgr = SYS_GET_KERNEL_VGPUMGR(pSys);
    KERNEL_PHYS_GPU_INFO *pPhysGpuInfo;
    NvU32 i;

    if (hypervisorIsVgxHyper())
    {
        NV2080_CTRL_VGPU_MGR_INTERNAL_SET_VGPU_MIG_TIMESLICE_MODE_PARAMS params = {0};
        RM_API *pRmApi = GPU_GET_PHYSICAL_RMAPI(pGpu);

        NV_ASSERT_OR_RETURN((pKernelVgpuMgr != NULL), NV_ERR_INVALID_ARGUMENT);

        NV_ASSERT_OR_RETURN((pGpu != NULL), NV_ERR_INVALID_ARGUMENT);

        NV_CHECK_OK_OR_RETURN(LEVEL_ERROR, kvgpumgrGetPgpuIndex(pKernelVgpuMgr, pGpu->gpuId, &i));

        pPhysGpuInfo = &(pKernelVgpuMgr->pgpuInfo[i]);
        NV_ASSERT_OR_RETURN((pPhysGpuInfo != NULL), NV_ERR_INVALID_ARGUMENT);

        pPhysGpuInfo->migTimeslicingModeEnabled = NV_FALSE;
        if ((pGpu->getProperty(pGpu, PDB_PROP_GPU_MIG_TIMESLICING_SUPPORTED)))
        {
            // Default disabled on VMWare as older ESXi hypervisor versions don't support
            // the MIG timeslicing. This support will be available with ESXi 9.x or later.
            pPhysGpuInfo->migTimeslicingModeEnabled = NV_TRUE;
        }
        params.bMigTimeslicingModeEnabled = pPhysGpuInfo->migTimeslicingModeEnabled;
        NV_CHECK_OK_OR_RETURN(LEVEL_ERROR, pRmApi->Control(pRmApi, pGpu->hInternalClient, pGpu->hInternalSubdevice,
                               NV2080_CTRL_CMD_VGPU_MGR_INTERNAL_SET_VGPU_MIG_TIMESLICE_MODE,
                               &params, sizeof(params)));
    }

    return NV_OK;
}
NvBool
kvgpumgrIsMigTimeslicingModeEnabled(OBJGPU *pGpu)
{
    OBJSYS *pSys = SYS_GET_INSTANCE();
    KernelVgpuMgr *pKernelVgpuMgr = SYS_GET_KERNEL_VGPUMGR(pSys);
    KERNEL_PHYS_GPU_INFO *pPhysGpuInfo;
    NvU32 i;
    NV_STATUS status;

    NV_CHECK_OR_ELSE(LEVEL_ERROR, (pKernelVgpuMgr != NULL), return NV_FALSE);

    NV_CHECK_OR_ELSE(LEVEL_ERROR, (pGpu != NULL), return NV_FALSE);

    NV_CHECK_OK_OR_ELSE(status, LEVEL_ERROR, kvgpumgrGetPgpuIndex(pKernelVgpuMgr, pGpu->gpuId, &i), return NV_FALSE);

    pPhysGpuInfo = &(pKernelVgpuMgr->pgpuInfo[i]);
    NV_CHECK_OR_ELSE(LEVEL_ERROR, (pPhysGpuInfo != NULL), return NV_FALSE);

    return pPhysGpuInfo->migTimeslicingModeEnabled;
}

NV_STATUS
kvgpumgrDetachGpu(NvU32 gpuPciId)
{
    NvU32 i , j;
    OBJSYS *pSys = SYS_GET_INSTANCE();
    KernelVgpuMgr *pKernelVgpuMgr = SYS_GET_KERNEL_VGPUMGR(pSys);
    KERNEL_PHYS_GPU_INFO *pPhysGpuInfo;
    VGPU_TYPE *pVgpuTypeInfo;
    KERNEL_VGPU_TYPE_PLACEMENT_INFO *pKernelVgpuTypePlacementInfo;
    VGPU_TYPE_PLACEMENT_INFO *pVgpuTypePlacementInfo;
    VGPU_INSTANCE_PLACEMENT_INFO *pVgpuInstancePlacementInfo;
    VGPU_TYPE_SUPPORTED_PLACEMENT_INFO *pVgpuTypeSupportedPlacementInfo;
    VGPU_INSTANCE_SUPPORTED_PLACEMENT_INFO *pVgpuInstanceSupportedPlacementInfo;
    REQUEST_VGPU_INFO_NODE *pRequestVgpu = NULL, *pRequestVgpuNext = NULL;

    NV_PRINTF(LEVEL_INFO, "Enter function\n");

    NV_ASSERT_OR_RETURN((pKernelVgpuMgr != NULL), NV_ERR_NOT_SUPPORTED);

    if (kvgpumgrGetPgpuIndex(pKernelVgpuMgr, gpuPciId, &i) != NV_OK)
        return NV_ERR_OBJECT_NOT_FOUND;

    pPhysGpuInfo = &(pKernelVgpuMgr->pgpuInfo[i]);
    pKernelVgpuTypePlacementInfo = &pPhysGpuInfo->kernelVgpuTypePlacementInfo;

    pPhysGpuInfo->createdVfMask = 0;
    pPhysGpuInfo->assignedSwizzIdMask = 0;
    pPhysGpuInfo->isAttached = NV_FALSE;
    pPhysGpuInfo->numCreatedVgpu = 0;

    for (i = 0; i < MAX_VGPU_TYPES_PER_PGPU; i++)
        pPhysGpuInfo->supportedTypeIds[i] = NVA081_CTRL_VGPU_CONFIG_INVALID_TYPE;

    for (i = 0; i < KMIGMGR_MAX_GPU_SWIZZID; i++)
        pPhysGpuInfo->assignedSwizzIdVgpuCount[i] = 0;
    /*
     * When heterogeneous vGPU mode is disabled, if vGPU supports homogeneous
     * placement ID, initially all homogeneous supported placement IDs are
     * creatable placement IDs.
     */
    if (pPhysGpuInfo->homogeneousPlacementSupported == NV_TRUE)
    {
        for (i = 0; i < pPhysGpuInfo->numVgpuTypes; i++)
        {
            pVgpuTypeInfo = pPhysGpuInfo->vgpuTypes[i];

            if (pVgpuTypeInfo == NULL)
                continue;

            pVgpuTypeSupportedPlacementInfo = &pVgpuTypeInfo->vgpuTypeSupportedPlacementInfo;
            pVgpuTypePlacementInfo = &pKernelVgpuTypePlacementInfo->vgpuTypePlacementInfo[i];

            for (j = 0; j < pVgpuTypeSupportedPlacementInfo->homogeneousPlacementCount; j++)
            {
                pVgpuInstanceSupportedPlacementInfo = &pVgpuTypeSupportedPlacementInfo->vgpuInstanceSupportedPlacementInfo[j];
                pVgpuInstancePlacementInfo = &pVgpuTypePlacementInfo->vgpuInstancePlacementInfo[j];
                pVgpuInstancePlacementInfo->creatablePlacementId = pVgpuInstanceSupportedPlacementInfo->homogeneousSupportedPlacementId;
            }
        }
        /* No vGPU instances running, so placement region is empty. */
        bitVectorClrAll(&pKernelVgpuTypePlacementInfo->usedPlacementRegionMap);
    }

    if (listCount(&(pKernelVgpuMgr->listRequestVgpuHead)) > 0)
    {
        for (pRequestVgpu = listHead(&(pKernelVgpuMgr->listRequestVgpuHead));
             pRequestVgpu != NULL;
             pRequestVgpu = pRequestVgpuNext)
        {
            pRequestVgpuNext = listNext(&(pKernelVgpuMgr->listRequestVgpuHead), pRequestVgpu);
            if (pRequestVgpu->gpuPciId == gpuPciId)
            {
                pKernelVgpuMgr->pHeap->eheapFree(pKernelVgpuMgr->pHeap, pRequestVgpu->vgpuId);
                listRemove(&(pKernelVgpuMgr->listRequestVgpuHead), pRequestVgpu);
            }
        }
    }

    listDestroy(&(pPhysGpuInfo->listHostVgpuDeviceHead));
    // Clear placement info for GI
    portMemSet(pPhysGpuInfo->kernelVgpuTypePlacementInfoGi, 0,
               sizeof(pPhysGpuInfo->kernelVgpuTypePlacementInfoGi));
    return NV_OK;
}

/*
 * @brief Sets Guest(VM) ID for the requested hostvgpudevice
 *
 * @param pParams               SET_GUEST_ID_PARAMS Pointer
 * @param pKernelHostVgpuDevice Device for which Vm ID need to be set
 * @param pGpu                  OBJGPU pointer
 *
 * @return NV_STATUS
 */
NV_STATUS
kvgpumgrRegisterGuestId(SET_GUEST_ID_PARAMS *pParams,
                        KERNEL_HOST_VGPU_DEVICE *pKernelHostVgpuDevice, OBJGPU *pGpu)
{
    NV_STATUS rmStatus = NV_OK;
    OBJSYS *pSys = SYS_GET_INSTANCE();
    KernelVgpuMgr *pKernelVgpuMgr = SYS_GET_KERNEL_VGPUMGR(pSys);

    if (pParams->action  == SET_GUEST_ID_ACTION_SET)
    {
        KERNEL_VGPU_GUEST *pVgpuGuest = NULL;

        for (pVgpuGuest = listHead(&(pKernelVgpuMgr->listVgpuGuestHead));
             pVgpuGuest != NULL;
             pVgpuGuest = listNext(&(pKernelVgpuMgr->listVgpuGuestHead), pVgpuGuest))
        {
            if (pParams->vmIdType == pVgpuGuest->guestVmInfo.vmIdType)
            {
                if (pParams->vmIdType == VM_ID_DOMAIN_ID)
                {
                    if (pParams->guestVmId.vmId == pVgpuGuest->guestVmInfo.guestVmId.vmId)
                    {
                        break;
                    }
                }
                else if (pParams->vmIdType == VM_ID_UUID)
                {
                    if (!portMemCmp(pParams->guestVmId.vmUuid,
                                  pVgpuGuest->guestVmInfo.guestVmId.vmUuid,
                                  VM_UUID_SIZE))
                    {
                        break;
                    }
                }
            }
        }

        if (pVgpuGuest == NULL)
        {
            pVgpuGuest = listAppendNew(&(pKernelVgpuMgr->listVgpuGuestHead));
            if (pVgpuGuest == NULL)
                return NV_ERR_NO_MEMORY;

            portMemSet(pVgpuGuest, 0, sizeof(KERNEL_VGPU_GUEST));
        }

        NV_ASSERT_OR_RETURN(pVgpuGuest->numVgpuDevices != MAX_VGPU_DEVICES_PER_VM,
                            NV_ERR_INSUFFICIENT_RESOURCES);

        pVgpuGuest->numVgpuDevices++;

        pKernelHostVgpuDevice->vgpuGuest = pVgpuGuest;

        pVgpuGuest->guestVmInfo.vmIdType = pParams->vmIdType;

        if (pParams->vmIdType == VM_ID_DOMAIN_ID)
        {
            pVgpuGuest->guestVmInfo.guestVmId.vmId = pParams->guestVmId.vmId;
        }
        else if (pParams->vmIdType == VM_ID_UUID)
        {
            portMemCopy(pVgpuGuest->guestVmInfo.guestVmId.vmUuid,
                        VM_UUID_SIZE,
                        pParams->guestVmId.vmUuid,
                        VM_UUID_SIZE);
        }
    }
    else if (pParams->action == SET_GUEST_ID_ACTION_UNSET)
    {
        KERNEL_VGPU_GUEST *pVgpuGuest = pKernelHostVgpuDevice->vgpuGuest;

        if (pVgpuGuest)
        {
            pVgpuGuest->numVgpuDevices--;
            if (pVgpuGuest->numVgpuDevices == 0)
            {
                listRemove(&(pKernelVgpuMgr->listVgpuGuestHead), pVgpuGuest);
            }
            pKernelHostVgpuDevice->vgpuGuest = NULL;
        }
        else
        {
            return NV_WARN_NULL_OBJECT;
        }
    }
    else
    {
        return NV_ERR_INVALID_ARGUMENT;
    }

    return rmStatus;
}

static NV_STATUS
_kvgpumgrVgpuAllocGPUInstanceSubscription
(
    OBJGPU *pGpu,
    NvU32 *phClient,
    NvU32 *phDevice,
    NvU32 swizzId
)
{
    NV_STATUS rmStatus;
    NVC637_ALLOCATION_PARAMETERS nvC637AllocParams;
    NvU32 hSubDevice;
    NvU32 hGPUInstanceSubscription;
    RM_API *pRmApi = rmapiGetInterface(RMAPI_GPU_LOCK_INTERNAL);

    NV_ASSERT_OR_RETURN(phClient != NULL, NV_ERR_INVALID_ARGUMENT);
    NV_ASSERT_OR_RETURN(phDevice != NULL, NV_ERR_INVALID_ARGUMENT);

    NV_ASSERT_OK_OR_RETURN(
        rmapiutilAllocClientAndDeviceHandles(pRmApi, pGpu, phClient, phDevice, &hSubDevice));

    // Allocate GPUInstanceSubscription handle
    NV_ASSERT_OK_OR_GOTO(
        rmStatus,
        serverutilGenResourceHandle(*phClient, &hGPUInstanceSubscription),
        cleanup);
    portMemSet(&nvC637AllocParams, 0, sizeof(nvC637AllocParams));
    nvC637AllocParams.swizzId = swizzId;

    // Allocate GPUInstanceSubscription object
    NV_CHECK_OK_OR_GOTO(
        rmStatus,
        LEVEL_ERROR,
        pRmApi->AllocWithHandle(pRmApi,
                                *phClient,
                                hSubDevice,
                                hGPUInstanceSubscription,
                                AMPERE_SMC_PARTITION_REF,
                                &nvC637AllocParams,
                                sizeof(nvC637AllocParams)),
        cleanup);

    return rmStatus;

cleanup:
    pRmApi->Free(pRmApi, *phClient, *phClient);
    return rmStatus;
}

NV_STATUS
kvgpumgrHeterogeneousGetChidOffset(NvU32 vgpuTypeId, NvU16 placementId,
                                   NvU32 numChannels, NvU64 *pChidOffset)
{
    VGPU_TYPE *pVgpuTypeInfo = NULL;
    NV_STATUS status;
    NvU32 i;
    VGPU_TYPE_SUPPORTED_PLACEMENT_INFO *pVgpuTypeSupportedPlacementInfo;
    VGPU_INSTANCE_SUPPORTED_PLACEMENT_INFO *pVgpuInstanceSupportedPlacementInfo;

    if ((status = kvgpumgrGetVgpuTypeInfo(vgpuTypeId, &pVgpuTypeInfo)) != NV_OK)
        return status;

    if (placementId == NVA081_PLACEMENT_ID_INVALID)
        return NV_ERR_INVALID_ARGUMENT;

    pVgpuTypeSupportedPlacementInfo = &pVgpuTypeInfo->vgpuTypeSupportedPlacementInfo;

    /* Channel count provided by plugin and calculated by RM should be same */
    if (numChannels != pVgpuTypeSupportedPlacementInfo->channelCount)
        return NV_ERR_INVALID_STATE;

    if (pVgpuTypeSupportedPlacementInfo->heterogeneousPlacementCount > MAX_VGPU_DEVICES_PER_PGPU)
        return NV_ERR_INVALID_INDEX;

    for (i = 0; i < pVgpuTypeSupportedPlacementInfo->heterogeneousPlacementCount; i++)
    {
        /*
         * For the placementId of the vGPU instance, find the corresponding
         * channel offset. heterogeneousSupportedPlacementId and
         * heterogeneousSupportedChidOffset share the same array index for a
         * give placement ID.
         */
        pVgpuInstanceSupportedPlacementInfo = &pVgpuTypeSupportedPlacementInfo->vgpuInstanceSupportedPlacementInfo[i];

        if (placementId == pVgpuInstanceSupportedPlacementInfo->heterogeneousSupportedPlacementId)
        {
            *pChidOffset = pVgpuInstanceSupportedPlacementInfo->heterogeneousSupportedChidOffset;
            return NV_OK;
        }
    }

    return NV_ERR_OBJECT_NOT_FOUND;
}

NV_STATUS
kvgpumgrHomogeneousGetChidOffset(NvU32 vgpuTypeId, NvU16 placementId,
                                 NvU32 numChannels, NvU64 *pChidOffset)
{
    VGPU_TYPE *pVgpuTypeInfo = NULL;
    NV_STATUS status;
    NvU32 i;
    VGPU_TYPE_SUPPORTED_PLACEMENT_INFO *pVgpuTypeSupportedPlacementInfo;
    VGPU_INSTANCE_SUPPORTED_PLACEMENT_INFO *pVgpuInstanceSupportedPlacementInfo;

    if ((status = kvgpumgrGetVgpuTypeInfo(vgpuTypeId, &pVgpuTypeInfo)) != NV_OK)
        return status;

    if (placementId == NVA081_PLACEMENT_ID_INVALID)
        return NV_ERR_INVALID_ARGUMENT;

    pVgpuTypeSupportedPlacementInfo = &pVgpuTypeInfo->vgpuTypeSupportedPlacementInfo;

    /* Channel count provided by plugin and calculated by RM should be same */
    if (numChannels != pVgpuTypeSupportedPlacementInfo->channelCount)
        return NV_ERR_INVALID_STATE;

    if (pVgpuTypeSupportedPlacementInfo->homogeneousPlacementCount > MAX_VGPU_DEVICES_PER_PGPU)
        return NV_ERR_INVALID_INDEX;

    for (i = 0; i < pVgpuTypeSupportedPlacementInfo->homogeneousPlacementCount; i++)
    {
        pVgpuInstanceSupportedPlacementInfo = &pVgpuTypeSupportedPlacementInfo->vgpuInstanceSupportedPlacementInfo[i];

        /*
         * For the placementId of the vGPU instance, find the corresponding
         * channel offset. homogeneousSupportedPlacementIds and
         * homogeneousSupportedChidOffsets share the same array index for a
         * give placement ID.
         */
        if (placementId == pVgpuInstanceSupportedPlacementInfo->homogeneousSupportedPlacementId)
        {
            *pChidOffset = pVgpuInstanceSupportedPlacementInfo->homogeneousSupportedChidOffset;
            return NV_OK;
        }
    }

    return NV_ERR_OBJECT_NOT_FOUND;
}

NV_STATUS
kvgpumgrGuestRegister(OBJGPU *pGpu,
                      NvU32 gfid,
                      NvU32 vgpuType,
                      NvU32 vmPid,
                      VM_ID_TYPE vmIdType,
                      VM_ID guestVmId,
                      NvHandle hPluginFBAllocationClient,
                      NvU32 numChannels,
                      NvU32 numPluginChannels,
                      NvU32 swizzId,
                      NvU32 vgpuDeviceInstanceId,
                      NvBool bDisableDefaultSmcExecPartRestore,
                      NvU16 placementId,
                      NvU8 *pVgpuDevName,
                      NvU32 accountingPid,
                      KERNEL_HOST_VGPU_DEVICE **ppKernelHostVgpuDevice)
{
    NV_STATUS                rmStatus                   = NV_OK;
    OBJSYS                  *pSys                       = SYS_GET_INSTANCE();
    KernelVgpuMgr           *pKernelVgpuMgr             = SYS_GET_KERNEL_VGPUMGR(pSys);
    SET_GUEST_ID_PARAMS      setGuestIDParams           = {0};
    NvU32                    i;
    KERNEL_PHYS_GPU_INFO    *pPhysGpuInfo;
    KERNEL_HOST_VGPU_DEVICE *pKernelHostVgpuDevice      = NULL;
    NvBool                   bMIGInUse                  = IS_MIG_IN_USE(pGpu);

    NvU32                    j;
    NvU32                    vgpuTypeIdx                = 0;
    REQUEST_VGPU_INFO_NODE  *pRequestVgpu               = NULL;

    NV_PRINTF(LEVEL_INFO, "Enter function\n");

    if (numPluginChannels > VGPU_MAX_PLUGIN_CHANNELS)
    {
        NV_PRINTF(LEVEL_ERROR, "failed, wrong number of plugin channels\n");
        return NV_ERR_INVALID_PARAMETER;
    }

    if (pGpu->getProperty(pGpu, PDB_PROP_GPU_IS_VIRTUALIZATION_MODE_HOST_VSGA))
    {
        NV_PRINTF(LEVEL_ERROR, "GPU is previously set in vSGA mode!\n");
        return NV_ERR_INVALID_DEVICE;
    }

    if (gpuIsSriovEnabled(pGpu))
    {
        GFID_ALLOC_STATUS gfidState;
        NV_ASSERT_OK_OR_RETURN(gpuGetGfidState(pGpu, gfid, &gfidState));
        NV_ASSERT_OR_RETURN((gfidState == GFID_FREE), NV_ERR_INSUFFICIENT_RESOURCES);
    }
    else
    {
        // MIG is supported only with SRIOV. Legacy mode doesn't support MIG
        NV_ASSERT_OR_RETURN(!bMIGInUse, NV_ERR_INVALID_STATE);
        numPluginChannels = 0;
    }

    if (kvgpumgrGetPgpuIndex(pKernelVgpuMgr, pGpu->gpuId, &i) != NV_OK)
        return NV_ERR_OBJECT_NOT_FOUND;

    pPhysGpuInfo = &(pKernelVgpuMgr->pgpuInfo[i]);

    // check if this vgpu type is allowed on physical gpu
    for (j = 0; j < pPhysGpuInfo->numVgpuTypes; j++)
    {
        if (vgpuType == pPhysGpuInfo->vgpuTypes[j]->vgpuTypeId)
        {
            vgpuTypeIdx = j;
            break;
        }
    }

    if (j == pPhysGpuInfo->numVgpuTypes)
    {
        return NV_ERR_NOT_SUPPORTED;
    }

    // On KVM, store a link to the corresponding REQUEST_VGPU_INFO_NODE in
    // KERNEL_HOST_VGPU_DEVICE and vice versa, so as to create a mapping between
    // vGPU device and the corresponding mdev device.
    if (osIsVgpuVfioPresent() == NV_OK)
    {
        for (pRequestVgpu = listHead(&pKernelVgpuMgr->listRequestVgpuHead);
             pRequestVgpu != NULL;
             pRequestVgpu = listNext(&pKernelVgpuMgr->listRequestVgpuHead, pRequestVgpu))
        {
            if (portMemCmp(pRequestVgpu->vgpuDevName, pVgpuDevName, VM_UUID_SIZE) == 0)
                break;
        }

        if (pRequestVgpu == NULL)
        {
            return NV_ERR_OBJECT_NOT_FOUND;
        }
    }

    /* On device-vm, swizzId is reserved during A084 object creation */
    if (IS_MIG_ENABLED(pGpu) &&
        (osIsVgpuDeviceVmPresent() == NV_OK)
        )

    {
        NvU32 partitionFlag;

        if (swizzId >= KMIGMGR_MAX_GPU_SWIZZID)
            return NV_ERR_INVALID_ARGUMENT;

        NV_ASSERT_OK_OR_RETURN(kvgpumgrGetPartitionFlag(vgpuType, &partitionFlag));

        NV_ASSERT_OK_OR_RETURN(kvgpumgrGetSwizzId(pGpu, pPhysGpuInfo, partitionFlag,
                               pPhysGpuInfo->vgpuTypes[vgpuTypeIdx], &swizzId));

    }

    /*
     * For heterogeneous vGPU mode, vGPU type is already validated based on placement ID
     * in NVA081_CTRL_CMD_VGPU_CONFIG_UPDATE_HETEROGENEOUS_INFO RmCtrl.
     * Both the RmCtrls are done before allocating the A084 object.
     */
    if (!IS_MIG_ENABLED(pGpu) && !pGpu->getProperty(pGpu, PDB_PROP_GPU_IS_VGPU_HETEROGENEOUS_MODE))
    {
        rmStatus = _kvgpumgrCheckVgpuTypeActivatable(pGpu, pPhysGpuInfo, pPhysGpuInfo->vgpuTypes[vgpuTypeIdx]);
        if (rmStatus != NV_OK)
            return rmStatus;

        rmStatus = kvgpumgrSetVgpuType(pGpu, pPhysGpuInfo, vgpuType, KMIGMGR_SWIZZID_INVALID);
        if (rmStatus != NV_OK)
            return rmStatus;
    }

    if (listCount(&(pPhysGpuInfo->listHostVgpuDeviceHead)) == MAX_VGPU_DEVICES_PER_PGPU)
    {
        rmStatus = NV_ERR_INSUFFICIENT_RESOURCES;
        goto failed;
    }

    // Register guest vgpu device instance
    pKernelHostVgpuDevice = listAppendNew(&(pPhysGpuInfo->listHostVgpuDeviceHead));
    if (pKernelHostVgpuDevice == NULL)
    {
        rmStatus = NV_ERR_NO_MEMORY;
        goto failed;
    }

    portMemSet(pKernelHostVgpuDevice, 0, sizeof(KERNEL_HOST_VGPU_DEVICE));

    pKernelHostVgpuDevice->vgpuType = vgpuType;

    NV_ASSERT_OR_RETURN(((vmIdType == VM_ID_DOMAIN_ID) ||
                         (vmIdType == VM_ID_UUID)),
                         NV_ERR_INVALID_ARGUMENT);

    setGuestIDParams.action   = SET_GUEST_ID_ACTION_SET;
    setGuestIDParams.vmPid    = vmPid;
    setGuestIDParams.vmIdType = vmIdType;
    if (vmIdType == VM_ID_DOMAIN_ID)
    {
        setGuestIDParams.guestVmId.vmId = guestVmId.vmId;
    }
    else if (vmIdType == VM_ID_UUID)
    {
        portMemCopy(setGuestIDParams.guestVmId.vmUuid,
                    VM_UUID_SIZE,
                    guestVmId.vmUuid,
                    VM_UUID_SIZE);
    }
    if (NV_OK != (rmStatus = kvgpumgrRegisterGuestId(&setGuestIDParams, pKernelHostVgpuDevice, pGpu)))
    {
        goto failed;
    }

    pKernelHostVgpuDevice->gfid                              = gfid;
    pKernelHostVgpuDevice->swizzId                           = swizzId;
    pKernelHostVgpuDevice->accountingPid                     = accountingPid;
    pKernelHostVgpuDevice->numPluginChannels                 = numPluginChannels;
    pKernelHostVgpuDevice->bDisableDefaultSmcExecPartRestore = bDisableDefaultSmcExecPartRestore;
    pKernelHostVgpuDevice->placementId                       = NVA081_PLACEMENT_ID_INVALID;

    if (IS_MIG_ENABLED(pGpu))
    {
        VGPU_TYPE *pVgpuTypeInfo;

        NV_ASSERT_OK_OR_RETURN(kvgpumgrGetVgpuTypeInfo(pPhysGpuInfo->vgpuTypes[vgpuTypeIdx]->vgpuTypeId,
                                                       &pVgpuTypeInfo));

        if (pVgpuTypeInfo->maxInstancePerGI > 1)
        {
            pKernelHostVgpuDevice->bDisableDefaultSmcExecPartRestore = NV_TRUE;
        }
    }

    if (osIsVgpuVfioPresent() == NV_OK)
    {
        pKernelHostVgpuDevice->pRequestVgpuInfoNode = pRequestVgpu;
        pRequestVgpu->pKernelHostVgpuDevice = pKernelHostVgpuDevice;
    }

    pKernelHostVgpuDevice->hPluginFBAllocationClient     = hPluginFBAllocationClient;
    pKernelHostVgpuDevice->vgpuDeviceGuestFbInfo.bValid  = NV_FALSE;
    pKernelHostVgpuDevice->vgpuDeviceGuestFbInfo.offset  = 0;
    pKernelHostVgpuDevice->vgpuDeviceGuestFbInfo.length  = 0;
    pKernelHostVgpuDevice->bOfflinedPageInfoValid        = NV_FALSE;

    pPhysGpuInfo->numActiveVgpu++;

    if (gpuIsSriovEnabled(pGpu))
    {
        // Set GFID in use
        gpuSetGfidUsage(pGpu, gfid, NV_TRUE);
    }

    // If MIG is enabled, then we need to set faultIds and per GPU instance VEID
    // offset in MMU_CFG registers. In MIG with memory partitioning, every
    // VM should be associated with a GPU instance, so set these params for
    // GPU instance engines
    //
    if (gpuIsSriovEnabled(pGpu) && bMIGInUse)
    {
        //
        // Alloc GPUInstanceSubscription object inside hostVgpuDeviceObject. This is
        // needed because hostVgpuDevice can be duped under any client and
        // will get freed after subdevice, so during hostVgpuDeviceDestruct
        // to get GPU instance reference, we need to keep the GPU instance alive
        //
        NV_ASSERT_OK_OR_GOTO(rmStatus, _kvgpumgrVgpuAllocGPUInstanceSubscription(pGpu,
                                                                                 &pKernelHostVgpuDevice->hMigClient,
                                                                                 &pKernelHostVgpuDevice->hMigDevice,
                                                                                 swizzId),
                             failed);
    }

    // Save/Restore default compute instance if not disabled during allocation
    if (gpuIsSriovEnabled(pGpu) && bMIGInUse && !pKernelHostVgpuDevice->bDisableDefaultSmcExecPartRestore)
    {
        KERNEL_MIG_GPU_INSTANCE *pKernelMIGGpuInstance;
        KernelMIGManager *pKernelMIGManager = GPU_GET_KERNEL_MIG_MANAGER(pGpu);

        // Get GPU instance
        NV_ASSERT_OK_OR_GOTO(rmStatus,
                             kmigmgrGetGPUInstanceInfo(pGpu, pKernelMIGManager,
                                                       swizzId, &pKernelMIGGpuInstance),
                             failed);
        //
        // Save any default compute instances as we may need to restore
        // them during destruct
        //
        NV_ASSERT_OK_OR_GOTO(rmStatus,
                             kmigmgrSaveComputeInstances(pGpu, pKernelMIGManager, pKernelMIGGpuInstance,
                                                         pKernelHostVgpuDevice->savedExecPartitions),
                             failed);

        // Determine if we have a VEID fragmentation between compute instances
        if (!kmigmgrIsPartitionVeidAllocationContiguous(pGpu, pKernelMIGManager, pKernelMIGGpuInstance))
        {
            // delete all compute instances and recreate then again
            NV_ASSERT_OK_OR_GOTO(rmStatus,
                                 kvgpuMgrRestoreSmcExecPart(pGpu,
                                                           pKernelHostVgpuDevice,
                                                           pKernelMIGGpuInstance),
                                 failed);
        }
    }

    *ppKernelHostVgpuDevice = pKernelHostVgpuDevice;

    return rmStatus;

failed:
    // Unset gfid usage
    if (gpuIsSriovEnabled(pGpu))
        gpuSetGfidUsage(pGpu, gfid, NV_FALSE);

    if (pKernelHostVgpuDevice != NULL)
    {
        if (pKernelHostVgpuDevice->vgpuGuest != NULL)
        {
            setGuestIDParams.action = SET_GUEST_ID_ACTION_UNSET;
            kvgpumgrRegisterGuestId(&setGuestIDParams, pKernelHostVgpuDevice, pGpu);
        }
        listRemove(&(pPhysGpuInfo->listHostVgpuDeviceHead), pKernelHostVgpuDevice);
    }

    if (osIsVgpuVfioPresent() != NV_OK && // not KVM
        swizzId != KMIGMGR_SWIZZID_INVALID)
        _kvgpumgrClearAssignedSwizzIdMask(pGpu, swizzId);

    return rmStatus;
}

NV_STATUS
kvgpumgrGuestUnregister(OBJGPU *pGpu, KERNEL_HOST_VGPU_DEVICE *pKernelHostVgpuDevice)
{
    NV_STATUS                    rmStatus        = NV_OK;
    OBJSYS                      *pSys            = SYS_GET_INSTANCE();
    KernelVgpuMgr               *pKernelVgpuMgr  = SYS_GET_KERNEL_VGPUMGR(pSys);
    RM_API                      *pRmApi          = rmapiGetInterface(RMAPI_GPU_LOCK_INTERNAL);
    KERNEL_PHYS_GPU_INFO        *pPhysGpuInfo;
    NvU32                        gfid            = pKernelHostVgpuDevice->gfid;
    SET_GUEST_ID_PARAMS          setGuestIDParams = {0};
    NvU32                        i;

    NV_PRINTF(LEVEL_INFO, "Enter function\n");

    if (gpuIsSriovEnabled(pGpu))
    {
        GFID_ALLOC_STATUS gfidState;

        // Sanity check on GFID
        NV_ASSERT_OK_OR_GOTO(rmStatus, gpuGetGfidState(pGpu, gfid, &gfidState), done);

        // Keep continuing in unregister path after sanity check
        if (gfidState == GFID_FREE)
        {
            NV_PRINTF(LEVEL_ERROR, "Request to unregister Invalid GFID\n");
            rmStatus = NV_ERR_INVALID_STATE;
            goto done;
        }
    }

    if (IS_MIG_IN_USE(pGpu) && !pKernelHostVgpuDevice->bDisableDefaultSmcExecPartRestore)
    {
        KernelMIGManager *pKernelMIGManager = GPU_GET_KERNEL_MIG_MANAGER(pGpu);
        KERNEL_MIG_GPU_INSTANCE *pKernelMIGGpuInstance;

        NV_ASSERT_OK_OR_GOTO(rmStatus, kmigmgrGetGPUInstanceInfo(pGpu, pKernelMIGManager,
                                                                pKernelHostVgpuDevice->swizzId,
                                                                &pKernelMIGGpuInstance),
                            done);

        NV_ASSERT_OK_OR_GOTO(rmStatus,
                            kvgpuMgrRestoreSmcExecPart(pGpu,
                                                        pKernelHostVgpuDevice,
                                                        pKernelMIGGpuInstance),
                            done);

        portMemSet(&pKernelHostVgpuDevice->savedExecPartitions, 0,
                sizeof(pKernelHostVgpuDevice->savedExecPartitions));
    }

    if (!gpuIsSriovEnabled(pGpu))
    {
        portMemFree(pKernelHostVgpuDevice->pGuestFbSegment);
        pKernelHostVgpuDevice->pGuestFbSegment = NULL;
    }

    if (pKernelHostVgpuDevice->hbmRegionList != NULL)
    {
        portMemFree(pKernelHostVgpuDevice->hbmRegionList);
        pKernelHostVgpuDevice->hbmRegionList = NULL;
        pKernelHostVgpuDevice->numValidHbmRegions = 0;
    }

    if (gpuIsSriovEnabled(pGpu))
    {
        if (IS_GSP_CLIENT(pGpu))
        {
            KernelFifo *pKernelFifo = GPU_GET_KERNEL_FIFO(pGpu);
            ENGINE_INFO tmpEngineInfo;
            Device *pMigDevice = NULL;

            portMemSet(&tmpEngineInfo, 0, sizeof(ENGINE_INFO));

            if (IS_MIG_IN_USE(pGpu))
            {
                RsClient *pClient;

                NV_ASSERT_OK_OR_GOTO(rmStatus,
                                     serverGetClientUnderLock(&g_resServ, pKernelHostVgpuDevice->hMigClient, &pClient),
                                     free_fifo_chids_exit);

                NV_ASSERT_OK_OR_GOTO(rmStatus,
                                     deviceGetByHandle(pClient, pKernelHostVgpuDevice->hMigDevice, &pMigDevice),
                                     free_fifo_chids_exit);
            }

            NV_ASSERT_OK_OR_GOTO(rmStatus,
                                 kfifoGetHostDeviceInfoTable_HAL(pGpu, pKernelFifo, &tmpEngineInfo, pMigDevice),
                                 free_fifo_chids_exit);

            vgpuMgrFreeSystemChannelIDs(pGpu,
                                        pKernelHostVgpuDevice->gfid,
                                        pKernelHostVgpuDevice->chidOffset,
                                        pKernelHostVgpuDevice->channelCount,
                                        pMigDevice,
                                        tmpEngineInfo.engineInfoListSize,
                                        tmpEngineInfo.engineInfoList);

free_fifo_chids_exit:
            portMemFree(tmpEngineInfo.engineInfoList);
            tmpEngineInfo.engineInfoList = NULL;
        }

        // Unset GFID usage
        gpuSetGfidUsage(pGpu, gfid, NV_FALSE);
    }

done:
    if (IS_MIG_IN_USE(pGpu))
    {
        pRmApi->Free(pRmApi, pKernelHostVgpuDevice->hMigClient, pKernelHostVgpuDevice->hMigClient);
        pKernelHostVgpuDevice->hMigClient = NV01_NULL_OBJECT;
        pKernelHostVgpuDevice->hMigDevice = NV01_NULL_OBJECT;
    }

    if (kvgpumgrGetPgpuIndex(pKernelVgpuMgr, pGpu->gpuId, &i) != NV_OK)
    {
        return NV_ERR_OBJECT_NOT_FOUND;
    }

    // Unset the guest id for host vGPU device
    setGuestIDParams.action = SET_GUEST_ID_ACTION_UNSET;
    if (NV_OK != (rmStatus = kvgpumgrRegisterGuestId(&setGuestIDParams,
                                                     pKernelHostVgpuDevice,
                                                     pGpu)))
    {
        // In case of Hyper-V, On VM reboot KERNEL_HOST_VGPU_DEVICE is not destroyed.
        // As guest id is unset during VM reboot by plugin explicitly. So we
        // will receive NV_WARN_NULL_OBJECT here and we need to ignore it.
        //
        // In case, plugin doesn't get a chance to unset the guest id e.g. in case
        // of crash or Force shutdown/Turn off then we need to get it unset here.
        if (hypervisorIsType(OS_HYPERVISOR_HYPERV) && rmStatus == NV_WARN_NULL_OBJECT)
        {
            rmStatus = NV_OK;
        }
        else
        {
            return rmStatus;
        }
    }

    pPhysGpuInfo = &(pKernelVgpuMgr->pgpuInfo[i]);

    NV_ASSERT_OR_RETURN((pPhysGpuInfo != NULL), NV_ERR_INVALID_ARGUMENT);

    pPhysGpuInfo->numActiveVgpu--;

    if (pPhysGpuInfo->numActiveVgpu == 0 && pPhysGpuInfo->numCreatedVgpu == 0)
    {
        kvgpumgrSetVgpuType(pGpu, pPhysGpuInfo, NVA081_CTRL_VGPU_CONFIG_INVALID_TYPE, pKernelHostVgpuDevice->swizzId);
    }

    if (pKernelHostVgpuDevice->pRequestVgpuInfoNode != NULL)
    {
        pKernelHostVgpuDevice->pRequestVgpuInfoNode->pKernelHostVgpuDevice = NULL;
    }

    if (osIsVgpuVfioPresent() != NV_OK) // Not KVM
    {
        if (pKernelHostVgpuDevice->swizzId != KMIGMGR_SWIZZID_INVALID)
            rmStatus = _kvgpumgrClearAssignedSwizzIdMask(pGpu, pKernelHostVgpuDevice->swizzId);
    }

    /* 
     * Accumulate the freed placement ID in creatable placement ID array.
     * In case of KVM, it will be happened at the time of mdev device removal. 
     */
    if (osIsVgpuVfioPresent() != NV_OK)
    {
        if (pKernelHostVgpuDevice->placementId != NVA081_PLACEMENT_ID_INVALID)
        {
            NvBool bHeterogeneousModeEnabled;

            NV_CHECK_OK_OR_RETURN(LEVEL_ERROR,
                kvgpuMgrGetHeterogeneousMode(pGpu,
                                             pKernelHostVgpuDevice->swizzId,
                                             &bHeterogeneousModeEnabled));

            if (bHeterogeneousModeEnabled)
            {
                rmStatus = _kvgpumgrUpdateHeterogeneousCreatablePlacementIds(pGpu,
                                    pKernelHostVgpuDevice->placementId,
                                    pKernelHostVgpuDevice->vgpuType, NV_FALSE,
                                    pKernelHostVgpuDevice->swizzId);
            }

            if (kvgpumgrCheckHomogeneousPlacementSupported(pGpu, pKernelHostVgpuDevice->swizzId) == NV_OK)
            {
                rmStatus = _kvgpumgrUpdateHomogeneousCreatablePlacementIds(pGpu,
                                        pKernelHostVgpuDevice->placementId,
                                        pKernelHostVgpuDevice->vgpuType, NV_FALSE,
                                        pKernelHostVgpuDevice->swizzId);
            }

            if (rmStatus != NV_OK)
            {
                NV_PRINTF(LEVEL_ERROR, "Failed to free placement ID: %d, error: 0x%x\n",
                          pKernelHostVgpuDevice->placementId, rmStatus);
            }
        }
    }

    listRemove(&(pPhysGpuInfo->listHostVgpuDeviceHead), pKernelHostVgpuDevice);

    return rmStatus;
}

//
// Helper function to check if pGPU is live migration capable.
//
NvBool
kvgpumgrCheckPgpuMigrationSupport(OBJGPU *pGpu)
{
    RM_API *pRmApi = GPU_GET_PHYSICAL_RMAPI(pGpu);
    NV2080_CTRL_VGPU_MGR_INTERNAL_GET_PGPU_MIGRATION_SUPPORT_PARAMS params = {0};
    NV_STATUS status;

    NV_CHECK_OK_OR_ELSE(status,
                        LEVEL_ERROR,
                        pRmApi->Control(pRmApi, pGpu->hInternalClient, pGpu->hInternalSubdevice,
                                        NV2080_CTRL_CMD_VGPU_MGR_INTERNAL_GET_PGPU_MIGRATION_SUPPORT,
                                        &params, sizeof(params)),
                        return NV_FALSE);

    return params.bIsMigrationSupported;
}

//
// Function to convert a NvU32 value to a ascii-encoded string.
// The value will be padded with 0 to use up the totalWidth. The value
// requiring space more that totalWidth will be truncated.
//
static NvU32
NvU32ToAsciiStr(NvU32 uval, NvU32 totalWidth, NvU8 *dest, NvBool invalid)
{
    NvU32 digitCount = 0;
    NvU8 nibble;

    if (totalWidth == 0)
        return (totalWidth);

    // Move the pointer to the end
    dest += (totalWidth - 1);

    // Start copying the data from the end in the reverse order
    while (digitCount < totalWidth)
    {
        digitCount++;
        if (invalid == NV_TRUE)
        {
            *dest-- = 'X';
            continue;
        }

        nibble = uval & 0xF;
        *dest-- = (nibble <= 9) ? (nibble + '0') : (nibble - 10 + 'A');
        uval = uval >> 4;
    }

    return (totalWidth);
}

NvU32 kvgpumgrGetPgpuDevIdEncoding(OBJGPU *pGpu, NvU8 *pgpuString,
                                   NvU32 strSize)
{
    NvU32 chipID = DRF_VAL(_PCI, _DEVID, _DEVICE, pGpu->idInfo.PCIDeviceID);
    NvU32 subID  = DRF_VAL(_PCI, _DEVID, _DEVICE, pGpu->idInfo.PCISubDeviceID);

    _get_chip_id_for_alias_pgpu(&chipID, &subID);

    if ((strSize < MAX_NVU32_TO_CONVERTED_STR_LEN)
        || (pgpuString == NULL))
    {
        return NV_U32_MAX;
    }

    // The encoding of the Dev ID is the value converted to string
    return (NvU32ToAsciiStr(chipID, DEVID_ENCODED_VALUE_SIZE, pgpuString, NV_FALSE));
}

NvU32 kvgpumgrGetPgpuSubdevIdEncoding(OBJGPU *pGpu, NvU8 *pgpuString,
                                      NvU32 strSize)
{
    NvU32 chipID = DRF_VAL(_PCI, _DEVID, _DEVICE, pGpu->idInfo.PCIDeviceID);
    NvU32 subID  = DRF_VAL(_PCI, _DEVID, _DEVICE, pGpu->idInfo.PCISubDeviceID);
    NvU32 bytes = 0;

    _get_chip_id_for_alias_pgpu(&chipID, &subID);

    if ((strSize < MAX_NVU32_TO_CONVERTED_STR_LEN)
        || (pgpuString == NULL))
    {
        return NV_U32_MAX;
    }

    // The encoding of the subdevice ID is its value converted to string
    bytes = NvU32ToAsciiStr(subID, SUBDEVID_ENCODED_VALUE_SIZE,
                                    pgpuString, NV_FALSE);
    return bytes;
}

NvU32 kvgpumgrGetPgpuFSEncoding(OBJGPU *pGpu, NvU8 *pgpuString,
                                NvU32 strSize)
{
    RM_API *pRmApi = GPU_GET_PHYSICAL_RMAPI(pGpu);
    NV2080_CTRL_VGPU_MGR_INTERNAL_GET_PGPU_FS_ENCODING_PARAMS params = {0};
    NV_STATUS status;

    params.pgpuStringSize = strSize;

    NV_CHECK_OK_OR_ELSE(status,
                          LEVEL_ERROR,
                          pRmApi->Control(pRmApi, pGpu->hInternalClient, pGpu->hInternalSubdevice,
                                          NV2080_CTRL_CMD_VGPU_MGR_INTERNAL_GET_PGPU_FS_ENCODING,
                                          &params, sizeof(params)),
                          return NV_U32_MAX);

    portMemCopy(pgpuString, strSize, params.pgpuString, params.pgpuStringSize);

    return params.pgpuStringSize;
}

NvU64
kvgpumgrGetEccAndPrReservedFb(OBJGPU *pGpu)
{
    Heap *pHeap = GPU_GET_HEAP(pGpu);
	const MEMORY_SYSTEM_STATIC_CONFIG *pMemorySystemConfig =
            kmemsysGetStaticConfig(pGpu, GPU_GET_KERNEL_MEMORY_SYSTEM(pGpu));
	NvU64 eccReservedFb = 0, pageRetirementReservedFb = 0, usableSizeFb;

    //
    // pageRetirementReservedFb is needed only in case of legacy, to account
    // for the pagepool used during page stitching
    //
    if (!gpuIsSriovEnabled(pGpu))
    {
        if (pGpu->getProperty(pGpu, PDB_PROP_GPU_ALLOW_PAGE_RETIREMENT))
        {
            if (IsPASCALorBetter(pGpu))
            {
                pageRetirementReservedFb = NV2080_CTRL_FB_OFFLINED_PAGES_MAX_PAGES * RM_PAGE_SIZE_HUGE;
            }
            else
            {
                pageRetirementReservedFb = NV2080_CTRL_FB_OFFLINED_PAGES_MAX_PAGES * RM_PAGE_SIZE_64K;
            }
        }
    }

    if (pMemorySystemConfig->bEnabledEccFBPA)
    {
        if ((pMemorySystemConfig->ramType != NV2080_CTRL_FB_INFO_RAM_TYPE_HBM1) &&
            (pMemorySystemConfig->ramType != NV2080_CTRL_FB_INFO_RAM_TYPE_HBM2) &&
            (pMemorySystemConfig->ramType != NV2080_CTRL_FB_INFO_RAM_TYPE_HBM3))
        {
            heapGetUsableSize(pHeap, &usableSizeFb);
            //
            // FB and ECC checksum calculation
            // eccReservedFb = totalFBEccOff / 16
            // i.e. totalFbEccOff = eccReservedFb * 16
            //
            // totalFbEccOff = totalFbEccOn + eccReservedFb
            // eccReservedFb * 16 = totalFbEccOn + eccReservedFb
            // eccReservedFb * 15 = totalFbEccOn
            // eccReservedFb = totalFbEccOn / 15 (totalFbEccOn is same as usableSizeFb)
            //
            eccReservedFb = usableSizeFb / 15;
        }
    }

	return pageRetirementReservedFb + eccReservedFb;
}

//
// A 32-bit variable is used to consolidate various GPU capabilities like
// ECC, SRIOV etc. The function sets the capabilities in the variable and
// converted to an ascii-encoded format.
//
NvU32 kvgpumgrGetPgpuCapEncoding(OBJGPU *pGpu, NvU8 *pgpuString, NvU32 strSize)
{
    RM_API *pRmApi = GPU_GET_PHYSICAL_RMAPI(pGpu);
    NvU32 pgpuCap = 0;

    if ((strSize < MAX_NVU32_TO_CONVERTED_STR_LEN) || (pgpuString == NULL))
    {
        return NV_U32_MAX;
    }

    {
        NV2080_CTRL_GPU_QUERY_ECC_STATUS_PARAMS eccStatus = {0};

        if (pRmApi->Control(pRmApi,
                            pGpu->hInternalClient,
                            pGpu->hInternalSubdevice,
                            NV2080_CTRL_CMD_GPU_QUERY_ECC_STATUS,
                            &eccStatus, sizeof(eccStatus)) == NV_OK)
        {
            NvU32 i;
            for (i = 0; i < NV2080_CTRL_GPU_ECC_UNIT_COUNT; i++)
            {
                if (eccStatus.units[i].enabled)
                {
                    pgpuCap |= PGPU_CAP_ECC_ON;
                    break;
                }
            }
        }
    }

    {
        NV0080_CTRL_GPU_GET_SRIOV_CAPS_PARAMS sriovCaps = {0};

        if (gpuGetSriovCaps_HAL(pGpu, &sriovCaps) == NV_OK)
        {
            if (sriovCaps.bSriovEnabled)
                pgpuCap |= PGPU_CAP_SRIOV_ON;
        }
    }

    return NvU32ToAsciiStr(pgpuCap, PGPU_CAP_ENCODED_VALUE_SIZE,
                           pgpuString, NV_FALSE);
}

/*
 * Get the user provide vGPU version range
 */
NV_STATUS
kvgpumgrGetHostVgpuVersion(NvU32 *user_min_supported_version,
                           NvU32 *user_max_supported_version)
{
    OBJSYS *pSys = SYS_GET_INSTANCE();
    KernelVgpuMgr *pKernelVgpuMgr = SYS_GET_KERNEL_VGPUMGR(pSys);

    *user_min_supported_version
        = pKernelVgpuMgr->user_min_supported_version;
    *user_max_supported_version
        = pKernelVgpuMgr->user_max_supported_version;
    return NV_OK;
}

/*
 * Set the user provide vGPU version range
 */
NV_STATUS
kvgpumgrSetHostVgpuVersion(NvU32 user_min_supported_version,
                           NvU32 user_max_supported_version)
{
    OBJSYS *pSys = SYS_GET_INSTANCE();
    KernelVgpuMgr *pKernelVgpuMgr = SYS_GET_KERNEL_VGPUMGR(pSys);
    NvU32 host_min_supported_version
        = GRIDSW_VERSION_EXTERNAL(NV_VGPU_MIN_SUPPORTED_GRIDSW_VERSION_EXTERNAL_MAJOR,
                                  NV_VGPU_MIN_SUPPORTED_GRIDSW_VERSION_EXTERNAL_MINOR);
    NvU32 host_max_supported_version
        = GRIDSW_VERSION_EXTERNAL(NV_VGPU_MAX_SUPPORTED_GRIDSW_VERSION_EXTERNAL_MAJOR,
                                  NV_VGPU_MAX_SUPPORTED_GRIDSW_VERSION_EXTERNAL_MINOR);

    /* Sanity check */
    if (user_min_supported_version > user_max_supported_version) {
        NV_PRINTF(LEVEL_ERROR,
                  "Maximum vGPU version (0x%x) being set is less than minimum version (0x%x)\n",
                  user_max_supported_version, user_min_supported_version);
        return NV_ERR_INVALID_ARGUMENT;
    }

    /* Sanity check: Fail in case the range being set by admin is a subset
     * of the range supported by the host driver
     */
    if ((user_min_supported_version < host_min_supported_version)
        || (user_max_supported_version > host_max_supported_version)) {

        NV_PRINTF(LEVEL_ERROR,
                  "vGPU version range being set (0x%x, 0x%x) is outside the range supported by host (0x%x, 0x%x)\n",
                  user_min_supported_version, user_max_supported_version,
                  host_min_supported_version, host_max_supported_version);

        return NV_ERR_INVALID_ARGUMENT;
    }

    NV_PRINTF(LEVEL_INFO,
              "vGPU version range enforced by user: (0x%x, 0x%x)\n",
              user_min_supported_version, user_max_supported_version);

    /* Save the information of the user specified range in the host */
    pKernelVgpuMgr->user_min_supported_version
        = user_min_supported_version;
    pKernelVgpuMgr->user_max_supported_version
        = user_max_supported_version;

    return NV_OK;
}

/*
 * Function to set swizzId is assigned to a vGPU device.
 */
static NV_STATUS
_kvgpumgrSetAssignedSwizzIdMask(OBJGPU       *pGpu,
                                VGPU_TYPE *vgpuTypeInfo,
                                NvU32         swizzId)
{
    OBJSYS *pSys = SYS_GET_INSTANCE();
    KernelVgpuMgr *pKernelVgpuMgr = SYS_GET_KERNEL_VGPUMGR(pSys);
    KERNEL_PHYS_GPU_INFO *pPhysGpuInfo;
    NvU32 i;
    NvU64 mask;
    NvBool bHeterogeneousModeEnabled = NV_FALSE;

    if (swizzId >= KMIGMGR_MAX_GPU_SWIZZID)
    {
        return NV_ERR_INVALID_ARGUMENT;
    }

    if (kvgpumgrGetPgpuIndex(pKernelVgpuMgr, pGpu->gpuId, &i) != NV_OK)
    {
        return NV_ERR_OBJECT_NOT_FOUND;
    }

    pPhysGpuInfo = &(pKernelVgpuMgr->pgpuInfo[i]);
    NV_ASSERT_OR_RETURN((pPhysGpuInfo != NULL), NV_ERR_INVALID_ARGUMENT);

    /* Validate that same ID is not already set and then set the ID */
    mask = NVBIT64(swizzId);

    if (IS_MIG_ENABLED(pGpu) && kvgpumgrIsMigTimeslicingModeEnabled(pGpu))
    {
        kvgpuMgrGetHeterogeneousMode(pGpu, swizzId, &bHeterogeneousModeEnabled);
    }

    if ((pPhysGpuInfo->assignedSwizzIdVgpuCount[swizzId] < vgpuTypeInfo->maxInstancePerGI) ||
         bHeterogeneousModeEnabled)
    {
        pPhysGpuInfo->assignedSwizzIdVgpuCount[swizzId]++;
        pPhysGpuInfo->assignedSwizzIdMask |= mask;
    }
    else
    {
        NV_PRINTF(LEVEL_ERROR, "SwizzID - %d max vgpu count is %d\n", swizzId, vgpuTypeInfo->maxInstancePerGI);
        DBG_BREAKPOINT();
        return NV_ERR_INSUFFICIENT_RESOURCES;
    }

    return NV_OK;
}

/*!
 * Function to mark swizzId is free to be used by other vGPU devices.
 */
NV_STATUS
_kvgpumgrClearAssignedSwizzIdMask(OBJGPU *pGpu,
                                  NvU32 swizzId)
{
    OBJSYS *pSys = SYS_GET_INSTANCE();
    KernelVgpuMgr *pKernelVgpuMgr = SYS_GET_KERNEL_VGPUMGR(pSys);
    KERNEL_PHYS_GPU_INFO *pPhysGpuInfo;
    NvU32 i;
    NvU64 mask;

    if (swizzId >= KMIGMGR_MAX_GPU_SWIZZID)
    {
        return NV_ERR_INVALID_ARGUMENT;
    }

    if (kvgpumgrGetPgpuIndex(pKernelVgpuMgr, pGpu->gpuId, &i) != NV_OK)
    {
        return NV_ERR_OBJECT_NOT_FOUND;
    }

    pPhysGpuInfo = &(pKernelVgpuMgr->pgpuInfo[i]);
    NV_ASSERT_OR_RETURN((pPhysGpuInfo != NULL), NV_ERR_INVALID_ARGUMENT);

    /* Validate that same ID is not already marked as free and then set the ID */
    mask = NVBIT64(swizzId);

    if (!(mask & pPhysGpuInfo->assignedSwizzIdMask))
    {
        NV_PRINTF(LEVEL_ERROR, "SwizzID - %d not in use\n", swizzId);
        DBG_BREAKPOINT();
        return NV_ERR_INVALID_STATE;
    }

    if (pPhysGpuInfo->assignedSwizzIdVgpuCount[swizzId] >= 1)
    {
        pPhysGpuInfo->assignedSwizzIdVgpuCount[swizzId]--;
    }
    else
    {
        NV_PRINTF(LEVEL_ERROR, "SwizzID - %d already has no assigned vGPU\n", swizzId);
        return NV_ERR_INVALID_STATE;
    }
    if (pPhysGpuInfo->assignedSwizzIdVgpuCount[swizzId] == 0)
    {
        pPhysGpuInfo->assignedSwizzIdMask &= ~mask;
    }

    return NV_OK;
}

NV_STATUS
kvgpumgrGetSwizzId(OBJGPU *pGpu,
                   KERNEL_PHYS_GPU_INFO *pPhysGpuInfo,
                   NvU32 partitionFlag,
                   VGPU_TYPE *vgpuTypeInfo,
                   NvU32 *swizzId)
{
    KernelMIGManager *pKernelMIGManager = GPU_GET_KERNEL_MIG_MANAGER(pGpu);
    NvU64 swizzIdInUseMask = 0;
    NvU32 id;
    NV_STATUS rmStatus = NV_OK;
    VGPU_TYPE *existingVgpuTypeInfo = NULL;
    NvBool bIsSwizzIdReserved = NV_FALSE;
    NvBool bFoundCreatableHeterogeneousType = NV_FALSE;

    swizzIdInUseMask = kmigmgrGetSwizzIdInUseMask(pGpu, pKernelMIGManager);

    // Determine valid swizzids not assigned to any vGPU device.
    for (id = 0; id < KMIGMGR_MAX_GPU_SWIZZID; id++)
    {
        //If specified GI is present, ignore other GIs
        if ((*swizzId != KMIGMGR_SWIZZID_INVALID) && (*swizzId != id))
            continue;

        if (NVBIT64(id) & swizzIdInUseMask)
        {
            KERNEL_MIG_GPU_INSTANCE *pKernelMIGGpuInstance;

            rmStatus = kmigmgrGetGPUInstanceInfo(pGpu, pKernelMIGManager, id, &pKernelMIGGpuInstance);
            if (rmStatus != NV_OK)
            {
                // Didn't find requested GPU instance
                NV_PRINTF(LEVEL_ERROR, "No valid GPU instance with SwizzId - %d found\n", id);
                return rmStatus;
            }

            if (pKernelMIGGpuInstance->partitionFlag == partitionFlag)
            {
                if (pPhysGpuInfo->migTimeslicingModeEnabled)
                {
                    NvBool bHeterogeneousModeEnabled;
                    NV_CHECK_OK_OR_RETURN(LEVEL_ERROR,
                        kvgpuMgrGetHeterogeneousMode(pGpu,
                                                     id,
                                                     &bHeterogeneousModeEnabled));
                    if (bHeterogeneousModeEnabled)
                    {
                        // search creatable heterogeneous type contains the current type
                        NvU32 *pSupportedTypeId = kvgpuMgrGetVgpuCreatableTypeIdFromSwizzId(pGpu,
                                                                                            pPhysGpuInfo, id);
                        for (NvU32 i = 0; i < pPhysGpuInfo->numVgpuTypes; i++)
                        {
                            if (pSupportedTypeId[i] == vgpuTypeInfo->vgpuTypeId)
                            {
                                bFoundCreatableHeterogeneousType = NV_TRUE;
                                break;
                            }
                        }
                        if (!bFoundCreatableHeterogeneousType)
                            continue;
                    }
                    else if (pPhysGpuInfo->assignedSwizzIdVgpuCount[id] > 0)
                    {
                        NvU32 vgpuTypeId;

                        if (_kvgpumgrGetCreatedVgpuOnGpuInstance(pGpu, id, &vgpuTypeId) == NV_OK)
                        {
                            NV_ASSERT_OK_OR_RETURN(kvgpumgrGetVgpuTypeInfo(vgpuTypeId, &existingVgpuTypeInfo));
                            if (kvgpumgrIsHeterogeneousVgpuTypeSupported() == NV_TRUE)
                            {
                                if (existingVgpuTypeInfo->profileSize != vgpuTypeInfo->profileSize)
                                    continue;
                            }
                            else
                            {
                                if (existingVgpuTypeInfo->vgpuTypeId != vgpuTypeInfo->vgpuTypeId)
                                     continue;
                            }
                        }
                    }
                }

                // Validate that same ID is not already set and then set the ID
                if ((pPhysGpuInfo->assignedSwizzIdVgpuCount[id] < vgpuTypeInfo->maxInstancePerGI) ||
                     bFoundCreatableHeterogeneousType)
                {
                    NV_ASSERT_OK_OR_RETURN(_kvgpumgrSetAssignedSwizzIdMask(pGpu, vgpuTypeInfo, pKernelMIGGpuInstance->swizzId));
                    *swizzId = pKernelMIGGpuInstance->swizzId;
                    bIsSwizzIdReserved = NV_TRUE;
                    break;
                }
            }
        }
    }

    if (bIsSwizzIdReserved == NV_FALSE)
    {
        return NV_ERR_INVALID_STATE;
    }

    return NV_OK;
}

NV_STATUS
kvgpumgrValidateSwizzId(OBJGPU *pGpu,
                        NvU32 vgpuTypeId,
                        NvU32 swizzId)
{
    KernelMIGManager *pKernelMIGManager = GPU_GET_KERNEL_MIG_MANAGER(pGpu);
    NvU32 partitionFlag                 = PARTITIONID_INVALID;
    NV_STATUS rmStatus                  = NV_OK;
    KERNEL_MIG_GPU_INSTANCE *pKernelMIGGpuInstance;
    VGPU_TYPE *vgpuTypeInfo = NULL;

    // Check if swizzId is valid.
    if (!kmigmgrIsSwizzIdInUse(pGpu, pKernelMIGManager, swizzId))
    {
        return NV_ERR_INVALID_ARGUMENT;
    }

    rmStatus = kmigmgrGetGPUInstanceInfo(pGpu, pKernelMIGManager, swizzId, &pKernelMIGGpuInstance);
    if (rmStatus != NV_OK)
    {
        // Didn't find requested GPU instance
        NV_PRINTF(LEVEL_ERROR, "No valid GPU instance with SwizzId - %d found\n", swizzId);
        return rmStatus;
    }

    NV_ASSERT_OK_OR_RETURN(
        kvgpumgrGetPartitionFlag(vgpuTypeId, &partitionFlag));

    // Check if swizzId is valid for vGPU type.
    if (pKernelMIGGpuInstance->partitionFlag == partitionFlag)
    {
        NV_ASSERT_OK_OR_RETURN(kvgpumgrGetVgpuTypeInfo(vgpuTypeId, &vgpuTypeInfo));
        NV_ASSERT_OK_OR_RETURN(_kvgpumgrSetAssignedSwizzIdMask(pGpu, vgpuTypeInfo, swizzId));
        return NV_OK;
    }

    return NV_ERR_INVALID_STATE;
}

NV_STATUS
kvgpumgrGetPartitionFlag(NvU32 vgpuTypeId, NvU32 *partitionFlag)
{
    NvU32 i;

    *partitionFlag = PARTITIONID_INVALID;

    for (i = 0; i < NV_ARRAY_ELEMENTS(vgpuSmcTypeIdMappings); i++)
    {
        if (vgpuSmcTypeIdMappings[i].vgpuTypeId == vgpuTypeId)
        {
            *partitionFlag = vgpuSmcTypeIdMappings[i].partitionFlag;
            return NV_OK;
        }
    }

    NV_PRINTF(LEVEL_ERROR, "Invalid SMC vGpu TypeId: 0x%x\n", vgpuTypeId);
    return NV_ERR_INVALID_ARGUMENT;
}

/*
 * Add or remove VF info to pgpuInfo of its PF
 * @param[in] gpuPciId          PCI ID of target PF
 * @param[in] cmd
 *      0/VGPU_CMD_PROCESS_VF_INFO.NV_VGPU_SAVE_VF_INFO     = Add VF info to VF list of target PF
 *      1/VGPU_CMD_PROCESS_VF_INFO.NV_VGPU_REMOVE_VF_INFO   = Remove VF info from VF list of target PF
 * @param[in] domain            Domain of VF to be stored
 * @param[in] bus               Bus no. of VF to be stored
 * @param[in] slot              Slot no. of VF to be stored
 * @param[in] function          Function of VF to be stored
 * @param[in] isMdevAttached    Flag to indicate if VF is registered with mdev
 * @param[out]vfPciInfo         Array of PCI information of VFs
 */
NV_STATUS
kvgpumgrProcessVfInfo(NvU32 gpuPciId, NvU8 cmd, NvU32 domain, NvU32 bus, NvU32 slot, NvU32 function, NvBool isMdevAttached, vgpu_vf_pci_info *vfPciInfo)
{
    NvU32 i;
    OBJSYS *pSys = SYS_GET_INSTANCE();
    KernelVgpuMgr *pKernelVgpuMgr = SYS_GET_KERNEL_VGPUMGR(pSys);
    NvU32 pgpuIndex;
    NV_STATUS status = NV_OK;
    KERNEL_PHYS_GPU_INFO *pPhysGpuInfo;
    vgpu_vf_pci_info *pVfInfo = NULL;

    if ((status = kvgpumgrGetPgpuIndex(pKernelVgpuMgr, gpuPciId, &pgpuIndex)) != NV_OK)
        return status;

    pPhysGpuInfo = &(pKernelVgpuMgr->pgpuInfo[pgpuIndex]);

    if (cmd == NV_VGPU_SAVE_VF_INFO)
    {
        /* Find existing entry for VF and set flags if found. */
        for (i = 0; i < MAX_VF_COUNT_PER_GPU; i++)
        {
            pVfInfo = &pPhysGpuInfo->vfPciInfo[i];

            if ((domain   == pVfInfo->domain) &&
                (bus      == pVfInfo->bus)    &&
                (slot     == pVfInfo->slot)   &&
                (function == pVfInfo->function))
            {
                pVfInfo->isNvidiaAttached = NV_TRUE;
                pVfInfo->isMdevAttached   = isMdevAttached;
                break;
            }
        }

        /*
         * If entry doesn't already exist, populate an empty slot and complain
         * if there isn't one.
         */
        if (i == MAX_VF_COUNT_PER_GPU)
        {
            for (i = 0; i < MAX_VF_COUNT_PER_GPU; i++)
            {
                pVfInfo = &pPhysGpuInfo->vfPciInfo[i];
                if ((!pVfInfo->isNvidiaAttached))
                {
                    pVfInfo->domain   = domain;
                    pVfInfo->bus      = bus;
                    pVfInfo->slot     = slot;
                    pVfInfo->function = function;

                    pVfInfo->isNvidiaAttached = NV_TRUE;
                    pVfInfo->isMdevAttached   = isMdevAttached;
                    break;
                }
            }

            if (i == MAX_VF_COUNT_PER_GPU)
            {
                NV_PRINTF(LEVEL_ERROR, "No free free slot to track VF PCI info\n");
                return NV_ERR_INVALID_STATE;
            }
        }
    }
    else if ((cmd == NV_VGPU_REMOVE_VF_PCI_INFO || cmd == NV_VGPU_REMOVE_VF_MDEV_INFO) && pPhysGpuInfo)
    {
        for (i = 0; i < MAX_VF_COUNT_PER_GPU; i++)
        {
            pVfInfo = &pPhysGpuInfo->vfPciInfo[i];

            if ((domain   == pVfInfo->domain) &&
                (bus      == pVfInfo->bus)    &&
                (slot     == pVfInfo->slot)   &&
                (function == pVfInfo->function))
            {
                if (cmd == NV_VGPU_REMOVE_VF_PCI_INFO)
                {
                    pVfInfo->isNvidiaAttached = NV_FALSE;
                }
                pVfInfo->isMdevAttached   = NV_FALSE;
                break;
            }
        }

        if (i == MAX_VF_COUNT_PER_GPU)
        {
            NV_PRINTF(LEVEL_ERROR, "Could not find entry for VF PCI info\n");
            return NV_ERR_MISSING_TABLE_ENTRY;
        }
    }
    else if (cmd == NV_VGPU_GET_VF_INFO && vfPciInfo)
    {
        for (i = 0; i < MAX_VF_COUNT_PER_GPU; i++)
        {
            pVfInfo = &pPhysGpuInfo->vfPciInfo[i];
            portMemCopy(&vfPciInfo[i], sizeof(vgpu_vf_pci_info), pVfInfo, sizeof(vgpu_vf_pci_info));
        }
    }
    else
    {
        NV_PRINTF(LEVEL_ERROR, "Requested invalid operation on VF info\n");
        status = NV_ERR_INVALID_ARGUMENT;
    }

    return status;
}

NV_STATUS
kvgpumgrEnumerateVgpuPerPgpu(OBJGPU *pGpu, NV2080_CTRL_VGPU_MGR_INTERNAL_ENUMERATE_VGPU_PER_PGPU_PARAMS *pParams)
{
    RM_API *pRmApi = GPU_GET_PHYSICAL_RMAPI(pGpu);

    NV_CHECK_OK_OR_RETURN(LEVEL_ERROR, pRmApi->Control(pRmApi, pGpu->hInternalClient, pGpu->hInternalSubdevice,
                                                       NV2080_CTRL_CMD_VGPU_MGR_INTERNAL_ENUMERATE_VGPU_PER_PGPU,
                                                       pParams, sizeof(*pParams)));

    return NV_OK;
}

NV_STATUS
kvgpumgrClearGuestVmInfo(OBJGPU *pGpu, KERNEL_HOST_VGPU_DEVICE *pKernelHostVgpuDevice)
{
    NV2080_CTRL_VGPU_MGR_INTERNAL_CLEAR_GUEST_VM_INFO_PARAMS params;
    RM_API *pRmApi = GPU_GET_PHYSICAL_RMAPI(pGpu);
    NV_STATUS status;

    portMemSet(&params, 0, sizeof(params));
    params.gfid = pKernelHostVgpuDevice->gfid;

    status = pRmApi->Control(pRmApi, pGpu->hInternalClient, pGpu->hInternalSubdevice,
                             NV2080_CTRL_CMD_VGPU_MGR_INTERNAL_CLEAR_GUEST_VM_INFO,
                             &params, sizeof(params));

    if (status != NV_OK)
        NV_PRINTF(LEVEL_ERROR, "Failed to clear guest vm info on GSP\n");

    return status;
}

NV_STATUS
kvgpumgrGetHostVgpuDeviceFromGfid(NvU32 gpuPciId, NvU32 gfid,
                                  KERNEL_HOST_VGPU_DEVICE** ppHostVgpuDevice)
{
    OBJSYS *pSys = SYS_GET_INSTANCE();
    KernelVgpuMgr *pKernelVgpuMgr = SYS_GET_KERNEL_VGPUMGR(pSys);
    NvU32 pgpuIndex;
    KERNEL_PHYS_GPU_INFO *pPgpuInfo;
    KERNEL_HOST_VGPU_DEVICE *pKernelHostVgpuDevice;

    NV_CHECK_OK_OR_RETURN(LEVEL_ERROR, kvgpumgrGetPgpuIndex(pKernelVgpuMgr, gpuPciId, &pgpuIndex));
    pPgpuInfo = &pKernelVgpuMgr->pgpuInfo[pgpuIndex];

    if (!pPgpuInfo->sriovEnabled)
        return NV_ERR_NOT_SUPPORTED;

    for (pKernelHostVgpuDevice = listHead(&(pPgpuInfo->listHostVgpuDeviceHead));
         pKernelHostVgpuDevice != NULL;
         pKernelHostVgpuDevice = listNext(&(pPgpuInfo->listHostVgpuDeviceHead), pKernelHostVgpuDevice))
    {
        if (pKernelHostVgpuDevice->gfid == gfid)
        {
            *ppHostVgpuDevice = pKernelHostVgpuDevice;
            return NV_OK;
        }
     }

    return NV_ERR_OBJECT_NOT_FOUND;
}

NV_STATUS
kvgpumgrGetVgpuFbUsage(OBJGPU *pGpu, NVA081_CTRL_VGPU_CONFIG_GET_VGPU_FB_USAGE_PARAMS *pParams)
{
    KERNEL_HOST_VGPU_DEVICE *pKernelHostVgpuDevice;
    NVA081_VGPU_FB_USAGE *pVgpuFbUsage;
    NvU32 i = 0;

    NV_PRINTF(LEVEL_INFO, "%s\n", __FUNCTION__);

    if (gpuIsSriovEnabled(pGpu))
    {
        NV2080_CTRL_VGPU_MGR_INTERNAL_GET_VGPU_FB_USAGE_PARAMS internalParams;
        RM_API *pRmApi = GPU_GET_PHYSICAL_RMAPI(pGpu);

        portMemSet(&internalParams, 0, sizeof(internalParams));

        NV_CHECK_OK_OR_RETURN(LEVEL_ERROR, pRmApi->Control(pRmApi, pGpu->hInternalClient, pGpu->hInternalSubdevice,
                                                           NV2080_CTRL_CMD_VGPU_MGR_INTERNAL_GET_VGPU_FB_USAGE,
                                                           &internalParams, sizeof(internalParams)));

        if (internalParams.vgpuCount > NV_ARRAY_ELEMENTS(pParams->vgpuFbUsage))
            return NV_ERR_INSUFFICIENT_RESOURCES;

        for (i = 0; i < internalParams.vgpuCount; i++)
        {
            NV_CHECK_OK_OR_RETURN(LEVEL_ERROR,
                 kvgpumgrGetHostVgpuDeviceFromGfid(pGpu->gpuId, internalParams.vgpuFbUsage[i].gfid, &pKernelHostVgpuDevice));

            pVgpuFbUsage = &pParams->vgpuFbUsage[i];

            portMemCopy(pVgpuFbUsage->vgpuUuid, VGPU_UUID_SIZE, pKernelHostVgpuDevice->vgpuUuid, VGPU_UUID_SIZE);

            pVgpuFbUsage->fbUsed = internalParams.vgpuFbUsage[i].fbUsed;
        }

        pParams->vgpuCount = internalParams.vgpuCount;
    }

    return NV_OK;
}

NV_STATUS
kvgpumgrSetVgpuEncoderCapacity(OBJGPU *pGpu, NvU8 *vgpuUuid, NvU32 encoderCapacity)
{
    OBJSYS *pSys = SYS_GET_INSTANCE();
    KernelVgpuMgr *pKernelVgpuMgr = SYS_GET_KERNEL_VGPUMGR(pSys);
    KERNEL_HOST_VGPU_DEVICE *pKernelHostVgpuDevice;
    KERNEL_PHYS_GPU_INFO *pPhysGpuInfo;
    NvU32 i;

    NV_PRINTF(LEVEL_INFO, "%s\n", __FUNCTION__);

    NV_CHECK_OK_OR_RETURN(LEVEL_ERROR, kvgpumgrGetPgpuIndex(pKernelVgpuMgr, pGpu->gpuId, &i));

    pPhysGpuInfo = &pKernelVgpuMgr->pgpuInfo[i];

    for (pKernelHostVgpuDevice = listHead(&(pPhysGpuInfo->listHostVgpuDeviceHead));
         pKernelHostVgpuDevice != NULL;
         pKernelHostVgpuDevice = listNext(&(pPhysGpuInfo->listHostVgpuDeviceHead), pKernelHostVgpuDevice))
    {
        if (portMemCmp(pKernelHostVgpuDevice->vgpuUuid, vgpuUuid, VGPU_UUID_SIZE) == 0)
            break;
    }

    if (pKernelHostVgpuDevice == NULL)
        return NV_ERR_OBJECT_NOT_FOUND;

    if (gpuIsSriovEnabled(pGpu))
    {
        NV2080_CTRL_VGPU_MGR_INTERNAL_SET_VGPU_ENCODER_CAPACITY_PARAMS params;
        RM_API *pRmApi = GPU_GET_PHYSICAL_RMAPI(pGpu);

        params.gfid = pKernelHostVgpuDevice->gfid;
        params.encoderCapacity = encoderCapacity;

        NV_CHECK_OK_OR_RETURN(LEVEL_ERROR,
                              pRmApi->Control(pRmApi, pGpu->hInternalClient, pGpu->hInternalSubdevice,
                                              NV2080_CTRL_CMD_VGPU_MGR_INTERNAL_SET_VGPU_ENCODER_CAPACITY,
                                              &params, sizeof(params)));
    }

    return NV_OK;
}

//
// Add vGPU info received on mdev_create sysfs call to REQUEST_VGPU_INFO_NODE
// list. REQUEST_VGPU_INFO_NODE is currently used only for vGPU on KVM.
//
// This funtion first checks whether the vGPU type is supported or not as
// only homegeneous vGPU types are supported currently. Also, this function
// only creates REQUEST_VGPU_INFO_NODE entry, actual vGPU will be created later
//
NV_STATUS
kvgpumgrCreateRequestVgpu(NvU32 gpuPciId, const NvU8 *pVgpuDevName,
                          NvU32 vgpuTypeId, NvU16 *vgpuId, NvU32 *gpuInstanceId,
                          NvU32 *placementId, NvU32 gpuPciBdf)
{
    OBJSYS *pSys = SYS_GET_INSTANCE();
    OBJGPU *pGpu = gpumgrGetGpuFromId(gpuPciId);
    KernelVgpuMgr *pKernelVgpuMgr = SYS_GET_KERNEL_VGPUMGR(pSys);
    OBJHYPERVISOR *pHypervisor = SYS_GET_HYPERVISOR(pSys);
    NvU32 pgpuIndex;
    NV_STATUS status = NV_OK;
    VGPU_TYPE *vgpuTypeInfo;
    KERNEL_PHYS_GPU_INFO *pPhysGpuInfo;
    REQUEST_VGPU_INFO_NODE *pRequestVgpu = NULL;
    NvU32 allocFlags = NVOS32_ALLOC_FLAGS_FORCE_MEM_GROWS_UP;
    NvU64 vgpuIdSize = 1, tmpVgpuId;
    NvU8 devfn = gpuDecodeDevice(gpuPciBdf);
    NvU32 swizzId = KMIGMGR_SWIZZID_INVALID;
    NvBool bHeterogeneousModeEnabled;

    if (pGpu == NULL)
    {
        NV_PRINTF(LEVEL_ERROR, "GPU %u is not initialized yet \n", gpuPciBdf);
        return NV_ERR_TIMEOUT_RETRY;
    }

    if ((status = kvgpumgrGetPgpuIndex(pKernelVgpuMgr, gpuPciId, &pgpuIndex)) != NV_OK)
        return status;

    pPhysGpuInfo = &(pKernelVgpuMgr->pgpuInfo[pgpuIndex]);

    if ((status = kvgpumgrGetVgpuTypeInfo(vgpuTypeId, &vgpuTypeInfo)) != NV_OK)
        return status;

    /*
     * For DriverVM, we will defer the createdVfMask validation later
     * during open call when the params provided.
     *
     */

    if (gpuIsSriovEnabled(pGpu) &&
        !(pHypervisor->getProperty(pHypervisor, PDB_PROP_HYPERVISOR_DRIVERVM_ENABLED)))
    {
        NvU8 fnId = devfn - pGpu->sriovState.firstVFOffset;

        NV_ASSERT_OR_RETURN((fnId < 64), NV_ERR_INVALID_ARGUMENT);

        if (pPhysGpuInfo->createdVfMask & NVBIT64(fnId))
            /* mdev device is already created on VF */
            return NV_ERR_INVALID_OPERATION;
    }

    if (IS_MIG_ENABLED(pGpu))
    {
        NvU32 partitionFlag = PARTITIONID_INVALID;

        NV_CHECK_OR_RETURN(LEVEL_INFO,
            IS_MIG_IN_USE(pGpu),
            NV_ERR_INVALID_OPERATION);

        NV_ASSERT_OK_OR_RETURN(
            kvgpumgrGetPartitionFlag(vgpuTypeInfo->vgpuTypeId, &partitionFlag));

        NV_ASSERT_OK_OR_RETURN(
            kvgpumgrGetSwizzId(pGpu, pPhysGpuInfo, partitionFlag, vgpuTypeInfo, &swizzId));
    }
    else
    {
        // Creation request for a MIG vgpuType
        if (vgpuTypeInfo->gpuInstanceSize)
            return NV_ERR_INVALID_OPERATION;

        status = kvgpumgrCheckVgpuTypeCreatable(pGpu, pPhysGpuInfo, vgpuTypeInfo);

        if (status != NV_OK)
            return status;
    }

    /* Initialize heap on first vGPU device creation */
    if (pKernelVgpuMgr->pHeap == NULL)
    {
        NvU64 tmpSize;
        pKernelVgpuMgr->pHeap = portMemAllocNonPaged(sizeof(OBJEHEAP));
        if (pKernelVgpuMgr->pHeap == NULL)
        {
            status = NV_ERR_NO_MEMORY;
            goto failed;
        }

        constructObjEHeap(pKernelVgpuMgr->pHeap, 1, 0xFFFF + 1, 0, 0);

        /* Verify if pHeap is allocated with required size */
        pKernelVgpuMgr->pHeap->eheapGetSize(pKernelVgpuMgr->pHeap, &tmpSize);
        if (tmpSize != 0xFFFF)
        {
            pKernelVgpuMgr->pHeap->eheapDestruct(pKernelVgpuMgr->pHeap);
            portMemFree(pKernelVgpuMgr->pHeap);
            pKernelVgpuMgr->pHeap = NULL;
            status = NV_ERR_INSUFFICIENT_RESOURCES;
            goto failed;
        }
    }

    status = pKernelVgpuMgr->pHeap->eheapAlloc(pKernelVgpuMgr->pHeap, 0xdeadbeef, &allocFlags,
                                         &tmpVgpuId, &vgpuIdSize, 1, 1,
                                         NULL, NULL, NULL);
    if (status != NV_OK)
    {
        NV_PRINTF(LEVEL_WARNING,
                  "Failed to allocate heap for vGPU ID 0x%x\n",
                  status);
        goto failed;
    }

    *vgpuId = (NvU16) tmpVgpuId;
    pRequestVgpu = listPrependNew(&(pKernelVgpuMgr->listRequestVgpuHead));
    if (pRequestVgpu == NULL)
    {
        status = NV_ERR_NO_MEMORY;
        goto failed;
    }

    portMemSet(pRequestVgpu, 0, sizeof(REQUEST_VGPU_INFO_NODE));

    portMemCopy(pRequestVgpu->vgpuDevName, VGPU_UUID_SIZE, pVgpuDevName, VGPU_UUID_SIZE);
    pRequestVgpu->gpuPciId = gpuPciId; /* For SRIOV, this is PF's gpuPciId */
    pRequestVgpu->gpuPciBdf = gpuPciBdf; /* For SRIOV, this is VF's gpuPciBdf */
    pRequestVgpu->vgpuId = *vgpuId;
    pRequestVgpu->placementId = NVA081_PLACEMENT_ID_INVALID;
    pRequestVgpu->vgpuTypeId = vgpuTypeId;

    if (IS_MIG_IN_USE(pGpu))
    {
        pRequestVgpu->swizzId = swizzId;
    }
    else
    {
        pRequestVgpu->swizzId = KMIGMGR_SWIZZID_INVALID;
    }
    if (!IS_MIG_ENABLED(pGpu) ||
        (IS_MIG_ENABLED(pGpu) && kvgpumgrIsMigTimeslicingModeEnabled(pGpu)))
    {
        NV_CHECK_OK_OR_RETURN(LEVEL_ERROR,
                            kvgpuMgrGetHeterogeneousMode(pGpu, swizzId, &bHeterogeneousModeEnabled));

        if (bHeterogeneousModeEnabled ||
            kvgpumgrCheckHomogeneousPlacementSupported(pGpu, pRequestVgpu->swizzId) == NV_OK)
        {
            NV_CHECK_OK_OR_GOTO(status, LEVEL_ERROR,
                kvgpumgrSetPlacementId(pRequestVgpu, pRequestVgpu->placementId), failed);
        }
    }

    kvgpumgrSetVgpuType(pGpu, pPhysGpuInfo, vgpuTypeId, pRequestVgpu->swizzId);
    pPhysGpuInfo->numCreatedVgpu++;

    if (gpuGetDevice(pGpu) != devfn)  /* SRIOV - VF */
        pPhysGpuInfo->createdVfMask |= NVBIT64(devfn - pGpu->sriovState.firstVFOffset);

    *placementId   = pRequestVgpu->placementId;
    *gpuInstanceId = pRequestVgpu->swizzId;

    return NV_OK;

failed:
    if (swizzId != KMIGMGR_SWIZZID_INVALID)
        _kvgpumgrClearAssignedSwizzIdMask(pGpu, swizzId);

    return status;
}

//
// Delete REQUEST_VGPU_INFO_NODE structure from list.
// REQUEST_VGPU_INFO_NODE is currently used only for vGPU on KVM.
//
NV_STATUS
kvgpumgrDeleteRequestVgpu(const NvU8 *pVgpuDevName, NvU16 vgpuId)
{
    OBJSYS *pSys = SYS_GET_INSTANCE();
    OBJGPU *pGpu = NULL;
    NV_STATUS status;
    NvU32 pgpuIndex;
    KernelVgpuMgr *pKernelVgpuMgr = SYS_GET_KERNEL_VGPUMGR(pSys);
    REQUEST_VGPU_INFO_NODE *pRequestVgpu = NULL;
    REQUEST_VGPU_INFO_NODE *pRequestVgpuNext = NULL;
    NvU8 devfn = 0;

    for (pRequestVgpu = listHead(&(pKernelVgpuMgr->listRequestVgpuHead));
         pRequestVgpu != NULL;
         pRequestVgpu = pRequestVgpuNext)
    {
        pRequestVgpuNext = listNext(&(pKernelVgpuMgr->listRequestVgpuHead), pRequestVgpu);
        if (portMemCmp(pVgpuDevName, pRequestVgpu->vgpuDevName, VGPU_UUID_SIZE) == 0)
        {
            if ((status = kvgpumgrGetPgpuIndex(pKernelVgpuMgr, pRequestVgpu->gpuPciId,
                                              &pgpuIndex)) != NV_OK)
                return status;

            pGpu  = gpumgrGetGpuFromId(pRequestVgpu->gpuPciId);
            devfn = gpuDecodeDevice(pRequestVgpu->gpuPciBdf);

            pKernelVgpuMgr->pgpuInfo[pgpuIndex].numCreatedVgpu--;

            if (pGpu && (gpuGetDevice(pGpu) != devfn)) // SRIOV - VF
                pKernelVgpuMgr->pgpuInfo[pgpuIndex].createdVfMask &= ~(NVBIT64(devfn - pGpu->sriovState.firstVFOffset));

            if (IS_MIG_ENABLED(pGpu))
                _kvgpumgrClearAssignedSwizzIdMask(pGpu, pRequestVgpu->swizzId);
            else if (pKernelVgpuMgr->pgpuInfo[pgpuIndex].numCreatedVgpu == 0)
                kvgpumgrSetVgpuType(pGpu, &pKernelVgpuMgr->pgpuInfo[pgpuIndex], NVA081_CTRL_VGPU_CONFIG_INVALID_TYPE,
                                    pRequestVgpu->swizzId);
            if (!IS_MIG_ENABLED(pGpu) ||
                (IS_MIG_ENABLED(pGpu) && kvgpumgrIsMigTimeslicingModeEnabled(pGpu)))
            {
                NvBool bHeterogeneousModeEnabled;

                NV_CHECK_OK_OR_RETURN(LEVEL_ERROR,
                    kvgpuMgrGetHeterogeneousMode(pGpu, pRequestVgpu->swizzId, &bHeterogeneousModeEnabled));
                if (bHeterogeneousModeEnabled)
                {
                    status = _kvgpumgrUpdateHeterogeneousCreatablePlacementIds(pGpu, pRequestVgpu->placementId,
                                                                            pRequestVgpu->vgpuTypeId, NV_FALSE,
                                                                            pRequestVgpu->swizzId);
                }
                else if (kvgpumgrCheckHomogeneousPlacementSupported(pGpu, pRequestVgpu->swizzId) == NV_OK)
                {
                    status = _kvgpumgrUpdateHomogeneousCreatablePlacementIds(pGpu, pRequestVgpu->placementId,
                                                                            pRequestVgpu->vgpuTypeId, NV_FALSE,
                                                                            pRequestVgpu->swizzId);
                }

                if (status != NV_OK)
                    NV_PRINTF(LEVEL_ERROR, "Failed to free up allocated placementId, status: 0x%x\n", status);
            }

            pKernelVgpuMgr->pHeap->eheapFree(pKernelVgpuMgr->pHeap, vgpuId);

            listRemove(&(pKernelVgpuMgr->listRequestVgpuHead), pRequestVgpu);

            return NV_OK;
        }
    }
    return NV_ERR_OBJECT_NOT_FOUND;
}

NV_STATUS kvgpumgrGetAvailableInstances(
    NvU32     *availInstances,
    OBJGPU    *pGpu,
    VGPU_TYPE *vgpuTypeInfo,
    NvU32      pgpuIndex,
    NvU8       devfn
)
{
    NV_STATUS rmStatus = NV_OK;
    OBJSYS *pSys = SYS_GET_INSTANCE();
    KernelVgpuMgr *pKernelVgpuMgr = SYS_GET_KERNEL_VGPUMGR(pSys);
    OBJHYPERVISOR *pHypervisor = SYS_GET_HYPERVISOR(pSys);

    *availInstances = 0;

    if (gpuIsSriovEnabled(pGpu))
    {
        if (!pHypervisor->getProperty(pHypervisor, PDB_PROP_HYPERVISOR_DRIVERVM_ENABLED))
        {
            NvU8 fnId = devfn - pGpu->sriovState.firstVFOffset;

            NV_ASSERT_OR_RETURN(fnId <= pGpu->sriovState.totalVFs, NV_ERR_INVALID_ARGUMENT);

            if ((pKernelVgpuMgr->pgpuInfo[pgpuIndex].createdVfMask) & (NVBIT64(fnId)))
                goto exit;
        }

        if (IS_MIG_ENABLED(pGpu))
        {
            if (IS_MIG_IN_USE(pGpu))
            {
                NvU64 swizzIdInUseMask = 0;
                NvU32 partitionFlag = PARTITIONID_INVALID;
                KernelMIGManager *pKernelMIGManager = GPU_GET_KERNEL_MIG_MANAGER(pGpu);
                NvU32 id;
                VGPU_TYPE *existingVgpuTypeInfo = NULL;

                swizzIdInUseMask = kmigmgrGetSwizzIdInUseMask(pGpu, pKernelMIGManager);

                if (!vgpuTypeInfo->gpuInstanceSize)
                {
                    // Query for a non MIG vgpuType
                    NV_PRINTF(LEVEL_INFO, "Query for a non MIG vGPU type \n");
                    rmStatus = NV_OK;
                    goto exit;
                }

                rmStatus = kvgpumgrGetPartitionFlag(vgpuTypeInfo->vgpuTypeId,
                                                    &partitionFlag);
                if (rmStatus != NV_OK)
                {
                    // Query for a non MIG vgpuType
                    NV_PRINTF(LEVEL_ERROR, "Failed to get partition flags.\n");
                    goto exit;
                }
    
                // Determine valid swizzids not assigned to any vGPU device.
                FOR_EACH_INDEX_IN_MASK(64, id, swizzIdInUseMask)
                {
                    KERNEL_MIG_GPU_INSTANCE *pKernelMIGGpuInstance;

                    rmStatus = kmigmgrGetGPUInstanceInfo(pGpu, pKernelMIGManager,
                                                         id, &pKernelMIGGpuInstance);
                    if (rmStatus != NV_OK)
                    {
                        // Didn't find requested GPU instance
                        NV_PRINTF(LEVEL_ERROR,
                                  "No valid GPU instance with SwizzId - %d found\n", id);
                        goto exit;
                    }

                    if (pKernelMIGGpuInstance->partitionFlag == partitionFlag)
                    {
                        if (pKernelVgpuMgr->pgpuInfo[pgpuIndex].assignedSwizzIdVgpuCount[id] > 0)
                        {
                            NvU32 vgpuTypeId;

                            if (_kvgpumgrGetCreatedVgpuOnGpuInstance(pGpu, id, &vgpuTypeId) == NV_OK)
                            {
                                NV_ASSERT_OK_OR_RETURN(kvgpumgrGetVgpuTypeInfo(vgpuTypeId, &existingVgpuTypeInfo));
                                if (kvgpumgrIsHeterogeneousVgpuTypeSupported() == NV_TRUE)
                                {
                                    if (existingVgpuTypeInfo->profileSize != vgpuTypeInfo->profileSize)
                                        continue;
                                }
                                else
                                {
                                    if (existingVgpuTypeInfo->vgpuTypeId != vgpuTypeInfo->vgpuTypeId)
                                         continue;
                                }
                            }
                        }
                       // Validate that same ID is not already set and VF is availablei
                       if (pKernelVgpuMgr->pgpuInfo[pgpuIndex].assignedSwizzIdVgpuCount[id] < vgpuTypeInfo->maxInstancePerGI)
                        {
                            *availInstances = 1;
                            break;
                        }
                    }
                }
                FOR_EACH_INDEX_IN_MASK_END;
            }
        }
        else
        {
            if (pKernelVgpuMgr->pgpuInfo[pgpuIndex].numCreatedVgpu < vgpuTypeInfo->maxInstance)
            {
                if (vgpuTypeInfo->gpuInstanceSize)
                {
                    // Query for a MIG vgpuType
                    NV_PRINTF(LEVEL_INFO, "Query for a MIG vGPU type\n");
                    rmStatus = NV_OK;
                    goto exit;
                }

                if (kvgpumgrCheckVgpuTypeCreatable(pGpu, &pKernelVgpuMgr->pgpuInfo[pgpuIndex],
                                                    vgpuTypeInfo) == NV_OK)
                    *availInstances = 1;
            }
        }
    }
    else
    {
        if (kvgpumgrCheckVgpuTypeCreatable(pGpu, &pKernelVgpuMgr->pgpuInfo[pgpuIndex], vgpuTypeInfo) == NV_OK)
            *availInstances = vgpuTypeInfo->maxInstance - pKernelVgpuMgr->pgpuInfo[pgpuIndex].numCreatedVgpu;
    }

exit:
    return rmStatus;
}

NV_STATUS
kvgpumgrGetHostVgpuDeviceFromVgpuDevName(NvU32 gpuPciId, const NvU8 *pVgpuDevName,
                                      KERNEL_HOST_VGPU_DEVICE **ppKernelHostVgpuDevice)
{
    OBJSYS *pSys = SYS_GET_INSTANCE();
    KernelVgpuMgr *pKernelVgpuMgr = SYS_GET_KERNEL_VGPUMGR(pSys);
    NvU32 pgpuIndex , rmStatus;
    KERNEL_PHYS_GPU_INFO *pPgpuInfo;
    KERNEL_HOST_VGPU_DEVICE *pKernelHostVgpuDevice;

    if (osIsVgpuVfioPresent() != NV_OK)
       return NV_ERR_NOT_SUPPORTED;

    if ((rmStatus = kvgpumgrGetPgpuIndex(pKernelVgpuMgr, gpuPciId, &pgpuIndex)) != NV_OK)
        return rmStatus;

    pPgpuInfo = &pKernelVgpuMgr->pgpuInfo[pgpuIndex];

    for (pKernelHostVgpuDevice = listHead(&(pPgpuInfo->listHostVgpuDeviceHead));
         pKernelHostVgpuDevice != NULL;
         pKernelHostVgpuDevice = listNext(&(pPgpuInfo->listHostVgpuDeviceHead), pKernelHostVgpuDevice))
    {
        if (pKernelHostVgpuDevice == NULL || pKernelHostVgpuDevice->pRequestVgpuInfoNode == NULL)
            return NV_ERR_INVALID_POINTER;

        if (portMemCmp(pKernelHostVgpuDevice->pRequestVgpuInfoNode->vgpuDevName,
                       pVgpuDevName, VM_UUID_SIZE) == 0)
        {
            *ppKernelHostVgpuDevice = pKernelHostVgpuDevice;
            return NV_OK;
        }
     }
    NV_PRINTF(LEVEL_ERROR, "Object not found\n");
    return NV_ERR_OBJECT_NOT_FOUND;
}

NV_STATUS
kvgpumgrGetHostVgpuDeviceFromVgpuUuid(NvU32 gpuPciId, NvU8 *vgpuUuid,
                                  KERNEL_HOST_VGPU_DEVICE **ppKernelHostVgpuDevice)
{
    OBJSYS *pSys = SYS_GET_INSTANCE();
    KernelVgpuMgr *pKernelVgpuMgr = SYS_GET_KERNEL_VGPUMGR(pSys);
    NvU32 pgpuIndex;
    KERNEL_PHYS_GPU_INFO *pPgpuInfo;
    KERNEL_HOST_VGPU_DEVICE *pKernelHostVgpuDevice;

    NV_CHECK_OK_OR_RETURN(LEVEL_ERROR, kvgpumgrGetPgpuIndex(pKernelVgpuMgr, gpuPciId, &pgpuIndex));

    pPgpuInfo = &pKernelVgpuMgr->pgpuInfo[pgpuIndex];

    for (pKernelHostVgpuDevice = listHead(&(pPgpuInfo->listHostVgpuDeviceHead));
         pKernelHostVgpuDevice != NULL;
         pKernelHostVgpuDevice = listNext(&(pPgpuInfo->listHostVgpuDeviceHead), pKernelHostVgpuDevice))
    {
        if (portMemCmp(pKernelHostVgpuDevice->vgpuUuid,
                             vgpuUuid, VM_UUID_SIZE) == 0)
        {
            *ppKernelHostVgpuDevice = pKernelHostVgpuDevice;
            return NV_OK;
        }
    }

    return NV_ERR_OBJECT_NOT_FOUND;
}

/*
 * On 1GB profiles on 48GB board, we do not deduct any reserved fb from vGPU FB length
 * as maxInstance is only 32, so guest complete 1GB of fb mem. In such cases with
 * heterogeneous vGPUs, for such 1G profiles on 48G board, some combinations with
 * other vGPU types do not work.
 * Due to this, A40-8Q at placement Id 0 cannot work with A40-1Q at placement Id 8 since
 * A40-1Q is occupying larger vGPU FB length and overlaps with A40-8Q's assigned FB.
 * Similar scenario happens for A40-8Q at placement Id 24 and A40-1Q at placement Id 32
 * So, we disable 1GB profiles in such cases at placement Id 8 and placement Id 32.
 */
static NvBool
_kvgpumgrIsPlacementValid(OBJGPU *pGpu, KERNEL_PHYS_GPU_INFO *pPgpuInfo, NvU32 vgpuTypeIndex,
                          NvU32 placementId)
{
    NvU64 vgpuFbLength;
    NvU32 denyListAdaHopper[] = {2, 3, 8, 9, 14, 15, 20, 21, 26, 27, 32, 33, 38, 39, 44, 45};
    NvU32 denyListAmpere[] = {8, 32};
    NvU32 i, length, *invalidPlacements;
    KERNEL_VGPU_TYPE_PLACEMENT_INFO *pKernelVgpuTypePlacementInfo;
    VGPU_TYPE_PLACEMENT_INFO *pVgpuTypePlacementInfo;

    if (IS_GSP_CLIENT(pGpu))
    {
        invalidPlacements = denyListAdaHopper;
        length = NV_ARRAY_ELEMENTS(denyListAdaHopper);
    }
    else
    {
        invalidPlacements = denyListAmpere;
        length = NV_ARRAY_ELEMENTS(denyListAmpere);
    }

    pKernelVgpuTypePlacementInfo = &pPgpuInfo->kernelVgpuTypePlacementInfo;
    pVgpuTypePlacementInfo = &pKernelVgpuTypePlacementInfo->vgpuTypePlacementInfo[vgpuTypeIndex];

    vgpuFbLength = pVgpuTypePlacementInfo->guestVmmuCount * gpuGetVmmuSegmentSize(pGpu);

    if (vgpuFbLength == VGPU_LOWER_FB_PROFILE_SIZE_1_GB)
    {
        for (i = 0; i < length; i++)
        {
           if (invalidPlacements[i] == placementId)
               return NV_FALSE;
        }
    }

    return NV_TRUE;
}

/*
 * This function recursively divides the placement region in 1/2, 1/4, 1/8, 1/16 partitions
 * and then calculates the placement IDs of each vGPU type. It will try to allocate
 * placement ID of 2 vGPU instances for a type in such a way that the placement ID
 * of first vGPU instance starts at the beginning of partition and the placement ID
 * of second vGPU instance ends at the end of partition.
 *
 * It finds the vGPU types whose 2 instances can be allocated in a 1/2, 1/4, 1/8, 1/16
 * partition and then allocates one at start of partition and other at end of partion.
 *
 * It does this recursively by taking input start and end as input alongwith the current
 * partition which specifies whether this is 1/1, 1/2, 1/4, 1/8, 1/16 partion.
 *
 * Similarly it also calculates the channel offsets for each supported placementId
 * for a vGPU type. The algorithm for calculating placementId and channel offset
 * is the same just that for channel offset the range is from 0 to 2048.
 *
 * For vGPU-GSP, we also need to define offsets in FB for GSP plugin heap.
 * For larger vGPU profiles with GSP plugin heap >=1 vMMU segment, the gsp heap is
 * placed at the end of vGPU FB.
 * For smaller profiles with gsp heap < 1 vMMU segment, heap is placed at end of
 * each 1/4th alignment (for 128MB vMMU segment size) OR 1/8th alighment (for 64MB
 * vMMU segment size). For such profiles the vMMU segment at the end of 1/4 or 1/8
 * partition is divided as per profile's heap size.
 *
 */
static void
_kvgpumgrSetHeterogeneousResources(OBJGPU *pGpu, KERNEL_PHYS_GPU_INFO *pPgpuInfo, NvU32 placementIdMin,
                                   NvU32 placementIdMax, NvU32 chidMin, NvU32 chidMax,
                                   NvU32 vmmuSegMin, NvU32 vmmuSegMax, NvU64 gspHeapOffsetMin,
                                   NvU64 gspHeapOffsetMax, NvU32 partitionCount,
                                   NvBool isLeftPartition)
{
    VGPU_TYPE *pVgpuTypeInfo;
    NvU32 heterogeneousMaxInstance, i;
    NvBool isCarveOutGspHeap = NV_FALSE;
    NvU64 newVmmuSegMin, newVmmuSegMax, newGspHeapOffsetMin, newGspHeapOffsetMax, vmmuSegSize;
    NvBool isDefineLeftPlacement, isDefineRightPlacement;
    VGPU_TYPE_SUPPORTED_PLACEMENT_INFO *pVgpuTypeSupportedPlacementInfo;
    VGPU_INSTANCE_SUPPORTED_PLACEMENT_INFO *pVgpuInstanceSupportedPlacementInfo;
    KERNEL_VGPU_TYPE_PLACEMENT_INFO *pKernelVgpuTypePlacementInfo;
    VGPU_TYPE_PLACEMENT_INFO *pVgpuTypePlacementInfo;
    VGPU_INSTANCE_PLACEMENT_INFO *pVgpuInstancePlacementInfo;

    /*
     * As max vGPU per GPU is 32, the minimum partition for 2 vGPUs will be 16.
     * So, if this is a 1/16 partition, don't recurse further
     */
    if (partitionCount > MAX_VGPU_DEVICES_PER_PGPU / 2)
        return;

    pKernelVgpuTypePlacementInfo = &pPgpuInfo->kernelVgpuTypePlacementInfo;

    vmmuSegSize = gpuGetVmmuSegmentSize(pGpu);
    for (i = 0; i < pPgpuInfo->numVgpuTypes; i++)
    {
        pVgpuTypeInfo = pPgpuInfo->vgpuTypes[i];

        pVgpuTypeSupportedPlacementInfo = &pVgpuTypeInfo->vgpuTypeSupportedPlacementInfo;
        pVgpuTypePlacementInfo = &pKernelVgpuTypePlacementInfo->vgpuTypePlacementInfo[i];

        if (pVgpuTypeSupportedPlacementInfo->placementSize == 0)
            continue;

        isDefineLeftPlacement = NV_FALSE;
        isDefineRightPlacement = NV_FALSE;

        heterogeneousMaxInstance = nvPrevPow2_U32(pVgpuTypeInfo->maxInstance);

        /*
         * If homogeneous maxInstances of vGPU type are in power-of-2 and can fit in this partition,
         * then define both the placements in this recursive call
         */
        if (partitionCount * 2 == pVgpuTypeInfo->maxInstance)
        {
            isDefineLeftPlacement = NV_TRUE;
            isDefineRightPlacement = NV_TRUE;

            /*
             * Check if the next recursive partition will be a smaller partition for which
             * we need to carve out GSP heap at the end of smaller partition.
             */
            if (IS_GSP_CLIENT(pGpu) && (isCarveOutGspHeap == NV_FALSE) &&
                (pVgpuTypeInfo->gspHeapSize == vmmuSegSize))
                isCarveOutGspHeap = NV_TRUE;
        }
        else if (((heterogeneousMaxInstance < pVgpuTypeInfo->maxInstance) &&
                  (heterogeneousMaxInstance == partitionCount)) ||
                 ((partitionCount == 1) && (pVgpuTypeInfo->maxInstance == 1)))
        {
            /*
             * If only one instance of vGPU type can fit in this partition, then define
             * placement depending on whether it is a left or right partition
             */
            isDefineLeftPlacement = isLeftPartition;
            isDefineRightPlacement = !isLeftPartition;
        }

        if (isDefineLeftPlacement == NV_TRUE)
        {
            /*
             * Fill placement ID / channel / FB offset of first vGPU in this partition
             * First vGPU is aligned to the starting of partition
             */
            if (_kvgpumgrIsPlacementValid(pGpu, pPgpuInfo, i, placementIdMin))
            {

                pVgpuInstanceSupportedPlacementInfo =
                    &pVgpuTypeSupportedPlacementInfo->vgpuInstanceSupportedPlacementInfo[pVgpuTypeSupportedPlacementInfo->heterogeneousPlacementCount];
                pVgpuInstancePlacementInfo =
                    &pVgpuTypePlacementInfo->vgpuInstancePlacementInfo[pVgpuTypeSupportedPlacementInfo->heterogeneousPlacementCount];

                pVgpuInstanceSupportedPlacementInfo->heterogeneousSupportedPlacementId = placementIdMin;
                pVgpuInstanceSupportedPlacementInfo->heterogeneousSupportedChidOffset = chidMin;

                if (!IS_GSP_CLIENT(pGpu))
                    pVgpuInstancePlacementInfo->heterogeneousSupportedVmmuOffset = vmmuSegMin;
                else
                {
                    /* If profile occupies > 1 vMMU segment for GSP heap, place heap at end of vGPU FB. */
                    if (pVgpuTypeInfo->gspHeapSize >= vmmuSegSize)
                    {
                        pVgpuInstancePlacementInfo->heterogeneousSupportedVmmuOffset = vmmuSegMin;
                        pVgpuInstancePlacementInfo->heterogeneousGspHeapOffset =
                            (vmmuSegMin + pVgpuTypePlacementInfo->guestVmmuCount) * vmmuSegSize;
                    }
                    else
                    {
                        /*
                         * If we're in smaller parition, GSP vMMU segment will be input to the function.
                         * Place the gsp heap of first vGPU at the starting of the GSP vMMU segment.
                         */
                        pVgpuInstancePlacementInfo->heterogeneousSupportedVmmuOffset = vmmuSegMin;
                        pVgpuInstancePlacementInfo->heterogeneousGspHeapOffset = gspHeapOffsetMin;
                    }

                    /* For 32:1 vGPU types, the GSP heap is placed between 2 consecutive guest vMMU segments */
                    if (heterogeneousMaxInstance == 32)
                    {
                        pVgpuInstancePlacementInfo->heterogeneousSupportedVmmuOffset = vmmuSegMin;
                        pVgpuInstancePlacementInfo->heterogeneousGspHeapOffset =
                            (vmmuSegMin + pVgpuTypePlacementInfo->guestVmmuCount) * vmmuSegSize;
                    }
                }

                pVgpuTypeSupportedPlacementInfo->heterogeneousPlacementCount++;
            }
        }

        if (isDefineRightPlacement == NV_TRUE)
        {
            /*
             * Fill placement ID / channel / FB offset of second vGPU in this partition
             * Second vGPU is aligned to the end of partition
             */
            if (_kvgpumgrIsPlacementValid(pGpu, pPgpuInfo, i,
                                          placementIdMax - pVgpuTypeSupportedPlacementInfo->placementSize))
            {

                pVgpuInstanceSupportedPlacementInfo =
                    &pVgpuTypeSupportedPlacementInfo->vgpuInstanceSupportedPlacementInfo[pVgpuTypeSupportedPlacementInfo->heterogeneousPlacementCount];
                pVgpuInstancePlacementInfo =
                    &pVgpuTypePlacementInfo->vgpuInstancePlacementInfo[pVgpuTypeSupportedPlacementInfo->heterogeneousPlacementCount];

                pVgpuInstanceSupportedPlacementInfo->heterogeneousSupportedPlacementId =
                    placementIdMax - pVgpuTypeSupportedPlacementInfo->placementSize;
                pVgpuInstanceSupportedPlacementInfo->heterogeneousSupportedChidOffset =
                    chidMax - pVgpuTypeSupportedPlacementInfo->channelCount;

                if (!IS_GSP_CLIENT(pGpu))
                    pVgpuInstancePlacementInfo->heterogeneousSupportedVmmuOffset = vmmuSegMax - pVgpuTypePlacementInfo->guestVmmuCount;
                else
                {
                    /* If profile occupies > 1 vMMU segment for GSP heap, place heap at end of vGPU FB. */
                    if (pVgpuTypeInfo->gspHeapSize >= vmmuSegSize)
                    {
                        pVgpuInstancePlacementInfo->heterogeneousSupportedVmmuOffset =
                            vmmuSegMax - pVgpuTypePlacementInfo->guestVmmuCount - (pVgpuTypeInfo->gspHeapSize / vmmuSegSize);
                        pVgpuInstancePlacementInfo->heterogeneousGspHeapOffset = (vmmuSegMax * vmmuSegSize) - pVgpuTypeInfo->gspHeapSize;
                    }
                    else
                    {
                        /*
                         * If we're in smaller parition, GSP vMMU segment will be input to the function.
                         * Place the gsp heap of first vGPU at the starting of the GSP vMMU segment.
                         */
                        pVgpuInstancePlacementInfo->heterogeneousSupportedVmmuOffset = vmmuSegMax - pVgpuTypePlacementInfo->guestVmmuCount;
                        pVgpuInstancePlacementInfo->heterogeneousGspHeapOffset = gspHeapOffsetMax - pVgpuTypeInfo->gspHeapSize;
                    }

                    /* For 32:1 vGPU types, the GSP heap is placed between 2 consecutive guest vMMU segments */
                    if (heterogeneousMaxInstance == 32)
                    {
                        pVgpuInstancePlacementInfo->heterogeneousSupportedVmmuOffset = vmmuSegMax - pVgpuTypePlacementInfo->guestVmmuCount;
                        pVgpuInstancePlacementInfo->heterogeneousGspHeapOffset =
                            ((vmmuSegMax - pVgpuTypePlacementInfo->guestVmmuCount) * vmmuSegSize) - pVgpuTypeInfo->gspHeapSize;
                    }
                }

                pVgpuTypeSupportedPlacementInfo->heterogeneousPlacementCount++;
            }
        }
    }

    /*
     * If the next recursive partition is for a left smaller partition which has GSP heap at
     * of start of partition, then update vmmuSegMin to reserve one segment at the
     * start of smaller partition. Also, init gsp min/max value for the reserved vMMU segment
     * at the start.
     */
    newVmmuSegMax = ((vmmuSegMin + vmmuSegMax) / 2);
    if (isCarveOutGspHeap)
    {
        NV_ASSERT((gspHeapOffsetMin == 0));

        newVmmuSegMin = vmmuSegMin + 1;
        newGspHeapOffsetMin = vmmuSegMin * vmmuSegSize;
        newGspHeapOffsetMax = newGspHeapOffsetMin + vmmuSegSize;
    }
    else
    {
        newVmmuSegMin = vmmuSegMin;
        newGspHeapOffsetMin = gspHeapOffsetMin;
        newGspHeapOffsetMax = (gspHeapOffsetMin + gspHeapOffsetMax) / 2;
    }

    /* Recursively call to get placment ID in left half of this partition */
    _kvgpumgrSetHeterogeneousResources(pGpu, pPgpuInfo, placementIdMin,
                             (placementIdMin + placementIdMax) / 2,
                             chidMin, (chidMin + chidMax) / 2, newVmmuSegMin,
                             newVmmuSegMax, newGspHeapOffsetMin, newGspHeapOffsetMax, partitionCount * 2,
                             NV_TRUE);

    /*
     * If the next recursive partition is for a right smaller partition which has GSP heap at
     * of end of partition, then update vmmuSegMax to reserve one segment at the
     * end of right partition. Also, init gsp min/max value for the reserved vMMU segment
     * at the end.
     */
    newVmmuSegMin = (vmmuSegMin + vmmuSegMax) / 2;
    if (isCarveOutGspHeap)
    {
        newVmmuSegMax = vmmuSegMax - 1;
        newGspHeapOffsetMin = newVmmuSegMax * vmmuSegSize;
        newGspHeapOffsetMax = newGspHeapOffsetMin + vmmuSegSize;
    }
    else
    {
        newVmmuSegMax = vmmuSegMax;
        newGspHeapOffsetMin = (gspHeapOffsetMin + gspHeapOffsetMax) / 2;;
        newGspHeapOffsetMax = gspHeapOffsetMax;
    }

    /* Recursively call to get placment ID in right half of this partition */
    _kvgpumgrSetHeterogeneousResources(pGpu, pPgpuInfo, (placementIdMin + placementIdMax) / 2,
                             placementIdMax, (chidMin + chidMax) / 2,
                             chidMax, newVmmuSegMin, newVmmuSegMax,
                             newGspHeapOffsetMin, newGspHeapOffsetMax, partitionCount * 2,
                             NV_FALSE);
}

/*
 * Determine the supported placement IDs for homogeneous vGPU profiles.
 */
static void
_kvgpumgrSetHomogeneousResources(OBJGPU *pGpu, KERNEL_PHYS_GPU_INFO *pPgpuInfo,
                                 NvU32 vmmuSegmentMin, NvU32 hostChannelCount)
{
    VGPU_TYPE *pVgpuTypeInfo;
    NvU32 i, j, chidMax, placementRegionIndexMax;
    VGPU_TYPE_SUPPORTED_PLACEMENT_INFO *pVgpuTypeSupportedPlacementInfo;
    VGPU_INSTANCE_SUPPORTED_PLACEMENT_INFO *pVgpuInstanceSupportedPlacementInfo;
    KERNEL_VGPU_TYPE_PLACEMENT_INFO *pKernelVgpuTypePlacementInfo;
    VGPU_TYPE_PLACEMENT_INFO *pVgpuTypePlacementInfo;
    VGPU_INSTANCE_PLACEMENT_INFO *pVgpuInstancePlacementInfo;

    if (!pGpu || !pPgpuInfo || !pPgpuInfo->homogeneousPlacementSupported) {
        return;
    }

    pKernelVgpuTypePlacementInfo = &pPgpuInfo->kernelVgpuTypePlacementInfo;

    for (i = 0; i < pPgpuInfo->numVgpuTypes; i++)
    {
        pVgpuTypeInfo = pPgpuInfo->vgpuTypes[i];
        chidMax = hostChannelCount;

        pVgpuTypeSupportedPlacementInfo = &pVgpuTypeInfo->vgpuTypeSupportedPlacementInfo;
        pVgpuTypePlacementInfo = &pKernelVgpuTypePlacementInfo->vgpuTypePlacementInfo[i];

        /*
         * Placement ID is not supported for vGPU MIG profiles. Hence, placementSize = 0
         * is already set for them in fn:kvgpumgrSetSupportedPlacementIds(). We use the
         * same below to skip those vGPU MIG profiles.
         */
        if (pVgpuTypeSupportedPlacementInfo->placementSize == 0)
            continue;

        placementRegionIndexMax = (pKernelVgpuTypePlacementInfo->placementRegionSize / pVgpuTypeSupportedPlacementInfo->placementSize);

        for (j = 0; j < placementRegionIndexMax; j++)
        {
            pVgpuInstanceSupportedPlacementInfo = &pVgpuTypeSupportedPlacementInfo->vgpuInstanceSupportedPlacementInfo[j];
            pVgpuInstancePlacementInfo = &pVgpuTypePlacementInfo->vgpuInstancePlacementInfo[j];

            pVgpuInstanceSupportedPlacementInfo->homogeneousSupportedPlacementId = j * pVgpuTypeSupportedPlacementInfo->placementSize;
            pVgpuInstancePlacementInfo->homogeneousSupportedVmmuOffset = vmmuSegmentMin + (j * pVgpuTypePlacementInfo->guestVmmuCount);

            pVgpuInstanceSupportedPlacementInfo->homogeneousSupportedChidOffset = chidMax - pVgpuTypeSupportedPlacementInfo->channelCount;
            chidMax -= pVgpuTypeSupportedPlacementInfo->channelCount;

            pVgpuTypeSupportedPlacementInfo->homogeneousPlacementCount++;

            if (pVgpuTypeSupportedPlacementInfo->homogeneousPlacementCount == pVgpuTypeInfo->maxInstance)
                break;
        }

        /*
         * Since, vGPU homogeneous mode is default, copy the supported placement IDs to
         * creatable placement IDs.
         */
        for (j = 0; j < pVgpuTypeSupportedPlacementInfo->homogeneousPlacementCount; j++)
        {
            pVgpuInstanceSupportedPlacementInfo = &pVgpuTypeSupportedPlacementInfo->vgpuInstanceSupportedPlacementInfo[j];

            pVgpuTypePlacementInfo->vgpuInstancePlacementInfo[j].creatablePlacementId =
                pVgpuInstanceSupportedPlacementInfo->homogeneousSupportedPlacementId;
        }
    }
}

/*
 * Calculate Max (worst case) Total GSP plugin heap size among vGPU types for Blackwell+ vGPUs.
 * In Blackwell+ with SMG,max total gsp plugin heap size will be calculated within the GI.
 * For Blackwell+ vGPUs, for 1:1 profile, GSP plugin heap size is 512MB.
 * However for N:1 profiles, total FB consumption for GSP plugin heap of
 * the running vGPUs can be more than 512MB.
 * E.g. 32:1 profile, 32MB heap size per instance = 1024MB total heap size.
 * So find out max or worst case Total GSP plugin heap size.
 */
static NvU64
_kvgpumgrGetMaxTotalGspHeapSize(OBJGPU *pGpu, KERNEL_PHYS_GPU_INFO *pPgpuInfo, NvU32 giPartitionFlag)
{
    NvU32 i;
    NvU64 totalGspHeapSize;
    NvU64 maxTotalGspHeapSize = 0;
    VGPU_TYPE *pVgpuTypeInfo;
    NvU32 placementRegionSize, placementSize;
    VGPU_TYPE_SUPPORTED_PLACEMENT_INFO *pVgpuTypeSupportedPlacementInfo;

    /* Calculate Max Total GSP Plugin heap size */
    for (i = 0; i < pPgpuInfo->numVgpuTypes; i++)
    {
        pVgpuTypeInfo = pPgpuInfo->vgpuTypes[i];
        if (pVgpuTypeInfo == NULL)
            break;

        if (IS_MIG_IN_USE(pGpu) && kvgpumgrIsMigTimeslicingModeEnabled(pGpu))
        {
            NvU32 partitionFlag;

            // Ignore non-MIG vgpu types
            if (pVgpuTypeInfo->gpuInstanceSize == 0)
                continue;
  
            kvgpumgrGetPartitionFlag(pVgpuTypeInfo->vgpuTypeId, &partitionFlag); 
            // Ignore vgpu types that does not belong to the current GI
            if (partitionFlag != giPartitionFlag)
                continue;

        }
        /* Ignore MIG vGPUs if MIG is not in use or mig timeslicing is not supported */
        else if (pVgpuTypeInfo->gpuInstanceSize != 0)
            continue;

        pVgpuTypeSupportedPlacementInfo = &pVgpuTypeInfo->vgpuTypeSupportedPlacementInfo;

        placementSize = pVgpuTypeSupportedPlacementInfo->placementSize;
        placementRegionSize = vgpuMgrGetPlacementRegionSize(pGpu, &pPgpuInfo->vgpuNominalFbSize, pVgpuTypeInfo);

        totalGspHeapSize = NV_ALIGN_UP((pVgpuTypeInfo->gspHeapSize * placementRegionSize / placementSize), RM_PAGE_SIZE_HUGE);

        if (maxTotalGspHeapSize < totalGspHeapSize)
            maxTotalGspHeapSize = totalGspHeapSize;
    }

    return maxTotalGspHeapSize;
}

static NvU32
_kvgpumgrGetMemDivFromPartitionFlag(NvU32 memSizeFlag)
{
    NvU32 memDiv = 0;

    switch (memSizeFlag)
    {
        case NV2080_CTRL_GPU_PARTITION_FLAG_MEMORY_SIZE_FULL:
            memDiv = 1;
            break;
        case NV2080_CTRL_GPU_PARTITION_FLAG_MEMORY_SIZE_HALF:
            memDiv = 2;
            break;
        case NV2080_CTRL_GPU_PARTITION_FLAG_MEMORY_SIZE_QUARTER:
            memDiv = 4;
            break;
        case NV2080_CTRL_GPU_PARTITION_FLAG_MEMORY_SIZE_EIGHTH:
            memDiv = 8;
            break;
        default:
            NV_ASSERT(0);
            NV_PRINTF(LEVEL_ERROR, "Unrecognized GPU mem partitioning flag 0x%x\n",
                      memSizeFlag);
            break;
    }
    return memDiv;
}

/* Set Resources for Heterogeneous Blackwell+ vGPUs */
static void
_kvgpumgrSetHeterogeneousResources_GB100(OBJGPU *pGpu,
                                         KERNEL_PHYS_GPU_INFO *pPgpuInfo,
                                         NvU32 vmmuSegmentMin,
                                         NvU32 totalVmmuCount,
                                         NvU32 hostChannelCount)
{
    NvU32 i, j;
    NvU64 vmmuSegSize;
    NvU32 guestMaxVmmuCount;
    NvU64 maxTotalGspHeapSize;
    NvU64 gspHeapStartOffset;
    NvU64 gspHeapOffset;
    VGPU_TYPE *pVgpuTypeInfo;
    KERNEL_VGPU_TYPE_PLACEMENT_INFO *pKernelVgpuTypePlacementInfo;
    VGPU_TYPE_PLACEMENT_INFO *pVgpuTypePlacementInfo;
    VGPU_INSTANCE_PLACEMENT_INFO *pVgpuInstancePlacementInfo;
    VGPU_TYPE_SUPPORTED_PLACEMENT_INFO *pVgpuTypeSupportedPlacementInfo;
    VGPU_INSTANCE_SUPPORTED_PLACEMENT_INFO *pVgpuInstanceSupportedPlacementInfo;

    /*
     * For Blackwell+ GPUs, for 1:1 profile, plugin GSP heap size is 512MB.
     * However for N:1 profiles, total FB consumption for GSP plugin heap of
     * the running vGPUs can be more than 512MB.
     * So for resource offset calculations, we consider max available FB for guests
     * after accounting for max total GSP heap size.
     */
    maxTotalGspHeapSize = _kvgpumgrGetMaxTotalGspHeapSize(pGpu, pPgpuInfo, 0);
    vmmuSegSize = gpuGetVmmuSegmentSize(pGpu);
    guestMaxVmmuCount = totalVmmuCount - ((maxTotalGspHeapSize + vmmuSegSize - 1) / vmmuSegSize);

    /* GSP Heap Start offset for N:1 */
    gspHeapStartOffset = (vmmuSegmentMin + guestMaxVmmuCount) * vmmuSegSize;

    pKernelVgpuTypePlacementInfo = &pPgpuInfo->kernelVgpuTypePlacementInfo;

    for (i = 0; i < pPgpuInfo->numVgpuTypes; i++)
    {
        NvU32 maxInstance;
        NvU32 placementRegionSize, placementId;

        pVgpuTypeInfo = pPgpuInfo->vgpuTypes[i];
        if (pVgpuTypeInfo == NULL)
            break;

        pVgpuTypeSupportedPlacementInfo = &pVgpuTypeInfo->vgpuTypeSupportedPlacementInfo;
        pVgpuTypePlacementInfo = &pKernelVgpuTypePlacementInfo->vgpuTypePlacementInfo[i];

        /* Skip MIG vGPUs profiles (pre-GB20x, MIG placementSize is set to 0) */
        if (pVgpuTypeSupportedPlacementInfo->placementSize == 0)
            continue;

        /* Set placementIds and resource offsets for each instance of this vGPU type */
        pVgpuTypeSupportedPlacementInfo->heterogeneousPlacementCount = 0;
        if (pGpu->getProperty(pGpu, PDB_PROP_GPU_MIG_TIMESLICING_SUPPORTED) && pVgpuTypeInfo->gpuInstanceSize != 0)
        {
            maxInstance = pVgpuTypeInfo->maxInstancePerGI;
        }
        else
        {
            maxInstance = pVgpuTypeInfo->maxInstance;
        }

        placementRegionSize = vgpuMgrGetPlacementRegionSize(pGpu, &pPgpuInfo->vgpuNominalFbSize, pVgpuTypeInfo);

        for (j = 0; j < maxInstance; j++)
        {
            pVgpuInstanceSupportedPlacementInfo =
                                &pVgpuTypeSupportedPlacementInfo->vgpuInstanceSupportedPlacementInfo[j];
            pVgpuInstancePlacementInfo = &pVgpuTypePlacementInfo->vgpuInstancePlacementInfo[j];

            pVgpuInstanceSupportedPlacementInfo->heterogeneousSupportedPlacementId =
                    (placementRegionSize * j) / maxInstance;
            placementId = pVgpuInstanceSupportedPlacementInfo->heterogeneousSupportedPlacementId;
            pVgpuInstanceSupportedPlacementInfo->heterogeneousSupportedChidOffset =
                                (hostChannelCount * placementId) / placementRegionSize;

            pVgpuTypeSupportedPlacementInfo->heterogeneousPlacementCount++;

           /*
            * Skip setting pVgpuInstancePlacementInfo here for MIG vGpu profiles as
            * it will get set during actual GI creation.
            */

            if (pVgpuTypeInfo->gpuInstanceSize) // MIG vGPU
                continue;

            if (maxInstance == 1)
            {
                pVgpuInstancePlacementInfo->heterogeneousSupportedVmmuOffset = vmmuSegmentMin;

                pVgpuInstancePlacementInfo->heterogeneousGspHeapOffset =
                                (vmmuSegmentMin + pVgpuTypePlacementInfo->guestVmmuCount) * vmmuSegSize;
            }
            else
            {
                pVgpuInstancePlacementInfo->heterogeneousSupportedVmmuOffset =
                                vmmuSegmentMin + (guestMaxVmmuCount * placementId) / placementRegionSize;

                gspHeapOffset = (maxTotalGspHeapSize * placementId) / placementRegionSize;
                pVgpuInstancePlacementInfo->heterogeneousGspHeapOffset =
                                gspHeapStartOffset + NV_ALIGN_DOWN(gspHeapOffset, RM_PAGE_SIZE_HUGE);
            }

        }
    }

    return;
}

/* Set Resources for Homogeneous Blackwell+ vGPUs */
static void
_kvgpumgrSetHomogeneousResources_GB100(OBJGPU *pGpu,
                                       KERNEL_PHYS_GPU_INFO *pPgpuInfo)
{
    NvU32 i, j;
    VGPU_TYPE *pVgpuTypeInfo;
    KERNEL_VGPU_TYPE_PLACEMENT_INFO *pKernelVgpuTypePlacementInfo;
    VGPU_TYPE_PLACEMENT_INFO *pVgpuTypePlacementInfo;
    VGPU_INSTANCE_PLACEMENT_INFO *pVgpuInstancePlacementInfo;
    VGPU_TYPE_SUPPORTED_PLACEMENT_INFO *pVgpuTypeSupportedPlacementInfo;
    VGPU_INSTANCE_SUPPORTED_PLACEMENT_INFO *pVgpuInstanceSupportedPlacementInfo;

    /*
     * For Blackwell+ GPUs, placementIds, max supported instances and resource offsets are
     * same for Homogeneous and Heterogeneous vGPUs.
     * So copy placementIds and resource offsets from Heterogeneous.
     */
    if (!pPgpuInfo->heterogeneousTimesliceSizesSupported)
    {
        pPgpuInfo->homogeneousPlacementSupported = NV_FALSE;
        NV_ASSERT_OR_RETURN_VOID(pPgpuInfo->homogeneousPlacementSupported);
    }

    pKernelVgpuTypePlacementInfo = &pPgpuInfo->kernelVgpuTypePlacementInfo;

    for (i = 0; i < pPgpuInfo->numVgpuTypes; i++)
    {
        pVgpuTypeInfo = pPgpuInfo->vgpuTypes[i];
        if (pVgpuTypeInfo == NULL)
            break;

        pVgpuTypeSupportedPlacementInfo = &pVgpuTypeInfo->vgpuTypeSupportedPlacementInfo;
        pVgpuTypePlacementInfo = &pKernelVgpuTypePlacementInfo->vgpuTypePlacementInfo[i];

        /* Skip MIG vGPUs profiles (pre-GB20x, MIG placementSize is set to 0) */
        if (pVgpuTypeSupportedPlacementInfo->placementSize == 0)
            continue;

        /* Copy placementIds and resource offsets for each instance of this vGPU type from Heterogeneous */
        pVgpuTypeSupportedPlacementInfo->homogeneousPlacementCount = 0;
        for (j = 0; j < pVgpuTypeSupportedPlacementInfo->heterogeneousPlacementCount; j++)
        {
            pVgpuInstanceSupportedPlacementInfo =
                                &pVgpuTypeSupportedPlacementInfo->vgpuInstanceSupportedPlacementInfo[j];
            pVgpuInstancePlacementInfo = &pVgpuTypePlacementInfo->vgpuInstancePlacementInfo[j];

            pVgpuInstanceSupportedPlacementInfo->homogeneousSupportedPlacementId =
                                pVgpuInstanceSupportedPlacementInfo->heterogeneousSupportedPlacementId;

            pVgpuInstanceSupportedPlacementInfo->homogeneousSupportedChidOffset =
                                pVgpuInstanceSupportedPlacementInfo->heterogeneousSupportedChidOffset;

            pVgpuInstancePlacementInfo->homogeneousSupportedVmmuOffset =
                                pVgpuInstancePlacementInfo->heterogeneousSupportedVmmuOffset;

            pVgpuTypeSupportedPlacementInfo->homogeneousPlacementCount++;
        }

        /*
         * Since, vGPU homogeneous mode is default, copy the supported placement IDs to
         * creatable placement IDs.
         */
        for (j = 0; j < pVgpuTypeSupportedPlacementInfo->homogeneousPlacementCount; j++)
        {
            pVgpuInstanceSupportedPlacementInfo =
                                &pVgpuTypeSupportedPlacementInfo->vgpuInstanceSupportedPlacementInfo[j];

            pVgpuTypePlacementInfo->vgpuInstancePlacementInfo[j].creatablePlacementId =
                                pVgpuInstanceSupportedPlacementInfo->homogeneousSupportedPlacementId;
        }
    }

    return;
}

static NvU32
_kvgpumgrGetplacementSize(NvU64 profileSize, NvU32 placementRegionSize)
{
    NvU32 placementSize = (NvU32) (profileSize / (1024 * 1024 * 1024));

    /*
       Report placementSize as 3 for 2GB profile on 96GB GPU (RTX PRO 6000 Blackwell SE)
       so that each vgpu gets atleast 64 channels because only 32 instances of 2GB are
       supported even though it can accommodate 48 vgpus.
    */    
    if (placementRegionSize == 96 && placementSize == 2)
        return 3;
    else
        return placementSize;
}

NV_STATUS
kvgpumgrSetSupportedPlacementIds(OBJGPU *pGpu)
{
    OBJSYS     *pSys     = SYS_GET_INSTANCE();
    KernelVgpuMgr *pKernelVgpuMgr = SYS_GET_KERNEL_VGPUMGR(pSys);
    KernelFifo *pKernelFifo = GPU_GET_KERNEL_FIFO(pGpu);
    MemoryManager *pMemoryManager = GPU_GET_MEMORY_MANAGER(pGpu);
    NV_STATUS rmStatus = NV_OK;
    KERNEL_PHYS_GPU_INFO *pPgpuInfo;
    NvU32 index, i, j;
    VGPU_TYPE *pVgpuTypeInfo;
    NvU32 hostChannelCount = 0, totalVmmuCount = 0;
    NvU64 totalReservedFb, vgpuReservedFb, guestFbLength, totalAvailableFb, totalRequiredFb;
    NvU64 gspHeapOffsetMax, gspHeapOffsetMin, pmaRegionLength, pmaBaseOffset, largestOffset;
    NvU64 vmmuOffsetMin, vmmuOffsetMax, vmmuSegmentMin, vmmuSegmentMax;
    Heap *pHeap = GPU_GET_HEAP(pGpu);
    VGPU_TYPE_SUPPORTED_PLACEMENT_INFO *pVgpuTypeSupportedPlacementInfo;
    VGPU_INSTANCE_SUPPORTED_PLACEMENT_INFO *pVgpuInstanceSupportedPlacementInfo;
    KERNEL_VGPU_TYPE_PLACEMENT_INFO *pKernelVgpuTypePlacementInfo;
    VGPU_TYPE_PLACEMENT_INFO *pVgpuTypePlacementInfo;

    NV_CHECK_OK_OR_RETURN(LEVEL_ERROR,
            kvgpumgrGetPgpuIndex(pKernelVgpuMgr, pGpu->gpuId, &index));

    pPgpuInfo = &(pKernelVgpuMgr->pgpuInfo[index]);

    pPgpuInfo->heterogeneousTimesliceSizesSupported = NV_FALSE;
    pPgpuInfo->homogeneousPlacementSupported = NV_FALSE;

    /* Heterogeneous vgpus enabled for SRIOV vGPUs on VMware, Device VM and KVM */
    if (gpuIsSriovEnabled(pGpu)
        )
    {
        if ((osIsVgpuVfioPresent() == NV_OK) ||
            (osIsVgpuDeviceVmPresent() == NV_OK))
        {
            pPgpuInfo->heterogeneousTimesliceSizesSupported = NV_TRUE;
        }
    }

    /*
     * We shall support Homogeneous vGPU placement ID feature only for SR-IOV vGPUs
     * with GSP. ADA vGPU qualifies this requirement. Hence, enable Homoogeneous vGPU
     * placement ID support for only ADA+.
     */
    if (IsADAorBetter(pGpu))
    {
        /*
         * Enable Homogeneous vgpu only on vGPUs on VMware, Device VM and KVM.
         * This will align it with Heterogeneous vgpu support.
         */
        if ((osIsVgpuVfioPresent() == NV_OK) ||
            (osIsVgpuDeviceVmPresent() == NV_OK))
        {
            pPgpuInfo->homogeneousPlacementSupported = NV_TRUE;
        }
    }

    if (!pPgpuInfo->heterogeneousTimesliceSizesSupported && !pPgpuInfo->homogeneousPlacementSupported)
        return rmStatus;

    /* skip placement Id calculation if already set */
    if (pPgpuInfo->isPlacementIdInfoSet)
        return rmStatus;

    hostChannelCount = kfifoChidMgrGetNumChannels(pGpu, pKernelFifo, pKernelFifo->ppChidMgr[0]);

    pKernelVgpuTypePlacementInfo = &pPgpuInfo->kernelVgpuTypePlacementInfo;

    for (i = 0; i < pPgpuInfo->numVgpuTypes; i++)
    {
        NvU32 placementRegionSize;
        pVgpuTypeInfo = pPgpuInfo->vgpuTypes[i];
        if (pVgpuTypeInfo == NULL)
            break;

        placementRegionSize = vgpuMgrGetPlacementRegionSize(pGpu, &pPgpuInfo->vgpuNominalFbSize, pVgpuTypeInfo);
        pVgpuTypeSupportedPlacementInfo = &pVgpuTypeInfo->vgpuTypeSupportedPlacementInfo;
        pVgpuTypePlacementInfo = &pKernelVgpuTypePlacementInfo->vgpuTypePlacementInfo[i];


        for (j = 0; j < MAX_VGPU_DEVICES_PER_PGPU; j++)
        {
            /* Set to NVA081_PLACEMENT_ID_INVALID initially */
            pVgpuInstanceSupportedPlacementInfo = &pVgpuTypeSupportedPlacementInfo->vgpuInstanceSupportedPlacementInfo[j];

            pVgpuInstanceSupportedPlacementInfo->heterogeneousSupportedPlacementId = NVA081_PLACEMENT_ID_INVALID;
            pVgpuInstanceSupportedPlacementInfo->homogeneousSupportedPlacementId   = NVA081_PLACEMENT_ID_INVALID;

        }

        /* Ignore MIG vGPUs if timeslice is not supported */
        if (pVgpuTypeInfo->gpuInstanceSize != 0 && !pGpu->getProperty(pGpu, PDB_PROP_GPU_MIG_TIMESLICING_SUPPORTED))
        {
            pVgpuTypeSupportedPlacementInfo->placementSize = 0;
            continue;
        }

        pVgpuTypeSupportedPlacementInfo->heterogeneousPlacementCount = 0;
        pVgpuTypeSupportedPlacementInfo->homogeneousPlacementCount = 0;
        pVgpuTypePlacementInfo->guestVmmuCount = 0;

        pVgpuTypeSupportedPlacementInfo->placementSize = _kvgpumgrGetplacementSize(pVgpuTypeInfo->profileSize, placementRegionSize);

        if (gpuIsSriovEnabled(pGpu))
        {
            NvU32 channelDivisor;

            if(pVgpuTypeInfo->gpuInstanceSize) // MIG vGPU
                channelDivisor = pVgpuTypeInfo->maxInstancePerGI;
            else
                channelDivisor = pVgpuTypeInfo->maxInstance;

            if (gpuIsNonPowerOf2ChannelCountSupported(pGpu))
                pVgpuTypeSupportedPlacementInfo->channelCount = (hostChannelCount * pVgpuTypeSupportedPlacementInfo->placementSize) /
                                                             placementRegionSize;
            else
                pVgpuTypeSupportedPlacementInfo->channelCount = nvPrevPow2_U32(hostChannelCount /
                                                             channelDivisor);

            /*
             * For sriov vGPUs with GSP disabled, guest FB length is only dependent on
             * profile fb length and RM reserved fb (including ECC). Calculate the
             * guest fb length here by reducing reserved FB from profile fb length
             * depending on max instance and aligning it to VMMU segment size.
             */
            if (!IS_GSP_CLIENT(pGpu))
            {
                totalReservedFb = memmgrGetVgpuHostRmReservedFb_HAL(pGpu,
                                    pMemoryManager, pVgpuTypeInfo->vgpuTypeId) +
                                  kvgpumgrGetEccAndPrReservedFb(pGpu);

                vgpuReservedFb = NV_ALIGN_UP((totalReservedFb / pVgpuTypeInfo->maxInstance),
                                     gpuGetVmmuSegmentSize(pGpu));

                /*
                 * For 1GB profiles on 48GB board, maxInstances is limited to 32. For such
                 * profiles, there is no per vgpu reserved fb carved out from guest fb as
                 * host has space for reserved fb.
                 */
                if (pVgpuTypeInfo->profileSize == VGPU_LOWER_FB_PROFILE_SIZE_1_GB)
                {
                    heapGetUsableSize(pHeap, &totalAvailableFb);

                    /* Total FB needed if vgpu reserved fb is 0 and not carved from guest fb */
                    totalRequiredFb = pVgpuTypeInfo->maxInstance *
                                      (VGPU_LOWER_FB_PROFILE_SIZE_1_GB + vgpuReservedFb);

                    if (totalRequiredFb < totalAvailableFb)
                        vgpuReservedFb = 0;
                }

                guestFbLength = pVgpuTypeInfo->profileSize - vgpuReservedFb;

                pVgpuTypePlacementInfo->guestVmmuCount = guestFbLength / gpuGetVmmuSegmentSize(pGpu);
            }
            else
            {
                if (IsGB100orBetter(pGpu))
                {
                    /* Get the guest FB from vGPU type */
                    /* TODO: Handle ECC enabled scenario */
                    pVgpuTypePlacementInfo->guestVmmuCount = pVgpuTypeInfo->fbLength / gpuGetVmmuSegmentSize(pGpu);
                }
                else
                {
                    /* Calculate the guest FB using similar calculation as done in vGPU plugin */
                    totalAvailableFb = pMemoryManager->Ram.fbTotalMemSizeMb << 20;
                    totalAvailableFb = NV_ALIGN_UP(totalAvailableFb, 8 * gpuGetVmmuSegmentSize(pGpu));

                    guestFbLength = (totalAvailableFb / pVgpuTypeInfo->maxInstance) -
                                    pVgpuTypeInfo->fbReservation - pVgpuTypeInfo->gspHeapSize;
                    guestFbLength = NV_MIN(guestFbLength, pVgpuTypeInfo->fbLength);
                    guestFbLength = NV_ALIGN_DOWN(guestFbLength, gpuGetVmmuSegmentSize(pGpu));

                    pVgpuTypePlacementInfo->guestVmmuCount = guestFbLength / gpuGetVmmuSegmentSize(pGpu);
                }
            }
        }

        if (pVgpuTypeInfo->maxInstance == 1)
        {
            pKernelVgpuTypePlacementInfo->placementRegionSize = 
                vgpuMgrGetPlacementRegionSize(pGpu, &pPgpuInfo->vgpuNominalFbSize, pVgpuTypeInfo);

            if (gpuIsSriovEnabled(pGpu))
            {
                if (pVgpuTypePlacementInfo->guestVmmuCount)
                {
                    totalVmmuCount = pVgpuTypePlacementInfo->guestVmmuCount;

                    if (IS_GSP_CLIENT(pGpu))
                        totalVmmuCount += (pVgpuTypeInfo->gspHeapSize / gpuGetVmmuSegmentSize(pGpu));
                }
            }
        }
    }

    /*
     * For SRIOV, the placement IDs are aligned to 1/2, 1/4, 1/8, 1/16 partitions
     * due to restrictions on channels assigned to VF being in power-of-2.
     */
    if (gpuIsSriovEnabled(pGpu))
    {
        pmaGetLargestFree(pHeap->pPmaObject, &pmaRegionLength,
                          &pmaBaseOffset, &largestOffset);

        vmmuOffsetMin = NV_ALIGN_UP(pmaBaseOffset + largestOffset, gpuGetVmmuSegmentSize(pGpu));
        vmmuOffsetMax = pmaBaseOffset + largestOffset + pmaRegionLength;

        totalRequiredFb = (totalVmmuCount * gpuGetVmmuSegmentSize(pGpu));
        if ((vmmuOffsetMax - vmmuOffsetMin) < totalRequiredFb)
        {
            NV_PRINTF(LEVEL_ERROR, "Required FB for heterogeneous vGPU (%llu) less "
                      "than available FB (%llu)\n", totalRequiredFb, pmaRegionLength);
            NV_ASSERT(0);

            // Disable heterogeneous vGPU support, but allow homogeneous placement ID vGPU support.
            pPgpuInfo->heterogeneousTimesliceSizesSupported = NV_FALSE;
        }

        vmmuSegmentMin = vmmuOffsetMin / gpuGetVmmuSegmentSize(pGpu);
        vmmuSegmentMax = vmmuSegmentMin + totalVmmuCount;

        gspHeapOffsetMin = 0;
        gspHeapOffsetMax = 0;

        if (pPgpuInfo->heterogeneousTimesliceSizesSupported)
        {
            if (IsGB100orBetter(pGpu))
                _kvgpumgrSetHeterogeneousResources_GB100(pGpu, pPgpuInfo, vmmuSegmentMin,
                                                         totalVmmuCount, hostChannelCount);
            else
                _kvgpumgrSetHeterogeneousResources(pGpu, pPgpuInfo, 0,
                                                   pKernelVgpuTypePlacementInfo->placementRegionSize,
                                                   0, hostChannelCount, vmmuSegmentMin, vmmuSegmentMax,
                                                   gspHeapOffsetMin, gspHeapOffsetMax, 1, NV_TRUE);
        }

        if (pPgpuInfo->homogeneousPlacementSupported)
        {
            if (IsGB100orBetter(pGpu))
                _kvgpumgrSetHomogeneousResources_GB100(pGpu, pPgpuInfo);
            else
                _kvgpumgrSetHomogeneousResources(pGpu, pPgpuInfo, vmmuSegmentMin, hostChannelCount);
        }
    }

    if (rmStatus == NV_OK)
        pPgpuInfo->isPlacementIdInfoSet = NV_TRUE;

    return rmStatus;
}

static NvBool
isPlacementOverlapping(NvU16 minId1, NvU16 maxId1, NvU16 id2, NvU16 size2)
{
    NvU16 max, min, size1;

    min = NV_MIN(minId1, id2);
    max = NV_MAX(maxId1, (id2 + size2));
    size1 = maxId1 - minId1;

    if ((size1 + size2) > (max - min))
        return NV_TRUE;

    return NV_FALSE;
}

static NvBool
isPlacementSubset(NvU16 min, NvU16 max, NvU16 id, NvU16 size)
{
    if ((min <= id) && (max >= (id + size - 1)))
        return NV_TRUE;

    return NV_FALSE;
}

static NV_STATUS
_kvgpumgrUpdateHomogeneousCreatablePlacementIds(OBJGPU *pGpu, NvU16 placementId,
                                                NvU32 vgpuTypeId, NvBool isAlloc,
                                                NvU32 swizzId)
{
    OBJSYS *pSys = SYS_GET_INSTANCE();
    KernelVgpuMgr *pKernelVgpuMgr = SYS_GET_KERNEL_VGPUMGR(pSys);
    KERNEL_PHYS_GPU_INFO *pPgpuInfo;
    VGPU_TYPE *pVgpuTypeInfo;
    NvU32 index, i , j, placementSize;
    NvU32 vgpuCount = 0;
    NvBool bIdFound = NV_FALSE, bPartitionMatch;;
    KERNEL_VGPU_TYPE_PLACEMENT_INFO *pKernelVgpuTypePlacementInfo;
    VGPU_TYPE_PLACEMENT_INFO *pVgpuTypePlacementInfo;
    VGPU_INSTANCE_PLACEMENT_INFO *pVgpuInstancePlacementInfo;
    VGPU_TYPE_SUPPORTED_PLACEMENT_INFO *pVgpuTypeSupportedPlacementInfo;
    VGPU_INSTANCE_SUPPORTED_PLACEMENT_INFO *pVgpuInstanceSupportedPlacementInfo;

    NV_CHECK_OK_OR_RETURN(LEVEL_ERROR,
            kvgpumgrGetPgpuIndex(pKernelVgpuMgr, pGpu->gpuId, &index));

    NV_CHECK_OK_OR_RETURN(LEVEL_ERROR,
            kvgpumgrGetVgpuTypeInfo(vgpuTypeId, &pVgpuTypeInfo));

    pPgpuInfo = &(pKernelVgpuMgr->pgpuInfo[index]);
    pKernelVgpuTypePlacementInfo = kvgpuMgrGetVgpuPlacementInfo(pGpu, pPgpuInfo, swizzId);
    NV_CHECK_OR_RETURN(LEVEL_ERROR, pKernelVgpuTypePlacementInfo != NULL, NV_ERR_OBJECT_NOT_FOUND);

    pVgpuTypeSupportedPlacementInfo = &pVgpuTypeInfo->vgpuTypeSupportedPlacementInfo;
    placementSize = pVgpuTypeSupportedPlacementInfo->placementSize;

    if (IS_MIG_IN_USE(pGpu))
    {
        NV_CHECK_OR_RETURN(LEVEL_ERROR, swizzId < KMIGMGR_MAX_GPU_SWIZZID, NV_ERR_INVALID_ARGUMENT);
        vgpuCount = pPgpuInfo->assignedSwizzIdVgpuCount[swizzId];
    }
    else
    {
        if (osIsVgpuVfioPresent() == NV_OK)
            vgpuCount = pPgpuInfo->numCreatedVgpu;
        else
            vgpuCount = pPgpuInfo->numActiveVgpu;
    }

    /* In Allocation path, check whether input placement ID is available to use. */
    if (isAlloc)
    {
        for (i = 0; i < pPgpuInfo->numVgpuTypes; i++)
        {
            pVgpuTypeInfo = pPgpuInfo->vgpuTypes[i];

            if (pVgpuTypeInfo == NULL)
                continue;

            pVgpuTypeSupportedPlacementInfo = &pVgpuTypeInfo->vgpuTypeSupportedPlacementInfo;
            pVgpuTypePlacementInfo = &pKernelVgpuTypePlacementInfo->vgpuTypePlacementInfo[i];

            if (pVgpuTypeSupportedPlacementInfo->placementSize != placementSize)
                continue;
            if (IS_MIG_IN_USE(pGpu))
            {
                if (!pVgpuTypeInfo->gpuInstanceSize) // Ignore non MIG vgpu types
                    continue;
                if (_kvgpumgrCheckPartitionFlag(pGpu, swizzId, pVgpuTypeInfo->vgpuTypeId, &bPartitionMatch) == NV_OK)
                {
                    if (!bPartitionMatch)
                        continue;
                }
            }
            else
            {
                if (pVgpuTypeInfo->gpuInstanceSize) // Ignore MIG vgpu types
                    continue;
            }
            for (j = 0; j < pVgpuTypeSupportedPlacementInfo->homogeneousPlacementCount; j++)
            {
                pVgpuInstancePlacementInfo = &pVgpuTypePlacementInfo->vgpuInstancePlacementInfo[j];

                if (pVgpuInstancePlacementInfo->creatablePlacementId == placementId)
                {
                    pVgpuInstancePlacementInfo->creatablePlacementId = NVA081_PLACEMENT_ID_INVALID;
                    bIdFound = NV_TRUE;
                    break;
                }
            }

            /* No creatable placement ID found */
            if (!bIdFound)
            {
                NV_PRINTF(LEVEL_ERROR, "Failed to find creatable placement id for vGPU type %d with placement id %d\n",
                    pVgpuTypeInfo->vgpuTypeId, placementId);
                return NV_ERR_INVALID_STATE;
            }
        }

        /*
         * If creatable placement ID found then, mark creatable placement IDs for other
         * vGPU Types with profile size not equal placement size of input to vgpuTypeId
         * as invalid.
         */
        for (i = 0; i < pPgpuInfo->numVgpuTypes; i++)
        {
            pVgpuTypeInfo = pPgpuInfo->vgpuTypes[i];

            if (pVgpuTypeInfo == NULL)
                continue;

            pVgpuTypeSupportedPlacementInfo = &pVgpuTypeInfo->vgpuTypeSupportedPlacementInfo;
            pVgpuTypePlacementInfo = &pKernelVgpuTypePlacementInfo->vgpuTypePlacementInfo[i];

            if (pVgpuTypeSupportedPlacementInfo->placementSize == placementSize)
                continue;

            if (IS_MIG_IN_USE(pGpu))
            {
                if (!pVgpuTypeInfo->gpuInstanceSize) // Ignore non MIG vgpu types
                    continue;
                if (_kvgpumgrCheckPartitionFlag(pGpu, swizzId, pVgpuTypeInfo->vgpuTypeId, &bPartitionMatch) == NV_OK)
                {
                    if (!bPartitionMatch)
                        continue;
                }
            }
            for (j = 0; j < pVgpuTypeSupportedPlacementInfo->homogeneousPlacementCount; j++)
            {
                pVgpuInstancePlacementInfo = &pVgpuTypePlacementInfo->vgpuInstancePlacementInfo[j];

                pVgpuInstancePlacementInfo->creatablePlacementId = NVA081_PLACEMENT_ID_INVALID;
            }
        }
    }
    /* In Free path, accumulate the freed placement ID */
    else
    {
        for (i = 0; i < pPgpuInfo->numVgpuTypes; i++)
        {
            pVgpuTypeInfo = pPgpuInfo->vgpuTypes[i];

            if (pVgpuTypeInfo == NULL)
                continue;

            pVgpuTypeSupportedPlacementInfo = &pVgpuTypeInfo->vgpuTypeSupportedPlacementInfo;
            pVgpuTypePlacementInfo = &pKernelVgpuTypePlacementInfo->vgpuTypePlacementInfo[i];

            /*
             * Restore all the creatable placement IDs with the supported placement IDs when
             * the last vGPU instance is freed.
             */
            if (vgpuCount == 0)
            {
                for (j = 0; j < pVgpuTypeSupportedPlacementInfo->homogeneousPlacementCount; j++)
                {
                    pVgpuInstancePlacementInfo = &pVgpuTypePlacementInfo->vgpuInstancePlacementInfo[j];
                    pVgpuInstanceSupportedPlacementInfo = &pVgpuTypeSupportedPlacementInfo->vgpuInstanceSupportedPlacementInfo[j];

                    pVgpuInstancePlacementInfo->creatablePlacementId = pVgpuInstanceSupportedPlacementInfo->homogeneousSupportedPlacementId;
                }
            }
            else
            {
                if (pVgpuTypeSupportedPlacementInfo->placementSize != placementSize)
                    continue;

                if (IS_MIG_IN_USE(pGpu))
                {
                    if (!pVgpuTypeInfo->gpuInstanceSize) // Ignore non MIG vgpu types
                        continue;
                    if (_kvgpumgrCheckPartitionFlag(pGpu, swizzId, pVgpuTypeInfo->vgpuTypeId, &bPartitionMatch) == NV_OK)
                    {
                        if (!bPartitionMatch)
                            continue;
                    }
                }
                else
                {
                    if (pVgpuTypeInfo->gpuInstanceSize)
                        continue;
                }

                /*
                 * Accumulate the freed placement ID for vGPU instance at its respective
                 * place where creatablePlacementIds slot is invalid.
                 */
                for (j = 0; j < pVgpuTypeSupportedPlacementInfo->homogeneousPlacementCount; j++)
                {
                    pVgpuInstancePlacementInfo = &pVgpuTypePlacementInfo->vgpuInstancePlacementInfo[j];
                    pVgpuInstanceSupportedPlacementInfo =
                        &pVgpuTypeSupportedPlacementInfo->vgpuInstanceSupportedPlacementInfo[j];

                    if ((pVgpuInstanceSupportedPlacementInfo->homogeneousSupportedPlacementId == placementId) &&
                        (pVgpuInstancePlacementInfo->creatablePlacementId == NVA081_PLACEMENT_ID_INVALID))
                    {
                        pVgpuInstancePlacementInfo->creatablePlacementId = placementId;
                        bIdFound = NV_TRUE;
                        break;
                    }
                }

                /* No creatable invalid placement ID found */
                if (!bIdFound)
                    return NV_ERR_INVALID_STATE;
            }
        }
    }

    return NV_OK;
}

static NV_STATUS
_kvgpumgrUpdateHeterogeneousCreatablePlacementIds(OBJGPU *pGpu, NvU16 placementId,
                                                  NvU32 vgpuTypeId, NvBool isAlloc,
                                                  NvU32 swizzId)
{
    OBJSYS *pSys = SYS_GET_INSTANCE();
    KernelVgpuMgr *pKernelVgpuMgr = SYS_GET_KERNEL_VGPUMGR(pSys);
    KERNEL_PHYS_GPU_INFO *pPgpuInfo;
    VGPU_TYPE *pVgpuTypeInfo;
    NvU32 index, i , j, placementSize, min = 0, max = 0, numCreatable = 0;
    NV_RANGE range;
    NvBool isVgpuTypeCreatable, bPartitionMatch;
    KERNEL_VGPU_TYPE_PLACEMENT_INFO *pKernelVgpuTypePlacementInfo;
    VGPU_TYPE_PLACEMENT_INFO *pVgpuTypePlacementInfo;
    VGPU_INSTANCE_PLACEMENT_INFO *pVgpuInstancePlacementInfo;
    VGPU_TYPE_SUPPORTED_PLACEMENT_INFO *pVgpuTypeSupportedPlacementInfo;
    VGPU_INSTANCE_SUPPORTED_PLACEMENT_INFO *pVgpuInstanceSupportedPlacementInfo;
    NvU32 *pSupportedTypeId;

    NV_CHECK_OK_OR_RETURN(LEVEL_ERROR,
            kvgpumgrGetPgpuIndex(pKernelVgpuMgr, pGpu->gpuId, &index));

    NV_CHECK_OK_OR_RETURN(LEVEL_ERROR,
            kvgpumgrGetVgpuTypeInfo(vgpuTypeId, &pVgpuTypeInfo));

    pPgpuInfo = &(pKernelVgpuMgr->pgpuInfo[index]);

    pKernelVgpuTypePlacementInfo = kvgpuMgrGetVgpuPlacementInfo(pGpu, pPgpuInfo, swizzId);
    NV_CHECK_OR_RETURN(LEVEL_ERROR, pKernelVgpuTypePlacementInfo != NULL, NV_ERR_OBJECT_NOT_FOUND);
    pSupportedTypeId = kvgpuMgrGetVgpuCreatableTypeIdFromSwizzId(pGpu, pPgpuInfo, swizzId);

    pVgpuTypeSupportedPlacementInfo = &pVgpuTypeInfo->vgpuTypeSupportedPlacementInfo;
    placementSize = pVgpuTypeSupportedPlacementInfo->placementSize;

    range = rangeMake(placementId, placementId + placementSize - 1);

    if (isAlloc)
    {
        /* Update the placment region bitmap to set (id + size) slots */
        bitVectorSetRange(&pKernelVgpuTypePlacementInfo->usedPlacementRegionMap, range);
        min = placementId;
        max = placementId + placementSize;
    }
    else
    {
        /* Update the placment region bitmap to unset (id + size) slots */
        bitVectorClrRange(&pKernelVgpuTypePlacementInfo->usedPlacementRegionMap, range);

        /* find the minimum slot before placement ID which is not occupied */
        min = placementId;
        for (j = placementId; j > 0;)
        {
            j--;

            if (bitVectorTest(&pKernelVgpuTypePlacementInfo->usedPlacementRegionMap, j))
                break;

            min = j;
        }

        /* find the maximum slot after placement ID + placementSize which is not occupied */
        max = placementId + placementSize - 1;
        for (j = placementId + placementSize; j < pKernelVgpuTypePlacementInfo->placementRegionSize; j++)
        {
            if (bitVectorTest(&pKernelVgpuTypePlacementInfo->usedPlacementRegionMap, j))
                break;

             max = j;
        }
    }

    for (i = 0; i < MAX_VGPU_TYPES_PER_PGPU; i++)
        pSupportedTypeId[i] = NVA081_CTRL_VGPU_CONFIG_INVALID_TYPE;

    /* Update creatable placement Ids of all vGPU type IDs */
    for (i = 0; i < pPgpuInfo->numVgpuTypes; i++)
    {
        pVgpuTypeInfo = pPgpuInfo->vgpuTypes[i];
        if (pVgpuTypeInfo == NULL)
            break;

        pVgpuTypeSupportedPlacementInfo = &pVgpuTypeInfo->vgpuTypeSupportedPlacementInfo;
        pVgpuTypePlacementInfo = &pKernelVgpuTypePlacementInfo->vgpuTypePlacementInfo[i];

        if (pVgpuTypeSupportedPlacementInfo->placementSize == 0)
            continue;

        if (IS_MIG_IN_USE(pGpu))
        {
            if (!pVgpuTypeInfo->gpuInstanceSize) // Ignore non MIG vgpu types
                continue;
            if (_kvgpumgrCheckPartitionFlag(pGpu, swizzId, pVgpuTypeInfo->vgpuTypeId, &bPartitionMatch) == NV_OK)
            {
                if (!bPartitionMatch)
                    continue;
            }
        }
        else
        {
            if (pVgpuTypeInfo->gpuInstanceSize) // Ignore MIG vgpu types
                continue;
        }
        isVgpuTypeCreatable = NV_FALSE;

        for (j = 0; j < pVgpuTypeSupportedPlacementInfo->heterogeneousPlacementCount; j++)
        {
            pVgpuInstanceSupportedPlacementInfo = &pVgpuTypeSupportedPlacementInfo->vgpuInstanceSupportedPlacementInfo[j];
            pVgpuInstancePlacementInfo = &pVgpuTypePlacementInfo->vgpuInstancePlacementInfo[j];

            /*
             * Creatable Placement IDs will be updated either on VM boot or VM shutdown.
             * 1. VM boot.
             *      Here, we compare every (id + size) of each vGPU type if it's partial/full
             *      overlapping with the (id + size) of the vGPU VM that is booting up.
             *      If yes, then mark that placement ID as INVALID.
             * 2. VM shutdown
             *      Here, we compare every INVALID id of each vGPU type and
             *      see if the corresponding supported (id + size) is a complete subset
             *      of (max - min) range of available placement slots.
             *      If it is subset, then the INVALID id is overwritten with
             *      corresponding supported id as it is now creatable.
             */
            if (pVgpuInstancePlacementInfo->creatablePlacementId != NVA081_PLACEMENT_ID_INVALID)
            {
                if (isAlloc && isPlacementOverlapping(min, max,
                                                      pVgpuInstancePlacementInfo->creatablePlacementId,
                                                      pVgpuTypeSupportedPlacementInfo->placementSize))
                {
                    pVgpuInstancePlacementInfo->creatablePlacementId = NVA081_PLACEMENT_ID_INVALID;
                }
            }
            else
            {
                if (!isAlloc && isPlacementSubset(min, max,
                                                  pVgpuInstanceSupportedPlacementInfo->heterogeneousSupportedPlacementId,
                                                  pVgpuTypeSupportedPlacementInfo->placementSize))
                {
                    pVgpuInstancePlacementInfo->creatablePlacementId =
                        pVgpuInstanceSupportedPlacementInfo->heterogeneousSupportedPlacementId;
                }
            }

            if (pVgpuInstancePlacementInfo->creatablePlacementId != NVA081_PLACEMENT_ID_INVALID)
                isVgpuTypeCreatable = NV_TRUE;
        }

        if (isVgpuTypeCreatable)
            pSupportedTypeId[numCreatable++] = pVgpuTypeInfo->vgpuTypeId;
    }

    return NV_OK;
}

NV_STATUS
kvgpumgrUpdateHeterogeneousInfo(OBJGPU *pGpu, Device *pDevice,
                                NvU32 vgpuTypeId, NvU16 *placementId,
                                NvU64 *guestFbLength, NvU64 *guestFbOffset,
                                NvU64 *gspHeapOffset, NvU64 *guestBar1PFOffset)
{
    OBJSYS *pSys = SYS_GET_INSTANCE();
    KernelVgpuMgr *pKernelVgpuMgr = SYS_GET_KERNEL_VGPUMGR(pSys);
    KERNEL_PHYS_GPU_INFO *pPgpuInfo;
    VGPU_TYPE *pVgpuTypeInfo;
    NvU32 index, i , j, swizzId = KMIGMGR_SWIZZID_INVALID;
    NvBool bIdFound = NV_FALSE;
    VGPU_TYPE_SUPPORTED_PLACEMENT_INFO *pVgpuTypeSupportedPlacementInfo;
    KERNEL_VGPU_TYPE_PLACEMENT_INFO *pKernelVgpuTypePlacementInfo;
    VGPU_TYPE_PLACEMENT_INFO *pVgpuTypePlacementInfo;
    VGPU_INSTANCE_PLACEMENT_INFO *pVgpuInstancePlacementInfo;

    NV_CHECK_OK_OR_RETURN(LEVEL_ERROR,
            kvgpumgrGetPgpuIndex(pKernelVgpuMgr, pGpu->gpuId, &index));

    pPgpuInfo = &(pKernelVgpuMgr->pgpuInfo[index]);

    if (IS_MIG_IN_USE(pGpu) && kvgpumgrIsMigTimeslicingModeEnabled(pGpu))
    {
        swizzId = kvgpuMgrGetSwizzIdFromDevice(pGpu, pDevice);
        pKernelVgpuTypePlacementInfo = kvgpuMgrGetVgpuPlacementInfo(pGpu, pPgpuInfo, swizzId);
        NV_CHECK_OR_RETURN(LEVEL_ERROR, pKernelVgpuTypePlacementInfo != NULL, NV_ERR_OBJECT_NOT_FOUND);
    }
    else
    {
        pKernelVgpuTypePlacementInfo = &pPgpuInfo->kernelVgpuTypePlacementInfo;
    }

    for (i = 0; i < pPgpuInfo->numVgpuTypes; i++)
    {
        pVgpuTypeInfo = pPgpuInfo->vgpuTypes[i];
        if (pVgpuTypeInfo == NULL)
            break;

        pVgpuTypeSupportedPlacementInfo = &pVgpuTypeInfo->vgpuTypeSupportedPlacementInfo;
        pVgpuTypePlacementInfo = &pKernelVgpuTypePlacementInfo->vgpuTypePlacementInfo[i];

        if (pVgpuTypeInfo->vgpuTypeId == vgpuTypeId)
        {
            for (j = 0; j < pVgpuTypeSupportedPlacementInfo->heterogeneousPlacementCount; j++)
            {
                pVgpuInstancePlacementInfo = &pVgpuTypePlacementInfo->vgpuInstancePlacementInfo[j];

                /*
                 * If hypervisor hasn't provided the placement ID as input, look
                 * for the first available placement ID for the input vGPU type.
                 * If hypervisor has provided the placement ID as input, just
                 * validate if it is a creatable placement ID.
                 */
                if (*placementId == NVA081_PLACEMENT_ID_INVALID)
                {
                    if (pVgpuInstancePlacementInfo->creatablePlacementId != NVA081_PLACEMENT_ID_INVALID)
                    {
                        *placementId = pVgpuInstancePlacementInfo->creatablePlacementId;
                        if (gpuIsSriovEnabled(pGpu))
                        {
                            *guestFbOffset = pVgpuInstancePlacementInfo->heterogeneousSupportedVmmuOffset *
                                             gpuGetVmmuSegmentSize(pGpu);
                            *guestFbLength = pVgpuTypePlacementInfo->guestVmmuCount *
                                             gpuGetVmmuSegmentSize(pGpu);
                            *gspHeapOffset = pVgpuInstancePlacementInfo->heterogeneousGspHeapOffset;
                            *guestBar1PFOffset = 0;
                        }
                        bIdFound = NV_TRUE;
                        break;
                    }
                }
                else
                {
                    if (pVgpuInstancePlacementInfo->creatablePlacementId == *placementId)
                    {
                        if (gpuIsSriovEnabled(pGpu))
                        {
                            *guestFbOffset = pVgpuInstancePlacementInfo->heterogeneousSupportedVmmuOffset *
                                             gpuGetVmmuSegmentSize(pGpu);
                            *guestFbLength = pVgpuTypePlacementInfo->guestVmmuCount *
                                             gpuGetVmmuSegmentSize(pGpu);
                            *gspHeapOffset = pVgpuInstancePlacementInfo->heterogeneousGspHeapOffset;
                            *guestBar1PFOffset    = 0;
                        }
                        bIdFound = NV_TRUE;
                        break;
                    }
                }
            }
        }
    }

    /* No creatable placement ID found */
    if (!bIdFound)
        return NV_ERR_INVALID_STATE;

    return _kvgpumgrUpdateHeterogeneousCreatablePlacementIds(pGpu, *placementId, vgpuTypeId, NV_TRUE, swizzId);
}

NV_STATUS
kvgpumgrUpdateHomogeneousInfo(OBJGPU *pGpu, Device *pDevice,
                              NvU32 vgpuTypeId, NvU16 *placementId,
                              NvU64 *guestFbLength, NvU64 *guestFbOffset)
{
    OBJSYS *pSys = SYS_GET_INSTANCE();
    KernelVgpuMgr *pKernelVgpuMgr = SYS_GET_KERNEL_VGPUMGR(pSys);
    KERNEL_PHYS_GPU_INFO *pPgpuInfo;
    VGPU_TYPE *pVgpuTypeInfo;
    NvU32 index, i , j, swizzId = KMIGMGR_SWIZZID_INVALID;
    NvBool bIdFound = NV_FALSE;
    KERNEL_VGPU_TYPE_PLACEMENT_INFO *pKernelVgpuTypePlacementInfo;
    VGPU_TYPE_PLACEMENT_INFO *pVgpuTypePlacementInfo;
    VGPU_INSTANCE_PLACEMENT_INFO *pVgpuInstancePlacementInfo;
    VGPU_TYPE_SUPPORTED_PLACEMENT_INFO *pVgpuTypeSupportedPlacementInfo;

    NV_CHECK_OR_RETURN(LEVEL_ERROR, pGpu != NULL, NV_ERR_NOT_SUPPORTED);

    NV_CHECK_OK_OR_RETURN(LEVEL_ERROR,
            kvgpumgrGetPgpuIndex(pKernelVgpuMgr, pGpu->gpuId, &index));

    pPgpuInfo = &(pKernelVgpuMgr->pgpuInfo[index]);

    if (IS_MIG_IN_USE(pGpu) && kvgpumgrIsMigTimeslicingModeEnabled(pGpu))
    {
        swizzId = kvgpuMgrGetSwizzIdFromDevice(pGpu, pDevice);
        pKernelVgpuTypePlacementInfo = kvgpuMgrGetVgpuPlacementInfo(pGpu, pPgpuInfo, swizzId);
        NV_CHECK_OR_RETURN(LEVEL_ERROR, pKernelVgpuTypePlacementInfo != NULL, NV_ERR_OBJECT_NOT_FOUND);
    }
    else
    {
        pKernelVgpuTypePlacementInfo = &pPgpuInfo->kernelVgpuTypePlacementInfo;
    }

    NV_CHECK_OK_OR_RETURN(LEVEL_ERROR, kvgpumgrCheckHomogeneousPlacementSupported(pGpu, swizzId));

    for (i = 0; i < pPgpuInfo->numVgpuTypes; i++)
    {
        pVgpuTypeInfo = pPgpuInfo->vgpuTypes[i];
        if (pVgpuTypeInfo == NULL)
            break;

        if (pVgpuTypeInfo->vgpuTypeId == vgpuTypeId)
        {
            pVgpuTypeSupportedPlacementInfo = &pVgpuTypeInfo->vgpuTypeSupportedPlacementInfo;
            pVgpuTypePlacementInfo = &pKernelVgpuTypePlacementInfo->vgpuTypePlacementInfo[i];

            for (j = 0; j < pVgpuTypeSupportedPlacementInfo->homogeneousPlacementCount; j++)
            {
                pVgpuInstancePlacementInfo = &pVgpuTypePlacementInfo->vgpuInstancePlacementInfo[j];

                /*
                 * If hypervisor hasn't provided the placement ID as input, look
                 * for the first available placement ID for the input vGPU type.
                 * If hypervisor has provided the placement ID as input, just
                 * validate if it is a creatable placement ID.
                 */
                if (*placementId == NVA081_PLACEMENT_ID_INVALID)
                {
                    if (pVgpuInstancePlacementInfo->creatablePlacementId != NVA081_PLACEMENT_ID_INVALID)
                    {
                        *placementId = pVgpuInstancePlacementInfo->creatablePlacementId;
                        *guestFbOffset = pVgpuInstancePlacementInfo->homogeneousSupportedVmmuOffset *
                                         gpuGetVmmuSegmentSize(pGpu);
                        *guestFbLength = pVgpuTypePlacementInfo->guestVmmuCount *
                                         gpuGetVmmuSegmentSize(pGpu);
                        bIdFound = NV_TRUE;
                        break;
                    }
                }
                else
                {
                    if (pVgpuInstancePlacementInfo->creatablePlacementId == *placementId)
                    {
                        *guestFbOffset = pVgpuInstancePlacementInfo->homogeneousSupportedVmmuOffset *
                                         gpuGetVmmuSegmentSize(pGpu);
                        *guestFbLength = pVgpuTypePlacementInfo->guestVmmuCount *
                                         gpuGetVmmuSegmentSize(pGpu);
                        bIdFound = NV_TRUE;
                        break;
                    }
                }
            }
        }
    }

    /* No creatable placement ID found */
    if (!bIdFound)
        return NV_ERR_INVALID_STATE;

    return _kvgpumgrUpdateHomogeneousCreatablePlacementIds(pGpu, *placementId, vgpuTypeId, NV_TRUE, swizzId);
}

NV_STATUS
kvgpumgrGetConfigEventInfoFromDb(NvHandle hClient,
                                 NvHandle hVgpuConfig,
                                 VGPU_CONFIG_EVENT_INFO_NODE **ppVgpuConfigEventInfoNode,
                                 NvU32 pgpuIndex)
{
    OBJSYS     *pSys     = SYS_GET_INSTANCE();
    KernelVgpuMgr *pKernelVgpuMgr = SYS_GET_KERNEL_VGPUMGR(pSys);
    VGPU_CONFIG_EVENT_INFO_NODE *pVgpuConfigEventInfoNodeTemp;

    if (!ppVgpuConfigEventInfoNode)
    {
        return NV_ERR_INVALID_ARGUMENT;
    }

    for (pVgpuConfigEventInfoNodeTemp = listHead(&(pKernelVgpuMgr->pgpuInfo[pgpuIndex].listVgpuConfigEventsHead));
         pVgpuConfigEventInfoNodeTemp != NULL;
         pVgpuConfigEventInfoNodeTemp = listNext(&(pKernelVgpuMgr->pgpuInfo[pgpuIndex].listVgpuConfigEventsHead), pVgpuConfigEventInfoNodeTemp))
    {
        if ((pVgpuConfigEventInfoNodeTemp->hClient == hClient) &&
            (pVgpuConfigEventInfoNodeTemp->hVgpuConfig == hVgpuConfig))
        {
            *ppVgpuConfigEventInfoNode = pVgpuConfigEventInfoNodeTemp;
            return NV_OK;
        }
    }

    *ppVgpuConfigEventInfoNode = NULL;
    return NV_ERR_OBJECT_NOT_FOUND;
}

NV_STATUS
kvgpuMgrRestoreSmcExecPart
(
    OBJGPU *pGpu,
    KERNEL_HOST_VGPU_DEVICE *pKernelHostVgpuDevice,
    KERNEL_MIG_GPU_INSTANCE *pKernelMIGGpuInstance
)
{
    NV_STATUS rmStatus = NV_OK;
    NvU32 CIIdx;
    NVC637_CTRL_EXEC_PARTITIONS_IMPORT_EXPORT_PARAMS *pExecPartImportParams = NULL;
    RM_API *pRmApi = rmapiGetInterface(RMAPI_GPU_LOCK_INTERNAL);

    //
    // Delete any compute instances which are alive in system. This may delete
    // instances existing in system even if they are same as default instance,
    // as we are not comparing instances
    //
    for (CIIdx = 0; CIIdx < NV_ARRAY_ELEMENTS(pKernelMIGGpuInstance->MIGComputeInstance); ++CIIdx)
    {
        NVC637_CTRL_EXEC_PARTITIONS_DELETE_PARAMS params;
        MIG_COMPUTE_INSTANCE *pMIGComputeInstance =
            &pKernelMIGGpuInstance->MIGComputeInstance[CIIdx];

        if (!pMIGComputeInstance->bValid)
            continue;

        portMemSet(&params, 0, sizeof(params));
        params.execPartCount = 1;
        params.execPartId[0] = CIIdx;

        //
        // There is no revert or failure in this call as we have to do best case
        // effort to delete and restore default compute instances
        //
        NV_ASSERT_OK(
            pRmApi->Control(pRmApi,
                            pKernelMIGGpuInstance->instanceHandles.hClient,
                            pKernelMIGGpuInstance->instanceHandles.hSubscription,
                            NVC637_CTRL_CMD_EXEC_PARTITIONS_DELETE,
                            &params,
                            sizeof(params)));
    }

    // Restore all saved compute instances
    pExecPartImportParams = portMemAllocNonPaged(sizeof(*pExecPartImportParams));
    NV_ASSERT_OR_RETURN(pExecPartImportParams != NULL, NV_ERR_NO_MEMORY);

    for (CIIdx = 0; CIIdx < NV_ARRAY_ELEMENTS(pKernelMIGGpuInstance->MIGComputeInstance); ++CIIdx)
    {
        GPUMGR_SAVE_COMPUTE_INSTANCE *pComputeInstanceSave = &pKernelHostVgpuDevice->savedExecPartitions[CIIdx];
        if (!pComputeInstanceSave->bValid)
            continue;

        portMemSet(pExecPartImportParams, 0, sizeof(*pExecPartImportParams));
        pExecPartImportParams->id = pComputeInstanceSave->id;
        pExecPartImportParams->bCreateCap = NV_TRUE;
        portMemCopy(&pExecPartImportParams->info, sizeof(pExecPartImportParams->info),
                          &pComputeInstanceSave->ciInfo, sizeof(pComputeInstanceSave->ciInfo));

        NV_ASSERT_OK_OR_GOTO(rmStatus,
                      pRmApi->Control(pRmApi,
                                      pKernelMIGGpuInstance->instanceHandles.hClient,
                                      pKernelMIGGpuInstance->instanceHandles.hSubscription,
                                      NVC637_CTRL_CMD_EXEC_PARTITIONS_IMPORT,
                                      pExecPartImportParams,
                                      sizeof(*pExecPartImportParams)),
                done);
    }

    //
    // Generate a subdevice event stating something has changed in GPU partition
    // config.
    //
    gpuNotifySubDeviceEvent(pGpu, NV2080_NOTIFIERS_SMC_CONFIG_UPDATE, NULL, 0, 0, 0);

done:
    portMemFree(pExecPartImportParams);
    return rmStatus;
}

/*
 * This function is currently supported only on Device VM.
 * Device VM use VM ID as identifier to set the VM Name.
 */
NV_STATUS
kvgpumgrGetHostVgpuDeviceFromVmId(NvU32 gpuPciId, NvU64 vmId,
                                  KERNEL_HOST_VGPU_DEVICE** ppHostVgpuDevice)
{
    OBJSYS *pSys = SYS_GET_INSTANCE();
    KernelVgpuMgr *pKernelVgpuMgr = SYS_GET_KERNEL_VGPUMGR(pSys);
    NvU32 pgpuIndex;
    KERNEL_PHYS_GPU_INFO *pPgpuInfo;
    KERNEL_HOST_VGPU_DEVICE *pKernelHostVgpuDevice;

    if (osIsVgpuDeviceVmPresent() != NV_OK)
        return NV_ERR_NOT_SUPPORTED;

    NV_CHECK_OK_OR_RETURN(LEVEL_ERROR, kvgpumgrGetPgpuIndex(pKernelVgpuMgr, gpuPciId, &pgpuIndex));
    pPgpuInfo = &pKernelVgpuMgr->pgpuInfo[pgpuIndex];

    if (!pPgpuInfo->sriovEnabled)
        return NV_ERR_NOT_SUPPORTED;

    for (pKernelHostVgpuDevice = listHead(&(pPgpuInfo->listHostVgpuDeviceHead));
         pKernelHostVgpuDevice != NULL;
         pKernelHostVgpuDevice = listNext(&(pPgpuInfo->listHostVgpuDeviceHead), pKernelHostVgpuDevice))
    {
        if (pKernelHostVgpuDevice->vgpuGuest->guestVmInfo.guestVmId.vmId == vmId)
        {
            *ppHostVgpuDevice = pKernelHostVgpuDevice;
            return NV_OK;
        }
     }

    return NV_ERR_OBJECT_NOT_FOUND;
}

NV_STATUS
kvgpumgrSetPlacementId(REQUEST_VGPU_INFO_NODE *pRequestVgpu, NvU32 placementId)
{
    OBJSYS                              *pSys           = SYS_GET_INSTANCE();
    KernelVgpuMgr                       *pKernelVgpuMgr = SYS_GET_KERNEL_VGPUMGR(pSys);
    OBJGPU                              *pGpu           = NULL;
    NV_STATUS                            rmStatus       = NV_OK;
    NvU32                                index, i , j, assignedPlacementId;
    KERNEL_PHYS_GPU_INFO                *pPgpuInfo;
    VGPU_TYPE                           *pVgpuTypeInfo;
    KERNEL_VGPU_TYPE_PLACEMENT_INFO     *pKernelVgpuTypePlacementInfo;
    NvU16                                placementCount;
    VGPU_TYPE_PLACEMENT_INFO            *pVgpuTypePlacementInfo;
    VGPU_INSTANCE_PLACEMENT_INFO        *pVgpuInstancePlacementInfo;
    VGPU_TYPE_SUPPORTED_PLACEMENT_INFO  *pVgpuTypeSupportedPlacementInfo;
    NvBool                               bIdFound = NV_FALSE, bHeterogeneousModeEnabled;

    // If the call is not coming from mdev-creation, verify the existing placement Id
    // against the placement Id to be assigned.
    if ((pRequestVgpu->placementId != NVA081_PLACEMENT_ID_INVALID) &&
        (pRequestVgpu->placementId == placementId))
        return NV_OK;

    pGpu = gpumgrGetGpuFromId(pRequestVgpu->gpuPciId);
    if (pGpu == NULL)
        return NV_ERR_NOT_SUPPORTED;

    // Check for placementId unsupported config.

    NV_CHECK_OK_OR_RETURN(LEVEL_ERROR,
            kvgpumgrGetPgpuIndex(pKernelVgpuMgr, pGpu->gpuId, &index));
 
    pPgpuInfo = &(pKernelVgpuMgr->pgpuInfo[index]);
    if (IS_MIG_IN_USE(pGpu) && kvgpumgrIsMigTimeslicingModeEnabled(pGpu))
    {
        pKernelVgpuTypePlacementInfo = kvgpuMgrGetVgpuPlacementInfo(pGpu, pPgpuInfo, pRequestVgpu->swizzId);
        NV_CHECK_OR_RETURN(LEVEL_ERROR, pKernelVgpuTypePlacementInfo != NULL, NV_ERR_OBJECT_NOT_FOUND);
    }
    else
    {
        pKernelVgpuTypePlacementInfo = &pPgpuInfo->kernelVgpuTypePlacementInfo;
    }
    NV_CHECK_OK_OR_RETURN(LEVEL_ERROR,
        kvgpuMgrGetHeterogeneousMode(pGpu, pRequestVgpu->swizzId, &bHeterogeneousModeEnabled));

    if (!bHeterogeneousModeEnabled &&
        (kvgpumgrCheckHomogeneousPlacementSupported(pGpu, pRequestVgpu->swizzId) != NV_OK))
    {
        pRequestVgpu->placementId = NVA081_PLACEMENT_ID_INVALID;
        return NV_ERR_NOT_SUPPORTED;
    }

    // Allocate either the first available placementId or the placementId specified by the user.
    for (i = 0; i < pPgpuInfo->numVgpuTypes; i++)
    {
        pVgpuTypeInfo = pPgpuInfo->vgpuTypes[i];
        if (pVgpuTypeInfo == NULL)
            break;
 
        pVgpuTypeSupportedPlacementInfo = &pVgpuTypeInfo->vgpuTypeSupportedPlacementInfo;
        pVgpuTypePlacementInfo = &pKernelVgpuTypePlacementInfo->vgpuTypePlacementInfo[i];
 
        if (pVgpuTypeInfo->vgpuTypeId == pRequestVgpu->vgpuTypeId)
        {
            placementCount = (bHeterogeneousModeEnabled) ?
                            pVgpuTypeSupportedPlacementInfo->heterogeneousPlacementCount :
                            pVgpuTypeSupportedPlacementInfo->homogeneousPlacementCount;

            for (j = 0; j < placementCount; j++)
            {
                pVgpuInstancePlacementInfo = &pVgpuTypePlacementInfo->vgpuInstancePlacementInfo[j];

                if (placementId == NVA081_PLACEMENT_ID_INVALID)
                {
                    if (pVgpuInstancePlacementInfo->creatablePlacementId != NVA081_PLACEMENT_ID_INVALID)
                    {
                        assignedPlacementId = pVgpuInstancePlacementInfo->creatablePlacementId;
                        bIdFound = NV_TRUE;
                        break;
                    }
                }
                else
                {
                    if (pVgpuInstancePlacementInfo->creatablePlacementId == placementId)
                    {
                        assignedPlacementId = placementId;
                        bIdFound = NV_TRUE;
                        break;
                    }
                }
            }
        }
    }

    // No creatable placement ID found
    if (!bIdFound)
        return NV_ERR_INVALID_STATE;

    // Reserve the new placementId then free up the pre-existing placementId if there is any
    if (bHeterogeneousModeEnabled)
    {
        rmStatus = _kvgpumgrUpdateHeterogeneousCreatablePlacementIds(pGpu,
                                    assignedPlacementId, pRequestVgpu->vgpuTypeId, NV_TRUE, pRequestVgpu->swizzId);
        if (rmStatus != NV_OK)
            return rmStatus;

        if (placementId != NVA081_PLACEMENT_ID_INVALID)
        {
            rmStatus = _kvgpumgrUpdateHeterogeneousCreatablePlacementIds(pGpu,
                                         pRequestVgpu->placementId,
                                         pRequestVgpu->vgpuTypeId, NV_FALSE,
                                         pRequestVgpu->swizzId);
        }
    }
    else
    {
        rmStatus = _kvgpumgrUpdateHomogeneousCreatablePlacementIds(pGpu,
                                    assignedPlacementId, pRequestVgpu->vgpuTypeId, NV_TRUE, pRequestVgpu->swizzId);
        if (rmStatus != NV_OK)
            return rmStatus;

        if (placementId != NVA081_PLACEMENT_ID_INVALID)
        {
            rmStatus = _kvgpumgrUpdateHomogeneousCreatablePlacementIds(pGpu,
                                         pRequestVgpu->placementId,
                                         pRequestVgpu->vgpuTypeId, NV_FALSE,
                                         pRequestVgpu->swizzId);
        }
    }

    if (rmStatus != NV_OK)
        return rmStatus;

    pRequestVgpu->placementId = assignedPlacementId;

    return rmStatus;
}

NV_STATUS
kvgpumgrSetGpuInstanceId(REQUEST_VGPU_INFO_NODE *pRequestVgpu, NvU32 swizzId)
{
    OBJSYS                  *pSys                   = SYS_GET_INSTANCE();
    OBJGPU                  *pGpu                   = NULL;
    KernelMIGManager        *pKernelMIGManager      = NULL;
    KernelVgpuMgr           *pKernelVgpuMgr         = SYS_GET_KERNEL_VGPUMGR(pSys);
    KERNEL_PHYS_GPU_INFO    *pPhysGpuInfo           = NULL;
    NvU64                    swizzIdInUseMask       = 0;
    NV_STATUS                rmStatus               = NV_OK;
    VGPU_TYPE               *vgpuTypeInfo           = NULL;
    NvU32                    partitionFlag          = PARTITIONID_INVALID;
    NvU32                    pgpuIndex;
    NvU32                    currPlacementId;

    if (pRequestVgpu->swizzId == swizzId)
        return NV_OK;

    pGpu = gpumgrGetGpuFromId(pRequestVgpu->gpuPciId);
    if (pGpu == NULL)
        return NV_ERR_NOT_SUPPORTED;

    pKernelMIGManager  = GPU_GET_KERNEL_MIG_MANAGER(pGpu);
    NV_CHECK_OR_RETURN(LEVEL_INFO, IS_MIG_ENABLED(pGpu), NV_ERR_INVALID_OPERATION);
    NV_CHECK_OR_RETURN(LEVEL_INFO, IS_MIG_IN_USE(pGpu), NV_ERR_INVALID_OPERATION);

    if ((rmStatus = kvgpumgrGetPgpuIndex(pKernelVgpuMgr, pRequestVgpu->gpuPciId, &pgpuIndex)) != NV_OK)
        return rmStatus;

    pPhysGpuInfo = &(pKernelVgpuMgr->pgpuInfo[pgpuIndex]);

    if ((rmStatus = kvgpumgrGetVgpuTypeInfo(pRequestVgpu->vgpuTypeId, &vgpuTypeInfo)) != NV_OK)
        return rmStatus;

    NV_ASSERT_OK_OR_RETURN(kvgpumgrGetPartitionFlag(vgpuTypeInfo->vgpuTypeId, &partitionFlag));

    // Store the placementid of the vGPU to be moved
    currPlacementId = pRequestVgpu->placementId;

    swizzIdInUseMask = kmigmgrGetSwizzIdInUseMask(pGpu, pKernelMIGManager);

    // Determine valid swizzids not assigned to any vGPU device.
    if (NVBIT64(swizzId) & swizzIdInUseMask)
    {
        KERNEL_MIG_GPU_INSTANCE *pKernelMIGGpuInstance;

        rmStatus = kmigmgrGetGPUInstanceInfo(pGpu, pKernelMIGManager, swizzId, &pKernelMIGGpuInstance);
        if (rmStatus != NV_OK)
        {
            // Didn't find requested GPU instance
            NV_PRINTF(LEVEL_ERROR, "No valid GPU instance with SwizzId - %d found\n", swizzId);
            return rmStatus;
        }

        if (pKernelMIGGpuInstance->partitionFlag == partitionFlag)
        {
            NvBool bHeterogeneousModeEnabledTargetGI = NV_FALSE;
            NvBool bHeterogeneousModeEnabledCurrentGI = NV_FALSE;
            /*
             * For platform that does not support mig + timeslice, do not check the existing vgpu on the target GI
             * For platform that has mig + timeslice, assuming heterogeneous vgpu types will be supported:
             *   - Check the target GI can support the vGPU type, then assign a new placement id for the vGPU type
             *   - If any steps above failed, abandon the migration and restore the placement id of the vGPU
             */
            if (kvgpumgrIsMigTimeslicingModeEnabled(pGpu))
            {
                VGPU_TYPE *pVgpuTypeInfo;
                NvBool bIdFound = NV_FALSE;
                VGPU_TYPE_SUPPORTED_PLACEMENT_INFO *pVgpuTypeSupportedPlacementInfo = NULL;
                VGPU_INSTANCE_PLACEMENT_INFO *pVgpuInstancePlacementInfo = NULL;
                VGPU_TYPE_PLACEMENT_INFO *pVgpuTypePlacementInfo = NULL;
                NvU32 placementCount = 0, i, j;
                KERNEL_VGPU_TYPE_PLACEMENT_INFO *pKernelVgpuTypePlacementInfo;

                // Check the current GI has heterogeneous mode enabled
                NV_CHECK_OK_OR_RETURN(LEVEL_ERROR,
                    kvgpuMgrGetHeterogeneousMode(pGpu, swizzId, &bHeterogeneousModeEnabledCurrentGI));
                // Check the target GI has heterogeneous mode enabled
                NV_CHECK_OK_OR_RETURN(LEVEL_ERROR,
                    kvgpuMgrGetHeterogeneousMode(pGpu, swizzId, &bHeterogeneousModeEnabledTargetGI));
                pKernelVgpuTypePlacementInfo =
                    kvgpuMgrGetVgpuPlacementInfo(pGpu, pPhysGpuInfo, swizzId);
                NV_CHECK_OR_RETURN(LEVEL_ERROR, pKernelVgpuTypePlacementInfo != NULL, NV_ERR_OBJECT_NOT_FOUND);
                // Check if target GI can support the vgpuType
                for (i = 0; i < pPhysGpuInfo->numVgpuTypes; i++)
                {
                    pVgpuTypeInfo = pPhysGpuInfo->vgpuTypes[i];
                    if (pVgpuTypeInfo == NULL)
                        break;

                    pVgpuTypeSupportedPlacementInfo = &pVgpuTypeInfo->vgpuTypeSupportedPlacementInfo;
                    pVgpuTypePlacementInfo = &pKernelVgpuTypePlacementInfo->vgpuTypePlacementInfo[i];
                    placementCount = bHeterogeneousModeEnabledTargetGI ?
                                    pVgpuTypeSupportedPlacementInfo->heterogeneousPlacementCount :
                                    pVgpuTypeSupportedPlacementInfo->homogeneousPlacementCount;
                    if (pVgpuTypeInfo->vgpuTypeId == pRequestVgpu->vgpuTypeId)
                    {
                        NV_CHECK_OR_GOTO(LEVEL_ERROR, pVgpuTypeSupportedPlacementInfo != NULL, error_exit);
                        NV_CHECK_OR_GOTO(LEVEL_ERROR, pVgpuTypePlacementInfo != NULL, error_exit);
                        // Check if target GI can support the placement id
                        for (j = 0; j < placementCount; j++)
                        {
                            pVgpuInstancePlacementInfo = &pVgpuTypePlacementInfo->vgpuInstancePlacementInfo[j];
                            if (pVgpuInstancePlacementInfo->creatablePlacementId != NVA081_PLACEMENT_ID_INVALID)
                            {
                                pRequestVgpu->placementId = pVgpuInstancePlacementInfo->creatablePlacementId;
                                bIdFound = NV_TRUE;
                                break;
                            }
                        }
                    }
                }
                if (!bIdFound)
                {
                    NV_PRINTF(LEVEL_ERROR, "Cannot find vGPU type %d placement id in GPU instance %d",
                        pRequestVgpu->vgpuTypeId, swizzId);
                    goto error_exit;
                }
            }
            else
            {
                // Validate that same ID is not already set and then set the ID
                if (pPhysGpuInfo->assignedSwizzIdVgpuCount[swizzId] >= vgpuTypeInfo->maxInstancePerGI )
                {
                    goto error_exit;
                }
            }

            // Clear the assigned swizzId of the current GI
            NV_ASSERT_OK_OR_RETURN(_kvgpumgrClearAssignedSwizzIdMask(pGpu, pRequestVgpu->swizzId));
            // Assign the swizzId to the vGPU device of the target GI
            NV_ASSERT_OK_OR_RETURN(_kvgpumgrSetAssignedSwizzIdMask(pGpu, vgpuTypeInfo,
                                    pKernelMIGGpuInstance->swizzId));

            /*
             * Now we confirmed the target GI can support the current vgpu, we can proceed to update the placement id
             * First, update the placement ID for the target GI
             * Then, free the placement ID for the current GI
             */
            if (kvgpumgrIsMigTimeslicingModeEnabled(pGpu))
            {
                if (bHeterogeneousModeEnabledTargetGI)
                {
                    // Update placement id of the target GI
                    NV_CHECK_OK_OR_GOTO(rmStatus, LEVEL_ERROR,
                        _kvgpumgrUpdateHeterogeneousCreatablePlacementIds(pGpu, pRequestVgpu->placementId,
                        pRequestVgpu->vgpuTypeId, NV_TRUE,
                        swizzId),
                        error_exit);
                }
                else
                {
                    // Update placement id of the target GI
                    NV_CHECK_OK_OR_GOTO(rmStatus, LEVEL_ERROR,
                        _kvgpumgrUpdateHomogeneousCreatablePlacementIds(pGpu, pRequestVgpu->placementId,
                        pRequestVgpu->vgpuTypeId, NV_TRUE,
                        swizzId),
                        error_exit);
                }

                if (bHeterogeneousModeEnabledCurrentGI)
                {
                    // Free the placement id for the vGPU from current GI
                    NV_CHECK_OK_OR_GOTO(rmStatus, LEVEL_ERROR,
                        _kvgpumgrUpdateHeterogeneousCreatablePlacementIds(pGpu, currPlacementId,
                        pRequestVgpu->vgpuTypeId, NV_FALSE,
                        pRequestVgpu->swizzId),
                        error_exit);
                }
                else
                {
                    // Free placement id for the vGPU from current GI
                    NV_CHECK_OK_OR_GOTO(rmStatus, LEVEL_ERROR,
                        _kvgpumgrUpdateHomogeneousCreatablePlacementIds(pGpu, currPlacementId,
                        pRequestVgpu->vgpuTypeId, NV_FALSE,
                        pRequestVgpu->swizzId),
                        error_exit);
                }
            }
        }
        else
            goto error_exit;

        // Assign swizzId to the REQUEST_VGPU_INFO_NODE
        pRequestVgpu->swizzId = pKernelMIGGpuInstance->swizzId;

        return NV_OK;
    }

error_exit:
    // Restore placementId
    pRequestVgpu->placementId = currPlacementId;
    return NV_ERR_INVALID_STATE;
}

NV_STATUS
kvgpuMgrGetPlacementResource(OBJGPU *pGpu, Device *pDevice,
                             const NvU8 *pVgpuDevName, NvU16 *placementId,
                             NvU64 *guestFbLength, NvU64 *guestFbOffset,
                             NvU64 *gspHeapOffset, NvU64 *guestBar1PFOffset)
{
    OBJSYS *pSys = SYS_GET_INSTANCE();
    KernelVgpuMgr *pKernelVgpuMgr = SYS_GET_KERNEL_VGPUMGR(pSys);
    REQUEST_VGPU_INFO_NODE *pRequestVgpu = NULL;
    NvU32 index, i , j;
    KERNEL_PHYS_GPU_INFO *pPgpuInfo;
    VGPU_TYPE *pVgpuTypeInfo;
    KERNEL_VGPU_TYPE_PLACEMENT_INFO *pKernelVgpuTypePlacementInfo;
    VGPU_INSTANCE_PLACEMENT_INFO *pVgpuInstancePlacementInfo;
    VGPU_TYPE_PLACEMENT_INFO *pVgpuTypePlacementInfo;
    VGPU_TYPE_SUPPORTED_PLACEMENT_INFO *pVgpuTypeSupportedPlacementInfo;
    NvBool bHeterogeneousModeEnabled;

    NV_CHECK_OK_OR_RETURN(LEVEL_ERROR, kvgpumgrGetPgpuIndex(pKernelVgpuMgr, pGpu->gpuId, &index));

    pPgpuInfo = &(pKernelVgpuMgr->pgpuInfo[index]);
    if (IS_MIG_IN_USE(pGpu) && kvgpumgrIsMigTimeslicingModeEnabled(pGpu))
    {
        NvU32 swizzId = kvgpuMgrGetSwizzIdFromDevice(pGpu, pDevice);
        pKernelVgpuTypePlacementInfo = kvgpuMgrGetVgpuPlacementInfo(pGpu, pPgpuInfo, swizzId);
        NV_CHECK_OR_RETURN(LEVEL_ERROR, pKernelVgpuTypePlacementInfo != NULL, NV_ERR_OBJECT_NOT_FOUND);
        NV_CHECK_OK_OR_RETURN(LEVEL_ERROR,
            kvgpuMgrGetHeterogeneousModePerGI(pGpu, swizzId, &bHeterogeneousModeEnabled));
    }
    else
    {
        pKernelVgpuTypePlacementInfo = &pPgpuInfo->kernelVgpuTypePlacementInfo;
        bHeterogeneousModeEnabled = pGpu->getProperty(pGpu, PDB_PROP_GPU_IS_VGPU_HETEROGENEOUS_MODE);
    }
    // Check the corresponding REQUEST_VGPU_INFO_NODE
    for (pRequestVgpu = listHead(&pKernelVgpuMgr->listRequestVgpuHead);
         pRequestVgpu != NULL;
         pRequestVgpu = listNext(&pKernelVgpuMgr->listRequestVgpuHead, pRequestVgpu))
    {
        if (portMemCmp(pRequestVgpu->vgpuDevName, pVgpuDevName, VM_UUID_SIZE) == 0)
            break;
    }

    if (pRequestVgpu == NULL)
        return NV_ERR_OBJECT_NOT_FOUND;

    // The placement resources corresponding to the placementId are obtained by iteration over vgpuTypes
    for (i = 0; i < pPgpuInfo->numVgpuTypes; i++)
    {
        pVgpuTypeInfo = pPgpuInfo->vgpuTypes[i];
        if (pVgpuTypeInfo == NULL)
            continue;

        pVgpuTypeSupportedPlacementInfo = &pVgpuTypeInfo->vgpuTypeSupportedPlacementInfo;
        pVgpuTypePlacementInfo = &pKernelVgpuTypePlacementInfo->vgpuTypePlacementInfo[i];

        if (pVgpuTypeInfo->vgpuTypeId != pRequestVgpu->vgpuTypeId)
            continue;

        if (pVgpuTypeSupportedPlacementInfo->placementSize == 0)
            continue;

        if (bHeterogeneousModeEnabled)
        {
            for (j = 0; j < pVgpuTypeSupportedPlacementInfo->heterogeneousPlacementCount; j++)
            {
                if (pVgpuTypeSupportedPlacementInfo->vgpuInstanceSupportedPlacementInfo[j].heterogeneousSupportedPlacementId ==
                    pRequestVgpu->placementId)
                {
                    pVgpuInstancePlacementInfo = &pVgpuTypePlacementInfo->vgpuInstancePlacementInfo[j];
                    *placementId        = pRequestVgpu->placementId;
                    if (gpuIsSriovEnabled(pGpu))
                    {
                        *guestFbOffset = pVgpuInstancePlacementInfo->heterogeneousSupportedVmmuOffset *
                                         gpuGetVmmuSegmentSize(pGpu);
                        *guestFbLength = pVgpuTypePlacementInfo->guestVmmuCount *
                                         gpuGetVmmuSegmentSize(pGpu);
                        *gspHeapOffset = pVgpuInstancePlacementInfo->heterogeneousGspHeapOffset;
                        *guestBar1PFOffset    = 0;
                    }

                    return NV_OK;
                }
            }
        }
        else
        {
            for (j = 0; j < pVgpuTypeSupportedPlacementInfo->homogeneousPlacementCount; j++)
            {
                if (pVgpuTypeSupportedPlacementInfo->vgpuInstanceSupportedPlacementInfo[j].homogeneousSupportedPlacementId ==
                    pRequestVgpu->placementId)
                {
                     pVgpuInstancePlacementInfo = &pVgpuTypePlacementInfo->vgpuInstancePlacementInfo[j];
                    *placementId        = pRequestVgpu->placementId;
                    *guestFbOffset      = pVgpuInstancePlacementInfo->homogeneousSupportedVmmuOffset *
                                                gpuGetVmmuSegmentSize(pGpu);
                    *guestFbLength      = pVgpuTypePlacementInfo->guestVmmuCount *
                                                gpuGetVmmuSegmentSize(pGpu);
                    *gspHeapOffset      = 0;
                    *guestBar1PFOffset  = 0;

                    return NV_OK;
                }
            }
        }
    }

    return NV_ERR_INVALID_STATE;
}
NV_STATUS
kvgpuMgrGetHeterogeneousMode(OBJGPU *pGpu, NvU32 swizzId, NvBool *bHeterogeneousModeEnabled)
{
    if (IS_MIG_IN_USE(pGpu))
    {
        if (pGpu->getProperty(pGpu, PDB_PROP_GPU_MIG_TIMESLICING_SUPPORTED) && kvgpumgrIsMigTimeslicingModeEnabled(pGpu))
            NV_CHECK_OK_OR_RETURN(LEVEL_ERROR,
                kvgpuMgrGetHeterogeneousModePerGI(pGpu, swizzId, bHeterogeneousModeEnabled));
        else
            *bHeterogeneousModeEnabled = NV_FALSE;
    }
    else
    {
        *bHeterogeneousModeEnabled = pGpu->getProperty(pGpu, PDB_PROP_GPU_IS_VGPU_HETEROGENEOUS_MODE);
    }
    return NV_OK;
}

KERNEL_VGPU_TYPE_PLACEMENT_INFO*
kvgpuMgrGetVgpuPlacementInfo
(
    OBJGPU *pGpu,
    KERNEL_PHYS_GPU_INFO *pPgpuInfo,
    NvU32 swizzId
)
{
    KERNEL_VGPU_TYPE_PLACEMENT_INFO *pKernelVgpuTypePlacementInfo = NULL;

    if (IS_MIG_IN_USE(pGpu) && kvgpumgrIsMigTimeslicingModeEnabled(pGpu))
    {
        for (NvU32 i = 0; i < MAX_GI_PARTITIONS_PER_GPU; i++)
        {
            if (pPgpuInfo->kernelVgpuTypePlacementInfoGi[i].swizzId == swizzId &&
                pPgpuInfo->kernelVgpuTypePlacementInfoGi[i].bUsed)
            {
                pKernelVgpuTypePlacementInfo = &pPgpuInfo->kernelVgpuTypePlacementInfoGi[i].kernelVgpuTypePlacementInfo;
                break;
            }
        }
    }
    else
    {
        pKernelVgpuTypePlacementInfo = &pPgpuInfo->kernelVgpuTypePlacementInfo;
    }
    return pKernelVgpuTypePlacementInfo;
}

NvU32*
kvgpuMgrGetVgpuCreatableTypeIdFromSwizzId
(
    OBJGPU *pGpu,
    KERNEL_PHYS_GPU_INFO *pPgpuInfo,
    NvU32 swizzId
)
{
    NvU32 *pSupportedTypeId = NULL;

    if (IS_MIG_IN_USE(pGpu) && kvgpumgrIsMigTimeslicingModeEnabled(pGpu))
    {
        for (NvU32 i = 0; i < MAX_GI_PARTITIONS_PER_GPU; i++)
        {
            if (pPgpuInfo->kernelVgpuTypePlacementInfoGi[i].swizzId == swizzId &&
                pPgpuInfo->kernelVgpuTypePlacementInfoGi[i].bUsed)
            {
                pSupportedTypeId = pPgpuInfo->kernelVgpuTypePlacementInfoGi[i].supportedTypeIds;
                break;
            }
        }
    }
    else
    {
        pSupportedTypeId = pPgpuInfo->supportedTypeIds;
    }
    return pSupportedTypeId;
}

NvU32
kvgpuMgrGetSwizzIdFromDevice
(
    OBJGPU *pGpu,
    Device *pDevice
)
{
    NvU32 swizzId = KMIGMGR_SWIZZID_INVALID;

    if (IS_MIG_IN_USE(pGpu))
    {
        KernelMIGManager *pKernelMIGManager = GPU_GET_KERNEL_MIG_MANAGER(pGpu);
        MIG_INSTANCE_REF ref;

        NV_CHECK_OK_OR_RETURN(LEVEL_ERROR, kmigmgrGetInstanceRefFromDevice(pGpu, pKernelMIGManager, pDevice, &ref));
        swizzId = ref.pKernelMIGGpuInstance->swizzId;
    }

    return swizzId;
}

NV_STATUS
kvgpuMgrReserveVgpuPlacementInfoPerGI
(
    OBJGPU  *pGpu,
    NvU32    swizzId
)
{
    OBJSYS     *pSys     = SYS_GET_INSTANCE();
    KernelVgpuMgr *pKernelVgpuMgr = SYS_GET_KERNEL_VGPUMGR(pSys);
    KERNEL_PHYS_GPU_INFO *pPgpuInfo;
    NvU32 index, i;

    NV_CHECK_OK_OR_RETURN(LEVEL_ERROR,
            kvgpumgrGetPgpuIndex(pKernelVgpuMgr, pGpu->gpuId, &index));

    pPgpuInfo = &(pKernelVgpuMgr->pgpuInfo[index]);

    for (i = 0; i < MAX_GI_PARTITIONS_PER_GPU; i++)
    {
        if (!pPgpuInfo->kernelVgpuTypePlacementInfoGi[i].bUsed)
        {
            pPgpuInfo->kernelVgpuTypePlacementInfoGi[i].swizzId = swizzId;
            pPgpuInfo->kernelVgpuTypePlacementInfoGi[i].bUsed = NV_TRUE;
            break;
        }
    }
    if (i == MAX_GI_PARTITIONS_PER_GPU)
    {
        return NV_ERR_OUT_OF_RANGE;
    }

    if (pPgpuInfo->heterogeneousTimesliceSizesSupported)
    {
        NV_CHECK_OK_OR_RETURN(LEVEL_ERROR,
            _kvgpuMgrSetHeterogeneousResourcesPerGI(pGpu, swizzId, pPgpuInfo));
    }
    if (pPgpuInfo->homogeneousPlacementSupported)
    {
        NV_CHECK_OK_OR_RETURN(LEVEL_ERROR,
            _kvgpumgrSetHomogeneousResourcesPerGI(pGpu, swizzId, pPgpuInfo));
    }
    return NV_OK;
}

NV_STATUS
kvgpuMgrClearVgpuPlacementInfoPerGI
(
    OBJGPU    *pGpu,
    NvU32      swizzId
)
{
    OBJSYS     *pSys     = SYS_GET_INSTANCE();
    KernelVgpuMgr *pKernelVgpuMgr = SYS_GET_KERNEL_VGPUMGR(pSys);
    KERNEL_PHYS_GPU_INFO *pPgpuInfo;
    NvU32 index, i;

    NV_CHECK_OK_OR_RETURN(LEVEL_ERROR,
            kvgpumgrGetPgpuIndex(pKernelVgpuMgr, pGpu->gpuId, &index));

    pPgpuInfo = &(pKernelVgpuMgr->pgpuInfo[index]);

    for (i = 0; i < MAX_GI_PARTITIONS_PER_GPU; i++)
    {
        if (pPgpuInfo->kernelVgpuTypePlacementInfoGi[i].swizzId == swizzId &&
            pPgpuInfo->kernelVgpuTypePlacementInfoGi[i].bUsed)
        {
            portMemSet(&pPgpuInfo->kernelVgpuTypePlacementInfoGi[i], 0,
                        sizeof(pPgpuInfo->kernelVgpuTypePlacementInfoGi[i]));
            break;
        }
    }
    if (i == MAX_GI_PARTITIONS_PER_GPU)
    {
        return NV_ERR_OUT_OF_RANGE;
    }

    return NV_OK;
}

NV_STATUS
kvgpuMgrSetHeterogeneousModePerGI
(
    OBJGPU    *pGpu,
    NvU32      swizzId,
    NvBool     bHeterogeneousModeEnabled
)
{
    OBJSYS     *pSys     = SYS_GET_INSTANCE();
    KernelVgpuMgr *pKernelVgpuMgr = SYS_GET_KERNEL_VGPUMGR(pSys);
    KERNEL_PHYS_GPU_INFO *pPgpuInfo;
    NvU32 index, i;

    NV_CHECK_OK_OR_RETURN(LEVEL_ERROR,
            kvgpumgrGetPgpuIndex(pKernelVgpuMgr, pGpu->gpuId, &index));

    pPgpuInfo = &(pKernelVgpuMgr->pgpuInfo[index]);

    for (i = 0; i < MAX_GI_PARTITIONS_PER_GPU; i++)
    {
        if (pPgpuInfo->kernelVgpuTypePlacementInfoGi[i].swizzId == swizzId &&
            pPgpuInfo->kernelVgpuTypePlacementInfoGi[i].bUsed)
        {
            pPgpuInfo->kernelVgpuTypePlacementInfoGi[i].bHeterogeneousModeEnabled = bHeterogeneousModeEnabled;
            break;
        }
    }
    if (i == MAX_GI_PARTITIONS_PER_GPU)
    {
        return NV_ERR_OUT_OF_RANGE;
    }

    return NV_OK;
}

NV_STATUS
kvgpuMgrGetHeterogeneousModePerGI
(
    OBJGPU    *pGpu,
    NvU32      swizzId,
    NvBool     *bHeterogeneousModeEnabled
)
{
    OBJSYS     *pSys     = SYS_GET_INSTANCE();
    KernelVgpuMgr *pKernelVgpuMgr = SYS_GET_KERNEL_VGPUMGR(pSys);
    KERNEL_PHYS_GPU_INFO *pPgpuInfo;
    NvU32 index, i;

    NV_CHECK_OK_OR_RETURN(LEVEL_ERROR,
            kvgpumgrGetPgpuIndex(pKernelVgpuMgr, pGpu->gpuId, &index));

    pPgpuInfo = &(pKernelVgpuMgr->pgpuInfo[index]);

    for (i = 0; i < MAX_GI_PARTITIONS_PER_GPU; i++)
    {
        if (pPgpuInfo->kernelVgpuTypePlacementInfoGi[i].swizzId == swizzId &&
            pPgpuInfo->kernelVgpuTypePlacementInfoGi[i].bUsed)
        {
            *bHeterogeneousModeEnabled = pPgpuInfo->kernelVgpuTypePlacementInfoGi[i].bHeterogeneousModeEnabled;
            break;
        }
    }
    if (i == MAX_GI_PARTITIONS_PER_GPU)
    {
        return NV_ERR_OUT_OF_RANGE;
    }

    return NV_OK;
}

static NV_STATUS
_kvgpuMgrSetHeterogeneousResourcesPerGI(OBJGPU *pGpu, NvU32 swizzId, KERNEL_PHYS_GPU_INFO *pPgpuInfo)
{
    OBJSYS     *pSys     = SYS_GET_INSTANCE();
    KernelVgpuMgr *pKernelVgpuMgr = SYS_GET_KERNEL_VGPUMGR(pSys);
    KERNEL_VGPU_TYPE_PLACEMENT_INFO *pKernelVgpuTypePlacementInfo = NULL, *pDefaultKernelVgpuTypePlacementInfo;
    NvU32 index, i, j, guestMaxVmmuCount, totalVmmuCount = 0, partitionFlag, numCreatable = 0, memDiv;
    NvU64 vmmuOffsetMin, vmmuSegmentMin, vmmuSegSize,
          maxTotalGspHeapSize, gspHeapStartOffset, gspHeapOffset,
          pmaRegionLength, pmaBaseOffset, largestOffset;
    VGPU_TYPE *pVgpuTypeInfo;
    VGPU_TYPE_PLACEMENT_INFO *pVgpuTypePlacementInfo, *pDefaultVgpuTypePlacementInfo;
    VGPU_TYPE_SUPPORTED_PLACEMENT_INFO *pVgpuTypeSupportedPlacementInfo;
    VGPU_INSTANCE_PLACEMENT_INFO *pVgpuInstancePlacementInfo;
    VGPU_INSTANCE_SUPPORTED_PLACEMENT_INFO *pVgpuInstanceSupportedPlacementInfo;
    KernelMIGManager *pKernelMIGManager = GPU_GET_KERNEL_MIG_MANAGER(pGpu);
    KERNEL_MIG_GPU_INSTANCE *pKernelMIGGpuInstance;
    Heap *pMemoryPartitionHeap  = NULL;
    Device *pDevice;
    RsClient *pClient;
    NvU32 *pSupportedTypeId;
    NvBool bFoundOne2OneProfile = 0;

    NV_CHECK_OR_RETURN(LEVEL_ERROR, pPgpuInfo != NULL, NV_ERR_INVALID_ARGUMENT);

    NV_CHECK_OK_OR_RETURN(LEVEL_ERROR,
            kvgpumgrGetPgpuIndex(pKernelVgpuMgr, pGpu->gpuId, &index));
    pPgpuInfo = &(pKernelVgpuMgr->pgpuInfo[index]);

    if (!gpuIsSriovEnabled(pGpu) || !IS_MIG_IN_USE(pGpu))
    {
        return NV_ERR_INVALID_ARGUMENT;
    }

    NV_ASSERT_OK_OR_RETURN(kmigmgrGetGPUInstanceInfo(pGpu, pKernelMIGManager, swizzId, &pKernelMIGGpuInstance));

    /*
     * Get the heap for the MIG partition and find memory range
     * vmmuOffsetMin = start of the GI memory aligned to VMMU segment size
     */
    NV_CHECK_OK_OR_RETURN(LEVEL_ERROR,
                    serverGetClientUnderLock(&g_resServ, pKernelMIGGpuInstance->instanceHandles.hClient, &pClient));

    NV_CHECK_OR_RETURN(LEVEL_ERROR,
                      (deviceGetByHandle(pClient, pKernelMIGGpuInstance->instanceHandles.hDevice, &pDevice) == NV_OK),
                      NV_ERR_INSUFFICIENT_PERMISSIONS);

    NV_CHECK_OR_RETURN(LEVEL_ERROR, (kmigmgrGetMemoryPartitionHeapFromDevice(pGpu,
                                       pKernelMIGManager, pDevice, &pMemoryPartitionHeap) == NV_OK),
                                       NV_ERR_INSUFFICIENT_PERMISSIONS);
    pmaGetLargestFree(pMemoryPartitionHeap->pPmaObject, &pmaRegionLength, &pmaBaseOffset, &largestOffset);
    vmmuOffsetMin = NV_ALIGN_UP((pmaBaseOffset + largestOffset), gpuGetVmmuSegmentSize(pGpu));
    vmmuSegmentMin = vmmuOffsetMin / gpuGetVmmuSegmentSize(pGpu);
    vmmuSegSize = gpuGetVmmuSegmentSize(pGpu);

    maxTotalGspHeapSize = _kvgpumgrGetMaxTotalGspHeapSize(pGpu, pPgpuInfo, pKernelMIGGpuInstance->partitionFlag);

    // Find GI placement info
    for (i = 0; i < MAX_GI_PARTITIONS_PER_GPU; i++)
    {
        if (pPgpuInfo->kernelVgpuTypePlacementInfoGi[i].bUsed &&
            pPgpuInfo->kernelVgpuTypePlacementInfoGi[i].swizzId == swizzId)
        {
            pKernelVgpuTypePlacementInfo =
                &pPgpuInfo->kernelVgpuTypePlacementInfoGi[i].kernelVgpuTypePlacementInfo;
            break;
        }
    }
    NV_ASSERT_OR_RETURN(pKernelVgpuTypePlacementInfo != NULL, NV_ERR_OBJECT_NOT_FOUND);
    memDiv = _kvgpumgrGetMemDivFromPartitionFlag(
                (DRF_VAL(2080_CTRL_GPU, _PARTITION_FLAG, _MEMORY_SIZE, pKernelMIGGpuInstance->partitionFlag)));
    if (memDiv != 0)
    {
        pKernelVgpuTypePlacementInfo->placementRegionSize =
            pPgpuInfo->kernelVgpuTypePlacementInfo.placementRegionSize /memDiv;
    }

    // Set initial creatable vgpu types for this GI
    pSupportedTypeId = kvgpuMgrGetVgpuCreatableTypeIdFromSwizzId(pGpu, pPgpuInfo, swizzId);

    pDefaultKernelVgpuTypePlacementInfo = &pPgpuInfo->kernelVgpuTypePlacementInfo;

    // Find the 1:1 profile of the current GI, calculate totalVmmuCount
    for (i = 0; i < pPgpuInfo->numVgpuTypes; i++)
    {
        pVgpuTypeInfo = pPgpuInfo->vgpuTypes[i];
        if (pVgpuTypeInfo == NULL)
            break;

        // Ignore non-MIG vgpu types
        if (pVgpuTypeInfo->gpuInstanceSize == 0)
            continue;

        kvgpumgrGetPartitionFlag(pVgpuTypeInfo->vgpuTypeId, &partitionFlag);

        // Ignore vgpu types that does not belong to the current GI
        if (partitionFlag != pKernelMIGGpuInstance->partitionFlag)
            continue;

        // Set initial creatable vgpu types for this GI
        pSupportedTypeId[numCreatable++] = pVgpuTypeInfo->vgpuTypeId;

        // Use 1:1 profile to get the totalVmmuCount
        if (!bFoundOne2OneProfile && pVgpuTypeInfo->maxInstancePerGI == 1)
        {
            /*
             * Global(Default) placement info will be aligned to the per GI placement info index
             * use its guestVmmuCount calculated from kvgpumgrSetSupportedPlacementIds, per GI placement info
             * is calcualted in the next loop
             */
            pDefaultVgpuTypePlacementInfo = &pDefaultKernelVgpuTypePlacementInfo->vgpuTypePlacementInfo[i];
            totalVmmuCount = pDefaultVgpuTypePlacementInfo->guestVmmuCount
                           + pVgpuTypeInfo->gspHeapSize / gpuGetVmmuSegmentSize(pGpu);
            bFoundOne2OneProfile = NV_TRUE;
        }
    }
    NV_ASSERT_OR_RETURN(bFoundOne2OneProfile, NV_ERR_INVALID_ARGUMENT);

    guestMaxVmmuCount = totalVmmuCount - ((maxTotalGspHeapSize + vmmuSegSize - 1) / vmmuSegSize);
    gspHeapStartOffset = (vmmuSegmentMin + guestMaxVmmuCount) * vmmuSegSize;

    // Iterate vgpu types belongs to current GI, then calculate placement resources
    for (i = 0; i < pPgpuInfo->numVgpuTypes; i++)
    {
        NvU32 placementRegionSize;
        pVgpuTypeInfo = pPgpuInfo->vgpuTypes[i];
        if (pVgpuTypeInfo == NULL)
            break;

        // Ignore non-MIG vgpu types
        if (pVgpuTypeInfo->gpuInstanceSize == 0)
            continue;

        kvgpumgrGetPartitionFlag(pVgpuTypeInfo->vgpuTypeId, &partitionFlag);

        // Ignore vgpu types that does not belong to the current GI
        if (partitionFlag != pKernelMIGGpuInstance->partitionFlag)
            continue;

        pVgpuTypeSupportedPlacementInfo = &pVgpuTypeInfo->vgpuTypeSupportedPlacementInfo;
        pVgpuTypePlacementInfo = &pKernelVgpuTypePlacementInfo->vgpuTypePlacementInfo[i];
        pDefaultVgpuTypePlacementInfo = &pDefaultKernelVgpuTypePlacementInfo->vgpuTypePlacementInfo[i];

        // Copy from Global(default) placement info calculated in kvgpumgrSetSupportedPlacementIds
        pVgpuTypePlacementInfo->guestVmmuCount = pDefaultVgpuTypePlacementInfo->guestVmmuCount;

        placementRegionSize = vgpuMgrGetPlacementRegionSize(pGpu, &pPgpuInfo->vgpuNominalFbSize, pVgpuTypeInfo);

        for (j = 0; j < pVgpuTypeSupportedPlacementInfo->heterogeneousPlacementCount; j++)
        {
            pVgpuInstancePlacementInfo = &pVgpuTypePlacementInfo->vgpuInstancePlacementInfo[j];
            pVgpuInstanceSupportedPlacementInfo =
                &pVgpuTypeSupportedPlacementInfo->vgpuInstanceSupportedPlacementInfo[j];

            if (pVgpuTypeInfo->maxInstancePerGI == 1)
            {
                pVgpuInstancePlacementInfo->heterogeneousSupportedVmmuOffset = vmmuSegmentMin;
                pVgpuInstancePlacementInfo->heterogeneousGspHeapOffset =
                            (vmmuSegmentMin + pVgpuTypePlacementInfo->guestVmmuCount) * vmmuSegSize;
            }
            else
            {
                NvU32 placementId = pVgpuInstanceSupportedPlacementInfo->heterogeneousSupportedPlacementId;

                pVgpuInstancePlacementInfo->heterogeneousSupportedVmmuOffset =
                            vmmuSegmentMin + ((guestMaxVmmuCount * placementId) / placementRegionSize);
                gspHeapOffset = (maxTotalGspHeapSize * placementId) / placementRegionSize;
                pVgpuInstancePlacementInfo->heterogeneousGspHeapOffset =
                            gspHeapStartOffset + NV_ALIGN_DOWN(gspHeapOffset, RM_PAGE_SIZE_HUGE);
            }
        }
    }
    return NV_OK;
}

static NV_STATUS
_kvgpumgrSetHomogeneousResourcesPerGI(OBJGPU *pGpu, NvU32 swizzId, KERNEL_PHYS_GPU_INFO *pPgpuInfo)
{
    KERNEL_VGPU_TYPE_PLACEMENT_INFO *pKernelVgpuTypePlacementInfo = NULL;
    VGPU_TYPE_SUPPORTED_PLACEMENT_INFO *pVgpuTypeSupportedPlacementInfo;
    VGPU_TYPE_PLACEMENT_INFO *pVgpuTypePlacementInfo;
    VGPU_TYPE *pVgpuTypeInfo;
    NvU32 i, j, partitionFlag;
    VGPU_INSTANCE_PLACEMENT_INFO *pVgpuInstancePlacementInfo;
    KernelMIGManager *pKernelMIGManager = GPU_GET_KERNEL_MIG_MANAGER(pGpu);
    KERNEL_MIG_GPU_INSTANCE *pKernelMIGGpuInstance;
    VGPU_INSTANCE_SUPPORTED_PLACEMENT_INFO *pVgpuInstanceSupportedPlacementInfo;

    NV_CHECK_OR_RETURN(LEVEL_ERROR, pPgpuInfo != NULL, NV_ERR_INVALID_ARGUMENT);

    NV_ASSERT_OK_OR_RETURN(kmigmgrGetGPUInstanceInfo(pGpu, pKernelMIGManager, swizzId, &pKernelMIGGpuInstance));

    if (!pPgpuInfo->heterogeneousTimesliceSizesSupported)
    {
        pPgpuInfo->homogeneousPlacementSupported = NV_FALSE;
        NV_CHECK_OR_RETURN(LEVEL_ERROR, pPgpuInfo->homogeneousPlacementSupported, NV_ERR_NOT_SUPPORTED);
    }

    for (i = 0; i < MAX_GI_PARTITIONS_PER_GPU; i++)
    {
        if (pPgpuInfo->kernelVgpuTypePlacementInfoGi[i].bUsed &&
            pPgpuInfo->kernelVgpuTypePlacementInfoGi[i].swizzId == swizzId)
        {
            pKernelVgpuTypePlacementInfo =
                &pPgpuInfo->kernelVgpuTypePlacementInfoGi[i].kernelVgpuTypePlacementInfo;
            break;
        }
    }
    NV_ASSERT_OR_RETURN(pKernelVgpuTypePlacementInfo != NULL, NV_ERR_OBJECT_NOT_FOUND);

    for (i = 0; i < pPgpuInfo->numVgpuTypes; i++)
    {
        pVgpuTypeInfo = pPgpuInfo->vgpuTypes[i];
        if (pVgpuTypeInfo == NULL)
            break;

        // Ignore non-MIG vgpu types
        if (pVgpuTypeInfo->gpuInstanceSize == 0)
            continue;

        kvgpumgrGetPartitionFlag(pVgpuTypeInfo->vgpuTypeId, &partitionFlag);

        // Ignore vgpu types that does not belong to the current GI
        if (partitionFlag != pKernelMIGGpuInstance->partitionFlag)
            continue;

        pVgpuTypeSupportedPlacementInfo = &pVgpuTypeInfo->vgpuTypeSupportedPlacementInfo;
        pVgpuTypePlacementInfo = &pKernelVgpuTypePlacementInfo->vgpuTypePlacementInfo[i];

        /*
         * For GB20x on MIG, heterogeneous and homogeneous placement count is the same, hence copy it over using
         * heterogenous placement count
         */
        for (j = 0; j < pVgpuTypeSupportedPlacementInfo->heterogeneousPlacementCount; j++)
        {
            pVgpuInstancePlacementInfo = &pVgpuTypePlacementInfo->vgpuInstancePlacementInfo[j];
            pVgpuInstanceSupportedPlacementInfo =
                &pVgpuTypeSupportedPlacementInfo->vgpuInstanceSupportedPlacementInfo[j];
            // Copy from heterogeneous vmmu offset
            pVgpuInstancePlacementInfo->homogeneousSupportedVmmuOffset =
                                pVgpuInstancePlacementInfo->heterogeneousSupportedVmmuOffset;
            /*
             * Initially, all supported placements are creatable
             * Copy from homogeneous placement count, homogeneous and heterogeneous placements are same
             */
            pVgpuInstancePlacementInfo->creatablePlacementId =
                pVgpuInstanceSupportedPlacementInfo->homogeneousSupportedPlacementId;
        }
    }
    return NV_OK;
}

