[WebGPU] Direct CPU->GPU buffer upload for UMA

[jchen10]

This eliminates the using of staging buffers when uploading data from CPU to GPU for UMA. The data is directly copied to the GPU buffer. The major challenge of enabling this feature is to avoid waiting for the mapping of the dest GPU buffer. The idea is to map the dest GPU buffers in advance, so that when the uploading starts, the dest GPU is always mapped. This is achieved by:

1. using map-at-creation for all storage buffers.
2. mapping the buffer once it's returned, and marking it as ready for reuse after the mapping is done.
3. checking to unmap the buffer whereever it's used by the command encoder.

[jchen10]

@xhcao PTAL

[fs-eire]

@jchen10 thank you for the optimization!

I checked the document of the feature. As it described, this feature allows to map the data on the storage buffer and makes it possible to get rid of the staging buffer for uploading/downloading.

For downloading it looks like this feature is a simple optimization. However for uploading the situation became a little bit complicated.

As already mentioned in the description of this PR, to enable this feature for uploading, we need the buffer to be in "mapped" state. If we call buffer.mapAsync in buffer upload, it will cause a GPU-CPU synchronization, which breaks the pipeline and may cause perf problem. However, creating all storage buffer with mapped=true may also cause other problems:

• Some buffers will never map to CPU in their whole lifecycle. Creating them as mapped and unmap them before use is unnecessary and will cause more memory usage (whether this impact performance or not need to run test and figure out by data)
• For some usage (for example, in wasm32), upload buffers are not mapped inside C++. Instead, the code opens holes using EM_ASM to directly upload data from outside of the wasm. Without this trick, many models may run out of 4GB memory and fail the wasm runtime.

I think the real problem is that the current implementation of ORT does not give the context information to allocator so BufferManager has no idea what is the usage of a buffer that it is creating. With this limitation, we have to choose either creating all buffers mapped or all buffers unmapped. Ideally ORT should give some context info so that buffer manager would do better job.

There may need more discussion about how to use this feature in ORT.

[jchen10]

@fs-eire Thanks for the comment. Absolutely it would be perfect if ORT could give more context info to the buffer manager. Let's investigate this further.
BTW, this optimization looks promising. On my TGL, it shows 5.1G->3.2G memory reduction without performance regression.
Without UMA:

model_benchmark.exe -i ..\models-genai\Phi-3.5-mini-instruct-onnx-web
Batch size: 1, prompt tokens: 17, tokens to generate: 128
Prompt processing (time to first token):
        avg (us):       286253
        avg (tokens/s): 59.388
        p50 (us):       294972
        stddev (us):    12417.6
        n:              5 * 17 token(s)
Token generation:
        avg (us):       56041.5
        avg (tokens/s): 17.8439
        p50 (us):       56055.8
        stddev (us):    1682.57
        n:              635 * 1 token(s)
Token sampling:
        avg (us):       15.64
        avg (tokens/s): 63938.6
        p50 (us):       15.8
        stddev (us):    0.364692
        n:              5 * 1 token(s)
E2E generation (entire generation loop):
        avg (ms):       7403.63
        p50 (ms):       7393.68
        stddev (ms):    20.0764
        n:              5
Peak working set size (bytes): 5530017792

With UMA:

model_benchmark.exe -i ..\models-genai\Phi-3.5-mini-instruct-onnx-web
Batch size: 1, prompt tokens: 17, tokens to generate: 128
Prompt processing (time to first token):
        avg (us):       275511
        avg (tokens/s): 61.7036
        p50 (us):       273323
        stddev (us):    6202.55
        n:              5 * 17 token(s)
Token generation:
        avg (us):       56323.8
        avg (tokens/s): 17.7545
        p50 (us):       56297.5
        stddev (us):    1842.57
        n:              635 * 1 token(s)
Token sampling:
        avg (us):       12
        avg (tokens/s): 83333.3
        p50 (us):       14.9
        stddev (us):    4.01933
        n:              5 * 1 token(s)
E2E generation (entire generation loop):
        avg (ms):       7428.74
        p50 (ms):       7417.44
        stddev (ms):    24.1045
        n:              5
Peak working set size (bytes): 3419979776

[xhcao]

In fact, when usage is usage & wgpu::BufferUsage::Storage, we will use UMA buffer. So remove IsUMABuffer and uma_buffers_, instead to use WGPUBufferUsage wgpuBufferGetUsage(WGPUBuffer buffer) == wgpu::BufferUsage::Storage where need to verify whether the buffer is UMA. If so, the code will be simple and clean.

[jchen10]

Good point. It would be simpler. Thanks!

[guschmue]

/azp run ONNX Runtime Web CI Pipeline,Windows GPU CI Pipeline,Linux Android Emulator QNN CI Pipeline

[azure-pipelines]

Azure Pipelines will not run the associated pipelines, because the pull request was updated after the run command was issued. Review the pull request again and issue a new run command.

[guschmue]

/azp run Linux CPU CI Pipeline,Linux CPU Minimal Build E2E CI Pipeline,Linux GPU CI Pipeline,Linux GPU TensorRT CI Pipeline, Linux OpenVINO CI Pipeline,Linux QNN CI Pipeline,MacOS CI Pipeline,Windows ARM64 QNN CI Pipeline,Windows CPU CI Pipeline

[guschmue]

/azp run Windows GPU TensorRT CI Pipeline,onnxruntime-binary-size-checks-ci-pipeline,orttraining-linux-ci-pipeline,orttraining-linux-gpu-ci-pipeline,orttraining-ortmodule-distributed,Windows x64 QNN CI Pipeline,Big Models

[azure-pipelines]

Azure Pipelines will not run the associated pipelines, because the pull request was updated after the run command was issued. Review the pull request again and issue a new run command.

[guschmue]

/azp run Windows GPU CUDA CI Pipeline,Windows GPU DML CI Pipeline,Windows GPU Doc Gen CI Pipeline, Win_TRT_Minimal_CUDA_Test_CI

[azure-pipelines]

Azure Pipelines will not run the associated pipelines, because the pull request was updated after the run command was issued. Review the pull request again and issue a new run command.

[azure-pipelines]

Azure Pipelines will not run the associated pipelines, because the pull request was updated after the run command was issued. Review the pull request again and issue a new run command.