/*******************************************************************************
    Copyright (c) 2015 NVIDIA Corporation

    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 "uvm_ioctl.h"
#include "uvm8_va_space.h"
#include "uvm8_va_range.h"
#include "uvm8_va_block.h"
#include "uvm8_api.h"
#include "uvm8_tracker.h"
#include "uvm8_gpu.h"

// Returns true if the interval [start, start + length - 1] is entirely covered
// by valid vmas. A vma is valid if the corresponding VM_SPECIAL flags are not
// set
//
// Locking: current->mm->mmap_sem must be locked
bool uvm_is_valid_vma_range(NvU64 start, NvU64 length)
{
    const NvU64 end = start + length;
    struct vm_area_struct *vma;

    uvm_assert_mmap_sem_locked(&current->mm->mmap_sem);

    vma = find_vma_intersection(current->mm, start, end);

    while (vma && (vma->vm_start <= start) && !(vma->vm_flags & VM_SPECIAL)) {
        if (vma->vm_end >= end)
            return true;
        start = vma->vm_end;
        vma = vma->vm_next;
    }

    return false;
}

NV_STATUS uvm_api_range_type_check(uvm_va_space_t *va_space, NvU64 base, NvU64 length)
{
    uvm_va_range_t *va_range, *va_range_last;
    const NvU64 last_address = base + length - 1;

    uvm_assert_rwsem_locked(&va_space->lock);

    if (uvm_api_range_invalid(base, length))
        return NV_ERR_INVALID_ADDRESS;

    // Check if passed interval overlaps with any VA range.
    if (uvm_va_space_range_empty(va_space, base, last_address)) {
        if (g_uvm_global.ats.enabled && uvm_is_valid_vma_range(base, length))
            return NV_WARN_NOTHING_TO_DO;
        else
            return NV_ERR_INVALID_ADDRESS;
    }

    va_range_last = NULL;

    uvm_for_each_managed_va_range_in_contig(va_range, va_space, base, last_address)
        va_range_last = va_range;

    // Check if passed interval overlaps with an unmanaged VA range, or a
    // sub-interval not tracked by a VA range
    if (!va_range_last || va_range_last->node.end < last_address)
        return NV_ERR_INVALID_ADDRESS;

    // Passed interval is fully covered by managed VA ranges
    return NV_OK;
}

static bool preferred_location_is_va_range_split_needed(uvm_va_range_t *va_range, void *data)
{
    uvm_processor_id_t processor_id;
    UVM_ASSERT(data);
    processor_id = *(uvm_processor_id_t*)data;
    return (processor_id != va_range->preferred_location);
}

static NV_STATUS preferred_location_unmap_remote_pages(uvm_va_block_t *va_block,
                                                       uvm_va_block_context_t *va_block_context)
{
    NV_STATUS status = NV_OK;
    NV_STATUS tracker_status;
    uvm_tracker_t local_tracker = UVM_TRACKER_INIT();
    uvm_va_range_t *va_range = va_block->va_range;
    uvm_processor_id_t preferred_location = va_range->preferred_location;
    const uvm_page_mask_t *mapped_mask;

    if (preferred_location == UVM_MAX_PROCESSORS || !uvm_processor_mask_test(&va_block->mapped, preferred_location))
        goto done;

    // Read duplication takes precedence over PreferredLocation. No mappings
    // need to be removed.
    if (uvm_va_range_is_read_duplicate(va_range))
        goto done;

    mapped_mask = uvm_va_block_map_mask_get(va_block, preferred_location);

    if (uvm_processor_mask_test(&va_block->resident, preferred_location)) {
        const uvm_page_mask_t *resident_mask = uvm_va_block_resident_mask_get(va_block, preferred_location);

        if (!uvm_page_mask_andnot(&va_block_context->caller_page_mask, mapped_mask, resident_mask))
            goto done;
    }
    else {
        uvm_page_mask_copy(&va_block_context->caller_page_mask, mapped_mask);
    }

    status = uvm_va_block_unmap(va_block,
                                va_block_context,
                                preferred_location,
                                uvm_va_block_region_from_block(va_block),
                                &va_block_context->caller_page_mask,
                                &local_tracker);

    tracker_status = uvm_tracker_add_tracker_safe(&va_block->tracker, &local_tracker);
    if (status == NV_OK)
        status = tracker_status;

done:
    uvm_tracker_deinit(&local_tracker);

    return status;
}

NV_STATUS uvm_va_block_set_preferred_location_locked(uvm_va_block_t *va_block,
                                                     uvm_va_block_context_t *va_block_context)
{
    uvm_assert_mutex_locked(&va_block->lock);

    uvm_va_block_mark_cpu_dirty(va_block);

    return preferred_location_unmap_remote_pages(va_block, va_block_context);
}

static NV_STATUS preferred_location_set(uvm_va_space_t *va_space,
                                        NvU64 base,
                                        NvU64 length,
                                        uvm_processor_id_t preferred_location,
                                        uvm_va_range_t **first_va_range_to_migrate)
{
    uvm_va_range_t *va_range, *va_range_last;
    const NvU64 last_address = base + length - 1;
    bool preferred_location_is_faultable_gpu = false;
    NV_STATUS status;

    uvm_assert_rwsem_locked_write(&va_space->lock);

    if (preferred_location != UVM_MAX_PROCESSORS) {
        UVM_ASSERT(first_va_range_to_migrate != NULL);

        *first_va_range_to_migrate = NULL;
        preferred_location_is_faultable_gpu = preferred_location != UVM_CPU_ID &&
                                              uvm_processor_mask_test(&va_space->faultable_processors,
                                                                      preferred_location);
    }

    status = uvm_va_space_split_span_as_needed(va_space,
                                               base,
                                               last_address + 1,
                                               preferred_location_is_va_range_split_needed,
                                               &preferred_location);
    if (status != NV_OK)
        return status;

    va_range_last = NULL;
    uvm_for_each_managed_va_range_in_contig(va_range, va_space, base, last_address) {
        bool found_non_migratable_interval = false;

        va_range_last = va_range;

        // If we didn't split the ends, check that they match
        if (va_range->node.start < base || va_range->node.end > last_address)
            UVM_ASSERT(va_range->preferred_location == preferred_location);

        if (preferred_location != UVM_MAX_PROCESSORS) {
            const NvU64 start = max(base, va_range->node.start);
            const NvU64 end = min(last_address, va_range->node.end);

            found_non_migratable_interval = !uvm_range_group_all_migratable(va_space, start, end);

            if (found_non_migratable_interval && preferred_location_is_faultable_gpu)
                return NV_ERR_INVALID_DEVICE;
        }

        status = uvm_va_range_set_preferred_location(va_range, preferred_location);
        if (status != NV_OK)
            return status;

        // Return the first VA range that needs to be migrated so the caller
        // function doesn't need to traverse the tree again
        if (found_non_migratable_interval && (*first_va_range_to_migrate == NULL))
            *first_va_range_to_migrate = va_range;
    }

    UVM_ASSERT(va_range_last && va_range_last->node.end >= last_address);

    return NV_OK;
}

NV_STATUS uvm_api_set_preferred_location(const UVM_SET_PREFERRED_LOCATION_PARAMS *params, struct file *filp)
{
    NV_STATUS status;
    NV_STATUS tracker_status;
    uvm_tracker_t local_tracker = UVM_TRACKER_INIT();
    uvm_va_space_t *va_space = uvm_va_space_get(filp);
    uvm_va_range_t *va_range = NULL;
    uvm_va_range_t *first_va_range_to_migrate = NULL;
    uvm_processor_id_t preferred_location_id;
    bool has_va_space_write_lock;
    const NvU64 start = params->requestedBase;
    const NvU64 length = params->length;
    const NvU64 end = start + length - 1;
    bool range_is_ats = false;

    UVM_ASSERT(va_space);

    uvm_down_read_mmap_sem(&current->mm->mmap_sem);
    uvm_va_space_down_write(va_space);
    has_va_space_write_lock = true;

    status = uvm_api_range_type_check(va_space, start, length);
    if (status != NV_OK) {
        if (status != NV_WARN_NOTHING_TO_DO)
            goto done;

        status = NV_OK;
        range_is_ats = true;
    }

    // If the CPU is the preferred location, we don't have to find the associated uvm_gpu_t
    if (uvm_uuid_is_cpu(&params->preferredLocation)) {
        preferred_location_id = UVM_CPU_ID;
    }
    else {
        // Translate preferredLocation into a live GPU ID, and check that this
        // GPU can address the virtual address range
        uvm_gpu_t *gpu = uvm_va_space_get_gpu_by_uuid(va_space, &params->preferredLocation);

        if (!gpu)
            status = NV_ERR_INVALID_DEVICE;
        else if (!uvm_gpu_can_address(gpu, end))
            status = NV_ERR_OUT_OF_RANGE;

        if (status != NV_OK)
            goto done;

        preferred_location_id = gpu->id;
    }

    UVM_ASSERT(status == NV_OK);

    // TODO: Bug 2098544: On ATS systems, honor the preferred location hint for
    // system memory ranges instead of ignoring it
    if (range_is_ats)
        goto done;

    status = preferred_location_set(va_space, start, length, preferred_location_id, &first_va_range_to_migrate);
    if (status != NV_OK)
        goto done;

    // No VA range to migrate, early exit
    if (!first_va_range_to_migrate)
        goto done;

    uvm_va_space_downgrade_write(va_space);
    has_va_space_write_lock = false;

    // No need to check for holes in the VA ranges span here, this was checked by preferred_location_set
    for (va_range = first_va_range_to_migrate; va_range; va_range = uvm_va_space_iter_next(va_range, end)) {
        uvm_range_group_range_iter_t iter;
        NvU64 cur_start = max(start, va_range->node.start);
        NvU64 cur_end = min(end, va_range->node.end);

        uvm_range_group_for_each_migratability_in(&iter, va_space, cur_start, cur_end) {
            if (!iter.migratable) {
                status = uvm_range_group_va_range_migrate(va_range, iter.start, iter.end, &local_tracker);
                if (status != NV_OK)
                    goto done;
            }
        }
    }

done:
    tracker_status = uvm_tracker_wait_deinit(&local_tracker);

    if (has_va_space_write_lock)
        uvm_va_space_up_write(va_space);
    else
        uvm_va_space_up_read(va_space);
    uvm_up_read_mmap_sem(&current->mm->mmap_sem);

    return status == NV_OK ? tracker_status : status;
}

NV_STATUS uvm_api_unset_preferred_location(const UVM_UNSET_PREFERRED_LOCATION_PARAMS *params, struct file *filp)
{
    NV_STATUS status;
    uvm_va_space_t *va_space = uvm_va_space_get(filp);

    UVM_ASSERT(va_space);

    uvm_down_read_mmap_sem(&current->mm->mmap_sem);
    uvm_va_space_down_write(va_space);

    status = uvm_api_range_type_check(va_space, params->requestedBase, params->length);

    if (status == NV_OK)
        status = preferred_location_set(va_space, params->requestedBase, params->length, UVM_MAX_PROCESSORS, NULL);
    else if (status == NV_WARN_NOTHING_TO_DO)
        status = NV_OK;

    uvm_va_space_up_write(va_space);
    uvm_up_read_mmap_sem(&current->mm->mmap_sem);
    return status;
}

static NV_STATUS va_block_set_accessed_by_locked(uvm_va_block_t *va_block,
                                                 uvm_va_block_context_t *va_block_context,
                                                 uvm_processor_id_t processor_id,
                                                 uvm_tracker_t *out_tracker)
{
    NV_STATUS status;
    NV_STATUS tracker_status;

    uvm_assert_mutex_locked(&va_block->lock);

    status = uvm_va_block_add_mappings(va_block,
                                       va_block_context,
                                       processor_id,
                                       uvm_va_block_region_from_block(va_block),
                                       NULL,
                                       UvmEventMapRemoteCausePolicy);

    tracker_status = uvm_tracker_add_tracker_safe(out_tracker, &va_block->tracker);

    return status == NV_OK ? tracker_status : status;
}

NV_STATUS uvm_va_block_set_accessed_by(uvm_va_block_t *va_block,
                                       uvm_va_block_context_t *va_block_context,
                                       uvm_processor_id_t processor_id)
{
    NV_STATUS status;
    uvm_tracker_t local_tracker = UVM_TRACKER_INIT();

    // Read duplication takes precedence over SetAccesedBy. Do not add mappings
    // if read duplication is enabled
    if (uvm_va_range_is_read_duplicate(va_block->va_range))
        return NV_OK;

    status = UVM_VA_BLOCK_LOCK_RETRY(va_block, NULL,
                                     va_block_set_accessed_by_locked(va_block,
                                                                     va_block_context,
                                                                     processor_id,
                                                                     &local_tracker));

    // TODO: Bug 1767224: Combine all accessed_by operations into single tracker
    if (status == NV_OK)
        status = uvm_tracker_wait(&local_tracker);

    uvm_tracker_deinit(&local_tracker);
    return status;
}

typedef struct
{
    uvm_processor_id_t processor_id;
    bool set_bit;
} accessed_by_split_params_t;

static bool accessed_by_is_va_range_split_needed(uvm_va_range_t *va_range, void *data)
{
    accessed_by_split_params_t *params = (accessed_by_split_params_t*)data;
    UVM_ASSERT(params);

    return (uvm_processor_mask_test(&va_range->accessed_by, params->processor_id) != params->set_bit);
}

static NV_STATUS accessed_by_set(uvm_va_space_t *va_space,
                                 NvU64 base,
                                 NvU64 length,
                                 const NvProcessorUuid *processor_uuid,
                                 bool set_bit)
{
    uvm_processor_id_t processor_id = UVM_MAX_PROCESSORS;
    uvm_va_range_t *va_range, *va_range_last;
    const NvU64 last_address = base + length - 1;
    bool range_is_sysmem = false;
    accessed_by_split_params_t split_params;
    NV_STATUS status;

    UVM_ASSERT(va_space);

    uvm_down_read_mmap_sem(&current->mm->mmap_sem);
    uvm_va_space_down_write(va_space);

    status = uvm_api_range_type_check(va_space, base, length);
    if (status != NV_OK) {
        if (status != NV_WARN_NOTHING_TO_DO)
            goto done;
        status = NV_OK;
        range_is_sysmem = true;
    }

    if (uvm_uuid_is_cpu(processor_uuid)) {
        processor_id = UVM_CPU_ID;
    }
    else {
        // Translate processor_uuid into a live GPU ID, and check that this GPU
        // can address the virtual address range
        uvm_gpu_t *gpu = uvm_va_space_get_gpu_by_uuid(va_space, processor_uuid);
        if (!gpu)
            status = NV_ERR_INVALID_DEVICE;
        else if (!uvm_gpu_can_address(gpu, last_address))
            status = NV_ERR_OUT_OF_RANGE;

        if (status != NV_OK)
            goto done;

        processor_id = gpu->id;
    }

    if (range_is_sysmem)
        goto done;

    split_params.processor_id = processor_id;
    split_params.set_bit = set_bit;
    status = uvm_va_space_split_span_as_needed(va_space,
                                               base,
                                               last_address + 1,
                                               accessed_by_is_va_range_split_needed,
                                               &split_params);
    if (status != NV_OK)
        goto done;

    va_range_last = NULL;
    uvm_for_each_managed_va_range_in_contig(va_range, va_space, base, last_address) {
        va_range_last = va_range;

        // If we didn't split the ends, check that they match
        if (va_range->node.start < base || va_range->node.end > last_address)
            UVM_ASSERT(uvm_processor_mask_test(&va_range->accessed_by, processor_id) == set_bit);

        if (set_bit) {
            status = uvm_va_range_set_accessed_by(va_range, processor_id);
            if (status != NV_OK)
                goto done;
        }
        else {
            uvm_va_range_unset_accessed_by(va_range, processor_id);
        }
    }

    UVM_ASSERT(va_range_last && va_range_last->node.end >= last_address);

done:
    uvm_va_space_up_write(va_space);
    uvm_up_read_mmap_sem(&current->mm->mmap_sem);

    return status;
}

NV_STATUS uvm_api_set_accessed_by(const UVM_SET_ACCESSED_BY_PARAMS *params, struct file *filp)
{
    uvm_va_space_t *va_space = uvm_va_space_get(filp);

    return accessed_by_set(va_space, params->requestedBase, params->length, &params->accessedByUuid, true);
}

NV_STATUS uvm_api_unset_accessed_by(const UVM_UNSET_ACCESSED_BY_PARAMS *params, struct file *filp)
{
    uvm_va_space_t *va_space = uvm_va_space_get(filp);

    return accessed_by_set(va_space, params->requestedBase, params->length, &params->accessedByUuid, false);
}

static NV_STATUS va_block_set_read_duplication_locked(uvm_va_block_t *va_block,
                                                      uvm_va_block_retry_t *va_block_retry,
                                                      uvm_va_block_context_t *va_block_context)
{
    uvm_processor_id_t src_id;

    uvm_assert_mutex_locked(&va_block->lock);

    for_each_id_in_mask(src_id, &va_block->resident) {
        NV_STATUS status;
        uvm_page_mask_t *resident_mask = uvm_va_block_resident_mask_get(va_block, src_id);

        // Calling uvm_va_block_make_resident_read_duplicate will break all
        // SetAccessedBy and remote mappings
        status = uvm_va_block_make_resident_read_duplicate(va_block,
                                                           va_block_retry,
                                                           va_block_context,
                                                           src_id,
                                                           uvm_va_block_region_from_block(va_block),
                                                           resident_mask,
                                                           NULL,
                                                           UVM_MAKE_RESIDENT_CAUSE_API_HINT);
        if (status != NV_OK)
            return status;
    }

    return NV_OK;
}

NV_STATUS uvm_va_block_set_read_duplication(uvm_va_block_t *va_block,
                                            uvm_va_block_context_t *va_block_context)
{
    NV_STATUS status;
    uvm_va_block_retry_t va_block_retry;

    status = UVM_VA_BLOCK_LOCK_RETRY(va_block, &va_block_retry,
                                     va_block_set_read_duplication_locked(va_block,
                                                                          &va_block_retry,
                                                                          va_block_context));

    return status;
}

static NV_STATUS va_block_unset_read_duplication_locked(uvm_va_block_t *va_block,
                                                        uvm_va_block_retry_t *va_block_retry,
                                                        uvm_va_block_context_t *va_block_context,
                                                        uvm_tracker_t *out_tracker)
{
    NV_STATUS status;
    uvm_processor_id_t processor_id;
    uvm_va_range_t *va_range = va_block->va_range;
    uvm_va_block_region_t block_region = uvm_va_block_region_from_block(va_block);
    uvm_page_mask_t *break_read_duplication_pages = &va_block_context->caller_page_mask;

    uvm_assert_mutex_locked(&va_block->lock);

    // 1- Iterate over all processors with resident copies to avoid migrations
    // and invalidate the rest of copies

    // If preferred_location is set and has resident copies, give it preference
    if (va_range->preferred_location != UVM_MAX_PROCESSORS &&
        uvm_processor_mask_test(&va_block->resident, va_range->preferred_location)) {
        uvm_page_mask_t *resident_mask = uvm_va_block_resident_mask_get(va_block, va_range->preferred_location);
        bool is_mask_empty = !uvm_page_mask_and(break_read_duplication_pages,
                                                &va_block->read_duplicated_pages,
                                                resident_mask);

        if (!is_mask_empty) {
            // make_resident breaks read duplication
            status = uvm_va_block_make_resident(va_block,
                                                va_block_retry,
                                                va_block_context,
                                                va_range->preferred_location,
                                                block_region,
                                                break_read_duplication_pages,
                                                NULL,
                                                UVM_MAKE_RESIDENT_CAUSE_API_HINT);
            if (status != NV_OK)
                return status;
        }
    }

    // Then iterate over the rest of processors
    for_each_id_in_mask(processor_id, &va_block->resident) {
        uvm_page_mask_t *resident_mask;
        bool is_mask_empty;

        if (processor_id == va_range->preferred_location)
            continue;

        resident_mask = uvm_va_block_resident_mask_get(va_block, processor_id);
        is_mask_empty = !uvm_page_mask_and(break_read_duplication_pages,
                                           &va_block->read_duplicated_pages,
                                           resident_mask);
        if (is_mask_empty)
            continue;

        // make_resident breaks read duplication
        status = uvm_va_block_make_resident(va_block,
                                            va_block_retry,
                                            va_block_context,
                                            processor_id,
                                            block_region,
                                            break_read_duplication_pages,
                                            NULL,
                                            UVM_MAKE_RESIDENT_CAUSE_API_HINT);
        if (status != NV_OK)
            return status;
    }

    // 2- Re-establish SetAccessedBy mappings
    for_each_id_in_mask(processor_id, &va_block->va_range->accessed_by) {
        status = va_block_set_accessed_by_locked(va_block,
                                                 va_block_context,
                                                 processor_id,
                                                 out_tracker);
        if (status != NV_OK)
            return status;
    }

    return NV_OK;
}

NV_STATUS uvm_va_block_unset_read_duplication(uvm_va_block_t *va_block,
                                              uvm_va_block_context_t *va_block_context)
{
    uvm_va_block_retry_t va_block_retry;
    NV_STATUS status = NV_OK;
    uvm_tracker_t local_tracker = UVM_TRACKER_INIT();

    // Restore all SetAccessedBy mappings
    status = UVM_VA_BLOCK_LOCK_RETRY(va_block, &va_block_retry,
                                     va_block_unset_read_duplication_locked(va_block,
                                                                            &va_block_retry,
                                                                            va_block_context,
                                                                            &local_tracker));
    if (status == NV_OK)
        status = uvm_tracker_wait(&local_tracker);

    uvm_tracker_deinit(&local_tracker);

    return status;
}

static bool read_duplication_is_va_range_split_needed(uvm_va_range_t *va_range, void *data)
{
    uvm_read_duplication_policy_t new_policy;

    UVM_ASSERT(data);

    new_policy = *(uvm_read_duplication_policy_t *)data;
    return va_range->read_duplication != new_policy;
}

static NV_STATUS read_duplication_set(uvm_va_space_t *va_space, NvU64 base, NvU64 length, bool enable)
{
    uvm_va_range_t *va_range, *va_range_last;
    const NvU64 last_address = base + length - 1;
    NV_STATUS status;
    uvm_read_duplication_policy_t new_policy;

    UVM_ASSERT(va_space);

    // We need mmap_sem as we may create CPU mappings
    uvm_down_read_mmap_sem(&current->mm->mmap_sem);
    uvm_va_space_down_write(va_space);

    status = uvm_api_range_type_check(va_space, base, length);
    if (status != NV_OK) {
        if (status == NV_WARN_NOTHING_TO_DO)
            status = NV_OK;

        goto done;
    }

    // Note that we never set the policy back to UNSET
    new_policy = enable ? UVM_READ_DUPLICATION_ENABLED : UVM_READ_DUPLICATION_DISABLED;

    status = uvm_va_space_split_span_as_needed(va_space,
                                               base,
                                               last_address + 1,
                                               read_duplication_is_va_range_split_needed,
                                               &new_policy);
    if (status != NV_OK)
        goto done;

    va_range_last = NULL;
    uvm_for_each_managed_va_range_in_contig(va_range, va_space, base, last_address) {
        va_range_last = va_range;

        // If we didn't split the ends, check that they match
        if (va_range->node.start < base || va_range->node.end > last_address)
            UVM_ASSERT(va_range->read_duplication == new_policy);

        // If the va_space cannot currently read duplicate, only change the user
        // state. All memory should already have read duplication unset.
        if (uvm_va_space_can_read_duplicate(va_space, NULL)) {

            // Handle SetAccessedBy mappings
            if (new_policy == UVM_READ_DUPLICATION_ENABLED) {
                status = uvm_va_range_set_read_duplication(va_range);
                if (status != NV_OK)
                    goto done;
            }
            else {
                // If unsetting read duplication fails, the return status is
                // not propagated back to the caller
                (void)uvm_va_range_unset_read_duplication(va_range);
            }
        }

        va_range->read_duplication = new_policy;
    }

    UVM_ASSERT(va_range_last && va_range_last->node.end >= last_address);

done:
    uvm_va_space_up_write(va_space);
    uvm_up_read_mmap_sem(&current->mm->mmap_sem);
    return status;
}

NV_STATUS uvm_api_enable_read_duplication(const UVM_ENABLE_READ_DUPLICATION_PARAMS *params, struct file *filp)
{
    uvm_va_space_t *va_space = uvm_va_space_get(filp);

    return read_duplication_set(va_space, params->requestedBase, params->length, true);
}

NV_STATUS uvm_api_disable_read_duplication(const UVM_DISABLE_READ_DUPLICATION_PARAMS *params, struct file *filp)
{
    uvm_va_space_t *va_space = uvm_va_space_get(filp);

    return read_duplication_set(va_space, params->requestedBase, params->length, false);
}

static NV_STATUS system_wide_atomics_set(uvm_va_space_t *va_space, const NvProcessorUuid *gpu_uuid, bool enable)
{
    NV_STATUS status = NV_OK;
    uvm_gpu_t *gpu;
    bool already_enabled;

    uvm_va_space_down_write(va_space);

    gpu = uvm_va_space_get_gpu_by_uuid(va_space, gpu_uuid);
    if (!gpu) {
        status = NV_ERR_INVALID_DEVICE;
        goto done;
    }

    if (gpu->scoped_atomics_supported) {
        status = NV_ERR_NOT_SUPPORTED;
        goto done;
    }

    if (!uvm_processor_mask_test(&va_space->faultable_processors, gpu->id)) {
        status = NV_ERR_NOT_SUPPORTED;
        goto done;
    }

    already_enabled = uvm_processor_mask_test(&va_space->system_wide_atomics_enabled_processors, gpu->id);
    if (enable && !already_enabled) {
        uvm_va_range_t *va_range;
        uvm_tracker_t local_tracker = UVM_TRACKER_INIT();
        uvm_va_block_context_t *va_block_context = uvm_va_space_block_context(va_space);
        NV_STATUS tracker_status;

        // Revoke atomic mappings from the calling GPU
        uvm_for_each_va_range(va_range, va_space) {
            uvm_va_block_t *va_block;

            if (va_range->type != UVM_VA_RANGE_TYPE_MANAGED)
                continue;

            for_each_va_block_in_va_range(va_range, va_block) {
                uvm_page_mask_t *non_resident_pages = &va_block_context->caller_page_mask;

                uvm_mutex_lock(&va_block->lock);

                if (!uvm_processor_mask_test(&va_block->mapped, gpu->id)) {
                    uvm_mutex_unlock(&va_block->lock);
                    continue;
                }

                uvm_page_mask_complement(non_resident_pages, &va_block->gpus[uvm_gpu_index(gpu->id)]->resident);

                status = uvm_va_block_revoke_prot(va_block,
                                                  va_block_context,
                                                  gpu->id,
                                                  uvm_va_block_region_from_block(va_block),
                                                  non_resident_pages,
                                                  UVM_PROT_READ_WRITE_ATOMIC,
                                                  &va_block->tracker);

                tracker_status = uvm_tracker_add_tracker_safe(&local_tracker, &va_block->tracker);

                uvm_mutex_unlock(&va_block->lock);

                if (status == NV_OK)
                    status = tracker_status;

                if (status != NV_OK) {
                    uvm_tracker_deinit(&local_tracker);
                    goto done;
                }
            }
        }
        status = uvm_tracker_wait_deinit(&local_tracker);

        uvm_processor_mask_set(&va_space->system_wide_atomics_enabled_processors, gpu->id);
    }
    else if (!enable && already_enabled) {
        // TODO: Bug 1767229: Promote write mappings to atomic
        uvm_processor_mask_clear(&va_space->system_wide_atomics_enabled_processors, gpu->id);
    }

done:
    uvm_va_space_up_write(va_space);
    return status;
}

NV_STATUS uvm_api_enable_system_wide_atomics(UVM_ENABLE_SYSTEM_WIDE_ATOMICS_PARAMS *params, struct file *filp)
{
    uvm_va_space_t *va_space = uvm_va_space_get(filp);

    return system_wide_atomics_set(va_space, &params->gpu_uuid, true);
}

NV_STATUS uvm_api_disable_system_wide_atomics(UVM_DISABLE_SYSTEM_WIDE_ATOMICS_PARAMS *params, struct file *filp)
{
    uvm_va_space_t *va_space = uvm_va_space_get(filp);

    return system_wide_atomics_set(va_space, &params->gpu_uuid, false);
}
