|
| 1 | +"""Run shipped UAV hook bodies and the tracking switch with CPU stand-ins. |
| 2 | +
|
| 3 | +Checks recording/forwarding and ordered thread isolation, not GPU restoration. |
| 4 | +""" |
| 5 | +import argparse |
| 6 | +from pathlib import Path |
| 7 | +import re |
| 8 | +import subprocess |
| 9 | +import tempfile |
| 10 | +from test_shutdown_routing import function |
| 11 | + |
| 12 | +ROOT = Path(__file__).resolve().parents[1] |
| 13 | + |
| 14 | + |
| 15 | +def main(): |
| 16 | + parser = argparse.ArgumentParser() |
| 17 | + parser.add_argument('--compiler', required=True) |
| 18 | + parser.add_argument('--driver') |
| 19 | + args = parser.parse_args() |
| 20 | + source = (ROOT/'OptiScaler/hooks/D3D12_Hooks.cpp').read_text(encoding='utf-8') |
| 21 | + cpp = r''' |
| 22 | +#include <vector> |
| 23 | +#include <unordered_map> |
| 24 | +#include <mutex> |
| 25 | +#include <thread> |
| 26 | +#include <iostream> |
| 27 | +#include <cstdlib> |
| 28 | +using UINT=unsigned;using D3D12_GPU_VIRTUAL_ADDRESS=unsigned long long; |
| 29 | +struct ID3D12GraphicsCommandList {}; |
| 30 | +enum class RootEntryType {Invalid,UAV}; |
| 31 | +struct RootState {RootEntryType type=RootEntryType::Invalid;D3D12_GPU_VIRTUAL_ADDRESS bufferLocation=0;}; |
| 32 | +// Count actual record attempts, including duplicate late -> early writes. |
| 33 | +struct CountingLock {int acquisitions=0;void lock(){++acquisitions;}void unlock(){}}; |
| 34 | +CountingLock rootStatesMutex; |
| 35 | +std::unordered_map<ID3D12GraphicsCommandList*,std::vector<RootState>> rootStates; |
| 36 | +struct D3D12Hooks {static void SetRootSignatureTracking(bool);}; |
| 37 | +int calls=0;ID3D12GraphicsCommandList* lastCmd=nullptr;UINT lastIndex=0; |
| 38 | +D3D12_GPU_VIRTUAL_ADDRESS lastAddress=0; |
| 39 | +void original(ID3D12GraphicsCommandList* cmd,UINT i,D3D12_GPU_VIRTUAL_ADDRESS address){ |
| 40 | + ++calls;lastCmd=cmd;lastIndex=i;lastAddress=address; |
| 41 | +} |
| 42 | +using HookFn=decltype(&original); |
| 43 | +struct Hook {HookFn o_earlyHook=original,o_lateHook=original;}; |
| 44 | +Hook s_SetComputeRootUnorderedAccessView,s_SetGraphicsRootUnorderedAccessView; |
| 45 | +''' |
| 46 | + names = ['isUpscalerActive', 'lateInProgressSetComputeRootUnorderedAccessView', |
| 47 | + 'lateInProgressSetGraphicsRootUnorderedAccessView'] |
| 48 | + for name in names: |
| 49 | + match = re.search(r'^static (?:thread_local )?bool '+name+r' = false;', source, re.M) |
| 50 | + assert match, name |
| 51 | + cpp += match[0]+'\n' |
| 52 | + cpp += function(source, 'void D3D12Hooks::SetRootSignatureTracking(')+'\n' |
| 53 | + for pipeline in ['Compute', 'Graphics']: |
| 54 | + for suffix in ['', 'Late']: |
| 55 | + body = function(source, 'static void hkSet'+pipeline+'RootUnorderedAccessView'+suffix+'(') |
| 56 | + # Substitute only the lock type to count duplicate recording; retain the hook body. |
| 57 | + cpp += body.replace('std::shared_mutex', 'CountingLock')+'\n' |
| 58 | + cpp += r''' |
| 59 | +int main(){int checks=0;auto check=[&](bool ok){++checks;if(!ok){std::cerr<<"FAIL "<<checks;std::exit(2);}}; |
| 60 | + ID3D12GraphicsCommandList a,b; |
| 61 | + for(bool graphics:{false,true}){ |
| 62 | + auto early=graphics?hkSetGraphicsRootUnorderedAccessView:hkSetComputeRootUnorderedAccessView; |
| 63 | + auto late=graphics?hkSetGraphicsRootUnorderedAccessViewLate:hkSetComputeRootUnorderedAccessViewLate; |
| 64 | + auto& hook=graphics?s_SetGraphicsRootUnorderedAccessView:s_SetComputeRootUnorderedAccessView; |
| 65 | + auto& nesting=graphics?lateInProgressSetGraphicsRootUnorderedAccessView:lateInProgressSetComputeRootUnorderedAccessView; |
| 66 | + hook.o_lateHook=early; |
| 67 | + for(bool latePath:{false,true}){ |
| 68 | + auto call=latePath?late:early; |
| 69 | + rootStates.clear();rootStates[&a].resize(2);calls=0;rootStatesMutex.acquisitions=0; |
| 70 | + D3D12Hooks::SetRootSignatureTracking(true); |
| 71 | + call(&a,1,123);check(rootStates[&a][1].type==RootEntryType::UAV); |
| 72 | + check(rootStates[&a][1].bufferLocation==123);check(calls==1); |
| 73 | + check(lastCmd==&a&&lastIndex==1&&lastAddress==123);check(rootStatesMutex.acquisitions==1); |
| 74 | + call(&a,1,0);check(rootStates[&a][1].bufferLocation==0);check(calls==2); |
| 75 | + call(&a,9,77);check(rootStates[&a].size()==2&&rootStates[&a][1].bufferLocation==0); |
| 76 | + check(calls==3&&lastIndex==9); |
| 77 | + call(nullptr,0,88);check(calls==4&&lastCmd==nullptr);check(!rootStates.contains(nullptr)); |
| 78 | + D3D12Hooks::SetRootSignatureTracking(false);int writes=rootStatesMutex.acquisitions; |
| 79 | + call(&a,1,99);check(calls==5);check(rootStatesMutex.acquisitions==writes); |
| 80 | + check(rootStates[&a][1].bufferLocation==0);check(!nesting); |
| 81 | + // Join orders access to all stand-ins, so the test itself has no data race. |
| 82 | + rootStates[&b].resize(1);bool workerEnabled=false; |
| 83 | + std::thread worker([&]{workerEnabled=!isUpscalerActive;call(&b,0,321);});worker.join(); |
| 84 | + check(workerEnabled);check(rootStates[&b][0].bufferLocation==321);check(isUpscalerActive); |
| 85 | + check(calls==6);D3D12Hooks::SetRootSignatureTracking(true);check(!isUpscalerActive); |
| 86 | + } |
| 87 | + nesting=true;rootStates[&a][0]={};calls=0;int writes=rootStatesMutex.acquisitions; |
| 88 | + early(&a,0,999);check(calls==1);check(rootStatesMutex.acquisitions==writes); |
| 89 | + check(rootStates[&a][0].type==RootEntryType::Invalid);nesting=false; |
| 90 | + } |
| 91 | + std::cout<<checks<<" D3D12 tracking checks passed (CPU stand-ins)\n"; |
| 92 | +} |
| 93 | +''' |
| 94 | + mutations = { |
| 95 | + 'global-tracking-switch': ('static thread_local bool isUpscalerActive', 'static bool isUpscalerActive'), |
| 96 | + 'compute-uav-inversion': ('if (!lateInProgressSetComputeRootUnorderedAccessView &&', |
| 97 | + 'if (lateInProgressSetComputeRootUnorderedAccessView &&'), |
| 98 | + 'graphics-uav-inversion': ('if (!lateInProgressSetGraphicsRootUnorderedAccessView &&', |
| 99 | + 'if (lateInProgressSetGraphicsRootUnorderedAccessView &&'), |
| 100 | + } |
| 101 | + with tempfile.TemporaryDirectory(prefix='aurora-d3d12-tracking-') as directory: |
| 102 | + path = Path(directory) |
| 103 | + cases = {'fixed': cpp} |
| 104 | + for name, (old, new) in mutations.items(): |
| 105 | + assert cpp.count(old) == 1, name |
| 106 | + cases[name] = cpp.replace(old, new) |
| 107 | + for name, text in cases.items(): |
| 108 | + test, exe = path/'test.cpp', path/'test.exe' |
| 109 | + test.write_text(text, encoding='utf-8') |
| 110 | + command = [args.compiler]+([args.driver] if args.driver else []) |
| 111 | + if Path(args.compiler).stem.lower() == 'cl': |
| 112 | + command += ['/nologo', '/EHsc', '/std:c++20', str(test), '/Fe:'+str(exe)] |
| 113 | + else: |
| 114 | + command += ['-std=c++20', str(test), '-o', str(exe)] |
| 115 | + subprocess.run(command, cwd=path, check=True) |
| 116 | + result = subprocess.run([str(exe)], capture_output=True, text=True, timeout=25) |
| 117 | + if name == 'fixed': |
| 118 | + assert result.returncode == 0, result |
| 119 | + print(result.stdout.strip()) |
| 120 | + else: |
| 121 | + assert result.returncode == 2 and 'FAIL ' in result.stderr, (name, result) |
| 122 | + print('Mutation rejected:', name) |
| 123 | + |
| 124 | + |
| 125 | +if __name__ == '__main__': |
| 126 | + main() |
0 commit comments