AMDGPU: add parallel restore of BO content to accelerate restore

[wweewrwer]

TL;DR:

This pull request extends CRIU to support parallel restore of AMDGPU buffer object content alongside other restore operations to accelerate the restoration.

The target issue:

In the current restore procedure of AMDGPU applications, the content of the AMDGPU buffer object (BO) is restored synchronously in CR_PLUGIN_HOOK__RESTORE_EXT_FILE. This procedure usually takes a significant amount of time, and during this time the target process cannot perform any other restore operations. However, this restoration has no logical dependencies with other restore operations. Parallelizing this part with other restore operations can speed up the restoration.

The parallel restore approach in this PR:

The core idea of these patch series is to offload the restore of the BO content from the target process to the main CRIU process (the main CRIU process refers to the parent process, and the target process refers to the child process created during the fork). To achieve this, we introduce a new hook, CR_PLUGIN_HOOK__RESTORE_ASYNCHRONOUS, in the main CRIU process. For the AMDGPU plugin, the target process will no longer restore BO contents in CR_PLUGIN_HOOK__RESTORE_EXT_FILE and just send the relevant BOs to the main CRIU process. the main CRIU process will receive the corresponding BOs in CR_PLUGIN_HOOK__RESTORE_ASYNCHRONOUS and begin the restoration. Meanwhile, the target process can continue with other parts of the restoration without being blocked by the BO content restoration. The full design of the idea can also be referred with the ACM SoCC'24 paper: On-demand and Parallel Checkpoint/Restore for GPU Applications.

Tests:

We evaluated the performance according to the following settings. The results show that parallel restore can speed up by 34.3% when images cached in the page cache, and 7.6% when restoring from disk.

Results:

	From disk	From page cache
Sequential restore	1728ms	254ms
Parallel restore	1596ms	167ms
Speed up	7.6%	34.3%

Settings:

CPU: Intel(R) Core(TM) i7-10700 CPU @ 2.90GHz

Memory: DDR4, 2x8GB

GPU: AMD MI50

Disk: 512GB, Samsung SSD 860

Docker image: rocm/pytorch:rocm5.6_ubuntu20.04_py3.8_pytorch_1.12.1

Example program：

example.py: a ResNet18 application. Enter 'y' to exit, or press any other key to perform inference.

import time
import os
import sys
import torch
import torchvision.models as models
import torchvision.transforms as transforms
torch.set_grad_enabled(False)

device = "cuda:0"

model = models.resnet18(weights='DEFAULT')
model = model.to(device)
model.eval()

batch_size = 1
channels = 3
height = 224
width = 224
input_tensor = torch.randn(batch_size, channels, height, width)
preprocess = transforms.Compose([
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])
input_tensor = preprocess(input_tensor)

while input()!="y":
    st = time.time()
    input_tensor = input_tensor.to(device)
    output = model(input_tensor)
    output = output.to("cpu")
    _, predicted_idx = torch.max(output, 1)
    torch.cuda.synchronize()
    ed = time.time()
    print("test time:",ed-st)
    sys.stdout.flush()

Steps：

0. Install CRIU

   Follow the standard CRIU installation process. Ensure you set the environment variable CRIU_LIBS_DIR to the plugins/amdgpu path.

1. Dump checkpoint image

   #In one shell
   python3 example.py
   #In another shell
   mkdir -p /tmp/criu-dump
   criu dump -t $(pgrep python3) -D /tmp/criu-dump -j --file-locks
   

2. Restore from disk

   Test for sequential restore:

   #Clear page cache
   sync; sudo sh -c "echo 3 > /proc/sys/vm/drop_caches" 
   criu restore -D /tmp/criu-dump -j --file-locks
   cat stats-restore | crit decode --pretty | grep restore_time
   

   Test for parallel restore:

   sync; sudo sh -c "echo 3 > /proc/sys/vm/drop_caches" 
   criu restore -D /tmp/criu-dump -j --file-locks --parallel
   cat stats-restore | crit decode --pretty | grep restore_time
   

3. Restore from page cache

   Install vmtouch for caching images:

   sudo apt install vmtouch
   

   Test:

   sync; sudo sh -c "echo 3 > /proc/sys/vm/drop_caches" 
   #Cache image in memory
   vmtouch -l criu-dump
   #Warm up environment 
   criu restore -D /tmp/criu-dump -j --file-locks
   #Begin to Test
   criu restore -D /tmp/criu-dump -j --file-locks
   cat stats-restore | crit decode --pretty | grep restore_time
   criu restore -D /tmp/criu-dump -j --file-locks --parallel
   cat stats-restore | crit decode --pretty | grep restore_time
   

[Ddnirvana]

Thanks for the above comments @avagin @rst0git , we are fixing and polishing the PR. Will update ASAP.

[rst0git]

@Ddnirvana @wweewrwer Thank you for your contributions! It might be good to also update the content of the following files to reflect these changes:

• https://github.com/checkpoint-restore/criu/blob/criu-dev/plugins/amdgpu/README.md
• https://github.com/checkpoint-restore/criu/blob/criu-dev/Documentation/criu-amdgpu-plugin.txt

[Ddnirvana]

@Ddnirvana @wweewrwer Thank you for your contributions! It might be good to also update the content of the following files to reflect these changes:

• https://github.com/checkpoint-restore/criu/blob/criu-dev/plugins/amdgpu/README.md
• https://github.com/checkpoint-restore/criu/blob/criu-dev/Documentation/criu-amdgpu-plugin.txt

@rst0git No problem. We will add proper description in the next version.

[dayatsin-amd]

Thank you @wweewrwer. Some minor nit picks, but overall the code looks good to me.

[wweewrwer]

