From 605a354baa1f055dcc517672378be8db1d8f105a Mon Sep 17 00:00:00 2001 From: Ti-Tai Wang Date: Tue, 30 Jun 2026 13:40:16 -0700 Subject: [PATCH 01/17] Reject CUDA BERT EmbedLayerNorm/SkipLayerNorm shapes exceeding 32-bit output indexing (#29264) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ### Summary The CUDA `EmbedLayerNormalization` and `SkipLayerNormalization` kernels compute output write offsets (`row_index * hidden_size`) using 32-bit arithmetic. For very large output tensors the element count can exceed `INT32_MAX`, at which point the offset is no longer representable in 32 bits. Every output write index in these kernels is a pure function of the launch grid and `hidden_size` — there is no data-dependent write indexing — so the maximum index is exactly `output_element_count - 1`, which the host knows from the input shapes before launch. This PR adds a **host-side guard** in each op's `ComputeInternal` that computes the output element count in 64-bit arithmetic and returns a clear error when it exceeds the supported 32-bit indexing range. ### Design - **`EmbedLayerNormalization`** (`embed_layer_norm.cc`): `output_element_count = (int64)batch_size * sequence_length * hidden_size`, guarded with `ORT_RETURN_IF_NOT(... <= INT32_MAX, ...)`. - **`SkipLayerNormalization`** (`skip_layer_norm.cc`): `output_element_count = input->Shape().Size()` (output shares the input shape), same guard. - Kernels are **unchanged** — they keep the original int32 indexing, so there is no extra register/occupancy cost in the hot path. This is pure host-side validation. ### Behavior This **rejects** (rather than silently attempting) single-op LayerNorm outputs larger than 2³¹ elements — a regime no real BERT-family model produces (it would require a multi-GB single-op activation). For all supported shapes there is no behavior or numeric change. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- onnxruntime/contrib_ops/cuda/bert/embed_layer_norm.cc | 11 +++++++++++ onnxruntime/contrib_ops/cuda/bert/skip_layer_norm.cc | 11 +++++++++++ 2 files changed, 22 insertions(+) diff --git a/onnxruntime/contrib_ops/cuda/bert/embed_layer_norm.cc b/onnxruntime/contrib_ops/cuda/bert/embed_layer_norm.cc index 864e2d1623923..20fafb1139dbb 100644 --- a/onnxruntime/contrib_ops/cuda/bert/embed_layer_norm.cc +++ b/onnxruntime/contrib_ops/cuda/bert/embed_layer_norm.cc @@ -1,6 +1,8 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. +#include + #include "core/providers/cuda/cuda_common.h" #include "contrib_ops/cpu/bert/embed_layer_norm_helper.h" #include "embed_layer_norm.h" @@ -61,6 +63,15 @@ Status EmbedLayerNorm::ComputeInternal(OpKernelContext* context) const { int sequence_length = static_cast(input_dims[1]); size_t element_size = sizeof(T); + // Element offsets into the output are 32-bit on device; reject shapes whose element count would + // exceed the 32-bit indexable range. The maximum output write index is + // batch_size * sequence_length * hidden_size - 1, so this guard covers every device write site. + const int64_t output_element_count = + static_cast(batch_size) * sequence_length * hidden_size; + ORT_RETURN_IF_NOT(output_element_count <= static_cast(std::numeric_limits::max()), + "EmbedLayerNormalization: output element count (", output_element_count, + ") exceeds the supported 32-bit indexing range."); + const bool broadcast_position_ids = (nullptr != position_ids && position_ids->Shape()[0] == 1); return LaunchEmbedLayerNormKernel( diff --git a/onnxruntime/contrib_ops/cuda/bert/skip_layer_norm.cc b/onnxruntime/contrib_ops/cuda/bert/skip_layer_norm.cc index 8557e326e5b15..865202801b2ce 100644 --- a/onnxruntime/contrib_ops/cuda/bert/skip_layer_norm.cc +++ b/onnxruntime/contrib_ops/cuda/bert/skip_layer_norm.cc @@ -1,6 +1,8 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. +#include + #include "core/providers/cuda/cuda_common.h" #include "core/common/narrow.h" #include "skip_layer_norm.h" @@ -77,6 +79,15 @@ Status SkipLayerNorm::ComputeInternal(OpKernelContext* ctx) const return Status::OK(); } + // Element offsets into the output are 32-bit on device; reject shapes whose element count would + // exceed the 32-bit indexable range. The output shares the input shape, so input->Shape().Size() + // is the output element count. The maximum output write index is + // row_count * hidden_size - 1 == output element count - 1, so this guard covers every device write site. + const int64_t output_element_count = input->Shape().Size(); + ORT_RETURN_IF_NOT(output_element_count <= static_cast(std::numeric_limits::max()), + "SkipLayerNormalization: output element count (", output_element_count, + ") exceeds the supported 32-bit indexing range."); + typedef typename ToCudaType::MappedType CudaT; const int skip_size = onnxruntime::narrow(skip->Shape().Size()); From a83b776e85342292de2c2a41dcab123754c0d333 Mon Sep 17 00:00:00 2001 From: Ti-Tai Wang Date: Tue, 30 Jun 2026 13:52:09 -0700 Subject: [PATCH 02/17] Fix optional-output guard in DecoderAttention/MultiHeadAttention shape inference (#29268) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ### Description The `DecoderAttention` and `MultiHeadAttention` shape-inference functions guarded population of their optional `present_key` (output 1) and `present_value` (output 2) outputs with `getNumOutputs() > 1`, but then write output index 2. `present_key` and `present_value` are produced as a both-or-neither pair, so this requires all three outputs (`> 2`) to be present before populating them — matching the existing `BaseGroupQueryAttention` (`>= 3`) and `EmbedLayerNorm` guards. It also adds an output-index range check in `InferenceContextImpl::getOutputType` so an output index beyond the declared output count fails inference cleanly instead of indexing past the end of the outputs container, mirroring the existing `DataPropagationContextImpl::getOutputType` and `getInputType` behavior. ### Motivation and Context A model that declares fewer outputs than the optional present outputs could previously drive shape inference to access an output index that was not declared. This makes the guard consistent with the other attention-family contrib ops. ### Changes - `onnxruntime/core/graph/contrib_ops/bert_defs.cc` — require all present outputs before populating `present_key`/`present_value` in `DecoderAttention` and `MultiHeadAttention`. - `onnxruntime/core/graph/graph.cc` — add an output-index range check in `InferenceContextImpl::getOutputType`. - `onnxruntime/test/contrib_ops/attention_optional_outputs_shape_inference_test.cc` — regression tests covering omitted optional present outputs, the 3-output positive cases, and the MHA/DMMHA two-output cases. - Adds a contrib-op shape-inference output-index safety skill doc plus a one-line coding-convention note. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --------- Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../SKILL.md | 237 ++++++++++++++++++ docs/Coding_Conventions_and_Standards.md | 1 + .../core/graph/contrib_ops/bert_defs.cc | 6 +- onnxruntime/core/graph/graph.cc | 4 + ...n_optional_outputs_shape_inference_test.cc | 230 +++++++++++++++++ 5 files changed, 475 insertions(+), 3 deletions(-) create mode 100644 .agents/skills/contrib-op-shape-inference-memory-safety/SKILL.md create mode 100644 onnxruntime/test/contrib_ops/attention_optional_outputs_shape_inference_test.cc diff --git a/.agents/skills/contrib-op-shape-inference-memory-safety/SKILL.md b/.agents/skills/contrib-op-shape-inference-memory-safety/SKILL.md new file mode 100644 index 0000000000000..e709455b60696 --- /dev/null +++ b/.agents/skills/contrib-op-shape-inference-memory-safety/SKILL.md @@ -0,0 +1,237 @@ +--- +name: contrib-op-shape-inference-memory-safety +description: "Audit and fix out-of-range output writes in ONNX Runtime operator shape-inference functions. Use when reviewing or fixing a contrib (or standard) op TypeAndShapeInference where a getNumOutputs() guard precedes a write to a higher output index - optional trailing outputs make a smaller output count schema-valid, so getOutputType(index) can run one past the declared outputs at Graph::Resolve." +--- + +# Contrib-Op Shape-Inference Output-Index Safety + +Reusable method for finding and fixing the bug class where an operator's +`TypeAndShapeInference` function guards an output write with `getNumOutputs() > N` but then +writes an output index **greater than** `N`. For a node that declares fewer outputs, the +written index is past the end of the inference context's output vector. + +> **Scope**: schema-level shape inference in `onnxruntime/core/graph/contrib_ops/*.cc` and +> `shape_inference_functions.cc`. This runs once during `Graph::Resolve` (model-load time), +> **EP-agnostic** - there is no per-EP (CPU/CUDA/ROCm) kernel duplicate of this code to +> chase. Op *kernels* allocate outputs via the bounds-safe `OpKernelContext::Output(index)` +> and are a separate concern. + +## 1. The pattern + +```cpp +// onnxruntime/core/graph/contrib_ops/bert_defs.cc (before) +propagateElemTypeFromInputToOutput(ctx, 0, 0); +if (ctx.getNumOutputs() > 1) { // guard says "> 1" + propagateElemTypeFromInputToOutput(ctx, 0, 1); + propagateElemTypeFromInputToOutput(ctx, 0, 2); // but writes index 2 +} +``` + +The guard `getNumOutputs() > 1` admits a node with **exactly 2 outputs** (indices 0, 1), yet +the body writes index **2**. The implication "`> 1` ⇒ index 2 exists" is false: `> 1` only +guarantees indices 0 and 1. + +### Why a smaller output count is valid + +Trailing outputs declared `OpSchema::Optional` **lower `min_output`**. ONNX derives +`min_output` = number of required outputs, `max_output` = total declared. The model checker +(`checker::check_node`) only enforces `min_output <= N <= max_output`. + +| Op | Output decls | min / max | A 2-output node? | +|---|---|---|---| +| `DecoderAttention` | out (req), new_key_cache (Opt), new_value_cache (Opt) | 1 / 3 | passes checker | +| `MultiHeadAttention` | out (req), present_key (Opt), present_value (Opt), qk (Opt) | 1 / 4 | passes checker | +| `DecoderMaskedMultiHeadAttention` | out (req) + 3 Optional | 1 / 4 | passes checker | + +So a node with `output=['out','present_key']` is schema-valid, passes the checker, and then +reaches the index-2 write. **A passing checker is not a guarantee the index is in range.** + +## 2. The sink (why the write is not caught) + +```cpp +// onnxruntime/core/graph/graph.cc - InferenceContextImpl +const TypeProto* getInputType(size_t index) const override { + return node_.InputDefs().at(index)->TypeAsProto(); // .at() -> bounds-checked +} +TypeProto* getOutputType(size_t index) override { + return &node_output_types_[index]; // operator[] -> NOT bounds-checked +} +``` + +- `node_output_types_` is sized to `node.OutputDefs().size()` in the `InferenceContextImpl` + ctor, so for a 2-output node it has 2 elements; `getOutputType(2)` returns one past the end. +- `getInputType` uses `.at()` (would throw on a bad index); `getOutputType` uses raw + `operator[]` (no check) - the asymmetry is the root cause. +- The call runs at `Graph::Resolve` → `InferAndVerifyTypeMatch` → `RunInferencing`. The + surrounding `ORT_TRY/ORT_CATCH(const std::exception&)` only catches *thrown* + `fail_shape_inference`; a raw out-of-range `operator[]` does not throw, so the catch does + not help. +- Because this is schema-level inference, it is **EP-independent** - no CUDA/ROCm copy. + +## 3. Audit technique — always sweep siblings + +Do not stop at the reported function. Grep **every** shape-inference guard and compare its +threshold against the **highest output index written before the next guard**. + +```bash +git grep -n 'getNumOutputs' -- \ + onnxruntime/core/graph/contrib_ops/*.cc \ + onnxruntime/core/graph/contrib_ops/shape_inference_functions.cc +``` + +For each `if (ctx.getNumOutputs() > N)` block, find the largest `index` passed to +`propagateElemTypeFromInputToOutput(ctx, _, index)` / `updateOutputShape(ctx, index, _)` / +`getOutputType(index)` inside it. **Rule: the guard must require strictly more outputs than +the highest index written** (write index `k` ⇒ guard must ensure `getNumOutputs() > k`). + +Correct exemplars already in the tree to copy: + +| Exemplar | Pattern | Why it is safe | +|---|---|---| +| `BaseGroupQueryAttention...` | `if (getNumOutputs() >= 3)` then writes idx 2 | guard covers highest index | +| `PagedAttention...` | nested `> 1` + inner `if (getNumOutputs() != 3) fail_shape_inference` | fails before any write | +| `EmbedLayerNormalizationShapeInference` | `> 2` then writes idx 2 | fixed by PR #28176 (precedent) | +| `SkipLayerNormalizationShapeInference` | each idx `k` guarded by `> k` | per-index guard | + +> **Gotcha — conditional writes can hide a vacuous audit.** A write may sit behind an inner +> condition (e.g. `hasInputShape(past_key_index)` before writing index 2). The site is still +> a bug, but you can only *observe* it when that inner condition is also satisfied. Keep this +> in mind both for the audit and for tests (§5). + +## 4. Fix patterns + +**Point fix (required): raise the guard to cover the highest index written.** + +```cpp +// before +if (ctx.getNumOutputs() > 1) { ... writes idx 2 ... } +// after +if (ctx.getNumOutputs() > 2) { // both present_key (idx 1) AND present_value (idx 2) + ... +} +``` + +Justify the threshold with the op's output semantics. For these attention ops the two trailing +outputs - `present_key` (idx 1) and `present_value` (idx 2) for `MultiHeadAttention`, +`new_key_cache` / `new_value_cache` for `DecoderAttention` (see the §1 table for each op's +exact output names) - are a **both-or-neither pair**: there is no valid configuration that +emits one without the other, so requiring all three outputs before populating indices 1 and 2 +is behavior-preserving. (`PagedAttention` encodes the same invariant via its nested `!= 3` +check.) + +**Defense-in-depth (recommended): bound the sink** so a future author cannot reintroduce the +class. + +```cpp +// onnxruntime/core/graph/graph.cc - InferenceContextImpl::getOutputType +TypeProto* getOutputType(size_t index) override { + if (index >= node_output_types_.size()) { + fail_type_inference("output index ", index, " is out of range; node has ", + node_output_types_.size(), " outputs"); + } + return &node_output_types_[index]; +} +``` + +This mirrors `getInputType`'s `.at()` and the existing bounds checks in the sibling +`DataPropagationContextImpl`. Placing it at the base layer transitively protects the NHWC and +quantization wrapper contexts. After the point fix this branch is unreachable through a normal +model (the guard already prevents the out-of-range index), so it is pure defense-in-depth. Its +failure mode is build-dependent: with exceptions enabled, `fail_type_inference` raises +`InferenceError` (a `std::exception`), caught by the existing `ORT_CATCH(const std::exception&)` +around `RunInferencing` and surfaced as a clean load-time error; under `ORT_NO_EXCEPTIONS` it is +**not** compiled out - ONNX's no-exceptions path prints the message to `std::cerr` and calls +`abort()`, a deterministic fail-fast (consistent with `getInputType`'s `.at()`, which likewise +terminates under no-exceptions). Either way the result is a controlled failure rather than an +out-of-range write. + +## 5. Test recipe + +Tests live in `onnxruntime/test/contrib_ops/*.cc` and are **auto-globbed** into the +`onnxruntime_provider_test` target by `cmake/onnxruntime_unittests.cmake` +(`test/contrib_ops/*.cc` pattern) - **no cmake edit needed** for a new file. See the +`ort-test` skill for the executable taxonomy (`onnxruntime_provider_test` vs +`onnxruntime_test_all`). + +Rules that make the regression test actually guard the fix: + +1. **Drive through `Model` + `Graph::Resolve`**, not ONNX's standalone `TestShapeInference`. + Only the full resolve path constructs the real `InferenceContextImpl` and hits the + `getOutputType` sink described in §2. A standalone ONNX shape-inference helper uses a + different context and **bypasses** the sink, so it cannot reproduce the bug. +2. **Negative tests must be NON-VACUOUS** - they must actually enter the write branch on + pre-fix source. If a write is gated by an inner condition (§3 gotcha), satisfy it: e.g. for + `MultiHeadAttention`/`DecoderMaskedMultiHeadAttention`, supply a **shaped `past_key`** + (and `past_sequence_length` / `past_present_share_buffer` as the op requires) so the + index-2 block runs. A negative test that only supplies `query` skips the block and passes + even on pre-fix source - regression-proof in name only. +3. **Add positive (all-outputs) cases**: a node with every output present must still infer the + trailing output types - proves the tightened guard did not over-restrict. +4. **Keep tests throw-free post-fix** so they are valid under `ORT_NO_EXCEPTIONS`. Any case + that is *expected* to `fail_shape_inference` (throws) must be excluded with + `#ifndef ORT_NO_EXCEPTIONS`. The "2 outputs must not go out of range" case is throw-free + after the point fix and is safe in all builds. + +**Verify the negative test is non-vacuous (sanitizer A/B)** - the most reliable way to prove a +negative test enters the previously-out-of-range branch: build the test at the **pre-fix** +commit with **AddressSanitizer** and confirm it flags the out-of-range output access; then +confirm it is clean after the fix. + +```bash +# Functional run (any Debug build): +cmake --build build/Linux/Debug --target onnxruntime_provider_test -j"$(nproc)" +./build/Linux/Debug/onnxruntime_provider_test \ + --gtest_filter='AttentionOptionalOutputsShapeInferenceTest.*' + +# A/B proof (isolated worktree at the pre-fix commit, CPU-only Debug + sanitizer): +git worktree add --detach ../ort-prefix-check ~1 +# copy the new test file in, then: +python3 tools/ci_build/build.py --build_dir build/asan --config Debug --parallel \ + --skip_tests --enable_address_sanitizer --skip_submodule_sync \ + --cmake_generator Ninja --target onnxruntime_provider_test +# Pre-fix: the negative tests fail (the sanitizer flags the out-of-range output access). +# Post-fix (cherry-pick the guard fix): all tests pass, no sanitizer report. +``` + +## 6. Process / wording conventions + +- Run **`lintrunner -a`** before pushing so the `CLANGFORMAT` / Python-format gate passes. See + the `ort-lint` skill. +- Use **correctness/robustness framing** in code, comments, commit messages, and the PR body + - describe the change as fixing an optional-output guard, not as a security fix. This + matches repo convention (compare `python-kwargs-setattr-security`) and keeps the PR neutral. + +## 7. Audit checklist (per-operator review) + +When reviewing or hardening any operator implementation or its shape inference: + +- [ ] Read the op's spec - ONNX standard op page, or for a contrib op its `OpSchema` + registration (`.Input/.Output/.Attr`, and `Optional`/`Variadic` markers). A local ONNX + checkout has the standard-op spec pages; contrib ops are defined only in ORT. +- [ ] Enumerate **all** inputs, attributes, and outputs, noting which are optional and the + resulting `min/max` input and output counts. +- [ ] Validate every input/attribute before indexing into it, to avoid out-of-range reads + (which can cascade into worse failures). Match each output-index write to a guard that + guarantees the index is in range (§3 rule). +- [ ] Prefer `ORT_RETURN_IF` / `ORT_RETURN_IF_NOT` for validation; use `ORT_ENFORCE` in + constructors. In shape inference use `fail_shape_inference` / `fail_type_inference`. +- [ ] Use `SafeInt<>` / `narrow<>()` for index and size arithmetic and casts to avoid overflow + or truncation that yields a wrong index. See `core/common/safeint.h` and + `docs/Coding_Conventions_and_Standards.md`. +- [ ] Ensure tests build and pass under **no-exceptions** builds; `#ifndef ORT_NO_EXCEPTIONS` + around any case expected to throw. +- [ ] Exclude EPs known not to support the op, with a comment explaining why. +- [ ] Check whether **other EPs (notably CUDA/ROCm)** implement the same op and whether the + same issue exists there. (For *shape inference* specifically, the logic is EP-agnostic + and single-source - confirm there is no kernel-side analogue.) + +## References + +- **PR #28176** - "Fix ... in EmbedLayerNormalizationShapeInference": the precedent that fixed + the identical `> 1` → `> 2` primitive in one site; the sibling attention sites were missed, + motivating the sweep in §3. +- **PR #29268** - this fix: guards corrected in `DecoderAttention` / `MultiHeadAttention` / + `DecoderMaskedMultiHeadAttention` shape inference, plus the `getOutputType` bounds check and + non-vacuous regression tests. +- Sibling skill: **`ort-test`** (test executables, `--gtest_filter`, contrib-op test layout); + **`ort-lint`** (`lintrunner -a`); **`ort-build`** (build flags, ASan). diff --git a/docs/Coding_Conventions_and_Standards.md b/docs/Coding_Conventions_and_Standards.md index 02af7ddaa49be..6dcf7e6dd43bd 100644 --- a/docs/Coding_Conventions_and_Standards.md +++ b/docs/Coding_Conventions_and_Standards.md @@ -109,6 +109,7 @@ void foo(gsl::span names) { * Use [SafeInt](https://github.com/dcleblanc/SafeInt) when calculating the size of memory to allocate to protect against overflow errors * `#include "core/common/safeint.h"` * search for `SafeInt` in the code for examples +* In operator shape inference, validate every output index against `getNumOutputs()` before writing it. Optional trailing outputs lower an op's `min_output`, so a node may legally declare fewer outputs than the schema's maximum; guard each optional output by the exact index it populates so a `getNumOutputs() > N` guard never writes an index greater than `N`. * The following C++ warnings should never be disabled in onnxruntime VC++ projects(Required by [Binskim](https://github.com/microsoft/binskim/blob/d9afb65c89a621411efded74c27999281d87867e/src/BinSkim.Rules/PERules/BA2007.EnableCriticalCompilerWarnings.cs)). 1. [4018](https://docs.microsoft.com/en-us/cpp/error-messages/compiler-warnings/compiler-warning-level-3-c4018) 'token' : signed/unsigned mismatch 2. [4146](https://docs.microsoft.com/en-us/cpp/error-messages/compiler-warnings/compiler-warning-level-2-c4146?view=msvc-160) unary minus operator applied to unsigned type, result still unsigned diff --git a/onnxruntime/core/graph/contrib_ops/bert_defs.cc b/onnxruntime/core/graph/contrib_ops/bert_defs.cc index 896774fb5c8d8..9fdcebf57081d 100644 --- a/onnxruntime/core/graph/contrib_ops/bert_defs.cc +++ b/onnxruntime/core/graph/contrib_ops/bert_defs.cc @@ -29,7 +29,7 @@ namespace contrib { void DecoderAttentionTypeAndShapeInference(ONNX_NAMESPACE::InferenceContext& ctx) { // Type inference ONNX_NAMESPACE::propagateElemTypeFromInputToOutput(ctx, 0, 0); - if (ctx.getNumOutputs() > 1) { + if (ctx.getNumOutputs() > 2) { // has new_key_cache and new_value_cache outputs; a pair, so present only when > 2 ONNX_NAMESPACE::propagateElemTypeFromInputToOutput(ctx, 0, 1); ONNX_NAMESPACE::propagateElemTypeFromInputToOutput(ctx, 0, 2); } @@ -38,7 +38,7 @@ void DecoderAttentionTypeAndShapeInference(ONNX_NAMESPACE::InferenceContext& ctx auto& query_shape = getInputShape(ctx, 0); updateOutputShape(ctx, 0, query_shape); } - if (ctx.getNumOutputs() > 1) { + if (ctx.getNumOutputs() > 2) { // has new_key_cache and new_value_cache outputs; a pair, so present only when > 2 if (hasInputShape(ctx, 6) && hasInputShape(ctx, 7)) { auto& cache_shape = getInputShape(ctx, 6); auto& cache_dims = cache_shape.dim(); @@ -199,7 +199,7 @@ void MultiHeadAttentionTypeAndShapeInference(ONNX_NAMESPACE::InferenceContext& c } } - if (ctx.getNumOutputs() > 1) { // has present output + if (ctx.getNumOutputs() > 2) { // has present_key and present_value outputs; a pair, so present only when > 2 if (hasInputShape(ctx, past_key_index)) { auto& past_shape = getInputShape(ctx, past_key_index); auto& past_dims = past_shape.dim(); diff --git a/onnxruntime/core/graph/graph.cc b/onnxruntime/core/graph/graph.cc index fe2df6a87d124..ff60f6cb5f741 100644 --- a/onnxruntime/core/graph/graph.cc +++ b/onnxruntime/core/graph/graph.cc @@ -2739,6 +2739,10 @@ class InferenceContextImpl : public ONNX_NAMESPACE::InferenceContext { } TypeProto* getOutputType(size_t index) override { + if (index >= node_output_types_.size()) { + fail_type_inference("output index ", index, " is out of range; node has ", + node_output_types_.size(), " outputs"); + } return &node_output_types_[index]; } diff --git a/onnxruntime/test/contrib_ops/attention_optional_outputs_shape_inference_test.cc b/onnxruntime/test/contrib_ops/attention_optional_outputs_shape_inference_test.cc new file mode 100644 index 0000000000000..d2d3a421a6560 --- /dev/null +++ b/onnxruntime/test/contrib_ops/attention_optional_outputs_shape_inference_test.cc @@ -0,0 +1,230 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +// Shape-inference correctness coverage for attention contrib ops with optional "present" outputs. +// DecoderAttention, MultiHeadAttention and DecoderMaskedMultiHeadAttention each expose an optional +// present_key (output 1) and present_value (output 2). These outputs are produced as a pair, so a +// node may declare either one output, or all three; declaring exactly two (present_key kept, +// present_value omitted) is also valid per the schemas. +// +// The "...Omitted" tests build each op with exactly two outputs and assert that Graph::Resolve() +// completes cleanly without referencing the absent third output. The "...AllPresentOutputs" tests +// build each op with all three outputs and assert that the present_key / present_value branch still +// runs and propagates their element types. Together they pin the guard to exactly "> 2": fewer +// outputs must not touch the missing one, and three outputs must still be inferred. +// +// These tests exercise only graph-load shape inference, which is execution-provider independent, so +// they run on the default CPU build with no provider-specific handling. The resolve path for these +// models is throw-free, so the tests are valid in builds compiled without exceptions +// (ORT_NO_EXCEPTIONS) and need no exception-specific guarding. + +#include "gtest/gtest.h" + +#include "core/graph/constants.h" +#include "core/graph/model.h" +#include "test/test_environment.h" +#include "test/unittest_util/graph_transform_test_builder.h" +#include "test/util/include/asserts.h" + +namespace onnxruntime { +namespace test { + +namespace { + +constexpr int kOnnxOpsetVersion = 17; + +// Builds a single-node model via add_node, resolves it (running the node's type/shape inference), +// asserts success, and runs the optional verifier against the resolved graph. +void BuildResolveAndVerify(const std::function& add_node, + const std::function& verify = nullptr) { + std::unordered_map domain_to_version; + domain_to_version[kOnnxDomain] = kOnnxOpsetVersion; + domain_to_version[kMSDomain] = 1; + + Model model("attention_optional_outputs", /*is_onnx_domain_only=*/false, ModelMetaData(), + PathString(), IOnnxRuntimeOpSchemaRegistryList(), domain_to_version, {}, + DefaultLoggingManager().DefaultLogger()); + + ModelTestBuilder builder(model.MainGraph()); + add_node(builder); + builder.SetGraphOutputs(); + + ASSERT_STATUS_OK(model.MainGraph().Resolve()); + if (verify) { + verify(model.MainGraph()); + } +} + +// Asserts that the given output NodeArg received a tensor element type from shape inference. +void ExpectInferredElemType(const NodeArg* output) { + ASSERT_NE(output, nullptr); + const ONNX_NAMESPACE::TypeProto* type = output->TypeAsProto(); + ASSERT_NE(type, nullptr); + EXPECT_TRUE(type->has_tensor_type()); + EXPECT_TRUE(type->tensor_type().has_elem_type()); +} + +} // namespace + +// MultiHeadAttention with present_key kept and present_value omitted (exactly two outputs). +// past_key (input 6), past_value (input 7) and past_sequence_length (input 8) are supplied with +// shapes so the present-output branch is active; with only two outputs declared, inference must not +// reference the absent present_value (output 2). Inputs 1-5 are optional and left empty. +TEST(AttentionOptionalOutputsShapeInferenceTest, MultiHeadAttentionPresentValueOmitted) { + BuildResolveAndVerify([](ModelTestBuilder& builder) { + NodeArg& empty = builder.graph_.GetOrCreateNodeArg("", nullptr); + NodeArg* query = builder.MakeInput(std::vector{2, 1, 4}); + NodeArg* past_key = builder.MakeInput(std::vector{2, 2, 3, 2}); + NodeArg* past_value = builder.MakeInput(std::vector{2, 2, 3, 2}); + NodeArg* past_sequence_length = builder.MakeInput(std::vector{1}); + NodeArg* output = builder.MakeOutput(std::nullopt); + NodeArg* present_key = builder.MakeOutput(std::nullopt); + std::vector inputs = {query, &empty, &empty, &empty, &empty, &empty, + past_key, past_value, past_sequence_length}; + Node& node = builder.AddNode("MultiHeadAttention", inputs, {output, present_key}, kMSDomain); + node.AddAttribute("num_heads", static_cast(2)); + }); +} + +// DecoderMaskedMultiHeadAttention with present_key kept and present_value omitted. +// past_key (input 5) and past_value (input 6) are supplied with shapes and past buffer sharing is +// enabled so the present-output branch is active; with only two outputs declared, inference must not +// reference the absent present_value (output 2). Inputs 1-4 are optional and left empty. +TEST(AttentionOptionalOutputsShapeInferenceTest, DecoderMaskedMultiHeadAttentionPresentValueOmitted) { + BuildResolveAndVerify([](ModelTestBuilder& builder) { + NodeArg& empty = builder.graph_.GetOrCreateNodeArg("", nullptr); + NodeArg* query = builder.MakeInput(std::vector{2, 1, 4}); + NodeArg* past_key = builder.MakeInput(std::vector{2, 2, 3, 2}); + NodeArg* past_value = builder.MakeInput(std::vector{2, 2, 3, 2}); + NodeArg* output = builder.MakeOutput(std::nullopt); + NodeArg* present_key = builder.MakeOutput(std::nullopt); + std::vector inputs = {query, &empty, &empty, &empty, &empty, past_key, past_value}; + Node& node = builder.AddNode("DecoderMaskedMultiHeadAttention", inputs, {output, present_key}, + kMSDomain); + node.AddAttribute("num_heads", static_cast(2)); + node.AddAttribute("past_present_share_buffer", static_cast(1)); + }); +} + +// DecoderAttention with new_key_cache kept and new_value_cache omitted (exactly two outputs). +TEST(AttentionOptionalOutputsShapeInferenceTest, DecoderAttentionNewValueCacheOmitted) { + BuildResolveAndVerify([](ModelTestBuilder& builder) { + // DecoderAttention requires inputs 0-4 and 8-11; inputs 5-7 are optional and left empty here. + NodeArg& empty = builder.graph_.GetOrCreateNodeArg("", nullptr); + NodeArg* query = builder.MakeInput(std::vector{2, 1, 4}); + NodeArg* key = builder.MakeInput(std::vector{2, 1, 4}); + NodeArg* q_weight = builder.MakeInput(std::vector{4, 4}); + NodeArg* kv_weight = builder.MakeInput(std::vector{4, 8}); + NodeArg* bias = builder.MakeInput(std::vector{12}); + NodeArg* static_kv = builder.MakeInput(std::vector{1}); + NodeArg* use_past = builder.MakeInput(std::vector{1}); + NodeArg* has_layer_state = builder.MakeInput(std::vector{1}); + NodeArg* has_key_padding_mask = builder.MakeInput(std::vector{1}); + + NodeArg* output = builder.MakeOutput(std::nullopt); + NodeArg* new_key_cache = builder.MakeOutput(std::nullopt); + + std::vector inputs = {query, key, q_weight, kv_weight, bias, &empty, &empty, &empty, + static_kv, use_past, has_layer_state, has_key_padding_mask}; + Node& node = builder.AddNode("DecoderAttention", inputs, {output, new_key_cache}, kMSDomain); + node.AddAttribute("num_heads", static_cast(2)); + }); +} + +// MultiHeadAttention with all three outputs: the present_key / present_value branch must still run. +TEST(AttentionOptionalOutputsShapeInferenceTest, MultiHeadAttentionAllPresentOutputs) { + NodeArg* present_key = nullptr; + NodeArg* present_value = nullptr; + BuildResolveAndVerify( + [&](ModelTestBuilder& builder) { + // present_key/present_value types are propagated from past_key (input 6) and past_value + // (input 7) when past buffer sharing is active, which the op detects from shaped past_key + // (input 6) and past_sequence_length (input 8). Leave inputs 1-5 empty to reach them. + NodeArg& empty = builder.graph_.GetOrCreateNodeArg("", nullptr); + NodeArg* query = builder.MakeInput(std::vector{2, 1, 4}); + NodeArg* past_key = builder.MakeInput(std::vector{2, 2, 3, 2}); + NodeArg* past_value = builder.MakeInput(std::vector{2, 2, 3, 2}); + NodeArg* past_sequence_length = builder.MakeInput(std::vector{1}); + + NodeArg* output = builder.MakeOutput(std::nullopt); + present_key = builder.MakeOutput(); + present_value = builder.MakeOutput(); + + std::vector inputs = {query, &empty, &empty, &empty, &empty, &empty, + past_key, past_value, past_sequence_length}; + Node& node = builder.AddNode("MultiHeadAttention", inputs, + {output, present_key, present_value}, kMSDomain); + node.AddAttribute("num_heads", static_cast(2)); + }, + [&](const Graph&) { + ExpectInferredElemType(present_key); + ExpectInferredElemType(present_value); + }); +} + +// DecoderMaskedMultiHeadAttention with all three outputs: shape inference must still populate them. +TEST(AttentionOptionalOutputsShapeInferenceTest, DecoderMaskedMultiHeadAttentionAllPresentOutputs) { + NodeArg* present_key = nullptr; + NodeArg* present_value = nullptr; + BuildResolveAndVerify( + [&](ModelTestBuilder& builder) { + // For this op past_key/past_value are inputs 5 and 6; leave inputs 1-4 empty to reach them. + NodeArg& empty = builder.graph_.GetOrCreateNodeArg("", nullptr); + NodeArg* query = builder.MakeInput(std::vector{2, 1, 4}); + NodeArg* past_key = builder.MakeInput(std::vector{2, 2, 3, 2}); + NodeArg* past_value = builder.MakeInput(std::vector{2, 2, 3, 2}); + + NodeArg* output = builder.MakeOutput(std::nullopt); + present_key = builder.MakeOutput(); + present_value = builder.MakeOutput(); + + std::vector inputs = {query, &empty, &empty, &empty, &empty, past_key, past_value}; + Node& node = builder.AddNode("DecoderMaskedMultiHeadAttention", inputs, + {output, present_key, present_value}, kMSDomain); + node.AddAttribute("num_heads", static_cast(2)); + node.AddAttribute("past_present_share_buffer", static_cast(1)); + }, + [&](const Graph&) { + ExpectInferredElemType(present_key); + ExpectInferredElemType(present_value); + }); +} + +// DecoderAttention with all three outputs: new_key_cache / new_value_cache types must be inferred. +TEST(AttentionOptionalOutputsShapeInferenceTest, DecoderAttentionAllCacheOutputs) { + NodeArg* new_key_cache = nullptr; + NodeArg* new_value_cache = nullptr; + BuildResolveAndVerify( + [&](ModelTestBuilder& builder) { + NodeArg& empty = builder.graph_.GetOrCreateNodeArg("", nullptr); + NodeArg* query = builder.MakeInput(std::vector{2, 1, 4}); + NodeArg* key = builder.MakeInput(std::vector{2, 1, 4}); + NodeArg* q_weight = builder.MakeInput(std::vector{4, 4}); + NodeArg* kv_weight = builder.MakeInput(std::vector{4, 8}); + NodeArg* bias = builder.MakeInput(std::vector{12}); + NodeArg* key_cache = builder.MakeInput(std::vector{2, 2, 3, 2}); + NodeArg* value_cache = builder.MakeInput(std::vector{2, 2, 3, 2}); + NodeArg* static_kv = builder.MakeInput(std::vector{1}); + NodeArg* use_past = builder.MakeInput(std::vector{1}); + NodeArg* has_layer_state = builder.MakeInput(std::vector{1}); + NodeArg* has_key_padding_mask = builder.MakeInput(std::vector{1}); + + NodeArg* output = builder.MakeOutput(std::nullopt); + new_key_cache = builder.MakeOutput(); + new_value_cache = builder.MakeOutput(); + + std::vector inputs = {query, key, q_weight, kv_weight, bias, &empty, key_cache, + value_cache, static_kv, use_past, has_layer_state, + has_key_padding_mask}; + Node& node = builder.AddNode("DecoderAttention", inputs, + {output, new_key_cache, new_value_cache}, kMSDomain); + node.AddAttribute("num_heads", static_cast(2)); + }, + [&](const Graph&) { + ExpectInferredElemType(new_key_cache); + ExpectInferredElemType(new_value_cache); + }); +} + +} // namespace test +} // namespace onnxruntime From 152d01f7e8eb227ddb9a118e2adfd6d857437cb1 Mon Sep 17 00:00:00 2001 From: Dmitri Smirnov Date: Tue, 30 Jun 2026 14:12:50 -0700 Subject: [PATCH 03/17] Handle non-trivially-copyable types in Loop/Scan output concatenation (#29397) This pull request improves the handling of string tensors in the CPU implementation of the Loop operator, ensuring correct memory management and copy semantics for non-trivially-copyable types like `std::string`. It also adds comprehensive unit tests to verify these behaviors, especially for cases involving string scan outputs and loop-carried variables. **Enhancements for string tensor support:** * Updated `ConcatenateCpuOutput` in `loop.cc` to properly detect string tensors and use `std::copy` for concatenation, ensuring correct handling of heap-allocated string payloads and avoiding unsafe byte-wise copying. [[1]](diffhunk://#diff-2c8478657254a53c4ce09684960c925593395336e00c43ef672d9427722e3ff7R276-L282) [[2]](diffhunk://#diff-2c8478657254a53c4ce09684960c925593395336e00c43ef672d9427722e3ff7R289-R304) * Modified `OutputIterator::ZeroOutCurrent` in `scan_utils.h` to skip zeroing for string tensors, as their elements are already default-constructed and cannot be safely set with `memset`. **New and extended unit tests for string tensor scenarios:** * Added tests to `loop_test.cc` covering: - String scan outputs, including long strings that exceed the small-string-optimization threshold and multi-element outputs. - Loop-carried string variables to ensure they are copied correctly. - Zero-iteration cases to confirm empty string scan outputs are handled without errors. --- .../core/providers/cpu/controlflow/loop.cc | 31 +- .../providers/cpu/controlflow/scan_utils.h | 9 +- .../providers/cpu/controlflow/loop_test.cc | 316 ++++++++++++++++++ 3 files changed, 343 insertions(+), 13 deletions(-) diff --git a/onnxruntime/core/providers/cpu/controlflow/loop.cc b/onnxruntime/core/providers/cpu/controlflow/loop.cc index 1baa9ddbfa0fa..41ab1bd0ff019 100644 --- a/onnxruntime/core/providers/cpu/controlflow/loop.cc +++ b/onnxruntime/core/providers/cpu/controlflow/loop.cc @@ -273,12 +273,15 @@ static Status ConcatenateCpuOutput(void* /*stream*/, void* output, size_t output_size_in_bytes) { const auto& first_output = per_iteration_output.front().Get(); const auto& per_iteration_shape = first_output.Shape(); - size_t bytes_per_iteration = first_output.SizeInBytes(); - - // we can't easily use a C++ template for the tensor element type, - // so use a span for some protection but work in bytes - gsl::span output_span = gsl::make_span(static_cast(output), - output_size_in_bytes); + const bool is_string = first_output.IsDataTypeString(); + const size_t bytes_per_iteration = first_output.SizeInBytes(); + const int64_t elements_per_iteration = first_output.Shape().Size(); + + // for the non-string path, create the output span once outside the loop + gsl::span output_span; + if (!is_string) { + output_span = gsl::make_span(static_cast(output), output_size_in_bytes); + } for (size_t i = 0, num_iterations = per_iteration_output.size(); i < num_iterations; ++i) { auto& ort_value = per_iteration_output[i]; @@ -290,10 +293,18 @@ static Status ConcatenateCpuOutput(void* /*stream*/, " Expected:", per_iteration_shape, " Got:", iteration_data.Shape()); } - auto src = gsl::make_span(static_cast(iteration_data.DataRaw()), - bytes_per_iteration); - auto dst = output_span.subspan(i * bytes_per_iteration, bytes_per_iteration); - gsl::copy(src, dst); + if (is_string) { + // std::string is not trivially copyable — move from the per-iteration tensors since they are + // discarded after concatenation + auto src = ort_value.GetMutable()->MutableDataAsSpan(); + auto* dst_begin = static_cast(output) + i * elements_per_iteration; + std::move(src.begin(), src.end(), dst_begin); + } else { + auto src = gsl::make_span(static_cast(iteration_data.DataRaw()), + bytes_per_iteration); + auto dst = output_span.subspan(i * bytes_per_iteration, bytes_per_iteration); + gsl::copy(src, dst); + } } return Status::OK(); diff --git a/onnxruntime/core/providers/cpu/controlflow/scan_utils.h b/onnxruntime/core/providers/cpu/controlflow/scan_utils.h index cf4cb2bd60084..38def699744c2 100644 --- a/onnxruntime/core/providers/cpu/controlflow/scan_utils.h +++ b/onnxruntime/core/providers/cpu/controlflow/scan_utils.h @@ -106,10 +106,13 @@ class OutputIterator { // set the output for the current iteration to zeros. used for short sequence lengths Status ZeroOutCurrent() { - auto status = Status::OK(); auto* tensor = (**this).GetMutable(); - status = zero_data_func_(tensor->MutableDataRaw(), tensor->SizeInBytes()); - return status; + if (tensor->IsDataTypeString()) { + // std::string is not trivially copyable — memset would corrupt the objects. + // The strings are already default-constructed (empty) from placement-new so nothing to do. + return Status::OK(); + } + return zero_data_func_(tensor->MutableDataRaw(), tensor->SizeInBytes()); } const OrtValue& GetOutput() const { diff --git a/onnxruntime/test/providers/cpu/controlflow/loop_test.cc b/onnxruntime/test/providers/cpu/controlflow/loop_test.cc index f968fc6fc2f2e..cbb7771d7f09e 100644 --- a/onnxruntime/test/providers/cpu/controlflow/loop_test.cc +++ b/onnxruntime/test/providers/cpu/controlflow/loop_test.cc @@ -1041,6 +1041,322 @@ TEST(Loop, IterationCountAsOutput) { test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kTensorrtExecutionProvider, kOpenVINOExecutionProvider}); } +// Verify that Loop correctly handles tensor(string) scan outputs. +// Strings are not trivially copyable so the concatenation path must use proper copy semantics. +// Uses strings exceeding the small-string-optimization threshold to exercise heap-allocated payloads. +TEST(Loop, StringScanOutput) { + auto create_subgraph = []() { + Model model("String scan output subgraph", false, DefaultLoggingManager().DefaultLogger()); + auto& graph = model.MainGraph(); + + std::vector inputs; + std::vector outputs; + + /* Subgraph produces a constant string tensor as a scan output each iteration. + + iter_num_in cond_in + (unused) | + [Identity] + | + cond_out + + [Constant] -> scan_output (string tensor, shape {1}) + */ + + // graph input types + TypeProto int64_scalar; + int64_scalar.mutable_tensor_type()->set_elem_type(TensorProto_DataType_INT64); + int64_scalar.mutable_tensor_type()->mutable_shape()->add_dim()->set_dim_value(1); + + TypeProto bool_scalar; + bool_scalar.mutable_tensor_type()->set_elem_type(TensorProto_DataType_BOOL); + bool_scalar.mutable_tensor_type()->mutable_shape()->add_dim()->set_dim_value(1); + + TypeProto string_tensor; + string_tensor.mutable_tensor_type()->set_elem_type(TensorProto_DataType_STRING); + string_tensor.mutable_tensor_type()->mutable_shape()->add_dim()->set_dim_value(1); + + // graph inputs + auto& iter_num_in = graph.GetOrCreateNodeArg("iter_num_in", &int64_scalar); + auto& cond_in = graph.GetOrCreateNodeArg("cond_in", &bool_scalar); + + // graph outputs + auto& cond_out = graph.GetOrCreateNodeArg("cond_out", &bool_scalar); + auto& scan_out = graph.GetOrCreateNodeArg("scan_out", &string_tensor); + + // cond_in -> cond_out + { + inputs = {&cond_in}; + outputs = {&cond_out}; + graph.AddNode("cond_identity", "Identity", "Forward cond", inputs, outputs); + } + + // Constant -> scan_out (string long enough to exceed SSO) + { + TensorProto value_tensor; + value_tensor.set_name("string_const"); + value_tensor.add_dims(1); + value_tensor.set_data_type(TensorProto_DataType_STRING); + // Use a string longer than typical SSO buffer (>22 chars) to ensure heap allocation + value_tensor.add_string_data("this_string_exceeds_sso_threshold_and_uses_heap_allocation"); + + auto& constant_node = graph.AddNode("string_constant", "Constant", "String constant", + {}, {&scan_out}); + constant_node.AddAttribute("value", value_tensor); + } + + graph.SetInputs({&iter_num_in, &cond_in}); + graph.SetOutputs({&cond_out, &scan_out}); + + auto status = graph.Resolve(); + EXPECT_EQ(status, Status::OK()); + + return graph.ToGraphProto(); + }; + + OpTester test("Loop", 11); + auto body = create_subgraph(); + test.AddAttribute("body", body); + test.AddInput("M", {1}, {3}); + test.AddInput("cond", {1}, {true}); + + // scan output: 3 iterations, each producing a {1} string tensor -> final shape {3, 1} + test.AddOutput("scan_out_final", {3, 1}, + {"this_string_exceeds_sso_threshold_and_uses_heap_allocation", + "this_string_exceeds_sso_threshold_and_uses_heap_allocation", + "this_string_exceeds_sso_threshold_and_uses_heap_allocation"}); + + // Only CPU EP supports string tensors. + std::vector> eps; + eps.push_back(DefaultCpuExecutionProvider()); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &eps); +} + +// Verify multi-element string scan output (shape {2} per iteration). +TEST(Loop, StringScanOutputMultiElement) { + auto create_subgraph = []() { + Model model("Multi-element string scan output", false, DefaultLoggingManager().DefaultLogger()); + auto& graph = model.MainGraph(); + + std::vector inputs; + std::vector outputs; + + // graph input types + TypeProto int64_scalar; + int64_scalar.mutable_tensor_type()->set_elem_type(TensorProto_DataType_INT64); + int64_scalar.mutable_tensor_type()->mutable_shape()->add_dim()->set_dim_value(1); + + TypeProto bool_scalar; + bool_scalar.mutable_tensor_type()->set_elem_type(TensorProto_DataType_BOOL); + bool_scalar.mutable_tensor_type()->mutable_shape()->add_dim()->set_dim_value(1); + + TypeProto string_tensor; + string_tensor.mutable_tensor_type()->set_elem_type(TensorProto_DataType_STRING); + string_tensor.mutable_tensor_type()->mutable_shape()->add_dim()->set_dim_value(2); + + // graph inputs + auto& iter_num_in = graph.GetOrCreateNodeArg("iter_num_in", &int64_scalar); + auto& cond_in = graph.GetOrCreateNodeArg("cond_in", &bool_scalar); + + // graph outputs + auto& cond_out = graph.GetOrCreateNodeArg("cond_out", &bool_scalar); + auto& scan_out = graph.GetOrCreateNodeArg("scan_out", &string_tensor); + + // cond_in -> cond_out + { + inputs = {&cond_in}; + outputs = {&cond_out}; + graph.AddNode("cond_identity", "Identity", "Forward cond", inputs, outputs); + } + + // Constant -> scan_out with 2 elements + { + TensorProto value_tensor; + value_tensor.set_name("string_const"); + value_tensor.add_dims(2); + value_tensor.set_data_type(TensorProto_DataType_STRING); + value_tensor.add_string_data("first_heap_allocated_string_that_exceeds_sso_buffer_size"); + value_tensor.add_string_data("second_heap_allocated_string_that_exceeds_sso_buffer_size"); + + auto& constant_node = graph.AddNode("string_constant", "Constant", "String constant", + {}, {&scan_out}); + constant_node.AddAttribute("value", value_tensor); + } + + graph.SetInputs({&iter_num_in, &cond_in}); + graph.SetOutputs({&cond_out, &scan_out}); + + auto status = graph.Resolve(); + EXPECT_EQ(status, Status::OK()); + + return graph.ToGraphProto(); + }; + + OpTester test("Loop", 11); + auto body = create_subgraph(); + test.AddAttribute("body", body); + test.AddInput("M", {1}, {2}); + test.AddInput("cond", {1}, {true}); + + // scan output: 2 iterations x {2} elements -> {2, 2} + test.AddOutput("scan_out_final", {2, 2}, + {"first_heap_allocated_string_that_exceeds_sso_buffer_size", + "second_heap_allocated_string_that_exceeds_sso_buffer_size", + "first_heap_allocated_string_that_exceeds_sso_buffer_size", + "second_heap_allocated_string_that_exceeds_sso_buffer_size"}); + + // Only CPU EP supports string tensors. + std::vector> eps; + eps.push_back(DefaultCpuExecutionProvider()); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &eps); +} + +// Verify Loop with a string loop-carried variable (uses IDataTransfer::CopyTensor path). +TEST(Loop, StringLoopCarriedVar) { + auto create_subgraph = []() { + Model model("String loop-carried var subgraph", false, DefaultLoggingManager().DefaultLogger()); + auto& graph = model.MainGraph(); + + std::vector inputs; + std::vector outputs; + + /* Subgraph passes through a string loop-carried variable unchanged. + + iter_num_in cond_in loop_var_in (string) + (unused) | | + [Identity] [Identity] + | | + cond_out loop_var_out + */ + + TypeProto int64_scalar; + int64_scalar.mutable_tensor_type()->set_elem_type(TensorProto_DataType_INT64); + int64_scalar.mutable_tensor_type()->mutable_shape()->add_dim()->set_dim_value(1); + + TypeProto bool_scalar; + bool_scalar.mutable_tensor_type()->set_elem_type(TensorProto_DataType_BOOL); + bool_scalar.mutable_tensor_type()->mutable_shape()->add_dim()->set_dim_value(1); + + TypeProto string_tensor; + string_tensor.mutable_tensor_type()->set_elem_type(TensorProto_DataType_STRING); + string_tensor.mutable_tensor_type()->mutable_shape()->add_dim()->set_dim_value(1); + + auto& iter_num_in = graph.GetOrCreateNodeArg("iter_num_in", &int64_scalar); + auto& cond_in = graph.GetOrCreateNodeArg("cond_in", &bool_scalar); + auto& loop_var_in = graph.GetOrCreateNodeArg("loop_var_in", &string_tensor); + + auto& cond_out = graph.GetOrCreateNodeArg("cond_out", &bool_scalar); + auto& loop_var_out = graph.GetOrCreateNodeArg("loop_var_out", &string_tensor); + + // cond_in -> cond_out + { + inputs = {&cond_in}; + outputs = {&cond_out}; + graph.AddNode("cond_identity", "Identity", "Forward cond", inputs, outputs); + } + + // loop_var_in -> loop_var_out + { + inputs = {&loop_var_in}; + outputs = {&loop_var_out}; + graph.AddNode("var_identity", "Identity", "Forward loop var", inputs, outputs); + } + + graph.SetInputs({&iter_num_in, &cond_in, &loop_var_in}); + graph.SetOutputs({&cond_out, &loop_var_out}); + + auto status = graph.Resolve(); + EXPECT_EQ(status, Status::OK()); + + return graph.ToGraphProto(); + }; + + OpTester test("Loop", 11); + auto body = create_subgraph(); + test.AddAttribute("body", body); + test.AddInput("M", {1}, {3}); + test.AddInput("cond", {1}, {true}); + test.AddInput("loop_var_init", {1}, + {"a_long_string_value_that_definitely_exceeds_the_sso_threshold"}); + + test.AddOutput("loop_var_final", {1}, + {"a_long_string_value_that_definitely_exceeds_the_sso_threshold"}); + + // Only CPU EP supports string tensors. + std::vector> eps; + eps.push_back(DefaultCpuExecutionProvider()); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &eps); +} + +// Verify Loop with zero trip count produces empty scan output for strings. +TEST(Loop, StringScanOutputZeroIterations) { + auto create_subgraph = []() { + Model model("String scan output zero iter", false, DefaultLoggingManager().DefaultLogger()); + auto& graph = model.MainGraph(); + + std::vector inputs; + std::vector outputs; + + TypeProto int64_scalar; + int64_scalar.mutable_tensor_type()->set_elem_type(TensorProto_DataType_INT64); + int64_scalar.mutable_tensor_type()->mutable_shape()->add_dim()->set_dim_value(1); + + TypeProto bool_scalar; + bool_scalar.mutable_tensor_type()->set_elem_type(TensorProto_DataType_BOOL); + bool_scalar.mutable_tensor_type()->mutable_shape()->add_dim()->set_dim_value(1); + + TypeProto string_tensor; + string_tensor.mutable_tensor_type()->set_elem_type(TensorProto_DataType_STRING); + string_tensor.mutable_tensor_type()->mutable_shape()->add_dim()->set_dim_value(1); + + auto& iter_num_in = graph.GetOrCreateNodeArg("iter_num_in", &int64_scalar); + auto& cond_in = graph.GetOrCreateNodeArg("cond_in", &bool_scalar); + + auto& cond_out = graph.GetOrCreateNodeArg("cond_out", &bool_scalar); + auto& scan_out = graph.GetOrCreateNodeArg("scan_out", &string_tensor); + + { + inputs = {&cond_in}; + outputs = {&cond_out}; + graph.AddNode("cond_identity", "Identity", "Forward cond", inputs, outputs); + } + + { + TensorProto value_tensor; + value_tensor.set_name("string_const"); + value_tensor.add_dims(1); + value_tensor.set_data_type(TensorProto_DataType_STRING); + value_tensor.add_string_data("never_produced_because_zero_iterations"); + + auto& constant_node = graph.AddNode("string_constant", "Constant", "String constant", + {}, {&scan_out}); + constant_node.AddAttribute("value", value_tensor); + } + + graph.SetInputs({&iter_num_in, &cond_in}); + graph.SetOutputs({&cond_out, &scan_out}); + + auto status = graph.Resolve(); + EXPECT_EQ(status, Status::OK()); + + return graph.ToGraphProto(); + }; + + OpTester test("Loop", 11); + auto body = create_subgraph(); + test.AddAttribute("body", body); + test.AddInput("M", {1}, {0}); + test.AddInput("cond", {1}, {true}); + + // Zero iterations -> scan output shape {0, 1} with no elements + test.AddOutput("scan_out_final", {0, 1}, {}); + + // Only CPU EP supports string tensors. + std::vector> eps; + eps.push_back(DefaultCpuExecutionProvider()); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &eps); +} + #if defined(USE_CUDA) // test that when part of the subgraph run on CUDA it executes successfully TEST(Loop, MixedExecutionProviders) { From 1a98f1586d8d3017251ffada7256084be07971c6 Mon Sep 17 00:00:00 2001 From: Dmitri Smirnov Date: Tue, 30 Jun 2026 14:20:44 -0700 Subject: [PATCH 04/17] Validate and inline external data in node tensor attributes during session initialization (#29250) This pull request introduces improved support and validation for external data in tensor attributes, particularly ensuring that external data in node attributes is properly validated and inlined, and that in-memory references are correctly rejected. Additionally, it introduces new tests to cover these scenarios, and refactors some utility functions in the test code for clarity and consistency. **External Data Handling and Validation:** * Updated `Graph::ConvertInitializersIntoOrtValues()` in `graph.cc` to: * Use `InlinedHashSet` for tracking validated external data paths for efficiency. * Add a new step that validates and inlines external data in node tensor attributes, ensuring that only file-based external data is accepted and in-memory references are rejected. This guarantees all execution providers have uniform access to attribute data. [[1]](diffhunk://#diff-e231a92b40d89409cc8e82436be0a15bc87ef95c93b303b9feaeab6e50c8835cL3842-R3842) [[2]](diffhunk://#diff-e231a92b40d89409cc8e82436be0a15bc87ef95c93b303b9feaeab6e50c8835cL3906-R3971) **Testing:** * Added comprehensive tests in `label_encoder_test.cc` to verify: * Valid external data in tensor attributes is loaded and inlined. * In-memory external data references in node attributes are rejected. * Duplicate key handling and singleton default tensor behavior. * Added a test in `tree_ensembler_test.cc` to ensure that in-memory external data references in node attributes are rejected, preventing invalid attribute usage. **Test Utility Refactoring:** * Refactored utility functions in `tree_ensembler_test.cc` and `treeregressor_test.cc`: * Renamed and standardized helper functions for array and string manipulations to improve code clarity (e.g., `MultiplyUpdateArray`, `MultiplyArraysValues`, `MultiplyUpdateArrayString`). [[1]](diffhunk://#diff-9bce50df70fddc092ce6fc5351812c621343f682a15ca178c34e6dd9415e9a39L47-R49) [[2]](diffhunk://#diff-9bce50df70fddc092ce6fc5351812c621343f682a15ca178c34e6dd9415e9a39L60-R62) [[3]](diffhunk://#diff-9bce50df70fddc092ce6fc5351812c621343f682a15ca178c34e6dd9415e9a39L90-R92) [[4]](diffhunk://#diff-9bce50df70fddc092ce6fc5351812c621343f682a15ca178c34e6dd9415e9a39L116-R127) [[5]](diffhunk://#diff-9bce50df70fddc092ce6fc5351812c621343f682a15ca178c34e6dd9415e9a39L174-R186) [[6]](diffhunk://#diff-9bce50df70fddc092ce6fc5351812c621343f682a15ca178c34e6dd9415e9a39L228-R230) [[7]](diffhunk://#diff-9bce50df70fddc092ce6fc5351812c621343f682a15ca178c34e6dd9415e9a39L243-R245) [[8]](diffhunk://#diff-08b3495816c68f145657ecff63d7b5f3d56813586ec62f7324c22977e70e336bR6-R13) [[9]](diffhunk://#diff-08b3495816c68f145657ecff63d7b5f3d56813586ec62f7324c22977e70e336bL23-R25) **Test File Includes:** * Added necessary includes for file utilities and path handling in test files to support new test scenarios involving external data files. [[1]](diffhunk://#diff-0db6b4a4d9a180daab3cc2eab5da4437a2a7a3a2e81f53119f5e6298bd3dc4a5R7-R8) [[2]](diffhunk://#diff-9bce50df70fddc092ce6fc5351812c621343f682a15ca178c34e6dd9415e9a39R6-R7) [[3]](diffhunk://#diff-08b3495816c68f145657ecff63d7b5f3d56813586ec62f7324c22977e70e336bR6-R13) --- onnxruntime/core/graph/graph.cc | 74 +++++++- .../providers/cpu/ml/label_encoder_test.cc | 174 ++++++++++++++++++ .../providers/cpu/ml/tree_ensembler_test.cc | 118 +++++++++--- .../providers/cpu/ml/treeregressor_test.cc | 150 ++++++++++----- 4 files changed, 442 insertions(+), 74 deletions(-) diff --git a/onnxruntime/core/graph/graph.cc b/onnxruntime/core/graph/graph.cc index ff60f6cb5f741..479a30a783f77 100644 --- a/onnxruntime/core/graph/graph.cc +++ b/onnxruntime/core/graph/graph.cc @@ -3843,7 +3843,7 @@ Status Graph::ConvertInitializersIntoOrtValues() { FindAllSubgraphs(all_subgraphs); const auto& model_path = GetModel().ModelPath(); - std::unordered_set validated_external_data_paths; + InlinedHashSet validated_external_data_paths; auto put_weights_maybe_in_memory_func = [&](Graph& graph) -> Status { // if we have any initializers that are not in memory, put them there. @@ -3907,7 +3907,77 @@ Status Graph::ConvertInitializersIntoOrtValues() { return Status::OK(); }; - return ForThisAndAllSubgraphs(all_subgraphs, put_weights_maybe_in_memory_func); + ORT_RETURN_IF_ERROR(ForThisAndAllSubgraphs(all_subgraphs, put_weights_maybe_in_memory_func)); + + // Validate and inline external data in node tensor attributes. + // In-memory references are rejected (no legitimate source creates them for attributes). + // File-based external data paths are validated, read from disk, and inlined as raw_data + // so all EPs (including plugins) can access attribute data uniformly. + auto inline_external_attr_tensors_func = [&](Graph& graph) -> Status { + // Helper: validate and inline a single external TensorProto. + auto inline_tensor = [&](ONNX_NAMESPACE::TensorProto& tensor_proto, + const Node& node, std::string_view attr_name) -> Status { + ORT_RETURN_IF(utils::HasExternalDataInMemory(tensor_proto), + "Node '", node.Name(), "' attribute '", attr_name, + "' contains an in-memory external data reference, which is not permitted ", + "for node attributes."); + + std::unique_ptr external_data_info; + { + auto create_status = + onnxruntime::ExternalDataInfo::Create(tensor_proto.external_data(), external_data_info); + if (!create_status.IsOK()) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "Node '", node.Name(), "' attribute '", attr_name, + "': ", create_status.ErrorMessage()); + } + } + const auto& location = external_data_info->GetRelPath(); + + if (validated_external_data_paths.count(location) == 0) { + auto path_status = utils::ValidateExternalDataPath(model_path, location); + if (!path_status.IsOK()) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "Node '", node.Name(), "' attribute '", attr_name, + "': ", path_status.ErrorMessage()); + } + validated_external_data_paths.insert(location); + } + + std::vector buffer; + auto unpack_status = utils::UnpackInitializerData(tensor_proto, model_path, buffer); + if (!unpack_status.IsOK()) { + return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, + "Node '", node.Name(), "' attribute '", attr_name, + "': ", unpack_status.ErrorMessage()); + } + + tensor_proto.clear_external_data(); + tensor_proto.set_data_location(ONNX_NAMESPACE::TensorProto_DataLocation_DEFAULT); + utils::SetRawDataInTensorProto(tensor_proto, buffer.data(), buffer.size()); + return Status::OK(); + }; + + for (auto& node : graph.Nodes()) { + for (auto& [attr_name, attr_proto] : node.GetMutableAttributes()) { + if (utils::HasTensor(attr_proto)) { + auto* tensor_proto = attr_proto.mutable_t(); + if (utils::HasExternalData(*tensor_proto)) { + ORT_RETURN_IF_ERROR(inline_tensor(*tensor_proto, node, attr_name)); + } + } else if (utils::HasTensors(attr_proto)) { + for (auto& tensor_proto : *attr_proto.mutable_tensors()) { + if (utils::HasExternalData(tensor_proto)) { + ORT_RETURN_IF_ERROR(inline_tensor(tensor_proto, node, attr_name)); + } + } + } + } + } + return Status::OK(); + }; + + return ForThisAndAllSubgraphs(all_subgraphs, inline_external_attr_tensors_func); } void Graph::SetName(const std::string& name) { diff --git a/onnxruntime/test/providers/cpu/ml/label_encoder_test.cc b/onnxruntime/test/providers/cpu/ml/label_encoder_test.cc index 034d206fec2f4..c68e6d73a56ef 100644 --- a/onnxruntime/test/providers/cpu/ml/label_encoder_test.cc +++ b/onnxruntime/test/providers/cpu/ml/label_encoder_test.cc @@ -4,6 +4,8 @@ #include "gtest/gtest.h" #include "core/framework/tensorprotoutils.h" #include "test/providers/provider_test_utils.h" +#include "test/util/include/file_util.h" +#include "core/platform/path_lib.h" namespace onnxruntime { namespace test { @@ -756,5 +758,177 @@ TEST(LabelEncoder, EmptyInputOpset4) { test.Run(); } +// External data in tensor attributes: file-based external data is validated and inlined +// during session initialization. In-memory references are rejected by the ONNX checker +// during Graph::Resolve() (it validates that external data locations are regular files). +// In no-exceptions builds, the ONNX checker's fail_check calls abort() so these tests cannot run. +#if !defined(ORT_NO_EXCEPTIONS) + +// RAII helper that creates a unique dummy binary file and removes it on destruction. +static std::pair CreateExternalDataFile(const void* data, size_t num_bytes) { + PathString filename(ORT_TSTR("ext_data_XXXXXX")); + FILE* fp = nullptr; + CreateTestFile(fp, filename); + ScopedFileDeleter deleter(filename); // ensure cleanup even if fwrite/fclose throws + size_t written = fwrite(data, 1, num_bytes, fp); + ORT_ENFORCE(written == num_bytes, "Failed to write external data file"); + ORT_ENFORCE(fclose(fp) == 0, "Failed to close external data file"); + return {ToUTF8String(filename), std::move(deleter)}; +} + +// Helper: create a TensorProto that references external data in the given file. +static ONNX_NAMESPACE::TensorProto MakeExternalInt64TensorProto(const std::string& name, + const std::string& filename, + int64_t num_elements) { + ONNX_NAMESPACE::TensorProto proto; + proto.set_name(name); + proto.set_data_type(ONNX_NAMESPACE::TensorProto_DataType_INT64); + proto.add_dims(num_elements); + proto.set_data_location(ONNX_NAMESPACE::TensorProto_DataLocation_EXTERNAL); + auto* loc = proto.add_external_data(); + loc->set_key("location"); + loc->set_value(filename); + auto* offset = proto.add_external_data(); + offset->set_key("offset"); + offset->set_value("0"); + auto* length = proto.add_external_data(); + length->set_key("length"); + length->set_value(std::to_string(num_elements * static_cast(sizeof(int64_t)))); + return proto; +} + +// Helper: create a TensorProto with in-memory external data reference (should be rejected). +static ONNX_NAMESPACE::TensorProto MakeInMemoryExternalTensorProto(const std::string& name, + int64_t num_elements) { + ONNX_NAMESPACE::TensorProto proto; + proto.set_name(name); + proto.set_data_type(ONNX_NAMESPACE::TensorProto_DataType_INT64); + proto.add_dims(num_elements); + proto.set_data_location(ONNX_NAMESPACE::TensorProto_DataLocation_EXTERNAL); + auto* loc = proto.add_external_data(); + loc->set_key("location"); + loc->set_value(ToUTF8String(onnxruntime::utils::kTensorProtoNativeEndianMemoryAddressTag)); + auto* offset = proto.add_external_data(); + offset->set_key("offset"); + offset->set_value("12345678"); + auto* length = proto.add_external_data(); + length->set_key("length"); + length->set_value(std::to_string(num_elements * static_cast(sizeof(int64_t)))); + return proto; +} + +// Valid external data in tensor attributes should be loaded and inlined during session initialization. +TEST(LabelEncoder, ExternalDataInKeysTensorOpset4) { + std::vector key_data{1, 2}; + auto [ext_path, ext_deleter] = CreateExternalDataFile(key_data.data(), key_data.size() * sizeof(int64_t)); + + OpTester test("LabelEncoder", 4, onnxruntime::kMLDomain); + test.AddAttribute("keys_tensor", MakeExternalInt64TensorProto("keys_tensor", ext_path, 2)); + + ONNX_NAMESPACE::TensorProto values_proto; + values_proto.set_name("values_tensor"); + values_proto.set_data_type(ONNX_NAMESPACE::TensorProto_DataType_INT64); + values_proto.add_dims(2); + values_proto.add_int64_data(10); + values_proto.add_int64_data(20); + test.AddAttribute("values_tensor", values_proto); + + ONNX_NAMESPACE::TensorProto default_proto; + default_proto.set_name("default_tensor"); + default_proto.set_data_type(ONNX_NAMESPACE::TensorProto_DataType_INT64); + default_proto.add_dims(1); + default_proto.add_int64_data(42); + test.AddAttribute("default_tensor", default_proto); + + test.AddInput("X", {1, 3}, {1, 2, 99}); + test.AddOutput("Y", {1, 3}, {10, 20, 42}); + + test.Run(); +} + +// In-memory external data references in node attributes are rejected during initialization. +TEST(LabelEncoder, RejectsInMemoryExternalDataInKeysTensorOpset4) { + OpTester test("LabelEncoder", 4, onnxruntime::kMLDomain); + test.AddAttribute("keys_tensor", MakeInMemoryExternalTensorProto("keys_tensor", 2)); + + ONNX_NAMESPACE::TensorProto values_proto; + values_proto.set_name("values_tensor"); + values_proto.set_data_type(ONNX_NAMESPACE::TensorProto_DataType_INT64); + values_proto.add_dims(2); + values_proto.add_int64_data(10); + values_proto.add_int64_data(20); + test.AddAttribute("values_tensor", values_proto); + + ONNX_NAMESPACE::TensorProto default_proto; + default_proto.set_name("default_tensor"); + default_proto.set_data_type(ONNX_NAMESPACE::TensorProto_DataType_INT64); + default_proto.add_dims(1); + default_proto.add_int64_data(0); + test.AddAttribute("default_tensor", default_proto); + + test.AddInput("X", {1, 2}, {1, 2}); + test.AddOutput("Y", {1, 2}, {10, 20}); + + // Error originates from the ONNX checker (checker::check_node) during Graph::Resolve(). + // There is no way to disable this check. + test.Run(OpTester::ExpectResult::kExpectFailure, "is not regular file"); +} + +#endif // !defined(ORT_NO_EXCEPTIONS) + +// Duplicate keys: emplace() keeps the first occurrence. Verify this behavior. +TEST(LabelEncoder, DuplicateKeysFirstWinsOpset4) { + std::vector dims{1, 3}; + + std::vector input{1, 2, 3}; + // key 1 maps to 10 (first), not 99 (second duplicate) + std::vector output{10, 20, 42}; + std::vector key_data{1, 2, 1}; // duplicate key 1 + std::vector value_data{10, 20, 99}; + + OpTester test("LabelEncoder", 4, onnxruntime::kMLDomain); + + test.AddAttribute("keys_int64s", key_data); + test.AddAttribute("values_int64s", value_data); + + ONNX_NAMESPACE::TensorProto default_proto; + default_proto.set_name("default_tensor"); + default_proto.set_data_type(ONNX_NAMESPACE::TensorProto_DataType_INT64); + default_proto.add_dims(1); + default_proto.add_int64_data(42); + test.AddAttribute("default_tensor", default_proto); + + test.AddInput("X", dims, input); + test.AddOutput("Y", dims, output); + + test.Run(); +} + +// Singleton 1D default_tensor (dims=[1]) — the ONNX spec requires this shape +TEST(LabelEncoder, SingletonDefaultTensorOpset4) { + std::vector dims{1, 3}; + + std::vector input{1, 2, 99}; + std::vector output{10, 20, -7}; + + OpTester test("LabelEncoder", 4, onnxruntime::kMLDomain); + + test.AddAttribute("keys_int64s", std::vector{1, 2}); + test.AddAttribute("values_int64s", std::vector{10, 20}); + + // 1D singleton default_tensor with dims=[1] + ONNX_NAMESPACE::TensorProto default_proto; + default_proto.set_name("default_tensor"); + default_proto.set_data_type(ONNX_NAMESPACE::TensorProto_DataType_INT64); + default_proto.add_dims(1); + default_proto.add_int64_data(-7); + test.AddAttribute("default_tensor", default_proto); + + test.AddInput("X", dims, input); + test.AddOutput("Y", dims, output); + + test.Run(); +} + } // namespace test } // namespace onnxruntime diff --git a/onnxruntime/test/providers/cpu/ml/tree_ensembler_test.cc b/onnxruntime/test/providers/cpu/ml/tree_ensembler_test.cc index 922fab7cd2e43..9f1822336b36f 100644 --- a/onnxruntime/test/providers/cpu/ml/tree_ensembler_test.cc +++ b/onnxruntime/test/providers/cpu/ml/tree_ensembler_test.cc @@ -3,6 +3,8 @@ #include "gtest/gtest.h" #include "test/providers/provider_test_utils.h" +#include "core/framework/tensorprotoutils.h" +#include "core/platform/path_lib.h" namespace onnxruntime { namespace test { @@ -44,7 +46,7 @@ static ONNX_NAMESPACE::TensorProto make_tensor(std::vector array, std:: } template -void _multiply_update_array(std::vector& data, int n, T inc = 0) { +void MultiplyUpdateArray(std::vector& data, int n, T inc = 0) { std::vector copy = data; data.resize(copy.size() * n); T cst = 0; @@ -57,7 +59,7 @@ void _multiply_update_array(std::vector& data, int n, T inc = 0) { } template -void _multiply_update_childnode(std::vector& childnodes, std::vector& childleafs, std::vector& otherchildleafs, int n) { +void MultiplyUpdateChildnode(std::vector& childnodes, std::vector& childleafs, std::vector& otherchildleafs, int n) { int64_t leafs_cnt = 0; int64_t nodes_cnt = childnodes.size(); for (auto& childleaf : childleafs) { @@ -87,7 +89,7 @@ void _multiply_update_childnode(std::vector& childnodes, std::vector& chil } template -void _multiply_arrays_values(std::vector& data, int64_t val) { +void MultiplyArraysValues(std::vector& data, int64_t val) { for (auto& curr : data) { curr *= val; } @@ -154,16 +156,16 @@ void GenTreeAndRunTest(const std::vector& X, const std::vector& Y, const i if (n_trees > 1) { // Multiplies the number of trees to test the parallelization by trees. - _multiply_update_array(tree_roots, n_trees, (int64_t)nodes_truenodeids.size()); - _multiply_update_array(nodes_featureids, n_trees); - _multiply_update_childnode(nodes_truenodeids, nodes_trueleafs, nodes_falseleafs, n_trees); - _multiply_update_childnode(nodes_falsenodeids, nodes_falseleafs, nodes_trueleafs, n_trees); - _multiply_update_array(nodes_trueleafs, n_trees); - _multiply_update_array(nodes_falseleafs, n_trees); - _multiply_update_array(leaf_targetids, n_trees); - _multiply_update_array(nodes_modes, n_trees); - _multiply_update_array(nodes_splits, n_trees); - _multiply_update_array(leaf_weights, n_trees); + MultiplyUpdateArray(tree_roots, n_trees, (int64_t)nodes_truenodeids.size()); + MultiplyUpdateArray(nodes_featureids, n_trees); + MultiplyUpdateChildnode(nodes_truenodeids, nodes_trueleafs, nodes_falseleafs, n_trees); + MultiplyUpdateChildnode(nodes_falsenodeids, nodes_falseleafs, nodes_trueleafs, n_trees); + MultiplyUpdateArray(nodes_trueleafs, n_trees); + MultiplyUpdateArray(nodes_falseleafs, n_trees); + MultiplyUpdateArray(leaf_targetids, n_trees); + MultiplyUpdateArray(nodes_modes, n_trees); + MultiplyUpdateArray(nodes_splits, n_trees); + MultiplyUpdateArray(leaf_weights, n_trees); } auto nodes_modes_as_tensor = make_tensor(nodes_modes, "nodes_modes"); @@ -212,17 +214,17 @@ void GenTreeAndRunTestWithSetMembership(const std::vector& X, const std::vect if (n_trees > 1) { // Multiplies the number of trees to test the parallelization by trees. - _multiply_update_array(tree_roots, n_trees, (int64_t)nodes_truenodeids.size()); - _multiply_update_array(nodes_featureids, n_trees); - _multiply_update_childnode(nodes_truenodeids, nodes_trueleafs, nodes_falseleafs, n_trees); - _multiply_update_childnode(nodes_falsenodeids, nodes_falseleafs, nodes_trueleafs, n_trees); - _multiply_update_array(nodes_trueleafs, n_trees); - _multiply_update_array(nodes_falseleafs, n_trees); - _multiply_update_array(leaf_targetids, n_trees); - _multiply_update_array(nodes_modes, n_trees); - _multiply_update_array(nodes_splits, n_trees); - _multiply_update_array(membership_values, n_trees); - _multiply_update_array(leaf_weights, n_trees); + MultiplyUpdateArray(tree_roots, n_trees, (int64_t)nodes_truenodeids.size()); + MultiplyUpdateArray(nodes_featureids, n_trees); + MultiplyUpdateChildnode(nodes_truenodeids, nodes_trueleafs, nodes_falseleafs, n_trees); + MultiplyUpdateChildnode(nodes_falsenodeids, nodes_falseleafs, nodes_trueleafs, n_trees); + MultiplyUpdateArray(nodes_trueleafs, n_trees); + MultiplyUpdateArray(nodes_falseleafs, n_trees); + MultiplyUpdateArray(leaf_targetids, n_trees); + MultiplyUpdateArray(nodes_modes, n_trees); + MultiplyUpdateArray(nodes_splits, n_trees); + MultiplyUpdateArray(membership_values, n_trees); + MultiplyUpdateArray(leaf_weights, n_trees); } auto nodes_modes_as_tensor = make_tensor(nodes_modes, "nodes_modes"); @@ -266,7 +268,7 @@ TEST(MLOpTest, TreeEnsembleDouble) { std::vector Y = {5.23f, 0.f, 5.23f, 0.f, 0.f, 12.12f}; GenTreeAndRunTest(X, Y, 1, 1); - _multiply_arrays_values(Y, 3); + MultiplyArraysValues(Y, 3); GenTreeAndRunTest(X, Y, 1, 3); } @@ -281,7 +283,7 @@ TEST(MLOpTest, TreeEnsembleSetMembership) { 0.f, 10.f, 0.f, 0.f}; GenTreeAndRunTestWithSetMembership(X, Y, 1, 1); - _multiply_arrays_values(Y, 5); + MultiplyArraysValues(Y, 5); GenTreeAndRunTestWithSetMembership(X, Y, 1, 5); } @@ -525,5 +527,69 @@ TEST(MLOpTest, TreeEnsembleIssue25400) { test.Run(); } +// In-memory external data references in node attributes are rejected by the ONNX checker +// during Graph::Resolve() (it validates that external data locations are regular files). +// In no-exceptions builds, the ONNX checker's fail_check calls abort() so these tests cannot run. +#if !defined(ORT_NO_EXCEPTIONS) + +// In-memory external data references in node attributes are rejected during Graph::Resolve(). +TEST(MLOpTest, TreeEnsembleRejectsInMemoryExternalDataInTensorAttribute) { + OpTester test("TreeEnsemble", 5, onnxruntime::kMLDomain); + + // nodes_splits with in-memory external data reference (should be rejected) + ONNX_NAMESPACE::TensorProto splits_proto; + splits_proto.set_name("nodes_splits"); + splits_proto.set_data_type(ONNX_NAMESPACE::TensorProto_DataType_FLOAT); + splits_proto.add_dims(1); + splits_proto.set_data_location(ONNX_NAMESPACE::TensorProto_DataLocation_EXTERNAL); + auto* loc = splits_proto.add_external_data(); + loc->set_key("location"); + loc->set_value(ToUTF8String(onnxruntime::utils::kTensorProtoNativeEndianMemoryAddressTag)); + auto* offset = splits_proto.add_external_data(); + offset->set_key("offset"); + offset->set_value("12345678"); + auto* length = splits_proto.add_external_data(); + length->set_key("length"); + length->set_value("4"); + test.AddAttribute("nodes_splits", splits_proto); + + // Minimal valid structure for remaining attributes + ONNX_NAMESPACE::TensorProto leaf_weights_proto; + leaf_weights_proto.set_name("leaf_weights"); + leaf_weights_proto.set_data_type(ONNX_NAMESPACE::TensorProto_DataType_FLOAT); + leaf_weights_proto.add_dims(2); + leaf_weights_proto.add_float_data(1.0f); + leaf_weights_proto.add_float_data(2.0f); + test.AddAttribute("leaf_weights", leaf_weights_proto); + + ONNX_NAMESPACE::TensorProto modes_proto; + modes_proto.set_name("nodes_modes"); + modes_proto.set_data_type(ONNX_NAMESPACE::TensorProto_DataType_UINT8); + modes_proto.add_dims(1); + modes_proto.add_int32_data(0); + test.AddAttribute("nodes_modes", modes_proto); + + test.AddAttribute("aggregate_function", static_cast(1)); + test.AddAttribute("leaf_targetids", std::vector{0, 0}); + test.AddAttribute("n_targets", static_cast(1)); + test.AddAttribute("nodes_falseleafs", std::vector{1}); + test.AddAttribute("nodes_falsenodeids", std::vector{1}); + test.AddAttribute("nodes_featureids", std::vector{0}); + test.AddAttribute("nodes_trueleafs", std::vector{1}); + test.AddAttribute("nodes_truenodeids", std::vector{0}); + test.AddAttribute("post_transform", static_cast(0)); + test.AddAttribute("tree_roots", std::vector{0}); + + std::vector X = {1.f}; + test.AddInput("X", {1, 1}, X); + test.AddOutput("Y", {1, 1}, {0.f}); + + // Error originates from the ONNX checker (checker::check_node) during Graph::Resolve(). + // There is no way to disable this check. + test.Run(OpTester::ExpectResult::kExpectFailure, "is not regular file"); +} + +#endif // !defined(ORT_NO_EXCEPTIONS) + } // namespace test } // namespace onnxruntime diff --git a/onnxruntime/test/providers/cpu/ml/treeregressor_test.cc b/onnxruntime/test/providers/cpu/ml/treeregressor_test.cc index a7892da3d676f..b6806042a5c07 100644 --- a/onnxruntime/test/providers/cpu/ml/treeregressor_test.cc +++ b/onnxruntime/test/providers/cpu/ml/treeregressor_test.cc @@ -3,12 +3,14 @@ #include "gtest/gtest.h" #include "test/providers/provider_test_utils.h" +#include "core/framework/tensorprotoutils.h" +#include "core/platform/path_lib.h" namespace onnxruntime { namespace test { template -void _multiply_update_array(std::vector& data, int n, T inc = 0) { +void MultiplyUpdateArray(std::vector& data, int n, T inc = 0) { std::vector copy = data; data.resize(copy.size() * n); T cst = 0; @@ -20,7 +22,7 @@ void _multiply_update_array(std::vector& data, int n, T inc = 0) { } } -void _multiply_update_array_string(std::vector& data, int n) { +void MultiplyUpdateArrayString(std::vector& data, int n) { std::vector copy = data; data.resize(copy.size() * n); for (int i = 0; i < n; ++i) { @@ -52,17 +54,17 @@ void GenTreeAndRunTest(int opsetml, const std::vector& X, const std::vector 1) { // Multiplies the number of trees to test the parallelization by trees. - _multiply_update_array(lefts, n_trees); - _multiply_update_array(rights, n_trees); - _multiply_update_array(treeids, n_trees, (int64_t)3); - _multiply_update_array(nodeids, n_trees); - _multiply_update_array(featureids, n_trees); - _multiply_update_array(thresholds, n_trees); - _multiply_update_array_string(modes, n_trees); - _multiply_update_array(target_treeids, n_trees, (int64_t)3); - _multiply_update_array(target_nodeids, n_trees); - _multiply_update_array(target_classids, n_trees); - _multiply_update_array(target_weights, n_trees); + MultiplyUpdateArray(lefts, n_trees); + MultiplyUpdateArray(rights, n_trees); + MultiplyUpdateArray(treeids, n_trees, (int64_t)3); + MultiplyUpdateArray(nodeids, n_trees); + MultiplyUpdateArray(featureids, n_trees); + MultiplyUpdateArray(thresholds, n_trees); + MultiplyUpdateArrayString(modes, n_trees); + MultiplyUpdateArray(target_treeids, n_trees, (int64_t)3); + MultiplyUpdateArray(target_nodeids, n_trees); + MultiplyUpdateArray(target_classids, n_trees); + MultiplyUpdateArray(target_weights, n_trees); } // add attributes @@ -146,17 +148,17 @@ void GenTreeAndRunTest_as_tensor(int opsetml, const std::vector& X, const std if (n_trees > 1) { // Multiplies the number of trees to test the parallelization by trees. - _multiply_update_array(lefts, n_trees); - _multiply_update_array(rights, n_trees); - _multiply_update_array(treeids, n_trees, (int64_t)3); - _multiply_update_array(nodeids, n_trees); - _multiply_update_array(featureids, n_trees); - _multiply_update_array(thresholds, n_trees); - _multiply_update_array_string(modes, n_trees); - _multiply_update_array(target_treeids, n_trees, (int64_t)3); - _multiply_update_array(target_nodeids, n_trees); - _multiply_update_array(target_classids, n_trees); - _multiply_update_array(target_weights, n_trees); + MultiplyUpdateArray(lefts, n_trees); + MultiplyUpdateArray(rights, n_trees); + MultiplyUpdateArray(treeids, n_trees, (int64_t)3); + MultiplyUpdateArray(nodeids, n_trees); + MultiplyUpdateArray(featureids, n_trees); + MultiplyUpdateArray(thresholds, n_trees); + MultiplyUpdateArrayString(modes, n_trees); + MultiplyUpdateArray(target_treeids, n_trees, (int64_t)3); + MultiplyUpdateArray(target_nodeids, n_trees); + MultiplyUpdateArray(target_classids, n_trees); + MultiplyUpdateArray(target_weights, n_trees); } // add attributes @@ -356,17 +358,17 @@ void GenTreeAndRunTest1(int opsetml, const std::string& aggFunction, bool one_ob if (n_trees > 1) { // Multiplies the number of trees to test the parallelization by trees. - _multiply_update_array(lefts, n_trees); - _multiply_update_array(rights, n_trees); - _multiply_update_array(treeids, n_trees, (int64_t)3); - _multiply_update_array(nodeids, n_trees); - _multiply_update_array(featureids, n_trees); - _multiply_update_array(thresholds, n_trees); - _multiply_update_array_string(modes, n_trees); - _multiply_update_array(target_treeids, n_trees, (int64_t)3); - _multiply_update_array(target_nodeids, n_trees); - _multiply_update_array(target_classids, n_trees); - _multiply_update_array(target_weights, n_trees); + MultiplyUpdateArray(lefts, n_trees); + MultiplyUpdateArray(rights, n_trees); + MultiplyUpdateArray(treeids, n_trees, (int64_t)3); + MultiplyUpdateArray(nodeids, n_trees); + MultiplyUpdateArray(featureids, n_trees); + MultiplyUpdateArray(thresholds, n_trees); + MultiplyUpdateArrayString(modes, n_trees); + MultiplyUpdateArray(target_treeids, n_trees, (int64_t)3); + MultiplyUpdateArray(target_nodeids, n_trees); + MultiplyUpdateArray(target_classids, n_trees); + MultiplyUpdateArray(target_weights, n_trees); } std::vector results; @@ -469,17 +471,17 @@ void GenTreeAndRunTest1_as_tensor(int opsetml, const std::string& aggFunction, b if (n_trees > 1) { // Multiplies the number of trees to test the parallelization by trees. - _multiply_update_array(lefts, n_trees); - _multiply_update_array(rights, n_trees); - _multiply_update_array(treeids, n_trees, (int64_t)3); - _multiply_update_array(nodeids, n_trees); - _multiply_update_array(featureids, n_trees); - _multiply_update_array(thresholds, n_trees); - _multiply_update_array_string(modes, n_trees); - _multiply_update_array(target_treeids, n_trees, (int64_t)3); - _multiply_update_array(target_nodeids, n_trees); - _multiply_update_array(target_classids, n_trees); - _multiply_update_array(target_weights, n_trees); + MultiplyUpdateArray(lefts, n_trees); + MultiplyUpdateArray(rights, n_trees); + MultiplyUpdateArray(treeids, n_trees, (int64_t)3); + MultiplyUpdateArray(nodeids, n_trees); + MultiplyUpdateArray(featureids, n_trees); + MultiplyUpdateArray(thresholds, n_trees); + MultiplyUpdateArrayString(modes, n_trees); + MultiplyUpdateArray(target_treeids, n_trees, (int64_t)3); + MultiplyUpdateArray(target_nodeids, n_trees); + MultiplyUpdateArray(target_classids, n_trees); + MultiplyUpdateArray(target_weights, n_trees); } std::vector results; @@ -1081,5 +1083,61 @@ TEST(MLOpTest, TreeEnsembleRegressorBaseValuesWrongSize) { test.Run(OpTester::ExpectResult::kExpectFailure, "base_values should have 0 or 2 values."); } +// In-memory external data references in node attributes are rejected by the ONNX checker +// during Graph::Resolve() (it validates that external data locations are regular files). +// In no-exceptions builds, the ONNX checker's fail_check calls abort() so these tests cannot run. +#if !defined(ORT_NO_EXCEPTIONS) + +TEST(MLOpTest, TreeEnsembleRegressorRejectsInMemoryExternalDataInTensorAttribute) { + OpTester test("TreeEnsembleRegressor", 3, onnxruntime::kMLDomain); + + // Minimal valid tree structure + std::vector lefts = {1, 0, 0}; + std::vector rights = {2, 0, 0}; + std::vector treeids = {0, 0, 0}; + std::vector nodeids = {0, 1, 2}; + std::vector featureids = {0, 0, 0}; + std::vector modes = {"BRANCH_LEQ", "LEAF", "LEAF"}; + + test.AddAttribute("nodes_truenodeids", lefts); + test.AddAttribute("nodes_falsenodeids", rights); + test.AddAttribute("nodes_treeids", treeids); + test.AddAttribute("nodes_nodeids", nodeids); + test.AddAttribute("nodes_featureids", featureids); + test.AddAttribute("nodes_modes", modes); + test.AddAttribute("target_treeids", std::vector{0, 0}); + test.AddAttribute("target_nodeids", std::vector{1, 2}); + test.AddAttribute("target_ids", std::vector{0, 0}); + test.AddAttribute("target_weights", std::vector{1.f, 2.f}); + test.AddAttribute("n_targets", static_cast(1)); + + // nodes_values_as_tensor with in-memory external data reference (should be rejected) + ONNX_NAMESPACE::TensorProto values_proto; + values_proto.set_name("nodes_values_as_tensor"); + values_proto.set_data_type(ONNX_NAMESPACE::TensorProto_DataType_FLOAT); + values_proto.add_dims(3); + values_proto.set_data_location(ONNX_NAMESPACE::TensorProto_DataLocation_EXTERNAL); + auto* loc = values_proto.add_external_data(); + loc->set_key("location"); + loc->set_value(ToUTF8String(utils::kTensorProtoNativeEndianMemoryAddressTag)); + auto* offset = values_proto.add_external_data(); + offset->set_key("offset"); + offset->set_value("12345678"); + auto* length = values_proto.add_external_data(); + length->set_key("length"); + length->set_value("12"); + test.AddAttribute("nodes_values_as_tensor", values_proto); + + std::vector X = {1.f}; + test.AddInput("X", {1, 1}, X); + test.AddOutput("Y", {1, 1}, {0.f}); + + // Error originates from the ONNX checker (checker::check_node) during Graph::Resolve(). + // There is no way to disable this check. + test.Run(OpTester::ExpectResult::kExpectFailure, "is not regular file"); +} + +#endif // !defined(ORT_NO_EXCEPTIONS) + } // namespace test } // namespace onnxruntime From 4f2ba9c4803eb3f87d6c0491e8fbebb97fbd5e10 Mon Sep 17 00:00:00 2001 From: mingyue Date: Tue, 30 Jun 2026 18:45:14 -0500 Subject: [PATCH 05/17] [VitisAI] Recognize native-endian in-memory address tag in process_ext_address (#29248) ### Description `vaip::process_ext_address` (in the VitisAI EP) decodes the in-memory "external data" address marker that ORT plants on initializers whose data lives in an in-process buffer. It only matched the little-endian tag `"*/_ORT_MEM_ADDR_/*"` (`kTensorProtoLittleEndianMemoryAddressTag`). This PR makes it also recognize the native-endian tag `"*/_ORT_NATIVE_ENDIAN_MEM_ADDR_/*"` (`kTensorProtoNativeEndianMemoryAddressTag`). ```cpp if (file == "*/_ORT_MEM_ADDR_/*" || file == "*/_ORT_NATIVE_ENDIAN_MEM_ADDR_/*") { auto addr = reinterpret_cast(offset); return {addr, size}; } ``` ### Motivation and Context Two existing changes combined to break VitisAI model compilation: 1. #27404 split the single in-memory marker into a little-endian and a native-endian variant, and switched `TensorToTensorProto(use_tensor_buffer=true)` to emit the **native-endian** tag by default. Both `HasExternalDataInMemory` and the data readers treat the two tags equivalently. 2. #28709 added an explicit `ORT_ENFORCE(!utils::HasExternalDataInMemory(tensor), ...)` in `Graph::Graph`, rejecting any in-memory address marker found in a deserialized model protobuf. Because `process_ext_address` still only recognized the little-endian tag, in-memory initializers (e.g. quantized weights moved out of the `TensorProto` during graph optimization) are no longer detected in `vaip::model_clone`. They fall through to the verbatim-copy branch and the native-endian marker is carried into the cloned model proto. When ORT then constructs the cloned `Graph`, the new enforce fires and aborts compilation: ``` Initializer 'onnx::Conv_1575_quantized' references an ORT in-memory address marker, which is not allowed in a model protobuf. ``` This regression is observed on 1.27.0, where both changes are present (1.25.x / 1.26.x have the native-endian tag but not the enforce, so the marker leaked silently without crashing). Both tags encode the buffer address in the `offset` field; on little-endian hosts (the only platform this EP targets) the byte layout is identical, so they are decoded the same way. Recognizing both tags restores the previous behavior of converting these initializers into a lightweight reference inside `model_clone`, so the in-memory marker never reaches `Graph::Graph`. --- .../core/providers/vitisai/imp/tensor_proto.cc | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/onnxruntime/core/providers/vitisai/imp/tensor_proto.cc b/onnxruntime/core/providers/vitisai/imp/tensor_proto.cc index 719ca8dd412bf..b19d9726f6c25 100644 --- a/onnxruntime/core/providers/vitisai/imp/tensor_proto.cc +++ b/onnxruntime/core/providers/vitisai/imp/tensor_proto.cc @@ -31,7 +31,17 @@ gsl::span process_ext_address(const ONNX_NAMESPACE::TensorProto& ten // checksum = (size_t)std::strtoull(data.mutable_value()->data(), &end, 10); } } - if (file == "*/_ORT_MEM_ADDR_/*") { + // ORT marks initializers whose data lives in an in-process memory buffer + // (e.g. weights moved out of the TensorProto during graph optimization) with + // one of two sentinels in the 'location' field. The original little-endian + // tag is still emitted for some paths, but newer ORT defaults to the + // native-endian tag (see TensorToTensorProto / kTensorProtoNativeEndianMemoryAddressTag). + // Both encode the buffer address in the 'offset' field; on little-endian + // hosts (the only platform this EP targets) the bytes are identical, so we + // decode them the same way. Recognizing only the little-endian tag would let + // native-endian-tagged initializers slip through model_clone and trip ORT's + // "in-memory address marker not allowed in a model protobuf" check. + if (file == "*/_ORT_MEM_ADDR_/*" || file == "*/_ORT_NATIVE_ENDIAN_MEM_ADDR_/*") { auto addr = reinterpret_cast(offset); return {addr, size}; } From 8e8190b5c89fca9a561f868c87243dbef278d355 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Tue, 30 Jun 2026 16:59:29 -0700 Subject: [PATCH 06/17] [CUDA] Make cuDNN optional at runtime (#29252) ## Summary Make cuDNN an optional runtime dependency for the CUDA Execution Provider and CUDA Plugin EP. The build still uses cuDNN headers, but provider binaries no longer directly depend on cuDNN shared libraries; cuDNN is loaded lazily when enabled and available, while no-cuDNN runs use native CUDA paths where available and report `NOT_IMPLEMENTED` for kernels that still require cuDNN. This also adds CI to validate no-cuDNN runtime environments. ## Key Changes | Area | Changes | |---|---| | CUDA EP runtime loading | Added a dynamic cuDNN loader and cuDNN symbol trampolines so CUDA provider binaries can avoid a direct cuDNN dependency. | | Provider options | Added `enable_cudnn`; removed `cudnn_path` from CUDA EP and CUDA Plugin EP provider configuration. | | CUDA Plugin EP | Wired optional cuDNN behavior through plugin EP config, kernel adapters, stream handles, and plugin utilities. | | Python preload behavior | Updated Python CUDA preload handling so cuDNN remains an optional dependency instead of an unconditional import/runtime requirement. | | Tests | Added/updated provider option coverage and CUDA Plugin EP no-cuDNN mode using `ORT_TEST_CUDA_PLUGIN_NO_CUDNN=1`. | | CI | Added Linux and Windows CUDA no-cuDNN workflows that build with cuDNN headers, exclude cuDNN from the runtime path, verify no direct cuDNN dependency, and run targeted tests. | | Documentation | Added `docs/CUDA_cuDNN_Optional_Design.md` and updated CUDA Plugin EP docs for no-cuDNN behavior and validation. | ## Testing Validated locally on Linux CUDA 13: - Rebuilt CUDA EP / CUDA Plugin EP with cuDNN headers available at build time. - Verified provider binaries have no direct cuDNN dependency: - `readelf -d ... | grep NEEDED | grep -i cudnn || echo "no cudnn DT_NEEDED"` - `ldd ... | grep -i cudnn || echo "no cudnn in ldd"` - Ran CUDA Plugin EP no-cuDNN validation: - `bash .env/cuda_130_plugin_no_cudnn.sh --test_plugin` - Result: `Ran 87 tests`, `OK (skipped=17)` Additional CI coverage is included for Linux and Windows no-cuDNN CUDA validation. --- .github/workflows/linux_cuda_no_cudnn.yml | 131 ++++ .github/workflows/windows_cuda_no_cudnn.yml | 262 +++++++ cmake/external/cudnn_frontend.cmake | 5 + cmake/onnxruntime_providers_cuda.cmake | 28 +- cmake/onnxruntime_providers_cuda_plugin.cmake | 9 +- cmake/onnxruntime_providers_nv.cmake | 2 + cmake/onnxruntime_providers_tensorrt.cmake | 2 + cmake/onnxruntime_python.cmake | 11 +- .../cudnn_frontend_win_dynamic_loading.patch | 52 ++ docs/CUDA_cuDNN_Optional_Design.md | 711 ++++++++++++++++++ docs/cuda_plugin_ep/QUICK_START.md | 21 + docs/cuda_plugin_ep/cuda_plugin_ep_design.md | 17 +- onnxruntime/__init__.py | 16 +- onnxruntime/core/providers/cuda/cuda_call.cc | 19 + .../providers/cuda/cuda_execution_provider.cc | 21 +- .../providers/cuda/cuda_execution_provider.h | 3 +- .../cuda/cuda_execution_provider_info.cc | 4 + .../cuda/cuda_execution_provider_info.h | 4 + onnxruntime/core/providers/cuda/cuda_kernel.h | 21 +- .../core/providers/cuda/cuda_stream_handle.cc | 15 +- .../core/providers/cuda/cudnn_fe_call.cc | 13 + .../core/providers/cuda/cudnn_loader.cc | 228 ++++++ .../core/providers/cuda/cudnn_loader.h | 47 ++ onnxruntime/core/providers/cuda/cudnn_stub.cc | 401 ++++++++++ .../core/providers/cuda/plugin/cuda_ep.cc | 4 +- .../core/providers/cuda/plugin/cuda_ep.h | 1 + .../providers/cuda/plugin/cuda_ep_factory.cc | 10 +- .../cuda/plugin/cuda_kernel_adapter.h | 63 +- .../providers/cuda/plugin/cuda_plugin_utils.h | 37 +- .../cuda/plugin/cuda_stream_plugin.cc | 14 +- .../cuda/plugin/cuda_stream_plugin.h | 3 +- .../test/python/onnxruntime_test_python.py | 2 + .../transformers/test_cuda_plugin_ep.py | 104 ++- 33 files changed, 2180 insertions(+), 101 deletions(-) create mode 100644 .github/workflows/linux_cuda_no_cudnn.yml create mode 100644 .github/workflows/windows_cuda_no_cudnn.yml create mode 100644 cmake/patches/cudnn_frontend/cudnn_frontend_win_dynamic_loading.patch create mode 100644 docs/CUDA_cuDNN_Optional_Design.md create mode 100644 onnxruntime/core/providers/cuda/cudnn_loader.cc create mode 100644 onnxruntime/core/providers/cuda/cudnn_loader.h create mode 100644 onnxruntime/core/providers/cuda/cudnn_stub.cc diff --git a/.github/workflows/linux_cuda_no_cudnn.yml b/.github/workflows/linux_cuda_no_cudnn.yml new file mode 100644 index 0000000000000..abe92589aa705 --- /dev/null +++ b/.github/workflows/linux_cuda_no_cudnn.yml @@ -0,0 +1,131 @@ +name: Linux CUDA No cuDNN CI + +on: + pull_request: + branches: [main, 'rel-*'] + paths: + - '.github/workflows/linux_cuda_no_cudnn.yml' + - 'cmake/onnxruntime_providers_cuda.cmake' + - 'cmake/onnxruntime_python.cmake' + - 'onnxruntime/__init__.py' + - 'onnxruntime/core/providers/cuda/**' + workflow_dispatch: + +concurrency: + group: ${{ github.workflow }}-${{ github.event_name == 'pull_request' && github.ref || github.sha }} + cancel-in-progress: true + +permissions: + contents: read + packages: write + attestations: write + id-token: write + +jobs: + build-linux-cuda-no-cudnn-x64-release: + name: Build Linux CUDA x64 Release without cuDNN link + uses: ./.github/workflows/reusable_linux_build.yml + with: + pool_name: "onnxruntime-github-Ubuntu2204-AMD-CPU" + build_config: Release + architecture: x64 + dockerfile_path: tools/ci_build/github/linux/docker/Dockerfile.manylinux2_28_cuda + docker_build_args: '--build-arg BASEIMAGE=onnxruntimebuildcache.azurecr.io/internal/azureml/onnxruntime/build/cuda13_x64_almalinux8_gcc14:20251107.1' + docker_image_repo: onnxruntimecuda13manylinuxbuild + extra_build_flags: '--use_binskim_compliant_compile_flags --build_wheel --parallel --nvcc_threads 4 --flash_nvcc_threads 4 --cuda_version=13.0 --cuda_home=/usr/local/cuda-13.0 --cudnn_home=/usr/local/cuda-13.0 --cmake_extra_defines CMAKE_CUDA_ARCHITECTURES=86 onnxruntime_BUILD_UNIT_TESTS=ON' + python_path_prefix: 'PATH=/opt/python/cp312-cp312/bin:$PATH' + run_tests: false + upload_build_output: true + execution_providers: 'cuda' + job_identifier: build-linux-cuda-no-cudnn-x64-release + secrets: + GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + + smoke-linux-cuda-no-cudnn-x64-release: + name: Smoke Linux CUDA x64 Release without cuDNN runtime use + needs: build-linux-cuda-no-cudnn-x64-release + runs-on: + - self-hosted + - "1ES.Pool=onnxruntime-github-linux-a10" + - "1ES.ImageOverride=onnxruntime-ubuntu2204-CUDA-A10-Test" + - "JobId=smoke-linux-cuda-no-cudnn-x64-release-${{ github.run_id }}-${{ github.run_number }}-${{ github.run_attempt }}" + permissions: + contents: read + packages: read + steps: + - name: Checkout code + uses: actions/checkout@v6 + + - uses: microsoft/onnxruntime-github-actions/build-docker-image@8bad63a3c05d448311dfa8e5f531171c97471aa1 + id: build_docker_image_step + with: + dockerfile: ${{ github.workspace }}/tools/ci_build/github/linux/docker/Dockerfile.manylinux2_28_cuda + image-name: ghcr.io/microsoft/onnxruntime/onnxruntimecuda13manylinuxbuild + build-args: '--build-arg BASEIMAGE=onnxruntimebuildcache.azurecr.io/internal/azureml/onnxruntime/build/cuda13_x64_almalinux8_gcc14:20251107.1' + push: true + azure-container-registry-name: onnxruntimebuildcache + env: + GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} + + - name: Download Build Artifact + uses: actions/download-artifact@v7 + with: + name: build-output-x64-Release + path: ${{ runner.temp }}/Release + + - name: Restore Executable Permissions + working-directory: ${{ runner.temp }}/Release + run: | + if [ -f perms.txt ]; then + while IFS= read -r file; do + if [ -f "$file" ]; then + chmod +x "$file" + fi + done < perms.txt + fi + + - name: Verify CUDA provider has no direct cuDNN dependency + run: | + docker run --rm --gpus all \ + -v "${{ runner.temp }}/Release:/build/Release" \ + "${{ steps.build_docker_image_step.outputs.full-image-name }}" \ + bash -lc 'set -euo pipefail + ldd /build/Release/Release/libonnxruntime_providers_cuda.so | tee /tmp/ldd.txt + ! grep -i cudnn /tmp/ldd.txt' + + - name: Run no-cuDNN CUDA EP smoke test + run: | + docker run --rm --gpus all \ + -v "${{ runner.temp }}/Release:/build/Release" \ + "${{ steps.build_docker_image_step.outputs.full-image-name }}" \ + bash -lc 'set -e + PATH=/opt/python/cp312-cp312/bin:$PATH + LD_LIBRARY_PATH=/usr/local/cuda-13.0/lib64:${LD_LIBRARY_PATH:-} + export PATH LD_LIBRARY_PATH + WHEEL_PATH=$(find /build/Release/Release/dist -type f -name "onnxruntime_gpu-*.whl" | head -n 1) + if [ -z "$WHEEL_PATH" ]; then + echo "No built onnxruntime GPU wheel found under /build/Release/Release/dist" >&2 + exit 1 + fi + echo "Installing $WHEEL_PATH" + python -m pip install --no-cache-dir --force-reinstall --no-deps numpy onnx "$WHEEL_PATH" + python - <<"PY" + import numpy as np + import onnx + import onnxruntime as ort + from onnx import TensorProto, helper + + x = helper.make_tensor_value_info("x", TensorProto.FLOAT, [2, 3]) + y = helper.make_tensor_value_info("y", TensorProto.FLOAT, [2, 3]) + node = helper.make_node("Add", ["x", "x"], ["y"]) + graph = helper.make_graph([node], "cuda_no_cudnn_smoke", [x], [y]) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 21)]) + model.ir_version = 10 + + providers = [("CUDAExecutionProvider", {"enable_cudnn": "0"})] + sess = ort.InferenceSession(model.SerializeToString(), providers=providers) + data = np.arange(6, dtype=np.float32).reshape(2, 3) + result = sess.run(None, {"x": data})[0] + np.testing.assert_allclose(result, data + data) + print("CUDA no-cuDNN smoke test passed") + PY' diff --git a/.github/workflows/windows_cuda_no_cudnn.yml b/.github/workflows/windows_cuda_no_cudnn.yml new file mode 100644 index 0000000000000..fe3269416bb73 --- /dev/null +++ b/.github/workflows/windows_cuda_no_cudnn.yml @@ -0,0 +1,262 @@ +name: Windows CUDA No cuDNN CI + +on: + pull_request: + branches: [main, 'rel-*'] + paths: + - '.github/workflows/windows_cuda_no_cudnn.yml' + - 'cmake/onnxruntime_providers_cuda.cmake' + - 'cmake/onnxruntime_providers_cuda_plugin.cmake' + - 'cmake/onnxruntime_python.cmake' + - 'docs/CUDA_cuDNN_Optional_Design.md' + - 'docs/cuda_plugin_ep/**' + - 'onnxruntime/__init__.py' + - 'onnxruntime/core/providers/cuda/**' + - 'onnxruntime/test/python/transformers/test_cuda_plugin_ep.py' + workflow_dispatch: + +concurrency: + group: ${{ github.workflow }}-${{ github.event_name == 'pull_request' && github.ref || github.sha }} + cancel-in-progress: true + +jobs: + build: + name: Windows CUDA Plugin EP Build without cuDNN + runs-on: [ + "self-hosted", + "1ES.Pool=onnxruntime-github-vs2022-latest", + "JobId=windows-cuda-plugin-no-cudnn-${{ github.run_id }}-${{ github.run_number }}-${{ github.run_attempt }}" + ] + steps: + - uses: actions/checkout@v6 + with: + fetch-depth: 0 + submodules: 'none' + + - uses: actions/setup-python@v6 + with: + python-version: '3.12' + architecture: x64 + + - name: Locate vcvarsall and Setup Env + uses: ./.github/actions/locate-vcvarsall-and-setup-env + with: + architecture: x64 + + - name: Install python modules + run: python -m pip install -r .\tools\ci_build\github\windows\python\requirements.txt + working-directory: ${{ github.workspace }} + shell: cmd + + - name: Download CUDA SDK v13.0 + working-directory: ${{ runner.temp }} + run: | + azcopy.exe cp --recursive "https://lotusscus.blob.core.windows.net/models/cuda_sdk/v13.0" . + dir + shell: pwsh + + - name: Prepare cuDNN SDK without runtime DLLs + working-directory: ${{ runner.temp }} + run: | + $cudnnSdkUri = "https://lotusscus.blob.core.windows.net/models/cudnn_sdk/$env:CUDNN_FOLDER" + azcopy.exe cp --recursive $cudnnSdkUri . + $cudnnRoot = Join-Path $env:RUNNER_TEMP $env:CUDNN_FOLDER + if (-not (Test-Path $cudnnRoot)) { + Write-Error "cuDNN SDK was not downloaded to the expected folder: $cudnnRoot" + exit 1 + } + Get-ChildItem -Path $cudnnRoot -Recurse -Include "cudnn*.dll" | Remove-Item -Force + if (Get-ChildItem -Path $cudnnRoot -Recurse -Include "cudnn*.dll" -ErrorAction SilentlyContinue) { + Write-Error "cuDNN runtime DLLs must not be present in the no-cuDNN build environment" + exit 1 + } + shell: pwsh + + - name: Add CUDA to PATH + shell: pwsh + run: | + Write-Host "Adding CUDA to PATH without adding any cuDNN directory" + Add-Content -Path $env:GITHUB_PATH -Value "$env:RUNNER_TEMP\v13.0\bin" + Add-Content -Path $env:GITHUB_PATH -Value "$env:RUNNER_TEMP\v13.0\bin\x64" + Add-Content -Path $env:GITHUB_PATH -Value "$env:RUNNER_TEMP\v13.0\extras\CUPTI\lib64" + + - name: Install CUDA Visual Studio integration + shell: pwsh + run: | + $sourceDir = "$env:RUNNER_TEMP\v13.0\extras\visual_studio_integration\MSBuildExtensions" + $targetDir = "${env:ProgramFiles}\Microsoft Visual Studio\2022\Enterprise\MSBuild\Microsoft\VC\v170\BuildCustomizations" + Copy-Item -Path "$sourceDir\*" -Destination $targetDir -Force + + - name: Set OnnxRuntimeBuildDirectory + shell: pwsh + run: | + $buildDir = Join-Path $env:RUNNER_TEMP "build" + echo "OnnxRuntimeBuildDirectory=$buildDir" >> $env:GITHUB_ENV + + - name: Build ONNX Runtime with CUDA Plugin EP and no cuDNN runtime path + working-directory: ${{ runner.temp }} + shell: pwsh + run: | + python.exe ${{ github.workspace }}\tools\ci_build\build.py ` + --update --build --config Release ` + --build_dir build ` + --skip_submodule_sync ` + --parallel ` + --nvcc_threads 4 ` + --flash_nvcc_threads 4 ` + --use_binskim_compliant_compile_flags ` + --cmake_generator "Visual Studio 17 2022" ` + --build_shared_lib ` + --build_wheel ` + --use_cuda ` + --cuda_home="$env:RUNNER_TEMP\v13.0" ` + --cudnn_home="$env:RUNNER_TEMP\$env:CUDNN_FOLDER" ` + --skip_tests ` + --use_vcpkg ` + --use_vcpkg_ms_internal_asset_cache ` + --enable_cuda_profiling ` + --cmake_extra_defines onnxruntime_QUICK_BUILD=ON ` + --cmake_extra_defines CMAKE_CUDA_ARCHITECTURES=86 ` + --cmake_extra_defines onnxruntime_BUILD_CUDA_EP_AS_PLUGIN=ON + + if ($lastExitCode -ne 0) { + exit $lastExitCode + } + + $outputDir = "${{ runner.temp }}\build\Release" + Write-Host "Cleaning up files from $outputDir..." + Remove-Item -Path "$outputDir\onnxruntime" -Recurse -Force -ErrorAction SilentlyContinue + Remove-Item -Path "$outputDir\pybind11" -Recurse -Force -ErrorAction SilentlyContinue + Remove-Item -Path "$outputDir\models" -Recurse -Force -ErrorAction SilentlyContinue + Remove-Item -Path "$outputDir\vcpkg_installed" -Recurse -Force -ErrorAction SilentlyContinue + Remove-Item -Path "$outputDir\_deps" -Recurse -Force -ErrorAction SilentlyContinue + Remove-Item -Path "$outputDir\CMakeCache.txt" -Force -ErrorAction SilentlyContinue + Remove-Item -Path "$outputDir\CMakeFiles" -Recurse -Force -ErrorAction SilentlyContinue + Remove-Item -Path $outputDir -Include "*.obj" -Recurse + + $cudnnArtifacts = Get-ChildItem -Path $outputDir -Recurse -Include "cudnn*.dll" -ErrorAction SilentlyContinue + if ($cudnnArtifacts) { + $cudnnArtifacts | ForEach-Object { Write-Host $_.FullName } + Write-Error "cuDNN runtime DLLs must not be present in no-cuDNN build artifacts" + exit 1 + } + + - name: Upload build artifacts + uses: actions/upload-artifact@v6 + with: + name: cuda-plugin-no-cudnn-build-artifacts + path: ${{ runner.temp }}\build + env: + DOTNET_SKIP_FIRST_TIME_EXPERIENCE: true + setVcvars: true + ALLOW_RELEASED_ONNX_OPSET_ONLY: '0' + ONNXRUNTIME_TEST_GPU_DEVICE_ID: '0' + AZCOPY_AUTO_LOGIN_TYPE: MSI + AZCOPY_MSI_CLIENT_ID: d712a4c7-a0cf-4e87-af75-31510eba0a8e + CUDNN_FOLDER: 9.14.0.64_cuda13 + + test: + name: Windows CUDA Plugin EP Test without cuDNN + needs: build + timeout-minutes: 120 + runs-on: [ + "self-hosted", + "1ES.Pool=onnxruntime-github-Win2022-GPU-A10", + "JobId=windows-cuda-plugin-no-cudnn-test-${{ github.run_id }}-${{ github.run_number }}-${{ github.run_attempt }}" + ] + steps: + - uses: actions/checkout@v6 + with: + fetch-depth: 0 + submodules: 'none' + + - name: Download build artifacts + uses: actions/download-artifact@v7 + with: + name: cuda-plugin-no-cudnn-build-artifacts + path: ${{ runner.temp }}\build + + - uses: actions/setup-python@v6 + with: + python-version: '3.12' + architecture: x64 + + - name: Locate vcvarsall and Setup Env + uses: ./.github/actions/locate-vcvarsall-and-setup-env + with: + architecture: x64 + + - name: Install python modules + run: python -m pip install -r .\tools\ci_build\github\windows\python\requirements.txt + working-directory: ${{ github.workspace }} + shell: cmd + + - name: Install torch for CPU only + run: python -m pip install torch + working-directory: ${{ github.workspace }} + shell: cmd + + - name: Download CUDA SDK v13.0 + working-directory: ${{ runner.temp }} + run: | + azcopy.exe cp --recursive "https://lotusscus.blob.core.windows.net/models/cuda_sdk/v13.0" . + dir + shell: pwsh + + - name: Add CUDA to PATH + shell: pwsh + run: | + Write-Host "Adding CUDA to PATH without adding any cuDNN directory" + Add-Content -Path $env:GITHUB_PATH -Value "$env:RUNNER_TEMP\v13.0\bin" + Add-Content -Path $env:GITHUB_PATH -Value "$env:RUNNER_TEMP\v13.0\bin\x64" + Add-Content -Path $env:GITHUB_PATH -Value "$env:RUNNER_TEMP\v13.0\extras\CUPTI\lib64" + + - name: Set OnnxRuntimeBuildDirectory + shell: pwsh + run: | + $buildDir = Join-Path $env:RUNNER_TEMP "build" + echo "OnnxRuntimeBuildDirectory=$buildDir" >> $env:GITHUB_ENV + + - name: Install ONNX Runtime Wheel + uses: ./.github/actions/install-onnxruntime-wheel + with: + whl-directory: ${{ runner.temp }}\build\Release\Release\dist + + - name: Verify GPU access + shell: pwsh + run: nvidia-smi + + - name: Verify CUDA plugin has no direct cuDNN dependency + shell: pwsh + run: | + $pluginPath = "${{ runner.temp }}\build\Release\Release\onnxruntime_providers_cuda_plugin.dll" + if (-not (Test-Path $pluginPath)) { + Write-Error "CUDA plugin EP library not found at $pluginPath" + exit 1 + } + + dumpbin /dependents $pluginPath | Tee-Object -FilePath $env:RUNNER_TEMP\cuda_plugin_dependents.txt + if (Select-String -Path $env:RUNNER_TEMP\cuda_plugin_dependents.txt -Pattern "cudnn" -SimpleMatch -Quiet) { + Write-Error "CUDA plugin EP has a direct cuDNN dependency" + exit 1 + } + + - name: Run CUDA Plugin EP Python Tests without cuDNN + working-directory: ${{ github.workspace }}\onnxruntime\test\python\transformers + shell: pwsh + run: | + $env:ORT_CUDA_PLUGIN_PATH = "${{ runner.temp }}\build\Release\Release\onnxruntime_providers_cuda_plugin.dll" + $env:ORT_TEST_CUDA_PLUGIN_EP = "1" + $env:ORT_TEST_CUDA_PLUGIN_NO_CUDNN = "1" + Write-Host "ORT_CUDA_PLUGIN_PATH=$env:ORT_CUDA_PLUGIN_PATH" + python test_cuda_plugin_ep.py + if ($lastExitCode -ne 0) { + exit $lastExitCode + } + env: + DOTNET_SKIP_FIRST_TIME_EXPERIENCE: true + setVcvars: true + ALLOW_RELEASED_ONNX_OPSET_ONLY: '0' + ONNXRUNTIME_TEST_GPU_DEVICE_ID: '0' + AZCOPY_AUTO_LOGIN_TYPE: MSI + AZCOPY_MSI_CLIENT_ID: d712a4c7-a0cf-4e87-af75-31510eba0a8e diff --git a/cmake/external/cudnn_frontend.cmake b/cmake/external/cudnn_frontend.cmake index 9af915a272b07..e8a84f99354f0 100644 --- a/cmake/external/cudnn_frontend.cmake +++ b/cmake/external/cudnn_frontend.cmake @@ -1,8 +1,13 @@ +# The cudnn_frontend shim only loads the CUDA runtime / driver libraries with +# their Linux names (libcudart.so.*, libcuda.so.1) when NV_CUDNN_FRONTEND_USE_DYNAMIC_LOADING +# is enabled. On Windows this always fails (dlopen is mapped to LoadLibrary), so patch the +# shim to use the Windows DLL names (cudart64_*.dll, nvcuda.dll). onnxruntime_fetchcontent_declare( cudnn_frontend URL ${DEP_URL_cudnn_frontend} URL_HASH SHA1=${DEP_SHA1_cudnn_frontend} + PATCH_COMMAND ${Patch_EXECUTABLE} --binary --ignore-whitespace -p1 < ${PROJECT_SOURCE_DIR}/patches/cudnn_frontend/cudnn_frontend_win_dynamic_loading.patch EXCLUDE_FROM_ALL ) diff --git a/cmake/onnxruntime_providers_cuda.cmake b/cmake/onnxruntime_providers_cuda.cmake index 2aa31276cc395..3a480033e77af 100644 --- a/cmake/onnxruntime_providers_cuda.cmake +++ b/cmake/onnxruntime_providers_cuda.cmake @@ -399,7 +399,9 @@ message( WARNING "To compile with NHWC ops enabled please compile against cuDNN 9 or newer." ) endif() endif() - target_link_libraries(${target} PRIVATE CUDA::cublasLt CUDA::cublas CUDNN::cudnn_all cudnn_frontend CUDA::curand CUDA::cufft CUDA::cudart CUDA::nvrtc CUDA::cuda_driver + target_compile_definitions(${target} PRIVATE NV_CUDNN_FRONTEND_USE_DYNAMIC_LOADING) + target_include_directories(${target} PRIVATE ${CUDNN_INCLUDE_DIR}) + target_link_libraries(${target} PRIVATE CUDA::cublasLt CUDA::cublas cudnn_frontend CUDA::curand CUDA::cufft CUDA::cudart CUDA::nvrtc CUDA::cuda_driver ${ABSEIL_LIBS} ${ONNXRUNTIME_PROVIDERS_SHARED} Boost::mp11 safeint_interface) endif() @@ -604,6 +606,30 @@ ) set_target_properties(onnxruntime_providers_cuda PROPERTIES PUBLIC_HEADER "${ONNXRUNTIME_CUDA_PROVIDER_PUBLIC_HEADERS}") + if(WIN32 AND NOT onnxruntime_CUDA_MINIMAL) + set(ORT_CUDNN_DLL_PATH "") + if(onnxruntime_CUDNN_HOME) + set(ORT_CUDNN_DLL_SEARCH_PATHS + "${onnxruntime_CUDNN_HOME}/bin/cudnn64_*.dll" + "${onnxruntime_CUDNN_HOME}/bin/x64/cudnn64_*.dll" + "${onnxruntime_CUDNN_HOME}/bin/${onnxruntime_CUDA_VERSION}/cudnn64_*.dll" + "${onnxruntime_CUDNN_HOME}/bin/${onnxruntime_CUDA_VERSION}/x64/cudnn64_*.dll" + ) + else() + set(ORT_CUDNN_DLL_SEARCH_PATHS "${onnxruntime_CUDA_HOME}/bin/cudnn64_*.dll") + endif() + foreach(search_path ${ORT_CUDNN_DLL_SEARCH_PATHS}) + file(GLOB ORT_CUDNN_DLL_PATH "${search_path}") + if(ORT_CUDNN_DLL_PATH) + break() + endif() + endforeach() + if(ORT_CUDNN_DLL_PATH) + add_custom_command(TARGET onnxruntime_providers_cuda POST_BUILD + COMMAND ${CMAKE_COMMAND} -E copy_if_different ${ORT_CUDNN_DLL_PATH} $ + ) + endif() + endif() install(TARGETS onnxruntime_providers_cuda PUBLIC_HEADER DESTINATION ${CMAKE_INSTALL_INCLUDEDIR}/onnxruntime/core/providers/cuda ARCHIVE DESTINATION ${CMAKE_INSTALL_LIBDIR} diff --git a/cmake/onnxruntime_providers_cuda_plugin.cmake b/cmake/onnxruntime_providers_cuda_plugin.cmake index 86e5579eb6761..f2c53aab31b39 100644 --- a/cmake/onnxruntime_providers_cuda_plugin.cmake +++ b/cmake/onnxruntime_providers_cuda_plugin.cmake @@ -348,12 +348,12 @@ endif() set(CUDA_PLUGIN_CUDNN_INCLUDE_DIR ${CUDNN_INCLUDE_DIR}) set(CUDA_PLUGIN_CUDNN_LIBRARY ${cudnn_LIBRARY}) -if(NOT CUDA_PLUGIN_CUDNN_INCLUDE_DIR OR NOT CUDA_PLUGIN_CUDNN_LIBRARY) - message(FATAL_ERROR "cuDNN not found (from main ORT search) for CUDA Plugin EP.") +if(NOT CUDA_PLUGIN_CUDNN_INCLUDE_DIR) + message(FATAL_ERROR "cuDNN headers not found (from main ORT search) for CUDA Plugin EP.") endif() message(STATUS "CUDA Plugin EP: cuDNN include: ${CUDA_PLUGIN_CUDNN_INCLUDE_DIR}") -message(STATUS "CUDA Plugin EP: cuDNN library: ${CUDA_PLUGIN_CUDNN_LIBRARY}") +message(STATUS "CUDA Plugin EP: cuDNN runtime library: ${CUDA_PLUGIN_CUDNN_LIBRARY}") # Include directories — only public ORT headers + CUDA toolkit + cuDNN + internal headers for adapter target_include_directories(onnxruntime_providers_cuda_plugin PRIVATE @@ -388,7 +388,6 @@ target_link_libraries(onnxruntime_providers_cuda_plugin PRIVATE CUDA::cufft CUDA::nvrtc CUDA::cuda_driver - CUDNN::cudnn_all cudnn_frontend Boost::mp11 safeint_interface @@ -403,6 +402,8 @@ target_link_libraries(onnxruntime_providers_cuda_plugin PRIVATE ${PROTOBUF_LIB} ) + target_compile_definitions(onnxruntime_providers_cuda_plugin PRIVATE NV_CUDNN_FRONTEND_USE_DYNAMIC_LOADING) + if (onnxruntime_ENABLE_CUDA_PROFILING) target_link_libraries(onnxruntime_providers_cuda_plugin PRIVATE CUDA::cupti) target_compile_definitions(onnxruntime_providers_cuda_plugin PRIVATE ENABLE_CUDA_PROFILING) diff --git a/cmake/onnxruntime_providers_nv.cmake b/cmake/onnxruntime_providers_nv.cmake index 92fc0ccadd667..ffb6fd236cf48 100644 --- a/cmake/onnxruntime_providers_nv.cmake +++ b/cmake/onnxruntime_providers_nv.cmake @@ -140,6 +140,8 @@ endif () "${ONNXRUNTIME_ROOT}/core/providers/shared_library/*.cc" "${ONNXRUNTIME_ROOT}/core/providers/cuda/cuda_stream_handle.h" "${ONNXRUNTIME_ROOT}/core/providers/cuda/cuda_stream_handle.cc" + "${ONNXRUNTIME_ROOT}/core/providers/cuda/cudnn_loader.h" + "${ONNXRUNTIME_ROOT}/core/providers/cuda/cudnn_loader.cc" "${ONNXRUNTIME_ROOT}/core/providers/cuda/cuda_graph.h" "${ONNXRUNTIME_ROOT}/core/providers/cuda/cuda_graph.cc" ) diff --git a/cmake/onnxruntime_providers_tensorrt.cmake b/cmake/onnxruntime_providers_tensorrt.cmake index 4184e0b049afc..073e26728c29e 100644 --- a/cmake/onnxruntime_providers_tensorrt.cmake +++ b/cmake/onnxruntime_providers_tensorrt.cmake @@ -161,6 +161,8 @@ "${ONNXRUNTIME_ROOT}/core/providers/shared_library/*.cc" "${ONNXRUNTIME_ROOT}/core/providers/cuda/cuda_stream_handle.h" "${ONNXRUNTIME_ROOT}/core/providers/cuda/cuda_stream_handle.cc" + "${ONNXRUNTIME_ROOT}/core/providers/cuda/cudnn_loader.h" + "${ONNXRUNTIME_ROOT}/core/providers/cuda/cudnn_loader.cc" "${ONNXRUNTIME_ROOT}/core/providers/cuda/cuda_graph.h" "${ONNXRUNTIME_ROOT}/core/providers/cuda/cuda_graph.cc" ) diff --git a/cmake/onnxruntime_python.cmake b/cmake/onnxruntime_python.cmake index de1d7559a1572..2b30ef0f239d5 100644 --- a/cmake/onnxruntime_python.cmake +++ b/cmake/onnxruntime_python.cmake @@ -339,16 +339,19 @@ if (WIN32) endif() endforeach() if(NOT CUDNN_DLL_PATH) - message(FATAL_ERROR "cuDNN not found in ${onnxruntime_CUDNN_HOME}") + message(STATUS "cuDNN not found in ${onnxruntime_CUDNN_HOME}. Python package metadata will record an empty cuDNN version.") endif() else() file(GLOB CUDNN_DLL_PATH "${onnxruntime_CUDA_HOME}/bin/cudnn64_*.dll") if (NOT CUDNN_DLL_PATH) - message(FATAL_ERROR "cuDNN not found in ${onnxruntime_CUDA_HOME}") + message(STATUS "cuDNN not found in ${onnxruntime_CUDA_HOME}. Python package metadata will record an empty cuDNN version.") endif() endif() - get_filename_component(CUDNN_DLL_NAME ${CUDNN_DLL_PATH} NAME_WE) - string(REPLACE "cudnn64_" "" CUDNN_VERSION "${CUDNN_DLL_NAME}") + set(CUDNN_VERSION "") + if(CUDNN_DLL_PATH) + get_filename_component(CUDNN_DLL_NAME ${CUDNN_DLL_PATH} NAME_WE) + string(REPLACE "cudnn64_" "" CUDNN_VERSION "${CUDNN_DLL_NAME}") + endif() if(NOT onnxruntime_CUDA_VERSION) set(onnxruntime_CUDA_VERSION ${CUDAToolkit_VERSION}) message("onnxruntime_CUDA_VERSION=${onnxruntime_CUDA_VERSION}") diff --git a/cmake/patches/cudnn_frontend/cudnn_frontend_win_dynamic_loading.patch b/cmake/patches/cudnn_frontend/cudnn_frontend_win_dynamic_loading.patch new file mode 100644 index 0000000000000..96242982a7468 --- /dev/null +++ b/cmake/patches/cudnn_frontend/cudnn_frontend_win_dynamic_loading.patch @@ -0,0 +1,52 @@ +--- a/include/cudnn_frontend_shim.h 2026-06-26 15:06:33.854043748 -0700 ++++ b/include/cudnn_frontend_shim.h 2026-06-26 15:06:33.966044710 -0700 +@@ -69,7 +69,14 @@ + dlerror(); + + // Attempt to open the cuda library ++#ifdef _WIN32 ++ // Reset the thread error state so the GetLastError()-based success check ++ // below is not tripped by a stale error from an earlier call. ++ SetLastError(0); ++ HMODULE handle = dlopen("nvcuda.dll", RTLD_NOW); ++#else + HMODULE handle = dlopen("libcuda.so.1", RTLD_NOW); ++#endif + const char *error = reinterpret_cast(dlerror()); + if (!handle || error) { + // If opening the library fails, throw an exception with the error message +@@ -85,13 +92,22 @@ + dlerror(); + + // List of potential libcudart libraries (Adding major version to support python package) ++#ifdef _WIN32 ++ constexpr const char *libs[] = {"cudart64_12.dll", "cudart64_13.dll"}; ++#else + constexpr const char *libs[] = {"libcudart.so.12", "libcudart.so.13"}; ++#endif + constexpr size_t num_libs = sizeof(libs) / sizeof(libs[0]); + + HMODULE lib_handle = nullptr; + int loaded_index = -1; + + for (size_t i = 0; i < num_libs; ++i) { ++#ifdef _WIN32 ++ // Reset the thread error state so the GetLastError()-based success check ++ // below is not tripped by a stale error from a previous failed attempt. ++ SetLastError(0); ++#endif + HMODULE handle = dlopen(libs[i], RTLD_NOW); + const char *error = reinterpret_cast(dlerror()); + +@@ -182,10 +200,10 @@ + #if defined NV_CUDNN_FRONTEND_USE_DYNAMIC_LOADING + + #define NV_FE_CALL_TO_CUDA(function_name, cuda_symbol, ...) \ +- return reinterpret_cast(get_cuda_symbol(CudaLibrary::CUDART, #cuda_symbol))(__VA_ARGS__); ++ return cuda_symbol(__VA_ARGS__); + #define NV_FE_CALL_TO_CU(function_name, cuda_symbol, ...) \ +- return reinterpret_cast(get_cuda_symbol(CudaLibrary::CUDA, #cuda_symbol))(__VA_ARGS__); ++ return cuda_symbol(__VA_ARGS__); + + #else + diff --git a/docs/CUDA_cuDNN_Optional_Design.md b/docs/CUDA_cuDNN_Optional_Design.md new file mode 100644 index 0000000000000..8d24b8db30186 --- /dev/null +++ b/docs/CUDA_cuDNN_Optional_Design.md @@ -0,0 +1,711 @@ +# Making cuDNN Optional for the CUDA Execution Provider + +Status: Draft / Proposal +Owner: (CUDA EP) +Scope: `onnxruntime/core/providers/cuda` (main static CUDA EP), CUDA Plugin EP +(`BUILD_CUDA_EP_AS_PLUGIN`), and the CUDA unit tests. **Out of scope:** TensorRT EP and +NV‑TensorRT‑RTX EP (they create and own their own cuDNN handles and inherently depend on +cuDNN/TensorRT). The existing build-time `USE_CUDA_MINIMAL` path is also out of scope; it is +used by TensorRT / NV‑TensorRT‑RTX integration and should remain available. + +--- + +## 1. Motivation + +Today the CUDA EP has a **hard link‑time and load‑time dependency** on cuDNN +(`libcudnn*.so` / `cudnn64_*.dll`) and on the header‑only `cudnn_frontend` library. If the +cuDNN shared libraries are not present on the machine, the ORT shared library +(`libonnxruntime.so` / `onnxruntime.dll`, or the CUDA provider DLL) **fails to load at all** — +even for models that use no cuDNN‑backed operators. + +cuDNN is large (hundreds of MB across its sub‑libraries) and is only needed by a subset of +operators (Conv, Pooling, BatchNorm, LRN, RNN/LSTM/GRU, the cuDNN reduction path, the cuDNN +softmax path, etc.). Many transformer / LLM models do not require any of these. + +**Goal:** Ship a *single* CUDA EP binary that + +1. loads and runs **without** cuDNN present, and +2. **lazily** attempts to load cuDNN at first use; if cuDNN is available, cuDNN‑backed + operators work exactly as today; if it is absent, those operators fail with a clear, + actionable error (Phase 1) and, incrementally, fall back to native CUDA kernels + (Phases 2 and 3). + +This is delivered in **three phases**: + +- **Phase 1 — Make cuDNN optional.** No new compute kernels. When cuDNN is missing, any op + that needs it throws a clear `NOT_IMPLEMENTED` ("cuDNN is required for operator X but was + not found") error at `Run` time. Everything that does not need cuDNN runs normally. +- **Phase 2 — Remove LLM‑relevant cuDNN dependencies.** Replace the cuDNN Softmax / + LogSoftmax and reduction paths with existing native CUDA / CUB paths. Phase 1 + Phase 2 is + the first milestone because LLM models may use these ops. +- **Phase 3 — Replace the remaining cuDNN‑backed NN ops.** Add native CUDA / CUTLASS / + Triton‑cubin fallbacks for Pooling, normalization, LRN, Conv, ConvTranspose, and FusedConv. + +--- + +## 2. Current state (as of this writing) + +### 2.1 How cuDNN is linked + +- `cmake/onnxruntime_providers_cuda.cmake`: + `target_link_libraries(... CUDNN::cudnn_all cudnn_frontend ...)` — a normal link dependency, + resolved at process load time. +- `cmake/onnxruntime_providers_cuda_plugin.cmake`: same (`CUDNN::cudnn_all`, `cudnn_frontend`). +- `cmake/onnxruntime_unittests.cmake`: links `cudnn_frontend` and includes `CUDNN_INCLUDE_DIR`. +- `cmake/deps.txt`: pins `cudnn_frontend` v1.x (header‑only C++ wrapper over the cuDNN v9 + *backend* API). +- `cmake/external/cudnn_frontend.cmake` fetches `cudnn_frontend`, sets `CUDNN_PATH` from + `onnxruntime_CUDNN_HOME`, disables its samples/tests/python bindings, and marks its headers + as system includes. In ORT this is a **compile-time header dependency**, not a runtime DLL. +- There is **no delay‑load** configured for cuDNN today (delay‑load is only used for DML / + WebGPU / a few Win32 API sets). + +### 2.2 How the handle is created and reached + +- The cuDNN handle is created **eagerly**: + - `CudaStream` constructor — `cudnnCreate(&cudnn_handle_)` / `cudnnSetStream(...)` + (`onnxruntime/core/providers/cuda/cuda_stream_handle.cc`). + - `CUDAExecutionProvider::PerThreadContext` — `cudnnCreate(&cudnn_handle_)` + (`onnxruntime/core/providers/cuda/cuda_execution_provider.cc`). + - CUDA Plugin EP — `cuda_stream_plugin.cc` and `cuda_kernel_adapter.h`. +- Kernels obtain the handle through: + - `CudaKernel::GetCudnnHandle(context)` → `stream->cudnn_handle_` + (`onnxruntime/core/providers/cuda/cuda_kernel.h`). + - `CUDAExecutionProvider::PerThreadDefaultCudnnHandle()`. + - The public `CudaContext` resource API (`cuda_context.h`, + `CudaResource::cudnn_handle_t`) — used by custom ops. + +### 2.3 Call sites and macros + +- `CUDNN_CALL` / `CUDNN_CALL_THROW` (in `shared_inc/cuda_call.h`) and + `CUDNN_RETURN_IF_ERROR` / `CUDNN2_RETURN_IF_ERROR` (in `cuda_common.h`). +- `CUDNN_FE_CALL` / `CUDNN_FE_RETURN_IF_ERROR` for the frontend (`cudnn_fe_call.*`). +- `CudaErrString` calls `cudnnGetErrorString` (`cuda_call.cc`). +- All of the above are already gated by `#ifndef USE_CUDA_MINIMAL` in shared CUDA code. That + build‑time path is used by TensorRT / NV‑TensorRT‑RTX related builds and is a useful + inventory of cuDNN‑touching code, but it is *not* the CUDA EP runtime behavior we want. + +### 2.4 Operators / components that depend on cuDNN + +| Area | Files | cuDNN usage | +|---|---|---| +| Conv / ConvTranspose (v9 graph) | `nn/conv.cc`, `nn/conv_transpose.cc` | `cudnn_frontend` graph API + `cudnnAddTensor` | +| Conv / ConvTranspose (legacy) | `nn/conv_8.h`, `nn/conv_transpose_8.h` | `cudnnConvolutionForward`, `cudnnConvolutionBackwardData`, algo search | +| FusedConv (contrib) | `contrib_ops/cuda/fused_conv.cc` | `cudnnConvolutionBiasActivationForward`, activation desc | +| Pooling | `nn/pool.cc` | `cudnnPoolingForward`, pooling desc | +| BatchNormalization | `nn/batch_norm.cc` | `cudnnBatchNormalizationForwardInference/Training` | +| InstanceNormalization | `nn/instance_norm.cc` | BatchNorm training helper | +| LRN | `nn/lrn.cc` | `cudnnLRNCrossChannelForward` | +| RNN / LSTM / GRU | `rnn/cudnn_rnn_base.*`, `rnn/{rnn,lstm,gru}.h` | `cudnnRNNForward`, RNN/dropout descriptors | +| Reductions | `reduction/reduction_ops.*` | `cudnnReduceTensor` (Reduce\*, ArgMax/ArgMin) | +| Softmax (cuDNN path) | `math/softmax_common.cc` | `cudnnSoftmaxForward/Backward` | +| Einsum | `math/einsum_utils/*` | passes cuDNN handle into helpers | +| Dropout descriptor | `cudnn_common.h` (`CudnnDropout`) | used by RNN | +| Tensor/Filter descriptors | `cudnn_common.*` | `cudnnCreate*Descriptor`, `cudnnSet*Descriptor` | +| Contrib attention (optional) | `contrib_ops/cuda/bert/group_query_attention.cc`, `quantization/attention_quantization.cc`, `math/bias_softmax.cc` | optional cuDNN flash attention / handle passthrough | + +> Note: Several of these ops already have, or can trivially get, a **non‑cuDNN** path +> (e.g. Softmax has a native warp/block kernel; reductions have CUB‑based paths; pooling and +> simple elementwise norms are straightforward). Softmax and reductions are Phase‑2 +> candidates; the remaining NN ops are Phase‑3 candidates. + +### 2.5 `cudnn_frontend` usage + +`cudnn_frontend` is currently used as a header-only C++ graph API wrapper around cuDNN's +backend API: + +- `onnxruntime/core/providers/cuda/nn/conv.h` includes `` and stores + `cudnn_frontend::graph::Graph`, `Tensor_attributes`, `Pointwise_attributes`, and variant + packs in `CudnnConvState`. +- `onnxruntime/core/providers/cuda/nn/conv.cc` builds cuDNN frontend graphs for v9 Conv, + optional bias fusion, optional activation fusion, heuristic selection, support checks, plan + building, workspace sizing, and graph execution. +- `onnxruntime/core/providers/cuda/nn/conv_transpose.cc` does the same for ConvTranspose + using `Conv_dgrad_attributes`. +- `onnxruntime/core/providers/cuda/cudnn_common.{h,cc}` defines `CudnnFeTensor`, a small ORT + helper that maps ORT tensor shapes/types into `cudnn_frontend::graph::Tensor_attributes`. +- `onnxruntime/core/providers/cuda/shared_inc/cudnn_fe_call.h` and + `onnxruntime/core/providers/cuda/cudnn_fe_call.cc` adapt `cudnn_frontend::error_t` into + ORT's `CudaCall` error-handling path. +- `onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc` has a cuDNN SDPA feature path + selected through kernel options, but it does not directly include the high-level frontend + graph headers in the same way Conv/ConvTranspose do. + +Important frontend detail: `cudnn_frontend` already has a dynamic-loading mode gated by +`NV_CUDNN_FRONTEND_USE_DYNAMIC_LOADING`. In that mode, `cudnn_frontend_shim.h` does **not** +link directly against `cudnnBackend*` symbols. Instead, it expects the embedding library to +define `cudnn_frontend::cudnn_dlhandle`, then resolves symbols from that handle with +`dlsym` / `GetProcAddress`. + +The frontend backend-symbol surface used by ORT's current graph path includes at least: + +- `cudnnGetVersion`, `cudnnGetErrorString`, and for cuDNN 9, `cudnnGetLastErrorString`. +- `cudnnBackendCreateDescriptor`, `cudnnBackendDestroyDescriptor`, + `cudnnBackendSetAttribute`, `cudnnBackendGetAttribute`, and `cudnnBackendFinalize`. +- `cudnnBackendExecute`. +- Version-gated helpers such as `cudnnBackendPopulateCudaGraph`, + `cudnnBackendUpdateCudaGraph`, and `cudnnGetExecutionPlanWorkspaceSize` if ORT starts using + frontend features that require them. + +Therefore, `cudnn_frontend` should stay in the build while ORT still has cuDNN frontend +Conv/ConvTranspose paths, but it should be compiled in dynamic-loading mode and wired to the +same ORT-owned cuDNN loader used by direct cuDNN calls. + +--- + +## 3. Design overview + +The core idea: **break the hard dependency by routing every cuDNN symbol through a thin, +lazily‑resolved trampoline layer**, plus an availability flag the EP and kernels consult. + +```mermaid +flowchart TD + K[CUDA kernel
e.g. Conv, Pool] -->|cudnnXxx(...)| S[cuDNN shim
trampolines] + FE[cudnn_frontend
header-only] -->|dynamic mode dlsym(cudnn_dlhandle)| L + S -->|first call: dlopen/LoadLibrary| L[cuDNN loader] + L -->|present| R[(real libcudnn*)] + L -->|absent| U[mark unavailable
return sentinel status] + K -.->|enable_cudnn && IsCudnnAvailable()?| L +``` + +### 3.1 The shim (no hard link) + +We **stop linking** `CUDNN::cudnn_all`. In its place we compile a generated translation unit, +`cudnn_stub.cc`, that **defines every direct cuDNN entry point ORT references**. Each +definition is a trampoline: + +```cpp +// Pseudocode for one entry +cudnnStatus_t cudnnConvolutionForward(cudnnHandle_t h, /* ... */) { + auto fn = CudnnLibrary::Get().convolution_forward; // resolved lazily + if (fn == nullptr) return CUDNN_STATUS_NOT_INITIALIZED; // cuDNN unavailable + return fn(h, /* ... */); +} +``` + +Because the trampolines have the **exact** cuDNN symbol names and signatures, ORT's direct +calls link against *our* definitions — so the final binary has **no `NEEDED`/import entry for +libcudnn** from those calls. + +`cudnn_frontend` is handled separately: compile it with +`NV_CUDNN_FRONTEND_USE_DYNAMIC_LOADING` and define `cudnn_frontend::cudnn_dlhandle` in ORT. +When the ORT loader successfully opens cuDNN, it sets that handle to the loaded cuDNN library +handle. `cudnn_frontend` then resolves `cudnnBackend*` symbols from the same handle. When +`enable_cudnn=0` or cuDNN is unavailable, ORT must guard all frontend graph-build/execute +entry points before calling frontend APIs so the frontend shim never tries to resolve symbols +from a null handle. + +This means `cudnn_frontend` stays a compile‑time‑only dependency. We still need its headers +and the cuDNN headers to build, but not cuDNN import libraries at link time. + +A single loader object owns the `dlopen`/`LoadLibrary` handles and the resolved function +pointers: + +```cpp +class CudnnLibrary { // onnxruntime/core/providers/cuda/cudnn_loader.{h,cc} + public: + static CudnnLibrary& Get(); // thread-safe singleton (std::call_once) + bool Available() const; // true iff all required libs + symbols resolved + // ... function pointer members, one per cuDNN entry ORT/frontend uses ... +}; +``` + +Loader responsibilities: + +- On first use, attempt to load the cuDNN runtime. For cuDNN **9**, this is a small set of + sub‑libraries (`libcudnn.so.9` umbrella plus, depending on packaging, + `libcudnn_graph`, `libcudnn_ops`, `libcudnn_cnn`, `libcudnn_engines_*`). On Windows the + corresponding `cudnn*64_9.dll` set. We load the umbrella `libcudnn` first; cuDNN itself + dlopens its sub‑libraries. +- Resolve each required symbol with `dlsym` / `GetProcAddress`. +- If the umbrella library or any **required** symbol is missing, set `available_ = false`. +- Do not accept a provider option that names a cuDNN runtime path. A provider-controlled + native library path is equivalent to a native code loading hook if provider options are + influenced by untrusted input. + +On Windows, cuDNN 9 is split into multiple DLLs. The loader should not rely on the process +working directory. Application or package code that needs an explicit directory should use a +trusted process-level preload mechanism, such as Python `preload_dlls(cudnn=True, +directory=...)`, before creating the session. On Linux, the C++ loader uses the default +dynamic loader search behavior for the umbrella `libcudnn.so.9`; deployment should provide +trusted library paths via the system loader, container image, package manager, or explicit +application preload. + +The loader must not run when the CUDA provider option `enable_cudnn=0` is set (see §3.3). +This keeps "force no cuDNN" tests deterministic even on machines where cuDNN is installed. + +**Symbol manifest.** The set of symbols is finite and enumerable (see §2.4 plus +`cudnn_common.*` and `cudnn_rnn_base.*`). We maintain direct ORT cuDNN calls as a single +header list (`cudnn_symbols.inc`, an X‑macro list) consumed by both the loader and the stub +generator so the two never drift. Frontend backend symbols are resolved by +`cudnn_frontend`'s own dynamic-loading shim from the same ORT-owned cuDNN handle; maintain a +separate frontend-symbol audit list for testing and version checks. + +> **Alternative considered (delay‑load only):** Windows `/DELAYLOAD:cudnn*.dll` gets us lazy +> load on Windows, but Linux has no equivalent, and `cudnn_frontend` requires explicit dynamic +> loading support to avoid backend API imports. Rejected because the request requires +> identical behavior on Linux and Windows from a single binary. The direct-call trampoline +> plus frontend dynamic-loading approach is uniform across both. + +### 3.2 Availability flag and handle lifecycle + +- `cudnnCreate` is only invoked through the shim. The eager `cudnnCreate` calls in + `CudaStream` / `PerThreadContext` / plugin stream become **conditional and non‑fatal**: + - Attempt `CudnnLibrary::Get().Available()`; if false, leave `cudnn_handle_ == nullptr` and + **do not throw**. + - If true, create the handle as today. +- `GetCudnnHandle(context)` returns `nullptr` when cuDNN is unavailable (it already returns a + raw handle; today it's never null). +- New helpers in `cuda_common.h`: + +```cpp +bool CudnnAvailable(const OpKernelContext* context); // provider option && runtime availability + +// For kernels: fail fast with a clear message if cuDNN is required but missing. +#define ORT_RETURN_IF_CUDNN_UNAVAILABLE(context, op_name) \ + ORT_RETURN_IF_ERROR(::onnxruntime::cuda::CheckCudnnAvailable(context, op_name)) +``` + +- `CudaErrString` must not call `cudnnGetErrorString` when cuDNN is unavailable + (route through the shim, which returns a static string in that case). + +### 3.3 CUDA provider option: `enable_cudnn` + +cuDNN can be disabled explicitly with a CUDA provider option: + +```text +enable_cudnn = 1 # default: try to load and use cuDNN when it is present +enable_cudnn = 0 # do not load cuDNN; force native CUDA paths / Phase-1 NOT_IMPLEMENTED +``` + +`enable_cudnn` is the policy switch. When it is `0`, ORT must not attempt to load cuDNN. ORT +intentionally does not provide a `cudnn_path` provider option because provider options can be +supplied by higher-level configuration systems, and allowing them to choose a native DLL/SO +path would create a library-loading security risk. + +Implementation details: + +- Add `constexpr const char* kEnableCudnn = "enable_cudnn"` in + `cuda::provider_option_names`. +- Add `bool enable_cudnn{true};` to `CUDAExecutionProviderInfo`. +- Parse it with `ProviderOptionsParser::AddAssignmentToReference(...)` in + `CUDAExecutionProviderInfo::FromProviderOptions(...)`. +- Emit it from `CUDAExecutionProviderInfo::ToProviderOptions(...)`. +- Include it in `std::hash` because it changes the EP behavior. +- Do **not** add a field to `OrtCUDAProviderOptionsV2` for Phase 1. That struct is public C + ABI surface; string-key provider options are sufficient and can be set through existing + provider-options APIs. +- Add an EP helper such as `CUDAExecutionProvider::IsCudnnEnabled()` or + `CudaKernel::IsCudnnEnabled()` so kernels can distinguish: + - cuDNN disabled by user (`enable_cudnn=0`), and + - cuDNN enabled but unavailable at runtime. + +The effective condition for cuDNN use is: + +```text +effective_cudnn_available = info.enable_cudnn && CudnnLibrary::Get().Available() +``` + +If `enable_cudnn=0`, ORT must not call `dlopen` / `LoadLibrary` for cuDNN and must not create +a cuDNN handle. If `enable_cudnn=1`, ORT uses trusted process-level library discovery: the +system loader search path, package/container deployment, or an explicit preload performed by +application code before session creation. + +### 3.4 Phase 1 fallback behavior (chosen: throw at Run time) + +Per the agreed design, in Phase 1 cuDNN‑dependent kernels **remain registered** but **fail +fast** when executed without cuDNN. The check is added at the top of each cuDNN op's +`ComputeInternal` (or centralized in shared base helpers such as `CudnnConvState` setup, +`CudnnRnnBase`, the reduction helper, etc.): + +```cpp +Status Conv::ComputeInternal(OpKernelContext* context) const { + ORT_RETURN_IF_CUDNN_UNAVAILABLE(context, "Conv"); + // ... existing cuDNN path ... +} +``` + +The guard should include the reason in the message: + +- `enable_cudnn=0`: "Operator 'Conv' on the CUDA EP requires cuDNN, but cuDNN was disabled + by the CUDA provider option 'enable_cudnn'." +- cuDNN missing: "Operator 'Conv' on the CUDA EP requires cuDNN, but cuDNN was not found at + runtime. Install cuDNN, or disable CUDA execution for this op/model." + +Rationale for "throw" over "don't register / fall back to CPU": + +- Keeps kernel registration tables identical regardless of runtime cuDNN presence (no + divergence between build/load configurations; simpler, lower‑risk). +- Produces a clear, attributable error instead of silent CPU fallback that can mask perf + cliffs. +- CPU fallback for individual nodes is still achievable by the user via EP assignment; we are + not removing that option, only not making it implicit. + +### 3.5 Builds in scope + +- **Main static CUDA EP** — primary target; shim + loader compiled in. +- **CUDA Plugin EP** (`BUILD_CUDA_EP_AS_PLUGIN`) — same shim/loader; the plugin's + `cuda_stream_plugin.cc` / `cuda_kernel_adapter.h` handle creation becomes conditional. +- **Unit tests** — link the shim instead of cuDNN; add tests that exercise both + cuDNN‑present and cuDNN‑absent behavior (the latter by forcing the loader into the + unavailable state, see §7). + +Python wheel packaging is unchanged by this design: cuDNN DLLs are not packed in the wheel +today, so Phase 1 is not introducing a new "CUDA-minimal" wheel flavor. The runtime loader +simply makes the existing package tolerant of environments where cuDNN is absent. + +TensorRT / NV‑RTX EPs are untouched and continue to link cuDNN as before. (If both a TRT EP +and the shimmed CUDA EP are in the same process, symbol collision must be avoided — see +§8 Risks.) + +--- + +## 4. Phase 1 — Make cuDNN optional (no new kernels) + +**Outcome:** ORT CUDA EP loads and runs without cuDNN. cuDNN‑backed ops throw a clear +`NOT_IMPLEMENTED` error when cuDNN is absent; everything else runs normally. When cuDNN *is* +present, behavior is byte‑for‑byte identical to today. + +### 4.1 Task breakdown + +1. **Symbol inventory & manifest.** + - Enumerate every direct `cudnn*` symbol referenced by ORT code (`cudnnCreate`, + descriptor APIs, legacy conv APIs, BN/LRN/pooling/reduction/softmax APIs, RNN APIs, + etc.). + - Capture direct calls as `cudnn_symbols.inc` (X‑macro: name, return type, signature). + - Separately audit the `cudnn_frontend` backend-symbol surface resolved by its dynamic + shim: `cudnnGetVersion`, `cudnnGetErrorString`, `cudnnBackend*` descriptor APIs, + `cudnnBackendExecute`, and version-gated graph helpers such as + `cudnnBackendPopulateCudaGraph`, `cudnnBackendUpdateCudaGraph`, and + `cudnnGetExecutionPlanWorkspaceSize`. + - *Verification:* link a probe binary that references only the direct-call manifest and + diff against `nm -D`/`dumpbin` of the real cuDNN to ensure completeness; add a frontend + graph-build/execute smoke test to prove `NV_CUDNN_FRONTEND_USE_DYNAMIC_LOADING` resolves + backend symbols through ORT's loaded cuDNN handle. + +2. **Loader (`cudnn_loader.{h,cc}`).** + - `dlopen`/`LoadLibrary` of the cuDNN umbrella lib with versioned name candidates + (`libcudnn.so.9`, `libcudnn.so`, `cudnn64_9.dll`, …). + - Do not accept a provider-supplied cuDNN path. Rely on trusted deployment/library-search + mechanisms or application-controlled preloading. + - On Windows, avoid relying on the process working directory. Python package users who need + an explicit directory should call `preload_dlls(cudnn=True, directory=...)` from trusted + application code before creating a CUDA EP session. + - Resolve all manifest symbols; populate function‑pointer table. + - `Available()` + thread‑safe one‑time init; report loader diagnostics without exposing a + provider-controlled library path option. + - Expose the raw library handle to `cudnn_frontend` dynamic-loading mode. + - Define and maintain `cudnn_frontend::cudnn_dlhandle` in one ORT translation unit when + `NV_CUDNN_FRONTEND_USE_DYNAMIC_LOADING` is enabled. Set it to the loader's cuDNN handle + after a successful load; keep it null when cuDNN is disabled or unavailable. + - Add an explicit "disabled" path: when the EP has `enable_cudnn=0`, skip loader + initialization entirely and report "disabled by provider option" to the error helper. + +3. **Stub/trampoline TU (`cudnn_stub.cc`).** + - Generate one trampoline per manifest entry forwarding to the loader's pointer; return a + sentinel `cudnnStatus_t` when unavailable. + - Handle the few non‑`cudnnStatus_t` entries (`cudnnGetErrorString`, `cudnnGetVersion`). + + The stubs must be compiled into the same target that currently links cuDNN. For Linux, + prefer hidden visibility for the stub definitions where possible to avoid exporting cuDNN + names from ORT provider binaries. The loader should use `RTLD_LOCAL` when opening cuDNN. + +4. **CMake changes.** + - Remove `CUDNN::cudnn_all` from `target_link_libraries` for the CUDA EP, plugin EP, and + tests; **keep** `CUDNN_INCLUDE_DIR` (headers) and `cudnn_frontend` (headers). + - Compile `cudnn_stub.cc` + `cudnn_loader.cc` into the EP. + - Compile CUDA EP targets that include `cudnn_frontend` with + `NV_CUDNN_FRONTEND_USE_DYNAMIC_LOADING`. This is required so frontend graph code uses + `dlsym` / `GetProcAddress` on `cudnn_frontend::cudnn_dlhandle` instead of creating + link-time imports for `cudnnBackend*` symbols. + - Keep `USE_CUDA_MINIMAL` working. It is used by TensorRT / NV‑TensorRT‑RTX related + builds and is not replaced by the optional-cuDNN runtime shim. + + Current CMake anchor points: + + - `cmake/onnxruntime_providers_cuda.cmake`: replace `CUDNN::cudnn_all` with the shim + sources/library while retaining `include(cudnn_frontend)` and the cuDNN include dirs; + add `NV_CUDNN_FRONTEND_USE_DYNAMIC_LOADING` for the provider target. + - `cmake/onnxruntime_providers_cuda_plugin.cmake`: same for + `onnxruntime_providers_cuda_plugin`. + - `cmake/onnxruntime_unittests.cmake`: unit tests should link against the shim path, not + against cuDNN import libraries, and should use the same frontend dynamic-loading define. + - `cmake/onnxruntime_python.cmake`: on Windows, the generated `version_info.py` currently + searches for `cudnn64_*.dll` and fails if it is missing. That fatal check must be + relaxed because cuDNN is optional; `cudnn_version` should be omitted, set to `None`, or + set to `"optional"` when no DLL is found. + +5. **Conditional handle creation.** + - `cuda_stream_handle.cc`, `cuda_execution_provider.cc`, `cuda_stream_plugin.cc`, + `cuda_kernel_adapter.h`: create the handle only if + `info.enable_cudnn && CudnnLibrary::Get().Available()`, otherwise leave it null and do + not throw. + - When `info.enable_cudnn` is false, skip the loader call entirely. + - External/custom-op resource behavior: `CudaResource::cudnn_handle_t` may be `nullptr`. + Any internal custom-op adapter that assumes a non-null handle must return the same clear + cuDNN-required error. + +6. **Guard all cuDNN op entry points.** + - Add `ORT_RETURN_IF_CUDNN_UNAVAILABLE(context, "")` or an equivalent helper to the + `ComputeInternal` of every op in the §2.4 table (centralize in shared bases where + possible: `CudnnRnnBase`, conv state setup, reduction helper, pooling, + batch/instance norm, LRN, cuDNN softmax path). + - For the cuDNN **softmax** and **reduction** paths that *already* have native + alternatives, prefer routing to the native path when cuDNN is absent instead of + throwing (Phase‑2 work; see §5). Otherwise throw. + - Guard frontend graph creation as well as frontend graph execution. `cudnn_frontend`'s + dynamic shim throws if it cannot resolve backend symbols, so ORT should fail with the + clearer provider-option / cuDNN-missing message before calling `validate()`, + `build_operation_graph()`, `create_execution_plans()`, `check_support()`, + `build_plans()`, or `execute()`. + +7. **Provider-option plumbing.** + - Add and parse `enable_cudnn` in `CUDAExecutionProviderInfo`. + - Return it from `GetProviderOptions()` / `ToProviderOptions()`. + - Include it in the EP hash. + - Add tests for parsing: `enable_cudnn` default true, `"0"` false, `"1"` true, invalid + values rejected. + +8. **Error‑string safety.** + - Make `CudaErrString` shim‑safe. + - Make `CudaErrString` report frontend dynamic-loading failures + without assuming cuDNN is available. + +9. **Docs & messaging.** + - Document the new behavior and `enable_cudnn`. + - Update Python package guidance for `onnxruntime.preload_dlls(cuda=True, cudnn=True, + directory=...)`: users can still preload a known cuDNN directory, but preloading is now + optional for CUDA EP load because the provider itself lazy-loads cuDNN. + - Update `onnxruntime/__init__.py` behavior as needed so missing cuDNN does not produce a + scary install warning by default in optional-cuDNN packages. If the user explicitly calls + `preload_dlls(cudnn=True)`, keep diagnostics useful and include the missing DLL name. + +10. **CI workflow for no-cuDNN builds.** + - Add a focused CI workflow/job that configures and builds the CUDA EP without cuDNN + import libraries available at link time. It should still provide cuDNN headers, because + the optional-cuDNN design remains source-compatible with cuDNN APIs and + `cudnn_frontend` headers. + - The job should verify the produced CUDA provider binary has no direct cuDNN runtime + dependency (`readelf -d` / `ldd` on Linux, `dumpbin /dependents` on Windows). + - Run at least a smoke test that imports ORT, initializes the CUDA EP, and executes a + non-cuDNN CUDA model with cuDNN runtime libraries absent from the runtime library path. + - Run a negative smoke test for one cuDNN-backed op, such as Conv, and assert the clear + `NOT_IMPLEMENTED` error rather than a dynamic-loader failure. + - Start with Linux CUDA CI, then add the equivalent Windows CUDA CI leg once the Windows + sub-DLL search behavior is implemented and stable. + +### 4.2 Acceptance criteria (Phase 1) + +- With cuDNN **removed** from the system: + - `libonnxruntime`/CUDA provider loads; a model with no cuDNN ops runs correctly on CUDA. + - A model with a cuDNN op (e.g. Conv) fails with the clear `NOT_IMPLEMENTED` message, not a + crash or loader error. +- With cuDNN **present** and `enable_cudnn=0`: + - ORT does not load cuDNN or create a cuDNN handle. + - Phase‑1 cuDNN ops fail with the "disabled by provider option" `NOT_IMPLEMENTED` message. + - Phase‑2 / Phase‑3 native fallback ops run through the native path once implemented. +- With cuDNN **present**: full existing test suite passes unchanged (no perf/accuracy + regression). +- Both main CUDA EP and plugin EP build and pass. +- A dedicated no-cuDNN CI job builds the CUDA EP without cuDNN import libraries, confirms the + provider binary has no direct cuDNN runtime dependency, and runs the no-cuDNN smoke tests. + +--- + +## 5. Phase 2 — Replace LLM‑relevant cuDNN paths + +**Outcome:** LLM‑focused CUDA workloads can run without cuDNN for common Softmax / +LogSoftmax and reduction patterns. This phase is part of the first milestone with Phase 1: +Phase 1 makes cuDNN optional, and Phase 2 removes the cuDNN dependency from ops that LLM +models may still use. + +### 5.1 Scope + +1. **Softmax / LogSoftmax** — native warp/block kernels already exist; make them the default + and drop the cuDNN path (or keep cuDNN only as an opt‑in fast path). +2. **Reductions / ArgMax / ArgMin** — CUB‑based implementations; remove the + `cudnnReduceTensor` dependency. + +### 5.2 Mechanism + +For these ops, prefer the native implementation regardless of cuDNN availability, unless a +specific cuDNN fast path is intentionally kept behind an opt‑in provider option: + +```text +use native CUDA / CUB implementation +optional: if provider option requests cuDNN fast path and cuDNN is available, use cuDNN +``` + +The important Phase‑2 property is that these ops must not require a non-null cuDNN handle. +If `enable_cudnn=0`, or if cuDNN is absent, they should still run through the native path. + +### 5.3 Acceptance criteria (Phase 2, per op) + +- Native path matches cuDNN within tolerance on the op's existing unit tests. +- With cuDNN absent, the op runs (no `NOT_IMPLEMENTED`). +- With cuDNN present, no regression in correctness; any retained cuDNN fast path is explicit + and test-covered. +- With `enable_cudnn=0`, no dynamic cuDNN load is attempted. + +--- + +## 6. Phase 3 — Replace remaining cuDNN‑backed NN ops + +**Outcome:** Broader CNN / vision-style CUDA workloads can run without cuDNN where practical. +These ops are outside the first Phase 1 + Phase 2 milestone because they are less central to +LLM workloads and, for convolution, much more expensive to replace well. + +### 6.1 Scope + +1. **Pooling** (Max/Average, global variants) — straightforward native kernels. +2. **BatchNormalization / InstanceNormalization (inference)** — elementwise affine over + precomputed stats; native kernel is simple. +3. **LRN** — native kernel. +4. **Conv / ConvTranspose / FusedConv** — the hard part. Options: + - implicit‑GEMM CUDA kernels for common cases, + - CUTLASS conv, + - precompiled Triton conv cubins, + - or im2col + existing GEMM as a correctness fallback. + Keep cuDNN as the preferred fast path when available; native kernel as fallback. + +RNN / LSTM / GRU are intentionally not part of the Phase‑3 scope for now. They are the +heaviest to replace and may remain cuDNN‑only with a clear `NOT_IMPLEMENTED` when cuDNN is +absent unless there is product demand. + +### 6.2 Mechanism + +For each Phase‑3 op, introduce a dispatch at `ComputeInternal`: + +```text +if info.enable_cudnn and CudnnLibrary::Get().Available() and (cuDNN path preferred): use cuDNN +else: use native CUDA / CUTLASS / Triton-cubin fallback +``` + +This preserves cuDNN performance where present while removing the hard requirement. +Triton cubins (see `docs/ORT_Use_Triton_Kernel.md`) are an option for fused/normalization +kernels; CUTLASS (already vendored) for conv/GEMM‑shaped work. + +### 6.3 Acceptance criteria (Phase 3, per op) + +- Native path matches cuDNN within tolerance on the op's existing unit tests. +- With cuDNN absent, the op runs (no `NOT_IMPLEMENTED`). +- With cuDNN present, no regression (cuDNN path still selected unless configured otherwise). +- With `enable_cudnn=0`, the op uses the native fallback and does not load cuDNN. + +--- + +## 7. Testing strategy + +- **Loader unit tests:** simulate cuDNN present vs absent. + - "Absent" is forced via a test hook on `CudnnLibrary` (e.g. an internal + `SetForceUnavailableForTest(true)` compiled only in test/internal builds), avoiding the + need to physically remove cuDNN in CI. + - "Disabled" is tested with CUDA provider option `enable_cudnn=0`; this should not touch + the dynamic loader at all. +- **Op‑level tests:** for each cuDNN op, assert the clear `NOT_IMPLEMENTED` error in the + forced‑absent mode (Phase 1), and correctness in present mode. +- **cuDNN frontend tests:** add a Conv / ConvTranspose test that exercises frontend graph + creation and execution with cuDNN present, and verifies that `enable_cudnn=0` fails before + `cudnn_frontend` attempts to resolve backend symbols. +- **Binary dependency tests:** inspect the CUDA provider binary (`ldd` / `readelf -d` on + Linux, `dumpbin /dependents` on Windows) and confirm there is no direct dependency on + `libcudnn*` / `cudnn64_*.dll` even though `cudnn_frontend` headers are compiled in. +- **Phase‑2 native path tests:** run Softmax / LogSoftmax and reduction tests with + `enable_cudnn=0` and with the forced‑absent loader hook. These should pass, not throw. +- **Phase‑3 native path tests:** as each Phase‑3 op is implemented, add the same + `enable_cudnn=0` / forced‑absent coverage for that op. +- **Load test:** a process that initializes the CUDA EP with the cuDNN libs unavailable and + runs a non‑cuDNN model end‑to‑end. +- **No-cuDNN CI workflow:** add a workflow/job, initially in Linux CUDA CI, that removes cuDNN + import libraries from the link/runtime environment while keeping cuDNN headers available. + It should build the CUDA EP, inspect dynamic dependencies, run the non-cuDNN smoke model, + and verify a cuDNN-backed op fails with ORT's `NOT_IMPLEMENTED` message. Add the Windows + equivalent after the Windows cuDNN sub-DLL loading path is covered. +- **Python preload tests:** verify `onnxruntime.preload_dlls(cudnn=True, directory=...)` + still preloads a user-provided cuDNN directory, while normal `import onnxruntime` and CUDA + EP initialization do not fail or print install guidance solely because cuDNN is missing. +- **Regression:** full existing CUDA suite with cuDNN present must stay green. +- CI legs that genuinely lack cuDNN can additionally validate the real (not forced) path. + +--- + +## 8. Risks and mitigations + +- **Symbol‑manifest completeness for direct cuDNN calls.** Missing a direct-call symbol in + `cudnn_symbols.inc` would cause an unresolved-symbol link error or a runtime trampoline + failure. *Mitigation:* the probe-binary diff in §4.1-1, and CI that builds with cuDNN + import libs unavailable. +- **`cudnn_frontend` dynamic-loading integration.** If + `NV_CUDNN_FRONTEND_USE_DYNAMIC_LOADING` is not applied consistently, frontend graph code may + still create imports for `cudnnBackend*` symbols. If `cudnn_frontend::cudnn_dlhandle` is not + defined/set by ORT, frontend calls will throw while resolving backend symbols. *Mitigation:* + apply the compile definition to every target that includes frontend headers, define the + global handle once in ORT, and add binary dependency plus Conv/ConvTranspose frontend smoke + tests. +- **`cudnn_frontend` may add new backend symbols across dependency bumps.** A future + `cudnn_frontend` version may resolve additional backend, CUDA, CUDART, NVRTC, or experimental + symbols. *Mitigation:* after every `cudnn_frontend` update, grep/audit + `cudnn_frontend_shim.h` and experimental shims, then update the frontend-symbol audit tests. +- **cuDNN sub‑library packaging differences (v9 split libs; distro/conda/pip layouts).** + *Mitigation:* rely on trusted deployment mechanisms for library discovery. Python users who + need a specific directory can call `onnxruntime.preload_dlls(cudnn=True, directory=...)` + from application code before creating the session; C++ deployments should use container, + package-manager, or system loader configuration. +- **Native library path provider options can become code-loading hooks.** If an attacker can + influence a provider option that names a DLL/SO directory, they can potentially cause ORT + to load attacker-controlled native code. *Mitigation:* do not expose a `cudnn_path` provider + option. Keep custom library discovery at trusted process/deployment layers instead. +- **Python preload behavior can conflict with optional cuDNN.** Today + `onnxruntime.preload_dlls(cudnn=True)` tries to load cuDNN and prints installation guidance + on failure. That is useful when the user requested preloading, but too alarming if cuDNN is + optional. *Mitigation:* keep explicit preloading available, but avoid invoking or requiring + cuDNN preload as part of normal optional-cuDNN package import / CUDA EP load. Update docs so + users who need a specific cuDNN directory can call `preload_dlls(cudnn=True, directory=...)` + before creating the session. +- **Python version metadata currently assumes cuDNN on Windows CUDA builds.** + `cmake/onnxruntime_python.cmake` fails if no `cudnn64_*.dll` is found when generating + `version_info.py`. *Mitigation:* make `cudnn_version` optional in that generated file. +- **Symbol collision when a cuDNN‑linking EP (TensorRT) and the shimmed CUDA EP coexist in + one process.** Our trampolines define real cuDNN symbol names; if TRT's cuDNN is also + loaded, the dynamic linker could bind either. *Mitigation:* keep the shim symbols with + internal/hidden visibility where possible and resolve the real library explicitly via the + loader (`RTLD_LOCAL`); document the constraint; consider a build option that keeps the + classic hard‑link behavior for TRT‑combined packages. +- **ABI/version skew** (built against cuDNN headers vN, loads runtime vM). *Mitigation:* + check `cudnnGetVersion` in the loader and refuse (mark unavailable) on incompatible major + versions. +- **Performance:** one extra indirect call per cuDNN entry — negligible relative to kernel + cost. + +--- + +## 9. Backward compatibility + +- Default behavior with cuDNN installed is unchanged (same kernels, same perf). +- The build still requires cuDNN **headers** (and `cudnn_frontend` headers) to compile; it no + longer requires cuDNN **import libraries** to link the in‑scope targets. +- `cudnn_frontend` remains a compile-time dependency while ORT uses frontend graph APIs for + Conv / ConvTranspose. It should not introduce a runtime cuDNN dependency when compiled with + `NV_CUDNN_FRONTEND_USE_DYNAMIC_LOADING` and wired to ORT's loader handle. +- The wheel package is unchanged: cuDNN DLLs are not packed into the wheel today, and this + design does not introduce a separate CUDA-minimal wheel. +- `onnxruntime.preload_dlls()` remains supported for users who want Python to preload CUDA / + cuDNN libraries from PyTorch, NVIDIA site packages, or an explicit directory. It becomes an + optional convenience path for cuDNN, not a requirement for importing ORT or initializing the + CUDA EP without cuDNN. +- No public API change is required for Phase 1. The custom‑op `CudaContext::cudnn_handle` + may now be `nullptr`; this is documented, and `FetchResource` returns null gracefully. + +--- + +## 10. Resolved decisions + +- Force-disabling cuDNN is a **CUDA provider option**, not a session option. Use + `enable_cudnn=0`. +- Providing a custom cuDNN runtime directory is **not** a CUDA provider option. Use trusted + deployment/library-search configuration, or Python `preload_dlls(cudnn=True, directory=...)` + from application code before creating the session. +- No new "CUDA-minimal" wheel is required for Phase 1. cuDNN DLLs are not packed in the wheel + today. +- Keep the existing `USE_CUDA_MINIMAL` build-time path. It is used by RTX/TensorRT-related EP + builds and is not replaced by the CUDA EP runtime shim. diff --git a/docs/cuda_plugin_ep/QUICK_START.md b/docs/cuda_plugin_ep/QUICK_START.md index 4055e1056b507..b5acab30748d9 100644 --- a/docs/cuda_plugin_ep/QUICK_START.md +++ b/docs/cuda_plugin_ep/QUICK_START.md @@ -15,6 +15,18 @@ build.bat --cmake_generator "Visual Studio 17 2022" --config Release --build_whe --cmake_extra_defines "onnxruntime_BUILD_CUDA_EP_AS_PLUGIN=ON" ``` +### Building and testing without cuDNN at runtime + +The CUDA Plugin EP build still requires cuDNN headers, but the plugin library must not have a hard runtime dependency on cuDNN. When cuDNN is not present, non-cuDNN kernels can still run. Kernels that still require cuDNN fail with `NOT_IMPLEMENTED` unless they have a native CUDA fallback. + +For local Linux CUDA 13 validation, use the no-cuDNN helper script. It keeps `CUDNN_HOME` available for headers, excludes cuDNN directories from `PATH` and `LD_LIBRARY_PATH`, verifies the plugin has no direct cuDNN dependency, and runs plugin tests in no-cuDNN mode: + +```bash +bash .env/cuda_130_plugin_no_cudnn.sh --build --test_plugin +``` + +The test mode sets `ORT_TEST_CUDA_PLUGIN_EP=1` and `ORT_TEST_CUDA_PLUGIN_NO_CUDNN=1`, which passes `enable_cudnn=0` to plugin sessions and skips plugin tests for operators that still require cuDNN, such as Conv, ConvTranspose, BatchNormalization, InstanceNormalization, LRN, ArgMax, reductions, Einsum, and cuDNN-backed pooling paths. + ## Minimum ONNX Runtime Version The plugin is compiled against the ONNX Runtime headers in this repository, but it is designed to load into an **older** ONNX Runtime runtime as well. The minimum compatible version is declared in [`plugin-ep-cuda/MIN_ONNXRUNTIME_VERSION`](../../plugin-ep-cuda/MIN_ONNXRUNTIME_VERSION) (currently **1.24.4**) and is the single source of truth: @@ -159,6 +171,15 @@ python test_cuda_plugin_ep.py The script validates plugin registration, device enumeration, provider options, operator coverage, and that key nodes are actually assigned to `CudaPluginExecutionProvider`. +To run the same focused test against a plugin build without cuDNN in the runtime search path: + +```bash +export ORT_TEST_CUDA_PLUGIN_NO_CUDNN=1 +export ORT_TEST_CUDA_PLUGIN_EP=1 +export ORT_CUDA_PLUGIN_PATH=/path/to/build/Release/libonnxruntime_providers_cuda_plugin.so +python test_cuda_plugin_ep.py +``` + ### Test against the minimum supported ORT version The plugin must keep working on the oldest supported ONNX Runtime (see [Minimum ONNX Runtime Version](#minimum-onnx-runtime-version)), not just the version it was built against. To validate this locally, install the minimum base runtime and run the same test against the freshly built plugin library: diff --git a/docs/cuda_plugin_ep/cuda_plugin_ep_design.md b/docs/cuda_plugin_ep/cuda_plugin_ep_design.md index 8f9a0d388d1af..b8286a5adf111 100644 --- a/docs/cuda_plugin_ep/cuda_plugin_ep_design.md +++ b/docs/cuda_plugin_ep/cuda_plugin_ep_design.md @@ -26,6 +26,21 @@ The ORT CUDA build produces four separate libraries: | `onnxruntime_providers_cuda` | `libonnxruntime_providers_cuda.so` | Shared module | In-tree CUDA EP (uses `SHARED_PROVIDER` bridge) | | `onnxruntime_providers_cuda_plugin` | `libonnxruntime_providers_cuda_plugin.so` | Shared module | Plugin CUDA EP (uses EP API adapters) | +### 2.1.1 Optional cuDNN Runtime Dependency + +The CUDA Plugin EP follows the in-tree CUDA EP's optional-cuDNN model. cuDNN headers are still required at build time, but the plugin shared library must not link directly to cuDNN or contain a cuDNN DLL/SO in its dynamic dependency table. cuDNN is loaded lazily through the ORT cuDNN loader when `enable_cudnn` is enabled and the runtime libraries are available through trusted process-level library discovery. + +The plugin exposes the same `enable_cudnn` provider option as the in-tree CUDA EP: + +```text +enable_cudnn = 1 # default: try to load and use cuDNN when available +enable_cudnn = 0 # do not load cuDNN; run native CUDA paths or fail cuDNN-required ops with NOT_IMPLEMENTED +``` + +There is intentionally no provider option for a custom cuDNN DLL/SO path. Provider options can flow from higher-level configuration systems, so allowing them to choose a native library path would create a code-loading security risk. Deployments that need a specific cuDNN directory should use trusted process-level mechanisms, such as the OS loader configuration, container image setup, or Python `preload_dlls(cudnn=True, directory=...)` before plugin registration. + +No-cuDNN plugin validation runs `test_cuda_plugin_ep.py` with `ORT_TEST_CUDA_PLUGIN_EP=1` and `ORT_TEST_CUDA_PLUGIN_NO_CUDNN=1`. That mode passes `enable_cudnn=0` to plugin sessions and skips tests for operators that still require cuDNN in the current implementation. Non-cuDNN operator coverage, plugin registration, device enumeration, graph assignment, CUDA graph, I/O binding, and profiling tests continue to run. + ### 2.2 Preprocessor Defines Each build target uses different preprocessor defines that control how framework types are resolved: @@ -618,7 +633,7 @@ The in-tree CUDA EP and shared provider bridge are compiled identically regardle ### 9.3 Plugin Independence -`libonnxruntime_providers_cuda_plugin.so` is **fully self-contained**. It does not depend on `libonnxruntime_providers_cuda.so` or `libonnxruntime_providers_shared.so` at load time. It statically links against `onnxruntime_framework`, `onnxruntime_graph`, `onnxruntime_common`, `onnxruntime_mlas`, `onnxruntime_flatbuffers`, and links dynamically against CUDA (`cudart`, `cublas`, `cublasLt`, `cufft`), cuDNN, and protobuf. Communication with the ORT runtime happens exclusively through the C API (`OrtApi`/`OrtEpApi`) passed at load time. +`libonnxruntime_providers_cuda_plugin.so` is **fully self-contained**. It does not depend on `libonnxruntime_providers_cuda.so` or `libonnxruntime_providers_shared.so` at load time. It statically links against `onnxruntime_framework`, `onnxruntime_graph`, `onnxruntime_common`, `onnxruntime_mlas`, `onnxruntime_flatbuffers`, and links dynamically against CUDA (`cudart`, `cublas`, `cublasLt`, `cufft`) and protobuf. cuDNN is loaded lazily only when enabled and available at runtime. Communication with the ORT runtime happens exclusively through the C API (`OrtApi`/`OrtEpApi`) passed at load time. ### 9.4 Build Outputs diff --git a/onnxruntime/__init__.py b/onnxruntime/__init__.py index 0a06156fe78d8..3b4f413a77b8f 100644 --- a/onnxruntime/__init__.py +++ b/onnxruntime/__init__.py @@ -330,14 +330,14 @@ def is_target_dll(path: str): def preload_dlls(cuda: bool = True, cudnn: bool = True, msvc: bool = True, directory=None): - """Preload CUDA 12.x+ and cuDNN 9.x DLLs in Windows or Linux, and MSVC runtime DLLs in Windows. + """Preload CUDA 12.x+ and optional cuDNN 9.x DLLs in Windows or Linux, and MSVC runtime DLLs in Windows. When the installed PyTorch is compatible (using same major version of CUDA and cuDNN), there is no need to call this function if `import torch` is done before `import onnxruntime`. Args: cuda (bool, optional): enable loading CUDA DLLs. Defaults to True. - cudnn (bool, optional): enable loading cuDNN DLLs. Defaults to True. + cudnn (bool, optional): enable loading cuDNN DLLs. Defaults to True. Missing cuDNN DLLs are ignored. msvc (bool, optional): enable loading MSVC DLLs in Windows. Defaults to True. directory(str, optional): a directory contains CUDA or cuDNN DLLs. It can be an absolute path, or a path relative to the directory of this file. @@ -430,6 +430,7 @@ def preload_dlls(cuda: bool = True, cudnn: bool = True, msvc: bool = True, direc # Try load DLLs from nvidia site packages. dll_paths = _get_nvidia_dll_paths(is_windows, cuda, cudnn) + optional_dll_filenames = {relative_path[-1] for relative_path in _get_nvidia_dll_paths(is_windows, False, cudnn)} loaded_dlls = [] for relative_path in dll_paths: dll_path = ( @@ -442,11 +443,8 @@ def preload_dlls(cuda: bool = True, cudnn: bool = True, msvc: bool = True, direc _ = ctypes.CDLL(dll_path) loaded_dlls.append(relative_path[-1]) except Exception as e: - print(f"Failed to load {dll_path}: {e}") - - # cuDNN DLLs that only exist in newer cuDNN releases (e.g. >= 9.23) and are - # optional for inference. Missing them on older cuDNN must not be treated as a failure. - _optional_dll_filenames = {"cudnn_engines_tensor_ir64_9.dll"} + if relative_path[-1] not in optional_dll_filenames: + print(f"Failed to load {dll_path}: {e}") # Try load DLLs with default path settings. has_failure = False @@ -456,9 +454,9 @@ def preload_dlls(cuda: bool = True, cudnn: bool = True, msvc: bool = True, direc try: _ = ctypes.CDLL(dll_filename) except Exception as e: - if dll_filename not in _optional_dll_filenames: + if dll_filename not in optional_dll_filenames: has_failure = True print(f"Failed to load {dll_filename}: {e}") if has_failure: - print("Please follow https://onnxruntime.ai/docs/install/#cuda-and-cudnn to install CUDA and CuDNN.") + print("Please follow https://onnxruntime.ai/docs/install/#cuda-and-cudnn to install CUDA.") diff --git a/onnxruntime/core/providers/cuda/cuda_call.cc b/onnxruntime/core/providers/cuda/cuda_call.cc index b9e909714f1c0..d27dbf266d183 100644 --- a/onnxruntime/core/providers/cuda/cuda_call.cc +++ b/onnxruntime/core/providers/cuda/cuda_call.cc @@ -9,6 +9,11 @@ #else #include #endif +#ifndef USE_CUDA_MINIMAL +#include "core/providers/cuda/cudnn_loader.h" +#endif + +#include #ifdef _WIN32 #else // POSIX @@ -89,6 +94,20 @@ std::conditional_t CudaCall( ERRTYPE retCode, const char* exprString, const char* libName, SUCCTYPE successCode, const char* msg, const char* file, const int line) { if (retCode != successCode) { +#ifndef USE_CUDA_MINIMAL + if constexpr (std::is_same_v) { + if (!cuda::CudnnLibrary::Get().Available()) { + auto status = ORT_MAKE_STATUS(ONNXRUNTIME, NOT_IMPLEMENTED, + "cuDNN is unavailable for CUDA Execution Provider: ", + cuda::CudnnLibrary::Get().Error()); + if constexpr (THRW) { + ORT_THROW(status.ErrorMessage()); + } else { + return status; + } + } + } +#endif try { #ifdef _WIN32 std::string hostname_str = GetEnvironmentVar("COMPUTERNAME"); diff --git a/onnxruntime/core/providers/cuda/cuda_execution_provider.cc b/onnxruntime/core/providers/cuda/cuda_execution_provider.cc index db24f43223613..86620ae3c60af 100755 --- a/onnxruntime/core/providers/cuda/cuda_execution_provider.cc +++ b/onnxruntime/core/providers/cuda/cuda_execution_provider.cc @@ -18,6 +18,7 @@ #include "core/providers/cuda/gpu_data_transfer.h" #include "core/providers/cuda/cuda_profiler.h" #include "core/providers/cuda/cuda_mempool_arena.h" +#include "core/providers/cuda/cudnn_loader.h" #include "core/session/onnxruntime_run_options_config_keys.h" #ifndef USE_CUDA_MINIMAL @@ -223,18 +224,19 @@ AllocatorPtr CUDAExecutionProvider::CreateCudaPinnedAllocator(const CUDAAllocato return CreateAllocator(pinned_memory_info); } -CUDAExecutionProvider::PerThreadContext::PerThreadContext(OrtDevice::DeviceId device_id, cudaStream_t stream, size_t /*gpu_mem_limit*/, - ArenaExtendStrategy /*arena_extend_strategy*/, CUDAExecutionProviderExternalAllocatorInfo /*external_allocator_info*/, - OrtArenaCfg* /*default_memory_arena_cfg*/) { +CUDAExecutionProvider::PerThreadContext::PerThreadContext(OrtDevice::DeviceId device_id, cudaStream_t stream, + const CUDAExecutionProviderInfo& info) { CUDA_CALL_THROW(cudaSetDevice(device_id)); #ifndef USE_CUDA_MINIMAL CUBLAS_CALL_THROW(cublasCreate(&cublas_handle_)); CUBLAS_CALL_THROW(cublasLtCreate(&cublas_lt_handle_)); CUBLAS_CALL_THROW(cublasSetStream(cublas_handle_, stream)); - CUDNN_CALL_THROW(cudnnCreate(&cudnn_handle_)); - CUDNN_CALL_THROW(cudnnSetStream(cudnn_handle_, stream)); - LOGS_DEFAULT(INFO) << "cuDNN version: " << cudnnGetVersion(); + if (info.enable_cudnn && cuda::CudnnLibrary::Get().Available()) { + CUDNN_CALL_THROW(cudnnCreate(&cudnn_handle_)); + CUDNN_CALL_THROW(cudnnSetStream(cudnn_handle_, stream)); + LOGS_DEFAULT(INFO) << "cuDNN version: " << cudnnGetVersion(); + } #endif cuda_graph_.SetStream(stream); } @@ -243,7 +245,9 @@ CUDAExecutionProvider::PerThreadContext::~PerThreadContext() { #ifndef USE_CUDA_MINIMAL ORT_IGNORE_RETURN_VALUE(CUBLAS_CALL(cublasDestroy(cublas_handle_))); ORT_IGNORE_RETURN_VALUE(CUBLAS_CALL(cublasLtDestroy(cublas_lt_handle_))); - ORT_IGNORE_RETURN_VALUE(CUDNN_CALL(cudnnDestroy(cudnn_handle_))); + if (cudnn_handle_ != nullptr) { + ORT_IGNORE_RETURN_VALUE(CUDNN_CALL(cudnnDestroy(cudnn_handle_))); + } #endif } @@ -446,8 +450,7 @@ CUDAExecutionProvider::PerThreadContext& CUDAExecutionProvider::GetPerThreadCont // get or create a context if (context_state_.retired_context_pool.empty()) { - context = std::make_shared(info_.device_id, stream_, info_.gpu_mem_limit, - info_.arena_extend_strategy, info_.external_allocator_info, info_.default_memory_arena_cfg); + context = std::make_shared(info_.device_id, stream_, info_); } else { context = context_state_.retired_context_pool.back(); context_state_.retired_context_pool.pop_back(); diff --git a/onnxruntime/core/providers/cuda/cuda_execution_provider.h b/onnxruntime/core/providers/cuda/cuda_execution_provider.h index 7b78457c6a120..3d617794b1fbf 100644 --- a/onnxruntime/core/providers/cuda/cuda_execution_provider.h +++ b/onnxruntime/core/providers/cuda/cuda_execution_provider.h @@ -142,8 +142,7 @@ class CUDAExecutionProvider : public IExecutionProvider { class PerThreadContext final { public: - PerThreadContext(OrtDevice::DeviceId device_id, cudaStream_t stream, size_t cuda_mem_limit, ArenaExtendStrategy arena_extend_strategy, - CUDAExecutionProviderExternalAllocatorInfo external_alloc_info, OrtArenaCfg* arena_cfg); + PerThreadContext(OrtDevice::DeviceId device_id, cudaStream_t stream, const CUDAExecutionProviderInfo& info); ~PerThreadContext(); ORT_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(PerThreadContext); diff --git a/onnxruntime/core/providers/cuda/cuda_execution_provider_info.cc b/onnxruntime/core/providers/cuda/cuda_execution_provider_info.cc index 25b77a268830b..7949537eb1181 100644 --- a/onnxruntime/core/providers/cuda/cuda_execution_provider_info.cc +++ b/onnxruntime/core/providers/cuda/cuda_execution_provider_info.cc @@ -27,6 +27,7 @@ constexpr const char* kGpuExternalEmptyCache = "gpu_external_empty_cache"; constexpr const char* kCudnnConvUseMaxWorkspace = "cudnn_conv_use_max_workspace"; constexpr const char* kEnableCudaGraph = "enable_cuda_graph"; constexpr const char* kCudnnConv1dPadToNc1d = "cudnn_conv1d_pad_to_nc1d"; +constexpr const char* kEnableCudnn = "enable_cudnn"; constexpr const char* kTunableOpEnable = "tunable_op_enable"; constexpr const char* kTunableOpTuningEnable = "tunable_op_tuning_enable"; constexpr const char* kTunableOpMaxTuningDurationMs = "tunable_op_max_tuning_duration_ms"; @@ -116,6 +117,7 @@ CUDAExecutionProviderInfo CUDAExecutionProviderInfo::FromProviderOptions(const P .AddAssignmentToReference(cuda::provider_option_names::kCudnnConvUseMaxWorkspace, info.cudnn_conv_use_max_workspace) .AddAssignmentToReference(cuda::provider_option_names::kEnableCudaGraph, info.enable_cuda_graph) .AddAssignmentToReference(cuda::provider_option_names::kCudnnConv1dPadToNc1d, info.cudnn_conv1d_pad_to_nc1d) + .AddAssignmentToReference(cuda::provider_option_names::kEnableCudnn, info.enable_cudnn) .AddValueParser( cuda::provider_option_names::kEnableSkipLayerNormStrictMode, [](const std::string& value_str) -> Status { @@ -176,6 +178,7 @@ ProviderOptions CUDAExecutionProviderInfo::ToProviderOptions(const CUDAExecution {cuda::provider_option_names::kCudnnConvUseMaxWorkspace, MakeStringWithClassicLocale(info.cudnn_conv_use_max_workspace)}, {cuda::provider_option_names::kEnableCudaGraph, MakeStringWithClassicLocale(info.enable_cuda_graph)}, {cuda::provider_option_names::kCudnnConv1dPadToNc1d, MakeStringWithClassicLocale(info.cudnn_conv1d_pad_to_nc1d)}, + {cuda::provider_option_names::kEnableCudnn, MakeStringWithClassicLocale(info.enable_cudnn)}, {cuda::provider_option_names::kTunableOpEnable, MakeStringWithClassicLocale(info.tunable_op.enable)}, {cuda::provider_option_names::kTunableOpTuningEnable, MakeStringWithClassicLocale(info.tunable_op.tuning_enable)}, {cuda::provider_option_names::kTunableOpMaxTuningDurationMs, MakeStringWithClassicLocale(info.tunable_op.max_tuning_duration_ms)}, @@ -200,6 +203,7 @@ ProviderOptions CUDAExecutionProviderInfo::ToProviderOptions(const OrtCUDAProvid {cuda::provider_option_names::kDoCopyInDefaultStream, MakeStringWithClassicLocale(info.do_copy_in_default_stream)}, {cuda::provider_option_names::kCudnnConvUseMaxWorkspace, MakeStringWithClassicLocale(info.cudnn_conv_use_max_workspace)}, {cuda::provider_option_names::kCudnnConv1dPadToNc1d, MakeStringWithClassicLocale(info.cudnn_conv1d_pad_to_nc1d)}, + {cuda::provider_option_names::kEnableCudnn, MakeStringWithClassicLocale(true)}, {cuda::provider_option_names::kTunableOpEnable, MakeStringWithClassicLocale(info.tunable_op_enable)}, {cuda::provider_option_names::kTunableOpTuningEnable, MakeStringWithClassicLocale(info.tunable_op_tuning_enable)}, {cuda::provider_option_names::kTunableOpMaxTuningDurationMs, MakeStringWithClassicLocale(info.tunable_op_max_tuning_duration_ms)}, diff --git a/onnxruntime/core/providers/cuda/cuda_execution_provider_info.h b/onnxruntime/core/providers/cuda/cuda_execution_provider_info.h index 2d7f5d2ad041c..8173828695aff 100644 --- a/onnxruntime/core/providers/cuda/cuda_execution_provider_info.h +++ b/onnxruntime/core/providers/cuda/cuda_execution_provider_info.h @@ -6,6 +6,7 @@ #include #include +#include #include "core/common/hash_combine.h" #include "core/framework/arena_extend_strategy.h" @@ -69,6 +70,8 @@ struct CUDAExecutionProviderInfo { // By default, for Conv1D, will pad [N,C,D] to [N,C,D,1], if turn on, will pad to [N,C,1,D]. bool cudnn_conv1d_pad_to_nc1d{false}; + bool enable_cudnn{true}; + cuda::TunableOpInfo tunable_op{}; bool prefer_nhwc{false}; @@ -113,6 +116,7 @@ struct std::hash<::onnxruntime::CUDAExecutionProviderInfo> { onnxruntime::HashCombine(info.gpu_mem_limit, value); onnxruntime::HashCombine(info.tunable_op.max_tuning_duration_ms, value); onnxruntime::HashCombine(info.sdpa_kernel, value); + onnxruntime::HashCombine(info.enable_cudnn, value); // Memory pointers onnxruntime::HashCombine(reinterpret_cast(info.user_compute_stream), value); diff --git a/onnxruntime/core/providers/cuda/cuda_kernel.h b/onnxruntime/core/providers/cuda/cuda_kernel.h index 1d891f204b9bd..85d629764da48 100644 --- a/onnxruntime/core/providers/cuda/cuda_kernel.h +++ b/onnxruntime/core/providers/cuda/cuda_kernel.h @@ -7,6 +7,7 @@ #include "core/providers/cuda/cuda_common.h" #include "core/providers/cuda/cuda_execution_provider.h" +#include "core/providers/cuda/cudnn_loader.h" #include "core/providers/cuda/cuda_fwd.h" #include #include "core/providers/cuda/cuda_stream_handle.h" @@ -129,7 +130,7 @@ class CudaKernel : public OpKernel { } inline cudnnHandle_t GetCudnnHandle(OpKernelContext* ctx) const { - return GetCudnnHandle(static_cast(ctx->GetComputeStream())); + return RequireCudnnHandle(GetCudnnHandle(static_cast(ctx->GetComputeStream()))); } static inline cudnnHandle_t GetCudnnHandle(onnxruntime::CudaStream* stream) { @@ -143,7 +144,7 @@ class CudaKernel : public OpKernel { } inline cudnnHandle_t GetCudnnHandleOrDefault(onnxruntime::Stream* stream) const { - return stream ? GetCudnnHandle(stream) : DefaultCudnnHandle(); + return stream ? RequireCudnnHandle(GetCudnnHandle(stream)) : DefaultCudnnHandle(); } inline cublasHandle_t GetCublasHandle(OpKernelContext* ctx) const { @@ -258,7 +259,21 @@ class CudaKernel : public OpKernel { } inline cudnnHandle_t DefaultCudnnHandle() const { - return provider_->PerThreadDefaultCudnnHandle(); + return RequireCudnnHandle(provider_->PerThreadDefaultCudnnHandle()); + } + + static inline cudnnHandle_t RequireCudnnHandle(cudnnHandle_t handle) { + if (handle == nullptr) { +#ifndef USE_CUDA_MINIMAL + ORT_THROW_IF_ERROR(ORT_MAKE_STATUS(ONNXRUNTIME, NOT_IMPLEMENTED, + "cuDNN is unavailable or disabled for CUDA Execution Provider: ", + cuda::CudnnLibrary::Get().Error())); +#else + ORT_THROW_IF_ERROR(ORT_MAKE_STATUS(ONNXRUNTIME, NOT_IMPLEMENTED, + "cuDNN is unavailable for CUDA Execution Provider in a CUDA minimal build.")); +#endif + } + return handle; } inline cudaStream_t DefaultCudaStream() const { diff --git a/onnxruntime/core/providers/cuda/cuda_stream_handle.cc b/onnxruntime/core/providers/cuda/cuda_stream_handle.cc index 39bc2ea35ed85..cd9b7c9d3bbde 100644 --- a/onnxruntime/core/providers/cuda/cuda_stream_handle.cc +++ b/onnxruntime/core/providers/cuda/cuda_stream_handle.cc @@ -3,6 +3,7 @@ #include "core/providers/cuda/cuda_resource.h" #include "core/providers/cuda/cuda_stream_handle.h" #include "core/providers/cuda/cuda_common.h" +#include "core/providers/cuda/cudnn_loader.h" #include "core/common/spin_pause.h" namespace onnxruntime { @@ -77,13 +78,17 @@ CudaStream::CudaStream(cudaStream_t stream, if (own_flag) { CUBLAS_CALL_THROW(cublasCreate(&cublas_handle_)); CUBLAS_CALL_THROW(cublasSetStream(cublas_handle_, stream)); - CUDNN_CALL_THROW(cudnnCreate(&cudnn_handle_)); - CUDNN_CALL_THROW(cudnnSetStream(cudnn_handle_, stream)); + if (ep_info_.enable_cudnn && cuda::CudnnLibrary::Get().Available()) { + CUDNN_CALL_THROW(cudnnCreate(&cudnn_handle_)); + CUDNN_CALL_THROW(cudnnSetStream(cudnn_handle_, stream)); + } } else { cublas_handle_ = external_cublas_handle; CUBLAS_CALL_THROW(cublasSetStream(cublas_handle_, stream)); cudnn_handle_ = external_cudnn_handle; - CUDNN_CALL_THROW(cudnnSetStream(cudnn_handle_, stream)); + if (cudnn_handle_ != nullptr) { + CUDNN_CALL_THROW(cudnnSetStream(cudnn_handle_, stream)); + } } #else (void)(external_cudnn_handle); @@ -96,7 +101,9 @@ CudaStream::~CudaStream() { #ifndef USE_CUDA_MINIMAL if (own_stream_) { cublasDestroy(cublas_handle_); - cudnnDestroy(cudnn_handle_); + if (cudnn_handle_ != nullptr) { + cudnnDestroy(cudnn_handle_); + } auto* handle = GetHandle(); if (handle) cudaStreamDestroy(static_cast(handle)); diff --git a/onnxruntime/core/providers/cuda/cudnn_fe_call.cc b/onnxruntime/core/providers/cuda/cudnn_fe_call.cc index 60d6b85544269..ce48c696425e8 100644 --- a/onnxruntime/core/providers/cuda/cudnn_fe_call.cc +++ b/onnxruntime/core/providers/cuda/cudnn_fe_call.cc @@ -11,6 +11,7 @@ #endif #if !defined(__CUDACC__) && !defined(USE_CUDA_MINIMAL) #include +#include "core/providers/cuda/cudnn_loader.h" #endif #ifdef _WIN32 #else // POSIX @@ -71,6 +72,18 @@ std::conditional_t CudaCall( ERRTYPE retCode, const char* exprString, const char* libName, SUCCTYPE successCode, const char* msg, const char* file, const int line) { if (retCode != successCode) { +#if !defined(__CUDACC__) && !defined(USE_CUDA_MINIMAL) + if (!cuda::CudnnLibrary::Get().Available()) { + auto status = ORT_MAKE_STATUS(ONNXRUNTIME, NOT_IMPLEMENTED, + "cuDNN is unavailable for CUDA Execution Provider: ", + cuda::CudnnLibrary::Get().Error()); + if constexpr (THRW) { + ORT_THROW(status.ErrorMessage()); + } else { + return status; + } + } +#endif try { #ifdef _WIN32 std::string hostname_str = GetEnvironmentVar("COMPUTERNAME"); diff --git a/onnxruntime/core/providers/cuda/cudnn_loader.cc b/onnxruntime/core/providers/cuda/cudnn_loader.cc new file mode 100644 index 0000000000000..df84410bb15a7 --- /dev/null +++ b/onnxruntime/core/providers/cuda/cudnn_loader.cc @@ -0,0 +1,228 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "core/providers/cuda/cudnn_loader.h" + +#ifndef USE_CUDA_MINIMAL + +#include + +#ifdef _WIN32 +#ifndef NOMINMAX +#define NOMINMAX +#endif +#include +#else +#include +#endif + +namespace { + +std::vector GetCandidateLibraryNames() { +#ifdef _WIN32 + constexpr const char* kCudnnLibraryName = "cudnn64_9.dll"; +#else + constexpr const char* kCudnnLibraryName = "libcudnn.so.9"; + constexpr const char* kCudnnUnversionedLibraryName = "libcudnn.so"; +#endif + + std::vector candidates; + candidates.push_back(kCudnnLibraryName); +#ifndef _WIN32 + candidates.push_back(kCudnnUnversionedLibraryName); +#endif + return candidates; +} + +#ifdef _WIN32 +// Search the directories listed in the PATH environment variable for the given +// library and, if found, load it by its full path. Loading by full path (rather +// than letting the loader search) preserves the historical PATH-based cuDNN +// discovery without ever loading from the current working directory, which the +// LOAD_LIBRARY_SEARCH_DEFAULT_DIRS-only search excludes for security reasons. +HMODULE LoadLibraryFromPathEnv(const std::string& candidate) { + DWORD length = GetEnvironmentVariableA("PATH", nullptr, 0); + if (length == 0) { + return nullptr; + } + + std::string path_value(length, '\0'); + length = GetEnvironmentVariableA("PATH", path_value.data(), length); + if (length == 0) { + return nullptr; + } + path_value.resize(length); + + size_t start = 0; + while (start <= path_value.size()) { + size_t end = path_value.find(';', start); + if (end == std::string::npos) { + end = path_value.size(); + } + + std::string dir = path_value.substr(start, end - start); + start = end + 1; + if (dir.empty()) { + continue; + } + + if (dir.back() != '\\' && dir.back() != '/') { + dir.push_back('\\'); + } + std::string full_path = dir + candidate; + + if (GetFileAttributesA(full_path.c_str()) == INVALID_FILE_ATTRIBUTES) { + continue; + } + + HMODULE handle = LoadLibraryExA(full_path.c_str(), nullptr, LOAD_LIBRARY_SEARCH_DEFAULT_DIRS); + if (handle != nullptr) { + return handle; + } + } + + return nullptr; +} +#endif + +void* LoadLibraryCandidate(const std::string& candidate, std::string& error) { +#ifdef _WIN32 + // Use LOAD_LIBRARY_SEARCH_DEFAULT_DIRS so cuDNN is resolved only from the + // application directory, %WINDIR%\System32, and directories added via + // AddDllDirectory/SetDefaultDllDirectories. This deliberately excludes the + // current working directory from the search order to avoid loading an + // attacker-controlled DLL from the process CWD. + HMODULE handle = LoadLibraryExA(candidate.c_str(), nullptr, LOAD_LIBRARY_SEARCH_DEFAULT_DIRS); + if (handle == nullptr) { + // Fall back to searching the directories listed in PATH (loading by full + // path), matching the pre-existing OS-loader behavior when cuDNN was a + // direct import dependency. The current working directory is never searched. + handle = LoadLibraryFromPathEnv(candidate); + } + if (handle == nullptr) { + error = "LoadLibrary failed for " + candidate + " with error " + std::to_string(GetLastError()); + } + return reinterpret_cast(handle); +#else + dlerror(); + void* handle = dlopen(candidate.c_str(), RTLD_NOW | RTLD_LOCAL); + if (handle == nullptr) { + const char* dl_error = dlerror(); + error = "dlopen failed for " + candidate + ": " + (dl_error != nullptr ? dl_error : "unknown error"); + } + return handle; +#endif +} + +void* GetLibrarySymbol(void* handle, const char* symbol, std::string& error) { +#ifdef _WIN32 + void* address = reinterpret_cast(GetProcAddress(reinterpret_cast(handle), symbol)); + if (address == nullptr) { + error = "GetProcAddress failed for " + std::string(symbol) + " with error " + std::to_string(GetLastError()); + } + return address; +#else + dlerror(); + void* address = dlsym(handle, symbol); + const char* dl_error = dlerror(); + if (address == nullptr || dl_error != nullptr) { + error = "dlsym failed for " + std::string(symbol) + ": " + (dl_error != nullptr ? dl_error : "unknown error"); + } + return address; +#endif +} + +} // namespace + +#if defined(NV_CUDNN_FRONTEND_USE_DYNAMIC_LOADING) +namespace cudnn_frontend { +#ifdef _WIN32 +HMODULE cudnn_dlhandle = nullptr; +#else +void* cudnn_dlhandle = nullptr; +#endif +} // namespace cudnn_frontend +#endif + +namespace onnxruntime::cuda { + +CudnnLibrary& CudnnLibrary::Get() { + static CudnnLibrary library; + return library; +} + +bool CudnnLibrary::Available() { + return EnsureLoaded(); +} + +std::string CudnnLibrary::Error() const { + std::lock_guard lock(mutex_); + return error_.empty() ? CudnnUnavailableErrorString() : error_; +} + +void* CudnnLibrary::Handle() { + return EnsureLoaded() ? handle_ : nullptr; +} + +bool CudnnLibrary::EnsureLoaded() { + std::lock_guard lock(mutex_); + if (load_attempted_) { + return available_; + } + + load_attempted_ = true; + std::string last_error; + for (const auto& candidate : GetCandidateLibraryNames()) { + handle_ = LoadLibraryCandidate(candidate, last_error); + if (handle_ != nullptr) { + available_ = true; + error_.clear(); +#if defined(NV_CUDNN_FRONTEND_USE_DYNAMIC_LOADING) +#ifdef _WIN32 + cudnn_frontend::cudnn_dlhandle = reinterpret_cast(handle_); +#else + cudnn_frontend::cudnn_dlhandle = handle_; +#endif +#endif + return true; + } + } + + available_ = false; + error_ = last_error.empty() ? "cuDNN library was not found" : last_error; + return false; +} + +void* CudnnLibrary::ResolveSymbol(const char* symbol) { + { + std::lock_guard lock(mutex_); + auto it = symbols_.find(symbol); + if (it != symbols_.end()) { + return it->second; + } + } + + if (!EnsureLoaded()) { + return nullptr; + } + + std::lock_guard lock(mutex_); + std::string symbol_error; + void* address = GetLibrarySymbol(handle_, symbol, symbol_error); + if (address == nullptr) { + available_ = false; + error_ = symbol_error; + return nullptr; + } + + symbols_.emplace(symbol, address); + return address; +} + +const char* CudnnUnavailableErrorString() { + return "cuDNN is not available. Install cuDNN, update the system library search path, or set enable_cudnn=0 to force native CUDA paths where available."; +} + +} // namespace onnxruntime::cuda + +#endif // USE_CUDA_MINIMAL diff --git a/onnxruntime/core/providers/cuda/cudnn_loader.h b/onnxruntime/core/providers/cuda/cudnn_loader.h new file mode 100644 index 0000000000000..8b6edc74e9441 --- /dev/null +++ b/onnxruntime/core/providers/cuda/cudnn_loader.h @@ -0,0 +1,47 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#ifndef USE_CUDA_MINIMAL + +#include +#include +#include + +#include "core/providers/cuda/cuda_pch.h" + +namespace onnxruntime::cuda { + +class CudnnLibrary { + public: + static CudnnLibrary& Get(); + + bool Available(); + std::string Error() const; + void* Handle(); + + template + T Resolve(const char* symbol) { + return reinterpret_cast(ResolveSymbol(symbol)); + } + + private: + CudnnLibrary() = default; + + bool EnsureLoaded(); + void* ResolveSymbol(const char* symbol); + + mutable std::mutex mutex_; + bool load_attempted_{false}; + bool available_{false}; + std::string error_; + void* handle_{nullptr}; + std::unordered_map symbols_; +}; + +const char* CudnnUnavailableErrorString(); + +} // namespace onnxruntime::cuda + +#endif // USE_CUDA_MINIMAL diff --git a/onnxruntime/core/providers/cuda/cudnn_stub.cc b/onnxruntime/core/providers/cuda/cudnn_stub.cc new file mode 100644 index 0000000000000..97a0c2e05299d --- /dev/null +++ b/onnxruntime/core/providers/cuda/cudnn_stub.cc @@ -0,0 +1,401 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "core/providers/cuda/cudnn_loader.h" + +#ifndef USE_CUDA_MINIMAL + +#define ORT_CUDNN_FORWARD_STATUS(name, ...) \ + using Fn = decltype(&name); \ + auto fn = onnxruntime::cuda::CudnnLibrary::Get().Resolve(#name); \ + return fn != nullptr ? fn(__VA_ARGS__) : CUDNN_STATUS_NOT_INITIALIZED + +extern "C" { + +size_t CUDNNWINAPI cudnnGetVersion(void) { + using Fn = decltype(&cudnnGetVersion); + auto fn = onnxruntime::cuda::CudnnLibrary::Get().Resolve("cudnnGetVersion"); + return fn != nullptr ? fn() : 0; +} + +const char* CUDNNWINAPI cudnnGetErrorString(cudnnStatus_t status) { + using Fn = decltype(&cudnnGetErrorString); + auto fn = onnxruntime::cuda::CudnnLibrary::Get().Resolve("cudnnGetErrorString"); + return fn != nullptr ? fn(status) : onnxruntime::cuda::CudnnUnavailableErrorString(); +} + +cudnnStatus_t CUDNNWINAPI cudnnCreate(cudnnHandle_t* handle) { + ORT_CUDNN_FORWARD_STATUS(cudnnCreate, handle); +} + +cudnnStatus_t CUDNNWINAPI cudnnDestroy(cudnnHandle_t handle) { + ORT_CUDNN_FORWARD_STATUS(cudnnDestroy, handle); +} + +cudnnStatus_t CUDNNWINAPI cudnnSetStream(cudnnHandle_t handle, cudaStream_t streamId) { + ORT_CUDNN_FORWARD_STATUS(cudnnSetStream, handle, streamId); +} + +cudnnStatus_t CUDNNWINAPI cudnnCreateTensorDescriptor(cudnnTensorDescriptor_t* tensorDesc) { + ORT_CUDNN_FORWARD_STATUS(cudnnCreateTensorDescriptor, tensorDesc); +} + +cudnnStatus_t CUDNNWINAPI cudnnDestroyTensorDescriptor(cudnnTensorDescriptor_t tensorDesc) { + ORT_CUDNN_FORWARD_STATUS(cudnnDestroyTensorDescriptor, tensorDesc); +} + +cudnnStatus_t CUDNNWINAPI cudnnSetTensorNdDescriptor(cudnnTensorDescriptor_t tensorDesc, cudnnDataType_t dataType, + int nbDims, const int dimA[], const int strideA[]) { + ORT_CUDNN_FORWARD_STATUS(cudnnSetTensorNdDescriptor, tensorDesc, dataType, nbDims, dimA, strideA); +} + +cudnnStatus_t CUDNNWINAPI cudnnSetTensor4dDescriptor(cudnnTensorDescriptor_t tensorDesc, cudnnTensorFormat_t format, + cudnnDataType_t dataType, int n, int c, int h, int w) { + ORT_CUDNN_FORWARD_STATUS(cudnnSetTensor4dDescriptor, tensorDesc, format, dataType, n, c, h, w); +} + +cudnnStatus_t CUDNNWINAPI cudnnGetTensorNdDescriptor(const cudnnTensorDescriptor_t tensorDesc, int nbDimsRequested, + cudnnDataType_t* dataType, int* nbDims, int dimA[], int strideA[]) { + ORT_CUDNN_FORWARD_STATUS(cudnnGetTensorNdDescriptor, tensorDesc, nbDimsRequested, dataType, nbDims, dimA, strideA); +} + +cudnnStatus_t CUDNNWINAPI cudnnCreateFilterDescriptor(cudnnFilterDescriptor_t* filterDesc) { + ORT_CUDNN_FORWARD_STATUS(cudnnCreateFilterDescriptor, filterDesc); +} + +cudnnStatus_t CUDNNWINAPI cudnnDestroyFilterDescriptor(cudnnFilterDescriptor_t filterDesc) { + ORT_CUDNN_FORWARD_STATUS(cudnnDestroyFilterDescriptor, filterDesc); +} + +cudnnStatus_t CUDNNWINAPI cudnnSetFilterNdDescriptor(cudnnFilterDescriptor_t filterDesc, cudnnDataType_t dataType, + cudnnTensorFormat_t format, int nbDims, const int filterDimA[]) { + ORT_CUDNN_FORWARD_STATUS(cudnnSetFilterNdDescriptor, filterDesc, dataType, format, nbDims, filterDimA); +} + +cudnnStatus_t CUDNNWINAPI cudnnSetFilter4dDescriptor(cudnnFilterDescriptor_t filterDesc, cudnnDataType_t dataType, + cudnnTensorFormat_t format, int k, int c, int h, int w) { + ORT_CUDNN_FORWARD_STATUS(cudnnSetFilter4dDescriptor, filterDesc, dataType, format, k, c, h, w); +} + +cudnnStatus_t CUDNNWINAPI cudnnCreateConvolutionDescriptor(cudnnConvolutionDescriptor_t* convDesc) { + ORT_CUDNN_FORWARD_STATUS(cudnnCreateConvolutionDescriptor, convDesc); +} + +cudnnStatus_t CUDNNWINAPI cudnnDestroyConvolutionDescriptor(cudnnConvolutionDescriptor_t convDesc) { + ORT_CUDNN_FORWARD_STATUS(cudnnDestroyConvolutionDescriptor, convDesc); +} + +cudnnStatus_t CUDNNWINAPI cudnnSetConvolutionNdDescriptor(cudnnConvolutionDescriptor_t convDesc, int arrayLength, + const int padA[], const int filterStrideA[], + const int dilationA[], cudnnConvolutionMode_t mode, + cudnnDataType_t computeType) { + ORT_CUDNN_FORWARD_STATUS(cudnnSetConvolutionNdDescriptor, convDesc, arrayLength, padA, filterStrideA, dilationA, mode, + computeType); +} + +cudnnStatus_t CUDNNWINAPI cudnnSetConvolutionGroupCount(cudnnConvolutionDescriptor_t convDesc, int groupCount) { + ORT_CUDNN_FORWARD_STATUS(cudnnSetConvolutionGroupCount, convDesc, groupCount); +} + +cudnnStatus_t CUDNNWINAPI cudnnSetConvolutionMathType(cudnnConvolutionDescriptor_t convDesc, + cudnnMathType_t mathType) { + ORT_CUDNN_FORWARD_STATUS(cudnnSetConvolutionMathType, convDesc, mathType); +} + +cudnnStatus_t CUDNNWINAPI cudnnCreateActivationDescriptor(cudnnActivationDescriptor_t* activationDesc) { + ORT_CUDNN_FORWARD_STATUS(cudnnCreateActivationDescriptor, activationDesc); +} + +cudnnStatus_t CUDNNWINAPI cudnnDestroyActivationDescriptor(cudnnActivationDescriptor_t activationDesc) { + ORT_CUDNN_FORWARD_STATUS(cudnnDestroyActivationDescriptor, activationDesc); +} + +cudnnStatus_t CUDNNWINAPI cudnnSetActivationDescriptor(cudnnActivationDescriptor_t activationDesc, + cudnnActivationMode_t mode, + cudnnNanPropagation_t reluNanOpt, double coef) { + ORT_CUDNN_FORWARD_STATUS(cudnnSetActivationDescriptor, activationDesc, mode, reluNanOpt, coef); +} + +cudnnStatus_t CUDNNWINAPI cudnnCreatePoolingDescriptor(cudnnPoolingDescriptor_t* poolingDesc) { + ORT_CUDNN_FORWARD_STATUS(cudnnCreatePoolingDescriptor, poolingDesc); +} + +cudnnStatus_t CUDNNWINAPI cudnnDestroyPoolingDescriptor(cudnnPoolingDescriptor_t poolingDesc) { + ORT_CUDNN_FORWARD_STATUS(cudnnDestroyPoolingDescriptor, poolingDesc); +} + +cudnnStatus_t CUDNNWINAPI cudnnSetPoolingNdDescriptor(cudnnPoolingDescriptor_t poolingDesc, + const cudnnPoolingMode_t mode, + const cudnnNanPropagation_t maxpoolingNanOpt, int nbDims, + const int windowDimA[], const int paddingA[], + const int strideA[]) { + ORT_CUDNN_FORWARD_STATUS(cudnnSetPoolingNdDescriptor, poolingDesc, mode, maxpoolingNanOpt, nbDims, windowDimA, + paddingA, strideA); +} + +cudnnStatus_t CUDNNWINAPI cudnnCreateLRNDescriptor(cudnnLRNDescriptor_t* normDesc) { + ORT_CUDNN_FORWARD_STATUS(cudnnCreateLRNDescriptor, normDesc); +} + +cudnnStatus_t CUDNNWINAPI cudnnDestroyLRNDescriptor(cudnnLRNDescriptor_t lrnDesc) { + ORT_CUDNN_FORWARD_STATUS(cudnnDestroyLRNDescriptor, lrnDesc); +} + +cudnnStatus_t CUDNNWINAPI cudnnSetLRNDescriptor(cudnnLRNDescriptor_t normDesc, unsigned lrnN, double lrnAlpha, + double lrnBeta, double lrnK) { + ORT_CUDNN_FORWARD_STATUS(cudnnSetLRNDescriptor, normDesc, lrnN, lrnAlpha, lrnBeta, lrnK); +} + +cudnnStatus_t CUDNNWINAPI cudnnCreateReduceTensorDescriptor(cudnnReduceTensorDescriptor_t* reduceTensorDesc) { + ORT_CUDNN_FORWARD_STATUS(cudnnCreateReduceTensorDescriptor, reduceTensorDesc); +} + +cudnnStatus_t CUDNNWINAPI cudnnDestroyReduceTensorDescriptor(cudnnReduceTensorDescriptor_t reduceTensorDesc) { + ORT_CUDNN_FORWARD_STATUS(cudnnDestroyReduceTensorDescriptor, reduceTensorDesc); +} + +cudnnStatus_t CUDNNWINAPI cudnnSetReduceTensorDescriptor(cudnnReduceTensorDescriptor_t reduceTensorDesc, + cudnnReduceTensorOp_t reduceTensorOp, + cudnnDataType_t reduceTensorCompType, + cudnnNanPropagation_t reduceTensorNanOpt, + cudnnReduceTensorIndices_t reduceTensorIndices, + cudnnIndicesType_t reduceTensorIndicesType) { + ORT_CUDNN_FORWARD_STATUS(cudnnSetReduceTensorDescriptor, reduceTensorDesc, reduceTensorOp, reduceTensorCompType, + reduceTensorNanOpt, reduceTensorIndices, reduceTensorIndicesType); +} + +cudnnStatus_t CUDNNWINAPI cudnnCreateRNNDescriptor(cudnnRNNDescriptor_t* rnnDesc) { + ORT_CUDNN_FORWARD_STATUS(cudnnCreateRNNDescriptor, rnnDesc); +} + +cudnnStatus_t CUDNNWINAPI cudnnDestroyRNNDescriptor(cudnnRNNDescriptor_t rnnDesc) { + ORT_CUDNN_FORWARD_STATUS(cudnnDestroyRNNDescriptor, rnnDesc); +} + +cudnnStatus_t CUDNNWINAPI cudnnSetRNNDescriptor_v8(cudnnRNNDescriptor_t rnnDesc, cudnnRNNAlgo_t algo, + cudnnRNNMode_t cellMode, cudnnRNNBiasMode_t biasMode, + cudnnDirectionMode_t dirMode, cudnnRNNInputMode_t inputMode, + cudnnDataType_t dataType, cudnnDataType_t mathPrec, + cudnnMathType_t mathType, int32_t inputSize, int32_t hiddenSize, + int32_t projSize, int32_t numLayers, + cudnnDropoutDescriptor_t dropoutDesc, uint32_t auxFlags) { + ORT_CUDNN_FORWARD_STATUS(cudnnSetRNNDescriptor_v8, rnnDesc, algo, cellMode, biasMode, dirMode, inputMode, dataType, + mathPrec, mathType, inputSize, hiddenSize, projSize, numLayers, dropoutDesc, auxFlags); +} + +cudnnStatus_t CUDNNWINAPI cudnnCreateRNNDataDescriptor(cudnnRNNDataDescriptor_t* rnnDataDesc) { + ORT_CUDNN_FORWARD_STATUS(cudnnCreateRNNDataDescriptor, rnnDataDesc); +} + +cudnnStatus_t CUDNNWINAPI cudnnDestroyRNNDataDescriptor(cudnnRNNDataDescriptor_t rnnDataDesc) { + ORT_CUDNN_FORWARD_STATUS(cudnnDestroyRNNDataDescriptor, rnnDataDesc); +} + +cudnnStatus_t CUDNNWINAPI cudnnSetRNNDataDescriptor(cudnnRNNDataDescriptor_t rnnDataDesc, cudnnDataType_t dataType, + cudnnRNNDataLayout_t layout, int maxSeqLength, int batchSize, + int vectorSize, const int seqLengthArray[], void* paddingFill) { + ORT_CUDNN_FORWARD_STATUS(cudnnSetRNNDataDescriptor, rnnDataDesc, dataType, layout, maxSeqLength, batchSize, + vectorSize, seqLengthArray, paddingFill); +} + +cudnnStatus_t CUDNNWINAPI cudnnCreateDropoutDescriptor(cudnnDropoutDescriptor_t* dropoutDesc) { + ORT_CUDNN_FORWARD_STATUS(cudnnCreateDropoutDescriptor, dropoutDesc); +} + +cudnnStatus_t CUDNNWINAPI cudnnDestroyDropoutDescriptor(cudnnDropoutDescriptor_t dropoutDesc) { + ORT_CUDNN_FORWARD_STATUS(cudnnDestroyDropoutDescriptor, dropoutDesc); +} + +cudnnStatus_t CUDNNWINAPI cudnnSetDropoutDescriptor(cudnnDropoutDescriptor_t dropoutDesc, cudnnHandle_t handle, + float dropout, void* states, size_t stateSizeInBytes, + unsigned long long seed) { + ORT_CUDNN_FORWARD_STATUS(cudnnSetDropoutDescriptor, dropoutDesc, handle, dropout, states, stateSizeInBytes, seed); +} + +cudnnStatus_t CUDNNWINAPI cudnnDropoutGetStatesSize(cudnnHandle_t handle, size_t* sizeInBytes) { + ORT_CUDNN_FORWARD_STATUS(cudnnDropoutGetStatesSize, handle, sizeInBytes); +} + +cudnnStatus_t CUDNNWINAPI cudnnDeriveBNTensorDescriptor(cudnnTensorDescriptor_t derivedBnDesc, + const cudnnTensorDescriptor_t xDesc, + cudnnBatchNormMode_t mode) { + ORT_CUDNN_FORWARD_STATUS(cudnnDeriveBNTensorDescriptor, derivedBnDesc, xDesc, mode); +} + +cudnnStatus_t CUDNNWINAPI cudnnAddTensor(cudnnHandle_t handle, const void* alpha, + const cudnnTensorDescriptor_t aDesc, const void* A, const void* beta, + const cudnnTensorDescriptor_t cDesc, void* C) { + ORT_CUDNN_FORWARD_STATUS(cudnnAddTensor, handle, alpha, aDesc, A, beta, cDesc, C); +} + +cudnnStatus_t CUDNNWINAPI cudnnActivationForward(cudnnHandle_t handle, cudnnActivationDescriptor_t activationDesc, + const void* alpha, const cudnnTensorDescriptor_t xDesc, const void* x, + const void* beta, const cudnnTensorDescriptor_t yDesc, void* y) { + ORT_CUDNN_FORWARD_STATUS(cudnnActivationForward, handle, activationDesc, alpha, xDesc, x, beta, yDesc, y); +} + +cudnnStatus_t CUDNNWINAPI cudnnPoolingForward(cudnnHandle_t handle, const cudnnPoolingDescriptor_t poolingDesc, + const void* alpha, const cudnnTensorDescriptor_t xDesc, const void* x, + const void* beta, const cudnnTensorDescriptor_t yDesc, void* y) { + ORT_CUDNN_FORWARD_STATUS(cudnnPoolingForward, handle, poolingDesc, alpha, xDesc, x, beta, yDesc, y); +} + +cudnnStatus_t CUDNNWINAPI cudnnLRNCrossChannelForward(cudnnHandle_t handle, cudnnLRNDescriptor_t normDesc, + cudnnLRNMode_t lrnMode, const void* alpha, + const cudnnTensorDescriptor_t xDesc, const void* x, + const void* beta, const cudnnTensorDescriptor_t yDesc, void* y) { + ORT_CUDNN_FORWARD_STATUS(cudnnLRNCrossChannelForward, handle, normDesc, lrnMode, alpha, xDesc, x, beta, yDesc, y); +} + +cudnnStatus_t CUDNNWINAPI cudnnSoftmaxForward(cudnnHandle_t handle, cudnnSoftmaxAlgorithm_t algo, + cudnnSoftmaxMode_t mode, const void* alpha, + const cudnnTensorDescriptor_t xDesc, const void* x, const void* beta, + const cudnnTensorDescriptor_t yDesc, void* y) { + ORT_CUDNN_FORWARD_STATUS(cudnnSoftmaxForward, handle, algo, mode, alpha, xDesc, x, beta, yDesc, y); +} + +cudnnStatus_t CUDNNWINAPI cudnnSoftmaxBackward(cudnnHandle_t handle, cudnnSoftmaxAlgorithm_t algo, + cudnnSoftmaxMode_t mode, const void* alpha, + const cudnnTensorDescriptor_t yDesc, const void* y, + const cudnnTensorDescriptor_t dyDesc, const void* dy, const void* beta, + const cudnnTensorDescriptor_t dxDesc, void* dx) { + ORT_CUDNN_FORWARD_STATUS(cudnnSoftmaxBackward, handle, algo, mode, alpha, yDesc, y, dyDesc, dy, beta, dxDesc, dx); +} + +cudnnStatus_t CUDNNWINAPI cudnnBatchNormalizationForwardInference( + cudnnHandle_t handle, cudnnBatchNormMode_t mode, const void* alpha, const void* beta, + const cudnnTensorDescriptor_t xDesc, const void* x, const cudnnTensorDescriptor_t yDesc, void* y, + const cudnnTensorDescriptor_t bnScaleBiasMeanVarDesc, const void* bnScale, const void* bnBias, + const void* estimatedMean, const void* estimatedVariance, double epsilon) { + ORT_CUDNN_FORWARD_STATUS(cudnnBatchNormalizationForwardInference, handle, mode, alpha, beta, xDesc, x, yDesc, y, + bnScaleBiasMeanVarDesc, bnScale, bnBias, estimatedMean, estimatedVariance, epsilon); +} + +cudnnStatus_t CUDNNWINAPI cudnnBatchNormalizationForwardTraining( + cudnnHandle_t handle, cudnnBatchNormMode_t mode, const void* alpha, const void* beta, + const cudnnTensorDescriptor_t xDesc, const void* x, const cudnnTensorDescriptor_t yDesc, void* y, + const cudnnTensorDescriptor_t bnScaleBiasMeanVarDesc, const void* bnScale, const void* bnBias, + double exponentialAverageFactor, void* resultRunningMean, void* resultRunningVariance, double epsilon, + void* resultSaveMean, void* resultSaveInvVariance) { + ORT_CUDNN_FORWARD_STATUS(cudnnBatchNormalizationForwardTraining, handle, mode, alpha, beta, xDesc, x, yDesc, y, + bnScaleBiasMeanVarDesc, bnScale, bnBias, exponentialAverageFactor, resultRunningMean, + resultRunningVariance, epsilon, resultSaveMean, resultSaveInvVariance); +} + +cudnnStatus_t CUDNNWINAPI cudnnConvolutionForward( + cudnnHandle_t handle, const void* alpha, const cudnnTensorDescriptor_t xDesc, const void* x, + const cudnnFilterDescriptor_t wDesc, const void* w, const cudnnConvolutionDescriptor_t convDesc, + cudnnConvolutionFwdAlgo_t algo, void* workSpace, size_t workSpaceSizeInBytes, const void* beta, + const cudnnTensorDescriptor_t yDesc, void* y) { + ORT_CUDNN_FORWARD_STATUS(cudnnConvolutionForward, handle, alpha, xDesc, x, wDesc, w, convDesc, algo, workSpace, + workSpaceSizeInBytes, beta, yDesc, y); +} + +cudnnStatus_t CUDNNWINAPI cudnnConvolutionBiasActivationForward( + cudnnHandle_t handle, const void* alpha1, const cudnnTensorDescriptor_t xDesc, const void* x, + const cudnnFilterDescriptor_t wDesc, const void* w, const cudnnConvolutionDescriptor_t convDesc, + cudnnConvolutionFwdAlgo_t algo, void* workSpace, size_t workSpaceSizeInBytes, const void* alpha2, + const cudnnTensorDescriptor_t zDesc, const void* z, const cudnnTensorDescriptor_t biasDesc, const void* bias, + const cudnnActivationDescriptor_t activationDesc, const cudnnTensorDescriptor_t yDesc, void* y) { + ORT_CUDNN_FORWARD_STATUS(cudnnConvolutionBiasActivationForward, handle, alpha1, xDesc, x, wDesc, w, convDesc, algo, + workSpace, workSpaceSizeInBytes, alpha2, zDesc, z, biasDesc, bias, activationDesc, yDesc, + y); +} + +cudnnStatus_t CUDNNWINAPI cudnnConvolutionBackwardData( + cudnnHandle_t handle, const void* alpha, const cudnnFilterDescriptor_t wDesc, const void* w, + const cudnnTensorDescriptor_t dyDesc, const void* dy, const cudnnConvolutionDescriptor_t convDesc, + cudnnConvolutionBwdDataAlgo_t algo, void* workSpace, size_t workSpaceSizeInBytes, const void* beta, + const cudnnTensorDescriptor_t dxDesc, void* dx) { + ORT_CUDNN_FORWARD_STATUS(cudnnConvolutionBackwardData, handle, alpha, wDesc, w, dyDesc, dy, convDesc, algo, + workSpace, workSpaceSizeInBytes, beta, dxDesc, dx); +} + +cudnnStatus_t CUDNNWINAPI cudnnFindConvolutionForwardAlgorithmEx( + cudnnHandle_t handle, const cudnnTensorDescriptor_t xDesc, const void* x, const cudnnFilterDescriptor_t wDesc, + const void* w, const cudnnConvolutionDescriptor_t convDesc, const cudnnTensorDescriptor_t yDesc, void* y, + const int requestedAlgoCount, int* returnedAlgoCount, cudnnConvolutionFwdAlgoPerf_t* perfResults, + void* workSpace, size_t workSpaceSizeInBytes) { + ORT_CUDNN_FORWARD_STATUS(cudnnFindConvolutionForwardAlgorithmEx, handle, xDesc, x, wDesc, w, convDesc, yDesc, y, + requestedAlgoCount, returnedAlgoCount, perfResults, workSpace, workSpaceSizeInBytes); +} + +cudnnStatus_t CUDNNWINAPI cudnnFindConvolutionBackwardDataAlgorithmEx( + cudnnHandle_t handle, const cudnnFilterDescriptor_t wDesc, const void* w, const cudnnTensorDescriptor_t dyDesc, + const void* dy, const cudnnConvolutionDescriptor_t convDesc, const cudnnTensorDescriptor_t dxDesc, void* dx, + const int requestedAlgoCount, int* returnedAlgoCount, cudnnConvolutionBwdDataAlgoPerf_t* perfResults, + void* workSpace, size_t workSpaceSizeInBytes) { + ORT_CUDNN_FORWARD_STATUS(cudnnFindConvolutionBackwardDataAlgorithmEx, handle, wDesc, w, dyDesc, dy, convDesc, dxDesc, + dx, requestedAlgoCount, returnedAlgoCount, perfResults, workSpace, workSpaceSizeInBytes); +} + +cudnnStatus_t CUDNNWINAPI cudnnGetConvolutionForwardAlgorithm_v7( + cudnnHandle_t handle, const cudnnTensorDescriptor_t srcDesc, const cudnnFilterDescriptor_t filterDesc, + const cudnnConvolutionDescriptor_t convDesc, const cudnnTensorDescriptor_t destDesc, const int requestedAlgoCount, + int* returnedAlgoCount, cudnnConvolutionFwdAlgoPerf_t* perfResults) { + ORT_CUDNN_FORWARD_STATUS(cudnnGetConvolutionForwardAlgorithm_v7, handle, srcDesc, filterDesc, convDesc, destDesc, + requestedAlgoCount, returnedAlgoCount, perfResults); +} + +cudnnStatus_t CUDNNWINAPI cudnnGetConvolutionForwardWorkspaceSize( + cudnnHandle_t handle, const cudnnTensorDescriptor_t xDesc, const cudnnFilterDescriptor_t wDesc, + const cudnnConvolutionDescriptor_t convDesc, const cudnnTensorDescriptor_t yDesc, cudnnConvolutionFwdAlgo_t algo, + size_t* sizeInBytes) { + ORT_CUDNN_FORWARD_STATUS(cudnnGetConvolutionForwardWorkspaceSize, handle, xDesc, wDesc, convDesc, yDesc, algo, + sizeInBytes); +} + +cudnnStatus_t CUDNNWINAPI cudnnGetReductionIndicesSize(cudnnHandle_t handle, + const cudnnReduceTensorDescriptor_t reduceTensorDesc, + const cudnnTensorDescriptor_t aDesc, + const cudnnTensorDescriptor_t cDesc, size_t* sizeInBytes) { + ORT_CUDNN_FORWARD_STATUS(cudnnGetReductionIndicesSize, handle, reduceTensorDesc, aDesc, cDesc, sizeInBytes); +} + +cudnnStatus_t CUDNNWINAPI cudnnGetReductionWorkspaceSize(cudnnHandle_t handle, + const cudnnReduceTensorDescriptor_t reduceTensorDesc, + const cudnnTensorDescriptor_t aDesc, + const cudnnTensorDescriptor_t cDesc, size_t* sizeInBytes) { + ORT_CUDNN_FORWARD_STATUS(cudnnGetReductionWorkspaceSize, handle, reduceTensorDesc, aDesc, cDesc, sizeInBytes); +} + +cudnnStatus_t CUDNNWINAPI cudnnReduceTensor(cudnnHandle_t handle, const cudnnReduceTensorDescriptor_t reduceTensorDesc, + void* indices, size_t indicesSizeInBytes, void* workspace, + size_t workspaceSizeInBytes, const void* alpha, + const cudnnTensorDescriptor_t aDesc, const void* A, const void* beta, + const cudnnTensorDescriptor_t cDesc, void* C) { + ORT_CUDNN_FORWARD_STATUS(cudnnReduceTensor, handle, reduceTensorDesc, indices, indicesSizeInBytes, workspace, + workspaceSizeInBytes, alpha, aDesc, A, beta, cDesc, C); +} + +cudnnStatus_t CUDNNWINAPI cudnnGetRNNTempSpaceSizes(cudnnHandle_t handle, cudnnRNNDescriptor_t rnnDesc, + cudnnForwardMode_t fwdMode, cudnnRNNDataDescriptor_t xDesc, + size_t* workSpaceSize, size_t* reserveSpaceSize) { + ORT_CUDNN_FORWARD_STATUS(cudnnGetRNNTempSpaceSizes, handle, rnnDesc, fwdMode, xDesc, workSpaceSize, + reserveSpaceSize); +} + +cudnnStatus_t CUDNNWINAPI cudnnGetRNNWeightParams( + cudnnHandle_t handle, cudnnRNNDescriptor_t rnnDesc, int32_t pseudoLayer, size_t weightSpaceSize, + const void* weightSpace, int32_t linLayerID, cudnnTensorDescriptor_t mDesc, void** mAddr, + cudnnTensorDescriptor_t bDesc, void** bAddr) { + ORT_CUDNN_FORWARD_STATUS(cudnnGetRNNWeightParams, handle, rnnDesc, pseudoLayer, weightSpaceSize, weightSpace, + linLayerID, mDesc, mAddr, bDesc, bAddr); +} + +cudnnStatus_t CUDNNWINAPI cudnnRNNForward( + cudnnHandle_t handle, cudnnRNNDescriptor_t rnnDesc, cudnnForwardMode_t fwdMode, const int32_t devSeqLengths[], + cudnnRNNDataDescriptor_t xDesc, const void* x, cudnnRNNDataDescriptor_t yDesc, void* y, + cudnnTensorDescriptor_t hDesc, const void* hx, void* hy, cudnnTensorDescriptor_t cDesc, const void* cx, void* cy, + size_t weightSpaceSize, const void* weightSpace, size_t workSpaceSize, void* workSpace, size_t reserveSpaceSize, + void* reserveSpace) { + ORT_CUDNN_FORWARD_STATUS(cudnnRNNForward, handle, rnnDesc, fwdMode, devSeqLengths, xDesc, x, yDesc, y, hDesc, hx, + hy, cDesc, cx, cy, weightSpaceSize, weightSpace, workSpaceSize, workSpace, + reserveSpaceSize, reserveSpace); +} + +} // extern "C" + +#undef ORT_CUDNN_FORWARD_STATUS + +#endif // USE_CUDA_MINIMAL diff --git a/onnxruntime/core/providers/cuda/plugin/cuda_ep.cc b/onnxruntime/core/providers/cuda/plugin/cuda_ep.cc index ac2ac04b26cfe..b4d7cf7a1ade2 100644 --- a/onnxruntime/core/providers/cuda/plugin/cuda_ep.cc +++ b/onnxruntime/core/providers/cuda/plugin/cuda_ep.cc @@ -3,6 +3,7 @@ #include "cuda_ep.h" #include "cuda_ep_factory.h" +#include "core/providers/cuda/cudnn_loader.h" #include "cuda_stream_plugin.h" #include "cuda_graph_plugin.h" #include "core/providers/cuda/plugin/cuda_kernel_adapter.h" @@ -205,6 +206,7 @@ CudaEp::CudaEp(CudaEpFactory& factory, const Config& config, const OrtLogger& lo adapter_config.cudnn_conv_algo = config_.cudnn_conv_algo; adapter_config.cudnn_conv_use_max_workspace = config_.cudnn_conv_use_max_workspace; adapter_config.cudnn_conv1d_pad_to_nc1d = config_.cudnn_conv1d_pad_to_nc1d; + adapter_config.enable_cudnn = config_.enable_cudnn; adapter_config.fuse_conv_bias = config_.fuse_conv_bias; adapter_config.sdpa_kernel = config_.sdpa_kernel; adapter_config.device_id = config_.device_id; @@ -401,7 +403,7 @@ OrtStatus* ORT_API_CALL CudaEp::CreateSyncStreamForDeviceImpl( return Ort::GetApi().CreateStatus(ORT_INVALID_ARGUMENT, error.c_str()); } - auto cuda_stream = std::make_unique(ep->factory_, device_id, this_ptr); + auto cuda_stream = std::make_unique(ep->factory_, device_id, ep->config_.enable_cudnn, this_ptr); if (ep->config_.has_user_compute_stream) { // A user-provided compute stream is honored for kernels regardless of whether CUDA graph diff --git a/onnxruntime/core/providers/cuda/plugin/cuda_ep.h b/onnxruntime/core/providers/cuda/plugin/cuda_ep.h index b05d4a2d3146d..68c515a028ac8 100644 --- a/onnxruntime/core/providers/cuda/plugin/cuda_ep.h +++ b/onnxruntime/core/providers/cuda/plugin/cuda_ep.h @@ -30,6 +30,7 @@ class CudaEp : public onnxruntime::ep::adapter::Ep { int cudnn_conv_algo = 0; ///< cuDNN convolution algorithm selection. bool cudnn_conv_use_max_workspace = true; ///< Use maximum workspace for cuDNN conv algo search. bool cudnn_conv1d_pad_to_nc1d = false; ///< Pad 1D convolutions to NC1D format. + bool enable_cudnn = true; ///< Enable runtime loading and use of cuDNN-backed kernels. bool fuse_conv_bias = false; ///< Enable cuDNN frontend conv+bias fusion. int sdpa_kernel = 0; ///< Attention backend bitmask override. bool enable_cuda_graph = false; ///< Enable CUDA graph capture and replay. diff --git a/onnxruntime/core/providers/cuda/plugin/cuda_ep_factory.cc b/onnxruntime/core/providers/cuda/plugin/cuda_ep_factory.cc index 8d4f7c7c06244..00d4ea5bef879 100644 --- a/onnxruntime/core/providers/cuda/plugin/cuda_ep_factory.cc +++ b/onnxruntime/core/providers/cuda/plugin/cuda_ep_factory.cc @@ -486,6 +486,7 @@ OrtStatus* ORT_API_CALL CudaEpFactory::CreateEpImpl( const std::string cudnn_conv1d_pad_key = ep_options_prefix + "cudnn_conv1d_pad_to_nc1d"; const std::string cudnn_conv_algo_key = ep_options_prefix + "cudnn_conv_algo"; const std::string cudnn_conv_algo_search_key = ep_options_prefix + "cudnn_conv_algo_search"; + const std::string enable_cudnn_key = ep_options_prefix + "enable_cudnn"; const std::string fuse_conv_bias_key = ep_options_prefix + "fuse_conv_bias"; const std::string sdpa_kernel_key = ep_options_prefix + "sdpa_kernel"; const std::string enable_cuda_graph_key = ep_options_prefix + "enable_cuda_graph"; @@ -514,6 +515,9 @@ OrtStatus* ORT_API_CALL CudaEpFactory::CreateEpImpl( {cudnn_conv_algo_search_key, cudnn_conv_algo_key, "ep.cuda.cudnn_conv_algo_search", "ep.cuda.cudnn_conv_algo", "cudnn_conv_algo_search", "cudnn_conv_algo"}, config.cudnn_conv_algo); + read_session_config_bool( + {enable_cudnn_key, "ep.cuda.enable_cudnn", "enable_cudnn"}, + config.enable_cudnn); read_session_config_bool( {fuse_conv_bias_key, "ep.cuda.fuse_conv_bias", "fuse_conv_bias"}, config.fuse_conv_bias); @@ -881,7 +885,11 @@ OrtStatus* ORT_API_CALL CudaEpFactory::CreateSyncStreamForDeviceImpl( auto* factory = static_cast(this_ptr); int req_device_id = factory->ep_api_.MemoryDevice_GetDeviceId(memory_device); - auto cuda_stream = std::make_unique(*factory, req_device_id, nullptr); + // Factory-level streams are not tied to a specific EP instance's enable_cudnn policy. Default cuDNN + // off here so this path never triggers a cuDNN load or handle creation; kernels that need cuDNN run + // on EP-owned streams created with the EP's actual enable_cudnn setting, and otherwise fall back to + // the per-thread default cuDNN handle. + auto cuda_stream = std::make_unique(*factory, req_device_id, false, nullptr); // Initialize CUDA handles (stream, cuBLAS, cuDNN) RETURN_IF_ERROR(cuda_stream->InitHandles()); diff --git a/onnxruntime/core/providers/cuda/plugin/cuda_kernel_adapter.h b/onnxruntime/core/providers/cuda/plugin/cuda_kernel_adapter.h index b6b563475b00b..ace200e9c2943 100644 --- a/onnxruntime/core/providers/cuda/plugin/cuda_kernel_adapter.h +++ b/onnxruntime/core/providers/cuda/plugin/cuda_kernel_adapter.h @@ -33,6 +33,7 @@ #include #include #include "core/providers/cuda/shared_inc/cuda_call.h" +#include "core/providers/cuda/cudnn_loader.h" #include "contrib_ops/cuda/bert/attention_kernel_options.h" #ifdef __CUDACC__ @@ -264,6 +265,11 @@ using ::onnxruntime::HandleNegativeAxis; { \ cudnnStatus_t _status = (expr); \ if (_status != CUDNN_STATUS_SUCCESS) { \ + if (!onnxruntime::cuda::CudnnLibrary::Get().Available()) { \ + return onnxruntime::common::Status(onnxruntime::common::ONNXRUNTIME, onnxruntime::common::NOT_IMPLEMENTED, \ + std::string("cuDNN is unavailable for CUDA Plugin Execution Provider: ") + \ + onnxruntime::cuda::CudnnLibrary::Get().Error()); \ + } \ return onnxruntime::common::Status(onnxruntime::common::ONNXRUNTIME, onnxruntime::common::FAIL, std::string("cuDNN error: ") + cudnnGetErrorString(_status)); \ } \ } @@ -520,6 +526,7 @@ struct CudaKernelAdapterRuntimeConfig { int cudnn_conv_algo = 0; bool cudnn_conv_use_max_workspace = true; bool cudnn_conv1d_pad_to_nc1d = false; + bool enable_cudnn = true; bool fuse_conv_bias = false; int sdpa_kernel = 0; int device_id = 0; @@ -587,6 +594,10 @@ inline DefaultCudaHandles& GetDefaultCudaHandlesForDevice(int device_id) { // Fallback handles are only used for code paths that need cuBLAS/cuDNN // without an active CudaSyncStream. Keep them thread-local so they are not // shared across callers that may use the libraries concurrently. + // + // Only cuBLAS/cuBLASLt are created here. The cuDNN fallback handle is created + // lazily by GetDefaultCudnnHandleForDevice() so that cuBLAS-only paths (and + // sessions with enable_cudnn=0) never trigger a cuDNN load. thread_local std::unordered_map handles_by_device; auto [it, inserted] = handles_by_device.try_emplace(device_id); if (inserted) { @@ -600,18 +611,7 @@ inline DefaultCudaHandles& GetDefaultCudaHandlesForDevice(int device_id) { handles_by_device.erase(it); ORT_THROW("Failed to create default cuBLAS handle for CUDA plugin device ", device_id); } - if (cudnnCreate(&it->second.cudnn) != CUDNN_STATUS_SUCCESS) { - cublasDestroy(it->second.cublas); - it->second.cublas = nullptr; - if (get_device_result == cudaSuccess) { - cudaSetDevice(prev_device); - } - handles_by_device.erase(it); - ORT_THROW("Failed to create default cuDNN handle for CUDA plugin device ", device_id); - } if (cublasLtCreate(&it->second.cublas_lt) != CUBLAS_STATUS_SUCCESS) { - cudnnDestroy(it->second.cudnn); - it->second.cudnn = nullptr; cublasDestroy(it->second.cublas); it->second.cublas = nullptr; if (get_device_result == cudaSuccess) { @@ -628,6 +628,31 @@ inline DefaultCudaHandles& GetDefaultCudaHandlesForDevice(int device_id) { return it->second; } +// Lazily creates the thread-local fallback cuDNN handle for the device. Callers +// must check enable_cudnn and CudnnLibrary::Available() before invoking this so +// that cuBLAS-only paths never trigger a cuDNN load. +inline cudnnHandle_t GetDefaultCudnnHandleForDevice(int device_id) { + DefaultCudaHandles& handles = GetDefaultCudaHandlesForDevice(device_id); + if (handles.cudnn != nullptr) { + return handles.cudnn; + } + + int prev_device = -1; + const cudaError_t get_device_result = cudaGetDevice(&prev_device); + PL_CUDA_CALL_THROW(cudaSetDevice(device_id)); + cudnnHandle_t cudnn = nullptr; + const cudnnStatus_t status = cudnnCreate(&cudnn); + if (get_device_result == cudaSuccess) { + cudaSetDevice(prev_device); + } + if (status != CUDNN_STATUS_SUCCESS) { + ORT_THROW("Failed to create default cuDNN handle for CUDA plugin device ", device_id); + } + + handles.cudnn = cudnn; + return cudnn; +} + inline const cudaDeviceProp& GetDevicePropForDevice(int device_id) { static std::mutex mutex; static std::unordered_map> props; @@ -729,6 +754,7 @@ inline void SetCudaKernelAdapterRuntimeConfigForProvider( config->cudnn_conv_algo = init_config.cudnn_conv_algo; config->cudnn_conv_use_max_workspace = init_config.cudnn_conv_use_max_workspace; config->cudnn_conv1d_pad_to_nc1d = init_config.cudnn_conv1d_pad_to_nc1d; + config->enable_cudnn = init_config.enable_cudnn; config->fuse_conv_bias = init_config.fuse_conv_bias; config->sdpa_kernel = init_config.sdpa_kernel; config->device_id = init_config.device_id; @@ -973,7 +999,12 @@ class CudaKernel : public OpKernel { inline cudaStream_t DefaultCudaStream() const { return Stream(static_cast(nullptr)); } inline cublasHandle_t DefaultCublasHandle() const { return detail::GetDefaultCudaHandlesForDevice(device_id_).cublas; } - inline cudnnHandle_t DefaultCudnnHandle() const { return detail::GetDefaultCudaHandlesForDevice(device_id_).cudnn; } + inline cudnnHandle_t DefaultCudnnHandle() const { + if (!runtime_config_->enable_cudnn || !onnxruntime::cuda::CudnnLibrary::Get().Available()) { + return nullptr; + } + return detail::GetDefaultCudnnHandleForDevice(device_id_); + } inline cublasLtHandle_t DefaultCublasLtHandle() const { return detail::GetDefaultCudaHandlesForDevice(device_id_).cublas_lt; } inline Status CopyTensor(const onnxruntime::Tensor& src, onnxruntime::Tensor& dst, onnxruntime::Stream& stream) const { @@ -1027,7 +1058,13 @@ class CudaKernel : public OpKernel { } handle = DefaultCudnnHandle(); - if (stream != nullptr) { + if (handle == nullptr) { + ORT_THROW_IF_ERROR(onnxruntime::common::Status( + onnxruntime::common::ONNXRUNTIME, onnxruntime::common::NOT_IMPLEMENTED, + std::string("cuDNN is unavailable or disabled for CUDA Plugin Execution Provider: ") + + onnxruntime::cuda::CudnnLibrary::Get().Error())); + } + if (handle != nullptr && stream != nullptr) { CUDNN_CALL_THROW(cudnnSetStream(handle, stream)); } return handle; diff --git a/onnxruntime/core/providers/cuda/plugin/cuda_plugin_utils.h b/onnxruntime/core/providers/cuda/plugin/cuda_plugin_utils.h index 6ad2bb0f53ccc..6f8e913ac9124 100644 --- a/onnxruntime/core/providers/cuda/plugin/cuda_plugin_utils.h +++ b/onnxruntime/core/providers/cuda/plugin/cuda_plugin_utils.h @@ -30,6 +30,8 @@ #include #include +#include "core/providers/cuda/cudnn_loader.h" + // Error handling macros #ifndef PL_CUDA_RETURN_IF_ERROR @@ -73,6 +75,14 @@ inline Ort::Status StatusFromCudnnError(cudnnStatus_t cudnn_err) { return Ort::Status{}; } + if (!onnxruntime::cuda::CudnnLibrary::Get().Available()) { + return Ort::Status{ + (std::string("cuDNN is unavailable for CUDA Plugin Execution Provider: ") + + onnxruntime::cuda::CudnnLibrary::Get().Error()) + .c_str(), + ORT_NOT_IMPLEMENTED}; + } + return Ort::Status{ (std::string("cuDNN error: ") + cudnnGetErrorString(cudnn_err)).c_str(), ORT_EP_FAIL}; @@ -110,16 +120,23 @@ inline bool TryGetCurrentCudaDevice(int& device_id) noexcept { #endif #ifndef PL_CUDNN_RETURN_IF_ERROR -#define PL_CUDNN_RETURN_IF_ERROR(cudnn_call_expr) \ - do { \ - cudnnStatus_t _cudnn_err = (cudnn_call_expr); \ - if (_cudnn_err != CUDNN_STATUS_SUCCESS) { \ - return Ort::GetApi().CreateStatus( \ - ORT_EP_FAIL, \ - (std::string("cuDNN error: ") + \ - cudnnGetErrorString(_cudnn_err)) \ - .c_str()); \ - } \ +#define PL_CUDNN_RETURN_IF_ERROR(cudnn_call_expr) \ + do { \ + cudnnStatus_t _cudnn_err = (cudnn_call_expr); \ + if (_cudnn_err != CUDNN_STATUS_SUCCESS) { \ + if (!onnxruntime::cuda::CudnnLibrary::Get().Available()) { \ + return Ort::GetApi().CreateStatus( \ + ORT_NOT_IMPLEMENTED, \ + (std::string("cuDNN is unavailable for CUDA Plugin Execution Provider: ") + \ + onnxruntime::cuda::CudnnLibrary::Get().Error()) \ + .c_str()); \ + } \ + return Ort::GetApi().CreateStatus( \ + ORT_EP_FAIL, \ + (std::string("cuDNN error: ") + \ + cudnnGetErrorString(_cudnn_err)) \ + .c_str()); \ + } \ } while (0) #endif diff --git a/onnxruntime/core/providers/cuda/plugin/cuda_stream_plugin.cc b/onnxruntime/core/providers/cuda/plugin/cuda_stream_plugin.cc index e497a9151c73d..80184dfe101f3 100644 --- a/onnxruntime/core/providers/cuda/plugin/cuda_stream_plugin.cc +++ b/onnxruntime/core/providers/cuda/plugin/cuda_stream_plugin.cc @@ -3,6 +3,7 @@ #include "cuda_stream_plugin.h" #include "cuda_ep_factory.h" +#include "core/providers/cuda/cudnn_loader.h" #include #include #include @@ -39,11 +40,12 @@ std::atomic& GetStreamMapGeneration() { // CudaSyncStream // --------------------------------------------------------------------------- -CudaSyncStream::CudaSyncStream(CudaEpFactory& factory, int device_id, +CudaSyncStream::CudaSyncStream(CudaEpFactory& factory, int device_id, bool enable_cudnn, const OrtEp* /*ep*/) : OrtSyncStreamImpl{}, factory_(factory), - device_id_(device_id) { + device_id_(device_id), + enable_cudnn_(enable_cudnn) { ort_version_supported = ORT_API_VERSION; GetHandle = GetHandleImpl; CreateNotification = CreateNotificationImpl; @@ -124,10 +126,10 @@ OrtStatus* CudaSyncStream::InitHandles() { if (status.IsOK()) { status = StatusFromCublasError(cublasSetStream(cublas_handle_, cuda_stream_)); } - if (status.IsOK()) { + if (status.IsOK() && enable_cudnn_ && onnxruntime::cuda::CudnnLibrary::Get().Available()) { status = StatusFromCudnnError(cudnnCreate(&cudnn_handle_)); } - if (status.IsOK()) { + if (status.IsOK() && cudnn_handle_ != nullptr) { status = StatusFromCudnnError(cudnnSetStream(cudnn_handle_, cuda_stream_)); } if (status.IsOK()) { @@ -192,10 +194,10 @@ OrtStatus* CudaSyncStream::InitHandlesWithUserStream(cudaStream_t user_stream) { if (status.IsOK()) { status = StatusFromCublasError(cublasSetStream(cublas_handle_, cuda_stream_)); } - if (status.IsOK()) { + if (status.IsOK() && enable_cudnn_ && onnxruntime::cuda::CudnnLibrary::Get().Available()) { status = StatusFromCudnnError(cudnnCreate(&cudnn_handle_)); } - if (status.IsOK()) { + if (status.IsOK() && cudnn_handle_ != nullptr) { status = StatusFromCudnnError(cudnnSetStream(cudnn_handle_, cuda_stream_)); } if (status.IsOK()) { diff --git a/onnxruntime/core/providers/cuda/plugin/cuda_stream_plugin.h b/onnxruntime/core/providers/cuda/plugin/cuda_stream_plugin.h index 1f1e5fade3de1..94d458149e2ca 100644 --- a/onnxruntime/core/providers/cuda/plugin/cuda_stream_plugin.h +++ b/onnxruntime/core/providers/cuda/plugin/cuda_stream_plugin.h @@ -28,7 +28,7 @@ class CudaEpFactory; /// or wraps an external stream for graph-mode registration/lifecycle tracking. class CudaSyncStream : public OrtSyncStreamImpl { public: - CudaSyncStream(CudaEpFactory& factory, int device_id, + CudaSyncStream(CudaEpFactory& factory, int device_id, bool enable_cudnn, const OrtEp* ep); ~CudaSyncStream(); @@ -70,6 +70,7 @@ class CudaSyncStream : public OrtSyncStreamImpl { CudaEpFactory& factory_; int device_id_; + bool enable_cudnn_ = true; cudaStream_t cuda_stream_ = nullptr; bool owns_stream_ = true; ///< False when wrapping an external stream (e.g., for CUDA graph). cublasHandle_t cublas_handle_ = nullptr; diff --git a/onnxruntime/test/python/onnxruntime_test_python.py b/onnxruntime/test/python/onnxruntime_test_python.py index bebd290bbe810..d60acfc8bc34f 100644 --- a/onnxruntime/test/python/onnxruntime_test_python.py +++ b/onnxruntime/test/python/onnxruntime_test_python.py @@ -439,6 +439,8 @@ def test_get_and_set_option_with_values(option_name, option_values): test_get_and_set_option_with_values("cudnn_conv_algo_search", ["DEFAULT", "EXHAUSTIVE", "HEURISTIC"]) + test_get_and_set_option_with_values("enable_cudnn", ["1", "0"]) + test_get_and_set_option_with_values("do_copy_in_default_stream", [0, 1]) test_get_and_set_option_with_values("tunable_op_enable", ["1", "0"]) diff --git a/onnxruntime/test/python/transformers/test_cuda_plugin_ep.py b/onnxruntime/test/python/transformers/test_cuda_plugin_ep.py index 3cb1392919398..51ee1a4568e95 100644 --- a/onnxruntime/test/python/transformers/test_cuda_plugin_ep.py +++ b/onnxruntime/test/python/transformers/test_cuda_plugin_ep.py @@ -27,6 +27,23 @@ TEST_SKIP = "SKIP" TEST_FAIL = "FAIL" EP_GRAPH_ASSIGNMENT_CONFIG_KEY = "session.record_ep_graph_assignment_info" +NO_CUDNN_PLUGIN_TEST = os.getenv("ORT_TEST_CUDA_PLUGIN_NO_CUDNN", "").upper() in {"1", "ON", "TRUE", "YES"} +requires_cudnn = unittest.skipIf(NO_CUDNN_PLUGIN_TEST, "test requires cuDNN-backed CUDA plugin kernels") +# Use the latest released ai.onnx opset so the model builders stay current as ONNX releases new opsets. +DEFAULT_ONNX_OPSET = max(v for (d, v) in helper.OP_SET_ID_VERSION_MAP if d == "ai.onnx") + + +def _make_released_opset_model(graph, producer_name="onnx-example"): + opset = OperatorSetIdProto() + opset.version = DEFAULT_ONNX_OPSET + return helper.make_model(graph, producer_name=producer_name, opset_imports=[opset]) + + +def _plugin_provider_options(extra_options=None): + options = {"enable_cudnn": "0"} if NO_CUDNN_PLUGIN_TEST else {} + if extra_options: + options.update(extra_options) + return options def require_cuda_plugin_ep(): @@ -140,7 +157,7 @@ def create_add_model(model_path): ], [helper.make_tensor_value_info("Y", TensorProto.FLOAT, [3, 2])], ) - model_def = helper.make_model(graph_def, producer_name="onnx-example") + model_def = _make_released_opset_model(graph_def) save(model_def, model_path) @@ -156,7 +173,7 @@ def create_matmul_model(model_path): ], [helper.make_tensor_value_info("Y", TensorProto.FLOAT, [3, 5])], ) - model_def = helper.make_model(graph_def, producer_name="onnx-example") + model_def = _make_released_opset_model(graph_def) save(model_def, model_path) @@ -181,7 +198,7 @@ def create_gemm_model(model_path, alpha=1.0, beta=1.0, transA=0, transB=0): ], [helper.make_tensor_value_info("Y", TensorProto.FLOAT, [m, n])], ) - model_def = helper.make_model(graph_def, producer_name="onnx-example") + model_def = _make_released_opset_model(graph_def) save(model_def, model_path) @@ -323,7 +340,7 @@ def run_operator_test( try: model_creator(model_path) sess_options = _create_session_options(session_config) - sess_options.add_provider_for_devices([target_device], {}) + sess_options.add_provider_for_devices([target_device], _plugin_provider_options()) sess = onnxrt.InferenceSession(model_path, sess_options=sess_options) active_providers = sess.get_providers() @@ -375,7 +392,7 @@ def run_provider_options_test(provider_options, expect_plugin_provider=True): model_path = tmp.name try: create_add_model(model_path) - providers = [(CUDA_PLUGIN_EP_NAME, provider_options), "CPUExecutionProvider"] + providers = [(CUDA_PLUGIN_EP_NAME, _plugin_provider_options(provider_options)), "CPUExecutionProvider"] sess = onnxrt.InferenceSession(model_path, sess_options=_create_session_options(), providers=providers) active_providers = sess.get_providers() assigned_nodes, assignment_info = _get_assigned_nodes(sess, CUDA_PLUGIN_EP_NAME) @@ -483,7 +500,7 @@ def _run_nhwc_model_test(target_device, op_name, model, feed_dict, expected_fn, try: save(model, model_path) sess_options = _create_session_options(_NHWC_CONFIG) - sess_options.add_provider_for_devices([target_device], {}) + sess_options.add_provider_for_devices([target_device], _plugin_provider_options()) sess = onnxrt.InferenceSession(model_path, sess_options=sess_options) assigned_nodes, assignment_info = _get_assigned_nodes(sess, CUDA_PLUGIN_EP_NAME) if not assigned_nodes: @@ -556,7 +573,7 @@ def _run_model_test( try: save(model, model_path) sess_options = _create_session_options() - sess_options.add_provider_for_devices([target_device], {}) + sess_options.add_provider_for_devices([target_device], _plugin_provider_options()) sess = onnxrt.InferenceSession(model_path, sess_options=sess_options) active_providers = sess.get_providers() assigned_nodes, assignment_info = _get_assigned_nodes(sess, ep_name) @@ -628,6 +645,7 @@ def test_registration_gemm(self): ) self.assertTrue(result, "Gemm plugin registration test failed") + @requires_cudnn def test_registration_conv(self): target_device = get_cuda_plugin_device() inputs = { @@ -658,7 +676,7 @@ def test_provider_options_second_device(self): model_path = tmp.name try: create_add_model(model_path) - providers = [(CUDA_PLUGIN_EP_NAME, {"device_id": "1"}), "CPUExecutionProvider"] + providers = [(CUDA_PLUGIN_EP_NAME, _plugin_provider_options({"device_id": "1"})), "CPUExecutionProvider"] sess = onnxrt.InferenceSession(model_path, sess_options=_create_session_options(), providers=providers) active_providers = sess.get_providers() @@ -687,6 +705,7 @@ def test_provider_options_second_device(self): # ---- NHWC layout tests ---- + @requires_cudnn def test_nhwc_conv(self): target_device = get_cuda_plugin_device() inputs = { @@ -703,6 +722,7 @@ def test_nhwc_conv(self): ) self.assertTrue(result, "Conv (NHWC) plugin test failed") + @requires_cudnn def test_nhwc_batch_normalization(self): target_device = get_cuda_plugin_device() inputs = {"X": np.random.rand(1, 3, 4, 4).astype(np.float32)} @@ -716,6 +736,7 @@ def test_nhwc_batch_normalization(self): ) self.assertTrue(result, "BatchNormalization (NHWC) plugin test failed") + @requires_cudnn def test_nhwc_maxpool(self): target_device = get_cuda_plugin_device() inputs = {"X": np.random.rand(1, 3, 4, 4).astype(np.float32)} @@ -729,6 +750,7 @@ def test_nhwc_maxpool(self): ) self.assertTrue(result, "MaxPool (NHWC) plugin test failed") + @requires_cudnn def test_nhwc_avgpool(self): target_device = get_cuda_plugin_device() inputs = {"X": np.random.rand(1, 3, 4, 4).astype(np.float32)} @@ -742,6 +764,7 @@ def test_nhwc_avgpool(self): ) self.assertTrue(result, "AveragePool (NHWC) plugin test failed") + @requires_cudnn def test_nhwc_conv_transpose(self): target_device = get_cuda_plugin_device() # ConvTranspose: input [1,2,4,4], weight [2,3,3,3] -> output [1,3,6,6] with stride=2, padding=1, output_padding=1 @@ -782,6 +805,7 @@ def expected_fn(feed): result = _run_nhwc_model_test(target_device, "ConvTranspose", model, {"X": x, "W": w}, expected_fn) self.assertEqual(result, TEST_PASS, "ConvTranspose (NHWC) plugin test failed") + @requires_cudnn def test_nhwc_global_max_pool(self): target_device = get_cuda_plugin_device() f_dtype = TensorProto.FLOAT @@ -800,6 +824,7 @@ def expected_fn(feed): result = _run_nhwc_model_test(target_device, "GlobalMaxPool", model, {"X": x}, expected_fn) self.assertEqual(result, TEST_PASS, "GlobalMaxPool (NHWC) plugin test failed") + @requires_cudnn def test_nhwc_global_average_pool(self): target_device = get_cuda_plugin_device() f_dtype = TensorProto.FLOAT @@ -867,6 +892,7 @@ def expected_fn(feed): result = _run_nhwc_model_test(target_device, "SpaceToDepth", model, {"X": x}, expected_fn) self.assertEqual(result, TEST_PASS, "SpaceToDepth (NHWC) plugin test failed") + @requires_cudnn def test_nhwc_lrn(self): target_device = get_cuda_plugin_device() f_dtype = TensorProto.FLOAT @@ -910,6 +936,7 @@ def expected_fn(feed): result = _run_nhwc_model_test(target_device, "GridSample", model, {"X": x, "grid": grid}, expected_fn) self.assertEqual(result, TEST_PASS, "GridSample (NHWC) plugin test failed") + @requires_cudnn def test_nhwc_conv_with_resource_accounting(self): # Smoke test for the NHWC two-pass partitioning flow combined with the resource # accountant (session.resource_cuda_partitioning_settings). The NHWC layout @@ -1592,26 +1619,32 @@ def test_plugin_ep_claims_key_ops(self): ), # second unary ("Sigmoid", "", 13, [("X", TensorProto.FLOAT, [2, 4])], [("Y", TensorProto.FLOAT, [2, 4])], None), - # cuDNN: ConvTranspose (Conv already tested by test_registration_conv) - ( - "ConvTranspose", - "", - 13, - [("X", TensorProto.FLOAT, [1, 2, 3, 3]), ("W", TensorProto.FLOAT, [2, 3, 3, 3])], - [("Y", TensorProto.FLOAT, [1, 3, 5, 5])], - None, - ), - # cuDNN: LRN (local response normalization) - ( - "LRN", - "", - 13, - [("X", TensorProto.FLOAT, [1, 2, 4, 4])], - [("Y", TensorProto.FLOAT, [1, 2, 4, 4])], - {"size": 3}, - ), ] + if not NO_CUDNN_PLUGIN_TEST: + probe_specs.extend( + [ + # cuDNN: ConvTranspose (Conv already tested by test_registration_conv) + ( + "ConvTranspose", + "", + 13, + [("X", TensorProto.FLOAT, [1, 2, 3, 3]), ("W", TensorProto.FLOAT, [2, 3, 3, 3])], + [("Y", TensorProto.FLOAT, [1, 3, 5, 5])], + None, + ), + # cuDNN: LRN (local response normalization) + ( + "LRN", + "", + 13, + [("X", TensorProto.FLOAT, [1, 2, 4, 4])], + [("Y", TensorProto.FLOAT, [1, 2, 4, 4])], + {"size": 3}, + ), + ] + ) + claimed = [] not_claimed = [] errors = [] @@ -1624,7 +1657,7 @@ def test_plugin_ep_claims_key_ops(self): save(model, model_path) sess_options = _create_session_options() sess_options.graph_optimization_level = onnxrt.GraphOptimizationLevel.ORT_DISABLE_ALL - sess_options.add_provider_for_devices([target_device], {}) + sess_options.add_provider_for_devices([target_device], _plugin_provider_options()) sess = onnxrt.InferenceSession(model_path, sess_options=sess_options) assigned_nodes, _ = _get_assigned_nodes(sess, CUDA_PLUGIN_EP_NAME) if assigned_nodes: @@ -1650,6 +1683,7 @@ def test_plugin_ep_claims_key_ops(self): # ---- Newly-included ops that previously lacked tests ---- + @requires_cudnn def test_op_einsum(self): """Test Einsum op (recently un-excluded from plugin build).""" target_device = get_cuda_plugin_device() @@ -1664,6 +1698,7 @@ def test_op_einsum(self): result = _run_model_test(target_device, "Einsum", model, feed, lambda f: f["A"] @ f["B"]) self.assertEqual(result, TEST_PASS, "Einsum test failed") + @requires_cudnn def test_op_einsum_batch(self): """Test Einsum op with batch matrix multiply.""" target_device = get_cuda_plugin_device() @@ -1793,6 +1828,7 @@ def test_op_flatten(self): result = _run_model_test(target_device, "Flatten", model, feed, lambda f: f["X"].reshape(2, 12)) self.assertEqual(result, TEST_PASS, "Flatten test failed") + @requires_cudnn def test_op_argmax(self): target_device = get_cuda_plugin_device() model = _make_simple_model( @@ -1863,6 +1899,7 @@ def expected(f): result = _run_model_test(target_device, "LayerNormalization", model, feed, expected) self.assertEqual(result, TEST_PASS, "LayerNormalization test failed") + @requires_cudnn def test_op_instance_normalization(self): target_device = get_cuda_plugin_device() n_channels = 3 @@ -1898,6 +1935,7 @@ def expected(f): result = _run_model_test(target_device, "InstanceNormalization", model, feed, expected) self.assertEqual(result, TEST_PASS, "InstanceNormalization test failed") + @requires_cudnn def test_op_conv_transpose(self): target_device = get_cuda_plugin_device() model = _make_simple_model( @@ -1920,6 +1958,7 @@ def expected(f): result = _run_model_test(target_device, "ConvTranspose", model, feed, expected) self.assertEqual(result, TEST_PASS, "ConvTranspose test failed") + @requires_cudnn def test_op_reduce_mean(self): target_device = get_cuda_plugin_device() model = _make_simple_model( @@ -1935,6 +1974,7 @@ def test_op_reduce_mean(self): ) self.assertEqual(result, TEST_PASS, "ReduceMean test failed") + @requires_cudnn def test_op_reduce_sum(self): target_device = get_cuda_plugin_device() model = _make_simple_model( @@ -2284,7 +2324,7 @@ def _create_cuda_graph_session(self, model_path, extra_session_config=None, prov if extra_session_config: for key, value in extra_session_config.items(): sess_options.add_session_config_entry(key, value) - provider_options = {"enable_cuda_graph": "1", **(provider_options or {})} + provider_options = _plugin_provider_options({"enable_cuda_graph": "1", **(provider_options or {})}) providers = [(CUDA_PLUGIN_EP_NAME, provider_options), "CPUExecutionProvider"] return onnxrt.InferenceSession(model_path, sess_options=sess_options, providers=providers) @@ -2470,7 +2510,7 @@ def test_cuda_graph_annotation_id(self): ], [helper.make_tensor_value_info("Y", TensorProto.FLOAT, ["M", "N"])], ) - model_def = helper.make_model(graph_def, producer_name="onnx-example") + model_def = _make_released_opset_model(graph_def) save(model_def, model_path) session = self._create_cuda_graph_session(model_path) @@ -2637,7 +2677,7 @@ def test_iobinding_add(self): try: create_add_model(model_path) sess_options = _create_session_options() - sess_options.add_provider_for_devices([target_device], {}) + sess_options.add_provider_for_devices([target_device], _plugin_provider_options()) sess = onnxrt.InferenceSession(model_path, sess_options=sess_options) assigned_nodes, assignment_info = _get_assigned_nodes(sess, CUDA_PLUGIN_EP_NAME) @@ -2678,7 +2718,7 @@ def test_iobinding_matmul(self): try: create_matmul_model(model_path) sess_options = _create_session_options() - sess_options.add_provider_for_devices([target_device], {}) + sess_options.add_provider_for_devices([target_device], _plugin_provider_options()) sess = onnxrt.InferenceSession(model_path, sess_options=sess_options) assigned_nodes, assignment_info = _get_assigned_nodes(sess, CUDA_PLUGIN_EP_NAME) @@ -2726,7 +2766,7 @@ def _run_profiling_test(self): try: create_matmul_model(model_path) sess_options = _create_session_options() - sess_options.add_provider_for_devices([target_device], {}) + sess_options.add_provider_for_devices([target_device], _plugin_provider_options()) profile_prefix = os.path.join(tempfile.gettempdir(), "cuda_plugin_ep_profiling_test") sess_options.enable_profiling = True From 76936d30de3d9a829a6a96dace77f9376684cdf1 Mon Sep 17 00:00:00 2001 From: Jambay Kinley Date: Tue, 30 Jun 2026 22:17:56 -0700 Subject: [PATCH 07/17] Expand model package with authoring tools and schema versioning (#28989) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ### Description Builds out the standalone `model_package/` C library so a single library covers the full lifecycle of an ONNX Runtime model package: inspection, authoring, content-addressed shared assets, commit, prune, and validation. The library remains free of any dependency on ONNX Runtime itself. **Public C API (`model_package/include/model_package.h`)** - Lifecycle: `ModelPackage_Open` / `ModelPackage_New` / `ModelPackage_Close`, with `ModelPackageOpenOptions` controlling external-path access, symlink following, and strict unknown-field handling. - Inspection: a POD `ModelPackageInfo` tree (`ModelPackage_Info`) plus by-name lookups for components, variants, and per-namespace `executor_info` entries, and round-trip JSON getters that preserve fields unknown to the current build. - Path resolution: `ModelPackage_ResolveStringRef` implements the package's resolution rules — relative paths anchored at a base directory, `sha256:[/sub/path]` for shared-asset content, and portable vs installed confinement (absolute paths and `..` segments only allowed under `layout: "installed"`). - Shared assets: SHA-256 directory hashing, `ModelPackage_AddSharedAsset` (with reproducible-build URI check and an optional `copy_in` staging mode), `ModelPackage_RemoveSharedAsset`, and `ModelPackage_ResolveAssetUri`. Assets under `/shared_assets/` are auto-discovered at `Open`. - Authoring: inline/external component setters, variant upsert/remove, per-namespace executor_info setters (inline and external), and package-level metadata / layout / `additional_metadata` setters. Mutations invalidate cached pointers in the mutated scope and its descendants. - Commit / prune / validate: `ModelPackage_Commit` writes the in-memory model to disk either in place or to a fresh `dest_root` ("save as"), with `PRESERVE` or `DENSE` write modes; `ModelPackage_Prune` reclaims unreferenced files under `shared_assets/` and tracked orphan variant/component directories left behind by removals; and `ModelPackage_Validate` runs schema, path-reachability, asset-rehash, and unknown-field checks and returns a JSON report. - Errors: opaque `ModelPackageStatus*` with a stable additive `ModelPackageErrorCode` enum (IO, schema, version, path confinement, asset missing, asset hash mismatch, not found, invalid arg, state). **Internal layout** The implementation is split into focused translation units: `manifest_parser`, `model_package_impl`, `authoring`, `commit_prune_validate`, `path_resolver`, `asset_hasher`, and an in-tree `sha256`. Shared error plumbing lives in `status_impl.h`. **ONNX Runtime integration (`onnxruntime/core/session/model_package/`)** The ORT-side glue is wired onto the library through the C inspection and path-resolution entry points. `model_package_context` now translates the library's info tree into ORT-internal structs and parses the `executor_info["ort"]` payload (`model_file`, `external_data`, `session_options`, `provider_options`). When a variant declares `external_data`, `CreateSessionForModelPackage` loads the model from a memory mapping and passes the resolved folder to the session via `kOrtSessionOptionsModelExternalInitializersFileFolderPath`, so external initializers — including those backed by a shared asset — are picked up at `Initialize` time. The experimental `OrtModelPackageApi_*_SinceV28` C entries introduced in #28990 are unchanged. **Documentation and tests** - `model_package/README.md` documents the on-disk layout, manifest and component schemas, shared-asset rules, path resolution, the authoring flow, and commit / prune / validate semantics. - `onnxruntime/core/session/model_package/README.md` documents the ORT consumer-side glue: the `executor_info["ort"]` schema, the variant selection algorithm, the session-creation contract, and the registered experimental C entries. - New library tests cover inspection, authoring, asset hashing, and commit (`model_package/tests/`). The ORT integration tests in `onnxruntime/test/autoep/test_model_package.cc` are reworked against the current C API surface. ### Motivation and Context ORT needs a single library that owns the model package format end to end — not just reading it, but producing it, validating it, and maintaining it on disk with content-addressed shared assets. Consolidating this behind one C API lets ORT, publisher tooling, and third-party consumers share the same parser, path-resolution rules, and on-disk invariants without each reimplementing them, and keeps the library independent of the ORT session runtime. --------- Co-authored-by: jambayk Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Co-authored-by: jambayk --- cmake/onnxruntime_session.cmake | 8 +- .../onnxruntime_experimental_c_api.inc | 28 + model_package/CMakeLists.txt | 50 +- model_package/README.md | 651 ++++++++++- model_package/include/model_package.h | 350 ++++++ model_package/include/model_package_api.h | 134 +-- model_package/src/api.cc | 243 ---- model_package/src/asset_hasher.cc | 97 ++ model_package/src/asset_hasher.h | 30 + model_package/src/authoring.cc | 593 ++++++++++ model_package/src/commit_prune_validate.cc | 769 +++++++++++++ model_package/src/manifest_parser.cc | 732 ++++++++++++ model_package/src/manifest_parser.h | 64 ++ model_package/src/model_package_impl.cc | 403 +++++++ model_package/src/model_package_impl.h | 161 +++ model_package/src/model_package_internal.h | 80 -- model_package/src/parser.cc | 595 ---------- model_package/src/parser.h | 27 - model_package/src/path_resolver.cc | 172 +++ model_package/src/path_resolver.h | 60 + model_package/src/sha256.cc | 229 ++++ model_package/src/sha256.h | 44 + model_package/src/status_impl.h | 30 + model_package/tests/test_asset_hashing.cc | 312 +++++ model_package/tests/test_authoring.cc | 525 +++++++++ model_package/tests/test_commit.cc | 504 ++++++++ model_package/tests/test_inspection.cc | 582 ++++++++++ .../core/session/model_package/README.md | 246 ++++ .../model_package/model_package_context.cc | 239 +++- .../model_package/model_package_context.h | 30 +- onnxruntime/core/session/model_package_api.cc | 44 +- onnxruntime/core/session/utils.cc | 77 +- onnxruntime/test/autoep/test_model_package.cc | 1020 +++++------------ 33 files changed, 7205 insertions(+), 1924 deletions(-) create mode 100644 model_package/include/model_package.h delete mode 100644 model_package/src/api.cc create mode 100644 model_package/src/asset_hasher.cc create mode 100644 model_package/src/asset_hasher.h create mode 100644 model_package/src/authoring.cc create mode 100644 model_package/src/commit_prune_validate.cc create mode 100644 model_package/src/manifest_parser.cc create mode 100644 model_package/src/manifest_parser.h create mode 100644 model_package/src/model_package_impl.cc create mode 100644 model_package/src/model_package_impl.h delete mode 100644 model_package/src/model_package_internal.h delete mode 100644 model_package/src/parser.cc delete mode 100644 model_package/src/parser.h create mode 100644 model_package/src/path_resolver.cc create mode 100644 model_package/src/path_resolver.h create mode 100644 model_package/src/sha256.cc create mode 100644 model_package/src/sha256.h create mode 100644 model_package/src/status_impl.h create mode 100644 model_package/tests/test_asset_hashing.cc create mode 100644 model_package/tests/test_authoring.cc create mode 100644 model_package/tests/test_commit.cc create mode 100644 model_package/tests/test_inspection.cc create mode 100644 onnxruntime/core/session/model_package/README.md diff --git a/cmake/onnxruntime_session.cmake b/cmake/onnxruntime_session.cmake index 6fb895cd1800c..21229c4a4f2de 100644 --- a/cmake/onnxruntime_session.cmake +++ b/cmake/onnxruntime_session.cmake @@ -13,10 +13,8 @@ file(GLOB onnxruntime_session_srcs CONFIGURE_DEPENDS # Standalone model package library (parsing/inspection with no ORT dependency). # Compiled as a static library and linked into onnxruntime_session. -# NOTE: ORT intentionally uses the library's internal C++ types directly (model_package::ParsePackage, -# model_package_internal.h) rather than going through its public C API (ModelPackage_*). This avoids -# double-wrapping (ORT C API -> standalone C API -> C++ internals). The public C API exists for -# external consumers (GenAI, FL) who link against the standalone library independently. +# ORT uses the standalone library's public C API (model_package.h) and translates +# the C handles into ORT-internal C++ types inside core/session/model_package/. set(MODEL_PACKAGE_LIB_DIR "${REPO_ROOT}/model_package") if(NOT onnxruntime_MINIMAL_BUILD) set(MODEL_PACKAGE_BUILD_SHARED OFF CACHE BOOL "" FORCE) @@ -59,7 +57,7 @@ onnxruntime_add_include_to_target(onnxruntime_session onnxruntime_common onnxrun target_link_libraries(onnxruntime_session PRIVATE onnxruntime_lora) if(TARGET model_package) target_link_libraries(onnxruntime_session PRIVATE model_package) - target_include_directories(onnxruntime_session PRIVATE ${MODEL_PACKAGE_LIB_DIR}/include ${MODEL_PACKAGE_LIB_DIR}/src) + target_include_directories(onnxruntime_session PRIVATE ${MODEL_PACKAGE_LIB_DIR}/include) endif() if(onnxruntime_ENABLE_INSTRUMENT) target_compile_definitions(onnxruntime_session PUBLIC ONNXRUNTIME_ENABLE_INSTRUMENT) diff --git a/include/onnxruntime/core/session/onnxruntime_experimental_c_api.inc b/include/onnxruntime/core/session/onnxruntime_experimental_c_api.inc index b1140485f7ff1..a8047566df63c 100644 --- a/include/onnxruntime/core/session/onnxruntime_experimental_c_api.inc +++ b/include/onnxruntime/core/session/onnxruntime_experimental_c_api.inc @@ -66,6 +66,7 @@ ORT_EXPERIMENTAL_API(28, OrtStatusPtr, OrtApi_ExperimentalApiTest, _Out_ int64_t // - OrtModelPackageApi_ModelPackage_GetVariantCount // - OrtModelPackageApi_ModelPackage_GetVariantNames // - OrtModelPackageApi_ModelPackage_GetVariantEpName +// - OrtModelPackageApi_ModelPackage_ResolveStringRef // 4) Select a component and resolve variant: // - OrtModelPackageApi_SelectComponent // 5) Query selected variant info (optional): @@ -212,6 +213,33 @@ ORT_EXPERIMENTAL_API(28, OrtStatusPtr, OrtModelPackageApi_ModelPackage_GetVarian _In_ const char* variant_name, _Outptr_result_maybenull_ const char** out_ep) +/** \brief Resolve a path reference declared inside the package against the model_package rules. + * + * Handles the path forms a package may use: + * - "sha256:" or "sha256:/": a content-addressed shared asset. Resolves to + * the asset's on-disk directory (honoring manifest shared_assets overrides), optionally + * joined with the confined tail. Errors if the asset is not declared/discovered. + * - any other value: a relative path resolved against `base_dir` (or the package root when + * `base_dir` is NULL), with portable-layout confinement (no absolute paths, no ".."). + * + * When `must_exist` is non-zero the resolved path must exist on disk. `out_path` is owned by + * `ctx` and remains valid until the next call to this function on the same context. + * + * \param[in] ctx The package context returned by OrtModelPackageApi_CreateModelPackageContext. + * \param[in] base_dir Base directory for relative inputs. May be NULL to use the package root. + * \param[in] input The path reference to resolve. + * \param[in] must_exist Non-zero to require that the resolved path exists on disk. + * \param[out] out_path Receives the resolved path string. + * + * \snippet{doc} snippets.dox OrtStatus Return Value + */ +ORT_EXPERIMENTAL_API(28, OrtStatusPtr, OrtModelPackageApi_ModelPackage_ResolveStringRef, + _In_ const OrtModelPackageContext* ctx, + _In_opt_ const char* base_dir, + _In_ const char* input, + _In_ int must_exist, + _Outptr_ const char** out_path) + /** \brief Select a component model and return an opaque component instance. * * The variant selection is also performed during this call based on the component diff --git a/model_package/CMakeLists.txt b/model_package/CMakeLists.txt index 326a1e541696a..4b296cdca96a0 100644 --- a/model_package/CMakeLists.txt +++ b/model_package/CMakeLists.txt @@ -52,8 +52,13 @@ endif() # ───────────────────────────────────────────────────────────────────────────── set(MODEL_PACKAGE_SOURCES - src/api.cc - src/parser.cc + src/asset_hasher.cc + src/authoring.cc + src/commit_prune_validate.cc + src/manifest_parser.cc + src/model_package_impl.cc + src/path_resolver.cc + src/sha256.cc ) if(MODEL_PACKAGE_BUILD_SHARED) @@ -91,6 +96,45 @@ install(TARGETS model_package RUNTIME DESTINATION ${CMAKE_INSTALL_BINDIR} ) -install(FILES include/model_package_api.h +install(FILES include/model_package_api.h include/model_package.h DESTINATION ${CMAKE_INSTALL_INCLUDEDIR} ) + +# ───────────────────────────────────────────────────────────────────────────── +# Tests +# ───────────────────────────────────────────────────────────────────────────── + +if(MODEL_PACKAGE_BUILD_TESTS) + enable_testing() + add_executable(test_inspection tests/test_inspection.cc) + target_link_libraries(test_inspection PRIVATE model_package nlohmann_json::nlohmann_json) + target_include_directories(test_inspection PRIVATE + ${CMAKE_CURRENT_SOURCE_DIR}/include + ${CMAKE_CURRENT_SOURCE_DIR}/src + ) + add_test(NAME inspection COMMAND test_inspection) + + add_executable(test_asset_hashing tests/test_asset_hashing.cc) + target_link_libraries(test_asset_hashing PRIVATE model_package nlohmann_json::nlohmann_json) + target_include_directories(test_asset_hashing PRIVATE + ${CMAKE_CURRENT_SOURCE_DIR}/include + ${CMAKE_CURRENT_SOURCE_DIR}/src + ) + add_test(NAME asset_hashing COMMAND test_asset_hashing) + + add_executable(test_authoring tests/test_authoring.cc) + target_link_libraries(test_authoring PRIVATE model_package nlohmann_json::nlohmann_json) + target_include_directories(test_authoring PRIVATE + ${CMAKE_CURRENT_SOURCE_DIR}/include + ${CMAKE_CURRENT_SOURCE_DIR}/src + ) + add_test(NAME authoring COMMAND test_authoring) + + add_executable(test_commit tests/test_commit.cc) + target_link_libraries(test_commit PRIVATE model_package nlohmann_json::nlohmann_json) + target_include_directories(test_commit PRIVATE + ${CMAKE_CURRENT_SOURCE_DIR}/include + ${CMAKE_CURRENT_SOURCE_DIR}/src + ) + add_test(NAME commit COMMAND test_commit) +endif() diff --git a/model_package/README.md b/model_package/README.md index 604720916d764..dcbb9252c71ed 100644 --- a/model_package/README.md +++ b/model_package/README.md @@ -1,78 +1,625 @@ # Model Package Library -A standalone C library for parsing and inspecting ONNX Runtime Model Packages. +A standalone C library for **reading, authoring, validating, and committing** +ONNX Runtime model packages. The library has no dependency on ONNX Runtime +itself, so any consumer (ORT, publisher tools, ...) can compile it in +without dragging in a session runtime. It is distributed and consumed as +**source** (see [Versioning and compatibility](#versioning-and-compatibility)). -**No dependency on ONNX Runtime.** This library can be consumed independently by any component (ORT, GenAI, FL, or external tools). +The library owns three things: -## What it does +1. The **on-disk layout** of a model package (directory + manifest + shared + assets). +2. The **schema** of `manifest.json` and `component.json`, including the + `executor_info` extension point. +3. The **resolution rules** for paths and content-addressed shared assets, + including portable vs installed confinement. -- Parses model package directory structures (`manifest.json`, `metadata.json`, `variant.json`) -- Provides read-only access to: - - Components and their variants - - EP compatibility declarations (opaque strings) - - Model file paths within variants - - Session/provider options per file - - Consumer metadata (opaque JSON) +It deliberately does **not** know about ONNX, execution providers, sessions, +or the JSON payload that lives under any `executor_info[""]` slot. +Each consumer owns its own slot and parses it itself. -## What it does NOT do +--- -- Variant selection (requires runtime EP factory validation → stays in ORT) -- Session creation (requires ORT `InferenceSession`) -- Any interpretation of `compatibility_string` tokens +## On-disk layout -## Building +A package is a directory containing a top-level `manifest.json`. Components +live under the package root, either declared inline in the manifest or as +external `component.json` files. Variants are directories under their +component. Shared assets are content-addressed directories under +`shared_assets/`. -```bash -cmake -B build -S . -cmake --build build +``` +package_root/ +├── manifest.json # required +├── decoder/ # external component (directory) +│ ├── component.json # required when external +│ └── cpu/ # variant_directory +│ ├── model.onnx +│ └── ort_info.json # executor_info["ort"], external form +├── encoder/ # inline component (no component.json) +│ └── cuda/ +│ └── model.onnx +└── shared_assets/ + └── sha256-<64hex>/ # content-addressed asset directory + ├── tokenizer.json + └── chat_template.jinja ``` -Options: -- `-DMODEL_PACKAGE_BUILD_SHARED=ON|OFF` — Build as shared (default) or static library -- `-DMODEL_PACKAGE_BUILD_TESTS=ON` — Build tests (default OFF) +- The package root must be a directory. A single file is **not** a package. +- A package has at least one component. A component has at least one variant. +- A variant always corresponds to a directory on disk (`variant_directory`). + Files inside that directory are referenced by `executor_info` payloads, not + by the manifest. +- `shared_assets/` is optional and only needs to exist if at least one + shared asset is published. -## C API Usage +### Portable vs installed layout -```c -#include "model_package_api.h" - -ModelPackageContext* ctx = NULL; -ModelPackageStatus* status = ModelPackage_CreateContext("/path/to/package", &ctx); -if (status != NULL) { - printf("Error: %s\n", ModelPackage_GetErrorMessage(status)); - ModelPackage_ReleaseStatus(status); - return; +`manifest.layout` declares how the package may use paths: + +- `"portable"` (default): every path is a `package_root`-relative POSIX path + with no `..` segments and no absolute paths. The package is self-contained + and movable. This is the format you ship. +- `"installed"`: absolute paths and `..` segments are allowed. This is for + packages that have been "installed" onto a system that links shared assets + to a system-wide cache, or that reference pre-existing files outside the + package root. + +The library enforces these rules at parse time. `ModelPackageOpenOptions. +allow_external_paths` can additionally relax portable confinement for read +operations, but the parser still rejects absolute paths inside the manifest +unless `layout == "installed"`. + +--- + +## `manifest.json` + +```jsonc +{ + "schema_version": "1.0", // required, "." (major gates compat) + "package_name": "phi-4-mini", // optional, free-form + "package_version":"4.0.0", // optional, free-form + "description": "Phi-4 mini reasoning model.", // optional + "layout": "portable", // optional: "portable" (default) | "installed" + + "components": { // required, at least one entry + "decoder": "decoder", // external — path relative to package_root + "encoder": { /* inline component body */ } + }, + + "shared_assets": { // optional + "sha256:<64hex>": "shared_assets/sha256-<64hex>" // optional path override + }, + + "additional_metadata": { /* free-form */ } // optional } +``` + +Field reference: + +| Field | Type | Required | Notes | +| -------------------- | --------------- | -------- | ----- | +| `schema_version` | string | yes | `"."` (e.g. `"1.0"`). The library accepts any package whose **major** is in its supported range and any **minor**; a major outside the range is an `ERR_VERSION`. A bare integer is accepted as `".0"`. Major gates compatibility; minor tells consumers which optional fields may be present. | +| `package_name` | string | no | Human label. Not used for resolution. | +| `package_version` | string | no | Human label. Not used for resolution. | +| `description` | string | no | Free-form. | +| `layout` | string | no | `"portable"` (default) or `"installed"`. | +| `components` | object | yes | Map of component name → component value. See below. | +| `shared_assets` | object | no | Map of `sha256:` URI → path override (string). | +| `additional_metadata`| any JSON value | no | Opaque to this library. Round-tripped verbatim. | + +By default the parser rejects unknown top-level keys (`strict_unknown_fields`, +on by default). Disable it via `ModelPackageOpenOptions` to round-trip +manifests authored against a newer schema. + +### Components -size_t count = 0; -ModelPackage_GetComponentCount(ctx, &count); +The value under `components[name]` is either: -for (size_t i = 0; i < count; i++) { - const char* name = NULL; - ModelPackage_GetComponentName(ctx, i, &name); - printf("Component: %s\n", name); +- **A string** — the path to an external component, resolved against + `package_root`. The path may be: + - **A directory.** The loader appends `component.json` and reads that + file. The filename is fixed in this form (must be exactly + `component.json`). + - **A file.** Loaded directly. The filename is not enforced and may be + anything (e.g. `decoder.json`). Useful when one directory holds + multiple component definitions. +- **A JSON object** — an inline component body matching the + [component schema](#componentjson) below. + +The component's "directory" is: + +- For an inline component, the package root itself. +- For an external component pointed at by a directory path, that directory. +- For an external component pointed at by a file path, the file's parent. + +Variant paths in the component body are resolved against this directory. + +### Shared assets + +`shared_assets[uri]` is an **override**: it says "the asset with this URI +lives at this path", overriding the default convention of +`/shared_assets/sha256-/`. Overrides are eagerly rejected +in portable layout when they would escape `package_root` (e.g. absolute paths, +`..` segments). + +Variants reference shared assets only by embedding `sha256:[/sub/path]` +strings inside their `executor_info` payloads. Consumers resolve those +references through [`ModelPackage_ResolveStringRef`](#path-resolution-rules). +The library never parses `executor_info` payloads, so it has no manifest-level +list of which variant uses which asset. + +--- + +## `component.json` + +When a component is external, `component.json` is the file referenced from +the manifest. When inline, the same body is embedded directly in +`manifest.components[name]`. + +```jsonc +{ + "component_name": "decoder", // optional, descriptive only + "variants": { // required, may be empty + "cpu": { /* variant body */ }, + "cuda": { /* variant body */ } + }, + "additional_metadata": { /* free-form */ } // optional } +``` + +Field reference: + +| Field | Type | Required | Notes | +| -------------------- | ------ | -------- | ----- | +| `component_name` | string | no | Sanity-checked as a string; not used for lookup. The map key in `components` wins. | +| `variants` | object | yes | Map of variant name → variant body. May be empty (placeholder component). | +| `additional_metadata`| any | no | Free-form. | -ModelPackage_ReleaseContext(ctx); +--- + +## Variant body + +A variant binds a single (EP, device, compatibility) triple to a single +on-disk directory plus zero or more per-consumer `executor_info` payloads. + +```jsonc +{ + "variant_directory": "cuda", // optional — defaults to variant name + "ep": "CUDAExecutionProvider", // optional + "device": "gpu", // optional ("cpu" | "gpu" | "npu") + "compatibility_string": "", // optional, opaque to library + "executor_info": { // optional + "ort": "ort_info.json", // string → external file + "other": { "filename": "model.onnx" } // object → inline JSON + }, + "additional_metadata": { /* free-form */ } // optional +} ``` -## Integration with ORT +Field reference: -ORT compiles this library as part of its build and wraps the C API through `OrtModelPackageApi`, adding: -- Variant selection via EP factory compatibility validation -- Session creation with merged options +| Field | Type | Required | Notes | +| ---------------------- | ---------------- | -------- | ----- | +| `variant_directory` | string | no | Path relative to the component directory. Defaults to the variant name. If declared but missing on disk, parse fails. | +| `ep` | string | no | Single ONNX Runtime EP name (e.g. `CPUExecutionProvider`). | +| `device` | string | no | Lower-case `cpu` / `gpu` / `npu`. ORT uses this for variant selection. | +| `compatibility_string` | string | no | Opaque to the library. ORT hands it to the EP's `ValidateCompiledModelCompatibilityInfo` callback. | +| `executor_info` | object | no | Map of consumer namespace → string (external file) or object (inline JSON). | +| `additional_metadata` | any | no | Free-form. | -## Package Format +#### `variant_directory` +- Always interpreted as a directory. +- Resolved against the **component directory** (not the package root). +- The library does not validate the directory's contents; consumers resolve + their own file references relative to it. + +#### `executor_info` + +This is the extension point that lets ORT and any future consumer share a +package without colliding. Keys are consumer namespaces; values are either: + +- **A string** — a path to a JSON file. Resolved against the variant + directory. The file must exist (in strict mode) and parse as JSON. +- **An inline JSON object** — embedded directly in the manifest. + +The library round-trips the payload but never interprets it. See +[`onnxruntime/core/session/model_package/README.md`](../onnxruntime/core/session/model_package/README.md) +for the `"ort"` namespace schema. + +Consumers can embed `sha256:[/sub/path]` references inside their +`executor_info` payload and resolve them through +`ModelPackage_ResolveStringRef`. The library does not maintain a per-variant +list of consumed assets; see [Shared assets](#shared-assets) for how URIs +enter the resolvable set. + +--- + +## Shared assets + +Shared assets are **directories** identified by a content hash. Two packages +that ship the same tokenizer will reuse the same asset directory on disk in +an installed layout, dedup-ing storage and downloads. + +### Canonical asset URI + +`ModelPackage_ComputeDirectoryHash(source_dir)` computes the canonical URI: + +1. Walk `source_dir` recursively, collecting regular files. Empty + subdirectories are ignored. +2. Reject symlinks (portability hazard). +3. For each file, compute `sha256(file_bytes)` → per-file hex digest. +4. Build a manifest text of lines ` \n` + sorted lexicographically by path. Paths use forward slashes, no leading + `./`. Non-ASCII paths must be NFC-normalized by the caller. +5. `asset_uri = "sha256:" + sha256(manifest_text)`, lowercase hex. + +The scheme hashes **both** file contents and file names, so renaming a file +inside an asset changes the URI. The on-disk directory name follows the +convention `sha256-` (dash, not colon) to keep the path filesystem-safe. + +### Default location + +`/shared_assets/sha256-/`. Override per-asset by adding an +entry to `manifest.shared_assets`. + +### How URIs enter the resolvable set + +At Open time the library populates the resolvable shared-asset table from +three sources, in order. Within each tier an already-seen URI is skipped: + +1. **Manifest overrides.** Every entry under `manifest.shared_assets` lands + first. These can also point at non-default paths (subject to the + layout's portability rules). +2. **On-disk discovery.** The library lists `/shared_assets/` + and admits each `sha256-` subdirectory it finds (sorted + lexicographically). The resolved path is the default + `/shared_assets/sha256-/`. A missing `shared_assets/` + directory is fine. +3. **Pending authoring stages.** Any `copy_in=true` source registered via + `ModelPackage_AddSharedAsset` is surfaced at its staged source path so + `ResolveStringRef` works before `Commit`. + +This means the manifest does not need to enumerate the assets that ship in +the conventional `shared_assets/` directory. The override list is only +needed when an asset lives outside the default convention. + +### Adding a shared asset programmatically + +```c +const char* uri = NULL; +ModelPackageStatus* st = ModelPackage_AddSharedAsset( + pkg, + "/path/to/tokenizer", // source_dir + NULL, // expected_uri_or_null (reproducible-build check) + /*copy_in=*/true, // stage for copy at Commit time + &uri); ``` -package_root/ -├── manifest.json # schema_version, components list -└── models/ - └── / - ├── metadata.json # variants + EP compatibility declarations - └── / - ├── variant.json # files list, consumer_metadata - └── model.onnx # (or other model files) + +`copy_in == false` stores an override path in the manifest and is rejected +eagerly in portable layout (the path is unlikely to be portable). `copy_in +== true` stages the source for copy when `ModelPackage_Commit()` runs. + +--- + +## Path resolution rules + +`ModelPackage_ResolveStringRef(pkg, base_dir, input, must_exist, &out)` is +the canonical path resolver. It accepts: + +| Input form | Resolution | +| --------------------------- | ---------- | +| `sha256:` | Returns the on-disk directory for that shared asset. Error if the asset isn't registered. | +| `sha256:/sub/path` | Returns `/sub/path`. The subpath is confined to the asset folder (no absolute, no `..`). | +| Relative path | Resolved against `base_dir` (or `package_root` when `base_dir` is NULL). | +| Absolute path / `..` segments | Allowed only in `installed` layout or when the package was opened with `allow_external_paths = true`. | + +In portable layout the resolver enforces that the resolved path stays +underneath `package_root`. Symlinks are followed by default +(`follow_symlinks`). + +`out_path` is a NUL-terminated thread-local pointer; copy it if it must +outlive the next `ResolveStringRef` call on the same thread. + +--- + +## C API quick tour + +All public entry points are declared in `include/model_package.h`. Reading a +package and walking the info tree: + +```c +#include "model_package.h" + +ModelPackage* pkg = NULL; +if (ModelPackageStatus* st = ModelPackage_Open("/path/to/pkg", NULL, &pkg)) { + fprintf(stderr, "open failed: %s\n", ModelPackageStatus_Message(st)); + ModelPackageStatus_Release(st); + return 1; +} + +const ModelPackageInfo* info = ModelPackage_Info(pkg); +printf("schema=%lld.%lld layout=%s\n", + (long long)info->schema_version_major, (long long)info->schema_version_minor, info->layout); +for (size_t i = 0; i < info->num_components; ++i) { + const ModelComponentInfo* c = &info->components[i]; + printf("component %s (%zu variants)\n", c->name, c->num_variants); + for (size_t v = 0; v < c->num_variants; ++v) { + const ModelVariantInfo* var = &c->variants[v]; + printf(" variant %s dir=%s ep=%s\n", + var->name, + var->variant_directory ? var->variant_directory : "(unset)", + var->ep ? var->ep : "(unset)"); + for (size_t e = 0; e < var->num_executor_infos; ++e) { + const ModelExecutorInfoEntry* ei = &var->executor_infos[e]; + printf(" executor_info[%s] = %s\n", ei->namespace_key, ei->json); + } + } +} + +ModelPackage_Close(pkg); +``` + +Authoring a new package from scratch: + +```c +ModelPackage* pkg = NULL; +ModelPackage_New(&pkg); +ModelPackage_SetMetadata(pkg, "phi-4-mini", "4.0.0", "Phi-4 mini."); + +ModelPackage_SetComponentInline(pkg, "decoder", "{\"variants\": {}}"); +ModelPackage_SetVariant(pkg, "decoder", "cpu", + "{\"variant_directory\":\"decoder/cpu\"," + " \"ep\":\"CPUExecutionProvider\"," + " \"device\":\"cpu\"}"); +ModelPackage_SetVariantExecutorInfoInline( + pkg, "decoder", "cpu", "ort", "{\"model_file\":\"model.onnx\"}"); + +const char* asset_uri = NULL; +ModelPackage_AddSharedAsset(pkg, "/src/tokenizer", NULL, /*copy_in=*/true, &asset_uri); +// asset_uri is owned by pkg; copy it if you need it past the next mutation. + +ModelPackage_Commit(pkg, "/path/to/new_pkg", MODEL_PACKAGE_WRITE_PRESERVE); +ModelPackage_Close(pkg); ``` -Single-component shorthand (metadata.json at root, no manifest.json) is also supported. +### Lifetime contract + +Every `const char*` and every `const ModelPackageInfo*` (plus sub-arrays) +returned by the read API is owned by the `ModelPackage` handle and remains +valid **until the next mutation of that scope** or until +`ModelPackage_Close()`. Any `Set*` / `Remove*` / `Add*` / `Commit` call +invalidates cached pointers in the mutated scope; re-read `Info()` after +mutating. + +`ModelPackage_AddSharedAsset`'s `out_uri` follows the same "valid until next +mutation" rule. + +`ModelPackage_ResolveStringRef` and `ModelPackage_ComputeDirectoryHash` +return pointers into a per-thread scratch slot; copy before the next call on +the same thread. + +### Commit modes + +`ModelPackage_Commit(pkg, dest, mode)`: + +- `dest == NULL` → in-place commit at `package_root`. +- `dest != NULL` → write a self-contained "save as". `dest` must be empty or + nonexistent. On success the package's root is updated to `dest`, so + subsequent in-place commits go there. + +`mode`: + +- `MODEL_PACKAGE_WRITE_PRESERVE` (default) — each component and + `executor_info` entry keeps its current inline-or-external shape. +- `MODEL_PACKAGE_WRITE_DENSE` — flatten every external component back inline + into `manifest.json`. Useful for single-file authoring inspection. + +### Prune + +`ModelPackage_Prune(pkg)` reclaims storage that the library itself manages: + +- Tracked orphan variant and component directories left behind by + `RemoveVariant`, `RemoveComponent`, `SetVariant`, or + `SetComponentExternal`. +- Stale `.tmp.` staging directories from interrupted commits, after + a short grace window. + +`Prune` deliberately never removes `shared_assets/sha256-/` directories. +Consumers freely embed `sha256:` references inside their own `executor_info` +payloads, and the library cannot prove an asset is unused without parsing +every consumer's namespace. Use `ModelPackage_RemoveSharedAsset(uri)` to +delete a shared asset explicitly when the caller knows it is unreferenced. + +Only paths registered through this API and strictly inside `package_root` +are touched. + +### Validate + +`ModelPackage_Validate(pkg, flags, &report_json)` runs a configurable set of +structural checks and returns a JSON report +`{"errors": [...], "warnings": [...]}`: + +| Flag | Checks | +| --------------------------------------- | ------ | +| `MODEL_PACKAGE_VALIDATE_SCHEMA` | Required keys, types, value ranges. | +| `MODEL_PACKAGE_VALIDATE_PATHS` | Every recorded path resolves under the configured layout. | +| `MODEL_PACKAGE_VALIDATE_ASSET_REHASH` | Recompute every asset directory hash and compare to its URI (slow). | +| `MODEL_PACKAGE_VALIDATE_UNKNOWN_FIELDS` | Surface unknown JSON fields as warnings. | +| `MODEL_PACKAGE_VALIDATE_ALL` | All of the above. | + +Errors cause a non-NULL status return; warnings alone return success. + +--- + +## Versioning and compatibility + +### Distributed as source + +The library is meant to be **vendored and compiled into each consumer's own +binary** (ORT, publisher tooling, third-party loaders). No prebuilt shared +library (`.so`/`.dll`) is published as the supported interface. + +A direct consequence is that the public POD structs in `model_package.h` have +**no binary boundary** to defend: within any single build there is exactly one +definition of every struct, so there is nothing for two separately-compiled +artifacts to disagree about. The library therefore carries **none** of the usual +ABI machinery — no per-struct `struct_size`/`cbSize`, no `abi_version`, no +library SOVERSION, and no offset `static_assert`s. Collections are exposed as +plain array members (`components`/`num_components`, `variants`/`num_variants`, +…) rather than count+index accessors, since accessors only earn their keep when +the library owns the struct stride across a binary boundary. + +The **only** compatibility contract is the on-disk data format, expressed by +`schema_version`. Everything a consumer needs to know about which fields and +objects a package may contain follows from that one value. + +### `schema_version` + +`schema_version` is a `"."` string in `manifest.json` (a bare +integer `N` is accepted and treated as `N.0`). It is parsed into +`ModelPackageInfo.schema_version_major` and `schema_version_minor`. + +- **major** — the data contract. Incremented only for a **breaking** change + (a field removed, renamed, retyped, or given new semantics). A consumer that + understands major *N* can read any `N.x` package. +- **minor** — additive evolution within a major. Incremented when a new + **optional** field or object is added. It never removes or reinterprets + anything, so it is fully backward- and forward-compatible within the major. + +Consumers should branch **solely on `schema_version_major` / `schema_version_minor`** +to decide which optional fields a package may carry — not on the presence or +absence of individual fields, and never on any library version. + +### What the parser enforces + +Each build declares the majors it understands as a closed range +(`kMinSupportedSchemaMajor … kMaxSupportedSchemaMajor` in `manifest_parser.cc`) +plus the highest minor it authored (`kMaxKnownSchemaMinor`): + +- **Unsupported major** → `ModelPackage_Open` fails with + `MODEL_PACKAGE_ERR_VERSION`. A consumer never silently misreads a package + whose contract it does not understand. +- **Any minor is accepted.** When the minor is **newer** than this build knows + (`minor > kMaxKnownSchemaMinor`), unknown-field strictness is relaxed for that + package so the additive fields a newer authoring tool wrote are **tolerated** + (read through, preserved on round-trip via the JSON getters) instead of + rejected. An older library can therefore load a newer-minor package and ignore + the fields it does not recognize. + +### Supporting a major version bump + +When a breaking change requires a new major, deployed packages do **not** have to +be rewritten and consumers do **not** have to upgrade in lockstep. The library is +designed to support a **range** of majors simultaneously: + +1. Bump `kMaxSupportedSchemaMajor` and add the new major's parse/serialize path, + keeping the existing major's path in place. The supported range now spans both. +2. Existing `N.x` packages keep loading unchanged through the old path; new + `(N+1).x` packages load through the new path. +3. Consumers branch on `schema_version_major` to pick the field set they read. + Code that only supports major *N* simply declines `(N+1).x` packages (the open + call returns `MODEL_PACKAGE_ERR_VERSION` for it) rather than misreading them. +4. A major is dropped from the supported range only when its packages are no + longer in circulation — an explicit, opt-in deprecation, never an implicit + break. + +This keeps already-published packages valid for as long as the library advertises +their major, which is the backward-compatibility guarantee external publishers +depend on. + +### How a major bump maps onto the structs + +A natural question is how a single C struct can represent two majors with +different fields. It can't — and it never has to, because **there is only one +struct definition in any given build**. The "old major" exists only as JSON on +disk; it is never a second C type in the consumer's binary. Since the library is +compiled from source, every consumer compiles exactly one definition of +`ModelPackageInfo`/`ModelVariantInfo`/etc. — the current one. Reconciling an +old-major package with that one definition is a **parse-time** job, not a +struct-layout one. + +The single struct is the **superset / newest** shape, and divergence between +majors is absorbed in three places: + +1. **Additive differences (common).** A field a new major added is present in the + struct and is simply `NULL`/`0`/empty when an older-major package lacks it — + the same mechanism as a minor bump. The consumer treats absence as "not + provided". + +2. **Parse-time normalization (preferred).** When a new major is added, its + parser path is added alongside the existing one, and **both populate the same + struct**. An older-major package is mapped up to the current in-memory model + (defaults filled, renamed fields mapped to their current names) before the + consumer sees it, so reads are uniform. `schema_version_major` then records the + *source* contract — useful for write-back and provenance — rather than + selecting a layout. + +3. **Non-migratable changes (rare).** A field whose *type* changes, or one + removed with no equivalent, cannot reuse the same name (C gives one field one + type). Add a new field for the new representation, populate the old field only + for old-major packages and the new field only for new-major packages, and let + the consumer branch on `schema_version_major`: + + ```c + // e.g. major 1 stored a single compatibility string; major 2 stores a list + const char* compatibility_string; // set when schema_version_major == 1 + const char* const* compatibilities; // set when schema_version_major == 2 + size_t num_compatibilities; + ``` + +**Escape hatch.** If a major bump is sweeping enough that the superset becomes +unwieldy, the standard move is **per-major typed structs** (e.g. a +`ModelPackageInfoV2` returned by a versioned accessor) — a deliberate API +expansion reserved for a wholesale redesign, not the default. In practice: prefer +normalizing old majors up to the newest struct at parse time; fall back to extra +nullable fields plus `schema_version_major` branching only when a change cannot be +auto-migrated. + +--- + +## What the library deliberately does NOT do + +- **Variant selection.** Picking which variant best matches the EPs the + caller has available requires EP factory introspection and is owned by the + executor. ORT's selector lives in + `onnxruntime/core/session/model_package/` and uses each EP's + `ValidateCompiledModelCompatibilityInfo` callback. +- **Session creation.** Building an `OrtSession` is ORT's job. +- **Interpreting `executor_info` payloads.** Each consumer namespace owns + its own slot. The library only validates that values are either strings + (paths) or objects. +- **Interpreting `compatibility_string`.** The format is owned by the EP + declared in `ep`. The library never parses it. + +--- + +## Building + +```bash +cmake -B build -S . [-DMODEL_PACKAGE_BUILD_TESTS=ON] +cmake --build build -j +ctest --test-dir build --output-on-failure # requires BUILD_TESTS=ON +``` + +CMake options: + +- `MODEL_PACKAGE_BUILD_SHARED` (default `ON`) — shared vs static. +- `MODEL_PACKAGE_BUILD_TESTS` (default `OFF`) — build the unit-test + executables (`test_asset_hashing`, `test_inspection`, `test_authoring`, + `test_commit`). + +The only build-time dependency is a vendored copy of nlohmann/json (header +only). + +--- + +## See also + +- `onnxruntime/core/session/model_package/README.md` — how ORT consumes this + library and the `executor_info["ort"]` schema. +- `model_package_redesign.md` in the `archive` repo — original design + rationale (extension fields, content addressing, portable vs installed, + shared-asset overrides). diff --git a/model_package/include/model_package.h b/model_package/include/model_package.h new file mode 100644 index 0000000000000..3b456852adb37 --- /dev/null +++ b/model_package/include/model_package.h @@ -0,0 +1,350 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +/// \file model_package.h +/// \brief Public C API for the ONNX Runtime Model Package library. +/// +/// A model package is a directory with a top-level `manifest.json` that +/// declares a set of components; each component declares a set of variants; +/// each variant points at a directory containing the model files and may +/// carry executor-specific configuration under per-namespace +/// `executor_info` entries. +/// +/// Error handling: functions that can fail return `ModelPackageStatus*`. +/// A `nullptr` return indicates success. Use `ModelPackageStatus_Message`, +/// `ModelPackageStatus_Code`, and `ModelPackageStatus_Release` to inspect and +/// release statuses. +/// +/// Object lifetime: every `const char*` and every `const ModelPackageInfo*` +/// (and its sub-arrays) returned by this API is owned by the `ModelPackage` +/// handle and remains valid until the next mutation of that scope or until +/// the package is closed. Mutations invalidate cached pointers in the mutated +/// scope and its descendants; callers must re-read `ModelPackage_Info()` +/// after any mutation. + +#pragma once + +#include +#include +#include + +#include "model_package_api.h" + +#ifdef __cplusplus +extern "C" { +#endif + +// The library is consumed as source (compiled into each consumer's own binary), +// so the structs below have no binary boundary to maintain: there is no +// struct_size/ABI versioning. Compatibility with on-disk packages is governed +// solely by `schema_version` (see ModelPackageInfo). + +// ───────────────────────────────────────────────────────────────────────────── +// Opaque handle +// ───────────────────────────────────────────────────────────────────────────── + +typedef struct ModelPackage ModelPackage; + +// ───────────────────────────────────────────────────────────────────────────── +// Status helpers +// ───────────────────────────────────────────────────────────────────────────── + +/// Get the error message from a status object. Returns NULL if `status` is NULL. +/// The returned string is owned by the status object. +MODEL_PACKAGE_API const char* ModelPackageStatus_Message(const ModelPackageStatus*); +/// Get the categorical error code. Returns `MODEL_PACKAGE_OK` when `status` is NULL. +MODEL_PACKAGE_API ModelPackageErrorCode ModelPackageStatus_Code(const ModelPackageStatus*); +/// Release a status object. Safe to call with NULL. +MODEL_PACKAGE_API void ModelPackageStatus_Release(ModelPackageStatus*); + +// ───────────────────────────────────────────────────────────────────────────── +// Lifecycle +// ───────────────────────────────────────────────────────────────────────────── + +typedef struct ModelPackageOpenOptions { + bool allow_external_paths; ///< default false; unlocks absolute paths and `..` segments + bool follow_symlinks; ///< default true + bool strict_unknown_fields; ///< default true; relax to round-trip newer schemas +} ModelPackageOpenOptions; + +/// Open an existing model package directory. `opts` may be NULL for defaults. +MODEL_PACKAGE_API ModelPackageStatus* ModelPackage_Open(const char* package_root, + const ModelPackageOpenOptions* opts, + ModelPackage** out); + +/// Create a new empty in-memory package for from-scratch authoring. +MODEL_PACKAGE_API ModelPackageStatus* ModelPackage_New(ModelPackage** out); + +/// Release a ModelPackage handle and all its caches. Safe on NULL. +MODEL_PACKAGE_API void ModelPackage_Close(ModelPackage* pkg); + +// ───────────────────────────────────────────────────────────────────────────── +// Data model — POD structs read from ModelPackage_Info() +// ───────────────────────────────────────────────────────────────────────────── + +typedef struct ModelExecutorInfoEntry { + const char* namespace_key; ///< executor namespace name (e.g. "ort") + const char* json; ///< canonical JSON value as string (object, array, etc.) +} ModelExecutorInfoEntry; + +typedef struct ModelVariantInfo { + const char* name; + /// Resolved absolute path to the variant's on-disk directory, or NULL when + /// no directory has been declared and the default location does not exist. + const char* variant_directory; + const char* ep; ///< NULL when unset + const char* device; ///< NULL when unset + const char* compatibility_string; ///< NULL when unset + const char* additional_metadata_json; ///< NULL when unset + size_t num_executor_infos; + const ModelExecutorInfoEntry* executor_infos; +} ModelVariantInfo; + +typedef struct ModelComponentInfo { + const char* name; + const char* additional_metadata_json; ///< NULL when unset + size_t num_variants; + const ModelVariantInfo* variants; +} ModelComponentInfo; + +typedef struct ModelSharedAssetInfo { + const char* uri; ///< "sha256:" + const char* resolved_path; ///< absolute on-disk directory path +} ModelSharedAssetInfo; + +typedef struct ModelPackageInfo { + int64_t schema_version_major; ///< parsed from on-disk "."; gates compatibility + int64_t schema_version_minor; ///< informational; indicates which optional fields may be present + const char* package_name; ///< NULL when unset + const char* package_version; ///< NULL when unset + const char* description; ///< NULL when unset + const char* layout; ///< "portable" or "installed" + const char* additional_metadata_json; ///< NULL when unset + size_t num_components; + const ModelComponentInfo* components; + size_t num_shared_assets; + const ModelSharedAssetInfo* shared_assets; +} ModelPackageInfo; + +/// Return the package-level info tree. Pointer is owned by the package and is +/// invalidated by any mutation. +MODEL_PACKAGE_API const ModelPackageInfo* ModelPackage_Info(const ModelPackage* pkg); + +// ───────────────────────────────────────────────────────────────────────────── +// Convenience lookups +// ───────────────────────────────────────────────────────────────────────────── + +/// Find a component by name. Returns NULL when not found. +MODEL_PACKAGE_API const ModelComponentInfo* ModelPackage_FindComponent(const ModelPackageInfo*, + const char* name); +/// Find a variant within a component by name. Returns NULL when not found. +MODEL_PACKAGE_API const ModelVariantInfo* ModelComponentInfo_FindVariant(const ModelComponentInfo*, + const char* name); +/// Find an executor_info entry by namespace. Returns NULL when not declared. +MODEL_PACKAGE_API const ModelExecutorInfoEntry* ModelVariantInfo_FindExecutorInfo( + const ModelVariantInfo*, const char* namespace_key); + +// ───────────────────────────────────────────────────────────────────────────── +// Round-trip JSON getters +// ───────────────────────────────────────────────────────────────────────────── + +/// Get the canonical schema-shaped JSON for the named component. Preserves +/// fields unknown to this build. The returned pointer is owned by the package. +MODEL_PACKAGE_API ModelPackageStatus* ModelPackage_GetComponentJson(const ModelPackage*, + const char* component_name, + const char** out_json); + +/// Get the canonical schema-shaped JSON for the named variant. +MODEL_PACKAGE_API ModelPackageStatus* ModelPackage_GetVariantJson(const ModelPackage*, + const char* component_name, + const char* variant_name, + const char** out_json); + +// ───────────────────────────────────────────────────────────────────────────── +// Asset resolution + hashing +// ───────────────────────────────────────────────────────────────────────────── + +/// Resolve a `sha256:` URI to an on-disk directory. Errors with +/// `MODEL_PACKAGE_ERR_ASSET_MISSING` when not resolvable. +MODEL_PACKAGE_API ModelPackageStatus* ModelPackage_ResolveAssetUri(const ModelPackage*, + const char* uri, + const char** out_path); + +/// Resolve a string reference using the model package's path resolution rules. +/// `input` may be: +/// - `sha256:` -> shared-asset folder +/// - `sha256:/sub/path` -> file or subdir inside a shared-asset folder +/// (sub/path is resolved with portable-mode +/// confinement under the asset folder: no +/// absolute, no `..`) +/// - relative path -> resolved against `base_dir` (or +/// `package_root` when `base_dir == NULL`), +/// confined to `package_root` in portable layout +/// - absolute path / `..` segments -> only allowed in installed layout, or in +/// any layout when the package was opened with +/// `ModelPackageOpenOptions.allow_external_paths` +/// +/// `must_exist` controls whether a missing target is `MODEL_PACKAGE_ERR_NOT_FOUND` +/// or the lexically-normalized path is returned anyway. +/// On success `*out_path` points to a NUL-terminated thread-local string; copy +/// it if you need it to outlive the next `ModelPackage_ResolveStringRef` call on +/// the same thread. +MODEL_PACKAGE_API ModelPackageStatus* ModelPackage_ResolveStringRef(const ModelPackage*, + const char* base_dir, + const char* input, + bool must_exist, + const char** out_path); + +/// Compute the canonical `sha256:` URI for a directory. On success, +/// `*out_uri` is set to a NUL-terminated string owned by an internal +/// thread-local slot; the caller must copy if it must outlive the next call +/// on the same thread. +MODEL_PACKAGE_API ModelPackageStatus* ModelPackage_ComputeDirectoryHash(const char* source_dir, + const char** out_uri); + +// ───────────────────────────────────────────────────────────────────────────── +// Authoring — mutation API +// ───────────────────────────────────────────────────────────────────────────── +// +// Each mutation invalidates info pointers in the mutated scope and its +// descendants. Strict unknown-field rejection follows the open-time option +// `strict_unknown_fields` (default true). + +/// Set or replace an inline component. `component_json` must be a JSON object +/// matching the component schema. An existing component with the same name is +/// replaced. +MODEL_PACKAGE_API ModelPackageStatus* ModelPackage_SetComponentInline(ModelPackage*, + const char* name, + const char* component_json); + +/// Set or replace an external component. `path` is recorded in the manifest +/// (relative to package_root, or absolute in installed layout). If the file +/// exists, it is loaded; otherwise the component is initialized empty +/// (`{"variants": {}}`). `path` may be a directory (resolves to +/// `/component.json`). +MODEL_PACKAGE_API ModelPackageStatus* ModelPackage_SetComponentExternal(ModelPackage*, + const char* name, + const char* path); + +/// Remove a component by name. No-op when the name is not present. +MODEL_PACKAGE_API ModelPackageStatus* ModelPackage_RemoveComponent(ModelPackage*, const char* name); + +/// Upsert a variant inside a component. `variant_json` must be a JSON object +/// matching the variant schema. The library does not validate that +/// `variant_directory` exists on disk; executors are responsible for resolving +/// their own file references at load time. +MODEL_PACKAGE_API ModelPackageStatus* ModelPackage_SetVariant(ModelPackage*, + const char* component_name, + const char* variant_name, + const char* variant_json); + +MODEL_PACKAGE_API ModelPackageStatus* ModelPackage_RemoveVariant(ModelPackage*, + const char* component_name, + const char* variant_name); + +MODEL_PACKAGE_API ModelPackageStatus* ModelPackage_SetVariantExecutorInfoInline(ModelPackage*, + const char* component, + const char* variant, + const char* ns, + const char* info_json); + +MODEL_PACKAGE_API ModelPackageStatus* ModelPackage_SetVariantExecutorInfoExternal(ModelPackage*, + const char* component, + const char* variant, + const char* ns, + const char* path); + +MODEL_PACKAGE_API ModelPackageStatus* ModelPackage_RemoveVariantExecutorInfo(ModelPackage*, + const char* component, + const char* variant, + const char* ns); + +/// Add a content-addressed shared asset. When `expected_uri_or_null` is +/// non-NULL, the computed URI must match (reproducible-build check). With +/// `copy_in == false`, an override path is stored in the manifest; this is +/// rejected eagerly in portable layout. With `copy_in == true`, the source +/// directory is staged for copy at `_Commit` time. +/// +/// `out_uri` is set to a NUL-terminated string owned by the package. The +/// pointer is only guaranteed to remain valid until the next mutation +/// (any ModelPackage_Set*, ModelPackage_Remove*, ModelPackage_AddSharedAsset, +/// or ModelPackage_Commit call), since those calls may rebuild the +/// shared-asset table or rehash the pending-copies map. Callers that need to +/// retain the URI must copy it into their own storage. +MODEL_PACKAGE_API ModelPackageStatus* ModelPackage_AddSharedAsset(ModelPackage*, + const char* source_dir, + const char* expected_uri_or_null, + bool copy_in, + const char** out_uri); + +MODEL_PACKAGE_API ModelPackageStatus* ModelPackage_RemoveSharedAsset(ModelPackage*, const char* uri); + +/// Set or clear package-level metadata. Any argument may be NULL to leave the +/// existing value untouched. Passing an empty string clears the field. +MODEL_PACKAGE_API ModelPackageStatus* ModelPackage_SetMetadata(ModelPackage*, + const char* name_or_null, + const char* version_or_null, + const char* description_or_null); + +/// Set the layout. Valid values: "portable" or "installed". +MODEL_PACKAGE_API ModelPackageStatus* ModelPackage_SetLayout(ModelPackage*, const char* layout); + +/// Set or clear `additional_metadata` at a given scope. +/// scope = "manifest" — component_or_null and variant_or_null must be NULL +/// scope = "component" — component_or_null is required, variant_or_null is NULL +/// scope = "variant" — component_or_null and variant_or_null are both required +/// `json_or_null == NULL` clears the field at that scope. +MODEL_PACKAGE_API ModelPackageStatus* ModelPackage_SetAdditionalMetadataJson(ModelPackage*, + const char* scope, + const char* component_or_null, + const char* variant_or_null, + const char* json_or_null); + +// ───────────────────────────────────────────────────────────────────────────── +// Commit / Prune / Validate +// ───────────────────────────────────────────────────────────────────────────── + +typedef enum { + MODEL_PACKAGE_WRITE_PRESERVE = 0, ///< each component/executor-info keeps its current shape + MODEL_PACKAGE_WRITE_DENSE = 1, ///< flatten all external components inline +} ModelPackageWriteMode; + +/// Persist the in-memory model to disk. `dest_root_or_null == NULL` commits +/// in-place at `package_root`. Otherwise `dest_root` must be empty or +/// nonexistent and the entire package is materialized there (self-contained +/// "save as"). On a successful dest_root commit, the package's root is +/// updated to `dest_root` so subsequent in-place commits go there. +MODEL_PACKAGE_API ModelPackageStatus* ModelPackage_Commit(ModelPackage*, + const char* dest_root_or_null, + ModelPackageWriteMode mode); + +/// Reclaim stale `.tmp.` staging directories under +/// `/shared_assets/` (left by interrupted commits, after a grace +/// window) and tracked orphan variant/component directories left behind by +/// RemoveVariant, RemoveComponent, SetVariant or SetComponentExternal. Only +/// paths registered through this API and inside `package_root` are touched. +/// Content-addressed shared-asset (`sha256-`) directories are never removed +/// — use ModelPackage_RemoveSharedAsset to reclaim those. +MODEL_PACKAGE_API ModelPackageStatus* ModelPackage_Prune(ModelPackage*); + +typedef enum { + MODEL_PACKAGE_VALIDATE_SCHEMA = 1 << 0, + MODEL_PACKAGE_VALIDATE_PATHS = 1 << 1, + MODEL_PACKAGE_VALIDATE_ASSET_REHASH = 1 << 2, + MODEL_PACKAGE_VALIDATE_UNKNOWN_FIELDS = 1 << 3, + MODEL_PACKAGE_VALIDATE_ALL = ~0, +} ModelPackageValidateFlags; + +/// Run structural and reachability checks. `*out_report_json` is set to a +/// JSON string owned by the package describing findings: +/// `{"errors": [{"code": "...", "message": "..."}, ...], +/// "warnings": [...]}` +/// Returns non-NULL status when any error-level finding fired; warnings alone +/// still return success. +MODEL_PACKAGE_API ModelPackageStatus* ModelPackage_Validate(ModelPackage*, + int flags, + const char** out_report_json); + +#ifdef __cplusplus +} // extern "C" +#endif diff --git a/model_package/include/model_package_api.h b/model_package/include/model_package_api.h index ca840c8a33e0e..36e678feed0f6 100644 --- a/model_package/include/model_package_api.h +++ b/model_package/include/model_package_api.h @@ -2,18 +2,15 @@ // Licensed under the MIT License. /// \file model_package_api.h -/// \brief Standalone C API for parsing and inspecting ONNX Runtime Model Packages. +/// \brief Core types shared by the model_package public API surface. /// -/// This library has no dependency on ONNX Runtime. It provides read-only access to -/// model package structure: components, variants, EP compatibility declarations, -/// model files, session/provider options, and consumer metadata. +/// This header defines the export macro, the opaque `ModelPackageStatus` type, +/// and the `ModelPackageErrorCode` enum used by every entry point in the +/// library. The actual API entry points live in `model_package.h`. /// -/// Error handling: Functions that can fail return `ModelPackageStatus*`. -/// A nullptr return indicates success. On failure, use `ModelPackage_GetErrorMessage()` -/// to retrieve the error string, and `ModelPackage_ReleaseStatus()` to free it. -/// -/// Lifetime: All `const char*` pointers returned by this API are owned by the -/// `ModelPackageContext` and remain valid until it is released. +/// Error handling: functions that can fail return `ModelPackageStatus*`. A +/// `nullptr` return indicates success. Use the `ModelPackageStatus_*` helpers +/// in `model_package.h` to inspect and release statuses. #pragma once @@ -45,110 +42,31 @@ extern "C" { #endif // ───────────────────────────────────────────────────────────────────────────── -// Opaque types +// Opaque status type // ───────────────────────────────────────────────────────────────────────────── /// Opaque status type. nullptr indicates success. typedef struct ModelPackageStatus ModelPackageStatus; -/// Opaque context holding a parsed model package. -typedef struct ModelPackageContext ModelPackageContext; - -// ───────────────────────────────────────────────────────────────────────────── -// Status API -// ───────────────────────────────────────────────────────────────────────────── - -/// Release a status object. Safe to call with nullptr. -MODEL_PACKAGE_API void ModelPackage_ReleaseStatus(ModelPackageStatus* status); - -/// Get the error message from a status object. Returns nullptr if status is nullptr. -/// The returned string is owned by the status object. -MODEL_PACKAGE_API const char* ModelPackage_GetErrorMessage(const ModelPackageStatus* status); - -// ───────────────────────────────────────────────────────────────────────────── -// Context lifecycle -// ───────────────────────────────────────────────────────────────────────────── - -/// Parse a model package from a directory path and create a context. -/// -/// \param[in] package_root_path Null-terminated UTF-8 path to the package root directory. -/// \param[out] out_context On success, receives the created context. Caller must release -/// via ModelPackage_ReleaseContext(). -/// \return nullptr on success, or a status object describing the error. -MODEL_PACKAGE_API ModelPackageStatus* ModelPackage_CreateContext( - const char* package_root_path, - ModelPackageContext** out_context); - -/// Release a model package context and all associated resources. -/// Safe to call with nullptr. -MODEL_PACKAGE_API void ModelPackage_ReleaseContext(ModelPackageContext* context); - -// ───────────────────────────────────────────────────────────────────────────── -// Package-level queries -// ───────────────────────────────────────────────────────────────────────────── - -/// Get the schema version declared in manifest.json. -MODEL_PACKAGE_API ModelPackageStatus* ModelPackage_GetSchemaVersion( - const ModelPackageContext* context, - int64_t* out_version); - -// ───────────────────────────────────────────────────────────────────────────── -// Component queries -// ───────────────────────────────────────────────────────────────────────────── - -/// Get the number of components in the package. -MODEL_PACKAGE_API ModelPackageStatus* ModelPackage_GetComponentCount( - const ModelPackageContext* context, - size_t* out_count); - -/// Get the name of a component by index. -/// -/// \param[in] context The package context. -/// \param[in] component_idx Zero-based index (must be < component count). -/// \param[out] out_name Receives a pointer to the component name string. -/// Lifetime is tied to the context. -MODEL_PACKAGE_API ModelPackageStatus* ModelPackage_GetComponentName( - const ModelPackageContext* context, - size_t component_idx, - const char** out_name); - -// ───────────────────────────────────────────────────────────────────────────── -// Variant queries // ───────────────────────────────────────────────────────────────────────────── - -/// Get the number of variants for a component. -MODEL_PACKAGE_API ModelPackageStatus* ModelPackage_GetVariantCount( - const ModelPackageContext* context, - const char* component_name, - size_t* out_count); - -/// Get the name of a variant by index. -MODEL_PACKAGE_API ModelPackageStatus* ModelPackage_GetVariantName( - const ModelPackageContext* context, - const char* component_name, - size_t variant_idx, - const char** out_name); - -/// Get the folder path for a variant (resolved absolute path). -MODEL_PACKAGE_API ModelPackageStatus* ModelPackage_GetVariantFolderPath( - const ModelPackageContext* context, - const char* component_name, - const char* variant_name, - const char** out_path); - -// ───────────────────────────────────────────────────────────────────────────── -// EP compatibility queries -// ───────────────────────────────────────────────────────────────────────────── - -/// Get the EP name declared for a variant. -/// -/// Each variant targets a single EP. When the variant does not declare an EP, -/// the returned pointer is set to nullptr. -MODEL_PACKAGE_API ModelPackageStatus* ModelPackage_GetVariantEpName( - const ModelPackageContext* context, - const char* component_name, - const char* variant_name, - const char** out_ep); +// Error codes +// ───────────────────────────────────────────────────────────────────────────── + +/// Categorical error codes attached to every non-OK ModelPackageStatus. +/// Stable additive enum: new codes will be appended at the end; existing +/// values will not be renumbered. +typedef enum ModelPackageErrorCode { + MODEL_PACKAGE_OK = 0, + MODEL_PACKAGE_ERR_IO = 1, ///< Filesystem read/write/sync failure. + MODEL_PACKAGE_ERR_SCHEMA = 2, ///< JSON value has wrong shape or wrong type. + MODEL_PACKAGE_ERR_VERSION = 3, ///< Unsupported schema_version. + MODEL_PACKAGE_ERR_PATH_CONFINEMENT = 4, ///< Path resolution escaped the allowed base. + MODEL_PACKAGE_ERR_ASSET_MISSING = 5, ///< Declared shared asset not resolvable. + MODEL_PACKAGE_ERR_ASSET_HASH_MISMATCH = 6, ///< Existing asset directory failed rehash. + MODEL_PACKAGE_ERR_NOT_FOUND = 7, ///< Named entity not present. + MODEL_PACKAGE_ERR_INVALID_ARG = 8, ///< Null pointer or otherwise invalid argument. + MODEL_PACKAGE_ERR_STATE = 9 ///< Operation not legal in current state. +} ModelPackageErrorCode; #ifdef __cplusplus } // extern "C" diff --git a/model_package/src/api.cc b/model_package/src/api.cc deleted file mode 100644 index 103bff8e1a4a3..0000000000000 --- a/model_package/src/api.cc +++ /dev/null @@ -1,243 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT License. - -#include "model_package_api.h" -#include "model_package_internal.h" -#include "parser.h" - -#include -#include - -// ───────────────────────────────────────────────────────────────────────────── -// Status implementation -// ───────────────────────────────────────────────────────────────────────────── - -struct ModelPackageStatus { - std::string message; -}; - -static ModelPackageStatus* MakeError(std::string msg) { - return new (std::nothrow) ModelPackageStatus{std::move(msg)}; -} - -// ───────────────────────────────────────────────────────────────────────────── -// Context is the public opaque type wrapping ContextImpl -// ───────────────────────────────────────────────────────────────────────────── - -struct ModelPackageContext { - model_package::ContextImpl impl; -}; - -// ───────────────────────────────────────────────────────────────────────────── -// ContextImpl lookup helpers -// ───────────────────────────────────────────────────────────────────────────── - -namespace model_package { - -const Component* ContextImpl::FindComponent(const char* name) const { - for (const auto& c : package_info.components) { - if (c.name == name) return &c; - } - return nullptr; -} - -const Variant* ContextImpl::FindVariant(const char* component_name, const char* variant_name) const { - const auto* comp = FindComponent(component_name); - if (!comp) return nullptr; - for (const auto& v : comp->variants) { - if (v.name == variant_name) return &v; - } - return nullptr; -} - -} // namespace model_package - -// ───────────────────────────────────────────────────────────────────────────── -// Validation macro -// ───────────────────────────────────────────────────────────────────────────── - -#define RETURN_IF_NULL(ptr, param_name) \ - do { \ - if ((ptr) == nullptr) \ - return MakeError(std::string(param_name) + " must not be null."); \ - } while (0) - -// ───────────────────────────────────────────────────────────────────────────── -// C API implementation -// ───────────────────────────────────────────────────────────────────────────── - -extern "C" { - -void ModelPackage_ReleaseStatus(ModelPackageStatus* status) { - delete status; -} - -const char* ModelPackage_GetErrorMessage(const ModelPackageStatus* status) { - if (status == nullptr) return nullptr; - return status->message.c_str(); -} - -ModelPackageStatus* ModelPackage_CreateContext( - const char* package_root_path, - ModelPackageContext** out_context) { - RETURN_IF_NULL(package_root_path, "package_root_path"); - RETURN_IF_NULL(out_context, "out_context"); - - *out_context = nullptr; - - auto ctx = std::make_unique(); - std::string error; - - if (!model_package::ParsePackage( - std::filesystem::path(std::string(package_root_path)), - ctx->impl.package_info, error)) { - return MakeError(std::move(error)); - } - - // Build component names cache. - ctx->impl.component_names_cache.clear(); - for (const auto& c : ctx->impl.package_info.components) { - ctx->impl.component_names_cache.push_back(c.name); - } - - // Build variant names cache. - for (const auto& c : ctx->impl.package_info.components) { - auto& names = ctx->impl.variant_names_cache[c.name]; - names.clear(); - for (const auto& v : c.variants) { - names.push_back(v.name); - } - } - - *out_context = ctx.release(); - return nullptr; -} - -void ModelPackage_ReleaseContext(ModelPackageContext* context) { - delete context; -} - -ModelPackageStatus* ModelPackage_GetSchemaVersion( - const ModelPackageContext* context, - int64_t* out_version) { - RETURN_IF_NULL(context, "context"); - RETURN_IF_NULL(out_version, "out_version"); - *out_version = context->impl.package_info.schema_version; - return nullptr; -} - -ModelPackageStatus* ModelPackage_GetComponentCount( - const ModelPackageContext* context, - size_t* out_count) { - RETURN_IF_NULL(context, "context"); - RETURN_IF_NULL(out_count, "out_count"); - *out_count = context->impl.package_info.components.size(); - return nullptr; -} - -ModelPackageStatus* ModelPackage_GetComponentName( - const ModelPackageContext* context, - size_t component_idx, - const char** out_name) { - RETURN_IF_NULL(context, "context"); - RETURN_IF_NULL(out_name, "out_name"); - - if (component_idx >= context->impl.component_names_cache.size()) { - return MakeError("component_idx out of range: " + std::to_string(component_idx)); - } - - *out_name = context->impl.component_names_cache[component_idx].c_str(); - return nullptr; -} - -ModelPackageStatus* ModelPackage_GetVariantCount( - const ModelPackageContext* context, - const char* component_name, - size_t* out_count) { - RETURN_IF_NULL(context, "context"); - RETURN_IF_NULL(component_name, "component_name"); - RETURN_IF_NULL(out_count, "out_count"); - - const auto* comp = context->impl.FindComponent(component_name); - if (!comp) { - return MakeError(std::string("Component not found: '") + component_name + "'."); - } - - *out_count = comp->variants.size(); - return nullptr; -} - -ModelPackageStatus* ModelPackage_GetVariantName( - const ModelPackageContext* context, - const char* component_name, - size_t variant_idx, - const char** out_name) { - RETURN_IF_NULL(context, "context"); - RETURN_IF_NULL(component_name, "component_name"); - RETURN_IF_NULL(out_name, "out_name"); - - auto it = context->impl.variant_names_cache.find(component_name); - if (it == context->impl.variant_names_cache.end()) { - return MakeError(std::string("Component not found: '") + component_name + "'."); - } - - if (variant_idx >= it->second.size()) { - return MakeError("variant_idx out of range: " + std::to_string(variant_idx)); - } - - *out_name = it->second[variant_idx].c_str(); - return nullptr; -} - -ModelPackageStatus* ModelPackage_GetVariantFolderPath( - const ModelPackageContext* context, - const char* component_name, - const char* variant_name, - const char** out_path) { - RETURN_IF_NULL(context, "context"); - RETURN_IF_NULL(component_name, "component_name"); - RETURN_IF_NULL(variant_name, "variant_name"); - RETURN_IF_NULL(out_path, "out_path"); - - const auto* variant = context->impl.FindVariant(component_name, variant_name); - if (!variant) { - return MakeError(std::string("Variant '") + variant_name + "' not found in component '" + - component_name + "'."); - } - - // Cache the path string for stable pointer. - std::string cache_key = std::string(component_name) + "/" + variant_name; - auto& cached = const_cast(context)->impl.folder_path_strings_cache[cache_key]; - if (cached.empty()) { - cached = variant->folder_path.string(); - } - *out_path = cached.c_str(); - return nullptr; -} - -ModelPackageStatus* ModelPackage_GetVariantEpName( - const ModelPackageContext* context, - const char* component_name, - const char* variant_name, - const char** out_ep) { - RETURN_IF_NULL(context, "context"); - RETURN_IF_NULL(component_name, "component_name"); - RETURN_IF_NULL(variant_name, "variant_name"); - - const auto* variant = context->impl.FindVariant(component_name, variant_name); - if (!variant) { - return MakeError(std::string("Variant '") + variant_name + "' not found in component '" + - component_name + "'."); - } - - if (out_ep) { - if (variant->ep_compatibility.ep.has_value()) { - *out_ep = variant->ep_compatibility.ep->c_str(); - } else { - *out_ep = nullptr; - } - } - return nullptr; -} - -} // extern "C" diff --git a/model_package/src/asset_hasher.cc b/model_package/src/asset_hasher.cc new file mode 100644 index 0000000000000..b41019d11a757 --- /dev/null +++ b/model_package/src/asset_hasher.cc @@ -0,0 +1,97 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "asset_hasher.h" + +#include +#include +#include +#include + +#include "sha256.h" +#include "status_impl.h" + +namespace fs = std::filesystem; + +namespace model_package { + +using model_package::MakeStatus; + +namespace { + +std::string ToPosix(const fs::path& rel) { + std::string s = rel.generic_string(); // generic_string uses '/' + // Strip leading "./" if any (lexical normalization edge case). + if (s.size() >= 2 && s[0] == '.' && s[1] == '/') s.erase(0, 2); + return s; +} + +} // namespace + +ModelPackageStatus* ComputeDirectoryAssetUri(const fs::path& source_dir, + std::string* out_uri) { + if (!out_uri) { + return MakeStatus(MODEL_PACKAGE_ERR_INVALID_ARG, "ComputeDirectoryAssetUri: out_uri is null."); + } + std::error_code ec; + if (!fs::exists(source_dir, ec) || !fs::is_directory(source_dir, ec)) { + return MakeStatus(MODEL_PACKAGE_ERR_NOT_FOUND, + "ComputeDirectoryAssetUri: '" + source_dir.string() + "' is not a directory."); + } + + // Collect (relative_posix_path, absolute_path) pairs. + std::vector> entries; + + auto walker = fs::recursive_directory_iterator( + source_dir, fs::directory_options::none, ec); + if (ec) { + return MakeStatus(MODEL_PACKAGE_ERR_IO, + "ComputeDirectoryAssetUri: cannot iterate '" + source_dir.string() + + "': " + ec.message()); + } + for (; walker != fs::recursive_directory_iterator(); walker.increment(ec)) { + if (ec) { + return MakeStatus(MODEL_PACKAGE_ERR_IO, + "ComputeDirectoryAssetUri: iteration error: " + ec.message()); + } + const fs::directory_entry& de = *walker; + if (de.is_symlink(ec)) { + return MakeStatus(MODEL_PACKAGE_ERR_SCHEMA, + "ComputeDirectoryAssetUri: symlink not allowed: '" + de.path().string() + "'."); + } + if (de.is_regular_file(ec)) { + fs::path rel = fs::relative(de.path(), source_dir, ec); + if (ec) { + return MakeStatus(MODEL_PACKAGE_ERR_IO, + "ComputeDirectoryAssetUri: relative path failed: " + ec.message()); + } + entries.emplace_back(ToPosix(rel), de.path()); + } else if (!de.is_directory(ec)) { + return MakeStatus(MODEL_PACKAGE_ERR_SCHEMA, + "ComputeDirectoryAssetUri: unsupported file kind: '" + + de.path().string() + "' (only regular files and directories allowed)."); + } + } + + std::sort(entries.begin(), entries.end(), + [](const auto& a, const auto& b) { return a.first < b.first; }); + + std::string manifest_text; + manifest_text.reserve(entries.size() * 96); + for (const auto& entry : entries) { + std::string file_hex = Sha256::HashFileHex(entry.second.string()); + if (file_hex.empty()) { + return MakeStatus(MODEL_PACKAGE_ERR_IO, + "ComputeDirectoryAssetUri: failed to hash file '" + entry.second.string() + "'."); + } + manifest_text.append(file_hex); + manifest_text.append(" "); + manifest_text.append(entry.first); + manifest_text.append("\n"); + } + + *out_uri = "sha256:" + Sha256::HashStringHex(manifest_text); + return nullptr; +} + +} // namespace model_package diff --git a/model_package/src/asset_hasher.h b/model_package/src/asset_hasher.h new file mode 100644 index 0000000000000..f9bd6eb1c5d9b --- /dev/null +++ b/model_package/src/asset_hasher.h @@ -0,0 +1,30 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +/// \file asset_hasher.h +/// \brief Directory Merkle hash for content-addressed shared assets. + +#pragma once + +#include +#include + +#include "model_package_api.h" + +namespace model_package { + +/// Compute the canonical asset URI for a directory: +/// 1. Walk recursively, collect regular files (ignore empty dirs). +/// 2. Reject symlinks (ERR_SCHEMA: portability hazard). +/// 3. For each file, compute sha256(file_bytes) → per-file hex. +/// 4. Build manifest text: ` \n` lines, +/// sorted lexicographically by path. Paths are POSIX (`/`), no leading +/// `./`. NFC normalization is the caller's responsibility for non-ASCII +/// paths; ASCII is identity. +/// 5. asset_uri = "sha256:" + sha256(manifest_text), lowercase hex. +/// +/// On success, *out_uri is set to the URI string. +ModelPackageStatus* ComputeDirectoryAssetUri(const std::filesystem::path& source_dir, + std::string* out_uri); + +} // namespace model_package diff --git a/model_package/src/authoring.cc b/model_package/src/authoring.cc new file mode 100644 index 0000000000000..4c7a9e9c1f6d5 --- /dev/null +++ b/model_package/src/authoring.cc @@ -0,0 +1,593 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +/// \file authoring.cc +/// \brief Mutation (authoring) API implementation. + +#include "model_package.h" + +#include +#include +#include +#include +#include +#include + +#include "asset_hasher.h" +#include "manifest_parser.h" +#include "model_package_impl.h" +#include "path_resolver.h" +#include "status_impl.h" + +namespace fs = std::filesystem; +namespace mp = model_package; +using model_package::MakeStatus; +using nlohmann::ordered_json; + +namespace { + +// Schema version stamped into newly authored packages, written as a "." +// string. Keep in sync with the parser's supported major + highest known minor +// (manifest_parser.cc: kMaxSupportedSchemaMajor / kMaxKnownSchemaMinor). +constexpr const char* kAuthoredSchemaVersion = "1.0"; + +ModelPackageStatus* NullArg(const char* name) { + return MakeStatus(MODEL_PACKAGE_ERR_INVALID_ARG, + std::string("model_package: '") + name + "' must not be null."); +} + +ModelPackageStatus* ParseJsonString(const char* json, const char* where, ordered_json* out) { + try { + *out = ordered_json::parse(json); + } catch (const ordered_json::parse_error& e) { + return MakeStatus(MODEL_PACKAGE_ERR_SCHEMA, + std::string(where) + ": JSON parse error: " + e.what()); + } + return nullptr; +} + +ModelPackageStatus* ExpectObject(const ordered_json& j, const char* where) { + if (!j.is_object()) { + return MakeStatus(MODEL_PACKAGE_ERR_SCHEMA, + std::string(where) + ": expected a JSON object."); + } + return nullptr; +} + +void RebuildComponentIndex(ModelPackage* pkg) { + pkg->component_index_by_name.clear(); + for (size_t i = 0; i < pkg->components.size(); ++i) { + pkg->component_index_by_name[pkg->components[i]->name] = i; + } +} + +mp::ComponentRecord* FindComponentRecord(ModelPackage* pkg, const std::string& name) { + auto it = pkg->component_index_by_name.find(name); + if (it == pkg->component_index_by_name.end()) return nullptr; + return pkg->components[it->second].get(); +} + +mp::VariantRecord* FindVariantRecord(mp::ComponentRecord* comp, const std::string& name) { + for (auto& v : comp->variants) { + if (v->name == name) return v.get(); + } + return nullptr; +} + +ModelPackageStatus* RefreshSharedAssetsHelper(ModelPackage* pkg) { + return mp::RefreshSharedAssets(pkg, mp::PathOptionsFor(pkg)); +} + +ModelPackageStatus* PostMutate(ModelPackage* pkg, bool refresh_assets = true) { + mp::DropViewCache(pkg); + if (refresh_assets) { + if (auto* s = RefreshSharedAssetsHelper(pkg)) return s; + } + if (auto* s = mp::RefreshPackageMetadata(pkg)) return s; + return mp::RefreshExecutorInfoCache(pkg, /*strict_missing_external=*/false); +} + +ordered_json& EnsureManifestComponentsObject(ModelPackage* pkg) { + if (!pkg->manifest.contains("components")) { + pkg->manifest["components"] = ordered_json::object(); + } + return pkg->manifest["components"]; +} + +} // namespace + +extern "C" { + +// ───────────────────────────────────────────────────────────────────────────── +// ModelPackage_New +// ───────────────────────────────────────────────────────────────────────────── + +ModelPackageStatus* ModelPackage_New(ModelPackage** out) { + if (!out) return NullArg("out"); + auto pkg = std::make_unique(); + pkg->manifest = ordered_json::object(); + // Authored at this build's schema version, written as a "." string. + pkg->manifest["schema_version"] = kAuthoredSchemaVersion; + pkg->manifest["layout"] = "portable"; + pkg->manifest["components"] = ordered_json::object(); + pkg->layout = "portable"; + pkg->strict_unknown_fields = true; + pkg->follow_symlinks = true; + pkg->allow_external_paths = false; + pkg->package_root = fs::path(); + if (auto* s = mp::RefreshPackageMetadata(pkg.get())) return s; + *out = pkg.release(); + return nullptr; +} + +// ───────────────────────────────────────────────────────────────────────────── +// Components +// ───────────────────────────────────────────────────────────────────────────── + +ModelPackageStatus* ModelPackage_SetComponentInline(ModelPackage* pkg, + const char* name, + const char* component_json) { + if (!pkg) return NullArg("pkg"); + if (!name) return NullArg("name"); + if (!component_json) return NullArg("component_json"); + + ordered_json body; + if (auto* s = ParseJsonString(component_json, + ("component '" + std::string(name) + "'").c_str(), &body)) return s; + if (auto* s = ExpectObject(body, ("component '" + std::string(name) + "'").c_str())) return s; + + auto opts = mp::PathOptionsFor(pkg); + auto rec = std::make_unique(); + rec->storage = mp::ComponentStorage::kInline; + rec->component_dir = pkg->package_root; + if (auto* s = mp::ParseComponentBody(pkg->package_root, opts, pkg->strict_unknown_fields, + name, body, pkg->package_root, rec.get())) return s; + + EnsureManifestComponentsObject(pkg)[name] = body; + + if (auto* existing = FindComponentRecord(pkg, name)) { + size_t idx = pkg->component_index_by_name[name]; + mp::RecordOrphanComponent(pkg, *pkg->components[idx]); + pkg->components[idx] = std::move(rec); + } else { + pkg->components.push_back(std::move(rec)); + } + RebuildComponentIndex(pkg); + return PostMutate(pkg); +} + +ModelPackageStatus* ModelPackage_SetComponentExternal(ModelPackage* pkg, + const char* name, + const char* path) { + if (!pkg) return NullArg("pkg"); + if (!name) return NullArg("name"); + if (!path) return NullArg("path"); + if (pkg->package_root.empty()) { + return MakeStatus(MODEL_PACKAGE_ERR_STATE, + "SetComponentExternal requires a package_root (use _Open or _Commit " + "with a dest_root first; or rely on _Commit(dest_root) to materialize)."); + } + + auto opts = mp::PathOptionsFor(pkg); + fs::path resolved; + // Allow the file/dir to not exist yet (we'll initialize empty). + if (auto* s = mp::ResolvePath(pkg->package_root, pkg->package_root, path, opts, + /*must_exist=*/false, &resolved)) return s; + std::error_code ec; + fs::path component_dir; + fs::path file_path; + if (fs::exists(resolved, ec) && fs::is_directory(resolved, ec)) { + file_path = resolved / "component.json"; + component_dir = resolved; + } else { + file_path = resolved; + component_dir = resolved.parent_path(); + } + ordered_json body; + if (fs::exists(file_path, ec)) { + std::ifstream f(file_path, std::ios::binary); + std::ostringstream buf; + buf << f.rdbuf(); + if (auto* s = ParseJsonString(buf.str().c_str(), + ("component '" + std::string(name) + "'").c_str(), &body)) return s; + } else { + body = ordered_json::object(); + body["variants"] = ordered_json::object(); + } + if (auto* s = ExpectObject(body, ("component '" + std::string(name) + "'").c_str())) return s; + + auto rec = std::make_unique(); + rec->storage = mp::ComponentStorage::kExternal; + rec->external_path = file_path; + rec->component_dir = component_dir; + if (auto* s = mp::ParseComponentBody(pkg->package_root, opts, pkg->strict_unknown_fields, + name, body, component_dir, rec.get())) return s; + + EnsureManifestComponentsObject(pkg)[name] = std::string(path); + + if (FindComponentRecord(pkg, name)) { + size_t idx = pkg->component_index_by_name[name]; + mp::RecordOrphanComponent(pkg, *pkg->components[idx]); + pkg->components[idx] = std::move(rec); + } else { + pkg->components.push_back(std::move(rec)); + } + RebuildComponentIndex(pkg); + return PostMutate(pkg); +} + +ModelPackageStatus* ModelPackage_RemoveComponent(ModelPackage* pkg, const char* name) { + if (!pkg) return NullArg("pkg"); + if (!name) return NullArg("name"); + auto it = pkg->component_index_by_name.find(name); + if (it == pkg->component_index_by_name.end()) return nullptr; + size_t idx = it->second; + mp::RecordOrphanComponent(pkg, *pkg->components[idx]); + pkg->components.erase(pkg->components.begin() + idx); + auto comps_it = pkg->manifest.find("components"); + if (comps_it != pkg->manifest.end() && comps_it->is_object()) { + comps_it->erase(name); + } + RebuildComponentIndex(pkg); + return PostMutate(pkg); +} + +// ───────────────────────────────────────────────────────────────────────────── +// Variants +// ───────────────────────────────────────────────────────────────────────────── + +ModelPackageStatus* ModelPackage_SetVariant(ModelPackage* pkg, + const char* component_name, + const char* variant_name, + const char* variant_json) { + if (!pkg) return NullArg("pkg"); + if (!component_name) return NullArg("component_name"); + if (!variant_name) return NullArg("variant_name"); + if (!variant_json) return NullArg("variant_json"); + auto* comp = FindComponentRecord(pkg, component_name); + if (!comp) { + return MakeStatus(MODEL_PACKAGE_ERR_NOT_FOUND, + std::string("SetVariant: component '") + component_name + "' not found."); + } + ordered_json body; + if (auto* s = ParseJsonString(variant_json, + ("variant '" + std::string(variant_name) + "'").c_str(), &body)) return s; + + auto vr = std::make_unique(); + auto opts = mp::PathOptionsFor(pkg); + if (auto* s = mp::ParseVariantBody(comp->component_dir, pkg->package_root, opts, + pkg->strict_unknown_fields, + variant_name, body, vr.get())) return s; + + // Update component.body["variants"][variant_name] + if (!comp->body.contains("variants") || !comp->body["variants"].is_object()) { + comp->body["variants"] = ordered_json::object(); + } + comp->body["variants"][variant_name] = body; + // If component is inline, mirror into manifest. + if (comp->storage == mp::ComponentStorage::kInline) { + pkg->manifest["components"][comp->name] = comp->body; + } + // Replace or append. + bool replaced = false; + for (auto& v : comp->variants) { + if (v->name == variant_name) { + mp::RecordOrphanVariantDir(pkg, *v); + v = std::move(vr); + replaced = true; + break; + } + } + if (!replaced) comp->variants.push_back(std::move(vr)); + + // Invalidate cached component JSON. + comp->component_json_cache.reset(); + return PostMutate(pkg); +} + +ModelPackageStatus* ModelPackage_RemoveVariant(ModelPackage* pkg, + const char* component_name, + const char* variant_name) { + if (!pkg) return NullArg("pkg"); + if (!component_name) return NullArg("component_name"); + if (!variant_name) return NullArg("variant_name"); + auto* comp = FindComponentRecord(pkg, component_name); + if (!comp) return nullptr; + auto pred = [&](const std::unique_ptr& v) { + if (v->name == variant_name) { + mp::RecordOrphanVariantDir(pkg, *v); + return true; + } + return false; + }; + comp->variants.erase(std::remove_if(comp->variants.begin(), comp->variants.end(), pred), + comp->variants.end()); + if (comp->body.contains("variants") && comp->body["variants"].is_object()) { + comp->body["variants"].erase(variant_name); + } + if (comp->storage == mp::ComponentStorage::kInline) { + pkg->manifest["components"][comp->name] = comp->body; + } + comp->component_json_cache.reset(); + return PostMutate(pkg); +} + +// ───────────────────────────────────────────────────────────────────────────── +// Variant executor_info +// ───────────────────────────────────────────────────────────────────────────── + +namespace { + +ModelPackageStatus* ReparseVariantInPlace(ModelPackage* pkg, + mp::ComponentRecord* comp, + mp::VariantRecord* var) { + auto opts = mp::PathOptionsFor(pkg); + auto rebuilt = std::make_unique(); + if (auto* s = mp::ParseVariantBody(comp->component_dir, pkg->package_root, opts, + pkg->strict_unknown_fields, + var->name, var->body, rebuilt.get())) return s; + *var = std::move(*rebuilt); + return nullptr; +} + +ModelPackageStatus* MutateExecutorInfo(ModelPackage* pkg, + const char* component, + const char* variant, + const char* namespace_, + const ordered_json* new_value /* null = remove */) { + if (!pkg) return NullArg("pkg"); + if (!component) return NullArg("component"); + if (!variant) return NullArg("variant"); + if (!namespace_) return NullArg("namespace"); + auto* comp = FindComponentRecord(pkg, component); + if (!comp) { + return MakeStatus(MODEL_PACKAGE_ERR_NOT_FOUND, + std::string("component '") + component + "' not found."); + } + auto* var = FindVariantRecord(comp, variant); + if (!var) { + return MakeStatus(MODEL_PACKAGE_ERR_NOT_FOUND, + std::string("variant '") + variant + "' not found in component '" + + component + "'."); + } + if (!var->body.contains("executor_info") || !var->body["executor_info"].is_object()) { + if (!new_value) return nullptr; // remove on absent -> nothing to do + var->body["executor_info"] = ordered_json::object(); + } + if (new_value) { + var->body["executor_info"][namespace_] = *new_value; + } else { + var->body["executor_info"].erase(namespace_); + if (var->body["executor_info"].empty()) { + var->body.erase("executor_info"); + } + } + comp->body["variants"][var->name] = var->body; + if (comp->storage == mp::ComponentStorage::kInline) { + pkg->manifest["components"][comp->name] = comp->body; + } + if (auto* s = ReparseVariantInPlace(pkg, comp, var)) return s; + comp->component_json_cache.reset(); + return PostMutate(pkg, /*refresh_assets=*/false); +} + +} // namespace + +ModelPackageStatus* ModelPackage_SetVariantExecutorInfoInline(ModelPackage* pkg, + const char* component, + const char* variant, + const char* namespace_, + const char* info_json) { + if (!info_json) return NullArg("info_json"); + ordered_json body; + if (auto* s = ParseJsonString(info_json, "executor_info", &body)) return s; + if (auto* s = ExpectObject(body, "executor_info inline value")) return s; + return MutateExecutorInfo(pkg, component, variant, namespace_, &body); +} + +ModelPackageStatus* ModelPackage_SetVariantExecutorInfoExternal(ModelPackage* pkg, + const char* component, + const char* variant, + const char* namespace_, + const char* path) { + if (!path) return NullArg("path"); + ordered_json body = std::string(path); + return MutateExecutorInfo(pkg, component, variant, namespace_, &body); +} + +ModelPackageStatus* ModelPackage_RemoveVariantExecutorInfo(ModelPackage* pkg, + const char* component, + const char* variant, + const char* namespace_) { + return MutateExecutorInfo(pkg, component, variant, namespace_, nullptr); +} + +// ───────────────────────────────────────────────────────────────────────────── +// Shared assets +// ───────────────────────────────────────────────────────────────────────────── + +ModelPackageStatus* ModelPackage_AddSharedAsset(ModelPackage* pkg, + const char* source_dir, + const char* expected_uri_or_null, + bool copy_in, + const char** out_uri) { + if (!pkg) return NullArg("pkg"); + if (!source_dir) return NullArg("source_dir"); + if (!out_uri) return NullArg("out_uri"); + *out_uri = nullptr; + + if (!copy_in && pkg->layout == "portable") { + return MakeStatus(MODEL_PACKAGE_ERR_STATE, + "AddSharedAsset: copy_in=false rejected in portable layout (the " + "path would point outside )."); + } + + std::string computed_uri; + if (auto* s = mp::ComputeDirectoryAssetUri(fs::path(source_dir), &computed_uri)) return s; + if (expected_uri_or_null) { + if (computed_uri != expected_uri_or_null) { + return MakeStatus(MODEL_PACKAGE_ERR_STATE, + std::string("AddSharedAsset: hash mismatch — computed ") + + computed_uri + ", expected " + expected_uri_or_null + "."); + } + } + + if (!pkg->manifest.contains("shared_assets") || !pkg->manifest["shared_assets"].is_object()) { + pkg->manifest["shared_assets"] = ordered_json::object(); + } + if (copy_in) { + // No manifest entry needed — the asset will be materialized at the default + // convention path on commit. LoadSharedAssets surfaces the staged source + // immediately so the URI shows up in ModelPackage_Info() before commit. + pkg->pending_shared_asset_copies[computed_uri] = fs::path(source_dir); + } else { + pkg->manifest["shared_assets"][computed_uri] = std::string(source_dir); + } + + if (auto* s = PostMutate(pkg)) return s; + + // Look up the record and return its URI. After PostMutate, the URI is + // always present in shared_assets_index_by_uri (either via the override + // path or via the pending-copy tier of LoadSharedAssets). + auto sit = pkg->shared_asset_index_by_uri.find(computed_uri); + if (sit == pkg->shared_asset_index_by_uri.end()) { + return MakeStatus(MODEL_PACKAGE_ERR_STATE, + std::string("AddSharedAsset: failed to register URI ") + computed_uri); + } + *out_uri = pkg->shared_assets[sit->second]->uri_cache.c_str(); + return nullptr; +} + +ModelPackageStatus* ModelPackage_RemoveSharedAsset(ModelPackage* pkg, const char* uri) { + if (!pkg) return NullArg("pkg"); + if (!uri) return NullArg("uri"); + std::string uri_str(uri); + if (pkg->manifest.contains("shared_assets") && pkg->manifest["shared_assets"].is_object()) { + pkg->manifest["shared_assets"].erase(uri_str); + if (pkg->manifest["shared_assets"].empty()) { + pkg->manifest.erase("shared_assets"); + } + } + pkg->pending_shared_asset_copies.erase(uri_str); + // Physically remove the on-disk directory at the default convention. If it + // stays on disk, the next RefreshSharedAssets would auto-discover it again + // and the removal would be a no-op. We only touch paths that live inside + // package_root. + if (!pkg->package_root.empty()) { + std::string dir_name = mp::DefaultSharedAssetDirName(uri_str); + if (!dir_name.empty()) { + std::filesystem::path on_disk = pkg->package_root / "shared_assets" / dir_name; + if (mp::IsInsidePackageRoot(pkg, on_disk)) { + std::error_code ec; + std::filesystem::remove_all(on_disk, ec); + } + } + } + return PostMutate(pkg); +} + +// ───────────────────────────────────────────────────────────────────────────── +// Package metadata +// ───────────────────────────────────────────────────────────────────────────── + +namespace { + +void SetOrClearString(ordered_json* obj, const char* key, const char* value) { + if (value == nullptr) return; // leave untouched + if (value[0] == '\0') { + obj->erase(key); + } else { + (*obj)[key] = std::string(value); + } +} + +} // namespace + +ModelPackageStatus* ModelPackage_SetMetadata(ModelPackage* pkg, + const char* name_or_null, + const char* version_or_null, + const char* description_or_null) { + if (!pkg) return NullArg("pkg"); + SetOrClearString(&pkg->manifest, "package_name", name_or_null); + SetOrClearString(&pkg->manifest, "package_version", version_or_null); + SetOrClearString(&pkg->manifest, "description", description_or_null); + return PostMutate(pkg, /*refresh_assets=*/false); +} + +ModelPackageStatus* ModelPackage_SetLayout(ModelPackage* pkg, const char* layout) { + if (!pkg) return NullArg("pkg"); + if (!layout) return NullArg("layout"); + std::string l(layout); + if (l != "portable" && l != "installed") { + return MakeStatus(MODEL_PACKAGE_ERR_SCHEMA, + "SetLayout: layout must be 'portable' or 'installed'."); + } + pkg->manifest["layout"] = l; + pkg->layout = l; + return PostMutate(pkg, /*refresh_assets=*/false); +} + +ModelPackageStatus* ModelPackage_SetAdditionalMetadataJson(ModelPackage* pkg, + const char* scope, + const char* component_or_null, + const char* variant_or_null, + const char* json_or_null) { + if (!pkg) return NullArg("pkg"); + if (!scope) return NullArg("scope"); + std::string s(scope); + ordered_json* target = nullptr; + mp::ComponentRecord* comp = nullptr; + mp::VariantRecord* var = nullptr; + if (s == "manifest") { + if (component_or_null || variant_or_null) { + return MakeStatus(MODEL_PACKAGE_ERR_INVALID_ARG, + "SetAdditionalMetadataJson: 'manifest' scope takes no component/variant."); + } + target = &pkg->manifest; + } else if (s == "component") { + if (!component_or_null) return NullArg("component"); + if (variant_or_null) { + return MakeStatus(MODEL_PACKAGE_ERR_INVALID_ARG, + "SetAdditionalMetadataJson: 'component' scope takes no variant."); + } + comp = FindComponentRecord(pkg, component_or_null); + if (!comp) { + return MakeStatus(MODEL_PACKAGE_ERR_NOT_FOUND, + std::string("component '") + component_or_null + "' not found."); + } + target = &comp->body; + } else if (s == "variant") { + if (!component_or_null) return NullArg("component"); + if (!variant_or_null) return NullArg("variant"); + comp = FindComponentRecord(pkg, component_or_null); + if (!comp) { + return MakeStatus(MODEL_PACKAGE_ERR_NOT_FOUND, + std::string("component '") + component_or_null + "' not found."); + } + var = FindVariantRecord(comp, variant_or_null); + if (!var) { + return MakeStatus(MODEL_PACKAGE_ERR_NOT_FOUND, + std::string("variant '") + variant_or_null + "' not found."); + } + target = &var->body; + } else { + return MakeStatus(MODEL_PACKAGE_ERR_INVALID_ARG, + "SetAdditionalMetadataJson: scope must be 'manifest', 'component', or 'variant'."); + } + if (json_or_null == nullptr) { + target->erase("additional_metadata"); + } else { + ordered_json body; + if (auto* st = ParseJsonString(json_or_null, "additional_metadata", &body)) return st; + (*target)["additional_metadata"] = body; + } + if (comp && comp->storage == mp::ComponentStorage::kInline) { + pkg->manifest["components"][comp->name] = comp->body; + } + if (comp) comp->component_json_cache.reset(); + if (var) var->additional_metadata_cache.reset(); + if (comp) comp->additional_metadata_cache.reset(); + return PostMutate(pkg, /*refresh_assets=*/false); +} + +} // extern "C" diff --git a/model_package/src/commit_prune_validate.cc b/model_package/src/commit_prune_validate.cc new file mode 100644 index 0000000000000..c603dd6bcbf00 --- /dev/null +++ b/model_package/src/commit_prune_validate.cc @@ -0,0 +1,769 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +/// \file commit_prune_validate.cc +/// \brief Commit, prune, and validate implementation. + +#include "model_package.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#ifndef _WIN32 +#include +#include +#include +#endif + +#include "asset_hasher.h" +#include "manifest_parser.h" +#include "model_package_impl.h" +#include "path_resolver.h" +#include "status_impl.h" + +namespace fs = std::filesystem; +namespace mp = model_package; +using model_package::MakeStatus; +using nlohmann::ordered_json; + +namespace { + +ModelPackageStatus* NullArg(const char* name) { + return MakeStatus(MODEL_PACKAGE_ERR_INVALID_ARG, + std::string("model_package: '") + name + "' must not be null."); +} + +// ───────────────────────────────────────────────────────────────────────────── +// fsync / random helpers (POSIX). Windows would substitute FlushFileBuffers + +// BCryptGenRandom; deferred to a follow-up. +// ───────────────────────────────────────────────────────────────────────────── + +std::string RandomSuffix() { + std::random_device rd; + uint64_t hi = (uint64_t(rd()) << 32) | rd(); + char buf[17]; + std::snprintf(buf, sizeof(buf), "%016llx", static_cast(hi)); + return buf; +} + +ModelPackageStatus* FsyncPath(const fs::path& p, bool is_dir) { +#ifdef _WIN32 + (void)p; + (void)is_dir; + return nullptr; +#else + int flags = is_dir ? (O_RDONLY | O_DIRECTORY) : O_RDONLY; + int fd = ::open(p.c_str(), flags); + if (fd < 0) { + // Best-effort: missing fsync targets are not fatal on tmpfs etc. + return nullptr; + } + if (::fsync(fd) != 0) { + int err = errno; + ::close(fd); + return MakeStatus(MODEL_PACKAGE_ERR_IO, + std::string("fsync '") + p.string() + "' failed: " + std::strerror(err)); + } + ::close(fd); + return nullptr; +#endif +} + +ModelPackageStatus* WriteFileAtomic(const fs::path& final_path, const std::string& bytes) { + fs::path tmp = final_path; + tmp += ".tmp." + RandomSuffix(); + { + std::ofstream f(tmp, std::ios::binary | std::ios::trunc); + if (!f) { + return MakeStatus(MODEL_PACKAGE_ERR_IO, + "Cannot open '" + tmp.string() + "' for writing."); + } + f.write(bytes.data(), static_cast(bytes.size())); + if (!f) { + return MakeStatus(MODEL_PACKAGE_ERR_IO, + "Write to '" + tmp.string() + "' failed."); + } + } + if (auto* s = FsyncPath(tmp, /*is_dir=*/false)) return s; + std::error_code ec; + fs::rename(tmp, final_path, ec); + if (ec) { + fs::remove(tmp, ec); + return MakeStatus(MODEL_PACKAGE_ERR_IO, + "Rename '" + tmp.string() + "' -> '" + final_path.string() + + "' failed: " + ec.message()); + } + if (auto* s = FsyncPath(final_path.parent_path(), /*is_dir=*/true)) return s; + return nullptr; +} + +ModelPackageStatus* CopyTreeNoFollow(const fs::path& src, const fs::path& dst) { + // Recursively copy `src` into `dst`. Refuses to follow symlinks (consistent + // with the directory hash semantics) so the on-disk bytes match the URI we + // already computed. + std::error_code ec; + fs::create_directories(dst, ec); + if (ec) return MakeStatus(MODEL_PACKAGE_ERR_IO, + "mkdir '" + dst.string() + "': " + ec.message()); + for (fs::recursive_directory_iterator it(src, fs::directory_options::none, ec), end; + it != end; it.increment(ec)) { + if (ec) return MakeStatus(MODEL_PACKAGE_ERR_IO, + "iterate '" + src.string() + "': " + ec.message()); + const auto& entry = *it; + fs::path rel = fs::relative(entry.path(), src, ec); + fs::path target = dst / rel; + if (entry.is_symlink()) { + return MakeStatus(MODEL_PACKAGE_ERR_SCHEMA, + "shared asset source contains a symlink: '" + entry.path().string() + "'."); + } + if (entry.is_directory()) { + fs::create_directories(target, ec); + if (ec) return MakeStatus(MODEL_PACKAGE_ERR_IO, + "mkdir '" + target.string() + "': " + ec.message()); + } else if (entry.is_regular_file()) { + fs::create_directories(target.parent_path(), ec); + fs::copy_file(entry.path(), target, fs::copy_options::overwrite_existing, ec); + if (ec) return MakeStatus(MODEL_PACKAGE_ERR_IO, + "copy '" + entry.path().string() + "' -> '" + + target.string() + "': " + ec.message()); + if (auto* s = FsyncPath(target, /*is_dir=*/false)) return s; + } else { + return MakeStatus(MODEL_PACKAGE_ERR_SCHEMA, + "unsupported file kind in shared asset source: '" + + entry.path().string() + "'."); + } + } + if (auto* s = FsyncPath(dst, /*is_dir=*/true)) return s; + return nullptr; +} + +ModelPackageStatus* CheckPortableConfinement(const fs::path& root, + const fs::path& candidate, + const std::string& where) { + std::error_code ec; + fs::path c = candidate.lexically_normal(); + fs::path r = root.lexically_normal(); + if (c.is_absolute()) { + // Confirm c is under r. + auto rel = fs::relative(c, r, ec); + // An empty relative path, or one whose first component is "..", escapes the root. + // (Checking only the first character would wrongly reject in-root dot-prefixed names + // such as ".hidden/component.json".) + if (ec || rel.empty() || rel.begin()->string() == "..") { + return MakeStatus(MODEL_PACKAGE_ERR_PATH_CONFINEMENT, + where + ": absolute path '" + c.string() + + "' escapes package_root '" + r.string() + "' (portable layout)."); + } + } else { + // Relative: a leading ".." escapes. + auto first = c.begin(); + if (first != c.end() && first->string() == "..") { + return MakeStatus(MODEL_PACKAGE_ERR_PATH_CONFINEMENT, + where + ": relative path '" + c.string() + + "' escapes package_root (portable layout)."); + } + } + return nullptr; +} + +// ───────────────────────────────────────────────────────────────────────────── +// Manifest serialization +// ───────────────────────────────────────────────────────────────────────────── + +std::string SerializeManifestForCommit(const ModelPackage* pkg) { + // Use the live in-memory manifest, but for external components, the + // ComponentRecord::body may have diverged from the string path. The manifest + // entry stays as the string in that case; the body is serialized separately + // into the external file. + return pkg->manifest.dump(2) + "\n"; +} + +ordered_json SerializeComponentBody(const mp::ComponentRecord* comp) { + return comp->body; +} + +// ───────────────────────────────────────────────────────────────────────────── +// In-place commit (PRESERVE / DENSE) +// ───────────────────────────────────────────────────────────────────────────── + +ModelPackageStatus* CheckDenseConstraints(ModelPackage* pkg) { + // Reject external executor_info in dense mode (dense flattens everything, + // but the in-memory model never loads external executor_info bodies, so we + // can't inline them surgically. ERR_STATE so the caller's intent is clear.) + for (const auto& comp : pkg->components) { + auto vit = comp->body.find("variants"); + if (vit == comp->body.end() || !vit->is_object()) continue; + for (auto v = vit->begin(); v != vit->end(); ++v) { + auto ei = v->find("executor_info"); + if (ei == v->end() || !ei->is_object()) continue; + for (auto e = ei->begin(); e != ei->end(); ++e) { + if (e->is_string()) { + return MakeStatus(MODEL_PACKAGE_ERR_STATE, + "WRITE_DENSE: component '" + comp->name + "' variant '" + + v.key() + "' has external executor_info '" + e.key() + + "' (string path). Convert to inline via " + "SetVariantExecutorInfoInline before dense commit."); + } + } + } + } + return nullptr; +} + +ModelPackageStatus* CommitSharedAssetsCopyIn(ModelPackage* pkg, const fs::path& root) { + if (pkg->pending_shared_asset_copies.empty()) return nullptr; + fs::path assets_root = root / "shared_assets"; + std::error_code ec; + fs::create_directories(assets_root, ec); + for (const auto& [uri, src] : pkg->pending_shared_asset_copies) { + std::string dir_name = mp::DefaultSharedAssetDirName(uri); + fs::path final_dir = assets_root / dir_name; + if (fs::exists(final_dir, ec)) continue; // already materialized — trust it. + fs::path stage_dir = assets_root / (dir_name + ".tmp." + RandomSuffix()); + if (auto* s = CopyTreeNoFollow(src, stage_dir)) { + fs::remove_all(stage_dir, ec); + return s; + } + // Re-hash staging to verify TOCTOU did not strike. + std::string verify_uri; + if (auto* s = mp::ComputeDirectoryAssetUri(stage_dir, &verify_uri)) { + fs::remove_all(stage_dir, ec); + return s; + } + if (verify_uri != uri) { + fs::remove_all(stage_dir, ec); + return MakeStatus(MODEL_PACKAGE_ERR_STATE, + "Shared asset source mutated during commit: expected " + + uri + ", staged " + verify_uri + "."); + } + fs::rename(stage_dir, final_dir, ec); + if (ec) { + fs::remove_all(stage_dir, ec); + return MakeStatus(MODEL_PACKAGE_ERR_IO, + "Rename shared asset dir '" + stage_dir.string() + "' -> '" + + final_dir.string() + "' failed: " + ec.message()); + } + if (auto* s = FsyncPath(assets_root, /*is_dir=*/true)) return s; + } + return nullptr; +} + +ModelPackageStatus* CommitExternalComponents(ModelPackage* pkg) { + // Write each external component's current in-memory body to its disk file. + // These are library-owned; for in-place PRESERVE commit we re-emit them + // every time (cheaper than tracking dirtiness). External executor_info + // files are opaque and intentionally left untouched. + for (const auto& comp : pkg->components) { + if (comp->storage != mp::ComponentStorage::kExternal) continue; + fs::path path = comp->external_path; + std::error_code ec; + fs::create_directories(path.parent_path(), ec); + std::string text = SerializeComponentBody(comp.get()).dump(2) + "\n"; + if (auto* s = WriteFileAtomic(path, text)) return s; + } + return nullptr; +} + +ModelPackageStatus* CommitInPlace(ModelPackage* pkg, ModelPackageWriteMode mode) { + if (pkg->package_root.empty()) { + return MakeStatus(MODEL_PACKAGE_ERR_STATE, + "Commit: package has no package_root. Use dest_root variant."); + } + std::error_code ec; + if (!fs::is_directory(pkg->package_root, ec)) { + fs::create_directories(pkg->package_root, ec); + if (ec) { + return MakeStatus(MODEL_PACKAGE_ERR_IO, + "Cannot create package_root '" + pkg->package_root.string() + + "': " + ec.message()); + } + } + + // Portable confinement pre-flight for external paths. + if (pkg->layout == "portable") { + for (const auto& comp : pkg->components) { + if (comp->storage == mp::ComponentStorage::kExternal) { + if (auto* s = CheckPortableConfinement(pkg->package_root, comp->external_path, + "component '" + comp->name + "'")) return s; + } + } + } + + // Dense mode: flatten external components into manifest before writing. + if (mode == MODEL_PACKAGE_WRITE_DENSE) { + if (auto* s = CheckDenseConstraints(pkg)) return s; + for (auto& comp : pkg->components) { + if (comp->storage == mp::ComponentStorage::kExternal) { + pkg->manifest["components"][comp->name] = comp->body; + // After commit, this becomes inline. + comp->storage = mp::ComponentStorage::kInline; + comp->external_path.clear(); + comp->component_dir = pkg->package_root; + } + } + } + + if (auto* s = CommitSharedAssetsCopyIn(pkg, pkg->package_root)) return s; + if (mode != MODEL_PACKAGE_WRITE_DENSE) { + if (auto* s = CommitExternalComponents(pkg)) return s; + } + + // Final manifest write. + fs::path manifest_path = pkg->package_root / "manifest.json"; + if (auto* s = WriteFileAtomic(manifest_path, SerializeManifestForCommit(pkg))) return s; + + pkg->pending_shared_asset_copies.clear(); + + // Re-derive shared assets + info view to pick up the materialized assets. + if (auto* s = mp::RefreshSharedAssets(pkg, mp::PathOptionsFor(pkg))) return s; + if (auto* s = mp::RefreshPackageMetadata(pkg)) return s; + mp::DropViewCache(pkg); + return nullptr; +} + +// ───────────────────────────────────────────────────────────────────────────── +// dest_root commit ("save as"): write to dest_root, then re-parse & swap. +// ───────────────────────────────────────────────────────────────────────────── + +ModelPackageStatus* CommitToDestRoot(ModelPackage* pkg, + const fs::path& dest_root, + ModelPackageWriteMode mode) { + std::error_code ec; + if (fs::exists(dest_root, ec)) { + if (!fs::is_directory(dest_root, ec)) { + return MakeStatus(MODEL_PACKAGE_ERR_STATE, + "Commit dest_root '" + dest_root.string() + "' exists and is not a directory."); + } + if (!fs::is_empty(dest_root, ec)) { + return MakeStatus(MODEL_PACKAGE_ERR_STATE, + "Commit dest_root '" + dest_root.string() + "' is not empty."); + } + } else { + fs::create_directories(dest_root, ec); + if (ec) { + return MakeStatus(MODEL_PACKAGE_ERR_IO, + "Cannot create dest_root '" + dest_root.string() + "': " + ec.message()); + } + } + + // Build a snapshot manifest mirroring `pkg->manifest`, then handle assets. + ordered_json manifest = pkg->manifest; + + // Dense mode constraints up-front. + if (mode == MODEL_PACKAGE_WRITE_DENSE) { + if (auto* s = CheckDenseConstraints(pkg)) return s; + for (const auto& comp : pkg->components) { + if (comp->storage == mp::ComponentStorage::kExternal) { + manifest["components"][comp->name] = comp->body; + } + } + } + + // Copy all shared assets into dest_root. Any manifest override entries are + // re-mapped to the default convention path under dest_root. + fs::path assets_root = dest_root / "shared_assets"; + // Gather source dirs for every URI we know about. + // 1. URIs already on disk (under current package_root) and not in pending: copy from there. + // 2. Pending copy_in sources: copy from staged source. + // 3. Manifest override entries: copy from the override path. + std::vector> to_copy; + for (const auto& rec : pkg->shared_assets) { + auto pit = pkg->pending_shared_asset_copies.find(rec->uri); + if (pit != pkg->pending_shared_asset_copies.end()) { + to_copy.emplace_back(rec->uri, pit->second); + } else { + to_copy.emplace_back(rec->uri, rec->resolved_path); + } + } + // Pending copies without a SharedAssetRecord shouldn't happen now that + // LoadSharedAssets surfaces pending copies, but stay defensive. + for (const auto& [uri, src] : pkg->pending_shared_asset_copies) { + if (pkg->shared_asset_index_by_uri.find(uri) == pkg->shared_asset_index_by_uri.end()) { + to_copy.emplace_back(uri, src); + } + } + // Only materialize shared_assets/ when something will actually land in it. + if (!to_copy.empty()) { + fs::create_directories(assets_root, ec); + } + + for (const auto& [uri, src] : to_copy) { + if (!fs::is_directory(src, ec)) { + return MakeStatus(MODEL_PACKAGE_ERR_NOT_FOUND, + "Commit dest_root: shared asset source '" + src.string() + + "' for " + uri + " is not a directory."); + } + std::string dir_name = mp::DefaultSharedAssetDirName(uri); + fs::path final_dir = assets_root / dir_name; + fs::path stage_dir = assets_root / (dir_name + ".tmp." + RandomSuffix()); + if (auto* s = CopyTreeNoFollow(src, stage_dir)) { + fs::remove_all(stage_dir, ec); + return s; + } + std::string verify_uri; + if (auto* s = mp::ComputeDirectoryAssetUri(stage_dir, &verify_uri)) { + fs::remove_all(stage_dir, ec); + return s; + } + if (verify_uri != uri) { + fs::remove_all(stage_dir, ec); + return MakeStatus(MODEL_PACKAGE_ERR_STATE, + "Shared asset hash mismatch during dest_root commit: expected " + + uri + ", staged " + verify_uri); + } + fs::rename(stage_dir, final_dir, ec); + if (ec) { + fs::remove_all(stage_dir, ec); + return MakeStatus(MODEL_PACKAGE_ERR_IO, "Rename failed: " + ec.message()); + } + } + // All assets now live at the default convention path; drop overrides. + manifest.erase("shared_assets"); + + // External components (PRESERVE mode): re-emit under dest_root using the same + // path string from the manifest. We treat the manifest string as relative to + // dest_root for portable mode; absolute paths are kept as-is iff the layout + // is installed. + if (mode == MODEL_PACKAGE_WRITE_PRESERVE) { + auto comps_it = manifest.find("components"); + if (comps_it != manifest.end() && comps_it->is_object()) { + for (auto e = comps_it->begin(); e != comps_it->end(); ++e) { + if (!e->is_string()) continue; + fs::path p(e->get()); + fs::path target; + if (p.is_absolute()) { + if (pkg->layout == "portable") { + return MakeStatus(MODEL_PACKAGE_ERR_PATH_CONFINEMENT, + "dest_root commit (portable): component '" + e.key() + + "' has absolute path '" + p.string() + "'."); + } + target = p; + } else { + target = dest_root / p; + std::error_code ec2; + fs::path normalized = target.lexically_normal(); + if (normalized.string().find(dest_root.lexically_normal().string()) != 0) { + return MakeStatus(MODEL_PACKAGE_ERR_PATH_CONFINEMENT, + "dest_root commit (portable): component '" + e.key() + + "' relative path '" + p.string() + "' escapes dest_root."); + } + target = normalized; + } + // Find the corresponding component body to write. + std::string ext_body; + for (const auto& comp : pkg->components) { + if (comp->name == e.key()) { + ext_body = comp->body.dump(2) + "\n"; + break; + } + } + std::error_code ec_md; + fs::create_directories(target.parent_path(), ec_md); + if (auto* s = WriteFileAtomic(target, ext_body)) return s; + } + } + } + + fs::path manifest_path = dest_root / "manifest.json"; + if (auto* s = WriteFileAtomic(manifest_path, manifest.dump(2) + "\n")) return s; + + // Re-parse the newly written package into a fresh state and swap in. + ModelPackageOpenOptions opts{}; + opts.allow_external_paths = pkg->allow_external_paths; + opts.follow_symlinks = pkg->follow_symlinks; + opts.strict_unknown_fields = pkg->strict_unknown_fields; + ModelPackage fresh; + if (auto* s = mp::ParsePackage(dest_root, opts, &fresh)) { + return s; + } + // Tear down the existing view cache for the old package, then swap. + mp::DropViewCache(pkg); + // Field-by-field swap (the opaque struct is non-trivial; std::swap of the + // struct works because all members are move/swap-friendly). + std::swap(pkg->package_root, fresh.package_root); + std::swap(pkg->manifest, fresh.manifest); + std::swap(pkg->layout, fresh.layout); + std::swap(pkg->components, fresh.components); + std::swap(pkg->shared_assets, fresh.shared_assets); + std::swap(pkg->component_index_by_name, fresh.component_index_by_name); + std::swap(pkg->shared_asset_index_by_uri, fresh.shared_asset_index_by_uri); + std::swap(pkg->package_name_cache, fresh.package_name_cache); + std::swap(pkg->package_version_cache, fresh.package_version_cache); + std::swap(pkg->description_cache, fresh.description_cache); + std::swap(pkg->layout_cache, fresh.layout_cache); + std::swap(pkg->additional_metadata_cache, fresh.additional_metadata_cache); + std::swap(pkg->schema_version_major, fresh.schema_version_major); + std::swap(pkg->schema_version_minor, fresh.schema_version_minor); + pkg->pending_shared_asset_copies.clear(); + pkg->info_cache.reset(); + + if (auto* s = mp::RefreshPackageMetadata(pkg)) return s; + return nullptr; +} + +// ───────────────────────────────────────────────────────────────────────────── +// Prune +// ───────────────────────────────────────────────────────────────────────────── + +constexpr std::chrono::seconds kPruneGrace{60}; + +bool IsTmpName(const fs::path& p) { + std::string name = p.filename().string(); + return name.find(".tmp.") != std::string::npos; +} + +bool IsOldEnough(const fs::path& p) { + std::error_code ec; + auto last = fs::last_write_time(p, ec); + if (ec) return false; + auto now = decltype(last)::clock::now(); + return (now - last) >= kPruneGrace; +} + +bool IsAncestorOrEqual(const fs::path& ancestor, const fs::path& descendant) { + // ancestor == descendant, or descendant lives under ancestor (boundary aware). + auto a = ancestor.lexically_normal().generic_string(); + auto d = descendant.lexically_normal().generic_string(); + if (d.size() < a.size()) return false; + if (d.compare(0, a.size(), a) != 0) return false; + return d.size() == a.size() || d[a.size()] == '/'; +} + +std::vector CollectLiveDirs(const ModelPackage* pkg) { + std::vector out; + for (const auto& c : pkg->components) { + if (c->storage == mp::ComponentStorage::kExternal) { + out.push_back(c->component_dir); + } + for (const auto& v : c->variants) { + if (v->resolved_directory.has_value()) { + out.push_back(*v->resolved_directory); + } + } + } + return out; +} + +// Drop entries we've handled (removed, or unsafe to touch). Entries that +// reference live state stay for a future Prune call. Tracked orphans don't +// wait on the kPruneGrace window: they were recorded by an in-session +// mutation, so there's no concurrent writer to protect against. The grace +// window is still applied to the shared_assets sweep below, which discovers +// candidates fresh from disk. +void SweepOrphanDirs(ModelPackage* pkg, + std::vector* pending, + const std::vector& live_dirs) { + pending->erase(std::remove_if(pending->begin(), pending->end(), [&](const fs::path& p) { + if (!mp::IsInsidePackageRoot(pkg, p)) return true; // outside our scope + std::error_code ec; + if (!fs::exists(p, ec)) return true; + // Skip if any live dir IS p or lives under it; deleting would damage live state. + for (const auto& live : live_dirs) { + if (IsAncestorOrEqual(p, live)) return false; + } + fs::remove_all(p, ec); + return true; + }), + pending->end()); +} + +} // namespace + +namespace model_package { + +bool IsInsidePackageRoot(const ModelPackage* pkg, const fs::path& p) { + if (pkg->package_root.empty()) return false; + return IsAncestorOrEqual(pkg->package_root, p); +} + +void RecordOrphanVariantDir(ModelPackage* pkg, const VariantRecord& v) { + if (!v.resolved_directory.has_value()) return; + if (!IsInsidePackageRoot(pkg, *v.resolved_directory)) return; + pkg->pending_orphan_variant_dirs.push_back(*v.resolved_directory); +} + +void RecordOrphanComponent(ModelPackage* pkg, const ComponentRecord& c) { + for (const auto& v : c.variants) { + RecordOrphanVariantDir(pkg, *v); + } + if (c.storage == ComponentStorage::kExternal && + IsInsidePackageRoot(pkg, c.component_dir)) { + pkg->pending_orphan_component_dirs.push_back(c.component_dir); + } +} + +} // namespace model_package + +extern "C" { + +ModelPackageStatus* ModelPackage_Commit(ModelPackage* pkg, + const char* dest_root_or_null, + ModelPackageWriteMode mode) { + if (!pkg) return NullArg("pkg"); + if (dest_root_or_null) { + return CommitToDestRoot(pkg, fs::path(dest_root_or_null), mode); + } + return CommitInPlace(pkg, mode); +} + +ModelPackageStatus* ModelPackage_Prune(ModelPackage* pkg) { + if (!pkg) return NullArg("pkg"); + if (pkg->package_root.empty()) return nullptr; + + // Shared assets are NEVER auto-pruned. The library cannot prove an asset is + // unused without parsing every consumer's executor_info payload, and a + // mistaken delete is worse than disk bloat for content-addressed dirs that + // dedupe naturally. Callers reclaim shared assets via explicit + // ModelPackage_RemoveSharedAsset(uri) (which still requires consumer-aware + // knowledge of what's reachable). + // + // Stale `.tmp.` staging dirs from interrupted commits are reclaimed + // here after a grace window: they belong to this library's own staging + // protocol and aren't user data. + std::error_code ec; + fs::path assets_root = pkg->package_root / "shared_assets"; + if (fs::is_directory(assets_root, ec)) { + for (const auto& entry : fs::directory_iterator(assets_root, ec)) { + if (ec) break; + if (!entry.is_directory()) continue; + if (!IsTmpName(entry.path())) continue; + if (!IsOldEnough(entry.path())) continue; + fs::remove_all(entry.path(), ec); + } + } + + // Tracked-orphan sweep: components before variants so a component_dir + // removal reclaims its child variant dirs in one shot. + std::vector live_dirs = CollectLiveDirs(pkg); + SweepOrphanDirs(pkg, &pkg->pending_orphan_component_dirs, live_dirs); + SweepOrphanDirs(pkg, &pkg->pending_orphan_variant_dirs, live_dirs); + + return nullptr; +} + +// ───────────────────────────────────────────────────────────────────────────── +// Validate +// ───────────────────────────────────────────────────────────────────────────── + +namespace { + +void AddFinding(ordered_json* arr, const std::string& code, const std::string& msg) { + ordered_json e = ordered_json::object(); + e["code"] = code; + e["message"] = msg; + arr->push_back(e); +} + +} // namespace + +ModelPackageStatus* ModelPackage_Validate(ModelPackage* pkg, int flags, + const char** out_report_json) { + if (!pkg) return NullArg("pkg"); + if (!out_report_json) return NullArg("out_report_json"); + *out_report_json = nullptr; + ordered_json report = ordered_json::object(); + report["errors"] = ordered_json::array(); + report["warnings"] = ordered_json::array(); + ordered_json* errors = &report["errors"]; + ordered_json* warnings = &report["warnings"]; + + std::error_code ec; + + // SCHEMA: re-validate the in-memory manifest by serializing then re-parsing + // into a scratch ModelPackage with strict mode. Validates schema for both + // committed and uncommitted state. + if (flags & MODEL_PACKAGE_VALIDATE_SCHEMA) { + // Re-run each component/variant through the parser to confirm shape. + for (const auto& comp : pkg->components) { + mp::ComponentRecord scratch; + auto opts = mp::PathOptionsFor(pkg); + if (auto* s = mp::ParseComponentBody(pkg->package_root, opts, + /*strict=*/true, + comp->name, comp->body, + comp->component_dir, &scratch)) { + AddFinding(errors, "SCHEMA", std::string("component '") + comp->name + "': " + ModelPackageStatus_Message(s)); + ModelPackageStatus_Release(s); + } + } + } + + // PATHS: each external component's path on disk; each shared-asset resolved_path exists. + if (flags & MODEL_PACKAGE_VALIDATE_PATHS) { + for (const auto& comp : pkg->components) { + if (comp->storage == mp::ComponentStorage::kExternal) { + if (!fs::exists(comp->external_path, ec)) { + AddFinding(warnings, "PATHS", + "component '" + comp->name + "' external file does not exist: " + + comp->external_path.string()); + } + } + } + for (const auto& rec : pkg->shared_assets) { + if (!fs::is_directory(rec->resolved_path, ec)) { + AddFinding(warnings, "PATHS", + "shared asset " + rec->uri + " resolved path is not a directory: " + + rec->resolved_path.string()); + } + } + } + + // ASSET_REHASH: re-hash each on-disk shared asset and compare to its URI. + if (flags & MODEL_PACKAGE_VALIDATE_ASSET_REHASH) { + for (const auto& rec : pkg->shared_assets) { + if (!fs::is_directory(rec->resolved_path, ec)) continue; // PATHS / REACH covers this. + std::string computed; + if (auto* s = mp::ComputeDirectoryAssetUri(rec->resolved_path, &computed)) { + AddFinding(errors, "ASSET_REHASH", + "shared asset " + rec->uri + ": hashing failed: " + + ModelPackageStatus_Message(s)); + ModelPackageStatus_Release(s); + continue; + } + if (computed != rec->uri) { + AddFinding(errors, "ASSET_REHASH", + "shared asset " + rec->uri + " on-disk hash differs: " + computed); + } + } + } + + // UNKNOWN_FIELDS: re-run with strict=true (only flags top-level / known scopes). + if (flags & MODEL_PACKAGE_VALIDATE_UNKNOWN_FIELDS) { + static const char* kKnown[] = { + "schema_version", "package_name", "package_version", "description", + "layout", "components", "shared_assets", "additional_metadata"}; + for (auto it = pkg->manifest.begin(); it != pkg->manifest.end(); ++it) { + bool found = false; + for (auto* k : kKnown) + if (it.key() == k) { + found = true; + break; + } + if (!found) { + AddFinding(warnings, "UNKNOWN_FIELDS", + "manifest contains unknown field '" + it.key() + "'."); + } + } + } + + pkg->last_validate_report = report.dump(2); + *out_report_json = pkg->last_validate_report->c_str(); + if (!errors->empty()) { + return MakeStatus(MODEL_PACKAGE_ERR_STATE, + "ModelPackage_Validate: " + std::to_string(errors->size()) + + " error(s) found. See out_report_json for details."); + } + return nullptr; +} + +} // extern "C" diff --git a/model_package/src/manifest_parser.cc b/model_package/src/manifest_parser.cc new file mode 100644 index 0000000000000..58b7ad16fdd7d --- /dev/null +++ b/model_package/src/manifest_parser.cc @@ -0,0 +1,732 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "manifest_parser.h" + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "path_resolver.h" +#include "status_impl.h" + +namespace fs = std::filesystem; + +namespace model_package { + +namespace { + +// The on-disk schema_version is a "." string (e.g. "1.0"). The major gates +// compatibility; the minor is informational and tells consumers which optional fields may +// be present. This build understands schema majors in [kMinSupportedSchemaMajor, +// kMaxSupportedSchemaMajor] and any minor: schema evolution within a major is additive and +// backward-compatible (newer minors only add optional fields), so a single parser reads +// every minor. A package whose major is below the minimum predates a breaking change and +// must be re-authored; one above the maximum was produced by a newer toolchain this build +// does not understand. kMaxKnownSchemaMinor is the highest minor this build authored/knows; +// a package with a higher minor is still accepted but may carry unknown optional fields, +// which are tolerated rather than rejected. +constexpr int64_t kMinSupportedSchemaMajor = 1; +constexpr int64_t kMaxSupportedSchemaMajor = 1; +constexpr int64_t kMaxKnownSchemaMinor = 0; +constexpr const char* kManifestFileName = "manifest.json"; +constexpr const char* kComponentFileName = "component.json"; + +constexpr const char* kSchemaVersionKey = "schema_version"; +constexpr const char* kPackageNameKey = "package_name"; +constexpr const char* kPackageVersionKey = "package_version"; +constexpr const char* kDescriptionKey = "description"; +constexpr const char* kLayoutKey = "layout"; +constexpr const char* kComponentsKey = "components"; +constexpr const char* kSharedAssetsKey = "shared_assets"; +constexpr const char* kAdditionalMetadataKey = "additional_metadata"; + +constexpr const char* kComponentNameKey = "component_name"; +constexpr const char* kVariantsKey = "variants"; + +constexpr const char* kVariantDirectoryKey = "variant_directory"; +constexpr const char* kEpKey = "ep"; +constexpr const char* kDeviceKey = "device"; +constexpr const char* kCompatibilityStringKey = "compatibility_string"; +constexpr const char* kExecutorInfoKey = "executor_info"; + +static const std::set kManifestKnownKeys = { + kSchemaVersionKey, + kPackageNameKey, + kPackageVersionKey, + kDescriptionKey, + kLayoutKey, + kComponentsKey, + kSharedAssetsKey, + kAdditionalMetadataKey, +}; + +static const std::set kComponentKnownKeys = { + kComponentNameKey, + kVariantsKey, + kAdditionalMetadataKey, +}; + +static const std::set kVariantKnownKeys = { + kVariantDirectoryKey, + kEpKey, + kDeviceKey, + kCompatibilityStringKey, + kExecutorInfoKey, + kAdditionalMetadataKey, +}; + +ModelPackageStatus* ReadFileToString(const fs::path& path, std::string* out) { + std::ifstream f(path, std::ios::binary); + if (!f) { + return MakeStatus(MODEL_PACKAGE_ERR_IO, + "Cannot open file: '" + path.string() + "': " + std::strerror(errno)); + } + std::ostringstream buf; + buf << f.rdbuf(); + *out = buf.str(); + return nullptr; +} + +ModelPackageStatus* ParseJsonFile(const fs::path& path, ordered_json* out) { + std::string contents; + if (auto* s = ReadFileToString(path, &contents)) return s; + try { + *out = ordered_json::parse(contents); + } catch (const ordered_json::parse_error& e) { + return MakeStatus(MODEL_PACKAGE_ERR_SCHEMA, + "Failed to parse JSON at '" + path.string() + "': " + e.what()); + } + return nullptr; +} + +ModelPackageStatus* ExpectObject(const ordered_json& j, const std::string& where) { + if (!j.is_object()) { + return MakeStatus(MODEL_PACKAGE_ERR_SCHEMA, where + ": expected a JSON object."); + } + return nullptr; +} + +ModelPackageStatus* CheckUnknownFields(const ordered_json& obj, + const std::set& known, + const std::string& where, + bool strict) { + if (!strict) return nullptr; + for (auto it = obj.begin(); it != obj.end(); ++it) { + if (known.find(it.key()) == known.end()) { + return MakeStatus(MODEL_PACKAGE_ERR_SCHEMA, + where + ": unknown field '" + it.key() + "'."); + } + } + return nullptr; +} + +ModelPackageStatus* ResolveVariantDirectory(const fs::path& component_dir, + const fs::path& package_root, + const ordered_json& variant_body, + const std::string& variant_name, + const PathResolverOptions& opts, + bool require_exists, + std::optional* out) { + auto it = variant_body.find(kVariantDirectoryKey); + bool explicitly_declared = (it != variant_body.end()); + std::string dir_input; + if (explicitly_declared) { + if (!it->is_string()) { + return MakeStatus(MODEL_PACKAGE_ERR_SCHEMA, + "variant '" + variant_name + "': variant_directory must be a string."); + } + dir_input = it->get(); + } else { + // Inferred default: missing-on-disk is fine; we just leave out unset. + dir_input = variant_name; + } + + fs::path resolved; + // Explicit value must exist; inferred default may not. + bool must_exist = require_exists || explicitly_declared; + auto* status = ResolvePath(component_dir, package_root, dir_input, opts, + must_exist, &resolved); + if (status) { + if (!must_exist && ModelPackageStatus_Code(status) == MODEL_PACKAGE_ERR_NOT_FOUND) { + ModelPackageStatus_Release(status); + *out = std::nullopt; + return nullptr; + } + return status; + } + std::error_code ec; + if (fs::exists(resolved, ec)) { + *out = resolved; + } else { + *out = std::nullopt; + } + return nullptr; +} + +ModelPackageStatus* ParseVariant(const fs::path& component_dir, + const fs::path& package_root, + const PathResolverOptions& opts, + bool strict, + const std::string& variant_name, + const ordered_json& variant_body, + VariantRecord* out); +ModelPackageStatus* ParseComponent(const fs::path& package_root, + const PathResolverOptions& opts, + bool strict, + const std::string& component_name, + const ordered_json& body, + const fs::path& component_dir, + ComponentRecord* out); +ModelPackageStatus* LoadSharedAssets(ModelPackage* pkg, const PathResolverOptions& opts); +ModelPackageStatus* PopulatePackageMetadata(ModelPackage* pkg); + +ModelPackageStatus* ParseVariant(const fs::path& component_dir, + const fs::path& package_root, + const PathResolverOptions& opts, + bool strict, + const std::string& variant_name, + const ordered_json& variant_body, + VariantRecord* out) { + if (auto* s = ExpectObject(variant_body, "variant '" + variant_name + "'")) return s; + if (auto* s = CheckUnknownFields(variant_body, kVariantKnownKeys, + "variant '" + variant_name + "'", strict)) + return s; + + out->name = variant_name; + out->body = variant_body; + out->name_cache = variant_name; + + auto stringopt = [&](const char* key, std::optional* dst) -> ModelPackageStatus* { + auto it = variant_body.find(key); + if (it == variant_body.end()) return nullptr; + if (!it->is_string()) { + return MakeStatus(MODEL_PACKAGE_ERR_SCHEMA, + std::string("variant '") + variant_name + "': '" + key + + "' must be a string."); + } + *dst = it->get(); + return nullptr; + }; + if (auto* s = stringopt(kEpKey, &out->ep_cache)) return s; + if (auto* s = stringopt(kDeviceKey, &out->device_cache)) return s; + if (auto* s = stringopt(kCompatibilityStringKey, &out->compatibility_string_cache)) return s; + + auto ei_it = variant_body.find(kExecutorInfoKey); + if (ei_it != variant_body.end()) { + if (!ei_it->is_object()) { + return MakeStatus(MODEL_PACKAGE_ERR_SCHEMA, + "variant '" + variant_name + "': executor_info must be an object."); + } + for (auto e = ei_it->begin(); e != ei_it->end(); ++e) { + if (!e->is_string() && !e->is_object()) { + return MakeStatus(MODEL_PACKAGE_ERR_SCHEMA, + "variant '" + variant_name + "': executor_info['" + e.key() + + "'] must be a string (path) or object (inline)."); + } + } + } + + // Resolve variant_directory if declared (records the resolved path when it + // exists on disk). We do NOT require the directory to exist here: executor + // semantics are not the library's concern, and executors must resolve their + // own file references against variant_directory at load time anyway. + std::optional resolved_dir; + auto* status = ResolveVariantDirectory(component_dir, package_root, variant_body, + variant_name, opts, + /*require_exists=*/false, &resolved_dir); + if (status) return status; + out->resolved_directory = resolved_dir; + out->resolved_directory_attempted = true; + if (resolved_dir.has_value()) { + out->resolved_directory_cache = resolved_dir->string(); + } + + return nullptr; +} + +ModelPackageStatus* ParseComponent(const fs::path& package_root, + const PathResolverOptions& opts, + bool strict, + const std::string& component_name, + const ordered_json& body, + const fs::path& component_dir, + ComponentRecord* out) { + if (auto* s = ExpectObject(body, "component '" + component_name + "'")) return s; + if (auto* s = CheckUnknownFields(body, kComponentKnownKeys, + "component '" + component_name + "'", strict)) + return s; + out->name = component_name; + out->name_cache = component_name; + out->component_dir = component_dir; + out->body = body; + + // Optional component_name override — for now we just sanity-check it. + auto cn_it = body.find(kComponentNameKey); + if (cn_it != body.end() && !cn_it->is_string()) { + return MakeStatus(MODEL_PACKAGE_ERR_SCHEMA, + "component '" + component_name + "': component_name must be a string."); + } + + auto variants_it = body.find(kVariantsKey); + if (variants_it == body.end()) { + return MakeStatus(MODEL_PACKAGE_ERR_SCHEMA, + "component '" + component_name + "': missing required 'variants' object."); + } + if (!variants_it->is_object()) { + return MakeStatus(MODEL_PACKAGE_ERR_SCHEMA, + "component '" + component_name + "': 'variants' must be an object."); + } + for (auto v = variants_it->begin(); v != variants_it->end(); ++v) { + auto vr = std::make_unique(); + if (auto* s = ParseVariant(component_dir, package_root, opts, strict, + v.key(), v.value(), vr.get())) { + return s; + } + out->variants.push_back(std::move(vr)); + } + return nullptr; +} + +ModelPackageStatus* LoadComponentForEntry(const fs::path& manifest_dir, + const fs::path& package_root, + const PathResolverOptions& opts, + bool strict, + const std::string& name, + const ordered_json& value, + std::unique_ptr* out) { + auto rec = std::make_unique(); + if (value.is_string()) { + rec->storage = ComponentStorage::kExternal; + fs::path resolved; + if (auto* s = ResolvePath(manifest_dir, package_root, value.get(), + opts, /*must_exist=*/true, &resolved)) { + return s; + } + std::error_code ec; + if (fs::is_directory(resolved, ec)) { + resolved /= kComponentFileName; + if (!fs::exists(resolved)) { + return MakeStatus(MODEL_PACKAGE_ERR_NOT_FOUND, + "component '" + name + "': directory has no '" + + kComponentFileName + "'."); + } + } + rec->external_path = resolved; + ordered_json body; + if (auto* s = ParseJsonFile(resolved, &body)) return s; + fs::path component_dir = resolved.parent_path(); + if (auto* s = ParseComponent(package_root, opts, strict, name, body, component_dir, rec.get())) { + return s; + } + } else if (value.is_object()) { + rec->storage = ComponentStorage::kInline; + if (auto* s = ParseComponent(package_root, opts, strict, name, value, manifest_dir, rec.get())) { + return s; + } + } else { + return MakeStatus(MODEL_PACKAGE_ERR_SCHEMA, + "component '" + name + "': value must be a string (path) or object (inline)."); + } + *out = std::move(rec); + return nullptr; +} + +ModelPackageStatus* LoadSharedAssets(ModelPackage* pkg, const PathResolverOptions& opts) { + // Source-of-truth ordering for the assembled shared_assets vector: + // 1. Manifest overrides (in declaration order). These specify a custom + // on-disk path for an asset URI (e.g. a system-wide cache or another + // location outside /shared_assets/). + // 2. Discovered sha256- subdirectories under /shared_assets/. + // These resolve to the default-convention path. Missing shared_assets/ is + // not an error: portable packages may not ship one yet, installed + // packages may route everything through overrides. + // 3. Pending copy_in assets from the authoring API that haven't been + // committed yet. These surface immediately so callers see the asset + // they just added; the staged source dir is reported as resolved_path + // until commit materializes it under shared_assets/. + // Within each tier, an URI that's already known is skipped. + std::vector ordered_uris; + std::unordered_map override_paths; + + auto sa_it = pkg->manifest.find(kSharedAssetsKey); + if (sa_it != pkg->manifest.end()) { + if (!sa_it->is_object()) { + return MakeStatus(MODEL_PACKAGE_ERR_SCHEMA, + "manifest: shared_assets must be an object."); + } + for (auto e = sa_it->begin(); e != sa_it->end(); ++e) { + const std::string uri = e.key(); + if (!IsSha256AssetUri(uri)) { + return MakeStatus(MODEL_PACKAGE_ERR_SCHEMA, + "manifest: shared_assets key '" + uri + "' is not a valid sha256: URI."); + } + if (!e->is_string()) { + return MakeStatus(MODEL_PACKAGE_ERR_SCHEMA, + "manifest: shared_assets['" + uri + "'] must be a string path."); + } + ordered_uris.push_back(uri); + override_paths.emplace(uri, e->get()); + } + } + std::set seen(ordered_uris.begin(), ordered_uris.end()); + + // Tier 2: discover sha256- dirs under /shared_assets/. + fs::path assets_root = pkg->package_root / "shared_assets"; + std::error_code ec; + if (!pkg->package_root.empty() && fs::is_directory(assets_root, ec)) { + std::vector discovered; + for (const auto& entry : fs::directory_iterator(assets_root, ec)) { + if (ec) break; + if (!entry.is_directory(ec)) continue; + std::string name = entry.path().filename().string(); + std::string uri = SharedAssetUriFromDirName(name); + if (uri.empty()) continue; // not a sha256- dir; ignore (.tmp staging, etc.) + if (!seen.insert(uri).second) continue; + discovered.push_back(std::move(uri)); + } + std::sort(discovered.begin(), discovered.end()); + for (auto& uri : discovered) ordered_uris.push_back(std::move(uri)); + } + + // Tier 3: pending copy_in entries. + for (const auto& [uri, src] : pkg->pending_shared_asset_copies) { + if (!seen.insert(uri).second) continue; + ordered_uris.push_back(uri); + } + + for (const auto& uri : ordered_uris) { + auto rec = std::make_unique(); + rec->uri = uri; + rec->uri_cache = uri; + auto override_it = override_paths.find(uri); + fs::path resolved; + if (override_it != override_paths.end()) { + if (auto* s = ResolvePath(pkg->package_root, pkg->package_root, override_it->second, + opts, /*must_exist=*/false, &resolved)) { + return s; + } + } else if (auto pending_it = pkg->pending_shared_asset_copies.find(uri); + pending_it != pkg->pending_shared_asset_copies.end() && + override_paths.find(uri) == override_paths.end()) { + // Pending copy_in with no override: surface the staged source until commit. + resolved = pending_it->second; + } else { + // Default convention: /shared_assets/sha256-/ + resolved = assets_root / DefaultSharedAssetDirName(uri); + } + rec->resolved_path = resolved; + rec->resolved_path_cache = resolved.string(); + pkg->shared_asset_index_by_uri.emplace(uri, pkg->shared_assets.size()); + pkg->shared_assets.push_back(std::move(rec)); + } + return nullptr; +} + +ModelPackageStatus* ParseSchemaVersion(ModelPackage* pkg) { + auto sv_it = pkg->manifest.find(kSchemaVersionKey); + if (sv_it == pkg->manifest.end()) { + return MakeStatus(MODEL_PACKAGE_ERR_SCHEMA, + "manifest: missing required 'schema_version'."); + } + + // schema_version is a "." string (e.g. "1.0"). A bare integer is accepted + // as shorthand for ".0". + int64_t major = 0; + int64_t minor = 0; + if (sv_it->is_string()) { + const std::string sv = sv_it->get(); + const size_t dot = sv.find('.'); + const std::string major_str = (dot == std::string::npos) ? sv : sv.substr(0, dot); + const std::string minor_str = (dot == std::string::npos) ? std::string("0") : sv.substr(dot + 1); + auto parse_part = [](const std::string& s, int64_t* out) -> bool { + if (s.empty() || s.find_first_not_of("0123456789") != std::string::npos) return false; + try { + *out = std::stoll(s); + } catch (const std::exception&) { + return false; + } + return true; + }; + if (dot != std::string::npos && minor_str.find('.') != std::string::npos) { + return MakeStatus(MODEL_PACKAGE_ERR_SCHEMA, + "manifest: 'schema_version' must be a \".\" string."); + } + if (!parse_part(major_str, &major) || !parse_part(minor_str, &minor)) { + return MakeStatus(MODEL_PACKAGE_ERR_SCHEMA, + "manifest: 'schema_version' must be a \".\" string."); + } + } else if (sv_it->is_number_integer() || sv_it->is_number_unsigned()) { + major = sv_it->get(); + minor = 0; + } else { + return MakeStatus(MODEL_PACKAGE_ERR_SCHEMA, + "manifest: 'schema_version' must be a \".\" string."); + } + + if (major < kMinSupportedSchemaMajor || major > kMaxSupportedSchemaMajor) { + std::string supported = (kMinSupportedSchemaMajor == kMaxSupportedSchemaMajor) + ? std::to_string(kMinSupportedSchemaMajor) + : std::to_string(kMinSupportedSchemaMajor) + "-" + + std::to_string(kMaxSupportedSchemaMajor); + return MakeStatus(MODEL_PACKAGE_ERR_VERSION, + "manifest: schema_version major " + std::to_string(major) + + " is not supported (this build supports major " + supported + ")."); + } + pkg->schema_version_major = major; + pkg->schema_version_minor = minor; + + // A package authored at a newer minor than this build knows may carry optional fields this + // build does not recognize. Those are additive and must be tolerated rather than rejected, + // so relax unknown-field strictness for a newer minor. + if (minor > kMaxKnownSchemaMinor) { + pkg->strict_unknown_fields = false; + } + return nullptr; +} + +ModelPackageStatus* PopulatePackageMetadata(ModelPackage* pkg) { + auto stropt = [&](const char* key, std::optional* dst) -> ModelPackageStatus* { + auto it = pkg->manifest.find(key); + if (it == pkg->manifest.end()) { + dst->reset(); + return nullptr; + } + if (!it->is_string()) { + return MakeStatus(MODEL_PACKAGE_ERR_SCHEMA, + std::string("manifest: '") + key + "' must be a string."); + } + *dst = it->get(); + return nullptr; + }; + if (auto* s = stropt(kPackageNameKey, &pkg->package_name_cache)) return s; + if (auto* s = stropt(kPackageVersionKey, &pkg->package_version_cache)) return s; + if (auto* s = stropt(kDescriptionKey, &pkg->description_cache)) return s; + + // layout: default "portable" + auto layout_it = pkg->manifest.find(kLayoutKey); + if (layout_it != pkg->manifest.end()) { + if (!layout_it->is_string()) { + return MakeStatus(MODEL_PACKAGE_ERR_SCHEMA, "manifest: 'layout' must be a string."); + } + pkg->layout = layout_it->get(); + if (pkg->layout != "portable" && pkg->layout != "installed") { + return MakeStatus(MODEL_PACKAGE_ERR_SCHEMA, + "manifest: 'layout' must be 'portable' or 'installed'."); + } + } else { + pkg->layout = "portable"; + } + pkg->layout_cache = pkg->layout; + + // additional_metadata: serialize as JSON string if present. + auto am_it = pkg->manifest.find(kAdditionalMetadataKey); + if (am_it != pkg->manifest.end()) { + pkg->additional_metadata_cache = am_it->dump(); + } else { + pkg->additional_metadata_cache.reset(); + } + return nullptr; +} + +} // namespace + +PathResolverOptions PathOptionsFor(const ModelPackage* pkg) { + PathResolverOptions o; + o.follow_symlinks = pkg->follow_symlinks; + o.allow_external_paths = pkg->allow_external_paths || (pkg->layout == "installed"); + return o; +} + +ModelPackageStatus* ParseVariantBody(const fs::path& component_dir, + const fs::path& package_root, + const PathResolverOptions& opts, + bool strict, + const std::string& variant_name, + const ordered_json& variant_body, + VariantRecord* out) { + return ParseVariant(component_dir, package_root, opts, strict, variant_name, variant_body, out); +} + +ModelPackageStatus* ParseComponentBody(const fs::path& package_root, + const PathResolverOptions& opts, + bool strict, + const std::string& component_name, + const ordered_json& body, + const fs::path& component_dir, + ComponentRecord* out) { + return ParseComponent(package_root, opts, strict, component_name, body, component_dir, out); +} + +ModelPackageStatus* RefreshPackageMetadata(ModelPackage* pkg) { + pkg->package_name_cache.reset(); + pkg->package_version_cache.reset(); + pkg->description_cache.reset(); + pkg->additional_metadata_cache.reset(); + return PopulatePackageMetadata(pkg); +} + +ModelPackageStatus* RefreshSharedAssets(ModelPackage* pkg, const PathResolverOptions& opts) { + pkg->shared_assets.clear(); + pkg->shared_asset_index_by_uri.clear(); + return LoadSharedAssets(pkg, opts); +} + +namespace { + +ModelPackageStatus* ResolveExecutorInfoEntry(const ModelPackage* pkg, + const VariantRecord& var, + const std::string& ns, + const ordered_json& entry, + bool strict_missing_external, + std::string* dst_json) { + if (entry.is_object()) { + *dst_json = entry.dump(); + return nullptr; + } + if (entry.is_string()) { + if (!var.resolved_directory.has_value()) { + if (!strict_missing_external) { + dst_json->clear(); + return nullptr; + } + return MakeStatus(MODEL_PACKAGE_ERR_NOT_FOUND, + "variant '" + var.name + "': executor_info['" + ns + + "'] points at an external file but the variant has no " + "resolved variant_directory to anchor it."); + } + PathResolverOptions opts = PathOptionsFor(pkg); + fs::path resolved; + if (auto* s = ResolvePath(*var.resolved_directory, pkg->package_root, + entry.get(), opts, + /*must_exist=*/strict_missing_external, &resolved)) { + if (!strict_missing_external) { + ModelPackageStatus_Release(s); + dst_json->clear(); + return nullptr; + } + return s; + } + std::ifstream f(resolved, std::ios::binary); + if (!f) { + if (!strict_missing_external) { + dst_json->clear(); + return nullptr; + } + return MakeStatus(MODEL_PACKAGE_ERR_IO, + "Cannot open executor_info file: '" + resolved.string() + "'."); + } + std::ostringstream buf; + buf << f.rdbuf(); + std::string contents = buf.str(); + try { + auto _ = ordered_json::parse(contents); + (void)_; + } catch (const std::exception& e) { + return MakeStatus(MODEL_PACKAGE_ERR_SCHEMA, + std::string("Failed to parse executor_info JSON at '") + + resolved.string() + "': " + e.what()); + } + *dst_json = std::move(contents); + return nullptr; + } + return MakeStatus(MODEL_PACKAGE_ERR_SCHEMA, + "variant '" + var.name + "': executor_info['" + ns + + "'] must be a string or object."); +} + +} // namespace + +ModelPackageStatus* RefreshExecutorInfoCache(ModelPackage* pkg, bool strict_missing_external) { + for (auto& comp : pkg->components) { + for (auto& vp : comp->variants) { + VariantRecord& var = *vp; + var.executor_info_resolved.clear(); + auto ei_it = var.body.find("executor_info"); + if (ei_it == var.body.end() || !ei_it->is_object()) continue; + var.executor_info_resolved.reserve(ei_it->size()); + for (auto e = ei_it->begin(); e != ei_it->end(); ++e) { + std::string body_json; + if (auto* s = ResolveExecutorInfoEntry(pkg, var, e.key(), e.value(), + strict_missing_external, &body_json)) { + return s; + } + var.executor_info_resolved.emplace_back(e.key(), std::move(body_json)); + } + } + } + return nullptr; +} + +ModelPackageStatus* ParsePackage(const fs::path& package_root, + const ModelPackageOpenOptions& opts, + ModelPackage* pkg) { + std::error_code ec; + if (!fs::exists(package_root, ec) || !fs::is_directory(package_root, ec)) { + return MakeStatus(MODEL_PACKAGE_ERR_IO, + "package_root '" + package_root.string() + "' is not a directory."); + } + pkg->package_root = fs::canonical(package_root, ec); + if (ec) pkg->package_root = package_root; + pkg->allow_external_paths = opts.allow_external_paths; + pkg->follow_symlinks = opts.follow_symlinks; + pkg->strict_unknown_fields = opts.strict_unknown_fields; + + fs::path manifest_path = pkg->package_root / kManifestFileName; + if (auto* s = ParseJsonFile(manifest_path, &pkg->manifest)) return s; + if (auto* s = ExpectObject(pkg->manifest, "manifest")) return s; + + // Validate the schema version first so an unsupported package fails fast, before any + // component/asset parsing. May relax pkg->strict_unknown_fields for a newer minor. + if (auto* s = ParseSchemaVersion(pkg)) return s; + + // Layout pre-read for path-resolver options. Done before strict-unknown + // check because we need the layout value to decide path-confinement. + PathResolverOptions presolve_opts; + presolve_opts.follow_symlinks = opts.follow_symlinks; + presolve_opts.allow_external_paths = opts.allow_external_paths; + { + auto layout_it = pkg->manifest.find(kLayoutKey); + if (layout_it != pkg->manifest.end() && layout_it->is_string() && + layout_it->get() == "installed") { + presolve_opts.allow_external_paths = true; + } + } + + if (auto* s = CheckUnknownFields(pkg->manifest, kManifestKnownKeys, "manifest", + pkg->strict_unknown_fields)) + return s; + + // Components. + auto comps_it = pkg->manifest.find(kComponentsKey); + if (comps_it == pkg->manifest.end()) { + return MakeStatus(MODEL_PACKAGE_ERR_SCHEMA, + "manifest: missing required 'components' object."); + } + if (!comps_it->is_object()) { + return MakeStatus(MODEL_PACKAGE_ERR_SCHEMA, "manifest: 'components' must be an object."); + } + for (auto e = comps_it->begin(); e != comps_it->end(); ++e) { + std::unique_ptr rec; + if (auto* s = LoadComponentForEntry(pkg->package_root, pkg->package_root, + presolve_opts, pkg->strict_unknown_fields, + e.key(), e.value(), &rec)) { + return s; + } + pkg->component_index_by_name.emplace(rec->name, pkg->components.size()); + pkg->components.push_back(std::move(rec)); + } + + if (auto* s = LoadSharedAssets(pkg, presolve_opts)) return s; + if (auto* s = PopulatePackageMetadata(pkg)) return s; + if (auto* s = RefreshExecutorInfoCache(pkg, /*strict_missing_external=*/true)) return s; + + return nullptr; +} + +} // namespace model_package diff --git a/model_package/src/manifest_parser.h b/model_package/src/manifest_parser.h new file mode 100644 index 0000000000000..d8266605440dc --- /dev/null +++ b/model_package/src/manifest_parser.h @@ -0,0 +1,64 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +/// \file manifest_parser.h +/// \brief Internal parser that reads a model package from disk into the +/// in-memory representation defined in model_package_impl.h. + +#pragma once + +#include "model_package_impl.h" +#include "path_resolver.h" + +namespace model_package { + +/// Parse the manifest at `/manifest.json` and all referenced +/// external component files, then populate `*pkg`. Caller owns `pkg`. +ModelPackageStatus* ParsePackage(const std::filesystem::path& package_root, + const ModelPackageOpenOptions& opts, + ModelPackage* pkg); + +/// Parse a single variant body into `out`. Used by authoring. +ModelPackageStatus* ParseVariantBody(const std::filesystem::path& component_dir, + const std::filesystem::path& package_root, + const PathResolverOptions& opts, + bool strict, + const std::string& variant_name, + const ordered_json& variant_body, + VariantRecord* out); + +/// Parse a single component body. `component_dir` is the directory used as the +/// base for the component's relative paths. +ModelPackageStatus* ParseComponentBody(const std::filesystem::path& package_root, + const PathResolverOptions& opts, + bool strict, + const std::string& component_name, + const ordered_json& body, + const std::filesystem::path& component_dir, + ComponentRecord* out); + +/// Re-derive package-level metadata (schema_version, package_name, version, +/// description, layout, additional_metadata) from `pkg->manifest` into the +/// package's stable string buffers. +ModelPackageStatus* RefreshPackageMetadata(ModelPackage* pkg); + +/// Re-derive `pkg->shared_assets` from `pkg->manifest.shared_assets` overrides, +/// plus any `sha256-` subdirectories discovered under +/// `/shared_assets/`, plus any pending copy_in entries staged via +/// the authoring API. Clears and replaces the existing shared_assets vector +/// and `shared_asset_index_by_uri`. +ModelPackageStatus* RefreshSharedAssets(ModelPackage* pkg, const PathResolverOptions& opts); + +/// Re-resolve every variant's executor_info entries into stable strings on the +/// VariantRecord (inline bodies dumped, external files loaded + JSON-parsed). +/// If `strict_missing_external` is true, missing external files are an error +/// (use at Open: the package is already published, files must be present); +/// if false, missing external files are recorded as an empty body (use during +/// authoring: callers may set the path before writing the file). Parse errors +/// on existing external files are always surfaced. +ModelPackageStatus* RefreshExecutorInfoCache(ModelPackage* pkg, bool strict_missing_external); + +/// Build PathResolverOptions appropriate for `pkg` (respects layout). +PathResolverOptions PathOptionsFor(const ModelPackage* pkg); + +} // namespace model_package diff --git a/model_package/src/model_package_impl.cc b/model_package/src/model_package_impl.cc new file mode 100644 index 0000000000000..cf383c659b1c7 --- /dev/null +++ b/model_package/src/model_package_impl.cc @@ -0,0 +1,403 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +/// \file model_package_impl.cc +/// \brief Implementation of the public C API declared in model_package.h. + +#include "model_package.h" + +#include +#include +#include +#include + +#include "asset_hasher.h" +#include "manifest_parser.h" +#include "model_package_impl.h" +#include "path_resolver.h" +#include "status_impl.h" + +namespace mp = model_package; +using mp::MakeStatus; + +namespace { + +ModelPackageStatus* NullArg(const char* name) { + return MakeStatus(MODEL_PACKAGE_ERR_INVALID_ARG, + std::string("model_package: '") + name + "' must not be null."); +} + +const char* OptStr(const std::optional& s) { + return s.has_value() ? s->c_str() : nullptr; +} + +} // namespace + +// ───────────────────────────────────────────────────────────────────────────── +// View cache materialization +// ───────────────────────────────────────────────────────────────────────────── + +namespace model_package { + +void DropViewCache(ModelPackage* pkg) { + if (!pkg) return; + pkg->info_cache.reset(); + for (auto& comp : pkg->components) { + comp->component_json_cache.reset(); + comp->additional_metadata_cache.reset(); + for (auto& var : comp->variants) { + var->variant_json_cache.reset(); + var->additional_metadata_cache.reset(); + } + } + pkg->additional_metadata_cache.reset(); +} + +const InfoViewCache& BuildOrGetViewCache(const ModelPackage* pkg) { + if (pkg->info_cache.has_value()) return *pkg->info_cache; + + pkg->info_cache.emplace(); + auto& cache = *pkg->info_cache; + const size_t num_components = pkg->components.size(); + + cache.executor_infos_storage.resize(num_components); + cache.variants_storage.resize(num_components); + cache.components.resize(num_components); + + for (size_t ci = 0; ci < num_components; ++ci) { + const auto& comp = *pkg->components[ci]; + const size_t num_variants = comp.variants.size(); + cache.executor_infos_storage[ci].clear(); + cache.variants_storage[ci].resize(num_variants); + + // Total executor_info entry count across this component's variants. + size_t total_execs = 0; + for (const auto& vp : comp.variants) { + total_execs += vp->executor_info_resolved.size(); + } + cache.executor_infos_storage[ci].reserve(total_execs); + + // First pass: append all executor_info entries so storage pointers stay + // stable for the second pass. + std::vector> ei_ranges(num_variants); + + for (size_t vi = 0; vi < num_variants; ++vi) { + const auto& var = *comp.variants[vi]; + size_t ei_begin = cache.executor_infos_storage[ci].size(); + // executor_info_resolved is populated eagerly by RefreshExecutorInfoCache + // (at Open and on every mutation); any parse/IO error surfaces there. + for (const auto& [ns_str, body_json] : var.executor_info_resolved) { + ModelExecutorInfoEntry entry{}; + entry.namespace_key = ns_str.c_str(); + entry.json = body_json.c_str(); + cache.executor_infos_storage[ci].push_back(entry); + } + ei_ranges[vi] = {ei_begin, cache.executor_infos_storage[ci].size()}; + } + + // Additional metadata strings live in the record-level cache; populate it + // lazily here as well. + for (size_t vi = 0; vi < num_variants; ++vi) { + auto& var = *comp.variants[vi]; + auto am_it = var.body.find("additional_metadata"); + if (am_it != var.body.end() && !var.additional_metadata_cache.has_value()) { + var.additional_metadata_cache = am_it->dump(); + } + } + if (auto am_it = comp.body.find("additional_metadata"); am_it != comp.body.end()) { + if (!comp.additional_metadata_cache.has_value()) { + comp.additional_metadata_cache = am_it->dump(); + } + } + + // Second pass: populate ModelVariantInfo entries pointing at the now-stable + // storage above. + for (size_t vi = 0; vi < num_variants; ++vi) { + const auto& var = *comp.variants[vi]; + ModelVariantInfo& vi_out = cache.variants_storage[ci][vi]; + vi_out = ModelVariantInfo{}; + vi_out.name = var.name_cache.c_str(); + vi_out.variant_directory = + var.resolved_directory_cache.has_value() ? var.resolved_directory_cache->c_str() : nullptr; + vi_out.ep = OptStr(var.ep_cache); + vi_out.device = OptStr(var.device_cache); + vi_out.compatibility_string = OptStr(var.compatibility_string_cache); + vi_out.additional_metadata_json = OptStr(var.additional_metadata_cache); + auto [ei_begin, ei_end] = ei_ranges[vi]; + vi_out.num_executor_infos = ei_end - ei_begin; + vi_out.executor_infos = + (vi_out.num_executor_infos > 0) ? &cache.executor_infos_storage[ci][ei_begin] : nullptr; + } + + ModelComponentInfo& ci_out = cache.components[ci]; + ci_out = ModelComponentInfo{}; + ci_out.name = comp.name_cache.c_str(); + ci_out.additional_metadata_json = OptStr(comp.additional_metadata_cache); + ci_out.num_variants = num_variants; + ci_out.variants = num_variants > 0 ? cache.variants_storage[ci].data() : nullptr; + } + + // Shared assets. + cache.shared_assets.resize(pkg->shared_assets.size()); + for (size_t i = 0; i < pkg->shared_assets.size(); ++i) { + const auto& rec = *pkg->shared_assets[i]; + ModelSharedAssetInfo& sa = cache.shared_assets[i]; + sa = ModelSharedAssetInfo{}; + sa.uri = rec.uri_cache.c_str(); + sa.resolved_path = rec.resolved_path_cache.c_str(); + } + + ModelPackageInfo& info = cache.info; + info = ModelPackageInfo{}; + info.schema_version_major = pkg->schema_version_major; + info.schema_version_minor = pkg->schema_version_minor; + info.package_name = OptStr(pkg->package_name_cache); + info.package_version = OptStr(pkg->package_version_cache); + info.description = OptStr(pkg->description_cache); + info.layout = pkg->layout_cache.c_str(); + info.additional_metadata_json = OptStr(pkg->additional_metadata_cache); + info.num_components = cache.components.size(); + info.components = cache.components.empty() ? nullptr : cache.components.data(); + info.num_shared_assets = cache.shared_assets.size(); + info.shared_assets = cache.shared_assets.empty() ? nullptr : cache.shared_assets.data(); + + return cache; +} + +} // namespace model_package + +// ───────────────────────────────────────────────────────────────────────────── +// Status helpers +// ───────────────────────────────────────────────────────────────────────────── + +extern "C" { + +const char* ModelPackageStatus_Message(const ModelPackageStatus* s) { + return s ? s->message.c_str() : nullptr; +} +ModelPackageErrorCode ModelPackageStatus_Code(const ModelPackageStatus* s) { + return s ? s->code : MODEL_PACKAGE_OK; +} +void ModelPackageStatus_Release(ModelPackageStatus* s) { + delete s; +} + +// ───────────────────────────────────────────────────────────────────────────── +// Lifecycle +// ───────────────────────────────────────────────────────────────────────────── + +ModelPackageStatus* ModelPackage_Open(const char* package_root, + const ModelPackageOpenOptions* opts, + ModelPackage** out) { + if (!package_root) return NullArg("package_root"); + if (!out) return NullArg("out"); + *out = nullptr; + + ModelPackageOpenOptions effective{}; + effective.allow_external_paths = false; + effective.follow_symlinks = true; + effective.strict_unknown_fields = true; + if (opts) { + effective.allow_external_paths = opts->allow_external_paths; + effective.follow_symlinks = opts->follow_symlinks; + effective.strict_unknown_fields = opts->strict_unknown_fields; + } + + auto pkg = std::make_unique(); + if (auto* s = mp::ParsePackage(std::filesystem::path(package_root), effective, pkg.get())) { + return s; + } + *out = pkg.release(); + return nullptr; +} + +void ModelPackage_Close(ModelPackage* pkg) { + if (!pkg) return; + mp::DropViewCache(pkg); + delete pkg; +} + +// ───────────────────────────────────────────────────────────────────────────── +// Info tree + convenience lookups +// ───────────────────────────────────────────────────────────────────────────── + +const ModelPackageInfo* ModelPackage_Info(const ModelPackage* pkg) { + if (!pkg) return nullptr; + return &mp::BuildOrGetViewCache(pkg).info; +} + +const ModelComponentInfo* ModelPackage_FindComponent(const ModelPackageInfo* info, + const char* name) { + if (!info || !name) return nullptr; + for (size_t i = 0; i < info->num_components; ++i) { + if (info->components[i].name && std::strcmp(info->components[i].name, name) == 0) { + return &info->components[i]; + } + } + return nullptr; +} + +const ModelVariantInfo* ModelComponentInfo_FindVariant(const ModelComponentInfo* comp, + const char* name) { + if (!comp || !name) return nullptr; + for (size_t i = 0; i < comp->num_variants; ++i) { + if (comp->variants[i].name && std::strcmp(comp->variants[i].name, name) == 0) { + return &comp->variants[i]; + } + } + return nullptr; +} + +const ModelExecutorInfoEntry* ModelVariantInfo_FindExecutorInfo(const ModelVariantInfo* var, + const char* namespace_key) { + if (!var || !namespace_key) return nullptr; + for (size_t i = 0; i < var->num_executor_infos; ++i) { + if (var->executor_infos[i].namespace_key && + std::strcmp(var->executor_infos[i].namespace_key, namespace_key) == 0) { + return &var->executor_infos[i]; + } + } + return nullptr; +} + +// ───────────────────────────────────────────────────────────────────────────── +// Shared assets +// ───────────────────────────────────────────────────────────────────────────── + +ModelPackageStatus* ModelPackage_ResolveAssetUri(const ModelPackage* pkg, + const char* uri, + const char** out_path) { + if (!pkg) return NullArg("pkg"); + if (!uri) return NullArg("uri"); + if (!out_path) return NullArg("out_path"); + *out_path = nullptr; + auto it = pkg->shared_asset_index_by_uri.find(uri); + if (it == pkg->shared_asset_index_by_uri.end()) { + return MakeStatus(MODEL_PACKAGE_ERR_ASSET_MISSING, + std::string("Asset URI not declared in this package: '") + uri + "'."); + } + *out_path = pkg->shared_assets[it->second]->resolved_path_cache.c_str(); + return nullptr; +} + +ModelPackageStatus* ModelPackage_ResolveStringRef(const ModelPackage* pkg, + const char* base_dir, + const char* input, + bool must_exist, + const char** out_path) { + if (!pkg) return NullArg("pkg"); + if (!input) return NullArg("input"); + if (!out_path) return NullArg("out_path"); + *out_path = nullptr; + static thread_local std::string slot; + + std::string uri_part, tail_part; + if (mp::TrySplitAssetUriPrefix(std::string(input), uri_part, tail_part)) { + auto asset_it = pkg->shared_asset_index_by_uri.find(uri_part); + if (asset_it == pkg->shared_asset_index_by_uri.end()) { + return MakeStatus(MODEL_PACKAGE_ERR_ASSET_MISSING, + std::string("Asset URI not declared in this package: '") + uri_part + "'."); + } + const std::string& asset_folder = pkg->shared_assets[asset_it->second]->resolved_path_cache; + if (tail_part.empty()) { + slot = asset_folder; + *out_path = slot.c_str(); + return nullptr; + } + // Tail is resolved with portable confinement under the asset folder: + // no absolute, no `..`. follow_symlinks mirrors the package setting. + mp::PathResolverOptions tail_opts; + tail_opts.allow_external_paths = false; + tail_opts.follow_symlinks = pkg->follow_symlinks; + std::filesystem::path resolved; + if (auto* s = mp::ResolvePath(asset_folder, asset_folder, tail_part, tail_opts, + must_exist, &resolved)) { + return s; + } + slot = resolved.string(); + *out_path = slot.c_str(); + return nullptr; + } + + std::filesystem::path base = base_dir ? std::filesystem::path(base_dir) : pkg->package_root; + std::filesystem::path resolved; + if (auto* s = mp::ResolvePath(base, pkg->package_root, std::string(input), + mp::PathOptionsFor(pkg), must_exist, &resolved)) { + return s; + } + slot = resolved.string(); + *out_path = slot.c_str(); + return nullptr; +} + +// ───────────────────────────────────────────────────────────────────────────── +// Round-trip JSON getters +// ───────────────────────────────────────────────────────────────────────────── + +ModelPackageStatus* ModelPackage_GetComponentJson(const ModelPackage* pkg, + const char* component_name, + const char** out_json) { + if (!pkg) return NullArg("pkg"); + if (!component_name) return NullArg("component_name"); + if (!out_json) return NullArg("out_json"); + *out_json = nullptr; + auto it = pkg->component_index_by_name.find(component_name); + if (it == pkg->component_index_by_name.end()) { + return MakeStatus(MODEL_PACKAGE_ERR_NOT_FOUND, + std::string("component '") + component_name + "' not found."); + } + auto& rec = pkg->components[it->second]; + if (!rec->component_json_cache.has_value()) { + rec->component_json_cache = rec->body.dump(); + } + *out_json = rec->component_json_cache->c_str(); + return nullptr; +} + +ModelPackageStatus* ModelPackage_GetVariantJson(const ModelPackage* pkg, + const char* component_name, + const char* variant_name, + const char** out_json) { + if (!pkg) return NullArg("pkg"); + if (!component_name) return NullArg("component_name"); + if (!variant_name) return NullArg("variant_name"); + if (!out_json) return NullArg("out_json"); + *out_json = nullptr; + auto it = pkg->component_index_by_name.find(component_name); + if (it == pkg->component_index_by_name.end()) { + return MakeStatus(MODEL_PACKAGE_ERR_NOT_FOUND, + std::string("component '") + component_name + "' not found."); + } + auto& comp = pkg->components[it->second]; + for (auto& var : comp->variants) { + if (var->name == variant_name) { + if (!var->variant_json_cache.has_value()) { + var->variant_json_cache = var->body.dump(); + } + *out_json = var->variant_json_cache->c_str(); + return nullptr; + } + } + return MakeStatus(MODEL_PACKAGE_ERR_NOT_FOUND, + std::string("variant '") + variant_name + "' not found in component '" + + component_name + "'."); +} + +// ───────────────────────────────────────────────────────────────────────────── +// Hashing utility +// ───────────────────────────────────────────────────────────────────────────── + +ModelPackageStatus* ModelPackage_ComputeDirectoryHash(const char* source_dir, + const char** out_uri) { + if (!source_dir) return NullArg("source_dir"); + if (!out_uri) return NullArg("out_uri"); + *out_uri = nullptr; + static thread_local std::string slot; + if (auto* s = mp::ComputeDirectoryAssetUri(std::filesystem::path(source_dir), &slot)) { + return s; + } + *out_uri = slot.c_str(); + return nullptr; +} + +} // extern "C" diff --git a/model_package/src/model_package_impl.h b/model_package/src/model_package_impl.h new file mode 100644 index 0000000000000..6770b9774e132 --- /dev/null +++ b/model_package/src/model_package_impl.h @@ -0,0 +1,161 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +/// \file model_package_impl.h +/// \brief Internal C++ representation of a ModelPackage handle. +/// +/// Records hold the parsed manifest data plus stable per-entity string buffers +/// so that all `const char*` exposed through the C API have package-owned +/// storage. The package builds an `InfoViewCache` lazily that materializes the +/// public POD struct tree returned by `ModelPackage_Info()`; any mutation +/// drops the cache so the next read produces a fresh tree. + +#pragma once + +#include +#include +#include +#include +#include +#include +#include + +#include + +#include "model_package.h" + +namespace model_package { + +using ordered_json = nlohmann::ordered_json; + +/// How the component's body is stored on disk relative to the manifest. +enum class ComponentStorage { + kInline, ///< body lives directly inside the manifest as an object + kExternal, ///< body lives in a separate file pointed to by a string +}; + +struct VariantRecord { + std::string name; + nlohmann::ordered_json body; ///< the full variant JSON object + + // Stable string buffers for ABI exposure. + std::string name_cache; + std::optional ep_cache; + std::optional device_cache; + std::optional compatibility_string_cache; + std::optional resolved_directory_cache; + mutable std::optional additional_metadata_cache; + mutable std::optional variant_json_cache; + + /// Resolved variant_directory for variants that have one. `std::nullopt` + /// means none was declared and the default location does not exist. + std::optional resolved_directory; + bool resolved_directory_attempted{false}; + + /// Pre-resolved executor_info entries. Populated eagerly at Open and + /// after any mutation that can touch executor_info. The first member is the + /// namespace key; the second is the serialized JSON body of that entry + /// (inline bodies are dumped, external file bodies are read + validated). + std::vector> executor_info_resolved; +}; + +struct ComponentRecord { + std::string name; + ComponentStorage storage{ComponentStorage::kInline}; + std::filesystem::path external_path; ///< valid iff storage == kExternal + std::filesystem::path component_dir; ///< base directory for relative paths inside this component + nlohmann::ordered_json body; ///< {"component_name": ..., "variants": {...}, "additional_metadata": {...}} + std::vector> variants; + + std::string name_cache; + mutable std::optional additional_metadata_cache; + mutable std::optional component_json_cache; +}; + +struct SharedAssetRecord { + std::string uri; ///< "sha256:" + std::filesystem::path resolved_path; + std::string uri_cache; + std::string resolved_path_cache; +}; + +/// Materialized POD-struct tree returned by ModelPackage_Info(). Owns all +/// backing storage (extra strings and array buffers) so pointers stay valid +/// until the next mutation drops the cache. +struct InfoViewCache { + // Per-variant arrays. Indexed [component_idx][variant_idx]. + std::vector> executor_infos_storage; + std::vector> variants_storage; + + std::vector components; + std::vector shared_assets; + ModelPackageInfo info{}; +}; + +} // namespace model_package + +// ───────────────────────────────────────────────────────────────────────────── +// Public opaque type (lives in the global namespace to match the C API) +// ───────────────────────────────────────────────────────────────────────────── + +struct ModelPackage { + std::filesystem::path package_root; + nlohmann::ordered_json manifest; ///< parsed manifest.json with declarations intact (component values stay in their original string-or-object form) + std::string layout; ///< "portable" | "installed" + + // Open-time options. + bool allow_external_paths{false}; + bool follow_symlinks{true}; + bool strict_unknown_fields{true}; + + // Package-level parsed data and stable string buffers. + int64_t schema_version_major{0}; + int64_t schema_version_minor{0}; + std::optional package_name_cache; + std::optional package_version_cache; + std::optional description_cache; + std::string layout_cache; + mutable std::optional additional_metadata_cache; + + std::vector> components; + std::vector> shared_assets; + + std::unordered_map component_index_by_name; + std::unordered_map shared_asset_index_by_uri; + + /// Authoring-time staging for copy_in=true shared assets that have not been + /// committed yet. Keyed by sha256: URI. + std::unordered_map pending_shared_asset_copies; + + /// Paths removed from the live tree, candidates for ModelPackage_Prune. + /// Populated by the authoring API; never by walking package_root. + std::vector pending_orphan_variant_dirs; + std::vector pending_orphan_component_dirs; + + /// Cache for the most recent ModelPackage_Validate report JSON. + mutable std::optional last_validate_report; + + /// Lazily built; dropped on any mutation. + mutable std::optional info_cache; +}; + +namespace model_package { + +/// Drop the materialized view cache. Call after any mutation that affects the +/// view tree. Safe on a cleared cache. +void DropViewCache(ModelPackage* pkg); + +/// Return the package's info view, building it lazily. +const InfoViewCache& BuildOrGetViewCache(const ModelPackage* pkg); + +/// Returns true iff `p` is `package_root` or lives under it (lexically). +bool IsInsidePackageRoot(const ModelPackage* pkg, const std::filesystem::path& p); + +/// Push the variant's resolved_directory onto the Prune candidates if it's +/// inside package_root. No-op if unresolved. +void RecordOrphanVariantDir(ModelPackage* pkg, const VariantRecord& v); + +/// Push every variant_dir of `c`, plus `c.component_dir` if external. +void RecordOrphanComponent(ModelPackage* pkg, const ComponentRecord& c); + +} // namespace model_package diff --git a/model_package/src/model_package_internal.h b/model_package/src/model_package_internal.h deleted file mode 100644 index 8d116b78f5880..0000000000000 --- a/model_package/src/model_package_internal.h +++ /dev/null @@ -1,80 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT License. - -/// \file model_package_internal.h -/// \brief Internal C++ types for the model package library. - -#pragma once - -#include -#include -#include -#include -#include - -namespace model_package { - -// ───────────────────────────────────────────────────────────────────────────── -// Data types -// ───────────────────────────────────────────────────────────────────────────── - -/// EP compatibility declaration for a variant (opaque to this library). -struct EpCompatibility { - std::optional ep; - std::optional device; - std::optional compatibility_string; -}; - -/// A single model file within a variant. -struct VariantFile { - std::string filename; - std::filesystem::path resolved_path; - - std::optional> session_options; - std::optional> provider_options; - std::optional> shared_files; -}; - -/// A variant of a component. -struct Variant { - std::string name; - std::filesystem::path folder_path; - // Single EP compatibility entry per variant (from metadata.json). - EpCompatibility ep_compatibility; - // Single model file entry (from variant.json). Empty when variant.json is absent. - std::optional file; - std::optional consumer_metadata_json; -}; - -/// A component in the model package. -struct Component { - std::string name; - std::vector variants; -}; - -/// Top-level model package descriptor. -struct PackageInfo { - int64_t schema_version{}; - std::filesystem::path root_path; - std::vector components; -}; - -// ───────────────────────────────────────────────────────────────────────────── -// Context implementation -// ───────────────────────────────────────────────────────────────────────────── - -/// Internal context holding parsed package data and C API caches. -struct ContextImpl { - PackageInfo package_info; - - // Caches for C API string access (stable pointers). - std::vector component_names_cache; - std::unordered_map> variant_names_cache; - std::unordered_map folder_path_strings_cache; - - // Lookup helpers. - const Component* FindComponent(const char* name) const; - const Variant* FindVariant(const char* component_name, const char* variant_name) const; -}; - -} // namespace model_package diff --git a/model_package/src/parser.cc b/model_package/src/parser.cc deleted file mode 100644 index 70d95b0297e38..0000000000000 --- a/model_package/src/parser.cc +++ /dev/null @@ -1,595 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT License. - -#include "parser.h" - -#include -#include -#include -#include -#include -#include - -#include "nlohmann/json.hpp" - -using json = nlohmann::json; - -namespace model_package { -namespace { - -// ───────────────────────────────────────────────────────────────────────────── -// JSON key constants -// ───────────────────────────────────────────────────────────────────────────── - -constexpr const char* kManifestFileName = "manifest.json"; -constexpr const char* kMetadataFileName = "metadata.json"; -constexpr const char* kVariantDescriptorFileName = "variant.json"; - -constexpr const char* kSchemaVersionKey = "schema_version"; -constexpr const char* kComponentsKey = "components"; -constexpr const char* kComponentNameKey = "component_name"; -constexpr const char* kVariantsKey = "variants"; - -constexpr const char* kEpKey = "ep"; -constexpr const char* kDeviceKey = "device"; -constexpr const char* kCompatibilityStringKey = "compatibility_string"; - -constexpr const char* kFilenameKey = "filename"; -constexpr const char* kSessionOptionsKey = "session_options"; -constexpr const char* kProviderOptionsKey = "provider_options"; -constexpr const char* kSharedFilesKey = "shared_files"; -constexpr const char* kConsumerMetadataKey = "consumer_metadata"; - -// ───────────────────────────────────────────────────────────────────────────── -// Internal schema types for deserialization -// ───────────────────────────────────────────────────────────────────────────── - -struct VariantMetadataSchema { - std::string filename; - std::optional> session_options; - std::optional> provider_options; - std::optional> shared_files; -}; - -struct EpCompatibilitySchema { - std::optional ep; - std::optional device; - std::optional compatibility_string; -}; - -struct VariantSchema { - EpCompatibilitySchema ep_info; -}; - -struct ComponentSchema { - std::optional component_name; - std::unordered_map variants; -}; - -struct ManifestSchema { - int64_t schema_version; - std::optional> components; -}; - -// ───────────────────────────────────────────────────────────────────────────── -// JSON helpers -// ───────────────────────────────────────────────────────────────────────────── - -std::string JsonScalarToString(const json& v, const char* key_name, const std::string& parent_key) { - if (v.is_string()) return v.get(); - if (v.is_number_integer()) return std::to_string(v.get()); - if (v.is_number_unsigned()) return std::to_string(v.get()); - if (v.is_number_float()) return v.dump(); - if (v.is_boolean()) return v.get() ? "true" : "false"; - - throw std::invalid_argument( - std::string("\"") + key_name + "\" under '" + parent_key + - "' must contain scalar (string/number/bool) values."); -} - -std::optional> ParseFlatOptionsObject( - const json& j, const char* key_name) { - if (!j.contains(key_name) || j[key_name].is_null()) { - return std::nullopt; - } - - const auto& obj = j[key_name]; - if (!obj.is_object()) { - throw std::invalid_argument(std::string("\"") + key_name + "\" must be an object."); - } - - std::unordered_map result; - result.reserve(obj.size()); - - for (auto it = obj.begin(); it != obj.end(); ++it) { - result.emplace(it.key(), JsonScalarToString(it.value(), key_name, it.key())); - } - - return result; -} - -std::optional ParseOptionalString(const json& j, const char* key_name) { - if (!j.contains(key_name) || j[key_name].is_null()) { - return std::nullopt; - } - - const auto& value = j[key_name]; - if (!value.is_string()) { - throw std::invalid_argument(std::string("\"") + key_name + "\" must be a string."); - } - return value.get(); -} - -// ───────────────────────────────────────────────────────────────────────────── -// nlohmann from_json overloads -// ───────────────────────────────────────────────────────────────────────────── - -void from_json(const json& j, EpCompatibilitySchema& c) { - if (!j.contains(kEpKey) || j[kEpKey].is_null()) { - throw std::invalid_argument(std::string("\"") + kEpKey + "\" is required in each ep_compatibility entry."); - } - if (!j[kEpKey].is_string()) { - throw std::invalid_argument(std::string("\"") + kEpKey + "\" must be a string."); - } - c.ep = j[kEpKey].get(); - if (c.ep->empty()) { - throw std::invalid_argument(std::string("\"") + kEpKey + "\" must be a non-empty string."); - } - - if (j.contains(kDeviceKey) && !j[kDeviceKey].is_null()) { - if (!j[kDeviceKey].is_string()) { - throw std::invalid_argument(std::string("\"") + kDeviceKey + "\" must be a string when present."); - } - c.device = j[kDeviceKey].get(); - } - c.compatibility_string = ParseOptionalString(j, kCompatibilityStringKey); -} - -void from_json(const json& j, VariantSchema& v) { - // EP fields (ep, device, compatibility_string) are now directly on the variant object. - // "ep" is required. - v.ep_info = j.get(); -} - -void from_json(const json& j, VariantMetadataSchema& v) { - v.filename = j.at(kFilenameKey).get(); - v.session_options = ParseFlatOptionsObject(j, kSessionOptionsKey); - v.provider_options = ParseFlatOptionsObject(j, kProviderOptionsKey); - v.shared_files = ParseFlatOptionsObject(j, kSharedFilesKey); -} - -void from_json(const json& j, ManifestSchema& m) { - m.schema_version = j.at(kSchemaVersionKey).get(); - - if (j.contains(kComponentsKey)) { - if (!j[kComponentsKey].is_array()) { - throw std::invalid_argument(std::string("\"") + kComponentsKey + "\" must be an array of strings"); - } - m.components = j[kComponentsKey].get>(); - } -} - -void from_json(const json& j, ComponentSchema& m) { - if (j.contains(kComponentNameKey) && j[kComponentNameKey].is_string()) { - m.component_name = j[kComponentNameKey].get(); - } - - m.variants = j.at(kVariantsKey).get>(); -} - -// ───────────────────────────────────────────────────────────────────────────── -// Parsing variants in declaration order (from the JSON object) -// ───────────────────────────────────────────────────────────────────────────── - -std::vector> ParseVariantsInOrder(const json& variants_obj) { - std::vector> result; - result.reserve(variants_obj.size()); - for (auto it = variants_obj.begin(); it != variants_obj.end(); ++it) { - result.emplace_back(it.key(), it.value().get()); - } - return result; -} - -// ───────────────────────────────────────────────────────────────────────────── -// Path validation -// ───────────────────────────────────────────────────────────────────────────── - -bool ValidatePathSegment(const std::string& segment, const char* segment_type, std::string& error) { - if (segment.empty()) { - error = std::string(segment_type) + " must not be empty."; - return false; - } - - if (std::filesystem::path(segment).is_absolute()) { - error = std::string(segment_type) + " must not be an absolute path: '" + segment + "'."; - return false; - } - - for (const auto& part : std::filesystem::path(segment)) { - if (part == "..") { - error = std::string(segment_type) + " must not contain '..' path components: '" + segment + "'."; - return false; - } - } - - return true; -} - -bool ValidatePathConfinement(const std::filesystem::path& resolved_path, - const std::filesystem::path& root, - const char* description, - std::string& error) { - auto normal_root = root.lexically_normal(); - auto normal_path = resolved_path.lexically_normal(); - - auto root_str = normal_root.string(); - auto path_str = normal_path.string(); - - if (path_str.size() < root_str.size() || - path_str.compare(0, root_str.size(), root_str) != 0 || - (path_str.size() > root_str.size() && path_str[root_str.size()] != std::filesystem::path::preferred_separator -#ifndef _WIN32 - && path_str[root_str.size()] != '/' -#endif - )) { - error = std::string(description) + " resolves outside the package root. Path: '" + - resolved_path.string() + "', Root: '" + root.string() + "'."; - return false; - } - - return true; -} - -// ───────────────────────────────────────────────────────────────────────────── -// Find single ONNX file in directory -// ───────────────────────────────────────────────────────────────────────────── - -bool FindSingleOnnxFile(const std::filesystem::path& search_dir, - std::filesystem::path& resolved_path, - std::string& error) { - std::vector onnx_files; - for (const auto& entry : std::filesystem::directory_iterator(search_dir)) { - if (!entry.is_regular_file()) continue; - - std::string ext = entry.path().extension().string(); - std::transform(ext.begin(), ext.end(), ext.begin(), - [](unsigned char c) { return static_cast(std::tolower(c)); }); - if (ext == ".onnx") { - onnx_files.push_back(entry.path()); - } - } - - if (onnx_files.empty()) { - error = "No ONNX model file found under " + search_dir.string(); - return false; - } - - if (onnx_files.size() > 1) { - error = "Multiple ONNX model files found under " + search_dir.string() + - ". Multiple ONNX files per variant are not supported yet."; - return false; - } - - resolved_path = onnx_files.front(); - return true; -} - -// ───────────────────────────────────────────────────────────────────────────── -// Parse variants from a single component -// ───────────────────────────────────────────────────────────────────────────── - -bool ParseVariantsFromComponent(const std::string& component_name, - const std::filesystem::path& component_root, - const json* variants_obj, - std::vector& out_variants, - std::string& error) { - if (variants_obj == nullptr) { - error = "Missing metadata variants for component: " + component_name; - return false; - } - - std::vector> variants; - try { - variants = ParseVariantsInOrder(*variants_obj); - } catch (const std::exception& ex) { - error = "Invalid metadata variant schema for component '" + component_name + "': " + ex.what(); - return false; - } - - for (const auto& [variant_name, variant_schema] : variants) { - if (!ValidatePathSegment(variant_name, "Variant name", error)) return false; - - const std::filesystem::path variant_root = component_root / variant_name; - if (!ValidatePathConfinement(variant_root, component_root, "Variant directory", error)) return false; - - const std::filesystem::path variant_descriptor_path = variant_root / kVariantDescriptorFileName; - - Variant variant_info{}; - variant_info.name = variant_name; - variant_info.folder_path = variant_root; - - // variant.json is optional. If present, it declares the file list, - // per-file session/provider options, and consumer metadata. - if (std::filesystem::exists(variant_descriptor_path)) { - std::ifstream vf(variant_descriptor_path, std::ios::binary); - if (!vf) { - error = "Failed to open variant.json at " + variant_descriptor_path.string(); - return false; - } - - json variant_doc; - try { - variant_doc = json::parse(vf); - } catch (const std::exception& ex) { - error = "variant.json at " + variant_descriptor_path.string() + " is not valid JSON: " + ex.what(); - return false; - } - - VariantMetadataSchema variant_metadata; - try { - variant_metadata = variant_doc.get(); - } catch (const std::exception& ex) { - error = "variant.json at " + variant_descriptor_path.string() + " has invalid schema: " + ex.what(); - return false; - } - - // consumer_metadata is a top-level optional field parsed separately from the schema struct. - if (variant_doc.contains(kConsumerMetadataKey) && variant_doc[kConsumerMetadataKey].is_object()) { - variant_info.consumer_metadata_json = variant_doc[kConsumerMetadataKey].dump(); - } - - if (!ValidatePathSegment(variant_metadata.filename, "File name", error)) return false; - - const std::filesystem::path candidate_path = variant_root / variant_metadata.filename; - if (!ValidatePathConfinement(candidate_path, variant_root, "Variant file path", error)) return false; - - if (!std::filesystem::exists(candidate_path)) { - error = "Variant '" + variant_name + "', file '" + variant_metadata.filename + - "' path does not exist: " + candidate_path.string(); - return false; - } - - std::filesystem::path resolved_model_path; - if (std::filesystem::is_regular_file(candidate_path)) { - resolved_model_path = candidate_path; - } else if (std::filesystem::is_directory(candidate_path)) { - if (!FindSingleOnnxFile(candidate_path, resolved_model_path, error)) return false; - } else { - error = "Variant '" + variant_name + "', file '" + variant_metadata.filename + - "' path is neither a file nor directory: " + candidate_path.string(); - return false; - } - - VariantFile file_info{}; - file_info.filename = variant_metadata.filename; - file_info.resolved_path = std::move(resolved_model_path); - file_info.session_options = variant_metadata.session_options; - file_info.provider_options = variant_metadata.provider_options; - file_info.shared_files = variant_metadata.shared_files; - - variant_info.file = std::move(file_info); - } - - // EP compatibility from metadata.json (single entry per variant) - variant_info.ep_compatibility.ep = variant_schema.ep_info.ep; - variant_info.ep_compatibility.device = variant_schema.ep_info.device; - variant_info.ep_compatibility.compatibility_string = variant_schema.ep_info.compatibility_string; - - out_variants.push_back(std::move(variant_info)); - } - - return true; -} - -} // namespace - -// ───────────────────────────────────────────────────────────────────────────── -// Public parser entry point -// ───────────────────────────────────────────────────────────────────────────── - -bool ParsePackage(const std::filesystem::path& package_root, - PackageInfo& out_package, - std::string& out_error) { - out_package = {}; - out_package.root_path = package_root; - - // Check for single-component mode: metadata.json at root - const auto root_metadata_path = package_root / kMetadataFileName; - if (std::filesystem::exists(root_metadata_path) && - std::filesystem::is_regular_file(root_metadata_path)) { - std::ifstream mf(root_metadata_path, std::ios::binary); - if (!mf) { - out_error = "Failed to open metadata.json at " + root_metadata_path.string(); - return false; - } - - json metadata_doc; - try { - metadata_doc = json::parse(mf); - } catch (const std::exception& ex) { - out_error = "metadata.json at " + root_metadata_path.string() + " is not valid JSON: " + ex.what(); - return false; - } - - ComponentSchema metadata_schema; - try { - metadata_schema = metadata_doc.get(); - } catch (const std::exception& ex) { - out_error = "metadata.json at " + root_metadata_path.string() + " has invalid schema: " + ex.what(); - return false; - } - - const std::string component_name = - metadata_schema.component_name.has_value() - ? *metadata_schema.component_name - : package_root.filename().string(); - - const json* variants_obj = &metadata_doc.at(kVariantsKey); - - Component component{}; - component.name = component_name; - - if (!ParseVariantsFromComponent(component_name, package_root, variants_obj, - component.variants, out_error)) { - return false; - } - - out_package.schema_version = 0; // Single-component mode doesn't have a manifest - out_package.components.push_back(std::move(component)); - return true; - } - - // Multi-component mode: manifest.json at root - const auto manifest_path = package_root / kManifestFileName; - if (!std::filesystem::exists(manifest_path)) { - out_error = "No manifest.json found at " + manifest_path.string(); - return false; - } - - std::ifstream f(manifest_path, std::ios::binary); - if (!f) { - out_error = "Failed to open manifest.json at " + manifest_path.string(); - return false; - } - - json doc; - try { - doc = json::parse(f); - } catch (const std::exception& ex) { - out_error = std::string("manifest.json is not valid JSON: ") + ex.what(); - return false; - } - - ManifestSchema manifest_schema; - try { - manifest_schema = doc.get(); - } catch (const std::exception& ex) { - out_error = std::string("manifest.json has invalid schema: ") + ex.what(); - return false; - } - - if (manifest_schema.schema_version != 1) { - out_error = "Unsupported schema_version in manifest.json: " + - std::to_string(manifest_schema.schema_version) + ". Expected 1."; - return false; - } - - out_package.schema_version = manifest_schema.schema_version; - - const bool has_components = manifest_schema.components.has_value(); - std::vector component_names; - std::unordered_map discovered_metadata_docs; - - if (has_components) { - component_names = *manifest_schema.components; - } else { - const auto models_dir = package_root / "models"; - if (!std::filesystem::exists(models_dir) || !std::filesystem::is_directory(models_dir)) { - out_error = "manifest.json missing \"components\" and no discoverable models directory at " + - models_dir.string(); - return false; - } - - for (const auto& entry : std::filesystem::directory_iterator(models_dir)) { - if (!entry.is_directory()) continue; - - const auto name = entry.path().filename().string(); - const auto metadata_path = entry.path() / kMetadataFileName; - if (!std::filesystem::exists(metadata_path)) continue; - - std::ifstream mf(metadata_path, std::ios::binary); - if (!mf) { - out_error = "Failed to open metadata.json at " + metadata_path.string(); - return false; - } - - json metadata_doc; - try { - metadata_doc = json::parse(mf); - (void)metadata_doc.get(); - } catch (const std::exception& ex) { - out_error = "metadata.json at " + metadata_path.string() + - " has invalid schema: " + std::string(ex.what()); - return false; - } - - discovered_metadata_docs.emplace(name, std::move(metadata_doc)); - component_names.push_back(name); - } - - if (component_names.empty()) { - out_error = - "manifest.json missing \"components\" and no component model folders with " - "metadata.json were found under " + - models_dir.string(); - return false; - } - } - - for (const auto& component_name : component_names) { - if (!ValidatePathSegment(component_name, "Component name", out_error)) return false; - - const auto component_root = package_root / "models" / component_name; - if (!ValidatePathConfinement(component_root, package_root, "Component directory", out_error)) return false; - - if (has_components && - (!std::filesystem::exists(component_root) || !std::filesystem::is_directory(component_root))) { - // Skip missing component directories (just warn — standalone library doesn't have logging, - // so we skip silently for now). - continue; - } - - json metadata_doc; - const json* variants_obj = nullptr; - const auto metadata_path = component_root / kMetadataFileName; - - if (!has_components) { - auto it_meta = discovered_metadata_docs.find(component_name); - if (it_meta != discovered_metadata_docs.end()) { - metadata_doc = it_meta->second; - variants_obj = &metadata_doc.at(kVariantsKey); - } - } else if (std::filesystem::exists(metadata_path)) { - std::ifstream mf(metadata_path, std::ios::binary); - if (mf) { - try { - metadata_doc = json::parse(mf); - (void)metadata_doc.get(); - variants_obj = &metadata_doc.at(kVariantsKey); - } catch (const std::exception&) { - // Ignore parse errors, fall through. - } - } - } - - if (!metadata_doc.is_null() && - metadata_doc.contains(kComponentNameKey) && - metadata_doc[kComponentNameKey].is_string()) { - const auto metadata_component_name = metadata_doc[kComponentNameKey].get(); - if (metadata_component_name != component_name) { - out_error = "metadata.json component_name '" + metadata_component_name + - "' does not match directory/manifest component name '" + component_name + "'."; - return false; - } - } - - Component component{}; - component.name = component_name; - - if (!ParseVariantsFromComponent(component_name, component_root, variants_obj, - component.variants, out_error)) { - return false; - } - - out_package.components.push_back(std::move(component)); - } - - if (out_package.components.empty()) { - out_error = "No valid component models were found under " + (package_root / "models").string(); - return false; - } - - return true; -} - -} // namespace model_package diff --git a/model_package/src/parser.h b/model_package/src/parser.h deleted file mode 100644 index ed3d22cb29d36..0000000000000 --- a/model_package/src/parser.h +++ /dev/null @@ -1,27 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT License. - -/// \file parser.h -/// \brief Model package JSON parser (internal). - -#pragma once - -#include -#include - -#include "model_package_internal.h" - -namespace model_package { - -/// Parse a model package from a directory. -/// Reads manifest.json, metadata.json per component, variant.json per variant. -/// -/// \param[in] package_root Path to the model package root directory. -/// \param[out] out_package On success, filled with the parsed package info. -/// \param[out] out_error On failure, filled with an error message. -/// \return true on success, false on error. -bool ParsePackage(const std::filesystem::path& package_root, - PackageInfo& out_package, - std::string& out_error); - -} // namespace model_package diff --git a/model_package/src/path_resolver.cc b/model_package/src/path_resolver.cc new file mode 100644 index 0000000000000..d62662d3ffcb1 --- /dev/null +++ b/model_package/src/path_resolver.cc @@ -0,0 +1,172 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "path_resolver.h" + +#include +#include +#include +#include + +#include "status_impl.h" + +namespace fs = std::filesystem; + +namespace model_package { + +namespace { + +bool IsHexLower(char c) { return (c >= '0' && c <= '9') || (c >= 'a' && c <= 'f'); } + +bool ContainsParentRefSegment(const fs::path& p) { + for (const auto& seg : p) { + if (seg == "..") return true; + } + return false; +} + +} // namespace + +bool IsSha256AssetUri(const std::string& uri) { + static constexpr const char* kPrefix = "sha256:"; + static constexpr size_t kPrefixLen = 7; + static constexpr size_t kHexLen = 64; + if (uri.size() != kPrefixLen + kHexLen) return false; + if (uri.compare(0, kPrefixLen, kPrefix) != 0) return false; + for (size_t i = kPrefixLen; i < uri.size(); ++i) { + if (!IsHexLower(uri[i])) return false; + } + return true; +} + +ModelPackageStatus* ResolvePath(const fs::path& base_dir, + const fs::path& package_root, + const std::string& input, + const PathResolverOptions& opts, + bool must_exist, + fs::path* out) { + if (!out) { + return model_package::MakeStatus(MODEL_PACKAGE_ERR_INVALID_ARG, + "ResolvePath: out must not be null."); + } + if (input.empty()) { + return model_package::MakeStatus(MODEL_PACKAGE_ERR_INVALID_ARG, + "ResolvePath: input must not be empty."); + } + + fs::path raw(input); + + if (!opts.allow_external_paths) { + if (raw.is_absolute() || raw.has_root_name()) { + return model_package::MakeStatus( + MODEL_PACKAGE_ERR_PATH_CONFINEMENT, + std::string("ResolvePath: absolute or drive-rooted path '") + input + + "' is not allowed in portable layout."); + } + if (ContainsParentRefSegment(raw)) { + return model_package::MakeStatus( + MODEL_PACKAGE_ERR_PATH_CONFINEMENT, + std::string("ResolvePath: '..' segments are not allowed in portable layout: '") + + input + "'."); + } + } + + fs::path joined = (raw.is_absolute() || raw.has_root_name()) ? raw : (base_dir / raw); + + std::error_code ec; + fs::path canonical; + bool exists_on_disk = fs::exists(joined, ec); + if (!exists_on_disk) { + if (must_exist) { + return model_package::MakeStatus( + MODEL_PACKAGE_ERR_NOT_FOUND, + std::string("ResolvePath: '") + joined.string() + "' does not exist."); + } + // Missing leaf (common during authoring/commit). When following symlinks, use + // weakly_canonical so any existing symlinks in the path prefix are still resolved; + // lexically_normal would leave a symlinked prefix unresolved and let it escape + // package_root undetected. Fall back to lexical normalization if that fails. + if (opts.follow_symlinks) { + canonical = fs::weakly_canonical(joined, ec); + if (ec) canonical = joined.lexically_normal(); + } else { + canonical = joined.lexically_normal(); + } + } else if (opts.follow_symlinks) { + canonical = fs::canonical(joined, ec); + if (ec) { + return model_package::MakeStatus( + MODEL_PACKAGE_ERR_IO, + std::string("ResolvePath: canonical('") + joined.string() + "') failed: " + ec.message()); + } + } else { + canonical = fs::weakly_canonical(joined, ec); + if (ec) { + canonical = joined.lexically_normal(); + } + } + + if (!opts.allow_external_paths && !package_root.empty()) { + // Confinement check: canonical must live under package_root's canonical form. This runs + // whether or not the leaf exists, so a not-yet-created path that resolves outside + // package_root (e.g. through a symlinked prefix) is still rejected. It is skipped when + // package_root is empty, which happens for in-memory authoring before a package has been + // anchored to a directory (there is no on-disk root to confine against yet); the + // absolute-path and ".." lexical checks above still apply in that case. + fs::path canonical_root = fs::weakly_canonical(package_root, ec); + if (ec) canonical_root = package_root.lexically_normal(); + + auto root_str = canonical_root.lexically_normal().string(); + auto can_str = canonical.lexically_normal().string(); + if (can_str.size() < root_str.size() || + can_str.compare(0, root_str.size(), root_str) != 0 || + (can_str.size() > root_str.size() && + can_str[root_str.size()] != fs::path::preferred_separator && + can_str[root_str.size()] != '/')) { + return model_package::MakeStatus( + MODEL_PACKAGE_ERR_PATH_CONFINEMENT, + std::string("ResolvePath: '") + can_str + + "' escapes package_root '" + root_str + "'."); + } + } + + *out = canonical; + return nullptr; +} + +bool TrySplitAssetUriPrefix(const std::string& input, std::string& uri, std::string& tail) { + static constexpr size_t kPrefixLen = 7; // "sha256:" + static constexpr size_t kHexLen = 64; + static constexpr size_t kUriLen = kPrefixLen + kHexLen; + if (input.size() < kUriLen) return false; + if (input.compare(0, kPrefixLen, "sha256:") != 0) return false; + for (size_t i = kPrefixLen; i < kUriLen; ++i) { + if (!IsHexLower(input[i])) return false; + } + if (input.size() == kUriLen) { + uri.assign(input); + tail.clear(); + return true; + } + if (input[kUriLen] != '/') return false; + uri.assign(input, 0, kUriLen); + tail.assign(input, kUriLen + 1, std::string::npos); + return true; +} + +std::string DefaultSharedAssetDirName(const std::string& uri) { + if (!IsSha256AssetUri(uri)) return {}; + return std::string(kSharedAssetOnDiskPrefix) + uri.substr(std::strlen("sha256:")); +} + +std::string SharedAssetUriFromDirName(const std::string& dir_name) { + const size_t prefix_len = std::strlen(kSharedAssetOnDiskPrefix); + if (dir_name.size() != prefix_len + 64) return {}; + if (dir_name.compare(0, prefix_len, kSharedAssetOnDiskPrefix) != 0) return {}; + for (size_t i = prefix_len; i < dir_name.size(); ++i) { + if (!IsHexLower(dir_name[i])) return {}; + } + return "sha256:" + dir_name.substr(prefix_len); +} + +} // namespace model_package diff --git a/model_package/src/path_resolver.h b/model_package/src/path_resolver.h new file mode 100644 index 0000000000000..f008897ff5bb0 --- /dev/null +++ b/model_package/src/path_resolver.h @@ -0,0 +1,60 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +/// \file path_resolver.h +/// \brief Path-resolution and confinement helpers. + +#pragma once + +#include +#include + +#include "model_package_api.h" // for ModelPackageStatus + +namespace model_package { + +struct PathResolverOptions { + bool allow_external_paths{false}; + bool follow_symlinks{true}; +}; + +/// Resolve a relative-or-absolute path string under a given base directory. +/// In portable mode (`allow_external_paths == false`): +/// - Reject absolute inputs (ERR_PATH_CONFINEMENT). +/// - Reject any path that, after canonicalization, escapes `package_root`. +/// - Reject `..` segments syntactically before resolution. +/// In installed mode: +/// - Absolute and `..` allowed. +/// - No confinement check. +/// +/// `must_exist` controls whether a missing target is an error (ERR_NOT_FOUND) +/// or whether the resolved (non-canonical) path is returned anyway. +/// Symlinks are followed when `follow_symlinks` is true. +ModelPackageStatus* ResolvePath(const std::filesystem::path& base_dir, + const std::filesystem::path& package_root, + const std::string& input, + const PathResolverOptions& opts, + bool must_exist, + std::filesystem::path* out); + +/// True if `uri` matches `^sha256:[0-9a-f]{64}$`. +bool IsSha256AssetUri(const std::string& uri); + +/// If `input` begins with a `sha256:` token followed by end-of-string or +/// '/', split into `uri` (the bare URI) and `tail` (substring after '/', or +/// empty). Returns true on a match, false otherwise. +bool TrySplitAssetUriPrefix(const std::string& input, std::string& uri, std::string& tail); + +/// Default on-disk directory name for a shared asset URI, i.e. the basename +/// under `/shared_assets/`. For `sha256:` this is +/// `sha256-`. Returns empty string if `uri` is not a valid sha256 URI. +std::string DefaultSharedAssetDirName(const std::string& uri); + +/// Inverse of `DefaultSharedAssetDirName`. If `dir_name` matches `sha256-` +/// returns the corresponding `sha256:` URI; otherwise returns empty string. +std::string SharedAssetUriFromDirName(const std::string& dir_name); + +/// Prefix shared by every default-convention shared-asset directory name. +constexpr const char* kSharedAssetOnDiskPrefix = "sha256-"; + +} // namespace model_package diff --git a/model_package/src/sha256.cc b/model_package/src/sha256.cc new file mode 100644 index 0000000000000..1ea26a555ad43 --- /dev/null +++ b/model_package/src/sha256.cc @@ -0,0 +1,229 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. +// +// Clean-room SHA-256 (FIPS 180-4) implementation. No external crypto deps. +// Intended for content-addressed asset hashing, not for cryptographic +// authentication. + +#include "sha256.h" + +#include +#include +#include +#include + +namespace model_package { + +namespace { + +constexpr uint32_t kInitState[8] = { + 0x6a09e667u, + 0xbb67ae85u, + 0x3c6ef372u, + 0xa54ff53au, + 0x510e527fu, + 0x9b05688cu, + 0x1f83d9abu, + 0x5be0cd19u, +}; + +constexpr uint32_t kRoundConstants[64] = { + 0x428a2f98u, + 0x71374491u, + 0xb5c0fbcfu, + 0xe9b5dba5u, + 0x3956c25bu, + 0x59f111f1u, + 0x923f82a4u, + 0xab1c5ed5u, + 0xd807aa98u, + 0x12835b01u, + 0x243185beu, + 0x550c7dc3u, + 0x72be5d74u, + 0x80deb1feu, + 0x9bdc06a7u, + 0xc19bf174u, + 0xe49b69c1u, + 0xefbe4786u, + 0x0fc19dc6u, + 0x240ca1ccu, + 0x2de92c6fu, + 0x4a7484aau, + 0x5cb0a9dcu, + 0x76f988dau, + 0x983e5152u, + 0xa831c66du, + 0xb00327c8u, + 0xbf597fc7u, + 0xc6e00bf3u, + 0xd5a79147u, + 0x06ca6351u, + 0x14292967u, + 0x27b70a85u, + 0x2e1b2138u, + 0x4d2c6dfcu, + 0x53380d13u, + 0x650a7354u, + 0x766a0abbu, + 0x81c2c92eu, + 0x92722c85u, + 0xa2bfe8a1u, + 0xa81a664bu, + 0xc24b8b70u, + 0xc76c51a3u, + 0xd192e819u, + 0xd6990624u, + 0xf40e3585u, + 0x106aa070u, + 0x19a4c116u, + 0x1e376c08u, + 0x2748774cu, + 0x34b0bcb5u, + 0x391c0cb3u, + 0x4ed8aa4au, + 0x5b9cca4fu, + 0x682e6ff3u, + 0x748f82eeu, + 0x78a5636fu, + 0x84c87814u, + 0x8cc70208u, + 0x90befffau, + 0xa4506cebu, + 0xbef9a3f7u, + 0xc67178f2u, +}; + +inline uint32_t Rotr(uint32_t x, int n) { return (x >> n) | (x << (32 - n)); } +inline uint32_t Ch(uint32_t x, uint32_t y, uint32_t z) { return (x & y) ^ (~x & z); } +inline uint32_t Maj(uint32_t x, uint32_t y, uint32_t z) { return (x & y) ^ (x & z) ^ (y & z); } +inline uint32_t Bsig0(uint32_t x) { return Rotr(x, 2) ^ Rotr(x, 13) ^ Rotr(x, 22); } +inline uint32_t Bsig1(uint32_t x) { return Rotr(x, 6) ^ Rotr(x, 11) ^ Rotr(x, 25); } +inline uint32_t Ssig0(uint32_t x) { return Rotr(x, 7) ^ Rotr(x, 18) ^ (x >> 3); } +inline uint32_t Ssig1(uint32_t x) { return Rotr(x, 17) ^ Rotr(x, 19) ^ (x >> 10); } + +} // namespace + +Sha256::Sha256() { + std::memcpy(state_, kInitState, sizeof(state_)); + bit_count_ = 0; + buffer_len_ = 0; +} + +void Sha256::Transform(const uint8_t block[64]) { + uint32_t w[64]; + for (int i = 0; i < 16; ++i) { + w[i] = (static_cast(block[i * 4]) << 24) | + (static_cast(block[i * 4 + 1]) << 16) | + (static_cast(block[i * 4 + 2]) << 8) | + (static_cast(block[i * 4 + 3])); + } + for (int i = 16; i < 64; ++i) { + w[i] = Ssig1(w[i - 2]) + w[i - 7] + Ssig0(w[i - 15]) + w[i - 16]; + } + + uint32_t a = state_[0], b = state_[1], c = state_[2], d = state_[3]; + uint32_t e = state_[4], f = state_[5], g = state_[6], h = state_[7]; + for (int i = 0; i < 64; ++i) { + uint32_t t1 = h + Bsig1(e) + Ch(e, f, g) + kRoundConstants[i] + w[i]; + uint32_t t2 = Bsig0(a) + Maj(a, b, c); + h = g; + g = f; + f = e; + e = d + t1; + d = c; + c = b; + b = a; + a = t1 + t2; + } + state_[0] += a; + state_[1] += b; + state_[2] += c; + state_[3] += d; + state_[4] += e; + state_[5] += f; + state_[6] += g; + state_[7] += h; +} + +void Sha256::Update(const void* data, size_t len) { + const uint8_t* p = static_cast(data); + bit_count_ += static_cast(len) * 8; + while (len > 0) { + size_t take = std::min(64 - buffer_len_, len); + std::memcpy(buffer_ + buffer_len_, p, take); + buffer_len_ += take; + p += take; + len -= take; + if (buffer_len_ == 64) { + Transform(buffer_); + buffer_len_ = 0; + } + } +} + +void Sha256::Final(uint8_t out[kDigestSize]) { + // Append 0x80, pad with zeros, append 64-bit big-endian length. + buffer_[buffer_len_++] = 0x80; + if (buffer_len_ > 56) { + std::memset(buffer_ + buffer_len_, 0, 64 - buffer_len_); + Transform(buffer_); + buffer_len_ = 0; + } + std::memset(buffer_ + buffer_len_, 0, 56 - buffer_len_); + uint64_t bc = bit_count_; + for (int i = 7; i >= 0; --i) { + buffer_[56 + i] = static_cast(bc & 0xff); + bc >>= 8; + } + Transform(buffer_); + for (int i = 0; i < 8; ++i) { + out[i * 4] = static_cast((state_[i] >> 24) & 0xff); + out[i * 4 + 1] = static_cast((state_[i] >> 16) & 0xff); + out[i * 4 + 2] = static_cast((state_[i] >> 8) & 0xff); + out[i * 4 + 3] = static_cast(state_[i] & 0xff); + } +} + +namespace { +constexpr char kHex[] = "0123456789abcdef"; +std::string ToHex(const uint8_t* bytes, size_t len) { + std::string s(len * 2, '0'); + for (size_t i = 0; i < len; ++i) { + s[i * 2] = kHex[(bytes[i] >> 4) & 0x0f]; + s[i * 2 + 1] = kHex[bytes[i] & 0x0f]; + } + return s; +} +} // namespace + +std::string Sha256::FinalHex() { + uint8_t out[kDigestSize]; + Final(out); + return ToHex(out, kDigestSize); +} + +std::string Sha256::HashBytesHex(const void* data, size_t len) { + Sha256 h; + h.Update(data, len); + return h.FinalHex(); +} + +std::string Sha256::HashStringHex(const std::string& s) { + return HashBytesHex(s.data(), s.size()); +} + +std::string Sha256::HashFileHex(const std::string& path) { + std::ifstream f(path, std::ios::binary); + if (!f) return std::string(); + Sha256 h; + char buf[8192]; + while (f) { + f.read(buf, sizeof(buf)); + std::streamsize n = f.gcount(); + if (n > 0) h.Update(buf, static_cast(n)); + } + return h.FinalHex(); +} + +} // namespace model_package diff --git a/model_package/src/sha256.h b/model_package/src/sha256.h new file mode 100644 index 0000000000000..da4125ecd80b0 --- /dev/null +++ b/model_package/src/sha256.h @@ -0,0 +1,44 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +/// \file sha256.h +/// \brief Minimal SHA-256 implementation used for content-addressed assets. +/// No external crypto dependency. + +#pragma once + +#include +#include +#include +#include + +namespace model_package { + +class Sha256 { + public: + static constexpr size_t kDigestSize = 32; + + Sha256(); + void Update(const void* data, size_t len); + void Update(const std::string& s) { Update(s.data(), s.size()); } + void Final(uint8_t out[kDigestSize]); + + /// Hex-encoded (lowercase) digest, 64 chars. + std::string FinalHex(); + + static std::string HashBytesHex(const void* data, size_t len); + static std::string HashStringHex(const std::string& s); + + /// Stream-hash a file by path. Returns the hex digest, or empty string on + /// IO error (caller should pre-check existence). + static std::string HashFileHex(const std::string& path); + + private: + void Transform(const uint8_t block[64]); + uint32_t state_[8]; + uint64_t bit_count_; + uint8_t buffer_[64]; + size_t buffer_len_; +}; + +} // namespace model_package diff --git a/model_package/src/status_impl.h b/model_package/src/status_impl.h new file mode 100644 index 0000000000000..6cc1c94238f98 --- /dev/null +++ b/model_package/src/status_impl.h @@ -0,0 +1,30 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +/// \file status_impl.h +/// \brief Internal representation of ModelPackageStatus, shared by all +/// implementation units in the model_package library. + +#pragma once + +#include +#include +#include + +#include "model_package_api.h" + +struct ModelPackageStatus { + ModelPackageErrorCode code{MODEL_PACKAGE_ERR_INVALID_ARG}; + std::string message; +}; + +namespace model_package { + +/// Allocate a new failure status. Returns nullptr if allocation fails (callers +/// should treat that as a generic error; we deliberately never throw out of the +/// C API). +inline ModelPackageStatus* MakeStatus(ModelPackageErrorCode code, std::string message) { + return new (std::nothrow) ModelPackageStatus{code, std::move(message)}; +} + +} // namespace model_package diff --git a/model_package/tests/test_asset_hashing.cc b/model_package/tests/test_asset_hashing.cc new file mode 100644 index 0000000000000..717c3fefea4b6 --- /dev/null +++ b/model_package/tests/test_asset_hashing.cc @@ -0,0 +1,312 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +/// \file test_asset_hashing.cc +/// \brief Tests for the directory Merkle hash and SHA-256 implementation. + +#include "model_package.h" +#include "model_package_api.h" +#include "sha256.h" + +#include +#include +#include +#include +#include +#include +#include + +namespace fs = std::filesystem; +using model_package::Sha256; + +namespace { + +int g_failed = 0; +int g_passed = 0; +const char* g_current = ""; + +#define CHECK(cond) \ + do { \ + if (!(cond)) { \ + std::fprintf(stderr, "[FAIL] %s line %d: CHECK(%s)\n", g_current, __LINE__, #cond); \ + return false; \ + } \ + } while (0) + +#define CHECK_OK(status) \ + do { \ + ModelPackageStatus* _s = (status); \ + if (_s != nullptr) { \ + std::fprintf(stderr, "[FAIL] %s line %d: expected OK, got: %s\n", \ + g_current, __LINE__, ModelPackageStatus_Message(_s)); \ + ModelPackageStatus_Release(_s); \ + return false; \ + } \ + } while (0) + +class Sandbox { + public: + Sandbox() { + std::random_device rd; + std::mt19937_64 g(rd()); + char buf[32]; + std::snprintf(buf, sizeof(buf), "mp_hash_%016llx", static_cast(g())); + root_ = fs::temp_directory_path() / buf; + fs::create_directories(root_); + } + ~Sandbox() { + std::error_code ec; + fs::remove_all(root_, ec); + } + Sandbox(const Sandbox&) = delete; + Sandbox& operator=(const Sandbox&) = delete; + const fs::path& root() const { return root_; } + void Write(const std::string& relpath, const std::string& contents) { + fs::path full = root_ / relpath; + fs::create_directories(full.parent_path()); + std::ofstream f(full, std::ios::binary); + f << contents; + } + + private: + fs::path root_; +}; + +// FIPS-180-4 known-answer test vectors. +bool test_sha256_known_vectors() { + CHECK(Sha256::HashStringHex("") == + "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855"); + CHECK(Sha256::HashStringHex("abc") == + "ba7816bf8f01cfea414140de5dae2223b00361a396177a9cb410ff61f20015ad"); + // Long message: 1,000,000 'a' characters. + std::string a_million(1000000, 'a'); + CHECK(Sha256::HashStringHex(a_million) == + "cdc76e5c9914fb9281a1c7e284d73e67f1809a48a497200e046d39ccc7112cd0"); + return true; +} + +bool test_sha256_incremental_matches_oneshot() { + std::string msg = "the quick brown fox jumps over the lazy dog"; + std::string oneshot = Sha256::HashStringHex(msg); + Sha256 h; + for (char c : msg) h.Update(&c, 1); + CHECK(h.FinalHex() == oneshot); + return true; +} + +bool test_directory_hash_basic() { + Sandbox s; + s.Write("a.txt", "alpha"); + s.Write("b.txt", "beta"); + + const char* uri = nullptr; + CHECK_OK(ModelPackage_ComputeDirectoryHash(s.root().c_str(), &uri)); + CHECK(uri != nullptr); + std::string u(uri); + CHECK(u.substr(0, 7) == "sha256:"); + CHECK(u.size() == 7 + 64); + return true; +} + +bool test_directory_hash_reproducible() { + Sandbox s1; + s1.Write("a.txt", "alpha"); + s1.Write("nested/b.txt", "beta"); + + Sandbox s2; + s2.Write("a.txt", "alpha"); + s2.Write("nested/b.txt", "beta"); + + const char* u1 = nullptr; + CHECK_OK(ModelPackage_ComputeDirectoryHash(s1.root().c_str(), &u1)); + std::string copy1(u1); + + const char* u2 = nullptr; + CHECK_OK(ModelPackage_ComputeDirectoryHash(s2.root().c_str(), &u2)); + CHECK(copy1 == std::string(u2)); + return true; +} + +bool test_directory_hash_name_change_differs() { + Sandbox s1; + s1.Write("a.txt", "alpha"); + + Sandbox s2; + s2.Write("b.txt", "alpha"); // same content, different name + + const char* u1 = nullptr; + const char* u2 = nullptr; + CHECK_OK(ModelPackage_ComputeDirectoryHash(s1.root().c_str(), &u1)); + std::string copy1(u1); + CHECK_OK(ModelPackage_ComputeDirectoryHash(s2.root().c_str(), &u2)); + CHECK(copy1 != std::string(u2)); + return true; +} + +bool test_directory_hash_swapped_names_differ() { + Sandbox s1; + s1.Write("a.txt", "alpha"); + s1.Write("b.txt", "beta"); + + Sandbox s2; + s2.Write("a.txt", "beta"); // swapped contents + s2.Write("b.txt", "alpha"); + + const char* u1 = nullptr; + const char* u2 = nullptr; + CHECK_OK(ModelPackage_ComputeDirectoryHash(s1.root().c_str(), &u1)); + std::string copy1(u1); + CHECK_OK(ModelPackage_ComputeDirectoryHash(s2.root().c_str(), &u2)); + CHECK(copy1 != std::string(u2)); + return true; +} + +bool test_directory_hash_content_change_differs() { + Sandbox s1; + s1.Write("a.txt", "alpha"); + Sandbox s2; + s2.Write("a.txt", "ALPHA"); + + const char* u1 = nullptr; + const char* u2 = nullptr; + CHECK_OK(ModelPackage_ComputeDirectoryHash(s1.root().c_str(), &u1)); + std::string copy1(u1); + CHECK_OK(ModelPackage_ComputeDirectoryHash(s2.root().c_str(), &u2)); + CHECK(copy1 != std::string(u2)); + return true; +} + +bool test_directory_hash_empty_dirs_ignored() { + Sandbox s1; + s1.Write("a.txt", "alpha"); + Sandbox s2; + s2.Write("a.txt", "alpha"); + fs::create_directories(s2.root() / "empty_subdir"); + + const char* u1 = nullptr; + const char* u2 = nullptr; + CHECK_OK(ModelPackage_ComputeDirectoryHash(s1.root().c_str(), &u1)); + std::string copy1(u1); + CHECK_OK(ModelPackage_ComputeDirectoryHash(s2.root().c_str(), &u2)); + CHECK(copy1 == std::string(u2)); + return true; +} + +bool test_directory_hash_rejects_symlink() { + Sandbox s; + s.Write("a.txt", "alpha"); + std::error_code ec; + fs::create_symlink("a.txt", s.root() / "a_link.txt", ec); + // If symlink creation isn't supported on this filesystem, skip the test + // (treat as pass — the rejection is the behavior under test). + if (ec) { + std::printf("[SKIP] %s (symlink unsupported)\n", g_current); + return true; + } + const char* uri = nullptr; + ModelPackageStatus* st = ModelPackage_ComputeDirectoryHash(s.root().c_str(), &uri); + CHECK(st != nullptr); + CHECK(ModelPackageStatus_Code(st) == MODEL_PACKAGE_ERR_SCHEMA); + ModelPackageStatus_Release(st); + return true; +} + +bool test_directory_hash_known_value_single_file() { + // Known-answer check: the directory URI hashes a manifest of " \n" + // lines, so compute the expected value the same way and compare. + Sandbox s; + s.Write("a.txt", "alpha"); + + std::string file_hex = Sha256::HashStringHex("alpha"); + std::string manifest = file_hex + " a.txt\n"; + std::string expected = "sha256:" + Sha256::HashStringHex(manifest); + + const char* uri = nullptr; + CHECK_OK(ModelPackage_ComputeDirectoryHash(s.root().c_str(), &uri)); + CHECK(std::string(uri) == expected); + return true; +} + +bool test_directory_hash_sorted_order_independent_of_walk() { + // Whether the OS walks "b.txt" before "a.txt" must not matter. + Sandbox s; + s.Write("a.txt", "alpha"); + s.Write("b.txt", "beta"); + s.Write("c.txt", "gamma"); + + // Compute expected manifest manually (sorted). + std::string hex_a = Sha256::HashStringHex("alpha"); + std::string hex_b = Sha256::HashStringHex("beta"); + std::string hex_c = Sha256::HashStringHex("gamma"); + std::string manifest = hex_a + " a.txt\n" + + hex_b + " b.txt\n" + + hex_c + " c.txt\n"; + std::string expected = "sha256:" + Sha256::HashStringHex(manifest); + + const char* uri = nullptr; + CHECK_OK(ModelPackage_ComputeDirectoryHash(s.root().c_str(), &uri)); + CHECK(std::string(uri) == expected); + return true; +} + +bool test_directory_hash_uses_forward_slash() { + Sandbox s; + s.Write("dir/sub/c.txt", "x"); + + std::string file_hex = Sha256::HashStringHex("x"); + // Path must be POSIX style in the manifest (forward slashes). + std::string manifest = file_hex + " dir/sub/c.txt\n"; + std::string expected = "sha256:" + Sha256::HashStringHex(manifest); + + const char* uri = nullptr; + CHECK_OK(ModelPackage_ComputeDirectoryHash(s.root().c_str(), &uri)); + CHECK(std::string(uri) == expected); + return true; +} + +bool test_missing_directory_errors() { + const char* uri = nullptr; + ModelPackageStatus* s = ModelPackage_ComputeDirectoryHash("/tmp/does_not_exist_xyzzy_zzz", &uri); + CHECK(s != nullptr); + CHECK(ModelPackageStatus_Code(s) == MODEL_PACKAGE_ERR_NOT_FOUND); + ModelPackageStatus_Release(s); + return true; +} + +struct Test { + const char* name; + bool (*fn)(); +}; + +const Test kTests[] = { + {"sha256_known_vectors", test_sha256_known_vectors}, + {"sha256_incremental_matches_oneshot", test_sha256_incremental_matches_oneshot}, + {"directory_hash_basic", test_directory_hash_basic}, + {"directory_hash_reproducible", test_directory_hash_reproducible}, + {"directory_hash_name_change_differs", test_directory_hash_name_change_differs}, + {"directory_hash_swapped_names_differ", test_directory_hash_swapped_names_differ}, + {"directory_hash_content_change_differs", test_directory_hash_content_change_differs}, + {"directory_hash_empty_dirs_ignored", test_directory_hash_empty_dirs_ignored}, + {"directory_hash_rejects_symlink", test_directory_hash_rejects_symlink}, + {"directory_hash_known_value_single_file", test_directory_hash_known_value_single_file}, + {"directory_hash_sorted_order_independent_of_walk", test_directory_hash_sorted_order_independent_of_walk}, + {"directory_hash_uses_forward_slash", test_directory_hash_uses_forward_slash}, + {"missing_directory_errors", test_missing_directory_errors}, +}; + +} // namespace + +int main() { + for (const auto& t : kTests) { + g_current = t.name; + bool ok = t.fn(); + if (ok) { + std::printf("[PASS] %s\n", t.name); + g_passed++; + } else { + g_failed++; + } + } + std::printf("\n=== %d passed, %d failed ===\n", g_passed, g_failed); + return g_failed == 0 ? 0 : 1; +} diff --git a/model_package/tests/test_authoring.cc b/model_package/tests/test_authoring.cc new file mode 100644 index 0000000000000..4f6808d966093 --- /dev/null +++ b/model_package/tests/test_authoring.cc @@ -0,0 +1,525 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +/// \file test_authoring.cc +/// \brief Authoring (mutation) API tests. + +#include "model_package.h" +#include "model_package_api.h" + +#include +#include +#include +#include +#include +#include + +namespace fs = std::filesystem; + +namespace { + +int g_failed = 0; +int g_passed = 0; +const char* g_current = ""; + +#define CHECK(cond) \ + do { \ + if (!(cond)) { \ + std::fprintf(stderr, "[FAIL] %s line %d: CHECK(%s)\n", g_current, __LINE__, #cond); \ + return false; \ + } \ + } while (0) + +#define CHECK_OK(status) \ + do { \ + ModelPackageStatus* _s = (status); \ + if (_s != nullptr) { \ + std::fprintf(stderr, "[FAIL] %s line %d: expected OK, got: %s\n", \ + g_current, __LINE__, ModelPackageStatus_Message(_s)); \ + ModelPackageStatus_Release(_s); \ + return false; \ + } \ + } while (0) + +#define CHECK_ERR(status, expected_code) \ + do { \ + ModelPackageStatus* _s = (status); \ + if (_s == nullptr) { \ + std::fprintf(stderr, "[FAIL] %s line %d: expected error %d, got OK\n", \ + g_current, __LINE__, (int)(expected_code)); \ + return false; \ + } \ + ModelPackageErrorCode _c = ModelPackageStatus_Code(_s); \ + if (_c != (expected_code)) { \ + std::fprintf(stderr, "[FAIL] %s line %d: expected error %d, got %d (%s)\n", \ + g_current, __LINE__, (int)(expected_code), (int)_c, \ + ModelPackageStatus_Message(_s)); \ + ModelPackageStatus_Release(_s); \ + return false; \ + } \ + ModelPackageStatus_Release(_s); \ + } while (0) + +class Sandbox { + public: + Sandbox() { + std::random_device rd; + std::mt19937_64 g(rd()); + char buf[32]; + std::snprintf(buf, sizeof(buf), "mp_auth_%016llx", static_cast(g())); + root_ = fs::temp_directory_path() / buf; + fs::create_directories(root_); + } + ~Sandbox() { + std::error_code ec; + fs::remove_all(root_, ec); + } + Sandbox(const Sandbox&) = delete; + Sandbox& operator=(const Sandbox&) = delete; + const fs::path& root() const { return root_; } + fs::path path(const std::string& rel) const { return root_ / rel; } + void Write(const std::string& rel, const std::string& contents) { + fs::path full = root_ / rel; + fs::create_directories(full.parent_path()); + std::ofstream f(full, std::ios::binary); + f << contents; + } + + private: + fs::path root_; +}; + +class PkgHandle { + public: + explicit PkgHandle(ModelPackage* p) : p_(p) {} + ~PkgHandle() { ModelPackage_Close(p_); } + PkgHandle(const PkgHandle&) = delete; + PkgHandle& operator=(const PkgHandle&) = delete; + ModelPackage* get() const { return p_; } + + private: + ModelPackage* p_; +}; + +// ───────────────────────────────────────────────────────────────────────────── +// ModelPackage_New +// ───────────────────────────────────────────────────────────────────────────── + +bool test_new_creates_empty_package() { + ModelPackage* raw = nullptr; + CHECK_OK(ModelPackage_New(&raw)); + CHECK(raw != nullptr); + PkgHandle p(raw); + const ModelPackageInfo* info = ModelPackage_Info(p.get()); + CHECK(info != nullptr); + CHECK(info->schema_version_major == 0); + CHECK(info->schema_version_minor == 0); + CHECK((info)->num_components == 0); + CHECK((info)->num_shared_assets == 0); + CHECK(std::string(info->layout) == "portable"); + return true; +} + +// ───────────────────────────────────────────────────────────────────────────── +// Component operations +// ───────────────────────────────────────────────────────────────────────────── + +bool test_set_component_inline_basic() { + ModelPackage* raw = nullptr; + CHECK_OK(ModelPackage_New(&raw)); + PkgHandle p(raw); + + CHECK_OK(ModelPackage_SetComponentInline(p.get(), "encoder", + R"({"variants": {}})")); + CHECK((ModelPackage_Info(p.get()))->num_components == 1); + const ModelComponentInfo* c = ModelPackage_FindComponent(ModelPackage_Info(p.get()), "encoder"); + CHECK(c != nullptr); + CHECK(std::string(c->name) == "encoder"); + CHECK((c)->num_variants == 0); + return true; +} + +bool test_set_component_inline_replaces_existing() { + ModelPackage* raw = nullptr; + CHECK_OK(ModelPackage_New(&raw)); + PkgHandle p(raw); + + CHECK_OK(ModelPackage_SetComponentInline(p.get(), "c", R"({"variants": {}})")); + CHECK_OK(ModelPackage_SetComponentInline(p.get(), "c", + R"({"variants": {"v1": {"variant_directory": "."}}})")); + CHECK((ModelPackage_Info(p.get()))->num_components == 1); + const ModelComponentInfo* c = ModelPackage_FindComponent(ModelPackage_Info(p.get()), "c"); + CHECK((c)->num_variants == 1); + return true; +} + +bool test_set_component_inline_rejects_unknown_field() { + ModelPackage* raw = nullptr; + CHECK_OK(ModelPackage_New(&raw)); + PkgHandle p(raw); + CHECK_ERR(ModelPackage_SetComponentInline(p.get(), "c", + R"({"variants": {}, "typo_field": 1})"), + MODEL_PACKAGE_ERR_SCHEMA); + CHECK((ModelPackage_Info(p.get()))->num_components == 0); + return true; +} + +bool test_set_component_inline_rejects_bad_json() { + ModelPackage* raw = nullptr; + CHECK_OK(ModelPackage_New(&raw)); + PkgHandle p(raw); + CHECK_ERR(ModelPackage_SetComponentInline(p.get(), "c", "not-json"), + MODEL_PACKAGE_ERR_SCHEMA); + return true; +} + +bool test_remove_component() { + ModelPackage* raw = nullptr; + CHECK_OK(ModelPackage_New(&raw)); + PkgHandle p(raw); + CHECK_OK(ModelPackage_SetComponentInline(p.get(), "a", R"({"variants": {}})")); + CHECK_OK(ModelPackage_SetComponentInline(p.get(), "b", R"({"variants": {}})")); + CHECK((ModelPackage_Info(p.get()))->num_components == 2); + CHECK_OK(ModelPackage_RemoveComponent(p.get(), "a")); + CHECK((ModelPackage_Info(p.get()))->num_components == 1); + const ModelPackageInfo* info = ModelPackage_Info(p.get()); + CHECK(ModelPackage_FindComponent(info, "a") == nullptr); + CHECK(ModelPackage_FindComponent(info, "b") != nullptr); + return true; +} + +bool test_remove_missing_component_is_noop() { + ModelPackage* raw = nullptr; + CHECK_OK(ModelPackage_New(&raw)); + PkgHandle p(raw); + CHECK_OK(ModelPackage_RemoveComponent(p.get(), "nope")); + return true; +} + +// ───────────────────────────────────────────────────────────────────────────── +// Variant operations +// ───────────────────────────────────────────────────────────────────────────── + +bool test_set_variant_upsert() { + ModelPackage* raw = nullptr; + CHECK_OK(ModelPackage_New(&raw)); + PkgHandle p(raw); + CHECK_OK(ModelPackage_SetComponentInline(p.get(), "c", R"({"variants": {}})")); + + CHECK_OK(ModelPackage_SetVariant(p.get(), "c", "v1", + R"({"variant_directory": ".", "ep": "CPU"})")); + const ModelComponentInfo* c = ModelPackage_FindComponent(ModelPackage_Info(p.get()), "c"); + CHECK((c)->num_variants == 1); + const ModelVariantInfo* v = ModelComponentInfo_FindVariant(c, "v1"); + CHECK(v != nullptr); + CHECK(std::string(v->ep) == "CPU"); + + // Upsert: change ep. + CHECK_OK(ModelPackage_SetVariant(p.get(), "c", "v1", + R"({"variant_directory": ".", "ep": "CUDA"})")); + c = ModelPackage_FindComponent(ModelPackage_Info(p.get()), "c"); + CHECK((c)->num_variants == 1); + v = ModelComponentInfo_FindVariant(c, "v1"); + CHECK(std::string(v->ep) == "CUDA"); + return true; +} + +bool test_set_variant_unknown_component_errors() { + ModelPackage* raw = nullptr; + CHECK_OK(ModelPackage_New(&raw)); + PkgHandle p(raw); + CHECK_ERR(ModelPackage_SetVariant(p.get(), "nope", "v1", R"({"variant_directory": "."})"), + MODEL_PACKAGE_ERR_NOT_FOUND); + return true; +} + +bool test_remove_variant() { + ModelPackage* raw = nullptr; + CHECK_OK(ModelPackage_New(&raw)); + PkgHandle p(raw); + CHECK_OK(ModelPackage_SetComponentInline(p.get(), "c", R"({"variants": {}})")); + CHECK_OK(ModelPackage_SetVariant(p.get(), "c", "v1", R"({"variant_directory": "."})")); + CHECK_OK(ModelPackage_RemoveVariant(p.get(), "c", "v1")); + const ModelComponentInfo* c = ModelPackage_FindComponent(ModelPackage_Info(p.get()), "c"); + CHECK((c)->num_variants == 0); + return true; +} + +// ───────────────────────────────────────────────────────────────────────────── +// Variant executor_info +// ───────────────────────────────────────────────────────────────────────────── + +bool test_set_executor_info_inline_and_remove() { + ModelPackage* raw = nullptr; + CHECK_OK(ModelPackage_New(&raw)); + PkgHandle p(raw); + CHECK_OK(ModelPackage_SetComponentInline(p.get(), "c", R"({"variants": {}})")); + CHECK_OK(ModelPackage_SetVariant(p.get(), "c", "v1", R"({"variant_directory": "."})")); + + CHECK_OK(ModelPackage_SetVariantExecutorInfoInline(p.get(), "c", "v1", "ort", + R"({"model": "m.onnx"})")); + const ModelVariantInfo* v = ModelComponentInfo_FindVariant( + ModelPackage_FindComponent(ModelPackage_Info(p.get()), "c"), "v1"); + const char* ej = nullptr; + const ModelExecutorInfoEntry* ei = ModelVariantInfo_FindExecutorInfo(v, "ort"); + ej = ei ? ei->json : nullptr; + CHECK(ej != nullptr); + CHECK(std::strstr(ej, "\"model\"") != nullptr); + + CHECK_OK(ModelPackage_RemoveVariantExecutorInfo(p.get(), "c", "v1", "ort")); + v = ModelComponentInfo_FindVariant(ModelPackage_FindComponent(ModelPackage_Info(p.get()), "c"), "v1"); + ei = ModelVariantInfo_FindExecutorInfo(v, "ort"); + ej = ei ? ei->json : nullptr; + CHECK(ei == nullptr); + CHECK(ej == nullptr); + return true; +} + +bool test_set_executor_info_external_records_path() { + ModelPackage* raw = nullptr; + CHECK_OK(ModelPackage_New(&raw)); + PkgHandle p(raw); + CHECK_OK(ModelPackage_SetComponentInline(p.get(), "c", R"({"variants": {}})")); + CHECK_OK(ModelPackage_SetVariant(p.get(), "c", "v1", R"({"variant_directory": "."})")); + CHECK_OK(ModelPackage_SetVariantExecutorInfoExternal(p.get(), "c", "v1", "ort", + "ort_info.json")); + return true; +} + +// ───────────────────────────────────────────────────────────────────────────── +// Package metadata +// ───────────────────────────────────────────────────────────────────────────── + +bool test_set_metadata() { + ModelPackage* raw = nullptr; + CHECK_OK(ModelPackage_New(&raw)); + PkgHandle p(raw); + CHECK_OK(ModelPackage_SetMetadata(p.get(), "mypkg", "1.0.0", "desc")); + const ModelPackageInfo* info = ModelPackage_Info(p.get()); + CHECK(std::string(info->package_name) == "mypkg"); + CHECK(std::string(info->package_version) == "1.0.0"); + CHECK(std::string(info->description) == "desc"); + + // Empty string clears. + CHECK_OK(ModelPackage_SetMetadata(p.get(), nullptr, "", nullptr)); + info = ModelPackage_Info(p.get()); + CHECK(info->package_version == nullptr); + CHECK(std::string(info->package_name) == "mypkg"); + return true; +} + +bool test_set_layout() { + ModelPackage* raw = nullptr; + CHECK_OK(ModelPackage_New(&raw)); + PkgHandle p(raw); + CHECK_OK(ModelPackage_SetLayout(p.get(), "installed")); + CHECK(std::string(ModelPackage_Info(p.get())->layout) == "installed"); + CHECK_ERR(ModelPackage_SetLayout(p.get(), "weird"), MODEL_PACKAGE_ERR_SCHEMA); + return true; +} + +bool test_set_additional_metadata_manifest_scope() { + ModelPackage* raw = nullptr; + CHECK_OK(ModelPackage_New(&raw)); + PkgHandle p(raw); + CHECK_OK(ModelPackage_SetAdditionalMetadataJson(p.get(), "manifest", nullptr, nullptr, + R"({"author":"jambayk"})")); + const ModelPackageInfo* info = ModelPackage_Info(p.get()); + CHECK(info->additional_metadata_json != nullptr); + CHECK(std::string(info->additional_metadata_json).find("jambayk") != std::string::npos); + + // Clear. + CHECK_OK(ModelPackage_SetAdditionalMetadataJson(p.get(), "manifest", nullptr, nullptr, nullptr)); + info = ModelPackage_Info(p.get()); + CHECK(info->additional_metadata_json == nullptr); + return true; +} + +bool test_set_additional_metadata_variant_scope() { + ModelPackage* raw = nullptr; + CHECK_OK(ModelPackage_New(&raw)); + PkgHandle p(raw); + CHECK_OK(ModelPackage_SetComponentInline(p.get(), "c", R"({"variants": {}})")); + CHECK_OK(ModelPackage_SetVariant(p.get(), "c", "v1", R"({"variant_directory": "."})")); + CHECK_OK(ModelPackage_SetAdditionalMetadataJson(p.get(), "variant", "c", "v1", + R"({"foo":"bar"})")); + const ModelVariantInfo* v = ModelComponentInfo_FindVariant( + ModelPackage_FindComponent(ModelPackage_Info(p.get()), "c"), "v1"); + CHECK(v != nullptr); + const char* md = v->additional_metadata_json; + CHECK(md != nullptr); + CHECK(std::string(md).find("foo") != std::string::npos); + return true; +} + +// ───────────────────────────────────────────────────────────────────────────── +// Shared assets — authoring +// ───────────────────────────────────────────────────────────────────────────── + +bool test_add_shared_asset_copy_in_true_portable_ok() { + Sandbox s; + s.Write("src/a.txt", "alpha"); + + ModelPackage* raw = nullptr; + CHECK_OK(ModelPackage_New(&raw)); + PkgHandle p(raw); + const char* uri = nullptr; + CHECK_OK(ModelPackage_AddSharedAsset(p.get(), (s.root() / "src").c_str(), + nullptr, /*copy_in=*/true, &uri)); + CHECK(uri != nullptr); + CHECK(std::string(uri).substr(0, 7) == "sha256:"); + return true; +} + +bool test_add_shared_asset_copy_in_false_portable_rejected() { + Sandbox s; + s.Write("src/a.txt", "alpha"); + + ModelPackage* raw = nullptr; + CHECK_OK(ModelPackage_New(&raw)); + PkgHandle p(raw); + const char* uri = nullptr; + CHECK_ERR(ModelPackage_AddSharedAsset(p.get(), (s.root() / "src").c_str(), + nullptr, /*copy_in=*/false, &uri), + MODEL_PACKAGE_ERR_STATE); + return true; +} + +bool test_add_shared_asset_copy_in_false_installed_ok() { + Sandbox s; + s.Write("src/a.txt", "alpha"); + + ModelPackage* raw = nullptr; + CHECK_OK(ModelPackage_New(&raw)); + PkgHandle p(raw); + CHECK_OK(ModelPackage_SetLayout(p.get(), "installed")); + const char* uri = nullptr; + CHECK_OK(ModelPackage_AddSharedAsset(p.get(), (s.root() / "src").c_str(), + nullptr, /*copy_in=*/false, &uri)); + CHECK(uri != nullptr); + // Surfaced as a manifest override -> shared_assets count should be 1. + CHECK((ModelPackage_Info(p.get()))->num_shared_assets == 1); + return true; +} + +bool test_add_shared_asset_expected_uri_mismatch_errors() { + Sandbox s; + s.Write("src/a.txt", "alpha"); + + ModelPackage* raw = nullptr; + CHECK_OK(ModelPackage_New(&raw)); + PkgHandle p(raw); + CHECK_OK(ModelPackage_SetLayout(p.get(), "installed")); + const char* uri = nullptr; + std::string bogus = "sha256:" + std::string(64, '0'); + CHECK_ERR(ModelPackage_AddSharedAsset(p.get(), (s.root() / "src").c_str(), + bogus.c_str(), /*copy_in=*/false, &uri), + MODEL_PACKAGE_ERR_STATE); + return true; +} + +bool test_remove_shared_asset() { + Sandbox s; + s.Write("src/a.txt", "alpha"); + + ModelPackage* raw = nullptr; + CHECK_OK(ModelPackage_New(&raw)); + PkgHandle p(raw); + CHECK_OK(ModelPackage_SetLayout(p.get(), "installed")); + const char* uri = nullptr; + CHECK_OK(ModelPackage_AddSharedAsset(p.get(), (s.root() / "src").c_str(), + nullptr, /*copy_in=*/false, &uri)); + std::string uri_copy(uri); + CHECK((ModelPackage_Info(p.get()))->num_shared_assets == 1); + CHECK_OK(ModelPackage_RemoveSharedAsset(p.get(), uri_copy.c_str())); + CHECK((ModelPackage_Info(p.get()))->num_shared_assets == 0); + return true; +} + +// ───────────────────────────────────────────────────────────────────────────── +// Round-trip through GetComponentJson / GetVariantJson +// ───────────────────────────────────────────────────────────────────────────── + +bool test_round_trip_component_json() { + ModelPackage* raw = nullptr; + CHECK_OK(ModelPackage_New(&raw)); + PkgHandle p(raw); + CHECK_OK(ModelPackage_SetComponentInline(p.get(), "c", + R"({"variants": {"v1": {"variant_directory": ".", "ep": "CPU"}}})")); + const char* j = nullptr; + CHECK_OK(ModelPackage_GetComponentJson(p.get(), "c", &j)); + CHECK(j != nullptr); + std::string s(j); + CHECK(s.find("\"variants\"") != std::string::npos); + CHECK(s.find("\"v1\"") != std::string::npos); + CHECK(s.find("\"CPU\"") != std::string::npos); + return true; +} + +// ───────────────────────────────────────────────────────────────────────────── +// View cache invalidation after mutation +// ───────────────────────────────────────────────────────────────────────────── + +bool test_view_cache_drops_on_remove() { + ModelPackage* raw = nullptr; + CHECK_OK(ModelPackage_New(&raw)); + PkgHandle p(raw); + CHECK_OK(ModelPackage_SetComponentInline(p.get(), "a", R"({"variants": {}})")); + CHECK_OK(ModelPackage_SetComponentInline(p.get(), "b", R"({"variants": {}})")); + const ModelComponentInfo* a = ModelPackage_FindComponent(ModelPackage_Info(p.get()), "a"); + CHECK(a != nullptr); + CHECK_OK(ModelPackage_RemoveComponent(p.get(), "a")); + // Old pointer was invalidated by the mutation; re-fetch and 'a' must now be gone. + const ModelPackageInfo* info = ModelPackage_Info(p.get()); + CHECK(ModelPackage_FindComponent(info, "a") == nullptr); + CHECK(ModelPackage_FindComponent(info, "b") != nullptr); + return true; +} + +struct Test { + const char* name; + bool (*fn)(); +}; + +const Test kTests[] = { + {"new_creates_empty_package", test_new_creates_empty_package}, + {"set_component_inline_basic", test_set_component_inline_basic}, + {"set_component_inline_replaces_existing", test_set_component_inline_replaces_existing}, + {"set_component_inline_rejects_unknown_field", test_set_component_inline_rejects_unknown_field}, + {"set_component_inline_rejects_bad_json", test_set_component_inline_rejects_bad_json}, + {"remove_component", test_remove_component}, + {"remove_missing_component_is_noop", test_remove_missing_component_is_noop}, + {"set_variant_upsert", test_set_variant_upsert}, + {"set_variant_unknown_component_errors", test_set_variant_unknown_component_errors}, + {"remove_variant", test_remove_variant}, + {"set_executor_info_inline_and_remove", test_set_executor_info_inline_and_remove}, + {"set_executor_info_external_records_path", test_set_executor_info_external_records_path}, + {"set_metadata", test_set_metadata}, + {"set_layout", test_set_layout}, + {"set_additional_metadata_manifest_scope", test_set_additional_metadata_manifest_scope}, + {"set_additional_metadata_variant_scope", test_set_additional_metadata_variant_scope}, + {"add_shared_asset_copy_in_true_portable_ok", test_add_shared_asset_copy_in_true_portable_ok}, + {"add_shared_asset_copy_in_false_portable_rejected", test_add_shared_asset_copy_in_false_portable_rejected}, + {"add_shared_asset_copy_in_false_installed_ok", test_add_shared_asset_copy_in_false_installed_ok}, + {"add_shared_asset_expected_uri_mismatch_errors", test_add_shared_asset_expected_uri_mismatch_errors}, + {"remove_shared_asset", test_remove_shared_asset}, + {"round_trip_component_json", test_round_trip_component_json}, + {"view_cache_drops_on_remove", test_view_cache_drops_on_remove}, +}; + +} // namespace + +int main() { + for (const auto& t : kTests) { + g_current = t.name; + bool ok = t.fn(); + if (ok) { + std::printf("[PASS] %s\n", t.name); + g_passed++; + } else { + g_failed++; + } + } + std::printf("\n=== %d passed, %d failed ===\n", g_passed, g_failed); + return g_failed == 0 ? 0 : 1; +} diff --git a/model_package/tests/test_commit.cc b/model_package/tests/test_commit.cc new file mode 100644 index 0000000000000..4ede82394e170 --- /dev/null +++ b/model_package/tests/test_commit.cc @@ -0,0 +1,504 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +/// \file test_commit.cc +/// \brief Commit, prune, and validate tests. + +#include "model_package.h" +#include "model_package_api.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace fs = std::filesystem; + +namespace { + +int g_failed = 0; +int g_passed = 0; +const char* g_current = ""; + +#define CHECK(cond) \ + do { \ + if (!(cond)) { \ + std::fprintf(stderr, "[FAIL] %s line %d: CHECK(%s)\n", g_current, __LINE__, #cond); \ + return false; \ + } \ + } while (0) + +#define CHECK_OK(status) \ + do { \ + ModelPackageStatus* _s = (status); \ + if (_s != nullptr) { \ + std::fprintf(stderr, "[FAIL] %s line %d: expected OK, got: %s\n", \ + g_current, __LINE__, ModelPackageStatus_Message(_s)); \ + ModelPackageStatus_Release(_s); \ + return false; \ + } \ + } while (0) + +#define CHECK_ERR(status, expected_code) \ + do { \ + ModelPackageStatus* _s = (status); \ + if (_s == nullptr) { \ + std::fprintf(stderr, "[FAIL] %s line %d: expected error %d, got OK\n", \ + g_current, __LINE__, (int)(expected_code)); \ + return false; \ + } \ + ModelPackageErrorCode _c = ModelPackageStatus_Code(_s); \ + if (_c != (expected_code)) { \ + std::fprintf(stderr, "[FAIL] %s line %d: expected error %d, got %d (%s)\n", \ + g_current, __LINE__, (int)(expected_code), (int)_c, \ + ModelPackageStatus_Message(_s)); \ + ModelPackageStatus_Release(_s); \ + return false; \ + } \ + ModelPackageStatus_Release(_s); \ + } while (0) + +class Sandbox { + public: + Sandbox() { + std::random_device rd; + std::mt19937_64 g(rd()); + char buf[32]; + std::snprintf(buf, sizeof(buf), "mp_commit_%016llx", static_cast(g())); + root_ = fs::temp_directory_path() / buf; + fs::create_directories(root_); + } + ~Sandbox() { + std::error_code ec; + fs::remove_all(root_, ec); + } + Sandbox(const Sandbox&) = delete; + Sandbox& operator=(const Sandbox&) = delete; + const fs::path& root() const { return root_; } + fs::path path(const std::string& rel) const { return root_ / rel; } + void Write(const std::string& rel, const std::string& contents) { + fs::path full = root_ / rel; + fs::create_directories(full.parent_path()); + std::ofstream f(full, std::ios::binary); + f << contents; + } + + private: + fs::path root_; +}; + +class PkgHandle { + public: + explicit PkgHandle(ModelPackage* p) : p_(p) {} + ~PkgHandle() { ModelPackage_Close(p_); } + PkgHandle(const PkgHandle&) = delete; + PkgHandle& operator=(const PkgHandle&) = delete; + ModelPackage* get() const { return p_; } + ModelPackage** outparam() { return &p_; } + + private: + ModelPackage* p_; +}; + +// Open a freshly-created in-memory package ready to commit at `root`. +// `root` must be empty/nonexistent for the subsequent dest_root commit. +PkgHandle MakeAuthoredPkgAt(const fs::path& /*root*/, + const std::string& layout = "portable") { + ModelPackage* raw = nullptr; + ModelPackage_New(&raw); + if (layout != "portable") ModelPackage_SetLayout(raw, layout.c_str()); + ModelPackage_SetComponentInline(raw, "encoder", R"({"variants": {}})"); + ModelPackage_SetVariant(raw, "encoder", "v1", R"({"ep": "CPU"})"); + return PkgHandle(raw); +} + +// ───────────────────────────────────────────────────────────────────────────── +// Commit (in-place, PRESERVE) +// ───────────────────────────────────────────────────────────────────────────── + +bool test_commit_inplace_basic_roundtrip() { + Sandbox s; + PkgHandle p = MakeAuthoredPkgAt(s.path("pkg")); + CHECK_OK(ModelPackage_Commit(p.get(), s.path("pkg").c_str(), MODEL_PACKAGE_WRITE_PRESERVE)); + // manifest.json exists. + CHECK(fs::is_regular_file(s.path("pkg") / "manifest.json")); + + // Reopen and confirm. + ModelPackage* re = nullptr; + CHECK_OK(ModelPackage_Open(s.path("pkg").c_str(), nullptr, &re)); + PkgHandle rep(re); + CHECK((ModelPackage_Info(rep.get()))->num_components == 1); + const ModelPackageInfo* info = ModelPackage_Info(rep.get()); + const ModelComponentInfo* c = ModelPackage_FindComponent(info, "encoder"); + CHECK(c != nullptr); + CHECK((c)->num_variants == 1); + const ModelVariantInfo* v = ModelComponentInfo_FindVariant(c, "v1"); + CHECK(std::string(v->ep) == "CPU"); + return true; +} + +bool test_commit_requires_package_root() { + ModelPackage* raw = nullptr; + CHECK_OK(ModelPackage_New(&raw)); + PkgHandle p(raw); + CHECK_ERR(ModelPackage_Commit(p.get(), nullptr, MODEL_PACKAGE_WRITE_PRESERVE), + MODEL_PACKAGE_ERR_STATE); + return true; +} + +bool test_commit_external_component_writes_file() { + Sandbox s; + // Author an inline package committed to disk first. + PkgHandle p = MakeAuthoredPkgAt(s.path("pkg")); + CHECK_OK(ModelPackage_Commit(p.get(), s.path("pkg").c_str(), MODEL_PACKAGE_WRITE_PRESERVE)); + + // Reopen, add an external component pointing at a file that doesn't exist yet. + ModelPackage* re = nullptr; + CHECK_OK(ModelPackage_Open(s.path("pkg").c_str(), nullptr, &re)); + PkgHandle rep(re); + CHECK_OK(ModelPackage_SetComponentExternal(rep.get(), "decoder", "decoder.json")); + CHECK_OK(ModelPackage_Commit(rep.get(), nullptr, MODEL_PACKAGE_WRITE_PRESERVE)); + CHECK(fs::is_regular_file(s.path("pkg") / "decoder.json")); + CHECK(fs::is_regular_file(s.path("pkg") / "manifest.json")); + + // Reopen yet again and verify external component round-trips. + ModelPackage* re2 = nullptr; + CHECK_OK(ModelPackage_Open(s.path("pkg").c_str(), nullptr, &re2)); + PkgHandle rep2(re2); + CHECK(ModelPackage_FindComponent(ModelPackage_Info(rep2.get()), "decoder") != nullptr); + return true; +} + +bool test_commit_pending_shared_asset_copy_in() { + Sandbox s; + s.Write("src_asset/m.onnx", "hello world"); + PkgHandle p = MakeAuthoredPkgAt(s.path("pkg")); + + const char* uri = nullptr; + CHECK_OK(ModelPackage_AddSharedAsset(p.get(), s.path("src_asset").c_str(), + nullptr, /*copy_in=*/true, &uri)); + std::string uri_copy(uri); + CHECK_OK(ModelPackage_Commit(p.get(), s.path("pkg").c_str(), + MODEL_PACKAGE_WRITE_PRESERVE)); + std::string hex = uri_copy.substr(7); + fs::path landed = s.path("pkg") / "shared_assets" / ("sha256-" + hex); + CHECK(fs::is_directory(landed)); + CHECK(fs::is_regular_file(landed / "m.onnx")); + return true; +} + +bool test_commit_dense_inlines_external_component() { + Sandbox s; + PkgHandle p = MakeAuthoredPkgAt(s.path("pkg")); + CHECK_OK(ModelPackage_Commit(p.get(), s.path("pkg").c_str(), MODEL_PACKAGE_WRITE_PRESERVE)); + CHECK_OK(ModelPackage_SetComponentExternal(p.get(), "decoder", "decoder.json")); + CHECK_OK(ModelPackage_Commit(p.get(), nullptr, MODEL_PACKAGE_WRITE_DENSE)); + // The dense commit should NOT have written decoder.json (component became inline). + CHECK(!fs::exists(s.path("pkg") / "decoder.json")); + // Manifest contains decoder as an inline object. + std::ifstream f(s.path("pkg") / "manifest.json"); + std::ostringstream oss; + oss << f.rdbuf(); + std::string m = oss.str(); + CHECK(m.find("\"decoder\"") != std::string::npos); + CHECK(m.find("\"variants\"") != std::string::npos); + return true; +} + +bool test_commit_dense_rejects_external_executor_info() { + Sandbox s; + PkgHandle p = MakeAuthoredPkgAt(s.path("pkg")); + CHECK_OK(ModelPackage_SetVariantExecutorInfoExternal( + p.get(), "encoder", "v1", "ort", "encoder/ort.json")); + CHECK_ERR(ModelPackage_Commit(p.get(), s.path("pkg").c_str(), MODEL_PACKAGE_WRITE_DENSE), + MODEL_PACKAGE_ERR_STATE); + return true; +} + +// ───────────────────────────────────────────────────────────────────────────── +// Commit (dest_root "save as") +// ───────────────────────────────────────────────────────────────────────────── + +bool test_commit_dest_root_self_contained() { + Sandbox s; + s.Write("src_asset/m.onnx", "alpha"); + PkgHandle p = MakeAuthoredPkgAt(s.path("orig")); + CHECK_OK(ModelPackage_Commit(p.get(), s.path("orig").c_str(), + MODEL_PACKAGE_WRITE_PRESERVE)); + + // Add an asset and commit as. + const char* uri = nullptr; + CHECK_OK(ModelPackage_AddSharedAsset(p.get(), s.path("src_asset").c_str(), + nullptr, /*copy_in=*/true, &uri)); + std::string uri_copy(uri); + fs::path saved = s.path("saved"); + CHECK_OK(ModelPackage_Commit(p.get(), saved.c_str(), MODEL_PACKAGE_WRITE_PRESERVE)); + CHECK(fs::is_regular_file(saved / "manifest.json")); + std::string hex = uri_copy.substr(7); + CHECK(fs::is_directory(saved / "shared_assets" / ("sha256-" + hex))); + + // After dest_root commit, in-memory state reflects the new root. + // (We can verify by mutating + committing in-place again.) + CHECK_OK(ModelPackage_SetMetadata(p.get(), "savedpkg", "1.0", nullptr)); + CHECK_OK(ModelPackage_Commit(p.get(), nullptr, MODEL_PACKAGE_WRITE_PRESERVE)); + // The most recent in-place commit should have landed at `saved`, not `orig`. + std::ifstream f(saved / "manifest.json"); + std::ostringstream oss; + oss << f.rdbuf(); + CHECK(oss.str().find("savedpkg") != std::string::npos); + return true; +} + +bool test_commit_dest_root_must_be_empty() { + Sandbox s; + PkgHandle p = MakeAuthoredPkgAt(s.path("pkg")); + CHECK_OK(ModelPackage_Commit(p.get(), s.path("pkg").c_str(), + MODEL_PACKAGE_WRITE_PRESERVE)); + s.Write("dest/something", "x"); + // Try to commit to non-empty dest. + CHECK_ERR(ModelPackage_Commit(p.get(), s.path("dest").c_str(), + MODEL_PACKAGE_WRITE_PRESERVE), + MODEL_PACKAGE_ERR_STATE); + return true; +} + +// ───────────────────────────────────────────────────────────────────────────── +// Prune +// ───────────────────────────────────────────────────────────────────────────── + +bool test_commit_dest_root_rehashes_existing_asset() { + Sandbox s; + s.Write("src_asset/m.onnx", "alpha"); + PkgHandle p = MakeAuthoredPkgAt(s.path("orig")); + const char* uri = nullptr; + CHECK_OK(ModelPackage_AddSharedAsset(p.get(), s.path("src_asset").c_str(), + nullptr, /*copy_in=*/true, &uri)); + std::string uri_copy(uri); + CHECK_OK(ModelPackage_Commit(p.get(), s.path("orig").c_str(), + MODEL_PACKAGE_WRITE_PRESERVE)); + + // Tamper with the landed sha256-/ dir under the existing package root. + std::string hex = uri_copy.substr(7); + fs::path landed = s.path("orig") / "shared_assets" / ("sha256-" + hex) / "m.onnx"; + { + std::ofstream f(landed, std::ios::binary); + f << "TAMPERED"; + } + + // CommitToDestRoot must rehash the source and refuse the mismatch. + CHECK_ERR(ModelPackage_Commit(p.get(), s.path("saved").c_str(), + MODEL_PACKAGE_WRITE_PRESERVE), + MODEL_PACKAGE_ERR_STATE); + return true; +} + +bool test_prune_never_touches_shared_assets() { + // Shared assets are content-addressed and only removed via explicit + // RemoveSharedAsset. Even an obviously orphan sha256-/ directory that + // matches no manifest entry must survive Prune. + Sandbox s; + PkgHandle p = MakeAuthoredPkgAt(s.path("pkg")); + CHECK_OK(ModelPackage_Commit(p.get(), s.path("pkg").c_str(), + MODEL_PACKAGE_WRITE_PRESERVE)); + + fs::path planted = s.path("pkg") / "shared_assets" / + ("sha256-" + std::string(64, 'a')); + fs::create_directories(planted); + // Backdate mtime to past grace window to make sure it isn't grace-protected. + auto old = fs::file_time_type::clock::now() - std::chrono::seconds(120); + std::error_code ec; + fs::last_write_time(planted, old, ec); + CHECK_OK(ModelPackage_Prune(p.get())); + CHECK(fs::is_directory(planted)); + return true; +} + +bool test_prune_reclaims_tracked_orphan_variant_dirs() { + Sandbox s; + PkgHandle p = MakeAuthoredPkgAt(s.path("pkg")); + CHECK_OK(ModelPackage_Commit(p.get(), s.path("pkg").c_str(), + MODEL_PACKAGE_WRITE_PRESERVE)); + // Now that package_root is anchored, materialize an on-disk variant dir and + // register it so subsequent removal records a tracked orphan. + fs::path victim = s.path("pkg") / "encoder" / "v1"; + fs::create_directories(victim); + CHECK_OK(ModelPackage_SetVariant(p.get(), "encoder", "v1", + R"({"ep":"CPU","variant_directory":"encoder/v1"})")); + CHECK_OK(ModelPackage_Commit(p.get(), nullptr, MODEL_PACKAGE_WRITE_PRESERVE)); + CHECK(fs::is_directory(victim)); + CHECK_OK(ModelPackage_RemoveVariant(p.get(), "encoder", "v1")); + CHECK_OK(ModelPackage_Commit(p.get(), nullptr, MODEL_PACKAGE_WRITE_PRESERVE)); + CHECK_OK(ModelPackage_Prune(p.get())); + CHECK(!fs::exists(victim)); + return true; +} + +bool test_prune_removes_stale_staging_dirs() { + Sandbox s; + PkgHandle p = MakeAuthoredPkgAt(s.path("pkg")); + CHECK_OK(ModelPackage_Commit(p.get(), s.path("pkg").c_str(), + MODEL_PACKAGE_WRITE_PRESERVE)); + + fs::path stage = s.path("pkg") / "shared_assets" / + ("sha256-" + std::string(64, 'c') + ".tmp.abcdef0123"); + fs::create_directories(stage); + auto old = fs::file_time_type::clock::now() - std::chrono::seconds(120); + std::error_code ec; + fs::last_write_time(stage, old, ec); + CHECK_OK(ModelPackage_Prune(p.get())); + CHECK(!fs::exists(stage)); + return true; +} + +// ───────────────────────────────────────────────────────────────────────────── +// Validate +// ───────────────────────────────────────────────────────────────────────────── + +bool test_validate_all_clean_package() { + Sandbox s; + PkgHandle p = MakeAuthoredPkgAt(s.path("pkg")); + CHECK_OK(ModelPackage_Commit(p.get(), s.path("pkg").c_str(), + MODEL_PACKAGE_WRITE_PRESERVE)); + const char* report = nullptr; + CHECK_OK(ModelPackage_Validate(p.get(), MODEL_PACKAGE_VALIDATE_ALL, &report)); + CHECK(report != nullptr); + CHECK(std::string(report).find("\"errors\": []") != std::string::npos); + return true; +} + +bool test_validate_paths_flags_missing_external() { + Sandbox s; + PkgHandle p = MakeAuthoredPkgAt(s.path("pkg")); + CHECK_OK(ModelPackage_Commit(p.get(), s.path("pkg").c_str(), + MODEL_PACKAGE_WRITE_PRESERVE)); + // Register an external component then delete the file behind the library's back. + CHECK_OK(ModelPackage_SetComponentExternal(p.get(), "decoder", "decoder.json")); + CHECK_OK(ModelPackage_Commit(p.get(), nullptr, MODEL_PACKAGE_WRITE_PRESERVE)); + std::error_code ec; + fs::remove(s.path("pkg") / "decoder.json", ec); + const char* report = nullptr; + CHECK_OK(ModelPackage_Validate(p.get(), MODEL_PACKAGE_VALIDATE_PATHS, &report)); + CHECK(std::string(report).find("PATHS") != std::string::npos); + return true; +} + +bool test_validate_asset_rehash_detects_mutation() { + Sandbox s; + s.Write("src_asset/m.onnx", "alpha"); + PkgHandle p = MakeAuthoredPkgAt(s.path("pkg")); + const char* uri = nullptr; + CHECK_OK(ModelPackage_AddSharedAsset(p.get(), s.path("src_asset").c_str(), + nullptr, /*copy_in=*/true, &uri)); + std::string uri_copy(uri); + CHECK_OK(ModelPackage_Commit(p.get(), s.path("pkg").c_str(), + MODEL_PACKAGE_WRITE_PRESERVE)); + // Mutate the on-disk shared asset directly. + std::string hex = uri_copy.substr(7); + fs::path landed = s.path("pkg") / "shared_assets" / ("sha256-" + hex) / "m.onnx"; + CHECK(fs::is_regular_file(landed)); + { + std::ofstream f(landed, std::ios::binary); + f << "MUTATED"; + } + const char* report = nullptr; + CHECK_ERR(ModelPackage_Validate(p.get(), MODEL_PACKAGE_VALIDATE_ASSET_REHASH, &report), + MODEL_PACKAGE_ERR_STATE); + CHECK(std::string(report).find("ASSET_REHASH") != std::string::npos); + return true; +} + +bool test_commit_accepts_unreferenced_shared_asset() { + // Shared assets no longer require an in-manifest reference: AddSharedAsset + // signals the user's intent to ship the asset, period. Commit materializes + // it under shared_assets/ at the default-convention path. + Sandbox s; + s.Write("src_asset/m.onnx", "alpha"); + PkgHandle p = MakeAuthoredPkgAt(s.path("pkg")); + const char* uri = nullptr; + CHECK_OK(ModelPackage_AddSharedAsset(p.get(), s.path("src_asset").c_str(), + nullptr, /*copy_in=*/true, &uri)); + std::string uri_copy(uri); + CHECK_OK(ModelPackage_Commit(p.get(), s.path("pkg").c_str(), + MODEL_PACKAGE_WRITE_PRESERVE)); + std::string hex = uri_copy.substr(7); + CHECK(fs::is_directory(s.path("pkg") / "shared_assets" / ("sha256-" + hex))); + // Same on dest_root path. + CHECK_OK(ModelPackage_Commit(p.get(), s.path("saved").c_str(), + MODEL_PACKAGE_WRITE_PRESERVE)); + CHECK(fs::is_directory(s.path("saved") / "shared_assets" / ("sha256-" + hex))); + return true; +} + +// ───────────────────────────────────────────────────────────────────────────── +// Atomicity hint: no stray .tmp.* under after successful commit +// ───────────────────────────────────────────────────────────────────────────── + +bool test_commit_leaves_no_temp_files() { + Sandbox s; + s.Write("src_asset/m.onnx", "alpha"); + PkgHandle p = MakeAuthoredPkgAt(s.path("pkg")); + CHECK_OK(ModelPackage_Commit(p.get(), s.path("pkg").c_str(), + MODEL_PACKAGE_WRITE_PRESERVE)); + const char* uri = nullptr; + CHECK_OK(ModelPackage_AddSharedAsset(p.get(), s.path("src_asset").c_str(), + nullptr, true, &uri)); + (void)uri; + CHECK_OK(ModelPackage_SetComponentExternal(p.get(), "decoder", "decoder.json")); + CHECK_OK(ModelPackage_Commit(p.get(), nullptr, + MODEL_PACKAGE_WRITE_PRESERVE)); + std::error_code ec; + for (auto& e : fs::recursive_directory_iterator(s.path("pkg"), ec)) { + if (e.path().filename().string().find(".tmp.") != std::string::npos) { + std::fprintf(stderr, " stray temp file: %s\n", e.path().c_str()); + return false; + } + } + return true; +} + +struct Test { + const char* name; + bool (*fn)(); +}; + +const Test kTests[] = { + {"commit_inplace_basic_roundtrip", test_commit_inplace_basic_roundtrip}, + {"commit_requires_package_root", test_commit_requires_package_root}, + {"commit_external_component_writes_file", test_commit_external_component_writes_file}, + {"commit_pending_shared_asset_copy_in", test_commit_pending_shared_asset_copy_in}, + {"commit_dense_inlines_external_component", test_commit_dense_inlines_external_component}, + {"commit_dense_rejects_external_executor_info", test_commit_dense_rejects_external_executor_info}, + {"commit_dest_root_self_contained", test_commit_dest_root_self_contained}, + {"commit_dest_root_must_be_empty", test_commit_dest_root_must_be_empty}, + {"commit_dest_root_rehashes_existing_asset", test_commit_dest_root_rehashes_existing_asset}, + {"prune_never_touches_shared_assets", test_prune_never_touches_shared_assets}, + {"prune_reclaims_tracked_orphan_variant_dirs", test_prune_reclaims_tracked_orphan_variant_dirs}, + {"prune_removes_stale_staging_dirs", test_prune_removes_stale_staging_dirs}, + {"validate_all_clean_package", test_validate_all_clean_package}, + {"validate_paths_flags_missing_external", test_validate_paths_flags_missing_external}, + {"validate_asset_rehash_detects_mutation", test_validate_asset_rehash_detects_mutation}, + {"commit_accepts_unreferenced_shared_asset", test_commit_accepts_unreferenced_shared_asset}, + {"commit_leaves_no_temp_files", test_commit_leaves_no_temp_files}, +}; + +} // namespace + +int main() { + for (const auto& t : kTests) { + g_current = t.name; + bool ok = t.fn(); + if (ok) { + std::printf("[PASS] %s\n", t.name); + g_passed++; + } else { + g_failed++; + } + } + std::printf("\n=== %d passed, %d failed ===\n", g_passed, g_failed); + return g_failed == 0 ? 0 : 1; +} diff --git a/model_package/tests/test_inspection.cc b/model_package/tests/test_inspection.cc new file mode 100644 index 0000000000000..0b51681bc7c80 --- /dev/null +++ b/model_package/tests/test_inspection.cc @@ -0,0 +1,582 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +/// \file test_inspection.cc +/// \brief Tests for the read-only inspection API (model_package.h). + +#include "model_package.h" +#include "model_package_api.h" + +#include +#include +#include +#include +#include +#include +#include + +namespace fs = std::filesystem; + +namespace { + +int g_failed = 0; +int g_passed = 0; +const char* g_current = ""; + +#define CHECK(cond) \ + do { \ + if (!(cond)) { \ + std::fprintf(stderr, "[FAIL] %s line %d: CHECK(%s)\n", g_current, __LINE__, #cond); \ + return false; \ + } \ + } while (0) + +#define CHECK_OK(status) \ + do { \ + ModelPackageStatus* _s = (status); \ + if (_s != nullptr) { \ + std::fprintf(stderr, "[FAIL] %s line %d: expected OK, got: %s\n", \ + g_current, __LINE__, ModelPackageStatus_Message(_s)); \ + ModelPackageStatus_Release(_s); \ + return false; \ + } \ + } while (0) + +#define CHECK_ERR(status, expected_code) \ + do { \ + ModelPackageStatus* _s = (status); \ + if (_s == nullptr) { \ + std::fprintf(stderr, "[FAIL] %s line %d: expected error %d, got OK\n", \ + g_current, __LINE__, (int)(expected_code)); \ + return false; \ + } \ + ModelPackageErrorCode _c = ModelPackageStatus_Code(_s); \ + if (_c != (expected_code)) { \ + std::fprintf(stderr, "[FAIL] %s line %d: expected error %d, got %d: %s\n", \ + g_current, __LINE__, (int)(expected_code), (int)_c, \ + ModelPackageStatus_Message(_s)); \ + ModelPackageStatus_Release(_s); \ + return false; \ + } \ + ModelPackageStatus_Release(_s); \ + } while (0) + +class Sandbox { + public: + Sandbox() { + std::random_device rd; + std::mt19937_64 g(rd()); + char buf[32]; + std::snprintf(buf, sizeof(buf), "mp_inspect_%016llx", static_cast(g())); + root_ = fs::temp_directory_path() / buf; + fs::create_directories(root_); + } + ~Sandbox() { + std::error_code ec; + fs::remove_all(root_, ec); + } + Sandbox(const Sandbox&) = delete; + Sandbox& operator=(const Sandbox&) = delete; + + const fs::path& root() const { return root_; } + + void Write(const std::string& relpath, const std::string& contents) { + fs::path full = root_ / relpath; + fs::create_directories(full.parent_path()); + std::ofstream f(full, std::ios::binary); + f << contents; + } + + void Touch(const std::string& relpath) { Write(relpath, ""); } + + private: + fs::path root_; +}; + +bool test_open_minimal_inline() { + Sandbox s; + s.Write("manifest.json", R"({ + "schema_version": 1, + "package_name": "test", + "components": { + "alpha": { + "variants": { + "cpu": {} + } + } + } + })"); + + ModelPackage* pkg = nullptr; + CHECK_OK(ModelPackage_Open(s.root().c_str(), nullptr, &pkg)); + CHECK(pkg != nullptr); + + const ModelPackageInfo* info = ModelPackage_Info(pkg); + CHECK(info != nullptr); + CHECK(info->schema_version_major == 1); + CHECK(info->schema_version_minor == 0); + CHECK(std::string(info->package_name) == "test"); + CHECK(std::string(info->layout) == "portable"); + CHECK((info)->num_components == 1); + CHECK((info)->num_shared_assets == 0); + CHECK(info->additional_metadata_json == nullptr); + + const ModelComponentInfo* c = &(info)->components[0]; + CHECK(c != nullptr); + CHECK(std::string(c->name) == "alpha"); + CHECK((c)->num_variants == 1); + + const ModelVariantInfo* v = &(c)->variants[0]; + CHECK(v != nullptr); + CHECK(std::string(v->name) == "cpu"); + CHECK(v->ep == nullptr); + CHECK(v->device == nullptr); + CHECK(v->compatibility_string == nullptr); + + ModelPackage_Close(pkg); + return true; +} + +bool test_open_full_inline_with_metadata() { + Sandbox s; + s.Write("manifest.json", R"({ + "schema_version": 1, + "package_name": "phi-4", + "package_version": "1.2.3", + "description": "demo", + "layout": "portable", + "additional_metadata": {"author": "team"}, + "components": { + "decoder": { + "additional_metadata": {"size": "small"}, + "variants": { + "cuda_fp16": { + "variant_directory": "decoder/cuda_fp16", + "ep": "CUDAExecutionProvider", + "device": "gpu", + "compatibility_string": "sm_80", + "additional_metadata": {"notes": "quantized"} + } + } + } + } + })"); + fs::create_directories(s.root() / "decoder" / "cuda_fp16"); + + ModelPackage* pkg = nullptr; + CHECK_OK(ModelPackage_Open(s.root().c_str(), nullptr, &pkg)); + const ModelPackageInfo* info = ModelPackage_Info(pkg); + CHECK(std::string(info->package_name) == "phi-4"); + CHECK(std::string(info->package_version) == "1.2.3"); + CHECK(std::string(info->description) == "demo"); + CHECK(info->additional_metadata_json != nullptr); + CHECK(std::string(info->additional_metadata_json).find("\"author\":\"team\"") != std::string::npos); + + const ModelComponentInfo* c = ModelPackage_FindComponent(info, "decoder"); + CHECK(c != nullptr); + const char* comp_meta = c->additional_metadata_json; + CHECK(comp_meta != nullptr); + CHECK(std::string(comp_meta).find("\"size\":\"small\"") != std::string::npos); + + const ModelVariantInfo* v = ModelComponentInfo_FindVariant(c, "cuda_fp16"); + CHECK(v != nullptr); + CHECK(std::string(v->ep) == "CUDAExecutionProvider"); + CHECK(std::string(v->device) == "gpu"); + CHECK(std::string(v->compatibility_string) == "sm_80"); + const char* var_meta = v->additional_metadata_json; + CHECK(var_meta != nullptr); + CHECK(std::string(var_meta).find("\"notes\":\"quantized\"") != std::string::npos); + + const char* resolved = v->variant_directory; + CHECK(resolved != nullptr); + CHECK(std::string(resolved).find("decoder/cuda_fp16") != std::string::npos); + + ModelPackage_Close(pkg); + return true; +} + +bool test_external_component_file() { + Sandbox s; + s.Write("manifest.json", R"({ + "schema_version": 1, + "components": { "decoder": "components/decoder.json" } + })"); + s.Write("components/decoder.json", R"({ + "variants": { "cpu": {} } + })"); + ModelPackage* pkg = nullptr; + CHECK_OK(ModelPackage_Open(s.root().c_str(), nullptr, &pkg)); + const ModelComponentInfo* c = ModelPackage_FindComponent(ModelPackage_Info(pkg), "decoder"); + CHECK(c != nullptr); + CHECK((c)->num_variants == 1); + ModelPackage_Close(pkg); + return true; +} + +bool test_external_component_directory() { + Sandbox s; + s.Write("manifest.json", R"({ + "schema_version": 1, + "components": { "decoder": "components/decoder" } + })"); + s.Write("components/decoder/component.json", R"({ + "variants": { "cpu": {} } + })"); + ModelPackage* pkg = nullptr; + CHECK_OK(ModelPackage_Open(s.root().c_str(), nullptr, &pkg)); + CHECK((ModelPackage_Info(pkg))->num_components == 1); + ModelPackage_Close(pkg); + return true; +} + +bool test_executor_info_inline_and_external() { + Sandbox s; + s.Write("manifest.json", R"({ + "schema_version": 1, + "components": { + "decoder": { + "variants": { + "cuda": { + "variant_directory": "v", + "executor_info": { + "ort": "ort_info.json", + "other": {"x": 1} + } + } + } + } + } + })"); + fs::create_directories(s.root() / "v"); + s.Write("v/ort_info.json", R"({"model_file":"model.onnx"})"); + + ModelPackage* pkg = nullptr; + CHECK_OK(ModelPackage_Open(s.root().c_str(), nullptr, &pkg)); + const ModelPackageInfo* info = ModelPackage_Info(pkg); + const ModelVariantInfo* v = + ModelComponentInfo_FindVariant(ModelPackage_FindComponent(info, "decoder"), "cuda"); + CHECK(v != nullptr); + + const ModelExecutorInfoEntry* ort_ei = ModelVariantInfo_FindExecutorInfo(v, "ort"); + const char* ort_json = ort_ei ? ort_ei->json : nullptr; + CHECK(ort_json != nullptr); + CHECK(std::string(ort_json).find("model.onnx") != std::string::npos); + + const ModelExecutorInfoEntry* other_ei = ModelVariantInfo_FindExecutorInfo(v, "other"); + const char* other_json = other_ei ? other_ei->json : nullptr; + CHECK(other_json != nullptr); + CHECK(std::string(other_json).find("\"x\":1") != std::string::npos); + + const ModelExecutorInfoEntry* missing_ei = ModelVariantInfo_FindExecutorInfo(v, "absent"); + const char* missing = missing_ei ? missing_ei->json : nullptr; + CHECK(missing_ei == nullptr); + CHECK(missing == nullptr); + + ModelPackage_Close(pkg); + return true; +} + +bool test_inline_executor_info_without_directory_accepted() { + // Library no longer requires variant_directory to exist for inline + // executor_info. Executors interpret their own payload. + Sandbox s; + s.Write("manifest.json", R"({ + "schema_version": 1, + "components": { + "decoder": { + "variants": { + "cuda": { + "executor_info": { "other": {"x": 1} } + } + } + } + } + })"); + ModelPackage* pkg = nullptr; + CHECK_OK(ModelPackage_Open(s.root().c_str(), nullptr, &pkg)); + ModelPackage_Close(pkg); + return true; +} + +bool test_path_confinement_rejects_external_paths() { + Sandbox s; + s.Write("manifest.json", R"({ + "schema_version": 1, + "components": { "x": "../escape.json" } + })"); + ModelPackage* pkg = nullptr; + CHECK_ERR(ModelPackage_Open(s.root().c_str(), nullptr, &pkg), MODEL_PACKAGE_ERR_PATH_CONFINEMENT); + return true; +} + +bool test_installed_layout_allows_absolute() { + // Build a package whose component lives outside its root. + Sandbox external; + external.Write("decoder.json", R"({"variants": {"cpu": {}}})"); + + Sandbox s; + std::string abs_comp = (external.root() / "decoder.json").string(); + // Escape backslashes for any platform that uses them — POSIX is fine as-is. + s.Write("manifest.json", std::string(R"({ + "schema_version": 1, + "layout": "installed", + "components": {"decoder": ")") + + abs_comp + R"("} + })"); + + ModelPackage* pkg = nullptr; + CHECK_OK(ModelPackage_Open(s.root().c_str(), nullptr, &pkg)); + CHECK((ModelPackage_Info(pkg))->num_components == 1); + ModelPackage_Close(pkg); + return true; +} + +bool test_shared_assets_resolve() { + Sandbox s; + s.Write("manifest.json", R"({ + "schema_version": 1, + "shared_assets": { + "sha256:aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa": "assets/a" + }, + "components": { + "x": { + "variants": { + "cpu": {} + } + } + } + })"); + fs::create_directories(s.root() / "assets" / "a"); + // Discovery: an on-disk sha256- dir without an override entry must + // surface alongside the explicit override. + fs::create_directories( + s.root() / "shared_assets" / + "sha256-bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb"); + + ModelPackage* pkg = nullptr; + CHECK_OK(ModelPackage_Open(s.root().c_str(), nullptr, &pkg)); + CHECK((ModelPackage_Info(pkg))->num_shared_assets == 2); + + const ModelSharedAssetInfo* a = &(ModelPackage_Info(pkg))->shared_assets[0]; + CHECK(a != nullptr); + CHECK(std::string(a->uri).find("aaaa") != std::string::npos); + CHECK(std::string(a->resolved_path).find("assets/a") != std::string::npos); + + const ModelSharedAssetInfo* b = &(ModelPackage_Info(pkg))->shared_assets[1]; + CHECK(b != nullptr); + CHECK(std::string(b->uri).find("bbbb") != std::string::npos); + // Default convention path: shared_assets/sha256- + CHECK(std::string(b->resolved_path).find("shared_assets/sha256-bb") != std::string::npos); + + // Resolve via API. + const char* path = nullptr; + CHECK_OK(ModelPackage_ResolveAssetUri(pkg, + "sha256:aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa", + &path)); + CHECK(std::string(path).find("assets/a") != std::string::npos); + + CHECK_ERR(ModelPackage_ResolveAssetUri(pkg, "sha256:not_a_known_one", &path), + MODEL_PACKAGE_ERR_ASSET_MISSING); + + ModelPackage_Close(pkg); + return true; +} + +bool test_unknown_field_rejected_strict() { + Sandbox s; + s.Write("manifest.json", R"({ + "schema_version": 1, + "components": { "x": {"variants": {"cpu": {"typo_field": 1}}} } + })"); + ModelPackage* pkg = nullptr; + CHECK_ERR(ModelPackage_Open(s.root().c_str(), nullptr, &pkg), MODEL_PACKAGE_ERR_SCHEMA); + return true; +} + +bool test_unknown_field_tolerated_lenient() { + Sandbox s; + s.Write("manifest.json", R"({ + "schema_version": 1, + "components": { "x": {"variants": {"cpu": {"typo_field": 1}}} } + })"); + ModelPackageOpenOptions opts{}; + opts.strict_unknown_fields = false; + opts.follow_symlinks = true; + ModelPackage* pkg = nullptr; + CHECK_OK(ModelPackage_Open(s.root().c_str(), &opts, &pkg)); + ModelPackage_Close(pkg); + return true; +} + +bool test_round_trip_getters_preserve_order() { + Sandbox s; + s.Write("manifest.json", R"({ + "schema_version": 1, + "components": { "decoder": {"variants": {"cuda": {"ep":"CUDAExecutionProvider","device":"gpu"}}} } + })"); + ModelPackage* pkg = nullptr; + CHECK_OK(ModelPackage_Open(s.root().c_str(), nullptr, &pkg)); + const char* comp_json = nullptr; + CHECK_OK(ModelPackage_GetComponentJson(pkg, "decoder", &comp_json)); + CHECK(comp_json != nullptr); + CHECK(std::string(comp_json).find("\"variants\":") != std::string::npos); + + const char* var_json = nullptr; + CHECK_OK(ModelPackage_GetVariantJson(pkg, "decoder", "cuda", &var_json)); + CHECK(var_json != nullptr); + // "ep" must appear before "device" — ordered_json preserves declaration order. + size_t ep_pos = std::string(var_json).find("\"ep\""); + size_t dev_pos = std::string(var_json).find("\"device\""); + CHECK(ep_pos != std::string::npos && dev_pos != std::string::npos && ep_pos < dev_pos); + ModelPackage_Close(pkg); + return true; +} + +bool test_round_trip_preserves_unknown_fields_lenient() { + Sandbox s; + s.Write("manifest.json", R"({ + "schema_version": 1, + "components": { "x": {"variants": {"cpu": {"future_field":"keepme"}}} } + })"); + ModelPackageOpenOptions opts{}; + opts.strict_unknown_fields = false; + opts.follow_symlinks = true; + ModelPackage* pkg = nullptr; + CHECK_OK(ModelPackage_Open(s.root().c_str(), &opts, &pkg)); + const char* var_json = nullptr; + CHECK_OK(ModelPackage_GetVariantJson(pkg, "x", "cpu", &var_json)); + CHECK(std::string(var_json).find("future_field") != std::string::npos); + ModelPackage_Close(pkg); + return true; +} + +bool test_missing_manifest() { + Sandbox s; + ModelPackage* pkg = nullptr; + CHECK_ERR(ModelPackage_Open(s.root().c_str(), nullptr, &pkg), MODEL_PACKAGE_ERR_IO); + return true; +} + +bool test_unsupported_schema_version() { + Sandbox s; + s.Write("manifest.json", R"({"schema_version": 99, "components": {}})"); + ModelPackage* pkg = nullptr; + CHECK_ERR(ModelPackage_Open(s.root().c_str(), nullptr, &pkg), MODEL_PACKAGE_ERR_VERSION); + return true; +} + +bool test_schema_version_string_and_minor() { + // "." string parses into the split fields. + { + Sandbox s; + s.Write("manifest.json", + R"({"schema_version": "1.0", "components": {"a": {"variants": {"cpu": {}}}}})"); + ModelPackage* pkg = nullptr; + CHECK_OK(ModelPackage_Open(s.root().c_str(), nullptr, &pkg)); + const ModelPackageInfo* info = ModelPackage_Info(pkg); + CHECK(info->schema_version_major == 1); + CHECK(info->schema_version_minor == 0); + ModelPackage_Close(pkg); + } + + // A newer minor than this build knows is accepted, and its unknown additive fields are + // tolerated rather than rejected even under the default strict mode. + { + Sandbox s; + s.Write("manifest.json", + R"({"schema_version": "1.7", "some_future_field": true, + "components": {"a": {"variants": {"cpu": {}}}}})"); + ModelPackage* pkg = nullptr; + CHECK_OK(ModelPackage_Open(s.root().c_str(), nullptr, &pkg)); + const ModelPackageInfo* info = ModelPackage_Info(pkg); + CHECK(info->schema_version_major == 1); + CHECK(info->schema_version_minor == 7); + ModelPackage_Close(pkg); + } + + // An unsupported major is rejected regardless of minor. + { + Sandbox s; + s.Write("manifest.json", R"({"schema_version": "2.0", "components": {}})"); + ModelPackage* pkg = nullptr; + CHECK_ERR(ModelPackage_Open(s.root().c_str(), nullptr, &pkg), MODEL_PACKAGE_ERR_VERSION); + } + + // A malformed schema_version string is a schema error. + { + Sandbox s; + s.Write("manifest.json", R"({"schema_version": "1.x", "components": {}})"); + ModelPackage* pkg = nullptr; + CHECK_ERR(ModelPackage_Open(s.root().c_str(), nullptr, &pkg), MODEL_PACKAGE_ERR_SCHEMA); + } + return true; +} + +bool test_invalid_sha256_uri_rejected() { + Sandbox s; + s.Write("manifest.json", R"({ + "schema_version": 1, + "shared_assets": { "sha256:notenough": "assets/a" }, + "components": {"x": {"variants": {"cpu": {}}}} + })"); + ModelPackage* pkg = nullptr; + CHECK_ERR(ModelPackage_Open(s.root().c_str(), nullptr, &pkg), MODEL_PACKAGE_ERR_SCHEMA); + return true; +} + +bool test_find_returns_null_on_missing() { + Sandbox s; + s.Write("manifest.json", R"({"schema_version":1,"components":{"a":{"variants":{"cpu":{}}}}})"); + ModelPackage* pkg = nullptr; + CHECK_OK(ModelPackage_Open(s.root().c_str(), nullptr, &pkg)); + const ModelPackageInfo* info = ModelPackage_Info(pkg); + CHECK(ModelPackage_FindComponent(info, "missing") == nullptr); + CHECK(ModelComponentInfo_FindVariant(ModelPackage_FindComponent(info, "a"), "missing") == nullptr); + ModelPackage_Close(pkg); + return true; +} + +struct Test { + const char* name; + bool (*fn)(); +}; + +const Test kTests[] = { + {"open_minimal_inline", test_open_minimal_inline}, + {"open_full_inline_with_metadata", test_open_full_inline_with_metadata}, + {"external_component_file", test_external_component_file}, + {"external_component_directory", test_external_component_directory}, + {"executor_info_inline_and_external", test_executor_info_inline_and_external}, + {"inline_executor_info_without_directory_accepted", + test_inline_executor_info_without_directory_accepted}, + {"path_confinement_rejects_external_paths", test_path_confinement_rejects_external_paths}, + {"installed_layout_allows_absolute", test_installed_layout_allows_absolute}, + {"shared_assets_resolve", test_shared_assets_resolve}, + {"unknown_field_rejected_strict", test_unknown_field_rejected_strict}, + {"unknown_field_tolerated_lenient", test_unknown_field_tolerated_lenient}, + {"round_trip_getters_preserve_order", test_round_trip_getters_preserve_order}, + {"round_trip_preserves_unknown_fields_lenient", + test_round_trip_preserves_unknown_fields_lenient}, + {"missing_manifest", test_missing_manifest}, + {"unsupported_schema_version", test_unsupported_schema_version}, + {"schema_version_string_and_minor", test_schema_version_string_and_minor}, + {"invalid_sha256_uri_rejected", test_invalid_sha256_uri_rejected}, + {"find_returns_null_on_missing", test_find_returns_null_on_missing}, +}; + +} // namespace + +int main() { + for (const auto& t : kTests) { + g_current = t.name; + bool ok = t.fn(); + if (ok) { + std::printf("[PASS] %s\n", t.name); + g_passed++; + } else { + g_failed++; + } + } + std::printf("\n=== %d passed, %d failed ===\n", g_passed, g_failed); + return g_failed == 0 ? 0 : 1; +} diff --git a/onnxruntime/core/session/model_package/README.md b/onnxruntime/core/session/model_package/README.md new file mode 100644 index 0000000000000..c4919219d7d40 --- /dev/null +++ b/onnxruntime/core/session/model_package/README.md @@ -0,0 +1,246 @@ +# ORT Model Package Integration + +This directory implements ONNX Runtime's consumer-side glue for the +standalone [`model_package` library](../../../../model_package/README.md): +loading packages, selecting variants against the runtime's execution +providers, and creating an `OrtSession` for the chosen variant. + +The package format, manifest schema, shared-asset rules, and the C +authoring/inspection API all live in `model_package/`. **This directory +adds three things on top**: + +1. The `executor_info["ort"]` payload schema (this is ORT's slot in the + variant body). +2. The variant selection algorithm, which queries each execution provider + factory and picks the highest-scoring variant. +3. The experimental `OrtModelPackageApi_*` C functions that wrap the library + and expose session creation. They are registered in + `include/onnxruntime/core/session/onnxruntime_experimental_c_api.inc` and + resolved by name through `OrtApi::GetExperimentalFunction`. + +ORT links the `model_package` library as a static archive; the library +itself never links against ORT. + +--- + +## Files + +| File | Responsibility | +| ------------------------------------- | -------------- | +| `model_package_context.h/.cc` | Translates the `model_package` library's C info tree into ORT-internal C++ structs (`ModelPackageInfo`, `ComponentInfo`, `VariantInfo`, `VariantModelInfo`). Parses the `executor_info["ort"]` payload. Owns `ModelPackageContext` (package-level) and `ModelPackageComponentContext` (per-component, with selected variant and provider list). | +| `model_package_options.h/.cc` | `ModelPackageOptions` snapshots EP intent (factories, devices, EP-name list) from an `OrtSessionOptions` at the moment `OrtModelPackageApi_CreateModelPackageOptionsFromSessionOptions_SinceV28` is called. Drives variant selection and provider construction. | +| `model_package_variant_selector.h/.cc`| `VariantSelector::SelectVariant` picks the best variant from a component given the EP list. Uses `OrtEpFactory::ValidateCompiledModelCompatibilityInfo`. | + +The C entry points themselves live in +`onnxruntime/core/session/model_package_api.cc` under +`namespace OrtExperimentalApis`. + +--- + +## `executor_info["ort"]` schema + +ORT's slot in `variant.executor_info` is a JSON object. All fields are +optional, but in practice `model_file` is required to load a session. + +```jsonc +{ + "model_file": "model.onnx", + "external_data": "weights", + "session_options": { "session.intra_op_thread_count": "4" }, + "provider_options": { "device_id": "0" } +} +``` + +| Field | Type | Required | Notes | +| ------------------ | ------ | -------- | ----- | +| `model_file` | string | yes (for session) | Path to the model file inside the variant. Resolved via `ModelPackage_ResolveStringRef`, anchored at the variant directory. Accepts relative paths, absolute paths or `..` segments (installed layout only), and `sha256:[/sub/path]` for shared-asset content. | +| `external_data` | string | no | Folder containing the model's external-initializers blobs. Wired into the session as ORT's external-initializers folder hint. Same resolution rules as `model_file`. | +| `session_options` | object | no | Map of `string -> string`. Merged on top of a fresh `OrtSessionOptions` when the caller passes `session_options == NULL` to `CreateSession`. Ignored when the caller supplies their own `OrtSessionOptions`. | +| `provider_options` | object | no | Map of `string -> string`. Merged into the variant's EP provider options on the default path. Ignored when the caller supplies their own `OrtSessionOptions`. | + +#### Inline vs external + +The slot follows the standard `executor_info` shape: the value may be either + +- a **string**, a path to a JSON file containing the body above (commonly + `ort_info.json` next to `model.onnx`), or +- an **object**, the body inlined into `component.json` / + `manifest.json`. + +Inline form keeps the package single-file. External form (the common case) +keeps the variant directory self-describing and survives `executor_info` +schema evolution without rewriting the manifest. + +The key under `executor_info` is the **executor namespace name** (`"ort"`), +not the EP. Other consumers use their own namespace key, so a single +variant can carry per-consumer payloads side by side. + +--- + +## Variant selection + +`ModelPackageOptions(env, session_options)` captures the **EP intent**: the +ordered list of execution providers registered on the session options, plus +their associated `OrtEpDevice` / `OrtHardwareDevice` / metadata. + +`VariantSelector::SelectVariant(component, ep_infos, &selected)` then walks +the component's variants and picks the best match: + +1. Use only the **first** EP from the captured list. (A policy may rank + several EPs; callers that need a specific EP should put it first. + Ranking across the full EP list is on the TODO list.) +2. For each variant, require `variant.ep == ep_info.ep_name`. +3. If `variant.device` is set (`"cpu"` / `"gpu"` / `"npu"`), require it to + match at least one of the EP's `OrtHardwareDevice` entries. +4. If both pass, call `OrtEpFactory::ValidateCompiledModelCompatibilityInfo` + with `variant.compatibility_string`. The EP returns an + `OrtCompiledModelCompatibility` enum which maps to a score: + + | Enum | Score | + | -------------------------------------------- | ----- | + | `EP_SUPPORTED_OPTIMAL` | 100 | + | `EP_SUPPORTED_PREFER_RECOMPILATION` | 50 | + | `EP_NOT_APPLICABLE` (or EP too old / no ABI) | 0 | + | `EP_UNSUPPORTED` | rejected | + +5. Pick the highest-scoring matching variant. Manifest declaration order + breaks ties. + +If no variant matches, `SelectComponent` fails with "No suitable model +variant found for the configured execution providers." + +ORT does **not** parse `compatibility_string`. The EP owns the format and +may encode multiple sub-targets (SoC ids, ISA flags, etc.) into the single +string internally; ORT only round-trips it through the EP callback. + +--- + +## Session creation contract + +`OrtModelPackageApi_CreateSession_SinceV28(env, component_ctx, session_options, &session)`. + +The `component_ctx` already knows which variant won selection and which +provider list it should use. Two paths: + +- **`session_options == NULL` (default).** ORT starts from a fresh + `OrtSessionOptions` and merges the variant's `session_options` / + `provider_options` from `executor_info["ort"]` on top. EPs declared in the + manifest are constructed and registered. This is what nearly all callers + want. + +- **`session_options != NULL` (advanced).** ORT uses the caller-supplied + `OrtSessionOptions` as-is. The manifest's `session_options` and + `provider_options` are **not** merged. Use this when you need custom EP + setup that does not round-trip through string options (shared CUDA + streams, shared QNN EP contexts, custom allocators, ...). The + `OrtSessionOptions` passed earlier to + `CreateModelPackageOptionsFromSessionOptions` only drives variant + selection / EP discovery; it is never silently re-applied here. + +In both modes, `external_data` from `executor_info["ort"]` is wired in as +ORT's external-initializers folder hint, so the model file can reference +weights stored next to (or shared by) the package. + +--- + +## C API surface + +The model package API is exposed via ONNX Runtime's +[experimental C API](../../../../docs/design/Experimental_C_API.md). Each +function is registered as a separate entry in +`include/onnxruntime/core/session/onnxruntime_experimental_c_api.inc` with +prefix `OrtModelPackageApi_` and version suffix `_SinceV28`. Consumers look +the functions up by name through `OrtApi::GetExperimentalFunction`, either +directly or via the typed C++ accessors in `Ort::Experimental::*` generated +from `onnxruntime_experimental_c_api.h`. + +The opaque handle types (`OrtModelPackageOptions`, `OrtModelPackageContext`, +`OrtModelPackageComponentContext`) are forward-declared at the top of +`onnxruntime_experimental_c_api.h`. + +Registered entries: + +| Function | Notes | +| ----------------------------------------------------- | ----- | +| `CreateModelPackageOptionsFromSessionOptions` | Snapshots EP intent. | +| `ReleaseModelPackageOptions` | | +| `CreateModelPackageContext` | Parses the manifest. | +| `ReleaseModelPackageContext` | | +| `ModelPackage_GetSchemaVersion` | | +| `ModelPackage_GetComponentCount` | | +| `ModelPackage_GetComponentNames` | | +| `ModelPackage_GetVariantCount` | | +| `ModelPackage_GetVariantNames` | | +| `ModelPackage_GetVariantEpName` | | +| `SelectComponent` | Resolves the best-matching variant. | +| `ReleaseModelPackageComponentContext` | | +| `ModelPackageComponent_GetSelectedVariantName` | | +| `ModelPackageComponent_GetSelectedVariantFolderPath` | | +| `CreateSession` | | + +> Experimental functions are not part of the stable ABI. Names, signatures +> and behaviour may change between releases until the surface is promoted +> to the stable `OrtApi`. Callers should null-check every lookup. + +Typical flow: + +```cpp +#include "onnxruntime_c_api.h" +#include "onnxruntime_experimental_c_api.h" + +const OrtApi* ort = OrtGetApiBase()->GetApi(ORT_API_VERSION); + +auto fn_create_opts = + Ort::Experimental::Get_OrtModelPackageApi_CreateModelPackageOptionsFromSessionOptions_SinceV28_Fn(ort); +auto fn_release_opts = + Ort::Experimental::Get_OrtModelPackageApi_ReleaseModelPackageOptions_SinceV28_Fn(ort); +auto fn_create_ctx = + Ort::Experimental::Get_OrtModelPackageApi_CreateModelPackageContext_SinceV28_Fn(ort); +auto fn_release_ctx = + Ort::Experimental::Get_OrtModelPackageApi_ReleaseModelPackageContext_SinceV28_Fn(ort); +auto fn_select = + Ort::Experimental::Get_OrtModelPackageApi_SelectComponent_SinceV28_Fn(ort); +auto fn_release_comp = + Ort::Experimental::Get_OrtModelPackageApi_ReleaseModelPackageComponentContext_SinceV28_Fn(ort); +auto fn_create_session = + Ort::Experimental::Get_OrtModelPackageApi_CreateSession_SinceV28_Fn(ort); + +OrtSessionOptions* so = nullptr; +ort->CreateSessionOptions(&so); +ort->SessionOptionsAppendExecutionProvider(so, "CUDAExecutionProvider", nullptr, nullptr, 0); + +OrtModelPackageOptions* mp_opts = nullptr; +fn_create_opts(env, so, &mp_opts); + +OrtModelPackageContext* ctx = nullptr; +fn_create_ctx(ORT_TSTR("/path/to/pkg"), &ctx); + +OrtModelPackageComponentContext* comp_ctx = nullptr; +fn_select(ctx, "decoder", mp_opts, &comp_ctx); + +OrtSession* session = nullptr; +fn_create_session(env, comp_ctx, nullptr, &session); + +ort->ReleaseSession(session); +fn_release_comp(comp_ctx); +fn_release_ctx(ctx); +fn_release_opts(mp_opts); +ort->ReleaseSessionOptions(so); +``` + +All `const char*` / `const ORTCHAR_T*` / array pointers returned by the API +are owned by the context that produced them and remain valid until the +context is released. + +--- + +## See also + +- [`model_package/README.md`](../../../../model_package/README.md): package + format, manifest/component schema, shared assets, path resolution, the + authoring C API, and the `executor_info` extension point. +- [`docs/design/Experimental_C_API.md`](../../../../docs/design/Experimental_C_API.md): + design and lifecycle rules for the experimental C API mechanism that + hosts these entries. +- `include/onnxruntime/core/session/onnxruntime_experimental_c_api.inc`: + the canonical list of `OrtModelPackageApi_*` entries. diff --git a/onnxruntime/core/session/model_package/model_package_context.cc b/onnxruntime/core/session/model_package/model_package_context.cc index ca4adb9c877a5..a0da46a10f88f 100644 --- a/onnxruntime/core/session/model_package/model_package_context.cc +++ b/onnxruntime/core/session/model_package/model_package_context.cc @@ -5,9 +5,7 @@ #include #include -#include #include -#include #include #include @@ -22,15 +20,23 @@ #include "core/session/provider_policy_context.h" #include "core/session/utils.h" -// We intentionally use the standalone model_package library's internal C++ types directly -// (model_package::ParsePackage, model_package_internal.h) rather than its public C API -// (ModelPackage_* functions). This avoids double-wrapping since ORT compiles the library in-tree. -// The public C API exists for external consumers (GenAI, FL) who link independently. -#include "model_package_internal.h" -#include "parser.h" +// Use the standalone model_package library's public C API. The library has no ORT +// dependency; ORT links it as a static archive (see cmake/onnxruntime_session.cmake) +// and translates the C handles into the ORT-internal C++ types defined in +// model_package_context.h here. +#include "model_package.h" namespace onnxruntime { +namespace { +// Deleter for the type-erased model_package handle held by ModelPackageContext. +void CloseModelPackageHandle(void* handle) { + if (handle != nullptr) { + ::ModelPackage_Close(static_cast<::ModelPackage*>(handle)); + } +} +} // namespace + namespace { Status FillOptionCachesFromMap( @@ -150,10 +156,9 @@ Status ModelPackageComponentContext::GetSelectedVariantFilePath(std::filesystem: "Selected variant index out of range for component: ", component_model_name_); const auto& selected_variant = component_model_info_.variants[selected_idx]; - ORT_RETURN_IF(!selected_variant.file.has_value(), + ORT_RETURN_IF(!selected_variant.file.has_value() || selected_variant.file->identifier.empty(), "Selected variant '", selected_variant.variant_name, - "' does not have a variant.json descriptor (or it lacks a 'filename' entry). " - "Component: ", + "' has no executor_info[\"ort\"][\"model_file\"]. Component: ", component_model_name_); out_path = selected_variant.file->model_file_path; @@ -345,54 +350,172 @@ Status ModelPackageComponentContext::GetSelectedVariantName(const std::string*& return Status::OK(); } -ModelPackageContext::ModelPackageContext(const std::filesystem::path& package_root) { - // Use the standalone model_package library for parsing. - model_package::PackageInfo pkg_info; - std::string error; - if (!model_package::ParsePackage(package_root, pkg_info, error)) { - ORT_THROW("Failed to parse model package: ", error); +Status ModelPackageComponentContext::GetSelectedVariantExternalDataFolder( + const std::string*& out_folder) const { + out_folder = nullptr; + const VariantInfo* selected_variant = nullptr; + ORT_RETURN_IF_ERROR(GetSelectedVariantInfo(selected_variant)); + ORT_RETURN_IF(selected_variant == nullptr, + "Selected variant is null for component: ", component_model_name_); + if (selected_variant->file.has_value() && + selected_variant->file->external_data_folder_path.has_value() && + !selected_variant->file->external_data_folder_path->empty()) { + out_folder = &(*selected_variant->file->external_data_folder_path); + } + return Status::OK(); +} + +ModelPackageContext::ModelPackageContext(const std::filesystem::path& package_root) + : package_handle_(nullptr, &CloseModelPackageHandle), package_root_(package_root) { + // Open the package via the model_package C API and keep the handle open for this context's + // lifetime (owned by package_handle_) so path references can be resolved later without + // reopening. The unique_ptr releases the handle even on exception paths during conversion. + ::ModelPackage* pkg = nullptr; + if (::ModelPackageStatus* st = ::ModelPackage_Open(package_root.string().c_str(), nullptr, &pkg)) { + std::string msg = ::ModelPackageStatus_Message(st) ? ::ModelPackageStatus_Message(st) : "unknown error"; + ::ModelPackageStatus_Release(st); + ORT_THROW("Failed to open model package at '", package_root.string(), "': ", msg); } + package_handle_.reset(pkg); - // Convert standalone library types to ORT internal types. - model_package_info_.schema_version = pkg_info.schema_version; + const ::ModelPackageInfo* pkg_info = ::ModelPackage_Info(pkg); + model_package_info_.schema_version = pkg_info ? pkg_info->schema_version_major : 0; model_package_info_.components.clear(); component_name_to_index_.clear(); - for (const auto& component : pkg_info.components) { - const auto& name = component.name; - size_t component_idx = model_package_info_.components.size(); - component_name_to_index_[name] = component_idx; + const size_t component_count = pkg_info ? pkg_info->num_components : 0; + for (size_t ci = 0; ci < component_count; ++ci) { + const ::ModelComponentInfo* component = &pkg_info->components[ci]; + + std::string component_name = component->name ? component->name : ""; + const size_t component_idx = model_package_info_.components.size(); + component_name_to_index_[component_name] = component_idx; ComponentInfo ort_component{}; - ort_component.component_name = name; + ort_component.component_name = component_name; ort_component.selected_variant_index.reset(); - for (const auto& variant : component.variants) { + const size_t variant_count = component->num_variants; + for (size_t vi = 0; vi < variant_count; ++vi) { + const ::ModelVariantInfo* variant = &component->variants[vi]; + VariantInfo ort_variant{}; - ort_variant.component_name = name; - ort_variant.variant_name = variant.name; - ort_variant.folder_path = variant.folder_path; - - // Convert EP compatibility (single entry per variant). - ort_variant.ep_compatibility.ep = variant.ep_compatibility.ep; - ort_variant.ep_compatibility.device = variant.ep_compatibility.device; - ort_variant.ep_compatibility.compatibility_string = variant.ep_compatibility.compatibility_string; - ort_variant.ep_compatibility.compiled_model_compatibility = OrtCompiledModelCompatibility_EP_NOT_APPLICABLE; - - // Convert file entry (single file per variant). - if (variant.file.has_value()) { + ort_variant.component_name = component_name; + ort_variant.variant_name = variant->name ? variant->name : ""; + + // Resolve the variant directory. Absence is treated as a soft signal; + // downstream callers that require a directory surface a clearer error + // at the point of use. + if (variant->variant_directory != nullptr) { + ort_variant.folder_path = std::filesystem::path(variant->variant_directory); + } + + // EP compatibility (single entry per variant). + if (variant->ep != nullptr) ort_variant.ep_compatibility.ep = std::string(variant->ep); + if (variant->device != nullptr) ort_variant.ep_compatibility.device = std::string(variant->device); + if (variant->compatibility_string != nullptr) + ort_variant.ep_compatibility.compatibility_string = std::string(variant->compatibility_string); + ort_variant.ep_compatibility.compiled_model_compatibility = + OrtCompiledModelCompatibility_EP_NOT_APPLICABLE; + + // Resolve the ORT executor_info from the manifest. + std::optional ort_obj; + if (const ::ModelExecutorInfoEntry* ei = + ::ModelVariantInfo_FindExecutorInfo(variant, "ort")) { + if (ei->json != nullptr && ei->json[0] != '\0') { + try { + ort_obj = json::parse(ei->json); + } catch (const std::exception& e) { + ORT_THROW("Failed to parse executor_info[\"ort\"] JSON for variant '", + ort_variant.variant_name, "' in component '", component_name, "': ", e.what()); + } + } + } + + if (ort_obj.has_value()) { + if (!ort_obj->is_object()) { + ORT_THROW("ORT variant configuration must be a JSON object for variant '", + ort_variant.variant_name, "' in component '", component_name, "'"); + } + VariantModelInfo ort_file{}; - ort_file.identifier = variant.file->filename; - ort_file.model_file_path = variant.file->resolved_path; - ort_file.session_options = variant.file->session_options; - ort_file.provider_options = variant.file->provider_options; - ort_file.shared_files = variant.file->shared_files; - ort_variant.file = std::move(ort_file); + + // Common resolver for ORT-side string refs (model_file, external_data). + // Delegates to ModelPackage_ResolveStringRef so accepted forms (relative, + // absolute, '..', sha256: URI, sha256: URI + subpath) and portable/installed + // confinement match the rest of the model_package library. + const std::string base_dir_str = ort_variant.folder_path.string(); + const char* base_dir = base_dir_str.empty() ? nullptr : base_dir_str.c_str(); + auto resolve_string_ref = [&](const char* field, const std::string& input, + bool must_exist) -> std::string { + const char* resolved = nullptr; + if (::ModelPackageStatus* st = ::ModelPackage_ResolveStringRef( + pkg, base_dir, input.c_str(), must_exist, &resolved)) { + std::string msg = ::ModelPackageStatus_Message(st) ? ::ModelPackageStatus_Message(st) + : "unknown error"; + ::ModelPackageStatus_Release(st); + ORT_THROW("Failed to resolve ORT variant '", field, "' = '", input, "' for variant '", + ort_variant.variant_name, "' in component '", component_name, "': ", msg); + } + return resolved ? std::string(resolved) : std::string{}; + }; + + if (auto it = ort_obj->find("model_file"); it != ort_obj->end()) { + if (!it->is_string()) { + ORT_THROW("ORT variant configuration: model_file must be a string for variant '", + ort_variant.variant_name, "' in component '", component_name, "'"); + } + const std::string model_file = it->get(); + ort_file.identifier = model_file; + ort_file.model_file_path = resolve_string_ref("model_file", model_file, + /*must_exist=*/false); + } + + auto fill_string_map = [&](const char* key, + std::optional>& dest) { + auto it = ort_obj->find(key); + if (it == ort_obj->end()) return; + if (!it->is_object()) { + ORT_THROW("ORT variant configuration: '", key, "' must be a JSON object for variant '", + ort_variant.variant_name, "' in component '", component_name, "'"); + } + std::unordered_map out; + out.reserve(it->size()); + for (auto kv = it->begin(); kv != it->end(); ++kv) { + if (!kv.value().is_string()) { + ORT_THROW("ORT variant configuration: '", key, "' entries must be strings for variant '", + ort_variant.variant_name, "' in component '", component_name, "'"); + } + out.emplace(kv.key(), kv.value().get()); + } + dest = std::move(out); + }; + fill_string_map("session_options", ort_file.session_options); + fill_string_map("provider_options", ort_file.provider_options); + + if (auto it = ort_obj->find("external_data"); it != ort_obj->end()) { + if (!it->is_string()) { + ORT_THROW("ORT variant configuration: external_data must be a string for variant '", + ort_variant.variant_name, "' in component '", component_name, "'"); + } + ort_file.external_data_folder_path = resolve_string_ref( + "external_data", it->get(), /*must_exist=*/false); + } + + if (!ort_file.identifier.empty() || ort_file.session_options.has_value() || + ort_file.provider_options.has_value() || ort_file.external_data_folder_path.has_value()) { + ort_variant.file = std::move(ort_file); + } } - // Consumer metadata. - if (variant.consumer_metadata_json.has_value()) { - ort_variant.consumer_metadata = nlohmann::json::parse(*variant.consumer_metadata_json); + // Variant-scope additional_metadata. + if (variant->additional_metadata_json != nullptr) { + try { + ort_variant.consumer_metadata = json::parse(variant->additional_metadata_json); + } catch (const std::exception& e) { + ORT_THROW("Failed to parse additional_metadata JSON for variant '", ort_variant.variant_name, + "' in component '", component_name, "': ", e.what()); + } } model_variant_infos_.push_back(ort_variant); @@ -405,7 +528,6 @@ ModelPackageContext::ModelPackageContext(const std::filesystem::path& package_ro // Create component names cache for quick lookup. component_names_cache_.clear(); component_names_cache_.reserve(model_package_info_.components.size()); - for (const auto& component : model_package_info_.components) { component_names_cache_.push_back(component.component_name); } @@ -415,6 +537,29 @@ size_t ModelPackageContext::GetComponentCount() const noexcept { return model_package_info_.components.size(); } +Status ModelPackageContext::ResolveStringRef(const std::string& base_dir, + const std::string& input, + bool must_exist, + const char*& out_path) const { + out_path = nullptr; + auto* pkg = static_cast<::ModelPackage*>(package_handle_.get()); + if (pkg == nullptr) { + return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "model package handle is not open"); + } + const char* resolved = nullptr; + if (::ModelPackageStatus* st = ::ModelPackage_ResolveStringRef( + pkg, base_dir.empty() ? nullptr : base_dir.c_str(), input.c_str(), must_exist, &resolved)) { + std::string msg = ::ModelPackageStatus_Message(st) ? ::ModelPackageStatus_Message(st) : "unknown error"; + ::ModelPackageStatus_Release(st); + return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Failed to resolve '", input, "' in model package: ", msg); + } + // Copy out of the library's transient buffer into a context-owned cache so the returned + // pointer stays valid until the next ResolveStringRef call. + resolve_string_ref_cache_ = resolved ? resolved : ""; + out_path = resolve_string_ref_cache_.c_str(); + return Status::OK(); +} + Status ModelPackageContext::GetComponentNames(gsl::span& out_names) const { out_names = gsl::span(component_names_cache_.data(), component_names_cache_.size()); diff --git a/onnxruntime/core/session/model_package/model_package_context.h b/onnxruntime/core/session/model_package/model_package_context.h index 1ecfffd7f74e0..a5ed5a04e917c 100644 --- a/onnxruntime/core/session/model_package/model_package_context.h +++ b/onnxruntime/core/session/model_package/model_package_context.h @@ -39,7 +39,11 @@ struct VariantModelInfo { // from variant.json file entry std::optional> session_options; std::optional> provider_options; - std::optional> shared_files; // logical_name -> checksum/path + + // Resolved folder containing the model's external initializer file, when + // executor_info.ort.external_data was set (path or sha256: URI). Empty + // otherwise. Used as the ORT external-initializers folder hint. + std::optional external_data_folder_path; }; // variant-level info (metadata.json + variant.json) @@ -124,6 +128,10 @@ class ModelPackageComponentContext { Status GetSelectedVariantName(const std::string*& out_name) const; + // Returns the resolved external_data folder for the selected variant, or + // nullptr-on-success if none was declared. Borrowed from VariantModelInfo. + Status GetSelectedVariantExternalDataFolder(const std::string*& out_folder) const; + std::vector>& MutableProviderList() { return provider_list_; } const std::vector& ExecutionDevices() const { return execution_devices_; } const std::vector& DevicesSelected() const { return devices_selected_; } @@ -184,7 +192,7 @@ class ModelPackageContext { gsl::span& out_variant_names) const; // Get the EP compatibility info declared on a variant. - // Lets callers (e.g. GenAI defaulting logic) inspect what EP a variant targets + // Lets callers inspect what EP a variant targets // before any EP has been resolved / before SelectComponent has been called. Status GetVariantEpCompatibility(const std::string& component_name, const std::string& variant_name, @@ -198,7 +206,25 @@ class ModelPackageContext { return model_variant_infos_; } + // Resolves a path reference from the package against the model_package library's rules: + // a "sha256:[/tail]" content-addressed shared-asset reference (honoring manifest + // overrides), or a plain relative path resolved against `base_dir` (empty base_dir falls + // back to the package root). When `must_exist` is true the resolved path must exist on + // disk. The returned pointer is owned by this context and stays valid until the next + // ResolveStringRef call. The underlying package handle is kept open for the context's + // lifetime so no reopen/reparse happens per call. + Status ResolveStringRef(const std::string& base_dir, const std::string& input, + bool must_exist, const char*& out_path) const; + private: + // The open model_package library handle, kept alive for this context's lifetime so path + // references can be resolved on demand. Stored type-erased (void*) to keep the + // model_package C header out of this ORT header; the deleter defined in the .cc closes it + // via ModelPackage_Close. + std::unique_ptr package_handle_; + std::filesystem::path package_root_{}; + mutable std::string resolve_string_ref_cache_{}; + ModelPackageInfo model_package_info_{}; std::vector model_variant_infos_; diff --git a/onnxruntime/core/session/model_package_api.cc b/onnxruntime/core/session/model_package_api.cc index 5fd25f9511eb9..aeae2ec1855d3 100644 --- a/onnxruntime/core/session/model_package_api.cc +++ b/onnxruntime/core/session/model_package_api.cc @@ -14,7 +14,6 @@ #include "core/session/model_package/model_package_context.h" #include "core/session/model_package/model_package_options.h" #include "core/session/utils.h" - #endif using namespace onnxruntime; @@ -400,13 +399,13 @@ ORT_API_STATUS_IMPL(OrtModelPackageApi_ModelPackage_GetVariantEpName_SinceV28, const onnxruntime::VariantEpCompatibilityInfo* info = nullptr; auto status = reinterpret_cast(ctx)->GetVariantEpCompatibility( component_name, variant_name, info); + if (!status.IsOK()) { + if (out_ep != nullptr) *out_ep = nullptr; + return onnxruntime::ToOrtStatus(status); + } if (out_ep != nullptr) { - if (status.IsOK() && info != nullptr && info->ep.has_value()) { - *out_ep = info->ep->c_str(); - } else { - *out_ep = nullptr; - } + *out_ep = (info != nullptr && info->ep.has_value()) ? info->ep->c_str() : nullptr; } return nullptr; #else @@ -419,6 +418,39 @@ ORT_API_STATUS_IMPL(OrtModelPackageApi_ModelPackage_GetVariantEpName_SinceV28, API_IMPL_END } +ORT_API_STATUS_IMPL(OrtModelPackageApi_ModelPackage_ResolveStringRef_SinceV28, + _In_ const OrtModelPackageContext* ctx, + _In_opt_ const char* base_dir, + _In_ const char* input, + _In_ int must_exist, + _Outptr_ const char** out_path) { + API_IMPL_BEGIN +#if !defined(ORT_MINIMAL_BUILD) + if (ctx == nullptr || input == nullptr || out_path == nullptr) { + return OrtApis::CreateStatus(ORT_INVALID_ARGUMENT, "ctx, input, and out_path must be non-null"); + } + *out_path = nullptr; + + const char* resolved = nullptr; + auto status = reinterpret_cast(ctx)->ResolveStringRef( + base_dir != nullptr ? std::string(base_dir) : std::string{}, std::string(input), + must_exist != 0, resolved); + if (!status.IsOK()) { + return onnxruntime::ToOrtStatus(status); + } + *out_path = resolved; + return nullptr; +#else + ORT_UNUSED_PARAMETER(ctx); + ORT_UNUSED_PARAMETER(base_dir); + ORT_UNUSED_PARAMETER(input); + ORT_UNUSED_PARAMETER(must_exist); + ORT_UNUSED_PARAMETER(out_path); + RETURN_NOT_IMPL_IN_MINIMAL_BUILD(); +#endif + API_IMPL_END +} + ORT_API_STATUS_IMPL(OrtModelPackageApi_ModelPackage_GetSchemaVersion_SinceV28, _In_ const OrtModelPackageContext* ctx, _Out_ int64_t* out_version) { diff --git a/onnxruntime/core/session/utils.cc b/onnxruntime/core/session/utils.cc index 330974aeed8d8..d196221ec55dc 100644 --- a/onnxruntime/core/session/utils.cc +++ b/onnxruntime/core/session/utils.cc @@ -981,16 +981,71 @@ OrtStatus* CreateSessionForModelPackage(_In_ const OrtSessionOptions* options, const std::filesystem::path& selected_model_path, onnxruntime::ModelPackageComponentContext& model_package_context, std::unique_ptr& sess) { - ORT_API_RETURN_IF_ERROR(CreateSessionAndLoadSingleModelImpl(options, env, - selected_model_path.c_str(), - /*model_data*/ nullptr, - /*model_data_length*/ 0, - sess)); + // When the variant declares an external_data folder (e.g. a shared asset + // under /shared_assets/sha256-/) we must switch to + // buffer load: ORT only honors session.model_external_initializers_file_folder_path + // when model_location_ is empty (see inference_session.cc). The mmap'd + // model buffer can be released right after Load; external initializers + // are read from the folder hint during Initialize. + const std::string* external_data_folder = nullptr; + ORT_API_RETURN_IF_STATUS_NOT_OK( + model_package_context.GetSelectedVariantExternalDataFolder(external_data_folder)); + + std::unique_ptr cloned_options; + const OrtSessionOptions* options_to_use = options; + onnxruntime::Env::MappedMemoryPtr mapped_model; + const void* model_data = nullptr; + size_t model_data_length = 0; + + if (external_data_folder != nullptr) { + cloned_options = options ? std::make_unique(*options) + : std::make_unique(); + ORT_API_RETURN_IF_STATUS_NOT_OK( + cloned_options->value.config_options.AddConfigEntry( + kOrtSessionOptionsModelExternalInitializersFileFolderPath, + external_data_folder->c_str())); + options_to_use = cloned_options.get(); + + size_t model_file_length = 0; + ORT_API_RETURN_IF_STATUS_NOT_OK( + onnxruntime::Env::Default().GetFileLength(selected_model_path.c_str(), model_file_length)); + if (model_file_length == 0) { + return OrtApis::CreateStatus( + ORT_FAIL, + ("model_package: selected variant model file is empty: " + selected_model_path.string()).c_str()); + } + ORT_API_RETURN_IF_STATUS_NOT_OK( + onnxruntime::Env::Default().MapFileIntoMemory(selected_model_path.c_str(), + /*offset=*/0, + model_file_length, + mapped_model)); + model_data = mapped_model.get(); + model_data_length = model_file_length; + } else if (options_to_use == nullptr) { + // No external_data and caller did not pass options: synthesize a default + // OrtSessionOptions so the downstream *options_to_use dereferences are safe. + cloned_options = std::make_unique(); + options_to_use = cloned_options.get(); + } + + if (model_data != nullptr) { + ORT_API_RETURN_IF_ERROR(CreateSessionAndLoadSingleModelImpl(options_to_use, env, + /*model_path*/ nullptr, + model_data, + model_data_length, + sess)); + } else { + ORT_API_RETURN_IF_ERROR(CreateSessionAndLoadSingleModelImpl(options_to_use, env, + selected_model_path.c_str(), + /*model_data*/ nullptr, + /*model_data_length*/ 0, + sess)); + } + mapped_model.reset(); - // Always rebuild providers from the effective session options (which include merged variant - // provider options). Providers created during EP selection used the original session options - // and would not reflect variant-specific provider options. - ORT_API_RETURN_IF_STATUS_NOT_OK(model_package_context.RebuildProviderListForSession(env, *options)); + // Providers were created earlier from the original options; rebuild now so + // any merged variant-specific provider options take effect. + ORT_API_RETURN_IF_STATUS_NOT_OK(model_package_context.RebuildProviderListForSession(env, *options_to_use)); auto& provider_list = model_package_context.MutableProviderList(); @@ -1000,10 +1055,10 @@ OrtStatus* CreateSessionForModelPackage(_In_ const OrtSessionOptions* options, } } - if (model_package_context.IsFromPolicy() && options != nullptr) { + if (model_package_context.IsFromPolicy()) { ProviderPolicyContext provider_policy_context; ORT_API_RETURN_IF_STATUS_NOT_OK(provider_policy_context.LogTelemetry( - *sess, *options, + *sess, *options_to_use, model_package_context.ExecutionDevices(), model_package_context.DevicesSelected())); } diff --git a/onnxruntime/test/autoep/test_model_package.cc b/onnxruntime/test/autoep/test_model_package.cc index c9065878da1ce..01a526da2a506 100644 --- a/onnxruntime/test/autoep/test_model_package.cc +++ b/onnxruntime/test/autoep/test_model_package.cc @@ -1,15 +1,18 @@ -// Copyright (c) Microsoft Corporation. +// Copyright (c) Microsoft Corporation. // Licensed under the MIT License. #include #include #include +#include #include #include #include +#include #include #include "gtest/gtest.h" +#include "nlohmann/json.hpp" #include "core/session/model_package/model_package_context.h" #include "core/session/onnxruntime_experimental_cxx_api.h" @@ -47,6 +50,8 @@ struct ModelPackageFns { ModelPackage_GetVariantNames{nullptr}; OrtExperimental_OrtModelPackageApi_ModelPackage_GetVariantEpName_SinceV28_Fn ModelPackage_GetVariantEpName{nullptr}; + OrtExperimental_OrtModelPackageApi_ModelPackage_ResolveStringRef_SinceV28_Fn + ModelPackage_ResolveStringRef{nullptr}; OrtExperimental_OrtModelPackageApi_SelectComponent_SinceV28_Fn SelectComponent{nullptr}; OrtExperimental_OrtModelPackageApi_ReleaseModelPackageComponentContext_SinceV28_Fn @@ -84,6 +89,8 @@ inline const ModelPackageFns& GetModelPackageFns() { Exp::Get_OrtModelPackageApi_ModelPackage_GetVariantNames_SinceV28_FnOrThrow(api); f.ModelPackage_GetVariantEpName = Exp::Get_OrtModelPackageApi_ModelPackage_GetVariantEpName_SinceV28_FnOrThrow(api); + f.ModelPackage_ResolveStringRef = + Exp::Get_OrtModelPackageApi_ModelPackage_ResolveStringRef_SinceV28_FnOrThrow(api); f.SelectComponent = Exp::Get_OrtModelPackageApi_SelectComponent_SinceV28_FnOrThrow(api); f.ReleaseModelPackageComponentContext = @@ -98,218 +105,121 @@ inline const ModelPackageFns& GetModelPackageFns() { }(); return fns; } -// ------------------------------------------------------------------ -// Helpers to build a test model package on disk -// ------------------------------------------------------------------ -std::filesystem::path CreateManifestJson(const std::filesystem::path& package_root, - std::string_view manifest_json) { - std::filesystem::path manifest_path = package_root / "manifest.json"; - std::filesystem::create_directories(package_root); - - std::ofstream os(manifest_path, std::ios::binary); - os << manifest_json; - return manifest_path; -} - -std::string MakeVariantJson(std::string_view filename) { - std::ostringstream oss; - oss << R"({ - "filename": ")" - << filename << R"(" - })"; - return oss.str(); -} - -void CreateVariantDescriptor(const std::filesystem::path& package_root, - std::string_view component_name, - std::string_view variant_name, - std::string_view variant_json) { - const auto variant_root = package_root / "models" / std::string(component_name) / std::string(variant_name); - std::filesystem::create_directories(variant_root); - - std::ofstream os(variant_root / "variant.json", std::ios::binary); - os << variant_json; -} +// ──────────────────────────────────────────────────────────────────────────── +// Fixture helpers for building model packages on disk. +// Every package is a single manifest.json at the package root that declares +// components/variants/executor_info inline. Variant directories live at +// `///` and contain the model file. +// ──────────────────────────────────────────────────────────────────────────── + +struct VariantSpec { + std::string variant_name; + std::string ep; // empty => omit + std::string device; // empty => omit + std::string compatibility_string; // empty => omit + std::filesystem::path source_model; // empty => no executor_info + std::optional> session_options; + std::optional> provider_options; +}; -std::filesystem::path CreateModelPackage( - const std::filesystem::path& package_root, - std::string_view manifest_json, - std::string_view component_name, - std::string_view variant_name_1, - std::string_view variant_name_2, - const std::filesystem::path& source_model_1, - const std::filesystem::path& source_model_2) { +// Build a single-component new-schema package on disk and return its root. +// `package_root` is wiped before writing. +std::filesystem::path BuildPackage(const std::filesystem::path& package_root, + const std::string& component_name, + const std::vector& variants) { std::error_code ec; std::filesystem::remove_all(package_root, ec); std::filesystem::create_directories(package_root); - CreateManifestJson(package_root, manifest_json); - - const auto models_root = package_root / "models" / std::string(component_name); - const auto variant1_dir = models_root / std::string(variant_name_1); - const auto variant2_dir = models_root / std::string(variant_name_2); - - std::filesystem::create_directories(variant1_dir); - std::filesystem::create_directories(variant2_dir); - - const auto variant1_model = variant1_dir / source_model_1.filename(); - const auto variant2_model = variant2_dir / source_model_2.filename(); - - std::filesystem::copy_file(source_model_1, variant1_model, std::filesystem::copy_options::overwrite_existing, ec); - std::filesystem::copy_file(source_model_2, variant2_model, std::filesystem::copy_options::overwrite_existing, ec); - - CreateVariantDescriptor(package_root, component_name, variant_name_1, - MakeVariantJson(source_model_1.filename().string())); - CreateVariantDescriptor(package_root, component_name, variant_name_2, - MakeVariantJson(source_model_2.filename().string())); - - return package_root; -} - -std::filesystem::path CreateComponentModelMetadata( - const std::filesystem::path& package_root, - std::string_view component_name, - std::string_view metadata_json) { - const auto component_root = package_root / "models" / std::string(component_name); - - std::filesystem::create_directories(component_root); - - const std::filesystem::path metadata_path = component_root / "metadata.json"; - std::ofstream os(metadata_path, std::ios::binary); - os << metadata_json; - - return component_root; -} - -std::string MakeManifestJson(std::string_view component_name) { - std::ostringstream oss; - oss << R"({ - "schema_version": 1, - "components": [")" - << component_name << R"("] - })"; - return oss.str(); -} - -std::string MakeMetadataJsonTwoVariants(std::string_view component_name, - std::string_view variant_name_1, - std::string_view variant_ep_1, - std::string_view variant_device_1, - std::string_view variant_compatibility_string_1, - std::string_view variant_name_2, - std::string_view variant_ep_2, - std::string_view variant_device_2, - std::string_view variant_compatibility_string_2) { - std::ostringstream oss; - oss << R"({ - "component_name": ")" - << component_name << R"(", - "variants": { - ")" - << variant_name_1 << R"(": { - "ep": ")" - << variant_ep_1 << R"(", - "device": ")" - << variant_device_1 << R"(", - "compatibility_string": ")" - << variant_compatibility_string_1 << R"(" - }, - ")" - << variant_name_2 << R"(": { - "ep": ")" - << variant_ep_2 << R"(", - "device": ")" - << variant_device_2 << R"(", - "compatibility_string": ")" - << variant_compatibility_string_2 << R"(" + using ojson = nlohmann::ordered_json; + ojson variants_obj = ojson::object(); + for (const auto& v : variants) { + const std::string variant_dir_rel = component_name + "/" + v.variant_name; + const auto variant_dir_abs = package_root / component_name / v.variant_name; + std::filesystem::create_directories(variant_dir_abs); + + ojson variant_obj = ojson::object(); + variant_obj["variant_directory"] = variant_dir_rel; + if (!v.ep.empty()) variant_obj["ep"] = v.ep; + if (!v.device.empty()) variant_obj["device"] = v.device; + if (!v.compatibility_string.empty()) variant_obj["compatibility_string"] = v.compatibility_string; + + if (!v.source_model.empty()) { + const std::string model_filename = v.source_model.filename().string(); + std::filesystem::copy_file(v.source_model, variant_dir_abs / model_filename, + std::filesystem::copy_options::overwrite_existing, ec); + + ojson ort_info = ojson::object(); + ort_info["model_file"] = model_filename; + if (v.session_options.has_value()) { + ojson so = ojson::object(); + for (const auto& kv : *v.session_options) so[kv.first] = kv.second; + ort_info["session_options"] = std::move(so); } - } - })"; - return oss.str(); -} - -std::filesystem::path CreateModelPackageApiTestPackage(bool multi_file_variant = false) { - const auto package_root = std::filesystem::temp_directory_path() / "ort_model_package_api_test"; - std::error_code ec; - std::filesystem::remove_all(package_root, ec); - - constexpr std::string_view manifest_json = R"({ - "schema_version": 1, - "components": ["model_1"] - })"; - - CreateModelPackage(package_root, manifest_json, - "model_1", "variant_1", "variant_2", - std::filesystem::path{"testdata/mul_1.onnx"}, std::filesystem::path{"testdata/mul_16.onnx"}); - - constexpr std::string_view metadata_json = R"({ - "component_name": "model_1", - "variants": { - "variant_1": { - "ep": "example_ep", - "device": "cpu", - "compatibility_string": "example_ep;version=0.1.0;ort_api_version=25;hardware_architecture=arch1" - }, - "variant_2": { - "ep": "example_ep", - "device": "npu", - "compatibility_string": "example_ep;version=0.1.0;ort_api_version=25;hardware_architecture=arch2" + if (v.provider_options.has_value()) { + ojson po = ojson::object(); + for (const auto& kv : *v.provider_options) po[kv.first] = kv.second; + ort_info["provider_options"] = std::move(po); } + ojson executor_info = ojson::object(); + executor_info["ort"] = std::move(ort_info); + variant_obj["executor_info"] = std::move(executor_info); } - })"; - - CreateComponentModelMetadata(package_root, "model_1", metadata_json); - if (!multi_file_variant) { - std::ofstream os(package_root / "models" / "model_1" / "variant_1" / "variant.json", std::ios::binary); - os << R"({ - "filename": "mul_1.onnx", - "session_options": { - "session.disable_prepacking": "1", - "session.intra_op.allow_spinning": "0" - }, - "provider_options": { - "backend_path": "example_backend", - "enable_htp": "1" - } - })"; - } else { - // Multi-file variants are no longer supported. For backward-compat testing, - // just write a single-file variant.json. - std::ofstream os(package_root / "models" / "model_1" / "variant_1" / "variant.json", std::ios::binary); - os << R"({ - "filename": "mul_1.onnx", - "session_options": { - "session.disable_prepacking": "1", - "session.intra_op.allow_spinning": "0" - }, - "provider_options": { - "backend_path": "example_backend", - "enable_htp": "1" - } - })"; + variants_obj[v.variant_name] = std::move(variant_obj); } - { - std::ofstream os(package_root / "models" / "model_1" / "variant_2" / "variant.json", std::ios::binary); - os << R"({ - "filename": "mul_16.onnx" - })"; - } + ojson component_obj = ojson::object(); + component_obj["variants"] = std::move(variants_obj); + + ojson components_obj = ojson::object(); + components_obj[component_name] = std::move(component_obj); + + ojson manifest = ojson::object(); + manifest["schema_version"] = "1.0"; + manifest["components"] = std::move(components_obj); + std::ofstream os(package_root / "manifest.json", std::ios::binary); + os << manifest.dump(2); return package_root; } +// Convenience: most tests use the same two-variant shape backed by mul_1.onnx / +// mul_16.onnx. `compat_1` and `compat_2` default to empty (no compatibility string). +std::filesystem::path BuildTwoVariantPackage(const std::filesystem::path& package_root, + std::string_view variant_name_1, + std::string_view device_1, + std::string_view compat_1, + const std::filesystem::path& model_1, + std::string_view variant_name_2, + std::string_view device_2, + std::string_view compat_2, + const std::filesystem::path& model_2, + std::string_view ep_name = "example_ep") { + std::vector variants; + variants.push_back(VariantSpec{std::string(variant_name_1), std::string(ep_name), std::string(device_1), std::string(compat_1), model_1, {}, {}}); + variants.push_back(VariantSpec{std::string(variant_name_2), std::string(ep_name), std::string(device_2), std::string(compat_2), model_2, {}, {}}); + return BuildPackage(package_root, "model_1", variants); +} + } // namespace -// ------------------------------------------------------------------ +// ──────────────────────────────────────────────────────────────────────────── // Model Package API tests -// ------------------------------------------------------------------ +// ──────────────────────────────────────────────────────────────────────────── TEST(ModelPackageApiTest, PackageContextQueries) { - const auto package_root = CreateModelPackageApiTestPackage(); + const auto package_root = std::filesystem::temp_directory_path() / "ort_model_package_api_test"; + BuildTwoVariantPackage(package_root, + "variant_1", "cpu", + "example_ep;version=0.1.0;ort_api_version=25;hardware_architecture=arch1", + "testdata/mul_1.onnx", + "variant_2", "npu", + "example_ep;version=0.1.0;ort_api_version=25;hardware_architecture=arch2", + "testdata/mul_16.onnx"); const auto& pkg_api = GetModelPackageFns(); + ASSERT_NE(pkg_api.CreateModelPackageContext, nullptr) << "Model package experimental API is not available"; auto context_deleter = [&pkg_api](OrtModelPackageContext* p) { if (p) pkg_api.ReleaseModelPackageContext(p); @@ -320,7 +230,6 @@ TEST(ModelPackageApiTest, PackageContextQueries) { ASSERT_ORTSTATUS_OK(pkg_api.CreateModelPackageContext(package_root.c_str(), &raw_context)); model_pkg_context.reset(raw_context); - // Query: component count + names size_t component_count = 0; ASSERT_ORTSTATUS_OK(pkg_api.ModelPackage_GetComponentCount(model_pkg_context.get(), &component_count)); ASSERT_EQ(component_count, 1u); @@ -334,7 +243,6 @@ TEST(ModelPackageApiTest, PackageContextQueries) { ASSERT_NE(component_names[0], nullptr); EXPECT_STREQ(component_names[0], "model_1"); - // Query: variant count + names size_t variant_count = 0; ASSERT_ORTSTATUS_OK(pkg_api.ModelPackage_GetVariantCount( model_pkg_context.get(), "model_1", &variant_count)); @@ -358,8 +266,82 @@ TEST(ModelPackageApiTest, PackageContextQueries) { std::filesystem::remove_all(package_root, ec); } +TEST(ModelPackageApiTest, ResolveStringRef) { + const auto package_root = std::filesystem::temp_directory_path() / "ort_model_package_resolve_test"; + std::vector variants; + variants.push_back(VariantSpec{"variant_1", "example_ep", "cpu", "", "testdata/mul_1.onnx", {}, {}}); + BuildPackage(package_root, "model_1", variants); + + // A content-addressed shared asset, discovered by convention at shared_assets/sha256-/. + const std::string digest(64, 'a'); + const auto asset_dir = package_root / "shared_assets" / ("sha256-" + digest); + std::filesystem::create_directories(asset_dir); + { + std::ofstream os(asset_dir / "asset.txt", std::ios::binary); + os << "hello"; + } + + const auto& pkg_api = GetModelPackageFns(); + ASSERT_NE(pkg_api.ModelPackage_ResolveStringRef, nullptr) << "Model package experimental API is not available"; + + auto context_deleter = [&pkg_api](OrtModelPackageContext* p) { + if (p) pkg_api.ReleaseModelPackageContext(p); + }; + std::unique_ptr ctx(nullptr, context_deleter); + OrtModelPackageContext* raw_context = nullptr; + ASSERT_ORTSTATUS_OK(pkg_api.CreateModelPackageContext(package_root.c_str(), &raw_context)); + ctx.reset(raw_context); + + const char* resolved = nullptr; + + // "sha256:" resolves to the shared asset directory (override/discovery aware). + ASSERT_ORTSTATUS_OK(pkg_api.ModelPackage_ResolveStringRef( + ctx.get(), nullptr, ("sha256:" + digest).c_str(), /*must_exist=*/1, &resolved)); + ASSERT_NE(resolved, nullptr); + EXPECT_EQ(std::filesystem::canonical(resolved), std::filesystem::canonical(asset_dir)); + + // "sha256:/" resolves the confined tail under the asset directory. + ASSERT_ORTSTATUS_OK(pkg_api.ModelPackage_ResolveStringRef( + ctx.get(), nullptr, ("sha256:" + digest + "/asset.txt").c_str(), /*must_exist=*/1, &resolved)); + ASSERT_NE(resolved, nullptr); + EXPECT_EQ(std::filesystem::canonical(resolved), std::filesystem::canonical(asset_dir / "asset.txt")); + + // A plain relative path resolves against base_dir. + const auto variant_dir = package_root / "model_1" / "variant_1"; + ASSERT_ORTSTATUS_OK(pkg_api.ModelPackage_ResolveStringRef( + ctx.get(), variant_dir.string().c_str(), "mul_1.onnx", /*must_exist=*/1, &resolved)); + ASSERT_NE(resolved, nullptr); + EXPECT_EQ(std::filesystem::canonical(resolved), std::filesystem::canonical(variant_dir / "mul_1.onnx")); + + // An undeclared sha256 asset is rejected even when must_exist is false. + const std::string missing_digest(64, 'b'); + OrtStatus* status = pkg_api.ModelPackage_ResolveStringRef( + ctx.get(), nullptr, ("sha256:" + missing_digest).c_str(), /*must_exist=*/0, &resolved); + EXPECT_NE(status, nullptr); + if (status != nullptr) Ort::GetApi().ReleaseStatus(status); + + std::error_code ec; + std::filesystem::remove_all(package_root, ec); +} + TEST(ModelPackageApiTest, SingleFileVariantInComponent_SelectComponentAndCreateSession) { - const auto package_root = CreateModelPackageApiTestPackage(); + const auto package_root = std::filesystem::temp_directory_path() / "ort_model_package_api_test"; + std::vector variants; + variants.push_back(VariantSpec{ + "variant_1", "example_ep", "cpu", + "example_ep;version=0.1.0;ort_api_version=25;hardware_architecture=arch1", + "testdata/mul_1.onnx", + std::unordered_map{ + {"session.disable_prepacking", "1"}, + {"session.intra_op.allow_spinning", "0"}, + }, + std::unordered_map{ + {"backend_path", "example_backend"}, + {"enable_htp", "1"}, + }}); + variants.push_back(VariantSpec{ + "variant_2", "example_ep", "npu", "example_ep;version=0.1.0;ort_api_version=25;hardware_architecture=arch2", "testdata/mul_16.onnx", {}, {}}); + BuildPackage(package_root, "model_1", variants); RegisteredEpDeviceUniquePtr example_ep; ASSERT_NO_FATAL_FAILURE(Utils::RegisterAndGetExampleEp(*ort_env, Utils::example_ep_info, example_ep)); @@ -370,6 +352,7 @@ TEST(ModelPackageApiTest, SingleFileVariantInComponent_SelectComponentAndCreateS session_options.AppendExecutionProvider_V2(*ort_env, {plugin_ep_device}, ep_options); const auto& pkg_api = GetModelPackageFns(); + ASSERT_NE(pkg_api.CreateModelPackageContext, nullptr) << "Model package experimental API is not available"; auto options_deleter = [&pkg_api](OrtModelPackageOptions* p) { if (p) pkg_api.ReleaseModelPackageOptions(p); @@ -428,156 +411,64 @@ TEST(ModelPackageApiTest, SingleFileVariantInComponent_SelectComponentAndCreateS } TEST(ModelPackageTest, LoadModelPackageAndRunInference_PluginEp_AppendV2) { - // Test Case 1: - // package_root is a model package directory which contains a manifest.json. - // This model package only contains one component model and it has a metadata.json. - // ORT should parse the manifest and the metadata.json to get model variants' constraints. - // ORT selects most suitable model variant based on constraints and then loads it to run inference successfully. - { - // Build model package on disk - const auto package_root = std::filesystem::temp_directory_path() / "ort_model_package_test"; - CreateModelPackage(package_root, MakeManifestJson("model_1"), - "model_1", "variant_1", "variant_2", - std::filesystem::path{"testdata/mul_1.onnx"}, std::filesystem::path{"testdata/mul_16.onnx"}); - - const std::string metadata_json = MakeMetadataJsonTwoVariants( - "model_1", - "variant_1", "example_ep", "cpu", "", - "variant_2", "example_ep", "npu", ""); - - CreateComponentModelMetadata(package_root, - "model_1", - metadata_json); - - // Register example EP and get its device - RegisteredEpDeviceUniquePtr example_ep; - ASSERT_NO_FATAL_FAILURE(Utils::RegisterAndGetExampleEp(*ort_env, Utils::example_ep_info, example_ep)); - Ort::ConstEpDevice plugin_ep_device(example_ep.get()); - - // Prepare session options with ExampleEP appended - Ort::SessionOptions session_options; - std::unordered_map ep_options; - session_options.AppendExecutionProvider_V2(*ort_env, {plugin_ep_device}, ep_options); - - // Create session from package root (directory) - // ORT should pick the variant_1 model since the constraints match the example EP device (device "cpu" matches) - Ort::Session session(*ort_env, package_root.c_str(), session_options); - - // Prepare input X - Ort::MemoryInfo memory_info = Ort::MemoryInfo::CreateCpu(OrtDeviceAllocator, OrtMemTypeCPU); - std::vector shape = {3, 2}; - std::vector input_data = {1.f, 2.f, 3.f, 4.f, 5.f, 6.f}; - Ort::Value input = Ort::Value::CreateTensor(memory_info, input_data.data(), input_data.size(), - shape.data(), shape.size()); - const char* input_names[] = {"X"}; - const char* output_names[] = {"Y"}; - std::vector inputs; - inputs.push_back(std::move(input)); - - // Run - auto outputs = session.Run(Ort::RunOptions{nullptr}, input_names, inputs.data(), inputs.size(), - output_names, 1); - ASSERT_EQ(outputs.size(), 1u); - const float* out = outputs[0].GetTensorData(); - gsl::span out_span(out, input_data.size()); - EXPECT_THAT(out_span, ::testing::ElementsAre(1.f, 4.f, 9.f, 16.f, 25.f, 36.f)); - - // Cleanup - std::error_code ec; - std::filesystem::remove_all(package_root, ec); - } - - // Test Case 2: - // package_root is a component model directory which contains a metadata.json. - // ORT should parse metadata.json to get model variants' constraints. - // ORT selects most suitable model variant based on constraints and then loads it to run inference successfully. - { - // Build model package on disk - const auto package_root = std::filesystem::temp_directory_path() / "ort_model_package_test"; + // package_root is a new-schema model package directory with one component and two variants. + // ORT parses the manifest, selects the variant whose device matches the registered EP (cpu), + // and loads/runs it successfully. + const auto package_root = std::filesystem::temp_directory_path() / "ort_model_package_test"; + BuildTwoVariantPackage(package_root, + "variant_1", "cpu", "", + "testdata/mul_1.onnx", + "variant_2", "npu", "", + "testdata/mul_16.onnx"); - CreateModelPackage(package_root, MakeManifestJson("model_1"), - "model_1", "variant_1", "variant_2", - std::filesystem::path{"testdata/mul_1.onnx"}, std::filesystem::path{"testdata/mul_16.onnx"}); + RegisteredEpDeviceUniquePtr example_ep; + ASSERT_NO_FATAL_FAILURE(Utils::RegisterAndGetExampleEp(*ort_env, Utils::example_ep_info, example_ep)); + Ort::ConstEpDevice plugin_ep_device(example_ep.get()); - const std::string metadata_json = MakeMetadataJsonTwoVariants( - "model_1", - "variant_1", "example_ep", "cpu", "", - "variant_2", "example_ep", "npu", ""); + Ort::SessionOptions session_options; + std::unordered_map ep_options; + session_options.AppendExecutionProvider_V2(*ort_env, {plugin_ep_device}, ep_options); - const auto component_model_root = CreateComponentModelMetadata(package_root, - "model_1", - metadata_json); + Ort::Session session(*ort_env, package_root.c_str(), session_options); - // Register example EP and get its device - RegisteredEpDeviceUniquePtr example_ep; - ASSERT_NO_FATAL_FAILURE(Utils::RegisterAndGetExampleEp(*ort_env, Utils::example_ep_info, example_ep)); - Ort::ConstEpDevice plugin_ep_device(example_ep.get()); + Ort::MemoryInfo memory_info = Ort::MemoryInfo::CreateCpu(OrtDeviceAllocator, OrtMemTypeCPU); + std::vector shape = {3, 2}; + std::vector input_data = {1.f, 2.f, 3.f, 4.f, 5.f, 6.f}; + Ort::Value input = Ort::Value::CreateTensor(memory_info, input_data.data(), input_data.size(), + shape.data(), shape.size()); + const char* input_names[] = {"X"}; + const char* output_names[] = {"Y"}; + std::vector inputs; + inputs.push_back(std::move(input)); - // Prepare session options with ExampleEP appended - Ort::SessionOptions session_options; - std::unordered_map ep_options; - session_options.AppendExecutionProvider_V2(*ort_env, {plugin_ep_device}, ep_options); + auto outputs = session.Run(Ort::RunOptions{nullptr}, input_names, inputs.data(), inputs.size(), + output_names, 1); + ASSERT_EQ(outputs.size(), 1u); + const float* out = outputs[0].GetTensorData(); + gsl::span out_span(out, input_data.size()); + EXPECT_THAT(out_span, ::testing::ElementsAre(1.f, 4.f, 9.f, 16.f, 25.f, 36.f)); - // Create session from component model root (directory) - // ORT should pick the variant_1 model since the constraints match the example EP device (device "cpu" matches) - Ort::Session session(*ort_env, component_model_root.c_str(), session_options); - - // Prepare input X - Ort::MemoryInfo memory_info = Ort::MemoryInfo::CreateCpu(OrtDeviceAllocator, OrtMemTypeCPU); - std::vector shape = {3, 2}; - std::vector input_data = {1.f, 2.f, 3.f, 4.f, 5.f, 6.f}; - Ort::Value input = Ort::Value::CreateTensor(memory_info, input_data.data(), input_data.size(), - shape.data(), shape.size()); - const char* input_names[] = {"X"}; - const char* output_names[] = {"Y"}; - std::vector inputs; - inputs.push_back(std::move(input)); - - // Run - auto outputs = session.Run(Ort::RunOptions{nullptr}, input_names, inputs.data(), inputs.size(), - output_names, 1); - ASSERT_EQ(outputs.size(), 1u); - const float* out = outputs[0].GetTensorData(); - gsl::span out_span(out, input_data.size()); - EXPECT_THAT(out_span, ::testing::ElementsAre(1.f, 4.f, 9.f, 16.f, 25.f, 36.f)); - - // Cleanup - std::error_code ec; - std::filesystem::remove_all(package_root, ec); - } + std::error_code ec; + std::filesystem::remove_all(package_root, ec); } TEST(ModelPackageTest, LoadModelPackageAndRunInference_PreferCpu) { - // Build model package on disk const auto package_root = std::filesystem::temp_directory_path() / "ort_model_package_test"; + BuildTwoVariantPackage(package_root, + "variant_1", "cpu", "", + "testdata/mul_1.onnx", + "variant_2", "npu", "", + "testdata/mul_16.onnx"); - CreateModelPackage(package_root, MakeManifestJson("model_1"), - "model_1", "variant_1", "variant_2", - std::filesystem::path{"testdata/mul_1.onnx"}, std::filesystem::path{"testdata/mul_16.onnx"}); - - const std::string metadata_json = MakeMetadataJsonTwoVariants( - "model_1", - "variant_1", "example_ep", "cpu", "", - "variant_2", "example_ep", "npu", ""); - - CreateComponentModelMetadata(package_root, - "model_1", - metadata_json); - - // Register example EP and get its device RegisteredEpDeviceUniquePtr example_ep; ASSERT_NO_FATAL_FAILURE(Utils::RegisterAndGetExampleEp(*ort_env, Utils::example_ep_info, example_ep)); Ort::ConstEpDevice plugin_ep_device(example_ep.get()); - // Prepare session options with ExampleEP appended Ort::SessionOptions session_options; session_options.SetEpSelectionPolicy(OrtExecutionProviderDevicePolicy_PREFER_CPU); - // Create session from package root (directory) - // ORT should pick the variant_1 model since the constraints match the example EP device (device "cpu" matches) Ort::Session session(*ort_env, package_root.c_str(), session_options); - // Prepare input X Ort::MemoryInfo memory_info = Ort::MemoryInfo::CreateCpu(OrtDeviceAllocator, OrtMemTypeCPU); std::vector shape = {3, 2}; std::vector input_data = {1.f, 2.f, 3.f, 4.f, 5.f, 6.f}; @@ -588,7 +479,6 @@ TEST(ModelPackageTest, LoadModelPackageAndRunInference_PreferCpu) { std::vector inputs; inputs.push_back(std::move(input)); - // Run auto outputs = session.Run(Ort::RunOptions{nullptr}, input_names, inputs.data(), inputs.size(), output_names, 1); ASSERT_EQ(outputs.size(), 1u); @@ -596,7 +486,6 @@ TEST(ModelPackageTest, LoadModelPackageAndRunInference_PreferCpu) { gsl::span out_span(out, input_data.size()); EXPECT_THAT(out_span, ::testing::ElementsAre(1.f, 4.f, 9.f, 16.f, 25.f, 36.f)); - // Cleanup std::error_code ec; std::filesystem::remove_all(package_root, ec); } @@ -610,7 +499,6 @@ TEST(ModelPackageTest, CheckCompiledModelCompatibilityInfo) { const ORTCHAR_T* output_model_file = ORT_TSTR("plugin_ep_compat_test.onnx"); std::filesystem::remove(output_model_file); - // Compile the model { Ort::SessionOptions session_options; std::unordered_map ep_options; @@ -630,153 +518,41 @@ TEST(ModelPackageTest, CheckCompiledModelCompatibilityInfo) { ASSERT_TRUE(std::filesystem::exists(output_model_file)); } - // Build model package on disk - const auto package_root = std::filesystem::temp_directory_path() / "ort_model_package_test"; - - CreateModelPackage(package_root, MakeManifestJson("model_1"), - "model_1", "variant_2", "variant_1", - std::filesystem::path{"testdata/mul_16.onnx"}, std::filesystem::path{"plugin_ep_compat_test.onnx"}); - // Build compat strings dynamically against current ORT_API_VERSION so the EP's ORT-version check - // doesn't short-circuit to PREFER_RECOMPILATION for both variants (which would make hardware_architecture - // irrelevant and the variant ranking collapse to a tie). With matching ORT versions, the arch differentiates: - // arch1 -> OPTIMAL, arch2 -> PREFER_RECOMPILATION; variant_1 must win. + // doesn't short-circuit to PREFER_RECOMPILATION for both variants. With matching ORT versions the + // hardware_architecture field differentiates: arch1 -> OPTIMAL, arch2 -> PREFER_RECOMPILATION, so + // variant_1 (mul_1) must win over variant_2 (mul_16). If variant_2 was picked, session init would + // fail with "No Op registered for Mul16". const std::string ort_api_version_str = std::to_string(ORT_API_VERSION); const std::string compat_arch2 = "example_ep;version=0.1.0;ort_api_version=" + ort_api_version_str + ";hardware_architecture=arch2"; const std::string compat_arch1 = "example_ep;version=0.1.0;ort_api_version=" + ort_api_version_str + ";hardware_architecture=arch1"; - const std::string metadata_json = MakeMetadataJsonTwoVariants( - "model_1", - "variant_2", "example_ep", "cpu", compat_arch2.c_str(), - "variant_1", "example_ep", "cpu", compat_arch1.c_str()); - CreateComponentModelMetadata(package_root, - "model_1", - metadata_json); + const auto package_root = std::filesystem::temp_directory_path() / "ort_model_package_test"; + BuildTwoVariantPackage(package_root, + "variant_2", "cpu", compat_arch2, + std::filesystem::path{"testdata/mul_16.onnx"}, + "variant_1", "cpu", compat_arch1, + std::filesystem::path{"plugin_ep_compat_test.onnx"}); - // Prepare session options with ExampleEP appended Ort::SessionOptions session_options; std::unordered_map ep_options; session_options.AppendExecutionProvider_V2(*ort_env, {plugin_ep_device}, ep_options); - // Create session from package root (directory) - // ORT should pick the variant_1 model since it has OrtCompiledModelCompatibility_EP_SUPPORTED_OPTIMAL for the example EP, - // while variant_2 is only OrtCompiledModelCompatibility_EP_SUPPORTED_PREFER_RECOMPILATION. - // If variant_2 was selected and loaded, i.e. mul_16.onnx, session initialization would fail with error "Error No Op registered for Mul16". Ort::Session session(*ort_env, package_root.c_str(), session_options); - // Cleanup std::error_code ec; std::filesystem::remove_all(package_root, ec); } -TEST(ModelPackageTest, LoadModelPackageAndRunInference_DiscoverComponentsFromModelsFolder) { - // manifest.json without "components"; discovery should scan models/* with metadata.json. - const auto package_root = std::filesystem::temp_directory_path() / "ort_model_package_discover_test"; - constexpr std::string_view manifest_json = R"({ - "schema_version": 1, - "model_name": "test_model" - })"; - - CreateModelPackage(package_root, manifest_json, - "model_1", "variant_1", "variant_2", - std::filesystem::path{"testdata/mul_1.onnx"}, std::filesystem::path{"testdata/mul_16.onnx"}); - - // Prepare component model with metadata and variants - const std::string component_name = "model_1"; - const std::string metadata_json = MakeMetadataJsonTwoVariants( - "model_1", - "variant_1", "example_ep", "cpu", "", - "variant_2", "example_ep", "npu", ""); - - // Create metadata.json under models/model_1 - const auto component_root = CreateComponentModelMetadata(package_root, - component_name, - metadata_json); - - // Add another component folder without metadata to ensure it's ignored - std::filesystem::create_directories(package_root / "models" / "ignored_component"); - - // Register example EP and get its device - RegisteredEpDeviceUniquePtr example_ep; - ASSERT_NO_FATAL_FAILURE(Utils::RegisterAndGetExampleEp(*ort_env, Utils::example_ep_info, example_ep)); - Ort::ConstEpDevice plugin_ep_device(example_ep.get()); - - // Prepare session options with ExampleEP appended - Ort::SessionOptions session_options; - std::unordered_map ep_options; - session_options.AppendExecutionProvider_V2(*ort_env, {plugin_ep_device}, ep_options); - - // Create session from package root (directory). Discovery should find model_1 via metadata.json, - // then pick variant_1 (device cpu) matching the example EP device. - // If variant_2 was selected and loaded, i.e. mul_16.onnx, session initialization would fail with error "Error No Op registered for Mul16". - Ort::Session session(*ort_env, package_root.c_str(), session_options); - - // Prepare input X - Ort::MemoryInfo memory_info = Ort::MemoryInfo::CreateCpu(OrtDeviceAllocator, OrtMemTypeCPU); - std::vector shape = {3, 2}; - std::vector input_data = {1.f, 2.f, 3.f, 4.f, 5.f, 6.f}; - Ort::Value input = Ort::Value::CreateTensor(memory_info, input_data.data(), input_data.size(), - shape.data(), shape.size()); - const char* input_names[] = {"X"}; - const char* output_names[] = {"Y"}; - std::vector inputs; - inputs.push_back(std::move(input)); - - // Run - auto outputs = session.Run(Ort::RunOptions{nullptr}, input_names, inputs.data(), inputs.size(), - output_names, 1); - ASSERT_EQ(outputs.size(), 1u); - const float* out = outputs[0].GetTensorData(); - gsl::span out_span(out, input_data.size()); - EXPECT_THAT(out_span, ::testing::ElementsAre(1.f, 4.f, 9.f, 16.f, 25.f, 36.f)); - - // Cleanup - std::error_code ec; - std::filesystem::remove_all(package_root, ec); -} - -TEST(ModelPackageTest, ParseVariantsFromRoot_PackageRootDirectory) { +TEST(ModelPackageTest, ParseVariantsFromPackageRoot) { const auto package_root = std::filesystem::temp_directory_path() / "ort_model_package_parse_from_package_root"; - std::error_code ec; - std::filesystem::remove_all(package_root, ec); - - // package_root is a model package directory (has manifest.json). - constexpr std::string_view manifest_json = R"({ - "schema_version": 1, - "components": ["model_1"] - })"; - - CreateModelPackage(package_root, manifest_json, - "model_1", "variant_1", "variant_2", - std::filesystem::path{"testdata/mul_1.onnx"}, std::filesystem::path{"testdata/mul_16.onnx"}); - - constexpr std::string_view metadata_json = R"({ - "component_name": "model_1", - "variants": { - "variant_1": { - "ep": "example_ep", - "device": "cpu" - }, - "variant_2": { - "ep": "example_ep", - "device": "npu" - } - } - })"; - - CreateComponentModelMetadata(package_root, "model_1", metadata_json); - - // New schema: per-variant descriptor in variant.json - { - std::ofstream os(package_root / "models" / "model_1" / "variant_1" / "variant.json", std::ios::binary); - os << R"({ "filename": "mul_1.onnx" })"; - } - { - std::ofstream os(package_root / "models" / "model_1" / "variant_2" / "variant.json", std::ios::binary); - os << R"({ "filename": "mul_16.onnx" })"; - } + BuildTwoVariantPackage(package_root, + "variant_1", "cpu", "", + std::filesystem::path{"testdata/mul_1.onnx"}, + "variant_2", "npu", "", + std::filesystem::path{"testdata/mul_16.onnx"}); ModelPackageContext ctx(package_root); const auto& variants = ctx.GetVariantInfos(); @@ -800,207 +576,19 @@ TEST(ModelPackageTest, ParseVariantsFromRoot_PackageRootDirectory) { EXPECT_EQ(v2->ep_compatibility.ep.value_or(""), "example_ep"); EXPECT_EQ(v2->ep_compatibility.device.value_or(""), "npu"); - std::filesystem::remove_all(package_root, ec); -} - -TEST(ModelPackageTest, ParseVariantsFromRoot_ComponentModelDirectory) { - const auto component_root = std::filesystem::temp_directory_path() / "ort_model_package_parse_from_component_root"; - std::error_code ec; - std::filesystem::remove_all(component_root, ec); - std::filesystem::create_directories(component_root); - - // package_root is a component model directory (has metadata.json, no manifest.json). - const auto variant_dir = component_root / "variant_1"; - std::filesystem::create_directories(variant_dir); - std::filesystem::copy_file("testdata/mul_1.onnx", variant_dir / "mul_1.onnx", - std::filesystem::copy_options::overwrite_existing, ec); - - constexpr std::string_view metadata_json = R"({ - "component_name": "model_1", - "variants": { - "variant_1": { - "ep": "example_ep", - "device": "cpu" - } - } - })"; - - { - std::ofstream os(component_root / "metadata.json", std::ios::binary); - os << metadata_json; - } - - { - std::ofstream os(variant_dir / "variant.json", std::ios::binary); - os << R"({ "filename": "mul_1.onnx" })"; - } - - ModelPackageContext ctx(component_root); - const auto& variants = ctx.GetVariantInfos(); - - ASSERT_EQ(variants.size(), 1u); - ASSERT_TRUE(variants[0].file.has_value()); - EXPECT_EQ(variants[0].file->model_file_path.filename().string(), "mul_1.onnx"); - - EXPECT_EQ(variants[0].ep_compatibility.ep.value_or(""), "example_ep"); - EXPECT_EQ(variants[0].ep_compatibility.device.value_or(""), "cpu"); - - std::filesystem::remove_all(component_root, ec); -} - -// ------------------------------------------------------------------ -// Tests for descriptor parser: enforced "ep" field in variant EP metadata. -// ------------------------------------------------------------------ -namespace { - -// Make a single-component, single-variant package on disk where metadata.json is written -// directly at the package root (the "single-component metadata flow" of the parser). -// In this flow variant EP metadata schema validation errors are propagated, instead of being -// swallowed by the manifest-driven discovery path which falls back to "Missing metadata variants". -// Returns the package_root. -std::filesystem::path MakeSingleComponentPackageWithMetadata(std::string_view subdir, - std::string_view metadata_json, - std::string_view variant_json = R"({"filename":"mul_1.onnx"})") { - const auto package_root = std::filesystem::temp_directory_path() / std::string(subdir); std::error_code ec; std::filesystem::remove_all(package_root, ec); - std::filesystem::create_directories(package_root); - - // Write metadata.json directly under package_root (no manifest, no models/ subdir). - { - std::ofstream os(package_root / "metadata.json", std::ios::binary); - os << metadata_json; - } - - // Variants live directly under package_root for the single-component flow. - const auto variant_dir = package_root / "variant_1"; - std::filesystem::create_directories(variant_dir); - std::filesystem::copy_file("testdata/mul_1.onnx", variant_dir / "mul_1.onnx", - std::filesystem::copy_options::overwrite_existing, ec); - - std::ofstream os(variant_dir / "variant.json", std::ios::binary); - os << variant_json; - - return package_root; } -} // namespace - -TEST(ModelPackageTest, ParserRejects_EpCompatibilityMissingEp) { - // The "ep" field is required in every variant descriptor. - // Omitting it must yield a parse error (not silently accept a wildcard / portable variant). - constexpr std::string_view metadata_json = R"({ - "component_name": "model_1", - "variants": { - "variant_1": { - "device": "cpu", - "compatibility_string": "anything" - } - } - })"; - const auto package_root = MakeSingleComponentPackageWithMetadata( - "ort_model_package_parser_missing_ep", metadata_json); - - try { - ModelPackageContext ctx(package_root); - FAIL() << "Expected exception for missing 'ep' field"; - } catch (const std::exception& e) { - EXPECT_NE(std::string(e.what()).find("ep"), std::string::npos) << e.what(); - } - - std::error_code ec; - std::filesystem::remove_all(package_root, ec); -} - -TEST(ModelPackageTest, ParserRejects_EpCompatibilityNullEp) { - constexpr std::string_view metadata_json = R"({ - "component_name": "model_1", - "variants": { - "variant_1": { - "ep": null, - "device": "cpu" - } - } - })"; - const auto package_root = MakeSingleComponentPackageWithMetadata( - "ort_model_package_parser_null_ep", metadata_json); - - try { - ModelPackageContext ctx(package_root); - FAIL() << "Expected exception for null 'ep' field"; - } catch (const std::exception& e) { - EXPECT_NE(std::string(e.what()).find("ep"), std::string::npos) << e.what(); - } - - std::error_code ec; - std::filesystem::remove_all(package_root, ec); -} - -TEST(ModelPackageTest, ParserRejects_EpCompatibilityEmptyEp) { - constexpr std::string_view metadata_json = R"({ - "component_name": "model_1", - "variants": { - "variant_1": { - "ep": "", - "device": "cpu" - } - } - })"; - const auto package_root = MakeSingleComponentPackageWithMetadata( - "ort_model_package_parser_empty_ep", metadata_json); - - try { - ModelPackageContext ctx(package_root); - FAIL() << "Expected exception for empty 'ep' field"; - } catch (const std::exception& e) { - EXPECT_NE(std::string(e.what()).find("ep"), std::string::npos) << e.what(); - } - - std::error_code ec; - std::filesystem::remove_all(package_root, ec); -} - -// ------------------------------------------------------------------ -// Tests for new pre-selection EP-compat traversal accessors. -// ------------------------------------------------------------------ TEST(ModelPackageApiTest, GetVariantEpName_ReturnsSingleEp) { const auto package_root = std::filesystem::temp_directory_path() / "ort_mp_pre_selection_ep_name"; - std::error_code ec; - std::filesystem::remove_all(package_root, ec); - - CreateManifestJson(package_root, MakeManifestJson("model_1")); - - const auto variant1_dir = package_root / "models" / "model_1" / "variant_1"; - const auto variant2_dir = package_root / "models" / "model_1" / "variant_2"; - std::filesystem::create_directories(variant1_dir); - std::filesystem::create_directories(variant2_dir); - std::filesystem::copy_file("testdata/mul_1.onnx", variant1_dir / "mul_1.onnx", - std::filesystem::copy_options::overwrite_existing, ec); - std::filesystem::copy_file("testdata/mul_1.onnx", variant2_dir / "mul_1.onnx", - std::filesystem::copy_options::overwrite_existing, ec); - - // Each variant declares a single EP. - constexpr std::string_view metadata_json = R"({ - "component_name": "model_1", - "variants": { - "variant_1": { - "ep": "example_ep", - "device": "cpu" - }, - "variant_2": { - "ep": "other_ep", - "device": "npu" - } - } - })"; - CreateComponentModelMetadata(package_root, "model_1", metadata_json); - - for (const auto& d : {variant1_dir, variant2_dir}) { - std::ofstream os(d / "variant.json", std::ios::binary); - os << R"({"filename":"mul_1.onnx"})"; - } + std::vector variants; + variants.push_back(VariantSpec{"variant_1", "example_ep", "cpu", "", "testdata/mul_1.onnx", {}, {}}); + variants.push_back(VariantSpec{"variant_2", "other_ep", "npu", "", "testdata/mul_1.onnx", {}, {}}); + BuildPackage(package_root, "model_1", variants); const auto& pkg_api = GetModelPackageFns(); + ASSERT_NE(pkg_api.CreateModelPackageContext, nullptr) << "Model package experimental API is not available"; auto context_deleter = [&pkg_api](OrtModelPackageContext* p) { if (p) pkg_api.ReleaseModelPackageContext(p); @@ -1010,14 +598,12 @@ TEST(ModelPackageApiTest, GetVariantEpName_ReturnsSingleEp) { ASSERT_ORTSTATUS_OK(pkg_api.CreateModelPackageContext(package_root.c_str(), &raw_ctx)); ctx.reset(raw_ctx); - // variant_1 targets example_ep const char* ep1 = nullptr; ASSERT_ORTSTATUS_OK(pkg_api.ModelPackage_GetVariantEpName( ctx.get(), "model_1", "variant_1", &ep1)); ASSERT_NE(ep1, nullptr); EXPECT_STREQ(ep1, "example_ep"); - // variant_2 targets other_ep const char* ep2 = nullptr; ASSERT_ORTSTATUS_OK(pkg_api.ModelPackage_GetVariantEpName( ctx.get(), "model_1", "variant_2", &ep2)); @@ -1028,47 +614,34 @@ TEST(ModelPackageApiTest, GetVariantEpName_ReturnsSingleEp) { ASSERT_ORTSTATUS_OK(pkg_api.ModelPackage_GetVariantEpName( ctx.get(), "model_1", "variant_1", nullptr)); + std::error_code ec; std::filesystem::remove_all(package_root, ec); } -// ------------------------------------------------------------------ -// ------------------------------------------------------------------ -// Test: variant selector tie-break is deterministic across repeated invocations. -// Two variants advertise compatibility for the same EP/device and EP returns the same -// validation score for both -- selection must be stable. -// ------------------------------------------------------------------ TEST(ModelPackageTest, VariantSelector_TieBreakIsDeterministic) { // Both variants point at the *same* model file (mul_1.onnx) so whichever wins works at runtime. // They advertise identical EP/device pairs and empty compatibility_string so the EP returns the - // same score (NOT_APPLICABLE) for both -- a tie. The fix in commit 27217da484 guarantees that - // ties resolve deterministically, i.e., selection is stable across repeated runs. + // same score (NOT_APPLICABLE) for both: ties must resolve deterministically across runs. RegisteredEpDeviceUniquePtr example_ep; ASSERT_NO_FATAL_FAILURE(Utils::RegisterAndGetExampleEp(*ort_env, Utils::example_ep_info, example_ep)); Ort::ConstEpDevice plugin_ep_device(example_ep.get()); - std::string first_selected_filename; + std::string first_selected_variant; for (int iter = 0; iter < 5; ++iter) { const auto package_root = std::filesystem::temp_directory_path() / "ort_mp_tie_break"; - std::error_code ec; - std::filesystem::remove_all(package_root, ec); - - CreateModelPackage(package_root, MakeManifestJson("model_1"), - "model_1", "variant_a", "variant_b", - std::filesystem::path{"testdata/mul_1.onnx"}, - std::filesystem::path{"testdata/mul_1.onnx"}); - - const std::string metadata_json = MakeMetadataJsonTwoVariants( - "model_1", - "variant_a", "example_ep", "cpu", "", - "variant_b", "example_ep", "cpu", ""); - CreateComponentModelMetadata(package_root, "model_1", metadata_json); + BuildTwoVariantPackage(package_root, + "variant_a", "cpu", "", + std::filesystem::path{"testdata/mul_1.onnx"}, + "variant_b", "cpu", "", + std::filesystem::path{"testdata/mul_1.onnx"}); Ort::SessionOptions session_options; std::unordered_map ep_options; session_options.AppendExecutionProvider_V2(*ort_env, {plugin_ep_device}, ep_options); const auto& pkg_api = GetModelPackageFns(); + ASSERT_NE(pkg_api.CreateModelPackageContext, nullptr) << "Model package experimental API is not available"; auto options_deleter = [&pkg_api](OrtModelPackageOptions* p) { if (p) pkg_api.ReleaseModelPackageOptions(p); }; auto context_deleter = [&pkg_api](OrtModelPackageContext* p) { if (p) pkg_api.ReleaseModelPackageContext(p); }; @@ -1095,61 +668,34 @@ TEST(ModelPackageTest, VariantSelector_TieBreakIsDeterministic) { ASSERT_ORTSTATUS_OK(pkg_api.ModelPackageComponent_GetSelectedVariantFolderPath(comp_ctx.get(), &selected_folder)); ASSERT_NE(selected_folder, nullptr); - // Path looks like .../models/model_1/ -- the folder name is the variant. + // Variant directories live at /model_1/; the leaf name is the variant. const auto selected_variant_dir = std::filesystem::path(selected_folder).filename().string(); ASSERT_TRUE(selected_variant_dir == "variant_a" || selected_variant_dir == "variant_b") << "unexpected variant dir: " << selected_variant_dir; if (iter == 0) { - first_selected_filename = selected_variant_dir; + first_selected_variant = selected_variant_dir; } else { - EXPECT_EQ(selected_variant_dir, first_selected_filename) + EXPECT_EQ(selected_variant_dir, first_selected_variant) << "tie-break selection drifted across runs (iter " << iter << ")"; } + std::error_code ec; std::filesystem::remove_all(package_root, ec); } } -// ------------------------------------------------------------------ -// Test: a variant's per-file `session_options` flow through OrtApis::AddSessionConfigEntry. -// We verify this by feeding a *known* typed key (session.intra_op_num_threads) a non-integer value: -// pre-change behavior would silently stuff it into AddConfigEntry and succeed; post-change -// behavior parses it via the typed dispatcher and fails CreateSession with a parse error. -// ------------------------------------------------------------------ TEST(ModelPackageTest, VariantSessionOptions_DispatchedThroughAddSessionConfigEntry) { + // Per-variant session_options assigns a typed key (session.intra_op_num_threads) a value that + // is not a valid integer. Routing this through OrtApis::AddSessionConfigEntry must reject it. const auto package_root = std::filesystem::temp_directory_path() / "ort_mp_session_options_dispatch"; - std::error_code ec; - std::filesystem::remove_all(package_root, ec); - - CreateManifestJson(package_root, MakeManifestJson("model_1")); - - const auto variant_dir = package_root / "models" / "model_1" / "variant_1"; - std::filesystem::create_directories(variant_dir); - std::filesystem::copy_file("testdata/mul_1.onnx", variant_dir / "mul_1.onnx", - std::filesystem::copy_options::overwrite_existing, ec); - - constexpr std::string_view metadata_json = R"({ - "component_name": "model_1", - "variants": { - "variant_1": { - "ep": "example_ep", "device": "cpu" - } - } - })"; - CreateComponentModelMetadata(package_root, "model_1", metadata_json); - - // Per-file session_options assigns a typed key (session.intra_op_num_threads) a value that is not a - // valid integer. Routing this through OrtApis::AddSessionConfigEntry (the new behavior) must reject it. - { - std::ofstream os(variant_dir / "variant.json", std::ios::binary); - os << R"({ - "filename": "mul_1.onnx", - "session_options": { - "session.intra_op_num_threads": "not_an_int" - } - })"; - } + std::vector variants; + variants.push_back(VariantSpec{ + "variant_1", "example_ep", "cpu", "", "testdata/mul_1.onnx", std::unordered_map{ + {"session.intra_op_num_threads", "not_an_int"}, + }, + {}}); + BuildPackage(package_root, "model_1", variants); RegisteredEpDeviceUniquePtr example_ep; ASSERT_NO_FATAL_FAILURE(Utils::RegisterAndGetExampleEp(*ort_env, Utils::example_ep_info, example_ep)); @@ -1160,6 +706,7 @@ TEST(ModelPackageTest, VariantSessionOptions_DispatchedThroughAddSessionConfigEn session_options.AppendExecutionProvider_V2(*ort_env, {plugin_ep_device}, ep_options); const auto& pkg_api = GetModelPackageFns(); + ASSERT_NE(pkg_api.CreateModelPackageContext, nullptr) << "Model package experimental API is not available"; auto options_deleter = [&pkg_api](OrtModelPackageOptions* p) { if (p) pkg_api.ReleaseModelPackageOptions(p); }; auto context_deleter = [&pkg_api](OrtModelPackageContext* p) { if (p) pkg_api.ReleaseModelPackageContext(p); }; @@ -1182,13 +729,9 @@ TEST(ModelPackageTest, VariantSessionOptions_DispatchedThroughAddSessionConfigEn ASSERT_ORTSTATUS_OK(pkg_api.SelectComponent(ctx.get(), "model_1", mp_opts.get(), &raw_comp_ctx)); comp_ctx.reset(raw_comp_ctx); - // CreateSession iterates the per-file session_options and dispatches each through OrtApis::AddSessionConfigEntry. - // The bad int value must surface as an error from this call. - // Pass nullptr for session_options so the metadata-merge path runs (it is skipped when the caller - // supplies their own session_options). + // Pass nullptr for session_options so the metadata-merge path runs. OrtSession* raw_session = nullptr; OrtStatus* st = pkg_api.CreateSession(*ort_env, comp_ctx.get(), /*session_options=*/nullptr, &raw_session); - // Clean up session first to avoid leaks if assertion fails. if (raw_session != nullptr) { Ort::GetApi().ReleaseSession(raw_session); raw_session = nullptr; @@ -1197,45 +740,32 @@ TEST(ModelPackageTest, VariantSessionOptions_DispatchedThroughAddSessionConfigEn const std::string err_msg = Ort::GetApi().GetErrorMessage(st); Ort::GetApi().ReleaseStatus(st); - // Message should mention either AddSessionConfigEntry or the typed-int parse failure. const bool mentions_dispatch = err_msg.find("AddSessionConfigEntry") != std::string::npos || err_msg.find("base-10 int32") != std::string::npos || err_msg.find("intra_op_num_threads") != std::string::npos; EXPECT_TRUE(mentions_dispatch) << "error did not mention typed dispatch: " << err_msg; + std::error_code ec; std::filesystem::remove_all(package_root, ec); } -// ------------------------------------------------------------------ -// Test: GetSelectedVariantFolderPath returns correct path even when variant.json is absent. -// ------------------------------------------------------------------ -TEST(ModelPackageApiTest, FolderPath_ReturnsCorrectPath_WhenVariantJsonAbsent) { - const auto package_root = std::filesystem::temp_directory_path() / "ort_mp_folder_path_no_variant_json"; +// GetSelectedVariantFolderPath returns the correct path even when the variant +// declares no executor_info (i.e., no `file` descriptor for the variant). +TEST(ModelPackageApiTest, FolderPath_ReturnsCorrectPath_WhenExecutorInfoAbsent) { + const auto package_root = std::filesystem::temp_directory_path() / "ort_mp_folder_path_no_executor_info"; + std::vector variants; + // No source_model => no executor_info is emitted for this variant. + VariantSpec only{"variant_1", "example_ep", "cpu", "", {}, {}, {}}; + variants.push_back(only); + BuildPackage(package_root, "model_1", variants); + + // Drop a model file in the variant directory so the package looks plausible on disk. std::error_code ec; - std::filesystem::remove_all(package_root, ec); - std::filesystem::create_directories(package_root); - - CreateManifestJson(package_root, MakeManifestJson("model_1")); - - const auto variant_dir = package_root / "models" / "model_1" / "variant_1"; - std::filesystem::create_directories(variant_dir); - - // Copy a model file but do NOT create variant.json - std::filesystem::copy_file("testdata/mul_1.onnx", variant_dir / "mul_1.onnx", + std::filesystem::copy_file("testdata/mul_1.onnx", + package_root / "model_1" / "variant_1" / "mul_1.onnx", std::filesystem::copy_options::overwrite_existing, ec); - constexpr std::string_view metadata_json = R"({ - "component_name": "model_1", - "variants": { - "variant_1": { - "ep": "example_ep", - "device": "cpu" - } - } - })"; - CreateComponentModelMetadata(package_root, "model_1", metadata_json); - RegisteredEpDeviceUniquePtr example_ep; ASSERT_NO_FATAL_FAILURE(Utils::RegisterAndGetExampleEp(*ort_env, Utils::example_ep_info, example_ep)); Ort::ConstEpDevice plugin_ep_device(example_ep.get()); @@ -1245,6 +775,7 @@ TEST(ModelPackageApiTest, FolderPath_ReturnsCorrectPath_WhenVariantJsonAbsent) { so.AppendExecutionProvider_V2(*ort_env, {plugin_ep_device}, ep_options); const auto& pkg_api = GetModelPackageFns(); + ASSERT_NE(pkg_api.CreateModelPackageContext, nullptr) << "Model package experimental API is not available"; OrtModelPackageOptions* raw_mp_opts = nullptr; ASSERT_ORTSTATUS_OK(pkg_api.CreateModelPackageOptionsFromSessionOptions(*ort_env, so, &raw_mp_opts)); @@ -1263,7 +794,6 @@ TEST(ModelPackageApiTest, FolderPath_ReturnsCorrectPath_WhenVariantJsonAbsent) { }; std::unique_ptr comp_ctx(raw_comp_ctx, component_context_deleter); - // GetSelectedVariantFolderPath should return the variant directory even without variant.json. const ORTCHAR_T* selected_folder = nullptr; ASSERT_ORTSTATUS_OK(pkg_api.ModelPackageComponent_GetSelectedVariantFolderPath(comp_ctx.get(), &selected_folder)); ASSERT_NE(selected_folder, nullptr); From 6f31cd68178f174648354a1e7a0a7b01bc30a85c Mon Sep 17 00:00:00 2001 From: Dmitri Smirnov Date: Wed, 1 Jul 2026 10:33:57 -0700 Subject: [PATCH 08/17] Harden ConvTranspose pad computation with SafeInt and consistency checks (#29446) This pull request strengthens validation and error handling for the ConvTranspose operator in ONNX Runtime, particularly when using explicit `output_shape` attributes. It adds comprehensive checks to prevent inconsistent or invalid configurations, improves arithmetic safety to guard against integer overflows, and introduces a suite of targeted unit tests to verify these behaviors. **Validation and Error Handling Improvements:** * Added stricter input validation in `ConvTransposeAttributes::ComputePadAndOutputShape` to ensure that all relevant parameters (`output_shape`, input size, stride, kernel, dilation, and output padding) are within valid ranges and to provide clear error messages when they are not. This includes checks that all values are positive and that output padding is non-negative. * Added a consistency check to verify that the explicit `output_shape` is compatible with the input dimensions and convolution parameters, preventing buffer overruns and logical inconsistencies. **Arithmetic Safety:** * Updated `ComputeTotalPad` to use `SafeInt` for all intermediate arithmetic, ensuring that integer overflows are detected and handled safely instead of producing undefined behavior. **Testing Enhancements:** * Added a comprehensive set of unit tests for `ConvTranspose` with explicit `output_shape`, including cases for invalid, inconsistent, and overflow-prone configurations, as well as valid edge cases (e.g., 1D, 2D, and 3D, large batch sizes, group > 1, and cases requiring padding). These tests verify that invalid configurations are rejected and that valid ones work as expected. These changes collectively improve the robustness, correctness, and maintainability of the ConvTranspose operator's implementation and its handling of explicit output shapes. --- onnxruntime/core/providers/common.h | 8 +- .../cpu/nn/conv_transpose_attributes.h | 54 ++- .../cpu/nn/conv_transpose_op_test.cc | 360 ++++++++++++++++++ 3 files changed, 417 insertions(+), 5 deletions(-) diff --git a/onnxruntime/core/providers/common.h b/onnxruntime/core/providers/common.h index aa20b88ef40cc..0e06dcb0f8144 100644 --- a/onnxruntime/core/providers/common.h +++ b/onnxruntime/core/providers/common.h @@ -155,9 +155,11 @@ inline Status ComputePadAndOutputShape(const int64_t in_dim, return Status::OK(); } -constexpr inline int64_t ComputeTotalPad(int64_t in_size, int64_t stride, int64_t adj, - int64_t kernel, int64_t dilation, int64_t out_size) { - return std::max(0, (in_size - 1) * stride + adj + (kernel - 1) * dilation + 1 - out_size); +inline int64_t ComputeTotalPad(int64_t in_size, int64_t stride, int64_t adj, + int64_t kernel, int64_t dilation, int64_t out_size) { + SafeInt safe_pad = (SafeInt(in_size) - 1) * stride + adj + + (SafeInt(kernel) - 1) * dilation + 1 - out_size; + return std::max(0, safe_pad); } inline void DistributePadding(AutoPadType pad_type, const int64_t& total_pad, diff --git a/onnxruntime/core/providers/cpu/nn/conv_transpose_attributes.h b/onnxruntime/core/providers/cpu/nn/conv_transpose_attributes.h index 1a14040215829..6b216080afdbb 100644 --- a/onnxruntime/core/providers/cpu/nn/conv_transpose_attributes.h +++ b/onnxruntime/core/providers/cpu/nn/conv_transpose_attributes.h @@ -305,14 +305,50 @@ struct ConvTransposeAttributes : public ConvAttributes { int64_t* out_size) const { // Output shape is explicitly provided - pad values will have to be computed if (*out_size != -1) { - if (*out_size < 0) { + if (*out_size <= 0) { return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, - "Explicit output size is negative: ", *out_size); + "Explicit output size must be positive. Got: ", *out_size); + } + if (in_size <= 0) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "Input spatial dimension must be positive. Got: ", in_size); + } + if (stride <= 0) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Stride must be positive. Got: ", stride); + } + if (kernel <= 0) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Kernel size must be positive. Got: ", kernel); + } + if (dilation <= 0) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Dilation must be positive. Got: ", dilation); + } + if (adj < 0) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "Output padding must be non-negative. Got: ", adj); } // total pad auto total_pad = ComputeTotalPad(in_size, stride, adj, kernel, dilation, *out_size); DistributePadding(pad_type, total_pad, *pad_head, *pad_tail); + + // Verify that the forward-conv re-derivation of input size from the output size and pads + // is consistent with the actual input size. Col2im re-derives the input spatial extent as: + // derived_in = (out_size + pad_head + pad_tail - dkernel) / stride + 1 + // If this exceeds in_size, Col2im would read past the col_buffer allocation. + // Note: derived_in < in_size is algebraically unreachable when adj >= 0 and total_pad >= 0, + // so this check is effectively one-sided (catches derived_in > in_size from oversized out_size). + SafeInt dkernel = (SafeInt(kernel) - 1) * dilation + 1; + int64_t derived_in = (SafeInt(*out_size) + *pad_head + *pad_tail - dkernel) / stride + 1; + if (derived_in != in_size) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "Explicit output_shape is inconsistent with input spatial dimensions" + " and convolution parameters. " + "Expected input size ", + derived_in, " but got ", in_size, + " (output_size=", *out_size, ", kernel=", kernel, + ", stride=", stride, ", dilation=", dilation, + ", output_padding=", adj, ")."); + } return Status::OK(); } @@ -346,6 +382,20 @@ struct ConvTransposeAttributes : public ConvAttributes { *out_size = SafeInt(in_size - 1) * stride + adj + SafeInt(kernel - 1) * dilation + 1 - *pad_head - *pad_tail; + + // Same consistency check as the explicit output_shape path: verify the forward-conv + // re-derivation of input size matches in_size. When output_padding (adj) >= stride + // (possible when dilation > stride passes the adj < max(stride, dilation) check), + // Col2im would compute a larger input extent and read past the col_buffer. + SafeInt dkernel2 = (SafeInt(kernel) - 1) * dilation + 1; + int64_t derived_in = (SafeInt(*out_size) + *pad_head + *pad_tail - dkernel2) / stride + 1; + if (derived_in != in_size) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "Computed output shape is inconsistent with input spatial dimensions. " + "output_padding (", + adj, ") may be too large for stride (", stride, + "). Expected input size ", derived_in, " but got ", in_size, "."); + } return Status::OK(); } }; diff --git a/onnxruntime/test/providers/cpu/nn/conv_transpose_op_test.cc b/onnxruntime/test/providers/cpu/nn/conv_transpose_op_test.cc index 7fff20784a19e..cfc9aee2adfbf 100644 --- a/onnxruntime/test/providers/cpu/nn/conv_transpose_op_test.cc +++ b/onnxruntime/test/providers/cpu/nn/conv_transpose_op_test.cc @@ -1729,5 +1729,365 @@ TEST(ConvTransposeTest, ConvTranspose_OutputPaddingExceedsStride) { {kTensorrtExecutionProvider, kQnnExecutionProvider, kDmlExecutionProvider, kWebGpuExecutionProvider}); } +// Test that an inconsistent explicit output_shape is rejected (output_shape too large +// relative to input spatial dimensions causes a pad/buffer size mismatch). +TEST(ConvTransposeTest, ConvTranspose_InconsistentOutputShape) { + OpTester test("ConvTranspose", 11); + test.AddShapeToTensorData(false); + // Input: 1x1x3x3, kernel 3x3, stride 1, no dilation. + // Natural output without padding = (3-1)*1 + 3 = 5. + // Setting output_shape to 100x100 is inconsistent. + test.AddAttribute("output_shape", std::vector{100, 100}); + test.AddInput("X", {1, 1, 3, 3}, std::vector(9, 1.0f)); + test.AddInput("W", {1, 1, 3, 3}, std::vector(9, 1.0f)); + test.AddOutput("Y", {0}, {}); + + // Exclude compiling EPs (TRT, QNN) that reject unsupported configs at partition time, + // DML which has its own validation, and CUDA/WebGPU which share ComputePadsAndOutputShape + // but may not be available in all builds. + test.Run(OpTester::ExpectResult::kExpectFailure, "inconsistent with input spatial dimensions", + {kTensorrtExecutionProvider, kQnnExecutionProvider, kDmlExecutionProvider, + kCudaExecutionProvider, kCudaNHWCExecutionProvider, kWebGpuExecutionProvider}); +} + +// Test that output_shape=0 is rejected. +TEST(ConvTransposeTest, ConvTranspose_ZeroOutputShape) { + OpTester test("ConvTranspose", 11); + test.AddShapeToTensorData(false); + test.AddAttribute("output_shape", std::vector{0, 0}); + test.AddInput("X", {1, 1, 3, 3}, std::vector(9, 1.0f)); + test.AddInput("W", {1, 1, 3, 3}, std::vector(9, 1.0f)); + test.AddOutput("Y", {0}, {}); + + test.Run(OpTester::ExpectResult::kExpectFailure, "output size must be positive", + {kTensorrtExecutionProvider, kQnnExecutionProvider, kDmlExecutionProvider, kWebGpuExecutionProvider}); +} + +// Test that an inconsistent 1D explicit output_shape is rejected. +TEST(ConvTransposeTest, ConvTranspose_1D_InconsistentOutputShape) { + OpTester test("ConvTranspose", 11); + test.AddShapeToTensorData(false); + // Input: 1x1x5, kernel_shape=3, stride=2, dilation=1. + // Natural (no-pad) output = (5-1)*2 + 3 = 11. output_shape=50 is way too large. + test.AddAttribute("kernel_shape", std::vector{3}); + test.AddAttribute("strides", std::vector{2}); + test.AddAttribute("output_shape", std::vector{50}); + test.AddInput("X", {1, 1, 5}, std::vector(5, 1.0f)); + test.AddInput("W", {1, 1, 3}, std::vector(3, 1.0f)); + test.AddOutput("Y", {0}, {}); + + test.Run(OpTester::ExpectResult::kExpectFailure, "inconsistent with input spatial dimensions", + {kTensorrtExecutionProvider, kQnnExecutionProvider, kDmlExecutionProvider, + kCudaExecutionProvider, kCudaNHWCExecutionProvider, kWebGpuExecutionProvider}); +} + +// Test that an inconsistent 3D explicit output_shape is rejected. +TEST(ConvTransposeTest, ConvTranspose_3D_InconsistentOutputShape) { + OpTester test("ConvTranspose", 11); + test.AddShapeToTensorData(false); + // Input: 1x1x2x2x2, kernel 2x2x2, stride 1, dilation 1. + // Natural output = (2-1)*1 + 2 = 3 per dim. output_shape=10x10x10 is too large. + test.AddAttribute("kernel_shape", std::vector{2, 2, 2}); + test.AddAttribute("output_shape", std::vector{10, 10, 10}); + test.AddInput("X", {1, 1, 2, 2, 2}, std::vector(8, 1.0f)); + test.AddInput("W", {1, 1, 2, 2, 2}, std::vector(8, 1.0f)); + test.AddOutput("Y", {0}, {}); + + // CUDA/WebGPU don't support 3D ConvTranspose in most builds. + test.Run(OpTester::ExpectResult::kExpectFailure, "inconsistent with input spatial dimensions", + {kTensorrtExecutionProvider, kQnnExecutionProvider, kDmlExecutionProvider, + kCudaExecutionProvider, kCudaNHWCExecutionProvider, kWebGpuExecutionProvider}); +} + +// Test that a valid 3D explicit output_shape with non-trivial padding works correctly. +TEST(ConvTransposeTest, ConvTranspose_3D_ValidOutputShape) { + ConvTransposeOpAttributes attrs = { + std::vector{2, 2, 2}, // kernel_shape + {}, // output_padding + std::vector{3, 3, 3}, // output_shape (natural no-pad output for 2x2x2 input, k=2, s=1) + std::vector{0, 0, 0, 0, 0, 0}, // pads + std::vector{1, 1, 1}, // strides + std::vector{1, 1, 1}, // dilations + 1, // group + "NOTSET" // auto_pad + }; + // Input 1x1x2x2x2 with all ones, kernel 1x1x2x2x2 with all ones. + // Output should be 1x1x3x3x3. Each output voxel sums overlapping kernel positions. + std::vector X(8, 1.0f); + std::vector W(8, 1.0f); + std::vector X_shape = {1, 1, 2, 2, 2}; + std::vector W_shape = {1, 1, 2, 2, 2}; + std::vector Y_shape = {1, 1, 3, 3, 3}; + // Corner=1, edge=2, face=4, center=8 (same as conv input for unit kernel). + std::vector expected_vals = { + 1.0f, 2.0f, 1.0f, 2.0f, 4.0f, 2.0f, 1.0f, 2.0f, 1.0f, + 2.0f, 4.0f, 2.0f, 4.0f, 8.0f, 4.0f, 2.0f, 4.0f, 2.0f, + 1.0f, 2.0f, 1.0f, 2.0f, 4.0f, 2.0f, 1.0f, 2.0f, 1.0f}; + TestConvTransposeOp(attrs, {X, W}, {X_shape, W_shape}, expected_vals, Y_shape, + OpTester::ExpectResult::kExpectSuccess, "", + {kTensorrtExecutionProvider, kCudaExecutionProvider, + kCudaNHWCExecutionProvider, kQnnExecutionProvider, + kDmlExecutionProvider, kWebGpuExecutionProvider}); +} + +// Test group > 1 with explicit output_shape. +TEST(ConvTransposeTest, ConvTranspose_2D_Group2_OutputShape) { + ConvTransposeOpAttributes attrs = { + std::vector{3, 3}, // kernel_shape + {}, // output_padding + std::vector{5, 5}, // output_shape: natural unpadded output for in=3, k=3, s=1 + std::vector{0, 0, 0, 0}, // pads + std::vector{1, 1}, // strides + std::vector{1, 1}, // dilations + 2, // group + "NOTSET" // auto_pad + }; + // X: 1x2x3x3 (2 input channels, group=2, so 1 channel per group) + // W: 2x1x3x3 (C=2, M/group=1, so output channels = 1*2 = 2) + std::vector X(18, 1.0f); + std::vector W(18, 1.0f); + std::vector X_shape = {1, 2, 3, 3}; + std::vector W_shape = {2, 1, 3, 3}; + std::vector Y_shape = {1, 2, 5, 5}; + // Each group produces a 5x5 output. With all-ones input (3x3) and all-ones kernel (3x3), + // it's the correlation of two 3x3 boxes producing the expected pattern. + std::vector expected_vals = { + 1.0f, 2.0f, 3.0f, 2.0f, 1.0f, + 2.0f, 4.0f, 6.0f, 4.0f, 2.0f, + 3.0f, 6.0f, 9.0f, 6.0f, 3.0f, + 2.0f, 4.0f, 6.0f, 4.0f, 2.0f, + 1.0f, 2.0f, 3.0f, 2.0f, 1.0f, + // Second group — identical + 1.0f, 2.0f, 3.0f, 2.0f, 1.0f, + 2.0f, 4.0f, 6.0f, 4.0f, 2.0f, + 3.0f, 6.0f, 9.0f, 6.0f, 3.0f, + 2.0f, 4.0f, 6.0f, 4.0f, 2.0f, + 1.0f, 2.0f, 3.0f, 2.0f, 1.0f}; + TestConvTransposeOp(attrs, {X, W}, {X_shape, W_shape}, expected_vals, Y_shape, + OpTester::ExpectResult::kExpectSuccess, "", + {kTensorrtExecutionProvider, kQnnExecutionProvider, + kOpenVINOExecutionProvider, kCudaNHWCExecutionProvider}); +} + +// Test with larger batch size and explicit output_shape. +TEST(ConvTransposeTest, ConvTranspose_2D_LargeBatch_OutputShape) { + ConvTransposeOpAttributes attrs = { + std::vector{2, 2}, // kernel_shape + {}, // output_padding + std::vector{3, 3}, // output_shape: (2-1)*1+2 = 3 + std::vector{0, 0, 0, 0}, // pads + std::vector{1, 1}, // strides + std::vector{1, 1}, // dilations + 1, // group + "NOTSET" // auto_pad + }; + // X: 4x1x2x2 (batch=4), W: 1x1x2x2 + std::vector X(16, 1.0f); + std::vector W = {1.0f, 1.0f, 1.0f, 1.0f}; + std::vector X_shape = {4, 1, 2, 2}; + std::vector W_shape = {1, 1, 2, 2}; + std::vector Y_shape = {4, 1, 3, 3}; + // Each batch image: ConvTranspose of 2x2 ones with 2x2 ones kernel → 3x3 + std::vector single_output = {1.0f, 2.0f, 1.0f, + 2.0f, 4.0f, 2.0f, + 1.0f, 2.0f, 1.0f}; + std::vector expected_vals; + for (int b = 0; b < 4; ++b) { + expected_vals.insert(expected_vals.end(), single_output.begin(), single_output.end()); + } + TestConvTransposeOp(attrs, {X, W}, {X_shape, W_shape}, expected_vals, Y_shape, + OpTester::ExpectResult::kExpectSuccess, "", + {kTensorrtExecutionProvider, kQnnExecutionProvider, kCudaNHWCExecutionProvider}); +} + +// Test that output_shape slightly larger than the natural maximum is caught (boundary case). +// Natural output = (3-1)*2 + 3 = 7. output_shape=9 requires negative padding and produces +// derived_in = (9-3)/2+1 = 4 != in_size=3. (output_shape=8 happens to pass due to integer +// division truncation: (8-3)/2+1 = 3 = in_size, so we use 9 for the true boundary.) +TEST(ConvTransposeTest, ConvTranspose_2D_Stride2_BoundaryOutputShape) { + OpTester test("ConvTranspose", 11); + test.AddShapeToTensorData(false); + test.AddAttribute("strides", std::vector{2, 2}); + test.AddAttribute("output_shape", std::vector{9, 9}); + test.AddInput("X", {1, 1, 3, 3}, std::vector(9, 1.0f)); + test.AddInput("W", {1, 1, 3, 3}, std::vector(9, 1.0f)); + test.AddOutput("Y", {0}, {}); + + test.Run(OpTester::ExpectResult::kExpectFailure, "inconsistent with input spatial dimensions", + {kTensorrtExecutionProvider, kQnnExecutionProvider, kDmlExecutionProvider, + kCudaExecutionProvider, kCudaNHWCExecutionProvider, kWebGpuExecutionProvider}); +} + +// Test that inconsistent output_shape with non-unit stride is caught. +TEST(ConvTransposeTest, ConvTranspose_2D_Stride2_InconsistentOutputShape) { + OpTester test("ConvTranspose", 11); + test.AddShapeToTensorData(false); + // Input: 1x1x3x3, kernel 3x3, stride 2, dilation 1. + // Natural output = (3-1)*2 + 3 = 7. output_shape=20x20 is too large. + test.AddAttribute("strides", std::vector{2, 2}); + test.AddAttribute("output_shape", std::vector{20, 20}); + test.AddInput("X", {1, 1, 3, 3}, std::vector(9, 1.0f)); + test.AddInput("W", {1, 1, 3, 3}, std::vector(9, 1.0f)); + test.AddOutput("Y", {0}, {}); + + test.Run(OpTester::ExpectResult::kExpectFailure, "inconsistent with input spatial dimensions", + {kTensorrtExecutionProvider, kQnnExecutionProvider, kDmlExecutionProvider, + kCudaExecutionProvider, kCudaNHWCExecutionProvider, kWebGpuExecutionProvider}); +} + +// Test that inconsistent output_shape with dilation > 1 is caught. +TEST(ConvTransposeTest, ConvTranspose_2D_Dilation_InconsistentOutputShape) { + OpTester test("ConvTranspose", 11); + test.AddShapeToTensorData(false); + // Input: 1x1x3x3, kernel 3x3, stride 1, dilation 2. + // dkernel = (3-1)*2+1 = 5. Natural output = (3-1)*1 + 5 = 7. + // output_shape=30x30 is too large. + test.AddAttribute("dilations", std::vector{2, 2}); + test.AddAttribute("output_shape", std::vector{30, 30}); + test.AddInput("X", {1, 1, 3, 3}, std::vector(9, 1.0f)); + test.AddInput("W", {1, 1, 3, 3}, std::vector(9, 1.0f)); + test.AddOutput("Y", {0}, {}); + + test.Run(OpTester::ExpectResult::kExpectFailure, "inconsistent with input spatial dimensions", + {kTensorrtExecutionProvider, kQnnExecutionProvider, kDmlExecutionProvider, + kCudaExecutionProvider, kCudaNHWCExecutionProvider, kWebGpuExecutionProvider}); +} + +// Test output_shape that is slightly smaller than natural (requiring positive padding). +// This is the normal legitimate use case for output_shape. +TEST(ConvTransposeTest, ConvTranspose_2D_OutputShape_RequiringPadding) { + ConvTransposeOpAttributes attrs = { + std::vector{3, 3}, // kernel_shape + {}, // output_padding + std::vector{3, 3}, // output_shape: smaller than natural (5), so pads will be added + std::vector{0, 0, 0, 0}, // pads (will be overwritten by computed pads) + std::vector{1, 1}, // strides + std::vector{1, 1}, // dilations + 1, // group + "NOTSET" // auto_pad + }; + // Input 1x1x3x3 all ones, kernel 1x1x3x3 all ones. + // Natural output = 5x5. With output_shape=3x3, total_pad = (3-1)*1+(3-1)*1+1-3 = 2 per dim. + // pad_head=1, pad_tail=1 (NOTSET → pad more on head). + // The result is the center 3x3 of the natural 5x5 output. + std::vector X(9, 1.0f); + std::vector W(9, 1.0f); + std::vector X_shape = {1, 1, 3, 3}; + std::vector W_shape = {1, 1, 3, 3}; + std::vector Y_shape = {1, 1, 3, 3}; + // Full 5x5 output would be: corner=1,edge=2,center area=3-9. + // With pad=1 on each side, we take center 3x3 of the 5x5 output, which equals: + // The center 3x3 of ConvTranspose(ones_3x3, ones_3x3, no padding). + // Full result: 1 2 3 2 1 / 2 4 6 4 2 / 3 6 9 6 3 / 2 4 6 4 2 / 1 2 3 2 1 + // Center 3x3: 4 6 4 / 6 9 6 / 4 6 4 + std::vector expected_vals = {4.0f, 6.0f, 4.0f, + 6.0f, 9.0f, 6.0f, + 4.0f, 6.0f, 4.0f}; + TestConvTransposeOp(attrs, {X, W}, {X_shape, W_shape}, expected_vals, Y_shape, + OpTester::ExpectResult::kExpectSuccess, "", + {kTensorrtExecutionProvider, kQnnExecutionProvider, + kOpenVINOExecutionProvider, kCudaNHWCExecutionProvider}); +} + +// Test that output_padding >= stride on the implicit path (no output_shape) is rejected +// when dilation > stride allows it past the adj < max(stride, dilation) pre-check. +TEST(ConvTransposeTest, ConvTranspose_ImplicitPath_OutputPaddingExceedsStride) { + OpTester test("ConvTranspose", 11); + test.AddShapeToTensorData(false); + // stride=2, dilation=3, kernel=3, output_padding=2. + // output_padding=2 < max(stride=2, dilation=3)=3, so it passes the ONNX-spec check. + // But adj=2 >= stride=2, causing Col2im to derive input size = in_size + 1 → OOB. + test.AddAttribute("strides", std::vector{2, 1}); + test.AddAttribute("dilations", std::vector{3, 1}); + test.AddAttribute("kernel_shape", std::vector{3, 1}); + test.AddAttribute("output_padding", std::vector{2, 0}); + test.AddInput("X", {1, 1, 3, 1}, std::vector(3, 1.0f)); + test.AddInput("W", {1, 1, 3, 1}, std::vector(3, 1.0f)); + test.AddOutput("Y", {0}, {}); + + test.Run(OpTester::ExpectResult::kExpectFailure, "inconsistent with input spatial dimensions", + {kTensorrtExecutionProvider, kQnnExecutionProvider, kDmlExecutionProvider, + kCudaExecutionProvider, kCudaNHWCExecutionProvider, kWebGpuExecutionProvider}); +} + +// Test that valid output_padding < stride on the implicit path works correctly. +// This exercises the adj term in the consistency check on the success side. +// Reuses the ConvTranspose_2D_outputpadding_strides2 configuration which is known correct. +TEST(ConvTransposeTest, ConvTranspose_2D_ValidOutputPadding_ConsistencyCheck) { + ConvTransposeOpAttributes attrs = { + std::vector{3, 3}, // kernel_shape + std::vector{1, 1}, // output_padding (adj=1 < stride=2) + {}, // output_shape (implicit) + std::vector{1, 1, 1, 1}, // pads + std::vector{2, 2}, // strides + std::vector{1, 1}, // dilations + 1, // group + "NOTSET" // auto_pad + }; + std::vector X_shape = {1, 1, 3, 3}; + std::vector X = {0.16857791f, -0.15161794f, 0.08540368f, + 0.1820628f, -0.21746576f, 0.08245695f, + 0.1431433f, -0.43156421f, 0.30591947f}; + std::vector W_shape = {1, 1, 3, 3}; + std::vector W = {-0.06230065f, 0.37932432f, -0.25388849f, + 0.33878803f, 0.43709868f, -0.22477469f, + 0.04118127f, -0.44696793f, 0.06373066f}; + std::vector Y_shape = {1, 1, 6, 6}; + std::vector expected_vals = { + 0.07368518f, -0.08925839f, -0.06627201f, 0.06301362f, 0.03732984f, -0.01919658f, + -0.00628807f, -0.02817563f, -0.01472169f, 0.04392925f, -0.00689478f, -0.01549204f, + 0.07957941f, -0.11459791f, -0.09505399f, 0.07681622f, 0.03604182f, -0.01853423f, + -0.0270785f, -0.00680824f, -0.06650258f, 0.08004665f, 0.07918708f, -0.0724144f, + 0.06256775f, -0.17838378f, -0.18863615f, 0.20064656f, 0.133717f, -0.06876295f, + -0.06398046f, -0.00864975f, 0.19289537f, -0.01490572f, -0.13673618f, 0.01949645f}; + TestConvTransposeOp(attrs, {X, W}, {X_shape, W_shape}, expected_vals, Y_shape, + OpTester::ExpectResult::kExpectSuccess, "", + {kTensorrtExecutionProvider, kQnnExecutionProvider, + kCudaNHWCExecutionProvider}); +} + +#if !defined(ORT_NO_EXCEPTIONS) +// Test that extreme attribute values causing arithmetic overflow are caught. +// SafeInt throws on overflow; in no-exceptions builds this aborts, so skip there. +TEST(ConvTransposeTest, ConvTranspose_OverflowInPadComputation) { + OpTester test("ConvTranspose", 11); + test.AddShapeToTensorData(false); + // dilation = (2^62 - 1), kernel = 3: (kernel-1) * dilation = 2 * (2^62 - 1) ≈ INT64_MAX → overflows + // when combined with the remaining terms. + constexpr int64_t kLargeDilation = (static_cast(1) << 62) - 1; // 4611686018427387903 + test.AddAttribute("strides", std::vector{1, 1}); + test.AddAttribute("dilations", std::vector{kLargeDilation, kLargeDilation}); + test.AddAttribute("kernel_shape", std::vector{3, 3}); + test.AddInput("X", {1, 1, 3, 3}, std::vector(9, 1.0f)); + test.AddInput("W", {1, 1, 3, 3}, std::vector(9, 1.0f)); + test.AddOutput("Y", {0}, {}); + + test.Run(OpTester::ExpectResult::kExpectFailure, "Integer overflow", + {kTensorrtExecutionProvider, kQnnExecutionProvider, kDmlExecutionProvider, + kCudaExecutionProvider, kCudaNHWCExecutionProvider, kWebGpuExecutionProvider}); +} + +// Test overflow in explicit output_shape path: large stride and dilation combined +// cause overflow in ComputeTotalPad. +TEST(ConvTransposeTest, ConvTranspose_OverflowInExplicitOutputShapePath) { + OpTester test("ConvTranspose", 11); + test.AddShapeToTensorData(false); + // stride = dilation = 2^61. (in_size-1)*stride = 2^62, (kernel-1)*dilation = 2^62. + // Sum ≈ 2^63 → overflows int64. + constexpr int64_t kLargeVal = static_cast(1) << 61; // 2305843009213693952 + test.AddAttribute("strides", std::vector{kLargeVal, 1}); + test.AddAttribute("dilations", std::vector{kLargeVal, 1}); + test.AddAttribute("kernel_shape", std::vector{3, 3}); + test.AddAttribute("output_shape", std::vector{5, 5}); + test.AddInput("X", {1, 1, 3, 3}, std::vector(9, 1.0f)); + test.AddInput("W", {1, 1, 3, 3}, std::vector(9, 1.0f)); + test.AddOutput("Y", {0}, {}); + + test.Run(OpTester::ExpectResult::kExpectFailure, "Integer overflow", + {kTensorrtExecutionProvider, kQnnExecutionProvider, kDmlExecutionProvider, + kCudaExecutionProvider, kCudaNHWCExecutionProvider, kWebGpuExecutionProvider}); +} +#endif // !defined(ORT_NO_EXCEPTIONS) + } // namespace test } // namespace onnxruntime From 36c6b7e1e6e0d4508f74d9cca7cbc583182cc970 Mon Sep 17 00:00:00 2001 From: Chi Lo <54722500+chilo-ms@users.noreply.github.com> Date: Wed, 1 Jul 2026 12:18:10 -0700 Subject: [PATCH 09/17] Fix brew install applesimutils failure by trusting wix/brew tap (#29450) Newer Homebrew versions refuse to load formulae from untrusted third-party taps. This adds `brew trust wix/brew` after tapping to allow the subsequent `brew install applesimutils` to succeed in React Native CI. **Changes:** - `.github/workflows/react_native.yml` - `tools/ci_build/github/azure-pipelines/templates/react-native-ci.yml` Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .github/workflows/react_native.yml | 1 + .../github/azure-pipelines/templates/react-native-ci.yml | 1 + 2 files changed, 2 insertions(+) diff --git a/.github/workflows/react_native.yml b/.github/workflows/react_native.yml index f93f7d7c2d08b..f9d5be61e8743 100644 --- a/.github/workflows/react_native.yml +++ b/.github/workflows/react_native.yml @@ -256,6 +256,7 @@ jobs: set -e -x npm install -g detox-cli brew tap wix/brew + brew trust wix/brew brew install applesimutils npm ci working-directory: ${{ github.workspace }}/js diff --git a/tools/ci_build/github/azure-pipelines/templates/react-native-ci.yml b/tools/ci_build/github/azure-pipelines/templates/react-native-ci.yml index 17f538c5f7ca3..55f5a20fb1d98 100644 --- a/tools/ci_build/github/azure-pipelines/templates/react-native-ci.yml +++ b/tools/ci_build/github/azure-pipelines/templates/react-native-ci.yml @@ -123,6 +123,7 @@ stages: - script: | brew tap wix/brew + brew trust wix/brew brew install applesimutils displayName: Install applesimutils From d64a76ccc8c326a23fe5a61678bd857bc5e4526e Mon Sep 17 00:00:00 2001 From: Ankit Maheshkar Date: Thu, 2 Jul 2026 01:22:28 +0530 Subject: [PATCH 10/17] [OVEP] OpenVINO Development Updates (#28954) ### Description This PR introduces minor development patches as listed below. - Fix the OVIR EpCtx model export fixing the changes introduced by https://github.com/microsoft/onnxruntime/pull/28725 - Add test for OVIR EpCtx Export on sample model. - Update OV Toolkit version to 2026.2 - Add test for OVEP Workload Type Test. --------- Signed-off-by: Jonathan Clohessy Signed-off-by: bfilipek Signed-off-by: dependabot[bot] Signed-off-by: Christian Bourjau Co-authored-by: Jaswanth Gannamaneni Co-authored-by: Ryan Metcalfe <107415876+RyanMetcalfeInt8@users.noreply.github.com> Co-authored-by: jatinwadhwa921 <110383850+jatinwadhwa921@users.noreply.github.com> Co-authored-by: derdeljan-msft Co-authored-by: Jonathan Clohessy Co-authored-by: Akshay Sonawane <111780983+apsonawane@users.noreply.github.com> Co-authored-by: Christopher Warrington Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> Co-authored-by: Ishwar Raut Co-authored-by: Gaurav Garg Co-authored-by: Xinpeng Dou <15529241576@163.com> Co-authored-by: Chi Lo <54722500+chilo-ms@users.noreply.github.com> Co-authored-by: adrastogi Co-authored-by: Aditya Rastogi Co-authored-by: qti-hungjuiw Co-authored-by: qti-yuduo Co-authored-by: Pradeep Sakhamoori Co-authored-by: Preetha Veeramalai Co-authored-by: Klimenko, Mikhail Co-authored-by: sfatimar Co-authored-by: Garth Long Co-authored-by: MayureshV1 <47039074+MayureshV1@users.noreply.github.com> Co-authored-by: Eric Crawford Co-authored-by: Vishnudas Thaniel S Co-authored-by: Javier Martinez Co-authored-by: Adam Pocock Co-authored-by: Changming Sun Co-authored-by: mingyue <131847423+mingyueliuh@users.noreply.github.com> Co-authored-by: Edward Chen <18449977+edgchen1@users.noreply.github.com> Co-authored-by: Susanta Bhattacharjee Co-authored-by: jatinwadhwa921 Co-authored-by: Jozef Wludzik Co-authored-by: Bartlomiej Filipek Co-authored-by: Kotomi-Du Co-authored-by: Rajeev Sekar Co-authored-by: liang Co-authored-by: n1harika Co-authored-by: Mayuresh M Varerkar Co-authored-by: Mikhail Dvoretckii Co-authored-by: bopeng1234 Co-authored-by: Copilot <198982749+Copilot@users.noreply.github.com> Co-authored-by: fs-eire <7679871+fs-eire@users.noreply.github.com> Co-authored-by: Wenqin Yang Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com> Co-authored-by: xieofxie Co-authored-by: hualxie Co-authored-by: Jiajia Qin Co-authored-by: Joshua Lochner Co-authored-by: Christian Bourjau Co-authored-by: Xiaofei Han Co-authored-by: Dmitri Smirnov Co-authored-by: chunghow-qti Co-authored-by: Guenther Schmuelling Co-authored-by: Jiawei Shao Co-authored-by: Tianlei Wu Co-authored-by: czekun Co-authored-by: Ryan Metcalfe Co-authored-by: Jaskaran Singh Nagi Co-authored-by: ai-fw-intg Co-authored-by: Rajeev Sekar Co-authored-by: RajeevSekar <117911837+RajeevSekar@users.noreply.github.com> Co-authored-by: Nazanin Beheshti --- .github/workflows/windows_openvino.yml | 8 +- cmake/onnxruntime_providers_openvino.cmake | 4 +- .../openvino/onnx_ctx_model_helper.cc | 17 +- .../openvino/ov_versions/capability.cc | 10 +- .../openvino/ov_versions/data_ops.cc | 10 +- .../providers/openvino/ov_versions/data_ops.h | 3 +- .../openvino/openvino_ep_context_test.cc | 230 ++++++++++++++++++ .../openvino_ep_workload_type_test.cc | 209 ++++++++++++++++ .../test/testdata/mul_1_ep_ctx_ovir.bin | Bin 0 -> 24 bytes .../test/testdata/mul_1_ep_ctx_ovir.onnx | Bin 0 -> 506 bytes .../test/testdata/mul_1_ep_ctx_ovir.xml | 56 +++++ .../x86_64/python/openvino/Dockerfile | 4 +- 12 files changed, 527 insertions(+), 24 deletions(-) create mode 100644 onnxruntime/test/providers/openvino/openvino_ep_workload_type_test.cc create mode 100644 onnxruntime/test/testdata/mul_1_ep_ctx_ovir.bin create mode 100644 onnxruntime/test/testdata/mul_1_ep_ctx_ovir.onnx create mode 100644 onnxruntime/test/testdata/mul_1_ep_ctx_ovir.xml diff --git a/.github/workflows/windows_openvino.yml b/.github/workflows/windows_openvino.yml index d87a4919a86fd..36cc1aebad8af 100644 --- a/.github/workflows/windows_openvino.yml +++ b/.github/workflows/windows_openvino.yml @@ -51,12 +51,12 @@ jobs: with: architecture: x64 - - name: Download OpenVINO Toolkit v2026.1.0 + - name: Download OpenVINO Toolkit v2026.2.1 env: - OpenVINOVersion: 2026.1.0 + OpenVINOVersion: 2026.2.1 shell: pwsh run: | - $Url ="https://storage.openvinotoolkit.org/repositories/openvino/packages/2026.1/windows_vc_mt/openvino_toolkit_windows_vc_mt_2026.1.0.21367.63e31528c62_x86_64.zip" + $Url ="https://storage.openvinotoolkit.org/repositories/openvino/packages/2026.2.1/windows_vc_mt/openvino_toolkit_windows_vc_mt_2026.2.1.21919.ede283a88e3_x86_64.zip" $OutputPath = "$env:RUNNER_TEMP\openvino.zip" $ExtractPath = "$env:RUNNER_TEMP\openvino-v$env:OpenVINOVersion" $TempExtractPath = "$env:RUNNER_TEMP\openvino_temp" @@ -99,7 +99,7 @@ jobs: shell: pwsh # Use $GITHUB_ENV to set the variable for subsequent steps run: | - $openVinoRootDir = Join-Path $env:RUNNER_TEMP "openvino-v2026.1.0" + $openVinoRootDir = Join-Path $env:RUNNER_TEMP "openvino-v2026.2.1" echo "OpenVINORootDir=$openVinoRootDir" >> $env:GITHUB_ENV - name: Print OpenVINORootDir after downloading OpenVINO diff --git a/cmake/onnxruntime_providers_openvino.cmake b/cmake/onnxruntime_providers_openvino.cmake index 882fc56d9a40b..1584d9f703575 100644 --- a/cmake/onnxruntime_providers_openvino.cmake +++ b/cmake/onnxruntime_providers_openvino.cmake @@ -13,8 +13,8 @@ # Header paths find_package(OpenVINO REQUIRED COMPONENTS Runtime ONNX) - if(OpenVINO_VERSION VERSION_LESS 2025.0) - message(FATAL_ERROR "OpenVINO 2025.0 and newer are supported. Please, use latest OpenVINO release") + if(OpenVINO_VERSION VERSION_LESS 2026.0) + message(FATAL_ERROR "OpenVINO 2026.0 and newer are supported. Please, use latest OpenVINO release") endif() if(OpenVINO_VERSION VERSION_GREATER_EQUAL 2024.4) diff --git a/onnxruntime/core/providers/openvino/onnx_ctx_model_helper.cc b/onnxruntime/core/providers/openvino/onnx_ctx_model_helper.cc index a21916be4c59b..b2519692a17f8 100644 --- a/onnxruntime/core/providers/openvino/onnx_ctx_model_helper.cc +++ b/onnxruntime/core/providers/openvino/onnx_ctx_model_helper.cc @@ -1,6 +1,7 @@ // Copyright (C) Intel Corporation // Licensed under the MIT License +#include #include #include #include @@ -243,10 +244,18 @@ std::shared_ptr EPCtxHandler::Initialize(const std::vectorDeserialize(ss); } } else { - ORT_THROW_IF_ERROR(utils::ValidateExternalDataPath(session_context.GetOutputModelPath(), - std::filesystem::path(ep_cache_context))); - std::filesystem::path ep_context_path = session_context.GetOutputModelPath().parent_path() / ep_cache_context; - if (ep_context_path.extension() != ".xml") { + const std::filesystem::path cache_context_path{ep_cache_context}; + // Compare the extension case-insensitively so that ".XML", ".Xml", etc. are also treated as XML. + std::string extension = cache_context_path.extension().string(); + std::transform(extension.begin(), extension.end(), extension.begin(), + [](unsigned char c) { return static_cast(std::tolower(c)); }); + const bool is_xml = (extension == ".xml"); + const std::filesystem::path& validation_base_path = is_xml + ? session_context.GetModelPath() + : session_context.GetOutputModelPath(); + ORT_THROW_IF_ERROR(utils::ValidateExternalDataPath(validation_base_path, cache_context_path)); + const std::filesystem::path ep_context_path = validation_base_path.parent_path() / cache_context_path; + if (!is_xml) { shared_context = shared_context_manager_->GetOrCreateSharedContext(ep_context_path); shared_context->Deserialize(); } diff --git a/onnxruntime/core/providers/openvino/ov_versions/capability.cc b/onnxruntime/core/providers/openvino/ov_versions/capability.cc index e932595f48565..d5bb021ccef44 100644 --- a/onnxruntime/core/providers/openvino/ov_versions/capability.cc +++ b/onnxruntime/core/providers/openvino/ov_versions/capability.cc @@ -41,16 +41,14 @@ GetCapability::GetCapability(const EPCtxHandler& ep_ctx_handler, npu_qdq_optimizer_enabled = true; // see data_ops.cc ~615 where we check for int16 types for gpu, this may change to a better approach later } -#if OPENVINO_VERSION_MAJOR == 2026 && OPENVINO_VERSION_MINOR == 1 +#if OPENVINO_VERSION_MAJOR == 2026 && OPENVINO_VERSION_MINOR == 2 + data_ops_ = std::make_unique(graph_viewer_, V_2026_2, device_type_, npu_qdq_optimizer_enabled); +#elif OPENVINO_VERSION_MAJOR == 2026 && OPENVINO_VERSION_MINOR == 1 data_ops_ = std::make_unique(graph_viewer_, V_2026_1, device_type_, npu_qdq_optimizer_enabled); #elif OPENVINO_VERSION_MAJOR == 2026 && OPENVINO_VERSION_MINOR == 0 data_ops_ = std::make_unique(graph_viewer_, V_2026_0, device_type_, npu_qdq_optimizer_enabled); -#elif OPENVINO_VERSION_MAJOR == 2025 && OPENVINO_VERSION_MINOR == 4 - data_ops_ = std::make_unique(graph_viewer_, V_2025_4, device_type_, npu_qdq_optimizer_enabled); -#elif OPENVINO_VERSION_MAJOR == 2025 && OPENVINO_VERSION_MINOR == 3 - data_ops_ = std::make_unique(graph_viewer_, V_2025_3, device_type_, npu_qdq_optimizer_enabled); #else - data_ops_ = std::make_unique(graph_viewer_, V_2026_1, device_type_, npu_qdq_optimizer_enabled); + data_ops_ = std::make_unique(graph_viewer_, V_2026_2, device_type_, npu_qdq_optimizer_enabled); #endif } diff --git a/onnxruntime/core/providers/openvino/ov_versions/data_ops.cc b/onnxruntime/core/providers/openvino/ov_versions/data_ops.cc index 37306d97b06ab..7aae81b987c0d 100644 --- a/onnxruntime/core/providers/openvino/ov_versions/data_ops.cc +++ b/onnxruntime/core/providers/openvino/ov_versions/data_ops.cc @@ -408,7 +408,7 @@ void DataOps::populate_op_mode_supported() { // populate unsupportedmode_t { - UnsupportedOpMode obj = {{V_2024_1, V_2024_2, V_2024_3, V_2024_4, V_2024_5, V_2024_6, V_2025_0, V_2025_1, V_2025_2, V_2025_3, V_2025_4, V_2026_0, V_2026_1}, + UnsupportedOpMode obj = {{V_2024_1, V_2024_2, V_2024_3, V_2024_4, V_2024_5, V_2024_6, V_2025_0, V_2025_1, V_2025_2, V_2025_3, V_2025_4, V_2026_0, V_2026_1, V_2026_2}, [this](const Node* node, const InitializedTensorSet&) { // If the Input of ReduceMax op is UINT8, it is rejected (Due to output mismatch) for (size_t i = 0; i < node->InputDefs().size(); i++) { @@ -425,7 +425,7 @@ void DataOps::populate_op_mode_supported() { { UnsupportedOpMode obj = {{V_2023_1, V_2023_2, V_2023_3, V_2024_0, V_2024_1, V_2024_2, V_2024_3, V_2024_4, V_2024_5, V_2024_6, V_2025_0, V_2025_1, - V_2025_2, V_2025_3, V_2025_4, V_2026_0, V_2026_1}, + V_2025_2, V_2025_3, V_2025_4, V_2026_0, V_2026_1, V_2026_2}, [this](const Node* node, const InitializedTensorSet&) { const auto& input_args = node->InputDefs(); const auto& input_arg = (input_args.size() > 1) ? input_args[1] : input_args[0]; @@ -445,7 +445,7 @@ void DataOps::populate_op_mode_supported() { { UnsupportedOpMode obj = {{V_2023_1, V_2023_2, V_2023_3, V_2024_0, V_2024_1, V_2024_2, V_2024_3, V_2024_4, V_2024_5, V_2024_6, V_2025_0, V_2025_1, - V_2025_2, V_2025_3, V_2025_4, V_2026_0, V_2026_1}, + V_2025_2, V_2025_3, V_2025_4, V_2026_0, V_2026_1, V_2026_2}, [this](const Node* node, const InitializedTensorSet&) { // If the operator is unsqueeze // If axes is an input, then we cannot produce a static graph. @@ -461,7 +461,7 @@ void DataOps::populate_op_mode_supported() { } { UnsupportedOpMode obj = {{V_2023_1, V_2023_2, V_2023_3, V_2024_0, V_2024_1, V_2024_2, V_2024_3, V_2024_4, V_2024_5, - V_2024_6, V_2025_0, V_2025_1, V_2025_2, V_2025_3, V_2025_4, V_2026_0, V_2026_1}, + V_2024_6, V_2025_0, V_2025_1, V_2025_2, V_2025_3, V_2025_4, V_2026_0, V_2026_1, V_2026_2}, [this](const Node* node, const InitializedTensorSet&) { // check for attributes auto& upsample_attr = node->GetAttributes(); @@ -492,7 +492,7 @@ void DataOps::populate_op_mode_supported() { { UnsupportedOpMode obj = {{V_2023_1, V_2023_2, V_2023_3, V_2024_0, V_2024_1, V_2024_2, V_2024_3, V_2024_4, V_2024_5, V_2024_6, V_2025_0, V_2025_1, - V_2025_2, V_2025_3, V_2025_4, V_2026_0, V_2026_1}, + V_2025_2, V_2025_3, V_2025_4, V_2026_0, V_2026_1, V_2026_2}, [this](const Node* node, const InitializedTensorSet&) { auto& attributes = node->GetAttributes(); if (attributes.count("coordinate_transformation_mode") > 0) { diff --git a/onnxruntime/core/providers/openvino/ov_versions/data_ops.h b/onnxruntime/core/providers/openvino/ov_versions/data_ops.h index af4bdc6efffff..1084e9204c8e5 100644 --- a/onnxruntime/core/providers/openvino/ov_versions/data_ops.h +++ b/onnxruntime/core/providers/openvino/ov_versions/data_ops.h @@ -40,7 +40,8 @@ enum versionNum { V_2025_3, V_2025_4, V_2026_0, - V_2026_1 + V_2026_1, + V_2026_2 }; using VersionNum = enum versionNum; diff --git a/onnxruntime/test/providers/openvino/openvino_ep_context_test.cc b/onnxruntime/test/providers/openvino/openvino_ep_context_test.cc index b5a8280a087a2..c085d17703496 100644 --- a/onnxruntime/test/providers/openvino/openvino_ep_context_test.cc +++ b/onnxruntime/test/providers/openvino/openvino_ep_context_test.cc @@ -1,12 +1,17 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. +#include #include +#include #include +#include +#include #include "core/framework/provider_options.h" #include "core/framework/tensor_shape.h" #include "core/common/float16.h" +#include "core/common/cpuid_info.h" #include "test/util/include/test_utils.h" #include "test/util/include/test/test_environment.h" @@ -26,6 +31,45 @@ using namespace onnxruntime::logging; extern std::unique_ptr ort_env; +namespace { + +// Returns true only on Intel CPUs. +// +// The OVIR EP context tests are gated on this because that path is currently +// validated only on Intel silicon. In particular, embed_mode = 0 dumps the +// compiled model to a separate .bin file and memory-maps it back on reload; +// this round-trip is unsupported on non-Intel CPUs (e.g. AMD), where it can +// crash. On those CPUs the OVIR EP context tests are skipped. +bool IsIntelCPU() { + return onnxruntime::CPUIDInfo::GetCPUIDInfo().GetCPUVendor() == "Intel"; +} + +// Runs a mul_1-style model (X[3,2] -> Y[3,2], Y = X * {1,2,3,4,5,6}) with +// X = all 2.0f and validates that Y == {2,4,6,8,10,12}. +void RunAndValidate(Ort::Session& session) { + const std::array input_shape = {3, 2}; + std::vector input_data(6, 2.0f); + Ort::MemoryInfo mem_info = Ort::MemoryInfo::CreateCpu(OrtAllocatorType::OrtDeviceAllocator, OrtMemTypeDefault); + Ort::Value input_tensor = Ort::Value::CreateTensor( + mem_info, input_data.data(), input_data.size(), input_shape.data(), input_shape.size()); + + const std::array input_names = {"X"}; + const std::array output_names = {"Y"}; + std::vector output_tensors(1); + + session.Run(Ort::RunOptions{nullptr}, input_names.data(), &input_tensor, 1, + output_names.data(), output_tensors.data(), 1); + + ASSERT_TRUE(output_tensors[0].IsTensor()); + ASSERT_EQ(output_tensors[0].GetTensorTypeAndShapeInfo().GetElementCount(), 6u); + + const float* out_data = output_tensors[0].GetTensorData(); + EXPECT_THAT(std::vector(out_data, out_data + 6), + ::testing::ElementsAre(2.f, 4.f, 6.f, 8.f, 10.f, 12.f)); +} + +} // namespace + class OVEPEPContextTests : public ::testing::Test { }; @@ -72,5 +116,191 @@ TEST_F(OVEPEPContextTests, OVEPEPContextFolderPath) { } } +// Runs an existing OVIR-encapsulated EP context model: "mul_1_ep_ctx_ovir.onnx" +// wraps a single EPContext node whose "ep_cache_context" points to a sibling +// OpenVINO IR (".xml" + ".bin"), so OVEP imports it via read_model()/ +// compile_model() instead of a pre-compiled blob. +// +// OVIR detection is filename-based (".onnx" -> ".xml"), so the model must be +// loaded from a path with the ".xml"/".bin" siblings next to it. +// +// CPU only. +class OVEPEPContextOVIRTests : public ::testing::Test { + protected: + void SetUp() override { + if (!IsIntelCPU()) { + GTEST_SKIP() << "OVIR EP context is only validated on Intel CPUs; skipping on non-Intel silicon."; + } + } + + static constexpr const char* kDevice = "CPU"; + static constexpr const ORTCHAR_T* kOvirModelPath = ORT_TSTR("testdata/mul_1_ep_ctx_ovir.onnx"); +}; + +TEST_F(OVEPEPContextOVIRTests, RunEpCtxOvirModel) { + ASSERT_TRUE(std::filesystem::exists(kOvirModelPath)) + << "Missing OVIR EP context model. Expected testdata/mul_1_ep_ctx_ovir.onnx " + "(with sibling .xml and .bin files)."; + + // Set up session options targeting CPU. + Ort::SessionOptions session_options; + std::unordered_map ov_options = {{"device_type", kDevice}}; + session_options.AppendExecutionProvider_OpenVINO_V2(ov_options); + + // Load the OVIR EP context model from disk (path-based load is required for + // OVIR encapsulation detection). + Ort::Session session(*ort_env, kOvirModelPath, session_options); + + RunAndValidate(session); +} + +// Negative / security test: an OVIR-encapsulated EP context model whose +// "ep_cache_context" attribute points outside the model directory via "../" +// traversal (e.g. "../../../etc/evil.xml") must be rejected at session-creation +// time rather than silently reading an arbitrary file off disk. + +TEST_F(OVEPEPContextOVIRTests, RejectsEpCacheContextPathTraversal) { + ASSERT_TRUE(std::filesystem::exists(kOvirModelPath)) + << "Missing OVIR EP context model. Expected testdata/mul_1_ep_ctx_ovir.onnx " + "(with sibling .xml and .bin files)."; + + // Load the known-good OVIR EP context model and rewrite its EPContext node so + // that ep_cache_context escapes the model directory. + ONNX_NAMESPACE::ModelProto model_proto; + ASSERT_STATUS_OK(Model::Load(kOvirModelPath, model_proto)); + + // Malicious relative path that escapes the model directory. The ".xml" + // extension routes validation through the OVIR ".xml" branch in + // EPCtxHandler::Initialize() (validated against the input model's directory), + // and "evil.xml" matches the "evil.onnx" output stem below so the node is also + // recognized as OVIR-encapsulated. + const std::string malicious_xml_path = "../../../etc/evil.xml"; + + bool patched = false; + for (auto& node : *model_proto.mutable_graph()->mutable_node()) { + if (node.op_type() != "EPContext") { + continue; + } + for (auto& attr : *node.mutable_attribute()) { + if (attr.name() == "embed_mode") { + attr.set_i(0); // force non-embed so the path (not an inline blob) is validated + } else if (attr.name() == "ep_cache_context") { + attr.set_s(malicious_xml_path); + patched = true; + } + } + } + ASSERT_TRUE(patched) << "Test model did not contain an EPContext ep_cache_context attribute to patch."; + + // Write the tampered model to a dedicated subfolder. The malicious ".xml" is + // intentionally never created on disk: validation must reject the path before + // any attempt to read it. + const std::filesystem::path out_dir = std::filesystem::path("testdata") / "ovir_epctx_path_traversal"; + std::filesystem::remove_all(out_dir); + std::filesystem::create_directories(out_dir); + const std::filesystem::path malicious_model = out_dir / "evil.onnx"; + { + std::ofstream ofs(malicious_model, std::ios::binary); + ASSERT_TRUE(ofs.is_open()) << "Failed to open " << malicious_model; + ASSERT_TRUE(model_proto.SerializeToOstream(&ofs)) << "Failed to serialize tampered model."; + } + + Ort::SessionOptions session_options; + std::unordered_map ov_options = {{"device_type", kDevice}}; + session_options.AppendExecutionProvider_OpenVINO_V2(ov_options); + + bool threw = false; + std::string error_message; + try { + Ort::Session session(*ort_env, malicious_model.c_str(), session_options); + } catch (const Ort::Exception& ex) { + threw = true; + error_message = ex.what(); + } + + std::filesystem::remove_all(out_dir); + + ASSERT_TRUE(threw) + << "Session creation should have rejected the path-traversal ep_cache_context, but it succeeded."; + EXPECT_THAT(error_message, ::testing::HasSubstr("escapes model directory")) + << "Expected a path-escape rejection. Actual error: " << error_message; +} + +// Generates an EP context model from the OVIR-encapsulated source model and +// then loads + runs the generated model, covering both EP context embed modes: +// embed_mode = 1: the compiled context is serialized INLINE into the .onnx. +// embed_mode = 0: the compiled context is dumped to a separate file and only +// its filename is stored in the .onnx EPContext node. +// +// embed_mode is only honored during generation (it is written into the +// EPContext node), so it must be exercised via a generate-then-run flow rather +// than by setting the option on a run-only session. +// +// CPU only. Parameter: embed_mode_enabled. +class OVEPOVIRModelsExportEPContextTests : public ::testing::TestWithParam { + protected: + void SetUp() override { + if (!IsIntelCPU()) { + GTEST_SKIP() << "OVIR EP context export is only validated on Intel CPUs; skipping on non-Intel silicon."; + } + } + + static constexpr const char* kDevice = "CPU"; + static constexpr const ORTCHAR_T* kOvirModelPath = ORT_TSTR("testdata/mul_1_ep_ctx_ovir.onnx"); +}; + +TEST_P(OVEPOVIRModelsExportEPContextTests, ExportEpCtxFromOVIRModel) { + const bool embed_mode = GetParam(); + + ASSERT_TRUE(std::filesystem::exists(kOvirModelPath)) + << "Missing OVIR EP context model. Expected testdata/mul_1_ep_ctx_ovir.onnx " + "(with sibling .xml and .bin files)."; + + // Generate the EP context model into a dedicated subfolder so that the + // separately-dumped blob (embed_mode = 0) doesn't collide with testdata. + const std::filesystem::path out_dir = + std::filesystem::path("testdata") / (std::string("ovir_epctx_export_embed_") + (embed_mode ? "on" : "off")); + std::filesystem::remove_all(out_dir); + std::filesystem::create_directories(out_dir); + const std::filesystem::path epctx_model = out_dir / "mul_1_ovir_epctx_export.onnx"; + + // --- Generate EP context model --- + { + Ort::SessionOptions session_options; + session_options.AddConfigEntry(kOrtSessionOptionEpContextEnable, "1"); + session_options.AddConfigEntry(kOrtSessionOptionEpContextFilePath, epctx_model.string().c_str()); + session_options.AddConfigEntry(kOrtSessionOptionEpContextEmbedMode, embed_mode ? "1" : "0"); + std::unordered_map ov_options = {{"device_type", kDevice}}; + session_options.AppendExecutionProvider_OpenVINO_V2(ov_options); + + // Creating the session triggers EP context export to epctx_model. + Ort::Session session(*ort_env, kOvirModelPath, session_options); + } + + ASSERT_TRUE(std::filesystem::exists(epctx_model)) + << "EP context model was not generated at " << epctx_model; + + // --- Load + run the generated EP context model --- + { + Ort::SessionOptions session_options; + std::unordered_map ov_options = {{"device_type", kDevice}}; + session_options.AppendExecutionProvider_OpenVINO_V2(ov_options); + + Ort::Session session(*ort_env, epctx_model.c_str(), session_options); + + RunAndValidate(session); + } + + std::filesystem::remove_all(out_dir); +} + +INSTANTIATE_TEST_SUITE_P( + OVEP_Tests, + OVEPOVIRModelsExportEPContextTests, + ::testing::Bool(), + [](const ::testing::TestParamInfo& info) { + return std::string("embed_") + (info.param ? "on" : "off"); + }); + } // namespace test } // namespace onnxruntime diff --git a/onnxruntime/test/providers/openvino/openvino_ep_workload_type_test.cc b/onnxruntime/test/providers/openvino/openvino_ep_workload_type_test.cc new file mode 100644 index 0000000000000..7d042b2c2e1ee --- /dev/null +++ b/onnxruntime/test/providers/openvino/openvino_ep_workload_type_test.cc @@ -0,0 +1,209 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include +#include +#include +#include +#include +#include + +#include "core/session/onnxruntime_cxx_api.h" +#include "gtest/gtest.h" + +extern std::unique_ptr ort_env; + +constexpr const ORTCHAR_T* kSqueezeNetModelUri = + ORT_TSTR("testdata/squeezenet/model.onnx"); + +class OVEPWorkloadTypeTests : public ::testing::Test { + protected: + // Check whether the NPU device can be registered at all. + static bool IsNPUAvailable() { + try { + Ort::SessionOptions opts; + std::unordered_map ov; + ov["device_type"] = "NPU"; + opts.AppendExecutionProvider_OpenVINO_V2(ov); + return true; + } catch (...) { + return false; + } + } + + // Allow NPU resources to be fully released between tests. + // Without this delay the NPU driver may fail to re-initialise. + // Skip the delay when the test was skipped (e.g. no NPU available), since no + // NPU resources were acquired and the sleep would only add build latency. + void TearDown() override { + if (IsSkipped()) { + return; + } + std::this_thread::sleep_for(std::chrono::milliseconds(50)); + } + + static Ort::Session CreateSqueezeNetSession( + Ort::SessionOptions& session_options, + std::unordered_map& ov_options) { + session_options.SetIntraOpNumThreads(1); + session_options.SetGraphOptimizationLevel( + GraphOptimizationLevel::ORT_ENABLE_ALL); + session_options.AppendExecutionProvider_OpenVINO_V2(ov_options); + + return Ort::Session(*ort_env, kSqueezeNetModelUri, session_options); + } + + // Run a single inference on the SqueezeNet session and return output + static std::vector RunSqueezeNet(Ort::Session& session, + const std::string& phase_label) { + Ort::AllocatorWithDefaultOptions allocator; + std::string input_name = + session.GetInputNameAllocated(0, allocator).get(); + std::string output_name = + session.GetOutputNameAllocated(0, allocator).get(); + const char* input_names[] = {input_name.c_str()}; + const char* output_names[] = {output_name.c_str()}; + + // SqueezeNet input: 1 × 3 × 224 × 224 = 150 528 floats + std::vector input_shape = {1, 3, 224, 224}; + constexpr size_t kInputSize = 1 * 3 * 224 * 224; + std::vector input_values(kInputSize, 1.0f); + + Ort::MemoryInfo mem_info = + Ort::MemoryInfo::CreateCpu(OrtArenaAllocator, OrtMemTypeDefault); + + Ort::Value input_tensor = Ort::Value::CreateTensor( + mem_info, input_values.data(), input_values.size(), + input_shape.data(), input_shape.size()); + + auto outputs = session.Run(Ort::RunOptions{nullptr}, input_names, + &input_tensor, 1, output_names, 1); + + EXPECT_EQ(outputs.size(), 1u) << phase_label; + if (outputs.empty()) return {}; + + auto type_shape = outputs[0].GetTensorTypeAndShapeInfo(); + std::vector out_shape = type_shape.GetShape(); + + // Expected: [1, 1000, 1, 1] + EXPECT_EQ(out_shape.size(), 4u) << phase_label; + if (out_shape.size() == 4u) { + EXPECT_EQ(out_shape[0], 1) << phase_label; + EXPECT_EQ(out_shape[1], 1000) << phase_label; + EXPECT_EQ(out_shape[2], 1) << phase_label; + EXPECT_EQ(out_shape[3], 1) << phase_label; + } + + size_t num_elements = type_shape.GetElementCount(); + EXPECT_EQ(num_elements, 1000u) << phase_label; + + const float* out_data = outputs[0].GetTensorData(); + std::vector result(out_data, out_data + num_elements); + + for (size_t i = 0; i < num_elements; ++i) { + EXPECT_TRUE(std::isfinite(result[i])) + << phase_label << " index " << i << " is not finite"; + } + + return result; + } + + // Compare two output vectors element-wise within a tolerance. + static void CompareOutputs(const std::vector& expected, + const std::vector& actual, + const std::string& label, + float tolerance = 1e-4f) { + ASSERT_EQ(expected.size(), actual.size()) << label << " size mismatch"; + for (size_t i = 0; i < expected.size(); ++i) { + EXPECT_NEAR(expected[i], actual[i], tolerance) + << label << " mismatch at index " << i; + } + } +}; + +namespace onnxruntime { +namespace test { + +// Test 1: Dynamic workload-type switching with consistency check +// Baseline (no workload type) → Efficient → Default +TEST_F(OVEPWorkloadTypeTests, OVEPWorkloadTypeDynamicSwitch) { + if (!IsNPUAvailable()) { + GTEST_SKIP() << "NPU device not available, skipping workload type test"; + } + + Ort::SessionOptions session_options; + std::unordered_map ov_options; + ov_options["device_type"] = "NPU"; + + Ort::Session session = CreateSqueezeNetSession(session_options, ov_options); + + const char* const keys[] = {"ep.dynamic.workload_type"}; + + // Phase 1: Baseline (no workload type set) + auto baseline_output = RunSqueezeNet(session, "Baseline"); + + // Phase 2: Switch to Efficient + const char* const eff_val[] = {"Efficient"}; + session.SetEpDynamicOptions(keys, eff_val, 1); + auto efficient_output = RunSqueezeNet(session, "Efficient"); + + // Phase 3: Switch to Default + const char* const def_val[] = {"Default"}; + session.SetEpDynamicOptions(keys, def_val, 1); + auto default_output = RunSqueezeNet(session, "Default"); + + // All modes should produce the same results + CompareOutputs(baseline_output, efficient_output, + "Baseline vs Efficient"); + CompareOutputs(baseline_output, default_output, + "Baseline vs Default"); +} + +// Test 2: Multiple inferences per workload mode +// Runs 10 inferences in each mode: +// Baseline × 10 → Efficient × 10 → Default × 10 +TEST_F(OVEPWorkloadTypeTests, OVEPWorkloadTypeMultipleInferencesPerMode) { + if (!IsNPUAvailable()) { + GTEST_SKIP() << "NPU device not available, skipping workload type test"; + } + + Ort::SessionOptions session_options; + std::unordered_map ov_options; + ov_options["device_type"] = "NPU"; + + Ort::Session session = CreateSqueezeNetSession(session_options, ov_options); + + const char* const keys[] = {"ep.dynamic.workload_type"}; + const char* const eff_val[] = {"Efficient"}; + const char* const def_val[] = {"Default"}; + + constexpr int kIterationsPerMode = 10; + + // Phase 1: Baseline – 10 runs without workload type + // Save the first run as the reference output. + auto reference_output = RunSqueezeNet(session, "Baseline iter 0"); + for (int i = 1; i < kIterationsPerMode; ++i) { + auto output = RunSqueezeNet(session, "Baseline iter " + std::to_string(i)); + CompareOutputs(reference_output, output, + "Baseline iter " + std::to_string(i) + " vs reference"); + } + + // Phase 2: Efficient – 10 runs + session.SetEpDynamicOptions(keys, eff_val, 1); + for (int i = 0; i < kIterationsPerMode; ++i) { + auto output = RunSqueezeNet(session, "Efficient iter " + std::to_string(i)); + CompareOutputs(reference_output, output, + "Efficient iter " + std::to_string(i) + " vs reference"); + } + + // Phase 3: Default – 10 runs + session.SetEpDynamicOptions(keys, def_val, 1); + for (int i = 0; i < kIterationsPerMode; ++i) { + auto output = RunSqueezeNet(session, "Default iter " + std::to_string(i)); + CompareOutputs(reference_output, output, + "Default iter " + std::to_string(i) + " vs reference"); + } +} + +} // namespace test +} // namespace onnxruntime diff --git a/onnxruntime/test/testdata/mul_1_ep_ctx_ovir.bin b/onnxruntime/test/testdata/mul_1_ep_ctx_ovir.bin new file mode 100644 index 0000000000000000000000000000000000000000..cc239b1d366323c00843f32cbdb3258718f23488 GIT binary patch literal 24 ZcmZQzXs~BsU~m8;AZ`HS1weej0RSmO1Rwwa literal 0 HcmV?d00001 diff --git a/onnxruntime/test/testdata/mul_1_ep_ctx_ovir.onnx b/onnxruntime/test/testdata/mul_1_ep_ctx_ovir.onnx new file mode 100644 index 0000000000000000000000000000000000000000..2c616129077efdd07e6dcd88c16917f57deecbca GIT binary patch literal 506 zcmZ9Iy-ve06oqjMDBP5iwjj!n7&0(a$bw2N9RnR0C|N+jVx{hdTEuqbxNYER`Uboj zkAMTML>WBzJNKTWb3N#NR^W>Od2EuWL;p%Ei-jHN5r$j+c=7Et08kEb#H172YjaZR zoI!vY;|U)l-o2#VQfxs@ zCPFQ3+NAW~0Z%3c%f-hc@Dw(4Juw`^gS!mX!_DBXUm%}Y=N20mHbPUn?SX%dQ^D~n zP4Yaf4&r-*Dtmi^kcSY$8KGz^tn~upbj;UG(D(zRre}G2cQ1KXdhZD zW>oRHvwLTB(A{G}wsRu=_KfQ{j$lRFYCv&im0W1j|Wl+y#VsbuP!y uUee6CCOD&i)Dt*d4aXuwBP1%Dq;9(bt4$qpG1|P+?A>a^$27610`d(5Q + + + + + + + 3 + 2 + + + + + + + + 3 + 2 + + + + + + + + 3 + 2 + + + 3 + 2 + + + + + 3 + 2 + + + + + + + 3 + 2 + + + + + + + + + + + diff --git a/tools/ci_build/github/linux/docker/inference/x86_64/python/openvino/Dockerfile b/tools/ci_build/github/linux/docker/inference/x86_64/python/openvino/Dockerfile index 03f351d942e70..17038fd4e080d 100644 --- a/tools/ci_build/github/linux/docker/inference/x86_64/python/openvino/Dockerfile +++ b/tools/ci_build/github/linux/docker/inference/x86_64/python/openvino/Dockerfile @@ -22,8 +22,8 @@ RUN dnf install -y --nodocs \ && dnf clean all \ && rm -rf /var/cache/dnf -ENV INTEL_OPENVINO_DIR=/opt/intel/openvino_2026.1.0 -ARG OPENVINO_PACKAGE_URL=https://storage.openvinotoolkit.org/repositories/openvino/packages/2026.1/linux/openvino_toolkit_rhel8_2026.1.0.21367.63e31528c62_x86_64.tgz +ENV INTEL_OPENVINO_DIR=/opt/intel/openvino_2026.2.1 +ARG OPENVINO_PACKAGE_URL=https://storage.openvinotoolkit.org/repositories/openvino/packages/2026.2.1/linux/openvino_toolkit_rhel8_2026.2.1.21919.ede283a88e3_x86_64.tgz ARG TEMP_DIR=/tmp/openvino_installer RUN mkdir -p ${TEMP_DIR} && \ From 0b5580da67c9342ef6b007cb5e376703b811afde Mon Sep 17 00:00:00 2001 From: Jambay Kinley Date: Wed, 1 Jul 2026 14:14:45 -0700 Subject: [PATCH 11/17] [CUDA] Batched small-M GEMV for MatMulNBits (4-bit/8-bit) (#29451) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ### Description Add a batched GEMV for the small-M range (M = 2..16) of CUDA `MatMulNBits` (4-bit and 8-bit), using the standard `[N, blocks, blob]` weight layout with **no prepacking**. Previously, `MatMulNBits` had a fast single-row GEMV for M = 1 (decode), but for M > 1 it fell back to full weight dequantization + cuBLAS, which dequantizes the entire weight matrix regardless of M. That fallback is therefore flat across small M and dominates latency at the row counts that matter for multi-row decode. The new half/bf16 path uses a `CtaM x CtaN` register-tiled kernel that streams the quantized weight once per block row and reuses each activation load across columns, so latency scales with M. M = 1 decode is unchanged. No prepacking is used, so there is no extra resident weight memory and no GEMM tactic profiling at session init. Also adds `profile_matmul_nbits.py` (same style as `profile_qmoe_gemv.py`, parseable with `parse_nsys.py`) and a docs experiment log. **Before (main: M > 1 dequant + cuBLAS) vs After (batched small-M GEMV)** — A100, block_size 32, fp16, average op latency in microseconds (lower is better). #### 4-bit (M = 2..16) Before: | matrix | K | N | M=1 | M=2 | M=4 | M=8 | M=16 | |---------|-------|--------|-------|--------|--------|--------|--------| | qkv | 4096 | 4096 | 25.9 | 77.2 | 69.3 | 69.5 | 69.9 | | o_proj | 4096 | 4096 | 23.2 | 69.2 | 69.1 | 69.3 | 69.7 | | gate_up | 4096 | 12288 | 39.3 | 172.0 | 172.1 | 172.1 | 172.4 | | down | 12288 | 4096 | 38.8 | 174.8 | 175.1 | 175.4 | 175.6 | | lm_head | 4096 | 151936 | 301.6 | 1868.0 | 1871.1 | 1877.8 | 1885.1 | After: | matrix | K | N | M=1 | M=2 | M=4 | M=8 | M=16 | |---------|-------|--------|-------|-------|-------|-------|--------| | qkv | 4096 | 4096 | 25.4 | 30.1 | 32.3 | 43.2 | 70.0 | | o_proj | 4096 | 4096 | 22.6 | 26.4 | 28.1 | 36.7 | 58.7 | | gate_up | 4096 | 12288 | 37.5 | 43.5 | 49.5 | 71.9 | 125.9 | | down | 12288 | 4096 | 37.9 | 47.4 | 52.9 | 76.8 | 150.1 | | lm_head | 4096 | 151936 | 300.7 | 329.9 | 424.1 | 635.3 | 1226.9 | Speedup (before / after): | matrix | M=2 | M=4 | M=8 | M=16 | |---------|-------|-------|-------|-------| | qkv | 2.56x | 2.15x | 1.61x | 1.00x | | o_proj | 2.62x | 2.46x | 1.89x | 1.19x | | gate_up | 3.95x | 3.48x | 2.39x | 1.37x | | down | 3.69x | 3.31x | 2.28x | 1.17x | | lm_head | 5.66x | 4.41x | 2.96x | 1.54x | #### 8-bit (M = 2..5) 8-bit weights are twice the bytes of 4-bit and the GEMV runs on CUDA cores, so it crosses over to the dequantize + cuBLAS (tensor-core) fallback at a lower M; the batched path covers M = 2..5 and M >= 6 keeps the fallback. Before: | matrix | K | N | M=1 | M=2 | M=3 | M=4 | M=5 | |---------|-------|--------|-------|--------|--------|--------|--------| | qkv | 4096 | 4096 | 36.2 | 80.6 | 72.9 | 72.8 | 73.3 | | o_proj | 4096 | 4096 | 31.1 | 73.2 | 72.7 | 73.6 | 73.4 | | gate_up | 4096 | 12288 | 63.5 | 184.4 | 184.2 | 184.1 | 184.3 | | down | 12288 | 4096 | 67.8 | 187.3 | 187.4 | 187.8 | 188.0 | | lm_head | 4096 | 151936 | 535.9 | 2025.0 | 2025.9 | 2028.1 | 2029.8 | After: | matrix | K | N | M=1 | M=2 | M=3 | M=4 | M=5 | |---------|-------|--------|-------|-------|-------|--------|--------| | qkv | 4096 | 4096 | 36.2 | 46.1 | 48.7 | 57.9 | 67.0 | | o_proj | 4096 | 4096 | 31.6 | 39.5 | 48.6 | 57.9 | 67.1 | | gate_up | 4096 | 12288 | 63.4 | 80.9 | 104.6 | 128.9 | 152.6 | | down | 12288 | 4096 | 68.0 | 96.2 | 119.3 | 146.4 | 172.0 | | lm_head | 4096 | 151936 | 536.3 | 647.0 | 896.9 | 1157.1 | 1420.1 | Speedup (before / after): | matrix | M=2 | M=3 | M=4 | M=5 | |---------|-------|-------|-------|-------| | qkv | 1.75x | 1.50x | 1.26x | 1.09x | | o_proj | 1.85x | 1.50x | 1.27x | 1.09x | | gate_up | 2.28x | 1.76x | 1.43x | 1.21x | | down | 1.95x | 1.57x | 1.28x | 1.09x | | lm_head | 3.13x | 2.26x | 1.75x | 1.43x | ### Motivation and Context The small-M (M = 2..16) regime is exactly what speculative decoding hits when verifying a block of draft tokens: each step runs the target model on a handful of rows rather than a single token. With the previous dispatch, that step paid the full-dequant + cuBLAS cost (flat ~172 us for an MLP projection, ~1.9 ms for the lm_head even at M = 2), which erased much of the speculative-decoding speedup over greedy decoding. Routing these row counts through a batched GEMV that scales with M (2.6-5.7x faster at M = 2, at or above parity through M = 16) restores the benefit, without the resident-memory and session-init costs of a prepacked weight layout. --------- Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../cuda/matmul_nbits_small_m_experiments.md | 127 ++++ .../cuda/quantization/matmul_4bits.cu | 645 +++++++++++++++++- .../cuda/quantization/matmul_8bits.cu | 243 ++++++- .../test/contrib_ops/matmul_4bits_test.cc | 41 ++ .../test/contrib_ops/matmul_8bits_test.cc | 48 ++ .../transformers/profile_matmul_nbits.py | 206 ++++++ 6 files changed, 1304 insertions(+), 6 deletions(-) create mode 100644 docs/contrib_ops/cuda/matmul_nbits_small_m_experiments.md create mode 100644 onnxruntime/test/python/transformers/profile_matmul_nbits.py diff --git a/docs/contrib_ops/cuda/matmul_nbits_small_m_experiments.md b/docs/contrib_ops/cuda/matmul_nbits_small_m_experiments.md new file mode 100644 index 0000000000000..046f17f81fa31 --- /dev/null +++ b/docs/contrib_ops/cuda/matmul_nbits_small_m_experiments.md @@ -0,0 +1,127 @@ +# MatMulNBits Small-M GEMV Profiling Experiments + +This file records CUDA `MatMulNBits` 4-bit (M = 2..16) and 8-bit (M = 2..5) +profiling results for the small-M range (e.g. multi-row decode or short +prefill) so future kernel and dispatch changes can be compared against a +stable baseline. + +> **Note**: These are **point-in-time** measurements captured on the specific GPU, driver, CUDA +> toolkit, and ORT build noted in the section header. Treat the numbers as a historical baseline +> for regression comparison, not as current performance guidance — re-run the benchmark script on +> your own hardware before drawing conclusions. + +## 2026-06-30 Baseline: A100 (SM80), Warmup 25, Repeat 200 + +### Setup + +- Machine/GPU: A100 (SM80). +- CUDA toolkit: 13.0. ONNX Runtime build: `build/Release`. +- Benchmark script: `onnxruntime/test/python/transformers/profile_matmul_nbits.py`. +- Warmup: 25 ORT runs before the measured `benchmark` NVTX range. Repeat: best of 10 trials x 200 runs. +- Weights: 4-bit, block_size 32, asymmetric (with zero points), per the standard `[N, blocks, blob]` layout. +- Matrices sized after a Qwen3-8B-class dense decoder + lm_head. + +Command template: + +```bash +# Host-timing table across all matrices: +python onnxruntime/test/python/transformers/profile_matmul_nbits.py --warmup 25 --repeat 200 + +# Single case: +python onnxruntime/test/python/transformers/profile_matmul_nbits.py --k 4096 --n 12288 --m 8 + +# Kernel-level via nsys + the repo parser: +nsys profile -t cuda,nvtx -o mnb --export=sqlite \ + python onnxruntime/test/python/transformers/profile_matmul_nbits.py --k 4096 --n 12288 --m 8 +python onnxruntime/test/python/transformers/parse_nsys.py mnb.sqlite --nvtx-range benchmark --pattern '%' +``` + +### Before / After (4-bit) + +`before` is the prior dispatch (M=1 single-row GEMV; M>1 falls back to weight dequantization + cuBLAS +GEMM, which is M-independent and dequantizes the full weight regardless of M). `after` is the batched +small-M GEMV (`MatMulFloat4BatchedKernel`) covering M=2..16. Values are average op latency in microseconds +(lower is better). + +Before (dequant + cuBLAS for M>1): + +| matrix | K | N | M=1 | M=2 | M=4 | M=8 | M=16 | +|---------|-------|--------|-------|--------|--------|--------|--------| +| qkv | 4096 | 4096 | 25.9 | 77.2 | 69.3 | 69.5 | 69.9 | +| o_proj | 4096 | 4096 | 23.2 | 69.2 | 69.1 | 69.3 | 69.7 | +| gate_up | 4096 | 12288 | 39.3 | 172.0 | 172.1 | 172.1 | 172.4 | +| down | 12288 | 4096 | 38.8 | 174.8 | 175.1 | 175.4 | 175.6 | +| lm_head | 4096 | 151936 | 301.6 | 1868.0 | 1871.1 | 1877.8 | 1885.1 | + +After (batched small-M GEMV for M=2..16): + +| matrix | K | N | M=1 | M=2 | M=4 | M=8 | M=16 | +|---------|-------|--------|-------|-------|-------|-------|--------| +| qkv | 4096 | 4096 | 25.4 | 30.1 | 32.3 | 43.2 | 70.0 | +| o_proj | 4096 | 4096 | 22.6 | 26.4 | 28.1 | 36.7 | 58.7 | +| gate_up | 4096 | 12288 | 37.5 | 43.5 | 49.5 | 71.9 | 125.9 | +| down | 12288 | 4096 | 37.9 | 47.4 | 52.9 | 76.8 | 150.1 | +| lm_head | 4096 | 151936 | 300.7 | 329.9 | 424.1 | 635.3 | 1226.9 | + +Speedup (before / after, >1 means the batched GEMV is faster): + +| matrix | M=2 | M=4 | M=8 | M=16 | +|---------|-------|-------|-------|-------| +| qkv | 2.56x | 2.15x | 1.61x | 1.00x | +| o_proj | 2.62x | 2.46x | 1.89x | 1.19x | +| gate_up | 3.95x | 3.48x | 2.39x | 1.37x | +| down | 3.69x | 3.31x | 2.28x | 1.17x | +| lm_head | 5.66x | 4.41x | 2.96x | 1.54x | + +### Before / After (8-bit) + +The 8-bit batched GEMV (`MatMulFloat8bKernelBatched`) covers M=2..5. 8-bit weights are twice the bytes of +4-bit and the GEMV runs on CUDA cores, so it crosses over to the dequantize + cuBLAS (tensor-core) fallback +at a lower M than the 4-bit path; M>=6 keeps the fallback. Values are average op latency in microseconds. + +Before (dequant + cuBLAS for M>1): + +| matrix | K | N | M=1 | M=2 | M=3 | M=4 | M=5 | +|---------|-------|--------|-------|--------|--------|--------|--------| +| qkv | 4096 | 4096 | 36.2 | 80.6 | 72.9 | 72.8 | 73.3 | +| o_proj | 4096 | 4096 | 31.1 | 73.2 | 72.7 | 73.6 | 73.4 | +| gate_up | 4096 | 12288 | 63.5 | 184.4 | 184.2 | 184.1 | 184.3 | +| down | 12288 | 4096 | 67.8 | 187.3 | 187.4 | 187.8 | 188.0 | +| lm_head | 4096 | 151936 | 535.9 | 2025.0 | 2025.9 | 2028.1 | 2029.8 | + +After (batched small-M GEMV for M=2..5): + +| matrix | K | N | M=1 | M=2 | M=3 | M=4 | M=5 | +|---------|-------|--------|-------|-------|-------|--------|--------| +| qkv | 4096 | 4096 | 36.2 | 46.1 | 48.7 | 57.9 | 67.0 | +| o_proj | 4096 | 4096 | 31.6 | 39.5 | 48.6 | 57.9 | 67.1 | +| gate_up | 4096 | 12288 | 63.4 | 80.9 | 104.6 | 128.9 | 152.6 | +| down | 12288 | 4096 | 68.0 | 96.2 | 119.3 | 146.4 | 172.0 | +| lm_head | 4096 | 151936 | 536.3 | 647.0 | 896.9 | 1157.1 | 1420.1 | + +Speedup (before / after, >1 means the batched GEMV is faster): + +| matrix | M=2 | M=3 | M=4 | M=5 | +|---------|-------|-------|-------|-------| +| qkv | 1.75x | 1.50x | 1.26x | 1.09x | +| o_proj | 1.85x | 1.50x | 1.27x | 1.09x | +| gate_up | 2.28x | 1.76x | 1.43x | 1.21x | +| down | 1.95x | 1.57x | 1.28x | 1.09x | +| lm_head | 3.13x | 2.26x | 1.75x | 1.43x | + +### Observations + +- The previous M>1 path dequantizes the entire weight matrix to a temporary buffer and calls cuBLAS, so + its latency is flat across M (e.g. `gate_up` ~172 us, `lm_head` ~1.9 ms even at M=2). The batched small-M + GEMV reads the quantized weight once and scales with M, giving 2.6-5.7x at M=2 (4-bit) / 1.8-3.1x at M=2 + (8-bit). +- 4-bit stays at or above parity through M=16; 8-bit wins through M=5 and falls back to dequant + cuBLAS + for M>=6, where the tensor-core GEMM beats the CUDA-core GEMV on the heavier 8-bit weights. +- M=1 decode is unchanged (same single-row GEMV in both builds). +- No prepacking is used, so there is no extra resident weight memory and no GEMM tactic profiling at + session init. + +### Next Experiments + +- Sweep block_size {16, 32, 64, 128} and bf16 activations. +- Add kernel-level (nsys) breakdowns to separate compute from launch/dispatch overhead. diff --git a/onnxruntime/contrib_ops/cuda/quantization/matmul_4bits.cu b/onnxruntime/contrib_ops/cuda/quantization/matmul_4bits.cu index b24e445401b99..1554cc2aecbea 100644 --- a/onnxruntime/contrib_ops/cuda/quantization/matmul_4bits.cu +++ b/onnxruntime/contrib_ops/cuda/quantization/matmul_4bits.cu @@ -5,6 +5,7 @@ #include #include #include +#include #include "core/providers/cuda/cu_inc/common.cuh" #include "core/providers/cuda/cuda_common.h" @@ -441,6 +442,631 @@ __global__ void __launch_bounds__(kWarpSize* kColsPerThreadBlock) MatMulFloatInt } } +// ===== Small-M batched GEMV (short prefill / small batch) ===== +// The single-row MatMulFloatInt4Kernel launches one block per output row (grid.y = m), so each row +// independently re-reads and re-dequantizes all of B; weight traffic and dequant work scale with M. +// For 2 <= M <= cap we instead dequantize each packed weight word once and accumulate it against +// CtaM activation rows held in registers, cutting weight traffic to ceil(M/CtaM)x. This is the same +// design used by TensorRT-LLM weightOnlyBatchedGemv / AWQ / llama.cpp MMVQ for small batch. +// +// Upper bound on M is kSmallMMax for all dtypes (measured on A100 vs the dequantize+cuBLAS fallback). +// half/bf16 run the register-tiled batched kernel, which stays ahead of the tensor-core GEMM through +// M<=16. float has no tensor-core GEMM fallback, so it uses the shared-memory small-M kernel over the +// same range. +constexpr int kSmallMMax = 16; +template +__host__ __device__ constexpr int SmallMCap() { + return kSmallMMax; +} + +// Holds the 8 dequantized weights produced from one packed 4-bit word, in the layout each +// AccumulateRow expects for its dtype. Splitting dequantization from the per-row multiply lets one +// dequantized weight feed all CtaM rows. +template +struct DequantizedEight; + +#if (!defined(__CUDA_ARCH__) || __CUDA_ARCH__ >= 530) && !defined(__HIPCC__) +template <> +struct DequantizedEight { + half2 v[4]; // order [04, 15, 26, 37], matching Convert8xInt4To8xHalfs + the prmt below +}; +__device__ __forceinline__ void DequantizeEight(uint32_t values_quant, half scale, uint8_t zp, DequantizedEight& d) { + half2 scale_half2 = {scale, scale}; + half zp_adjust = -scale * __short2half_rn(zp); + half2 zp_adjust2 = {zp_adjust, zp_adjust}; + half2 elements[4]; + Convert8xInt4To8xHalfs(values_quant, elements); + d.v[0] = elements[0] * scale_half2 + zp_adjust2; + d.v[1] = elements[1] * scale_half2 + zp_adjust2; + d.v[2] = elements[2] * scale_half2 + zp_adjust2; + d.v[3] = elements[3] * scale_half2 + zp_adjust2; +} +__device__ __forceinline__ void AccumulateRow(const DequantizedEight& d, const half* a, half* sums) { + uint4 vec_a = *(reinterpret_cast(a)); + constexpr uint32_t kLowHalf2 = 0x5410; + constexpr uint32_t kHighHalf2 = 0x7632; + uint4 vp; + asm volatile("prmt.b32 %0, %1, %2, %3;\n" : "=r"(vp.x) : "r"(vec_a.x), "r"(vec_a.z), "r"(kLowHalf2)); + asm volatile("prmt.b32 %0, %1, %2, %3;\n" : "=r"(vp.y) : "r"(vec_a.x), "r"(vec_a.z), "r"(kHighHalf2)); + asm volatile("prmt.b32 %0, %1, %2, %3;\n" : "=r"(vp.z) : "r"(vec_a.y), "r"(vec_a.w), "r"(kLowHalf2)); + asm volatile("prmt.b32 %0, %1, %2, %3;\n" : "=r"(vp.w) : "r"(vec_a.y), "r"(vec_a.w), "r"(kHighHalf2)); + half2* s = reinterpret_cast(sums); + s[0] = s[0] + d.v[0] * (*(reinterpret_cast(&vp.x))); + s[1] = s[1] + d.v[1] * (*(reinterpret_cast(&vp.y))); + s[2] = s[2] + d.v[2] * (*(reinterpret_cast(&vp.z))); + s[3] = s[3] + d.v[3] * (*(reinterpret_cast(&vp.w))); +} +#else +template <> +struct DequantizedEight { + half v[8]; +}; +__device__ __forceinline__ void DequantizeEight(uint32_t values_quant, half scale, uint8_t zp, DequantizedEight& d) { + half zp_adjust = -scale * __short2half_rn(zp); +#pragma unroll + for (int i = 0; i < 8; i++) { + d.v[i] = __uint2half_rn((values_quant >> (4 * i)) & 0xF) * scale + zp_adjust; + } +} +__device__ __forceinline__ void AccumulateRow(const DequantizedEight& d, const half* a, half* sums) { +#pragma unroll + for (int i = 0; i < 8; i++) { + sums[i] += d.v[i] * a[i]; + } +} +#endif + +template <> +struct DequantizedEight { + float v[8]; +}; +__device__ __forceinline__ void DequantizeEight(uint32_t values_quant, float scale, uint8_t zp, DequantizedEight& d) { + float zp_adjust = -scale * zp; +#pragma unroll + for (int i = 0; i < 8; i++) { + d.v[i] = float((values_quant >> (4 * i)) & 0xF) * scale + zp_adjust; + } +} +__device__ __forceinline__ void AccumulateRow(const DequantizedEight& d, const float* a, float* sums) { + float4 a0 = *(reinterpret_cast(a)); + float4 a1 = *(reinterpret_cast(a + 4)); + sums[0] += d.v[0] * a0.x; + sums[1] += d.v[1] * a0.y; + sums[2] += d.v[2] * a0.z; + sums[3] += d.v[3] * a0.w; + sums[4] += d.v[4] * a1.x; + sums[5] += d.v[5] * a1.y; + sums[6] += d.v[6] * a1.z; + sums[7] += d.v[7] * a1.w; +} + +template <> +struct DequantizedEight { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800 + __nv_bfloat162 v[4]; +#else + nv_bfloat16 v[8]; +#endif +}; +__device__ __forceinline__ void DequantizeEight(uint32_t values_quant, nv_bfloat16 scale, uint8_t zp, DequantizedEight& d) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800 + __nv_bfloat162 scale_bf162 = __bfloat162bfloat162(scale); + nv_bfloat16 zp_adjust = -scale * __uint2bfloat16_rn(zp); + __nv_bfloat162 zp_adjust2 = __bfloat162bfloat162(zp_adjust); + __nv_bfloat162 elements[4]; + Convert8xInt4To8xBF16s(values_quant, elements); + d.v[0] = __hfma2(elements[0], scale_bf162, zp_adjust2); + d.v[1] = __hfma2(elements[1], scale_bf162, zp_adjust2); + d.v[2] = __hfma2(elements[2], scale_bf162, zp_adjust2); + d.v[3] = __hfma2(elements[3], scale_bf162, zp_adjust2); +#endif +} +__device__ __forceinline__ void AccumulateRow(const DequantizedEight& d, const nv_bfloat16* a, nv_bfloat16* sums) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800 + uint4 vec_a = *(reinterpret_cast(a)); + constexpr uint32_t kLowHalf2 = 0x5410; + constexpr uint32_t kHighHalf2 = 0x7632; + uint4 vp; + asm volatile("prmt.b32 %0, %1, %2, %3;\n" : "=r"(vp.x) : "r"(vec_a.x), "r"(vec_a.z), "r"(kLowHalf2)); + asm volatile("prmt.b32 %0, %1, %2, %3;\n" : "=r"(vp.y) : "r"(vec_a.x), "r"(vec_a.z), "r"(kHighHalf2)); + asm volatile("prmt.b32 %0, %1, %2, %3;\n" : "=r"(vp.z) : "r"(vec_a.y), "r"(vec_a.w), "r"(kLowHalf2)); + asm volatile("prmt.b32 %0, %1, %2, %3;\n" : "=r"(vp.w) : "r"(vec_a.y), "r"(vec_a.w), "r"(kHighHalf2)); + __nv_bfloat162* s = reinterpret_cast<__nv_bfloat162*>(sums); + s[0] = __hfma2(d.v[0], *reinterpret_cast<__nv_bfloat162*>(&vp.x), s[0]); + s[1] = __hfma2(d.v[1], *reinterpret_cast<__nv_bfloat162*>(&vp.y), s[1]); + s[2] = __hfma2(d.v[2], *reinterpret_cast<__nv_bfloat162*>(&vp.z), s[2]); + s[3] = __hfma2(d.v[3], *reinterpret_cast<__nv_bfloat162*>(&vp.w), s[3]); +#endif +} + +// ---- Small-M batched GEMV (half/bf16): CtaM x CtaN register tiling -------------------------- +// 2-wide accumulator type (half2 / bf162). +template +struct Acc2; +#if (!defined(__CUDA_ARCH__) || __CUDA_ARCH__ >= 530) && !defined(__HIPCC__) +template <> +struct Acc2 { + using type = half2; +}; +#endif +template <> +struct Acc2 { + using type = __nv_bfloat162; +}; + +// Four 2-wide weight lanes in NATURAL element order [01,23,45,67], so a naturally-loaded activation +// (uint4 reinterpreted as four half2) can be multiply-accumulated with no per-activation permute. +template +struct WPack { + typename Acc2::type v[4]; +}; + +// DequantizeEight emits [04,15,26,37] (the order of Convert8xInt4To8xHalfs); repack to natural order +// once per column. Doing the prmt on the (CtaN) weights instead of the (CtaM) activations cuts the +// permute count by CtaM/CtaN, which dominates at small M. +#if (!defined(__CUDA_ARCH__) || __CUDA_ARCH__ >= 530) && !defined(__HIPCC__) +__device__ __forceinline__ WPack PackNatural(const DequantizedEight& d) { + uint32_t d0 = *reinterpret_cast(&d.v[0]); + uint32_t d1 = *reinterpret_cast(&d.v[1]); + uint32_t d2 = *reinterpret_cast(&d.v[2]); + uint32_t d3 = *reinterpret_cast(&d.v[3]); + constexpr uint32_t kLo = 0x5410; // (x0,x1) of two half2 -> elements 0,1 + constexpr uint32_t kHi = 0x7632; // (y0,y1) -> elements 4,5 + WPack w; + uint32_t t; + asm volatile("prmt.b32 %0, %1, %2, %3;\n" : "=r"(t) : "r"(d0), "r"(d1), "r"(kLo)); + w.v[0] = *reinterpret_cast(&t); + asm volatile("prmt.b32 %0, %1, %2, %3;\n" : "=r"(t) : "r"(d2), "r"(d3), "r"(kLo)); + w.v[1] = *reinterpret_cast(&t); + asm volatile("prmt.b32 %0, %1, %2, %3;\n" : "=r"(t) : "r"(d0), "r"(d1), "r"(kHi)); + w.v[2] = *reinterpret_cast(&t); + asm volatile("prmt.b32 %0, %1, %2, %3;\n" : "=r"(t) : "r"(d2), "r"(d3), "r"(kHi)); + w.v[3] = *reinterpret_cast(&t); + return w; +} +__device__ __forceinline__ void DotAccum(const WPack& w, const half2* a4, half2& acc) { + acc = __hfma2(w.v[0], a4[0], acc); + acc = __hfma2(w.v[1], a4[1], acc); + acc = __hfma2(w.v[2], a4[2], acc); + acc = __hfma2(w.v[3], a4[3], acc); +} +__device__ __forceinline__ float HorizontalAdd(half2 acc) { + return static_cast(acc.x) + static_cast(acc.y); +} +#endif + +__device__ __forceinline__ WPack PackNatural(const DequantizedEight& d) { + WPack w; +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800 + uint32_t d0 = *reinterpret_cast(&d.v[0]); + uint32_t d1 = *reinterpret_cast(&d.v[1]); + uint32_t d2 = *reinterpret_cast(&d.v[2]); + uint32_t d3 = *reinterpret_cast(&d.v[3]); + constexpr uint32_t kLo = 0x5410; + constexpr uint32_t kHi = 0x7632; + uint32_t t; + asm volatile("prmt.b32 %0, %1, %2, %3;\n" : "=r"(t) : "r"(d0), "r"(d1), "r"(kLo)); + w.v[0] = *reinterpret_cast<__nv_bfloat162*>(&t); + asm volatile("prmt.b32 %0, %1, %2, %3;\n" : "=r"(t) : "r"(d2), "r"(d3), "r"(kLo)); + w.v[1] = *reinterpret_cast<__nv_bfloat162*>(&t); + asm volatile("prmt.b32 %0, %1, %2, %3;\n" : "=r"(t) : "r"(d0), "r"(d1), "r"(kHi)); + w.v[2] = *reinterpret_cast<__nv_bfloat162*>(&t); + asm volatile("prmt.b32 %0, %1, %2, %3;\n" : "=r"(t) : "r"(d2), "r"(d3), "r"(kHi)); + w.v[3] = *reinterpret_cast<__nv_bfloat162*>(&t); +#endif + return w; +} +__device__ __forceinline__ void DotAccum(const WPack& w, const __nv_bfloat162* a4, __nv_bfloat162& acc) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800 + acc = __hfma2(w.v[0], a4[0], acc); + acc = __hfma2(w.v[1], a4[1], acc); + acc = __hfma2(w.v[2], a4[2], acc); + acc = __hfma2(w.v[3], a4[3], acc); +#endif +} +__device__ __forceinline__ float HorizontalAdd(__nv_bfloat162 acc) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800 + return static_cast(acc.x) + static_cast(acc.y); +#else + return 0.f; +#endif +} + +// Each warp computes CtaN output columns x CtaM rows. Lanes split K (8 elems/lane/iter); per-row +// activations are loaded once (uint4) and reused across CtaN columns, and the int4->half order permute +// is applied once per column weight (not per row). A single 2-wide accumulator per (row,column) keeps +// CtaM=m up to 16 in registers with the weight streamed exactly once. The launch bound pins >=3 blocks +// per SM so the CtaN=4 tiling (which minimizes activation L2 traffic) keeps enough occupancy to hide +// memory latency. Standard MatMulNBits [N, blocks, blob] layout, no prepacking; scales/zp from global. +template +__global__ void __launch_bounds__(kWarpSize* kColsPerThreadBlock, 3) MatMulFloat4BatchedKernel( + T* output, + const T* a_data, + const uint8_t* b_data_quant, + const T* scales_data, + const uint8_t* zero_points, + int m, + int n, + int k, + int blocks_per_K) { + using AccT = typename Acc2::type; + const int lane_id = threadIdx.x; + const int warp_id = WarpUniform(threadIdx.y); + const int col_base = (blockIdx.x * kColsPerThreadBlock + warp_id) * CtaN; + const int m_base = blockIdx.y * CtaM; + const int valid = m - m_base; + constexpr int k_per_iter = kWarpSize * kElementsPerThreadPerIteration; // 256 + const int zp_blocks = (blocks_per_K + 1) / 2; + + const T* a_base = a_data + static_cast(m_base) * k + (lane_id << 3); + const uint8_t* b_ptr[CtaN]; +#pragma unroll + for (int c = 0; c < CtaN; c++) { + b_ptr[c] = b_data_quant + static_cast(col_base + c) * blocks_per_K * (block_size / 2) + lane_id * 4; + } + + AccT acc[CtaM][CtaN]; +#pragma unroll + for (int r = 0; r < CtaM; r++) { +#pragma unroll + for (int c = 0; c < CtaN; c++) { + acc[r][c] = AccT{}; + } + } + + int k_id = 0; + int t_meta_k = lane_id * 8 / block_size; + constexpr int kWork = CtaM * CtaN; + constexpr int kMainUnroll = (kWork >= 20) ? 1 : (kWork >= 12) ? 2 + : 4; + +#define BATCHED_BODY(i) \ + do { \ + WPack w[CtaN]; \ + const int bk = t_meta_k + k_per_iter / block_size * (i); \ + _Pragma("unroll") for (int c = 0; c < CtaN; c++) { \ + uint32_t value = *(reinterpret_cast(b_ptr[c] + k_per_iter / 2 * (i))); \ + T scale = scales_data[static_cast(col_base + c) * blocks_per_K + bk]; \ + uint8_t zp = 8; \ + if constexpr (has_zero_point) { \ + uint8_t zpb = zero_points[static_cast(col_base + c) * zp_blocks + (bk >> 1)]; \ + zp = (bk & 1) ? (zpb >> 4) : (zpb & 0x0f); \ + } \ + DequantizedEight d; \ + DequantizeEight(value, scale, zp, d); \ + w[c] = PackNatural(d); \ + } \ + _Pragma("unroll") for (int r = 0; r < CtaM; r++) { \ + if (r < valid) { \ + AccT a4[4]; \ + *reinterpret_cast(a4) = *reinterpret_cast( \ + a_base + static_cast(r) * k + k_id + (i) * k_per_iter); \ + _Pragma("unroll") for (int c = 0; c < CtaN; c++) { \ + DotAccum(w[c], a4, acc[r][c]); \ + } \ + } \ + } \ + } while (false) + +#define BATCHED_UNROLL(unroll_size) \ + do { \ + constexpr int kUnroll = unroll_size; \ + constexpr int kUnrollStep = kUnroll * k_per_iter; \ + const int k_unroll_bound = k - k % kUnrollStep; \ + for (; k_id < k_unroll_bound; k_id += kUnrollStep) { \ + _Pragma("unroll") for (int i = 0; i < kUnroll; i++) { \ + BATCHED_BODY(i); \ + } \ + _Pragma("unroll") for (int c = 0; c < CtaN; c++) { \ + b_ptr[c] += k_per_iter / 2 * kUnroll; \ + } \ + t_meta_k += k_per_iter / block_size * kUnroll; \ + } \ + } while (false) + + BATCHED_UNROLL(kMainUnroll); + BATCHED_UNROLL(1); +#undef BATCHED_UNROLL + + if (k_id + lane_id * 8 < k) { + WPack w[CtaN]; + const int bk = t_meta_k; +#pragma unroll + for (int c = 0; c < CtaN; c++) { + uint32_t value = *(reinterpret_cast(b_ptr[c])); + T scale = scales_data[static_cast(col_base + c) * blocks_per_K + bk]; + uint8_t zp = 8; + if constexpr (has_zero_point) { + uint8_t zpb = zero_points[static_cast(col_base + c) * zp_blocks + (bk >> 1)]; + zp = (bk & 1) ? (zpb >> 4) : (zpb & 0x0f); + } + DequantizedEight d; + DequantizeEight(value, scale, zp, d); + w[c] = PackNatural(d); + } +#pragma unroll + for (int r = 0; r < CtaM; r++) { + if (r < valid) { + AccT a4[4]; + *reinterpret_cast(a4) = *reinterpret_cast(a_base + static_cast(r) * k + k_id); +#pragma unroll + for (int c = 0; c < CtaN; c++) { + DotAccum(w[c], a4, acc[r][c]); + } + } + } + } +#undef BATCHED_BODY + +#pragma unroll + for (int r = 0; r < CtaM; r++) { + if (r >= valid) continue; +#pragma unroll + for (int c = 0; c < CtaN; c++) { + float sum = HorizontalAdd(acc[r][c]); + for (int i = kWarpSize / 2; i > 0; i = i / 2) { + sum += WARP_SHFL_DOWN(sum, i); + } + if (lane_id == 0) { + output[static_cast(m_base + r) * n + (col_base + c)] = static_cast(sum); + } + } + } +} + +// Batched GEMV: block computes CtaM rows x kColsPerThreadBlock columns. Grid is +// (ceil(N/kColsPerThreadBlock), ceil(M/CtaM)). Mirrors MatMulFloatInt4Kernel's shared-memory scale/zp +// staging and packed-weight indexing; the only change is looping CtaM rows per dequantized word. +template +__global__ void __launch_bounds__(kWarpSize* kColsPerThreadBlock) MatMulFloatInt4KernelSmallM( + T* output, + const T* a_data, + const uint8_t* b_data_quant, + const T* scales_data, + const uint8_t* zero_points, + int m, + int n, + int k, + int blocks_per_K) { + const int n_block_id = blockIdx.x; + const int m_base = blockIdx.y * CtaM; + const int lane_id = threadIdx.x; + const int warp_id = WarpUniform(threadIdx.y); + const int n_id = n_block_id * kColsPerThreadBlock + warp_id; + constexpr int k_per_iter = kWarpSize * kElementsPerThreadPerIteration; + + extern __shared__ char shared_buffer[]; + T* b_scale_vec = (T*)shared_buffer; + int offset = n_block_id * kColsPerThreadBlock * blocks_per_K; + for (int i = warp_id * kWarpSize + lane_id; i < kColsPerThreadBlock * blocks_per_K; i += kColsPerThreadBlock * kWarpSize) { + b_scale_vec[i] = scales_data[offset + i]; + } + + uint8_t* b_zp_vec; + (void)b_zp_vec; + if constexpr (has_zero_point) { + b_zp_vec = reinterpret_cast(b_scale_vec + kColsPerThreadBlock * blocks_per_K); + const int b_zp_k = (blocks_per_K + 1) / 2; + int zp_offset = n_block_id * kColsPerThreadBlock * b_zp_k; + for (int i = warp_id * kWarpSize + lane_id; i < kColsPerThreadBlock * b_zp_k; i += kColsPerThreadBlock * kWarpSize) { + b_zp_vec[2 * i] = (zero_points[zp_offset + i] & 0x0f); + b_zp_vec[2 * i + 1] = (zero_points[zp_offset + i] >> 4); + } + b_zp_vec += warp_id * b_zp_k * 2; + } + __syncthreads(); + + const int valid = m - m_base; + const T* a_row[CtaM]; +#pragma unroll + for (int r = 0; r < CtaM; r++) { + a_row[r] = a_data + static_cast(m_base + r) * k + (lane_id << 3); + } + b_scale_vec += warp_id * blocks_per_K; + + T sums[CtaM][8]; +#pragma unroll + for (int r = 0; r < CtaM; r++) { +#pragma unroll + for (int j = 0; j < 8; j++) { + sums[r][j] = static_cast(0); + } + } + + int k_id = 0; + int t_meta_k = lane_id * 8 / block_size; + b_data_quant += n_id * blocks_per_K * (block_size / 2) + lane_id * 4; + +#define SmallMUnRoll(unroll_size) \ + do { \ + constexpr int kUnroll = unroll_size; \ + constexpr int kUnrollStep = kUnroll * k_per_iter; \ + const int k_unroll_bound = k - k % kUnrollStep; \ + for (; k_id < k_unroll_bound; k_id += kUnrollStep) { \ + _Pragma("unroll") for (int i = 0; i < kUnroll; i++) { \ + uint32_t value = *(reinterpret_cast(b_data_quant + k_per_iter / 2 * i)); \ + T scale = b_scale_vec[t_meta_k + k_per_iter / block_size * i]; \ + uint8_t zp = 8; \ + if constexpr (has_zero_point) { \ + zp = b_zp_vec[t_meta_k + k_per_iter / block_size * i]; \ + } \ + DequantizedEight d; \ + DequantizeEight(value, scale, zp, d); \ + _Pragma("unroll") for (int r = 0; r < CtaM; r++) { \ + if (r < valid) AccumulateRow(d, a_row[r] + k_id + i * k_per_iter, sums[r]); \ + } \ + } \ + b_data_quant += k_per_iter / 2 * kUnroll; \ + t_meta_k += k_per_iter / block_size * kUnroll; \ + } \ + } while (false) + + SmallMUnRoll(16); + SmallMUnRoll(4); + SmallMUnRoll(1); +#undef SmallMUnRoll + + if (k_id + lane_id * 8 < k) { + uint32_t value = *(reinterpret_cast(b_data_quant)); + T scale = b_scale_vec[t_meta_k]; + uint8_t zp = 8; + if constexpr (has_zero_point) { + zp = b_zp_vec[t_meta_k]; + } + DequantizedEight d; + DequantizeEight(value, scale, zp, d); +#pragma unroll + for (int r = 0; r < CtaM; r++) { + if (r < valid) AccumulateRow(d, a_row[r] + k_id, sums[r]); + } + } + +#pragma unroll + for (int r = 0; r < CtaM; r++) { + if (r >= valid) continue; + float sum = static_cast(sums[r][0] + sums[r][1] + sums[r][2] + sums[r][3] + + sums[r][4] + sums[r][5] + sums[r][6] + sums[r][7]); + for (int i = kWarpSize / 2; i > 0; i = i / 2) { + sum += WARP_SHFL_DOWN(sum, i); + } + if (lane_id == 0) { + output[static_cast(m_base + r) * n + n_id] = sum; + } + } +} + +// Launches the batched small-M kernel for 2 <= m <= SmallMCap(). Returns false if m is out of range +// so the caller can fall back to the single-row (m==1) or dequant+GEMM (large m) paths. +template +bool TryMatMulSmallM4Bits( + T* output, + const T* a_data, + const uint8_t* b_data_quant, + const T* scales_data, + const uint8_t* zero_points, + int m, + int n, + int k, + int block_size, + size_t shared_mem_size, + cudaStream_t stream) { + if (m < 2 || m > SmallMCap()) { + return false; + } + const int cta_m = (m <= 2) ? 2 : 4; + dim3 threads(GPU_WARP_SIZE_HOST, kColsPerThreadBlock); + dim3 blocks((n + kColsPerThreadBlock - 1) / kColsPerThreadBlock, (m + cta_m - 1) / cta_m); + +#define SmallMDispatch(BS, CM) \ + if (nullptr != zero_points) { \ + MatMulFloatInt4KernelSmallM<<>>( \ + output, a_data, b_data_quant, scales_data, zero_points, m, n, k, (k + BS - 1) / BS); \ + } else { \ + MatMulFloatInt4KernelSmallM<<>>( \ + output, a_data, b_data_quant, scales_data, zero_points, m, n, k, (k + BS - 1) / BS); \ + } +#define SmallMDispatchBlock(CM) \ + if (16 == block_size) { \ + SmallMDispatch(16, CM) \ + } else if (32 == block_size) { \ + SmallMDispatch(32, CM) \ + } else if (64 == block_size) { \ + SmallMDispatch(64, CM) \ + } else if (128 == block_size) { \ + SmallMDispatch(128, CM) \ + } else { \ + return false; \ + } + + if (cta_m == 2) { + SmallMDispatchBlock(2) + } else { + SmallMDispatchBlock(4) + } + +#undef SmallMDispatchBlock +#undef SmallMDispatch + return true; +} + +// Small-M launcher (half/bf16): picks CtaM >= m from {2,4,8,16} so a single block streams the +// weight once, and CtaN columns/warp (largest of {4,2,1} dividing N/kCols) to reuse each activation +// load across columns. Returns false for float or out-of-range m so the caller falls back. +template +bool TryMatMulBatched4Bits( + T* output, + const T* a_data, + const uint8_t* b_data_quant, + const T* scales_data, + const uint8_t* zero_points, + int m, + int n, + int k, + int block_size, + size_t /*shared_mem_size*/, + cudaStream_t stream) { + if constexpr (std::is_same::value) { + return false; + } else { + if (m < 2 || m > kSmallMMax) { + return false; + } + // CtaM = smallest of {2,4,8,16} >= m, streaming the weight once per block row. Non-power-of-2 CtaM + // (10/12/14) compile to materially slower code, so the row-tile is rounded up (M>8 uses CtaM=16). + // CtaN = 2 columns/warp where N allows (halves activation L2 traffic); CtaN=4 lost to register + // pressure so it is not used by default. + const int cta_m = (m <= 2) ? 2 : (m <= 4) ? 4 + : (m <= 8) ? 8 + : 16; + const int cta_n = (n % (kColsPerThreadBlock * 2) == 0) ? 2 : 1; + dim3 threads(GPU_WARP_SIZE_HOST, kColsPerThreadBlock); + dim3 blocks(n / (kColsPerThreadBlock * cta_n), (m + cta_m - 1) / cta_m); + +#define BatchedDispatch(BS, CM, CN) \ + if (nullptr != zero_points) { \ + MatMulFloat4BatchedKernel<<>>( \ + output, a_data, b_data_quant, scales_data, zero_points, m, n, k, (k + BS - 1) / BS); \ + } else { \ + MatMulFloat4BatchedKernel<<>>( \ + output, a_data, b_data_quant, scales_data, zero_points, m, n, k, (k + BS - 1) / BS); \ + } +#define BatchedDispatchN(CM, CN) \ + if (16 == block_size) { \ + BatchedDispatch(16, CM, CN) \ + } else if (32 == block_size) { \ + BatchedDispatch(32, CM, CN) \ + } else if (64 == block_size) { \ + BatchedDispatch(64, CM, CN) \ + } else if (128 == block_size) { \ + BatchedDispatch(128, CM, CN) \ + } else { \ + return false; \ + } +#define BatchedDispatchM(CN) \ + switch (cta_m) { \ + case 2: \ + BatchedDispatchN(2, CN) break; \ + case 4: \ + BatchedDispatchN(4, CN) break; \ + case 8: \ + BatchedDispatchN(8, CN) break; \ + default: \ + BatchedDispatchN(16, CN) break; \ + } + + if (cta_n == 2) { + BatchedDispatchM(2) + } else { + BatchedDispatchM(1) + } + +#undef BatchedDispatchM +#undef BatchedDispatchN +#undef BatchedDispatch + return true; + } +} + template bool TryMatMul4Bits( T* output, @@ -455,7 +1081,7 @@ bool TryMatMul4Bits( int block_size, size_t shared_mem_per_block, cudaStream_t stream) { - if (n % kColsPerThreadBlock != 0 || k % 8 != 0 || m > 1) { + if (n % kColsPerThreadBlock != 0 || k % 8 != 0 || m > SmallMCap()) { return false; } @@ -477,6 +1103,15 @@ bool TryMatMul4Bits( return false; } + // The register-tiled batched path (half/bf16, 2 <= m <= cap) launches with no shared memory, so try + // it before the shared-memory budget gate that only constrains the shared-memory kernels below. + if (m >= 2) { + if (TryMatMulBatched4Bits(output, a_data, b_data_quant, scales_data, zero_points, + m, n, k, block_size, 0, stream)) { + return true; + } + } + dim3 blocks((n + kColsPerThreadBlock - 1) / kColsPerThreadBlock, m); dim3 threads(GPU_WARP_SIZE_HOST, kColsPerThreadBlock); int blocks_per_K = (k + block_size - 1) / block_size; @@ -486,6 +1121,14 @@ bool TryMatMul4Bits( return false; } + // Float, and any half/bf16 shape the batched path rejected, falls back to the shared-memory small-M + // kernel. m == 1 uses the single-row kernel below; larger m used the dequantize + cuBLAS path + // (returned false above). + if (m >= 2) { + return TryMatMulSmallM4Bits(output, a_data, b_data_quant, scales_data, zero_points, + m, n, k, block_size, shared_mem_size, stream); + } + #define MatMulFloatInt4KernelDispatch(block_size) \ if (nullptr != zero_points) { \ MatMulFloatInt4Kernel<<>>( \ diff --git a/onnxruntime/contrib_ops/cuda/quantization/matmul_8bits.cu b/onnxruntime/contrib_ops/cuda/quantization/matmul_8bits.cu index 2788675765b62..495922185c939 100644 --- a/onnxruntime/contrib_ops/cuda/quantization/matmul_8bits.cu +++ b/onnxruntime/contrib_ops/cuda/quantization/matmul_8bits.cu @@ -175,6 +175,185 @@ __device__ __forceinline__ void AccumulateEightElements8b( #endif } +// ===== Small-M batched GEMV (8-bit) ===== +// Same idea as the 4-bit batched kernel: each warp computes CtaN columns x CtaM rows. Lanes split K +// (8 elems/lane/iter via one uint64 weight load); each column's 8 weights are dequantized once per row +// group, per-row activations are loaded once and reused across CtaN columns, and a single float +// accumulator per (row, column) keeps register pressure low while the weight is streamed once. The +// launch bound pins >=3 blocks per SM. Standard [N, blocks, blob] layout, no prepacking; the 8-bit path +// accumulates in float, so this stays dtype-agnostic (float/half/bf16). +constexpr int kSmallMMax8 = 5; +template +__host__ __device__ constexpr int SmallMCap8() { + return kSmallMMax8; +} + +template +__device__ __forceinline__ float ToFloat8b(T v); +template <> +__device__ __forceinline__ float ToFloat8b(float v) { return v; } +template <> +__device__ __forceinline__ float ToFloat8b(half v) { return __half2float(v); } +template <> +__device__ __forceinline__ float ToFloat8b(nv_bfloat16 v) { return __bfloat162float(v); } + +template +__device__ __forceinline__ void DequantizeEight8b(uint64_t values_quant, T scale, uint8_t zp, float dq[8]) { + float scale_f = ToFloat8b(scale); + float zp_f = static_cast(zp); +#pragma unroll + for (int i = 0; i < 8; ++i) { + uint8_t q = (values_quant >> (i * 8)) & 0xFF; + dq[i] = (static_cast(q) - zp_f) * scale_f; + } +} + +// Load 8 consecutive activations (one lane's iteration tile) into float[8]. +template +__device__ __forceinline__ void LoadEightActivations8b(const T* a, float out[8]) { +#pragma unroll + for (int j = 0; j < 8; ++j) { + out[j] = ToFloat8b(a[j]); + } +} + +template +__global__ void __launch_bounds__(kWarpSize* kColsPerThreadBlock, 3) MatMulFloat8bKernelBatched( + T* output, + const T* a_data, + const uint8_t* b_data_quant, + const T* scales_data, + const uint8_t* zero_points, + int m, + int n, + int k, + int blocks_per_K) { + const int lane_id = threadIdx.x; + const int warp_id = threadIdx.y; + const int col_base = (blockIdx.x * kColsPerThreadBlock + warp_id) * CtaN; + const int m_base = blockIdx.y * CtaM; + const int valid = m - m_base; + constexpr int k_per_iter = kWarpSize * kElementsPerThreadPerIteration; // 256 + const int lane_offset = lane_id * kElementsPerThreadPerIteration; // lane_id * 8 + + const T* a_base = a_data + static_cast(m_base) * k + lane_offset; + const uint8_t* b_ptr[CtaN]; +#pragma unroll + for (int c = 0; c < CtaN; ++c) { + b_ptr[c] = b_data_quant + static_cast(col_base + c) * blocks_per_K * block_size + lane_offset; + } + + float acc[CtaM][CtaN]; +#pragma unroll + for (int r = 0; r < CtaM; ++r) { +#pragma unroll + for (int c = 0; c < CtaN; ++c) { + acc[r][c] = 0.0f; + } + } + + int k_id = 0; + int t_meta_k = lane_offset / block_size; + constexpr int kWork = CtaM * CtaN; + constexpr int kMainUnroll = (kWork >= 20) ? 1 : (kWork >= 12) ? 2 + : 4; + +#define BATCHED8_BODY(i) \ + do { \ + float dq[CtaN][8]; \ + const int bk = t_meta_k + k_per_iter / block_size * (i); \ + _Pragma("unroll") for (int c = 0; c < CtaN; ++c) { \ + uint64_t value = *reinterpret_cast(b_ptr[c] + k_per_iter * (i)); \ + T scale = scales_data[static_cast(col_base + c) * blocks_per_K + bk]; \ + uint8_t zp = kDefaultZeroPoint; \ + if constexpr (has_zero_point) { \ + zp = zero_points[static_cast(col_base + c) * blocks_per_K + bk]; \ + } \ + DequantizeEight8b(value, scale, zp, dq[c]); \ + } \ + _Pragma("unroll") for (int r = 0; r < CtaM; ++r) { \ + if (r < valid) { \ + float av[8]; \ + LoadEightActivations8b(a_base + static_cast(r) * k + k_id + (i) * k_per_iter, av); \ + _Pragma("unroll") for (int c = 0; c < CtaN; ++c) { \ + float s = 0.0f; \ + _Pragma("unroll") for (int j = 0; j < 8; ++j) { \ + s = fmaf(av[j], dq[c][j], s); \ + } \ + acc[r][c] += s; \ + } \ + } \ + } \ + } while (false) + +#define BATCHED8_UNROLL(unroll_size) \ + do { \ + constexpr int kUnroll = unroll_size; \ + constexpr int kUnrollStep = kUnroll * k_per_iter; \ + const int k_unroll_bound = k - k % kUnrollStep; \ + for (; k_id < k_unroll_bound; k_id += kUnrollStep) { \ + _Pragma("unroll") for (int i = 0; i < kUnroll; ++i) { \ + BATCHED8_BODY(i); \ + } \ + _Pragma("unroll") for (int c = 0; c < CtaN; ++c) { \ + b_ptr[c] += kUnrollStep; \ + } \ + t_meta_k += k_per_iter / block_size * kUnroll; \ + } \ + } while (false) + + BATCHED8_UNROLL(kMainUnroll); + BATCHED8_UNROLL(1); +#undef BATCHED8_UNROLL + + if (lane_offset + k_id < k) { + float dq[CtaN][8]; + const int bk = t_meta_k; +#pragma unroll + for (int c = 0; c < CtaN; ++c) { + uint64_t value = *reinterpret_cast(b_ptr[c]); + T scale = scales_data[static_cast(col_base + c) * blocks_per_K + bk]; + uint8_t zp = kDefaultZeroPoint; + if constexpr (has_zero_point) { + zp = zero_points[static_cast(col_base + c) * blocks_per_K + bk]; + } + DequantizeEight8b(value, scale, zp, dq[c]); + } +#pragma unroll + for (int r = 0; r < CtaM; ++r) { + if (r < valid) { + float av[8]; + LoadEightActivations8b(a_base + static_cast(r) * k + k_id, av); +#pragma unroll + for (int c = 0; c < CtaN; ++c) { + float s = 0.0f; +#pragma unroll + for (int j = 0; j < 8; ++j) { + s = fmaf(av[j], dq[c][j], s); + } + acc[r][c] += s; + } + } + } + } +#undef BATCHED8_BODY + +#pragma unroll + for (int r = 0; r < CtaM; ++r) { + if (r >= valid) continue; +#pragma unroll + for (int c = 0; c < CtaN; ++c) { + float sum = acc[r][c]; + for (int off = kWarpSize / 2; off > 0; off /= 2) { + sum += WARP_SHFL_DOWN(sum, off); + } + if (lane_id == 0) { + output[static_cast(m_base + r) * n + (col_base + c)] = static_cast(sum); + } + } + } +} + // --- CUDA Kernel: MatMulFloat8bKernel (Optimized for m=1) --- // Computes C(1, N) = A(1, K) x B(K, N) // B(K, N) is quantized with 8 bits and block_size bs, stored as [N, K/bs, bs] @@ -346,17 +525,17 @@ bool TryMatMul8Bits( const uint8_t* b_data_quant, // Input B Quantized [N, K/bs, bs] const T* scales_data, // Scales [N, K/bs] const uint8_t* zero_points, // Zero Points [N, K/bs] (can be nullptr) - int m, // Rows of A and C (MUST be 1) + int m, // Rows of A and C (1 single-row, 2..cap batched) int n, // Columns of B and C int k, // Columns of A / Rows of B int block_size, // Quantization block size for B size_t shared_mem_per_block, // Available shared memory per block cudaStream_t stream) { // Constraints Check - // m must be 1 (since this kernel is optimized for m=1) // N must be a multiple of kColsPerThreadBlock (8) for warps to align with columns. // K must be a multiple of kElementsPerThreadPerIteration (8) for full uint64_t reads/processing. - if (m != 1 || n % kColsPerThreadBlock != 0 || k % kElementsPerThreadPerIteration != 0) { + // m up to the small-M cap is handled (m==1 single-row, 2..cap batched); larger m falls back. + if (m < 1 || m > SmallMCap8() || n % kColsPerThreadBlock != 0 || k % kElementsPerThreadPerIteration != 0) { return false; } @@ -380,7 +559,61 @@ bool TryMatMul8Bits( // Calculate K / block_size (no rounding needed due to k % block_size == 0 check) int blocks_per_K = k / block_size; - // --- Shared Memory Calculation --- + // 2 <= m <= cap: batched GEMV. CtaM = smallest of {2,4,8} >= m streams the weight once per block row + // (m=5 uses CtaM=8 and skips the unused rows); CtaN = 2 columns/warp where N allows (reuses each + // activation load across columns). One float accumulator per (row, column) keeps register pressure + // low. 8-bit weights are twice the bytes of 4-bit and the GEMV runs on CUDA cores, so it crosses over + // to the dequantize + cuBLAS (tensor-core) fallback at a lower M than the 4-bit path; the cap keeps + // only the row counts where it is faster on every matrix shape. This path launches with no shared + // memory, so it runs before the shared-memory budget gate that only constrains the m==1 kernel below. + if (m >= 2) { + const int cta_m = (m <= 2) ? 2 : (m <= 4) ? 4 + : 8; + const int cta_n = (n % (kColsPerThreadBlock * 2) == 0) ? 2 : 1; + dim3 batched_blocks(n / (kColsPerThreadBlock * cta_n), (m + cta_m - 1) / cta_m); +#define MatMulFloat8bBatchedDispatch(bs, cm, cn) \ + if (nullptr != zero_points) { \ + MatMulFloat8bKernelBatched<<>>( \ + output, a_data, b_data_quant, scales_data, zero_points, m, n, k, blocks_per_K); \ + } else { \ + MatMulFloat8bKernelBatched<<>>( \ + output, a_data, b_data_quant, scales_data, nullptr, m, n, k, blocks_per_K); \ + } +#define MatMulFloat8bBatchedDispatchN(cm, cn) \ + if (16 == block_size) { \ + MatMulFloat8bBatchedDispatch(16, cm, cn) \ + } else if (32 == block_size) { \ + MatMulFloat8bBatchedDispatch(32, cm, cn) \ + } else if (64 == block_size) { \ + MatMulFloat8bBatchedDispatch(64, cm, cn) \ + } else if (128 == block_size) { \ + MatMulFloat8bBatchedDispatch(128, cm, cn) \ + } else if (256 == block_size) { \ + MatMulFloat8bBatchedDispatch(256, cm, cn) \ + } else { \ + return false; \ + } +#define MatMulFloat8bBatchedDispatchM(cn) \ + switch (cta_m) { \ + case 2: \ + MatMulFloat8bBatchedDispatchN(2, cn) break; \ + case 4: \ + MatMulFloat8bBatchedDispatchN(4, cn) break; \ + default: \ + MatMulFloat8bBatchedDispatchN(8, cn) break; \ + } + if (cta_n == 2) { + MatMulFloat8bBatchedDispatchM(2) + } else { + MatMulFloat8bBatchedDispatchM(1) + } +#undef MatMulFloat8bBatchedDispatchM +#undef MatMulFloat8bBatchedDispatchN +#undef MatMulFloat8bBatchedDispatch + return true; + } + + // --- Shared Memory Calculation (m == 1) --- // Memory for scales + optional zero points for the columns handled by the block size_t scale_zp_shared_mem = (sizeof(T) + (zero_points != nullptr ? sizeof(uint8_t) : 0)) * static_cast(blocks_per_K) * kColsPerThreadBlock; @@ -395,7 +628,7 @@ bool TryMatMul8Bits( return false; } - // --- Kernel Launch --- + // --- Kernel Launch (m == 1) --- // Macro to simplify dispatching based on block size and presence of zero_points #define MatMulFloat8bKernelM1Dispatch(bs) \ if (nullptr != zero_points) { \ diff --git a/onnxruntime/test/contrib_ops/matmul_4bits_test.cc b/onnxruntime/test/contrib_ops/matmul_4bits_test.cc index bedf035d320f8..4d5a1c5c17c0b 100644 --- a/onnxruntime/test/contrib_ops/matmul_4bits_test.cc +++ b/onnxruntime/test/contrib_ops/matmul_4bits_test.cc @@ -869,6 +869,47 @@ TEST(MatMulNBits, Fp16_Int4_NoZeroPoint) { } } +// Exercises the CUDA small-M batched GEMV tiles: CtaM in {2,4,8,16} (with M values that are not a +// multiple of CtaM so the row-skip path runs) and CtaN in {1,2} (N divisible / not divisible by 16). +TEST(MatMulNBits, Fp16_Int4_SmallMBatchedTiles) { + constexpr float abs_error = 0.1f; + constexpr bool zp_is_4bit = true; + for (auto block_size : {32, 128}) { + for (auto m : {3, 4, 5, 8, 12, 16}) { + for (auto has_zeropoint : {false, true}) { + RunTest(m, 256, 1024, block_size, has_zeropoint, zp_is_4bit, abs_error); // N % 16 == 0 -> CtaN=2 + RunTest(m, 24, 1024, block_size, has_zeropoint, zp_is_4bit, abs_error); // N % 16 != 0 -> CtaN=1 + } + } + } +} + +TEST(MatMulNBits, BFloat16_Int4_SmallMBatchedTiles) { + if (!HasCudaEnvironment(800)) { + GTEST_SKIP() << "Skipping BFloat16 tests on CUDA < 8.0"; + } + + constexpr float abs_error = 0.1f; + for (auto block_size : {32, 128}) { + for (auto m : {3, 4, 5, 8, 12, 16}) { + for (auto n : {256, 24}) { // N=256 -> CtaN=2, N=24 -> CtaN=1 + for (auto has_zeropoint : {false, true}) { + TestOptions opts{}; + opts.M = m, opts.N = n, opts.K = 1024; + opts.block_size = block_size; + opts.has_zero_point = has_zeropoint; + opts.zp_is_4bit = true; + opts.output_abs_error = abs_error; + opts.output_rel_error = 0.02f; + std::vector> eps; + eps.push_back(DefaultCudaExecutionProvider()); + RunTest(opts, std::move(eps)); + } + } + } + } +} + TEST(MatMulNBits, Fp16_Int4_GptOssRouterShapeNoZeroPoint) { constexpr float abs_error = 0.1f; diff --git a/onnxruntime/test/contrib_ops/matmul_8bits_test.cc b/onnxruntime/test/contrib_ops/matmul_8bits_test.cc index 411e83536c190..c3cd8b388dd8e 100644 --- a/onnxruntime/test/contrib_ops/matmul_8bits_test.cc +++ b/onnxruntime/test/contrib_ops/matmul_8bits_test.cc @@ -682,6 +682,54 @@ TEST(MatMulNBits, BFloat16_Int8_Chunked_BFloat16ZeroPoint) { RunTest8Bits(opts); } } + +// Exercises the CUDA small-M batched GEMV tiles for 8-bit: CtaM in {2,4,8} (M=3,5 hit the row-skip +// path) and CtaN in {1,2} (N divisible / not divisible by 16). 8-bit caps the batched path at M=5. +TEST(MatMulNBits, Fp16_Int8_SmallMBatchedTiles) { + constexpr float abs_error = 0.1f; + constexpr float rel_error = 0.02f; + for (auto block_size : {32, 128}) { + for (auto m : {2, 3, 4, 5}) { + for (auto n : {256, 24}) { // N=256 -> CtaN=2, N=24 -> CtaN=1 + for (auto has_zeropoint : {false, true}) { + TestOptions8Bits opts{}; + opts.M = m, opts.N = n, opts.K = 1024; + opts.block_size = block_size; + opts.has_zero_point = has_zeropoint; + opts.zp_is_typed = false; + opts.output_abs_error = abs_error; + opts.output_rel_error = rel_error; + RunTest8Bits(opts); + } + } + } + } +} + +TEST(MatMulNBits, BFloat16_Int8_SmallMBatchedTiles) { + if (!HasCudaEnvironment(800)) { + GTEST_SKIP() << "Skipping BFloat16 tests on CUDA < 8.0"; + } + + constexpr float abs_error = 0.1f; + constexpr float rel_error = 0.02f; + for (auto block_size : {32, 128}) { + for (auto m : {2, 3, 4, 5}) { + for (auto n : {256, 24}) { + for (auto has_zeropoint : {false, true}) { + TestOptions8Bits opts{}; + opts.M = m, opts.N = n, opts.K = 1024; + opts.block_size = block_size; + opts.has_zero_point = has_zeropoint; + opts.zp_is_typed = false; + opts.output_abs_error = abs_error; + opts.output_rel_error = rel_error; + RunTest8Bits(opts); + } + } + } + } +} #endif #if !defined(USE_CUDA) && !defined(USE_WEBGPU) diff --git a/onnxruntime/test/python/transformers/profile_matmul_nbits.py b/onnxruntime/test/python/transformers/profile_matmul_nbits.py new file mode 100644 index 0000000000000..75edc892ede1f --- /dev/null +++ b/onnxruntime/test/python/transformers/profile_matmul_nbits.py @@ -0,0 +1,206 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- + +""" +Profiling script for the CUDA MatMulNBits (4-bit / 8-bit weight-only) op. + +Times the op across representative decoder weight matrices and row counts M +(M=1 decode and M>1 small-batch). Mirrors profile_qmoe_gemv.py: +warmup + measured runs wrapped in NVTX ranges so the results can be parsed +with parse_nsys.py for kernel-level timing, while also printing host-observed +average latency for a quick end-to-end view. + +Usage: + # Host-timing table across all cases: + python profile_matmul_nbits.py --warmup 25 --repeat 200 + + # Single case: + python profile_matmul_nbits.py --k 4096 --n 4096 --m 8 --block-size 32 --bits 4 --dtype fp16 + + # Kernel-level via nsys + the repo parser: + nsys profile -t cuda,nvtx -o mnb --export=sqlite \ + python profile_matmul_nbits.py --k 4096 --n 4096 --m 8 + python parse_nsys.py mnb.sqlite --nvtx-range benchmark --pattern '%' +""" + +import argparse +import json +import time +from contextlib import nullcontext + +import numpy as np +import torch +from onnx import TensorProto, helper + +import onnxruntime + +try: + import ml_dtypes + + bfloat16 = ml_dtypes.bfloat16 +except ImportError: + bfloat16 = None + +try: + import nvtx + + has_nvtx = True +except ImportError: + has_nvtx = False + nvtx = None + +RESULT_PREFIX = "MATMUL_NBITS_RESULT " + +_OT = {"fp16": TensorProto.FLOAT16, "bf16": TensorProto.BFLOAT16} +_TT = {"fp16": torch.float16, "bf16": torch.bfloat16} +_ELEM = {torch.float16: TensorProto.FLOAT16, torch.bfloat16: TensorProto.BFLOAT16} + +# Representative decoder weight matrices (K = in features, N = out features), +# sized after a Qwen3-8B-class dense decoder + lm_head. +DEFAULT_CASES = [ + ("qkv", 4096, 4096), + ("o_proj", 4096, 4096), + ("gate_up", 4096, 12288), + ("down", 12288, 4096), + ("lm_head", 4096, 151936), +] +DEFAULT_MS = [1, 2, 4, 8, 16] + + +def _nvtx_range(name, color="green"): + if not has_nvtx: + return nullcontext() + return nvtx.annotate(name, color=color) + + +def build_model(k, n, block_size, bits, onnx_dtype, with_zero_point=True): + rng = np.random.default_rng(0) + n_blocks = (k + block_size - 1) // block_size + blob = block_size // (8 // bits) + b = rng.integers(0, 256, size=(n, n_blocks, blob), dtype=np.uint8) + scales_f32 = rng.random(n * n_blocks).astype(np.float32) * 0.02 + 0.01 + if onnx_dtype == TensorProto.FLOAT16: + scales = scales_f32.astype(np.float16) + elif onnx_dtype == TensorProto.BFLOAT16: + if bfloat16 is None: + raise RuntimeError("ml_dtypes is required for bf16 (pip install ml_dtypes)") + scales = scales_f32.astype(bfloat16) + else: + scales = scales_f32 + inits = [ + helper.make_tensor("B", TensorProto.UINT8, list(b.shape), b.tobytes(), raw=True), + helper.make_tensor("scales", onnx_dtype, list(scales.shape), scales.tobytes(), raw=True), + ] + inputs = ["A", "B", "scales"] + if with_zero_point: + # 4-bit zero points are packed two per byte; 8-bit zero points are one byte per block. + zp_count = n * ((n_blocks + 1) // 2) if bits == 4 else n * n_blocks + zp = rng.integers(0, 256, size=(zp_count,), dtype=np.uint8) + inits.append(helper.make_tensor("zero_points", TensorProto.UINT8, list(zp.shape), zp.tobytes(), raw=True)) + inputs.append("zero_points") + node = helper.make_node( + "MatMulNBits", + inputs, + ["Y"], + domain="com.microsoft", + K=k, + N=n, + bits=bits, + block_size=block_size, + accuracy_level=0, + ) + graph = helper.make_graph( + [node], + "mnb", + [helper.make_tensor_value_info("A", onnx_dtype, ["M", k])], + [helper.make_tensor_value_info("Y", onnx_dtype, ["M", n])], + inits, + ) + model = helper.make_model( + graph, opset_imports=[helper.make_opsetid("com.microsoft", 1), helper.make_opsetid("", 17)] + ) + return model.SerializeToString() + + +def make_session(model): + so = onnxruntime.SessionOptions() + so.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_DISABLE_ALL + so.log_severity_level = 3 + return onnxruntime.InferenceSession(model, so, providers=["CUDAExecutionProvider"]) + + +def run_case(name, m, k, n, block_size, bits, dtype, warmup, repeat): + onnx_dtype = _OT[dtype] + torch_dtype = _TT[dtype] + sess = make_session(build_model(k, n, block_size, bits, onnx_dtype)) + a = np.random.default_rng(m).random((m, k)).astype(np.float32) * 0.02 - 0.01 + at = torch.from_numpy(a).to(torch_dtype).cuda().contiguous() + y = torch.empty((m, n), dtype=torch_dtype, device="cuda") + io = sess.io_binding() + io.bind_input("A", "cuda", 0, _ELEM[torch_dtype], list(at.shape), at.data_ptr()) + io.bind_output("Y", "cuda", 0, _ELEM[torch_dtype], list(y.shape), y.data_ptr()) + + with _nvtx_range("warmup", "yellow"): + for _ in range(warmup): + sess.run_with_iobinding(io) + torch.cuda.synchronize() + + best = float("inf") + trials = 10 + with _nvtx_range("benchmark", "green"): + for _ in range(trials): + torch.cuda.synchronize() + t0 = time.perf_counter() + for _ in range(repeat): + sess.run_with_iobinding(io) + torch.cuda.synchronize() + best = min(best, (time.perf_counter() - t0) / repeat) + us = best * 1e6 + result = { + "case": name, + "m": m, + "k": k, + "n": n, + "block_size": block_size, + "bits": bits, + "dtype": dtype, + "avg_us": round(us, 2), + } + print(RESULT_PREFIX + json.dumps(result)) + return result + + +def main(): + p = argparse.ArgumentParser(description="Profile CUDA MatMulNBits") + p.add_argument("--k", type=int, help="in features (single-case mode)") + p.add_argument("--n", type=int, help="out features (single-case mode)") + p.add_argument("--m", type=int, help="rows (single-case mode)") + p.add_argument("--block-size", type=int, default=32) + p.add_argument("--bits", type=int, default=4) + p.add_argument("--dtype", default="fp16", choices=["fp16", "bf16"]) + p.add_argument("--warmup", type=int, default=25) + p.add_argument("--repeat", type=int, default=200) + p.add_argument("--ms", type=int, nargs="+", default=DEFAULT_MS, help="row counts to sweep in table mode") + args = p.parse_args() + + if args.k and args.n and args.m: + run_case("custom", args.m, args.k, args.n, args.block_size, args.bits, args.dtype, args.warmup, args.repeat) + return + + rows = {} + for name, k, n in DEFAULT_CASES: + rows[name] = { + m: run_case(name, m, k, n, args.block_size, args.bits, args.dtype, args.warmup, args.repeat)["avg_us"] + for m in args.ms + } + + print(f"\n MatMulNBits {args.bits}-bit block{args.block_size} {args.dtype} (avg us, lower is better)") + print(" " + f"{'matrix':10s} {'K':>6} {'N':>7} " + " ".join(f"M={m:<5}" for m in args.ms)) + for name, k, n in DEFAULT_CASES: + print(" " + f"{name:10s} {k:>6} {n:>7} " + " ".join(f"{rows[name][m]:7.1f}" for m in args.ms)) + + +if __name__ == "__main__": + main() From d7592544248c90bef67476e6a1c18a15c0dd1bf1 Mon Sep 17 00:00:00 2001 From: Fei Chen Date: Fri, 3 Jul 2026 05:53:56 +0800 Subject: [PATCH 12/17] WebGPU: Add indirect dispatch for flash attention graph capture to support Gemma4 (#29236) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary This adds indirect dispatch support for WebGPU flash attention to enable graph capture for Gemma4. When graph capture is on, \`total_seqlen\` is GPU-resident and cannot be read on the CPU, so dispatch group sizes must be computed on GPU. CopyKVCache normally prepares the indirect dispatch buffer as a side effect. For Gemma4 kv_empty layers (shared KV layers that skip CopyKVCache), a dedicated \`PrepareIndirectDispatchProgram\` fills the indirect buffer instead — avoiding \`dispatch(0)\` crashes. - Adds \`PrepareIndirectDispatchProgram\`: a single-thread GPU kernel that reads \`total_sequence_length_input[0]\` and writes 3 x uint32 dispatch dims to \`indirect_buffer\`, matching the logic in \`CopyKVCacheProgram\` - \`use_indirect_dispatch\` is gated on \`context.IsGraphCaptureEnabled()\` — indirect dispatch is only needed when \`total_seqlen\` is GPU-resident; when graph capture is off, CPU-side dispatch sizing works correctly even for kv_empty layers - \`PrepareIndirectDispatchProgram\` takes \`total_sequence_length_input\` (the global max, batch-safe) rather than \`seqlen_k\` (per-batch index 0, unsafe for batch > 1) ## Changes - \`flash_attention.cc\`: - New \`PrepareIndirectDispatchProgram::GenerateShaderCode\`: reads \`total_sequence_length_input[0]\`, computes tile count, calls \`populate_indirect_dispatch_buffer\` — identical logic to \`CopyKVCacheProgram\`'s indirect dispatch block - \`use_indirect_dispatch\` gated on \`seqlen_k != nullptr && total_seqlen != nullptr && context.IsGraphCaptureEnabled()\` - kv_empty path calls \`PrepareIndirectDispatchProgram\` when \`use_indirect_dispatch\` is true (i.e., under graph capture) - \`flash_attention.h\`: Added \`PrepareIndirectDispatchProgram\` class declaration ## Test plan - [x] Verified with Gemma4 INT4 model: ~95-105 tok/s with GC=ON vs ~75 tok/s with GC=OFF - [x] Tested multiple prompts (short/long) with graph capture enabled — no crashes, correct output - [x] Verified warmup (capture run) followed by replay produces consistent results - [x] Unit test: \`WebGPU_SharedKV_IndirectDispatchForGraphCapture\` — exercises the full ORT graph capture simulation path end-to-end: - Model built via the Graph API (opset 17) with all GQA inputs declared as proper graph inputs - All inputs allocated as GPU-resident tensors via \`InferenceSession::GetAllocator\` + \`IOBinding\`; \`total_seqlen\` is a real \`WGPUBuffer\` so \`PrepareIndirectDispatchProgram\` reads it from GPU - WebGPU EP registered with \`kEnableGraphCapture=ON\`; first \`Run\` captures, second \`Run\` replays - Replay output copied GPU→CPU and compared against a CPU reference to verify correctness - Covers the kv_empty branch specifically: \`key\`/\`value\` have \`sequence_length=0\`, triggering the \`PrepareIndirectDispatchProgram\` code path (CopyKVCache is skipped for kv_empty layers) --------- Co-authored-by: Claude Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> --- .../webgpu/bert/flash_attention.cc | 28 +++ .../contrib_ops/webgpu/bert/flash_attention.h | 14 ++ .../webgpu/bert/group_query_attention.cc | 13 +- .../group_query_attention_op_test.cc | 199 ++++++++++++++++++ 4 files changed, 248 insertions(+), 6 deletions(-) diff --git a/onnxruntime/contrib_ops/webgpu/bert/flash_attention.cc b/onnxruntime/contrib_ops/webgpu/bert/flash_attention.cc index ff6938c0e9bf4..ce4018956f5c0 100644 --- a/onnxruntime/contrib_ops/webgpu/bert/flash_attention.cc +++ b/onnxruntime/contrib_ops/webgpu/bert/flash_attention.cc @@ -159,6 +159,17 @@ Status CopyKVCacheProgram::GenerateShaderCode(ShaderHelper& shader) const { return Status::OK(); } +Status PrepareIndirectDispatchProgram::GenerateShaderCode(ShaderHelper& shader) const { + shader.AddInput("total_sequence_length_input", ShaderUsage::None); + shader.AddOutput("indirect_buffer", ShaderUsage::None); + shader.AdditionalImplementation() << kPopulateIndirectDispatchBufferFn; + shader.MainFunctionBody() + << " let global_total_seq_length = u32(total_sequence_length_input[0]);\n" + << " let num_total_seq_length_tile = (global_total_seq_length + uniforms.tile_size - 1u) / uniforms.tile_size;\n" + << " populate_indirect_dispatch_buffer(num_total_seq_length_tile, uniforms.num_heads * uniforms.num_q_tiles, uniforms.batch_size);\n"; + return Status::OK(); +} + Status CopyKVCache(onnxruntime::webgpu::ComputeContext& context, const WebgpuAttentionParameters& parameters, const Tensor* K, const Tensor* past_key, Tensor* present_key, const Tensor* V, const Tensor* past_value, Tensor* present_value, @@ -547,6 +558,23 @@ Status ApplyFlashAttention(const Tensor* Q, const Tensor* K, const Tensor* V, co present_key = const_cast(past_key); present_value = const_cast(past_value); } + + // CopyKVCache normally prepares the indirect dispatch buffer. For kv_empty layers + // CopyKVCache is skipped, so we prepare it here. Only needed under graph capture + // because that is when total_seqlen is GPU-resident and CPU-side dispatch sizing + // is unavailable. + if (use_indirect_dispatch) { + PrepareIndirectDispatchProgram program; + program.AddInput({total_seqlen, ProgramTensorMetadataDependency::None}); + program.AddOutput({indirect_buffer_ptr, ProgramTensorMetadataDependency::None}); + program.SetDispatchGroupSize(1) + .SetWorkgroupSize(1) + .AddUniformVariables({{tile_size}, + {static_cast(parameters.num_heads_)}, + {num_q_tiles}, + {static_cast(parameters.batch_size_)}}); + ORT_RETURN_IF_ERROR(context.RunProgram(program)); + } } // When TurboQuant is active, create u32 tensor views over present/past KV cache buffers. diff --git a/onnxruntime/contrib_ops/webgpu/bert/flash_attention.h b/onnxruntime/contrib_ops/webgpu/bert/flash_attention.h index 02f25a753fff8..f917f50c7fd72 100644 --- a/onnxruntime/contrib_ops/webgpu/bert/flash_attention.h +++ b/onnxruntime/contrib_ops/webgpu/bert/flash_attention.h @@ -70,6 +70,20 @@ class CopyKVCacheProgram final : public Program { bool use_seqlen_k_; }; +class PrepareIndirectDispatchProgram final : public Program { + public: + PrepareIndirectDispatchProgram() + : Program{"PrepareIndirectDispatch"} {} + + Status GenerateShaderCode(ShaderHelper& sh) const override; + + WEBGPU_PROGRAM_DEFINE_UNIFORM_VARIABLES( + {"tile_size", ProgramUniformVariableDataType::Uint32}, + {"num_heads", ProgramUniformVariableDataType::Uint32}, + {"num_q_tiles", ProgramUniformVariableDataType::Uint32}, + {"batch_size", ProgramUniformVariableDataType::Uint32}); +}; + class FlashAttentionProgram final : public Program { public: FlashAttentionProgram(const std::string& kernel_name, diff --git a/onnxruntime/contrib_ops/webgpu/bert/group_query_attention.cc b/onnxruntime/contrib_ops/webgpu/bert/group_query_attention.cc index 879a63a2482ca..7cc00555b6789 100644 --- a/onnxruntime/contrib_ops/webgpu/bert/group_query_attention.cc +++ b/onnxruntime/contrib_ops/webgpu/bert/group_query_attention.cc @@ -353,7 +353,13 @@ Status GroupQueryAttention::ComputeInternal(onnxruntime::webgpu::ComputeContext& past_value->DataRaw() == present_value->DataRaw(); ORT_ENFORCE(parameters.total_sequence_length_ <= parameters.seqlen_present_kv_cache_, "Total sequence length cannot be greater than the existing KV cache length."); - ORT_ENFORCE(!context.IsGraphCaptureEnabled() || parameters.past_present_share_buffer_, + // kv_sequence_length==0 fast path: K/V inputs are empty (shared KV layer). + // Skip all K/V processing; only apply RoPE to Q if needed. + // Use past_key/past_value directly as the KV context. + const bool kv_empty = (parameters.kv_sequence_length_ == 0); + // kv_empty layers (e.g. Gemma4 layers 15-34) reuse KV from another layer so + // past/present cannot share the same buffer — exempt them from this check. + ORT_ENFORCE(!context.IsGraphCaptureEnabled() || kv_empty || parameters.past_present_share_buffer_, "Graph capture requires past/present KV cache to share the same buffer (static KV cache)."); Tensor qSplit; @@ -363,11 +369,6 @@ Status GroupQueryAttention::ComputeInternal(onnxruntime::webgpu::ComputeContext& Tensor qRotary; Tensor kRotary; - // kv_sequence_length==0 fast path: K/V inputs are empty (shared KV layer). - // Skip all K/V processing; only apply RoPE to Q if needed. - // Use past_key/past_value directly as the KV context. - const bool kv_empty = (parameters.kv_sequence_length_ == 0); - // Use a sliding window if the total sequence exceeds the window's length. bool use_sliding_window = (local_window_size_ != -1 && local_window_size_ < parameters.total_sequence_length_); bool will_use_flash_attention = false; diff --git a/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc b/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc index ef7a04bc67d4f..16af4a4e489aa 100644 --- a/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc +++ b/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc @@ -14,7 +14,12 @@ #include "test/providers/provider_test_utils.h" #include "test/util/include/default_providers.h" #ifdef USE_WEBGPU +#include "core/graph/model.h" #include "core/providers/webgpu/webgpu_provider_options.h" +#include "core/session/inference_session.h" +#include "core/session/IOBinding.h" +#include "test/test_environment.h" +#include "test/unittest_util/framework_test_utils.h" #endif namespace onnxruntime { @@ -2562,6 +2567,200 @@ TEST(GroupQueryAttentionTest, WebGPU_SharedKV_SlidingWindow) { tester.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers); } +#ifdef USE_WEBGPU +// WebGPU graph capture test for kv_empty (Gemma4 shared-KV) layers. +// +// When graph capture is enabled, total_seqlen is GPU-resident and +// PrepareIndirectDispatchProgram must compute dispatch sizes on GPU. +// This test exercises the full ORT graph capture path by allocating all +// inputs as GPU tensors via IOBinding, running capture then replay, and +// verifying the replay output matches the CPU reference. +TEST(GroupQueryAttentionTest, WebGPU_SharedKV_IndirectDispatchForGraphCapture) { + constexpr int batch_size = 1; + constexpr int q_seq_len = 1; + constexpr int past_seq_len = 32; + constexpr int num_heads = 2; + constexpr int kv_num_heads = 1; + constexpr int head_size = 8; + constexpr int hidden_size = num_heads * head_size; + constexpr int kv_hidden_size = kv_num_heads * head_size; + + // Build a GQA model directly using the Graph API so graph inputs are + // properly declared (OpTester::BuildModel doesn't call graph.Resolve or + // SetInputs, producing a proto that InferenceSession::Load rejects). + std::unique_ptr p_model; + { + std::unordered_map domain_to_version; + domain_to_version[kOnnxDomain] = 17; + domain_to_version[kMSDomain] = 1; + p_model = std::make_unique( + "gqa_gc_test", true, ModelMetaData(), PathString(), + IOnnxRuntimeOpSchemaRegistryList(), domain_to_version, + std::vector{}, + DefaultLoggingManager().DefaultLogger(), + ModelOptions(true, true)); + onnxruntime::Graph& graph = p_model->MainGraph(); + + ONNX_NAMESPACE::TypeProto tp_float, tp_int32; + tp_float.mutable_tensor_type()->set_elem_type(ONNX_NAMESPACE::TensorProto_DataType_FLOAT); + tp_int32.mutable_tensor_type()->set_elem_type(ONNX_NAMESPACE::TensorProto_DataType_INT32); + + std::vector in_defs = { + &graph.GetOrCreateNodeArg("query", &tp_float), + &graph.GetOrCreateNodeArg("key", &tp_float), + &graph.GetOrCreateNodeArg("value", &tp_float), + &graph.GetOrCreateNodeArg("past_key", &tp_float), + &graph.GetOrCreateNodeArg("past_value", &tp_float), + &graph.GetOrCreateNodeArg("seqlens_k", &tp_int32), + &graph.GetOrCreateNodeArg("total_sequence_length", &tp_int32), + // optional inputs 8-12 are absent — use empty NodeArgs + &graph.GetOrCreateNodeArg("", nullptr), + &graph.GetOrCreateNodeArg("", nullptr), + &graph.GetOrCreateNodeArg("", nullptr), + &graph.GetOrCreateNodeArg("", nullptr), + &graph.GetOrCreateNodeArg("", nullptr), + }; + std::vector out_defs = { + &graph.GetOrCreateNodeArg("output", &tp_float), + }; + + auto& node = graph.AddNode("gqa", "GroupQueryAttention", "GQA", + in_defs, out_defs, nullptr, kMSDomain); + node.AddAttribute("num_heads", static_cast(num_heads)); + node.AddAttribute("kv_num_heads", static_cast(kv_num_heads)); + + ORT_THROW_IF_ERROR(graph.Resolve()); + } + + std::string model_data; + p_model->ToProto().SerializeToString(&model_data); + + // Create InferenceSession with graph capture EP. + SessionOptions so; + InferenceSession session{so, GetEnvironment()}; + ConfigOptions config_options{}; + ORT_THROW_IF_ERROR(config_options.AddConfigEntry(webgpu::options::kEnableGraphCapture, + webgpu::options::kEnableGraphCapture_ON)); + auto webgpu_ep = WebGpuExecutionProviderWithOptions(config_options); + if (!webgpu_ep) { + GTEST_SKIP() << "WebGPU EP not available"; + } + IExecutionProvider* ep_ptr = webgpu_ep.get(); + ORT_THROW_IF_ERROR(session.RegisterExecutionProvider(std::move(webgpu_ep))); + std::istringstream model_stream(model_data); + ORT_THROW_IF_ERROR(session.Load(model_stream)); + ORT_THROW_IF_ERROR(session.Initialize()); + + // Get GPU allocator from session. + OrtMemoryInfo gpu_mem_info(WEBGPU_BUFFER, OrtAllocatorType::OrtDeviceAllocator, + OrtDevice(OrtDevice::GPU, OrtDevice::MemType::DEFAULT, OrtDevice::VendorIds::NONE, 0)); + auto gpu_alloc = session.GetAllocator(gpu_mem_info); + AllocatorPtr cpu_alloc = TestCPUExecutionProvider()->CreatePreferredAllocators()[0]; + + // Capture-run input data (set A). + std::vector query_data(batch_size * q_seq_len * hidden_size); + std::vector past_key_data(batch_size * kv_num_heads * past_seq_len * head_size); + std::vector past_value_data(batch_size * kv_num_heads * past_seq_len * head_size); + for (size_t i = 0; i < query_data.size(); i++) query_data[i] = 0.1f * static_cast(i % 7 + 1); + for (size_t i = 0; i < past_key_data.size(); i++) past_key_data[i] = 0.2f * static_cast(i % 5 + 1); + for (size_t i = 0; i < past_value_data.size(); i++) past_value_data[i] = 0.3f * static_cast(i % 3 + 1); + + // Replay-run input data (set B) — different values to verify replay actually uses new data. + std::vector query_data2(batch_size * q_seq_len * hidden_size); + std::vector past_key_data2(batch_size * kv_num_heads * past_seq_len * head_size); + std::vector past_value_data2(batch_size * kv_num_heads * past_seq_len * head_size); + for (size_t i = 0; i < query_data2.size(); i++) query_data2[i] = 0.5f * static_cast(i % 11 + 1); + for (size_t i = 0; i < past_key_data2.size(); i++) past_key_data2[i] = 0.4f * static_cast(i % 7 + 1); + for (size_t i = 0; i < past_value_data2.size(); i++) past_value_data2[i] = 0.6f * static_cast(i % 4 + 1); + + // Helper: create GPU OrtValue by copying from CPU data (nullptr = zero-size tensor). + auto make_gpu_value = [&](const void* data, MLDataType dtype, TensorShape shape) -> OrtValue { + Tensor gpu_tensor(dtype, shape, gpu_alloc); + if (data != nullptr && shape.Size() > 0) { + Tensor cpu_tensor(dtype, shape, const_cast(data), cpu_alloc->Info()); + ORT_THROW_IF_ERROR(ep_ptr->GetDataTransfer()->CopyTensor(cpu_tensor, gpu_tensor)); + } + OrtValue val; + Tensor::InitOrtValue(std::move(gpu_tensor), val); + return val; + }; + + // Helper: overwrite an existing GPU tensor's data in-place from CPU (same buffer, new values). + // Under graph capture the WGPUBuffer pointer is baked in; rebinding is not allowed. + auto update_gpu_value = [&](const OrtValue& gpu_val, const void* data, MLDataType dtype, TensorShape shape) { + if (shape.Size() > 0) { + Tensor cpu_tensor(dtype, shape, const_cast(data), cpu_alloc->Info()); + ORT_THROW_IF_ERROR(ep_ptr->GetDataTransfer()->CopyTensor(cpu_tensor, const_cast(gpu_val.Get()))); + } + }; + + auto q_val = make_gpu_value(query_data.data(), DataTypeImpl::GetType(), + TensorShape{batch_size, q_seq_len, hidden_size}); + auto key_val = make_gpu_value(nullptr, DataTypeImpl::GetType(), + TensorShape{batch_size, 0, kv_hidden_size}); + auto value_val = make_gpu_value(nullptr, DataTypeImpl::GetType(), + TensorShape{batch_size, 0, kv_hidden_size}); + auto pk_val = make_gpu_value(past_key_data.data(), DataTypeImpl::GetType(), + TensorShape{batch_size, kv_num_heads, past_seq_len, head_size}); + auto pv_val = make_gpu_value(past_value_data.data(), DataTypeImpl::GetType(), + TensorShape{batch_size, kv_num_heads, past_seq_len, head_size}); + std::vector seqlens_k_data = {past_seq_len - 1}; + auto sk_val = make_gpu_value(seqlens_k_data.data(), DataTypeImpl::GetType(), + TensorShape{batch_size}); + std::vector total_seqlen_data = {past_seq_len}; + auto ts_val = make_gpu_value(total_seqlen_data.data(), DataTypeImpl::GetType(), + TensorShape{1}); + + // Preallocate GPU output. + Tensor gpu_out_tensor(DataTypeImpl::GetType(), + TensorShape{batch_size, q_seq_len, hidden_size}, gpu_alloc); + OrtValue out_val; + Tensor::InitOrtValue(std::move(gpu_out_tensor), out_val); + + // Bind inputs and output. + std::unique_ptr io_binding; + ORT_THROW_IF_ERROR(session.NewIOBinding(&io_binding)); + ORT_THROW_IF_ERROR(io_binding->BindInput("query", q_val)); + ORT_THROW_IF_ERROR(io_binding->BindInput("key", key_val)); + ORT_THROW_IF_ERROR(io_binding->BindInput("value", value_val)); + ORT_THROW_IF_ERROR(io_binding->BindInput("past_key", pk_val)); + ORT_THROW_IF_ERROR(io_binding->BindInput("past_value", pv_val)); + ORT_THROW_IF_ERROR(io_binding->BindInput("seqlens_k", sk_val)); + ORT_THROW_IF_ERROR(io_binding->BindInput("total_sequence_length", ts_val)); + ORT_THROW_IF_ERROR(io_binding->BindOutput("output", out_val)); + ORT_THROW_IF_ERROR(io_binding->SynchronizeInputs()); + + // Run 1: capture. + RunOptions run_options; + ORT_THROW_IF_ERROR(session.Run(run_options, *io_binding)); + + // Run 2: replay with different inputs (set B) written into the same GPU buffers. + // Rebinding is not allowed under graph capture — the WGPUBuffer pointers are baked + // in at capture time, so new data must be copied into the existing buffers. + update_gpu_value(q_val, query_data2.data(), DataTypeImpl::GetType(), + TensorShape{batch_size, q_seq_len, hidden_size}); + update_gpu_value(pk_val, past_key_data2.data(), DataTypeImpl::GetType(), + TensorShape{batch_size, kv_num_heads, past_seq_len, head_size}); + update_gpu_value(pv_val, past_value_data2.data(), DataTypeImpl::GetType(), + TensorShape{batch_size, kv_num_heads, past_seq_len, head_size}); + ORT_THROW_IF_ERROR(session.Run(run_options, *io_binding)); + + // Copy replay output GPU -> CPU and compare against CPU reference for set B. + const int output_size = batch_size * q_seq_len * hidden_size; + auto& gpu_out = io_binding->GetOutputs()[0].Get(); + Tensor cpu_out_tensor(DataTypeImpl::GetType(), gpu_out.Shape(), cpu_alloc); + ORT_THROW_IF_ERROR(ep_ptr->GetDataTransfer()->CopyTensor(gpu_out, cpu_out_tensor)); + std::vector webgpu_output(cpu_out_tensor.Data(), + cpu_out_tensor.Data() + output_size); + + auto cpu_output = RunGQASharedKV( + batch_size, q_seq_len, past_seq_len, query_data2, past_key_data2, past_value_data2, + num_heads, kv_num_heads, head_size, GqaTargetEp::kCpu); + + ExpectOutputsMatch(webgpu_output, cpu_output, 0.05f, "SharedKV_IndirectDispatchForGraphCapture_vs_CPU"); +} +#endif // USE_WEBGPU + // --------------------------------------------------------------------------- // Batched right-padded packed-QKV prefill with do_rotary. // From 99ba5db1816e39f4b9ca0b88668e387d21c7eb8f Mon Sep 17 00:00:00 2001 From: Chi Lo <54722500+chilo-ms@users.noreply.github.com> Date: Thu, 2 Jul 2026 16:52:12 -0700 Subject: [PATCH 13/17] Fix copy-paste typos in DynamicQuantizeLSTM zero-point and scale validation (#29462) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ### Description Two copy-paste typos in  DynamicQuantizeLSTM::Compute  ( dynamic_quantize_lstm.cc ) cause the recurrence-weight zero-point and scale to be validated against the wrong tensor's shape: 1. L181:  R_zp_shape = w_zp->Shape()  →  r_zp->Shape() .  ZeroPointCheck  then iterates  w_zp 's element count over the smaller  r_zp  tensor, reading past it (OOB read). 2. L188:  WeightCheck(W_scale_shape, R_scale)  →  WeightCheck(R_scale_shape, R_scale) . The recurrence scale shape is validated against the input scale shape instead of its own. ### Motivation and Context Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../contrib_ops/cpu/quantization/dynamic_quantize_lstm.cc | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/onnxruntime/contrib_ops/cpu/quantization/dynamic_quantize_lstm.cc b/onnxruntime/contrib_ops/cpu/quantization/dynamic_quantize_lstm.cc index 2094af78f40b7..ba3c8a2ada30b 100644 --- a/onnxruntime/contrib_ops/cpu/quantization/dynamic_quantize_lstm.cc +++ b/onnxruntime/contrib_ops/cpu/quantization/dynamic_quantize_lstm.cc @@ -178,14 +178,14 @@ Status DynamicQuantizeLSTM::Compute(OpKernelContext* context) const { const Tensor* r_zp = context->Input(11); const TensorShape& W_zp_shape = w_zp->Shape(); - const TensorShape& R_zp_shape = w_zp->Shape(); + const TensorShape& R_zp_shape = r_zp->Shape(); const TensorShape& W_scale_shape = w_scale->Shape(); const TensorShape& R_scale_shape = r_scale->Shape(); WeightCheck(W_zp_shape, W_zero_point); WeightCheck(R_zp_shape, R_zero_point); WeightCheck(W_scale_shape, W_scale); - WeightCheck(W_scale_shape, R_scale); + WeightCheck(R_scale_shape, R_scale); const bool is_W_signed = (W != nullptr) ? W->IsDataType() : is_W_signed_; const bool is_R_signed = (R != nullptr) ? R->IsDataType() : is_R_signed_; From 3eb4a3564d6c7ec548558bb9cc21d9a8c4f20e7f Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Thu, 2 Jul 2026 16:53:06 -0700 Subject: [PATCH 14/17] Fix GQA out-of-bounds read in past KV buffer (CPU) (#29447) ### Description Fixes an out-of-bounds (OOB) read in the CPU `GroupQueryAttention` (GQA) operator. During token generation (decode), `ConcatStateChunkGQA` copies `seqlens_k + 1 - sequence_length` rows out of the past key/value buffers. The existing validation only bounded the *present* buffer write (`seqlens_k < present_kv_seqlen`), but the present buffer can be larger than the past buffer when `total_sequence_length` exceeds the past sequence length. A large `seqlens_k` combined with a small past buffer therefore read past the end of the past key/value tensors. ### Motivation and Context `present_kv_seqlen = max(total_sequence_length, past_sequence_length)`, so the pre-existing `seqlens_k < present_kv_seqlen` check does not bound the past-side read when `total_sequence_length > past_sequence_length`. With a crafted `seqlens_k`, the decode path in `ConcatStateChunkGQA` reads `seqlens_k + 1 - sequence_length` rows from the smaller past buffer, causing an OOB read. ### Key Changes | File | Change | |---|---| | `onnxruntime/contrib_ops/cpu/bert/group_query_attention.cc` | Add a per-batch bound in the `seqlens_k` validation block: for the decode case (`past_kv_seqlen > 0 && kv_sequence_length != 0 && !is_first_prompt`), reject when `seqlens_k + 1 - sequence_length > past_kv_seqlen`. | | `onnxruntime/test/contrib_ops/group_query_attention_op_test.cc` | Add regression test `SeqlensKExceedsPastBuffer_OOBRead` exercising a large `seqlens_k` against a small past buffer. | Shared KV (`kv_sequence_length == 0`) appends no new KV and its past read is already bounded by the present-buffer check together with the `total_sequence_length <= seqlen_past_kv_cache` enforcement in the apply-attention paths, so it needs no additional check. ### Testing Notes - Build and run the GQA tests: ``` cmake --build build/Linux/Debug --target onnxruntime_provider_test -j$(nproc) ./build/Linux/Debug/onnxruntime_provider_test --gtest_filter="*GroupQueryAttention*" ``` - All 41 CPU GQA tests pass, including the new `SeqlensKExceedsPastBuffer_OOBRead` regression and the existing shared-KV CPU cases. - `lintrunner` reports no issues on the changed files. --- .../cpu/bert/group_query_attention.cc | 20 +++++++++++++ .../group_query_attention_op_test.cc | 29 +++++++++++++++++++ 2 files changed, 49 insertions(+) diff --git a/onnxruntime/contrib_ops/cpu/bert/group_query_attention.cc b/onnxruntime/contrib_ops/cpu/bert/group_query_attention.cc index debda282eb4f1..1080498831474 100644 --- a/onnxruntime/contrib_ops/cpu/bert/group_query_attention.cc +++ b/onnxruntime/contrib_ops/cpu/bert/group_query_attention.cc @@ -145,6 +145,7 @@ Status GroupQueryAttention::Compute(OpKernelContext* context) const { // Validate seqlens_k values before they are used as GEMM dimensions to prevent OOB access. { const int32_t* seqlens_k_data = seqlens_k->Data(); + const int past_kv_seqlen = parameters.seqlen_past_kv_cache; for (int b = 0; b < batch_size; b++) { if (seqlens_k_data[b] < 0 || seqlens_k_data[b] >= present_kv_seqlen) { return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, @@ -156,6 +157,25 @@ Status GroupQueryAttention::Compute(OpKernelContext* context) const { "seqlens_k[", b, "] = ", seqlens_k_data[b], " is too small for sequence_length ", sequence_length); } + // Bound the number of past KV rows copied out of the past key/value buffers during + // token generation (decode). ConcatStateChunkGQA copies (seqlens_k + 1 - sequence_length) + // rows from the past buffer (sized past_kv_seqlen). The present-buffer check above does + // not bound this past-side read, because the present buffer can be larger than the past + // buffer when total_sequence_length exceeds the past sequence length. A large seqlens_k + // combined with a small past buffer would otherwise read past the end of the past tensors. + // Shared KV (kv_sequence_length == 0) appends no new KV and its past read is already + // bounded by the present-buffer check together with the total_sequence_length <= + // seqlen_past_kv_cache enforcement in the apply-attention paths, so it needs no check here. + if (past_key != nullptr && past_value != nullptr && parameters.kv_sequence_length != 0 && + !parameters.is_first_prompt) { + const int64_t past_rows = static_cast(seqlens_k_data[b]) + 1 - sequence_length; + if (past_rows > past_kv_seqlen) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "seqlens_k[", b, "] = ", seqlens_k_data[b], " requires ", past_rows, + " past KV rows, which exceeds the past buffer sequence length ", + past_kv_seqlen, "."); + } + } } } int q_hidden_size = parameters.hidden_size; diff --git a/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc b/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc index 16af4a4e489aa..64c3f206b8a93 100644 --- a/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc +++ b/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc @@ -374,6 +374,35 @@ TEST(GroupQueryAttentionTest, NonPromptSeqlensKUnderflow_OOB) { /*past_seq_len=*/4); } +// Regression: present buffer large enough (total_seq_len passes the present-buffer check), +// but the past buffer is much smaller. ConcatStateChunkGQA would copy +// (seqlens_k + 1 - sequence_length) rows out of the small past buffer, reading past its end. +TEST(GroupQueryAttentionTest, SeqlensKExceedsPastBuffer_OOBRead) { + // present_kv_seqlen = max(total_seq_len=100, past_seq_len=2) = 100, so seqlens_k=50 passes the + // present-buffer check, but past_seqlen = 51 - 1 = 50 rows >> past buffer (2 rows) => OOB read. + RunGQASeqlensKTest( + /*seqlens_k_data=*/{50}, + /*total_seq_len=*/100, + /*batch_size=*/1, + /*sequence_length=*/1, + OpTester::ExpectResult::kExpectFailure, + "exceeds the past buffer sequence length", + /*provide_past=*/true, + /*past_seq_len=*/2); +} + +TEST(GroupQueryAttentionTest, SeqlensKExceedsEmptyPastBuffer_OOBRead) { + RunGQASeqlensKTest( + /*seqlens_k_data=*/{50}, + /*total_seq_len=*/100, + /*batch_size=*/1, + /*sequence_length=*/1, + OpTester::ExpectResult::kExpectFailure, + "exceeds the past buffer sequence length", + /*provide_past=*/true, + /*past_seq_len=*/0); +} + // INT32_MAX seqlens_k: rejected by the >= present_kv_seqlen check. TEST(GroupQueryAttentionTest, Int32MaxSeqlensK_OOB) { RunGQASeqlensKTest( From d7efadab308ed9bde763d8b014e30c2e5664b614 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Thu, 2 Jul 2026 18:10:30 -0700 Subject: [PATCH 15/17] Fix negative-axis handling in ExpandDims shape inference (#29448) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit # PR: Fix negative-axis handling in ExpandDims shape inference ## Description The `com.microsoft.ExpandDims` type/shape-inference function mishandled negative `axis` values, which could lead to an out-of-bounds read during graph resolution (`Graph::Resolve`). This PR corrects the axis normalization and makes the scalar `axis` read robust to both tensor encodings, and adds a regression test that exercises the shape-inference path. ## Summary of Changes ### Fix | File | Change | |------|--------| | `onnxruntime/core/graph/contrib_ops/contrib_defs.cc` | Normalize a negative `axis` against the output rank (`rank + axis + 1`) instead of the off-by-two `rank + axis - 1`, so the insertion index stays within `[0, rank]`. Read the scalar `axis` via the existing `ParseScalar` helper, which handles both `raw_data` and `int32_data` encodings and validates the element count. | ### Test | File | Change | |------|--------| | `onnxruntime/test/contrib_ops/expand_dims_test.cc` | Add `ExpandDimsTest.NegativeAxisConstInitializerShapeInference` plus a `RunExpandDimsConstAxisTest` helper that supplies `axis` as a constant initializer so the operator's shape-inference function is exercised (the existing tests pass `axis` as a runtime input, which skips that path). | ## Details - For an output of `rank + 1` dimensions, a negative `axis` must be normalized as `axis + (rank + 1)`. The previous `rank + axis - 1` formula produced a negative insertion index for the most-negative valid axes, which was then used to index the protobuf dimension list out of bounds. - The axis value was previously read with `int32_data()[0]`. When the value is stored as `raw_data` (the common encoding for serialized models and the one produced by the test harness), `int32_data()` is empty and the access is out of bounds. `ParseScalar` decodes either encoding and validates the count. ## Testing - Built `onnxruntime_provider_test` and ran the ExpandDims suite: `./onnxruntime_provider_test --gtest_filter="ExpandDimsTest.*"` — all 6 tests pass. - Confirmed the new regression test fails (process aborts) without the fix and passes with it. - Existing positive/negative out-of-range and kernel tests are unchanged. ## Checklist - [x] Tests added/updated - [ ] Documentation updated (if applicable) - [x] No breaking changes - [ ] CI passes --------- Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> --- .../core/graph/contrib_ops/contrib_defs.cc | 14 ++++- .../test/contrib_ops/expand_dims_test.cc | 52 +++++++++++++++++++ 2 files changed, 64 insertions(+), 2 deletions(-) diff --git a/onnxruntime/core/graph/contrib_ops/contrib_defs.cc b/onnxruntime/core/graph/contrib_ops/contrib_defs.cc index f3f2f521ecab2..6242edb64d26b 100644 --- a/onnxruntime/core/graph/contrib_ops/contrib_defs.cc +++ b/onnxruntime/core/graph/contrib_ops/contrib_defs.cc @@ -2002,11 +2002,21 @@ ONNX_MS_OPERATOR_SET_SCHEMA(ExpandDims, 1, const ONNX_NAMESPACE::TensorProto* axis_initializer = ctx.getInputData(1); if (!axis_initializer) return; - const int axis = axis_initializer->int32_data()[0]; + // Read the scalar axis robustly. ParseScalar handles both raw_data and + // int32_data encodings and validates the element count, avoiding an + // out-of-bounds read when the value is stored as raw_data. A present but + // malformed initializer (wrong element type or not a single scalar) is a + // model error, so fail shape inference rather than silently skipping it. + int axis = 0; + if (!ParseScalar(axis_initializer, axis)) { + fail_shape_inference("Input axis must be a single int32 scalar initializer"); + } if (axis > rank || axis < -rank - 1) { fail_shape_inference("Input axis is invalid: ", axis); } - int pos = axis >= 0 ? axis : rank + axis - 1; + // The output has rank + 1 dimensions, so a negative axis is normalized + // against the output rank: pos = axis + (rank + 1). + int pos = axis >= 0 ? axis : rank + axis + 1; ONNX_NAMESPACE::TensorShapeProto output_shape; for (int i = 0; i < pos; ++i) { output_shape.add_dim(); diff --git a/onnxruntime/test/contrib_ops/expand_dims_test.cc b/onnxruntime/test/contrib_ops/expand_dims_test.cc index cb81690ba954a..3a0dfed04eb6f 100644 --- a/onnxruntime/test/contrib_ops/expand_dims_test.cc +++ b/onnxruntime/test/contrib_ops/expand_dims_test.cc @@ -51,6 +51,39 @@ void RunExpandDimsExpectedFailureTest(const std::vector& input_shape, c test.AddOutput("Y", input_shape, input_data); test.Run(BaseTester::ExpectResult::kExpectFailure, expected_failure_message); } + +// Runs ExpandDims with the axis supplied as a constant initializer. This causes the +// operator's shape-inference function (which only runs when the axis value is known at +// graph-resolution time via getInputData) to be exercised, in addition to the kernel. +void RunExpandDimsConstAxisTest(const std::vector& input_shape, const int32_t expand_axis, + const std::vector& expected_output_shape) { + SCOPED_TRACE(MakeString("input_shape: ", TensorShape(input_shape), ", expand_axis: ", expand_axis)); + + const auto input_size = ShapeSize(input_shape); + const auto input_data = GenerateInput(input_size); + + OpTester test("ExpandDims", 1, onnxruntime::kMSDomain); + test.AddInput("X", input_shape, input_data); + test.AddInput("axis", {}, {expand_axis}, /*is_initializer=*/true); + test.AddOutput("Y", expected_output_shape, input_data); + test.Run(); +} + +void RunExpandDimsMalformedConstAxisTest(const std::vector& input_shape, + const std::vector& expand_axes, + const std::string& expected_failure_message) { + SCOPED_TRACE(MakeString("input_shape: ", TensorShape(input_shape))); + + const auto input_size = ShapeSize(input_shape); + const auto input_data = GenerateInput(input_size); + + OpTester test("ExpandDims", 1, onnxruntime::kMSDomain); + test.AddInput("X", input_shape, input_data); + test.AddInput("axis", {static_cast(expand_axes.size())}, expand_axes, + /*is_initializer=*/true); + test.AddOutput("Y", input_shape, input_data); + test.Run(BaseTester::ExpectResult::kExpectFailure, expected_failure_message); +} } // namespace TEST(ExpandDimsTest, Basic) { @@ -69,6 +102,25 @@ TEST(ExpandDimsTest, MinAxis) { RunExpandDimsTest({}, -1, {1}); } +// Regression test for a heap out-of-bounds read in the ExpandDims shape-inference function. +// When the axis is a constant initializer, shape inference runs during graph resolution and +// normalizes a negative axis against the output rank. The previous formula produced a negative +// dimension index for the most-negative axes (e.g. pos == -2), indexing before the start of the +// protobuf RepeatedPtrField backing store. These cases must resolve to valid output shapes. +TEST(ExpandDimsTest, NegativeAxisConstInitializerShapeInference) { + RunExpandDimsConstAxisTest({2, 3}, -1, {2, 3, 1}); + RunExpandDimsConstAxisTest({2, 3}, -2, {2, 1, 3}); + RunExpandDimsConstAxisTest({2, 3}, -3, {1, 2, 3}); // minimum valid axis; previously read out of bounds + RunExpandDimsConstAxisTest({}, -1, {1}); // scalar, minimum valid axis; previously read out of bounds +} + +#ifndef ORT_NO_EXCEPTIONS +TEST(ExpandDimsTest, MalformedConstInitializerAxisFailsShapeInference) { + RunExpandDimsMalformedConstAxisTest({2, 3}, {0, 1}, + "Input axis must be a single int32 scalar initializer"); +} +#endif // !ORT_NO_EXCEPTIONS + TEST(ExpandDimsTest, PositiveAxisOutOfRange) { RunExpandDimsExpectedFailureTest({2, 3}, 3, "Axis must be within range [-3, 2]. Axis is 3"); RunExpandDimsExpectedFailureTest({}, 1, "Axis must be within range [-1, 0]. Axis is 1"); From af0588b4947b039da2fd9bf8bed3a2bcd18c4022 Mon Sep 17 00:00:00 2001 From: Chi Lo <54722500+chilo-ms@users.noreply.github.com> Date: Fri, 3 Jul 2026 06:58:34 -0700 Subject: [PATCH 16/17] [EP ABI] Add API to select the best compiled model compatibility info from candidate strings (#28387) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ### Description This PR adds a new EP API, `SelectBestModelCandidate`, that selects the best model variant from a set of candidates described by key-value metadata. The EP evaluates each candidate's metadata against the given hardware device and optional session options, and returns the index of the best match. **Key design points:** - Each candidate is an `OrtKeyValuePairs` representing one model variant. This future-proofs the API — additional metadata keys can be added over time without changing the function signature. - **Single-model variants (simple case):** The KVP contains a single `ep_compatibility_info` key with the compatibility string from the ONNX model metadata. - **Multi-model variants:** When a variant bundles multiple sub-models (e.g., prefill + decode in a GenAI scenario), the KVP uses indexed keys so the EP can inspect each sub-model independently: - `num_models` — number of sub-models (e.g., "3") - `.ep_compatibility_info` — compatibility string for sub-model *i* (required per sub-model) - `.role` — role of sub-model *i*, e.g., "prefill", "decode" (optional) - `.future_meaningful_info` — additional EP-meaningful metadata for sub-model i (optional) A basic EP implementation validates all `.ep_compatibility_info` entries. An advanced implementation can also consider `role` or other metadata for smarter ranking. - This approach delegates variant selection entirely to the EP, which has the domain knowledge to handle structurally mismatched variants (different sub-model counts, different roles, etc.) without ORT needing to understand model roles or compute aggregated scores. ### Motivation and Context The existing `ValidateCompiledModelCompatibilityInfo()` alone is not sufficient for some EPs to determine the best compatible model when there are multiple candidates. For example, an EP may support multiple compilation modes (e.g., "speed optimized" vs "memory optimized") that produce different compatibility strings. The EP can implement this function to evaluate the candidate metadata and select the best compatible variant based on its own criteria and the target device. --------- Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../core/session/onnxruntime_ep_c_api.h | 51 +++++++++++++ .../session/plugin_ep/ep_factory_internal.cc | 1 + .../session/plugin_ep/ep_factory_internal.h | 9 +++ .../plugin_ep/ep_factory_internal_impl.h | 18 +++++ .../plugin_ep/forward_to_factory_impl.h | 11 +++ .../library/example_plugin_ep/ep_factory.cc | 74 +++++++++++++++++++ .../library/example_plugin_ep/ep_factory.h | 8 ++ .../test/autoep/test_ep_compatibility.cc | 67 +++++++++++++++++ 8 files changed, 239 insertions(+) create mode 100644 onnxruntime/test/autoep/test_ep_compatibility.cc diff --git a/include/onnxruntime/core/session/onnxruntime_ep_c_api.h b/include/onnxruntime/core/session/onnxruntime_ep_c_api.h index 34a57b11e1748..f8329d4e8dc61 100644 --- a/include/onnxruntime/core/session/onnxruntime_ep_c_api.h +++ b/include/onnxruntime/core/session/onnxruntime_ep_c_api.h @@ -3044,6 +3044,57 @@ struct OrtEpFactory { */ ORT_API2_STATUS(DeinitGraphicsInterop, _In_ OrtEpFactory* this_ptr, _In_ const OrtEpDevice* ep_device); + + /** \brief Select the best model variant candidate from metadata. + * + * Evaluates each candidate's metadata against the given hardware device and optional session options, + * and returns the index of the best match. + * + * Each candidate is an OrtKeyValuePairs representing one model variant. The KVP uses indexed keys + * so that the EP can inspect each model's metadata independently. A variant always has num_models >= 1. + * + * Required and optional keys: + * - "num_models" — number of models in this variant (>= 1) (required) + * - ".ep_compatibility_info" — compatibility string for model i (required per model) + * - ".role" — role/purpose of model i (e.g., "prefill", "decode") (optional) + * - ".future_meaningful_info" — additional EP-meaningful metadata for model i (optional) + * + * where is a zero-based index (e.g., "0.ep_compatibility_info", "1.ep_compatibility_info"). + * + * The implementer should loop from 0 to num_models - 1 and validate each ".ep_compatibility_info" entry. + * An advanced implementation may additionally consider "role" or other metadata when ranking candidates. + * + * **Why this function exists:** + * + * The existing ValidateCompiledModelCompatibilityInfo() alone is not sufficient for some EPs to determine the best + * compatible model when there are multiple candidates. For example, an EP may support multiple compilation modes + * (e.g., "speed optimized" vs "memory optimized") that produce different compatibility strings. The EP can implement + * this function to evaluate the candidate metadata and select the best compatible variant based on its own criteria, + * the target device, and the session options. + * + * If all candidates are unsupported, this function succeeds and sets `selected_index` to SIZE_MAX. + * + * \note The implementer should validate each ".ep_compatibility_info" in the candidate (e.g., by calling + * ValidateCompiledModelCompatibilityInfo for each one) before determining the best match. + * + * \param[in] this_ptr The OrtEpFactory instance. + * \param[in] device The target hardware device that the EP would run on. Must map to this EP. + * \param[in] candidates Array of OrtKeyValuePairs pointers (one per model variant). + * \param[in] num_candidates Number of candidates (i.e., number of model variants to evaluate). + * \param[in] session_options Optional session options to consider when selecting the best candidate. + * May be nullptr if no session-level preferences are relevant. + * \param[out] selected_index Selected candidate index, or SIZE_MAX if all unsupported. + * + * \snippet{doc} snippets.dox OrtStatus Return Value + * + * \since Version 1.28. + */ + ORT_API2_STATUS(SelectBestModelCandidate, _In_ OrtEpFactory* this_ptr, + _In_ const OrtHardwareDevice* device, + _In_reads_(num_candidates) const OrtKeyValuePairs* const* candidates, + _In_ size_t num_candidates, + _In_opt_ const OrtSessionOptions* session_options, + _Out_ size_t* selected_index); }; #ifdef __cplusplus diff --git a/onnxruntime/core/session/plugin_ep/ep_factory_internal.cc b/onnxruntime/core/session/plugin_ep/ep_factory_internal.cc index ca39d6e750088..fb8a90f25ce13 100644 --- a/onnxruntime/core/session/plugin_ep/ep_factory_internal.cc +++ b/onnxruntime/core/session/plugin_ep/ep_factory_internal.cc @@ -36,6 +36,7 @@ EpFactoryInternal::EpFactoryInternal(std::unique_ptr impl OrtEpFactory::GetHardwareDeviceIncompatibilityDetails = Forward::GetHardwareDeviceIncompatibilityDetails; OrtEpFactory::InitGraphicsInterop = Forward::InitGraphicsInterop; OrtEpFactory::DeinitGraphicsInterop = Forward::DeinitGraphicsInterop; + OrtEpFactory::SelectBestModelCandidate = Forward::SelectBestModelCandidate; } InternalExecutionProviderFactory::InternalExecutionProviderFactory(EpFactoryInternal& ep_factory, diff --git a/onnxruntime/core/session/plugin_ep/ep_factory_internal.h b/onnxruntime/core/session/plugin_ep/ep_factory_internal.h index 9ac883da06465..cdfe5347185b7 100644 --- a/onnxruntime/core/session/plugin_ep/ep_factory_internal.h +++ b/onnxruntime/core/session/plugin_ep/ep_factory_internal.h @@ -106,6 +106,15 @@ class EpFactoryInternal : public OrtEpFactory { return impl_->DeinitGraphicsInterop(ep_device); } + OrtStatus* SelectBestModelCandidate(_In_ const OrtHardwareDevice* device, + _In_reads_(num_candidates) const OrtKeyValuePairs* const* candidates, + _In_ size_t num_candidates, + _In_opt_ const OrtSessionOptions* session_options, + _Out_ size_t* selected_index) noexcept { + return impl_->SelectBestModelCandidate(device, candidates, num_candidates, + session_options, selected_index); + } + // Function ORT calls to release an EP instance. void ReleaseEp(OrtEp* /*ep*/) noexcept { // we never create an OrtEp so we should never be trying to release one diff --git a/onnxruntime/core/session/plugin_ep/ep_factory_internal_impl.h b/onnxruntime/core/session/plugin_ep/ep_factory_internal_impl.h index a6140b2ac260f..cbadabfd8c3eb 100644 --- a/onnxruntime/core/session/plugin_ep/ep_factory_internal_impl.h +++ b/onnxruntime/core/session/plugin_ep/ep_factory_internal_impl.h @@ -122,6 +122,24 @@ class EpFactoryInternalImpl { return nullptr; } + virtual OrtStatus* SelectBestModelCandidate( + _In_ const OrtHardwareDevice* device, + _In_reads_(num_candidates) const OrtKeyValuePairs* const* candidates, + _In_ size_t num_candidates, + _In_opt_ const OrtSessionOptions* session_options, + _Out_ size_t* selected_index) noexcept { + ORT_UNUSED_PARAMETER(device); + ORT_UNUSED_PARAMETER(session_options); + if (candidates == nullptr || num_candidates == 0 || selected_index == nullptr) { + return OrtApis::CreateStatus(ORT_INVALID_ARGUMENT, + "Invalid arguments to SelectBestModelCandidate."); + } + + // Default implementation: all candidates are unsupported, sets `selected_index` to SIZE_MAX. + *selected_index = SIZE_MAX; + return nullptr; + } + // Function ORT calls to release an EP instance. void ReleaseEp(OrtEp* ep); diff --git a/onnxruntime/core/session/plugin_ep/forward_to_factory_impl.h b/onnxruntime/core/session/plugin_ep/forward_to_factory_impl.h index d4d93152c4f80..9dbf25a8251a4 100644 --- a/onnxruntime/core/session/plugin_ep/forward_to_factory_impl.h +++ b/onnxruntime/core/session/plugin_ep/forward_to_factory_impl.h @@ -106,6 +106,17 @@ struct ForwardToFactoryImpl { return static_cast(this_ptr)->DeinitGraphicsInterop(ep_device); } + static OrtStatus* ORT_API_CALL SelectBestModelCandidate( + OrtEpFactory* this_ptr, + _In_ const OrtHardwareDevice* device, + _In_reads_(num_candidates) const OrtKeyValuePairs* const* candidates, + size_t num_candidates, + _In_opt_ const OrtSessionOptions* session_options, + size_t* selected_index) noexcept { + return static_cast(this_ptr)->SelectBestModelCandidate( + device, candidates, num_candidates, session_options, selected_index); + } + static void ORT_API_CALL ReleaseEp(OrtEpFactory* this_ptr, OrtEp* ep) noexcept { static_cast(this_ptr)->ReleaseEp(ep); } diff --git a/onnxruntime/test/autoep/library/example_plugin_ep/ep_factory.cc b/onnxruntime/test/autoep/library/example_plugin_ep/ep_factory.cc index a323ec0e8c15e..2fefaeffb5d34 100644 --- a/onnxruntime/test/autoep/library/example_plugin_ep/ep_factory.cc +++ b/onnxruntime/test/autoep/library/example_plugin_ep/ep_factory.cc @@ -4,6 +4,7 @@ #include "ep_factory.h" #include +#include #include "ep.h" #include "ep_allocator.h" @@ -14,6 +15,22 @@ #include "core/session/onnxruntime_ep_device_ep_metadata_keys.h" #include "core/session/onnxruntime_session_options_config_keys.h" +namespace { +int CompatibilityRank(OrtCompiledModelCompatibility c) { + switch (c) { + case OrtCompiledModelCompatibility_EP_SUPPORTED_OPTIMAL: + return 3; + case OrtCompiledModelCompatibility_EP_SUPPORTED_PREFER_RECOMPILATION: + return 2; + case OrtCompiledModelCompatibility_EP_NOT_APPLICABLE: + return 1; + case OrtCompiledModelCompatibility_EP_UNSUPPORTED: + default: + return 0; + } +} +} // namespace + ExampleEpFactory::ExampleEpFactory(const char* ep_name, ApiPtrs apis, const OrtLogger& default_logger) : OrtEpFactory{}, ApiPtrs(apis), @@ -47,6 +64,7 @@ ExampleEpFactory::ExampleEpFactory(const char* ep_name, ApiPtrs apis, const OrtL GetNumCustomOpDomains = GetNumCustomOpDomainsImpl; GetCustomOpDomains = GetCustomOpDomainsImpl; ValidateCompiledModelCompatibilityInfo = ValidateCompiledModelCompatibilityInfoImpl; + SelectBestModelCandidate = SelectBestModelCandidateImpl; // setup the OrtMemoryInfo instances required by the EP. // We pretend the device the EP is running on is GPU. @@ -537,3 +555,59 @@ OrtStatus* ORT_API_CALL ExampleEpFactory::ValidateCompiledModelCompatibilityInfo *model_compatibility = OrtCompiledModelCompatibility_EP_SUPPORTED_OPTIMAL; return nullptr; } + +OrtStatus* ORT_API_CALL ExampleEpFactory::SelectBestModelCandidateImpl( + OrtEpFactory* this_ptr, + const OrtHardwareDevice* device, + const OrtKeyValuePairs* const* candidates, + size_t num_candidates, + const OrtSessionOptions* /*session_options*/, + size_t* selected_index) noexcept { + auto& factory = *static_cast(this_ptr); + + if (selected_index == nullptr) { + return factory.ort_api.CreateStatus(ORT_INVALID_ARGUMENT, "selected_index cannot be nullptr"); + } + + *selected_index = std::numeric_limits::max(); + + if (candidates == nullptr || num_candidates == 0) { + return factory.ort_api.CreateStatus(ORT_INVALID_ARGUMENT, "candidates cannot be nullptr or empty"); + } + + const OrtHardwareDevice* devices[] = {device}; + + int best_rank = -1; + size_t best_idx = std::numeric_limits::max(); + + for (size_t i = 0; i < num_candidates; ++i) { + const char* compatibility_info = + factory.ort_api.GetKeyValue(candidates[i], "ep_compatibility_info"); + + if (compatibility_info == nullptr) { + return factory.ort_api.CreateStatus(ORT_INVALID_ARGUMENT, + "candidate metadata is missing required key: ep_compatibility_info"); + } + + OrtCompiledModelCompatibility compatibility = OrtCompiledModelCompatibility_EP_UNSUPPORTED; + OrtStatus* status = ValidateCompiledModelCompatibilityInfoImpl( + this_ptr, devices, 1, compatibility_info, &compatibility); + if (status != nullptr) { + return status; + } + + const int rank = CompatibilityRank(compatibility); + if (rank > best_rank) { + best_rank = rank; + best_idx = i; + } + } + + if (best_rank <= CompatibilityRank(OrtCompiledModelCompatibility_EP_UNSUPPORTED)) { + *selected_index = std::numeric_limits::max(); + } else { + *selected_index = best_idx; + } + + return nullptr; +} diff --git a/onnxruntime/test/autoep/library/example_plugin_ep/ep_factory.h b/onnxruntime/test/autoep/library/example_plugin_ep/ep_factory.h index b8b4a1a434e58..74c64b069fe70 100644 --- a/onnxruntime/test/autoep/library/example_plugin_ep/ep_factory.h +++ b/onnxruntime/test/autoep/library/example_plugin_ep/ep_factory.h @@ -110,6 +110,14 @@ class ExampleEpFactory : public OrtEpFactory, public ApiPtrs { const char* compatibility_info, OrtCompiledModelCompatibility* model_compatibility) noexcept; + static OrtStatus* ORT_API_CALL SelectBestModelCandidateImpl( + OrtEpFactory* this_ptr, + const OrtHardwareDevice* device, + const OrtKeyValuePairs* const* candidates, + size_t num_candidates, + const OrtSessionOptions* session_options, + size_t* selected_index) noexcept; + const std::string ep_name_; // EP name const std::string vendor_{"Contoso"}; // EP vendor name const uint32_t vendor_id_{0xB357}; // EP vendor ID diff --git a/onnxruntime/test/autoep/test_ep_compatibility.cc b/onnxruntime/test/autoep/test_ep_compatibility.cc new file mode 100644 index 0000000000000..437ebc70ef08f --- /dev/null +++ b/onnxruntime/test/autoep/test_ep_compatibility.cc @@ -0,0 +1,67 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include +#include + +#include "gtest/gtest.h" + +#include "core/session/onnxruntime_c_api.h" +#include "core/session/onnxruntime_cxx_api.h" +#include "core/session/abi_devices.h" +#include "test/autoep/test_autoep_utils.h" + +using namespace onnxruntime; + +TEST(EpCompatibilitySelectBestTest, SelectBestModelCandidate_UsesHardwareDevices) { + Ort::Env env{ORT_LOGGING_LEVEL_WARNING, "EpCompatSelectBestExampleEp"}; + + onnxruntime::test::RegisteredEpDeviceUniquePtr example_ep; + onnxruntime::test::Utils::RegisterAndGetExampleEp( + env, onnxruntime::test::Utils::example_ep_info, example_ep); + + ASSERT_NE(example_ep.get(), nullptr); + + const OrtEpDevice* ep_device = example_ep.get(); + OrtEpFactory* factory = ep_device->GetMutableFactory(); + ASSERT_NE(factory, nullptr); + ASSERT_NE(factory->SelectBestModelCandidate, nullptr); + + const std::string ep_name = factory->GetName(factory); + + const std::string optimal = + ep_name + ";version=0.1.0;ort_api_version=" + std::to_string(ORT_API_VERSION) + ";hardware_architecture=arch1"; + + const std::string prefer_recompile = + ep_name + ";version=9.9.9;ort_api_version=" + std::to_string(ORT_API_VERSION) + ";hardware_architecture=arch1"; + + const std::string unsupported = + "SomeOtherEp;version=0.1.0;ort_api_version=" + std::to_string(ORT_API_VERSION) + ";hardware_architecture=arch1"; + + const OrtApi* api = OrtGetApiBase()->GetApi(ORT_API_VERSION); + ASSERT_NE(api, nullptr); + + OrtKeyValuePairs* kvp0 = nullptr; + OrtKeyValuePairs* kvp1 = nullptr; + OrtKeyValuePairs* kvp2 = nullptr; + api->CreateKeyValuePairs(&kvp0); + api->CreateKeyValuePairs(&kvp1); + api->CreateKeyValuePairs(&kvp2); + api->AddKeyValuePair(kvp0, "ep_compatibility_info", unsupported.c_str()); + api->AddKeyValuePair(kvp1, "ep_compatibility_info", prefer_recompile.c_str()); + api->AddKeyValuePair(kvp2, "ep_compatibility_info", optimal.c_str()); + + const OrtKeyValuePairs* candidates[] = {kvp0, kvp1, kvp2}; + + size_t selected_index = std::numeric_limits::max(); + OrtStatus* st = factory->SelectBestModelCandidate( + factory, ep_device->device, candidates, 3, nullptr, &selected_index); + + ASSERT_EQ(st, nullptr) << (st ? api->GetErrorMessage(st) : ""); + + EXPECT_EQ(selected_index, 2u); + + api->ReleaseKeyValuePairs(kvp0); + api->ReleaseKeyValuePairs(kvp1); + api->ReleaseKeyValuePairs(kvp2); +} \ No newline at end of file From abe7ba9231568ee878187964e78e90e0948b3c78 Mon Sep 17 00:00:00 2001 From: Edward Chen <18449977+edgchen1@users.noreply.github.com> Date: Fri, 3 Jul 2026 08:38:51 -0700 Subject: [PATCH 17/17] Shard Windows ASan test runs (#29503) ### Description Use GoogleTest sharding (`GTEST_TOTAL_SHARDS` and `GTEST_SHARD_INDEX`) to split up unit tests into multiple runs in Windows ASan CI build. ### Motivation and Context AddressSanitizer runs out of memory. Splitting the tests into multiple runs seems to mitigate it. --------- Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --- .github/workflows/windows_build_x64_asan.yml | 78 ++++++++++++++++++-- 1 file changed, 72 insertions(+), 6 deletions(-) diff --git a/.github/workflows/windows_build_x64_asan.yml b/.github/workflows/windows_build_x64_asan.yml index 116938b5129db..1be3bd84fcd73 100644 --- a/.github/workflows/windows_build_x64_asan.yml +++ b/.github/workflows/windows_build_x64_asan.yml @@ -21,6 +21,11 @@ jobs: ] timeout-minutes: 300 + env: + # Number of GoogleTest shards to split the test run across. + # Try increasing this if AddressSanitizer runs out of memory. + SHARD_COUNT: 4 + steps: - name: Checkout repository uses: actions/checkout@v6 @@ -38,10 +43,71 @@ jobs: with: architecture: x64 - - name: Build and Test (Combined) - shell: cmd + - name: Write shared build options + shell: pwsh + run: | + $ErrorActionPreference = 'Stop' + + $buildDir = Join-Path $env:RUNNER_TEMP 'build' + $optsFile = Join-Path $env:RUNNER_TEMP 'asan_build_options.txt' + + # Common build.py options shared by the build and test phases below. + @( + '--config Debug' + "--build_dir `"$buildDir`"" + '--skip_submodule_sync' + '--parallel' + '--use_vcpkg' + '--use_vcpkg_ms_internal_asset_cache' + '--cmake_generator "Visual Studio 17 2022"' + '--disable_memleak_checker' + '--enable_address_sanitizer' + ) | Set-Content -Path $optsFile -Encoding ascii + + Write-Host "Wrote build options to ${optsFile}:" + + Get-Content $optsFile | Write-Host + + - name: Build + shell: pwsh + run: | + $ErrorActionPreference = 'Stop' + + $buildPy = Join-Path $env:GITHUB_WORKSPACE 'tools\ci_build\build.py' + $optsFile = Join-Path $env:RUNNER_TEMP 'asan_build_options.txt' + $reqFile = Join-Path $env:GITHUB_WORKSPACE 'tools\ci_build\github\windows\python\requirements.txt' + + python -m pip install -r $reqFile + if ($LASTEXITCODE -ne 0) { throw "pip install failed ($LASTEXITCODE)" } + + # '@' must be quoted so PowerShell does not treat it as the splat operator. + python $buildPy "@$optsFile" --update --build + if ($LASTEXITCODE -ne 0) { throw "Build failed ($LASTEXITCODE)" } + + - name: Test + shell: pwsh run: | - @echo off - echo %PATH% - python -m pip install -r "%GITHUB_WORKSPACE%\tools\ci_build/github/windows\python\requirements.txt" - python "%GITHUB_WORKSPACE%\tools\ci_build\build.py" --config Debug --build_dir "%RUNNER_TEMP%\build" --skip_submodule_sync --parallel --test_parallel 4 --use_vcpkg --use_vcpkg_ms_internal_asset_cache --cmake_generator "Visual Studio 17 2022" --disable_memleak_checker --enable_address_sanitizer + $ErrorActionPreference = 'Stop' + + $buildPy = Join-Path $env:GITHUB_WORKSPACE 'tools\ci_build\build.py' + $optsFile = Join-Path $env:RUNNER_TEMP 'asan_build_options.txt' + $shardCount = [int]$env:SHARD_COUNT + if ($shardCount -lt 1) { throw "SHARD_COUNT must be >= 1 (got $shardCount)" } + + # GoogleTest reads the GTEST_TOTAL_SHARDS and GTEST_SHARD_INDEX env vars. + # Each ctest-launched test binary runs only its shard. + $env:GTEST_TOTAL_SHARDS = $shardCount + + for ($i = 0; $i -lt $shardCount; $i++) { + Write-Host "=== Running test shard $($i + 1) of $shardCount ===" + + $env:GTEST_SHARD_INDEX = $i + + # NOTE: The ctest commands produce a test output file with `--gtest_output=xml:...` that will be + # overwritten by subsequent shard runs. This is fine for now since nothing consumes the output files. + # If the output files are required, they can be renamed after each shard run so they don't get overwritten. + + # '@' must be quoted so PowerShell does not treat it as the splat operator. + python $buildPy "@$optsFile" --test --test_parallel 4 + if ($LASTEXITCODE -ne 0) { throw "Test shard $($i + 1) of $shardCount failed ($LASTEXITCODE)" } + }