Make memory hooks work with static libcudart, by using bistro hooks for cuda driver API by yosefe · Pull Request #6038 · openucx/ucx

yosefe · 2020-12-16T21:22:15Z

Why

Support Cuda memory allocate/release hooks also for applications which are linked to libcudart statically, or use directly the driver API.
Fixed #3210

How

Use binary instrumentation (bistro) hooks for Cuda Driver APIs, which work regardless of how a function is called (from .plt or by function pointer).

Before this PR, if an application was statically linked to libcudart, we could not install hooks on Cuda memory events: cuda runtime calls (e.g cudaFree) are not going through .plt, and libcudart is calling driver API (e.g cuMemFree) by function pointer.
After this PR, the calls from cudart to cuda driver API can be intercepter.

If any other library is intercepting the same function by .plt hooks, it can work alongside with bistro hooks. However, if the other library is also using binary patching, in will probably be incompatible with the UCX hooks.

Binary instrumentation hook method:

Copy minimal number of bytes of target function prologue to a newly allocated buffer with execute permissions.
After these code bytes, place an indirect jump instruction to the rest of the target function code.

If there are not enough position independent instructions, such as push <reg> or sub <imm>, <reg>, to accommodate
for the indirect jump (3), this method will fail.
Place an indirect jump in the target function start to the UCM replacement
UCM replacement will treat the executable buffer (1) as the "original" function which should be used from the replacement (2)

memory allocate/release
in the application
      +
      |
      |       target function
      |       (e.g cuMemAlloc)               UCM replacement
      +----->-----------------+              (e.g ucm_cuMemAlloc)
             |jmp to UCM hook +-------------->+-----------+
             |                | (3)           | dispatch  +-------------> [ rcache ]
             |                |               | UCM event |
             +----------------+               |           |
             | unmodified     +<---+          +-----------+
             | code of target |    |          | call orig | (4)
             | function       |    |      +---+ function  |
             |                |    |      |   +-----------+
             |                |    |      |
             |                |    |      |   executable buffer
             |                |    |      |   (e.g ucm_orig_cuMemAlloc)
             |                |    |      +-->+-----------+
             |                |    |          | copy of   | (1)
             |                |    |          | function  |
             |                |    |          | prologue  |
             |                |    |          +-----------+
             |                |    +----------+ jump back | (2)
             +----------------+               +-----------+

hoopoepg · 2020-12-16T21:28:23Z

+        }
+
+        /* check the opcode */
+        if (((rex == 0) || rex == 0x41) && ((opcode & 0xF0) == 0x50)) {


((rex == 0) || rex == 0x41) is this correct?

rex == 0 is pushing IA32 register and rex==0x41 is pushing rN register

hoopoepg · 2020-12-16T21:34:52Z

+            }
+        }
+
+        /* unsupported instruction - bail */


bail? not fail?

yes, "bail" means exit

hoopoepg · 2020-12-16T21:36:19Z

+
+    pthread_mutex_lock(&mutex);
+
+    if (page == MAP_FAILED) {


this condition is always true

hoopoepg · 2020-12-16T21:37:59Z

+        }
+    }
+
+    if (size + offset > ucm_get_page_size()) {


offset is 0 here... and wrap size + offset by ()

hoopoepg · 2020-12-16T21:38:42Z

+
+    /* allocate space in the page */
+    result  = UCS_PTR_BYTE_OFFSET(page, offset);
+    offset += size;


it seems body of this function was in loop :)

forgot to add 'static'

brminich · 2020-12-17T06:39:09Z

+ * Set library function call hook using Binary Instrumentation
+ * method (BISTRO): replace function body by user defined call
+ *
+ * @param symbol function name to replace


missing @param func_ptr

brminich · 2020-12-17T06:53:16Z

    ucm_event_enter();

-    ucm_trace("ucm_cuMemFree(dptr=%p)",(void*)dptr);
+    ucm_diag("ucm_cuMemFree(dptr=%p)",(void*)dptr);


space after ,

brminich · 2020-12-17T06:53:55Z

+
+    ucm_event_enter();
+
+    ucm_diag("ucm_cuMemFree_v2(dptr=%p)",(void*)dptr);


space after ,

