412 struct npu_context {  struct
420 	void (*release_cb)(struct npu_context *context, void *priv);  argument
532 #define mn_to_npu_context(x) container_of(x, struct npu_context, mn)
557 static void acquire_atsd_reg(struct npu_context *npu_context,  in acquire_atsd_reg()  argument
574 			npdev = READ_ONCE(npu_context->npdev[i][j]);  in acquire_atsd_reg()
618 static void mmio_invalidate(struct npu_context *npu_context, int va,  in mmio_invalidate()  argument
622 	unsigned long pid = npu_context->mm->context.id;  in mmio_invalidate()
624 	if (npu_context->nmmu_flush)  in mmio_invalidate()
630 		flush_all_mm(npu_context->mm);  in mmio_invalidate()
636 	acquire_atsd_reg(npu_context, mmio_atsd_reg);  in mmio_invalidate()
660 	struct npu_context *npu_context = mn_to_npu_context(mn);  in pnv_npu2_mn_release()  local
663 	if (npu_context->release_cb)  in pnv_npu2_mn_release()
664 		npu_context->release_cb(npu_context, npu_context->priv);  in pnv_npu2_mn_release()
670 	mmio_invalidate(npu_context, 0, 0, true);  in pnv_npu2_mn_release()
678 	struct npu_context *npu_context = mn_to_npu_context(mn);  in pnv_npu2_mn_change_pte()  local
680 	mmio_invalidate(npu_context, 1, address, true);  in pnv_npu2_mn_change_pte()
687 	struct npu_context *npu_context = mn_to_npu_context(mn);  in pnv_npu2_mn_invalidate_range()  local
695 		mmio_invalidate(npu_context, 0, 0, true);  in pnv_npu2_mn_invalidate_range()
698 			mmio_invalidate(npu_context, 1, address, false);  in pnv_npu2_mn_invalidate_range()
701 		mmio_invalidate(npu_context, 1, address, true);  in pnv_npu2_mn_invalidate_range()
726 struct npu_context *pnv_npu2_init_context(struct pci_dev *gpdev,  in pnv_npu2_init_context()
728 			void (*cb)(struct npu_context *, void *),  in pnv_npu2_init_context()  argument
737 	struct npu_context *npu_context;  in pnv_npu2_init_context()  local
785 	npu_context = mm->context.npu_context;  in pnv_npu2_init_context()
786 	if (npu_context) {  in pnv_npu2_init_context()
787 		if (npu_context->release_cb != cb ||  in pnv_npu2_init_context()
788 			npu_context->priv != priv) {  in pnv_npu2_init_context()
796 		WARN_ON(!kref_get_unless_zero(&npu_context->kref));  in pnv_npu2_init_context()
800 	if (!npu_context) {  in pnv_npu2_init_context()
808 		npu_context = kzalloc(sizeof(struct npu_context), GFP_KERNEL);  in pnv_npu2_init_context()
809 		if (npu_context) {  in pnv_npu2_init_context()
810 			kref_init(&npu_context->kref);  in pnv_npu2_init_context()
811 			npu_context->mm = mm;  in pnv_npu2_init_context()
812 			npu_context->mn.ops = &nv_nmmu_notifier_ops;  in pnv_npu2_init_context()
813 			rc = __mmu_notifier_register(&npu_context->mn, mm);  in pnv_npu2_init_context()
817 			kfree(npu_context);  in pnv_npu2_init_context()
824 		mm->context.npu_context = npu_context;  in pnv_npu2_init_context()
827 	npu_context->release_cb = cb;  in pnv_npu2_init_context()
828 	npu_context->priv = priv;  in pnv_npu2_init_context()
838 	WRITE_ONCE(npu_context->npdev[npu->index][nvlink_index], npdev);  in pnv_npu2_init_context()
845 		npu_context->nmmu_flush = false;  in pnv_npu2_init_context()
848 		npu_context->nmmu_flush = true;  in pnv_npu2_init_context()
850 	return npu_context;  in pnv_npu2_init_context()
856 	struct npu_context *npu_context =  in pnv_npu2_release_context()  local
857 		container_of(kref, struct npu_context, kref);  in pnv_npu2_release_context()
859 	if (!npu_context->nmmu_flush)  in pnv_npu2_release_context()
860 		mm_context_remove_copro(npu_context->mm);  in pnv_npu2_release_context()
862 	npu_context->mm->context.npu_context = NULL;  in pnv_npu2_release_context()
869 void pnv_npu2_destroy_context(struct npu_context *npu_context,  in pnv_npu2_destroy_context()  argument
891 	WRITE_ONCE(npu_context->npdev[npu->index][nvlink_index], NULL);  in pnv_npu2_destroy_context()
892 	opal_npu_destroy_context(nphb->opal_id, npu_context->mm->context.id,  in pnv_npu2_destroy_context()
895 	removed = kref_put(&npu_context->kref, pnv_npu2_release_context);  in pnv_npu2_destroy_context()
903 		mmu_notifier_unregister(&npu_context->mn,  in pnv_npu2_destroy_context()
904 					npu_context->mm);  in pnv_npu2_destroy_context()
906 		kfree(npu_context);  in pnv_npu2_destroy_context()
915 int pnv_npu2_handle_fault(struct npu_context *context, uintptr_t *ea,  in pnv_npu2_handle_fault()