linux/drivers/gpu/drm/xe/xe_tlb_inval_job.c

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

#include "xe_assert.h"
#include "xe_dep_job_types.h"
#include "xe_dep_scheduler.h"
#include "xe_exec_queue.h"
#include "xe_gt_printk.h"
#include "xe_gt_types.h"
#include "xe_page_reclaim.h"
#include "xe_tlb_inval.h"
#include "xe_tlb_inval_job.h"
#include "xe_migrate.h"
#include "xe_pm.h"
#include "xe_vm.h"

/** struct xe_tlb_inval_job - TLB invalidation job */
struct xe_tlb_inval_job {
	/** @dep: base generic dependency Xe job */
	struct xe_dep_job dep;
	/** @tlb_inval: TLB invalidation client */
	struct xe_tlb_inval *tlb_inval;
	/** @q: exec queue issuing the invalidate */
	struct xe_exec_queue *q;
	/** @vm: VM which TLB invalidation is being issued for */
	struct xe_vm *vm;
	/** @prl: Embedded copy of page reclaim list */
	struct xe_page_reclaim_list prl;
	/** @refcount: ref count of this job */
	struct kref refcount;
	/**
	 * @fence: dma fence to indicate completion. 1 way relationship - job
	 * can safely reference fence, fence cannot safely reference job.
	 */
	struct dma_fence *fence;
	/** @start: Start address to invalidate */
	u64 start;
	/** @end: End address to invalidate */
	u64 end;
	/** @type: GT type */
	int type;
	/** @fence_armed: Fence has been armed */
	bool fence_armed;
};

static struct dma_fence *xe_tlb_inval_job_run(struct xe_dep_job *dep_job)
{
	struct xe_tlb_inval_job *job =
		container_of(dep_job, typeof(*job), dep);
	struct xe_tlb_inval_fence *ifence =
		container_of(job->fence, typeof(*ifence), base);
	struct drm_suballoc *prl_sa = NULL;

	if (xe_page_reclaim_list_valid(&job->prl)) {
		prl_sa = xe_page_reclaim_create_prl_bo(job->tlb_inval, &job->prl, ifence);
		if (IS_ERR(prl_sa))
			prl_sa = NULL; /* Indicate fall back PPC flush with NULL */
	}

	xe_tlb_inval_range(job->tlb_inval, ifence, job->start,
			   job->end, job->vm->usm.asid, prl_sa);

	return job->fence;
}

static void xe_tlb_inval_job_free(struct xe_dep_job *dep_job)
{
	struct xe_tlb_inval_job *job =
		container_of(dep_job, typeof(*job), dep);

	/* Pairs with get in xe_tlb_inval_job_push */
	xe_tlb_inval_job_put(job);
}

static const struct xe_dep_job_ops dep_job_ops = {
	.run_job = xe_tlb_inval_job_run,
	.free_job = xe_tlb_inval_job_free,
};

/**
 * xe_tlb_inval_job_create() - TLB invalidation job create
 * @q: exec queue issuing the invalidate
 * @tlb_inval: TLB invalidation client
 * @dep_scheduler: Dependency scheduler for job
 * @vm: VM which TLB invalidation is being issued for
 * @start: Start address to invalidate
 * @end: End address to invalidate
 * @type: GT type
 *
 * Create a TLB invalidation job and initialize internal fields. The caller is
 * responsible for releasing the creation reference.
 *
 * Return: TLB invalidation job object on success, ERR_PTR failure
 */
struct xe_tlb_inval_job *
xe_tlb_inval_job_create(struct xe_exec_queue *q, struct xe_tlb_inval *tlb_inval,
			struct xe_dep_scheduler *dep_scheduler,
			struct xe_vm *vm, u64 start, u64 end, int type)
{
	struct xe_tlb_inval_job *job;
	struct drm_sched_entity *entity =
		xe_dep_scheduler_entity(dep_scheduler);
	struct xe_tlb_inval_fence *ifence;
	int err;

