adev->gfx.funcs->switch_partition_mode(xcp_mgr->adev,
                                                       num_xcc_per_xcp);
 
-       if (adev->nbio.funcs->set_compute_partition_mode)
-               adev->nbio.funcs->set_compute_partition_mode(adev, mode);
-
        /* Init info about new xcps */
        *num_xcps = num_xcc / num_xcc_per_xcp;
        amdgpu_xcp_init(xcp_mgr, *num_xcps, mode);
 
 static int gfx_v9_4_3_switch_compute_partition(struct amdgpu_device *adev,
                                                int num_xccs_per_xcp)
 {
-       int i, num_xcc;
-       u32 tmp = 0;
-
-       num_xcc = NUM_XCC(adev->gfx.xcc_mask);
+       int ret;
 
-       for (i = 0; i < num_xcc; i++) {
-               tmp = REG_SET_FIELD(tmp, CP_HYP_XCP_CTL, NUM_XCC_IN_XCP,
-                                   num_xccs_per_xcp);
-               tmp = REG_SET_FIELD(tmp, CP_HYP_XCP_CTL, VIRTUAL_XCC_ID,
-                                   i % num_xccs_per_xcp);
-               WREG32_SOC15(GC, GET_INST(GC, i), regCP_HYP_XCP_CTL, tmp);
-       }
+       ret = psp_spatial_partition(&adev->psp, NUM_XCC(adev->gfx.xcc_mask) /
+                                                       num_xccs_per_xcp);
+       if (ret)
+               return ret;
 
        adev->gfx.num_xcc_per_xcp = num_xccs_per_xcp;
 
-       return 0;
+       return ret;
 }
 
 static int gfx_v9_4_3_ih_to_xcc_inst(struct amdgpu_device *adev, int ih_node)