Reject foreign allocation handles in the DirectX EP allocator (PLAT-207667) - #106
Reject foreign allocation handles in the DirectX EP allocator (PLAT-207667)#106danieyan-amd wants to merge 2 commits into
Conversation
…07667)
Amuse crashes on the AMDGPU execution device (Medusa/MDS1, EP 1.8.99.99) both
when generating with SD3 Medium and when installing SD1.5. Crash dump analysis
shows two distinct fault sites in directx-backend.dll, both reached the same way:
[this+0xf8] (m_allocator) -> DecodeDataHandle(ptr) -> deref result
* +0x587af mov rax,[rcx+8] reading the vtable slot for AddRef()
* +0x4f5aa mov rax,[rax+0x18] reading a member on the GetResource() path
Depending on the garbage read, this surfaces as 0xC0000005 or, when the bogus
vtable slot is a non-null invalid target, as a Control Flow Guard fastfail
(0xC0000409, FAST_FAIL_GUARD_ICALL_CHECK_FAILURE).
Root cause: DecodeDataHandle only null-checks and then static_casts. A pointer
that this allocator never produced - notably one from the factory's CpuAllocator
stub, which is plain malloc with no D3D12 resource - is non-null, so it passes
the check and every subsequent field read faults. The QueryInterface
discrimination already present in one caller does not help, because its fallback
path still blind-casts.
Fix, two layers:
1. The allocators track the handles they hand out and DecodeDataHandle rejects
anything else with E_INVALIDARG. Validation is a registry lookup rather than
an owner check, because reading GetOwner() would itself dereference the
untrusted pointer - the exact fault at +0x4f5aa. The registry is process-wide
so a handle that legitimately crosses allocator instances (sessions sharing
tensors) is still accepted.
2. ProviderFactory::CreateAllocatorImpl no longer falls through to the
CpuAllocator stub for GPU memory once an EP exists; it returns ORT_EP_FAIL so
session creation fails recoverably instead of handing out malloc'd pointers.
This is deliberately scoped to m_ep_raw != nullptr: ORT also requests a GPU
allocator during factory registration, before any EP exists, and failing there
breaks registration outright (observed while testing). The stub returned on
that path is harmless now that DecodeDataHandle validates.
Verified by building the DirectX backend locally (USE_DML, ORT 1.27) and running
a Conv/Relu/Add/ReduceMean model through the rebuilt EP: it registers, is
selected, and infers correctly, so the validation rejects no legitimate handle.
Note: DecodeDataHandle now takes a lock and a hash lookup per call. If this shows
up on GPU-path profiles, move the registry to a shared_mutex or a lock-free set.
DecodeDataHandle now rejects handles the allocator never produced, but three of its callers dereference the result with no exception boundary, so the throw escaped a noexcept frame and terminated the process (0xC0000409) instead of crashing on the bad pointer (0xC0000005). Both are fatal; neither is diagnosable. Guard those three so an invalid handle degrades to the result each caller already handles: GetABIDataInterface yields a null interface, TryGetPooledAllocationId returns 0 (the answer it already gives for an unpooled handle), and GetAllocationFromDataPointer returns nullptr - it had a dormant '!alloc_info' branch for exactly this case. Verified on a local repro (SD3.5 vae_decoder through the DirectX EP, which reproduces the reported access violation 3/3): shipped build crashes 0xC0000005, this build returns ORT_FAIL 'DML operator Compute failed: op=Transpose HR=0x80070057' 3/3 with the process intact.
|
danieyan-amd seems not to be a GitHub user. You need a GitHub account to be able to sign the CLA. If you have already a GitHub account, please add the email address used for this commit to your account. You have signed the CLA already but the status is still pending? Let us recheck it. |
Update — local repro found, fix verified end-to-endThe crash reproduces on ordinary hardware, no Amuse and no APU involved. Registering the shipped That is the same function, same size and same fault position as the crash reported on the mov rcx, [rbx+0xf8] ; m_allocator
call <DecodeDataHandle>
mov rbx, rax
test rax, rax / je ; null-checked, and it passes
mov rcx, [rax] ; vtable
mov rax, [rcx+8] ; <== FAULT (AddRef slot)A simple static Conv/Relu/Add/ReduceMean graph does not trigger it; it needs a real model, which Second commit: the rejection had to be made non-fatalWith only the first commit the process still died, just differently —
Each degrades to a value the caller already handles — Result
Open question for reviewersThe remaining question is why a non-EP pointer reaches Repro harness (parent/child so a crash is captured as an exit code rather than killing the run) is |
Root cause of the functional failure, traced — and it isn't what the code comment saysI instrumented a throwaway build (env-gated, since reverted) to log every Three things fall out of it:
Because The factory promises GPU memory and returns a CPU allocator, because it is asked before it can What the trace rules out
The second row is the one that matters for this PR: it confirms an exact-match registry is the The third row is worth flagging to whoever owns this file: the comment in Two consequences for this PR1. My 2. The cure is a separate change and needs an owner. The factory already holds the
Until that lands, this PR is containment: 8 crashes eliminated, 0 regressions, 0 models restored to Regression evidenceA simple static Conv/Relu/Add model runs clean on both shipped and patched builds ( Separately, SD1.5 |
|
Is this intentionally in draft status? |
Yes it is, im working on a fix for this jira |
Root cause found — it is ORT's memory-pattern planner, and there is a no-code-change workaroundFurther instrumentation identified the actual mechanism. It is not what this PR's description (or the The mechanism
return allocInfo.Detach(); // PluginDmlAllocationInfo*, not a data addressOffsetting that pointer produces something that is not an object.
This is why the original ticket says the app "crashes after generating an image". Run 1 records Proof, and a workaround that needs no code changeSame build, same model, one session option:
Across the model set with
All eight previously-crashing models run correctly on the unpatched binary. Users can be unblocked What this means for this PRIt is still worth merging, and it is still correct — a wild pointer dereference reachable from Two honest caveats:
I'd suggest (2) as the real fix and am happy to prepare it. The proper long-term answer is for the Unrelated defect found along the waySD1.5 |
Summary
Fixes a crash in the AMD GPU EP's DirectML backend where a pointer the DML allocator never produced is decoded as a
PluginDmlAllocationInfo*and dereferenced. Reported as PLAT-207667 (Amuse on Medusa/MDS1); it reproduces both when generating with SD3 Medium and when installing SD1.5.Root cause
DmlBucketizedBufferAllocator::DecodeDataHandleonly null-checks and thenstatic_casts. For GPU/DEFAULT memtype the factory can hand back aCpuAllocatorstub whoseAllocis plainmalloc; that pointer is non-null, so it passes the check and is then treated as aPluginDmlAllocationInfo— every subsequent field read faults.dml_factory.ccalready documents this exact scenario in a comment ("blind-cast to aPluginDmlAllocationInfo*and AddRef'd through a garbage vtable -> 0xC0000005"); this change makes it unreachable. I confirmed at runtime that the fallthrough which produces the stub really is taken.Crash dumps show two fault sites, one bug:
directx-backend.dll+0x587afmov rax,[rcx+8]— vtable slot forAddRef0xC0000005directx-backend.dll+0x4f5aamov rax,[rax+0x18]— member read on theGetResourcepath0xC0000005A third crash is
0xC0000409FAST_FAIL_GUARD_ICALL_CHECK_FAILURE— the same defect when the garbage vtable slot happens to be a non-null invalid target, so Control Flow Guard blocks the indirect call.Not fixed in any shipped build: the newest pipeline build (
msftpipeline/20260828-10D-RC9, 1.8.99.26) still has the same null-check-and-cast stub, andmainis unpatched.Changes
dml_bucketized_buffer_allocator.{h,cc}— the allocators track the handles they hand out;DecodeDataHandlerejects anything else withE_INVALIDARG. Validation is a registry lookup rather than an owner check, because readingGetOwner()would itself dereference the untrusted pointer — the exact fault at+0x4f5aa. The registry is process-wide so a handle that legitimately crosses allocator instances (sessions sharing tensors) still decodes.dml_factory.cc—CreateAllocatorImplno longer substitutes a CPU allocator for GPU memory once an EP exists; it returnsORT_EP_FAILso session creation fails recoverably instead of handing outmalloc'd pointers. Deliberately scoped tom_ep_raw != nullptr: ORT also requests a GPU allocator during factory registration, before any EP exists, and failing there breaks registration outright (observed while testing).Testing
Built the DirectX backend locally (
--use_dml, ORT 1.27) and ran a Conv/Relu/Add/ReduceMean model through the rebuilt EP: it registers, is selected, and infers correctly — so the validation rejects no legitimate handle. Also verified in the binary thatDecodeDataHandlenow performs the lookup rather than the bare cast.End-to-end validation on MDS1 still needs a signed build — that is the remaining step.
Reviewer notes
DecodeDataHandlenow takes a lock and a hash lookup per call. If this shows up on GPU-path profiles, move the registry to ashared_mutexor a lock-free set.E_INVALIDARG) rather than adding a new one, so the exception-safety contract is unchanged. Most call sites are insideORT_TRY; a few helpers (GetABIDataInterface,TryGetPooledAllocationId,GetAllocationFromDataPointer) rely on their callers — worth a second pair of eyes.AMDGPU_DirectXorDirectMLexecution device; only theAMDGPUumbrella path crashes.