	xe_assert(vm->xe, type == XE_EXEC_QUEUE_TLB_INVAL_MEDIA_GT ||
		  type == XE_EXEC_QUEUE_TLB_INVAL_PRIMARY_GT);

	job = kmalloc_obj(*job);
	if (!job)
		return ERR_PTR(-ENOMEM);

	job->q = q;
	job->vm = vm;
	job->tlb_inval = tlb_inval;
	job->start = start;
	job->end = end;
	job->fence_armed = false;
	xe_page_reclaim_list_init(&job->prl);
	job->dep.ops = &dep_job_ops;
	job->type = type;
	kref_init(&job->refcount);
	xe_exec_queue_get(q);	/* Pairs with put in xe_tlb_inval_job_destroy */
	xe_vm_get(vm);		/* Pairs with put in xe_tlb_inval_job_destroy */

	ifence = kmalloc_obj(*ifence);
	if (!ifence) {
		err = -ENOMEM;
		goto err_job;
	}
	job->fence = &ifence->base;

	err = drm_sched_job_init(&job->dep.drm, entity, 1, NULL,
				 q->xef ? q->xef->drm->client_id : 0);
	if (err)
		goto err_fence;

	/* Pairs with put in xe_tlb_inval_job_destroy */
	xe_pm_runtime_get_noresume(gt_to_xe(q->gt));

	return job;

err_fence:
	kfree(ifence);
err_job:
	xe_vm_put(vm);
	xe_exec_queue_put(q);
	kfree(job);

	return ERR_PTR(err);
}

/**
 * xe_tlb_inval_job_add_page_reclaim() - Embed PRL into a TLB job
 * @job: TLB invalidation job that may trigger reclamation
 * @prl: Page reclaim list populated during unbind
 *
 * Copies @prl into the job and takes an extra reference to the entry page so
 * ownership can transfer to the TLB fence when the job is pushed.
 */
void xe_tlb_inval_job_add_page_reclaim(struct xe_tlb_inval_job *job,
				       struct xe_page_reclaim_list *prl)
{
	struct xe_device *xe = gt_to_xe(job->q->gt);

	xe_gt_WARN_ON(job->q->gt, !xe->info.has_page_reclaim_hw_assist);
	job->prl = *prl;
	/* Pair with put in job_destroy */
	xe_page_reclaim_entries_get(job->prl.entries);
}

static void xe_tlb_inval_job_destroy(struct kref *ref)
{
	struct xe_tlb_inval_job *job = container_of(ref, typeof(*job),
						    refcount);
	struct xe_tlb_inval_fence *ifence =
		container_of(job->fence, typeof(*ifence), base);
	struct xe_exec_queue *q = job->q;
	struct xe_device *xe = gt_to_xe(q->gt);
	struct xe_vm *vm = job->vm;

	/* BO creation retains a copy (if used), so no longer needed */
	xe_page_reclaim_entries_put(job->prl.entries);

	if (!job->fence_armed)
		kfree(ifence);
	else
		/* Ref from xe_tlb_inval_fence_init */
		dma_fence_put(job->fence);

	drm_sched_job_cleanup(&job->dep.drm);
	kfree(job);
	xe_vm_put(vm);		/* Pairs with get from xe_tlb_inval_job_create */
	xe_exec_queue_put(q);	/* Pairs with get from xe_tlb_inval_job_create */
	xe_pm_runtime_put(xe);	/* Pairs with get from xe_tlb_inval_job_create */
}

/**
 * xe_tlb_inval_job_alloc_dep() - TLB invalidation job alloc dependency
 * @job: TLB invalidation job to alloc dependency for
 *
 * Allocate storage for a dependency in the TLB invalidation fence. This
 * function should be called at most once per job and must be paired with
 * xe_tlb_inval_job_push being called with a real fence.
 *
 * Return: 0 on success, -errno on failure
 */
int xe_tlb_inval_job_alloc_dep(struct xe_tlb_inval_job *job)
{
	xe_assert(gt_to_xe(job->q->gt), !xa_load(&job->dep.drm.dependencies, 0));
	might_alloc(GFP_KERNEL);

	return drm_sched_job_add_dependency(&job->dep.drm,
					    dma_fence_get_stub());
}

