In order to facilitate vfio-ap device migration:

* The function pointers to the callbacks used to facilitate migration of a
  vfio-ap device must be plugged in. These function pointers are specified
  via the fields contained in a struct vfio_migration_ops instance:

  	migration_set_state: Set the migration state of the vfio-ap device
  	migration_get_state: Get the migration state of the vfio-ap device
  	migration_get_data_size: Get the estimated data number bytes of
  				 data that will be required to complete
  				 STOP_COPY
* The migration_flags field of the struct vfio_device instance is set to
  VFIO_MIGRATION_STOP_COPY to indicate that STOP, STOP_COPY and RESUMING
  are the migration states supported for vfio-ap device migration.

* The initial migration state of the vfio-ap device is set to RUNNING

* The boolean flag indicating that all data sent from userspace during the
  RESUMINMG phase has been received is initialized to false.

Signed-off-by: Anthony Krowiak <akrowiak@linux.ibm.com>
---
 drivers/s390/crypto/vfio_ap_migration.c | 24 ++++++++++++++++++++++++
 drivers/s390/crypto/vfio_ap_ops.c       | 10 ++++++++++
 drivers/s390/crypto/vfio_ap_private.h   |  1 +
 3 files changed, 35 insertions(+)

diff --git a/drivers/s390/crypto/vfio_ap_migration.c b/drivers/s390/crypto/vfio_ap_migration.c
index 4af734aeee17..924c25585987 100644
--- a/drivers/s390/crypto/vfio_ap_migration.c
+++ b/drivers/s390/crypto/vfio_ap_migration.c
@@ -217,6 +217,30 @@ static int vfio_ap_release_mig_file(struct inode *file_inode, struct file *filp)
 	return 0;
 }
 
+static const struct vfio_migration_ops vfio_ap_migration_ops;
+
+/**
+ * vfio_ap_setup_vfio_device_migration:
+ *
+ * Plug the object containing the function pointers to the migration callbacks
+ * used to facilitate migration into the vfio device and set the flags
+ * indicating migration states STOP, STOP_COPY and RESUMING are supported.
+ *
+ * @matrix_mdev: pointer to object containing the vfio device (vdev)
+ */
+void vfio_ap_setup_vfio_device_migration(struct ap_matrix_mdev *matrix_mdev)
+{
+	/*
+	 * Live guest migration is not supported for SE guests, so there is no
+	 * need to initialize the migration capabilities.
+	 */
+	if (ap_is_se_guest())
+		return;
+
+	matrix_mdev->vdev.migration_flags = VFIO_MIGRATION_STOP_COPY;
+	matrix_mdev->vdev.mig_ops = &vfio_ap_migration_ops;
+}
+
 /**
  * validate_stop_copy_read_parms: Validate the input parameters to the
  *                                vfio_ap_stop_copy_read function
diff --git a/drivers/s390/crypto/vfio_ap_ops.c b/drivers/s390/crypto/vfio_ap_ops.c
index 8d70089eb9c9..52de74dc2c6c 100644
--- a/drivers/s390/crypto/vfio_ap_ops.c
+++ b/drivers/s390/crypto/vfio_ap_ops.c
@@ -811,9 +811,19 @@ static int vfio_ap_mdev_probe(struct mdev_device *mdev)
 
 	matrix_mdev = vfio_alloc_device(ap_matrix_mdev, vdev, &mdev->dev,
 					&vfio_ap_matrix_dev_ops);
+
 	if (IS_ERR(matrix_mdev))
 		return PTR_ERR(matrix_mdev);
 
+	/*
+	 * Set up vfio-ap device migration before calling
+	 * vfio_register_emulated_iommu_dev; otherwise, the VFIO core
+	 * will see mig_ops as NULL during the registration. This could
+	 * prevent the VFIO core from properly setting up migration
+	 * infrastructure like debugfs entries.
+	 */
+	vfio_ap_setup_vfio_device_migration(matrix_mdev);
+
 	ret = vfio_register_emulated_iommu_dev(&matrix_mdev->vdev);
 	if (ret)
 		goto err_put_vdev;
diff --git a/drivers/s390/crypto/vfio_ap_private.h b/drivers/s390/crypto/vfio_ap_private.h
index c6e577064e7e..19772e59fdc6 100644
--- a/drivers/s390/crypto/vfio_ap_private.h
+++ b/drivers/s390/crypto/vfio_ap_private.h
@@ -173,6 +173,7 @@ void vfio_ap_on_scan_complete(struct ap_config_info *new_config_info,
 			      struct ap_config_info *old_config_info);
 
 int vfio_ap_init_migration_data(struct ap_matrix_mdev *matrix_mdev);
+void vfio_ap_setup_vfio_device_migration(struct ap_matrix_mdev *matrix_mdev);
 void vfio_ap_release_migration_data(struct ap_matrix_mdev *matrix_mdev);
 void vfio_ap_reset_migration_state(struct ap_matrix_mdev *matrix_mdev);
 
-- 
2.53.0