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.