brminich · 2020-12-17T06:54:47Z

    ret = ucm_orig_cuMemAlloc(dptr, size);
    if (ret == CUDA_SUCCESS) {
-        ucm_trace("ucm_cuMemAlloc(dptr=%p size:%lu)",(void *)*dptr, size);
+        ucm_diag("ucm_cuMemAlloc(dptr=%p size:%lu)",(void *)*dptr, size);


would not diag be too flooding?

brminich · 2020-12-17T06:55:09Z

    ret = ucm_orig_cuMemAlloc(dptr, size);
    if (ret == CUDA_SUCCESS) {
-        ucm_trace("ucm_cuMemAlloc(dptr=%p size:%lu)",(void *)*dptr, size);
+        ucm_diag("ucm_cuMemAlloc(dptr=%p size:%lu)",(void *)*dptr, size);


space after ,

brminich · 2020-12-17T06:55:23Z

+
+    ret = ucm_orig_cuMemAlloc_v2(dptr, size);
+    if (ret == CUDA_SUCCESS) {
+        ucm_diag("ucm_cuMemAlloc_v2(dptr=%p size:%lu)",(void *)*dptr, size);


space after ,

brminich · 2020-12-17T06:55:58Z

+                                      ElementSizeBytes);
    if (ret == CUDA_SUCCESS) {
-        ucm_trace("ucm_cuMemHostGetDevicePointer(pdptr=%p p=%p)",(void *)*pdptr, p);
+        ucm_trace("ucm_cuMemAllocPitch_v2(dptr=%p size:%lu)",(void *)*dptr,


space after ,

brminich · 2020-12-17T06:57:52Z

+                continue;
+            }
+
+            ucm_debug("failed to install bistro hook for '%s', fallback to reloc",


better print status here than in previous ucs_debug, where it is always OK

yosefe · 2020-12-28T13:33:38Z

bot:pipe:retest

yosefe · 2020-12-29T21:42:54Z

/azp run

yosefe · 2020-12-29T21:42:59Z

bot:pipe:retest

azure-pipelines · 2020-12-29T21:43:07Z

Azure Pipelines successfully started running 1 pipeline(s).

yosefe · 2020-12-30T11:56:22Z

/azp run

azure-pipelines · 2020-12-30T11:56:34Z

Azure Pipelines successfully started running 1 pipeline(s).

yosefe · 2020-12-30T17:44:04Z

@Akshay-Venkatesh @bureddy can you pls take a look?

bureddy · 2020-12-30T22:16:12Z

- *               optional, may be NULL
+ * @param func_ptr  Pointer to function to patch.
+ * @param hook      User-defined function-replacer.
+ * @param symbol    Function name to replace.


should we we add @PARAM orig_func_p documentation here?

bureddy · 2020-12-31T02:22:15Z

+UCM_CUDA_FREE_FUNC(cuMemFree, CUresult, CUdeviceptr)
+UCM_CUDA_FREE_FUNC(cuMemFree_v2, CUresult, CUdeviceptr)
+UCM_CUDA_FREE_FUNC(cuMemFreeHost, CUresult, void*)
+UCM_CUDA_FREE_FUNC(cuMemFreeHost_v2, CUresult, void*)


cuMemFreeHost is also dispatching UCS_MEMORY_TYPE_CUDA memtype free event. it is not an issue today as we don't care if it not find in the mem type cache. but, should we differentiate cuMemFree and cuMemFreeHost?

bureddy · 2020-12-31T02:40:02Z

+    if (status == UCS_OK) {
+        ucm_debug("installed cuda memory bistro hooks on driver API");
+        ucm_cudamem_installed = 1;
+    } else {


we also have to install reloc for driver funcs if bistro fails. application might be using driver API directly.
can we update test_cuda_hook.c with driver api test as well?

yosefe · 2021-01-01T15:27:15Z

@bureddy fixed

Akshay-Venkatesh · 2021-01-02T18:13:35Z

@Akshay-Venkatesh @bureddy can you pls take a look?

Some high-level questions/requests:

Can you please explain how/why this solves the static linking issue? "How" section in the PR seems to describe the low-level mechanism of binary instrumentation. Specifying current limitations with cudart static linking (such as use of function pointers in cudart), how current hooks in ucm/cuda fails to intercept when cudart static is used, and how this PR addresses issues could be useful for future reference.
Should we document any limitations with current PR (such as memory allocated before ucp_init)?
Does this PR address the case when there are other libraries intercept cu* functions?

yosefe · 2021-01-02T20:37:11Z

Can you please explain how/why this solves the static linking issue? "How" section in the PR seems to describe the low-level mechanism of binary instrumentation. Specifying current limitations with cudart static linking (such as use of function pointers in cudart), how current hooks in ucm/cuda fails to intercept when cudart static is used, and how this PR addresses issues could be useful for future reference.

Added

Should we document any limitations with current PR (such as memory allocated before ucp_init)?

Memory before ucp_init is not a problem: we set suspicious regions to UNKNOWN and force detection during runtime

Does this PR address the case when there are other libraries intercept cu* functions?

Added note: if other library is hooking .plt it would work, if using bistro - no

bureddy · 2021-01-02T21:35:37Z

-    return ret;
+    event.mem_type.address  = (void*)ptr;
+    event.mem_type.size     = length;
+    event.mem_type.mem_type = UCS_MEMORY_TYPE_CUDA;


should this be "event.mem_type.mem_type = mem_type" ?

bureddy · 2021-01-02T22:30:21Z

-cudaError_t ucm_cudaHostGetDevicePointer(void **pDevice, void *pHost, unsigned int flags)
-{
-    cudaError_t ret;
+        if (ucm_mmap_hook_mode() == UCM_MMAP_HOOK_BISTRO) {


@yosefe should we detach cuda hooks from ucm_mmap_hook_mode() flag? may be we have a case where we need to to disable malloc bistro hooks but not cuda hooks ( for ex: on summit with spectrum mpi).
how about controlling cuda hooks with new flag "ucm_cuda_hook_mode: and "remove ucm_global_opts.enable_cuda_reloc" ?

bureddy · 2021-01-03T06:23:01Z


-    if (!ucm_global_opts.enable_cuda_reloc) {
+    if (ucm_cuda_hook_mode() == UCM_MMAP_HOOK_NONE) {
        ucm_debug("cudamem relocations are disabled by configuration");


s/relocations/hooks?

bureddy · 2021-01-03T06:29:10Z

+   " reloc  - Use ELF relocation table to set hooks.\n"
+#if UCM_BISTRO_HOOKS
+   " bistro - Use binary instrumentation to set hooks. In this mode, calls to\n"
+   "          statically a statically Cuda runtime can be hooked on the the\n"


seems wording is not clear

bureddy · 2021-01-03T06:51:59Z

+        status = ucm_cuda_install_hooks(ucm_cuda_runtime_funcs, &used_reloc);
+        if (status == UCS_OK) {
+            ucm_debug("installed cuda memory %s hooks on runtime API",
+                      used_reloc ? "reloc" : "bistro");


seems "ucm_cuda_install_hooks" can do bistro for some funcs if it is possible and reloc for others. is this debug msg accurate?

In order to catch cuda memory events even when Cuda is statically linked, use bistro.

yosefe added the WIP-DNM Work in progress / Do not review label Dec 16, 2020

hoopoepg reviewed Dec 16, 2020

View reviewed changes

yosefe force-pushed the topic/cuda-bistro-hooks branch from 76d2862 to ee105b1 Compare December 16, 2020 22:05

brminich reviewed Dec 17, 2020

View reviewed changes

yosefe force-pushed the topic/cuda-bistro-hooks branch from ee105b1 to 79c12b7 Compare December 19, 2020 12:16

yosefe removed the WIP-DNM Work in progress / Do not review label Dec 19, 2020

yosefe force-pushed the topic/cuda-bistro-hooks branch from 79c12b7 to d87b139 Compare December 19, 2020 12:47

yosefe force-pushed the topic/cuda-bistro-hooks branch 3 times, most recently from 8595659 to 1b69f10 Compare December 28, 2020 11:19

yosefe force-pushed the topic/cuda-bistro-hooks branch from 1b69f10 to 80693f4 Compare December 30, 2020 12:18

bureddy reviewed Dec 31, 2020

View reviewed changes

yosefe mentioned this pull request Dec 31, 2020

UCP/CUDA/ARCH: Catch mis-detected host memory by SEGV handler #5431

Closed

yosefe changed the title ~~UCM/CUDA: Add bistro hooks for Cuda~~ Make memory hooks work with static libcudart, by using bistro hooks for cuda driver API Jan 2, 2021

bureddy reviewed Jan 2, 2021

View reviewed changes

bureddy reviewed Jan 3, 2021

View reviewed changes

bureddy approved these changes Jan 3, 2021

View reviewed changes

yosefe added 2 commits January 4, 2021 10:25

UCM/BISTRO/CUDA: Add bisro hooks for Cuda

27444eb

In order to catch cuda memory events even when Cuda is statically linked, use bistro.

JENKINS: Fix distributed task workers count

705eb62

yosefe force-pushed the topic/cuda-bistro-hooks branch from 9502681 to 705eb62 Compare January 4, 2021 08:25

yosefe merged commit c79072c into openucx:master Jan 4, 2021

yosefe deleted the topic/cuda-bistro-hooks branch January 4, 2021 22:27

yosefe mentioned this pull request Jan 5, 2021

UCM/BISTRO: Improve x86 parser to recognize some mov flavors #6095

Merged

This was referenced Apr 25, 2024

OpenMPI 4.0.2 configured with UCX and CUDA-10.1-update2 crashes on DGX2/V100 platform #4307

Closed

ucx + cupy-core unexpectedly pulls in cuda-cudart conda-forge/ucx-split-feedstock#172

Closed


		ucm_event_enter();

		ucm_diag("ucm_cuMemFree_v2(dptr=%p)",(void*)dptr);

Conversation

yosefe commented Dec 16, 2020 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Why

How

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

yosefe commented Dec 28, 2020

Uh oh!

yosefe commented Dec 29, 2020

Uh oh!

yosefe commented Dec 29, 2020

Uh oh!

azure-pipelines Bot commented Dec 29, 2020

Uh oh!

yosefe commented Dec 30, 2020

Uh oh!

azure-pipelines Bot commented Dec 30, 2020

Uh oh!

yosefe commented Dec 30, 2020

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

yosefe commented Jan 1, 2021

Uh oh!

Akshay-Venkatesh commented Jan 2, 2021

Uh oh!

yosefe commented Jan 2, 2021

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

yosefe commented Dec 16, 2020 •

edited

Loading