@rst0git @avagin @dayatsin-amd hi maintainers, thanks for your prior reviews and comments. We have fixed all the issues, as the following:

1. Use the proper APIs to allocate (xmalloc, etc.)
2. Enable the optimizations by default
3. Change the name of hook
4. Fix the issues to run in Podman containers
5. Other fixes (line width, comments, etc.)
6. Add descriptions in README to explain the optimizations.

Please let us know if you have any further comments

[dayatsin-amd]

Thank you @wweewrwer

[rst0git]

@wweewrwer Would you be able to merge the fixup commits into the previous commits using git rebase?
https://github.com/checkpoint-restore/criu/blob/criu-dev/CONTRIBUTING.md#submit-your-work-upstream

[wweewrwer]

@wweewrwer Would you be able to merge the fixup commits into the previous commits using git rebase? https://github.com/checkpoint-restore/criu/blob/criu-dev/CONTRIBUTING.md#submit-your-work-upstream

@rst0git Thanks for your comment! I have merged the fixup commits into the previous commits using git rebase. Please let me know if you have any further comments.

[wweewrwer]

@rst0git @avagin
Dear maintainers,

We have pushed the V4 version of the PR, completing all mentioned issues since the last version. Specifically, we: (1) support multiple commands (from a single process), (2) support multiple processes restore, and (3) fix other minor issues mentioned.

Details:

• Replaced UDP with TCP to distinguish messages between different processes and commands.
• Multiple-command support: Instead of receiving the command only once, the hook function now launches a dedicated thread to receive commands indefinitely until all tasks finish their restore stage. The main thread in this hook uses restore_wait_inprogress_tasks to determine when tasks have finished. Once completed, it sends an exit command to the parallel restore thread to stop receiving commands.
• Multi-process support: In the case of multiple processes, they are restored in parallel (with different processes) by default, which will not benefit from the parallel optimizations. Therefore, we introduce a flag (called parallel_disabled) to only enable the optimization for single-process (which is the common case) as a fast path, and fallback to original restore otherwise.
• Multi-GPU parallel restore support: In the original restore, when a process has multiple GPUs, the content on each GPU is restored in parallel. In this version, we have supported multi-GPU parallel restore utilizing the original design.
• Other issues: Big thanks to Andrei and Radostin for other issues and suggestions, which are all fixed accordingly.

We have performed all the tests with the above changes. The PR can still bring 31% decrease for the restore latency in the case of single process, and achieves the same results for mutlti-process scenarios.

Please let me know if you have any further comments.

[wweewrwer]

@rst0git @avagin Just a friendly reminder about the updates in this PR (in case maintainers miss the prior notifications)

[avagin]

Have you investigated other approaches of parallel restoring of BO? For example, it is possible to fork a thread and restoring BO asynchronously in context of its process. In this case, two BO will be restored concurrently.

[wweewrwer]

Have you investigated other approaches of parallel restoring of BO? For example, it is possible to fork a thread and restoring BO asynchronously in context of its process. In this case, two BO will be restored concurrently.

Yes, we have investigated the approach of forking a thread in the background, but it cannot work as it conflicts with the restore logic of CRIU.

Specifically, when CRIU tries to restore its memory state, it will unmap all old mappings. However, some mappings may be needed by the background thread for BO restoring. Therefore, a thread can only run in parallel with shorter procedures (possibly before entering the restorer blob), while offloading the restore of BO content to a new process (in this PR) can be parallelized with almost the entire restore procedure.

Below figure shows the issues that BO restore must be finished before the CPU memory state restore:

[wweewrwer]

@rst0git @avagin Dear maintainers/reviewers, just want to know if there are any further issues/concerns about the latest version?

[dayatsin-amd]

I have requested this PR to be validated on a multi-GPU set-up internally at AMD. Can you give us a few days to confirm there is no regression.

Thank you for this patch!

[wweewrwer]

I have requested this PR to be validated on a multi-GPU set-up internally at AMD. Can you give us a few days to confirm there is no regression.

Thank you for this patch!

Sure. Thanks!

[Ddnirvana]

I have requested this PR to be validated on a multi-GPU set-up internally at AMD. Can you give us a few days to confirm there is no regression.

Thank you for this patch!

Dear David @dayatsin-amd , just want to know are there any progress/results about the internal regression test. Thank you again for the assistance and happy new year btw :)

[avagin]

Have you investigated other approaches of parallel restoring of BO? For example, it is possible to fork a thread and restoring BO asynchronously in context of its process. In this case, two BO will be restored concurrently.

Yes, we have investigated the approach of forking a thread in the background, but it cannot work as it conflicts with the restore logic of CRIU.

Specifically, when CRIU tries to restore its memory state, it will unmap all old mappings. However, some mappings may be needed by the background thread for BO restoring. Therefore, a thread can only run in parallel with shorter procedures (possibly before entering the restorer blob), while offloading the restore of BO content to a new process (in this PR) can be parallelized with almost the entire restore procedure.

Everything what is happening in the restore blob should be fast. All mappings are restored before switching into the restore blob. There, the restored mappings are just remapped to proper addresses. I am still not convinced that the idea of restoring buffer objects from the main process is really what we need here.. I can miss something, but I want to see a clear explanation with numbers why the proposed solution is a valuable one.

Additionally, I see two potential issues:

• Sequential Restoration: This change seems to introduce a new bottleneck by restoring buffer objects sequentially. Could this cause performance problems for workloads with many buffer objects across multiple processes? It would be helpful to understand how this approach scales.
• Plugin Hook Execution: Running the plugin hook in the main CRIU process for an extended period and making it dependent on other processes is problematic. This deviates from the expectation that multiple plugins should operate independently with equal capabilities.