fix context switch on versal (Xilinx#2376)

src/runtime_src/core/pcie/driver/linux/xocl/userpf/common.h

@@ -128,6 +128,8 @@ struct xocl_dev	{
 
 	uint64_t			mig_cache_expire_secs;
 	ktime_t				mig_cache_expires;
+
+	struct mem_topology		*mem_topo;
 };
 
 /**

src/runtime_src/core/pcie/driver/linux/xocl/userpf/xocl_drm.c

@@ -595,7 +595,7 @@ int xocl_cleanup_mem(struct xocl_drm *drm_p)
 	return 0;
 }
 
-int xocl_init_mem(struct xocl_drm *drm_p, struct mem_topology *new_topo)
+int xocl_init_mem(struct xocl_drm *drm_p)
 {
 	size_t length = 0;
 	size_t mm_size = 0, mm_stat_size = 0;
@@ -618,10 +618,7 @@ int xocl_init_mem(struct xocl_drm *drm_p, struct mem_topology *new_topo)
 		reserved2 = 0x1000000;
 	}
 
-	if (XOCL_DSA_IS_VERSAL(drm_p->xdev))
-		topo = new_topo;
-	else
-		topo = XOCL_MEM_TOPOLOGY(drm_p->xdev);
+	topo = XOCL_MEM_TOPOLOGY(drm_p->xdev);
 
 	if (topo == NULL)
 		return 0;

src/runtime_src/core/pcie/driver/linux/xocl/userpf/xocl_drv.c

@@ -1023,6 +1023,8 @@ void xocl_userpf_remove(struct pci_dev *pdev)
 		vfree(xdev->core.dyn_subdev_store);
 	if (xdev->ulp_blob)
 		vfree(xdev->ulp_blob);
+	if (xdev->mem_topo)
+		vfree(xdev->mem_topo);
 	mutex_destroy(&xdev->core.lock);
 	mutex_destroy(&xdev->dev_lock);
 	mutex_destroy(&xdev->wq_lock);

src/runtime_src/core/pcie/driver/linux/xocl/userpf/xocl_ioctl.c

@@ -385,10 +385,15 @@ xocl_read_axlf_helper(struct xocl_drm *drm_p, struct drm_xocl_axlf *axlf_ptr)
 	}
 
 	/* Switching the xclbin, make sure none of the buffers are used. */
-	if (!preserve_mem && !XOCL_DSA_IS_VERSAL(xdev)) {
+	if (!preserve_mem) {
 		err = xocl_cleanup_mem(drm_p);
 		if (err)
 			goto done;
+
+		if (XOCL_DSA_IS_VERSAL(xdev)) {
+			vfree(xdev->mem_topo);
+			xdev->mem_topo = NULL;
+		}
 	}
 
 	if (!XOCL_DSA_IS_VERSAL(xdev)) {
@@ -411,16 +416,17 @@ xocl_read_axlf_helper(struct xocl_drm *drm_p, struct drm_xocl_axlf *axlf_ptr)
 			(void) xocl_pci_rbar_refresh(xdev->core.pdev,
 					xdev->p2p_bar_idx);
 		}
+	} else if (!preserve_mem) {
+		xdev->mem_topo = vmalloc(size);
+		if (!xdev->mem_topo) {
+			err = -ENOMEM;
+			goto done;
+		}
+		memcpy(xdev->mem_topo, new_topology, size);
 	}
 
-	if (XOCL_DSA_IS_VERSAL(xdev)) {
-		if (drm_p->mm == NULL) {
-			rc = xocl_init_mem(drm_p, new_topology);
-			if (err == 0)
-				err = rc;
-		}
-	} else if (!preserve_mem) {
-		rc = xocl_init_mem(drm_p, NULL);
+	if (!preserve_mem) {
+		rc = xocl_init_mem(drm_p);
 		if (err == 0)
 			err = rc;
 	}

src/runtime_src/core/pcie/driver/linux/xocl/xocl_drm.h

@@ -86,7 +86,7 @@ int xocl_mm_insert_node(struct xocl_drm *drm_p, u32 ddr,
 void *xocl_drm_init(xdev_handle_t xdev);
 void xocl_drm_fini(struct xocl_drm *drm_p);
 uint32_t xocl_get_shared_ddr(struct xocl_drm *drm_p, struct mem_data *m_data);
-int xocl_init_mem(struct xocl_drm *drm_p, struct mem_topology *topo);
+int xocl_init_mem(struct xocl_drm *drm_p);
 int xocl_cleanup_mem(struct xocl_drm *drm_p);
 
 #if LINUX_VERSION_CODE >= KERNEL_VERSION(4, 13, 0)

src/runtime_src/core/pcie/driver/linux/xocl/xocl_drv.h

@@ -575,7 +575,9 @@ struct xocl_mb_scheduler_funcs {
 	-ENODEV)
 
 #define XOCL_MEM_TOPOLOGY(xdev)						\
-	((struct mem_topology *)xocl_icap_get_data(xdev, MEMTOPO_AXLF))
+	(XOCL_DSA_IS_VERSAL(xdev) ?					\
+	((struct mem_topology *)(((struct xocl_dev *)(xdev))->mem_topo)) : \
+	((struct mem_topology *)xocl_icap_get_data(xdev, MEMTOPO_AXLF)))
 #define XOCL_IP_LAYOUT(xdev)						\
 	((struct ip_layout *)xocl_icap_get_data(xdev, IPLAYOUT_AXLF))
 #define XOCL_XCLBIN_ID(xdev)						\

src/runtime_src/core/pcie/linux/shim.cpp

@@ -1609,7 +1609,8 @@ int xclLoadXclBin(xclDeviceHandle handle, const xclBin *buffer)
       START_DEVICE_PROFILING_CB(handle);
     }
     if (!ret && xrt_core::config::get_ert() &&
-      xclbin::get_axlf_section(buffer, PDI) &&
+      (xclbin::get_axlf_section(buffer, PDI) ||
+      xclbin::get_axlf_section(buffer, BITSTREAM_PARTIAL_PDI)) &&
       xrt_core::config::get_pdi_load())
         ret = xrt_core::scheduler::loadXclbinToPS(handle, buffer);
     return ret;