/**
 * xe_tlb_inval_job_push() - TLB invalidation job push
 * @job: TLB invalidation job to push
 * @m: The migration object being used
 * @fence: Dependency for TLB invalidation job
 *
 * Pushes a TLB invalidation job for execution, using @fence as a dependency.
 * Storage for @fence must be preallocated with xe_tlb_inval_job_alloc_dep
 * prior to this call if @fence is not signaled. Takes a reference to the job’s
 * finished fence, which the caller is responsible for releasing, and return it
 * to the caller. This function is safe to be called in the path of reclaim.
 *
 * Return: Job's finished fence on success, cannot fail
 */
struct dma_fence *xe_tlb_inval_job_push(struct xe_tlb_inval_job *job,
					struct xe_migrate *m,
					struct dma_fence *fence)
{
	struct xe_tlb_inval_fence *ifence =
		container_of(job->fence, typeof(*ifence), base);

	if (!dma_fence_is_signaled(fence)) {
		void *ptr;

		/*
		 * Can be in path of reclaim, hence the preallocation of fence
		 * storage in xe_tlb_inval_job_alloc_dep. Verify caller did
		 * this correctly.
		 */
		xe_assert(gt_to_xe(job->q->gt),
			  xa_load(&job->dep.drm.dependencies, 0) ==
			  dma_fence_get_stub());

		dma_fence_get(fence);	/* ref released once dependency processed by scheduler */
		ptr = xa_store(&job->dep.drm.dependencies, 0, fence,
			       GFP_ATOMIC);
		xe_assert(gt_to_xe(job->q->gt), !xa_is_err(ptr));
	}

	xe_tlb_inval_job_get(job);	/* Pairs with put in free_job */
	job->fence_armed = true;

	/*
	 * We need the migration lock to protect the job's seqno and the spsc
	 * queue, only taken on migration queue, user queues protected dma-resv
	 * VM lock.
	 */
	xe_migrate_job_lock(m, job->q);

	/* Creation ref pairs with put in xe_tlb_inval_job_destroy */
	xe_tlb_inval_fence_init(job->tlb_inval, ifence, false);
	dma_fence_get(job->fence);	/* Pairs with put in DRM scheduler */

	drm_sched_job_arm(&job->dep.drm);
	/*
	 * caller ref, get must be done before job push as it could immediately
	 * signal and free.
	 */
	dma_fence_get(&job->dep.drm.s_fence->finished);
	drm_sched_entity_push_job(&job->dep.drm);

	/* Let the upper layers fish this out */
	xe_exec_queue_tlb_inval_last_fence_set(job->q, job->vm,
					       &job->dep.drm.s_fence->finished,
					       job->type);

	xe_migrate_job_unlock(m, job->q);

	/*
	 * Not using job->fence, as it has its own dma-fence context, which does
	 * not allow TLB invalidation fences on the same queue, GT tuple to
	 * be squashed in dma-resv/DRM scheduler. Instead, we use the DRM scheduler
	 * context and job's finished fence, which enables squashing.
	 */
	return &job->dep.drm.s_fence->finished;
}

/**
 * xe_tlb_inval_job_get() - Get a reference to TLB invalidation job
 * @job: TLB invalidation job object
 *
 * Increment the TLB invalidation job's reference count
 */
void xe_tlb_inval_job_get(struct xe_tlb_inval_job *job)
{
	kref_get(&job->refcount);
}

/**
 * xe_tlb_inval_job_put() - Put a reference to TLB invalidation job
 * @job: TLB invalidation job object
 *
 * Decrement the TLB invalidation job's reference count, call
 * xe_tlb_inval_job_destroy when reference count == 0. Skips decrement if
 * input @job is NULL or IS_ERR.
 */
void xe_tlb_inval_job_put(struct xe_tlb_inval_job *job)
{
	if (!IS_ERR_OR_NULL(job))
		kref_put(&job->refcount, xe_tlb_inval_job_destroy);
}