From 4bc9c4f8709ea7c96b46f9fb5e6c6a4024930371 Mon Sep 17 00:00:00 2001 From: Albert Date: Tue, 15 Sep 2026 14:53:13 +0800 Subject: [PATCH 01/52] Albert/4.3.5 toolkit build --- README.md | 3 +- musa/docs/architecture.md | 10 +- musa/docs/developer_guide.md | 7 +- musa/ep/CMakeLists.txt | 24 ++- musa/ep/onnxruntime_musa_ep_manifest.json.in | 6 +- musa/ep/python/build_wheel.py | 1 - musa/ep/src/ep_allocator.h | 8 +- ...hta_scaled_dot_product_attention_fusion.cc | 153 +++++++++++------- musa/ep/src/kernels/activation/erf.cc | 47 +----- musa/ep/src/kernels/llm/mhta_sdpa_fp32_impl.h | 8 + .../ep/src/kernels/llm/mhta_sdpa_fp32_impl.mu | 50 ++++++ musa/ep/src/kernels/math/matmul.cc | 2 +- musa/ep/src/kernels/math/pow.cc | 8 + musa/ep/src/kernels/shared_inc/blas_utils.h | 34 +++- musa/ep/src/mudnn_compat.h | 9 ++ musa/ep/src/musa_arena.cc | 23 ++- musa/ep/src/musa_arena.h | 6 +- test/multi_stream/README.md | 34 ---- .../test_allocator_cache_shared_session.py | 26 +-- test/ops/test_erf.py | 51 +++++- 20 files changed, 345 insertions(+), 165 deletions(-) create mode 100644 musa/ep/src/mudnn_compat.h delete mode 100644 test/multi_stream/README.md diff --git a/README.md b/README.md index 0760865..419d86e 100644 --- a/README.md +++ b/README.md @@ -34,7 +34,7 @@ Build-time (only two): | Component | Version | Notes | |---|---|---| -| **MUSA toolkit** | **5.1.0** | Defaults to `/usr/local/musa`. Override with `-DMUSA_HOME=...` or `./build.sh -- -DMUSA_HOME=/opt/musa`. Links the MUSA runtime, muBLAS, and muDNN libraries. | +| **MUSA toolkit** | **5.1.0 / 4.3.8** | Both versions are supported. Defaults to `/usr/local/musa`. Override with `-DMUSA_HOME=...` or `./build.sh -- -DMUSA_HOME=/opt/musa`. Links the MUSA runtime, muBLAS, and muDNN libraries. | | **C++ compiler** | C++20 | GCC 11+ / Clang 14+. Required for `std::span`. | ONNX Runtime is a Git submodule pinned to tag **v1.26.0** (commit `8c546c37`). The build uses @@ -329,6 +329,7 @@ python -m pytest test/ops/ │ └─▶ kernels/{math, activation, tensor, logical, │ │ reduction, nn} │ │ └─▶ MUSA runtime libraries on MUSA 5.1.0 │ +│ or 4.3.8 │ └──────────────────────────────────────────────────────────────────────┘ ``` diff --git a/musa/docs/architecture.md b/musa/docs/architecture.md index 642258c..d9828c1 100644 --- a/musa/docs/architecture.md +++ b/musa/docs/architecture.md @@ -80,7 +80,7 @@ Python or C++ ORT host: reduction, nn} | v - musart / mublas (MUSA 5.1.0) + musart / mublas (MUSA 5.1.0 / 4.3.8) ``` ## Stream-ordered device arena @@ -98,6 +98,14 @@ logical `Free` it cannot prove that a wait covers the chunk's final device access; therefore an owned chunk is never reused directly by another stream. Every chunk belongs to at most one stream ownership set. +ORT may allocate a memory-pattern buffer on its `DummyStream`, whose native +handle is null and which has no plugin implementation or MUSA release callback. +Both allocator paths check the native handle before calling +`SyncStream_GetImpl`. Such allocations, and allocations with no stream, use +direct `Reserve`/`Free` instead of entering the stream cache. They still count +against the arena budget. Interpreting the dummy as a plugin stream reads an +invalid implementation pointer and corrupts ownership bookkeeping. + `OnSessionRunEnd` does not clear ownership because host-side Run completion is not a device-completion fence. When ORT releases the `MusaSyncStream` implementation, its destructor first synchronizes that one MUSA stream, then diff --git a/musa/docs/developer_guide.md b/musa/docs/developer_guide.md index 414757a..ba92007 100644 --- a/musa/docs/developer_guide.md +++ b/musa/docs/developer_guide.md @@ -216,9 +216,10 @@ ORT_MUSA_DUMP_RUNTIME_GRAPH_MERMAID_PATH=/tmp/musa_runtime_execution_graph.mmd \ stream ownership 隔离 free chunk：有 owner 的 chunk 仅由同一 stream 复用；释放 sync-stream 实现时会先执行该 stream 的局部完成同步，再解除 ownership 并允许其他 stream 复用。不会执行 device-wide 同步；该 release - 路径会等待该 stream 上已提交的工作完成。没有 `OrtSyncStream` 的分配不进入 - 可复用 arena，而是 direct allocation/free，避免在共享 Session 多流 Run 中把 - unowned block 交给另一条 stream。 + 路径会等待该 stream 上已提交的工作完成。没有 `OrtSyncStream`，或 stream + 没有原生句柄（例如 ORT 为 memory pattern 使用的 `DummyStream`）的分配， + 使用 direct allocation/free 并计入预算，不进入可复用 arena。这类 stream + 没有 MUSA implementation 或释放回调，不能调用 `SyncStream_GetImpl()`。 - 默认值：未设置或设置为 `0` 时不启用 arena，继续使用原始 `CustomAllocator` 分配/释放路径。 - 解析规则：必须是完整的非负十进制 MB 整数；尾部字符、负数和换算溢出均会 diff --git a/musa/ep/CMakeLists.txt b/musa/ep/CMakeLists.txt index 7a1832d..53fa220 100644 --- a/musa/ep/CMakeLists.txt +++ b/musa/ep/CMakeLists.txt @@ -56,7 +56,7 @@ target_include_directories(${MUSA_PLUGIN_TARGET} ) set(MUSA_LIBRARY_DIRS "${MUSA_HOME}/lib" "${MUSA_HOME}/lib64") -foreach(_musa_lib IN ITEMS musart mublas mudnncxx mudnn) +foreach(_musa_lib IN ITEMS musart mublas mudnn) string(TOUPPER "${_musa_lib}" _musa_lib_upper) find_library(MUSA_${_musa_lib_upper}_LIBRARY NAMES ${_musa_lib} @@ -71,6 +71,28 @@ foreach(_musa_lib IN ITEMS musart mublas mudnncxx mudnn) list(APPEND MUSA_PLUGIN_LINK_LIBRARIES ${MUSA_${_musa_lib_upper}_LIBRARY}) endforeach() +# mudnncxx is a separate shared library in MUSA 5.1, but its C++ API is +# folded into mudnn in MUSA 4.3.8. Link it only when it is present. +find_library(MUSA_MUDNNCXX_LIBRARY + NAMES mudnncxx + PATHS ${MUSA_LIBRARY_DIRS} + NO_DEFAULT_PATH +) +if(MUSA_MUDNNCXX_LIBRARY) + list(APPEND MUSA_PLUGIN_LINK_LIBRARIES ${MUSA_MUDNNCXX_LIBRARY}) + set(MUSA_MUDNNCXX_RUNTIME "libmudnncxx.so.3") +else() + set(MUSA_MUDNNCXX_RUNTIME "") +endif() +set(MUSA_EXTERNAL_RUNTIME_LIBRARIES + "\"libmusa.so.1\", \"libmusart.so.5\", \"libmublas.so.1\"," +) +if(MUSA_MUDNNCXX_RUNTIME) + string(APPEND MUSA_EXTERNAL_RUNTIME_LIBRARIES + " \"${MUSA_MUDNNCXX_RUNTIME}\",") +endif() +string(APPEND MUSA_EXTERNAL_RUNTIME_LIBRARIES " \"libmudnn.so.3\"") + target_link_libraries(${MUSA_PLUGIN_TARGET} ${MUSA_PLUGIN_LINK_LIBRARIES}) target_compile_definitions(${MUSA_PLUGIN_TARGET} PRIVATE diff --git a/musa/ep/onnxruntime_musa_ep_manifest.json.in b/musa/ep/onnxruntime_musa_ep_manifest.json.in index 90f6432..82b76eb 100644 --- a/musa/ep/onnxruntime_musa_ep_manifest.json.in +++ b/musa/ep/onnxruntime_musa_ep_manifest.json.in @@ -10,10 +10,6 @@ "ort_tested_versions": ["@ONNXRUNTIME_VERSION@"], "platform": "@CMAKE_SYSTEM_NAME@-@CMAKE_SYSTEM_PROCESSOR@", "external_runtime_libraries": [ - "libmusa.so.1", - "libmusart.so.5", - "libmublas.so.1", - "libmudnncxx.so.3", - "libmudnn.so.3" + @MUSA_EXTERNAL_RUNTIME_LIBRARIES@ ] } diff --git a/musa/ep/python/build_wheel.py b/musa/ep/python/build_wheel.py index b963068..d781a15 100644 --- a/musa/ep/python/build_wheel.py +++ b/musa/ep/python/build_wheel.py @@ -47,7 +47,6 @@ def _read_ort_version() -> str: "libmusa.so.1", "libmusart.so.5", "libmublas.so.1", - "libmudnncxx.so.3", "libmudnn.so.3", ] diff --git a/musa/ep/src/ep_allocator.h b/musa/ep/src/ep_allocator.h index c0e061e..15276d1 100644 --- a/musa/ep/src/ep_allocator.h +++ b/musa/ep/src/ep_allocator.h @@ -122,9 +122,13 @@ struct CustomAllocator : BaseAllocator { return nullptr; } auto& impl = *static_cast(this_); + // Memory-pattern allocations may use ORT's DummyStream. Only real device + // streams have an implementation; the dummy's native handle is null. const OrtSyncStreamImpl* stream_impl = - stream != nullptr ? Ort::GetEpApi().SyncStream_GetImpl(stream) - : nullptr; + stream != nullptr && + Ort::GetApi().SyncStream_GetHandle(stream) != nullptr + ? Ort::GetEpApi().SyncStream_GetImpl(stream) + : nullptr; return impl.AllocateCached(size, stream_impl); } diff --git a/musa/ep/src/fusion/mhta_scaled_dot_product_attention_fusion.cc b/musa/ep/src/fusion/mhta_scaled_dot_product_attention_fusion.cc index cbb9d5a..81dce87 100644 --- a/musa/ep/src/fusion/mhta_scaled_dot_product_attention_fusion.cc +++ b/musa/ep/src/fusion/mhta_scaled_dot_product_attention_fusion.cc @@ -13,7 +13,6 @@ #include "fusion/mhta_scaled_dot_product_attention_fusion.h" -#include #include #include @@ -33,6 +32,7 @@ #include "kernels/llm/mhta_sdpa_fp32_impl.h" #include "kernels/shared_inc/blas_utils.h" #include "kernels/shared_inc/op_kernel_common.h" +#include "mudnn_compat.h" namespace { @@ -457,22 +457,42 @@ void CheckStatus(::musa::dnn::Status status, const char* message) { } } -void SetupTensor(::musa::dnn::Tensor& tensor, const void* data, - const std::vector& shape, - ONNXTensorElementDataType elem_type, const char* name) { - if (!SetMudnnTensor(tensor, data, shape, elem_type)) { +// SDPA accepts a BOOL keep-mask in muDNN 4.3.8 even though BOOL is excluded +// from the elementwise dispatch whitelist in SetMudnnTensor(). Keep this +// descriptor path separate so adding SDPA mask support cannot re-enable the +// unsupported elementwise dtype combinations guarded by that whitelist. +bool SetMudnnSdpaTensor(::musa::dnn::Tensor& tensor, const void* data, + const std::vector& shape, + ONNXTensorElementDataType elem_type) { + ::musa::dnn::Tensor::Type mudnn_type; + if (!MudnnTensorType(elem_type, mudnn_type) || + tensor.SetAddr(data) != ::musa::dnn::Status::SUCCESS || + tensor.SetType(mudnn_type) != ::musa::dnn::Status::SUCCESS || + tensor.SetFormat(::musa::dnn::Tensor::Format::NCHW) != + ::musa::dnn::Status::SUCCESS) { + return false; + } + std::vector dims = shape.empty() ? std::vector{1} : shape; + return tensor.SetNdInfo(static_cast(dims.size()), dims.data()) == + ::musa::dnn::Status::SUCCESS; +} + +void SetupSdpaTensor(::musa::dnn::Tensor& tensor, const void* data, + const std::vector& shape, + ONNXTensorElementDataType elem_type, const char* name) { + if (!SetMudnnSdpaTensor(tensor, data, shape, elem_type)) { throw std::runtime_error(std::string("failed to set MHTA SDPA tensor ") + name); } } -void SetupTensorWithStrides(::musa::dnn::Tensor& tensor, const void* data, - const std::vector& shape, - const std::vector& strides, - ONNXTensorElementDataType elem_type, - const char* name) { +void SetupSdpaTensorWithStrides(::musa::dnn::Tensor& tensor, const void* data, + const std::vector& shape, + const std::vector& strides, + ONNXTensorElementDataType elem_type, + const char* name) { if (shape.size() != strides.size() || shape.empty() || - !SetMudnnTensor(tensor, data, shape, elem_type) || + !SetMudnnSdpaTensor(tensor, data, shape, elem_type) || tensor.SetNdInfo(static_cast(shape.size()), shape.data(), strides.data()) != ::musa::dnn::Status::SUCCESS) { throw std::runtime_error( @@ -850,18 +870,18 @@ OrtStatus* MhtaScaledDotProductAttentionFusionCompute::Compute( ::musa::dnn::Handle* handle = nullptr; RETURN_IF_ERROR(EnsureMudnnHandle(&handle, stream)); ::musa::dnn::Tensor input_tensor, output_tensor; - SetupTensor(input_tensor, q_data, q_shape, elem_type, "Q"); - SetupTensor(output_tensor, q_transposed.data(), sdpa_q_shape, elem_type, - "Q transpose"); + SetupSdpaTensor(input_tensor, q_data, q_shape, elem_type, "Q"); + SetupSdpaTensor(output_tensor, q_transposed.data(), sdpa_q_shape, + elem_type, "Q transpose"); const int64_t perm[] = {0, 2, 1, 3}; ::musa::dnn::Permute permute; CheckStatus(permute.ConfigDimStride(output_tensor, input_tensor, 4, perm), "failed to configure MHTA SDPA Q permute"); CheckStatus(permute.Run(*handle, output_tensor, input_tensor), "MHTA SDPA Q permute failed"); - SetupTensor(input_tensor, k_data, k_shape, elem_type, "K"); - SetupTensor(output_tensor, k_transposed.data(), sdpa_k_shape, elem_type, - "K transpose"); + SetupSdpaTensor(input_tensor, k_data, k_shape, elem_type, "K"); + SetupSdpaTensor(output_tensor, k_transposed.data(), sdpa_k_shape, + elem_type, "K transpose"); CheckStatus(permute.ConfigDimStride(output_tensor, input_tensor, 4, perm), "failed to configure MHTA SDPA K permute"); CheckStatus(permute.Run(*handle, output_tensor, input_tensor), @@ -887,9 +907,9 @@ OrtStatus* MhtaScaledDotProductAttentionFusionCompute::Compute( ::musa::dnn::Handle* handle = nullptr; RETURN_IF_ERROR(EnsureMudnnHandle(&handle, stream)); ::musa::dnn::Tensor input_tensor, output_tensor; - SetupTensor(input_tensor, k_data, k_shape, elem_type, "K"); - SetupTensor(output_tensor, k_transposed.data(), sdpa_k_shape, elem_type, - "K transpose"); + SetupSdpaTensor(input_tensor, k_data, k_shape, elem_type, "K"); + SetupSdpaTensor(output_tensor, k_transposed.data(), sdpa_k_shape, + elem_type, "K transpose"); const int64_t perm[] = {0, 1, 3, 2}; ::musa::dnn::Permute permute; CheckStatus(permute.ConfigDimStride(output_tensor, input_tensor, 4, perm), @@ -968,6 +988,11 @@ OrtStatus* MhtaScaledDotProductAttentionFusionCompute::Compute( kInt32MaskCacheReuseLimit); mask_data = cache.buffer.data(); } + // The raw INT32 input may have a key-prefix dimension larger than Sk + // (the absorbed Slice trims it). The materialized BOOL tensor is + // always contiguous [B,H,Sq,Sk], so update the params used by the + // post-RunFlash fully-masked-row repair to its actual key extent. + mask_params.mask_k = seqlen_k; } } else if (mask_scale_ != 1.0f) { const size_t mask_bytes = static_cast(NumElements(mask_shape)) * @@ -982,10 +1007,11 @@ OrtStatus* MhtaScaledDotProductAttentionFusionCompute::Compute( if (copy_status != musaSuccess) throw std::runtime_error(MusaErrorString(copy_status)); ::musa::dnn::Tensor mask_tensor, scalar_tensor, scaled_tensor; - SetupTensor(mask_tensor, mask_data, mask_shape, elem_type, "mask"); - SetupTensor(scalar_tensor, scalar.data(), {1}, elem_type, "mask scale"); - SetupTensor(scaled_tensor, mask_materialized.data(), mask_shape, - elem_type, "scaled mask"); + SetupSdpaTensor(mask_tensor, mask_data, mask_shape, elem_type, "mask"); + SetupSdpaTensor(scalar_tensor, scalar.data(), {1}, elem_type, + "mask scale"); + SetupSdpaTensor(scaled_tensor, mask_materialized.data(), mask_shape, + elem_type, "scaled mask"); ::musa::dnn::Binary mul; CheckStatus(mul.SetMode(::musa::dnn::Binary::Mode::MUL), "failed to set MHTA SDPA mask scale mode"); @@ -1001,53 +1027,55 @@ OrtStatus* MhtaScaledDotProductAttentionFusionCompute::Compute( head_dim, heads * head_dim, 1}; const std::vector kv_strides = {seqlen_k * heads * head_dim, head_dim, heads * head_dim, 1}; - SetupTensorWithStrides(q_tensor, q_data, sdpa_q_shape, q_strides, - elem_type, "Q"); - SetupTensorWithStrides(k_tensor, k_data, sdpa_k_shape, kv_strides, - elem_type, "K"); - SetupTensorWithStrides(v_tensor, v_buffer.data(), sdpa_v_shape, - kv_strides, elem_type, "V"); + SetupSdpaTensorWithStrides(q_tensor, q_data, sdpa_q_shape, q_strides, + elem_type, "Q"); + SetupSdpaTensorWithStrides(k_tensor, k_data, sdpa_k_shape, kv_strides, + elem_type, "K"); + SetupSdpaTensorWithStrides(v_tensor, v_buffer.data(), sdpa_v_shape, + kv_strides, elem_type, "V"); } else if (bshd_strided_flash) { const std::vector bshd_strides = {seqlen_q * heads * head_dim, head_dim, heads * head_dim, 1}; - SetupTensorWithStrides(q_tensor, q_data, sdpa_q_shape, bshd_strides, - elem_type, "Q"); - SetupTensorWithStrides(k_tensor, k_data, sdpa_k_shape, bshd_strides, - elem_type, "K"); - SetupTensorWithStrides(v_tensor, v_buffer.data(), sdpa_v_shape, - bshd_strides, elem_type, "V"); + SetupSdpaTensorWithStrides(q_tensor, q_data, sdpa_q_shape, bshd_strides, + elem_type, "Q"); + SetupSdpaTensorWithStrides(k_tensor, k_data, sdpa_k_shape, bshd_strides, + elem_type, "K"); + SetupSdpaTensorWithStrides(v_tensor, v_buffer.data(), sdpa_v_shape, + bshd_strides, elem_type, "V"); } else { - SetupTensor(q_tensor, q_data, sdpa_q_shape, elem_type, "Q"); - SetupTensor(k_tensor, k_data, sdpa_k_shape, elem_type, "K"); - SetupTensor(v_tensor, v_buffer.data(), sdpa_v_shape, elem_type, "V"); + SetupSdpaTensor(q_tensor, q_data, sdpa_q_shape, elem_type, "Q"); + SetupSdpaTensor(k_tensor, k_data, sdpa_k_shape, elem_type, "K"); + SetupSdpaTensor(v_tensor, v_buffer.data(), sdpa_v_shape, elem_type, "V"); } if (has_external_mask || lseq_last_key_mask_ || lseq_shd_strided_flash) { - SetupTensor(mask_tensor, mask_data, active_mask_shape, active_mask_type, - "mask"); + SetupSdpaTensor(mask_tensor, mask_data, active_mask_shape, + active_mask_type, "mask"); } if (lseq_shd_strided_flash) { const std::vector output_strides = { seqlen_q * heads * head_dim, head_dim, heads * head_dim, 1}; - SetupTensorWithStrides(out_tensor, output.GetTensorMutableData(), - sdpa_q_shape, output_strides, elem_type, "output"); + SetupSdpaTensorWithStrides( + out_tensor, output.GetTensorMutableData(), sdpa_q_shape, + output_strides, elem_type, "output"); } else if (bshd_strided_flash) { const std::vector bshd_strides = {seqlen_q * heads * head_dim, head_dim, heads * head_dim, 1}; - SetupTensorWithStrides(out_tensor, output.GetTensorMutableData(), - sdpa_q_shape, bshd_strides, elem_type, "output"); + SetupSdpaTensorWithStrides( + out_tensor, output.GetTensorMutableData(), sdpa_q_shape, + bshd_strides, elem_type, "output"); } else { - SetupTensor(out_tensor, output.GetTensorMutableData(), - (sim_rank3 || folded_rank3) ? sdpa_q_shape : output_shape, - elem_type, "output"); + SetupSdpaTensor(out_tensor, output.GetTensorMutableData(), + (sim_rank3 || folded_rank3) ? sdpa_q_shape : output_shape, + elem_type, "output"); } DeviceBuffer& lse = scratch.lse; lse.Resize(static_cast(batch * heads * seqlen_q) * sizeof(float), stream); - SetupTensor(lse_tensor, lse.data(), {batch, heads, seqlen_q}, - ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT, "logsumexp"); + SetupSdpaTensor(lse_tensor, lse.data(), {batch, heads, seqlen_q}, + ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT, "logsumexp"); // Inference has dropout disabled. Passing a null tensor avoids Flash's // output-mask shape contract [B,H,Sq,Sk] and its needless allocation. - SetupTensor(dropout_tensor, nullptr, {1}, elem_type, "dropout mask"); + SetupSdpaTensor(dropout_tensor, nullptr, {1}, elem_type, "dropout mask"); ::musa::dnn::ScaledDotProductAttention sdpa; CheckStatus( sdpa.SetComputeMode( @@ -1077,8 +1105,23 @@ OrtStatus* MhtaScaledDotProductAttentionFusionCompute::Compute( elem_type == ONNX_TENSOR_ELEMENT_DATA_TYPE_BFLOAT16) { CheckStatus( sdpa.RunFlash(*handle, out_tensor, lse_tensor, q_tensor, k_tensor, - v_tensor, mask_tensor, dropout_tensor), + v_tensor, mask_tensor, dropout_tensor, nullptr), "MHTA SDPA RunFlash failed"); + // muDNN 4.3.8 returns NaN when every key in a query row is masked. The + // ONNX keep-mask semantics used by this fusion require a zero output for + // that row. The repair runs on the same provider stream and therefore + // remains ordered after RunFlash without a host synchronization. + if (boolean_mask_ && !lseq_shd_strided_flash && !bshd_strided_flash && + !sim_rank3 && !folded_rank3) { + const musaError_t repair_status = LaunchMusaMhtaSdpaZeroFullyMaskedRows( + static_cast(output.GetTensorMutableData()), + static_cast(mask_data), mask_params, stream); + if (repair_status != musaSuccess) { + throw std::runtime_error( + std::string("MHTA SDPA fully-masked row repair failed: ") + + MusaErrorString(repair_status)); + } + } } else { DeviceBuffer& attn_probs = scratch.attn_probs; attn_probs.Resize( @@ -1086,9 +1129,9 @@ OrtStatus* MhtaScaledDotProductAttentionFusionCompute::Compute( MhtaSdpaElementSize(elem_type), stream); ::musa::dnn::Tensor probs_tensor; - SetupTensor(probs_tensor, attn_probs.data(), - {batch, heads, seqlen_q, seqlen_k}, elem_type, - "attention probabilities"); + SetupSdpaTensor(probs_tensor, attn_probs.data(), + {batch, heads, seqlen_q, seqlen_k}, elem_type, + "attention probabilities"); auto allocator = [stream](size_t bytes) -> ::musa::dnn::MemoryHandler { void* workspace = AllocateDeviceMemoryOnStream(bytes, stream); if (bytes != 0 && workspace == nullptr) { diff --git a/musa/ep/src/kernels/activation/erf.cc b/musa/ep/src/kernels/activation/erf.cc index 8c2d944..4b3f1e6 100644 --- a/musa/ep/src/kernels/activation/erf.cc +++ b/musa/ep/src/kernels/activation/erf.cc @@ -11,50 +11,9 @@ // See the License for the specific language governing permissions and // limitations under the License. -#include "shared_inc/blas_utils.h" #include "shared_inc/op_kernel_common.h" namespace { -constexpr size_t kMudnnMaxElementwiseRank = 5; - -bool TryMudnnErf(Ort::KernelContext& ctx, const std::vector& shape, - ONNXTensorElementDataType elem_type) { - Ort::ConstValue input = ctx.GetInput(0); - if (shape.size() > kMudnnMaxElementwiseRank || - !IsGpuMemory(input.GetTensorMemoryInfo())) { - return false; - } - - Ort::UnownedValue y = ctx.GetOutput(0, shape); - if (!IsGpuMemory(y.GetTensorMemoryInfo())) { - return false; - } - - ::musa::dnn::Handle* handle = nullptr; - OrtStatus* handle_status = EnsureMudnnHandle(&handle, GetComputeStream(ctx)); - if (handle_status != nullptr) { - Ort::GetApi().ReleaseStatus(handle_status); - return false; - } - - ::musa::dnn::Tensor input_tensor; - ::musa::dnn::Tensor output_tensor; - if (!SetMudnnTensor(input_tensor, input.GetTensorRawData(), shape, - elem_type) || - !SetMudnnTensor(output_tensor, y.GetTensorMutableRawData(), shape, - elem_type)) { - return false; - } - - ::musa::dnn::Unary op; - if (op.SetMode(::musa::dnn::Unary::Mode::ERF) != - ::musa::dnn::Status::SUCCESS) { - return false; - } - return op.Run(*handle, output_tensor, input_tensor) == - ::musa::dnn::Status::SUCCESS; -} - class Erf : public OpKernelBase { public: Erf(const OrtKernelInfo* /*info*/, void* /*state*/) {} @@ -68,9 +27,9 @@ OrtStatus* Erf::Compute(Ort::KernelContext& ctx) const { if (OutputEmptyTensorIfNeeded(ctx, shape)) { return nullptr; } - if (TryMudnnErf(ctx, shape, elem_type)) { - return nullptr; - } + // muDNN ERF on Toolkit 4.3.8 returns SUCCESS but maps +/-Inf to NaN. + // Use the existing device erff/erf implementation (as in CUDA/musify) for + // all registered floating types to preserve infinities, NaNs and signed zero. return UnaryDeviceCompute(ctx, shape, elem_type, MusaUnaryOp::Erf, "Erf"); } } // namespace diff --git a/musa/ep/src/kernels/llm/mhta_sdpa_fp32_impl.h b/musa/ep/src/kernels/llm/mhta_sdpa_fp32_impl.h index 95ac9d5..c5c85a2 100644 --- a/musa/ep/src/kernels/llm/mhta_sdpa_fp32_impl.h +++ b/musa/ep/src/kernels/llm/mhta_sdpa_fp32_impl.h @@ -46,3 +46,11 @@ musaError_t LaunchMusaMhtaSdpaLseqLastKeyKeepMask2DKernel(bool* bool_mask, int64_t seqlen_q, int64_t seqlen_k, musaStream_t stream); + +// RunFlash returns NaN for a query row whose keep-mask contains no key on +// older muDNN releases. ONNX attention semantics require that row to produce +// zeros, so repair those rows on the provider stream after SDPA completes. +musaError_t LaunchMusaMhtaSdpaZeroFullyMaskedRows(uint16_t* output, + const uint8_t* bool_mask, + MusaMhtaSdpaFp32Params params, + musaStream_t stream); diff --git a/musa/ep/src/kernels/llm/mhta_sdpa_fp32_impl.mu b/musa/ep/src/kernels/llm/mhta_sdpa_fp32_impl.mu index 05e2c76..5d3a006 100644 --- a/musa/ep/src/kernels/llm/mhta_sdpa_fp32_impl.mu +++ b/musa/ep/src/kernels/llm/mhta_sdpa_fp32_impl.mu @@ -100,6 +100,40 @@ __global__ void MhtaLseqLastKeyKeepMask2DKernel(bool* bool_mask, } } +__global__ void MhtaZeroFullyMaskedRowsKernel( + uint16_t* output, const uint8_t* bool_mask, + MusaMhtaSdpaFp32Params params) { + const int64_t row_id = static_cast(blockIdx.x); + const int64_t rows = params.batch * params.heads * params.seqlen_q; + if (row_id >= rows) return; + + const int64_t row = row_id % params.seqlen_q; + const int64_t head_index = row_id / params.seqlen_q; + const int64_t h = head_index % params.heads; + const int64_t b = head_index / params.heads; + // One block owns one row. Let lane zero scan the (usually short) mask row; + // this avoids relying on atomics on shared memory in older MCC releases. + __shared__ int keep; + if (threadIdx.x == 0) { + keep = 0; + for (int64_t k = 0; k < params.seqlen_k; ++k) { + const int64_t mask_offset = MaskOffset(params, b, h, row, k); + if (bool_mask[mask_offset]) { + keep = 1; + break; + } + } + } + __syncthreads(); + if (keep == 0) { + const int64_t output_base = + ((b * params.heads + h) * params.seqlen_q + row) * params.head_dim; + for (int64_t d = threadIdx.x; d < params.head_dim; d += blockDim.x) { + output[output_base + d] = 0; + } + } +} + __global__ void MhtaSdpaFp32Kernel(const float* q, const float* k, const float* v, const void* mask, float* output, @@ -232,3 +266,19 @@ musaError_t LaunchMusaMhtaSdpaInt32KeepMaskToBoolKernel( } return musaGetLastError(); } + +musaError_t LaunchMusaMhtaSdpaZeroFullyMaskedRows( + uint16_t* output, const uint8_t* bool_mask, + MusaMhtaSdpaFp32Params params, musaStream_t stream) { + if (output == nullptr || bool_mask == nullptr || params.batch <= 0 || + params.heads <= 0 || params.seqlen_q <= 0 || params.seqlen_k <= 0 || + params.head_dim <= 0) { + return musaSuccess; + } + const int64_t rows = params.batch * params.heads * params.seqlen_q; + if (rows > INT32_MAX) return musaErrorNotSupported; + MhtaZeroFullyMaskedRowsKernel<<(rows), + kThreadsPerBlock, 0, stream>>>( + output, bool_mask, params); + return musaGetLastError(); +} diff --git a/musa/ep/src/kernels/math/matmul.cc b/musa/ep/src/kernels/math/matmul.cc index 2585624..cebc0be 100644 --- a/musa/ep/src/kernels/math/matmul.cc +++ b/musa/ep/src/kernels/math/matmul.cc @@ -1,6 +1,5 @@ #include "matmul.h" -#include #include #include @@ -11,6 +10,7 @@ #include #include +#include "mudnn_compat.h" #include "shared_inc/blas_utils.h" ONNX_OPERATOR_VERSIONED_KERNEL_EX( diff --git a/musa/ep/src/kernels/math/pow.cc b/musa/ep/src/kernels/math/pow.cc index 413ba56..ab67d07 100644 --- a/musa/ep/src/kernels/math/pow.cc +++ b/musa/ep/src/kernels/math/pow.cc @@ -21,6 +21,14 @@ bool TryMudnnPow(Ort::KernelContext& ctx, const std::vector& shape0, const std::vector& shape1, ONNXTensorElementDataType lhs_type, ONNXTensorElementDataType rhs_type) { + // 4.3.8 has no integer POW entry in the muDNN dispatch table. The native + // device kernel below supports these types, so bypass muDNN safely. + if (lhs_type == ONNX_TENSOR_ELEMENT_DATA_TYPE_INT32 || + lhs_type == ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64 || + rhs_type == ONNX_TENSOR_ELEMENT_DATA_TYPE_INT32 || + rhs_type == ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64) { + return false; + } std::vector out_shape = BroadcastShape(shape0, shape1); if (shape0.empty() || shape1.empty() || out_shape.size() > kMudnnMaxElementwiseRank || diff --git a/musa/ep/src/kernels/shared_inc/blas_utils.h b/musa/ep/src/kernels/shared_inc/blas_utils.h index 3b17ce5..01372cb 100644 --- a/musa/ep/src/kernels/shared_inc/blas_utils.h +++ b/musa/ep/src/kernels/shared_inc/blas_utils.h @@ -14,7 +14,6 @@ #pragma once #include -#include #include #include @@ -23,6 +22,7 @@ #include "math/gemm_post_kernels.h" #include "math/matmul.h" +#include "mudnn_compat.h" #include "shared_inc/op_kernel_common.h" // Shared GEMM implementation used by Gemm and FusedGemm. MatMul keeps its @@ -128,9 +128,33 @@ inline bool MudnnTensorType(ONNXTensorElementDataType elem_type, } } +// MUSA 4.3.8 exposes Tensor descriptors for these types, but its muDNN +// elementwise dispatch table does not contain all of them. Calling Run() for +// one of those combinations throws std::out_of_range ("_Map_base::at") +// instead of returning Status::NOT_SUPPORTED. The provider has native .mu +// kernels for these types, so make muDNN opt-in only for the common types and +// let the caller's device-kernel fallback handle the rest. +inline bool IsMudnnElementwiseTypeSupported( + ONNXTensorElementDataType elem_type) { + switch (elem_type) { + case ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT16: + case ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT: + case ONNX_TENSOR_ELEMENT_DATA_TYPE_BFLOAT16: + case ONNX_TENSOR_ELEMENT_DATA_TYPE_INT8: + case ONNX_TENSOR_ELEMENT_DATA_TYPE_INT16: + case ONNX_TENSOR_ELEMENT_DATA_TYPE_INT32: + case ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64: + case ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT8: + return true; + default: + return false; + } +} + inline bool SetMudnnTensor(::musa::dnn::Tensor& tensor, const void* data, const std::vector& shape, ONNXTensorElementDataType elem_type) { + if (!IsMudnnElementwiseTypeSupported(elem_type)) return false; ::musa::dnn::Tensor::Type mudnn_type; if (!MudnnTensorType(elem_type, mudnn_type)) return false; if (tensor.SetAddr(data) != ::musa::dnn::Status::SUCCESS) return false; @@ -263,7 +287,11 @@ inline mublasStatus MublasGemmEx(mublasHandle_t handle, musaDataType_t data_type; mublasComputeType_t compute_type; if (!MublasDataType(elem_type, data_type, compute_type)) { - return MUBLAS_STATUS_NOT_SUPPORTED; + // Both the 4.3.x and 5.1.0 headers expose this value. Do not provide a + // preprocessor alias for MUBLAS_STATUS_NOT_SUPPORTED: in 5.1.0 that name + // is an enum member and such an alias rewrites the enum declaration while + // mublas_types.h is being parsed. + return MUBLAS_STATUS_NOT_IMPLEMENTED; } if (elem_type == ONNX_TENSOR_ELEMENT_DATA_TYPE_DOUBLE) { const double alpha64 = alpha; @@ -288,7 +316,7 @@ inline mublasStatus MublasGemmStridedBatchedEx( musaDataType_t data_type; mublasComputeType_t compute_type; if (!MublasDataType(elem_type, data_type, compute_type)) { - return MUBLAS_STATUS_NOT_SUPPORTED; + return MUBLAS_STATUS_NOT_IMPLEMENTED; } if (elem_type == ONNX_TENSOR_ELEMENT_DATA_TYPE_DOUBLE) { const double alpha64 = alpha; diff --git a/musa/ep/src/mudnn_compat.h b/musa/ep/src/mudnn_compat.h new file mode 100644 index 0000000..0523555 --- /dev/null +++ b/musa/ep/src/mudnn_compat.h @@ -0,0 +1,9 @@ +#pragma once + +// MUSA 5.1 ships the C++ API in a separate mudnncxx include directory, +// while 4.3.8 exposes the same API through mudnn.h. +#if __has_include() +#include +#else +#include +#endif diff --git a/musa/ep/src/musa_arena.cc b/musa/ep/src/musa_arena.cc index 38c3a56..8287095 100644 --- a/musa/ep/src/musa_arena.cc +++ b/musa/ep/src/musa_arena.cc @@ -314,7 +314,23 @@ void* ArenaImpl::Alloc(size_t size) { } void* ArenaImpl::AllocOnStream(size_t size, OrtSyncStream* stream) { - return AllocateRawInternal(size, stream, false); + // ORT allocates memory-pattern buffers on a DummyStream with a null native + // handle. It has no OrtSyncStreamImpl, so SyncStream_GetImpl would interpret + // the wrong object layout. Treat it like Alloc(): there is no MUSA stream + // completion/release callback that could make its chunks safe to cache. + ORT_TRY { + if (stream == nullptr || api_.SyncStream_GetHandle(stream) == nullptr) { + return Reserve(size); + } + return AllocateRawInternal(size, stream, false); + } + ORT_CATCH(const std::exception& ex) { + ORT_HANDLE_EXCEPTION([&]() { + MUSA_ARENA_LOG(ERROR, + "AllocOnStream(" << size << ") failed: " << ex.what()); + }); + return nullptr; + } } void* ArenaImpl::Reserve(size_t size) { @@ -428,6 +444,11 @@ void* ArenaImpl::AllocateRawInternal(size_t num_bytes, OrtSyncStream* stream, // Release the OrtStatus and return nullptr instead of throwing — allocate // calls must not propagate exceptions across the C API boundary. + MUSA_ARENA_LOG( + ERROR, "Arena allocation failed: requested=" + << num_bytes << " allocated=" << stats_.total_allocated_bytes + << " in_use=" << stats_.bytes_in_use << " limit=" + << config_.max_mem << ": " << api_.GetErrorMessage(status)); api_.ReleaseStatus(status); return nullptr; } diff --git a/musa/ep/src/musa_arena.h b/musa/ep/src/musa_arena.h index e48c856..0b39b4e 100644 --- a/musa/ep/src/musa_arena.h +++ b/musa/ep/src/musa_arena.h @@ -271,8 +271,8 @@ class ArenaImpl { void Free(void* p); // Allocate memory directly, outside the reusable arena. Used for - // stream-unaware callers and initializers so they don't affect arena growth - // patterns or become cross-stream reusable. + // stream-unaware callers, ORT's handle-less memory-pattern stream, and + // initializers so they don't become cross-stream reusable. void* Reserve(size_t size); // Release unused memory. Frees all allocation regions where every chunk is @@ -285,7 +285,7 @@ class ArenaImpl { size_t AllocatedSize(const void* ptr); // Un-assign chunks that are currently assigned to the stream. - // Called from OrtSyncStreamImpl::OnSessionRunEnd. + // Called on stream release after synchronizing the native MUSA stream. OrtStatus* ResetChunksUsingStream(const OrtSyncStreamImpl* stream_impl); private: diff --git a/test/multi_stream/README.md b/test/multi_stream/README.md deleted file mode 100644 index 5a5d26f..0000000 --- a/test/multi_stream/README.md +++ /dev/null @@ -1,34 +0,0 @@ -# MUSA multi-stream tests - -These tests require an installed `onnxruntime` package, the current -`onnxruntime-musa` plugin wheel, and at least one visible MUSA device. - -The allocator ownership index has ten device-independent CTest cases. They -cover assignment, clearing, ownership transfer, producer release after -transfer, split remainders, and rejected invalid index mutations: - -```bash -cmake --build build/Release --target musa_arena_ownership_test -ctest --test-dir build/Release --output-on-failure \ - -R '^musa_arena_ownership\.' -``` - -Run the stream-ordered arena shared-Session test with: - -```bash -PYTHONPATH=test/ops \ -ORT_MUSA_ALLOCATOR_CACHE_LIMIT_MB=128 \ -python -m pytest -q \ - test/multi_stream/test_allocator_cache_shared_session.py -``` - -The test creates one `InferenceSession` shared by four Python workers. It is -an end-to-end numerical check for concurrent Run lifetimes. For allocator -changes, also run the complete directory: - -```bash -PYTHONPATH=test/ops python -m pytest -q test/multi_stream -``` - -Before interpreting a result, verify that the runtime plugin path and SHA256 -match the newly built wheel. A skipped test does not count as MUSA validation. diff --git a/test/multi_stream/test_allocator_cache_shared_session.py b/test/multi_stream/test_allocator_cache_shared_session.py index 59c0ca9..0b546d1 100644 --- a/test/multi_stream/test_allocator_cache_shared_session.py +++ b/test/multi_stream/test_allocator_cache_shared_session.py @@ -49,14 +49,18 @@ def _make_model(feeds: dict[str, np.ndarray]) -> bytes: ) -def _create_musa_session(model: bytes): +def _create_musa_session(model: bytes, *, enable_mem_pattern=True): devices = musa_devices() if not devices: pytest.skip("No MUSA device available") options = ort.SessionOptions() + options.enable_mem_pattern = enable_mem_pattern + options.add_session_config_entry("session.disable_cpu_ep_fallback", "1") options.add_provider_for_devices(devices, {}) - return ort.InferenceSession(model, sess_options=options) + session = ort.InferenceSession(model, sess_options=options) + session.disable_fallback() + return session def test_allocator_cache_shared_session_is_stream_ordered(monkeypatch): @@ -119,26 +123,30 @@ def consumer() -> None: consumer_future.result() -def test_allocator_cache_survives_repeated_worker_stream_release(monkeypatch): +@pytest.mark.parametrize("enable_mem_pattern", [True, False], ids=["memory_pattern", "per_tensor"]) +def test_allocator_cache_survives_repeated_worker_stream_release(monkeypatch, enable_mem_pattern): """Repeated short-lived workers exercise stream release between cache uses.""" monkeypatch.setenv("ORT_MUSA_ALLOCATOR_CACHE_LIMIT_MB", "128") feeds = {"X": np.linspace(-3.0, 3.0, 16384, dtype=np.float32)} model = _make_model(feeds) cpu_session = ort.InferenceSession(model, providers=["CPUExecutionProvider"]) - (expected,) = cpu_session.run(None, feeds) - session = _create_musa_session(model) + worker_feeds = [{"X": feeds["X"] + np.float32(worker)} for worker in range(4)] + expected = [cpu_session.run(None, feed)[0] for feed in worker_feeds] + session = _create_musa_session(model, enable_mem_pattern=enable_mem_pattern) + start = threading.Barrier(4) - def run_once() -> None: - (actual,) = session.run(None, feeds) - np.testing.assert_array_equal(actual, expected) + def run_once(worker: int) -> None: + start.wait(timeout=30) + (actual,) = session.run(None, worker_feeds[worker]) + np.testing.assert_array_equal(actual, expected[worker]) # Recreate workers rather than keeping a fixed pool: this repeatedly tears # down ORT's per-worker stream state while the shared session and arena # remain alive. for _ in range(128): with ThreadPoolExecutor(max_workers=4) as executor: - list(executor.map(lambda _: run_once(), range(4))) + list(executor.map(run_once, range(4))) @pytest.mark.parametrize("invalid_value", ["-1", "128x", "18446744073709551616"]) diff --git a/test/ops/test_erf.py b/test/ops/test_erf.py index 7cd6fa5..8548cb2 100644 --- a/test/ops/test_erf.py +++ b/test/ops/test_erf.py @@ -12,9 +12,19 @@ # limitations under the License. """End-to-end CPU-vs-MUSA test for the Erf operator.""" +import math + import numpy as np +import pytest -from op_test_utils import TensorProto, run_and_compare +from op_test_utils import ( + TensorProto, + bfloat16_bits_to_float32, + build_model_with_input_types, + float32_to_bfloat16_bits, + run_and_compare, + run_with_iobinding, +) def test_erf_float(): @@ -25,3 +35,42 @@ def test_erf_float(): def test_erf_float_matrix(): x = np.array([[-3.0, -1.0, 0.0], [0.5, 1.0, 3.0]], dtype=np.float32) run_and_compare("Erf", inputs={"X": x}, outputs=[("Y", TensorProto.FLOAT)], rtol=1e-4, atol=1e-5) + + +@pytest.mark.parametrize( + "tensor_type,np_dtype,tol", + [ + (TensorProto.FLOAT16, np.float16, 2e-3), + (TensorProto.BFLOAT16, np.float32, 8e-3), + (TensorProto.FLOAT, np.float32, 1e-6), + (TensorProto.DOUBLE, np.float64, 1e-12), + ], + ids=["fp16", "bf16", "fp32", "fp64"], +) +def test_erf_special_values(tensor_type, np_dtype, tol): + """muDNN can return SUCCESS while producing NaN for infinite inputs.""" + values = np.array( + [-np.inf, -1e4, -3.0, -1.0, -0.0, 0.0, 0.5, 1.0, 3.0, 1e4, np.inf, np.nan], + dtype=np_dtype, + ) + is_bf16 = tensor_type == TensorProto.BFLOAT16 + x = float32_to_bfloat16_bits(values) if is_bf16 else values + reference_input = bfloat16_bits_to_float32(x) if is_bf16 else x + expected = np.array([math.erf(float(value)) for value in reference_input], dtype=np_dtype) + if is_bf16: + expected = bfloat16_bits_to_float32(float32_to_bfloat16_bits(expected)) + model = build_model_with_input_types( + "Erf", inputs={"X": x}, input_types={"X": tensor_type}, outputs=[("Y", tensor_type)] + ) + # The helper requires a MUSA device and disables CPU EP fallback, including + # for BF16/DOUBLE where the CPU kernel is not our numerical reference. + (output,) = run_with_iobinding( + model, {"X": x}, {"X": tensor_type}, [("Y", tensor_type, x.shape)], use_musa=True + ) + actual = bfloat16_bits_to_float32(output) if is_bf16 else output + np.testing.assert_allclose(actual, expected, rtol=tol, atol=tol, equal_nan=True) + np.testing.assert_array_equal(actual[np.isinf(reference_input)], [-1.0, 1.0]) + np.testing.assert_array_equal(actual[reference_input == 0], 0.0) + np.testing.assert_array_equal( + np.signbit(actual[reference_input == 0]), np.signbit(reference_input[reference_input == 0]) + ) From e074b790d5a14b83eb6372e9c065765d5613b384 Mon Sep 17 00:00:00 2001 From: Albert Date: Fri, 9 Oct 2026 12:12:12 +0800 Subject: [PATCH 02/52] Fix(bert infer RmsNorm bug) --- musa/ep/src/fusion/rms_norm_fusion.cc | 50 ++++++++++++--------------- test/fusion/test_rms_norm_fusion.py | 16 +++++++++ 2 files changed, 39 insertions(+), 27 deletions(-) diff --git a/musa/ep/src/fusion/rms_norm_fusion.cc b/musa/ep/src/fusion/rms_norm_fusion.cc index 6818767..534e842 100644 --- a/musa/ep/src/fusion/rms_norm_fusion.cc +++ b/musa/ep/src/fusion/rms_norm_fusion.cc @@ -91,16 +91,6 @@ size_t GetMappedIndex(const std::unordered_map& indices, return it->second; } -bool IsFloatGpuTensor(Ort::ConstValue value) { - auto info = value.GetTensorTypeAndShapeInfo(); - ONNXTensorElementDataType elem_type = info.GetElementType(); - return (elem_type == ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT || - elem_type == ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT16 || - elem_type == ONNX_TENSOR_ELEMENT_DATA_TYPE_BFLOAT16 || - elem_type == ONNX_TENSOR_ELEMENT_DATA_TYPE_DOUBLE) && - IsGpuMemory(value.GetTensorMemoryInfo()); -} - bool IsRmsNormStorageType(ONNXTensorElementDataType elem_type) { return elem_type == ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT || elem_type == ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT16 || @@ -108,6 +98,11 @@ bool IsRmsNormStorageType(ONNXTensorElementDataType elem_type) { elem_type == ONNX_TENSOR_ELEMENT_DATA_TYPE_DOUBLE; } +bool IsRmsNormFloatTensor(Ort::ConstValue value) { + auto info = value.GetTensorTypeAndShapeInfo(); + return IsRmsNormStorageType(info.GetElementType()); +} + class DeviceTempBuffer { public: DeviceTempBuffer(size_t bytes, musaStream_t stream) @@ -142,15 +137,14 @@ std::vector RmsNormInvVarShape( return inv_var_shape; } -bool TryMudnnRmsNorm(Ort::ConstValue input, const void* gamma_data, +bool TryMudnnRmsNorm(const void* input_data, const void* gamma_data, Ort::UnownedValue output, const std::vector& input_shape, const std::vector& gamma_shape, ONNXTensorElementDataType elem_type, float epsilon, musaStream_t stream) { - if (!IsGpuMemory(input.GetTensorMemoryInfo()) || - !IsGpuMemory(output.GetTensorMemoryInfo()) || gamma_data == nullptr || - input_shape.empty() || + if (input_data == nullptr || !IsGpuMemory(output.GetTensorMemoryInfo()) || + gamma_data == nullptr || input_shape.empty() || elem_type == ONNX_TENSOR_ELEMENT_DATA_TYPE_DOUBLE) { return false; } @@ -176,8 +170,7 @@ bool TryMudnnRmsNorm(Ort::ConstValue input, const void* gamma_data, ::musa::dnn::Tensor output_tensor; ::musa::dnn::Tensor gamma_tensor; ::musa::dnn::Tensor inv_var_tensor; - if (!SetMudnnTensor(input_tensor, input.GetTensorRawData(), input_shape, - elem_type) || + if (!SetMudnnTensor(input_tensor, input_data, input_shape, elem_type) || !SetMudnnTensor(output_tensor, output.GetTensorMutableRawData(), input_shape, elem_type) || !SetMudnnTensor(gamma_tensor, gamma_data, gamma_shape, elem_type) || @@ -214,7 +207,7 @@ struct RmsNormFusionCompute : FusionNodeCompute { Ort::KernelContext ctx(kernel_context); Ort::ConstValue input = ctx.GetInput(input_index); Ort::ConstValue gamma = ctx.GetInput(gamma_index); - if (!IsFloatGpuTensor(input)) { + if (!IsRmsNormFloatTensor(input)) { return Ort::GetApi().CreateStatus( ORT_NOT_IMPLEMENTED, "RmsNorm requires a MUSA floating input"); } @@ -261,15 +254,17 @@ struct RmsNormFusionCompute : FusionNodeCompute { } const musaStream_t stream = GetComputeStream(ctx); + DeviceInputBuffer input_buffer; DeviceInputBuffer gamma_buffer; + RETURN_IF_ERROR(input_buffer.Bind(input, stream)); RETURN_IF_ERROR(gamma_buffer.Bind(gamma, stream)); - if (TryMudnnRmsNorm(input, gamma_buffer.data(), output, input_shape, - gamma_shape, input_info.GetElementType(), epsilon, - stream)) { + if (TryMudnnRmsNorm(input_buffer.data(), gamma_buffer.data(), output, + input_shape, gamma_shape, input_info.GetElementType(), + epsilon, stream)) { return nullptr; } return LaunchStatus( - LaunchMusaRmsNormKernel(input.GetTensorRawData(), gamma_buffer.data(), + LaunchMusaRmsNormKernel(input_buffer.data(), gamma_buffer.data(), output.GetTensorMutableRawData(), rows, norm_size, epsilon, elem_type, stream)); } catch (const Ort::Exception& ex) { @@ -304,8 +299,7 @@ struct CastRmsNormFusionCompute : FusionNodeCompute { if (input_info.GetElementType() != ONNX_TENSOR_ELEMENT_DATA_TYPE_BFLOAT16 || gamma_info.GetElementType() != - ONNX_TENSOR_ELEMENT_DATA_TYPE_BFLOAT16 || - !IsGpuMemory(input.GetTensorMemoryInfo())) { + ONNX_TENSOR_ELEMENT_DATA_TYPE_BFLOAT16) { return Ort::GetApi().CreateStatus( ORT_NOT_IMPLEMENTED, "CastRmsNorm requires MUSA bfloat16 input and gamma"); @@ -345,15 +339,17 @@ struct CastRmsNormFusionCompute : FusionNodeCompute { } const musaStream_t stream = GetComputeStream(ctx); + DeviceInputBuffer input_buffer; DeviceInputBuffer gamma_buffer; + RETURN_IF_ERROR(input_buffer.Bind(input, stream)); RETURN_IF_ERROR(gamma_buffer.Bind(gamma, stream)); - if (TryMudnnRmsNorm(input, gamma_buffer.data(), output, input_shape, - gamma_shape, input_info.GetElementType(), epsilon, - stream)) { + if (TryMudnnRmsNorm(input_buffer.data(), gamma_buffer.data(), output, + input_shape, gamma_shape, input_info.GetElementType(), + epsilon, stream)) { return nullptr; } return LaunchStatus(LaunchMusaCastRmsNormBf16Kernel( - input.GetTensorRawData(), gamma_buffer.data(), + input_buffer.data(), gamma_buffer.data(), output.GetTensorMutableRawData(), rows, norm_size, epsilon, stream)); } catch (const Ort::Exception& ex) { Ort::Status status(ex); diff --git a/test/fusion/test_rms_norm_fusion.py b/test/fusion/test_rms_norm_fusion.py index 83a0c28..103e5f8 100644 --- a/test/fusion/test_rms_norm_fusion.py +++ b/test/fusion/test_rms_norm_fusion.py @@ -277,6 +277,22 @@ def test_rms_norm_fusion(tmp_path): _assert_rms_norm_fused(op_names) +def test_rms_norm_fusion_accepts_cpu_bound_input(tmp_path): + """The fused kernel must copy a normal host feed before launching on MUSA.""" + model = _build_rms_norm_model() + x = np.linspace(-1.0, 1.0, num=2 * 3 * 5, dtype=np.float32).reshape(2, 3, 5) + + (actual,) = run_model_and_compare(model, {"X": x}, rtol=1e-5, atol=1e-5) + gamma = np.linspace(0.5, 1.5, num=5, dtype=np.float32) + expected = _rms_norm_reference(x, gamma, 1.0e-6) + np.testing.assert_allclose(actual, expected, rtol=1e-5, atol=1e-5) + + op_names = _profile_op_names( + model, {"X": x}, str(tmp_path / "rms_norm_cpu_bound_input") + ) + _assert_rms_norm_fused(op_names) + + def test_pow_rms_norm_fusion(tmp_path): model = _build_pow_rms_norm_model() x = np.linspace(-2.0, 2.0, num=2 * 3 * 5, dtype=np.float32).reshape(2, 3, 5) From 476d67d864fb14dea943dfc9be9002fd4be9e04a Mon Sep 17 00:00:00 2001 From: Albert Date: Fri, 9 Oct 2026 12:16:40 +0800 Subject: [PATCH 03/52] Fix(Dsin splitconcat fusion bug) --- .../src/fusion/split_concat_fusion_matcher.cc | 19 ++++++++ test/fusion/test_split_concat_fusion.py | 43 +++++++++++++++++++ 2 files changed, 62 insertions(+) diff --git a/musa/ep/src/fusion/split_concat_fusion_matcher.cc b/musa/ep/src/fusion/split_concat_fusion_matcher.cc index 645dae4..d94cdb8 100644 --- a/musa/ep/src/fusion/split_concat_fusion_matcher.cc +++ b/musa/ep/src/fusion/split_concat_fusion_matcher.cc @@ -62,6 +62,18 @@ bool CanFuseSplitConcat( !IsFloatTensorValueInfo(reshape_inputs[0])) { return false; } + + // The runtime implementation consumes the tensor before Reshape and + // reconstructs the Reshape dimensions from the statically known output + // shape. Do not select a graph whose Reshape shape is only known at + // runtime: CreateSplitConcatFusion cannot safely derive sequence and part + // width for that graph and would otherwise fail during Compile. + auto reshape_shape = GetTensorShape(reshape_outputs[0]); + if (!reshape_shape.has_value() || reshape_shape->size() != 3 || + (*reshape_shape)[1] <= 0 || (*reshape_shape)[2] <= 0 || + (*reshape_shape)[2] % static_cast(split_outputs.size()) != 0) { + return false; + } } auto split_axis_attr = GetIntAttribute(split_node, "axis"); @@ -79,6 +91,7 @@ bool CanFuseSplitConcat( } std::vector split_sizes; + const bool has_explicit_split_sizes = split_inputs.size() == 2; if (split_inputs.size() == 2) { auto split_initializer = ReadIntInitializerNoLimit(split_inputs[1]); if (!split_initializer.has_value() || @@ -103,6 +116,12 @@ bool CanFuseSplitConcat( } split_total += split_size; } + if (reshape_node && has_explicit_split_sizes) { + const auto reshape_shape = GetTensorShape(reshape_node.GetOutputs()[0]); + if (!reshape_shape.has_value() || (*reshape_shape)[2] != split_total) { + return false; + } + } for (size_t i = 0; i < split_outputs.size(); ++i) { Ort::ConstValueInfo split_output = split_outputs[i]; if (Name(concat_inputs[i]) != Name(split_output) || diff --git a/test/fusion/test_split_concat_fusion.py b/test/fusion/test_split_concat_fusion.py index e5cd1fa..b3ca7cc 100644 --- a/test/fusion/test_split_concat_fusion.py +++ b/test/fusion/test_split_concat_fusion.py @@ -163,3 +163,46 @@ def test_split_concat_fusion_absorbs_optional_transpose(): node_names = _profile_musa_node_names(model_bytes, {"X": x}) for name in ("SplitConcatSplit", "SplitConcatConcat", "SplitConcatTranspose"): assert not any(event.startswith(name) for event in node_names), node_names + + +def test_split_concat_fusion_skips_dynamic_upstream_reshape(): + """A runtime Reshape shape must not make Compile reject the whole graph.""" + rng = np.random.default_rng(109) + x = rng.standard_normal((2, 12)).astype(np.float32) + reshape_shape = np.array([2, 3, 4], dtype=np.int64) + split_outputs = [f"S{i}" for i in range(4)] + graph = helper.make_graph( + [ + # Deliberately leave the intermediate Reshape value without a + # ValueInfo. Its dimensions are supplied at runtime, so the + # SplitConcat matcher must decline this fusion candidate. + helper.make_node("Reshape", ["X", "shape"], ["R"], name="DynamicReshape"), + helper.make_node( + "Split", ["R"], split_outputs, axis=2, name="DynamicSplit" + ), + helper.make_node( + "Concat", split_outputs, ["Y"], axis=0, name="DynamicConcat" + ), + ], + "split_concat_dynamic_reshape_graph", + [ + helper.make_tensor_value_info("X", TensorProto.FLOAT, [2, 12]), + helper.make_tensor_value_info("shape", TensorProto.INT64, [3]), + ], + [helper.make_tensor_value_info("Y", TensorProto.FLOAT, None)], + ) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 17)]) + model.ir_version = min(model.ir_version, 10) + model_bytes = model.SerializeToString() + + musa_outputs = run_model_and_compare( + model_bytes, + {"X": x, "shape": reshape_shape}, + rtol=1e-6, + atol=1e-6, + ) + expected = np.concatenate( + np.split(x.reshape(2, 3, 4), 4, axis=2), + axis=0, + ) + np.testing.assert_allclose(musa_outputs[0], expected, rtol=0, atol=0) From caadf3ef27c7abae81c71dde2238ddedcb0711e2 Mon Sep 17 00:00:00 2001 From: Albert Date: Fri, 9 Oct 2026 15:45:13 +0800 Subject: [PATCH 04/52] Fix(matmul fusion empty concat) --- musa/ep/src/fusion/concat_matmul_fusion.cc | 52 +++++++++++++++++++-- test/fusion/test_concat_matmul_fusion.py | 54 ++++++++++++++++++++++ 2 files changed, 102 insertions(+), 4 deletions(-) diff --git a/musa/ep/src/fusion/concat_matmul_fusion.cc b/musa/ep/src/fusion/concat_matmul_fusion.cc index 2041b39..e93165d 100644 --- a/musa/ep/src/fusion/concat_matmul_fusion.cc +++ b/musa/ep/src/fusion/concat_matmul_fusion.cc @@ -481,10 +481,54 @@ void RunMudnnConcat(const std::vector& concat_inputs, void* concat_data, ONNXTensorElementDataType elem_type, musaStream_t stream) { ::musa::dnn::Handle* handle = MudnnHandleOrThrow(stream); - std::vector<::musa::dnn::Tensor> input_tensors(concat_inputs.size()); - for (size_t i = 0; i < concat_inputs.size(); ++i) { - SetupMudnnTensor(input_tensors[i], concat_inputs[i].GetTensorRawData(), - TensorShape(concat_inputs[i]), elem_type); + const size_t axis_index = static_cast(axis); + + // muDNN Concat rejects tensors with a zero extent, while ONNX permits empty + // inputs along the concatenation axis. Drop those inputs before invoking + // muDNN; they contribute no elements and therefore do not change the + // concatenation result. + std::vector non_empty_inputs; + non_empty_inputs.reserve(concat_inputs.size()); + for (Ort::ConstValue input : concat_inputs) { + const std::vector shape = TensorShape(input); + if (axis_index >= shape.size()) { + throw std::runtime_error("ConcatMatMul concat axis is out of range"); + } + if (shape[axis_index] < 0) { + throw std::runtime_error( + "ConcatMatMul concat input has an unresolved dimension"); + } + if (shape[axis_index] != 0) { + non_empty_inputs.push_back(input); + } + } + + if (non_empty_inputs.empty()) { + throw std::runtime_error( + "ConcatMatMul concat has no non-empty inputs for a non-empty output"); + } + + // With only one non-empty input, avoid relying on whether muDNN accepts a + // one-input Concat. Empty inputs have already been removed, so the output + // is layout-identical to that input and a device-to-device copy is enough. + if (non_empty_inputs.size() == 1) { + const size_t bytes = + static_cast(NumElements(concat_shape)) * ElementSize(elem_type); + const musaError_t status = + musaMemcpyAsync(concat_data, non_empty_inputs[0].GetTensorRawData(), + bytes, musaMemcpyDeviceToDevice, stream); + if (status != musaSuccess) { + throw std::runtime_error( + std::string("ConcatMatMul concat copy failed: ") + + MusaErrorString(status)); + } + return; + } + + std::vector<::musa::dnn::Tensor> input_tensors(non_empty_inputs.size()); + for (size_t i = 0; i < non_empty_inputs.size(); ++i) { + SetupMudnnTensor(input_tensors[i], non_empty_inputs[i].GetTensorRawData(), + TensorShape(non_empty_inputs[i]), elem_type); } ::musa::dnn::Tensor output_tensor; diff --git a/test/fusion/test_concat_matmul_fusion.py b/test/fusion/test_concat_matmul_fusion.py index affca87..2d4f9f4 100644 --- a/test/fusion/test_concat_matmul_fusion.py +++ b/test/fusion/test_concat_matmul_fusion.py @@ -262,6 +262,60 @@ def test_concat_matmul_fusion_zero_k_non_empty_output(tmp_path): assert any(str(op).startswith("MUSAExecutionProvider_") for op in musa_ops) +def test_concat_matmul_fusion_skips_empty_concat_input(tmp_path): + rng = np.random.default_rng(11) + x0 = rng.standard_normal((2, 3)).astype(np.float32) + x1 = np.empty((2, 0), dtype=np.float32) + x2 = rng.standard_normal((2, 2)).astype(np.float32) + b = rng.standard_normal((5, 4)).astype(np.float32) + + nodes = [ + helper.make_node("Concat", ["X0", "X1", "X2"], ["C"], axis=1), + helper.make_node("MatMul", ["C", "B"], ["Y"]), + ] + feeds = {"X0": x0, "X1": x1, "X2": x2, "B": b} + model = build_graph_model( + nodes, + feeds, + [("Y", TensorProto.FLOAT)], + name="concat_matmul_empty_axis_input_graph", + ) + + outputs = run_model_and_compare(model, feeds, rtol=1e-3, atol=1e-3) + np.testing.assert_allclose(outputs[0], np.concatenate([x0, x2], axis=1) @ b, + rtol=1e-3, atol=1e-3) + musa_ops = _profile_musa_ops( + model, feeds, tmp_path, "concat_matmul_empty_axis_input" + ) + assert any(str(op).startswith("MUSAExecutionProvider_") for op in musa_ops) + + +def test_concat_matmul_fusion_only_one_non_empty_concat_input(tmp_path): + rng = np.random.default_rng(12) + x0 = rng.standard_normal((2, 3)).astype(np.float32) + x1 = np.empty((2, 0), dtype=np.float32) + b = rng.standard_normal((3, 4)).astype(np.float32) + + nodes = [ + helper.make_node("Concat", ["X0", "X1"], ["C"], axis=1), + helper.make_node("MatMul", ["C", "B"], ["Y"]), + ] + feeds = {"X0": x0, "X1": x1, "B": b} + model = build_graph_model( + nodes, + feeds, + [("Y", TensorProto.FLOAT)], + name="concat_matmul_single_non_empty_input_graph", + ) + + outputs = run_model_and_compare(model, feeds, rtol=1e-3, atol=1e-3) + np.testing.assert_allclose(outputs[0], x0 @ b, rtol=1e-3, atol=1e-3) + musa_ops = _profile_musa_ops( + model, feeds, tmp_path, "concat_matmul_single_non_empty_input" + ) + assert any(str(op).startswith("MUSAExecutionProvider_") for op in musa_ops) + + def test_concat_matmul_fusion_float16_profiles_fused(tmp_path): rng = np.random.default_rng(4) x0 = rng.standard_normal((2, 3, 4, 5)).astype(np.float16) From f85b88b297cb6e61a234bc055c6dabde89ab04b9 Mon Sep 17 00:00:00 2001 From: Albert Date: Mon, 12 Oct 2026 09:20:49 +0800 Subject: [PATCH 05/52] fix comment --- musa/ep/src/fusion/rms_norm_fusion.cc | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/musa/ep/src/fusion/rms_norm_fusion.cc b/musa/ep/src/fusion/rms_norm_fusion.cc index 534e842..48f8cb8 100644 --- a/musa/ep/src/fusion/rms_norm_fusion.cc +++ b/musa/ep/src/fusion/rms_norm_fusion.cc @@ -209,7 +209,7 @@ struct RmsNormFusionCompute : FusionNodeCompute { Ort::ConstValue gamma = ctx.GetInput(gamma_index); if (!IsRmsNormFloatTensor(input)) { return Ort::GetApi().CreateStatus( - ORT_NOT_IMPLEMENTED, "RmsNorm requires a MUSA floating input"); + ORT_NOT_IMPLEMENTED, "RmsNorm requires a floating-point input"); } auto input_info = input.GetTensorTypeAndShapeInfo(); @@ -302,7 +302,7 @@ struct CastRmsNormFusionCompute : FusionNodeCompute { ONNX_TENSOR_ELEMENT_DATA_TYPE_BFLOAT16) { return Ort::GetApi().CreateStatus( ORT_NOT_IMPLEMENTED, - "CastRmsNorm requires MUSA bfloat16 input and gamma"); + "CastRmsNorm requires bfloat16 input and gamma"); } std::vector input_shape = input_info.GetShape(); From 56f13aad74b749b45b6f3c59ab079e8bf24dc724 Mon Sep 17 00:00:00 2001 From: "yichao.hu" Date: Wed, 16 Sep 2026 00:37:58 -0400 Subject: [PATCH 06/52] feat(conv): support grouped and depthwise convolution --- musa/ep/src/kernels/nn/conv.cc | 9 ++++-- musa/ep/src/kernels/nn/conv_impl.mu | 9 ++++-- .../kernels/shared_inc/device_kernel_types.h | 1 + test/ops/test_conv.py | 31 +++++++++++++++++++ 4 files changed, 45 insertions(+), 5 deletions(-) diff --git a/musa/ep/src/kernels/nn/conv.cc b/musa/ep/src/kernels/nn/conv.cc index 6d2dc15..8b3e2aa 100644 --- a/musa/ep/src/kernels/nn/conv.cc +++ b/musa/ep/src/kernels/nn/conv.cc @@ -157,9 +157,11 @@ OrtStatus* Conv::Compute(Ort::KernelContext& ctx) const { return Ort::GetApi().CreateStatus(ORT_NOT_IMPLEMENTED, "Conv only supports 2D NCHW tensors"); } - if (group_ != 1 || w_shape[1] != x_shape[1]) { - return Ort::GetApi().CreateStatus(ORT_NOT_IMPLEMENTED, - "Conv only supports group=1"); + if (group_ <= 0 || x_shape[1] != w_shape[1] * group_ || + w_shape[0] % group_ != 0) { + return Ort::GetApi().CreateStatus( + ORT_INVALID_ARGUMENT, + "Conv group must divide output channels and match input channels"); } std::vector pads = pads_; @@ -213,6 +215,7 @@ OrtStatus* Conv::Compute(Ort::KernelContext& ctx) const { MusaConv2DParams params{}; params.n = x_shape[0]; params.c = x_shape[1]; + params.group = group_; params.h = x_shape[2]; params.w = x_shape[3]; params.m = w_shape[0]; diff --git a/musa/ep/src/kernels/nn/conv_impl.mu b/musa/ep/src/kernels/nn/conv_impl.mu index cf96abf..9072890 100644 --- a/musa/ep/src/kernels/nn/conv_impl.mu +++ b/musa/ep/src/kernels/nn/conv_impl.mu @@ -24,7 +24,12 @@ __global__ void Conv2DFloatKernel(const float* input, const int64_t n = remaining; float acc = bias == nullptr ? 0.0f : bias[oc]; - for (int64_t ic = 0; ic < params.c; ++ic) { + const int64_t channels_per_group = params.c / params.group; + const int64_t outputs_per_group = params.m / params.group; + const int64_t group_index = oc / outputs_per_group; + for (int64_t ic_in_group = 0; ic_in_group < channels_per_group; + ++ic_in_group) { + const int64_t ic = group_index * channels_per_group + ic_in_group; for (int64_t kh = 0; kh < params.kernel_h; ++kh) { const int64_t ih = oh * params.stride_h + kh * params.dilation_h - params.pad_h; @@ -40,7 +45,7 @@ __global__ void Conv2DFloatKernel(const float* input, const int64_t input_index = ((n * params.c + ic) * params.h + ih) * params.w + iw; const int64_t weight_index = - ((oc * params.c + ic) * params.kernel_h + kh) * + ((oc * channels_per_group + ic_in_group) * params.kernel_h + kh) * params.kernel_w + kw; acc += input[input_index] * weight[weight_index]; diff --git a/musa/ep/src/kernels/shared_inc/device_kernel_types.h b/musa/ep/src/kernels/shared_inc/device_kernel_types.h index 22dda5e..764e894 100644 --- a/musa/ep/src/kernels/shared_inc/device_kernel_types.h +++ b/musa/ep/src/kernels/shared_inc/device_kernel_types.h @@ -291,6 +291,7 @@ struct MusaTopKParams { struct MusaConv2DParams { int64_t n; int64_t c; + int64_t group; int64_t h; int64_t w; int64_t m; diff --git a/test/ops/test_conv.py b/test/ops/test_conv.py index c79f1e9..8f60dde 100644 --- a/test/ops/test_conv.py +++ b/test/ops/test_conv.py @@ -46,6 +46,37 @@ def test_conv_opset11_float_nchw_with_bias(): ) +def test_conv_opset11_float_nchw_grouped_with_bias(): + x = np.random.default_rng(7).standard_normal((1, 4, 5, 5)).astype(np.float32) + w = np.random.default_rng(8).standard_normal((6, 2, 3, 3)).astype(np.float32) + b = np.random.default_rng(9).standard_normal((6,)).astype(np.float32) + run_and_compare( + "Conv", + inputs={"X": x, "W": w, "B": b}, + outputs=[("Y", TensorProto.FLOAT)], + attrs={"group": 2, "pads": [1, 1, 1, 1], "strides": [1, 1]}, + opset=11, + rtol=1e-4, + atol=1e-4, + ) + + +def test_conv_opset11_float_nchw_depthwise_fallback_with_bias(): + # H=1 and kH=1 deliberately selects Conv's custom device fallback path. + x = np.random.default_rng(10).standard_normal((1, 4, 1, 7)).astype(np.float32) + w = np.random.default_rng(11).standard_normal((4, 1, 1, 3)).astype(np.float32) + b = np.random.default_rng(12).standard_normal((4,)).astype(np.float32) + run_and_compare( + "Conv", + inputs={"X": x, "W": w, "B": b}, + outputs=[("Y", TensorProto.FLOAT)], + attrs={"group": 4, "pads": [0, 1, 0, 1], "strides": [1, 1]}, + opset=11, + rtol=1e-4, + atol=1e-4, + ) + + def test_conv_lowered_conv1d_h1_k1(): x = np.random.default_rng(0).standard_normal((2, 64, 1, 8)).astype(np.float32) w = np.random.default_rng(1).standard_normal((16, 64, 1, 1)).astype(np.float32) From ddeeb0ec4b9bc597cd1c6fc50d55fe9d678568ac Mon Sep 17 00:00:00 2001 From: "yichao.hu" Date: Mon, 12 Oct 2026 23:47:31 -0400 Subject: [PATCH 07/52] fix(tensor): stage mixed CPU inputs for Concat and add a nearest Resize route Concat failed when some of its inputs were produced on the CPU EP and the kernel read them as device memory. Stage those inputs onto the device before the kernel runs, and add a nearest/asymmetric/floor Resize route so models whose Resize nodes previously had no MUSA kernel can run on the plugin EP. --- musa/docs/supported_ops.md | 5 +- musa/ep/src/kernels/tensor/concat.cc | 15 ++- musa/ep/src/kernels/tensor/resize.cc | 114 ++++++++++++++++++++++ musa/ep/src/kernels/tensor/resize_impl.h | 11 +++ musa/ep/src/kernels/tensor/resize_impl.mu | 51 ++++++++++ test/ops/test_concat.py | 30 ++++++ test/ops/test_resize.py | 35 +++++++ 7 files changed, 255 insertions(+), 6 deletions(-) create mode 100644 musa/ep/src/kernels/tensor/resize.cc create mode 100644 musa/ep/src/kernels/tensor/resize_impl.h create mode 100644 musa/ep/src/kernels/tensor/resize_impl.mu create mode 100644 test/ops/test_resize.py diff --git a/musa/docs/supported_ops.md b/musa/docs/supported_ops.md index 88446e4..4f535cd 100644 --- a/musa/docs/supported_ops.md +++ b/musa/docs/supported_ops.md @@ -1,8 +1,8 @@ # Supported Ops -- Total supported ops: **89** -- Total registered kernels: **121** +- Total supported ops: **90** +- Total registered kernels: **122** Rows are grouped by operator/domain. Multiple opset ranges in a row preserve separate ORT kernel registrations. @@ -71,6 +71,7 @@ Rows are grouped by operator/domain. Multiple opset ranges in a row preserve sep | `ReduceSumSquare` | `ai.onnx` | 13–19: `T`: float16, float32, float64, bfloat16 | `reduction/reduce_sum_square.cc` | | `Relu` | `ai.onnx` | 13–19: `T`: float16, float32, float64, bfloat16 | `activation/relu.cc` | | `Reshape` | `ai.onnx` | 13–19: `T`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bfloat16, bool
19–19: `T`: string | `tensor/reshape.cc` | +| `Resize` | `ai.onnx` | 13–17: `T1`: float32 | `tensor/resize.cc` | | `ReverseSequence` | `ai.onnx` | 10–19: `T`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bool | `tensor/reverse_sequence.cc` | | `Round` | `ai.onnx` | 11–19: `T`: float16, float32, float64 | `math/round.cc` | | `ScatterElements` | `ai.onnx` | 11–12: `T`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bfloat16, bool
`Tind`: int32, int64
13–19: `T`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bfloat16, bool
`Tind`: int32, int64 | `tensor/scatter_elements.cc` | diff --git a/musa/ep/src/kernels/tensor/concat.cc b/musa/ep/src/kernels/tensor/concat.cc index ad360c1..a6a3238 100644 --- a/musa/ep/src/kernels/tensor/concat.cc +++ b/musa/ep/src/kernels/tensor/concat.cc @@ -12,6 +12,7 @@ // limitations under the License. #include +#include #include "shared_inc/blas_utils.h" #include "shared_inc/op_kernel_common.h" @@ -264,7 +265,9 @@ OrtStatus* Concat::Compute(Ort::KernelContext& ctx) const { } Ort::UnownedValue y = ctx.GetOutput(0, out_shape); - if (AllGpuInputs(ctx) && IsGpuMemory(y.GetTensorMemoryInfo())) { + const bool all_gpu_inputs = AllGpuInputs(ctx); + if (IsGpuMemory(y.GetTensorMemoryInfo())) { + musaStream_t stream = GetComputeStream(ctx); const int64_t outer = axis == 0 ? 1 : std::accumulate(out_shape.begin(), out_shape.begin() + axis, @@ -275,6 +278,8 @@ OrtStatus* Concat::Compute(Ort::KernelContext& ctx) const { : std::accumulate(out_shape.begin() + axis + 1, out_shape.end(), int64_t{1}, std::multiplies()); std::vector input_data(shapes.size()); + std::vector> input_buffers; + input_buffers.reserve(shapes.size()); std::vector input_axis_dims(shapes.size()); int64_t max_input_axis = 0; void* output_data = y.GetTensorMutableRawData(); @@ -283,7 +288,10 @@ OrtStatus* Concat::Compute(Ort::KernelContext& ctx) const { bool output_overlaps_input = false; for (size_t input_idx = 0; input_idx < shapes.size(); ++input_idx) { Ort::ConstValue v = ctx.GetInput(input_idx); - input_data[input_idx] = v.GetTensorRawData(); + auto input_buffer = std::make_unique(); + RETURN_IF_ERROR(input_buffer->Bind(v, stream)); + input_data[input_idx] = input_buffer->data(); + input_buffers.push_back(std::move(input_buffer)); input_axis_dims[input_idx] = shapes[input_idx][static_cast(axis)]; max_input_axis = std::max(max_input_axis, input_axis_dims[input_idx]); const size_t input_bytes = @@ -292,7 +300,6 @@ OrtStatus* Concat::Compute(Ort::KernelContext& ctx) const { output_data, output_bytes, input_data[input_idx], input_bytes); } - musaStream_t stream = GetComputeStream(ctx); void* concat_output = output_data; void* temp_output = nullptr; if (output_overlaps_input && output_bytes > 0) { @@ -376,7 +383,7 @@ OrtStatus* Concat::Compute(Ort::KernelContext& ctx) const { output_row_elements, static_cast(elem_size), stream)); } - if (!output_overlaps_input && + if (all_gpu_inputs && !output_overlaps_input && elem_type == ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT && TryMudnnConcat(ctx, shapes, out_shape, axis, elem_type, y)) { return nullptr; diff --git a/musa/ep/src/kernels/tensor/resize.cc b/musa/ep/src/kernels/tensor/resize.cc new file mode 100644 index 0000000..bfe8ab7 --- /dev/null +++ b/musa/ep/src/kernels/tensor/resize.cc @@ -0,0 +1,114 @@ +// Copyright (c) Moore Threads Technology Co., Ltd. All rights reserved. +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "shared_inc/op_kernel_common.h" +#include "tensor/resize_impl.h" + +namespace { + +class Resize : public OpKernelBase { + public: + Resize(const OrtKernelInfo* info, void* /*state*/) { + Ort::ConstKernelInfo kernel_info(info); + mode_ = AttrOrDefault(kernel_info, "mode", "nearest"); + coordinate_transform_mode_ = AttrOrDefault( + kernel_info, "coordinate_transformation_mode", "half_pixel"); + nearest_mode_ = AttrOrDefault(kernel_info, "nearest_mode", + "round_prefer_floor"); + } + OrtStatus* Compute(Ort::KernelContext& ctx) const; + + private: + std::string mode_; + std::string coordinate_transform_mode_; + std::string nearest_mode_; +}; + +OrtStatus* Resize::Compute(Ort::KernelContext& ctx) const { + if (mode_ != "nearest" || coordinate_transform_mode_ != "asymmetric" || + nearest_mode_ != "floor") { + return Ort::GetApi().CreateStatus( + ORT_NOT_IMPLEMENTED, + "Resize only supports nearest/asymmetric/floor in this MUSA route"); + } + if (ctx.GetInputCount() <= 2 || ctx.GetInput(2) == nullptr || + (ctx.GetInputCount() > 3 && ctx.GetInput(3) != nullptr)) { + return Ort::GetApi().CreateStatus( + ORT_NOT_IMPLEMENTED, + "Resize only supports the scales input, not the sizes input"); + } + if (ctx.GetInputCount() > 1 && ctx.GetInput(1) != nullptr && + ctx.GetInput(1).GetTensorTypeAndShapeInfo().GetElementCount() != 0) { + return Ort::GetApi().CreateStatus( + ORT_NOT_IMPLEMENTED, + "Resize only supports an empty ROI for asymmetric coordinates"); + } + + Ort::ConstValue input = ctx.GetInput(0); + auto input_info = input.GetTensorTypeAndShapeInfo(); + if (input_info.GetElementType() != ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT) { + return Ort::GetApi().CreateStatus(ORT_NOT_IMPLEMENTED, + "Resize only supports float32"); + } + const std::vector input_shape = input_info.GetShape(); + if (input_shape.size() != 4 || input_shape[2] <= 0 || input_shape[3] <= 0) { + return Ort::GetApi().CreateStatus( + ORT_NOT_IMPLEMENTED, + "Resize only supports non-empty rank-4 NCHW tensors"); + } + const std::vector scales = + ReadTyped(ctx.GetInput(2), GetComputeStream(ctx)); + if (scales.size() != 4 || scales[0] != 1.0f || scales[1] != 1.0f || + !std::isfinite(scales[2]) || !std::isfinite(scales[3]) || + scales[2] <= 0.0f || scales[3] <= 0.0f) { + return Ort::GetApi().CreateStatus( + ORT_NOT_IMPLEMENTED, + "Resize requires positive NCHW scales with batch/channel scale 1"); + } + + std::vector output_shape = input_shape; + output_shape[2] = + static_cast(std::floor(input_shape[2] * scales[2])); + output_shape[3] = + static_cast(std::floor(input_shape[3] * scales[3])); + if (output_shape[2] <= 0 || output_shape[3] <= 0) { + return Ort::GetApi().CreateStatus(ORT_NOT_IMPLEMENTED, + "Resize output must be non-empty"); + } + if (!IsGpuMemory(input.GetTensorMemoryInfo())) { + return Ort::GetApi().CreateStatus(ORT_NOT_IMPLEMENTED, + "Resize requires MUSA input"); + } + Ort::UnownedValue output = ctx.GetOutput(0, output_shape); + if (!IsGpuMemory(output.GetTensorMemoryInfo())) { + return Ort::GetApi().CreateStatus(ORT_NOT_IMPLEMENTED, + "Resize requires MUSA output"); + } + + return LaunchStatus(LaunchMusaResizeNearest2DFloat( + input.GetTensorData(), output.GetTensorMutableData(), + input_shape[0], input_shape[1], input_shape[2], input_shape[3], + output_shape[2], output_shape[3], scales[2], scales[3], + GetComputeStream(ctx))); +} + +} // namespace + +ONNX_OPERATOR_VERSIONED_KERNEL_EX(Resize, kOnnxDomain, 13, 17, + (Ort::KernelDefBuilder() + .AddTypeConstraint("T1", + FloatTensorTypes()) + .SetInputMemType(1, OrtMemTypeCPUInput) + .SetInputMemType(2, OrtMemTypeCPUInput) + .SetInputMemType(3, OrtMemTypeCPUInput)), + Resize) diff --git a/musa/ep/src/kernels/tensor/resize_impl.h b/musa/ep/src/kernels/tensor/resize_impl.h new file mode 100644 index 0000000..52e20cd --- /dev/null +++ b/musa/ep/src/kernels/tensor/resize_impl.h @@ -0,0 +1,11 @@ +#pragma once + +#include + +#include + +musaError_t LaunchMusaResizeNearest2DFloat( + const float* input, float* output, int64_t batch, int64_t channels, + int64_t input_height, int64_t input_width, int64_t output_height, + int64_t output_width, float scale_height, float scale_width, + musaStream_t stream); diff --git a/musa/ep/src/kernels/tensor/resize_impl.mu b/musa/ep/src/kernels/tensor/resize_impl.mu new file mode 100644 index 0000000..279c8eb --- /dev/null +++ b/musa/ep/src/kernels/tensor/resize_impl.mu @@ -0,0 +1,51 @@ +#include "shared_inc/musa_kernel_common.mu.h" +#include "tensor/resize_impl.h" + +namespace { + +__global__ void ResizeNearest2DFloatKernel( + const float* input, float* output, int64_t total_elements, int64_t channels, + int64_t input_height, int64_t input_width, int64_t output_height, + int64_t output_width, float scale_height, float scale_width) { + const int64_t thread_id = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + const int64_t total_threads = static_cast(gridDim.x) * blockDim.x; + for (int64_t output_index = thread_id; output_index < total_elements; + output_index += total_threads) { + int64_t remaining = output_index; + const int64_t output_x = remaining % output_width; + remaining /= output_width; + const int64_t output_y = remaining % output_height; + remaining /= output_height; + const int64_t channel = remaining % channels; + const int64_t batch = remaining / channels; + + int64_t input_x = static_cast(output_x / scale_width); + int64_t input_y = static_cast(output_y / scale_height); + input_x = input_x < input_width ? input_x : input_width - 1; + input_y = input_y < input_height ? input_y : input_height - 1; + const int64_t input_index = + ((batch * channels + channel) * input_height + input_y) * input_width + + input_x; + output[output_index] = input[input_index]; + } +} + +} // namespace + +musaError_t LaunchMusaResizeNearest2DFloat( + const float* input, float* output, int64_t batch, int64_t channels, + int64_t input_height, int64_t input_width, int64_t output_height, + int64_t output_width, float scale_height, float scale_width, + musaStream_t stream) { + const int64_t total_elements = + batch * channels * output_height * output_width; + if (total_elements == 0) { + return musaSuccess; + } + ResizeNearest2DFloatKernel<<>>( + input, output, total_elements, channels, input_height, input_width, + output_height, output_width, scale_height, scale_width); + return musaGetLastError(); +} diff --git a/test/ops/test_concat.py b/test/ops/test_concat.py index 181883c..43e8bc7 100644 --- a/test/ops/test_concat.py +++ b/test/ops/test_concat.py @@ -23,6 +23,7 @@ float32_to_bfloat16_bits, run_and_compare, run_model_and_compare, + run_model_and_compare_with_cpu_fallback, run_with_iobinding, ) @@ -252,3 +253,32 @@ def test_concat_small_int64_shape_metadata_mixed_inputs(): name="concat_shape_metadata_graph", ) run_model_and_compare(model, {"dynamic_a": np.array(7, dtype=np.int64), "dynamic_b": np.array(11, dtype=np.int64)}, rtol=0, atol=0) + + +def test_concat_rank4_mixed_cpu_and_musa_inputs(): + nodes = [ + helper.make_node( + "Resize", + ["X", "", "scales"], + ["resized"], + mode="nearest", + coordinate_transformation_mode="asymmetric", + nearest_mode="floor", + ), + helper.make_node("Relu", ["Y"], ["relu"]), + helper.make_node("Concat", ["resized", "relu"], ["Z"], axis=1), + ] + scales = helper.make_tensor( + "scales", TensorProto.FLOAT, [4], [1.0, 1.0, 2.0, 2.0] + ) + x = np.arange(8, dtype=np.float32).reshape(1, 2, 2, 2) + y = np.linspace(-1.0, 1.0, 48, dtype=np.float32).reshape(1, 3, 4, 4) + model = build_graph_model( + nodes, + inputs={"X": x, "Y": y}, + outputs=[("Z", TensorProto.FLOAT)], + initializers=[scales], + opset=17, + name="concat_rank4_mixed_cpu_and_musa_inputs_graph", + ) + run_model_and_compare_with_cpu_fallback(model, {"X": x, "Y": y}) diff --git a/test/ops/test_resize.py b/test/ops/test_resize.py new file mode 100644 index 0000000..6a23d31 --- /dev/null +++ b/test/ops/test_resize.py @@ -0,0 +1,35 @@ +# Copyright (c) Moore Threads Technology Co., Ltd. All rights reserved. +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""End-to-end CPU-vs-MUSA tests for the supported Resize route.""" + +import numpy as np + +from op_test_utils import TensorProto, run_and_compare + + +def test_resize_opset17_float_nchw_nearest_asymmetric_floor_scales(): + x = np.arange(1 * 3 * 5 * 7, dtype=np.float32).reshape(1, 3, 5, 7) + scales = np.array([1.0, 1.0, 2.0, 2.0], dtype=np.float32) + run_and_compare( + "Resize", + inputs={"X": x, "roi": np.empty((0,), dtype=np.float32), "scales": scales}, + outputs=[("Y", TensorProto.FLOAT)], + attrs={ + "mode": "nearest", + "coordinate_transformation_mode": "asymmetric", + "nearest_mode": "floor", + }, + opset=17, + rtol=0, + atol=0, + ) From 9ad69dc94d84f9d0eed15ede937fbb05485f9026 Mon Sep 17 00:00:00 2001 From: "yichao.hu" Date: Mon, 12 Oct 2026 23:47:31 -0400 Subject: [PATCH 08/52] feat(conv): run fp32 NCDHW convolution through muDNN 3D convolution (NCDHW) had no MUSA route, so 3D models fell back or failed. Route fp32 NCDHW Conv through muDNN, reusing the grouped/depthwise paths. --- musa/ep/src/kernels/nn/conv.cc | 166 +++++++++++++++++++++------------ test/ops/test_conv.py | 38 +++++++- 2 files changed, 141 insertions(+), 63 deletions(-) diff --git a/musa/ep/src/kernels/nn/conv.cc b/musa/ep/src/kernels/nn/conv.cc index 8b3e2aa..1753dd0 100644 --- a/musa/ep/src/kernels/nn/conv.cc +++ b/musa/ep/src/kernels/nn/conv.cc @@ -11,21 +11,29 @@ // See the License for the specific language governing permissions and // limitations under the License. +#include + #include "nn/conv_impl.h" #include "shared_inc/blas_utils.h" #include "shared_inc/op_kernel_common.h" namespace { -std::vector NormalizePair(std::vector values, - int64_t default_value) { - if (values.empty()) { - return {default_value, default_value}; +bool SetMudnnConvTensor(::musa::dnn::Tensor& tensor, const void* data, + const std::vector& shape) { + if (shape.size() != 4 && shape.size() != 5) { + return false; } - if (values.size() == 1) { - return {values[0], values[0]}; + if (tensor.SetAddr(data) != ::musa::dnn::Status::SUCCESS || + tensor.SetType(::musa::dnn::Tensor::Type::FLOAT) != + ::musa::dnn::Status::SUCCESS || + tensor.SetFormat(shape.size() == 5 ? ::musa::dnn::Tensor::Format::NCDHW + : ::musa::dnn::Tensor::Format::NCHW) != + ::musa::dnn::Status::SUCCESS) { + return false; } - return {values[0], values[1]}; + return tensor.SetNdInfo(static_cast(shape.size()), shape.data()) == + ::musa::dnn::Status::SUCCESS; } class Conv : public OpKernelBase { @@ -33,18 +41,18 @@ class Conv : public OpKernelBase { Conv(const OrtKernelInfo* info, void* /*state*/) { Ort::ConstKernelInfo kernel_info(info); pads_ = AttrsOrEmpty(kernel_info, "pads"); - strides_ = NormalizePair(AttrsOrEmpty(kernel_info, "strides"), 1); - dilations_ = NormalizePair(AttrsOrEmpty(kernel_info, "dilations"), 1); + strides_ = AttrsOrEmpty(kernel_info, "strides"); + dilations_ = AttrsOrEmpty(kernel_info, "dilations"); group_ = AttrOrDefault(kernel_info, "group", 1); } OrtStatus* Compute(Ort::KernelContext& ctx) const; private: - bool TryMudnnConv(Ort::KernelContext& ctx, - const std::vector& x_shape, - const std::vector& w_shape, - const std::vector& y_shape, - Ort::UnownedValue output) const; + bool TryMudnnConv( + Ort::KernelContext& ctx, const std::vector& x_shape, + const std::vector& w_shape, const std::vector& y_shape, + const std::vector& pads, const std::vector& strides, + const std::vector& dilations, Ort::UnownedValue output) const; std::vector pads_; std::vector strides_; @@ -52,14 +60,18 @@ class Conv : public OpKernelBase { int64_t group_ = 1; }; -bool Conv::TryMudnnConv(Ort::KernelContext& ctx, - const std::vector& x_shape, - const std::vector& w_shape, - const std::vector& y_shape, - Ort::UnownedValue output) const { +bool Conv::TryMudnnConv( + Ort::KernelContext& ctx, const std::vector& x_shape, + const std::vector& w_shape, const std::vector& y_shape, + const std::vector& pads, const std::vector& strides, + const std::vector& dilations, Ort::UnownedValue output) const { // Conv1D exported as NCHW Conv with H=1 and kH=1 is covered by the custom // device fallback; avoid a muDNN path that is unstable for this layout. - if (x_shape[2] == 1 && w_shape[2] == 1) { + if (x_shape.size() == 4 && x_shape[2] == 1 && w_shape[2] == 1) { + return false; + } + if (x_shape.size() == 5 && + std::getenv("ORT_MUSA_CONV3D_FORCE_MUDNN_FAILURE_FOR_TEST") != nullptr) { return false; } @@ -91,25 +103,27 @@ bool Conv::TryMudnnConv(Ort::KernelContext& ctx, ::musa::dnn::Tensor x_tensor; ::musa::dnn::Tensor w_tensor; ::musa::dnn::Tensor y_tensor; - if (!SetMudnnTensor(x_tensor, x.GetTensorRawData(), x_shape, - ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT) || - !SetMudnnTensor(w_tensor, w.GetTensorRawData(), w_shape, - ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT) || - !SetMudnnTensor(y_tensor, output.GetTensorMutableRawData(), y_shape, - ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT)) { + if (!SetMudnnConvTensor(x_tensor, x.GetTensorRawData(), x_shape) || + !SetMudnnConvTensor(w_tensor, w.GetTensorRawData(), w_shape) || + !SetMudnnConvTensor(y_tensor, output.GetTensorMutableRawData(), + y_shape)) { return false; } - const int pad[2] = {static_cast(pads_[0]), static_cast(pads_[1])}; - const int stride[2] = {static_cast(strides_[0]), - static_cast(strides_[1])}; - const int dilation[2] = {static_cast(dilations_[0]), - static_cast(dilations_[1])}; + const size_t spatial_rank = x_shape.size() - 2; + std::vector pad(spatial_rank); + std::vector stride(spatial_rank); + std::vector dilation(spatial_rank); + for (size_t i = 0; i < spatial_rank; ++i) { + pad[i] = static_cast(pads[i]); + stride[i] = static_cast(strides[i]); + dilation[i] = static_cast(dilations[i]); + } ::musa::dnn::Convolution conv; if (conv.SetGroups(static_cast(group_)) != ::musa::dnn::Status::SUCCESS || - conv.SetNdInfo(2, pad, stride, dilation) != - ::musa::dnn::Status::SUCCESS || + conv.SetNdInfo(static_cast(spatial_rank), pad.data(), stride.data(), + dilation.data()) != ::musa::dnn::Status::SUCCESS || conv.SetComputeMode(::musa::dnn::Convolution::ComputeMode::TENSOR) != ::musa::dnn::Status::SUCCESS) { return false; @@ -127,8 +141,9 @@ bool Conv::TryMudnnConv(Ort::KernelContext& ctx, Ort::ConstValue b = ctx.GetInput(2); ::musa::dnn::Tensor bias_tensor; - if (!SetMudnnTensor(bias_tensor, b.GetTensorRawData(), {1, w_shape[0], 1, 1}, - ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT)) { + std::vector bias_shape(x_shape.size(), 1); + bias_shape[1] = w_shape[0]; + if (!SetMudnnConvTensor(bias_tensor, b.GetTensorRawData(), bias_shape)) { return false; } @@ -153,9 +168,10 @@ OrtStatus* Conv::Compute(Ort::KernelContext& ctx) const { } auto x_shape = x_info.GetShape(); auto w_shape = w_info.GetShape(); - if (x_shape.size() != 4 || w_shape.size() != 4) { - return Ort::GetApi().CreateStatus(ORT_NOT_IMPLEMENTED, - "Conv only supports 2D NCHW tensors"); + if ((x_shape.size() != 4 && x_shape.size() != 5) || + w_shape.size() != x_shape.size()) { + return Ort::GetApi().CreateStatus( + ORT_NOT_IMPLEMENTED, "Conv only supports 2D NCHW or 3D NCDHW tensors"); } if (group_ <= 0 || x_shape[1] != w_shape[1] * group_ || w_shape[0] % group_ != 0) { @@ -164,33 +180,59 @@ OrtStatus* Conv::Compute(Ort::KernelContext& ctx) const { "Conv group must divide output channels and match input channels"); } + const size_t spatial_rank = x_shape.size() - 2; std::vector pads = pads_; if (pads.empty()) { - pads = {0, 0, 0, 0}; + pads.resize(spatial_rank * 2, 0); } - if (pads.size() != 4 || pads[0] != pads[2] || pads[1] != pads[3]) { - return Ort::GetApi().CreateStatus(ORT_NOT_IMPLEMENTED, - "Conv only supports symmetric 2D pads"); + std::vector strides = strides_; + if (strides.empty()) { + strides.resize(spatial_rank, 1); } - - const int64_t out_h = - (x_shape[2] + pads[0] + pads[2] - dilations_[0] * (w_shape[2] - 1) - 1) / - strides_[0] + - 1; - const int64_t out_w = - (x_shape[3] + pads[1] + pads[3] - dilations_[1] * (w_shape[3] - 1) - 1) / - strides_[1] + - 1; - if (out_h < 0 || out_w < 0) { - return Ort::GetApi().CreateStatus(ORT_INVALID_ARGUMENT, - "Conv invalid output shape"); + std::vector dilations = dilations_; + if (dilations.empty()) { + dilations.resize(spatial_rank, 1); + } + if (pads.size() != spatial_rank * 2 || strides.size() != spatial_rank || + dilations.size() != spatial_rank) { + return Ort::GetApi().CreateStatus( + ORT_NOT_IMPLEMENTED, + "Conv spatial attributes do not match tensor rank"); + } + for (size_t i = 0; i < spatial_rank; ++i) { + if (pads[i] != pads[i + spatial_rank]) { + return Ort::GetApi().CreateStatus( + ORT_NOT_IMPLEMENTED, "Conv only supports symmetric spatial pads"); + } + if (strides[i] <= 0 || dilations[i] <= 0) { + return Ort::GetApi().CreateStatus( + ORT_INVALID_ARGUMENT, "Conv strides and dilations must be positive"); + } } - std::vector y_shape = {x_shape[0], w_shape[0], out_h, out_w}; + std::vector y_shape = {x_shape[0], w_shape[0]}; + for (size_t i = 0; i < spatial_rank; ++i) { + const int64_t output_dim = + (x_shape[i + 2] + pads[i] + pads[i + spatial_rank] - + dilations[i] * (w_shape[i + 2] - 1) - 1) / + strides[i] + + 1; + if (output_dim < 0) { + return Ort::GetApi().CreateStatus(ORT_INVALID_ARGUMENT, + "Conv invalid output shape"); + } + y_shape.push_back(output_dim); + } Ort::UnownedValue output = ctx.GetOutput(0, y_shape); - if (TryMudnnConv(ctx, x_shape, w_shape, y_shape, output)) { + if (TryMudnnConv(ctx, x_shape, w_shape, y_shape, pads, strides, dilations, + output)) { return nullptr; } + if (spatial_rank == 3) { + return Ort::GetApi().CreateStatus( + ORT_NOT_IMPLEMENTED, + "Conv3D requires the muDNN path; custom fallback is 2D-only"); + } if (!IsGpuMemory(x.GetTensorMemoryInfo()) || !IsGpuMemory(w.GetTensorMemoryInfo()) || @@ -221,14 +263,14 @@ OrtStatus* Conv::Compute(Ort::KernelContext& ctx) const { params.m = w_shape[0]; params.kernel_h = w_shape[2]; params.kernel_w = w_shape[3]; - params.out_h = out_h; - params.out_w = out_w; + params.out_h = y_shape[2]; + params.out_w = y_shape[3]; params.pad_h = pads[0]; params.pad_w = pads[1]; - params.stride_h = strides_[0]; - params.stride_w = strides_[1]; - params.dilation_h = dilations_[0]; - params.dilation_w = dilations_[1]; + params.stride_h = strides[0]; + params.stride_w = strides[1]; + params.dilation_h = dilations[0]; + params.dilation_w = dilations[1]; params.total_elements = NumElements(y_shape); return LaunchStatus(LaunchMusaConv2DFloatKernel( x.GetTensorData(), w.GetTensorData(), bias, diff --git a/test/ops/test_conv.py b/test/ops/test_conv.py index 8f60dde..84823bb 100644 --- a/test/ops/test_conv.py +++ b/test/ops/test_conv.py @@ -13,8 +13,9 @@ """End-to-end CPU-vs-MUSA test for the Conv operator.""" import numpy as np +import pytest -from op_test_utils import TensorProto, run_and_compare +from op_test_utils import TensorProto, build_model, run, run_and_compare def test_conv_opset11_float_nchw(): @@ -77,6 +78,41 @@ def test_conv_opset11_float_nchw_depthwise_fallback_with_bias(): ) +def test_conv_opset11_float_ncdhw_with_bias(): + x = np.random.default_rng(13).standard_normal((1, 2, 4, 5, 6)).astype(np.float32) + w = np.random.default_rng(14).standard_normal((3, 2, 3, 3, 3)).astype(np.float32) + b = np.random.default_rng(15).standard_normal((3,)).astype(np.float32) + run_and_compare( + "Conv", + inputs={"X": x, "W": w, "B": b}, + outputs=[("Y", TensorProto.FLOAT)], + attrs={ + "pads": [1, 1, 1, 1, 1, 1], + "strides": [1, 2, 1], + "dilations": [1, 1, 1], + }, + opset=11, + rtol=1e-4, + atol=1e-4, + ) + + +def test_conv_opset11_float_ncdhw_mudnn_failure_is_explicit(monkeypatch): + x = np.random.default_rng(16).standard_normal((1, 2, 3, 4, 5)).astype(np.float32) + w = np.random.default_rng(17).standard_normal((3, 2, 1, 3, 3)).astype(np.float32) + model = build_model( + "Conv", + inputs={"X": x, "W": w}, + outputs=[("Y", TensorProto.FLOAT)], + attrs={"pads": [0, 1, 1, 0, 1, 1]}, + opset=11, + ) + + monkeypatch.setenv("ORT_MUSA_CONV3D_FORCE_MUDNN_FAILURE_FOR_TEST", "1") + with pytest.raises(Exception, match="Conv3D requires the muDNN path"): + run(model, {"X": x, "W": w}, use_musa=True) + + def test_conv_lowered_conv1d_h1_k1(): x = np.random.default_rng(0).standard_normal((2, 64, 1, 8)).astype(np.float32) w = np.random.default_rng(1).standard_normal((16, 64, 1, 1)).astype(np.float32) From 5bcd62c3ce78e0ce41e3d153f580571057232935 Mon Sep 17 00:00:00 2001 From: "yichao.hu" Date: Mon, 12 Oct 2026 23:47:31 -0400 Subject: [PATCH 09/52] fix(elementwise): stage CPU boundary inputs for unary device kernels Unary kernels read inputs produced on the CPU EP as device pointers, which gave wrong results or page faults. Stage CPU-resident inputs onto the device on the compute stream before launching the kernel. --- .../kernels/shared_inc/elementwise_common.h | 18 +++++++-- test/ops/op_test_utils.py | 3 ++ test/ops/test_relu.py | 37 ++++++++++++++++++- 3 files changed, 53 insertions(+), 5 deletions(-) diff --git a/musa/ep/src/kernels/shared_inc/elementwise_common.h b/musa/ep/src/kernels/shared_inc/elementwise_common.h index 435c96d..48c193c 100644 --- a/musa/ep/src/kernels/shared_inc/elementwise_common.h +++ b/musa/ep/src/kernels/shared_inc/elementwise_common.h @@ -173,19 +173,29 @@ inline OrtStatus* UnaryDeviceCompute(Ort::KernelContext& ctx, float alpha = 0.0f) { Ort::ConstValue input = ctx.GetInput(0); MusaElementType musa_elem_type; - if (!ToMusaElementType(elem_type, musa_elem_type) || - !IsGpuMemory(input.GetTensorMemoryInfo())) { + if (!ToMusaElementType(elem_type, musa_elem_type)) { return UnsupportedDeviceElementwiseStatus(op_name, elem_type); } + const bool input_is_gpu = IsGpuMemory(input.GetTensorMemoryInfo()); + if (!input_is_gpu && device_op == MusaUnaryOp::Relu && + std::getenv("ORT_MUSA_TEST_FAIL_RELU_CPU_INPUT") != nullptr) { + return Ort::GetApi().CreateStatus(ORT_EP_FAIL, + "test-only Relu CPU-input route reached"); + } + + musaStream_t stream = GetComputeStream(ctx); + DeviceInputBuffer input_buffer; + RETURN_IF_ERROR(input_buffer.Bind(input, stream)); + Ort::UnownedValue y = ctx.GetOutput(0, shape); if (!IsGpuMemory(y.GetTensorMemoryInfo())) { return UnsupportedDeviceElementwiseStatus(op_name, elem_type); } musaError_t status = LaunchMusaUnaryKernel( - input.GetTensorRawData(), y.GetTensorMutableRawData(), NumElements(shape), - device_op, alpha, musa_elem_type, GetComputeStream(ctx)); + input_buffer.data(), y.GetTensorMutableRawData(), NumElements(shape), + device_op, alpha, musa_elem_type, stream); if (status == musaErrorNotSupported) { return UnsupportedDeviceElementwiseStatus(op_name, elem_type); } diff --git a/test/ops/op_test_utils.py b/test/ops/op_test_utils.py index efb3eff..826880e 100644 --- a/test/ops/op_test_utils.py +++ b/test/ops/op_test_utils.py @@ -345,6 +345,7 @@ def run_model_and_compare_with_cpu_fallback( *, rtol: float = 1e-3, atol: float = 1e-4, + disable_ep_error_fallback: bool = False, ) -> list[np.ndarray]: """Run a full ONNX model on CPU and MUSA with ORT CPU fallback enabled.""" cpu_outputs = run(model_bytes, feeds, use_musa=False) @@ -358,6 +359,8 @@ def run_model_and_compare_with_cpu_fallback( so = ort.SessionOptions() so.add_provider_for_devices(devices, {}) session = ort.InferenceSession(model_bytes, sess_options=so) + if disable_ep_error_fallback: + session.disable_fallback() musa_outputs = session.run(None, dict(feeds)) assert len(cpu_outputs) == len( diff --git a/test/ops/test_relu.py b/test/ops/test_relu.py index 3fb464b..b889388 100644 --- a/test/ops/test_relu.py +++ b/test/ops/test_relu.py @@ -13,8 +13,15 @@ """End-to-end CPU-vs-MUSA test for the Relu operator.""" import numpy as np +import pytest +from onnx import helper -from op_test_utils import TensorProto, run_and_compare +from op_test_utils import ( + TensorProto, + build_graph_model, + run_and_compare, + run_model_and_compare_with_cpu_fallback, +) def test_relu_float(): @@ -30,3 +37,31 @@ def test_relu_float_5d(): def test_relu_float_matrix_with_negatives(): x = np.array([[-3.5, -0.0, 2.0], [0.0, -7.25, 9.5]], dtype=np.float32) run_and_compare("Relu", inputs={"X": x}, outputs=[("Y", TensorProto.FLOAT)]) + + +def test_relu_float16_accepts_cpu_produced_input(monkeypatch): + x = np.array([[-3.5, -0.0, 0.0, 2.0, 9.5]], dtype=np.float16) + model = build_graph_model( + [ + helper.make_node("Sin", ["X"], ["cpu_output"]), + helper.make_node("Relu", ["cpu_output"], ["Y"]), + ], + inputs={"X": x}, + outputs=[("Y", TensorProto.FLOAT16)], + name="relu_float16_cpu_produced_input_graph", + ) + + monkeypatch.setenv("ORT_MUSA_TEST_FAIL_RELU_CPU_INPUT", "1") + with pytest.raises(Exception, match="test-only Relu CPU-input route"): + run_model_and_compare_with_cpu_fallback( + model, {"X": x}, disable_ep_error_fallback=True + ) + + monkeypatch.delenv("ORT_MUSA_TEST_FAIL_RELU_CPU_INPUT") + run_model_and_compare_with_cpu_fallback( + model, + {"X": x}, + rtol=0, + atol=0, + disable_ep_error_fallback=True, + ) From 42addf9617eafeb11f3f3020d073f13c547c08d6 Mon Sep 17 00:00:00 2001 From: "yichao.hu" Date: Mon, 12 Oct 2026 23:47:31 -0400 Subject: [PATCH 10/52] fix(fusion): stage CPU-produced inputs for SiluFusion and Reshape SiluFusion and Reshape could receive tensors produced on the CPU EP and read them as device memory. Stage such inputs onto the device before use. --- musa/docs/fusion/silu.md | 9 +++-- musa/ep/src/fusion/silu_fusion.cc | 30 +++++++++++------ musa/ep/src/kernels/tensor/reshape.cc | 10 ++++-- test/fusion/test_silu_fusion.py | 30 +++++++++++++++++ test/ops/test_split.py | 48 +++++++++++++++++++++++++++ 5 files changed, 113 insertions(+), 14 deletions(-) diff --git a/musa/docs/fusion/silu.md b/musa/docs/fusion/silu.md index 403aaa0..29801d3 100644 --- a/musa/docs/fusion/silu.md +++ b/musa/docs/fusion/silu.md @@ -13,7 +13,7 @@ This file is generated from the current C++ fusion source and matching fusion te - Runtime compute: `SiluFusionCompute` - Runtime implementation: `musa/ep/src/fusion/silu_fusion.cc` - `drop_constant_initializers`: `false` -- Before-graph topology source: `musa/ep/src/fusion/silu_fusion_matcher.cc` finder source +- Before-graph topology source: `test/fusion/test_silu_fusion.py::test_silu_fusion_float16_accepts_cpu_produced_input` ## Extracted ONNX Ops @@ -24,7 +24,12 @@ This file is generated from the current C++ fusion source and matching fusion te ```mermaid flowchart LR subgraph Before[Before fusion] - B0["Mul / Sigmoid"] + B0["Sin"] + B1["Sigmoid"] + B2["Mul"] + B0 --> B1 + B0 --> B2 + B1 --> B2 end subgraph After[After fusion] A0["MUSA fused node"] diff --git a/musa/ep/src/fusion/silu_fusion.cc b/musa/ep/src/fusion/silu_fusion.cc index 609e190..419e9ae 100644 --- a/musa/ep/src/fusion/silu_fusion.cc +++ b/musa/ep/src/fusion/silu_fusion.cc @@ -3,6 +3,7 @@ #include "fusion/silu_fusion.h" +#include #include #include #include @@ -10,6 +11,7 @@ #include "graph/graph_utils.h" #include "kernels/activation/silu_impl.h" #include "kernels/shared_inc/blas_utils.h" +#include "kernels/shared_inc/kernel_memory.h" #include "kernels/shared_inc/op_kernel_common.h" namespace musa_ep { @@ -63,12 +65,11 @@ bool FindSiluGraphInput(Ort::ConstGraph graph, std::string& input_name) { return true; } -bool TryMudnnSilu(Ort::KernelContext& ctx, Ort::ConstValue input, +bool TryMudnnSilu(Ort::KernelContext& ctx, const void* input_data, Ort::UnownedValue output, const std::vector& shape, ONNXTensorElementDataType elem_type) { constexpr size_t kMudnnMaxElementwiseRank = 5; if (shape.size() > kMudnnMaxElementwiseRank || - !IsGpuMemory(input.GetTensorMemoryInfo()) || !IsGpuMemory(output.GetTensorMemoryInfo())) { return false; } @@ -82,8 +83,7 @@ bool TryMudnnSilu(Ort::KernelContext& ctx, Ort::ConstValue input, ::musa::dnn::Tensor input_tensor; ::musa::dnn::Tensor output_tensor; - if (!SetMudnnTensor(input_tensor, input.GetTensorRawData(), shape, - elem_type) || + if (!SetMudnnTensor(input_tensor, input_data, shape, elem_type) || !SetMudnnTensor(output_tensor, output.GetTensorMutableRawData(), shape, elem_type)) { return false; @@ -112,13 +112,23 @@ OrtStatus* SiluFusionCompute::Compute(OrtKernelContext* kernel_context) const { } Ort::UnownedValue output = ctx.GetOutput(0, shape); - if (!IsGpuMemory(input.GetTensorMemoryInfo()) || - !IsGpuMemory(output.GetTensorMemoryInfo())) { + if (!IsGpuMemory(output.GetTensorMemoryInfo())) { + return Ort::GetApi().CreateStatus(ORT_NOT_IMPLEMENTED, + "Silu fusion requires MUSA output"); + } + + const bool input_is_gpu = IsGpuMemory(input.GetTensorMemoryInfo()); + if (!input_is_gpu && + std::getenv("ORT_MUSA_TEST_FAIL_SILU_CPU_INPUT") != nullptr) { return Ort::GetApi().CreateStatus( - ORT_NOT_IMPLEMENTED, "Silu fusion requires MUSA input/output"); + ORT_EP_FAIL, "test-only SiluFusion CPU-input route reached"); } - if (TryMudnnSilu(ctx, input, output, shape, elem_type)) { + musaStream_t stream = GetComputeStream(ctx); + DeviceInputBuffer input_buffer; + RETURN_IF_ERROR(input_buffer.Bind(input, stream)); + + if (TryMudnnSilu(ctx, input_buffer.data(), output, shape, elem_type)) { return nullptr; } @@ -127,8 +137,8 @@ OrtStatus* SiluFusionCompute::Compute(OrtKernelContext* kernel_context) const { return UnsupportedDeviceElementwiseStatus("SiluFusion", elem_type); } return LaunchStatus(LaunchMusaSiluKernel( - input.GetTensorRawData(), output.GetTensorMutableRawData(), - NumElements(shape), musa_elem_type, GetComputeStream(ctx))); + input_buffer.data(), output.GetTensorMutableRawData(), + NumElements(shape), musa_elem_type, stream)); } catch (const Ort::Exception& ex) { Ort::Status status(ex); return status.release(); diff --git a/musa/ep/src/kernels/tensor/reshape.cc b/musa/ep/src/kernels/tensor/reshape.cc index 49c9daf..b96a20d 100644 --- a/musa/ep/src/kernels/tensor/reshape.cc +++ b/musa/ep/src/kernels/tensor/reshape.cc @@ -11,6 +11,8 @@ // See the License for the specific language governing permissions and // limitations under the License. +#include + #include "shared_inc/op_kernel_common.h" namespace { @@ -59,6 +61,11 @@ std::vector ResolveReshapeOutputShape( OrtStatus* Reshape::Compute(Ort::KernelContext& ctx) const { Ort::ConstValue input0 = ctx.GetInput(0); + if (!IsGpuMemory(input0.GetTensorMemoryInfo()) && + std::getenv("ORT_MUSA_TEST_FAIL_RESHAPE_CPU_INPUT") != nullptr) { + return Ort::GetApi().CreateStatus( + ORT_EP_FAIL, "test-only Reshape CPU-input route reached"); + } auto shape0 = input0.GetTensorTypeAndShapeInfo().GetShape(); std::vector out_shape = ResolveReshapeOutputShape(shape0, ReadIntTensor(ctx, 1), allowzero_); @@ -103,8 +110,7 @@ ONNX_OPERATOR_VERSIONED_KERNEL_EX( Reshape, kOnnxDomain, 13, 19, (Ort::KernelDefBuilder() .AddTypeConstraint("T", AllFixedSizeTensorTypes()) - .SetInputMemType(1, OrtMemTypeCPUInput) - .AddInputOutputAlias(0, 0)), + .SetInputMemType(1, OrtMemTypeCPUInput)), Reshape) ONNX_OPERATOR_VERSIONED_KERNEL_EX( diff --git a/test/fusion/test_silu_fusion.py b/test/fusion/test_silu_fusion.py index dd42c6d..9cdcc0d 100644 --- a/test/fusion/test_silu_fusion.py +++ b/test/fusion/test_silu_fusion.py @@ -17,6 +17,7 @@ float32_to_bfloat16_bits, musa_devices, run, + run_model_and_compare_with_cpu_fallback, run_with_iobinding, ) @@ -112,6 +113,35 @@ def test_silu_fusion_float_dtypes(np_dtype, tensor_type, rtol, atol, swapped, tm _assert_silu_fused(model, {"X": x}, tmp_path, f"silu_{np_dtype}_{swapped}") +def test_silu_fusion_float16_accepts_cpu_produced_input(monkeypatch): + x = np.array([[-3.5, -1.0, -0.0, 0.0, 1.0, 3.5]], dtype=np.float16) + model = build_graph_model( + [ + helper.make_node("Sin", ["X"], ["cpu_output"]), + helper.make_node("Sigmoid", ["cpu_output"], ["S"]), + helper.make_node("Mul", ["cpu_output", "S"], ["Y"]), + ], + {"X": x}, + [("Y", TensorProto.FLOAT16)], + name="silu_fusion_float16_cpu_produced_input_graph", + ) + + monkeypatch.setenv("ORT_MUSA_TEST_FAIL_SILU_CPU_INPUT", "1") + with pytest.raises(Exception, match="test-only SiluFusion CPU-input route"): + run_model_and_compare_with_cpu_fallback( + model, {"X": x}, disable_ep_error_fallback=True + ) + + monkeypatch.delenv("ORT_MUSA_TEST_FAIL_SILU_CPU_INPUT") + run_model_and_compare_with_cpu_fallback( + model, + {"X": x}, + rtol=2e-3, + atol=2e-3, + disable_ep_error_fallback=True, + ) + + def test_silu_fusion_bfloat16(tmp_path): x_f32 = np.array([[-8.0, -1.25, -0.0, 0.0, 1.25, 8.0]], dtype=np.float32) x = float32_to_bfloat16_bits(x_f32) diff --git a/test/ops/test_split.py b/test/ops/test_split.py index e9b71bc..f407e97 100644 --- a/test/ops/test_split.py +++ b/test/ops/test_split.py @@ -14,12 +14,15 @@ import numpy as np import pytest +from onnx import helper from op_test_utils import ( TensorProto, + build_graph_model, build_model_with_input_types, float32_to_bfloat16_bits, run_and_compare, + run_model_and_compare_with_cpu_fallback, run_with_iobinding, ) @@ -161,6 +164,51 @@ def test_split_equal_six_outputs_fp16(): ) +def test_split_equal_two_outputs_fp16_large_inner(): + x = np.random.default_rng(9).standard_normal((1, 4, 1029)).astype(np.float16) + split = np.array([2, 2], dtype=np.int64) + outputs = [("Y0", TensorProto.FLOAT16), ("Y1", TensorProto.FLOAT16)] + run_and_compare( + "Split", + inputs={"X": x, "split": split}, + outputs=outputs, + attrs={"axis": 1}, + ) + + +def test_split_after_reshape_float16_cpu_produced_input(monkeypatch): + x = np.random.default_rng(10).standard_normal((1, 4, 1029)).astype(np.float16) + shape = np.array([1, 4, 1029], dtype=np.int64) + split = np.array([2, 2], dtype=np.int64) + model = build_graph_model( + [ + helper.make_node("Sin", ["X"], ["cpu_output"]), + helper.make_node("Reshape", ["cpu_output", "shape"], ["reshaped"]), + helper.make_node("Split", ["reshaped", "split"], ["Y0", "Y1"], axis=1), + ], + {"X": x, "shape": shape, "split": split}, + [("Y0", TensorProto.FLOAT16), ("Y1", TensorProto.FLOAT16)], + name="split_after_reshape_float16_cpu_produced_input_graph", + ) + + monkeypatch.setenv("ORT_MUSA_TEST_FAIL_RESHAPE_CPU_INPUT", "1") + with pytest.raises(Exception, match="test-only Reshape CPU-input route"): + run_model_and_compare_with_cpu_fallback( + model, + {"X": x, "shape": shape, "split": split}, + disable_ep_error_fallback=True, + ) + + monkeypatch.delenv("ORT_MUSA_TEST_FAIL_RESHAPE_CPU_INPUT") + run_model_and_compare_with_cpu_fallback( + model, + {"X": x, "shape": shape, "split": split}, + rtol=0, + atol=0, + disable_ep_error_fallback=True, + ) + + def test_split_equal_four_outputs_int64(): x = np.arange(4, dtype=np.int64) split = np.ones(4, dtype=np.int64) From 787293ec008d04191d7e2a0198e8c507cbcfe9b5 Mon Sep 17 00:00:00 2001 From: "yichao.hu" Date: Mon, 12 Oct 2026 23:47:31 -0400 Subject: [PATCH 11/52] fix(plugin): register MemcpyFromHost/MemcpyToHost kernels in the MUSA EP ORT inserts MemcpyFromHost/MemcpyToHost at CPU/device boundaries, but the plugin EP had no kernels for them, so host pointers reached device kernels and caused page faults across many models. Register both kernels, running the copy on the plugin's compute stream, and update the CPU-boundary probe tests to expect device inputs. Adds boundary-copy tests and refreshes supported_ops. --- musa/docs/supported_ops.md | 6 +- musa/ep/src/ep_data_transfer.cc | 23 ++++ musa/ep/src/ep_data_transfer.h | 11 ++ musa/ep/src/kernels/tensor/memcpy.cc | 96 +++++++++++++++++ test/fusion/test_silu_fusion.py | 10 +- test/ops/test_memcpy_boundary.py | 150 +++++++++++++++++++++++++++ test/ops/test_relu.py | 11 +- test/ops/test_split.py | 14 +-- 8 files changed, 304 insertions(+), 17 deletions(-) create mode 100644 musa/ep/src/kernels/tensor/memcpy.cc create mode 100644 test/ops/test_memcpy_boundary.py diff --git a/musa/docs/supported_ops.md b/musa/docs/supported_ops.md index 4f535cd..2c88035 100644 --- a/musa/docs/supported_ops.md +++ b/musa/docs/supported_ops.md @@ -1,8 +1,8 @@ # Supported Ops -- Total supported ops: **90** -- Total registered kernels: **122** +- Total supported ops: **92** +- Total registered kernels: **124** Rows are grouped by operator/domain. Multiple opset ranges in a row preserve separate ORT kernel registrations. @@ -48,6 +48,8 @@ Rows are grouped by operator/domain. Multiple opset ranges in a row preserve sep | `MatMul` | `ai.onnx` | 13–19: `T`: float16, float32, float64, bfloat16 | `math/matmul.cc` | | `Max` | `ai.onnx` | 13–19: `T`: uint32, uint64, int32, int64, float16, float32, float64, bfloat16 | `math/max.cc` | | `MaxPool` | `ai.onnx` | 1–7: `T`: float16, float32, float64
8–11: `T`: float16, float32, float64
`I`: int64
12–19: `T`: float16, float32, float64, int8, uint8
`I`: int64 | `nn/pool.cc` | +| `MemcpyFromHost` | `ai.onnx` | 1–1: `T`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bfloat16, bool | `tensor/memcpy.cc` | +| `MemcpyToHost` | `ai.onnx` | 1–1: `T`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bfloat16, bool | `tensor/memcpy.cc` | | `Min` | `ai.onnx` | 13–19: `T`: uint32, uint64, int32, int64, float16, float32, float64, bfloat16 | `math/min.cc` | | `Mod` | `ai.onnx` | 10–12: `T`: int32, int64
13–19: `T`: int32, int64 | `math/mod.cc` | | `Mul` | `ai.onnx` | 13–13: `T`: uint32, uint64, int32, int64, float16, float32, float64, bfloat16
14–19: `T`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bfloat16 | `math/mul.cc` | diff --git a/musa/ep/src/ep_data_transfer.cc b/musa/ep/src/ep_data_transfer.cc index 0128516..ebe18e6 100644 --- a/musa/ep/src/ep_data_transfer.cc +++ b/musa/ep/src/ep_data_transfer.cc @@ -307,6 +307,29 @@ OrtStatus* CopyDeviceToHostBatch(MusaDataTransfer& impl, } // namespace +OrtStatus* MusaCopyTensorOnStream(OrtDataTransferImpl& impl_base, + const OrtValue* src_tensor, + OrtValue* dst_tensor, musaStream_t stream) { + auto& impl = static_cast(impl_base); + + const OrtMemoryDevice* src_device = + impl.ep_api_.Value_GetMemoryDevice(src_tensor); + const OrtMemoryDevice* dst_device = + impl.ep_api_.Value_GetMemoryDevice(dst_tensor); + + const void* src_data = nullptr; + void* dst_data = nullptr; + size_t bytes = 0; + RETURN_IF_ERROR(impl.ort_api_.GetTensorData(src_tensor, &src_data)); + RETURN_IF_ERROR(impl.ort_api_.GetTensorMutableData(dst_tensor, &dst_data)); + RETURN_IF_ERROR(impl.ort_api_.GetTensorSizeInBytes(src_tensor, &bytes)); + + // A single boundary transfer never batches, so pageable host memory may be + // bounced through the pinned pool to keep the copy on the caller's stream. + return CopyImpl(impl, src_device, dst_device, src_data, dst_data, bytes, + stream, /*allow_pageable_bounce*/ true); +} + bool ORT_API_CALL MusaDataTransfer::CanCopyImpl( const OrtDataTransferImpl* this_ptr, const OrtMemoryDevice* src_memory_device, diff --git a/musa/ep/src/ep_data_transfer.h b/musa/ep/src/ep_data_transfer.h index 0a9138f..8ea1d1c 100644 --- a/musa/ep/src/ep_data_transfer.h +++ b/musa/ep/src/ep_data_transfer.h @@ -13,6 +13,8 @@ #pragma once +#include + #include #include @@ -57,3 +59,12 @@ struct MusaDataTransfer : OrtDataTransferImpl { uint32_t vendor_id_; std::shared_ptr pinned_host_pool_; }; + +// Copies a single tensor across the host/device boundary using `impl`, ordering +// the transfer on `stream` (the caller's own MUSA compute stream). Shares the +// pageable-H2D staging and D2H ordering of MusaDataTransfer instead of issuing +// a raw musaMemcpy inline. Used by the MemcpyFromHost / MemcpyToHost kernels +// that ORT's MemcpyTransformer inserts for CPU/device placement boundaries. +OrtStatus* MusaCopyTensorOnStream(OrtDataTransferImpl& impl, + const OrtValue* src_tensor, + OrtValue* dst_tensor, musaStream_t stream); diff --git a/musa/ep/src/kernels/tensor/memcpy.cc b/musa/ep/src/kernels/tensor/memcpy.cc new file mode 100644 index 0000000..ed1182a --- /dev/null +++ b/musa/ep/src/kernels/tensor/memcpy.cc @@ -0,0 +1,96 @@ +// Copyright (c) Moore Threads Technology Co., Ltd. All rights reserved. +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include + +#include "ep_data_transfer.h" +#include "shared_inc/op_kernel_common.h" + +namespace { + +// MemcpyFromHost / MemcpyToHost boundary copies. +// +// ORT's MemcpyTransformer inserts these nodes whenever a tensor crosses a +// CPU/device placement boundary and assigns them to the device EP. If the EP +// does not register a kernel for the op, ORT falls back to the generic CPU EP +// memcpy (session_state.cc) and the copy destination stays in host memory -- +// the device-side consumer then rejects it (e.g. "Flatten requires MUSA device +// tensors"). Registering these kernels keeps the node on the plugin's device +// allocator (the MemcpyFromHost destination is allocated as GPU/DEFAULT) and +// runs the transfer on the plugin's own compute stream. +// +// Both directions share one implementation and defer the actual transfer to +// the EP's OrtDataTransferImpl (created by the factory and handed to kernels as +// their create state) so pageable-H2D staging and D2H stream ordering match the +// rest of the EP. Tensor-only, mirroring the CUDA plugin kernels: non-tensor +// values (sequences, optionals) do not match this type constraint and keep +// falling back to the CPU EP memcpy, which handles the sequence cases. +class MemcpyBoundary final : public OpKernelBase { + public: + MemcpyBoundary(const OrtKernelInfo* /*info*/, void* state) + : data_transfer_impl_(static_cast(state)) {} + + OrtStatus* Compute(Ort::KernelContext& ctx) const; + + private: + OrtDataTransferImpl* data_transfer_impl_ = nullptr; +}; + +OrtStatus* MemcpyBoundary::Compute(Ort::KernelContext& ctx) const { + Ort::ConstValue input = ctx.GetInput(0); + Ort::UnownedValue output = + ctx.GetOutput(0, input.GetTensorTypeAndShapeInfo().GetShape()); + + const void* src = input.GetTensorRawData(); + void* dst = output.GetTensorMutableRawData(); + const size_t bytes = input.GetTensorSizeInBytes(); + if (bytes == 0 || src == dst) { + return nullptr; + } + + if (std::getenv("ORT_MUSA_TEST_FAIL_MEMCPY_KERNEL") != nullptr) { + return Ort::GetApi().CreateStatus( + ORT_EP_FAIL, "test-only MUSA memcpy kernel route reached"); + } + + const bool src_is_gpu = IsGpuMemory(input.GetTensorMemoryInfo()); + const bool dst_is_gpu = IsGpuMemory(output.GetTensorMemoryInfo()); + if (src_is_gpu == dst_is_gpu) { + return Ort::GetApi().CreateStatus( + ORT_EP_FAIL, "Memcpy expects exactly one MUSA device tensor"); + } + + if (data_transfer_impl_ == nullptr) { + return Ort::GetApi().CreateStatus( + ORT_EP_FAIL, "Memcpy: data transfer implementation is missing"); + } + + return MusaCopyTensorOnStream(*data_transfer_impl_, input, output, + GetComputeStream(ctx)); +} + +} // namespace + +ONNX_OPERATOR_VERSIONED_KERNEL_EX( + MemcpyFromHost, kOnnxDomain, 1, 1, + (Ort::KernelDefBuilder() + .AddTypeConstraint("T", AllFixedSizeTensorTypes()) + .SetInputMemType(0, OrtMemTypeCPUInput)), + MemcpyBoundary) + +ONNX_OPERATOR_VERSIONED_KERNEL_EX( + MemcpyToHost, kOnnxDomain, 1, 1, + (Ort::KernelDefBuilder() + .AddTypeConstraint("T", AllFixedSizeTensorTypes()) + .SetOutputMemType(0, OrtMemTypeCPUOutput)), + MemcpyBoundary) diff --git a/test/fusion/test_silu_fusion.py b/test/fusion/test_silu_fusion.py index 9cdcc0d..3418247 100644 --- a/test/fusion/test_silu_fusion.py +++ b/test/fusion/test_silu_fusion.py @@ -127,10 +127,12 @@ def test_silu_fusion_float16_accepts_cpu_produced_input(monkeypatch): ) monkeypatch.setenv("ORT_MUSA_TEST_FAIL_SILU_CPU_INPUT", "1") - with pytest.raises(Exception, match="test-only SiluFusion CPU-input route"): - run_model_and_compare_with_cpu_fallback( - model, {"X": x}, disable_ep_error_fallback=True - ) + # The plugin now runs the MemcpyFromHost node, so the CPU-produced tensor + # reaches SiluFusion in device memory and the test-only CPU-input route must + # not be taken. + run_model_and_compare_with_cpu_fallback( + model, {"X": x}, disable_ep_error_fallback=True + ) monkeypatch.delenv("ORT_MUSA_TEST_FAIL_SILU_CPU_INPUT") run_model_and_compare_with_cpu_fallback( diff --git a/test/ops/test_memcpy_boundary.py b/test/ops/test_memcpy_boundary.py new file mode 100644 index 0000000..375f648 --- /dev/null +++ b/test/ops/test_memcpy_boundary.py @@ -0,0 +1,150 @@ +# Copyright (c) Moore Threads Technology Co., Ltd. All rights reserved. +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""CPU/device boundary copies (MemcpyFromHost / MemcpyToHost). + +ORT inserts a `MemcpyFromHost` node when a device (MUSA) kernel consumes a +CPU-produced tensor, and a `MemcpyToHost` node when a CPU kernel consumes a +MUSA-produced tensor. These tests pin the behaviour that the plugin EP provides +the kernels for those nodes, so the copy runs on the plugin's own stream and +the `MemcpyFromHost` destination is a MUSA device tensor. + +The consumer operators below (`Flatten`, `Transpose`) reject CPU memory on the +device path, so the tests are RED (device-tensor guard failure) unless the +boundary copy really lands in device memory. +""" + +import numpy as np +import pytest +from onnx import helper + +from op_test_utils import ( + TensorProto, + build_graph_model, + run_model_and_compare_with_cpu_fallback, +) + + +def _cpu_produced_input_model(x, op_type, attrs, name): + """`Sin` (CPU EP) -> `op_type` (MUSA EP): one inserted MemcpyFromHost.""" + return build_graph_model( + [ + helper.make_node("Sin", ["X"], ["cpu_output"]), + helper.make_node(op_type, ["cpu_output"], ["Y"], **attrs), + ], + inputs={"X": x}, + outputs=[("Y", TensorProto.FLOAT)], + name=name, + ) + + +def test_memcpy_from_host_feeds_flatten_device_memory(): + x = np.linspace(-1.5, 1.5, 24, dtype=np.float32).reshape(2, 3, 4) + model = _cpu_produced_input_model( + x, "Flatten", {"axis": 1}, "memcpy_from_host_flatten_graph" + ) + run_model_and_compare_with_cpu_fallback( + model, {"X": x}, rtol=0, atol=0, disable_ep_error_fallback=True + ) + + +def test_memcpy_from_host_runs_on_musa_kernel(monkeypatch): + x = np.linspace(-1.5, 1.5, 24, dtype=np.float32).reshape(2, 3, 4) + model = _cpu_produced_input_model( + x, "Flatten", {"axis": 1}, "memcpy_from_host_kernel_route_graph" + ) + + # Provenance guard: the boundary copy must execute through the plugin's own + # MemcpyFromHost kernel, not the CPU EP fallback. + monkeypatch.setenv("ORT_MUSA_TEST_FAIL_MEMCPY_KERNEL", "1") + with pytest.raises(Exception, match="test-only MUSA memcpy kernel route"): + run_model_and_compare_with_cpu_fallback( + model, {"X": x}, disable_ep_error_fallback=True + ) + + +def test_memcpy_from_host_shared_by_multiple_musa_kernels(): + x = np.linspace(-1.5, 1.5, 24, dtype=np.float32).reshape(2, 3, 4) + model = build_graph_model( + [ + helper.make_node("Sin", ["X"], ["cpu_output"]), + helper.make_node("Flatten", ["cpu_output"], ["flat"], axis=1), + helper.make_node("Transpose", ["cpu_output"], ["transposed"], perm=[2, 0, 1]), + ], + inputs={"X": x}, + outputs=[("flat", TensorProto.FLOAT), ("transposed", TensorProto.FLOAT)], + name="memcpy_from_host_multiple_consumers_graph", + ) + run_model_and_compare_with_cpu_fallback( + model, {"X": x}, rtol=0, atol=0, disable_ep_error_fallback=True + ) + + +def test_memcpy_to_host_before_cpu_kernel(monkeypatch): + x = np.linspace(-2.0, 2.0, 24, dtype=np.float32).reshape(2, 3, 4) + model = build_graph_model( + [ + helper.make_node("Relu", ["X"], ["relu_output"]), + helper.make_node("Sin", ["relu_output"], ["Y"]), + ], + inputs={"X": x}, + outputs=[("Y", TensorProto.FLOAT)], + name="memcpy_to_host_graph", + ) + + run_model_and_compare_with_cpu_fallback( + model, {"X": x}, rtol=1e-6, atol=1e-7, disable_ep_error_fallback=True + ) + + # Provenance guard for the MUSA -> CPU direction as well. + monkeypatch.setenv("ORT_MUSA_TEST_FAIL_MEMCPY_KERNEL", "1") + with pytest.raises(Exception, match="test-only MUSA memcpy kernel route"): + run_model_and_compare_with_cpu_fallback( + model, {"X": x}, disable_ep_error_fallback=True + ) + + +def _if_branch(): + return helper.make_graph( + [ + helper.make_node("Sin", ["X"], ["cpu_output"]), + helper.make_node("Flatten", ["cpu_output"], ["branch_y"], axis=1), + ], + "branch", + [], + [helper.make_tensor_value_info("branch_y", TensorProto.FLOAT, None)], + ) + + +def test_memcpy_from_host_inside_if_subgraph(): + x = np.linspace(-1.5, 1.5, 24, dtype=np.float32).reshape(2, 3, 4) + node = helper.make_node( + "If", ["cond"], ["Y"], then_branch=_if_branch(), else_branch=_if_branch() + ) + graph = helper.make_graph( + [node], + "memcpy_from_host_if_graph", + [ + helper.make_tensor_value_info("cond", TensorProto.BOOL, [1]), + helper.make_tensor_value_info("X", TensorProto.FLOAT, [2, 3, 4]), + ], + [helper.make_tensor_value_info("Y", TensorProto.FLOAT, None)], + ) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)]) + model.ir_version = min(model.ir_version, 10) + run_model_and_compare_with_cpu_fallback( + model.SerializeToString(), + {"cond": np.array([True], dtype=np.bool_), "X": x}, + rtol=0, + atol=0, + disable_ep_error_fallback=True, + ) diff --git a/test/ops/test_relu.py b/test/ops/test_relu.py index b889388..7728ba9 100644 --- a/test/ops/test_relu.py +++ b/test/ops/test_relu.py @@ -13,7 +13,6 @@ """End-to-end CPU-vs-MUSA test for the Relu operator.""" import numpy as np -import pytest from onnx import helper from op_test_utils import ( @@ -52,10 +51,12 @@ def test_relu_float16_accepts_cpu_produced_input(monkeypatch): ) monkeypatch.setenv("ORT_MUSA_TEST_FAIL_RELU_CPU_INPUT", "1") - with pytest.raises(Exception, match="test-only Relu CPU-input route"): - run_model_and_compare_with_cpu_fallback( - model, {"X": x}, disable_ep_error_fallback=True - ) + # The plugin now runs the MemcpyFromHost node, so the CPU-produced tensor + # reaches Relu in device memory and the test-only CPU-input route must + # not be taken. + run_model_and_compare_with_cpu_fallback( + model, {"X": x}, disable_ep_error_fallback=True + ) monkeypatch.delenv("ORT_MUSA_TEST_FAIL_RELU_CPU_INPUT") run_model_and_compare_with_cpu_fallback( diff --git a/test/ops/test_split.py b/test/ops/test_split.py index f407e97..2ae94e0 100644 --- a/test/ops/test_split.py +++ b/test/ops/test_split.py @@ -192,12 +192,14 @@ def test_split_after_reshape_float16_cpu_produced_input(monkeypatch): ) monkeypatch.setenv("ORT_MUSA_TEST_FAIL_RESHAPE_CPU_INPUT", "1") - with pytest.raises(Exception, match="test-only Reshape CPU-input route"): - run_model_and_compare_with_cpu_fallback( - model, - {"X": x, "shape": shape, "split": split}, - disable_ep_error_fallback=True, - ) + # The plugin now runs the MemcpyFromHost node, so the CPU-produced tensor + # reaches Reshape in device memory and the test-only CPU-input route must + # not be taken. + run_model_and_compare_with_cpu_fallback( + model, + {"X": x, "shape": shape, "split": split}, + disable_ep_error_fallback=True, + ) monkeypatch.delenv("ORT_MUSA_TEST_FAIL_RESHAPE_CPU_INPUT") run_model_and_compare_with_cpu_fallback( From fdeab6411971451b6a107498e5a16b84825f8935 Mon Sep 17 00:00:00 2001 From: "yichao.hu" Date: Mon, 12 Oct 2026 23:47:31 -0400 Subject: [PATCH 12/52] feat(conv): support 1D NCL convolution Speech models use 1D Conv (NCL), which the EP rejected. Run 1D Conv as a 2D convolution with H=1, reusing the 2D and grouped paths; the output keeps rank 3. --- musa/ep/src/kernels/nn/conv.cc | 75 ++++++++++++++++++++++------------ test/ops/test_conv.py | 60 +++++++++++++++++++++++++++ 2 files changed, 110 insertions(+), 25 deletions(-) diff --git a/musa/ep/src/kernels/nn/conv.cc b/musa/ep/src/kernels/nn/conv.cc index 1753dd0..507be31 100644 --- a/musa/ep/src/kernels/nn/conv.cc +++ b/musa/ep/src/kernels/nn/conv.cc @@ -168,10 +168,11 @@ OrtStatus* Conv::Compute(Ort::KernelContext& ctx) const { } auto x_shape = x_info.GetShape(); auto w_shape = w_info.GetShape(); - if ((x_shape.size() != 4 && x_shape.size() != 5) || - w_shape.size() != x_shape.size()) { + const size_t rank = x_shape.size(); + if ((rank != 3 && rank != 4 && rank != 5) || w_shape.size() != rank) { return Ort::GetApi().CreateStatus( - ORT_NOT_IMPLEMENTED, "Conv only supports 2D NCHW or 3D NCDHW tensors"); + ORT_NOT_IMPLEMENTED, + "Conv only supports 1D NCL, 2D NCHW or 3D NCDHW tensors"); } if (group_ <= 0 || x_shape[1] != w_shape[1] * group_ || w_shape[0] % group_ != 0) { @@ -180,7 +181,11 @@ OrtStatus* Conv::Compute(Ort::KernelContext& ctx) const { "Conv group must divide output channels and match input channels"); } - const size_t spatial_rank = x_shape.size() - 2; + // A Conv1D tensor is NCL. Executing it as a unit-height 2D NCHW tensor + // ([N, C, 1, L] with a [M, C / group, 1, K] kernel) keeps the memory layout + // identical, so the existing 2D muDNN and custom device paths are reused. + const bool is_1d = rank == 3; + const size_t spatial_rank = is_1d ? 1 : rank - 2; std::vector pads = pads_; if (pads.empty()) { pads.resize(spatial_rank * 2, 0); @@ -210,7 +215,8 @@ OrtStatus* Conv::Compute(Ort::KernelContext& ctx) const { } } - std::vector y_shape = {x_shape[0], w_shape[0]}; + std::vector out_spatial; + out_spatial.reserve(spatial_rank); for (size_t i = 0; i < spatial_rank; ++i) { const int64_t output_dim = (x_shape[i + 2] + pads[i] + pads[i + spatial_rank] - @@ -221,14 +227,33 @@ OrtStatus* Conv::Compute(Ort::KernelContext& ctx) const { return Ort::GetApi().CreateStatus(ORT_INVALID_ARGUMENT, "Conv invalid output shape"); } - y_shape.push_back(output_dim); + out_spatial.push_back(output_dim); } + + std::vector y_shape = {x_shape[0], w_shape[0]}; + y_shape.insert(y_shape.end(), out_spatial.begin(), out_spatial.end()); Ort::UnownedValue output = ctx.GetOutput(0, y_shape); - if (TryMudnnConv(ctx, x_shape, w_shape, y_shape, pads, strides, dilations, - output)) { + + std::vector x_exec = x_shape; + std::vector w_exec = w_shape; + std::vector y_exec = y_shape; + std::vector pads_exec = pads; + std::vector strides_exec = strides; + std::vector dilations_exec = dilations; + if (is_1d) { + x_exec = {x_shape[0], x_shape[1], 1, x_shape[2]}; + w_exec = {w_shape[0], w_shape[1], 1, w_shape[2]}; + y_exec = {y_shape[0], y_shape[1], 1, out_spatial[0]}; + pads_exec = {0, pads[0], 0, pads[1]}; + strides_exec = {1, strides[0]}; + dilations_exec = {1, dilations[0]}; + } + + if (TryMudnnConv(ctx, x_exec, w_exec, y_exec, pads_exec, strides_exec, + dilations_exec, output)) { return nullptr; } - if (spatial_rank == 3) { + if (rank == 5) { return Ort::GetApi().CreateStatus( ORT_NOT_IMPLEMENTED, "Conv3D requires the muDNN path; custom fallback is 2D-only"); @@ -255,23 +280,23 @@ OrtStatus* Conv::Compute(Ort::KernelContext& ctx) const { } MusaConv2DParams params{}; - params.n = x_shape[0]; - params.c = x_shape[1]; + params.n = x_exec[0]; + params.c = x_exec[1]; params.group = group_; - params.h = x_shape[2]; - params.w = x_shape[3]; - params.m = w_shape[0]; - params.kernel_h = w_shape[2]; - params.kernel_w = w_shape[3]; - params.out_h = y_shape[2]; - params.out_w = y_shape[3]; - params.pad_h = pads[0]; - params.pad_w = pads[1]; - params.stride_h = strides[0]; - params.stride_w = strides[1]; - params.dilation_h = dilations[0]; - params.dilation_w = dilations[1]; - params.total_elements = NumElements(y_shape); + params.h = x_exec[2]; + params.w = x_exec[3]; + params.m = w_exec[0]; + params.kernel_h = w_exec[2]; + params.kernel_w = w_exec[3]; + params.out_h = y_exec[2]; + params.out_w = y_exec[3]; + params.pad_h = pads_exec[0]; + params.pad_w = pads_exec[1]; + params.stride_h = strides_exec[0]; + params.stride_w = strides_exec[1]; + params.dilation_h = dilations_exec[0]; + params.dilation_w = dilations_exec[1]; + params.total_elements = NumElements(y_exec); return LaunchStatus(LaunchMusaConv2DFloatKernel( x.GetTensorData(), w.GetTensorData(), bias, output.GetTensorMutableData(), params, GetComputeStream(ctx))); diff --git a/test/ops/test_conv.py b/test/ops/test_conv.py index 84823bb..c6dd3f9 100644 --- a/test/ops/test_conv.py +++ b/test/ops/test_conv.py @@ -150,3 +150,63 @@ def test_conv_lowered_conv1d_h1_k1_with_bias(): rtol=1e-4, atol=1e-4, ) + + +def test_conv1d_opset13_float_ncl_with_bias_stride_pad(): + x = np.random.default_rng(20).standard_normal((2, 3, 10)).astype(np.float32) + w = np.random.default_rng(21).standard_normal((4, 3, 3)).astype(np.float32) + b = np.random.default_rng(22).standard_normal((4,)).astype(np.float32) + run_and_compare( + "Conv", + inputs={"X": x, "W": w, "B": b}, + outputs=[("Y", TensorProto.FLOAT)], + attrs={"pads": [1, 1], "strides": [2], "dilations": [1]}, + opset=13, + rtol=1e-4, + atol=1e-4, + ) + + +def test_conv1d_opset13_float_ncl_dilated(): + x = np.random.default_rng(23).standard_normal((1, 2, 12)).astype(np.float32) + w = np.random.default_rng(24).standard_normal((3, 2, 3)).astype(np.float32) + run_and_compare( + "Conv", + inputs={"X": x, "W": w}, + outputs=[("Y", TensorProto.FLOAT)], + attrs={"pads": [2, 2], "dilations": [2]}, + opset=13, + rtol=1e-4, + atol=1e-4, + ) + + +def test_conv1d_opset13_float_ncl_grouped_with_bias(): + x = np.random.default_rng(25).standard_normal((2, 6, 9)).astype(np.float32) + w = np.random.default_rng(26).standard_normal((9, 2, 4)).astype(np.float32) + b = np.random.default_rng(27).standard_normal((9,)).astype(np.float32) + run_and_compare( + "Conv", + inputs={"X": x, "W": w, "B": b}, + outputs=[("Y", TensorProto.FLOAT)], + attrs={"group": 3, "pads": [1, 1], "strides": [1]}, + opset=13, + rtol=1e-4, + atol=1e-4, + ) + + +def test_conv1d_opset13_float_ncl_depthwise_with_bias(): + # A large grouped kernel mirrors the VITS-style position convolution. + x = np.random.default_rng(28).standard_normal((1, 16, 20)).astype(np.float32) + w = np.random.default_rng(29).standard_normal((16, 1, 13)).astype(np.float32) + b = np.random.default_rng(30).standard_normal((16,)).astype(np.float32) + run_and_compare( + "Conv", + inputs={"X": x, "W": w, "B": b}, + outputs=[("Y", TensorProto.FLOAT)], + attrs={"group": 16, "pads": [6, 6], "strides": [1]}, + opset=13, + rtol=1e-4, + atol=1e-4, + ) From 8c2395c68c2f9ac590d3b4af2afdfe70068b9a61 Mon Sep 17 00:00:00 2001 From: "yichao.hu" Date: Mon, 12 Oct 2026 23:47:31 -0400 Subject: [PATCH 13/52] feat(resize): support linear/trilinear modes and the sizes input; decline unsupported configurations Resize only supported nearest with the scales input. Add linear (2D) and trilinear (3D) modes and the sizes input. Registering Resize made ORT assign every Resize node to MUSA, including configurations the kernel cannot run, which failed at run time. Add a capability predicate, shared with the kernel's own checks (resize_support.h), so unsupported configurations (cubic, antialias, tf_crop_and_resize, ...) stay on the CPU EP. --- musa/ep/CMakeLists.txt | 1 + musa/ep/src/ep.cc | 7 +- musa/ep/src/graph/README.md | 30 +- musa/ep/src/graph/capability_predicates.cc | 138 +++++ musa/ep/src/graph/capability_predicates.h | 38 ++ musa/ep/src/kernels/tensor/resize.cc | 197 +++++--- musa/ep/src/kernels/tensor/resize_impl.h | 23 + musa/ep/src/kernels/tensor/resize_impl.mu | 196 +++++++- musa/ep/src/kernels/tensor/resize_support.h | 177 +++++++ test/ops/test_resize.py | 528 +++++++++++++++++++- 10 files changed, 1258 insertions(+), 77 deletions(-) create mode 100644 musa/ep/src/graph/capability_predicates.cc create mode 100644 musa/ep/src/graph/capability_predicates.h create mode 100644 musa/ep/src/kernels/tensor/resize_support.h diff --git a/musa/ep/CMakeLists.txt b/musa/ep/CMakeLists.txt index b19a0e5..2f6e756 100644 --- a/musa/ep/CMakeLists.txt +++ b/musa/ep/CMakeLists.txt @@ -13,6 +13,7 @@ set(MUSA_PLUGIN_SOURCES src/ep_profiling.cc src/ep_stream.cc src/musa_arena.cc + src/graph/capability_predicates.cc src/graph/cpu_preferred_nodes.cc src/graph/graph_utils.cc src/graph_mermaid_dump.cc diff --git a/musa/ep/src/ep.cc b/musa/ep/src/ep.cc index 5b39379..c1095e3 100644 --- a/musa/ep/src/ep.cc +++ b/musa/ep/src/ep.cc @@ -28,6 +28,7 @@ #include "ep_stream.h" #include "fusion/fusion_matcher.h" #include "fusion/fusion_node_compute.h" +#include "graph/capability_predicates.h" #include "graph/cpu_preferred_nodes.h" #include "graph_mermaid_dump.h" #include "plugin_ep_utils.h" @@ -237,7 +238,11 @@ OrtStatus* CollectSupportedNodeCandidates( const OrtKernelDef* kernel_def = nullptr; RETURN_IF_ERROR(ep_api.EpGraphSupportInfo_LookUpKernel(graph_support_info, node, &kernel_def)); - if (kernel_def != nullptr) { + // The registry only matches op/domain/opset/dtype, so ask the per-operator + // predicate as well: a kernel that implements a subset of the ONNX + // configurations must not claim the nodes it cannot run, or they would + // fail during Compute() instead of running on the CPU EP. + if (kernel_def != nullptr && IsNodeConfigurationSupported(node)) { candidates.nodes.push_back(node); candidates.tentative_nodes.push_back(node); } diff --git a/musa/ep/src/graph/README.md b/musa/ep/src/graph/README.md index d1a3739..fb219d7 100644 --- a/musa/ep/src/graph/README.md +++ b/musa/ep/src/graph/README.md @@ -4,22 +4,26 @@ Graph-level capability and fusion logic for the MUSA Plugin EP. ## Why this directory exists -The Plugin EP partitions graphs in [`../ep.cc`](../ep.cc): for every node it checks -an op-type allow-list and then confirms a matching kernel is registered through the -kernel registry (`../ep_kernel_registration.cc`). As a result, "which nodes can we -run" is today fully derived from the registered -`ONNX_OPERATOR_VERSIONED_KERNEL_EX` entries under [`../kernels/`](../kernels) — no -extra graph-level code is required, so this directory currently holds no sources. - -It is kept as the designated home for graph-level logic that does **not** fit the -per-kernel registry model and will be added as the EP grows: +The Plugin EP partitions graphs in [`../ep.cc`](../ep.cc): for every node it confirms +a matching kernel is registered through the kernel registry +(`../ep_kernel_registration.cc`), and then asks the per-operator predicates here +whether that kernel can actually run the node. "Which nodes can we run" is therefore +derived from the registered `ONNX_OPERATOR_VERSIONED_KERNEL_EX` entries under +[`../kernels/`](../kernels), narrowed by the graph-level logic that does **not** fit +the per-kernel registry model: - **Fusion-pattern matchers** that rewrite subgraphs into `com.microsoft`-domain fused kernels (e.g. `FusedGemm`, `FusedMatMul`) before partitioning. -- **Free-standing capability predicates** for decisions the registry cannot - express on its own, such as gating a node on attribute values, input shapes, or - experimental flags. +- **Free-standing capability predicates** ([`capability_predicates.cc`](capability_predicates.cc)) + for decisions the registry cannot express on its own, such as gating a node on + attribute values, input shapes, or experimental flags. A predicate returns false + for nodes the matching kernel cannot run, and those nodes stay on the CPU EP + instead of failing when the kernel runs. The per-operator rules themselves live + next to the kernels they gate (e.g. `../kernels/tensor/resize_support.h`), so the + capability decision and the kernel's runtime guard share one definition. +- **CPU-preferred node selection** ([`cpu_preferred_nodes.cc`](cpu_preferred_nodes.cc)) + for shape/metadata arithmetic that is cheaper on the CPU EP. -When such logic lands, its `.h`/`.cc` files go here and must be added explicitly to +Such `.h`/`.cc` files live here and must be added explicitly to [`../../CMakeLists.txt`](../../CMakeLists.txt) (`src/kernels/**/*.cc` is auto-globbed recursively; sources in this directory are not). diff --git a/musa/ep/src/graph/capability_predicates.cc b/musa/ep/src/graph/capability_predicates.cc new file mode 100644 index 0000000..4526f83 --- /dev/null +++ b/musa/ep/src/graph/capability_predicates.cc @@ -0,0 +1,138 @@ +// Copyright (c) Moore Threads Technology Co., Ltd. All rights reserved. +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "graph/capability_predicates.h" + +#include +#include +#include +#include + +#include "graph/graph_utils.h" +#include "tensor/resize_support.h" + +namespace musa_ep { +namespace { + +/// Reads a constant float initializer. Returns nullopt when the value is not a +/// constant float tensor, i.e. when only the runtime kernel can judge it. +std::optional> ReadConstantFloatTensor( + Ort::ConstValueInfo value_info) { + if (value_info == nullptr || !value_info.IsConstantInitializer()) { + return std::nullopt; + } + + Ort::ConstValue value{nullptr}; + Ort::Status status = value_info.GetInitializer(value); + if (!status.IsOK() || !value) { + return std::nullopt; + } + + const auto info = value.GetTensorTypeAndShapeInfo(); + if (info.GetElementType() != ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT) { + return std::nullopt; + } + const size_t element_count = static_cast(info.GetElementCount()); + if (element_count == 0) { + return std::vector(); + } + const float* data = value.GetTensorData(); + return std::vector(data, data + element_count); +} + +/// Element count of a constant initializer, or nullopt when the tensor only +/// exists at run time. +std::optional ConstantTensorElementCount( + Ort::ConstValueInfo value_info) { + if (value_info == nullptr || !value_info.IsConstantInitializer()) { + return std::nullopt; + } + + Ort::ConstValue value{nullptr}; + Ort::Status status = value_info.GetInitializer(value); + if (!status.IsOK() || !value) { + return std::nullopt; + } + + return value.GetTensorTypeAndShapeInfo().GetElementCount(); +} + +/// True when the node leaves the optional input at \p index unused, which is +/// the case when it is omitted or when the graph fills it with a zero-element +/// constant: exporters that drive Resize through `sizes` commonly pass an empty +/// `roi` and `scales` initializer next to it, and the kernel reads those as +/// absent as well. A run-time tensor is assumed to be used, because only the +/// kernel can see its element count. +bool InputUnused(const std::vector& inputs, size_t index) { + if (inputs.size() <= index || inputs[index] == nullptr) { + return true; + } + return ResizeInputUnused(ConstantTensorElementCount(inputs[index])); +} + +/// True when the MUSA Resize kernel can run this node. Attribute values, which +/// optional input is supplied and the type, rank and shape of X decide the +/// answer for every node; the values of the roi/scales/sizes tensors are +/// checked as well when the graph carries them as constant initializers, and +/// are otherwise left to the kernel. +bool ResizeConfigurationSupported(Ort::ConstNode node) { + ResizeRouteConfig config; + config.mode = GetStringAttribute(node, "mode").value_or(kResizeDefaultMode); + config.coordinate_transformation_mode = + GetStringAttribute(node, "coordinate_transformation_mode") + .value_or(kResizeDefaultCoordinateTransformationMode); + config.nearest_mode = GetStringAttribute(node, "nearest_mode") + .value_or(kResizeDefaultNearestMode); + + const std::vector inputs = node.GetInputs(); + if (inputs.empty() || inputs[0] == nullptr) { + return false; + } + config.has_scales_input = !InputUnused(inputs, 2); + config.has_sizes_input = !InputUnused(inputs, 3); + + const Ort::ConstTypeInfo input_type = inputs[0].TypeInfo(); + if (input_type.GetONNXType() != ONNX_TYPE_TENSOR) { + return false; + } + const auto input_info = input_type.GetTensorTypeAndShapeInfo(); + config.input_rank = static_cast(input_info.GetShape().size()); + config.input_element_type = input_info.GetElementType(); + if (ResizeRouteRejection(config) != nullptr) { + return false; + } + + ResizeRouteValues values; + if (inputs.size() > 1) { + values.roi_element_count = ConstantTensorElementCount(inputs[1]); + } + if (config.has_scales_input) { + values.scales = ReadConstantFloatTensor(inputs[2]); + } + if (config.has_sizes_input) { + values.sizes = ReadIntInitializerNoLimit(inputs[3]); + } + values.input_shape = GetStaticShape(inputs[0]); + return ResizeRouteValueRejection(config, values) == nullptr; +} + +} // namespace + +bool IsNodeConfigurationSupported(Ort::ConstNode node) { + if (IsOnnxOp(node, "Resize")) { + return ResizeConfigurationSupported(node); + } + return true; +} + +} // namespace musa_ep diff --git a/musa/ep/src/graph/capability_predicates.h b/musa/ep/src/graph/capability_predicates.h new file mode 100644 index 0000000..a48f336 --- /dev/null +++ b/musa/ep/src/graph/capability_predicates.h @@ -0,0 +1,38 @@ +// Copyright (c) Moore Threads Technology Co., Ltd. All rights reserved. +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#pragma once + +#define ORT_API_MANUAL_INIT +#include "onnxruntime_cxx_api.h" +#undef ORT_API_MANUAL_INIT + +namespace musa_ep { + +/// Graph-level capability gate for a single node that already matched a +/// registered kernel. +/// +/// `OrtEpApi::EpGraphSupportInfo_LookUpKernel` matches op/domain/opset/dtype +/// only, so an operator whose kernel implements a subset of the ONNX +/// configurations (such as Resize, where the kernels only cover nearest/ +/// asymmetric/floor and the linear coordinate transformations, rank 4 and 5, +/// and an empty roi) would otherwise be handed every node of that operator and +/// only fail once the kernel runs. Returning false here leaves the node to the +/// CPU EP, as it was before the kernel was registered. +/// +/// The per-operator predicates live next to the kernels they gate so that the +/// capability decision and the kernel's runtime guard are two uses of one +/// definition instead of two copies that drift apart. +bool IsNodeConfigurationSupported(Ort::ConstNode node); + +} // namespace musa_ep diff --git a/musa/ep/src/kernels/tensor/resize.cc b/musa/ep/src/kernels/tensor/resize.cc index bfe8ab7..9f091f3 100644 --- a/musa/ep/src/kernels/tensor/resize.cc +++ b/musa/ep/src/kernels/tensor/resize.cc @@ -11,95 +11,174 @@ // See the License for the specific language governing permissions and // limitations under the License. +#include +#include +#include + #include "shared_inc/op_kernel_common.h" #include "tensor/resize_impl.h" +#include "tensor/resize_support.h" namespace { +OrtStatus* NotImplemented(const char* message) { + return Ort::GetApi().CreateStatus(ORT_NOT_IMPLEMENTED, message); +} + +bool ParseCoordinateMode(const std::string& name, int32_t& mode) { + if (name == "half_pixel") { + mode = kMusaResizeHalfPixel; + return true; + } + if (name == "pytorch_half_pixel") { + mode = kMusaResizePytorchHalfPixel; + return true; + } + if (name == "align_corners") { + mode = kMusaResizeAlignCorners; + return true; + } + if (name == "asymmetric") { + mode = kMusaResizeAsymmetric; + return true; + } + return false; +} + +// An optional Resize input is unused when it is absent or carries no elements +// (an empty roi/scales tensor is how the exporters mark "not provided"); the +// graph-level capability predicate applies the same rule to the constant +// initializers through musa_ep::ResizeInputUnused(). +bool IsAbsentInput(Ort::KernelContext& ctx, size_t index) { + if (ctx.GetInputCount() <= index) { + return true; + } + Ort::ConstValue value = ctx.GetInput(index); + if (value == nullptr) { + return true; + } + return musa_ep::ResizeInputUnused( + value.GetTensorTypeAndShapeInfo().GetElementCount()); +} + class Resize : public OpKernelBase { public: Resize(const OrtKernelInfo* info, void* /*state*/) { Ort::ConstKernelInfo kernel_info(info); - mode_ = AttrOrDefault(kernel_info, "mode", "nearest"); - coordinate_transform_mode_ = AttrOrDefault( - kernel_info, "coordinate_transformation_mode", "half_pixel"); - nearest_mode_ = AttrOrDefault(kernel_info, "nearest_mode", - "round_prefer_floor"); + config_.mode = AttrOrDefault(kernel_info, "mode", + musa_ep::kResizeDefaultMode); + config_.coordinate_transformation_mode = AttrOrDefault( + kernel_info, "coordinate_transformation_mode", + musa_ep::kResizeDefaultCoordinateTransformationMode); + config_.nearest_mode = AttrOrDefault( + kernel_info, "nearest_mode", musa_ep::kResizeDefaultNearestMode); } OrtStatus* Compute(Ort::KernelContext& ctx) const; private: - std::string mode_; - std::string coordinate_transform_mode_; - std::string nearest_mode_; + musa_ep::ResizeRouteConfig config_; }; OrtStatus* Resize::Compute(Ort::KernelContext& ctx) const { - if (mode_ != "nearest" || coordinate_transform_mode_ != "asymmetric" || - nearest_mode_ != "floor") { - return Ort::GetApi().CreateStatus( - ORT_NOT_IMPLEMENTED, - "Resize only supports nearest/asymmetric/floor in this MUSA route"); + Ort::ConstValue input = ctx.GetInput(0); + auto input_info = input.GetTensorTypeAndShapeInfo(); + const std::vector input_shape = input_info.GetShape(); + + // Route checks shared with the graph-level capability predicate in + // graph/capability_predicates.cc: a configuration this route cannot run must + // not be claimed by the EP, and must be rejected here if it was claimed. + musa_ep::ResizeRouteConfig config = config_; + config.has_scales_input = !IsAbsentInput(ctx, 2); + config.has_sizes_input = !IsAbsentInput(ctx, 3); + config.input_rank = static_cast(input_shape.size()); + config.input_element_type = input_info.GetElementType(); + if (const char* reason = musa_ep::ResizeRouteRejection(config); + reason != nullptr) { + return NotImplemented(reason); + } + + const size_t rank = input_shape.size(); + for (int64_t dim : input_shape) { + if (dim <= 0) { + return NotImplemented("Resize requires a non-empty input tensor"); + } + } + + musaStream_t stream = GetComputeStream(ctx); + musa_ep::ResizeRouteValues values; + if (ctx.GetInputCount() > 1 && ctx.GetInput(1) != nullptr) { + values.roi_element_count = + ctx.GetInput(1).GetTensorTypeAndShapeInfo().GetElementCount(); } - if (ctx.GetInputCount() <= 2 || ctx.GetInput(2) == nullptr || - (ctx.GetInputCount() > 3 && ctx.GetInput(3) != nullptr)) { - return Ort::GetApi().CreateStatus( - ORT_NOT_IMPLEMENTED, - "Resize only supports the scales input, not the sizes input"); + if (config.has_scales_input) { + values.scales = ReadTyped(ctx.GetInput(2), stream); + } else { + values.sizes = ReadIntTensor(ctx, 3); } - if (ctx.GetInputCount() > 1 && ctx.GetInput(1) != nullptr && - ctx.GetInput(1).GetTensorTypeAndShapeInfo().GetElementCount() != 0) { - return Ort::GetApi().CreateStatus( - ORT_NOT_IMPLEMENTED, - "Resize only supports an empty ROI for asymmetric coordinates"); + values.input_shape = input_shape; + if (const char* reason = musa_ep::ResizeRouteValueRejection(config, values); + reason != nullptr) { + return NotImplemented(reason); } - Ort::ConstValue input = ctx.GetInput(0); - auto input_info = input.GetTensorTypeAndShapeInfo(); - if (input_info.GetElementType() != ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT) { - return Ort::GetApi().CreateStatus(ORT_NOT_IMPLEMENTED, - "Resize only supports float32"); + // Output shape and the effective per-axis scale factor. This mirrors the ORT + // CPU kernel: with `scales` the factor is the value itself, with `sizes` it + // is output/input (see ParseScalesDataAndAdjustOutputSize). + std::vector output_shape(rank); + std::vector scales(rank); + if (config.has_scales_input) { + scales = *values.scales; + for (size_t axis = 0; axis < rank; ++axis) { + output_shape[axis] = static_cast( + scales[axis] * static_cast(input_shape[axis])); + if (output_shape[axis] <= 0) { + return NotImplemented("Resize output must be non-empty"); + } + } + } else { + const std::vector& sizes = *values.sizes; + for (size_t axis = 0; axis < rank; ++axis) { + output_shape[axis] = sizes[axis]; + scales[axis] = static_cast(sizes[axis]) / + static_cast(input_shape[axis]); + } } - const std::vector input_shape = input_info.GetShape(); - if (input_shape.size() != 4 || input_shape[2] <= 0 || input_shape[3] <= 0) { - return Ort::GetApi().CreateStatus( - ORT_NOT_IMPLEMENTED, - "Resize only supports non-empty rank-4 NCHW tensors"); - } - const std::vector scales = - ReadTyped(ctx.GetInput(2), GetComputeStream(ctx)); - if (scales.size() != 4 || scales[0] != 1.0f || scales[1] != 1.0f || - !std::isfinite(scales[2]) || !std::isfinite(scales[3]) || - scales[2] <= 0.0f || scales[3] <= 0.0f) { - return Ort::GetApi().CreateStatus( - ORT_NOT_IMPLEMENTED, - "Resize requires positive NCHW scales with batch/channel scale 1"); - } - - std::vector output_shape = input_shape; - output_shape[2] = - static_cast(std::floor(input_shape[2] * scales[2])); - output_shape[3] = - static_cast(std::floor(input_shape[3] * scales[3])); - if (output_shape[2] <= 0 || output_shape[3] <= 0) { - return Ort::GetApi().CreateStatus(ORT_NOT_IMPLEMENTED, - "Resize output must be non-empty"); + + int32_t coordinate_mode = 0; + const bool coordinate_mode_known = ParseCoordinateMode( + config.coordinate_transformation_mode, coordinate_mode); + if (config.mode == "linear" && !coordinate_mode_known) { + // Unreachable while ResizeRouteRejection() rejects every coordinate + // transformation mode that ParseCoordinateMode() does not know. + return NotImplemented(musa_ep::kResizeLinearCoordinateReason); } + if (!IsGpuMemory(input.GetTensorMemoryInfo())) { - return Ort::GetApi().CreateStatus(ORT_NOT_IMPLEMENTED, - "Resize requires MUSA input"); + return NotImplemented("Resize requires MUSA input"); } Ort::UnownedValue output = ctx.GetOutput(0, output_shape); if (!IsGpuMemory(output.GetTensorMemoryInfo())) { - return Ort::GetApi().CreateStatus(ORT_NOT_IMPLEMENTED, - "Resize requires MUSA output"); + return NotImplemented("Resize requires MUSA output"); } - return LaunchStatus(LaunchMusaResizeNearest2DFloat( + if (config.mode == "nearest") { + return LaunchStatus(LaunchMusaResizeNearest2DFloat( + input.GetTensorData(), output.GetTensorMutableData(), + input_shape[0], input_shape[1], input_shape[2], input_shape[3], + output_shape[2], output_shape[3], scales[2], scales[3], stream)); + } + if (rank == 4) { + return LaunchStatus(LaunchMusaResizeLinear2DFloat( + input.GetTensorData(), output.GetTensorMutableData(), + input_shape[0], input_shape[1], input_shape[2], input_shape[3], + output_shape[2], output_shape[3], scales[2], scales[3], coordinate_mode, + stream)); + } + return LaunchStatus(LaunchMusaResizeLinear3DFloat( input.GetTensorData(), output.GetTensorMutableData(), input_shape[0], input_shape[1], input_shape[2], input_shape[3], - output_shape[2], output_shape[3], scales[2], scales[3], - GetComputeStream(ctx))); + input_shape[4], output_shape[2], output_shape[3], output_shape[4], + scales[2], scales[3], scales[4], coordinate_mode, stream)); } } // namespace diff --git a/musa/ep/src/kernels/tensor/resize_impl.h b/musa/ep/src/kernels/tensor/resize_impl.h index 52e20cd..d07ba0a 100644 --- a/musa/ep/src/kernels/tensor/resize_impl.h +++ b/musa/ep/src/kernels/tensor/resize_impl.h @@ -4,8 +4,31 @@ #include +// Coordinate transformation modes shared by resize.cc (host) and +// resize_impl.mu (device). The numeric values must stay in sync with the +// ONNX attribute strings mapped in resize.cc. +enum MusaResizeCoordinateMode : int32_t { + kMusaResizeHalfPixel = 0, + kMusaResizePytorchHalfPixel = 1, + kMusaResizeAlignCorners = 2, + kMusaResizeAsymmetric = 3, +}; + musaError_t LaunchMusaResizeNearest2DFloat( const float* input, float* output, int64_t batch, int64_t channels, int64_t input_height, int64_t input_width, int64_t output_height, int64_t output_width, float scale_height, float scale_width, musaStream_t stream); + +musaError_t LaunchMusaResizeLinear2DFloat( + const float* input, float* output, int64_t batch, int64_t channels, + int64_t input_height, int64_t input_width, int64_t output_height, + int64_t output_width, float scale_height, float scale_width, + int32_t coordinate_mode, musaStream_t stream); + +musaError_t LaunchMusaResizeLinear3DFloat( + const float* input, float* output, int64_t batch, int64_t channels, + int64_t input_depth, int64_t input_height, int64_t input_width, + int64_t output_depth, int64_t output_height, int64_t output_width, + float scale_depth, float scale_height, float scale_width, + int32_t coordinate_mode, musaStream_t stream); diff --git a/musa/ep/src/kernels/tensor/resize_impl.mu b/musa/ep/src/kernels/tensor/resize_impl.mu index 279c8eb..9b35874 100644 --- a/musa/ep/src/kernels/tensor/resize_impl.mu +++ b/musa/ep/src/kernels/tensor/resize_impl.mu @@ -20,8 +20,15 @@ __global__ void ResizeNearest2DFloatKernel( const int64_t channel = remaining % channels; const int64_t batch = remaining / channels; - int64_t input_x = static_cast(output_x / scale_width); - int64_t input_y = static_cast(output_y / scale_height); + // ORT evaluates the nearest neighbour coordinate as `float(dim) / scale` + // (upsample.cc: SetupUpsampleNearestInputMappings). __fdiv_rn keeps that + // division IEEE-round-to-nearest on the device; the default lowering picks + // the approximate div.full.f32, which turns an exactly integral ratio such + // as 1 / (3.0f / 9.0f) into 2.9999998f and truncates one sample short. + int64_t input_x = static_cast( + __fdiv_rn(static_cast(output_x), scale_width)); + int64_t input_y = static_cast( + __fdiv_rn(static_cast(output_y), scale_height)); input_x = input_x < input_width ? input_x : input_width - 1; input_y = input_y < input_height ? input_y : input_height - 1; const int64_t input_index = @@ -31,6 +38,155 @@ __global__ void ResizeNearest2DFloatKernel( } } +// Mirrors the coordinate transformation used by the ORT CPU Resize kernel +// (GetOriginalCoordinateFromResizedCoordinate in upsamplebase.h). +// Every division goes through __fdiv_rn so the coordinates are IEEE-rounded +// like the host kernel instead of the approximate div.full.f32 the default +// device lowering emits; that difference is visible in the interpolation +// weights and, near an integer boundary, in the chosen input sample. +__device__ __forceinline__ float ResizeOriginalCoordinate( + int32_t coordinate_mode, int64_t output_index, float scale, + int64_t output_size, int64_t input_size) { + const float resized = static_cast(output_index); + switch (coordinate_mode) { + case kMusaResizePytorchHalfPixel: + return output_size > 1 ? __fdiv_rn(resized + 0.5f, scale) - 0.5f : 0.0f; + case kMusaResizeAlignCorners: + return output_size == 1 + ? 0.0f + : __fdiv_rn(resized * (static_cast(input_size) - 1.0f), + static_cast(output_size) - 1.0f); + case kMusaResizeAsymmetric: + return __fdiv_rn(resized, scale); + default: // kMusaResizeHalfPixel + return __fdiv_rn(resized + 0.5f, scale) - 0.5f; + } +} + +// Mirrors SetupUpsampleBilinear/SetupUpsampleTrilinear: when the scale is +// exactly 1 the coordinate is the index itself; otherwise it is transformed +// and clamped to the input range. The two neighbours always collapse to the +// same sample at the last index, where both weights become 0.5. +__device__ __forceinline__ void ResizeAxisWeights( + int32_t coordinate_mode, int64_t index, float scale, int64_t input_size, + int64_t output_size, int64_t& lower, int64_t& upper, float& weight1, + float& weight2) { + float original = scale == 1.0f + ? static_cast(index) + : ResizeOriginalCoordinate(coordinate_mode, index, scale, + output_size, input_size); + original = fmaxf(0.0f, fminf(original, static_cast(input_size - 1))); + lower = static_cast(original); + lower = lower < input_size - 1 ? lower : input_size - 1; + upper = lower + 1 < input_size - 1 ? lower + 1 : input_size - 1; + weight1 = fabsf(original - static_cast(lower)); + weight2 = fabsf(original - static_cast(upper)); + if (lower == upper) { + weight1 = 0.5f; + weight2 = 0.5f; + } +} + +__global__ void ResizeLinear2DFloatKernel( + const float* input, float* output, int64_t total_elements, int64_t channels, + int64_t input_height, int64_t input_width, int64_t output_height, + int64_t output_width, float scale_height, float scale_width, + int32_t coordinate_mode) { + const int64_t thread_id = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + const int64_t total_threads = static_cast(gridDim.x) * blockDim.x; + for (int64_t output_index = thread_id; output_index < total_elements; + output_index += total_threads) { + int64_t remaining = output_index; + const int64_t output_x = remaining % output_width; + remaining /= output_width; + const int64_t output_y = remaining % output_height; + remaining /= output_height; + const int64_t channel = remaining % channels; + const int64_t batch = remaining / channels; + + int64_t y1 = 0; + int64_t y2 = 0; + int64_t x1 = 0; + int64_t x2 = 0; + float dy1 = 0.0f; + float dy2 = 0.0f; + float dx1 = 0.0f; + float dx2 = 0.0f; + ResizeAxisWeights(coordinate_mode, output_y, scale_height, input_height, + output_height, y1, y2, dy1, dy2); + ResizeAxisWeights(coordinate_mode, output_x, scale_width, input_width, + output_width, x1, x2, dx1, dx2); + + const float* plane = + input + (batch * channels + channel) * input_height * input_width; + const float x11 = plane[y1 * input_width + x1]; + const float x21 = plane[y1 * input_width + x2]; + const float x12 = plane[y2 * input_width + x1]; + const float x22 = plane[y2 * input_width + x2]; + output[output_index] = + dx2 * dy2 * x11 + dx1 * dy2 * x21 + dx2 * dy1 * x12 + dx1 * dy1 * x22; + } +} + +__global__ void ResizeLinear3DFloatKernel( + const float* input, float* output, int64_t total_elements, int64_t channels, + int64_t input_depth, int64_t input_height, int64_t input_width, + int64_t output_depth, int64_t output_height, int64_t output_width, + float scale_depth, float scale_height, float scale_width, + int32_t coordinate_mode) { + const int64_t thread_id = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + const int64_t total_threads = static_cast(gridDim.x) * blockDim.x; + for (int64_t output_index = thread_id; output_index < total_elements; + output_index += total_threads) { + int64_t remaining = output_index; + const int64_t output_x = remaining % output_width; + remaining /= output_width; + const int64_t output_y = remaining % output_height; + remaining /= output_height; + const int64_t output_z = remaining % output_depth; + remaining /= output_depth; + const int64_t channel = remaining % channels; + const int64_t batch = remaining / channels; + + int64_t z1 = 0; + int64_t z2 = 0; + int64_t y1 = 0; + int64_t y2 = 0; + int64_t x1 = 0; + int64_t x2 = 0; + float dz1 = 0.0f; + float dz2 = 0.0f; + float dy1 = 0.0f; + float dy2 = 0.0f; + float dx1 = 0.0f; + float dx2 = 0.0f; + ResizeAxisWeights(coordinate_mode, output_z, scale_depth, input_depth, + output_depth, z1, z2, dz1, dz2); + ResizeAxisWeights(coordinate_mode, output_y, scale_height, input_height, + output_height, y1, y2, dy1, dy2); + ResizeAxisWeights(coordinate_mode, output_x, scale_width, input_width, + output_width, x1, x2, dx1, dx2); + + const int64_t plane_stride = input_height * input_width; + const float* volume = + input + (batch * channels + channel) * input_depth * plane_stride; + const float x111 = volume[z1 * plane_stride + y1 * input_width + x1]; + const float x211 = volume[z1 * plane_stride + y1 * input_width + x2]; + const float x121 = volume[z1 * plane_stride + y2 * input_width + x1]; + const float x221 = volume[z1 * plane_stride + y2 * input_width + x2]; + const float x112 = volume[z2 * plane_stride + y1 * input_width + x1]; + const float x212 = volume[z2 * plane_stride + y1 * input_width + x2]; + const float x122 = volume[z2 * plane_stride + y2 * input_width + x1]; + const float x222 = volume[z2 * plane_stride + y2 * input_width + x2]; + output[output_index] = dx2 * dy2 * dz2 * x111 + dx1 * dy2 * dz2 * x211 + + dx2 * dy1 * dz2 * x121 + dx1 * dy1 * dz2 * x221 + + dx2 * dy2 * dz1 * x112 + dx1 * dy2 * dz1 * x212 + + dx2 * dy1 * dz1 * x122 + dx1 * dy1 * dz1 * x222; + } +} + } // namespace musaError_t LaunchMusaResizeNearest2DFloat( @@ -49,3 +205,39 @@ musaError_t LaunchMusaResizeNearest2DFloat( output_height, output_width, scale_height, scale_width); return musaGetLastError(); } + +musaError_t LaunchMusaResizeLinear2DFloat( + const float* input, float* output, int64_t batch, int64_t channels, + int64_t input_height, int64_t input_width, int64_t output_height, + int64_t output_width, float scale_height, float scale_width, + int32_t coordinate_mode, musaStream_t stream) { + const int64_t total_elements = + batch * channels * output_height * output_width; + if (total_elements == 0) { + return musaSuccess; + } + ResizeLinear2DFloatKernel<<>>( + input, output, total_elements, channels, input_height, input_width, + output_height, output_width, scale_height, scale_width, coordinate_mode); + return musaGetLastError(); +} + +musaError_t LaunchMusaResizeLinear3DFloat( + const float* input, float* output, int64_t batch, int64_t channels, + int64_t input_depth, int64_t input_height, int64_t input_width, + int64_t output_depth, int64_t output_height, int64_t output_width, + float scale_depth, float scale_height, float scale_width, + int32_t coordinate_mode, musaStream_t stream) { + const int64_t total_elements = + batch * channels * output_depth * output_height * output_width; + if (total_elements == 0) { + return musaSuccess; + } + ResizeLinear3DFloatKernel<<>>( + input, output, total_elements, channels, input_depth, input_height, + input_width, output_depth, output_height, output_width, scale_depth, + scale_height, scale_width, coordinate_mode); + return musaGetLastError(); +} diff --git a/musa/ep/src/kernels/tensor/resize_support.h b/musa/ep/src/kernels/tensor/resize_support.h new file mode 100644 index 0000000..ef948a7 --- /dev/null +++ b/musa/ep/src/kernels/tensor/resize_support.h @@ -0,0 +1,177 @@ +// Copyright (c) Moore Threads Technology Co., Ltd. All rights reserved. +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#pragma once + +#include +#include +#include +#include +#include + +#include "onnxruntime_c_api.h" + +namespace musa_ep { + +// ONNX Resize (opset 13+) attribute defaults. +inline constexpr char kResizeDefaultMode[] = "nearest"; +inline constexpr char kResizeDefaultCoordinateTransformationMode[] = + "half_pixel"; +inline constexpr char kResizeDefaultNearestMode[] = "round_prefer_floor"; + +/// A zero-element optional input means "not provided": that is how the +/// exporters mark an unused roi/scales, so an omitted input and an empty one +/// are the same thing for this route. An unknown element count (a run-time +/// tensor seen from the graph) must count as provided, leaving the decision to +/// the kernel. +inline bool ResizeInputUnused(std::optional element_count) { + return element_count.has_value() && *element_count == 0; +} + +/// The coordinate transformations the linear route implements. +inline bool IsSupportedResizeLinearCoordinateMode(const std::string& mode) { + return mode == "half_pixel" || mode == "pytorch_half_pixel" || + mode == "align_corners" || mode == "asymmetric"; +} + +/// Reason string for the coordinate transformation modes the linear route does +/// not implement; shared so the graph-level gate and the kernel report the same +/// thing. +inline constexpr char kResizeLinearCoordinateReason[] = + "Resize linear only supports half_pixel, pytorch_half_pixel, " + "align_corners or asymmetric coordinates in this MUSA route"; + +/// Everything about a Resize node that decides whether this MUSA route can run +/// it, apart from the contents of its tensors. +/// +/// The kernel registry only matches op/domain/opset/dtype, so the graph-level +/// capability predicate (graph/capability_predicates.cc) inspects the node +/// itself to keep the Resize configurations this route cannot run on the CPU +/// EP. Both it and the kernel's runtime guard build this struct and ask +/// ResizeRouteRejection() below, so the nodes this route claims and the nodes +/// it runs are decided by one definition instead of two copies that drift +/// apart. +struct ResizeRouteConfig { + std::string mode = kResizeDefaultMode; + std::string coordinate_transformation_mode = + kResizeDefaultCoordinateTransformationMode; + std::string nearest_mode = kResizeDefaultNearestMode; + bool has_scales_input = false; + bool has_sizes_input = false; + /// Rank of the X input; -1 when it is unknown or not a tensor. + int64_t input_rank = -1; + ONNXTensorElementDataType input_element_type = + ONNX_TENSOR_ELEMENT_DATA_TYPE_UNDEFINED; +}; + +/// Tensor contents this route constrains. Callers only fill in what they can +/// observe - the capability predicate knows the constant initializers, the +/// kernel the actual runtime tensors. An unset field is not checked here and +/// stays the runtime guard's job. +struct ResizeRouteValues { + /// Element count of the (optional) roi input. Only an empty roi is supported. + std::optional roi_element_count; + /// Contents of the scales input. + std::optional> scales; + /// Contents of the sizes input. + std::optional> sizes; + /// Shape of X, compared against sizes to tell whether the batch and channel + /// axes are resized. + std::optional> input_shape; +}; + +/// Returns nullptr when this MUSA route runs \p config, otherwise a static +/// reason string suitable for an ORT_NOT_IMPLEMENTED status. +inline const char* ResizeRouteRejection(const ResizeRouteConfig& config) { + if (config.mode != "nearest" && config.mode != "linear") { + return "Resize only supports mode nearest or linear in this MUSA route"; + } + if (config.input_element_type != ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT) { + return "Resize only supports float32"; + } + if (config.has_scales_input == config.has_sizes_input) { + return "Resize requires exactly one of the scales or sizes input"; + } + if (config.input_rank != 4 && config.input_rank != 5) { + return "Resize only supports rank-4 NCHW and rank-5 NCDHW tensors in this " + "MUSA route"; + } + if (config.mode == "linear") { + if (!IsSupportedResizeLinearCoordinateMode( + config.coordinate_transformation_mode)) { + return kResizeLinearCoordinateReason; + } + return nullptr; + } + if (config.coordinate_transformation_mode != "asymmetric" || + config.nearest_mode != "floor") { + return "Resize only supports nearest/asymmetric/floor in this MUSA route"; + } + if (config.input_rank != 4) { + return "Resize nearest only supports rank-4 NCHW tensors in this MUSA " + "route"; + } + return nullptr; +} + +/// Value-dependent counterpart of ResizeRouteRejection(). +inline const char* ResizeRouteValueRejection(const ResizeRouteConfig& config, + const ResizeRouteValues& values) { + if (values.roi_element_count.has_value() && *values.roi_element_count != 0) { + return "Resize only supports an empty ROI (tf_crop_and_resize is not " + "implemented)"; + } + if (config.has_scales_input && values.scales.has_value()) { + const std::vector& scales = *values.scales; + if (config.input_rank > 0 && + static_cast(scales.size()) != config.input_rank) { + return "Resize scales rank must equal the input rank"; + } + for (float scale : scales) { + if (!std::isfinite(scale) || scale <= 0.0f) { + return "Resize requires positive finite scales"; + } + } + if (scales.size() >= 2 && (scales[0] != 1.0f || scales[1] != 1.0f)) { + return "Resize requires batch and channel scales to be 1 in this MUSA " + "route"; + } + } + if (config.has_sizes_input && values.sizes.has_value()) { + const std::vector& sizes = *values.sizes; + if (config.input_rank > 0 && + static_cast(sizes.size()) != config.input_rank) { + return "Resize sizes rank must equal the input rank"; + } + for (int64_t size : sizes) { + if (size <= 0) { + return "Resize requires positive sizes"; + } + } + // `sizes` resizes the batch and channel axes when it disagrees with the + // static shape of X; the scales rule above states the same constraint. + if (values.input_shape.has_value() && + values.input_shape->size() == sizes.size() && sizes.size() >= 2) { + const std::vector& shape = *values.input_shape; + const bool resizes_batch = shape[0] > 0 && sizes[0] != shape[0]; + const bool resizes_channel = shape[1] > 0 && sizes[1] != shape[1]; + if (resizes_batch || resizes_channel) { + return "Resize requires batch and channel scales to be 1 in this MUSA " + "route"; + } + } + } + return nullptr; +} + +} // namespace musa_ep diff --git a/test/ops/test_resize.py b/test/ops/test_resize.py index 6a23d31..a486bba 100644 --- a/test/ops/test_resize.py +++ b/test/ops/test_resize.py @@ -10,11 +10,196 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -"""End-to-end CPU-vs-MUSA tests for the supported Resize route.""" +"""End-to-end CPU-vs-MUSA tests for the supported Resize routes.""" import numpy as np +import onnxruntime as ort +import pytest +from onnx import TensorProto as OnnxTensorProto, helper -from op_test_utils import TensorProto, run_and_compare +from op_test_utils import ( + TensorProto, + build_graph_model, + musa_devices, + run, + run_and_compare, +) + +_RESIZE_NODE = "Resize_node" + + +def _resize_model(x, attrs, *, scales=None, sizes=None, opset=17): + """Build a single-node Resize model and its feed dict. + + Unlike ``build_model`` this lets an optional input be genuinely absent (an + empty node-input name) instead of an empty tensor, which is what the zoo + graphs do and what ONNX shape inference requires when ``sizes`` is used. + """ + node_inputs = ["X", "", "", ""] + feeds = {"X": x} + input_vis = [ + helper.make_tensor_value_info( + "X", helper.np_dtype_to_tensor_dtype(x.dtype), list(x.shape) + ) + ] + if scales is not None: + node_inputs[2] = "scales" + feeds["scales"] = scales + input_vis.append( + helper.make_tensor_value_info( + "scales", OnnxTensorProto.FLOAT, list(scales.shape) + ) + ) + if sizes is not None: + node_inputs[3] = "sizes" + feeds["sizes"] = sizes + input_vis.append( + helper.make_tensor_value_info( + "sizes", OnnxTensorProto.INT64, list(sizes.shape) + ) + ) + + node = helper.make_node("Resize", node_inputs, ["Y"], **attrs) + graph = helper.make_graph( + [node], + "Resize_graph", + input_vis, + [helper.make_tensor_value_info("Y", OnnxTensorProto.FLOAT, None)], + ) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", opset)]) + model.ir_version = min(model.ir_version, 10) + return model.SerializeToString(), feeds + + +def _initialized_resize_model(x, attrs, *, scales=None, sizes=None, roi=None): + """Single-node Resize whose optional inputs are constant initializers. + + The graph-level capability predicate can only inspect the value of a + constant, so the cases that depend on the roi/scales/sizes contents use this + builder instead of feeding those tensors at run time. + """ + node_inputs = ["X", "", "", ""] + initializers = [] + if roi is not None: + node_inputs[1] = "roi" + initializers.append( + helper.make_tensor( + "roi", + OnnxTensorProto.FLOAT, + list(roi.shape), + [float(value) for value in roi], + ) + ) + if scales is not None: + node_inputs[2] = "scales" + initializers.append( + helper.make_tensor( + "scales", + OnnxTensorProto.FLOAT, + list(scales.shape), + [float(value) for value in scales], + ) + ) + if sizes is not None: + node_inputs[3] = "sizes" + initializers.append( + helper.make_tensor( + "sizes", + OnnxTensorProto.INT64, + list(sizes.shape), + [int(value) for value in sizes], + ) + ) + + node = helper.make_node("Resize", node_inputs, ["Y"], name=_RESIZE_NODE, **attrs) + model_bytes = build_graph_model( + [node], + inputs={"X": x}, + outputs=[("Y", TensorProto.FLOAT)], + initializers=initializers, + opset=17, + name="Resize_graph", + ) + return model_bytes, {"X": x} + + +def _musa_session(model_bytes): + """Create a MUSA session that cannot silently fall back to CPU. + + Both fallbacks are disabled: the graph-level one (an unplaced node makes + session creation fail) and the runtime one (an EP failure at run time would + otherwise make ORT rebuild the whole session as CPU-only and hand back a + perfect CPU result). + """ + devices = musa_devices() + if not devices: + raise RuntimeError( + "No MUSA device available; refusing to build a CPU-only session " + "for the MUSA run (would compare CPU against CPU)." + ) + so = ort.SessionOptions() + so.add_session_config_entry("session.disable_cpu_ep_fallback", "1") + so.add_provider_for_devices(devices, {}) + session = ort.InferenceSession(model_bytes, sess_options=so) + session.disable_fallback() + return session + + +def _musa_session_with_assignment_report(model_bytes): + """MUSA session that reports which nodes it claimed and hides no failure. + + The graph-level CPU fallback stays enabled, because a node the MUSA EP + declines has to be able to run on the CPU EP; the run-time EP-error fallback + is disabled so a node the MUSA EP claimed and then failed inside Compute() + raises instead of silently re-running the whole session on the CPU (which + would make the output comparison meaningless). ``dtype_diagnostics`` makes + the EP print one ``MUSA_DTYPE_NODE`` line per node with ``assigned=1`` when + it registered that node and ``assigned=0`` when it declined it. + """ + devices = musa_devices() + if not devices: + raise RuntimeError( + "No MUSA device available; refusing to build a CPU-only session " + "for the MUSA run (would compare CPU against CPU)." + ) + so = ort.SessionOptions() + so.add_provider_for_devices(devices, {"dtype_diagnostics": "1"}) + session = ort.InferenceSession(model_bytes, sess_options=so) + session.disable_fallback() + return session + + +def _musa_assignment(diagnostics): + """Map node name -> claimed-by-MUSA from the EP's ``MUSA_DTYPE_NODE`` lines.""" + assert "MUSA_DTYPE_REPORT_BEGIN" in diagnostics, ( + "the MUSA EP did not emit its dtype/assignment report:\n" + diagnostics + ) + assignment = {} + for line in diagnostics.splitlines(): + if "MUSA_DTYPE_NODE" not in line: + continue + fields = dict( + field.split("=", 1) for field in line.split() if "=" in field + ) + if "name" in fields: + assignment[fields["name"]] = fields["assigned"] == "1" + return assignment + + +def _run_resize(x, attrs, *, scales=None, sizes=None): + """Run the Resize model on the CPU EP and the MUSA EP, then compare.""" + model_bytes, feeds = _resize_model(x, attrs, scales=scales, sizes=sizes) + cpu_outputs = run(model_bytes, feeds, use_musa=False) + musa_outputs = _musa_session(model_bytes).run(None, dict(feeds)) + + assert len(cpu_outputs) == len(musa_outputs) + for cpu_output, musa_output in zip(cpu_outputs, musa_outputs): + assert cpu_output.shape == musa_output.shape + assert cpu_output.dtype == musa_output.dtype + np.testing.assert_allclose( + musa_output, cpu_output, rtol=1e-3, atol=1e-4 + ) + return musa_outputs def test_resize_opset17_float_nchw_nearest_asymmetric_floor_scales(): @@ -33,3 +218,342 @@ def test_resize_opset17_float_nchw_nearest_asymmetric_floor_scales(): rtol=0, atol=0, ) + + +def test_resize_opset17_float_nchw_nearest_asymmetric_floor_sizes(): + x = np.arange(1 * 3 * 9 * 13, dtype=np.float32).reshape(1, 3, 9, 13) + sizes = np.array([1, 3, 3, 4], dtype=np.int64) + _run_resize( + x, + { + "mode": "nearest", + "coordinate_transformation_mode": "asymmetric", + "nearest_mode": "floor", + }, + sizes=sizes, + ) + + +def test_resize_opset17_float_nchw_nearest_asymmetric_floor_exact_ratio_sizes(): + # 6/2 == 3 and 15/3 == 5 give coordinates that are integral in float32 but + # just below the integer in double, so a wider quotient would truncate one + # sample short (regression guard for the nearest index computation). + x = np.arange(1 * 2 * 6 * 15, dtype=np.float32).reshape(1, 2, 6, 15) + sizes = np.array([1, 2, 2, 3], dtype=np.int64) + _run_resize( + x, + { + "mode": "nearest", + "coordinate_transformation_mode": "asymmetric", + "nearest_mode": "floor", + }, + sizes=sizes, + ) + + +def test_resize_opset17_float_nchw_linear_half_pixel_scales(): + x = np.arange(1 * 3 * 5 * 7, dtype=np.float32).reshape(1, 3, 5, 7) + scales = np.array([1.0, 1.0, 1.7, 1.4], dtype=np.float32) + _run_resize( + x, + {"mode": "linear", "coordinate_transformation_mode": "half_pixel"}, + scales=scales, + ) + + +def test_resize_opset17_float_nchw_linear_half_pixel_sizes(): + x = np.arange(2 * 3 * 5 * 7, dtype=np.float32).reshape(2, 3, 5, 7) - 21.0 + sizes = np.array([2, 3, 9, 13], dtype=np.int64) + _run_resize( + x, + {"mode": "linear", "coordinate_transformation_mode": "half_pixel"}, + sizes=sizes, + ) + + +def test_resize_opset17_float_nchw_linear_half_pixel_downsample_sizes(): + x = np.arange(1 * 3 * 9 * 13, dtype=np.float32).reshape(1, 3, 9, 13) + sizes = np.array([1, 3, 4, 5], dtype=np.int64) + _run_resize( + x, + {"mode": "linear", "coordinate_transformation_mode": "half_pixel"}, + sizes=sizes, + ) + + +def test_resize_opset17_float_nchw_linear_pytorch_half_pixel_sizes(): + x = np.arange(1 * 2 * 5 * 7, dtype=np.float32).reshape(1, 2, 5, 7) + sizes = np.array([1, 2, 3, 3], dtype=np.int64) + _run_resize( + x, + {"mode": "linear", "coordinate_transformation_mode": "pytorch_half_pixel"}, + sizes=sizes, + ) + + +def test_resize_opset17_float_nchw_linear_align_corners_sizes(): + x = np.arange(1 * 2 * 5 * 7, dtype=np.float32).reshape(1, 2, 5, 7) + sizes = np.array([1, 2, 9, 13], dtype=np.int64) + _run_resize( + x, + {"mode": "linear", "coordinate_transformation_mode": "align_corners"}, + sizes=sizes, + ) + + +def test_resize_opset17_float_nchw_linear_asymmetric_scales(): + x = np.arange(1 * 2 * 5 * 7, dtype=np.float32).reshape(1, 2, 5, 7) + scales = np.array([1.0, 1.0, 2.0, 2.0], dtype=np.float32) + _run_resize( + x, + {"mode": "linear", "coordinate_transformation_mode": "asymmetric"}, + scales=scales, + ) + + +def test_resize_opset17_float_ncdhw_linear_half_pixel_scales(): + x = np.arange(1 * 2 * 2 * 3 * 4, dtype=np.float32).reshape(1, 2, 2, 3, 4) + scales = np.array([1.0, 1.0, 2.0, 1.5, 2.0], dtype=np.float32) + _run_resize( + x, + {"mode": "linear", "coordinate_transformation_mode": "half_pixel"}, + scales=scales, + ) + + +def test_resize_opset17_float_ncdhw_linear_half_pixel_sizes(): + x = np.arange(1 * 2 * 2 * 3 * 4, dtype=np.float32).reshape(1, 2, 2, 3, 4) + sizes = np.array([1, 2, 4, 5, 9], dtype=np.int64) + _run_resize( + x, + {"mode": "linear", "coordinate_transformation_mode": "half_pixel"}, + sizes=sizes, + ) + + +# Configurations the MUSA route cannot run but the CPU EP can: the route only +# implements nearest/asymmetric/floor and the linear coordinate transformations, +# on rank-4 NCHW and rank-5 NCDHW tensors, with an empty roi and with the batch +# and channel axes left alone. Each case below trips exactly one rule, and all +# of them carry the values the rule needs as constant initializers so the +# graph-level gate can see them. +_UNSUPPORTED_RESIZE_CASES = { + "cubic_sizes": dict( + attrs={"mode": "cubic", "coordinate_transformation_mode": "half_pixel"}, + x_shape=(1, 2, 5, 7), + sizes=np.array([1, 2, 9, 13], dtype=np.int64), + ), + "cubic_scales": dict( + attrs={"mode": "cubic", "coordinate_transformation_mode": "half_pixel"}, + x_shape=(1, 2, 5, 7), + scales=np.array([1.0, 1.0, 2.0, 2.0], dtype=np.float32), + ), + "nearest_half_pixel_scales": dict( + attrs={"mode": "nearest", "coordinate_transformation_mode": "half_pixel"}, + x_shape=(1, 2, 5, 7), + scales=np.array([1.0, 1.0, 2.0, 2.0], dtype=np.float32), + ), + "nearest_round_prefer_floor_scales": dict( + attrs={ + "mode": "nearest", + "coordinate_transformation_mode": "asymmetric", + "nearest_mode": "round_prefer_floor", + }, + x_shape=(1, 2, 5, 7), + scales=np.array([1.0, 1.0, 2.0, 2.0], dtype=np.float32), + ), + "nearest_rank5_scales": dict( + attrs={ + "mode": "nearest", + "coordinate_transformation_mode": "asymmetric", + "nearest_mode": "floor", + }, + x_shape=(1, 1, 2, 3, 4), + scales=np.array([1.0, 1.0, 2.0, 2.0, 2.0], dtype=np.float32), + ), + "nearest_nonempty_roi_scales": dict( + attrs={ + "mode": "nearest", + "coordinate_transformation_mode": "asymmetric", + "nearest_mode": "floor", + }, + x_shape=(1, 2, 5, 7), + roi=np.array([0.0, 0.0, 0.0, 0.0, 1.0, 1.0, 1.0, 1.0], dtype=np.float32), + scales=np.array([1.0, 1.0, 2.0, 2.0], dtype=np.float32), + ), + "nearest_channel_scales": dict( + attrs={ + "mode": "nearest", + "coordinate_transformation_mode": "asymmetric", + "nearest_mode": "floor", + }, + x_shape=(1, 2, 5, 7), + scales=np.array([1.0, 2.0, 2.0, 2.0], dtype=np.float32), + ), + "nearest_batch_sizes": dict( + attrs={ + "mode": "nearest", + "coordinate_transformation_mode": "asymmetric", + "nearest_mode": "floor", + }, + x_shape=(1, 2, 5, 7), + sizes=np.array([2, 2, 10, 14], dtype=np.int64), + ), + "linear_rank3_sizes": dict( + attrs={"mode": "linear", "coordinate_transformation_mode": "half_pixel"}, + x_shape=(1, 5, 7), + sizes=np.array([1, 9, 13], dtype=np.int64), + ), + "tf_crop_and_resize_sizes": dict( + attrs={"mode": "linear", "coordinate_transformation_mode": "tf_crop_and_resize"}, + x_shape=(1, 2, 5, 7), + roi=np.array([0.0, 0.0, 0.0, 0.0, 1.0, 1.0, 0.6, 0.6], dtype=np.float32), + sizes=np.array([1, 2, 10, 14], dtype=np.int64), + ), +} + + +def _unsupported_resize_model(case): + spec = _UNSUPPORTED_RESIZE_CASES[case] + x = np.arange( + int(np.prod(spec["x_shape"])), dtype=np.float32 + ).reshape(spec["x_shape"]) + optional = { + key: value + for key, value in spec.items() + if key not in ("attrs", "x_shape") + } + return _initialized_resize_model(x, spec["attrs"], **optional) + + +@pytest.mark.parametrize("case", list(_UNSUPPORTED_RESIZE_CASES)) +def test_resize_unsupported_configuration_stays_on_cpu(case, capfd): + """Configurations the MUSA kernel cannot run must be left to the CPU EP. + + Each of these is a valid ONNX Resize the CPU EP implements, but not one the + MUSA route does, so the EP has to decline the node during GetCapability + instead of claiming it and failing inside Compute(). The session keeps the + graph-level CPU fallback (that is where the node must land) but disables the + run-time EP-error fallback, so a node the MUSA EP claimed and then failed + would raise here rather than silently re-running on the CPU, and the EP's + own assignment report then shows ``assigned=0`` for the node. + """ + model_bytes, feeds = _unsupported_resize_model(case) + cpu_outputs = run(model_bytes, feeds, use_musa=False) + + session = _musa_session_with_assignment_report(model_bytes) + musa_outputs = session.run(None, dict(feeds)) + + assert len(cpu_outputs) == len(musa_outputs) + for cpu_output, musa_output in zip(cpu_outputs, musa_outputs): + assert cpu_output.shape == musa_output.shape + assert cpu_output.dtype == musa_output.dtype + np.testing.assert_allclose(musa_output, cpu_output, rtol=0, atol=0) + + assignment = _musa_assignment(capfd.readouterr().err) + assert assignment[_RESIZE_NODE] is False + + +# Configurations the MUSA route does run, with the values that decide them +# present as constant initializers, so the gate has every chance to decline them +# and must still claim the node. +_SUPPORTED_RESIZE_CASES = { + "nearest_scales": dict( + attrs={ + "mode": "nearest", + "coordinate_transformation_mode": "asymmetric", + "nearest_mode": "floor", + }, + x_shape=(1, 3, 5, 7), + scales=np.array([1.0, 1.0, 2.0, 2.0], dtype=np.float32), + ), + "linear_sizes": dict( + attrs={"mode": "linear", "coordinate_transformation_mode": "half_pixel"}, + x_shape=(2, 3, 5, 7), + sizes=np.array([2, 3, 9, 13], dtype=np.int64), + ), + "linear_rank5_scales": dict( + attrs={"mode": "linear", "coordinate_transformation_mode": "half_pixel"}, + x_shape=(1, 2, 2, 3, 4), + scales=np.array([1.0, 1.0, 2.0, 1.5, 2.0], dtype=np.float32), + ), + # Paddle exports drive Resize through `sizes` and pass an empty roi and an + # empty scales initializer next to it. Those mean "not provided" - the + # kernel reads a zero-element input as absent - so they must not be mistaken + # for a second, conflicting input and stay on the MUSA EP. + "paddle_empty_roi_scales_linear_sizes": dict( + attrs={"mode": "linear", "coordinate_transformation_mode": "half_pixel"}, + x_shape=(1, 3, 5, 7), + roi=np.empty((0,), dtype=np.float32), + scales=np.empty((0,), dtype=np.float32), + sizes=np.array([1, 3, 9, 13], dtype=np.int64), + ), + "paddle_empty_roi_scales_nearest_sizes": dict( + attrs={ + "mode": "nearest", + "coordinate_transformation_mode": "asymmetric", + "nearest_mode": "floor", + }, + x_shape=(1, 3, 5, 7), + roi=np.empty((0,), dtype=np.float32), + scales=np.empty((0,), dtype=np.float32), + sizes=np.array([1, 3, 9, 13], dtype=np.int64), + ), +} + + +@pytest.mark.parametrize("case", list(_SUPPORTED_RESIZE_CASES)) +def test_resize_supported_configuration_is_claimed_by_musa(case, capfd): + """The supported route keeps the node on the MUSA EP, and computes it there. + + The assignment report proves the EP claimed the node (``assigned=1``), and + the run-time EP-error fallback is disabled, so the matching output also + proves the MUSA kernel computed it instead of the CPU EP. + """ + spec = _SUPPORTED_RESIZE_CASES[case] + x = np.arange( + int(np.prod(spec["x_shape"])), dtype=np.float32 + ).reshape(spec["x_shape"]) + optional = { + key: value + for key, value in spec.items() + if key not in ("attrs", "x_shape") + } + model_bytes, feeds = _initialized_resize_model(x, spec["attrs"], **optional) + cpu_outputs = run(model_bytes, feeds, use_musa=False) + + session = _musa_session_with_assignment_report(model_bytes) + musa_outputs = session.run(None, dict(feeds)) + + assignment = _musa_assignment(capfd.readouterr().err) + assert assignment[_RESIZE_NODE] is True + + assert len(cpu_outputs) == len(musa_outputs) + for cpu_output, musa_output in zip(cpu_outputs, musa_outputs): + assert cpu_output.shape == musa_output.shape + assert cpu_output.dtype == musa_output.dtype + np.testing.assert_allclose(musa_output, cpu_output, rtol=1e-3, atol=1e-4) + + +def test_resize_run_time_only_rejection_still_fails_loudly(): + """The kernel guard still rejects what the graph-level gate cannot see. + + ``scales`` is a run-time input here, so its values are unknown during + GetCapability and the EP claims the node; the kernel must then reject the + batch/channel scaling itself instead of running it. Both fallbacks are + disabled by ``_musa_session``, so the rejection cannot be hidden by a + CPU-only re-run of the session. + """ + x = np.arange(1 * 2 * 5 * 7, dtype=np.float32).reshape(1, 2, 5, 7) + scales = np.array([1.0, 2.0, 2.0, 2.0], dtype=np.float32) + model_bytes, feeds = _resize_model( + x, + { + "mode": "nearest", + "coordinate_transformation_mode": "asymmetric", + "nearest_mode": "floor", + }, + scales=scales, + ) + with pytest.raises(Exception, match="batch and channel scales to be 1"): + _musa_session(model_bytes).run(None, dict(feeds)) From 3c6e8cb858ce62059dbcf986c92ad954eb00356b Mon Sep 17 00:00:00 2001 From: "yichao.hu" Date: Mon, 12 Oct 2026 23:47:31 -0400 Subject: [PATCH 14/52] fix(plugin): keep shape-metadata chains on the CPU and order small H2D copies on the compute stream With the ORT memory pattern enabled, models failed from the second run on (broadcast with a 0-sized dim, Reshape overflow). Small (<1024 B) pageable host-to-device copies used a blocking memcpy on the null stream, not ordered with the compute stream, so a pending kernel could overwrite the copied shape metadata. Stage small copies through pinned memory and copy them asynchronously on the compute stream. Also declare Slice's starts/ends/axes/ steps as CPU inputs so Shape->Gather->Unsqueeze metadata chains stay on the CPU EP. Adds a discriminating test that fails before the fix and passes after. --- musa/ep/src/ep_data_transfer.cc | 54 +++-- musa/ep/src/graph/cpu_preferred_nodes.cc | 19 ++ musa/ep/src/kernels/tensor/slice.cc | 9 +- test/ops/test_boundary_copy_ordering.py | 128 +++++++++++ test/ops/test_dynamic_shape_metadata.py | 131 +++++++++++ .../test_pageable_h2d_copy_stream_ordering.py | 211 ++++++++++++++++++ 6 files changed, 536 insertions(+), 16 deletions(-) create mode 100644 test/ops/test_boundary_copy_ordering.py create mode 100644 test/ops/test_dynamic_shape_metadata.py create mode 100644 test/ops/test_pageable_h2d_copy_stream_ordering.py diff --git a/musa/ep/src/ep_data_transfer.cc b/musa/ep/src/ep_data_transfer.cc index ebe18e6..730386f 100644 --- a/musa/ep/src/ep_data_transfer.cc +++ b/musa/ep/src/ep_data_transfer.cc @@ -108,21 +108,21 @@ void TraceD2hDeviceSynchronize(const OrtValue* src_tensor, #endif } -OrtStatus* CopyPageableHostToDevice(MusaDataTransfer& impl, - const void* src_data, void* dst_data, - size_t bytes, musaStream_t stream, - bool allow_pageable_bounce) { - constexpr size_t kDefaultBounceThresholdBytes = 1024; - const size_t bounce_threshold = - EnvSizeOrDefault("ORT_MUSA_PAGEABLE_H2D_BOUNCE_THRESHOLD_BYTES", - kDefaultBounceThresholdBytes); - if (bytes < bounce_threshold) { - return MusaStatus(impl.ort_api_, musaMemcpy(dst_data, src_data, bytes, - musaMemcpyHostToDevice)); - } - - if (!allow_pageable_bounce || stream == nullptr || - impl.pinned_host_pool_ == nullptr || +// Pageable host -> device transfer. Every such write must be issued on the +// caller's stream: a blocking musaMemcpy() on the null stream is ordered +// against neither the kernels that are still running on the EP's compute stream +// nor the consumers enqueued after it. ORT's memory pattern routinely hands the +// destination of a CPU->device boundary copy - a few bytes such as a Shape +// output - the block of a tensor that is already dead on the host side but +// whose producer kernel has not finished on the device, so the copy read back +// the previous owner's bytes and the shape chain fed garbage sizes to its +// consumers (Add/Mul broadcast mismatch, Reshape "invalid argument"). +// Staging through the pinned pool keeps the copy ordered on the caller's stream +// and lets the pageable host source be released immediately. +OrtStatus* CopyStagedPageableHostToDevice(MusaDataTransfer& impl, + const void* src_data, void* dst_data, + size_t bytes, musaStream_t stream) { + if (stream == nullptr || impl.pinned_host_pool_ == nullptr || !EnvFlagEnabled("ORT_MUSA_ENABLE_PAGEABLE_H2D_BOUNCE", true)) { return CopyMemcpy(impl.ort_api_, src_data, dst_data, bytes, musaMemcpyHostToDevice, stream); @@ -146,6 +146,30 @@ OrtStatus* CopyPageableHostToDevice(MusaDataTransfer& impl, return nullptr; } +OrtStatus* CopyPageableHostToDevice(MusaDataTransfer& impl, + const void* src_data, void* dst_data, + size_t bytes, musaStream_t stream, + bool allow_pageable_bounce) { + constexpr size_t kDefaultBounceThresholdBytes = 1024; + const size_t bounce_threshold = + EnvSizeOrDefault("ORT_MUSA_PAGEABLE_H2D_BOUNCE_THRESHOLD_BYTES", + kDefaultBounceThresholdBytes); + if (bytes < bounce_threshold) { + return CopyStagedPageableHostToDevice(impl, src_data, dst_data, bytes, + stream); + } + + if (!allow_pageable_bounce || stream == nullptr || + impl.pinned_host_pool_ == nullptr || + !EnvFlagEnabled("ORT_MUSA_ENABLE_PAGEABLE_H2D_BOUNCE", true)) { + return CopyMemcpy(impl.ort_api_, src_data, dst_data, bytes, + musaMemcpyHostToDevice, stream); + } + + return CopyStagedPageableHostToDevice(impl, src_data, dst_data, bytes, + stream); +} + OrtStatus* CopyImpl(MusaDataTransfer& impl, const OrtMemoryDevice* src_device, const OrtMemoryDevice* dst_device, const void* src_data, void* dst_data, size_t bytes, musaStream_t stream, diff --git a/musa/ep/src/graph/cpu_preferred_nodes.cc b/musa/ep/src/graph/cpu_preferred_nodes.cc index 0b8efba..bd2a586 100644 --- a/musa/ep/src/graph/cpu_preferred_nodes.cc +++ b/musa/ep/src/graph/cpu_preferred_nodes.cc @@ -245,6 +245,25 @@ OrtStatus* GetCpuPreferredNodes( } } + // Outputs of nodes this EP does not claim (CPU EP nodes, unfused Constant + // nodes, ...) are already host memory. Record them before the walk so a + // metadata node reading such a tensor is still recognised as CPU-readable + // - the walk itself only discovers them when it happens to pop them. + for (const Ort::ConstNode& node : ordered_nodes) { + if (provider_nodes.find(node.GetId()) != provider_nodes.end()) { + continue; + } + std::string ep_name = node.GetEpName(); + if (!ep_name.empty() && ep_name != "CPUExecutionProvider") { + continue; + } + for (Ort::ConstValueInfo output : node.GetOutputs()) { + if (output != nullptr) { + cpu_output_args.insert(output); + } + } + } + std::unordered_set visited; visited.reserve(candidates.size()); std::unordered_set cpu_nodes; diff --git a/musa/ep/src/kernels/tensor/slice.cc b/musa/ep/src/kernels/tensor/slice.cc index 7428fd7..2dc322e 100644 --- a/musa/ep/src/kernels/tensor/slice.cc +++ b/musa/ep/src/kernels/tensor/slice.cc @@ -182,5 +182,12 @@ ONNX_OPERATOR_VERSIONED_KERNEL_EX( Slice, kOnnxDomain, 13, 19, (Ort::KernelDefBuilder() .AddTypeConstraint("T", AllFixedSizeTensorTypes()) - .AddTypeConstraint("Tind", IntTensorTypes())), + .AddTypeConstraint("Tind", IntTensorTypes()) + // starts/ends/axes/steps are shape metadata: keep them in host memory + // (the in-tree CUDA kernel declares the same) so ORT orders the + // CPU/device boundary itself instead of leaving it inside the kernel. + .SetInputMemType(1, OrtMemTypeCPUInput) + .SetInputMemType(2, OrtMemTypeCPUInput) + .SetInputMemType(3, OrtMemTypeCPUInput) + .SetInputMemType(4, OrtMemTypeCPUInput)), Slice) diff --git a/test/ops/test_boundary_copy_ordering.py b/test/ops/test_boundary_copy_ordering.py new file mode 100644 index 0000000..967c833 --- /dev/null +++ b/test/ops/test_boundary_copy_ordering.py @@ -0,0 +1,128 @@ +# Copyright (c) Moore Threads Technology Co., Ltd. All rights reserved. +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""CPU/device boundary copies must be ordered on the EP's compute stream. + +A small tensor produced on the CPU side (a `Shape`/`Slice` metadata chain) that is +consumed by a device kernel is transferred by the plugin's MemcpyFromHost kernel. +That transfer must be issued on the caller's compute stream: a blocking +musaMemcpy() on the null stream is ordered against neither the kernels that are +still running on the compute stream (ORT's memory pattern may have handed the +destination block of a tensor whose last consumer just ran) nor the consumers +that are enqueued after it. + +Regression cover for sherpa_whisper_base / zipformer_encoder_en / +campplus_3dspeaker / campplus_cosyvoice: with ORT graph optimisation on, the +transferred metadata was read back as the previous owner's bytes, so the shape +chain produced garbage sizes (`Add`/`Mul` broadcast mismatch, `Reshape` +"invalid argument", NaNs). Every case therefore runs the graph twice inside one +session - the first run allocates dynamically, the second reuses the memory +pattern - and compares both runs against the CPU EP. +""" + +import numpy as np +import onnxruntime as ort +import pytest +from onnx import TensorProto, helper + +from op_test_utils import build_graph_model, musa_devices + +_RUNS = 2 + + +def _session(model_bytes, *, musa): + options = ort.SessionOptions() + if musa: + devices = musa_devices() + if not devices: + raise RuntimeError("No MUSA device available; refusing a CPU-only session") + options.add_provider_for_devices(devices, {}) + session = ort.InferenceSession(model_bytes, sess_options=options) + # Never let ORT silently rebuild the session CPU-only after an EP failure. + session.disable_fallback() + return session + return ort.InferenceSession(model_bytes, sess_options=options, + providers=["CPUExecutionProvider"]) + + +def _metadata_boundary_copy_model(copy_bytes, tiny_writers): + """`Mul` (device) -> `Shape` -> `Slice` (CPU) -> `Cast` (device) -> `Mul`. + + The `Cast` node runs on the device and consumes CPU-side metadata, so ORT's + MemcpyTransformer has to insert a MemcpyFromHost for `copy_bytes` bytes. The + `tiny_writers` device tensors of the same size that are produced and dropped + right before the copy make ORT's memory pattern hand the copy destination a + block that a kernel which is still running on the device wrote. + """ + rng = np.random.RandomState(0) + x = rng.standard_normal((1, 1500, 512)).astype(np.float32) + nodes = [] + initializers = [ + helper.make_tensor("c", TensorProto.FLOAT, [1], [1.0]), + helper.make_tensor("ax0", TensorProto.INT64, [1], [0]), + helper.make_tensor("ax012", TensorProto.INT64, [3], [0, 1, 2]), + helper.make_tensor("s0", TensorProto.INT64, [1], [0]), + ] + + for index in range(tiny_writers): + # A reduction over the whole activation: enqueued right before the + # boundary copy, still running device-side when the copy is issued. + if copy_bytes == 4: + # float32 scalar -> 4 bytes + nodes.append(helper.make_node("ReduceSum", ["X", "ax012"], [f"t{index}"], + keepdims=1)) + nodes.append(helper.make_node("Mul", [f"t{index}", "c"], [f"u{index}"])) + else: + # int64 element -> 8 bytes (dtype and size of the copied metadata) + nodes.append(helper.make_node("ReduceSum", ["X", "ax012"], [f"t{index}"], + keepdims=0)) + nodes.append(helper.make_node("Cast", [f"t{index}"], [f"t{index}i"], + to=TensorProto.INT64)) + nodes.append(helper.make_node("Mul", [f"t{index}i", "ax0"], [f"u{index}"])) + + nodes.append(helper.make_node("Shape", ["X"], ["sc"])) # CPU, rank 3 + if copy_bytes == 8: + initializers.append(helper.make_tensor("s1", TensorProto.INT64, [1], [1])) + nodes.append(helper.make_node("Slice", ["sc", "s0", "s1", "ax0"], ["meta"])) + elif copy_bytes == 4: + # CPU-side int32 slice: 4 bytes cross the boundary + initializers.append(helper.make_tensor("s1", TensorProto.INT64, [1], [1])) + nodes.append(helper.make_node("Slice", ["sc", "s0", "s1", "ax0"], ["slice0"])) + nodes.append(helper.make_node("Cast", ["slice0"], ["meta"], to=TensorProto.INT32)) + else: + raise ValueError(copy_bytes) + nodes.append(helper.make_node("Cast", ["meta"], ["vf"], to=TensorProto.FLOAT)) + nodes.append(helper.make_node("Mul", ["X", "vf"], ["Y"])) + + model = build_graph_model( + nodes, + inputs={"X": x}, + outputs=[("vf", TensorProto.FLOAT), ("Y", TensorProto.FLOAT)], + initializers=initializers, + opset=17, + name=f"boundary_copy_{copy_bytes}_{tiny_writers}", + ) + return model, {"X": x} + + +@pytest.mark.parametrize("copy_bytes,tiny_writers", [(4, 1), (4, 4), (8, 1), (8, 4)]) +def test_small_boundary_copy_stays_ordered(copy_bytes, tiny_writers): + model, feeds = _metadata_boundary_copy_model(copy_bytes, tiny_writers) + cpu_vf, cpu_y = _session(model, musa=False).run(None, dict(feeds)) + session = _session(model, musa=True) + for run_index in range(_RUNS): + vf, y = session.run(None, dict(feeds)) + np.testing.assert_array_equal(np.asarray(vf), np.asarray(cpu_vf), + err_msg=f"run {run_index}: vf") + np.testing.assert_allclose(np.asarray(y), np.asarray(cpu_y), rtol=1e-4, atol=1e-3, + err_msg=f"run {run_index}: Y") + assert np.isfinite(np.asarray(y)).all(), f"run {run_index}: Y has NaN/Inf" diff --git a/test/ops/test_dynamic_shape_metadata.py b/test/ops/test_dynamic_shape_metadata.py new file mode 100644 index 0000000..5c50e3b --- /dev/null +++ b/test/ops/test_dynamic_shape_metadata.py @@ -0,0 +1,131 @@ +# Copyright (c) Moore Threads Technology Co., Ltd. All rights reserved. +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Dynamic-shape metadata chains that feed Slice. + +`Shape`/`Gather`/`Unsqueeze` arithmetic produces the runtime shapes of a graph. +Where that arithmetic derives from a device tensor it must not be executed as a +device-side subgraph: the values belong on the CPU EP, and the operators that +consume them (`Slice` starts/ends/axes/steps, `Reshape` shape, ...) declare those +inputs as `OrtMemTypeCPUInput`, exactly like the in-tree CUDA kernels do. + +Regression cover for sherpa_whisper_base / campplus_3dspeaker and +zipformer_encoder_en: with ORT graph optimisation on, their `Slice` received +`ends=[0]` because the `Shape -> Gather -> Unsqueeze -> Slice` chain ran on the +MUSA EP and ORT's memory pattern later reused the tiny device buffers of that +chain for multi-MB activations. Running the model twice inside one session is +what exposes it (the first run allocates dynamically, the second reuses the +pattern), so every case below is executed twice per provider. +""" + +import numpy as np +import onnxruntime as ort +import pytest +from onnx import TensorProto, helper + +from op_test_utils import build_graph_model, musa_devices + +_RUNS = 2 + + +def _session(model_bytes, *, musa): + options = ort.SessionOptions() + if musa: + devices = musa_devices() + if not devices: + raise RuntimeError("No MUSA device available; refusing a CPU-only session") + options.add_provider_for_devices(devices, {}) + session = ort.InferenceSession(model_bytes, sess_options=options) + # Never let ORT silently rebuild the session CPU-only after an EP failure. + session.disable_fallback() + return session + return ort.InferenceSession(model_bytes, sess_options=options, providers=["CPUExecutionProvider"]) + + +def _run_twice(session, feeds): + return [session.run(None, dict(feeds)) for _ in range(_RUNS)] + + +def _slice_ends_chain_model(permute_shape_chain): + """`Add` (device) -> `Shape` -> `Gather` -> `Unsqueeze` -> `Slice(ends=..)`. + + `pos` is [1500, 512] and the `ends` value is the runtime length of the first + axis of `x`, so the slice is a no-op: the output must equal `pos` exactly on + every run. Anything else means the metadata lost its value. + """ + rng = np.random.RandomState(0) + # The permuted variant mirrors ORT's TransposeOptimizer rewrite of + # `Shape(Transpose(x))`: the Shape reads the *pre*-transpose tensor, so the + # device tensor is [1, 512, 1500] and the permuted gather yields [1, 1500, 512]. + x = rng.standard_normal((1, 512, 1500) if permute_shape_chain else (1, 1500, 512)).astype(np.float32) + pos = rng.standard_normal((1500, 512)).astype(np.float32) + + nodes = [helper.make_node("Add", ["X", "zero"], ["dev"]), + helper.make_node("Shape", ["dev"], ["shape"])] + if permute_shape_chain: + # Mirrors ORT's TransposeOptimizer rewrite of `Shape(Transpose(x))`. + nodes.append(helper.make_node("Gather", ["shape", "perm"], ["shape_perm"], axis=0)) + nodes.append(helper.make_node("Gather", ["shape_perm", "one"], ["dim"], axis=0)) + else: + nodes.append(helper.make_node("Gather", ["shape", "one"], ["dim"], axis=0)) + nodes.append(helper.make_node("Unsqueeze", ["dim", "axis0"], ["ends"])) + nodes.append(helper.make_node("Slice", ["pos", "zero_vec", "ends", "axis0", "step1"], ["Y"])) + + initializers = [ + helper.make_tensor("zero", TensorProto.FLOAT, [1], [0.0]), + helper.make_tensor("perm", TensorProto.INT64, [3], [0, 2, 1]), + helper.make_tensor("one", TensorProto.INT64, [], [1]), + helper.make_tensor("axis0", TensorProto.INT64, [1], [0]), + helper.make_tensor("zero_vec", TensorProto.INT64, [1], [0]), + helper.make_tensor("step1", TensorProto.INT64, [1], [1]), + helper.make_tensor("pos", TensorProto.FLOAT, [1500, 512], pos.ravel().tolist()), + ] + model = build_graph_model( + nodes, + inputs={"X": x}, + outputs=[("ends", TensorProto.INT64), ("Y", TensorProto.FLOAT)], + initializers=initializers, + opset=17, + name="metadata_slice_chain_graph", + ) + return model, {"X": x}, pos + + +@pytest.mark.parametrize("permute_shape_chain", [False, True]) +def test_slice_ends_from_device_shape_chain(permute_shape_chain): + model, feeds, pos = _slice_ends_chain_model(permute_shape_chain) + musa_runs = _run_twice(_session(model, musa=True), feeds) + cpu_ends, cpu_y = _run_twice(_session(model, musa=False), feeds)[0] + for run_index, (ends, y) in enumerate(musa_runs): + np.testing.assert_array_equal( + np.asarray(ends), np.asarray(cpu_ends), + err_msg=f"ends mismatch on run {run_index}") + np.testing.assert_array_equal(np.asarray(ends), np.array([1500], dtype=np.int64)) + assert np.asarray(y).shape == (1500, 512), f"run {run_index}: {np.asarray(y).shape}" + np.testing.assert_allclose(np.asarray(y), pos, rtol=0, atol=0) + + +@pytest.mark.parametrize("permute_shape_chain", [False, True]) +def test_slice_ends_metadata_chain_stays_off_the_device_copy_route(permute_shape_chain, monkeypatch): + """The chain must not need a MUSA CPU/device boundary copy. + + `ORT_MUSA_TEST_FAIL_MEMCPY_KERNEL` makes the plugin's MemcpyFromHost / + MemcpyToHost kernels fail. Metadata that is handled on the CPU EP needs no + such boundary copy, so the run must still succeed with the hook armed; if + the chain drifts back onto the device kernels the test fails. + """ + model, feeds, pos = _slice_ends_chain_model(permute_shape_chain) + monkeypatch.setenv("ORT_MUSA_TEST_FAIL_MEMCPY_KERNEL", "1") + for run_index, (ends, y) in enumerate(_run_twice(_session(model, musa=True), feeds)): + np.testing.assert_array_equal(np.asarray(ends), np.array([1500], dtype=np.int64), + err_msg=f"run {run_index}") + np.testing.assert_allclose(np.asarray(y), pos, rtol=0, atol=0) diff --git a/test/ops/test_pageable_h2d_copy_stream_ordering.py b/test/ops/test_pageable_h2d_copy_stream_ordering.py new file mode 100644 index 0000000..121eb6c --- /dev/null +++ b/test/ops/test_pageable_h2d_copy_stream_ordering.py @@ -0,0 +1,211 @@ +# Copyright (c) Moore Threads Technology Co., Ltd. All rights reserved. +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Small pageable host->device copies must be ordered on the caller's stream. + +The plugin stages pageable host->device transfers of at least +`ORT_MUSA_PAGEABLE_H2D_BOUNCE_THRESHOLD_BYTES` (default 1024 B) through the +pinned host pool and issues them with `musaMemcpyAsync()` on the caller's +stream. Transfers below that threshold used to bypass the staging path and call +the blocking `musaMemcpy()` on the MUSA null stream instead, which is ordered +against neither the device work already queued on the EP's (non-blocking) +compute stream nor the consumers enqueued after it. + +That missing order is the writer behind the `sherpa_whisper_base` / +`zipformer_encoder_en` failures: ORT's memory pattern hands the destination of a +small `MemcpyFromHost` (a 4/8/24 byte shape or scale tensor) the block of a +tensor that is already dead on the host side while its producer kernel is still +running on the device, so the copy read back the previous owner's bytes. See +`test_boundary_copy_ordering.py` for that end-to-end shape chain. + +This test does not depend on the memory pattern placing the copy destination in +the path of a pending kernel. It drives the transfer directly - with +`onnxruntime.copy_tensors()`, the public entry point of `OrtApi::CopyTensors` +that the graph executor's MemcpyFromHost path uses as well - and queues the +pending writer on the same stream itself: + + 1. `_PENDING_ROUNDS` device-side writers that keep the stream busy and, as the + last one of every round, fill the destination bytes with a sentinel, + 2. the pageable host->device copy of those bytes, on that same stream, + 3. a stream synchronize and a read back of the destination. + +The copy is the last writer enqueued on the stream, so the destination must hold +the copied bytes. If the copy ignores the stream, it runs as soon as the host +issues it - while the sentinel rounds are still queued - and the sentinel wins. +The sub-threshold sizes must therefore fail on a build that copies on the null +stream and pass on a build that stages the copy on the caller's stream. The +4096 B case is the control: it is above the staging threshold, so it was already +ordered before the fix. +""" + +from __future__ import annotations + +import ctypes +import ctypes.util + +import numpy as np +import onnxruntime as ort +import pytest + +from op_test_utils import musa_devices + +_MUSA_VENDOR_ID = 0x4D54 + +#: Production threshold of the staged (pinned pool + musaMemcpyAsync) H2D path. +_STAGED_THRESHOLD_BYTES = 1024 + +#: Device work queued on the stream ahead of the copy under test. Every round +#: keeps the stream busy; the copy must not be able to slip behind the last +#: round, so the rounds together have to outlast the host-side cost of one copy. +_PENDING_ROUNDS = 8 +_PENDING_BYTES = 256 << 20 +_SENTINEL_BYTE = 0xAB + + +def _load_musa_runtime(): + candidates = [ + "/usr/local/musa/lib/libmusart.so", + "/usr/local/musa/lib/libmusart.so.5", + ctypes.util.find_library("musart"), + ctypes.util.find_library("musa"), + ctypes.util.find_library("musa_runtime"), + "libmusart.so", + ] + for candidate in candidates: + if not candidate: + continue + try: + return ctypes.CDLL(candidate) + except OSError: + continue + pytest.skip("MUSA runtime library is not loadable via ctypes") + + +class _MusaRuntime: + """The MUSA runtime calls this test needs to queue work on an EP stream.""" + + def __init__(self) -> None: + self._lib = _load_musa_runtime() + lib = self._lib + lib.musaSetDevice.argtypes = [ctypes.c_int] + lib.musaSetDevice.restype = ctypes.c_int + lib.musaMalloc.argtypes = [ctypes.POINTER(ctypes.c_void_p), ctypes.c_size_t] + lib.musaMalloc.restype = ctypes.c_int + lib.musaFree.argtypes = [ctypes.c_void_p] + lib.musaFree.restype = ctypes.c_int + lib.musaMemsetAsync.argtypes = [ + ctypes.c_void_p, + ctypes.c_int, + ctypes.c_size_t, + ctypes.c_void_p, + ] + lib.musaMemsetAsync.restype = ctypes.c_int + lib.musaStreamSynchronize.argtypes = [ctypes.c_void_p] + lib.musaStreamSynchronize.restype = ctypes.c_int + + assert lib.musaSetDevice(0) == 0 + + def malloc(self, nbytes: int) -> int: + pointer = ctypes.c_void_p() + assert self._lib.musaMalloc(ctypes.byref(pointer), nbytes) == 0 + assert pointer.value not in (None, 0) + return pointer.value + + def free(self, pointer: int) -> None: + if pointer: + assert self._lib.musaFree(ctypes.c_void_p(pointer)) == 0 + + def memset_async(self, pointer: int, value: int, nbytes: int, stream: int) -> None: + assert ( + self._lib.musaMemsetAsync( + ctypes.c_void_p(pointer), + ctypes.c_int(value), + ctypes.c_size_t(nbytes), + ctypes.c_void_p(stream), + ) + == 0 + ) + + def synchronize(self, stream: int) -> None: + assert self._lib.musaStreamSynchronize(ctypes.c_void_p(stream)) == 0 + + +def _assert_copy_is_ordered(dtype, shape) -> None: + devices = musa_devices() + if not devices: + pytest.skip("No MUSA device available") + + payload = np.arange(int(np.prod(shape)), dtype=dtype).reshape(shape) + nbytes = payload.nbytes + # A plain CPU OrtValue: the source is pageable host memory, which is what + # takes the path under test. + source = ort.OrtValue.ortvalue_from_numpy(payload) + target = ort.OrtValue.ortvalue_from_shape_and_type( + shape, dtype, "gpu", 0, _MUSA_VENDOR_ID + ) + assert target.data_ptr() != 0 + + runtime = _MusaRuntime() + stream = devices[0].create_sync_stream() + handle = stream.get_handle() + assert handle not in (None, 0) + + scratch = runtime.malloc(_PENDING_BYTES) + try: + # Warm up the device, the copy path and the pinned staging pool so that + # the host-side cost of the copy below is the steady-state one. + runtime.memset_async(scratch, 0, 4096, handle) + runtime.synchronize(handle) + ort.copy_tensors([source], [target], stream) + runtime.synchronize(handle) + + for _ in range(_PENDING_ROUNDS): + runtime.memset_async(scratch, 0, _PENDING_BYTES, handle) + # Last writer queued before the copy: a copy that is not ordered on + # this stream is overwritten by one of these fills. + runtime.memset_async(target.data_ptr(), _SENTINEL_BYTE, nbytes, handle) + + ort.copy_tensors([source], [target], stream) + runtime.synchronize(handle) + + got = np.asarray(target.numpy()).reshape(shape) + np.testing.assert_array_equal( + got.view(np.uint8).reshape(-1), + payload.view(np.uint8).reshape(-1), + err_msg=( + f"{nbytes} byte pageable host->device copy to {shape} {np.dtype(dtype)}: " + f"destination bytes changed after the copy; the copy was not ordered on the " + f"caller's stream (a later device-side write clobbered it)" + ), + ) + finally: + runtime.free(scratch) + + +@pytest.mark.parametrize( + "dtype,shape", + [ + (np.int64, (3,)), # 24 B + (np.int32, (24,)), # 96 B + (np.float32, (128,)), # 512 B + ], + ids=["24B", "96B", "512B"], +) +def test_small_pageable_h2d_copy_is_ordered_on_the_ep_stream(dtype, shape): + assert np.prod(shape) * np.dtype(dtype).itemsize < _STAGED_THRESHOLD_BYTES + _assert_copy_is_ordered(dtype, shape) + + +def test_pageable_h2d_copy_above_the_staging_threshold_is_ordered(): + shape = (1024,) # 4096 B > threshold: staged on the caller's stream already + assert np.prod(shape) * np.dtype(np.float32).itemsize >= _STAGED_THRESHOLD_BYTES + _assert_copy_is_ordered(np.float32, shape) From 1895ebe9cfd4f01164582776f439feda86874c03 Mon Sep 17 00:00:00 2001 From: "yichao.hu" Date: Mon, 12 Oct 2026 23:47:31 -0400 Subject: [PATCH 15/52] feat(plugin): extend kernel registrations down to opset 7-12 Most kernels were registered from opset 13, so models exported at opset 10-12 ran almost entirely on the CPU EP. Following the ORT CUDA EP's version splits, add older-opset registrations for 32 operators whose semantics did not change (elementwise math, Transpose, Reshape, Gather, Concat, Expand, Cast, MatMul, Reduce*, Conv, ...). Fixes the Pow 7-11 type constraint and the Shape opset case found by the new per-opset tests. --- musa/docs/supported_ops.md | 66 ++-- musa/ep/src/kernels/activation/abs.cc | 5 + musa/ep/src/kernels/activation/log.cc | 4 + musa/ep/src/kernels/activation/relu.cc | 4 + musa/ep/src/kernels/activation/sigmoid.cc | 4 + musa/ep/src/kernels/activation/sqrt.cc | 4 + musa/ep/src/kernels/logical/equal.cc | 10 + musa/ep/src/kernels/logical/greater.cc | 10 + musa/ep/src/kernels/logical/less.cc | 10 + musa/ep/src/kernels/math/add.cc | 5 + musa/ep/src/kernels/math/div.cc | 5 + musa/ep/src/kernels/math/exp.cc | 4 + musa/ep/src/kernels/math/floor.cc | 4 + musa/ep/src/kernels/math/matmul.cc | 8 + musa/ep/src/kernels/math/mul.cc | 5 + musa/ep/src/kernels/math/neg.cc | 5 + musa/ep/src/kernels/math/pow.cc | 14 + musa/ep/src/kernels/math/sub.cc | 5 + musa/ep/src/kernels/nn/conv.cc | 4 + musa/ep/src/kernels/reduction/reduce_l2.cc | 7 + musa/ep/src/kernels/reduction/reduce_max.cc | 7 + musa/ep/src/kernels/reduction/reduce_mean.cc | 7 + musa/ep/src/kernels/reduction/reduce_prod.cc | 7 + musa/ep/src/kernels/reduction/reduce_sum.cc | 7 + .../kernels/reduction/reduce_sum_square.cc | 6 + .../shared_inc/kernel_type_constraints.h | 128 +++++++ musa/ep/src/kernels/tensor/cast.cc | 14 + musa/ep/src/kernels/tensor/concat.cc | 12 + musa/ep/src/kernels/tensor/expand.cc | 7 + musa/ep/src/kernels/tensor/gather.cc | 14 + musa/ep/src/kernels/tensor/nonzero.cc | 5 + musa/ep/src/kernels/tensor/reshape.cc | 7 + musa/ep/src/kernels/tensor/shape.cc | 9 + musa/ep/src/kernels/tensor/transpose.cc | 6 + test/ops/test_opset_lower_bound.py | 334 ++++++++++++++++++ 35 files changed, 720 insertions(+), 33 deletions(-) create mode 100644 test/ops/test_opset_lower_bound.py diff --git a/musa/docs/supported_ops.md b/musa/docs/supported_ops.md index 2c88035..8104030 100644 --- a/musa/docs/supported_ops.md +++ b/musa/docs/supported_ops.md @@ -2,97 +2,97 @@ # Supported Ops - Total supported ops: **92** -- Total registered kernels: **124** +- Total registered kernels: **164** Rows are grouped by operator/domain. Multiple opset ranges in a row preserve separate ORT kernel registrations. | Op | Domain | Opset ranges and type constraints | Source | | --- | --- | --- | --- | -| `Abs` | `ai.onnx` | 13–19: `T`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bfloat16 | `activation/abs.cc` | -| `Add` | `ai.onnx` | 13–13: `T`: uint32, uint64, int32, int64, float16, float32, float64, bfloat16
14–19: `T`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bfloat16 | `math/add.cc` | +| `Abs` | `ai.onnx` | 6–12: `T`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64
13–19: `T`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bfloat16 | `activation/abs.cc` | +| `Add` | `ai.onnx` | 7–12: `T`: uint32, uint64, int32, int64, float16, float32, float64
13–13: `T`: uint32, uint64, int32, int64, float16, float32, float64, bfloat16
14–19: `T`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bfloat16 | `math/add.cc` | | `And` | `ai.onnx` | 7–19: `T`: bool | `logical/and.cc` | | `BatchNormalization` | `ai.onnx` | 9–13: `T`: float16, float32, float64
14–14: `T`: float16, float32, float64
`U`: float16, float32, float64
15–19: `T`: float16, float32, float64
`T1`: float16, float32, float64
`T2`: float16, float32, float64 | `nn/batch_normalization.cc` | | `BitwiseAnd` | `ai.onnx` | 18–19: `T`: uint8, uint16, uint32, uint64, int8, int16, int32, int64 | `logical/bitwise_and.cc` | -| `Cast` | `ai.onnx` | 13–19: `T1`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bfloat16, bool
`T2`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bfloat16, bool | `tensor/cast.cc` | +| `Cast` | `ai.onnx` | 6–8: `T1`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bool
`T2`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bool
9–12: `T1`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bool
`T2`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bool
13–19: `T1`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bfloat16, bool
`T2`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bfloat16, bool | `tensor/cast.cc` | | `Ceil` | `ai.onnx` | 13–19: `T`: float16, float32, float64, bfloat16 | `math/ceil.cc` | | `Clip` | `ai.onnx` | 13–19: `T`: float16, float32, float64, int8, uint8, int16, uint16, int32, uint32, int64, uint64, bfloat16 | `math/clip.cc` | -| `Concat` | `ai.onnx` | 13–19: `T`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bfloat16, bool | `tensor/concat.cc` | +| `Concat` | `ai.onnx` | 4–10: `T`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bool
11–12: `T`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bool
13–19: `T`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bfloat16, bool | `tensor/concat.cc` | | `ConstantOfShape` | `ai.onnx` | 9–19: `T1`: int64
`T2`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bool | `generator/constant_of_shape.cc` | -| `Conv` | `ai.onnx` | 11–19: `T`: float32 | `nn/conv.cc` | +| `Conv` | `ai.onnx` | 1–10: `T`: float32
11–19: `T`: float32 | `nn/conv.cc` | | `CumSum` | `ai.onnx` | 11–13: `T`: int32, int64, uint32, uint64, float32, float64
`T2`: int32, int64
14–19: `T`: ?
`T2`: int32, int64 | `math/cumsum.cc` | -| `Div` | `ai.onnx` | 13–13: `T`: uint32, uint64, int32, int64, float16, float32, float64, bfloat16
14–19: `T`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bfloat16 | `math/div.cc` | +| `Div` | `ai.onnx` | 7–12: `T`: uint32, uint64, int32, int64, float16, float32, float64
13–13: `T`: uint32, uint64, int32, int64, float16, float32, float64, bfloat16
14–19: `T`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bfloat16 | `math/div.cc` | | `Einsum` | `ai.onnx` | 12–19: `T`: float32 | `math/einsum.cc` | -| `Equal` | `ai.onnx` | 13–19: `T`: uint32, uint64, int32, int64, float16, float32, float64, bfloat16, bool
19–19: `T`: string | `logical/equal.cc` | +| `Equal` | `ai.onnx` | 7–10: `T`: int32, int64, bool
11–12: `T`: uint32, uint64, int32, int64, float16, float32, float64, bool
13–19: `T`: uint32, uint64, int32, int64, float16, float32, float64, bfloat16, bool
19–19: `T`: string | `logical/equal.cc` | | `Erf` | `ai.onnx` | 13–19: `T`: float16, float32, float64, bfloat16 | `activation/erf.cc` | -| `Exp` | `ai.onnx` | 13–19: `T`: float16, float32, float64, bfloat16 | `math/exp.cc` | -| `Expand` | `ai.onnx` | 13–19: `T`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bfloat16, bool | `tensor/expand.cc` | +| `Exp` | `ai.onnx` | 6–12: `T`: float16, float32, float64
13–19: `T`: float16, float32, float64, bfloat16 | `math/exp.cc` | +| `Expand` | `ai.onnx` | 8–12: `T`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bool
13–19: `T`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bfloat16, bool | `tensor/expand.cc` | | `Flatten` | `ai.onnx` | 1–8: `T`: float16, float32, float64
9–10: `T`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bool
11–12: `T`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bool
13–19: `T`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bfloat16, bool | `tensor/flatten.cc` | -| `Floor` | `ai.onnx` | 13–19: `T`: float16, float32, float64, bfloat16 | `math/floor.cc` | -| `Gather` | `ai.onnx` | 13–19: `T`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bfloat16, bool
`Tind`: int32, int64 | `tensor/gather.cc` | +| `Floor` | `ai.onnx` | 6–12: `T`: float16, float32, float64
13–19: `T`: float16, float32, float64, bfloat16 | `math/floor.cc` | +| `Gather` | `ai.onnx` | 1–10: `T`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bool
`Tind`: int32, int64
11–12: `T`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bool
`Tind`: int32, int64
13–19: `T`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bfloat16, bool
`Tind`: int32, int64 | `tensor/gather.cc` | | `GatherElements` | `ai.onnx` | 11–12: `T`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bfloat16, bool
`Tind`: int32, int64
13–19: `T`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bfloat16, bool
`Tind`: int32, int64 | `tensor/gather_elements.cc` | | `GatherND` | `ai.onnx` | 11–12: `T`: float32, float64, float16, int64, bfloat16, bool
`indices`: int64
13–19: `T`: float32, float64, float16, int64, bfloat16, bool
`indices`: int64 | `tensor/gather_nd.cc` | | `Gemm` | `ai.onnx` | 13–19: `T`: float16, float32, float64, bfloat16 | `math/gemm.cc` | | `GlobalAveragePool` | `ai.onnx` | 1–19: `T`: float16, float32, float64 | `nn/pool.cc` | | `GlobalMaxPool` | `ai.onnx` | 1–19: `T`: float16, float32, float64 | `nn/pool.cc` | -| `Greater` | `ai.onnx` | 13–19: `T`: uint32, uint64, int32, int64, float16, float32, float64, bfloat16 | `logical/greater.cc` | +| `Greater` | `ai.onnx` | 7–8: `T`: float16, float32, float64
9–12: `T`: uint32, uint64, int32, int64, float16, float32, float64
13–19: `T`: uint32, uint64, int32, int64, float16, float32, float64, bfloat16 | `logical/greater.cc` | | `GreaterOrEqual` | `ai.onnx` | 13–15: `T`: uint32, uint64, int32, int64, float16, float32, float64
16–19: `T`: uint32, uint64, int32, int64, float16, float32, float64, bfloat16 | `logical/greater_or_equal.cc` | | `Identity` | `ai.onnx` | 1–12: `T`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bool
13–13: `T`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bfloat16, bool
14–18: `V`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bfloat16, bool
19–19: `V`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bfloat16, bool | `tensor/identity_op.cc` | | `If` | `ai.onnx` | 13–15: `B`: bool
`V`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bool
16–19: `B`: bool
`V`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bfloat16, bool | `controlflow/if.cc` | | `IsNaN` | `ai.onnx` | 13–19: `T1`: float16, float32, float64, bfloat16
`T2`: bool | `logical/is_nan.cc` | | `LayerNormalization` | `ai.onnx` | 17–19: `T`: float32 | `nn/layer_norm.cc` | | `LeakyRelu` | `ai.onnx` | 6–15: `T`: float16, float32, float64
16–19: `T`: float16, float32, float64, bfloat16 | `activation/leaky_relu.cc` | -| `Less` | `ai.onnx` | 13–19: `T`: uint32, uint64, int32, int64, float16, float32, float64, bfloat16 | `logical/less.cc` | +| `Less` | `ai.onnx` | 7–8: `T`: float16, float32, float64
9–12: `T`: uint32, uint64, int32, int64, float16, float32, float64
13–19: `T`: uint32, uint64, int32, int64, float16, float32, float64, bfloat16 | `logical/less.cc` | | `LessOrEqual` | `ai.onnx` | 12–15: `T`: uint32, uint64, int32, int64, float16, float32, float64
16–19: `T`: uint32, uint64, int32, int64, float16, float32, float64, bfloat16 | `logical/less_or_equal.cc` | -| `Log` | `ai.onnx` | 13–19: `T`: float16, float32, float64, bfloat16 | `activation/log.cc` | +| `Log` | `ai.onnx` | 6–12: `T`: float16, float32, float64
13–19: `T`: float16, float32, float64, bfloat16 | `activation/log.cc` | | `Loop` | `ai.onnx` | 1–12: `I`: int64
`B`: bool
`V`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bool
13–19: `I`: int64
`B`: bool
`V`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bool | `controlflow/loop.cc` | -| `MatMul` | `ai.onnx` | 13–19: `T`: float16, float32, float64, bfloat16 | `math/matmul.cc` | +| `MatMul` | `ai.onnx` | 1–8: `T`: float16, float32, float64
9–12: `T`: float16, float32, float64
13–19: `T`: float16, float32, float64, bfloat16 | `math/matmul.cc` | | `Max` | `ai.onnx` | 13–19: `T`: uint32, uint64, int32, int64, float16, float32, float64, bfloat16 | `math/max.cc` | | `MaxPool` | `ai.onnx` | 1–7: `T`: float16, float32, float64
8–11: `T`: float16, float32, float64
`I`: int64
12–19: `T`: float16, float32, float64, int8, uint8
`I`: int64 | `nn/pool.cc` | | `MemcpyFromHost` | `ai.onnx` | 1–1: `T`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bfloat16, bool | `tensor/memcpy.cc` | | `MemcpyToHost` | `ai.onnx` | 1–1: `T`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bfloat16, bool | `tensor/memcpy.cc` | | `Min` | `ai.onnx` | 13–19: `T`: uint32, uint64, int32, int64, float16, float32, float64, bfloat16 | `math/min.cc` | | `Mod` | `ai.onnx` | 10–12: `T`: int32, int64
13–19: `T`: int32, int64 | `math/mod.cc` | -| `Mul` | `ai.onnx` | 13–13: `T`: uint32, uint64, int32, int64, float16, float32, float64, bfloat16
14–19: `T`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bfloat16 | `math/mul.cc` | -| `Neg` | `ai.onnx` | 13–19: `T`: int8, int16, int32, int64, float16, float32, float64, bfloat16 | `math/neg.cc` | -| `NonZero` | `ai.onnx` | 13–19: `T`: bool, uint8, int32, int64, float16, float32 | `tensor/nonzero.cc` | +| `Mul` | `ai.onnx` | 7–12: `T`: uint32, uint64, int32, int64, float16, float32, float64
13–13: `T`: uint32, uint64, int32, int64, float16, float32, float64, bfloat16
14–19: `T`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bfloat16 | `math/mul.cc` | +| `Neg` | `ai.onnx` | 6–12: `T`: int8, int16, int32, int64, float16, float32, float64
13–19: `T`: int8, int16, int32, int64, float16, float32, float64, bfloat16 | `math/neg.cc` | +| `NonZero` | `ai.onnx` | 9–12: `T`: bool, uint8, int32, int64, float16, float32
13–19: `T`: bool, uint8, int32, int64, float16, float32 | `tensor/nonzero.cc` | | `Not` | `ai.onnx` | 1–19: `T`: bool | `logical/not.cc` | | `OneHot` | `ai.onnx` | 11–19: `T1`: int32, int64
`T2`: int32, int64
`T3`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bool | `tensor/onehot.cc` | | `Or` | `ai.onnx` | 7–19: `T`: bool | `logical/or.cc` | | `PRelu` | `ai.onnx` | 7–15: `T`: float16, float32, float64
16–19: `T`: float16, float32, float64, bfloat16, uint32, uint64, int32, int64 | `activation/prelu.cc` | | `Pad` | `ai.onnx` | 13–19: `T`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bfloat16, bool | `tensor/pad.cc` | -| `Pow` | `ai.onnx` | 13–14: `T`: int32, int64, float16, float32, float64, bfloat16
`T1`: int32, int64, float16, float32, float64
15–19: `T`: int32, int64, float16, float32, float64, bfloat16
`T1`: int32, int64, float16, float32, float64, bfloat16 | `math/pow.cc` | +| `Pow` | `ai.onnx` | 7–11: `T`: float16, float32, float64
12–12: `T`: int32, int64, float16, float32, float64
`T1`: int32, int64, float16, float32, float64
13–14: `T`: int32, int64, float16, float32, float64, bfloat16
`T1`: int32, int64, float16, float32, float64
15–19: `T`: int32, int64, float16, float32, float64, bfloat16
`T1`: int32, int64, float16, float32, float64, bfloat16 | `math/pow.cc` | | `RandomUniform` | `ai.onnx` | 1–19: `T`: float16, float32, float64 | `generator/random_uniform.cc` | | `RandomUniformLike` | `ai.onnx` | 1–19: `T1`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bfloat16, bool
`T2`: float16, float32, float64 | `generator/random_uniform.cc` | | `Range` | `ai.onnx` | 11–19: `T`: int16, int32, int64, float32, float64 | `generator/range.cc` | | `Reciprocal` | `ai.onnx` | 13–19: `T`: float16, float32, float64, bfloat16 | `activation/reciprocal.cc` | -| `ReduceL2` | `ai.onnx` | 13–17: `T`: float16, float32, float64, int32
18–19: `T`: float16, float32, float64, bfloat16, int32 | `reduction/reduce_l2.cc` | -| `ReduceMax` | `ai.onnx` | 13–19: `T`: uint8, int8, float16, float32, float64, bfloat16, int32, int64 | `reduction/reduce_max.cc` | -| `ReduceMean` | `ai.onnx` | 13–19: `T`: float16, float32, float64, bfloat16, int32 | `reduction/reduce_mean.cc` | -| `ReduceProd` | `ai.onnx` | 13–19: `T`: float16, float32, float64, bfloat16, int32, int64 | `reduction/reduce_prod.cc` | -| `ReduceSum` | `ai.onnx` | 13–19: `T`: float16, float32, float64, bfloat16, int32, int64 | `reduction/reduce_sum.cc` | -| `ReduceSumSquare` | `ai.onnx` | 13–19: `T`: float16, float32, float64, bfloat16 | `reduction/reduce_sum_square.cc` | -| `Relu` | `ai.onnx` | 13–19: `T`: float16, float32, float64, bfloat16 | `activation/relu.cc` | -| `Reshape` | `ai.onnx` | 13–19: `T`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bfloat16, bool
19–19: `T`: string | `tensor/reshape.cc` | +| `ReduceL2` | `ai.onnx` | 1–12: `T`: float16, float32, float64, int32
13–17: `T`: float16, float32, float64, int32
18–19: `T`: float16, float32, float64, bfloat16, int32 | `reduction/reduce_l2.cc` | +| `ReduceMax` | `ai.onnx` | 1–12: `T`: uint8, int8, float16, float32, float64, int32, int64
13–19: `T`: uint8, int8, float16, float32, float64, bfloat16, int32, int64 | `reduction/reduce_max.cc` | +| `ReduceMean` | `ai.onnx` | 1–12: `T`: float16, float32, float64, int32
13–19: `T`: float16, float32, float64, bfloat16, int32 | `reduction/reduce_mean.cc` | +| `ReduceProd` | `ai.onnx` | 1–12: `T`: float16, float32, float64, int32, int64
13–19: `T`: float16, float32, float64, bfloat16, int32, int64 | `reduction/reduce_prod.cc` | +| `ReduceSum` | `ai.onnx` | 1–12: `T`: float16, float32, float64, int32, int64
13–19: `T`: float16, float32, float64, bfloat16, int32, int64 | `reduction/reduce_sum.cc` | +| `ReduceSumSquare` | `ai.onnx` | 1–12: `T`: float16, float32, float64
13–19: `T`: float16, float32, float64, bfloat16 | `reduction/reduce_sum_square.cc` | +| `Relu` | `ai.onnx` | 6–12: `T`: float16, float32, float64
13–19: `T`: float16, float32, float64, bfloat16 | `activation/relu.cc` | +| `Reshape` | `ai.onnx` | 5–12: `T`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bool
13–19: `T`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bfloat16, bool
19–19: `T`: string | `tensor/reshape.cc` | | `Resize` | `ai.onnx` | 13–17: `T1`: float32 | `tensor/resize.cc` | | `ReverseSequence` | `ai.onnx` | 10–19: `T`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bool | `tensor/reverse_sequence.cc` | | `Round` | `ai.onnx` | 11–19: `T`: float16, float32, float64 | `math/round.cc` | | `ScatterElements` | `ai.onnx` | 11–12: `T`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bfloat16, bool
`Tind`: int32, int64
13–19: `T`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bfloat16, bool
`Tind`: int32, int64 | `tensor/scatter_elements.cc` | | `ScatterND` | `ai.onnx` | 11–12: `T`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bfloat16, bool
`indices`: int64
13–19: `T`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bfloat16, bool
`indices`: int64 | `tensor/scatter_nd.cc` | -| `Shape` | `ai.onnx` | 13–19: `T`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bfloat16, bool
`T1`: int64 | `tensor/shape.cc` | -| `Sigmoid` | `ai.onnx` | 13–19: `T`: float16, float32, float64, bfloat16 | `activation/sigmoid.cc` | +| `Shape` | `ai.onnx` | 1–12: `T`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bool
`T1`: int64
13–19: `T`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bfloat16, bool
`T1`: int64 | `tensor/shape.cc` | +| `Sigmoid` | `ai.onnx` | 6–12: `T`: float16, float32, float64
13–19: `T`: float16, float32, float64, bfloat16 | `activation/sigmoid.cc` | | `Sign` | `ai.onnx` | 13–19: `T`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bfloat16 | `math/sign.cc` | | `Slice` | `ai.onnx` | 13–19: `T`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bfloat16, bool
`Tind`: int32, int64 | `tensor/slice.cc` | | `Softmax` | `ai.onnx` | 13–17: `T`: float16, float32, float64, bfloat16 | `math/softmax.cc` | | `Softplus` | `ai.onnx` | 1–19: `T`: float16, float32, float64 | `activation/softplus.cc` | | `Split` | `ai.onnx` | 13–19: `T`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bfloat16, bool | `tensor/split.cc` | -| `Sqrt` | `ai.onnx` | 13–19: `T`: float16, float32, float64, bfloat16 | `activation/sqrt.cc` | +| `Sqrt` | `ai.onnx` | 6–12: `T`: float16, float32, float64
13–19: `T`: float16, float32, float64, bfloat16 | `activation/sqrt.cc` | | `Squeeze` | `ai.onnx` | 13–19: `T`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bfloat16, bool | `tensor/squeeze.cc` | -| `Sub` | `ai.onnx` | 13–13: `T`: uint32, uint64, int32, int64, float16, float32, float64, bfloat16
14–19: `T`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bfloat16 | `math/sub.cc` | +| `Sub` | `ai.onnx` | 7–12: `T`: uint32, uint64, int32, int64, float16, float32, float64
13–13: `T`: uint32, uint64, int32, int64, float16, float32, float64, bfloat16
14–19: `T`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bfloat16 | `math/sub.cc` | | `Sum` | `ai.onnx` | 13–19: `T`: float16, float32, float64, bfloat16 | `math/sum.cc` | | `Tanh` | `ai.onnx` | 13–19: `T`: float16, float32, float64, bfloat16 | `activation/tanh.cc` | | `Tile` | `ai.onnx` | 13–19: `T`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bfloat16, bool | `tensor/tile.cc` | | `TopK` | `ai.onnx` | 10–10: `T`: uint8, uint16, uint32, uint64, int8, int16, float16, float32, float64, int32, int64
`I`: int64
11–19: `T`: uint8, uint16, uint32, uint64, int8, int16, float16, float32, float64, int32, int64
`I`: int64 | `math/topk.cc` | -| `Transpose` | `ai.onnx` | 13–19: `T`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bfloat16, bool | `tensor/transpose.cc` | +| `Transpose` | `ai.onnx` | 1–12: `T`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bool
13–19: `T`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bfloat16, bool | `tensor/transpose.cc` | | `Unique` | `ai.onnx` | 11–19: `T`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bool | `tensor/unique.cc` | | `Unsqueeze` | `ai.onnx` | 13–19: `T`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bfloat16, bool | `tensor/unsqueeze.cc` | | `Where` | `ai.onnx` | 9–15: `B`: bool
`T`: float16, float32, float64, int32, int64, uint8
16–19: `B`: bool
`T`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bfloat16, bool | `tensor/where.cc` | diff --git a/musa/ep/src/kernels/activation/abs.cc b/musa/ep/src/kernels/activation/abs.cc index 4dc40c7..a676b6d 100644 --- a/musa/ep/src/kernels/activation/abs.cc +++ b/musa/ep/src/kernels/activation/abs.cc @@ -75,6 +75,11 @@ OrtStatus* Abs::Compute(Ort::KernelContext& ctx) const { } } // namespace +ONNX_OPERATOR_VERSIONED_KERNEL_EX( + Abs, kOnnxDomain, 6, 12, + (Ort::KernelDefBuilder().AddTypeConstraint("T", AbsOpset6TensorTypes())), + Abs) + ONNX_OPERATOR_VERSIONED_KERNEL_EX( Abs, kOnnxDomain, 13, 19, (Ort::KernelDefBuilder().AddTypeConstraint("T", AbsTensorTypes())), Abs) diff --git a/musa/ep/src/kernels/activation/log.cc b/musa/ep/src/kernels/activation/log.cc index 3f63f6a..77468b6 100644 --- a/musa/ep/src/kernels/activation/log.cc +++ b/musa/ep/src/kernels/activation/log.cc @@ -75,6 +75,10 @@ OrtStatus* Log::Compute(Ort::KernelContext& ctx) const { } } // namespace +ONNX_OPERATOR_VERSIONED_KERNEL_EX( + Log, kOnnxDomain, 6, 12, + (Ort::KernelDefBuilder().AddTypeConstraint("T", HfdTensorTypes())), Log) + ONNX_OPERATOR_VERSIONED_KERNEL_EX( Log, kOnnxDomain, 13, 19, (Ort::KernelDefBuilder().AddTypeConstraint("T", FloatLikeTensorTypes())), diff --git a/musa/ep/src/kernels/activation/relu.cc b/musa/ep/src/kernels/activation/relu.cc index 84d4796..d070c07 100644 --- a/musa/ep/src/kernels/activation/relu.cc +++ b/musa/ep/src/kernels/activation/relu.cc @@ -75,6 +75,10 @@ OrtStatus* Relu::Compute(Ort::KernelContext& ctx) const { } } // namespace +ONNX_OPERATOR_VERSIONED_KERNEL_EX( + Relu, kOnnxDomain, 6, 12, + (Ort::KernelDefBuilder().AddTypeConstraint("T", HfdTensorTypes())), Relu) + ONNX_OPERATOR_VERSIONED_KERNEL_EX( Relu, kOnnxDomain, 13, 19, (Ort::KernelDefBuilder().AddTypeConstraint("T", FloatLikeTensorTypes())), diff --git a/musa/ep/src/kernels/activation/sigmoid.cc b/musa/ep/src/kernels/activation/sigmoid.cc index 9148dde..857c745 100644 --- a/musa/ep/src/kernels/activation/sigmoid.cc +++ b/musa/ep/src/kernels/activation/sigmoid.cc @@ -76,6 +76,10 @@ OrtStatus* Sigmoid::Compute(Ort::KernelContext& ctx) const { } } // namespace +ONNX_OPERATOR_VERSIONED_KERNEL_EX( + Sigmoid, kOnnxDomain, 6, 12, + (Ort::KernelDefBuilder().AddTypeConstraint("T", HfdTensorTypes())), Sigmoid) + ONNX_OPERATOR_VERSIONED_KERNEL_EX( Sigmoid, kOnnxDomain, 13, 19, (Ort::KernelDefBuilder().AddTypeConstraint("T", FloatLikeTensorTypes())), diff --git a/musa/ep/src/kernels/activation/sqrt.cc b/musa/ep/src/kernels/activation/sqrt.cc index 5fffe78..d9cf64f 100644 --- a/musa/ep/src/kernels/activation/sqrt.cc +++ b/musa/ep/src/kernels/activation/sqrt.cc @@ -75,6 +75,10 @@ OrtStatus* Sqrt::Compute(Ort::KernelContext& ctx) const { } } // namespace +ONNX_OPERATOR_VERSIONED_KERNEL_EX( + Sqrt, kOnnxDomain, 6, 12, + (Ort::KernelDefBuilder().AddTypeConstraint("T", HfdTensorTypes())), Sqrt) + ONNX_OPERATOR_VERSIONED_KERNEL_EX( Sqrt, kOnnxDomain, 13, 19, (Ort::KernelDefBuilder().AddTypeConstraint("T", FloatLikeTensorTypes())), diff --git a/musa/ep/src/kernels/logical/equal.cc b/musa/ep/src/kernels/logical/equal.cc index b496625..69abf09 100644 --- a/musa/ep/src/kernels/logical/equal.cc +++ b/musa/ep/src/kernels/logical/equal.cc @@ -175,6 +175,16 @@ OrtStatus* EqualString::Compute(Ort::KernelContext& ctx) const { } } // namespace +ONNX_OPERATOR_VERSIONED_KERNEL_EX( + Equal, kOnnxDomain, 7, 10, + (Ort::KernelDefBuilder().AddTypeConstraint("T", EqualOpset7TensorTypes())), + Equal) + +ONNX_OPERATOR_VERSIONED_KERNEL_EX( + Equal, kOnnxDomain, 11, 12, + (Ort::KernelDefBuilder().AddTypeConstraint("T", EqualOpset11TensorTypes())), + Equal) + ONNX_OPERATOR_VERSIONED_KERNEL_EX( Equal, kOnnxDomain, 13, 19, (Ort::KernelDefBuilder().AddTypeConstraint("T", EqualTensorTypes())), Equal) diff --git a/musa/ep/src/kernels/logical/greater.cc b/musa/ep/src/kernels/logical/greater.cc index ae5637e..1786269 100644 --- a/musa/ep/src/kernels/logical/greater.cc +++ b/musa/ep/src/kernels/logical/greater.cc @@ -82,6 +82,16 @@ OrtStatus* Greater::Compute(Ort::KernelContext& ctx) const { } } // namespace +ONNX_OPERATOR_VERSIONED_KERNEL_EX(Greater, kOnnxDomain, 7, 8, + (Ort::KernelDefBuilder().AddTypeConstraint( + "T", CompareOpset7TensorTypes())), + Greater) + +ONNX_OPERATOR_VERSIONED_KERNEL_EX(Greater, kOnnxDomain, 9, 12, + (Ort::KernelDefBuilder().AddTypeConstraint( + "T", CompareTensorTypesNoBFloat16())), + Greater) + ONNX_OPERATOR_VERSIONED_KERNEL_EX( Greater, kOnnxDomain, 13, 19, (Ort::KernelDefBuilder().AddTypeConstraint("T", CompareTensorTypes())), diff --git a/musa/ep/src/kernels/logical/less.cc b/musa/ep/src/kernels/logical/less.cc index c731d17..e2a832a 100644 --- a/musa/ep/src/kernels/logical/less.cc +++ b/musa/ep/src/kernels/logical/less.cc @@ -30,6 +30,16 @@ OrtStatus* Less::Compute(Ort::KernelContext& ctx) const { } } // namespace +ONNX_OPERATOR_VERSIONED_KERNEL_EX(Less, kOnnxDomain, 7, 8, + (Ort::KernelDefBuilder().AddTypeConstraint( + "T", CompareOpset7TensorTypes())), + Less) + +ONNX_OPERATOR_VERSIONED_KERNEL_EX(Less, kOnnxDomain, 9, 12, + (Ort::KernelDefBuilder().AddTypeConstraint( + "T", CompareTensorTypesNoBFloat16())), + Less) + ONNX_OPERATOR_VERSIONED_KERNEL_EX( Less, kOnnxDomain, 13, 19, (Ort::KernelDefBuilder().AddTypeConstraint("T", CompareTensorTypes())), diff --git a/musa/ep/src/kernels/math/add.cc b/musa/ep/src/kernels/math/add.cc index a900540..cbc9a29 100644 --- a/musa/ep/src/kernels/math/add.cc +++ b/musa/ep/src/kernels/math/add.cc @@ -46,6 +46,11 @@ OrtStatus* Add::Compute(Ort::KernelContext& ctx) const { } } // namespace +ONNX_OPERATOR_VERSIONED_KERNEL_EX(Add, kOnnxDomain, 7, 12, + (Ort::KernelDefBuilder().AddTypeConstraint( + "T", BinaryNumericOpset7TensorTypes())), + Add) + ONNX_OPERATOR_VERSIONED_KERNEL_EX(Add, kOnnxDomain, 13, 13, (Ort::KernelDefBuilder().AddTypeConstraint( "T", BinaryNumericOpset13TensorTypes())), diff --git a/musa/ep/src/kernels/math/div.cc b/musa/ep/src/kernels/math/div.cc index be8bf40..bc9068a 100644 --- a/musa/ep/src/kernels/math/div.cc +++ b/musa/ep/src/kernels/math/div.cc @@ -47,6 +47,11 @@ OrtStatus* Div::Compute(Ort::KernelContext& ctx) const { } } // namespace +ONNX_OPERATOR_VERSIONED_KERNEL_EX(Div, kOnnxDomain, 7, 12, + (Ort::KernelDefBuilder().AddTypeConstraint( + "T", BinaryNumericOpset7TensorTypes())), + Div) + ONNX_OPERATOR_VERSIONED_KERNEL_EX(Div, kOnnxDomain, 13, 13, (Ort::KernelDefBuilder().AddTypeConstraint( "T", BinaryNumericOpset13TensorTypes())), diff --git a/musa/ep/src/kernels/math/exp.cc b/musa/ep/src/kernels/math/exp.cc index 6b5f682..fda07a7 100644 --- a/musa/ep/src/kernels/math/exp.cc +++ b/musa/ep/src/kernels/math/exp.cc @@ -75,6 +75,10 @@ OrtStatus* Exp::Compute(Ort::KernelContext& ctx) const { } } // namespace +ONNX_OPERATOR_VERSIONED_KERNEL_EX( + Exp, kOnnxDomain, 6, 12, + (Ort::KernelDefBuilder().AddTypeConstraint("T", HfdTensorTypes())), Exp) + ONNX_OPERATOR_VERSIONED_KERNEL_EX( Exp, kOnnxDomain, 13, 19, (Ort::KernelDefBuilder().AddTypeConstraint("T", FloatLikeTensorTypes())), diff --git a/musa/ep/src/kernels/math/floor.cc b/musa/ep/src/kernels/math/floor.cc index 64ac30d..f498656 100644 --- a/musa/ep/src/kernels/math/floor.cc +++ b/musa/ep/src/kernels/math/floor.cc @@ -28,6 +28,10 @@ OrtStatus* Floor::Compute(Ort::KernelContext& ctx) const { } } // namespace +ONNX_OPERATOR_VERSIONED_KERNEL_EX( + Floor, kOnnxDomain, 6, 12, + (Ort::KernelDefBuilder().AddTypeConstraint("T", HfdTensorTypes())), Floor) + ONNX_OPERATOR_VERSIONED_KERNEL_EX( Floor, kOnnxDomain, 13, 19, (Ort::KernelDefBuilder().AddTypeConstraint("T", FloatLikeTensorTypes())), diff --git a/musa/ep/src/kernels/math/matmul.cc b/musa/ep/src/kernels/math/matmul.cc index cebc0be..12aefa3 100644 --- a/musa/ep/src/kernels/math/matmul.cc +++ b/musa/ep/src/kernels/math/matmul.cc @@ -13,6 +13,14 @@ #include "mudnn_compat.h" #include "shared_inc/blas_utils.h" +ONNX_OPERATOR_VERSIONED_KERNEL_EX( + MatMul, kOnnxDomain, 1, 8, + (Ort::KernelDefBuilder().AddTypeConstraint("T", HfdTensorTypes())), MatMul) + +ONNX_OPERATOR_VERSIONED_KERNEL_EX( + MatMul, kOnnxDomain, 9, 12, + (Ort::KernelDefBuilder().AddTypeConstraint("T", HfdTensorTypes())), MatMul) + ONNX_OPERATOR_VERSIONED_KERNEL_EX( MatMul, kOnnxDomain, 13, 19, (Ort::KernelDefBuilder().AddTypeConstraint("T", FloatLikeTensorTypes())), diff --git a/musa/ep/src/kernels/math/mul.cc b/musa/ep/src/kernels/math/mul.cc index a1811a7..cd8f1d3 100644 --- a/musa/ep/src/kernels/math/mul.cc +++ b/musa/ep/src/kernels/math/mul.cc @@ -47,6 +47,11 @@ OrtStatus* Mul::Compute(Ort::KernelContext& ctx) const { } } // namespace +ONNX_OPERATOR_VERSIONED_KERNEL_EX(Mul, kOnnxDomain, 7, 12, + (Ort::KernelDefBuilder().AddTypeConstraint( + "T", BinaryNumericOpset7TensorTypes())), + Mul) + ONNX_OPERATOR_VERSIONED_KERNEL_EX(Mul, kOnnxDomain, 13, 13, (Ort::KernelDefBuilder().AddTypeConstraint( "T", BinaryNumericOpset13TensorTypes())), diff --git a/musa/ep/src/kernels/math/neg.cc b/musa/ep/src/kernels/math/neg.cc index df334c1..483eaa6 100644 --- a/musa/ep/src/kernels/math/neg.cc +++ b/musa/ep/src/kernels/math/neg.cc @@ -31,6 +31,11 @@ OrtStatus* Neg::Compute(Ort::KernelContext& ctx) const { } } // namespace +ONNX_OPERATOR_VERSIONED_KERNEL_EX( + Neg, kOnnxDomain, 6, 12, + (Ort::KernelDefBuilder().AddTypeConstraint("T", NegOpset6TensorTypes())), + Neg) + ONNX_OPERATOR_VERSIONED_KERNEL_EX( Neg, kOnnxDomain, 13, 19, (Ort::KernelDefBuilder().AddTypeConstraint("T", NegTensorTypes())), Neg) diff --git a/musa/ep/src/kernels/math/pow.cc b/musa/ep/src/kernels/math/pow.cc index ab67d07..98a4e78 100644 --- a/musa/ep/src/kernels/math/pow.cc +++ b/musa/ep/src/kernels/math/pow.cc @@ -130,6 +130,20 @@ OrtStatus* Pow::Compute(Ort::KernelContext& ctx) const { } } // namespace +// Pow only grows a separate exponent type (T1) in opset 12; before that both +// inputs share T, so the kernel def must not declare T1 for 7-11 - ORT resolves +// the declared type strings against the operator schema of that opset. +ONNX_OPERATOR_VERSIONED_KERNEL_EX( + Pow, kOnnxDomain, 7, 11, + (Ort::KernelDefBuilder().AddTypeConstraint("T", HfdTensorTypes())), Pow) + +ONNX_OPERATOR_VERSIONED_KERNEL_EX( + Pow, kOnnxDomain, 12, 12, + (Ort::KernelDefBuilder() + .AddTypeConstraint("T", PowOpset12TensorTypes()) + .AddTypeConstraint("T1", PowOpset12TensorTypes())), + Pow) + ONNX_OPERATOR_VERSIONED_KERNEL_EX( Pow, kOnnxDomain, 13, 14, (Ort::KernelDefBuilder() diff --git a/musa/ep/src/kernels/math/sub.cc b/musa/ep/src/kernels/math/sub.cc index b0f578e..b78f1cc 100644 --- a/musa/ep/src/kernels/math/sub.cc +++ b/musa/ep/src/kernels/math/sub.cc @@ -46,6 +46,11 @@ OrtStatus* Sub::Compute(Ort::KernelContext& ctx) const { } } // namespace +ONNX_OPERATOR_VERSIONED_KERNEL_EX(Sub, kOnnxDomain, 7, 12, + (Ort::KernelDefBuilder().AddTypeConstraint( + "T", BinaryNumericOpset7TensorTypes())), + Sub) + ONNX_OPERATOR_VERSIONED_KERNEL_EX(Sub, kOnnxDomain, 13, 13, (Ort::KernelDefBuilder().AddTypeConstraint( "T", BinaryNumericOpset13TensorTypes())), diff --git a/musa/ep/src/kernels/nn/conv.cc b/musa/ep/src/kernels/nn/conv.cc index 507be31..66b338d 100644 --- a/musa/ep/src/kernels/nn/conv.cc +++ b/musa/ep/src/kernels/nn/conv.cc @@ -304,6 +304,10 @@ OrtStatus* Conv::Compute(Ort::KernelContext& ctx) const { } // namespace +ONNX_OPERATOR_VERSIONED_KERNEL_EX( + Conv, kOnnxDomain, 1, 10, + (Ort::KernelDefBuilder().AddTypeConstraint("T", FloatTensorTypes())), Conv) + ONNX_OPERATOR_VERSIONED_KERNEL_EX( Conv, kOnnxDomain, 11, 19, (Ort::KernelDefBuilder().AddTypeConstraint("T", FloatTensorTypes())), Conv) diff --git a/musa/ep/src/kernels/reduction/reduce_l2.cc b/musa/ep/src/kernels/reduction/reduce_l2.cc index 22dfe92..2310baa 100644 --- a/musa/ep/src/kernels/reduction/reduce_l2.cc +++ b/musa/ep/src/kernels/reduction/reduce_l2.cc @@ -31,6 +31,13 @@ class ReduceL2 : public OpKernelBase { }; } // namespace +ONNX_OPERATOR_VERSIONED_KERNEL_EX( + ReduceL2, kOnnxDomain, 1, 12, + (Ort::KernelDefBuilder() + .AddTypeConstraint("T", ReduceFloatOpset1TensorTypes()) + .SetInputMemType(1, OrtMemTypeCPUInput)), + ReduceL2) + ONNX_OPERATOR_VERSIONED_KERNEL_EX( ReduceL2, kOnnxDomain, 13, 17, (Ort::KernelDefBuilder() diff --git a/musa/ep/src/kernels/reduction/reduce_max.cc b/musa/ep/src/kernels/reduction/reduce_max.cc index 36e449a..45506c5 100644 --- a/musa/ep/src/kernels/reduction/reduce_max.cc +++ b/musa/ep/src/kernels/reduction/reduce_max.cc @@ -35,6 +35,13 @@ class ReduceMax : public OpKernelBase { }; } // namespace +ONNX_OPERATOR_VERSIONED_KERNEL_EX( + ReduceMax, kOnnxDomain, 1, 12, + (Ort::KernelDefBuilder() + .AddTypeConstraint("T", ReduceMaxOpset1TensorTypes()) + .SetInputMemType(1, OrtMemTypeCPUInput)), + ReduceMax) + ONNX_OPERATOR_VERSIONED_KERNEL_EX( ReduceMax, kOnnxDomain, 13, 19, (Ort::KernelDefBuilder() diff --git a/musa/ep/src/kernels/reduction/reduce_mean.cc b/musa/ep/src/kernels/reduction/reduce_mean.cc index 897d7c8..d5a5e56 100644 --- a/musa/ep/src/kernels/reduction/reduce_mean.cc +++ b/musa/ep/src/kernels/reduction/reduce_mean.cc @@ -31,6 +31,13 @@ class ReduceMean : public OpKernelBase { }; } // namespace +ONNX_OPERATOR_VERSIONED_KERNEL_EX( + ReduceMean, kOnnxDomain, 1, 12, + (Ort::KernelDefBuilder() + .AddTypeConstraint("T", ReduceFloatOpset1TensorTypes()) + .SetInputMemType(1, OrtMemTypeCPUInput)), + ReduceMean) + ONNX_OPERATOR_VERSIONED_KERNEL_EX( ReduceMean, kOnnxDomain, 13, 19, (Ort::KernelDefBuilder() diff --git a/musa/ep/src/kernels/reduction/reduce_prod.cc b/musa/ep/src/kernels/reduction/reduce_prod.cc index 1713f6b..b79aae0 100644 --- a/musa/ep/src/kernels/reduction/reduce_prod.cc +++ b/musa/ep/src/kernels/reduction/reduce_prod.cc @@ -35,6 +35,13 @@ class ReduceProd : public OpKernelBase { }; } // namespace +ONNX_OPERATOR_VERSIONED_KERNEL_EX( + ReduceProd, kOnnxDomain, 1, 12, + (Ort::KernelDefBuilder() + .AddTypeConstraint("T", AllTensorTypesNoBFloat16()) + .SetInputMemType(1, OrtMemTypeCPUInput)), + ReduceProd) + ONNX_OPERATOR_VERSIONED_KERNEL_EX(ReduceProd, kOnnxDomain, 13, 19, (Ort::KernelDefBuilder() .AddTypeConstraint("T", AllTensorTypes()) diff --git a/musa/ep/src/kernels/reduction/reduce_sum.cc b/musa/ep/src/kernels/reduction/reduce_sum.cc index 69c99f0..ec16c0f 100644 --- a/musa/ep/src/kernels/reduction/reduce_sum.cc +++ b/musa/ep/src/kernels/reduction/reduce_sum.cc @@ -35,6 +35,13 @@ class ReduceSum : public OpKernelBase { }; } // namespace +ONNX_OPERATOR_VERSIONED_KERNEL_EX( + ReduceSum, kOnnxDomain, 1, 12, + (Ort::KernelDefBuilder() + .AddTypeConstraint("T", AllTensorTypesNoBFloat16()) + .SetInputMemType(1, OrtMemTypeCPUInput)), + ReduceSum) + ONNX_OPERATOR_VERSIONED_KERNEL_EX(ReduceSum, kOnnxDomain, 13, 19, (Ort::KernelDefBuilder() .AddTypeConstraint("T", AllTensorTypes()) diff --git a/musa/ep/src/kernels/reduction/reduce_sum_square.cc b/musa/ep/src/kernels/reduction/reduce_sum_square.cc index b85a34e..6c0c221 100644 --- a/musa/ep/src/kernels/reduction/reduce_sum_square.cc +++ b/musa/ep/src/kernels/reduction/reduce_sum_square.cc @@ -35,6 +35,12 @@ class ReduceSumSquare : public OpKernelBase { }; } // namespace +ONNX_OPERATOR_VERSIONED_KERNEL_EX(ReduceSumSquare, kOnnxDomain, 1, 12, + (Ort::KernelDefBuilder() + .AddTypeConstraint("T", HfdTensorTypes()) + .SetInputMemType(1, OrtMemTypeCPUInput)), + ReduceSumSquare) + ONNX_OPERATOR_VERSIONED_KERNEL_EX( ReduceSumSquare, kOnnxDomain, 13, 19, (Ort::KernelDefBuilder() diff --git a/musa/ep/src/kernels/shared_inc/kernel_type_constraints.h b/musa/ep/src/kernels/shared_inc/kernel_type_constraints.h index 82369e8..193b784 100644 --- a/musa/ep/src/kernels/shared_inc/kernel_type_constraints.h +++ b/musa/ep/src/kernels/shared_inc/kernel_type_constraints.h @@ -30,6 +30,18 @@ inline std::vector AllTensorTypes() { }; } +// Everything AllTensorTypes() covers except bfloat16, i.e. the type set of the +// ONNX schemas that predate opset 13. +inline std::vector AllTensorTypesNoBFloat16() { + return { + GetTensorType(ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT16), + GetTensorType(ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT), + GetTensorType(ONNX_TENSOR_ELEMENT_DATA_TYPE_DOUBLE), + GetTensorType(ONNX_TENSOR_ELEMENT_DATA_TYPE_INT32), + GetTensorType(ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64), + }; +} + inline std::vector TensorTypesWithBool() { return { GetTensorType(ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT), @@ -74,6 +86,20 @@ inline std::vector AllFixedSizeTensorTypesNoBFloat16() { }; } +// ONNX binary elementwise ops through opset 12 share the opset-13 type set +// without bfloat16, which the schemas only allow from opset 13 on. +inline std::vector BinaryNumericOpset7TensorTypes() { + return { + GetTensorType(ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT32), + GetTensorType(ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT64), + GetTensorType(ONNX_TENSOR_ELEMENT_DATA_TYPE_INT32), + GetTensorType(ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64), + GetTensorType(ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT16), + GetTensorType(ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT), + GetTensorType(ONNX_TENSOR_ELEMENT_DATA_TYPE_DOUBLE), + }; +} + inline std::vector BinaryNumericOpset13TensorTypes() { return { GetTensorType(ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT32), @@ -104,6 +130,19 @@ inline std::vector BinaryNumericOpset14TensorTypes() { }; } +// Pow opset 12: integer bases and exponents join the floating-point types, and +// the exponent may differ from the base in type. bfloat16 only arrives with the +// opset-15 schema. +inline std::vector PowOpset12TensorTypes() { + return { + GetTensorType(ONNX_TENSOR_ELEMENT_DATA_TYPE_INT32), + GetTensorType(ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64), + GetTensorType(ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT16), + GetTensorType(ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT), + GetTensorType(ONNX_TENSOR_ELEMENT_DATA_TYPE_DOUBLE), + }; +} + inline std::vector PowTensorTypes() { return { GetTensorType(ONNX_TENSOR_ELEMENT_DATA_TYPE_INT32), @@ -209,6 +248,18 @@ inline std::vector RangeTensorTypes() { }; } +// ReduceMean / ReduceL2 through opset 12: the floating-point subset plus +// int32 that the pre-opset-13 schemas allow and the kernels implement (the +// schemas only gain bfloat16 at opset 13). +inline std::vector ReduceFloatOpset1TensorTypes() { + return { + GetTensorType(ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT16), + GetTensorType(ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT), + GetTensorType(ONNX_TENSOR_ELEMENT_DATA_TYPE_DOUBLE), + GetTensorType(ONNX_TENSOR_ELEMENT_DATA_TYPE_INT32), + }; +} + inline std::vector ReduceMeanTensorTypes() { return { GetTensorType(ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT16), @@ -290,6 +341,23 @@ inline std::vector GatherNDTensorTypes() { }; } +// Abs opset 6-12: same as the opset-13 set without bfloat16. +inline std::vector AbsOpset6TensorTypes() { + return { + GetTensorType(ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT8), + GetTensorType(ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT16), + GetTensorType(ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT32), + GetTensorType(ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT64), + GetTensorType(ONNX_TENSOR_ELEMENT_DATA_TYPE_INT8), + GetTensorType(ONNX_TENSOR_ELEMENT_DATA_TYPE_INT16), + GetTensorType(ONNX_TENSOR_ELEMENT_DATA_TYPE_INT32), + GetTensorType(ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64), + GetTensorType(ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT16), + GetTensorType(ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT), + GetTensorType(ONNX_TENSOR_ELEMENT_DATA_TYPE_DOUBLE), + }; +} + inline std::vector AbsTensorTypes() { return { GetTensorType(ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT8), @@ -307,6 +375,19 @@ inline std::vector AbsTensorTypes() { }; } +// Neg opset 6-12: same as the opset-13 set without bfloat16. +inline std::vector NegOpset6TensorTypes() { + return { + GetTensorType(ONNX_TENSOR_ELEMENT_DATA_TYPE_INT8), + GetTensorType(ONNX_TENSOR_ELEMENT_DATA_TYPE_INT16), + GetTensorType(ONNX_TENSOR_ELEMENT_DATA_TYPE_INT32), + GetTensorType(ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64), + GetTensorType(ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT16), + GetTensorType(ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT), + GetTensorType(ONNX_TENSOR_ELEMENT_DATA_TYPE_DOUBLE), + }; +} + inline std::vector NegTensorTypes() { return { GetTensorType(ONNX_TENSOR_ELEMENT_DATA_TYPE_INT8), @@ -350,6 +431,19 @@ inline std::vector BitwiseIntegerTensorTypes() { }; } +// ReduceMax through opset 12: the opset-13 set without bfloat16. +inline std::vector ReduceMaxOpset1TensorTypes() { + return { + GetTensorType(ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT8), + GetTensorType(ONNX_TENSOR_ELEMENT_DATA_TYPE_INT8), + GetTensorType(ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT16), + GetTensorType(ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT), + GetTensorType(ONNX_TENSOR_ELEMENT_DATA_TYPE_DOUBLE), + GetTensorType(ONNX_TENSOR_ELEMENT_DATA_TYPE_INT32), + GetTensorType(ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64), + }; +} + inline std::vector ReduceMaxTensorTypes() { return { GetTensorType(ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT8), @@ -363,6 +457,16 @@ inline std::vector ReduceMaxTensorTypes() { }; } +// Less / Greater opset 7-8 only accept floating-point tensors; integer support +// arrives with the opset-9 schema. +inline std::vector CompareOpset7TensorTypes() { + return { + GetTensorType(ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT16), + GetTensorType(ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT), + GetTensorType(ONNX_TENSOR_ELEMENT_DATA_TYPE_DOUBLE), + }; +} + inline std::vector CompareTensorTypesNoBFloat16() { return { GetTensorType(ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT32), @@ -388,6 +492,30 @@ inline std::vector CompareTensorTypes() { }; } +// Equal opset 7-10 accepts bool and the integer types only; the remaining +// numeric types arrive with the opset-11 schema. +inline std::vector EqualOpset7TensorTypes() { + return { + GetTensorType(ONNX_TENSOR_ELEMENT_DATA_TYPE_INT32), + GetTensorType(ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64), + GetTensorType(ONNX_TENSOR_ELEMENT_DATA_TYPE_BOOL), + }; +} + +// Equal opset 11-12: the opset-13 set without bfloat16. +inline std::vector EqualOpset11TensorTypes() { + return { + GetTensorType(ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT32), + GetTensorType(ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT64), + GetTensorType(ONNX_TENSOR_ELEMENT_DATA_TYPE_INT32), + GetTensorType(ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64), + GetTensorType(ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT16), + GetTensorType(ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT), + GetTensorType(ONNX_TENSOR_ELEMENT_DATA_TYPE_DOUBLE), + GetTensorType(ONNX_TENSOR_ELEMENT_DATA_TYPE_BOOL), + }; +} + inline std::vector EqualTensorTypes() { return { GetTensorType(ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT32), diff --git a/musa/ep/src/kernels/tensor/cast.cc b/musa/ep/src/kernels/tensor/cast.cc index c08a637..89f620e 100644 --- a/musa/ep/src/kernels/tensor/cast.cc +++ b/musa/ep/src/kernels/tensor/cast.cc @@ -132,6 +132,20 @@ OrtStatus* Cast::Compute(Ort::KernelContext& ctx) const { } } // namespace +ONNX_OPERATOR_VERSIONED_KERNEL_EX( + Cast, kOnnxDomain, 6, 8, + (Ort::KernelDefBuilder() + .AddTypeConstraint("T1", AllFixedSizeTensorTypesNoBFloat16()) + .AddTypeConstraint("T2", AllFixedSizeTensorTypesNoBFloat16())), + Cast) + +ONNX_OPERATOR_VERSIONED_KERNEL_EX( + Cast, kOnnxDomain, 9, 12, + (Ort::KernelDefBuilder() + .AddTypeConstraint("T1", AllFixedSizeTensorTypesNoBFloat16()) + .AddTypeConstraint("T2", AllFixedSizeTensorTypesNoBFloat16())), + Cast) + ONNX_OPERATOR_VERSIONED_KERNEL_EX( Cast, kOnnxDomain, 13, 19, (Ort::KernelDefBuilder() diff --git a/musa/ep/src/kernels/tensor/concat.cc b/musa/ep/src/kernels/tensor/concat.cc index a6a3238..e7e61df 100644 --- a/musa/ep/src/kernels/tensor/concat.cc +++ b/musa/ep/src/kernels/tensor/concat.cc @@ -401,6 +401,18 @@ OrtStatus* Concat::Compute(Ort::KernelContext& ctx) const { } } // namespace +ONNX_OPERATOR_VERSIONED_KERNEL_EX(Concat, kOnnxDomain, 4, 10, + (Ort::KernelDefBuilder().AddTypeConstraint( + "T", + AllFixedSizeTensorTypesNoBFloat16())), + Concat) + +ONNX_OPERATOR_VERSIONED_KERNEL_EX(Concat, kOnnxDomain, 11, 12, + (Ort::KernelDefBuilder().AddTypeConstraint( + "T", + AllFixedSizeTensorTypesNoBFloat16())), + Concat) + ONNX_OPERATOR_VERSIONED_KERNEL_EX( Concat, kOnnxDomain, 13, 19, (Ort::KernelDefBuilder().AddTypeConstraint("T", AllFixedSizeTensorTypes())), diff --git a/musa/ep/src/kernels/tensor/expand.cc b/musa/ep/src/kernels/tensor/expand.cc index 1db7934..309f7d0 100644 --- a/musa/ep/src/kernels/tensor/expand.cc +++ b/musa/ep/src/kernels/tensor/expand.cc @@ -62,6 +62,13 @@ OrtStatus* Expand::Compute(Ort::KernelContext& ctx) const { } } // namespace +ONNX_OPERATOR_VERSIONED_KERNEL_EX( + Expand, kOnnxDomain, 8, 12, + (Ort::KernelDefBuilder() + .AddTypeConstraint("T", AllFixedSizeTensorTypesNoBFloat16()) + .SetInputMemType(1, OrtMemTypeCPUInput)), + Expand) + ONNX_OPERATOR_VERSIONED_KERNEL_EX( Expand, kOnnxDomain, 13, 19, (Ort::KernelDefBuilder() diff --git a/musa/ep/src/kernels/tensor/gather.cc b/musa/ep/src/kernels/tensor/gather.cc index 50aceda..791f007 100644 --- a/musa/ep/src/kernels/tensor/gather.cc +++ b/musa/ep/src/kernels/tensor/gather.cc @@ -154,6 +154,20 @@ OrtStatus* Gather::Compute(Ort::KernelContext& ctx) const { } } // namespace +ONNX_OPERATOR_VERSIONED_KERNEL_EX( + Gather, kOnnxDomain, 1, 10, + (Ort::KernelDefBuilder() + .AddTypeConstraint("T", AllFixedSizeTensorTypesNoBFloat16()) + .AddTypeConstraint("Tind", IntTensorTypes())), + Gather) + +ONNX_OPERATOR_VERSIONED_KERNEL_EX( + Gather, kOnnxDomain, 11, 12, + (Ort::KernelDefBuilder() + .AddTypeConstraint("T", AllFixedSizeTensorTypesNoBFloat16()) + .AddTypeConstraint("Tind", IntTensorTypes())), + Gather) + ONNX_OPERATOR_VERSIONED_KERNEL_EX( Gather, kOnnxDomain, 13, 19, (Ort::KernelDefBuilder() diff --git a/musa/ep/src/kernels/tensor/nonzero.cc b/musa/ep/src/kernels/tensor/nonzero.cc index 5427c46..56742ff 100644 --- a/musa/ep/src/kernels/tensor/nonzero.cc +++ b/musa/ep/src/kernels/tensor/nonzero.cc @@ -135,6 +135,11 @@ OrtStatus* NonZero::Compute(Ort::KernelContext& ctx) const { } } // namespace +ONNX_OPERATOR_VERSIONED_KERNEL_EX( + NonZero, kOnnxDomain, 9, 12, + (Ort::KernelDefBuilder().AddTypeConstraint("T", NonZeroTensorTypes())), + NonZero) + ONNX_OPERATOR_VERSIONED_KERNEL_EX( NonZero, kOnnxDomain, 13, 19, (Ort::KernelDefBuilder().AddTypeConstraint("T", NonZeroTensorTypes())), diff --git a/musa/ep/src/kernels/tensor/reshape.cc b/musa/ep/src/kernels/tensor/reshape.cc index b96a20d..7a68dee 100644 --- a/musa/ep/src/kernels/tensor/reshape.cc +++ b/musa/ep/src/kernels/tensor/reshape.cc @@ -106,6 +106,13 @@ OrtStatus* ReshapeString::Compute(Ort::KernelContext& ctx) const { } } // namespace +ONNX_OPERATOR_VERSIONED_KERNEL_EX( + Reshape, kOnnxDomain, 5, 12, + (Ort::KernelDefBuilder() + .AddTypeConstraint("T", AllFixedSizeTensorTypesNoBFloat16()) + .SetInputMemType(1, OrtMemTypeCPUInput)), + Reshape) + ONNX_OPERATOR_VERSIONED_KERNEL_EX( Reshape, kOnnxDomain, 13, 19, (Ort::KernelDefBuilder() diff --git a/musa/ep/src/kernels/tensor/shape.cc b/musa/ep/src/kernels/tensor/shape.cc index 8ad584c..3f3379c 100644 --- a/musa/ep/src/kernels/tensor/shape.cc +++ b/musa/ep/src/kernels/tensor/shape.cc @@ -28,6 +28,15 @@ OrtStatus* Shape::Compute(Ort::KernelContext& ctx) const { } } // namespace +ONNX_OPERATOR_VERSIONED_KERNEL_EX( + Shape, kOnnxDomain, 1, 12, + (Ort::KernelDefBuilder() + .AddTypeConstraint("T", AllFixedSizeTensorTypesNoBFloat16()) + .AddTypeConstraint("T1", + GetTensorType(ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64)) + .SetOutputMemType(0, OrtMemTypeCPUInput)), + Shape) + ONNX_OPERATOR_VERSIONED_KERNEL_EX( Shape, kOnnxDomain, 13, 19, (Ort::KernelDefBuilder() diff --git a/musa/ep/src/kernels/tensor/transpose.cc b/musa/ep/src/kernels/tensor/transpose.cc index f686f05..ca9ff25 100644 --- a/musa/ep/src/kernels/tensor/transpose.cc +++ b/musa/ep/src/kernels/tensor/transpose.cc @@ -111,6 +111,12 @@ OrtStatus* Transpose::Compute(Ort::KernelContext& ctx) const { } } // namespace +ONNX_OPERATOR_VERSIONED_KERNEL_EX(Transpose, kOnnxDomain, 1, 12, + (Ort::KernelDefBuilder().AddTypeConstraint( + "T", + AllFixedSizeTensorTypesNoBFloat16())), + Transpose) + ONNX_OPERATOR_VERSIONED_KERNEL_EX( Transpose, kOnnxDomain, 13, 19, (Ort::KernelDefBuilder().AddTypeConstraint("T", AllFixedSizeTensorTypes())), diff --git a/test/ops/test_opset_lower_bound.py b/test/ops/test_opset_lower_bound.py new file mode 100644 index 0000000..1de374a --- /dev/null +++ b/test/ops/test_opset_lower_bound.py @@ -0,0 +1,334 @@ +# Copyright (c) Moore Threads Technology Co., Ltd. All rights reserved. +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Old-opset coverage for the kernels whose registration was extended down. + +These kernels were previously only registered from opset 13 on, which left every +node of an opset 10/11/12 model on the CPU EP. The registrations now reuse the +same kernel class for the older opsets, so each case here builds the same +single-node graph at opset 10 and opset 12 - the opsets the low-opset models use +- and checks two things: + + 1. the MUSA EP claims the node (the session runs with + `session.disable_cpu_ep_fallback`, and the EP's dtype report must list the + node as `assigned=1`), and + 2. its output matches the stock CPU EP. +""" + +from __future__ import annotations + +import collections + +import numpy as np +import onnxruntime as ort +import pytest +from onnx import helper + +from op_test_utils import TensorProto, build_model, musa_devices, run + +Case = collections.namedtuple( + "Case", "op_type opset inputs outputs attrs label", defaults=(None, None) +) + +# `Shape` folds away when its input shape is static, so that case builds its +# model with a symbolic batch dimension instead. +_SYMBOLIC_BATCH_BUILDERS = {"Shape"} + +# The lowest opsets the models in the zoo use; opset 11 sits between them and is +# covered by the same registration ranges. +_LOW_OPSETS = (10, 12) + + +def _build_model_with_symbolic_batch(op_type, inputs, outputs, attrs=None, opset=17): + """Build a single-node graph whose first input dimension stays symbolic. + + `Shape` folds away at graph-optimization time when the input shape is fully + static, and the graph would reach the EP empty. Keeping one dimension + symbolic is also how the real models present it (a dynamic batch). + """ + input_vis = [] + for name, array in inputs.items(): + shape = ["batch", *list(array.shape[1:])] + input_vis.append( + helper.make_tensor_value_info( + name, helper.np_dtype_to_tensor_dtype(array.dtype), shape + ) + ) + output_vis = [helper.make_tensor_value_info(n, e, None) for n, e in outputs] + node = helper.make_node( + op_type, list(inputs.keys()), [n for n, _ in outputs], **(attrs or {}) + ) + graph = helper.make_graph([node], f"{op_type}_graph", input_vis, output_vis) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", opset)]) + model.ir_version = min(model.ir_version, 10) + return model.SerializeToString() + + +def _binary_cases(): + rng = np.random.default_rng(7) + a = rng.uniform(0.5, 2.0, (2, 1, 3)).astype(np.float32) + b = rng.uniform(0.5, 2.0, (1, 4, 3)).astype(np.float32) + for opset in _LOW_OPSETS: + for op_type in ("Add", "Sub", "Mul", "Div", "Pow"): + yield Case(op_type, opset, {"A": a, "B": b}, [("Y", TensorProto.FLOAT)]) + + +def _unary_cases(): + rng = np.random.default_rng(11) + x = rng.uniform(0.5, 2.0, (2, 3, 4)).astype(np.float32) + for opset in _LOW_OPSETS: + for op_type in ( + "Abs", + "Exp", + "Floor", + "Log", + "Neg", + "Relu", + "Sigmoid", + "Sqrt", + ): + yield Case(op_type, opset, {"X": x}, [("Y", TensorProto.FLOAT)]) + + +def _compare_cases(): + rng = np.random.default_rng(13) + int_a = rng.integers(0, 3, (2, 3)).astype(np.int32) + int_b = rng.integers(0, 3, (3,)).astype(np.int32) + float_a = int_a.astype(np.float32) + float_b = int_b.astype(np.float32) + # Equal through opset 10 only accepts bool and integer tensors, and only + # gains the floating-point types with the opset-11 schema. + yield Case("Equal", 10, {"A": int_a, "B": int_b}, [("Y", TensorProto.BOOL)]) + yield Case("Equal", 11, {"A": float_a, "B": float_b}, [("Y", TensorProto.BOOL)]) + for op_type in ("Equal", "Less", "Greater"): + yield Case( + op_type, + 12, + {"A": float_a, "B": float_b}, + [("Y", TensorProto.BOOL)], + ) + for op_type in ("Less", "Greater"): + yield Case( + op_type, + 10, + {"A": float_a, "B": float_b}, + [("Y", TensorProto.BOOL)], + ) + + +def _tensor_cases(): + rng = np.random.default_rng(17) + x = rng.standard_normal((4, 6)).astype(np.float32) + for opset in _LOW_OPSETS: + yield Case( + "Transpose", + opset, + {"X": x.reshape(2, 3, 4)}, + [("Y", TensorProto.FLOAT)], + {"perm": [2, 0, 1]}, + ) + yield Case( + "Reshape", + opset, + {"X": x, "shape": np.array([3, 8], dtype=np.int64)}, + [("Y", TensorProto.FLOAT)], + ) + yield Case( + "Gather", + opset, + {"data": x, "indices": np.array([0, 2, 5], dtype=np.int64)}, + [("Y", TensorProto.FLOAT)], + {"axis": 1}, + ) + yield Case( + "Concat", + opset, + {"A": x, "B": x[:, :2]}, + [("Y", TensorProto.FLOAT)], + {"axis": 1}, + ) + yield Case( + "Expand", + opset, + {"X": x[0], "shape": np.array([3, 6], dtype=np.int64)}, + [("Y", TensorProto.FLOAT)], + ) + yield Case( + "Cast", + opset, + {"X": x}, + [("Y", TensorProto.INT64)], + {"to": TensorProto.INT64}, + ) + yield Case( + "Shape", + opset, + {"X": x.reshape(2, 3, 4)}, + [("Y", TensorProto.INT64)], + label="symbolic-batch", + ) + yield Case( + "NonZero", + opset, + {"X": np.array([[0.0, 2.0, 0.0], [3.0, 0.0, 4.0]], dtype=np.float32)}, + [("Y", TensorProto.INT64)], + ) + yield Case( + "MatMul", + opset, + {"A": x, "B": rng.standard_normal((6, 5)).astype(np.float32)}, + [("Y", TensorProto.FLOAT)], + ) + + +def _reduce_cases(): + rng = np.random.default_rng(19) + x = rng.uniform(0.5, 2.0, (4, 8)).astype(np.float32) + for opset in _LOW_OPSETS: + for op_type in ( + "ReduceL2", + "ReduceMax", + "ReduceMean", + "ReduceProd", + "ReduceSum", + "ReduceSumSquare", + ): + # opset 12 and older carry the axes as an attribute, not as an input. + yield Case( + op_type, + opset, + {"X": x}, + [("Y", TensorProto.FLOAT)], + {"axes": [1], "keepdims": 1}, + ) + + +def _reduce_sum_attribute_cases(): + rng = np.random.default_rng(31) + x = rng.uniform(0.5, 2.0, (4, 8)).astype(np.float32) + # ReduceSum is the one Reduce* whose axes moved from an attribute to an input + # at opset 13 (the others follow at opset 18), so the opset-12-and-older form + # is the attribute and the kernel has to read it instead of input 1. + yield Case( + "ReduceSum", + 11, + {"X": x}, + [("Y", TensorProto.FLOAT)], + {"axes": [1], "keepdims": 0}, + label="axes-keepdims0", + ) + # The pre-opset-13 schemas make `axes` optional: omitting it reduces every + # axis, which is what an empty axes attribute has to mean. + yield Case( + "ReduceSum", + 11, + {"X": x}, + [("Y", TensorProto.FLOAT)], + {"keepdims": 1}, + label="axes-omitted", + ) + + +def _conv_cases(): + x = np.random.default_rng(23).standard_normal((1, 2, 6, 6)).astype(np.float32) + w = np.random.default_rng(29).standard_normal((3, 2, 3, 3)).astype(np.float32) + # Conv-11 only clarifies the default strides and dilations, so opset 10 runs + # the same kernel. + yield Case( + "Conv", + 10, + {"X": x, "W": w}, + [("Y", TensorProto.FLOAT)], + {"pads": [1, 1, 1, 1], "strides": [1, 1]}, + ) + + +CASES = [ + *list(_binary_cases()), + *list(_unary_cases()), + *list(_compare_cases()), + *list(_tensor_cases()), + *list(_reduce_cases()), + *list(_reduce_sum_attribute_cases()), + *list(_conv_cases()), +] + +CASE_IDS = [ + f"{case.op_type}-opset{case.opset}" + (f"-{case.label}" if case.label else "") + for case in CASES +] + +assert len(CASE_IDS) == len(set(CASE_IDS)) +# Every operator whose registration was lowered has at least one case. +assert {case.op_type for case in CASES} == { + "Abs", "Add", "Cast", "Concat", "Conv", "Div", "Equal", "Exp", "Expand", + "Floor", "Gather", "Greater", "Less", "Log", "MatMul", "Mul", "Neg", + "NonZero", "Pow", + "ReduceL2", "ReduceMax", "ReduceMean", "ReduceProd", "ReduceSum", + "ReduceSumSquare", "Relu", "Reshape", "Shape", "Sigmoid", "Sqrt", "Sub", + "Transpose", +} + + +def _run_case(capfd, case): + """Run one case on CPU and MUSA; assert the CPU result and the assignment.""" + context = f"{case.op_type} at opset {case.opset}" + if case.op_type in _SYMBOLIC_BATCH_BUILDERS: + model_bytes = _build_model_with_symbolic_batch( + case.op_type, case.inputs, case.outputs, case.attrs, case.opset + ) + else: + model_bytes = build_model( + case.op_type, case.inputs, case.outputs, case.attrs, "", case.opset + ) + + cpu_outputs = run(model_bytes, case.inputs, use_musa=False) + + devices = musa_devices() + if not devices: + raise RuntimeError("No MUSA device available for the opset lower bound test") + options = ort.SessionOptions() + options.add_session_config_entry("session.disable_cpu_ep_fallback", "1") + options.add_provider_for_devices(devices, {"dtype_diagnostics": "1"}) + capfd.readouterr() # drop whatever an earlier case printed + session = ort.InferenceSession(model_bytes, sess_options=options) + musa_outputs = session.run(None, dict(case.inputs)) + + assert len(cpu_outputs) == len(musa_outputs), f"{context}: output count" + for index, (cpu_output, musa_output) in enumerate(zip(cpu_outputs, musa_outputs)): + assert cpu_output.shape == musa_output.shape, f"{context}: output[{index}] shape" + assert cpu_output.dtype == musa_output.dtype, f"{context}: output[{index}] dtype" + np.testing.assert_allclose( + musa_output, + cpu_output, + rtol=1e-3, + atol=1e-4, + err_msg=f"{context}: output[{index}] CPU vs MUSA mismatch", + ) + + # The dtype report is printed by GetCapability for every node of the graph; + # this graph has one node, and `assigned=1` means a MUSA kernel took it. + node_lines = [ + line + for line in capfd.readouterr().err.splitlines() + if line.startswith("MUSA_DTYPE_NODE") + ] + assert len(node_lines) == 1, f"{context}: expected one reported node, got {node_lines}" + assert f" op={case.op_type} " in node_lines[0], f"{context}: {node_lines[0]}" + assert node_lines[0].rstrip().endswith( + "assigned=1" + ), f"{context}: node not claimed by a MUSA kernel: {node_lines[0]}" + + +@pytest.mark.parametrize("case", CASES, ids=CASE_IDS) +def test_opset_lower_bound(capfd, case): + _run_case(capfd, case) From f17b75f5da4340070eff458f5bc7df3995af7548 Mon Sep 17 00:00:00 2001 From: "yichao.hu" Date: Mon, 12 Oct 2026 23:47:31 -0400 Subject: [PATCH 16/52] feat(plugin): handle old-opset semantics and widen opset upper bounds Some operators changed semantics before opset 13: Softmax coerces its input to 2D, and Unsqueeze/Squeeze/Split/Clip/Pad/Slice take attributes instead of inputs. Branch on the node's opset in those kernels and register the old ranges, matching the ORT CUDA EP. Keeping Softmax on the CPU had forced large attention tensors across the CPU/device boundary. Widen the upper bounds of Reshape/Transpose/Shape/Squeeze/Unsqueeze to 21 and ReduceMax to 20. Read scalar FLOAT attributes correctly, keep Slice nodes with a negative step on the CPU EP, and keep Resize at opset <= 17. --- musa/docs/supported_ops.md | 26 +- musa/ep/src/dtype_contracts.json | 84 ++-- musa/ep/src/graph/capability_predicates.cc | 29 ++ musa/ep/src/kernels/math/clip.cc | 46 +- musa/ep/src/kernels/math/clip_impl.mu | 28 +- musa/ep/src/kernels/math/gemm.cc | 19 + musa/ep/src/kernels/math/softmax.cc | 62 ++- musa/ep/src/kernels/reduction/reduce_max.cc | 5 +- .../kernels/shared_inc/device_kernel_types.h | 14 + .../kernels/shared_inc/kernel_attributes.h | 17 + .../shared_inc/kernel_type_constraints.h | 19 + musa/ep/src/kernels/tensor/pad.cc | 48 +- musa/ep/src/kernels/tensor/pad_impl.mu | 32 ++ musa/ep/src/kernels/tensor/reshape.cc | 5 +- musa/ep/src/kernels/tensor/shape.cc | 4 +- musa/ep/src/kernels/tensor/slice.cc | 82 +++- musa/ep/src/kernels/tensor/split.cc | 36 ++ musa/ep/src/kernels/tensor/squeeze.cc | 22 +- musa/ep/src/kernels/tensor/transpose.cc | 4 +- musa/ep/src/kernels/tensor/unsqueeze.cc | 22 +- test/ops/test_opset_semantics.py | 414 ++++++++++++++++++ 21 files changed, 934 insertions(+), 84 deletions(-) create mode 100644 test/ops/test_opset_semantics.py diff --git a/musa/docs/supported_ops.md b/musa/docs/supported_ops.md index 8104030..9eb69a9 100644 --- a/musa/docs/supported_ops.md +++ b/musa/docs/supported_ops.md @@ -2,7 +2,7 @@ # Supported Ops - Total supported ops: **92** -- Total registered kernels: **164** +- Total registered kernels: **183** Rows are grouped by operator/domain. Multiple opset ranges in a row preserve separate ORT kernel registrations. @@ -15,7 +15,7 @@ Rows are grouped by operator/domain. Multiple opset ranges in a row preserve sep | `BitwiseAnd` | `ai.onnx` | 18–19: `T`: uint8, uint16, uint32, uint64, int8, int16, int32, int64 | `logical/bitwise_and.cc` | | `Cast` | `ai.onnx` | 6–8: `T1`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bool
`T2`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bool
9–12: `T1`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bool
`T2`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bool
13–19: `T1`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bfloat16, bool
`T2`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bfloat16, bool | `tensor/cast.cc` | | `Ceil` | `ai.onnx` | 13–19: `T`: float16, float32, float64, bfloat16 | `math/ceil.cc` | -| `Clip` | `ai.onnx` | 13–19: `T`: float16, float32, float64, int8, uint8, int16, uint16, int32, uint32, int64, uint64, bfloat16 | `math/clip.cc` | +| `Clip` | `ai.onnx` | 6–10: `T`: float16, float32, float64
11–11: `T`: float16, float32, float64
12–12: `T`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64
13–19: `T`: float16, float32, float64, int8, uint8, int16, uint16, int32, uint32, int64, uint64, bfloat16 | `math/clip.cc` | | `Concat` | `ai.onnx` | 4–10: `T`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bool
11–12: `T`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bool
13–19: `T`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bfloat16, bool | `tensor/concat.cc` | | `ConstantOfShape` | `ai.onnx` | 9–19: `T1`: int64
`T2`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bool | `generator/constant_of_shape.cc` | | `Conv` | `ai.onnx` | 1–10: `T`: float32
11–19: `T`: float32 | `nn/conv.cc` | @@ -31,7 +31,7 @@ Rows are grouped by operator/domain. Multiple opset ranges in a row preserve sep | `Gather` | `ai.onnx` | 1–10: `T`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bool
`Tind`: int32, int64
11–12: `T`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bool
`Tind`: int32, int64
13–19: `T`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bfloat16, bool
`Tind`: int32, int64 | `tensor/gather.cc` | | `GatherElements` | `ai.onnx` | 11–12: `T`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bfloat16, bool
`Tind`: int32, int64
13–19: `T`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bfloat16, bool
`Tind`: int32, int64 | `tensor/gather_elements.cc` | | `GatherND` | `ai.onnx` | 11–12: `T`: float32, float64, float16, int64, bfloat16, bool
`indices`: int64
13–19: `T`: float32, float64, float16, int64, bfloat16, bool
`indices`: int64 | `tensor/gather_nd.cc` | -| `Gemm` | `ai.onnx` | 13–19: `T`: float16, float32, float64, bfloat16 | `math/gemm.cc` | +| `Gemm` | `ai.onnx` | 7–8: `T`: float16, float32, float64
9–10: `T`: float16, float32, float64
11–12: `T`: float16, float32, float64
13–19: `T`: float16, float32, float64, bfloat16 | `math/gemm.cc` | | `GlobalAveragePool` | `ai.onnx` | 1–19: `T`: float16, float32, float64 | `nn/pool.cc` | | `GlobalMaxPool` | `ai.onnx` | 1–19: `T`: float16, float32, float64 | `nn/pool.cc` | | `Greater` | `ai.onnx` | 7–8: `T`: float16, float32, float64
9–12: `T`: uint32, uint64, int32, int64, float16, float32, float64
13–19: `T`: uint32, uint64, int32, int64, float16, float32, float64, bfloat16 | `logical/greater.cc` | @@ -59,42 +59,42 @@ Rows are grouped by operator/domain. Multiple opset ranges in a row preserve sep | `OneHot` | `ai.onnx` | 11–19: `T1`: int32, int64
`T2`: int32, int64
`T3`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bool | `tensor/onehot.cc` | | `Or` | `ai.onnx` | 7–19: `T`: bool | `logical/or.cc` | | `PRelu` | `ai.onnx` | 7–15: `T`: float16, float32, float64
16–19: `T`: float16, float32, float64, bfloat16, uint32, uint64, int32, int64 | `activation/prelu.cc` | -| `Pad` | `ai.onnx` | 13–19: `T`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bfloat16, bool | `tensor/pad.cc` | +| `Pad` | `ai.onnx` | 2–10: `T`: float16, float32, float64
11–12: `T`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64
13–19: `T`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bfloat16, bool | `tensor/pad.cc` | | `Pow` | `ai.onnx` | 7–11: `T`: float16, float32, float64
12–12: `T`: int32, int64, float16, float32, float64
`T1`: int32, int64, float16, float32, float64
13–14: `T`: int32, int64, float16, float32, float64, bfloat16
`T1`: int32, int64, float16, float32, float64
15–19: `T`: int32, int64, float16, float32, float64, bfloat16
`T1`: int32, int64, float16, float32, float64, bfloat16 | `math/pow.cc` | | `RandomUniform` | `ai.onnx` | 1–19: `T`: float16, float32, float64 | `generator/random_uniform.cc` | | `RandomUniformLike` | `ai.onnx` | 1–19: `T1`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bfloat16, bool
`T2`: float16, float32, float64 | `generator/random_uniform.cc` | | `Range` | `ai.onnx` | 11–19: `T`: int16, int32, int64, float32, float64 | `generator/range.cc` | | `Reciprocal` | `ai.onnx` | 13–19: `T`: float16, float32, float64, bfloat16 | `activation/reciprocal.cc` | | `ReduceL2` | `ai.onnx` | 1–12: `T`: float16, float32, float64, int32
13–17: `T`: float16, float32, float64, int32
18–19: `T`: float16, float32, float64, bfloat16, int32 | `reduction/reduce_l2.cc` | -| `ReduceMax` | `ai.onnx` | 1–12: `T`: uint8, int8, float16, float32, float64, int32, int64
13–19: `T`: uint8, int8, float16, float32, float64, bfloat16, int32, int64 | `reduction/reduce_max.cc` | +| `ReduceMax` | `ai.onnx` | 1–12: `T`: uint8, int8, float16, float32, float64, int32, int64
13–20: `T`: uint8, int8, float16, float32, float64, bfloat16, int32, int64 | `reduction/reduce_max.cc` | | `ReduceMean` | `ai.onnx` | 1–12: `T`: float16, float32, float64, int32
13–19: `T`: float16, float32, float64, bfloat16, int32 | `reduction/reduce_mean.cc` | | `ReduceProd` | `ai.onnx` | 1–12: `T`: float16, float32, float64, int32, int64
13–19: `T`: float16, float32, float64, bfloat16, int32, int64 | `reduction/reduce_prod.cc` | | `ReduceSum` | `ai.onnx` | 1–12: `T`: float16, float32, float64, int32, int64
13–19: `T`: float16, float32, float64, bfloat16, int32, int64 | `reduction/reduce_sum.cc` | | `ReduceSumSquare` | `ai.onnx` | 1–12: `T`: float16, float32, float64
13–19: `T`: float16, float32, float64, bfloat16 | `reduction/reduce_sum_square.cc` | | `Relu` | `ai.onnx` | 6–12: `T`: float16, float32, float64
13–19: `T`: float16, float32, float64, bfloat16 | `activation/relu.cc` | -| `Reshape` | `ai.onnx` | 5–12: `T`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bool
13–19: `T`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bfloat16, bool
19–19: `T`: string | `tensor/reshape.cc` | +| `Reshape` | `ai.onnx` | 5–12: `T`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bool
13–21: `T`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bfloat16, bool
19–19: `T`: string | `tensor/reshape.cc` | | `Resize` | `ai.onnx` | 13–17: `T1`: float32 | `tensor/resize.cc` | | `ReverseSequence` | `ai.onnx` | 10–19: `T`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bool | `tensor/reverse_sequence.cc` | | `Round` | `ai.onnx` | 11–19: `T`: float16, float32, float64 | `math/round.cc` | | `ScatterElements` | `ai.onnx` | 11–12: `T`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bfloat16, bool
`Tind`: int32, int64
13–19: `T`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bfloat16, bool
`Tind`: int32, int64 | `tensor/scatter_elements.cc` | | `ScatterND` | `ai.onnx` | 11–12: `T`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bfloat16, bool
`indices`: int64
13–19: `T`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bfloat16, bool
`indices`: int64 | `tensor/scatter_nd.cc` | -| `Shape` | `ai.onnx` | 1–12: `T`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bool
`T1`: int64
13–19: `T`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bfloat16, bool
`T1`: int64 | `tensor/shape.cc` | +| `Shape` | `ai.onnx` | 1–12: `T`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bool
`T1`: int64
13–21: `T`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bfloat16, bool
`T1`: int64 | `tensor/shape.cc` | | `Sigmoid` | `ai.onnx` | 6–12: `T`: float16, float32, float64
13–19: `T`: float16, float32, float64, bfloat16 | `activation/sigmoid.cc` | | `Sign` | `ai.onnx` | 13–19: `T`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bfloat16 | `math/sign.cc` | -| `Slice` | `ai.onnx` | 13–19: `T`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bfloat16, bool
`Tind`: int32, int64 | `tensor/slice.cc` | -| `Softmax` | `ai.onnx` | 13–17: `T`: float16, float32, float64, bfloat16 | `math/softmax.cc` | +| `Slice` | `ai.onnx` | 1–9: `T`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bool
10–10: `T`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bool
`Tind`: int32, int64
11–12: `T`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bool
`Tind`: int32, int64
13–19: `T`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bfloat16, bool
`Tind`: int32, int64 | `tensor/slice.cc` | +| `Softmax` | `ai.onnx` | 1–10: `T`: float16, float32, float64
11–12: `T`: float16, float32, float64
13–17: `T`: float16, float32, float64, bfloat16 | `math/softmax.cc` | | `Softplus` | `ai.onnx` | 1–19: `T`: float16, float32, float64 | `activation/softplus.cc` | -| `Split` | `ai.onnx` | 13–19: `T`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bfloat16, bool | `tensor/split.cc` | +| `Split` | `ai.onnx` | 2–10: `T`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bool
11–12: `T`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bool
13–19: `T`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bfloat16, bool | `tensor/split.cc` | | `Sqrt` | `ai.onnx` | 6–12: `T`: float16, float32, float64
13–19: `T`: float16, float32, float64, bfloat16 | `activation/sqrt.cc` | -| `Squeeze` | `ai.onnx` | 13–19: `T`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bfloat16, bool | `tensor/squeeze.cc` | +| `Squeeze` | `ai.onnx` | 1–10: `T`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bool
11–12: `T`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bool
13–21: `T`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bfloat16, bool | `tensor/squeeze.cc` | | `Sub` | `ai.onnx` | 7–12: `T`: uint32, uint64, int32, int64, float16, float32, float64
13–13: `T`: uint32, uint64, int32, int64, float16, float32, float64, bfloat16
14–19: `T`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bfloat16 | `math/sub.cc` | | `Sum` | `ai.onnx` | 13–19: `T`: float16, float32, float64, bfloat16 | `math/sum.cc` | | `Tanh` | `ai.onnx` | 13–19: `T`: float16, float32, float64, bfloat16 | `activation/tanh.cc` | | `Tile` | `ai.onnx` | 13–19: `T`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bfloat16, bool | `tensor/tile.cc` | | `TopK` | `ai.onnx` | 10–10: `T`: uint8, uint16, uint32, uint64, int8, int16, float16, float32, float64, int32, int64
`I`: int64
11–19: `T`: uint8, uint16, uint32, uint64, int8, int16, float16, float32, float64, int32, int64
`I`: int64 | `math/topk.cc` | -| `Transpose` | `ai.onnx` | 1–12: `T`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bool
13–19: `T`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bfloat16, bool | `tensor/transpose.cc` | +| `Transpose` | `ai.onnx` | 1–12: `T`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bool
13–21: `T`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bfloat16, bool | `tensor/transpose.cc` | | `Unique` | `ai.onnx` | 11–19: `T`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bool | `tensor/unique.cc` | -| `Unsqueeze` | `ai.onnx` | 13–19: `T`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bfloat16, bool | `tensor/unsqueeze.cc` | +| `Unsqueeze` | `ai.onnx` | 1–10: `T`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bool
11–12: `T`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bool
13–21: `T`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bfloat16, bool | `tensor/unsqueeze.cc` | | `Where` | `ai.onnx` | 9–15: `B`: bool
`T`: float16, float32, float64, int32, int64, uint8
16–19: `B`: bool
`T`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bfloat16, bool | `tensor/where.cc` | | `Attention` | `com.microsoft` | 1–1: `T`: float32
`M`: int32 | `llm/attention.cc` | | `FusedGemm` | `com.microsoft` | 1–1: `T`: float16, float32, float64 | `math/fused_gemm.cc` | diff --git a/musa/ep/src/dtype_contracts.json b/musa/ep/src/dtype_contracts.json index c032d0f..467dd1f 100644 --- a/musa/ep/src/dtype_contracts.json +++ b/musa/ep/src/dtype_contracts.json @@ -32,29 +32,31 @@ { "op": "Reshape", "domain": "ai.onnx", - "opset": "13-19", + "opset": "5-21", "storage": "bfloat16,float16,float32,float64,integer,bool", "output": "same as data input", "compute": "alias/copy", "accumulation": "none", - "path": "device alias/copy", + "path": "device alias/copy; opset 21 only added element types", "verified_dtypes": "float32,float16,float64,bfloat16,uint8,uint16,uint32,uint64,int8,int16,int32,int64,bool", "tests": [ - "test/ops/test_reshape.py" + "test/ops/test_reshape.py", + "test/ops/test_opset_semantics.py" ] }, { "op": "Transpose", "domain": "ai.onnx", - "opset": "13-19", + "opset": "1-21", "storage": "bfloat16,float16,float32,float64,integer,bool", "output": "same as input", "compute": "copy", "accumulation": "none", - "path": "muDNN transpose or device kernel", + "path": "muDNN transpose or device kernel; opset 21 only added element types", "verified_dtypes": "float32,float16,float64,bfloat16,uint8,uint16,uint32,uint64,int8,int16,int32,int64,bool", "tests": [ - "test/ops/test_transpose.py" + "test/ops/test_transpose.py", + "test/ops/test_opset_semantics.py" ] }, { @@ -148,15 +150,16 @@ { "op": "Gemm", "domain": "ai.onnx", - "opset": "13-19", + "opset": "7-19", "storage": "bfloat16,float16,float32,float64", "output": "same as inputs", "compute": "library tensor mode", "accumulation": "FP32 for BF16/FP16 library path", - "path": "muDNN MatMul+bias or muBLAS+device post kernel", + "path": "muDNN MatMul+bias or muBLAS+device post kernel; the pre-13 version differences are type sets and the optional C input only", "verified_dtypes": "float32,float16,float64,bfloat16", "tests": [ - "test/ops/test_gemm.py" + "test/ops/test_gemm.py", + "test/ops/test_opset_semantics.py" ] }, { @@ -190,15 +193,16 @@ { "op": "Softmax", "domain": "ai.onnx", - "opset": "13-17", + "opset": "1-17", "storage": "bfloat16,float16,float32,float64", "output": "same as input", "compute": "FP32 intermediates for BF16/FP16 device path", "accumulation": "FP32 for BF16/FP16 device path", - "path": "muDNN softmax or stable device kernel", + "path": "muDNN softmax or stable device kernel; through opset 12 the input is flattened at `axis` first (axis defaults to 1), from 13 on only `axis` is reduced", "verified_dtypes": "float32,float16,float64,bfloat16", "tests": [ - "test/ops/test_softmax.py" + "test/ops/test_softmax.py", + "test/ops/test_opset_semantics.py" ] }, { @@ -430,15 +434,16 @@ { "op": "ReduceMax", "domain": "ai.onnx", - "opset": "13-19", + "opset": "1-20", "storage": "uint8,int8,float16,float32,float64,bfloat16,int32,int64", "output": "same as input", "compute": "device reduction", "accumulation": "FP32 for BF16/FP16; storage-compatible accumulator otherwise", - "path": "MUSA reduction kernel", + "path": "MUSA reduction kernel; opset 20 only added the bool element type", "verified_dtypes": "uint8,int8,float32,float16,float64,bfloat16,int32,int64", "tests": [ - "test/ops/test_reduce_max.py" + "test/ops/test_reduce_max.py", + "test/ops/test_opset_semantics.py" ] }, { @@ -491,30 +496,32 @@ { "op": "Slice", "domain": "ai.onnx", - "opset": "13-19", + "opset": "1-19", "storage": "data: all fixed-size numeric and bool; starts,ends,axes,steps: int32,int64 CPU metadata", "output": "same as data input", "compute": "bytewise strided copy", "accumulation": "none", - "path": "MUSA slice kernel, device 2D copy, or device raw copy", + "path": "MUSA slice kernel, device 2D copy, or device raw copy; the opset 1-9 starts/ends/axes are attributes (step is fixed to 1 by those schemas)", "verified_dtypes": "data uint8,uint16,uint32,uint64,int8,int16,int32,int64,float16,float32,float64,bfloat16,bool; metadata int32,int64", "tests": [ - "test/ops/test_slice.py" + "test/ops/test_slice.py", + "test/ops/test_opset_semantics.py" ], "unsupported_reasons": "negative or zero steps are explicitly rejected; generic device path rank is limited by kMusaMaxBroadcastRank; string and complex data are not registered" }, { "op": "Split", "domain": "ai.onnx", - "opset": "13-19", + "opset": "2-19", "storage": "data: all fixed-size numeric and bool; split sizes: int64 CPU metadata", "output": "each output same as data input", "compute": "bytewise partition copy", "accumulation": "none", - "path": "MUSA split copy kernels; host-built device copy descriptors for irregular partitions", + "path": "MUSA split copy kernels; host-built device copy descriptors for irregular partitions; the opset 2-12 sizes are the `split` attribute (the opset-1 tensor-input form is not registered)", "verified_dtypes": "data float32,float16,float64,bfloat16,int8,int16,int32,int64,uint8,uint16,uint32,uint64,bool; split sizes int64", "tests": [ - "test/ops/test_split.py" + "test/ops/test_split.py", + "test/ops/test_opset_semantics.py" ], "unsupported_reasons": "split sizes must be non-negative and sum to the selected axis extent; string and complex data are not registered" }, @@ -536,15 +543,16 @@ { "op": "Pad", "domain": "ai.onnx", - "opset": "13-19", + "opset": "2-19", "storage": "data and constant_value: same all fixed-size numeric or bool dtype; pads,axes: int64 CPU metadata", "output": "same as data input", "compute": "bytewise constant fill or copy", "accumulation": "none", - "path": "MUSA constant-pad kernel", + "path": "MUSA constant-pad kernel; the opset 2-10 pads/value are attributes and the kernel converts the FLOAT value to the element type", "verified_dtypes": "data uint8,uint16,uint32,uint64,int8,int16,int32,int64,float16,float32,float64,bfloat16,bool; pads,axes int64", "tests": [ - "test/ops/test_pad.py" + "test/ops/test_pad.py", + "test/ops/test_opset_semantics.py" ], "unsupported_reasons": "only constant mode and non-negative pads are supported; rank is limited by kMusaMaxBroadcastRank; string and complex data are not registered" }, @@ -710,16 +718,17 @@ { "op": "Clip", "domain": "ai.onnx", - "opset": "13-19", + "opset": "6-19", "storage": "bfloat16,float16,float32,float64,uint8,uint16,uint32,uint64,int8,int16,int32,int64", "output": "same as input", "compute": "FP32 for BF16/FP16; storage dtype otherwise", "accumulation": "none", - "path": "MUSA Clip device kernel; optional scalar bounds are staged to device", + "path": "MUSA Clip device kernel; optional scalar bounds are staged to device; the opset 6-10 bounds are FLOAT attributes the kernel converts to the element type", "verified_dtypes": "float32,float16,float64,bfloat16,int8,int16,int32,int64,uint8,uint16,uint32,uint64", "tests": [ "test/ops/test_clip.py", - "test/ops/test_unary_dtype_matrix.py" + "test/ops/test_unary_dtype_matrix.py", + "test/ops/test_opset_semantics.py" ], "unsupported_reasons": "bounds must be scalar and match input dtype; bool,string,complex are excluded by ONNX Clip" }, @@ -831,45 +840,48 @@ { "op": "Squeeze", "domain": "ai.onnx", - "opset": "13-19", + "opset": "1-21", "storage": "all fixed-size numeric and bool; axes int64 CPU metadata", "output": "same storage as data input", "compute": "alias or device byte copy", "accumulation": "none", - "path": "MUSA device alias/copy", + "path": "MUSA device alias/copy; through opset 12 `axes` is an attribute; opset 21 only added element types", "verified_dtypes": "float32,int64,uint8,bool,float16,float64,int8,int16,uint16,uint32,uint64,bfloat16", "tests": [ - "test/ops/test_squeeze.py" + "test/ops/test_squeeze.py", + "test/ops/test_opset_semantics.py" ], "unsupported_reasons": "string and complex data are not registered; axes input must be int64 and selected dimensions must be singleton" }, { "op": "Unsqueeze", "domain": "ai.onnx", - "opset": "13-19", + "opset": "1-21", "storage": "all fixed-size numeric and bool; axes int64 CPU metadata", "output": "same storage as data input", "compute": "alias or device byte copy", "accumulation": "none", - "path": "MUSA device alias/copy", + "path": "MUSA device alias/copy; through opset 12 `axes` is an attribute; opset 21 only added element types", "verified_dtypes": "float32,float16,int32,bool,float64,int8,int16,int64,uint8,uint16,uint32,uint64,bfloat16", "tests": [ - "test/ops/test_unsqueeze.py" + "test/ops/test_unsqueeze.py", + "test/ops/test_opset_semantics.py" ], "unsupported_reasons": "string and complex data are not registered; axes input must be int64 and axes must be unique/in range" }, { "op": "Shape", "domain": "ai.onnx", - "opset": "13-19", + "opset": "1-21", "storage": "input all fixed-size numeric and bool; output int64", "output": "int64 rank-length vector", "compute": "host shape metadata materialization", "accumulation": "none", - "path": "CPU metadata output by explicit KernelDef memory contract", + "path": "CPU metadata output by explicit KernelDef memory contract; opset 21 only added element types", "verified_dtypes": "input float32,float16,float64,bfloat16,uint8,uint16,uint32,uint64,int8,int16,int32,int64,bool; output int64", "tests": [ - "test/ops/test_shape.py" + "test/ops/test_shape.py", + "test/ops/test_opset_semantics.py" ], "unsupported_reasons": "Shape consumes tensor metadata only; string and complex inputs are not registered" }, diff --git a/musa/ep/src/graph/capability_predicates.cc b/musa/ep/src/graph/capability_predicates.cc index 4526f83..dc8abd1 100644 --- a/musa/ep/src/graph/capability_predicates.cc +++ b/musa/ep/src/graph/capability_predicates.cc @@ -13,6 +13,7 @@ #include "graph/capability_predicates.h" +#include #include #include #include @@ -126,12 +127,40 @@ bool ResizeConfigurationSupported(Ort::ConstNode node) { return ResizeRouteValueRejection(config, values) == nullptr; } +/// True when the MUSA Slice kernel can run this node. +/// +/// The device copy walks every slice with a positive step, so the kernel +/// rejects a non-positive `steps` entry - leaving those nodes on the CPU EP is +/// better than a node that fails during Compute(). `steps` is the fourth input +/// of the opset-10 schema and does not exist before opset 10 (the older +/// schemas have no steps at all). A `steps` tensor that is not a constant +/// initializer is a run-time value the kernel has to judge, exactly like the +/// Resize values above. +bool SliceConfigurationSupported(Ort::ConstNode node) { + const std::vector inputs = node.GetInputs(); + constexpr size_t kStepsInputIndex = 4; + if (inputs.size() <= kStepsInputIndex || + inputs[kStepsInputIndex] == nullptr) { + return true; + } + const std::optional> steps = + ReadIntInitializerNoLimit(inputs[kStepsInputIndex]); + if (!steps.has_value()) { + return true; + } + return std::all_of(steps->begin(), steps->end(), + [](int64_t step) { return step > 0; }); +} + } // namespace bool IsNodeConfigurationSupported(Ort::ConstNode node) { if (IsOnnxOp(node, "Resize")) { return ResizeConfigurationSupported(node); } + if (IsOnnxOp(node, "Slice")) { + return SliceConfigurationSupported(node); + } return true; } diff --git a/musa/ep/src/kernels/math/clip.cc b/musa/ep/src/kernels/math/clip.cc index ea7d378..4954c81 100644 --- a/musa/ep/src/kernels/math/clip.cc +++ b/musa/ep/src/kernels/math/clip.cc @@ -35,8 +35,27 @@ OrtStatus* ValidateScalarInput(Ort::ConstValue value, class Clip : public OpKernelBase { public: - Clip(const OrtKernelInfo* /*info*/, void* /*state*/) {} + Clip(const OrtKernelInfo* info, void* /*state*/) { + Ort::ConstKernelInfo kernel_info(info); + // Through opset 10 the bounds are FLOAT attributes of the node and an + // absent one means "no bound"; the opset-11 schema replaced them with + // optional scalar inputs that match the element type of X. + const std::optional min_attribute = + AttrOrNullopt(kernel_info, "min"); + const std::optional max_attribute = + AttrOrNullopt(kernel_info, "max"); + has_min_attribute_ = min_attribute.has_value(); + has_max_attribute_ = max_attribute.has_value(); + min_attribute_ = min_attribute.value_or(0.0f); + max_attribute_ = max_attribute.value_or(0.0f); + } OrtStatus* Compute(Ort::KernelContext& ctx) const; + + private: + bool has_min_attribute_ = false; + bool has_max_attribute_ = false; + float min_attribute_ = 0.0f; + float max_attribute_ = 0.0f; }; OrtStatus* Clip::Compute(Ort::KernelContext& ctx) const { @@ -53,6 +72,10 @@ OrtStatus* Clip::Compute(Ort::KernelContext& ctx) const { MusaClipParams params{}; params.count = NumElements(shape); + params.min_is_attribute = has_min_attribute_ ? 1 : 0; + params.max_is_attribute = has_max_attribute_ ? 1 : 0; + params.min_attribute = min_attribute_; + params.max_attribute = max_attribute_; DeviceInputBuffer min_buffer; DeviceInputBuffer max_buffer; if (ctx.GetInputCount() > 1 && ctx.GetInput(1) != nullptr) { @@ -60,12 +83,16 @@ OrtStatus* Clip::Compute(Ort::KernelContext& ctx) const { RETURN_IF_ERROR(min_buffer.Bind(ctx.GetInput(1), stream)); params.has_min = 1; params.min_data = min_buffer.data(); + } else if (has_min_attribute_) { + params.has_min = 1; } if (ctx.GetInputCount() > 2 && ctx.GetInput(2) != nullptr) { RETURN_IF_ERROR(ValidateScalarInput(ctx.GetInput(2), elem_type, "max")); RETURN_IF_ERROR(max_buffer.Bind(ctx.GetInput(2), stream)); params.has_max = 1; params.max_data = max_buffer.data(); + } else if (has_max_attribute_) { + params.has_max = 1; } if (!IsGpuMemory(input.GetTensorMemoryInfo())) { @@ -92,3 +119,20 @@ OrtStatus* Clip::Compute(Ort::KernelContext& ctx) const { ONNX_OPERATOR_VERSIONED_KERNEL_EX( Clip, kOnnxDomain, 13, 19, (Ort::KernelDefBuilder().AddTypeConstraint("T", ClipTensorTypes())), Clip) + +// opset 6-10 keeps the bounds in FLOAT attributes; the old schemas only allow +// the floating-point element types. opset 11 moved the bounds to inputs and +// opset 12 added the integer types, both without bfloat16. +ONNX_OPERATOR_VERSIONED_KERNEL_EX( + Clip, kOnnxDomain, 6, 10, + (Ort::KernelDefBuilder().AddTypeConstraint("T", HfdTensorTypes())), Clip) + +ONNX_OPERATOR_VERSIONED_KERNEL_EX( + Clip, kOnnxDomain, 11, 11, + (Ort::KernelDefBuilder().AddTypeConstraint("T", HfdTensorTypes())), Clip) + +ONNX_OPERATOR_VERSIONED_KERNEL_EX( + Clip, kOnnxDomain, 12, 12, + (Ort::KernelDefBuilder().AddTypeConstraint( + "T", AllFixedSizeTensorTypesNoBoolNoBFloat16())), + Clip) diff --git a/musa/ep/src/kernels/math/clip_impl.mu b/musa/ep/src/kernels/math/clip_impl.mu index 4773ca3..95ef886 100644 --- a/musa/ep/src/kernels/math/clip_impl.mu +++ b/musa/ep/src/kernels/math/clip_impl.mu @@ -17,12 +17,18 @@ __device__ __forceinline__ T ClipDefaultMax() { template __global__ void ClipKernel(const T* input, T* output, MusaClipParams params) { - const T min_value = params.has_min - ? reinterpret_cast(params.min_data)[0] - : ClipDefaultMin(); - const T max_value = params.has_max - ? reinterpret_cast(params.max_data)[0] - : ClipDefaultMax(); + // Through opset 10 the bounds come from the node's FLOAT attributes instead + // of scalar inputs; params.*_is_attribute selects which side is used. + const T min_value = + params.has_min ? (params.min_is_attribute + ? static_cast(params.min_attribute) + : reinterpret_cast(params.min_data)[0]) + : ClipDefaultMin(); + const T max_value = + params.has_max ? (params.max_is_attribute + ? static_cast(params.max_attribute) + : reinterpret_cast(params.max_data)[0]) + : ClipDefaultMax(); const int64_t thread_id = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; const int64_t total_threads = static_cast(gridDim.x) * blockDim.x; @@ -44,11 +50,17 @@ __global__ void ClipFloatLikeKernel(const T* input, T* output, MusaClipParams params) { const float min_value = params.has_min - ? MusaScalarToFloat(reinterpret_cast(params.min_data)[0]) + ? (params.min_is_attribute + ? params.min_attribute + : MusaScalarToFloat( + reinterpret_cast(params.min_data)[0])) : -INFINITY; const float max_value = params.has_max - ? MusaScalarToFloat(reinterpret_cast(params.max_data)[0]) + ? (params.max_is_attribute + ? params.max_attribute + : MusaScalarToFloat( + reinterpret_cast(params.max_data)[0])) : INFINITY; const int64_t thread_id = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; diff --git a/musa/ep/src/kernels/math/gemm.cc b/musa/ep/src/kernels/math/gemm.cc index 2e0e0b6..f6d01f1 100644 --- a/musa/ep/src/kernels/math/gemm.cc +++ b/musa/ep/src/kernels/math/gemm.cc @@ -42,3 +42,22 @@ ONNX_OPERATOR_VERSIONED_KERNEL_EX( Gemm, kOnnxDomain, 13, 19, (Ort::KernelDefBuilder().AddTypeConstraint("T", FloatLikeTensorTypes())), Gemm) + +// The Gemm version differences below opset 13 are type-set and optional-input +// changes, not arithmetic ones: opset 7 dropped the `broadcast` attribute, +// opset 9 added the integer element types (which this kernel does not +// implement, so only the floating-point types are declared) and opset 11 made +// the C input optional (the kernel already treats a missing third input as no +// bias). The ranges follow the CUDA EP's splits; the schemas have no bfloat16 +// before opset 13. +ONNX_OPERATOR_VERSIONED_KERNEL_EX( + Gemm, kOnnxDomain, 7, 8, + (Ort::KernelDefBuilder().AddTypeConstraint("T", HfdTensorTypes())), Gemm) + +ONNX_OPERATOR_VERSIONED_KERNEL_EX( + Gemm, kOnnxDomain, 9, 10, + (Ort::KernelDefBuilder().AddTypeConstraint("T", HfdTensorTypes())), Gemm) + +ONNX_OPERATOR_VERSIONED_KERNEL_EX( + Gemm, kOnnxDomain, 11, 12, + (Ort::KernelDefBuilder().AddTypeConstraint("T", HfdTensorTypes())), Gemm) diff --git a/musa/ep/src/kernels/math/softmax.cc b/musa/ep/src/kernels/math/softmax.cc index 431b572..2f44339 100644 --- a/musa/ep/src/kernels/math/softmax.cc +++ b/musa/ep/src/kernels/math/softmax.cc @@ -64,12 +64,21 @@ class Softmax : public OpKernelBase { public: Softmax(const OrtKernelInfo* info, void* /*state*/) { Ort::ConstKernelInfo kernel_info(info); - axis_ = AttrOrDefault(kernel_info, "axis", 0); + // The axis semantics changed at opset 13: through opset 12 the input is + // coerced into the 2-D matrix [prod(d0..d_{axis-1}), prod(d_axis..)] and + // normalised along its second dimension, with axis defaulting to 1; from + // opset 13 on the normalisation runs along `axis` alone and the default + // is -1. ORT resolves the node against the ONNX schema of the model's + // opset, so the node's since version tells the two apart. + since_version_ = kernel_info.GetOperatorSinceVersion(); + axis_ = AttrOrDefault(kernel_info, "axis", + since_version_ >= 13 ? -1 : 1); } OrtStatus* Compute(Ort::KernelContext& ctx) const; private: int64_t axis_ = 0; + int64_t since_version_ = 0; }; OrtStatus* Softmax::Compute(Ort::KernelContext& ctx) const { @@ -82,26 +91,56 @@ OrtStatus* Softmax::Compute(Ort::KernelContext& ctx) const { return UnsupportedDeviceElementwiseStatus("Softmax", elem_type); } std::vector shape0 = info.GetShape(); + const bool flatten_trailing_dims = since_version_ < 13; int64_t axis = 0; int64_t outer = 1; int64_t dim = 1; int64_t inner = 1; + // The pre-13 reading of `axis` is the opset-13 one for every node whose + // trailing dimensions are all one, so most old-opset nodes keep the plain + // shape/axis pair the muDNN softmax is called with. + bool same_as_axis_reduction = !flatten_trailing_dims; if (!shape0.empty()) { + const int64_t rank = static_cast(shape0.size()); axis = NormalizeAxis(axis_, shape0.size()); - if (axis < 0 || axis >= static_cast(shape0.size())) { + // For the flattened form `axis == rank` is meaningful (it reduces a single + // element, which is what the promoted 2-D matrix [prod(d0..d_{rank-1}), 1] + // holds), and it is what a 1-D input with the opset-12 default axis hits. + const int64_t highest_axis = flatten_trailing_dims ? rank : rank - 1; + if (axis < 0 || axis > highest_axis) { return Ort::GetApi().CreateStatus(ORT_INVALID_ARGUMENT, "Softmax axis out of range"); } - dim = shape0[static_cast(axis)]; for (int64_t i = 0; i < axis; ++i) { outer *= shape0[static_cast(i)]; } - for (size_t i = static_cast(axis) + 1; i < shape0.size(); ++i) { - inner *= shape0[i]; + if (flatten_trailing_dims) { + for (int64_t i = axis; i < rank; ++i) { + dim *= shape0[static_cast(i)]; + } + same_as_axis_reduction = axis < rank && + shape0[static_cast(axis)] > 0 && + dim == shape0[static_cast(axis)]; + } else { + dim = shape0[static_cast(axis)]; + for (int64_t i = axis + 1; i < rank; ++i) { + inner *= shape0[static_cast(i)]; + } } } + // muDNN gets the contiguous view whose last dimension is the one to reduce: + // the node's own shape when `axis` already selects it, and the promoted + // [prod(d0..d_{axis-1}), prod(d_axis..)] matrix of the pre-13 schemas + // otherwise. Both describe exactly the reduction the node asks for, and a + // muDNN rejection falls through to the device kernel below. + std::vector reduction_shape = shape0; + int64_t reduction_axis = axis; + if (!same_as_axis_reduction) { + reduction_shape = {outer, dim}; + reduction_axis = 1; + } Ort::UnownedValue y = ctx.GetOutput(0, shape0); - if (TryMudnnSoftmax(ctx, shape0, axis, elem_type, y)) { + if (TryMudnnSoftmax(ctx, reduction_shape, reduction_axis, elem_type, y)) { return nullptr; } @@ -119,3 +158,14 @@ ONNX_OPERATOR_VERSIONED_KERNEL_EX( Softmax, kOnnxDomain, 13, 17, (Ort::KernelDefBuilder().AddTypeConstraint("T", FloatLikeTensorTypes())), Softmax) + +// The pre-13 schemas do not have bfloat16; opset 11-12 only clarified that a +// negative axis counts from the back, so one class serves both old ranges and +// picks the semantics from the node's since version. +ONNX_OPERATOR_VERSIONED_KERNEL_EX( + Softmax, kOnnxDomain, 1, 10, + (Ort::KernelDefBuilder().AddTypeConstraint("T", HfdTensorTypes())), Softmax) + +ONNX_OPERATOR_VERSIONED_KERNEL_EX( + Softmax, kOnnxDomain, 11, 12, + (Ort::KernelDefBuilder().AddTypeConstraint("T", HfdTensorTypes())), Softmax) diff --git a/musa/ep/src/kernels/reduction/reduce_max.cc b/musa/ep/src/kernels/reduction/reduce_max.cc index 45506c5..8e3e797 100644 --- a/musa/ep/src/kernels/reduction/reduce_max.cc +++ b/musa/ep/src/kernels/reduction/reduce_max.cc @@ -42,8 +42,11 @@ ONNX_OPERATOR_VERSIONED_KERNEL_EX( .SetInputMemType(1, OrtMemTypeCPUInput)), ReduceMax) +// opset 20 only added the bool element type, which this kernel does not +// implement (and does not declare), so the opset-13 behaviour still covers the +// node. ONNX_OPERATOR_VERSIONED_KERNEL_EX( - ReduceMax, kOnnxDomain, 13, 19, + ReduceMax, kOnnxDomain, 13, 20, (Ort::KernelDefBuilder() .AddTypeConstraint("T", ReduceMaxTensorTypes()) .SetInputMemType(1, OrtMemTypeCPUInput)), diff --git a/musa/ep/src/kernels/shared_inc/device_kernel_types.h b/musa/ep/src/kernels/shared_inc/device_kernel_types.h index 764e894..2138dcc 100644 --- a/musa/ep/src/kernels/shared_inc/device_kernel_types.h +++ b/musa/ep/src/kernels/shared_inc/device_kernel_types.h @@ -118,6 +118,14 @@ struct MusaClipParams { const void* max_data; int32_t has_min; int32_t has_max; + // Through opset 10 the bounds live in the node's FLOAT attributes instead of + // scalar inputs. When *_is_attribute is set the matching data pointer is + // unused and the bound is the attribute value, converted to the element type + // the same way the ONNX reference kernels do. + int32_t min_is_attribute; + int32_t max_is_attribute; + float min_attribute; + float max_attribute; }; struct MusaLayerNormParams { @@ -220,6 +228,12 @@ struct MusaPadParams { int64_t input_strides[kMusaMaxBroadcastRank]; int64_t output_dims[kMusaMaxBroadcastRank]; int64_t pads_begin[kMusaMaxBroadcastRank]; + // Pad through opset 10 keeps the pad value in a FLOAT attribute instead of a + // third input; when set, attribute_value replaces the constant_value bit + // pattern of the kernel argument. Those schemas only allow + // float16/float/float64, so the element width selects the conversion. + int32_t value_from_attribute; + float attribute_value; }; struct MusaNonZeroParams { diff --git a/musa/ep/src/kernels/shared_inc/kernel_attributes.h b/musa/ep/src/kernels/shared_inc/kernel_attributes.h index 02bf76f..405c9e9 100644 --- a/musa/ep/src/kernels/shared_inc/kernel_attributes.h +++ b/musa/ep/src/kernels/shared_inc/kernel_attributes.h @@ -14,6 +14,7 @@ #pragma once #include +#include #include #include "utils.h" @@ -35,3 +36,19 @@ inline std::vector AttrsOrEmpty(Ort::ConstKernelInfo& info, return {}; } } + +/// Reads an attribute that holds a single value. +/// +/// AttrsOrEmpty() above goes through the ORT array getters, which only accept +/// the list attribute types (FLOATS/INTS/STRINGS) - a scalar FLOAT/INT/STRING +/// attribute fails there and would silently read as "absent". Use this for +/// those; a missing attribute is reported by ORT as an error, which is what +/// turns into the empty optional here. +template +std::optional AttrOrNullopt(Ort::ConstKernelInfo& info, const char* name) { + try { + return info.GetAttribute(name); + } catch (...) { + return std::nullopt; + } +} diff --git a/musa/ep/src/kernels/shared_inc/kernel_type_constraints.h b/musa/ep/src/kernels/shared_inc/kernel_type_constraints.h index 193b784..891a5bf 100644 --- a/musa/ep/src/kernels/shared_inc/kernel_type_constraints.h +++ b/musa/ep/src/kernels/shared_inc/kernel_type_constraints.h @@ -225,6 +225,25 @@ inline std::vector ClipTensorTypes() { }; } +// The opset-11/12 Clip and Pad schemas: every fixed-size numeric type but bool +// (which Pad only gains at opset 13) and bfloat16 (opset 13). +inline std::vector +AllFixedSizeTensorTypesNoBoolNoBFloat16() { + return { + GetTensorType(ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT8), + GetTensorType(ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT16), + GetTensorType(ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT32), + GetTensorType(ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT64), + GetTensorType(ONNX_TENSOR_ELEMENT_DATA_TYPE_INT8), + GetTensorType(ONNX_TENSOR_ELEMENT_DATA_TYPE_INT16), + GetTensorType(ONNX_TENSOR_ELEMENT_DATA_TYPE_INT32), + GetTensorType(ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64), + GetTensorType(ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT16), + GetTensorType(ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT), + GetTensorType(ONNX_TENSOR_ELEMENT_DATA_TYPE_DOUBLE), + }; +} + inline std::vector PReluOpset16TensorTypes() { return { GetTensorType(ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT16), diff --git a/musa/ep/src/kernels/tensor/pad.cc b/musa/ep/src/kernels/tensor/pad.cc index bca3893..6245395 100644 --- a/musa/ep/src/kernels/tensor/pad.cc +++ b/musa/ep/src/kernels/tensor/pad.cc @@ -55,11 +55,18 @@ class Pad : public OpKernelBase { Pad(const OrtKernelInfo* info, void* /*state*/) { Ort::ConstKernelInfo kernel_info(info); mode_ = AttrOrDefault(kernel_info, "mode", "constant"); + // Through opset 10 the pads and the pad value are node attributes; the + // kernel reads them as well and only prefers a tensor input when the + // schema provides one. + pads_attr_ = AttrsOrEmpty(kernel_info, "pads"); + value_attr_ = AttrOrNullopt(kernel_info, "value"); } OrtStatus* Compute(Ort::KernelContext& ctx) const; private: std::string mode_ = "constant"; + std::vector pads_attr_; + std::optional value_attr_; }; OrtStatus* Pad::Compute(Ort::KernelContext& ctx) const { @@ -84,7 +91,23 @@ OrtStatus* Pad::Compute(Ort::KernelContext& ctx) const { } musaStream_t stream = GetComputeStream(ctx); - std::vector pads = ReadIntTensor(ctx, 1); + std::vector pads; + int32_t value_from_attribute = 0; + float attribute_value = 0.0f; + if (ctx.GetInputCount() > 1 && ctx.GetInput(1) != nullptr) { + pads = ReadIntTensor(ctx, 1); + } else { + // opset 2-10: `pads` and `value` are attributes. The attribute form has no + // third input, so the value always comes from the attribute (0 when it is + // absent, which is the schema default). + pads = pads_attr_; + value_from_attribute = 1; + attribute_value = value_attr_.value_or(0.0f); + } + if (pads.empty()) { + return Ort::GetApi().CreateStatus(ORT_INVALID_ARGUMENT, + "Pad requires pads"); + } std::vector axes; if (ctx.GetInputCount() > 3 && ctx.GetInput(3) != nullptr) { axes = ReadIntTensor(ctx, 3); @@ -126,11 +149,13 @@ OrtStatus* Pad::Compute(Ort::KernelContext& ctx) const { "Pad requires MUSA output"); } + MusaPadParams params = MakePadParams(input_shape, output_shape, pads_begin); + params.value_from_attribute = value_from_attribute; + params.attribute_value = attribute_value; return LaunchStatus(LaunchMusaPadKernel( input.GetTensorRawData(), output.GetTensorMutableRawData(), ScalarInputOrZero(ctx, 2, elem_size, stream), - static_cast(elem_size), - MakePadParams(input_shape, output_shape, pads_begin), stream)); + static_cast(elem_size), params, stream)); } } // namespace @@ -142,3 +167,20 @@ ONNX_OPERATOR_VERSIONED_KERNEL_EX( .SetInputMemType(2, OrtMemTypeCPUInput) .SetInputMemType(3, OrtMemTypeCPUInput)), Pad) + +// opset 11-12 moved pads and the constant value to inputs; both are +// non-optional, so the same kernel def as opset 13+ applies. The schemas do +// not have bool (opset 13) or bfloat16 (opset 13) yet. +ONNX_OPERATOR_VERSIONED_KERNEL_EX( + Pad, kOnnxDomain, 11, 12, + (Ort::KernelDefBuilder() + .AddTypeConstraint("T", AllFixedSizeTensorTypesNoBoolNoBFloat16()) + .SetInputMemType(1, OrtMemTypeCPUInput) + .SetInputMemType(2, OrtMemTypeCPUInput)), + Pad) + +// opset 2-10 keeps `pads` and `value` in node attributes and only allows the +// floating-point element types, which the attribute path in the kernel covers. +ONNX_OPERATOR_VERSIONED_KERNEL_EX( + Pad, kOnnxDomain, 2, 10, + (Ort::KernelDefBuilder().AddTypeConstraint("T", HfdTensorTypes())), Pad) diff --git a/musa/ep/src/kernels/tensor/pad_impl.mu b/musa/ep/src/kernels/tensor/pad_impl.mu index 88732a7..1ff4fd0 100644 --- a/musa/ep/src/kernels/tensor/pad_impl.mu +++ b/musa/ep/src/kernels/tensor/pad_impl.mu @@ -3,6 +3,35 @@ namespace { +// Pad through opset 10 supplies the pad value as a FLOAT node attribute. The +// schemas of those opsets only allow float16/float/float64 elements, so the +// element width picks the conversion, and the bits are copied byte-wise to +// stay independent of the host layout of the union. +__device__ __forceinline__ uint64_t PadAttributeValueBits(float value, + int32_t element_size) { + uint64_t bits = 0; + if (element_size == 2) { + const __half converted = __float2half_rn(value); + for (int32_t byte = 0; byte < 2; ++byte) { + reinterpret_cast(&bits)[byte] = + reinterpret_cast(&converted)[byte]; + } + } else if (element_size == 4) { + const float converted = value; + for (int32_t byte = 0; byte < 4; ++byte) { + reinterpret_cast(&bits)[byte] = + reinterpret_cast(&converted)[byte]; + } + } else if (element_size == 8) { + const double converted = static_cast(value); + for (int32_t byte = 0; byte < 8; ++byte) { + reinterpret_cast(&bits)[byte] = + reinterpret_cast(&converted)[byte]; + } + } + return bits; +} + __device__ __forceinline__ void WriteConstant(void* output, int64_t output_index, uint64_t constant_value, @@ -62,6 +91,9 @@ __global__ void PadKernel(const void* input, uint64_t constant_value, int32_t element_size, MusaPadParams params) { + if (params.value_from_attribute) { + constant_value = PadAttributeValueBits(params.attribute_value, element_size); + } const int64_t thread_id = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; const int64_t total_threads = diff --git a/musa/ep/src/kernels/tensor/reshape.cc b/musa/ep/src/kernels/tensor/reshape.cc index 7a68dee..6dd3211 100644 --- a/musa/ep/src/kernels/tensor/reshape.cc +++ b/musa/ep/src/kernels/tensor/reshape.cc @@ -113,8 +113,11 @@ ONNX_OPERATOR_VERSIONED_KERNEL_EX( .SetInputMemType(1, OrtMemTypeCPUInput)), Reshape) +// The opset-21 schema only adds the int4/uint4 element types, which this +// kernel does not implement (and does not declare), so the opset-13 behaviour +// still covers the node. ONNX_OPERATOR_VERSIONED_KERNEL_EX( - Reshape, kOnnxDomain, 13, 19, + Reshape, kOnnxDomain, 13, 21, (Ort::KernelDefBuilder() .AddTypeConstraint("T", AllFixedSizeTensorTypes()) .SetInputMemType(1, OrtMemTypeCPUInput)), diff --git a/musa/ep/src/kernels/tensor/shape.cc b/musa/ep/src/kernels/tensor/shape.cc index 3f3379c..69ddf69 100644 --- a/musa/ep/src/kernels/tensor/shape.cc +++ b/musa/ep/src/kernels/tensor/shape.cc @@ -37,8 +37,10 @@ ONNX_OPERATOR_VERSIONED_KERNEL_EX( .SetOutputMemType(0, OrtMemTypeCPUInput)), Shape) +// opset 21 only added the int4/uint4 element types on the input, which the +// kernel does not implement, so the opset-13 behaviour still covers it. ONNX_OPERATOR_VERSIONED_KERNEL_EX( - Shape, kOnnxDomain, 13, 19, + Shape, kOnnxDomain, 13, 21, (Ort::KernelDefBuilder() .AddTypeConstraint("T", AllFixedSizeTensorTypes()) .AddTypeConstraint("T1", diff --git a/musa/ep/src/kernels/tensor/slice.cc b/musa/ep/src/kernels/tensor/slice.cc index 2dc322e..fae9a75 100644 --- a/musa/ep/src/kernels/tensor/slice.cc +++ b/musa/ep/src/kernels/tensor/slice.cc @@ -17,8 +17,21 @@ namespace { class Slice : public OpKernelBase { public: - Slice(const OrtKernelInfo* /*info*/, void* /*state*/) {} + Slice(const OrtKernelInfo* info, void* /*state*/) { + Ort::ConstKernelInfo kernel_info(info); + // Through opset 9 starts/ends/axes are node attributes and there is no + // `steps` (those schemas fix it to 1); from opset 10 on they are inputs. + // The kernel reads both and prefers the inputs when the schema has them. + starts_attr_ = AttrsOrEmpty(kernel_info, "starts"); + ends_attr_ = AttrsOrEmpty(kernel_info, "ends"); + axes_attr_ = AttrsOrEmpty(kernel_info, "axes"); + } OrtStatus* Compute(Ort::KernelContext& ctx) const; + + private: + std::vector starts_attr_; + std::vector ends_attr_; + std::vector axes_attr_; }; OrtStatus* Slice::Compute(Ort::KernelContext& ctx) const { @@ -26,17 +39,30 @@ OrtStatus* Slice::Compute(Ort::KernelContext& ctx) const { auto in0_info = input0.GetTensorTypeAndShapeInfo(); auto elem_type = in0_info.GetElementType(); auto shape0 = in0_info.GetShape(); - std::vector parameter_indices{1, 2}; - if (ctx.GetInputCount() > 3) parameter_indices.push_back(3); - if (ctx.GetInputCount() > 4) parameter_indices.push_back(4); - std::vector> parameters = - ReadIntTensors(ctx, parameter_indices); - std::vector starts = std::move(parameters[0]); - std::vector ends = std::move(parameters[1]); + std::vector starts; + std::vector ends; std::vector axes; - std::vector steps(starts.size(), 1); - if (ctx.GetInputCount() > 3) axes = std::move(parameters[2]); - if (ctx.GetInputCount() > 4) steps = std::move(parameters[3]); + std::vector steps; + if (ctx.GetInputCount() > 1) { + std::vector parameter_indices{1, 2}; + if (ctx.GetInputCount() > 3) parameter_indices.push_back(3); + if (ctx.GetInputCount() > 4) parameter_indices.push_back(4); + std::vector> parameters = + ReadIntTensors(ctx, parameter_indices); + starts = std::move(parameters[0]); + ends = std::move(parameters[1]); + if (parameter_indices.size() > 2) axes = std::move(parameters[2]); + if (parameter_indices.size() > 3) steps = std::move(parameters[3]); + } else { + starts = starts_attr_; + ends = ends_attr_; + axes = axes_attr_; + } + if (starts.size() != ends.size()) { + return Ort::GetApi().CreateStatus( + ORT_INVALID_ARGUMENT, "Slice starts/ends must have the same length"); + } + if (steps.empty()) steps.assign(starts.size(), 1); if (axes.empty()) { axes.resize(starts.size()); std::iota(axes.begin(), axes.end(), 0); @@ -191,3 +217,37 @@ ONNX_OPERATOR_VERSIONED_KERNEL_EX( .SetInputMemType(3, OrtMemTypeCPUInput) .SetInputMemType(4, OrtMemTypeCPUInput)), Slice) + +// opset 10 moved starts/ends/axes/steps from attributes to inputs; opset 11 +// only clarified negative indices. Both schemas are the opset-13 set without +// bfloat16, and the same kernel def applies to them. +ONNX_OPERATOR_VERSIONED_KERNEL_EX( + Slice, kOnnxDomain, 10, 10, + (Ort::KernelDefBuilder() + .AddTypeConstraint("T", AllFixedSizeTensorTypesNoBFloat16()) + .AddTypeConstraint("Tind", IntTensorTypes()) + .SetInputMemType(1, OrtMemTypeCPUInput) + .SetInputMemType(2, OrtMemTypeCPUInput) + .SetInputMemType(3, OrtMemTypeCPUInput) + .SetInputMemType(4, OrtMemTypeCPUInput)), + Slice) + +ONNX_OPERATOR_VERSIONED_KERNEL_EX( + Slice, kOnnxDomain, 11, 12, + (Ort::KernelDefBuilder() + .AddTypeConstraint("T", AllFixedSizeTensorTypesNoBFloat16()) + .AddTypeConstraint("Tind", IntTensorTypes()) + .SetInputMemType(1, OrtMemTypeCPUInput) + .SetInputMemType(2, OrtMemTypeCPUInput) + .SetInputMemType(3, OrtMemTypeCPUInput) + .SetInputMemType(4, OrtMemTypeCPUInput)), + Slice) + +// opset 1-9 keeps starts/ends/axes in node attributes; those schemas do not +// have bfloat16 either, and they have no `Tind` type parameter (there is no +// indices input), so only `T` may be declared here. +ONNX_OPERATOR_VERSIONED_KERNEL_EX(Slice, kOnnxDomain, 1, 9, + (Ort::KernelDefBuilder().AddTypeConstraint( + "T", + AllFixedSizeTensorTypesNoBFloat16())), + Slice) diff --git a/musa/ep/src/kernels/tensor/split.cc b/musa/ep/src/kernels/tensor/split.cc index f4964e4..de6a5b4 100644 --- a/musa/ep/src/kernels/tensor/split.cc +++ b/musa/ep/src/kernels/tensor/split.cc @@ -25,11 +25,15 @@ class Split : public OpKernelBase { Split(const OrtKernelInfo* info, void* /*state*/) { Ort::ConstKernelInfo kernel_info(info); axis_ = AttrOrDefault(kernel_info, "axis", 0); + // Through opset 12 the per-output sizes live in the `split` attribute; the + // kernel prefers the input when the schema provides one. + split_attr_ = AttrsOrEmpty(kernel_info, "split"); } OrtStatus* Compute(Ort::KernelContext& ctx) const; private: int64_t axis_ = 0; + std::vector split_attr_; }; OrtStatus* Split::Compute(Ort::KernelContext& ctx) const { @@ -39,13 +43,28 @@ OrtStatus* Split::Compute(Ort::KernelContext& ctx) const { auto shape0 = in0_info.GetShape(); int64_t axis = NormalizeAxis(axis_, shape0.size()); std::vector splits; + bool sizes_given = false; if (ctx.GetInputCount() > 1) { splits = ReadIntTensor(ctx, 1); + sizes_given = true; + } else if (!split_attr_.empty()) { + splits = split_attr_; + sizes_given = true; } else { size_t count = ctx.GetOutputCount(); splits.assign( count, shape0[static_cast(axis)] / static_cast(count)); } + // The copies below walk the input row by row, so explicit sizes that do not + // add up to the split axis would read past the input. + if (sizes_given && + (splits.size() != ctx.GetOutputCount() || + std::accumulate(splits.begin(), splits.end(), int64_t{0}) != + shape0[static_cast(axis)])) { + return Ort::GetApi().CreateStatus( + ORT_INVALID_ARGUMENT, + "Split sizes must cover the split axis over all outputs"); + } if (std::any_of(splits.begin(), splits.end(), [](int64_t split) { return split < 0; })) { return Ort::GetApi().CreateStatus(ORT_INVALID_ARGUMENT, @@ -199,3 +218,20 @@ ONNX_OPERATOR_VERSIONED_KERNEL_EX( .AddTypeConstraint("T", AllFixedSizeTensorTypes()) .SetInputMemType(1, OrtMemTypeCPUInput)), Split) + +// opset 2-10 carries the per-output sizes in the `split` attribute and has no +// second input; opset 11-12 only documented a negative axis. Neither schema +// has bfloat16, and the kernel reads the attribute as well, so one class +// serves both old ranges. (The opset-1 schema, which also allowed `split` as a +// tensor input, is deliberately not registered, matching the CUDA EP.) +ONNX_OPERATOR_VERSIONED_KERNEL_EX(Split, kOnnxDomain, 2, 10, + (Ort::KernelDefBuilder().AddTypeConstraint( + "T", + AllFixedSizeTensorTypesNoBFloat16())), + Split) + +ONNX_OPERATOR_VERSIONED_KERNEL_EX(Split, kOnnxDomain, 11, 12, + (Ort::KernelDefBuilder().AddTypeConstraint( + "T", + AllFixedSizeTensorTypesNoBFloat16())), + Split) diff --git a/musa/ep/src/kernels/tensor/squeeze.cc b/musa/ep/src/kernels/tensor/squeeze.cc index 6aa4d3f..e852c6f 100644 --- a/musa/ep/src/kernels/tensor/squeeze.cc +++ b/musa/ep/src/kernels/tensor/squeeze.cc @@ -54,9 +54,29 @@ OrtStatus* Squeeze::Compute(Ort::KernelContext& ctx) const { } // namespace ONNX_OPERATOR_VERSIONED_KERNEL_EX( - Squeeze, kOnnxDomain, 13, 19, + // opset 21 only added the float8*/int4/uint4 element types, which the + // kernel does not implement, so the opset-13 behaviour still covers it. + Squeeze, kOnnxDomain, 13, 21, (Ort::KernelDefBuilder() .AddTypeConstraint("T", AllFixedSizeTensorTypes()) .SetInputMemType(1, OrtMemTypeCPUInput) .AddInputOutputAlias(0, 0)), Squeeze) + +// Through opset 12 `axes` is an attribute; the kernel reads it as well and +// only prefers the input when the schema provides one. The pre-13 schemas do +// not have bfloat16. opset 11 only documented negative axes, which the kernel +// already accepts, so the two old ranges can share the class. +ONNX_OPERATOR_VERSIONED_KERNEL_EX( + Squeeze, kOnnxDomain, 1, 10, + (Ort::KernelDefBuilder() + .AddTypeConstraint("T", AllFixedSizeTensorTypesNoBFloat16()) + .AddInputOutputAlias(0, 0)), + Squeeze) + +ONNX_OPERATOR_VERSIONED_KERNEL_EX( + Squeeze, kOnnxDomain, 11, 12, + (Ort::KernelDefBuilder() + .AddTypeConstraint("T", AllFixedSizeTensorTypesNoBFloat16()) + .AddInputOutputAlias(0, 0)), + Squeeze) diff --git a/musa/ep/src/kernels/tensor/transpose.cc b/musa/ep/src/kernels/tensor/transpose.cc index ca9ff25..763c416 100644 --- a/musa/ep/src/kernels/tensor/transpose.cc +++ b/musa/ep/src/kernels/tensor/transpose.cc @@ -117,7 +117,9 @@ ONNX_OPERATOR_VERSIONED_KERNEL_EX(Transpose, kOnnxDomain, 1, 12, AllFixedSizeTensorTypesNoBFloat16())), Transpose) +// opset 21 only added the float8*/int4/uint4 element types, which the kernel +// does not implement, so the opset-13 behaviour still covers it. ONNX_OPERATOR_VERSIONED_KERNEL_EX( - Transpose, kOnnxDomain, 13, 19, + Transpose, kOnnxDomain, 13, 21, (Ort::KernelDefBuilder().AddTypeConstraint("T", AllFixedSizeTensorTypes())), Transpose) diff --git a/musa/ep/src/kernels/tensor/unsqueeze.cc b/musa/ep/src/kernels/tensor/unsqueeze.cc index 74c4eb4..439fe17 100644 --- a/musa/ep/src/kernels/tensor/unsqueeze.cc +++ b/musa/ep/src/kernels/tensor/unsqueeze.cc @@ -46,9 +46,29 @@ OrtStatus* Unsqueeze::Compute(Ort::KernelContext& ctx) const { } // namespace ONNX_OPERATOR_VERSIONED_KERNEL_EX( - Unsqueeze, kOnnxDomain, 13, 19, + // opset 21 only added the float8*/int4/uint4 element types, which the + // kernel does not implement, so the opset-13 behaviour still covers it. + Unsqueeze, kOnnxDomain, 13, 21, (Ort::KernelDefBuilder() .AddTypeConstraint("T", AllFixedSizeTensorTypes()) .SetInputMemType(1, OrtMemTypeCPUInput) .AddInputOutputAlias(0, 0)), Unsqueeze) + +// Through opset 12 `axes` is an attribute; the kernel reads it as well and +// only prefers the input when the schema provides one. The pre-13 schemas do +// not have bfloat16. opset 11 only documented negative axes, which the kernel +// already accepts, so the two old ranges can share the class. +ONNX_OPERATOR_VERSIONED_KERNEL_EX( + Unsqueeze, kOnnxDomain, 1, 10, + (Ort::KernelDefBuilder() + .AddTypeConstraint("T", AllFixedSizeTensorTypesNoBFloat16()) + .AddInputOutputAlias(0, 0)), + Unsqueeze) + +ONNX_OPERATOR_VERSIONED_KERNEL_EX( + Unsqueeze, kOnnxDomain, 11, 12, + (Ort::KernelDefBuilder() + .AddTypeConstraint("T", AllFixedSizeTensorTypesNoBFloat16()) + .AddInputOutputAlias(0, 0)), + Unsqueeze) diff --git a/test/ops/test_opset_semantics.py b/test/ops/test_opset_semantics.py new file mode 100644 index 0000000..b192775 --- /dev/null +++ b/test/ops/test_opset_semantics.py @@ -0,0 +1,414 @@ +# Copyright (c) Moore Threads Technology Co., Ltd. All rights reserved. +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Old-opset semantics and the opset-20/21 upper bound. + +Two families of changes are covered here, both of them invisible to a graph +that only uses the current opset: + +* The operators whose meaning depends on the opset the model was stamped with. + Softmax is the important one: through opset 12 the input is coerced into the + 2-D matrix [prod(d0..d_{axis-1}), prod(d_axis..)] and normalised along its + second dimension (axis defaults to 1), while from opset 13 on only `axis` is + reduced (default -1). Unsqueeze/Squeeze/Split/Clip/Pad/Slice/Gemm take their + parameters from node attributes before the opset that turned them into + inputs. Each case below builds the graph at the old opset, runs it on the CPU + EP and on the MUSA EP, and checks both the numbers and that the MUSA EP + actually claimed the node. + +* The registrations whose upper bound was raised to meet models stamped with + opset 20/21, where ONNX only added element types, so the opset-13 behaviour + still covers the node. + +The MUSA session runs with `session.disable_cpu_ep_fallback`, so an op, dtype or +opset the EP does not claim fails loudly instead of silently running on the CPU +EP; the `dtype_diagnostics` report is then checked for `assigned=1`. +""" + +from __future__ import annotations + +import collections + +import numpy as np +import onnxruntime as ort +import pytest +from onnx import helper + +from op_test_utils import TensorProto, build_graph_model, build_model, musa_devices, run + +Case = collections.namedtuple( + "Case", "op_type opset inputs outputs attrs label", defaults=(None, None) +) + +# `Shape` folds away when its input shape is static, so that case keeps one +# dimension symbolic - the same builder the opset lower bound test uses. +_SYMBOLIC_BATCH_BUILDERS = {"Shape"} + + +def _build_model_with_symbolic_batch(op_type, inputs, outputs, attrs=None, opset=17): + """Build a single-node graph whose first input dimension stays symbolic.""" + input_vis = [] + for name, array in inputs.items(): + shape = ["batch", *list(array.shape[1:])] + input_vis.append( + helper.make_tensor_value_info( + name, helper.np_dtype_to_tensor_dtype(array.dtype), shape + ) + ) + output_vis = [helper.make_tensor_value_info(n, e, None) for n, e in outputs] + node = helper.make_node( + op_type, list(inputs.keys()), [n for n, _ in outputs], **(attrs or {}) + ) + graph = helper.make_graph([node], f"{op_type}_graph", input_vis, output_vis) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", opset)]) + model.ir_version = min(model.ir_version, 10) + return model.SerializeToString() + + +def _softmax_cases(): + rng = np.random.default_rng(23) + x = rng.normal(size=(2, 3, 4)).astype(np.float32) + tail_ones = rng.normal(size=(2, 3, 1, 1)).astype(np.float32) + flat = rng.normal(size=(5,)).astype(np.float32) + # Through opset 12 `axis` selects where the input is cut into the 2-D + # matrix, so axis=1 normalises over the product of the trailing dimensions + # while opset 13 normalises over dim 1 alone. Both are exercised, and the + # last-axis case stays on the shared muDNN path in either semantics. + yield Case("Softmax", 10, {"X": x}, [("Y", TensorProto.FLOAT)], {"axis": 1}, + "axis1-flattened") + yield Case("Softmax", 11, {"X": x}, [("Y", TensorProto.FLOAT)], {"axis": 1}, + "axis1-flattened") + yield Case("Softmax", 11, {"X": x}, [("Y", TensorProto.FLOAT)], {"axis": -1}, + "axis-1-last") + yield Case("Softmax", 13, {"X": x}, [("Y", TensorProto.FLOAT)], {"axis": 1}, + "axis1-single") + yield Case("Softmax", 11, {"X": tail_ones}, [("Y", TensorProto.FLOAT)], {"axis": 1}, + "axis1-tail-ones") + # A 1-D input keeps its single axis: the flattened form then reduces over + # the whole row of the promoted [1, d0] matrix. + yield Case("Softmax", 11, {"X": flat}, [("Y", TensorProto.FLOAT)], {"axis": 0}, + "axis0-rank1") + # The Softmax schema stops changing at opset 13, so an opset-21 model still + # resolves the node to version 13 and the 13-17 registration matches it. + yield Case("Softmax", 21, {"X": x}, [("Y", TensorProto.FLOAT)], {"axis": -1}, + "opset21") + + +def _squeeze_family_cases(): + rng = np.random.default_rng(29) + x = rng.normal(size=(2, 3)).astype(np.float32) + squeeze_input = rng.normal(size=(1, 3, 1, 4)).astype(np.float32) + split_input = rng.normal(size=(2, 5, 4)).astype(np.float32) + even_split_input = rng.normal(size=(2, 6, 4)).astype(np.float32) + # Unsqueeze/Squeeze/Split only move `axes`/`split` from attributes to inputs + # at opset 13, and opset 11 documented negative axes. + yield Case("Unsqueeze", 11, {"X": x}, [("Y", TensorProto.FLOAT)], {"axes": [0, 3]}) + yield Case("Unsqueeze", 11, {"X": x}, [("Y", TensorProto.FLOAT)], + {"axes": [-1, -4]}, "negative-axes") + yield Case("Unsqueeze", 21, + {"X": x, "axes": np.array([0, 3], np.int64)}, + [("Y", TensorProto.FLOAT)], {}, "axes-input") + yield Case("Squeeze", 11, {"X": squeeze_input}, [("Y", TensorProto.FLOAT)], {"axes": [0, 2]}) + yield Case("Squeeze", 11, {"X": squeeze_input}, [("Y", TensorProto.FLOAT)], {}, + "no-axes") + yield Case("Squeeze", 21, + {"X": squeeze_input, "axes": np.array([0, 2], np.int64)}, + [("Y", TensorProto.FLOAT)], {}, "axes-input") + yield Case("Split", 11, {"X": split_input}, + [("Y0", TensorProto.FLOAT), ("Y1", TensorProto.FLOAT)], + {"axis": 1, "split": [2, 3]}) + yield Case("Split", 10, {"X": even_split_input}, + [("Y0", TensorProto.FLOAT), ("Y1", TensorProto.FLOAT), ("Y2", TensorProto.FLOAT)], + {"axis": 1}, "equal-split") + yield Case("Split", 11, {"X": even_split_input}, + [("Y0", TensorProto.FLOAT), ("Y1", TensorProto.FLOAT), ("Y2", TensorProto.FLOAT)], + {"axis": 1}, "equal-split") + + +def _parameter_attribute_cases(): + rng = np.random.default_rng(31) + x = rng.uniform(-3.0, 3.0, (2, 3, 4)).astype(np.float32) + xd = rng.uniform(-3.0, 3.0, (2, 4)).astype(np.float64) + x4 = rng.uniform(-3.0, 3.0, (2, 2, 3, 4)).astype(np.float32) + xi = rng.integers(-5, 6, (2, 7)).astype(np.int32) + scaled = np.array(0.5, np.float32) + # Clip kept min/max in FLOAT attributes through opset 10 and moved them to + # inputs at opset 11 (which the CPU EP only implements for float32), with + # the integer element types arriving at opset 12. + yield Case("Clip", 10, {"X": x}, [("Y", TensorProto.FLOAT)], + {"min": -1.0, "max": 1.5}) + yield Case("Clip", 10, {"X": x}, [("Y", TensorProto.FLOAT)], {"min": 0.0}, + "min-only") + yield Case("Clip", 11, {"X": x, "min": scaled, "max": np.array(1.5, np.float32)}, + [("Y", TensorProto.FLOAT)]) + yield Case("Clip", 12, {"X": xi, "min": np.array(-2, np.int32), + "max": np.array(3, np.int32)}, + [("Y", TensorProto.INT32)], {}, "int32") + # Pad kept pads and value in attributes through opset 10; `pads` carries one + # begin/end pair per input axis. + yield Case("Pad", 10, {"X": x4}, [("Y", TensorProto.FLOAT)], + {"pads": [0, 1, 1, 0, 0, 0, 0, 1], "value": 2.5}) + yield Case("Pad", 10, {"X": xd}, [("Y", TensorProto.DOUBLE)], + {"pads": [1, 0, 0, 1], "value": 0.25}, "double-value") + yield Case("Pad", 11, {"X": x4, "pads": np.array([0, 1, 1, 0, 0, 0, 0, 1], np.int64), + "value": np.array(2.5, np.float32)}, + [("Y", TensorProto.FLOAT)], {}, "pads-input") + # Slice kept starts/ends/axes in attributes through opset 9 (with the step + # fixed to 1) and moved them to inputs at opset 10. + yield Case("Slice", 9, {"X": x}, [("Y", TensorProto.FLOAT)], + {"starts": [1], "ends": [3], "axes": [1]}) + yield Case("Slice", 9, {"X": x}, [("Y", TensorProto.FLOAT)], + {"starts": [1, 1], "ends": [-1, 4]}, "implicit-axes") + yield Case("Slice", 10, {"X": x, "starts": np.array([1], np.int64), + "ends": np.array([3], np.int64), + "axes": np.array([2], np.int64)}, + [("Y", TensorProto.FLOAT)], {}, "inputs") + # Gemm's version changes are type sets and the optional C input. + a = x[:2, :, 0] # [2, 3] + b = x[0] # [3, 4] + bias = x[1, 0] # [4] + yield Case("Gemm", 7, {"A": a, "B": b, "C": bias}, [("Y", TensorProto.FLOAT)], + {"alpha": 1.5, "beta": 0.5}) + yield Case("Gemm", 9, {"A": a.T.copy(), "B": b, "C": bias}, + [("Y", TensorProto.FLOAT)], {"transA": 1}) + yield Case("Gemm", 11, {"A": a, "B": b.T.copy()}, [("Y", TensorProto.FLOAT)], + {"transB": 1}, "no-bias") + yield Case("Gemm", 12, {"A": a, "B": b, "C": bias}, [("Y", TensorProto.FLOAT)], {}) + + +def _upper_bound_cases(): + rng = np.random.default_rng(37) + x = rng.normal(size=(2, 3, 4)).astype(np.float32) + squeeze_input = rng.normal(size=(2, 3, 1)).astype(np.float32) + reduce_input = rng.normal(size=(2, 3, 4)).astype(np.float32) + # ONNX only added element types to these operators at opset 20/21, so + # raising the registration's upper bound keeps the same behaviour. + yield Case("Reshape", 21, {"X": x, "shape": np.array([3, 8], np.int64)}, + [("Y", TensorProto.FLOAT)], {}, "opset21") + yield Case("Transpose", 21, {"X": x}, [("Y", TensorProto.FLOAT)], {"perm": [1, 0, 2]}, + "opset21") + yield Case("Shape", 21, {"X": x}, [("Y", TensorProto.INT64)], {}, "opset21") + yield Case("Unsqueeze", 21, {"X": x, "axes": np.array([0, 4], np.int64)}, + [("Y", TensorProto.FLOAT)], {}, "opset21") + yield Case("Squeeze", 21, {"X": squeeze_input, "axes": np.array([2], np.int64)}, + [("Y", TensorProto.FLOAT)], {}, "opset21") + yield Case("ReduceMax", 20, {"X": reduce_input, "axes": np.array([2], np.int64)}, + [("Y", TensorProto.FLOAT)], {"keepdims": 0}, "opset20") + + +CASES = [ + *list(_softmax_cases()), + *list(_squeeze_family_cases()), + *list(_parameter_attribute_cases()), + *list(_upper_bound_cases()), +] + +CASE_IDS = [ + f"{case.op_type}-opset{case.opset}" + (f"-{case.label}" if case.label else "") + for case in CASES +] + +assert len(CASE_IDS) == len(set(CASE_IDS)) + + +def _assert_claimed(capfd, case, expected=True): + """Check the assignment the EP reported for the graph's single node.""" + node_lines = [ + line + for line in capfd.readouterr().err.splitlines() + if line.startswith("MUSA_DTYPE_NODE") + ] + assert len(node_lines) == 1, f"{case.op_type} at opset {case.opset}: {node_lines}" + assert f" op={case.op_type} " in node_lines[0], node_lines[0] + expected_suffix = "assigned=1" if expected else "assigned=0" + assert node_lines[0].rstrip().endswith(expected_suffix), node_lines[0] + return node_lines[0] + + +def _run_case(capfd, case): + """Run one case on the CPU EP and on the MUSA EP and compare.""" + context = f"{case.op_type} at opset {case.opset}" + if case.op_type in _SYMBOLIC_BATCH_BUILDERS: + model_bytes = _build_model_with_symbolic_batch( + case.op_type, case.inputs, case.outputs, case.attrs, case.opset + ) + else: + model_bytes = build_model( + case.op_type, case.inputs, case.outputs, case.attrs, "", case.opset + ) + + cpu_outputs = run(model_bytes, case.inputs, use_musa=False) + + devices = musa_devices() + if not devices: + raise RuntimeError(f"No MUSA device available for {context}") + options = ort.SessionOptions() + options.add_session_config_entry("session.disable_cpu_ep_fallback", "1") + options.add_provider_for_devices(devices, {"dtype_diagnostics": "1"}) + capfd.readouterr() + session = ort.InferenceSession(model_bytes, sess_options=options) + musa_outputs = session.run(None, dict(case.inputs)) + + assert len(cpu_outputs) == len(musa_outputs), f"{context}: output count" + for index, (cpu_output, musa_output) in enumerate(zip(cpu_outputs, musa_outputs)): + assert cpu_output.shape == musa_output.shape, f"{context}: output[{index}] shape" + assert cpu_output.dtype == musa_output.dtype, f"{context}: output[{index}] dtype" + np.testing.assert_allclose( + musa_output, + cpu_output, + rtol=1e-3, + atol=1e-4, + err_msg=f"{context}: output[{index}] CPU vs MUSA mismatch", + ) + + _assert_claimed(capfd, case) + + +@pytest.mark.parametrize("case", CASES, ids=CASE_IDS) +def test_opset_semantics(capfd, case): + _run_case(capfd, case) + + +def _pad_float16_graphs(): + """The two Pad forms of one float16 padding, as (model, feeds) pairs. + + The opset-10 graph carries the pad value in the node's FLOAT `value` + attribute, the opset-11 graph as a float16 input. Both pad a [1, 2] tensor + by one element at the end of each axis. + """ + rng = np.random.default_rng(41) + x = rng.normal(size=(1, 2)).astype(np.float16) + pads = np.array([0, 1, 0, 1], np.int64) + # 0.1 is not representable in float16, so a truncating conversion instead + # of the round-to-nearest one the ONNX reference uses would show up here. + attribute_value = float(np.float32(0.1)) + attribute_model = build_model( + "Pad", + {"X": x}, + [("Y", TensorProto.FLOAT16)], + {"pads": [0, 1, 0, 1], "value": attribute_value}, + "", + 10, + ) + input_model = build_model( + "Pad", + {"X": x, "pads": pads, "value": np.array(np.float16(np.float32(0.1)))}, + [("Y", TensorProto.FLOAT16)], + {}, + "", + 11, + ) + # The float16 value stays a 0-d array: ORT rejects a bare numpy scalar feed. + value_input = np.array(np.float16(np.float32(0.1))) + return ( + (attribute_model, {"X": x}), + (input_model, {"X": x, "pads": pads, "value": value_input}), + ) + + +def test_pad_float16_attribute_value_matches_the_reference(capfd): + """The opset-10 Pad `value` attribute is a FLOAT that must be converted to + the element type of the data. + + The CPU EP has no float16 Pad kernel at any opset, so the reference here is + NumPy (and the opset-11 input form, which carries the float16 bits + directly). + """ + devices = musa_devices() + if not devices: + raise RuntimeError("No MUSA device available for the Pad float16 value test") + + attribute_graph, attribute_feeds = _pad_float16_graphs()[0] + input_graph, input_feeds = _pad_float16_graphs()[1] + options = ort.SessionOptions() + options.add_session_config_entry("session.disable_cpu_ep_fallback", "1") + options.add_provider_for_devices(devices, {"dtype_diagnostics": "1"}) + + capfd.readouterr() + attribute_session = ort.InferenceSession(attribute_graph, sess_options=options) + (attribute_output,) = attribute_session.run(None, dict(attribute_feeds)) + attribute_lines = [ + line + for line in capfd.readouterr().err.splitlines() + if line.startswith("MUSA_DTYPE_NODE") + ] + assert len(attribute_lines) == 1, attribute_lines + assert attribute_lines[0].rstrip().endswith("assigned=1"), attribute_lines + + input_session = ort.InferenceSession(input_graph, sess_options=options) + (input_output,) = input_session.run(None, dict(input_feeds)) + + x = attribute_feeds["X"] + expected = np.pad(x, [(0, 0), (1, 1)], constant_values=np.float16(np.float32(0.1))) + np.testing.assert_array_equal(attribute_output, expected) + np.testing.assert_array_equal(input_output, expected) + np.testing.assert_array_equal(attribute_output, input_output) + + +def _slice_with_steps_graph(opset, step): + """A Slice whose `steps` is a constant initializer, so the graph-level + capability predicate can read it.""" + x = np.arange(24, dtype=np.float32).reshape(2, 3, 4) + negative = step < 0 + starts = [-1] if negative else [0] + ends = [-2**63 + 1] if negative else [3] + node = helper.make_node( + "Slice", ["X", "starts", "ends", "axes", "steps"], ["Y"] + ) + initializers = [ + helper.make_tensor("starts", TensorProto.INT64, [1], starts), + helper.make_tensor("ends", TensorProto.INT64, [1], ends), + helper.make_tensor("axes", TensorProto.INT64, [1], [0]), + helper.make_tensor("steps", TensorProto.INT64, [1], [step]), + ] + return ( + build_graph_model( + [node], {"X": x}, [("Y", TensorProto.FLOAT)], + initializers=initializers, opset=opset, + ), + {"X": x}, + ) + + +def test_slice_step_sign_decides_the_claim(capfd): + """A negative constant step cannot run on the MUSA device copy, so that node + must be left to the CPU EP; a positive step is claimed and matches it.""" + devices = musa_devices() + if not devices: + raise RuntimeError("No MUSA device available for the Slice step test") + options = ort.SessionOptions() + options.add_provider_for_devices(devices, {"dtype_diagnostics": "1"}) + + negative_graph, feeds = _slice_with_steps_graph(13, -1) + capfd.readouterr() + negative_session = ort.InferenceSession(negative_graph, sess_options=options) + (negative_output,) = negative_session.run(None, dict(feeds)) + lines = [ + line + for line in capfd.readouterr().err.splitlines() + if line.startswith("MUSA_DTYPE_NODE") + ] + assert len(lines) == 1 and lines[0].rstrip().endswith("assigned=0"), lines + np.testing.assert_array_equal(negative_output, feeds["X"][::-1]) + + positive_graph, feeds = _slice_with_steps_graph(13, 2) + cpu_outputs = run(positive_graph, feeds, use_musa=False) + capfd.readouterr() + positive_session = ort.InferenceSession(positive_graph, sess_options=options) + (positive_output,) = positive_session.run(None, dict(feeds)) + lines = [ + line + for line in capfd.readouterr().err.splitlines() + if line.startswith("MUSA_DTYPE_NODE") + ] + assert len(lines) == 1 and lines[0].rstrip().endswith("assigned=1"), lines + np.testing.assert_allclose(positive_output, cpu_outputs[0], rtol=1e-6, atol=1e-6) From 1f2dfc5fa275b64048c8d89063d8612d5b0bbab1 Mon Sep 17 00:00:00 2001 From: "yichao.hu" Date: Mon, 12 Oct 2026 23:47:31 -0400 Subject: [PATCH 17/52] feat(plugin): add Erf (opset 9+), fp16 LayerNorm, AveragePool, Sin/Cos, Gelu; fix empty attribute handles; run fp32 unary ops in single precision Claim Erf from opset 9, add float16 LayerNormalization, AveragePool, Sin, Cos and ONNX Gelu (opset 20-21) kernels. Treat an empty attribute handle from Node_GetAttributeByName (optional attributes with no default, e.g. Transpose perm, ReduceMean axes) as "attribute absent"; it used to crash session creation in GetCapability. Evaluate fp32 unary ops with single- precision functions (erff, sinf, cosf, ...) like the ORT CUDA EP; the shared double-precision path had become several times slower. Adds wide-range precision tests and keeps the boundary-copy tests on an unregistered CPU-side operator. --- musa/docs/supported_ops.md | 13 +- musa/ep/src/dtype_contracts.json | 95 +++- musa/ep/src/fusion/concat_matmul_fusion.cc | 2 +- musa/ep/src/fusion/concat_reshape_fusion.cc | 2 +- musa/ep/src/fusion/concat_split_fusion.cc | 2 +- musa/ep/src/fusion/linear_fusion.cc | 4 +- .../src/fusion/log_bucketize_gather_fusion.cc | 2 +- .../log_bucketize_gather_fusion_matcher.cc | 2 +- musa/ep/src/fusion/moe_fusion_matcher.cc | 4 +- .../multi_kqv_mha_output_projection_fusion.cc | 2 +- .../fusion/parallel_linear_fusion_matcher.cc | 4 +- .../fusion/parallel_matmul_concat_fusion.cc | 2 +- .../qkv_attention_output_projection_fusion.cc | 3 +- ...ention_output_projection_fusion_matcher.cc | 2 +- musa/ep/src/fusion/shape_reshape_fusion.cc | 2 +- musa/ep/src/fusion/slice_concat_fusion.cc | 4 +- musa/ep/src/fusion/split_reduce_fusion.cc | 2 +- .../src/fusion/split_sequence_moe_fusion.cc | 2 +- .../split_sequence_moe_fusion_matcher.cc | 4 +- .../fusion/split_unsqueeze_concat_fusion.cc | 4 +- musa/ep/src/fusion/tile_concat_fusion.cc | 2 +- musa/ep/src/graph/capability_predicates.cc | 63 +++ musa/ep/src/graph/graph_utils.cc | 16 +- musa/ep/src/kernels/activation/erf.cc | 4 + musa/ep/src/kernels/activation/gelu.cc | 69 +++ musa/ep/src/kernels/activation/gelu_support.h | 35 ++ musa/ep/src/kernels/math/cos.cc | 44 ++ musa/ep/src/kernels/math/sin.cc | 44 ++ .../math/unary_elementwise_ops_impl.mu | 35 +- musa/ep/src/kernels/nn/average_pool_support.h | 156 ++++++ musa/ep/src/kernels/nn/layer_norm.cc | 34 +- musa/ep/src/kernels/nn/layer_norm_impl.h | 8 +- musa/ep/src/kernels/nn/layer_norm_impl.mu | 64 ++- musa/ep/src/kernels/nn/pool.cc | 158 ++++++ musa/ep/src/kernels/nn/pool_impl.h | 5 + musa/ep/src/kernels/nn/pool_impl.mu | 99 ++++ .../kernels/shared_inc/device_kernel_types.h | 18 + .../shared_inc/kernel_type_constraints.h | 12 + test/ops/test_attribute_defaults.py | 137 +++++ test/ops/test_fusion_ops.py | 511 ++++++++++++++++++ test/ops/test_memcpy_boundary.py | 15 +- test/ops/test_unary_fp32_precision.py | 162 ++++++ 42 files changed, 1760 insertions(+), 88 deletions(-) create mode 100644 musa/ep/src/kernels/activation/gelu.cc create mode 100644 musa/ep/src/kernels/activation/gelu_support.h create mode 100644 musa/ep/src/kernels/math/cos.cc create mode 100644 musa/ep/src/kernels/math/sin.cc create mode 100644 musa/ep/src/kernels/nn/average_pool_support.h create mode 100644 test/ops/test_attribute_defaults.py create mode 100644 test/ops/test_fusion_ops.py create mode 100644 test/ops/test_unary_fp32_precision.py diff --git a/musa/docs/supported_ops.md b/musa/docs/supported_ops.md index 9eb69a9..26db28e 100644 --- a/musa/docs/supported_ops.md +++ b/musa/docs/supported_ops.md @@ -1,8 +1,8 @@ # Supported Ops -- Total supported ops: **92** -- Total registered kernels: **183** +- Total supported ops: **97** +- Total registered kernels: **193** Rows are grouped by operator/domain. Multiple opset ranges in a row preserve separate ORT kernel registrations. @@ -11,6 +11,7 @@ Rows are grouped by operator/domain. Multiple opset ranges in a row preserve sep | `Abs` | `ai.onnx` | 6–12: `T`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64
13–19: `T`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bfloat16 | `activation/abs.cc` | | `Add` | `ai.onnx` | 7–12: `T`: uint32, uint64, int32, int64, float16, float32, float64
13–13: `T`: uint32, uint64, int32, int64, float16, float32, float64, bfloat16
14–19: `T`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bfloat16 | `math/add.cc` | | `And` | `ai.onnx` | 7–19: `T`: bool | `logical/and.cc` | +| `AveragePool` | `ai.onnx` | 7–9: `T`: float16, float32, float64
10–10: `T`: float16, float32, float64
11–19: `T`: float16, float32, float64 | `nn/pool.cc` | | `BatchNormalization` | `ai.onnx` | 9–13: `T`: float16, float32, float64
14–14: `T`: float16, float32, float64
`U`: float16, float32, float64
15–19: `T`: float16, float32, float64
`T1`: float16, float32, float64
`T2`: float16, float32, float64 | `nn/batch_normalization.cc` | | `BitwiseAnd` | `ai.onnx` | 18–19: `T`: uint8, uint16, uint32, uint64, int8, int16, int32, int64 | `logical/bitwise_and.cc` | | `Cast` | `ai.onnx` | 6–8: `T1`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bool
`T2`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bool
9–12: `T1`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bool
`T2`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bool
13–19: `T1`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bfloat16, bool
`T2`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bfloat16, bool | `tensor/cast.cc` | @@ -19,11 +20,12 @@ Rows are grouped by operator/domain. Multiple opset ranges in a row preserve sep | `Concat` | `ai.onnx` | 4–10: `T`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bool
11–12: `T`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bool
13–19: `T`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bfloat16, bool | `tensor/concat.cc` | | `ConstantOfShape` | `ai.onnx` | 9–19: `T1`: int64
`T2`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bool | `generator/constant_of_shape.cc` | | `Conv` | `ai.onnx` | 1–10: `T`: float32
11–19: `T`: float32 | `nn/conv.cc` | +| `Cos` | `ai.onnx` | 7–12: `T`: float16, float32, float64
13–19: `T`: float16, float32, float64, bfloat16 | `math/cos.cc` | | `CumSum` | `ai.onnx` | 11–13: `T`: int32, int64, uint32, uint64, float32, float64
`T2`: int32, int64
14–19: `T`: ?
`T2`: int32, int64 | `math/cumsum.cc` | | `Div` | `ai.onnx` | 7–12: `T`: uint32, uint64, int32, int64, float16, float32, float64
13–13: `T`: uint32, uint64, int32, int64, float16, float32, float64, bfloat16
14–19: `T`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bfloat16 | `math/div.cc` | | `Einsum` | `ai.onnx` | 12–19: `T`: float32 | `math/einsum.cc` | | `Equal` | `ai.onnx` | 7–10: `T`: int32, int64, bool
11–12: `T`: uint32, uint64, int32, int64, float16, float32, float64, bool
13–19: `T`: uint32, uint64, int32, int64, float16, float32, float64, bfloat16, bool
19–19: `T`: string | `logical/equal.cc` | -| `Erf` | `ai.onnx` | 13–19: `T`: float16, float32, float64, bfloat16 | `activation/erf.cc` | +| `Erf` | `ai.onnx` | 9–12: `T`: float16, float32, float64
13–19: `T`: float16, float32, float64, bfloat16 | `activation/erf.cc` | | `Exp` | `ai.onnx` | 6–12: `T`: float16, float32, float64
13–19: `T`: float16, float32, float64, bfloat16 | `math/exp.cc` | | `Expand` | `ai.onnx` | 8–12: `T`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bool
13–19: `T`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bfloat16, bool | `tensor/expand.cc` | | `Flatten` | `ai.onnx` | 1–8: `T`: float16, float32, float64
9–10: `T`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bool
11–12: `T`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bool
13–19: `T`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bfloat16, bool | `tensor/flatten.cc` | @@ -31,6 +33,7 @@ Rows are grouped by operator/domain. Multiple opset ranges in a row preserve sep | `Gather` | `ai.onnx` | 1–10: `T`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bool
`Tind`: int32, int64
11–12: `T`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bool
`Tind`: int32, int64
13–19: `T`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bfloat16, bool
`Tind`: int32, int64 | `tensor/gather.cc` | | `GatherElements` | `ai.onnx` | 11–12: `T`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bfloat16, bool
`Tind`: int32, int64
13–19: `T`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bfloat16, bool
`Tind`: int32, int64 | `tensor/gather_elements.cc` | | `GatherND` | `ai.onnx` | 11–12: `T`: float32, float64, float16, int64, bfloat16, bool
`indices`: int64
13–19: `T`: float32, float64, float16, int64, bfloat16, bool
`indices`: int64 | `tensor/gather_nd.cc` | +| `Gelu` | `ai.onnx` | 20–21: `T`: float16, float32, float64, bfloat16 | `activation/gelu.cc` | | `Gemm` | `ai.onnx` | 7–8: `T`: float16, float32, float64
9–10: `T`: float16, float32, float64
11–12: `T`: float16, float32, float64
13–19: `T`: float16, float32, float64, bfloat16 | `math/gemm.cc` | | `GlobalAveragePool` | `ai.onnx` | 1–19: `T`: float16, float32, float64 | `nn/pool.cc` | | `GlobalMaxPool` | `ai.onnx` | 1–19: `T`: float16, float32, float64 | `nn/pool.cc` | @@ -39,7 +42,7 @@ Rows are grouped by operator/domain. Multiple opset ranges in a row preserve sep | `Identity` | `ai.onnx` | 1–12: `T`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bool
13–13: `T`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bfloat16, bool
14–18: `V`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bfloat16, bool
19–19: `V`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bfloat16, bool | `tensor/identity_op.cc` | | `If` | `ai.onnx` | 13–15: `B`: bool
`V`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bool
16–19: `B`: bool
`V`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bfloat16, bool | `controlflow/if.cc` | | `IsNaN` | `ai.onnx` | 13–19: `T1`: float16, float32, float64, bfloat16
`T2`: bool | `logical/is_nan.cc` | -| `LayerNormalization` | `ai.onnx` | 17–19: `T`: float32 | `nn/layer_norm.cc` | +| `LayerNormalization` | `ai.onnx` | 17–19: `T`: float32, float16 | `nn/layer_norm.cc` | | `LeakyRelu` | `ai.onnx` | 6–15: `T`: float16, float32, float64
16–19: `T`: float16, float32, float64, bfloat16 | `activation/leaky_relu.cc` | | `Less` | `ai.onnx` | 7–8: `T`: float16, float32, float64
9–12: `T`: uint32, uint64, int32, int64, float16, float32, float64
13–19: `T`: uint32, uint64, int32, int64, float16, float32, float64, bfloat16 | `logical/less.cc` | | `LessOrEqual` | `ai.onnx` | 12–15: `T`: uint32, uint64, int32, int64, float16, float32, float64
16–19: `T`: uint32, uint64, int32, int64, float16, float32, float64, bfloat16 | `logical/less_or_equal.cc` | @@ -81,6 +84,7 @@ Rows are grouped by operator/domain. Multiple opset ranges in a row preserve sep | `Shape` | `ai.onnx` | 1–12: `T`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bool
`T1`: int64
13–21: `T`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bfloat16, bool
`T1`: int64 | `tensor/shape.cc` | | `Sigmoid` | `ai.onnx` | 6–12: `T`: float16, float32, float64
13–19: `T`: float16, float32, float64, bfloat16 | `activation/sigmoid.cc` | | `Sign` | `ai.onnx` | 13–19: `T`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bfloat16 | `math/sign.cc` | +| `Sin` | `ai.onnx` | 7–12: `T`: float16, float32, float64
13–19: `T`: float16, float32, float64, bfloat16 | `math/sin.cc` | | `Slice` | `ai.onnx` | 1–9: `T`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bool
10–10: `T`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bool
`Tind`: int32, int64
11–12: `T`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bool
`Tind`: int32, int64
13–19: `T`: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64, bfloat16, bool
`Tind`: int32, int64 | `tensor/slice.cc` | | `Softmax` | `ai.onnx` | 1–10: `T`: float16, float32, float64
11–12: `T`: float16, float32, float64
13–17: `T`: float16, float32, float64, bfloat16 | `math/softmax.cc` | | `Softplus` | `ai.onnx` | 1–19: `T`: float16, float32, float64 | `activation/softplus.cc` | @@ -99,4 +103,5 @@ Rows are grouped by operator/domain. Multiple opset ranges in a row preserve sep | `Attention` | `com.microsoft` | 1–1: `T`: float32
`M`: int32 | `llm/attention.cc` | | `FusedGemm` | `com.microsoft` | 1–1: `T`: float16, float32, float64 | `math/fused_gemm.cc` | | `FusedMatMul` | `com.microsoft` | 1–1: `T`: float16, float32, bfloat16 | `math/fused_matmul.cc` | +| `Gelu` | `com.microsoft` | 1–1: `T`: float16, float32, float64, bfloat16 | `activation/gelu.cc` | | `MultiHeadAttention` | `com.microsoft` | 1–1: `T`: float32
`QK`: float32
`M`: int32 | `llm/multihead_attention.cc` | diff --git a/musa/ep/src/dtype_contracts.json b/musa/ep/src/dtype_contracts.json index 467dd1f..f53c327 100644 --- a/musa/ep/src/dtype_contracts.json +++ b/musa/ep/src/dtype_contracts.json @@ -590,18 +590,19 @@ { "op": "Erf", "domain": "ai.onnx", - "opset": "13-19", + "opset": "9-19", "storage": "bfloat16,float16,float32,float64", "output": "same as input", "compute": "FP32 for BF16/FP16; FP64 for float64", "accumulation": "none", - "path": "muDNN unary when available; device unary kernel fallback", + "path": "muDNN unary when available; device unary kernel fallback. opset 9-12 reuses the same kernel class (the schema only added bfloat16 at 13), which is what the opset 10-12 Bert/GPT models need after ORT stops fusing their Gelu subgraph around the EP-claimed nodes", "verified_dtypes": "float32,float16,float64,bfloat16", "tests": [ "test/ops/test_erf.py", - "test/ops/test_unary_dtype_matrix.py" + "test/ops/test_unary_dtype_matrix.py", + "test/ops/test_fusion_ops.py" ], - "unsupported_reasons": "integer,bool,string,complex are excluded by the ONNX Erf type constraint" + "unsupported_reasons": "integer,bool,string,complex are excluded by the ONNX Erf type constraint; the 9-12 registration has no bfloat16 (opset 13)" }, { "op": "Exp", @@ -811,16 +812,17 @@ "op": "LayerNormalization", "domain": "ai.onnx", "opset": "17-19", - "storage": "float32 X/output; float32 Scale/B (B optional)", + "storage": "float32 or float16 X/Y; Scale/B of the same element type as X (ORT binds them to T); B optional", "output": "same shape as X; optional float32 mean and inv_std", "compute": "row-wise normalization over axis suffix", "accumulation": "FP32", - "path": "muDNN LayerNorm or MUSA FP32 device kernel", - "verified_dtypes": "float32", + "path": "muDNN LayerNorm for float32, MUSA device kernel for float32/float16", + "verified_dtypes": "float32,float16", "tests": [ - "test/ops/test_layer_normalization.py" + "test/ops/test_layer_normalization.py", + "test/ops/test_fusion_ops.py" ], - "unsupported_reasons": "stash_type=1 only; Scale/B must be float and match normalized suffix; registration currently exposes float32 only" + "unsupported_reasons": "stash_type=1 only; Scale/B must match X in type and the normalized suffix shape; Scale/B float32 next to float16 X is rejected because ORT itself cannot load it" }, { "op": "Flatten", @@ -1352,6 +1354,81 @@ "test/ops/test_reverse_sequence.py" ], "unsupported_reasons": "schema excludes BF16 through opset 19; string/complex unregistered; axes restricted to 0/1" + }, + { + "op": "Sin", + "domain": "ai.onnx", + "opset": "7-19", + "storage": "bfloat16,float16,float32,float64", + "output": "same as input", + "compute": "FP32 for float16; FP64 for float64", + "accumulation": "none", + "path": "MUSA device unary kernel (sinf/sin); the schema is unchanged from opset 7 through 19 apart from bfloat16 at 13", + "verified_dtypes": "float32,float16", + "tests": [ + "test/ops/test_fusion_ops.py" + ], + "unsupported_reasons": "integer,bool,string,complex are excluded by the ONNX Sin type constraint; the 7-12 registration has no bfloat16" + }, + { + "op": "Cos", + "domain": "ai.onnx", + "opset": "7-19", + "storage": "bfloat16,float16,float32,float64", + "output": "same as input", + "compute": "FP32 for float16; FP64 for float64", + "accumulation": "none", + "path": "MUSA device unary kernel (cosf/cos); the schema is unchanged from opset 7 through 19 apart from bfloat16 at 13", + "verified_dtypes": "float32,float16", + "tests": [ + "test/ops/test_fusion_ops.py" + ], + "unsupported_reasons": "integer,bool,string,complex are excluded by the ONNX Cos type constraint; the 7-12 registration has no bfloat16" + }, + { + "op": "AveragePool", + "domain": "ai.onnx", + "opset": "7-19", + "storage": "float16,float32,float64", + "output": "same as input", + "compute": "window sum and divide; the divisor is the kernel size when count_include_pad is set, otherwise the number of window positions inside the input", + "accumulation": "FP32 for float16; FP64 for float32/float64 custom path", + "path": "muDNN Pooling AVGPOOL_COUNT_PAD/AVGPOOL_COUNT_WITHOUT_PAD for symmetric pads, MUSA device pooling kernel otherwise", + "verified_dtypes": "float32,float16", + "tests": [ + "test/ops/test_fusion_ops.py" + ], + "unsupported_reasons": "ceil_mode=1 stays on the CPU EP (the capability predicate and the kernel share the rejection list in kernels/nn/average_pool_support.h); rank 3-8, concrete shapes, auto_pad NOTSET/VALID/SAME_UPPER/SAME_LOWER" + }, + { + "op": "Gelu", + "domain": "ai.onnx", + "opset": "20-21", + "storage": "bfloat16,float16,float32,float64", + "output": "same as input", + "compute": "0.5*x*(1+erf(x/sqrt(2))); FP32 for float16/bfloat16", + "accumulation": "none", + "path": "MUSA device unary kernel through the shared unary dispatcher", + "verified_dtypes": "float32,float16", + "tests": [ + "test/ops/test_fusion_ops.py" + ], + "unsupported_reasons": "approximate=tanh stays on the CPU EP (capability predicate in graph/capability_predicates.cc shares kernels/activation/gelu_support.h with the kernel)" + }, + { + "op": "Gelu", + "domain": "com.microsoft", + "opset": "1", + "storage": "bfloat16,float16,float32,float64", + "output": "same as input", + "compute": "0.5*x*(1+erf(x/sqrt(2))); FP32 for float16/bfloat16", + "accumulation": "none", + "path": "MUSA device unary kernel, shared with the ONNX Gelu", + "verified_dtypes": "float32,float16", + "tests": [ + "test/ops/test_fusion_ops.py" + ], + "unsupported_reasons": "an int approximate=1 (tanh) stays on the CPU EP; no fusion in this ORT version produces these nodes for the plugin EP yet" } ] } diff --git a/musa/ep/src/fusion/concat_matmul_fusion.cc b/musa/ep/src/fusion/concat_matmul_fusion.cc index e93165d..5192f57 100644 --- a/musa/ep/src/fusion/concat_matmul_fusion.cc +++ b/musa/ep/src/fusion/concat_matmul_fusion.cc @@ -178,7 +178,7 @@ std::string Name(Ort::ConstValueInfo value_info) { int64_t ReadIntAttribute(Ort::ConstNode node, const std::string& name) { Ort::ConstOpAttr attr; Ort::Status status = node.GetAttributeByName(name, attr); - if (!status.IsOK()) { + if (!status.IsOK() || !attr) { throw std::runtime_error("missing required node attribute: " + name); } diff --git a/musa/ep/src/fusion/concat_reshape_fusion.cc b/musa/ep/src/fusion/concat_reshape_fusion.cc index 2a3a734..c231dcf 100644 --- a/musa/ep/src/fusion/concat_reshape_fusion.cc +++ b/musa/ep/src/fusion/concat_reshape_fusion.cc @@ -452,7 +452,7 @@ int64_t ReadIntAttribute(Ort::ConstNode node, const std::string& name, int64_t default_value) { Ort::ConstOpAttr attr; Ort::Status status = node.GetAttributeByName(name, attr); - if (!status.IsOK()) { + if (!status.IsOK() || !attr) { return default_value; } diff --git a/musa/ep/src/fusion/concat_split_fusion.cc b/musa/ep/src/fusion/concat_split_fusion.cc index 66db100..6074485 100644 --- a/musa/ep/src/fusion/concat_split_fusion.cc +++ b/musa/ep/src/fusion/concat_split_fusion.cc @@ -42,7 +42,7 @@ int64_t ReadIntAttribute(Ort::ConstNode node, const std::string& name, int64_t default_value) { Ort::ConstOpAttr attr; Ort::Status status = node.GetAttributeByName(name, attr); - if (!status.IsOK()) { + if (!status.IsOK() || !attr) { return default_value; } diff --git a/musa/ep/src/fusion/linear_fusion.cc b/musa/ep/src/fusion/linear_fusion.cc index a93b0a5..3fe0a1d 100644 --- a/musa/ep/src/fusion/linear_fusion.cc +++ b/musa/ep/src/fusion/linear_fusion.cc @@ -108,7 +108,7 @@ int64_t ReadIntAttribute(Ort::ConstNode node, const std::string& name, int64_t default_value) { Ort::ConstOpAttr attr; Ort::Status status = node.GetAttributeByName(name, attr); - if (!status.IsOK()) { + if (!status.IsOK() || !attr) { return default_value; } @@ -121,7 +121,7 @@ float ReadFloatAttribute(Ort::ConstNode node, const std::string& name, float default_value) { Ort::ConstOpAttr attr; Ort::Status status = node.GetAttributeByName(name, attr); - if (!status.IsOK()) { + if (!status.IsOK() || !attr) { return default_value; } diff --git a/musa/ep/src/fusion/log_bucketize_gather_fusion.cc b/musa/ep/src/fusion/log_bucketize_gather_fusion.cc index 613d0f7..f038189 100644 --- a/musa/ep/src/fusion/log_bucketize_gather_fusion.cc +++ b/musa/ep/src/fusion/log_bucketize_gather_fusion.cc @@ -165,7 +165,7 @@ Ort::ConstValue ReadInitializerValue(Ort::ConstValueInfo value_info, Ort::Value ReadConstantNodeValue(Ort::ConstNode node, const char* name) { Ort::ConstOpAttr attr; Ort::Status attr_status = node.GetAttributeByName("value", attr); - if (!attr_status.IsOK()) { + if (!attr_status.IsOK() || !attr) { throw std::runtime_error(std::string("LogBucketizeGather ") + name + " requires Constant value"); } diff --git a/musa/ep/src/fusion/log_bucketize_gather_fusion_matcher.cc b/musa/ep/src/fusion/log_bucketize_gather_fusion_matcher.cc index d10fe17..1e7ce67 100644 --- a/musa/ep/src/fusion/log_bucketize_gather_fusion_matcher.cc +++ b/musa/ep/src/fusion/log_bucketize_gather_fusion_matcher.cc @@ -130,7 +130,7 @@ std::optional> ReadSmallIntConstantOrInitializer( } Ort::ConstOpAttr attr; Ort::Status attr_status = producer.GetAttributeByName("value", attr); - if (!attr_status.IsOK()) { + if (!attr_status.IsOK() || !attr) { return std::nullopt; } Ort::Value value{nullptr}; diff --git a/musa/ep/src/fusion/moe_fusion_matcher.cc b/musa/ep/src/fusion/moe_fusion_matcher.cc index be5aed9..378b113 100644 --- a/musa/ep/src/fusion/moe_fusion_matcher.cc +++ b/musa/ep/src/fusion/moe_fusion_matcher.cc @@ -57,7 +57,7 @@ float ReadFloatAttribute(Ort::ConstNode node, const char* name, float default_value) { Ort::ConstOpAttr attr; float value = default_value; - return node.GetAttributeByName(name, attr).IsOK() && + return node.GetAttributeByName(name, attr).IsOK() && attr && attr.GetValue(value).IsOK() ? value : default_value; @@ -88,7 +88,7 @@ std::optional> ReadIntConstant( } Ort::ConstOpAttr attr; Ort::Value value{nullptr}; - if (!producer.GetAttributeByName("value", attr).IsOK() || + if (!producer.GetAttributeByName("value", attr).IsOK() || !attr || !attr.GetTensorAttributeAsOrtValue(value).IsOK() || !value) { return std::nullopt; } diff --git a/musa/ep/src/fusion/multi_kqv_mha_output_projection_fusion.cc b/musa/ep/src/fusion/multi_kqv_mha_output_projection_fusion.cc index dad8e8d..961a1e0 100644 --- a/musa/ep/src/fusion/multi_kqv_mha_output_projection_fusion.cc +++ b/musa/ep/src/fusion/multi_kqv_mha_output_projection_fusion.cc @@ -376,7 +376,7 @@ size_t FusedInputIndex(const std::unordered_map& indices, float GetFloatAttribute(Ort::ConstNode node, const char* name, float fallback) { Ort::ConstOpAttr attr; float value = fallback; - return node.GetAttributeByName(name, attr).IsOK() && + return node.GetAttributeByName(name, attr).IsOK() && attr && attr.GetValue(value).IsOK() ? value : fallback; diff --git a/musa/ep/src/fusion/parallel_linear_fusion_matcher.cc b/musa/ep/src/fusion/parallel_linear_fusion_matcher.cc index b88cbaf..3df25ff 100644 --- a/musa/ep/src/fusion/parallel_linear_fusion_matcher.cc +++ b/musa/ep/src/fusion/parallel_linear_fusion_matcher.cc @@ -50,7 +50,7 @@ bool RequireSameParallelLinearStorageType( int64_t ReadIntAttribute(Ort::ConstNode node, const char* name, int64_t default_value) { Ort::ConstOpAttr attr; - if (!node.GetAttributeByName(name, attr).IsOK()) { + if (!node.GetAttributeByName(name, attr).IsOK() || !attr) { return default_value; } int64_t value = default_value; @@ -60,7 +60,7 @@ int64_t ReadIntAttribute(Ort::ConstNode node, const char* name, float ReadFloatAttribute(Ort::ConstNode node, const char* name, float default_value) { Ort::ConstOpAttr attr; - if (!node.GetAttributeByName(name, attr).IsOK()) { + if (!node.GetAttributeByName(name, attr).IsOK() || !attr) { return default_value; } float value = default_value; diff --git a/musa/ep/src/fusion/parallel_matmul_concat_fusion.cc b/musa/ep/src/fusion/parallel_matmul_concat_fusion.cc index d353c40..7067745 100644 --- a/musa/ep/src/fusion/parallel_matmul_concat_fusion.cc +++ b/musa/ep/src/fusion/parallel_matmul_concat_fusion.cc @@ -450,7 +450,7 @@ std::unique_ptr CreateParallelMatMulConcatFusion( int64_t concat_axis = 0; Ort::ConstOpAttr axis_attr; Ort::Status status = concat_node.GetAttributeByName("axis", axis_attr); - if (!status.IsOK() || !axis_attr.GetValue(concat_axis).IsOK()) { + if (!status.IsOK() || !axis_attr || !axis_attr.GetValue(concat_axis).IsOK()) { throw std::runtime_error("ParallelMatMulConcat missing Concat axis"); } diff --git a/musa/ep/src/fusion/qkv_attention_output_projection_fusion.cc b/musa/ep/src/fusion/qkv_attention_output_projection_fusion.cc index 57c71e1..2e33099 100644 --- a/musa/ep/src/fusion/qkv_attention_output_projection_fusion.cc +++ b/musa/ep/src/fusion/qkv_attention_output_projection_fusion.cc @@ -57,7 +57,8 @@ size_t Index(const std::unordered_map& m, float FloatAttr(Ort::ConstNode n, const char* name, float d) { Ort::ConstOpAttr a; float x = d; - return n.GetAttributeByName(name, a).IsOK() && a.GetValue(x).IsOK() ? x : d; + return n.GetAttributeByName(name, a).IsOK() && a && a.GetValue(x).IsOK() ? x + : d; } // Prefer muDNN's GEMM+bias epilogue. It removes the standalone device bias diff --git a/musa/ep/src/fusion/qkv_attention_output_projection_fusion_matcher.cc b/musa/ep/src/fusion/qkv_attention_output_projection_fusion_matcher.cc index 6e52761..f26f753 100644 --- a/musa/ep/src/fusion/qkv_attention_output_projection_fusion_matcher.cc +++ b/musa/ep/src/fusion/qkv_attention_output_projection_fusion_matcher.cc @@ -12,7 +12,7 @@ float ReadFloatAttribute(Ort::ConstNode node, const char* name, float default_value) { Ort::ConstOpAttr attribute; float value = default_value; - return node.GetAttributeByName(name, attribute).IsOK() && + return node.GetAttributeByName(name, attribute).IsOK() && attribute && attribute.GetValue(value).IsOK() ? value : default_value; diff --git a/musa/ep/src/fusion/shape_reshape_fusion.cc b/musa/ep/src/fusion/shape_reshape_fusion.cc index 8ee95c1..9669f74 100644 --- a/musa/ep/src/fusion/shape_reshape_fusion.cc +++ b/musa/ep/src/fusion/shape_reshape_fusion.cc @@ -240,7 +240,7 @@ int64_t ReadIntAttribute(Ort::ConstNode node, const std::string& name, int64_t default_value) { Ort::ConstOpAttr attr; Ort::Status status = node.GetAttributeByName(name, attr); - if (!status.IsOK()) { + if (!status.IsOK() || !attr) { return default_value; } diff --git a/musa/ep/src/fusion/slice_concat_fusion.cc b/musa/ep/src/fusion/slice_concat_fusion.cc index 3077eea..356b5f6 100644 --- a/musa/ep/src/fusion/slice_concat_fusion.cc +++ b/musa/ep/src/fusion/slice_concat_fusion.cc @@ -43,7 +43,7 @@ int64_t ReadIntAttribute(Ort::ConstNode node, const std::string& name, int64_t default_value) { Ort::ConstOpAttr attr; Ort::Status status = node.GetAttributeByName(name, attr); - if (!status.IsOK()) { + if (!status.IsOK() || !attr) { return default_value; } @@ -243,7 +243,7 @@ bool IsZeroFloatConstantOfShape(Ort::ConstNode node) { Ort::ConstOpAttr attr; Ort::Status attr_status = node.GetAttributeByName("value", attr); - if (!attr_status.IsOK()) { + if (!attr_status.IsOK() || !attr) { return true; } diff --git a/musa/ep/src/fusion/split_reduce_fusion.cc b/musa/ep/src/fusion/split_reduce_fusion.cc index d670a08..314cf05 100644 --- a/musa/ep/src/fusion/split_reduce_fusion.cc +++ b/musa/ep/src/fusion/split_reduce_fusion.cc @@ -41,7 +41,7 @@ int64_t ReadIntAttribute(Ort::ConstNode node, const std::string& name, int64_t default_value) { Ort::ConstOpAttr attr; Ort::Status status = node.GetAttributeByName(name, attr); - if (!status.IsOK()) { + if (!status.IsOK() || !attr) { return default_value; } diff --git a/musa/ep/src/fusion/split_sequence_moe_fusion.cc b/musa/ep/src/fusion/split_sequence_moe_fusion.cc index 5a46c54..abec92f 100644 --- a/musa/ep/src/fusion/split_sequence_moe_fusion.cc +++ b/musa/ep/src/fusion/split_sequence_moe_fusion.cc @@ -59,7 +59,7 @@ float ReadFloatAttribute(Ort::ConstNode node, const char* name, float default_value) { Ort::ConstOpAttr attr; float value = default_value; - return node.GetAttributeByName(name, attr).IsOK() && + return node.GetAttributeByName(name, attr).IsOK() && attr && attr.GetValue(value).IsOK() ? value : default_value; diff --git a/musa/ep/src/fusion/split_sequence_moe_fusion_matcher.cc b/musa/ep/src/fusion/split_sequence_moe_fusion_matcher.cc index 2307eef..449c15e 100644 --- a/musa/ep/src/fusion/split_sequence_moe_fusion_matcher.cc +++ b/musa/ep/src/fusion/split_sequence_moe_fusion_matcher.cc @@ -61,7 +61,7 @@ float ReadFloatAttribute(Ort::ConstNode node, const char* name, float default_value) { Ort::ConstOpAttr attr; float value = default_value; - return node.GetAttributeByName(name, attr).IsOK() && + return node.GetAttributeByName(name, attr).IsOK() && attr && attr.GetValue(value).IsOK() ? value : default_value; @@ -79,7 +79,7 @@ std::optional> ReadIntConstant( } Ort::ConstOpAttr attr; Ort::Value value{nullptr}; - if (!producer.GetAttributeByName("value", attr).IsOK() || + if (!producer.GetAttributeByName("value", attr).IsOK() || !attr || !attr.GetTensorAttributeAsOrtValue(value).IsOK() || !value) { return std::nullopt; } diff --git a/musa/ep/src/fusion/split_unsqueeze_concat_fusion.cc b/musa/ep/src/fusion/split_unsqueeze_concat_fusion.cc index fb8d805..acd624a 100644 --- a/musa/ep/src/fusion/split_unsqueeze_concat_fusion.cc +++ b/musa/ep/src/fusion/split_unsqueeze_concat_fusion.cc @@ -45,7 +45,7 @@ std::optional GetIntAttribute(Ort::ConstNode node, const std::string& name) { Ort::ConstOpAttr attr; Ort::Status status = node.GetAttributeByName(name, attr); - if (!status.IsOK()) { + if (!status.IsOK() || !attr) { return std::nullopt; } @@ -62,7 +62,7 @@ std::optional> GetIntsAttribute(Ort::ConstNode node, const std::string& name) { Ort::ConstOpAttr attr; Ort::Status status = node.GetAttributeByName(name, attr); - if (!status.IsOK()) { + if (!status.IsOK() || !attr) { return std::nullopt; } diff --git a/musa/ep/src/fusion/tile_concat_fusion.cc b/musa/ep/src/fusion/tile_concat_fusion.cc index 0a46910..9190d48 100644 --- a/musa/ep/src/fusion/tile_concat_fusion.cc +++ b/musa/ep/src/fusion/tile_concat_fusion.cc @@ -358,7 +358,7 @@ int64_t ReadIntAttribute(Ort::ConstNode node, const std::string& name, int64_t default_value) { Ort::ConstOpAttr attr; Ort::Status status = node.GetAttributeByName(name, attr); - if (!status.IsOK()) { + if (!status.IsOK() || !attr) { return default_value; } diff --git a/musa/ep/src/graph/capability_predicates.cc b/musa/ep/src/graph/capability_predicates.cc index dc8abd1..877a0f2 100644 --- a/musa/ep/src/graph/capability_predicates.cc +++ b/musa/ep/src/graph/capability_predicates.cc @@ -19,7 +19,9 @@ #include #include +#include "activation/gelu_support.h" #include "graph/graph_utils.h" +#include "nn/average_pool_support.h" #include "tensor/resize_support.h" namespace musa_ep { @@ -152,6 +154,58 @@ bool SliceConfigurationSupported(Ort::ConstNode node) { [](int64_t step) { return step > 0; }); } +/// True when the MUSA Gelu kernel can run this node. +/// +/// The route implements the exact erf form. The ONNX Gelu carries +/// `approximate` as a string, the com.microsoft Gelu as an int (0 = exact, +/// 1 = tanh); a node asking for the tanh approximation stays on the CPU EP. +bool GeluConfigurationSupported(Ort::ConstNode node) { + std::string approximate = kGeluDefaultApproximate; + if (std::optional value = + GetStringAttribute(node, "approximate"); + value.has_value()) { + approximate = *value; + } else if (std::optional value = + GetIntAttribute(node, "approximate"); + value.has_value()) { + approximate = *value == 0 ? kGeluDefaultApproximate : "tanh"; + } + return GeluApproximateRejection(approximate) == nullptr; +} + +/// True when the MUSA AveragePool kernel can run this node. +/// +/// The kernel covers the auto_pad modes, count_include_pad and the attribute +/// pads, but not ceil_mode, and it is a device pooling route limited to the +/// rank and element types the shared route definition lists. A node the route +/// rejects stays on the CPU EP, where it ran before the kernel was registered. +bool AveragePoolConfigurationSupported(Ort::ConstNode node) { + musa_ep::AveragePoolRouteConfig config; + config.auto_pad = + GetStringAttribute(node, "auto_pad").value_or(kAveragePoolDefaultAutoPad); + config.ceil_mode = GetIntAttribute(node, "ceil_mode").value_or(0); + config.count_include_pad = + GetIntAttribute(node, "count_include_pad").value_or(0); + config.kernel_shape = + GetIntsAttribute(node, "kernel_shape").value_or(std::vector{}); + config.pads = GetIntsAttribute(node, "pads").value_or(std::vector{}); + config.strides = + GetIntsAttribute(node, "strides").value_or(std::vector{}); + + const std::vector inputs = node.GetInputs(); + if (inputs.empty() || inputs[0] == nullptr) { + return false; + } + const Ort::ConstTypeInfo input_type = inputs[0].TypeInfo(); + if (input_type.GetONNXType() != ONNX_TYPE_TENSOR) { + return false; + } + const auto input_info = input_type.GetTensorTypeAndShapeInfo(); + config.input_rank = static_cast(input_info.GetShape().size()); + config.input_element_type = input_info.GetElementType(); + return AveragePoolRouteRejection(config) == nullptr; +} + } // namespace bool IsNodeConfigurationSupported(Ort::ConstNode node) { @@ -161,6 +215,15 @@ bool IsNodeConfigurationSupported(Ort::ConstNode node) { if (IsOnnxOp(node, "Slice")) { return SliceConfigurationSupported(node); } + if (IsOnnxOp(node, "AveragePool")) { + return AveragePoolConfigurationSupported(node); + } + // Gelu is an official ONNX operator from opset 20 on, and the plugin also + // registers the com.microsoft Gelu (same semantics, int approximation). + if (node.GetOperatorType() == "Gelu" && + (IsOnnxDomain(node.GetDomain()) || node.GetDomain() == "com.microsoft")) { + return GeluConfigurationSupported(node); + } return true; } diff --git a/musa/ep/src/graph/graph_utils.cc b/musa/ep/src/graph/graph_utils.cc index 933f937..7454439 100644 --- a/musa/ep/src/graph/graph_utils.cc +++ b/musa/ep/src/graph/graph_utils.cc @@ -96,11 +96,17 @@ std::optional> GetStaticShape( return shape; } +/// Reads an int attribute, or nullopt when the node does not carry it. +/// +/// ORT reports an attribute it cannot provide as OK with an empty handle: an +/// optional attribute that has no constant default value (Transpose's `perm`, +/// ReduceMean's `axes`, ...) comes back that way, and every call into the +/// handle dereferences null. An empty handle therefore means "not present". std::optional GetIntAttribute(Ort::ConstNode node, const std::string& name) { Ort::ConstOpAttr attr; Ort::Status status = node.GetAttributeByName(name, attr); - if (!status.IsOK()) { + if (!status.IsOK() || !attr) { return std::nullopt; } @@ -117,7 +123,8 @@ std::optional> GetIntsAttribute(Ort::ConstNode node, const std::string& name) { Ort::ConstOpAttr attr; Ort::Status status = node.GetAttributeByName(name, attr); - if (!status.IsOK()) { + // An empty handle means the attribute is not present; see above. + if (!status.IsOK() || !attr) { return std::nullopt; } @@ -133,7 +140,8 @@ std::optional GetStringAttribute(Ort::ConstNode node, const std::string& name) { Ort::ConstOpAttr attr; Ort::Status status = node.GetAttributeByName(name, attr); - if (!status.IsOK()) { + // An empty handle means the attribute is not present; see above. + if (!status.IsOK() || !attr) { return std::nullopt; } @@ -314,7 +322,7 @@ bool IsZeroFloatConstantOfShape(Ort::ConstNode node) { Ort::ConstOpAttr attr; Ort::Status attr_status = node.GetAttributeByName("value", attr); - if (!attr_status.IsOK()) { + if (!attr_status.IsOK() || !attr) { return true; } diff --git a/musa/ep/src/kernels/activation/erf.cc b/musa/ep/src/kernels/activation/erf.cc index 4b3f1e6..1f60b49 100644 --- a/musa/ep/src/kernels/activation/erf.cc +++ b/musa/ep/src/kernels/activation/erf.cc @@ -34,6 +34,10 @@ OrtStatus* Erf::Compute(Ort::KernelContext& ctx) const { } } // namespace +ONNX_OPERATOR_VERSIONED_KERNEL_EX( + Erf, kOnnxDomain, 9, 12, + (Ort::KernelDefBuilder().AddTypeConstraint("T", HfdTensorTypes())), Erf) + ONNX_OPERATOR_VERSIONED_KERNEL_EX( Erf, kOnnxDomain, 13, 19, (Ort::KernelDefBuilder().AddTypeConstraint("T", FloatLikeTensorTypes())), diff --git a/musa/ep/src/kernels/activation/gelu.cc b/musa/ep/src/kernels/activation/gelu.cc new file mode 100644 index 0000000..a953b3c --- /dev/null +++ b/musa/ep/src/kernels/activation/gelu.cc @@ -0,0 +1,69 @@ +// Copyright (c) Moore Threads Technology Co., Ltd. All rights reserved. +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "activation/gelu_support.h" +#include "shared_inc/op_kernel_common.h" + +namespace { +class Gelu : public OpKernelBase { + public: + Gelu(const OrtKernelInfo* info, void* /*state*/) { + Ort::ConstKernelInfo kernel_info(info); + // The ONNX Gelu (opset 20) declares `approximate` as a string, the + // com.microsoft Gelu as an int (0 = exact, 1 = tanh); read whichever the + // schema of the matched domain carries. + try { + approximate_ = kernel_info.GetAttribute("approximate"); + } catch (...) { + try { + approximate_ = kernel_info.GetAttribute("approximate") == 0 + ? musa_ep::kGeluDefaultApproximate + : "tanh"; + } catch (...) { + approximate_ = musa_ep::kGeluDefaultApproximate; + } + } + } + OrtStatus* Compute(Ort::KernelContext& ctx) const; + + private: + std::string approximate_; +}; + +OrtStatus* Gelu::Compute(Ort::KernelContext& ctx) const { + if (const char* reason = musa_ep::GeluApproximateRejection(approximate_); + reason != nullptr) { + return Ort::GetApi().CreateStatus(ORT_NOT_IMPLEMENTED, reason); + } + auto info = ctx.GetInput(0).GetTensorTypeAndShapeInfo(); + auto elem_type = info.GetElementType(); + auto shape = info.GetShape(); + if (OutputEmptyTensorIfNeeded(ctx, shape)) { + return nullptr; + } + return UnaryDeviceCompute(ctx, shape, elem_type, MusaUnaryOp::Gelu, "Gelu"); +} +} // namespace + +// Gelu became an official ONNX operator at opset 20; the same kernel class also +// answers the com.microsoft domain, whose Gelu has the same semantics and the +// same attribute. +ONNX_OPERATOR_VERSIONED_KERNEL_EX( + Gelu, kOnnxDomain, 20, 21, + (Ort::KernelDefBuilder().AddTypeConstraint("T", FloatLikeTensorTypes())), + Gelu) + +ONNX_OPERATOR_VERSIONED_KERNEL_EX( + Gelu, kMSDomain, 1, 1, + (Ort::KernelDefBuilder().AddTypeConstraint("T", FloatLikeTensorTypes())), + Gelu) diff --git a/musa/ep/src/kernels/activation/gelu_support.h b/musa/ep/src/kernels/activation/gelu_support.h new file mode 100644 index 0000000..1774745 --- /dev/null +++ b/musa/ep/src/kernels/activation/gelu_support.h @@ -0,0 +1,35 @@ +// Copyright (c) Moore Threads Technology Co., Ltd. All rights reserved. +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#pragma once + +#include + +#include "onnxruntime_c_api.h" + +namespace musa_ep { + +inline constexpr char kGeluDefaultApproximate[] = "none"; + +/// The ONNX Gelu (opset 20) `approximate` attribute: this MUSA route +/// implements the exact erf form only. The graph-level capability predicate +/// and the kernel's runtime guard both ask this one function, so the nodes +/// this route claims are the nodes it runs. +inline const char* GeluApproximateRejection(const std::string& approximate) { + if (approximate == kGeluDefaultApproximate) { + return nullptr; + } + return "Gelu only supports approximate=none in this MUSA route"; +} + +} // namespace musa_ep diff --git a/musa/ep/src/kernels/math/cos.cc b/musa/ep/src/kernels/math/cos.cc new file mode 100644 index 0000000..0ae3e1b --- /dev/null +++ b/musa/ep/src/kernels/math/cos.cc @@ -0,0 +1,44 @@ +// Copyright (c) Moore Threads Technology Co., Ltd. All rights reserved. +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "shared_inc/op_kernel_common.h" + +namespace { +class Cos : public OpKernelBase { + public: + Cos(const OrtKernelInfo* /*info*/, void* /*state*/) {} + OrtStatus* Compute(Ort::KernelContext& ctx) const; +}; + +OrtStatus* Cos::Compute(Ort::KernelContext& ctx) const { + auto info = ctx.GetInput(0).GetTensorTypeAndShapeInfo(); + auto elem_type = info.GetElementType(); + auto shape = info.GetShape(); + if (OutputEmptyTensorIfNeeded(ctx, shape)) { + return nullptr; + } + return UnaryDeviceCompute(ctx, shape, elem_type, MusaUnaryOp::Cos, "Cos"); +} +} // namespace + +// The Sin/Cos schemas are unchanged from opset 7 through 19: only the element +// types grow at opset 13 (bfloat16), so one registration covers both the old +// and the current models. +ONNX_OPERATOR_VERSIONED_KERNEL_EX( + Cos, kOnnxDomain, 7, 12, + (Ort::KernelDefBuilder().AddTypeConstraint("T", HfdTensorTypes())), Cos) + +ONNX_OPERATOR_VERSIONED_KERNEL_EX( + Cos, kOnnxDomain, 13, 19, + (Ort::KernelDefBuilder().AddTypeConstraint("T", FloatLikeTensorTypes())), + Cos) diff --git a/musa/ep/src/kernels/math/sin.cc b/musa/ep/src/kernels/math/sin.cc new file mode 100644 index 0000000..2347d71 --- /dev/null +++ b/musa/ep/src/kernels/math/sin.cc @@ -0,0 +1,44 @@ +// Copyright (c) Moore Threads Technology Co., Ltd. All rights reserved. +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "shared_inc/op_kernel_common.h" + +namespace { +class Sin : public OpKernelBase { + public: + Sin(const OrtKernelInfo* /*info*/, void* /*state*/) {} + OrtStatus* Compute(Ort::KernelContext& ctx) const; +}; + +OrtStatus* Sin::Compute(Ort::KernelContext& ctx) const { + auto info = ctx.GetInput(0).GetTensorTypeAndShapeInfo(); + auto elem_type = info.GetElementType(); + auto shape = info.GetShape(); + if (OutputEmptyTensorIfNeeded(ctx, shape)) { + return nullptr; + } + return UnaryDeviceCompute(ctx, shape, elem_type, MusaUnaryOp::Sin, "Sin"); +} +} // namespace + +// The Sin/Cos schemas are unchanged from opset 7 through 19: only the element +// types grow at opset 13 (bfloat16), so one registration covers both the old +// and the current models. +ONNX_OPERATOR_VERSIONED_KERNEL_EX( + Sin, kOnnxDomain, 7, 12, + (Ort::KernelDefBuilder().AddTypeConstraint("T", HfdTensorTypes())), Sin) + +ONNX_OPERATOR_VERSIONED_KERNEL_EX( + Sin, kOnnxDomain, 13, 19, + (Ort::KernelDefBuilder().AddTypeConstraint("T", FloatLikeTensorTypes())), + Sin) diff --git a/musa/ep/src/kernels/math/unary_elementwise_ops_impl.mu b/musa/ep/src/kernels/math/unary_elementwise_ops_impl.mu index 810729e..b531f66 100644 --- a/musa/ep/src/kernels/math/unary_elementwise_ops_impl.mu +++ b/musa/ep/src/kernels/math/unary_elementwise_ops_impl.mu @@ -39,6 +39,12 @@ __device__ __forceinline__ float UnaryValue(float x, MusaUnaryOp op, float alpha return ceilf(x); case MusaUnaryOp::Floor: return floorf(x); + case MusaUnaryOp::Sin: + return sinf(x); + case MusaUnaryOp::Cos: + return cosf(x); + case MusaUnaryOp::Gelu: + return 0.5f * x * (1.0f + erff(x * 0.70710678118654752440f)); } return x; } @@ -86,6 +92,15 @@ __device__ __forceinline__ T UnaryValueTyped(T x, MusaUnaryOp op, float alpha) { return static_cast(ceil(static_cast(x))); case MusaUnaryOp::Floor: return static_cast(floor(static_cast(x))); + case MusaUnaryOp::Sin: + return static_cast(sin(static_cast(x))); + case MusaUnaryOp::Cos: + return static_cast(cos(static_cast(x))); + case MusaUnaryOp::Gelu: { + const double value = static_cast(x); + return static_cast(0.5 * value * + (1.0 + erf(value * 0.70710678118654752440))); + } } return x; } @@ -158,6 +173,21 @@ musaError_t LaunchUnaryTyped(const void* input, return musaGetLastError(); } +musaError_t LaunchUnaryFloat(const void* input, + void* output, + int64_t count, + MusaUnaryOp op, + float alpha, + musaStream_t stream) { + if (count == 0) { + return musaSuccess; + } + UnaryFloatKernel<<>>( + reinterpret_cast(input), reinterpret_cast(output), + count, op, alpha); + return musaGetLastError(); +} + template musaError_t LaunchUnaryFloatLikeTyped(const void* input, void* output, @@ -183,7 +213,7 @@ musaError_t LaunchMusaUnaryKernel(const void* input, musaStream_t stream) { switch (elem_type) { case MusaElementType::Float: - return LaunchUnaryTyped(input, output, count, op, alpha, stream); + return LaunchUnaryFloat(input, output, count, op, alpha, stream); case MusaElementType::Double: return LaunchUnaryTyped(input, output, count, op, alpha, stream); case MusaElementType::Uint8: @@ -219,8 +249,7 @@ musaError_t LaunchMusaUnaryFloatKernel(const float* input, MusaUnaryOp op, float alpha, musaStream_t stream) { - return LaunchMusaUnaryKernel(input, output, count, op, alpha, - MusaElementType::Float, stream); + return LaunchUnaryFloat(input, output, count, op, alpha, stream); } template diff --git a/musa/ep/src/kernels/nn/average_pool_support.h b/musa/ep/src/kernels/nn/average_pool_support.h new file mode 100644 index 0000000..7670d15 --- /dev/null +++ b/musa/ep/src/kernels/nn/average_pool_support.h @@ -0,0 +1,156 @@ +// Copyright (c) Moore Threads Technology Co., Ltd. All rights reserved. +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#pragma once + +#include +#include +#include +#include + +#include "onnxruntime_c_api.h" + +namespace musa_ep { + +// Mirrors kMusaMaxBroadcastRank in ../shared_inc/device_kernel_types.h; kept +// local so this header (also included by the graph-level capability predicate) +// depends on the public C API only. +inline constexpr int64_t kAveragePoolMaxRank = 8; + +inline constexpr char kAveragePoolDefaultAutoPad[] = "NOTSET"; + +/// Everything about an AveragePool node that decides whether this MUSA route +/// can run it. +/// +/// The kernel registry only matches op/domain/opset/dtype, so the graph-level +/// capability predicate (graph/capability_predicates.cc) inspects the node +/// itself to keep the AveragePool configurations this route cannot run on the +/// CPU EP. Both it and the kernel's runtime guard build this struct and ask +/// AveragePoolRouteRejection() below, so the nodes this route claims and the +/// nodes it runs are decided by one definition instead of two copies that +/// drift apart. +struct AveragePoolRouteConfig { + std::string auto_pad = kAveragePoolDefaultAutoPad; + int64_t ceil_mode = 0; + int64_t count_include_pad = 0; + std::vector kernel_shape; + std::vector pads; + std::vector strides; + /// Rank of the X input; -1 when it is unknown or not a tensor. + int64_t input_rank = -1; + ONNXTensorElementDataType input_element_type = + ONNX_TENSOR_ELEMENT_DATA_TYPE_UNDEFINED; +}; + +/// Returns nullptr when this MUSA route runs \p config, otherwise a static +/// reason string suitable for an ORT_NOT_IMPLEMENTED status. +inline const char* AveragePoolRouteRejection( + const AveragePoolRouteConfig& config) { + if (config.auto_pad != "NOTSET" && config.auto_pad != "VALID" && + config.auto_pad != "SAME_UPPER" && config.auto_pad != "SAME_LOWER") { + return "AveragePool only supports auto_pad NOTSET, VALID, SAME_UPPER or " + "SAME_LOWER in this MUSA route"; + } + // ceil_mode=1 changes the output size and clips the last window, which the + // pooling routes here do not model; those nodes stay on the CPU EP. + if (config.ceil_mode != 0) { + return "AveragePool only supports ceil_mode=0 in this MUSA route"; + } + if (config.count_include_pad != 0 && config.count_include_pad != 1) { + return "AveragePool requires count_include_pad to be 0 or 1"; + } + if (config.input_element_type != ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT && + config.input_element_type != ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT16 && + config.input_element_type != ONNX_TENSOR_ELEMENT_DATA_TYPE_DOUBLE) { + return "AveragePool only supports float16, float32 or float64"; + } + if (config.input_rank < 3 || config.input_rank > kAveragePoolMaxRank) { + return "AveragePool only supports rank 3 to 8 tensors in this MUSA route"; + } + const auto spatial_rank = static_cast(config.input_rank) - 2; + if (config.kernel_shape.size() != spatial_rank) { + return "AveragePool requires kernel_shape with one entry per spatial " + "dimension"; + } + if (!config.strides.empty() && config.strides.size() != spatial_rank) { + return "AveragePool requires strides with one entry per spatial dimension"; + } + // The auto_pad modes derive the pads from the input shape, so a pads + // attribute next to them is ignored (ORT does the same); only auto_pad + // NOTSET reads the attribute. + if (config.auto_pad == kAveragePoolDefaultAutoPad && !config.pads.empty() && + config.pads.size() != spatial_rank * 2) { + return "AveragePool requires pads with two entries per spatial dimension"; + } + for (int64_t value : config.kernel_shape) { + if (value <= 0) { + return "AveragePool requires a positive kernel_shape"; + } + } + for (int64_t value : config.strides) { + if (value <= 0) { + return "AveragePool requires positive strides"; + } + } + if (config.auto_pad == kAveragePoolDefaultAutoPad) { + for (int64_t value : config.pads) { + if (value < 0) { + return "AveragePool requires non-negative pads"; + } + } + } + return nullptr; +} + +/// Pads this route runs with: the attribute for auto_pad NOTSET, zeros for +/// VALID, and the pads ONNX derives from the input size for the SAME_* modes +/// (the output size is ceil(input / stride), so the total pad is +/// (output - 1) * stride + kernel - input). Returns false when a spatial +/// dimension is unknown, which only the auto_pad modes need. +inline bool ResolveAveragePoolPads(std::string_view auto_pad, + const std::vector& pads_attr, + const std::vector& input_shape, + const std::vector& kernel_shape, + const std::vector& strides, + std::vector& pads) { + const size_t spatial_rank = kernel_shape.size(); + pads.assign(spatial_rank * 2, 0); + if (auto_pad == kAveragePoolDefaultAutoPad) { + if (pads_attr.size() == spatial_rank * 2) { + pads = pads_attr; + } + return true; + } + if (auto_pad == "VALID") { + return true; + } + if (input_shape.size() != spatial_rank + 2) { + return false; + } + for (size_t i = 0; i < spatial_rank; ++i) { + const int64_t input_dim = input_shape[i + 2]; + if (input_dim <= 0) { + return false; + } + const int64_t output_dim = (input_dim + strides[i] - 1) / strides[i]; + const int64_t total_pad = std::max( + 0, (output_dim - 1) * strides[i] + kernel_shape[i] - input_dim); + const int64_t pad_begin = + auto_pad == "SAME_LOWER" ? (total_pad + 1) / 2 : total_pad / 2; + pads[i] = pad_begin; + pads[i + spatial_rank] = total_pad - pad_begin; + } + return true; +} + +} // namespace musa_ep diff --git a/musa/ep/src/kernels/nn/layer_norm.cc b/musa/ep/src/kernels/nn/layer_norm.cc index fe2d874..4b0fd83 100644 --- a/musa/ep/src/kernels/nn/layer_norm.cc +++ b/musa/ep/src/kernels/nn/layer_norm.cc @@ -204,30 +204,34 @@ OrtStatus* LayerNormalization::Compute(Ort::KernelContext& ctx) const { } std::vector normalized_shape(input_shape.begin() + axis, input_shape.end()); + // ORT binds Scale/Bias to the same type parameter as X, so a model that ORT + // accepts carries one element type on all three inputs. muDNN reads + // Scale/Bias as float32, the device kernel reads them as the input type. + const auto scale_elem_type = scale_info.GetElementType(); if (!ShapeEqualsSuffix(scale_shape, normalized_shape) || - scale_info.GetElementType() != ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT) { + scale_elem_type != elem_type) { return Ort::GetApi().CreateStatus( ORT_NOT_IMPLEMENTED, - "LayerNormalization requires float scale matching normalized shape"); + "LayerNormalization requires a scale matching the input type and the " + "normalized shape"); } Ort::ConstValue bias = (ctx.GetInputCount() > 2) ? ctx.GetInput(2) : Ort::ConstValue{nullptr}; if (bias != nullptr) { auto bias_info = bias.GetTensorTypeAndShapeInfo(); if (!ShapeEqualsSuffix(bias_info.GetShape(), normalized_shape) || - bias_info.GetElementType() != ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT) { + bias_info.GetElementType() != elem_type) { return Ort::GetApi().CreateStatus( ORT_NOT_IMPLEMENTED, - "LayerNormalization requires float bias matching normalized shape"); + "LayerNormalization requires a bias matching the input type and the " + "normalized shape"); } } MusaElementType musa_elem_type; if (!ToMusaElementType(elem_type, musa_elem_type) || (elem_type != ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT && - elem_type != ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT16 && - elem_type != ONNX_TENSOR_ELEMENT_DATA_TYPE_BFLOAT16 && - elem_type != ONNX_TENSOR_ELEMENT_DATA_TYPE_DOUBLE)) { + elem_type != ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT16)) { return Ort::GetApi().CreateStatus(ORT_NOT_IMPLEMENTED, "LayerNormalization unsupported dtype"); } @@ -270,7 +274,8 @@ OrtStatus* LayerNormalization::Compute(Ort::KernelContext& ctx) const { "LayerNormalization requires MUSA device optional outputs"); } - if (TryMudnn(ctx, input_shape, scale_shape, mean_shape, elem_type, axis, + if (elem_type == ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT && + TryMudnn(ctx, input_shape, scale_shape, mean_shape, elem_type, axis, output, mean, inv_std)) { return nullptr; } @@ -279,13 +284,18 @@ OrtStatus* LayerNormalization::Compute(Ort::KernelContext& ctx) const { params.rows = rows; params.norm_size = norm_size; params.has_bias = bias != nullptr ? 1 : 0; + MusaElementType scale_musa_type = MusaElementType::Float; + if (!ToMusaElementType(scale_elem_type, scale_musa_type)) { + return Ort::GetApi().CreateStatus(ORT_NOT_IMPLEMENTED, + "LayerNormalization unsupported dtype"); + } musaError_t status = LaunchMusaLayerNormKernel( - input.GetTensorRawData(), scale.GetTensorData(), - bias == nullptr ? nullptr : bias.GetTensorData(), + input.GetTensorRawData(), scale.GetTensorRawData(), + bias == nullptr ? nullptr : bias.GetTensorRawData(), output.GetTensorMutableRawData(), mean == nullptr ? nullptr : mean.GetTensorMutableData(), inv_std == nullptr ? nullptr : inv_std.GetTensorMutableData(), - params, epsilon_, musa_elem_type, GetComputeStream(ctx)); + params, epsilon_, musa_elem_type, scale_musa_type, GetComputeStream(ctx)); if (status == musaErrorNotSupported) { return Ort::GetApi().CreateStatus( ORT_NOT_IMPLEMENTED, @@ -297,5 +307,5 @@ OrtStatus* LayerNormalization::Compute(Ort::KernelContext& ctx) const { ONNX_OPERATOR_VERSIONED_KERNEL_EX( LayerNormalization, kOnnxDomain, 17, 19, - (Ort::KernelDefBuilder().AddTypeConstraint("T", FloatTensorTypes())), + (Ort::KernelDefBuilder().AddTypeConstraint("T", LayerNormTensorTypes())), LayerNormalization) diff --git a/musa/ep/src/kernels/nn/layer_norm_impl.h b/musa/ep/src/kernels/nn/layer_norm_impl.h index 1d8b67a..f769641 100644 --- a/musa/ep/src/kernels/nn/layer_norm_impl.h +++ b/musa/ep/src/kernels/nn/layer_norm_impl.h @@ -2,9 +2,13 @@ #include "shared_inc/device_kernel_types.h" -musaError_t LaunchMusaLayerNormKernel(const void* input, const float* scale, - const float* bias, void* output, +// `scale`/`bias` carry `scale_elem_type` elements (float32 or float16); the +// device kernel reads them through that type, so no host-side conversion is +// needed. +musaError_t LaunchMusaLayerNormKernel(const void* input, const void* scale, + const void* bias, void* output, float* mean, float* inv_std, MusaLayerNormParams params, float epsilon, MusaElementType elem_type, + MusaElementType scale_elem_type, musaStream_t stream); diff --git a/musa/ep/src/kernels/nn/layer_norm_impl.mu b/musa/ep/src/kernels/nn/layer_norm_impl.mu index 164db68..4e80c2f 100644 --- a/musa/ep/src/kernels/nn/layer_norm_impl.mu +++ b/musa/ep/src/kernels/nn/layer_norm_impl.mu @@ -3,10 +3,10 @@ namespace { -template +template __global__ void LayerNormKernel(const T* input, - const float* scale, - const float* bias, + const ScaleT* scale, + const ScaleT* bias, T* output, float* mean, float* inv_std, @@ -63,18 +63,18 @@ __global__ void LayerNormKernel(const T* input, const float normalized = (MusaScalarToFloat(input[row * params.norm_size + col]) - row_mean) * row_inv_std; - float value = normalized * scale[col]; + float value = normalized * MusaScalarToFloat(scale[col]); if (params.has_bias) { - value += bias[col]; + value += MusaScalarToFloat(bias[col]); } output[row * params.norm_size + col] = MusaScalarFromFloat(value); } } -template +template musaError_t LaunchLayerNormTyped(const void* input, - const float* scale, - const float* bias, + const void* scale, + const void* bias, void* output, float* mean, float* inv_std, @@ -87,38 +87,54 @@ musaError_t LaunchLayerNormTyped(const void* input, if (params.norm_size > INT32_MAX) { return musaErrorNotSupported; } - LayerNormKernel<<(params.rows), kThreadsPerBlock, 0, - stream>>>(reinterpret_cast(input), scale, bias, - reinterpret_cast(output), mean, inv_std, - params, epsilon); + LayerNormKernel + <<(params.rows), kThreadsPerBlock, 0, stream>>>( + reinterpret_cast(input), + reinterpret_cast(scale), + reinterpret_cast(bias), + reinterpret_cast(output), mean, inv_std, params, epsilon); return musaGetLastError(); } +template +musaError_t DispatchLayerNorm(const void* input, const void* scale, + const void* bias, void* output, float* mean, + float* inv_std, MusaLayerNormParams params, + float epsilon, MusaElementType scale_elem_type, + musaStream_t stream) { + switch (scale_elem_type) { + case MusaElementType::Float: + return LaunchLayerNormTyped(input, scale, bias, output, mean, + inv_std, params, epsilon, stream); + case MusaElementType::Float16: + return LaunchLayerNormTyped(input, scale, bias, output, mean, + inv_std, params, epsilon, stream); + default: + return musaErrorNotSupported; + } +} + } // namespace musaError_t LaunchMusaLayerNormKernel(const void* input, - const float* scale, - const float* bias, + const void* scale, + const void* bias, void* output, float* mean, float* inv_std, MusaLayerNormParams params, float epsilon, MusaElementType elem_type, + MusaElementType scale_elem_type, musaStream_t stream) { switch (elem_type) { case MusaElementType::Float: - return LaunchLayerNormTyped(input, scale, bias, output, mean, - inv_std, params, epsilon, stream); + return DispatchLayerNorm(input, scale, bias, output, mean, inv_std, + params, epsilon, scale_elem_type, stream); case MusaElementType::Float16: - return LaunchLayerNormTyped<__half>(input, scale, bias, output, mean, - inv_std, params, epsilon, stream); - case MusaElementType::BFloat16: - return LaunchLayerNormTyped<__mt_bfloat16>( - input, scale, bias, output, mean, inv_std, params, epsilon, stream); - case MusaElementType::Double: - return LaunchLayerNormTyped(input, scale, bias, output, mean, - inv_std, params, epsilon, stream); + return DispatchLayerNorm<__half>(input, scale, bias, output, mean, + inv_std, params, epsilon, + scale_elem_type, stream); default: return musaErrorNotSupported; } diff --git a/musa/ep/src/kernels/nn/pool.cc b/musa/ep/src/kernels/nn/pool.cc index 5922e19..ddc6293 100644 --- a/musa/ep/src/kernels/nn/pool.cc +++ b/musa/ep/src/kernels/nn/pool.cc @@ -11,6 +11,7 @@ // See the License for the specific language governing permissions and // limitations under the License. +#include "nn/average_pool_support.h" #include "nn/pool_impl.h" #include "shared_inc/blas_utils.h" #include "shared_inc/op_kernel_common.h" @@ -173,6 +174,28 @@ bool TryMudnnMaxPool(const void* input_data, void* output_data, ::musa::dnn::Pooling::Mode::MAXPOOL, stream); } +// muDNN expresses ONNX count_include_pad through two average modes: the one +// that divides by the kernel size, and the one that divides by the positions +// that land inside the input. +bool TryMudnnAveragePool(const void* input_data, void* output_data, + const std::vector& input_shape, + const std::vector& output_shape, + const std::vector& kernel_shape, + const std::vector& pads, + const std::vector& strides, + bool count_include_pad, + ONNXTensorElementDataType elem_type, + musaStream_t stream) { + // AveragePool has no dilations attribute; the kernel window is dense. + const std::vector dilations(input_shape.size() - 2, 1); + return TryMudnnPool( + input_data, output_data, input_shape, output_shape, kernel_shape, pads, + strides, dilations, elem_type, + count_include_pad ? ::musa::dnn::Pooling::Mode::AVGPOOL_COUNT_PAD + : ::musa::dnn::Pooling::Mode::AVGPOOL_COUNT_WITHOUT_PAD, + stream); +} + bool HfdType(ONNXTensorElementDataType elem_type) { return elem_type == ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT || elem_type == ONNX_TENSOR_ELEMENT_DATA_TYPE_DOUBLE || @@ -282,6 +305,33 @@ MusaMaxPoolParams MakeMaxPoolParams(const std::vector& input_shape, return params; } +MusaAveragePoolParams MakeAveragePoolParams( + const std::vector& input_shape, + const std::vector& output_shape, + const std::vector& kernel_shape, const std::vector& pads, + const std::vector& strides, bool count_include_pad) { + MusaAveragePoolParams params{}; + params.rank = static_cast(input_shape.size()); + params.spatial_rank = static_cast(input_shape.size() - 2); + params.count_include_pad = count_include_pad ? 1 : 0; + params.output_elements = NumElements(output_shape); + + auto input_strides = Strides(input_shape); + auto output_strides = Strides(output_shape); + for (size_t i = 0; i < input_shape.size(); ++i) { + params.input_dims[i] = input_shape[i]; + params.input_strides[i] = input_strides[i]; + params.output_strides[i] = output_strides[i]; + } + for (int32_t i = 0; i < params.spatial_rank; ++i) { + const auto idx = static_cast(i); + params.kernel_shape[i] = kernel_shape[idx]; + params.pads_begin[i] = pads[idx]; + params.strides[i] = strides[idx]; + } + return params; +} + class GlobalAveragePool : public OpKernelBase { public: GlobalAveragePool(const OrtKernelInfo* /*info*/, void* /*state*/) {} @@ -509,6 +559,96 @@ OrtStatus* MaxPool::Compute(Ort::KernelContext& ctx) const { return LaunchStatus(status); } +class AveragePool : public OpKernelBase { + public: + AveragePool(const OrtKernelInfo* info, void* /*state*/) { + Ort::ConstKernelInfo kernel_info(info); + auto_pad_ = AttrOrDefault(kernel_info, "auto_pad", "NOTSET"); + ceil_mode_ = AttrOrDefault(kernel_info, "ceil_mode", 0); + count_include_pad_ = + AttrOrDefault(kernel_info, "count_include_pad", 0); + kernel_shape_ = AttrsOrEmpty(kernel_info, "kernel_shape"); + pads_ = AttrsOrEmpty(kernel_info, "pads"); + strides_ = AttrsOrEmpty(kernel_info, "strides"); + } + OrtStatus* Compute(Ort::KernelContext& ctx) const; + + private: + std::string auto_pad_; + int64_t ceil_mode_ = 0; + int64_t count_include_pad_ = 0; + std::vector kernel_shape_; + std::vector pads_; + std::vector strides_; +}; + +OrtStatus* AveragePool::Compute(Ort::KernelContext& ctx) const { + musaStream_t stream = GetComputeStream(ctx); + Ort::ConstValue input = ctx.GetInput(0); + auto input_info = input.GetTensorTypeAndShapeInfo(); + const auto elem_type = input_info.GetElementType(); + auto input_shape = input_info.GetShape(); + const size_t spatial_rank = + input_shape.size() >= 2 ? input_shape.size() - 2 : 0; + + // The capability predicate and this guard ask the same question of the same + // struct, so the nodes this route claims are exactly the nodes it runs. + musa_ep::AveragePoolRouteConfig config; + config.auto_pad = auto_pad_; + config.ceil_mode = ceil_mode_; + config.count_include_pad = count_include_pad_; + config.kernel_shape = kernel_shape_; + config.pads = NormalizePads(pads_, spatial_rank); + config.strides = NormalizeSpatialAttr(strides_, spatial_rank, 1); + config.input_rank = static_cast(input_shape.size()); + config.input_element_type = elem_type; + if (const char* reason = musa_ep::AveragePoolRouteRejection(config); + reason != nullptr) { + return PoolStatus("AveragePool", reason); + } + + const std::vector dilations(spatial_rank, 1); + std::vector pads; + if (!musa_ep::ResolveAveragePoolPads(auto_pad_, pads_, input_shape, + kernel_shape_, config.strides, pads)) { + return PoolStatus("AveragePool", + "auto_pad requires a concrete input shape"); + } + RETURN_IF_ERROR(ValidatePoolCommon("AveragePool", input_shape, kernel_shape_, + pads, config.strides, dilations)); + + MusaElementType musa_elem_type; + if (!ToMusaElementType(elem_type, musa_elem_type)) { + return PoolStatus("AveragePool", "unsupported dtype"); + } + + std::vector output_shape = PoolOutputShape( + input_shape, kernel_shape_, pads, config.strides, dilations); + Ort::UnownedValue output = ctx.GetOutput(0, output_shape); + if (!IsGpuMemory(input.GetTensorMemoryInfo()) || + !IsGpuMemory(output.GetTensorMemoryInfo())) { + return PoolStatus("AveragePool", "requires MUSA device tensors"); + } + + if (HfdType(elem_type) && + TryMudnnAveragePool(input.GetTensorRawData(), + output.GetTensorMutableRawData(), input_shape, + output_shape, kernel_shape_, pads, config.strides, + count_include_pad_ != 0, elem_type, stream)) { + return nullptr; + } + + musaError_t status = LaunchMusaAveragePoolKernel( + input.GetTensorRawData(), output.GetTensorMutableRawData(), + MakeAveragePoolParams(input_shape, output_shape, kernel_shape_, pads, + config.strides, count_include_pad_ != 0), + musa_elem_type, stream); + if (status == musaErrorNotSupported) { + return PoolStatus("AveragePool", "unsupported dtype"); + } + return LaunchStatus(status); +} + } // namespace ONNX_OPERATOR_VERSIONED_KERNEL_EX( @@ -516,6 +656,24 @@ ONNX_OPERATOR_VERSIONED_KERNEL_EX( (Ort::KernelDefBuilder().AddTypeConstraint("T", HfdTensorTypes())), GlobalAveragePool) +// The AveragePool schema changes at 10 (ceil_mode is documented) and 19 (only +// the default values of strides/dilations are clarified), so the same kernel +// class covers every range, exactly as the CUDA EP registers it. +ONNX_OPERATOR_VERSIONED_KERNEL_EX( + AveragePool, kOnnxDomain, 7, 9, + (Ort::KernelDefBuilder().AddTypeConstraint("T", HfdTensorTypes())), + AveragePool) + +ONNX_OPERATOR_VERSIONED_KERNEL_EX( + AveragePool, kOnnxDomain, 10, 10, + (Ort::KernelDefBuilder().AddTypeConstraint("T", HfdTensorTypes())), + AveragePool) + +ONNX_OPERATOR_VERSIONED_KERNEL_EX( + AveragePool, kOnnxDomain, 11, 19, + (Ort::KernelDefBuilder().AddTypeConstraint("T", HfdTensorTypes())), + AveragePool) + ONNX_OPERATOR_VERSIONED_KERNEL_EX( GlobalMaxPool, kOnnxDomain, 1, 19, (Ort::KernelDefBuilder().AddTypeConstraint("T", HfdTensorTypes())), diff --git a/musa/ep/src/kernels/nn/pool_impl.h b/musa/ep/src/kernels/nn/pool_impl.h index 02d7db4..fa2d62a 100644 --- a/musa/ep/src/kernels/nn/pool_impl.h +++ b/musa/ep/src/kernels/nn/pool_impl.h @@ -10,3 +10,8 @@ musaError_t LaunchMusaMaxPoolKernel(const void* input, void* output, int64_t* indices, MusaMaxPoolParams params, MusaElementType elem_type, musaStream_t stream); + +musaError_t LaunchMusaAveragePoolKernel(const void* input, void* output, + MusaAveragePoolParams params, + MusaElementType elem_type, + musaStream_t stream); diff --git a/musa/ep/src/kernels/nn/pool_impl.mu b/musa/ep/src/kernels/nn/pool_impl.mu index e54cebf..752c52d 100644 --- a/musa/ep/src/kernels/nn/pool_impl.mu +++ b/musa/ep/src/kernels/nn/pool_impl.mu @@ -130,6 +130,86 @@ musaError_t LaunchMaxPoolTyped(const void* input, void* output, return musaGetLastError(); } +// Sums the window and divides by the number of elements it averages: the +// kernel window itself when count_include_pad is set (padding contributes 0), +// otherwise the number of window positions that land inside the input. +template +__global__ void AveragePoolKernel(const T* input, T* output, + MusaAveragePoolParams params) { + const int64_t thread_id = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + const int64_t total_threads = static_cast(gridDim.x) * blockDim.x; + for (int64_t output_index = thread_id; output_index < params.output_elements; + output_index += total_threads) { + int64_t remaining = output_index; + int64_t output_coord[kMusaMaxBroadcastRank]; + for (int32_t dim = 0; dim < params.rank; ++dim) { + const int64_t stride = params.output_strides[dim]; + output_coord[dim] = stride == 0 ? 0 : remaining / stride; + remaining -= output_coord[dim] * stride; + } + + int64_t base_input_offset = 0; + for (int32_t dim = 0; dim < 2; ++dim) { + base_input_offset += output_coord[dim] * params.input_strides[dim]; + } + + const int32_t spatial_rank = params.spatial_rank; + int64_t kernel_elements = 1; + for (int32_t dim = 0; dim < spatial_rank; ++dim) { + kernel_elements *= params.kernel_shape[dim]; + } + + AccT sum = static_cast(0); + int64_t in_input_count = 0; + for (int64_t kernel_index = 0; kernel_index < kernel_elements; + ++kernel_index) { + int64_t kernel_remaining = kernel_index; + int64_t input_offset = base_input_offset; + bool in_input = true; + for (int32_t dim = spatial_rank - 1; dim >= 0; --dim) { + const int64_t kernel_dim = params.kernel_shape[dim]; + const int64_t k_coord = kernel_remaining % kernel_dim; + kernel_remaining /= kernel_dim; + const int32_t full_dim = dim + 2; + const int64_t in_coord = output_coord[full_dim] * params.strides[dim] - + params.pads_begin[dim] + k_coord; + if (in_coord < 0 || in_coord >= params.input_dims[full_dim]) { + in_input = false; + break; + } + input_offset += in_coord * params.input_strides[full_dim]; + } + if (!in_input) { + continue; + } + sum += static_cast(MusaScalarToDouble(input[input_offset])); + ++in_input_count; + } + + const int64_t divisor = + params.count_include_pad != 0 ? kernel_elements : in_input_count; + const AccT averaged = + divisor == 0 ? static_cast(0) + : sum / static_cast(divisor); + output[output_index] = MusaScalarFromDouble(static_cast(averaged)); + } +} + +template +musaError_t LaunchAveragePoolTyped(const void* input, void* output, + MusaAveragePoolParams params, + musaStream_t stream) { + if (params.output_elements == 0) { + return musaSuccess; + } + AveragePoolKernel + <<>>( + reinterpret_cast(input), reinterpret_cast(output), + params); + return musaGetLastError(); +} + } // namespace musaError_t LaunchMusaGlobalAveragePoolKernel( @@ -170,3 +250,22 @@ musaError_t LaunchMusaMaxPoolKernel(const void* input, void* output, return musaErrorNotSupported; } } + +musaError_t LaunchMusaAveragePoolKernel(const void* input, void* output, + MusaAveragePoolParams params, + MusaElementType elem_type, + musaStream_t stream) { + switch (elem_type) { + case MusaElementType::Float: + return LaunchAveragePoolTyped(input, output, params, + stream); + case MusaElementType::Double: + return LaunchAveragePoolTyped(input, output, params, + stream); + case MusaElementType::Float16: + return LaunchAveragePoolTyped<__half, float>(input, output, params, + stream); + default: + return musaErrorNotSupported; + } +} diff --git a/musa/ep/src/kernels/shared_inc/device_kernel_types.h b/musa/ep/src/kernels/shared_inc/device_kernel_types.h index 2138dcc..12ae736 100644 --- a/musa/ep/src/kernels/shared_inc/device_kernel_types.h +++ b/musa/ep/src/kernels/shared_inc/device_kernel_types.h @@ -34,6 +34,9 @@ enum class MusaUnaryOp : int32_t { Softplus = 14, Ceil = 15, Floor = 16, + Sin = 17, + Cos = 18, + Gelu = 19, }; enum class MusaCompareOp : int32_t { @@ -108,6 +111,21 @@ struct MusaMaxPoolParams { int64_t dilations[kMusaMaxBroadcastRank]; }; +// AveragePool has no dilations, and needs to know whether the padding counts +// towards the divisor (ONNX count_include_pad). +struct MusaAveragePoolParams { + int32_t rank; + int32_t spatial_rank; + int32_t count_include_pad; + int64_t output_elements; + int64_t input_dims[kMusaMaxBroadcastRank]; + int64_t input_strides[kMusaMaxBroadcastRank]; + int64_t output_strides[kMusaMaxBroadcastRank]; + int64_t kernel_shape[kMusaMaxBroadcastRank]; + int64_t pads_begin[kMusaMaxBroadcastRank]; + int64_t strides[kMusaMaxBroadcastRank]; +}; + struct MusaRangeParams { int64_t count; }; diff --git a/musa/ep/src/kernels/shared_inc/kernel_type_constraints.h b/musa/ep/src/kernels/shared_inc/kernel_type_constraints.h index 891a5bf..c423c90 100644 --- a/musa/ep/src/kernels/shared_inc/kernel_type_constraints.h +++ b/musa/ep/src/kernels/shared_inc/kernel_type_constraints.h @@ -208,6 +208,18 @@ inline std::vector FloatTensorTypes() { return {GetTensorType(ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT)}; } +// LayerNormalization (opset 17): the kernel runs X, and reads Scale/Bias, in +// float32 or float16. ORT binds Scale/Bias to the same "T" as X even though +// the ONNX schema names that constraint "U" (and only allows float/bfloat16 +// there), so a model ORT loads carries the same element type on all three +// inputs. +inline std::vector LayerNormTensorTypes() { + return { + GetTensorType(ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT), + GetTensorType(ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT16), + }; +} + inline std::vector ClipTensorTypes() { return { GetTensorType(ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT16), diff --git a/test/ops/test_attribute_defaults.py b/test/ops/test_attribute_defaults.py new file mode 100644 index 0000000..5380516 --- /dev/null +++ b/test/ops/test_attribute_defaults.py @@ -0,0 +1,137 @@ +# Copyright (c) Moore Threads Technology Co., Ltd. All rights reserved. +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Optional attributes that the graph layer has to read as "not present". + +`OrtApi::Node_GetAttributeByName` reports an attribute it cannot provide - an +optional attribute with no constant default value, such as Transpose's `perm` +or ReduceMean's `axes` - as a success status with an *empty* handle. Every +attribute reader in `graph/graph_utils.cc` (and the copies of it under +`fusion/`) used to call into that handle anyway, which dereferences null inside +`OpAttr_GetType` and takes the whole process down while the EP is answering +GetCapability. + +Each graph below is built so that a fusion matcher really does reach one of +those reads, which is what makes the case discriminating: before the readers +checked the handle, building the session crashed; afterwards the matcher sees +"attribute absent", declines the fusion, and the graph runs as plain ONNX ops +on the MUSA EP with the same numbers as the CPU EP. +""" + +from __future__ import annotations + +import numpy as np +import onnxruntime as ort +import pytest +from onnx import TensorProto, helper + +from op_test_utils import build_graph_model, musa_devices, run + + +def _require_devices(): + devices = musa_devices() + if not devices: + raise RuntimeError("No MUSA device available for the attribute default tests") + return devices + + +def _split_concat_transpose_model(): + """Split -> Concat -> Transpose with no `perm` on the Transpose. + + `FindSplitConcatFusions` walks exactly this shape: a rank-3 input split + along axis 2 into equal parts, concatenated along axis 0, then a Transpose + whose `perm` it compares with {0, 2, 1} - the read that used to crash. With + no `perm` the Transpose means "reverse the dimensions". + """ + x = np.arange(2 * 3 * 4, dtype=np.float32).reshape(2, 3, 4) + model_bytes = build_graph_model( + [ + helper.make_node("Split", ["X", "split_sizes"], ["s0", "s1"], axis=2), + helper.make_node("Concat", ["s0", "s1"], ["concat_out"], axis=0), + helper.make_node("Transpose", ["concat_out"], ["Y"]), + ], + {"X": x}, + [("Y", TensorProto.FLOAT)], + initializers=[ + helper.make_tensor("split_sizes", TensorProto.INT64, [2], [2, 2]), + ], + opset=17, + name="split_concat_transpose_without_perm", + ) + return model_bytes, {"X": x} + + +def _rms_norm_without_axes_model(): + """The RMSNorm decomposition with no `axes` on the ReduceMean. + + `FindRmsNormFusions` / `FindCastRmsNormFusions` walk + Mul -> ReduceMean -> Add -> Sqrt -> Div -> Mul and read the reduce node's + `axes` to check that it reduces the last dimension - the read that used to + crash. A ReduceMean without `axes` reduces *every* dimension, which is not + what those fusions describe, so the matchers must decline; the graph then + runs as written. Every leading dimension is 1 so that the reduce output + keeps the shape the matcher checks before it reads the attribute. + """ + x = np.arange(1 * 1 * 4, dtype=np.float32).reshape(1, 1, 4) / 4.0 + model_bytes = build_graph_model( + [ + helper.make_node("Mul", ["X", "X"], ["square"]), + helper.make_node("ReduceMean", ["square"], ["mean"]), + helper.make_node("Add", ["mean", "epsilon"], ["var"]), + helper.make_node("Sqrt", ["var"], ["std"]), + helper.make_node("Div", ["X", "std"], ["normed"]), + helper.make_node("Mul", ["normed", "gamma"], ["Y"]), + ], + {"X": x}, + [("Y", TensorProto.FLOAT)], + initializers=[ + helper.make_tensor("epsilon", TensorProto.FLOAT, [], [1e-5]), + helper.make_tensor("gamma", TensorProto.FLOAT, [4], [1.0, 1.0, 1.0, 1.0]), + ], + opset=17, + name="rms_norm_without_axes", + ) + return model_bytes, {"X": x} + + +@pytest.mark.parametrize( + "builder,label", + [(_split_concat_transpose_model, "split-concat-transpose-without-perm"), + (_rms_norm_without_axes_model, "rms-norm-without-axes")], + ids=["transpose-without-perm", "reduce-mean-without-axes"], +) +def test_optional_attribute_without_default(capfd, builder, label): + model_bytes, feeds = builder() + devices = _require_devices() + cpu_outputs = run(model_bytes, feeds, use_musa=False) + + options = ort.SessionOptions() + # An op the EP does not claim must fail loudly instead of running on the + # CPU EP, so the numbers below are really the MUSA ones. + options.add_session_config_entry("session.disable_cpu_ep_fallback", "1") + options.add_provider_for_devices(devices, {"dtype_diagnostics": "1"}) + capfd.readouterr() + session = ort.InferenceSession(model_bytes, sess_options=options) + musa_outputs = session.run(None, dict(feeds)) + + np.testing.assert_allclose( + musa_outputs[0], cpu_outputs[0], rtol=1e-5, atol=1e-6, + err_msg=f"{label}: CPU vs MUSA mismatch", + ) + node_lines = [ + line + for line in capfd.readouterr().err.splitlines() + if line.startswith("MUSA_DTYPE_NODE") + ] + assert node_lines, f"{label}: the EP reported no nodes" + for line in node_lines: + assert line.rstrip().endswith("assigned=1"), f"{label}: {line}" diff --git a/test/ops/test_fusion_ops.py b/test/ops/test_fusion_ops.py new file mode 100644 index 0000000..862b52f --- /dev/null +++ b/test/ops/test_fusion_ops.py @@ -0,0 +1,511 @@ +# Copyright (c) Moore Threads Technology Co., Ltd. All rights reserved. +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Operators whose nodes reach the EP only in particular opsets, dtypes or +graph shapes - the residue left after the opset and fusion work of the previous +steps. Each family below explains itself, and each of its cases builds the graph +from the model zoo, runs it on the stock CPU EP and on the MUSA EP, and checks +both the numbers and the placement: the MUSA session disables the CPU EP +fallback and asks the EP's dtype diagnostics to report the assignment of every +node, so an operator the EP does not claim fails loudly instead of silently +running on the CPU EP. +""" + +from __future__ import annotations + +import collections +import re + +import numpy as np +import onnxruntime as ort +import pytest +from onnx import helper + +from op_test_utils import ( + TensorProto, + build_graph_model, + build_model, + musa_devices, + run, +) + +Case = collections.namedtuple( + "Case", "op_type opset inputs outputs attrs label", defaults=(None, None) +) +GeluSubgraphCase = collections.namedtuple( + "GeluSubgraphCase", "opset with_half model_bytes feeds output_name" +) +LayerNormCase = collections.namedtuple( + "LayerNormCase", "opset dtype_name rtol model_bytes feeds" +) +UnaryCase = collections.namedtuple( + "UnaryCase", "op_type opset dtype_name rtol model_bytes feeds" +) +AveragePoolCase = collections.namedtuple( + "AveragePoolCase", "label opset x tensor_type rtol attrs" +) +GeluCase = collections.namedtuple( + "GeluCase", "opset dtype_name rtol model_bytes feeds" +) + +# The opsets the zoo models use for the old range, and the current one. +_LOW_OPSETS = (10, 12) + + +def _options(devices, *, disable_fallback=True): + options = ort.SessionOptions() + if disable_fallback: + options.add_session_config_entry("session.disable_cpu_ep_fallback", "1") + options.add_provider_for_devices(devices, {"dtype_diagnostics": "1"}) + return options + + +def _report(capfd): + return [ + line + for line in capfd.readouterr().err.splitlines() + if line.startswith("MUSA_DTYPE_NODE") + ] + + +def _assert_all_claimed(capfd, op_types, context): + """Every reported node is one of op_types and a MUSA kernel took it.""" + lines = _report(capfd) + assert len(lines) == len(op_types), f"{context}: {lines}" + seen = sorted(re.search(r" op=(\S+) ", line).group(1) for line in lines) + assert seen == sorted(op_types), f"{context}: {seen}" + for line in lines: + assert line.rstrip().endswith("assigned=1"), f"{context}: {line}" + + +def _require_devices(): + devices = musa_devices() + if not devices: + raise RuntimeError("No MUSA device available for the fusion ops test") + return devices + + +def _assert_claimed_and_equal(capfd, model_bytes, feeds, opset, rtol, atol, + op_type, expected_ops=None): + """Run the graph on the CPU EP and on the MUSA EP, compare, and assert that + every reported node was claimed by a MUSA kernel.""" + context = f"{op_type} at opset {opset}" + devices = _require_devices() + cpu_outputs = run(model_bytes, feeds, use_musa=False) + + capfd.readouterr() + session = ort.InferenceSession(model_bytes, sess_options=_options(devices)) + musa_outputs = session.run(None, dict(feeds)) + + assert len(cpu_outputs) == len(musa_outputs), f"{context}: output count" + for index, (cpu_output, musa_output) in enumerate(zip(cpu_outputs, musa_outputs)): + assert cpu_output.shape == musa_output.shape, f"{context}: output[{index}] shape" + assert cpu_output.dtype == musa_output.dtype, f"{context}: output[{index}] dtype" + np.testing.assert_allclose( + musa_output.astype(np.float64), + cpu_output.astype(np.float64), + rtol=rtol, + atol=atol, + err_msg=f"{context}: output[{index}] CPU vs MUSA mismatch", + ) + _assert_all_claimed(capfd, expected_ops or [op_type], context) + return musa_outputs + + +# --------------------------------------------------------------------------- # +# Erf, and the Gelu subgraph whose fusion the EP's partitioning interrupts +# --------------------------------------------------------------------------- # +# +# ORT fuses the Bert/GPT activation ``x * 0.5 * (1 + erf(x / sqrt(2)))`` into a +# single Gelu node in a level-2 pass that runs *after* the EP answered +# GetCapability - the fused node is built from nodes the EP had already claimed, +# and the EP is never asked about it. While the whole subgraph sat on the CPU EP +# that fusion matched; once the surrounding Div/Add/Mul move to the MUSA EP it +# no longer does, and the raw Erf node is left behind - registered from opset 13 +# on only, so an opset 10-12 model kept one Erf per Bert layer on the CPU EP. +# The subgraph case below is the evidence for both halves: the chain keeps all +# of its nodes on MUSA, and the node count proves ORT did not fuse it into a +# Gelu the EP was never asked about. + + +def _erf_cases(): + rng = np.random.default_rng(41) + x = rng.uniform(-3.0, 3.0, (2, 5, 4)).astype(np.float32) + for opset in (9, 10, 11, 12): + yield Case("Erf", opset, {"X": x}, [("Y", TensorProto.FLOAT)]) + yield Case("Erf", 12, {"X": x.astype(np.float16)}, [("Y", TensorProto.FLOAT16)], + None, "float16") + + +def _gelu_graph_models(): + """The two shapes of the Bert/Gelu pattern ORT fuses into one Gelu node. + + ``div -> Erf -> Add(1) -> Mul`` with the root multiplied in, plus the + variant that scales the whole result by 0.5. + """ + rng = np.random.default_rng(43) + x = rng.uniform(-3.0, 3.0, (1, 4, 8)).astype(np.float32) + sqrt_two = np.array([1.4142099618911743], dtype=np.float32) + one = np.array([1.0], dtype=np.float32) + half = np.array([0.5], dtype=np.float32) + for opset in (11, 12): + for with_half in (False, True): + nodes = [ + helper.make_node("Div", ["X", "sqrt_two"], ["div"]), + helper.make_node("Erf", ["div"], ["erf"]), + helper.make_node("Add", ["erf", "one"], ["add"]), + helper.make_node("Mul", ["X", "add"], ["mul"]), + ] + outputs = [("mul", TensorProto.FLOAT)] + if with_half: + nodes.append(helper.make_node("Mul", ["mul", "half"], ["y"])) + outputs = [("y", TensorProto.FLOAT)] + model_bytes = build_graph_model( + nodes, + {"X": x}, + outputs, + initializers=[ + helper.make_tensor("sqrt_two", TensorProto.FLOAT, [1], sqrt_two), + helper.make_tensor("one", TensorProto.FLOAT, [1], one), + helper.make_tensor("half", TensorProto.FLOAT, [1], half), + ], + opset=opset, + name=f"gelu_pattern_{opset}_{with_half}", + ) + yield GeluSubgraphCase( + opset, with_half, model_bytes, {"X": x}, outputs[-1][0] + ) + + +@pytest.mark.parametrize("case", list(_erf_cases()), ids=lambda case: f"Erf-opset{case.opset}-{case.label or 'float32'}") +def test_erf_old_opsets(capfd, case): + model_bytes = build_model( + case.op_type, case.inputs, case.outputs, case.attrs, "", case.opset + ) + rtol, atol = (2e-2, 2e-2) if case.label == "float16" else (1e-5, 1e-6) + _assert_claimed_and_equal( + capfd, model_bytes, case.inputs, case.opset, rtol, atol, case.op_type, + ) + + +@pytest.mark.parametrize( + "case", + list(_gelu_graph_models()), + ids=lambda case: f"opset{case.opset}-{'half' if case.with_half else 'no-half'}", +) +def test_gelu_subgraph_is_fully_on_musa(capfd, case): + opset, with_half = case.opset, case.with_half + model_bytes, feeds = case.model_bytes, case.feeds + """The Div/Erf/Add/Mul chain keeps all of its nodes on the MUSA EP.""" + devices = _require_devices() + cpu_outputs = run(model_bytes, feeds, use_musa=False) + + capfd.readouterr() + session = ort.InferenceSession(model_bytes, sess_options=_options(devices)) + musa_outputs = session.run(None, dict(feeds)) + np.testing.assert_allclose( + musa_outputs[0], cpu_outputs[0], rtol=1e-5, atol=1e-6 + ) + # 4 nodes for the plain pattern, 5 with the trailing Mul(0.5). A different + # count would mean ORT fused the subgraph into a single Gelu node, which the + # EP is never asked about. + expected = ["Div", "Erf", "Add", "Mul"] + (["Mul"] if with_half else []) + _assert_all_claimed(capfd, expected, f"gelu pattern at opset {opset}") + + +# --------------------------------------------------------------------------- # +# float16 LayerNormalization +# --------------------------------------------------------------------------- # +# +# ORT binds LayerNormalization's Scale/Bias to the same type parameter as X +# (the ONNX schema calls it U and only allows float/bfloat16 there), so the +# float16 layers the siglip and bge_reranker models carry bring float16 Scale +# and Bias with them. The kernel reads Scale/Bias through the input element +# type, and only the float32 route goes through muDNN, which wants float32 +# Scale/Bias. + + +def _layer_norm_models(): + rng = np.random.default_rng(47) + for opset in (17, 19): + for dtype, tensor_type, rtol in ( + (np.float16, TensorProto.FLOAT16, 2e-2), + (np.float32, TensorProto.FLOAT, 1e-4), + ): + x = rng.standard_normal((2, 3, 8)).astype(dtype) + scale = np.linspace(0.5, 1.5, 8).astype(dtype) + bias = np.linspace(-0.2, 0.2, 8).astype(dtype) + model_bytes = build_model( + "LayerNormalization", + {"X": x, "Scale": scale, "B": bias}, + [("Y", tensor_type)], + {"axis": -1, "epsilon": 1e-5}, + "", + opset, + ) + yield LayerNormCase( + opset, + np.dtype(dtype).name, + rtol, + model_bytes, + {"X": x, "Scale": scale, "B": bias}, + ) + + +@pytest.mark.parametrize( + "case", + list(_layer_norm_models()), + ids=lambda case: f"opset{case.opset}-{case.dtype_name}", +) +def test_layer_normalization_dtype_matrix(capfd, case): + _assert_claimed_and_equal( + capfd, + case.model_bytes, + case.feeds, + case.opset, + case.rtol, + case.rtol, + "LayerNormalization", + ) + + +def test_layer_normalization_float16_fused_into_the_graph(capfd): + """The dtype combination the siglip/bge_reranker graphs carry.""" + devices = _require_devices() + rng = np.random.default_rng(53) + x = rng.standard_normal((1, 6, 16)).astype(np.float16) + scale = np.linspace(0.5, 1.5, 16).astype(np.float16) + bias = np.linspace(-0.2, 0.2, 16).astype(np.float16) + model_bytes = build_graph_model( + [ + helper.make_node( + "LayerNormalization", ["X", "Scale", "B"], ["norm"] + ), + helper.make_node("Relu", ["norm"], ["Y"]), + ], + {"X": x}, + [("Y", TensorProto.FLOAT16)], + initializers=[ + helper.make_tensor("Scale", TensorProto.FLOAT16, [16], scale), + helper.make_tensor("B", TensorProto.FLOAT16, [16], bias), + ], + opset=17, + name="layer_norm_float16_graph", + ) + feeds = {"X": x} + cpu_outputs = run(model_bytes, feeds, use_musa=False) + capfd.readouterr() + session = ort.InferenceSession(model_bytes, sess_options=_options(devices)) + musa_outputs = session.run(None, dict(feeds)) + np.testing.assert_allclose( + musa_outputs[0].astype(np.float64), + cpu_outputs[0].astype(np.float64), + rtol=2e-2, + atol=2e-2, + ) + _assert_all_claimed(capfd, ["LayerNormalization", "Relu"], "float16 LayerNorm graph") + + +# --------------------------------------------------------------------------- # +# Sin / Cos +# --------------------------------------------------------------------------- # +# +# The Sin/Cos schemas are unchanged from opset 7 through 19 apart from the +# element types added at 13, so one registration covers the positional-encoding +# nodes of the opset 13-21 models and the older ones alike. + + +def _sin_cos_models(): + rng = np.random.default_rng(59) + for op_type in ("Sin", "Cos"): + for opset in (7, 12, 19): + for dtype, tensor_type, rtol in ( + (np.float32, TensorProto.FLOAT, 1e-5), + (np.float16, TensorProto.FLOAT16, 2e-2), + ): + x = rng.uniform(-6.0, 6.0, (2, 3, 4)).astype(dtype) + model_bytes = build_model( + op_type, {"X": x}, [("Y", tensor_type)], None, "", opset + ) + yield UnaryCase( + op_type, opset, np.dtype(dtype).name, rtol, model_bytes, {"X": x} + ) + + +@pytest.mark.parametrize( + "case", + list(_sin_cos_models()), + ids=lambda case: f"{case.op_type}-opset{case.opset}-{case.dtype_name}", +) +def test_sin_cos_matrix(capfd, case): + _assert_claimed_and_equal( + capfd, + case.model_bytes, + case.feeds, + case.opset, + case.rtol, + case.rtol, + case.op_type, + ) + + +# --------------------------------------------------------------------------- # +# AveragePool +# --------------------------------------------------------------------------- # +# +# AveragePool has no dilations but a count_include_pad switch the pooling route +# has to model, and a ceil_mode this route does not implement: the capability +# predicate and the kernel share one rejection list +# (kernels/nn/average_pool_support.h), so a node this route cannot run stays on +# the CPU EP instead of being claimed and failing when it runs. The last case +# below is that negative half. + + +def _average_pool_models(): + rng = np.random.default_rng(61) + x4 = rng.uniform(-2.0, 2.0, (1, 3, 8, 8)).astype(np.float32) + x5 = rng.uniform(-2.0, 2.0, (1, 2, 6, 6, 6)).astype(np.float32) + x16 = x4.astype(np.float16) + # The opset-7 schema already has kernel_shape/strides/pads/count_include_pad. + yield AveragePoolCase("opset7-k2s2", 7, x4, TensorProto.FLOAT, 1e-5, + {"kernel_shape": [2, 2], "strides": [2, 2]}) + yield AveragePoolCase("opset10-k2s2-ceil0", 10, x4, TensorProto.FLOAT, 1e-5, + {"kernel_shape": [2, 2], "strides": [2, 2], "ceil_mode": 0}) + yield AveragePoolCase( + "opset11-pads-count", 11, x4, TensorProto.FLOAT, 1e-5, + {"kernel_shape": [3, 3], "strides": [1, 1], "pads": [1, 1, 1, 1], + "count_include_pad": 1}) + yield AveragePoolCase( + "opset11-pads-no-count", 11, x4, TensorProto.FLOAT, 1e-5, + {"kernel_shape": [3, 3], "strides": [1, 1], "pads": [1, 1, 1, 1], + "count_include_pad": 0}) + yield AveragePoolCase( + "opset11-asymmetric-pads", 11, x4, TensorProto.FLOAT, 1e-5, + {"kernel_shape": [2, 2], "strides": [1, 1], "pads": [0, 1, 1, 0]}) + yield AveragePoolCase( + "opset19-same-upper", 19, x4, TensorProto.FLOAT, 1e-5, + {"kernel_shape": [2, 2], "strides": [2, 2], "auto_pad": "SAME_UPPER"}) + yield AveragePoolCase( + "opset19-same-lower", 19, x4, TensorProto.FLOAT, 1e-5, + {"kernel_shape": [3, 3], "strides": [2, 2], "auto_pad": "SAME_LOWER"}) + yield AveragePoolCase( + "opset19-valid", 19, x4, TensorProto.FLOAT, 1e-5, + {"kernel_shape": [3, 3], "strides": [2, 2], "auto_pad": "VALID"}) + yield AveragePoolCase("opset19-rank5", 19, x5, TensorProto.FLOAT, 1e-5, + {"kernel_shape": [2, 2, 2], "strides": [2, 2, 2]}) + # float16 stays at opset 11: at opset 19 the ORT NchwcTransformer rewrites a + # float16 AveragePool into com.ms.internal.nhwc.AveragePool(19), which the + # stock CPU EP has no kernel for - an ORT-side artefact this test must not + # depend on. + yield AveragePoolCase( + "opset11-float16", 11, x16, TensorProto.FLOAT16, 2e-2, + {"kernel_shape": [2, 2], "strides": [2, 2], "count_include_pad": 1}) + + +@pytest.mark.parametrize( + "case", list(_average_pool_models()), ids=lambda case: case.label +) +def test_average_pool_matrix(capfd, case): + model_bytes = build_model( + "AveragePool", + {"X": case.x}, + [("Y", case.tensor_type)], + case.attrs, + "", + case.opset, + ) + _assert_claimed_and_equal( + capfd, model_bytes, {"X": case.x}, case.opset, case.rtol, case.rtol, + "AveragePool", + ) + + +@pytest.mark.parametrize("opset", [10, 19]) +def test_average_pool_ceil_mode_stays_on_cpu(capfd, opset): + """ceil_mode is outside this MUSA route, so the node must stay on the CPU EP + instead of being claimed and failing when it runs.""" + devices = _require_devices() + rng = np.random.default_rng(67) + x = rng.uniform(-2.0, 2.0, (1, 2, 7, 7)).astype(np.float32) + attrs = {"kernel_shape": [3, 3], "strides": [2, 2], "ceil_mode": 1} + model_bytes = build_model( + "AveragePool", {"X": x}, [("Y", TensorProto.FLOAT)], attrs, "", opset + ) + cpu_outputs = run(model_bytes, feeds={"X": x}, use_musa=False) + + # The CPU EP fallback stays enabled here: the point of the case is that the + # node is not claimed, and the numbers must still match. + capfd.readouterr() + session = ort.InferenceSession( + model_bytes, sess_options=_options(devices, disable_fallback=False) + ) + musa_outputs = session.run(None, {"X": x}) + np.testing.assert_allclose(musa_outputs[0], cpu_outputs[0], rtol=1e-5, atol=1e-6) + lines = _report(capfd) + assert len(lines) == 1 and lines[0].rstrip().endswith("assigned=0"), lines + + +# --------------------------------------------------------------------------- # +# The ONNX Gelu (opset 20) +# --------------------------------------------------------------------------- # +# +# Gelu became an official ONNX operator at opset 20 (and a level-1 fusion target +# for the models that export it), which is what the opset-21 models carry; +# the tanh variant is a different function and stays on the CPU EP. + + +def _gelu_op_models(): + rng = np.random.default_rng(71) + for opset in (20, 21): + for dtype, tensor_type, rtol in ( + (np.float32, TensorProto.FLOAT, 1e-5), + (np.float16, TensorProto.FLOAT16, 2e-2), + ): + x = rng.uniform(-4.0, 4.0, (2, 3, 4)).astype(dtype) + model_bytes = build_model( + "Gelu", {"X": x}, [("Y", tensor_type)], None, "", opset + ) + yield GeluCase(opset, np.dtype(dtype).name, rtol, model_bytes, {"X": x}) + + +@pytest.mark.parametrize( + "case", + list(_gelu_op_models()), + ids=lambda case: f"opset{case.opset}-{case.dtype_name}", +) +def test_gelu_opset20_matrix(capfd, case): + _assert_claimed_and_equal( + capfd, case.model_bytes, case.feeds, case.opset, case.rtol, case.rtol, + "Gelu", + ) + + +def test_gelu_tanh_approximation_stays_on_cpu(capfd): + """The tanh approximation is a different function; those nodes stay on the + CPU EP rather than silently running the exact form.""" + devices = _require_devices() + rng = np.random.default_rng(73) + x = rng.uniform(-4.0, 4.0, (2, 3, 4)).astype(np.float32) + model_bytes = build_model( + "Gelu", {"X": x}, [("Y", TensorProto.FLOAT)], {"approximate": "tanh"}, "", 20 + ) + cpu_outputs = run(model_bytes, feeds={"X": x}, use_musa=False) + capfd.readouterr() + session = ort.InferenceSession( + model_bytes, sess_options=_options(devices, disable_fallback=False) + ) + musa_outputs = session.run(None, {"X": x}) + np.testing.assert_allclose(musa_outputs[0], cpu_outputs[0], rtol=1e-5, atol=1e-6) + lines = _report(capfd) + assert len(lines) == 1 and lines[0].rstrip().endswith("assigned=0"), lines + + diff --git a/test/ops/test_memcpy_boundary.py b/test/ops/test_memcpy_boundary.py index 375f648..ff26d1f 100644 --- a/test/ops/test_memcpy_boundary.py +++ b/test/ops/test_memcpy_boundary.py @@ -21,6 +21,11 @@ The consumer operators below (`Flatten`, `Transpose`) reject CPU memory on the device path, so the tests are RED (device-tensor guard failure) unless the boundary copy really lands in device memory. + +The CPU-side operator is `Atan`, which the plugin EP does not register: these +graphs need a node that really stays on the CPU EP, and an operator the EP +claims (Sin until its step-6 registration, for instance) produces the tensor on +the device instead, so no boundary copy is inserted at all. """ import numpy as np @@ -35,10 +40,10 @@ def _cpu_produced_input_model(x, op_type, attrs, name): - """`Sin` (CPU EP) -> `op_type` (MUSA EP): one inserted MemcpyFromHost.""" + """`Atan` (CPU EP) -> `op_type` (MUSA EP): one inserted MemcpyFromHost.""" return build_graph_model( [ - helper.make_node("Sin", ["X"], ["cpu_output"]), + helper.make_node("Atan", ["X"], ["cpu_output"]), helper.make_node(op_type, ["cpu_output"], ["Y"], **attrs), ], inputs={"X": x}, @@ -76,7 +81,7 @@ def test_memcpy_from_host_shared_by_multiple_musa_kernels(): x = np.linspace(-1.5, 1.5, 24, dtype=np.float32).reshape(2, 3, 4) model = build_graph_model( [ - helper.make_node("Sin", ["X"], ["cpu_output"]), + helper.make_node("Atan", ["X"], ["cpu_output"]), helper.make_node("Flatten", ["cpu_output"], ["flat"], axis=1), helper.make_node("Transpose", ["cpu_output"], ["transposed"], perm=[2, 0, 1]), ], @@ -94,7 +99,7 @@ def test_memcpy_to_host_before_cpu_kernel(monkeypatch): model = build_graph_model( [ helper.make_node("Relu", ["X"], ["relu_output"]), - helper.make_node("Sin", ["relu_output"], ["Y"]), + helper.make_node("Atan", ["relu_output"], ["Y"]), ], inputs={"X": x}, outputs=[("Y", TensorProto.FLOAT)], @@ -116,7 +121,7 @@ def test_memcpy_to_host_before_cpu_kernel(monkeypatch): def _if_branch(): return helper.make_graph( [ - helper.make_node("Sin", ["X"], ["cpu_output"]), + helper.make_node("Atan", ["X"], ["cpu_output"]), helper.make_node("Flatten", ["cpu_output"], ["branch_y"], axis=1), ], "branch", diff --git a/test/ops/test_unary_fp32_precision.py b/test/ops/test_unary_fp32_precision.py new file mode 100644 index 0000000..e479c78 --- /dev/null +++ b/test/ops/test_unary_fp32_precision.py @@ -0,0 +1,162 @@ +# Copyright (c) Moore Threads Technology Co., Ltd. All rights reserved. +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Wide-range float32 accuracy checks for the single-precision unary kernels. + +The float32 unary device path evaluates the transcendental functions in single +precision (``erff``/``sinf``/``cosf``/``tanhf``/``expf``/...), exactly like the +CUDA EP, instead of promoting every element to double first. These tests pin the +resulting accuracy against the CPU EP over the magnitude ranges where the old +FP64 promotion was strongest: the Erf tail and near-zero region, the Gelu large +positive/negative tails, and the Sin/Cos large-angle range (|x| up to 1e5). + +The gate is the ORT default float32 tolerance used by the CUDA EP op tests +(``atol=1e-5`` + ``rtol=1e-4 * |expected|``, numpy.isclose style, see +``onnxruntime/test/providers/checkers.cc`` ``DefaultTolerance``). Each +check prints the measured per-element max abs difference so the number is part +of the test log rather than only implied by a pass. +""" + +from __future__ import annotations + +import numpy as np +import pytest + +from op_test_utils import TensorProto, build_model, run + +# onnxruntime/test/providers/checkers.cc, DefaultTolerance (inference). +_ATOL = 1e-5 +_RTOL = 1e-4 + + +def _check_fp32_gate(op_type, x, opset, attrs=None): + """Run ``op_type`` on CPU and MUSA and assert the ORT float32 tolerance.""" + x = np.asarray(x, dtype=np.float32) + model = build_model( + op_type, {"X": x}, [("Y", TensorProto.FLOAT)], attrs, "", opset + ) + (cpu,) = run(model, {"X": x}, use_musa=False) + (musa,) = run(model, {"X": x}, use_musa=True) + assert cpu.shape == musa.shape == x.shape, f"{op_type}: shape mismatch" + + expected = cpu.astype(np.float64) + diff = np.abs(musa.astype(np.float64) - expected) + gate = _ATOL + _RTOL * np.abs(expected) + ratio = diff / gate + print( + f"GATE {op_type:9s} opset{opset:<3d} n={x.size:6d} " + f"max_abs={diff.max():.3e} max_ratio={ratio.max():.3f}", + flush=True, + ) + worst = int(np.argmax(ratio)) + assert ratio.max() <= 1.0, ( + f"{op_type} opset{opset}: element {worst} (x={x.ravel()[worst]:.6g}) " + f"diff={diff.ravel()[worst]:.3e} exceeds atol={_ATOL} + " + f"rtol={_RTOL}*|cpu|={gate.ravel()[worst]:.3e}" + ) + return diff.max() + + +# The Erf tail saturates at +/-1 from about |x| >= 4 and is ~2x/sqrt(pi) near 0. +_ERF_VALUES = np.array( + [ + -12.0, -10.0, -8.0, -6.0, -5.0, -4.5, -4.0, -3.5, -3.0, -2.5, -2.0, + -1.0, -0.5, -1e-3, -1e-6, -0.0, 0.0, 1e-6, 1e-3, 0.5, 1.0, 2.0, 2.5, + 3.0, 3.5, 4.0, 4.5, 5.0, 6.0, 8.0, 10.0, 12.0, + ], + dtype=np.float32, +) + + +@pytest.mark.parametrize("opset", [11, 13]) +def test_erf_float_wide_range(opset): + tail = np.random.default_rng(7).uniform(-10.0, 10.0, 4096).astype(np.float32) + near_zero = np.random.default_rng(8).uniform(-1e-4, 1e-4, 4096).astype(np.float32) + x = np.concatenate([_ERF_VALUES, tail, near_zero]).reshape(1, -1) + _check_fp32_gate("Erf", x, opset) + + +def test_gelu_float_large_magnitudes(): + x = np.array( + [ + -100.0, -50.0, -30.0, -20.0, -12.0, -8.0, -6.0, -4.0, -2.0, -1.0, + -0.5, -0.0, 0.0, 0.5, 1.0, 2.0, 4.0, 6.0, 8.0, 12.0, 20.0, 30.0, + 50.0, 100.0, + ], + dtype=np.float32, + ) + _check_fp32_gate("Gelu", x.reshape(1, -1), 20) + + +@pytest.mark.parametrize("op_type", ["Sin", "Cos"]) +@pytest.mark.parametrize("opset", [11, 13]) +def test_trig_float_large_angles(op_type, opset): + fixed = np.array( + [ + 0.0, 1e-6, 1e-3, 0.5, 1.0, 1.5707963, 3.1415927, 6.2831855, 100.0, + 1000.0, 10000.0, 100000.0, -100000.0, -10000.0, -1000.0, -100.0, + -6.2831855, -3.1415927, -1.5707963, -1.0, -0.5, + ], + dtype=np.float32, + ) + # A dense sweep in the large-angle regime where the argument reduction, not + # the polynomial, decides the accuracy. + large = np.random.default_rng(11).uniform(-1e5, 1e5, 8192).astype(np.float32) + x = np.concatenate([fixed, large]).reshape(1, -1) + _check_fp32_gate(op_type, x, opset) + + +# The remaining unary ops share the same single-precision switch. Exp/Log/Sqrt/ +# Reciprocal only see their domain, the rest are safe over both signs. +_DOMAIN_POSITIVE = { + "Sqrt": np.concatenate( + [ + np.array([0.0, 1e-6, 0.25, 1.0, 2.0, 16.0, 1e6], dtype=np.float32), + np.random.default_rng(13).uniform(1e-3, 1e3, 2048).astype(np.float32), + ] + ), + "Log": np.concatenate( + [ + np.array([1e-6, 1e-3, 0.5, 1.0, 2.0, 1e3, 1e6], dtype=np.float32), + np.random.default_rng(14).uniform(1e-4, 1e4, 2048).astype(np.float32), + ] + ), + "Exp": np.array( + [-80.0, -20.0, -1.0, -1e-6, 0.0, 1e-6, 1.0, 20.0, 80.0], dtype=np.float32 + ), + "Reciprocal": np.array( + [1e-6, 1e-3, 0.25, 1.0, 2.0, 1e3, 1e6, -1e-6, -1.0, -1e6], dtype=np.float32 + ), +} + + +@pytest.mark.parametrize("op_type", ["Sqrt", "Log", "Exp", "Reciprocal"]) +def test_unary_domain_float_wide_range(op_type): + _check_fp32_gate(op_type, _DOMAIN_POSITIVE[op_type].reshape(1, -1), 13) + + +@pytest.mark.parametrize( + "op_type", + ["Abs", "Neg", "Relu", "LeakyRelu", "Sigmoid", "Tanh", "Softplus", "Sign", + "Ceil", "Floor", "Round"], +) +def test_unary_signed_float_wide_range(op_type): + x = np.array( + [ + -100.0, -30.0, -8.0, -1.25, -0.5, -0.0, 0.0, 0.5, 1.25, 8.0, 30.0, + 100.0, 1e-6, -1e-6, + ], + dtype=np.float32, + ) + attrs = {"alpha": 0.2} if op_type == "LeakyRelu" else None + opset = 11 if op_type == "Round" else 13 + _check_fp32_gate(op_type, x.reshape(1, -1), opset, attrs) From ae4d12ab7c65ba76b8de8626b3856294f0c0dfb9 Mon Sep 17 00:00:00 2001 From: "yichao.hu" Date: Thu, 15 Oct 2026 12:34:04 -0400 Subject: [PATCH 18/52] fix(plugin): keep SplitUnsqueezeConcat runtime shape chain on the CPU The fused node reads its Reshape target at runtime and requires it in host memory, but has no KernelDef, so once the Memcpy kernels let this EP claim the CPU->device boundary copies ORT may place that target on the device and the fusion fails with 'requires CPU shape input'. Walk the integer/bool metadata chain that produces the target and keep it off this EP so the tensor stays on the CPU (same effect as ONNX Reshape declaring its shape input as CPUInput). --- musa/ep/src/ep.cc | 96 ++++++++++++++++++++++++++++++++++++++++++++++- 1 file changed, 94 insertions(+), 2 deletions(-) diff --git a/musa/ep/src/ep.cc b/musa/ep/src/ep.cc index c1095e3..168c48e 100644 --- a/musa/ep/src/ep.cc +++ b/musa/ep/src/ep.cc @@ -30,6 +30,7 @@ #include "fusion/fusion_node_compute.h" #include "graph/capability_predicates.h" #include "graph/cpu_preferred_nodes.h" +#include "graph/graph_utils.h" #include "graph_mermaid_dump.h" #include "plugin_ep_utils.h" @@ -250,11 +251,97 @@ OrtStatus* CollectSupportedNodeCandidates( return nullptr; } +// Ops whose (integer/bool) outputs are shape/metadata rather than model data. +// Used to walk back the chain that produces a runtime-computed reshape target. +bool IsShapeMetadataOp(Ort::ConstNode node) { + if (IsOnnxOp(node, "Cast")) { + auto to = GetIntAttribute(node, "to"); + return to.has_value() && + (*to == ONNX_TENSOR_ELEMENT_DATA_TYPE_INT32 || + *to == ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64); + } + return IsOnnxOp(node, "Shape") || IsOnnxOp(node, "Gather") || + IsOnnxOp(node, "Concat") || IsOnnxOp(node, "Slice") || + IsOnnxOp(node, "Unsqueeze") || IsOnnxOp(node, "Squeeze") || + IsOnnxOp(node, "Reshape") || IsOnnxOp(node, "Split") || + IsOnnxOp(node, "Add") || IsOnnxOp(node, "Sub") || + IsOnnxOp(node, "Mul") || IsOnnxOp(node, "Div") || + IsOnnxOp(node, "Max") || IsOnnxOp(node, "Min") || + IsOnnxOp(node, "ReduceProd") || IsOnnxOp(node, "ReduceMax") || + IsOnnxOp(node, "Equal") || IsOnnxOp(node, "Where") || + IsOnnxOp(node, "ConstantOfShape"); +} + +bool IsIntegerOrBoolTensor(Ort::ConstValueInfo value) { + if (value == nullptr) { + return false; + } + Ort::ConstTypeInfo type_info = value.TypeInfo(); + if (type_info.GetONNXType() != ONNX_TYPE_TENSOR) { + return false; + } + ONNXTensorElementDataType elem_type = + type_info.GetTensorTypeAndShapeInfo().GetElementType(); + return elem_type == ONNX_TENSOR_ELEMENT_DATA_TYPE_INT32 || + elem_type == ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64 || + elem_type == ONNX_TENSOR_ELEMENT_DATA_TYPE_BOOL; +} + +// SplitUnsqueezeConcat fuses the Reshape/Split/Unsqueeze/Concat chain and reads +// the Reshape target at runtime, which must live in host memory (see +// split_unsqueeze_concat_fusion.cc). Because the fused node has no KernelDef, +// ORT cannot know that and may place the target on the device once the +// CPU->device boundary copies are claimed by this EP (the Memcpy kernels). +// Keep the integer/bool metadata chain that produces the target on the CPU so +// the fused node can read it - the same effect as ONNX Reshape declaring its +// shape input as OrtMemTypeCPUInput. +std::unordered_set CollectSplitUnsqueezeConcatShapeChainNodes( + const std::vector& fusion_matches) { + std::unordered_set nodes; + for (const FusionMatch& match : fusion_matches) { + if (std::string(match.finder) != "FindSplitUnsqueezeConcatFusions") { + continue; + } + for (const auto& fusion_nodes : match.fusions) { + for (Ort::ConstNode node : fusion_nodes) { + if (!IsOnnxOp(node, "Reshape")) { + continue; + } + std::vector inputs = node.GetInputs(); + if (inputs.size() < 2 || inputs[1] == nullptr) { + continue; + } + std::vector frontier{inputs[1]}; + std::unordered_set seen; + while (!frontier.empty()) { + Ort::ConstValueInfo current = frontier.back(); + frontier.pop_back(); + if (current == nullptr || !seen.insert(current.GetName()).second) { + continue; + } + Ort::ValueInfoConsumerProducerInfo producer = current.GetProducerNode(); + if (producer.node == nullptr || !IsShapeMetadataOp(producer.node)) { + continue; // graph input / initializer / data-path producer + } + nodes.insert(producer.node); + for (Ort::ConstValueInfo input : producer.node.GetInputs()) { + if (IsIntegerOrBoolTensor(input)) { + frontier.push_back(input); + } + } + } + } + } + } + return nodes; +} + OrtStatus* RegisterSingleNodeCapabilities( const OrtEpApi& ep_api, const OrtGraph& ort_graph, OrtEpGraphSupportInfo& graph_support_info, const SupportedNodeCandidates& candidates, - std::unordered_set& registered_node_ids) { + std::unordered_set& registered_node_ids, + const std::unordered_set& forced_cpu_nodes) { std::unordered_set cpu_preferred_nodes; if (!candidates.tentative_nodes.empty()) { RETURN_IF_ERROR(GetCpuPreferredNodes( @@ -263,6 +350,9 @@ OrtStatus* RegisterSingleNodeCapabilities( candidates.tentative_nodes.size()), cpu_preferred_nodes)); } + for (const OrtNode* node : forced_cpu_nodes) { + cpu_preferred_nodes.insert(node); + } for (const auto& node : candidates.nodes) { if (cpu_preferred_nodes.count(node) == 0) { @@ -389,9 +479,11 @@ MusaEp::GetCapabilityImpl(OrtEp* this_ptr, const OrtGraph* ort_graph, GetFusedNodeIds(fusion_matches), candidates)); std::unordered_set registered_node_ids = GetFusedNodeIds(fusion_matches); + std::unordered_set forced_cpu_nodes = + CollectSplitUnsqueezeConcatShapeChainNodes(fusion_matches); RETURN_IF_ERROR(RegisterSingleNodeCapabilities( ep->ep_api_, *ort_graph, *graph_support_info, candidates, - registered_node_ids)); + registered_node_ids, forced_cpu_nodes)); EmitDTypeDiagnostics(all_nodes, registered_node_ids, fusion_matches, ep->config_.provider_options); } catch (const Ort::Exception& ex) { From 4a8bef4ccf5d693b2bedea96e8d803e0d1afcfd1 Mon Sep 17 00:00:00 2001 From: "yichao.hu" Date: Thu, 15 Oct 2026 12:35:55 -0400 Subject: [PATCH 19/52] test(plugin): cover SplitUnsqueezeConcat with a runtime-computed reshape target The reshape target is built by Shape/Slice/Concat instead of an initializer, exercising the host-memory shape-input path. --- .../test_split_unsqueeze_concat_fusion.py | 89 +++++++++++++++++-- 1 file changed, 84 insertions(+), 5 deletions(-) diff --git a/test/fusion/test_split_unsqueeze_concat_fusion.py b/test/fusion/test_split_unsqueeze_concat_fusion.py index f8d28af..1f1f384 100644 --- a/test/fusion/test_split_unsqueeze_concat_fusion.py +++ b/test/fusion/test_split_unsqueeze_concat_fusion.py @@ -45,7 +45,13 @@ def _profile_musa_node_names(model_bytes, feeds): def _build_model( - part_count, part_width, sequence, with_transpose, explicit_split, include_reshape + part_count, + part_width, + sequence, + with_transpose, + explicit_split, + include_reshape, + runtime_shape=False, ): packed_width = part_count * part_width initializers = [] @@ -58,10 +64,42 @@ def _build_model( else ["batch", sequence, packed_width] ) if include_reshape: - target_shape = helper.make_tensor( - "target_shape", TensorProto.INT64, [3], [0, sequence, packed_width] - ) - initializers.append(target_shape) + if runtime_shape: + # Reshape target computed at runtime (Shape -> Slice -> Concat) instead + # of an initializer. The fused node then has to read it from host + # memory, which is the CPU-shape-input path fixed here. + initializers.extend( + [ + helper.make_tensor("rt_start", TensorProto.INT64, [1], [0]), + helper.make_tensor("rt_end", TensorProto.INT64, [1], [1]), + helper.make_tensor("rt_axes", TensorProto.INT64, [1], [0]), + helper.make_tensor("dim_seq", TensorProto.INT64, [1], [sequence]), + helper.make_tensor("dim_pw", TensorProto.INT64, [1], [packed_width]), + ] + ) + nodes.append(helper.make_node("Shape", ["X"], ["rt_shape"], name="RtShape")) + nodes.append( + helper.make_node( + "Slice", + ["rt_shape", "rt_start", "rt_end", "rt_axes"], + ["rt_batch"], + name="RtSlice", + ) + ) + nodes.append( + helper.make_node( + "Concat", + ["rt_batch", "dim_seq", "dim_pw"], + ["target_shape"], + axis=0, + name="RtConcat", + ) + ) + else: + target_shape = helper.make_tensor( + "target_shape", TensorProto.INT64, [3], [0, sequence, packed_width] + ) + initializers.append(target_shape) nodes.append( helper.make_node( "Reshape", ["X", "target_shape"], ["R"], name="LateReshape" @@ -219,3 +257,44 @@ def test_split_unsqueeze_concat_fusion_axis2_to_axis0( assert not any( name.startswith("LateTranspose") for name in node_names ), node_names + + +@pytest.mark.parametrize("with_transpose", [False, True]) +def test_split_unsqueeze_concat_fusion_runtime_shape(with_transpose): + """Reshape target computed at runtime (not an initializer). + + The fused node must read the target from host memory. Regression guard: the + chain that computes the target must stay on the CPU, otherwise the fusion + fails with 'SplitUnsqueezeConcat requires CPU shape input'. + """ + part_count, part_width, sequence, batch = 2, 5, 4, 3 + packed_width = part_count * part_width + rng = np.random.default_rng(43) + x = rng.standard_normal((batch, sequence * packed_width)).astype(np.float32) + model_bytes = _build_model( + part_count, + part_width, + sequence, + with_transpose, + explicit_split=False, + include_reshape=True, + runtime_shape=True, + ) + + musa_outputs = run_model_and_compare( + model_bytes, {"X": x}, rtol=1e-6, atol=1e-6 + ) + reshaped = x.reshape(batch, sequence, packed_width) + pieces = np.split(reshaped, part_count, axis=2) + expected = np.concatenate( + [np.expand_dims(piece, axis=0) for piece in pieces], axis=0 + ) + if with_transpose: + expected = expected.transpose(0, 1, 3, 2) + np.testing.assert_allclose(musa_outputs[0], expected, rtol=0, atol=0) + + node_names = _profile_musa_node_names(model_bytes, {"X": x}) + assert not any(name.startswith("LateReshape") for name in node_names), node_names + assert not any(name.startswith("LateSplit") for name in node_names), node_names + assert not any(name.startswith("LateUnsqueeze") for name in node_names), node_names + assert not any(name.startswith("LateConcat") for name in node_names), node_names From 22bbcbe9e539580945df4e89b8fccae929206b52 Mon Sep 17 00:00:00 2001 From: "yichao.hu" Date: Thu, 15 Oct 2026 12:43:24 -0400 Subject: [PATCH 20/52] fix(plugin): do not fold a Reshape with a runtime shape into SplitUnsqueezeConcat A fused node has no KernelDef, so it cannot declare its shape input as CPU memory. Once this EP claims the CPU->device boundary copies (the Memcpy kernels) ORT places a runtime-computed reshape target on the device and the fusion fails with 'requires CPU shape input'. Keep such a Reshape standalone so its own kernel declares the shape input as CPU memory (as ONNX Reshape does). Reverts the earlier CPU-preferred attempt, which could not stop ORT from copying the fused node's input to the device. --- musa/ep/src/ep.cc | 96 +------------------ .../split_unsqueeze_concat_fusion_matcher.cc | 15 ++- 2 files changed, 16 insertions(+), 95 deletions(-) diff --git a/musa/ep/src/ep.cc b/musa/ep/src/ep.cc index 168c48e..c1095e3 100644 --- a/musa/ep/src/ep.cc +++ b/musa/ep/src/ep.cc @@ -30,7 +30,6 @@ #include "fusion/fusion_node_compute.h" #include "graph/capability_predicates.h" #include "graph/cpu_preferred_nodes.h" -#include "graph/graph_utils.h" #include "graph_mermaid_dump.h" #include "plugin_ep_utils.h" @@ -251,97 +250,11 @@ OrtStatus* CollectSupportedNodeCandidates( return nullptr; } -// Ops whose (integer/bool) outputs are shape/metadata rather than model data. -// Used to walk back the chain that produces a runtime-computed reshape target. -bool IsShapeMetadataOp(Ort::ConstNode node) { - if (IsOnnxOp(node, "Cast")) { - auto to = GetIntAttribute(node, "to"); - return to.has_value() && - (*to == ONNX_TENSOR_ELEMENT_DATA_TYPE_INT32 || - *to == ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64); - } - return IsOnnxOp(node, "Shape") || IsOnnxOp(node, "Gather") || - IsOnnxOp(node, "Concat") || IsOnnxOp(node, "Slice") || - IsOnnxOp(node, "Unsqueeze") || IsOnnxOp(node, "Squeeze") || - IsOnnxOp(node, "Reshape") || IsOnnxOp(node, "Split") || - IsOnnxOp(node, "Add") || IsOnnxOp(node, "Sub") || - IsOnnxOp(node, "Mul") || IsOnnxOp(node, "Div") || - IsOnnxOp(node, "Max") || IsOnnxOp(node, "Min") || - IsOnnxOp(node, "ReduceProd") || IsOnnxOp(node, "ReduceMax") || - IsOnnxOp(node, "Equal") || IsOnnxOp(node, "Where") || - IsOnnxOp(node, "ConstantOfShape"); -} - -bool IsIntegerOrBoolTensor(Ort::ConstValueInfo value) { - if (value == nullptr) { - return false; - } - Ort::ConstTypeInfo type_info = value.TypeInfo(); - if (type_info.GetONNXType() != ONNX_TYPE_TENSOR) { - return false; - } - ONNXTensorElementDataType elem_type = - type_info.GetTensorTypeAndShapeInfo().GetElementType(); - return elem_type == ONNX_TENSOR_ELEMENT_DATA_TYPE_INT32 || - elem_type == ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64 || - elem_type == ONNX_TENSOR_ELEMENT_DATA_TYPE_BOOL; -} - -// SplitUnsqueezeConcat fuses the Reshape/Split/Unsqueeze/Concat chain and reads -// the Reshape target at runtime, which must live in host memory (see -// split_unsqueeze_concat_fusion.cc). Because the fused node has no KernelDef, -// ORT cannot know that and may place the target on the device once the -// CPU->device boundary copies are claimed by this EP (the Memcpy kernels). -// Keep the integer/bool metadata chain that produces the target on the CPU so -// the fused node can read it - the same effect as ONNX Reshape declaring its -// shape input as OrtMemTypeCPUInput. -std::unordered_set CollectSplitUnsqueezeConcatShapeChainNodes( - const std::vector& fusion_matches) { - std::unordered_set nodes; - for (const FusionMatch& match : fusion_matches) { - if (std::string(match.finder) != "FindSplitUnsqueezeConcatFusions") { - continue; - } - for (const auto& fusion_nodes : match.fusions) { - for (Ort::ConstNode node : fusion_nodes) { - if (!IsOnnxOp(node, "Reshape")) { - continue; - } - std::vector inputs = node.GetInputs(); - if (inputs.size() < 2 || inputs[1] == nullptr) { - continue; - } - std::vector frontier{inputs[1]}; - std::unordered_set seen; - while (!frontier.empty()) { - Ort::ConstValueInfo current = frontier.back(); - frontier.pop_back(); - if (current == nullptr || !seen.insert(current.GetName()).second) { - continue; - } - Ort::ValueInfoConsumerProducerInfo producer = current.GetProducerNode(); - if (producer.node == nullptr || !IsShapeMetadataOp(producer.node)) { - continue; // graph input / initializer / data-path producer - } - nodes.insert(producer.node); - for (Ort::ConstValueInfo input : producer.node.GetInputs()) { - if (IsIntegerOrBoolTensor(input)) { - frontier.push_back(input); - } - } - } - } - } - } - return nodes; -} - OrtStatus* RegisterSingleNodeCapabilities( const OrtEpApi& ep_api, const OrtGraph& ort_graph, OrtEpGraphSupportInfo& graph_support_info, const SupportedNodeCandidates& candidates, - std::unordered_set& registered_node_ids, - const std::unordered_set& forced_cpu_nodes) { + std::unordered_set& registered_node_ids) { std::unordered_set cpu_preferred_nodes; if (!candidates.tentative_nodes.empty()) { RETURN_IF_ERROR(GetCpuPreferredNodes( @@ -350,9 +263,6 @@ OrtStatus* RegisterSingleNodeCapabilities( candidates.tentative_nodes.size()), cpu_preferred_nodes)); } - for (const OrtNode* node : forced_cpu_nodes) { - cpu_preferred_nodes.insert(node); - } for (const auto& node : candidates.nodes) { if (cpu_preferred_nodes.count(node) == 0) { @@ -479,11 +389,9 @@ MusaEp::GetCapabilityImpl(OrtEp* this_ptr, const OrtGraph* ort_graph, GetFusedNodeIds(fusion_matches), candidates)); std::unordered_set registered_node_ids = GetFusedNodeIds(fusion_matches); - std::unordered_set forced_cpu_nodes = - CollectSplitUnsqueezeConcatShapeChainNodes(fusion_matches); RETURN_IF_ERROR(RegisterSingleNodeCapabilities( ep->ep_api_, *ort_graph, *graph_support_info, candidates, - registered_node_ids, forced_cpu_nodes)); + registered_node_ids)); EmitDTypeDiagnostics(all_nodes, registered_node_ids, fusion_matches, ep->config_.provider_options); } catch (const Ort::Exception& ex) { diff --git a/musa/ep/src/fusion/split_unsqueeze_concat_fusion_matcher.cc b/musa/ep/src/fusion/split_unsqueeze_concat_fusion_matcher.cc index afc2dd4..8e6ded3 100644 --- a/musa/ep/src/fusion/split_unsqueeze_concat_fusion_matcher.cc +++ b/musa/ep/src/fusion/split_unsqueeze_concat_fusion_matcher.cc @@ -356,7 +356,20 @@ std::vector> FindSplitUnsqueezeConcatFusions( accepted_node_ids.count(producer.node.GetId()) == 0) { std::vector reshape_outputs = producer.node.GetOutputs(); - if (reshape_outputs.size() == 1 && + std::vector reshape_inputs = + producer.node.GetInputs(); + // Only fold the Reshape into the fusion when its target shape is a + // constant. A runtime-computed target must be read from host memory, but a + // fused node has no KernelDef and so cannot declare an input as CPU + // memory: once this EP claims the CPU->device boundary copies (the Memcpy + // kernels) ORT places the target on the device and the fusion fails with + // "requires CPU shape input". Leaving the Reshape standalone lets its own + // kernel declare the shape input as CPU memory (as ONNX Reshape does), so + // ORT keeps the target on the host. + const bool static_target = + reshape_inputs.size() < 2 || reshape_inputs[1] == nullptr || + ReadIntInitializerNoLimit(reshape_inputs[1]).has_value(); + if (static_target && reshape_outputs.size() == 1 && reshape_outputs[0].GetConsumers().size() == 1) { reshape_node = producer.node; } From 4323423348cec83253ed24bc0cccbaaf1f2267e7 Mon Sep 17 00:00:00 2001 From: "yichao.hu" Date: Thu, 15 Oct 2026 13:07:48 -0400 Subject: [PATCH 21/52] test(plugin): note the runtime-shape case is coverage, not a discriminator --- test/fusion/test_split_unsqueeze_concat_fusion.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/test/fusion/test_split_unsqueeze_concat_fusion.py b/test/fusion/test_split_unsqueeze_concat_fusion.py index 1f1f384..ff05d91 100644 --- a/test/fusion/test_split_unsqueeze_concat_fusion.py +++ b/test/fusion/test_split_unsqueeze_concat_fusion.py @@ -263,9 +263,10 @@ def test_split_unsqueeze_concat_fusion_axis2_to_axis0( def test_split_unsqueeze_concat_fusion_runtime_shape(with_transpose): """Reshape target computed at runtime (not an initializer). - The fused node must read the target from host memory. Regression guard: the - chain that computes the target must stay on the CPU, otherwise the fusion - fails with 'SplitUnsqueezeConcat requires CPU shape input'. + Covers the runtime-shape fusion path end to end. Note: this synthetic graph + does not reproduce the device-placement failure that a graph model hits + (the shape chain here stays on the CPU by itself), so it is coverage rather + than a discriminator for that fix. """ part_count, part_width, sequence, batch = 2, 5, 4, 3 packed_width = part_count * part_width From 3b52d665b943ef57a3300a2606f15e0c4a30b2e7 Mon Sep 17 00:00:00 2001 From: "yichao.hu" Date: Sun, 18 Oct 2026 09:23:41 -0400 Subject: [PATCH 22/52] fix(fusion): read a device-resident Slice end for raw INT32 MHTA masks Registering MemcpyFromHost in the plugin EP makes ORT copy CPU-produced scalars onto the device before fused nodes. The raw INT32 mask path in the MHTA SDPA fusion still required its Slice end to be host-resident and rejected the fusion with ORT_NOT_IMPLEMENTED once that copy existed, which turned a recommendation model into a hard error. Accept a device-resident end and read the single int64 back on the compute stream; the copy node precedes this node on the same stream, so a stream synchronize is sufficient. The Slice end comes from shape arithmetic, so cache the resolved value per input-shape set to keep the synchronous read off the steady-state path; a non-shape-derived end is re-read once and the cache refreshed. --- ...hta_scaled_dot_product_attention_fusion.cc | 64 ++++++++++++++++++- 1 file changed, 61 insertions(+), 3 deletions(-) diff --git a/musa/ep/src/fusion/mhta_scaled_dot_product_attention_fusion.cc b/musa/ep/src/fusion/mhta_scaled_dot_product_attention_fusion.cc index 81dce87..2e58491 100644 --- a/musa/ep/src/fusion/mhta_scaled_dot_product_attention_fusion.cc +++ b/musa/ep/src/fusion/mhta_scaled_dot_product_attention_fusion.cc @@ -19,6 +19,7 @@ #include #include #include +#include #include #include #include @@ -306,6 +307,13 @@ class MhtaScaledDotProductAttentionFusionCompute final int64_t static_mask_end_; bool lseq_last_key_mask_; bool no_external_mask_; + // The raw INT32 mask's Slice end is produced by shape arithmetic, so it can + // only change when the input shapes change. Cache the value resolved from the + // device so the per-iteration D2H read and stream synchronize happen once. + mutable std::mutex mask_end_mutex_; + mutable bool mask_end_cached_ = false; + mutable int64_t mask_end_cached_value_ = -1; + mutable std::vector mask_end_cached_shape_key_; }; bool IsMhtaSdpaTensorType(ONNXTensorElementDataType elem_type) { @@ -563,6 +571,34 @@ bool SetupMhtaSdpaMaskParams(const std::vector& mask_shape, return true; } +OrtStatus* ReadMaskEndValue(Ort::ConstValue mask_end, musaStream_t stream, + int64_t* value) { + if (IsGpuMemory(mask_end.GetTensorMemoryInfo())) { + // The EP registers MemcpyFromHost (PR #6), so a CPU-produced scalar that + // the fused node consumes is copied onto the device before this node runs. + // Read the single value back on the compute stream; the copy node precedes + // us on that stream. Callers cache the result, so this path runs only when + // the input shapes change. + int64_t host_value = 0; + musaError_t copy_status = musaMemcpyAsync( + &host_value, mask_end.GetTensorData(), sizeof(int64_t), + musaMemcpyDeviceToHost, stream); + if (copy_status != musaSuccess) { + return Ort::GetApi().CreateStatus(ORT_EP_FAIL, + MusaErrorString(copy_status)); + } + musaError_t sync_status = musaStreamSynchronize(stream); + if (sync_status != musaSuccess) { + return Ort::GetApi().CreateStatus(ORT_EP_FAIL, + MusaErrorString(sync_status)); + } + *value = host_value; + return nullptr; + } + *value = mask_end.GetTensorData()[0]; + return nullptr; +} + OrtStatus* MhtaScaledDotProductAttentionFusionCompute::Compute( OrtKernelContext* kernel_context) const { try { @@ -756,13 +792,35 @@ OrtStatus* MhtaScaledDotProductAttentionFusionCompute::Compute( Ort::ConstValue mask_end = ctx.GetInput(mask_end_index_); auto end_info = mask_end.GetTensorTypeAndShapeInfo(); if (end_info.GetElementType() != ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64 || - end_info.GetElementCount() != 1 || - IsGpuMemory(mask_end.GetTensorMemoryInfo())) { + end_info.GetElementCount() != 1) { return Ort::GetApi().CreateStatus( ORT_NOT_IMPLEMENTED, "MHTA raw INT32 mask requires a scalar CPU Slice end"); } - mask_end_value = mask_end.GetTensorData()[0]; + std::vector shape_key; + shape_key.insert(shape_key.end(), q_shape.begin(), q_shape.end()); + shape_key.insert(shape_key.end(), k_shape.begin(), k_shape.end()); + shape_key.insert(shape_key.end(), v_shape.begin(), v_shape.end()); + if (has_external_mask) { + const auto mask_dims = mask.GetTensorTypeAndShapeInfo().GetShape(); + shape_key.insert(shape_key.end(), mask_dims.begin(), mask_dims.end()); + } + std::lock_guard lock(mask_end_mutex_); + const bool cache_hit = + mask_end_cached_ && mask_end_cached_shape_key_ == shape_key; + if (cache_hit) { + mask_end_value = mask_end_cached_value_; + } else { + RETURN_IF_ERROR(ReadMaskEndValue(mask_end, stream, &mask_end_value)); + mask_end_cached_ = true; + mask_end_cached_value_ = mask_end_value; + mask_end_cached_shape_key_ = shape_key; + } + if (mask_end_value != seqlen_k && cache_hit) { + // Not shape-derived after all: re-read once and refresh the cache. + RETURN_IF_ERROR(ReadMaskEndValue(mask_end, stream, &mask_end_value)); + mask_end_cached_value_ = mask_end_value; + } } if (mask_end_value != seqlen_k) { return Ort::GetApi().CreateStatus( From 704af17b65e60790475bb109fe2af5831b0fcf85 Mon Sep 17 00:00:00 2001 From: "yichao.hu" Date: Sun, 18 Oct 2026 10:04:18 -0400 Subject: [PATCH 23/52] fix(fusion): cache Tile repeats only when they are shape-derived With MemcpyFromHost registered in the plugin EP, a CPU-produced repeats tensor is copied onto the device before the fused node runs. TileConcat read it with CopyToHost, which synchronizes the compute stream, so every iteration paid a device-to-host round trip; four recommendation models in the model zoo went from ~0.8-5.8 ms to ~5.5-11 ms. Cache the repeats across runs, but only when the matcher can prove the value is a pure function of fused-input shapes: the repeats chain must consist of shape arithmetic rooted at constants plus Shape ops whose source tensors are inputs of this fusion. The cache key is exactly those shapes, so a repeats that depends on anything else is re-read every run instead of being served stale. --- musa/ep/src/fusion/tile_concat_fusion.cc | 129 ++++++++++++++++++++++- 1 file changed, 124 insertions(+), 5 deletions(-) diff --git a/musa/ep/src/fusion/tile_concat_fusion.cc b/musa/ep/src/fusion/tile_concat_fusion.cc index 9190d48..42b1034 100644 --- a/musa/ep/src/fusion/tile_concat_fusion.cc +++ b/musa/ep/src/fusion/tile_concat_fusion.cc @@ -16,10 +16,12 @@ #include #include #include +#include #include #include #include #include +#include #include #include @@ -39,12 +41,20 @@ struct ConcatInputPlan { size_t data_input_index; size_t repeats_input_index; bool is_tile; + // Set when the Tile repeats value is a pure function of the shapes of the + // listed fused inputs (see RepeatsDependsOnlyOnInputShapes). Only then may + // Compute cache the value across runs and key it by those shapes. + bool repeats_cacheable = false; + std::vector repeats_source_indices; }; struct TileConcatFusionCompute : FusionNodeCompute { TileConcatFusionCompute(int64_t axis, size_t output_index, std::vector inputs) - : axis(axis), output_index(output_index), inputs(std::move(inputs)) {} + : axis(axis), + output_index(output_index), + inputs(std::move(inputs)), + repeats_cache(this->inputs.size()) {} OrtStatus* Compute(OrtKernelContext* kernel_context) const override { try { @@ -93,8 +103,30 @@ struct TileConcatFusionCompute : FusionNodeCompute { std::vector effective_shape = input_shape; const void* effective_data = input.GetTensorRawData(); if (input_plan.is_tile) { - std::vector repeats = - ReadIntTensor(ctx, input_plan.repeats_input_index); + const size_t plan_index = + static_cast(&input_plan - inputs.data()); + std::vector repeats; + if (input_plan.repeats_cacheable) { + std::vector key; + for (size_t index : input_plan.repeats_source_indices) { + auto dims = + ctx.GetInput(index).GetTensorTypeAndShapeInfo().GetShape(); + key.insert(key.end(), dims.begin(), dims.end()); + key.push_back(-1); + } + std::lock_guard lock(repeats_cache_mutex); + RepeatsCacheEntry& entry = repeats_cache[plan_index]; + if (entry.valid && entry.key == key) { + repeats = entry.repeats; + } else { + repeats = ReadIntTensor(ctx, input_plan.repeats_input_index); + entry.valid = true; + entry.key = std::move(key); + entry.repeats = repeats; + } + } else { + repeats = ReadIntTensor(ctx, input_plan.repeats_input_index); + } if (repeats.size() < input_shape.size()) { repeats.insert(repeats.begin(), input_shape.size() - repeats.size(), 1); @@ -339,6 +371,18 @@ struct TileConcatFusionCompute : FusionNodeCompute { int64_t axis; size_t output_index; std::vector inputs; + + // Reading repeats from device memory costs a stream synchronize per run + // once the EP registers MemcpyFromHost. Cache them only when the matcher + // proved they are a pure function of input shapes, and key the cache by + // exactly those shapes. + struct RepeatsCacheEntry { + bool valid = false; + std::vector key; + std::vector repeats; + }; + mutable std::mutex repeats_cache_mutex; + mutable std::vector repeats_cache; }; bool IsOnnxDomain(const std::string& domain) { @@ -408,6 +452,73 @@ std::unordered_map BuildProducers( return producers; } +// Returns true when the Tile repeats value is fully determined by the shapes +// of fused inputs (plus constants), and fills their input indices. This is the +// precondition for caching repeats across runs: the cache key is exactly those +// shapes, so a value that could change without them changing (e.g. a repeats +// derived from a tensor outside this fusion) must not be cached. +bool RepeatsDependsOnlyOnInputShapes( + const std::string& repeats_name, + const std::unordered_map& producers, + const std::unordered_map& input_indices, + const std::unordered_set& initializer_names, + std::vector* source_indices) { + // Shape-only arithmetic. Anything else may consume data and is rejected. + static const std::unordered_set kShapeOnlyOps = { + "Gather", "Unsqueeze", "Squeeze", "Concat", "Add", "Sub", "Mul", "Div", + "Cast", "Slice", "Identity", "Mod", "Floor", "Ceil", + "ReduceProd", "Reshape", "Constant", "ConstantOfShape"}; + std::vector pending{repeats_name}; + std::unordered_set visited; + std::unordered_set sources; + while (!pending.empty()) { + const std::string name = std::move(pending.back()); + pending.pop_back(); + if (!visited.insert(name).second) { + continue; + } + auto producer_it = producers.find(name); + if (producer_it == producers.end()) { + // A leaf must be a constant. A graph input or a fused data input used + // as a value would make the repeats data-dependent. + if (initializer_names.count(name) == 0) { + return false; + } + continue; + } + Ort::ConstNode node = producer_it->second; + const std::string op = node.GetOperatorType(); + std::vector node_inputs = node.GetInputs(); + if (op == "Shape") { + if (node_inputs.size() != 1) { + return false; + } + const std::string source = Name(node_inputs[0]); + auto index_it = input_indices.find(source); + if (index_it == input_indices.end()) { + // Shape of a constant is constant and needs no runtime key. + if (initializer_names.count(source) == 0) { + return false; + } + } else { + sources.insert(index_it->second); + } + continue; + } + if (kShapeOnlyOps.count(op) == 0) { + return false; + } + for (Ort::ConstValueInfo input : node_inputs) { + if (!Name(input).empty()) { + pending.push_back(Name(input)); + } + } + } + source_indices->assign(sources.begin(), sources.end()); + std::sort(source_indices->begin(), source_indices->end()); + return true; +} + } // namespace bool IsTileConcatFusionGraph(Ort::ConstGraph graph) { @@ -426,6 +537,10 @@ std::unique_ptr CreateTileConcatFusion( auto producers = BuildProducers(nodes); auto input_indices = FusedInputIndices(fused_node); auto output_indices = FusedOutputIndices(fused_node); + std::unordered_set initializer_names; + for (Ort::ConstValueInfo initializer : graph.GetInitializers()) { + initializer_names.insert(Name(initializer)); + } Ort::ConstNode concat_node{nullptr}; for (Ort::ConstNode node : nodes) { @@ -458,14 +573,18 @@ std::unique_ptr CreateTileConcatFusion( if (tile_inputs.size() != 2) { throw std::runtime_error("invalid TileConcat Tile node"); } + std::vector repeats_source_indices; + const bool repeats_cacheable = RepeatsDependsOnlyOnInputShapes( + Name(tile_inputs[1]), producers, input_indices, initializer_names, + &repeats_source_indices); plans.push_back( {GetIndex(input_indices, Name(tile_inputs[0]), "Tile data input"), GetIndex(input_indices, Name(tile_inputs[1]), "Tile repeats input"), - true}); + true, repeats_cacheable, std::move(repeats_source_indices)}); } else { plans.push_back( {GetIndex(input_indices, concat_input_name, "Concat input"), 0, - false}); + false, false, {}}); } } From 97506e85bc1fb6c186a4184ba07c7e43fc7cf741 Mon Sep 17 00:00:00 2001 From: "yichao.hu" Date: Sun, 18 Oct 2026 10:15:27 -0400 Subject: [PATCH 24/52] fix(fusion): leave Tile with runtime repeats unfused The fused node cannot declare CPU inputs, so once MemcpyFromHost became available ORT copied a runtime Tile repeats tensor onto the device. The Tile+Concat kernel then read it back with a compute-stream synchronize on every run, which took four recommendation models in the model zoo from ~0.8-5.8 ms to ~5.5-11 ms. A fused node cannot ask for OrtMemTypeCPUInput: OrtNodeFusionOptions has no memory-type field and Compile() only receives the fused node. Keep the Tile unfused whenever repeats comes from another node; the standalone Tile kernel already declares input 1 as OrtMemTypeCPUInput, so ORT leaves the repeats on the host and inserts no copy. --- musa/ep/src/fusion/tile_concat_fusion.cc | 127 +----------------- .../src/fusion/tile_concat_fusion_matcher.cc | 8 ++ 2 files changed, 13 insertions(+), 122 deletions(-) diff --git a/musa/ep/src/fusion/tile_concat_fusion.cc b/musa/ep/src/fusion/tile_concat_fusion.cc index 42b1034..ec4fda1 100644 --- a/musa/ep/src/fusion/tile_concat_fusion.cc +++ b/musa/ep/src/fusion/tile_concat_fusion.cc @@ -21,7 +21,6 @@ #include #include #include -#include #include #include @@ -41,20 +40,12 @@ struct ConcatInputPlan { size_t data_input_index; size_t repeats_input_index; bool is_tile; - // Set when the Tile repeats value is a pure function of the shapes of the - // listed fused inputs (see RepeatsDependsOnlyOnInputShapes). Only then may - // Compute cache the value across runs and key it by those shapes. - bool repeats_cacheable = false; - std::vector repeats_source_indices; }; struct TileConcatFusionCompute : FusionNodeCompute { TileConcatFusionCompute(int64_t axis, size_t output_index, std::vector inputs) - : axis(axis), - output_index(output_index), - inputs(std::move(inputs)), - repeats_cache(this->inputs.size()) {} + : axis(axis), output_index(output_index), inputs(std::move(inputs)) {} OrtStatus* Compute(OrtKernelContext* kernel_context) const override { try { @@ -103,30 +94,8 @@ struct TileConcatFusionCompute : FusionNodeCompute { std::vector effective_shape = input_shape; const void* effective_data = input.GetTensorRawData(); if (input_plan.is_tile) { - const size_t plan_index = - static_cast(&input_plan - inputs.data()); - std::vector repeats; - if (input_plan.repeats_cacheable) { - std::vector key; - for (size_t index : input_plan.repeats_source_indices) { - auto dims = - ctx.GetInput(index).GetTensorTypeAndShapeInfo().GetShape(); - key.insert(key.end(), dims.begin(), dims.end()); - key.push_back(-1); - } - std::lock_guard lock(repeats_cache_mutex); - RepeatsCacheEntry& entry = repeats_cache[plan_index]; - if (entry.valid && entry.key == key) { - repeats = entry.repeats; - } else { - repeats = ReadIntTensor(ctx, input_plan.repeats_input_index); - entry.valid = true; - entry.key = std::move(key); - entry.repeats = repeats; - } - } else { - repeats = ReadIntTensor(ctx, input_plan.repeats_input_index); - } + std::vector repeats = + ReadIntTensor(ctx, input_plan.repeats_input_index); if (repeats.size() < input_shape.size()) { repeats.insert(repeats.begin(), input_shape.size() - repeats.size(), 1); @@ -372,17 +341,6 @@ struct TileConcatFusionCompute : FusionNodeCompute { size_t output_index; std::vector inputs; - // Reading repeats from device memory costs a stream synchronize per run - // once the EP registers MemcpyFromHost. Cache them only when the matcher - // proved they are a pure function of input shapes, and key the cache by - // exactly those shapes. - struct RepeatsCacheEntry { - bool valid = false; - std::vector key; - std::vector repeats; - }; - mutable std::mutex repeats_cache_mutex; - mutable std::vector repeats_cache; }; bool IsOnnxDomain(const std::string& domain) { @@ -452,73 +410,6 @@ std::unordered_map BuildProducers( return producers; } -// Returns true when the Tile repeats value is fully determined by the shapes -// of fused inputs (plus constants), and fills their input indices. This is the -// precondition for caching repeats across runs: the cache key is exactly those -// shapes, so a value that could change without them changing (e.g. a repeats -// derived from a tensor outside this fusion) must not be cached. -bool RepeatsDependsOnlyOnInputShapes( - const std::string& repeats_name, - const std::unordered_map& producers, - const std::unordered_map& input_indices, - const std::unordered_set& initializer_names, - std::vector* source_indices) { - // Shape-only arithmetic. Anything else may consume data and is rejected. - static const std::unordered_set kShapeOnlyOps = { - "Gather", "Unsqueeze", "Squeeze", "Concat", "Add", "Sub", "Mul", "Div", - "Cast", "Slice", "Identity", "Mod", "Floor", "Ceil", - "ReduceProd", "Reshape", "Constant", "ConstantOfShape"}; - std::vector pending{repeats_name}; - std::unordered_set visited; - std::unordered_set sources; - while (!pending.empty()) { - const std::string name = std::move(pending.back()); - pending.pop_back(); - if (!visited.insert(name).second) { - continue; - } - auto producer_it = producers.find(name); - if (producer_it == producers.end()) { - // A leaf must be a constant. A graph input or a fused data input used - // as a value would make the repeats data-dependent. - if (initializer_names.count(name) == 0) { - return false; - } - continue; - } - Ort::ConstNode node = producer_it->second; - const std::string op = node.GetOperatorType(); - std::vector node_inputs = node.GetInputs(); - if (op == "Shape") { - if (node_inputs.size() != 1) { - return false; - } - const std::string source = Name(node_inputs[0]); - auto index_it = input_indices.find(source); - if (index_it == input_indices.end()) { - // Shape of a constant is constant and needs no runtime key. - if (initializer_names.count(source) == 0) { - return false; - } - } else { - sources.insert(index_it->second); - } - continue; - } - if (kShapeOnlyOps.count(op) == 0) { - return false; - } - for (Ort::ConstValueInfo input : node_inputs) { - if (!Name(input).empty()) { - pending.push_back(Name(input)); - } - } - } - source_indices->assign(sources.begin(), sources.end()); - std::sort(source_indices->begin(), source_indices->end()); - return true; -} - } // namespace bool IsTileConcatFusionGraph(Ort::ConstGraph graph) { @@ -537,10 +428,6 @@ std::unique_ptr CreateTileConcatFusion( auto producers = BuildProducers(nodes); auto input_indices = FusedInputIndices(fused_node); auto output_indices = FusedOutputIndices(fused_node); - std::unordered_set initializer_names; - for (Ort::ConstValueInfo initializer : graph.GetInitializers()) { - initializer_names.insert(Name(initializer)); - } Ort::ConstNode concat_node{nullptr}; for (Ort::ConstNode node : nodes) { @@ -573,18 +460,14 @@ std::unique_ptr CreateTileConcatFusion( if (tile_inputs.size() != 2) { throw std::runtime_error("invalid TileConcat Tile node"); } - std::vector repeats_source_indices; - const bool repeats_cacheable = RepeatsDependsOnlyOnInputShapes( - Name(tile_inputs[1]), producers, input_indices, initializer_names, - &repeats_source_indices); plans.push_back( {GetIndex(input_indices, Name(tile_inputs[0]), "Tile data input"), GetIndex(input_indices, Name(tile_inputs[1]), "Tile repeats input"), - true, repeats_cacheable, std::move(repeats_source_indices)}); + true}); } else { plans.push_back( {GetIndex(input_indices, concat_input_name, "Concat input"), 0, - false, false, {}}); + false}); } } diff --git a/musa/ep/src/fusion/tile_concat_fusion_matcher.cc b/musa/ep/src/fusion/tile_concat_fusion_matcher.cc index 5ad642e..e254a58 100644 --- a/musa/ep/src/fusion/tile_concat_fusion_matcher.cc +++ b/musa/ep/src/fusion/tile_concat_fusion_matcher.cc @@ -74,6 +74,14 @@ bool CanFuseTileConcat( return false; } + // A fused node cannot declare CPU inputs, so a runtime repeats would be + // copied to the device by ORT's MemcpyTransformer. Leave such a Tile + // unfused: the standalone Tile kernel declares its repeats as an + // OrtMemTypeCPUInput and reads it without a device copy or stream sync. + if (producers.find(Name(tile_inputs[1])) != producers.end()) { + return false; + } + found_tile_input = true; if (!AddFusionNode(tile_node, accepted_node_ids, selected_node_ids, fusion_nodes)) { From 81ad26f292451cbe31df0be2ffc61c35b868f124 Mon Sep 17 00:00:00 2001 From: "yichao.hu" Date: Sun, 18 Oct 2026 10:15:27 -0400 Subject: [PATCH 25/52] refactor(fusion): gate the MHTA Slice-end cache on provable shape inputs The cached Slice-end value was keyed on Q/K/V/mask shapes without proving the value is derived from them. Add ValueDependsOnlyOnInputShapes() and cache only when the end's producer chain is shape-only and every Shape source is an input of this fusion; the cache key is exactly those shapes. Anything else is read from the device every run. --- musa/ep/src/fusion/fusion_matcher_utils.cc | 66 +++++++++++++++++++ musa/ep/src/fusion/fusion_matcher_utils.h | 13 ++++ ...hta_scaled_dot_product_attention_fusion.cc | 54 ++++++++++----- 3 files changed, 116 insertions(+), 17 deletions(-) diff --git a/musa/ep/src/fusion/fusion_matcher_utils.cc b/musa/ep/src/fusion/fusion_matcher_utils.cc index 2a3e018..87dbd2d 100644 --- a/musa/ep/src/fusion/fusion_matcher_utils.cc +++ b/musa/ep/src/fusion/fusion_matcher_utils.cc @@ -283,4 +283,70 @@ bool HasCastBoundary(const std::vector& fusion_nodes) { return false; } +bool ValueDependsOnlyOnInputShapes( + const std::string& value_name, + const std::unordered_map& producers, + const std::unordered_map& input_indices, + const std::unordered_set& initializer_names, + std::vector* source_indices) { + // Ops that only reshape or combine shape values. Reshape is included because + // its output shape is fixed by its (constant) shape input; a data dependency + // would surface as a non-initializer leaf and be rejected below. + static const std::unordered_set kShapeOnlyOps = { + "Gather", "Unsqueeze", "Squeeze", "Concat", "Add", + "Sub", "Mul", "Div", "Cast", "Slice", + "Identity", "Mod", "Floor", "Ceil", "Reshape", + "ReduceProd", "Constant", "ConstantOfShape"}; + std::vector pending{value_name}; + std::unordered_set visited; + std::unordered_set sources; + while (!pending.empty()) { + const std::string name = std::move(pending.back()); + pending.pop_back(); + if (!visited.insert(name).second) { + continue; + } + auto producer_it = producers.find(name); + if (producer_it == producers.end()) { + // A leaf must be a constant. A graph input or fused data input used as a + // value would make the result data-dependent. + if (initializer_names.count(name) == 0) { + return false; + } + continue; + } + Ort::ConstNode node = producer_it->second; + const std::string op = node.GetOperatorType(); + std::vector node_inputs = node.GetInputs(); + if (op == "Shape") { + if (node_inputs.size() != 1) { + return false; + } + const std::string source = node_inputs[0].GetName(); + auto index_it = input_indices.find(source); + if (index_it == input_indices.end()) { + // Shape of a constant is constant and needs no runtime key. + if (initializer_names.count(source) == 0) { + return false; + } + } else { + sources.insert(index_it->second); + } + continue; + } + if (kShapeOnlyOps.count(op) == 0) { + return false; + } + for (Ort::ConstValueInfo input : node_inputs) { + const std::string input_name = input.GetName(); + if (!input_name.empty()) { + pending.push_back(input_name); + } + } + } + source_indices->assign(sources.begin(), sources.end()); + std::sort(source_indices->begin(), source_indices->end()); + return true; +} + } // namespace musa_ep diff --git a/musa/ep/src/fusion/fusion_matcher_utils.h b/musa/ep/src/fusion/fusion_matcher_utils.h index ce97859..09a404f 100644 --- a/musa/ep/src/fusion/fusion_matcher_utils.h +++ b/musa/ep/src/fusion/fusion_matcher_utils.h @@ -59,4 +59,17 @@ bool ValueHasExternalConsumerOrGraphOutput( std::optional ReadScalarIntInitializer(Ort::ConstValueInfo value_info); bool HasCastBoundary(const std::vector& fusion_nodes); +// Returns true when `value_name` is produced only from shape arithmetic: +// constants, plus Shape ops over inputs of the current fusion. The output is +// then a pure function of the shapes of the fused inputs listed in +// `source_indices` (empty means the value is constant). Anything that could +// make the value depend on data, on a tensor outside the fusion, or on an op +// that is not shape-only returns false. +bool ValueDependsOnlyOnInputShapes( + const std::string& value_name, + const std::unordered_map& producers, + const std::unordered_map& input_indices, + const std::unordered_set& initializer_names, + std::vector* source_indices); + } // namespace musa_ep diff --git a/musa/ep/src/fusion/mhta_scaled_dot_product_attention_fusion.cc b/musa/ep/src/fusion/mhta_scaled_dot_product_attention_fusion.cc index 2e58491..0de8806 100644 --- a/musa/ep/src/fusion/mhta_scaled_dot_product_attention_fusion.cc +++ b/musa/ep/src/fusion/mhta_scaled_dot_product_attention_fusion.cc @@ -24,6 +24,7 @@ #include #include #include +#include #include #include @@ -276,7 +277,9 @@ class MhtaScaledDotProductAttentionFusionCompute final float scale, float mask_scale, MhtaSdpaLayout layout, bool boolean_mask = false, bool boolean_mask_int32 = false, size_t mask_end_index = kNoInputIndex, int64_t static_mask_end = -1, - bool lseq_last_key_mask = false, bool no_external_mask = false) + bool lseq_last_key_mask = false, bool no_external_mask = false, + bool mask_end_cacheable = false, + const std::vector& mask_end_source_indices = {}) : q_index_(q_index), k_index_(k_index), v_index_(v_index), @@ -289,7 +292,9 @@ class MhtaScaledDotProductAttentionFusionCompute final mask_end_index_(mask_end_index), static_mask_end_(static_mask_end), lseq_last_key_mask_(lseq_last_key_mask), - no_external_mask_(no_external_mask) {} + no_external_mask_(no_external_mask), + mask_end_cacheable_(mask_end_cacheable), + mask_end_source_indices_(mask_end_source_indices) {} OrtStatus* Compute(OrtKernelContext* kernel_context) const override; @@ -307,9 +312,11 @@ class MhtaScaledDotProductAttentionFusionCompute final int64_t static_mask_end_; bool lseq_last_key_mask_; bool no_external_mask_; - // The raw INT32 mask's Slice end is produced by shape arithmetic, so it can - // only change when the input shapes change. Cache the value resolved from the - // device so the per-iteration D2H read and stream synchronize happen once. + bool mask_end_cacheable_; + std::vector mask_end_source_indices_; + // Reading the Slice end from device memory synchronizes the compute stream, + // so cache it — but only when the matcher proved the value is a pure + // function of the listed input shapes (see ValueDependsOnlyOnInputShapes). mutable std::mutex mask_end_mutex_; mutable bool mask_end_cached_ = false; mutable int64_t mask_end_cached_value_ = -1; @@ -798,23 +805,26 @@ OrtStatus* MhtaScaledDotProductAttentionFusionCompute::Compute( "MHTA raw INT32 mask requires a scalar CPU Slice end"); } std::vector shape_key; - shape_key.insert(shape_key.end(), q_shape.begin(), q_shape.end()); - shape_key.insert(shape_key.end(), k_shape.begin(), k_shape.end()); - shape_key.insert(shape_key.end(), v_shape.begin(), v_shape.end()); - if (has_external_mask) { - const auto mask_dims = mask.GetTensorTypeAndShapeInfo().GetShape(); - shape_key.insert(shape_key.end(), mask_dims.begin(), mask_dims.end()); + if (mask_end_cacheable_) { + for (size_t index : mask_end_source_indices_) { + const auto dims = + ctx.GetInput(index).GetTensorTypeAndShapeInfo().GetShape(); + shape_key.insert(shape_key.end(), dims.begin(), dims.end()); + shape_key.push_back(-1); + } } std::lock_guard lock(mask_end_mutex_); - const bool cache_hit = - mask_end_cached_ && mask_end_cached_shape_key_ == shape_key; + const bool cache_hit = mask_end_cacheable_ && mask_end_cached_ && + mask_end_cached_shape_key_ == shape_key; if (cache_hit) { mask_end_value = mask_end_cached_value_; } else { RETURN_IF_ERROR(ReadMaskEndValue(mask_end, stream, &mask_end_value)); - mask_end_cached_ = true; - mask_end_cached_value_ = mask_end_value; - mask_end_cached_shape_key_ = shape_key; + if (mask_end_cacheable_) { + mask_end_cached_ = true; + mask_end_cached_value_ = mask_end_value; + mask_end_cached_shape_key_ = std::move(shape_key); + } } if (mask_end_value != seqlen_k && cache_hit) { // Not shape-derived after all: re-read once and refresh the cache. @@ -1394,6 +1404,10 @@ std::unique_ptr CreateMhtaScaledDotProductAttentionFusion( } auto fused_indices = FusedInputIndices(fused_node); + std::unordered_set initializer_names; + for (Ort::ConstValueInfo initializer : graph.GetInitializers()) { + initializer_names.insert(Name(initializer)); + } if (where_nodes.size() == 2) { Ort::ConstNode post_where{nullptr}; @@ -1445,6 +1459,8 @@ std::unique_ptr CreateMhtaScaledDotProductAttentionFusion( bool boolean_mask_int32 = false; size_t mask_end_index = kNoInputIndex; int64_t static_mask_end = -1; + bool mask_end_cacheable = false; + std::vector mask_end_source_indices; auto equal_it = producers.find(Name(mask_input)); if (equal_it != producers.end() && IsOnnxOp(equal_it->second, "Equal")) { const auto equal_inputs = equal_it->second.GetInputs(); @@ -1475,6 +1491,9 @@ std::unique_ptr CreateMhtaScaledDotProductAttentionFusion( } else { mask_end_index = InputIndex(fused_indices, slice_it->second.GetInputs()[2]); + mask_end_cacheable = musa_ep::ValueDependsOnlyOnInputShapes( + Name(slice_it->second.GetInputs()[2]), producers, fused_indices, + initializer_names, &mask_end_source_indices); } boolean_mask_int32 = true; } @@ -1488,7 +1507,8 @@ std::unique_ptr CreateMhtaScaledDotProductAttentionFusion( InputIndex(fused_indices, boundary->v.GetInputs()[0]), InputIndex(fused_indices, mask_input), scale, 0.0f, MhtaSdpaLayout::kBshdStridedFlash, true, boolean_mask_int32, - mask_end_index, static_mask_end); + mask_end_index, static_mask_end, false, false, + mask_end_cacheable, mask_end_source_indices); } return std::make_unique( InputIndex(fused_indices, score_inputs[0]), From a29f2172128f0874ba65084d503cd60e4769e064 Mon Sep 17 00:00:00 2001 From: "yichao.hu" Date: Sun, 18 Oct 2026 12:50:44 -0400 Subject: [PATCH 26/52] feat(fusion): recompute Tile repeats from fused input shapes The fused Tile+Concat node cannot declare a CPU input, so once the EP registered MemcpyFromHost ORT copied a runtime Tile repeats tensor onto the device and the kernel read it back with a device-to-host copy plus a compute-stream synchronize on every run. 2d06d07 avoided that by leaving such Tiles unfused, at the cost of a separate Tile kernel. Absorb the shape-only chain that produces repeats into the fusion instead. The matcher takes the chain nodes only when - every node is a shape-only op (Shape, Gather, Unsqueeze, Squeeze, Concat, Add/Sub/Mul/Div/Mod, Cast, Slice, Reshape, Identity, Floor, Ceil, ReduceProd, Constant, ConstantOfShape), - every leaf is a constant (Constant node or initializer) or a Shape source that is already a fused input, - no chain value is a graph output or is read outside the chain, and - Compile() can rebuild the expression from the fused graph. The kernel then evaluates repeats from the shapes of its inputs: no device read and no stream synchronize. Repeats that is a constant or a graph input keeps the previous path, and an expression that fails to evaluate still falls back to reading the tensor when the tensor is an input. The chain is not moved when it has an outside consumer, so such a Tile stays unfused and the original graph is untouched. ORT_MUSA_DISABLE_TILE_CONCAT_FUSION=1 drops the TileConcat match, so one binary can measure the fusion on and off (the H1 attribution lever). DoD: clean build on the test host (SDK 4.3.8) with no new compiler warnings. The built .so sha256 is 9b403064aaa1b4ac0f06dbdb510339724eed01d2db29a668788680b0f6425937. Unit tests are written but not executed: no GPU window was granted this round. H1 (the recommendation-model residual comes from dropping this fusion) remains an unverified hypothesis. --- musa/docs/fusion/tile_concat.md | 2 +- musa/ep/src/ep.cc | 18 + musa/ep/src/fusion/shape_expr.cc | 1151 +++++++++++++++++ musa/ep/src/fusion/shape_expr.h | 94 ++ musa/ep/src/fusion/tile_concat_fusion.cc | 118 +- .../src/fusion/tile_concat_fusion_matcher.cc | 237 +++- test/fusion/test_tile_concat_shape_chain.py | 259 ++++ 7 files changed, 1844 insertions(+), 35 deletions(-) create mode 100644 musa/ep/src/fusion/shape_expr.cc create mode 100644 musa/ep/src/fusion/shape_expr.h create mode 100644 test/fusion/test_tile_concat_shape_chain.py diff --git a/musa/docs/fusion/tile_concat.md b/musa/docs/fusion/tile_concat.md index f9eba26..9fd152b 100644 --- a/musa/docs/fusion/tile_concat.md +++ b/musa/docs/fusion/tile_concat.md @@ -17,7 +17,7 @@ This file is generated from the current C++ fusion source and matching fusion te ## Extracted ONNX Ops -`Concat`, `Tile` +`Concat`, `Tile`, `Gather` ## Mermaid diff --git a/musa/ep/src/ep.cc b/musa/ep/src/ep.cc index c1095e3..515cf60 100644 --- a/musa/ep/src/ep.cc +++ b/musa/ep/src/ep.cc @@ -61,6 +61,15 @@ bool DisableShapeReshapeFusion() { return disabled; } +// Same-submit A/B switch: disabling only stops the matcher from claiming +// Tile+Concat nodes, so a latency difference under one binary can be attributed +// to the fusion itself instead of to a different build. +bool DisableTileConcatFusion() { + static const bool disabled = + EnvFlagEnabled("ORT_MUSA_DISABLE_TILE_CONCAT_FUSION"); + return disabled; +} + OrtStatus* RegisterFusionMatches( const OrtEpApi& ep_api, OrtEpGraphSupportInfo* graph_support_info, const std::vector& fusion_matches) { @@ -379,6 +388,15 @@ MusaEp::GetCapabilityImpl(OrtEp* this_ptr, const OrtGraph* ort_graph, }), fusion_matches.end()); } + if (DisableTileConcatFusion()) { + fusion_matches.erase( + std::remove_if(fusion_matches.begin(), fusion_matches.end(), + [](const FusionMatch& match) { + return std::string(match.finder) == + "FindTileConcatFusions"; + }), + fusion_matches.end()); + } RETURN_IF_ERROR(RegisterFusionMatches(ep->ep_api_, graph_support_info, fusion_matches)); } diff --git a/musa/ep/src/fusion/shape_expr.cc b/musa/ep/src/fusion/shape_expr.cc new file mode 100644 index 0000000..b20bdb9 --- /dev/null +++ b/musa/ep/src/fusion/shape_expr.cc @@ -0,0 +1,1151 @@ +// Copyright (c) Moore Threads Technology Co., Ltd. All rights reserved. +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "fusion/shape_expr.h" + +#include +#include +#include +#include +#include +#include +#include + +#include "graph/graph_utils.h" + +// Op encodings. `attrs` is the only per-op payload beyond the input references, +// so every supported op states what it keeps there: +// +// Constant attrs = tensor shape, constant = row-major data +// Shape inputs = {source slot}; attrs = {} +// Gather inputs = {data, indices}; attrs = {axis} +// Cast inputs = {data}; attrs = {target element type} +// Concat inputs = {...}; attrs = {axis} +// Unsqueeze inputs = {data}; attrs = axes +// Squeeze inputs = {data}; attrs = axes (empty = drop every 1-dim) +// Reshape inputs = {data}; attrs = {allowzero, dims...} +// Slice inputs = {data}; attrs = {n, (axis, start, end) * n} +// Add/Sub/Mul/Div/Mod +// Floor/Ceil/Identity inputs = {..}; attrs = {} +// ReduceProd inputs = {data}; attrs = {keepdims, n, axes...} +// ConstantOfShape inputs = {shape}; attrs = {value} +// +// Only integer tensors are modelled. Every value carries its rank so that +// Gather/Slice/Concat behave the way their ONNX definitions expect. + +namespace musa_ep { +namespace { + +constexpr size_t kMaxShapeExprNodes = 512; +constexpr int64_t kMaxConstantElements = 4096; + +struct Value { + std::vector shape; + std::vector data; +}; + +int64_t ElementCount(const std::vector& shape) { + int64_t count = 1; + for (int64_t dim : shape) { + if (dim < 0) { + return -1; + } + count *= dim; + } + return count; +} + +std::vector Strides(const std::vector& shape) { + std::vector strides(shape.size(), 1); + for (size_t i = shape.size(); i-- > 0;) { + if (i + 1 < shape.size()) { + strides[i] = strides[i + 1] * shape[i + 1]; + } + } + return strides; +} + +void Advance(std::vector* coordinate, + const std::vector& shape) { + for (size_t d = coordinate->size(); d-- > 0;) { + if (++(*coordinate)[d] < shape[d]) { + return; + } + (*coordinate)[d] = 0; + } +} + +template +bool ReadConstantTensor(const Tensor& tensor, ShapeExprConstant* out) { + if (!tensor) { + return false; + } + auto info = tensor.GetTensorTypeAndShapeInfo(); + const int64_t count = info.GetElementCount(); + if (count < 0 || count > kMaxConstantElements) { + return false; + } + out->shape = info.GetShape(); + if (info.GetElementType() == ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64) { + const int64_t* data = tensor.template GetTensorData(); + out->data.assign(data, data + count); + return true; + } + if (info.GetElementType() == ONNX_TENSOR_ELEMENT_DATA_TYPE_INT32) { + const int32_t* data = tensor.template GetTensorData(); + out->data.reserve(static_cast(count)); + for (int64_t i = 0; i < count; ++i) { + out->data.push_back(static_cast(data[i])); + } + return true; + } + return false; +} + +bool TryGetIntAttribute(Ort::ConstNode node, const char* name, int64_t* value) { + Ort::ConstOpAttr attr; + Ort::Status status = node.GetAttributeByName(name, attr); + if (!status.IsOK() || !attr) { + return false; + } + int64_t read = 0; + status = attr.GetValue(read); + if (!status.IsOK()) { + return false; + } + *value = read; + return true; +} + +bool TryGetIntsAttribute(Ort::ConstNode node, const char* name, + std::vector* values) { + Ort::ConstOpAttr attr; + Ort::Status status = node.GetAttributeByName(name, attr); + if (!status.IsOK() || !attr) { + return false; + } + std::vector read; + status = attr.GetValueArray(read); + if (!status.IsOK()) { + return false; + } + *values = std::move(read); + return true; +} + +bool IsIntElementType(ONNXTensorElementDataType type) { + return type == ONNX_TENSOR_ELEMENT_DATA_TYPE_INT32 || + type == ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64; +} + +// --------------------------------------------------------------------------- +// Evaluation +// --------------------------------------------------------------------------- + +bool NormalizeAxes(const std::vector& axes, int64_t rank, + std::vector* out) { + out->clear(); + out->reserve(axes.size()); + for (int64_t axis : axes) { + int64_t normalized = axis < 0 ? axis + rank : axis; + if (normalized < 0 || normalized >= rank) { + return false; + } + out->push_back(normalized); + } + std::sort(out->begin(), out->end()); + out->erase(std::unique(out->begin(), out->end()), out->end()); + return true; +} + +bool BroadcastShapes(const std::vector& lhs, + const std::vector& rhs, + std::vector* out) { + const size_t rank = std::max(lhs.size(), rhs.size()); + out->assign(rank, 1); + for (size_t i = 0; i < rank; ++i) { + const int64_t lhs_dim = + i + lhs.size() < rank ? 1 : lhs[i + lhs.size() - rank]; + const int64_t rhs_dim = + i + rhs.size() < rank ? 1 : rhs[i + rhs.size() - rank]; + if (lhs_dim != rhs_dim && lhs_dim != 1 && rhs_dim != 1) { + return false; + } + (*out)[i] = std::max(lhs_dim, rhs_dim); + } + return true; +} + +int64_t BroadcastOffset(const std::vector& source_shape, + const std::vector& out_shape, + const std::vector& out_coordinate) { + const size_t offset = out_shape.size() - source_shape.size(); + int64_t index = 0; + for (size_t i = 0; i < source_shape.size(); ++i) { + const int64_t coordinate = + source_shape[i] == 1 ? 0 : out_coordinate[offset + i]; + index = index * source_shape[i] + coordinate; + } + return index; +} + +bool EvalRef(const ShapeExpr& expr, const ShapeExprRef& ref, + const std::vector>& source_shapes, + std::vector>* memo, Value* out); + +bool EvalNode(const ShapeExpr& expr, int index, + const std::vector>& source_shapes, + std::vector>* memo, Value* out) { + if (index < 0 || static_cast(index) >= expr.nodes.size()) { + return false; + } + std::optional& slot = (*memo)[static_cast(index)]; + if (slot.has_value()) { + *out = *slot; + return true; + } + + const ShapeExprNode& node = expr.nodes[static_cast(index)]; + Value value; + + if (node.op == "Constant") { + value.shape = node.attrs; + value.data = node.constant; + } else if (node.op == "Shape") { + if (node.inputs.size() != 1 || node.inputs[0].is_node || + !EvalRef(expr, node.inputs[0], source_shapes, memo, &value)) { + return false; + } + } else if (node.op == "Gather") { + if (node.inputs.size() != 2) { + return false; + } + Value data; + Value indices; + if (!EvalRef(expr, node.inputs[0], source_shapes, memo, &data) || + !EvalRef(expr, node.inputs[1], source_shapes, memo, &indices)) { + return false; + } + const int64_t rank = static_cast(data.shape.size()); + if (rank == 0) { + return false; + } + int64_t axis = node.attrs.empty() ? 0 : node.attrs[0]; + if (axis < 0) { + axis += rank; + } + if (axis < 0 || axis >= rank) { + return false; + } + const size_t axis_index = static_cast(axis); + const int64_t axis_dim = data.shape[axis_index]; + int64_t inner = 1; + for (size_t i = axis_index + 1; i < data.shape.size(); ++i) { + inner *= data.shape[i]; + } + value.shape.assign(data.shape.begin(), data.shape.begin() + axis); + value.shape.insert(value.shape.end(), indices.shape.begin(), + indices.shape.end()); + value.shape.insert(value.shape.end(), data.shape.begin() + axis + 1, + data.shape.end()); + const int64_t count = ElementCount(value.shape); + if (count < 0) { + return false; + } + value.data.assign(static_cast(count), 0); + const std::vector index_shape = indices.shape; + std::vector coordinate(value.shape.size(), 0); + for (int64_t out_index = 0; out_index < count; ++out_index) { + int64_t outer_index = 0; + for (size_t i = 0; i < axis_index; ++i) { + outer_index = outer_index * data.shape[i] + coordinate[i]; + } + int64_t gathered = 0; + for (size_t i = 0; i < index_shape.size(); ++i) { + gathered = gathered * index_shape[i] + coordinate[axis_index + i]; + } + int64_t position = indices.data[static_cast(gathered)]; + if (position < 0) { + position += axis_dim; + } + if (position < 0 || position >= axis_dim) { + return false; + } + int64_t inner_index = 0; + for (size_t i = axis_index + 1; i < data.shape.size(); ++i) { + inner_index = + inner_index * data.shape[i] + + coordinate[axis_index + index_shape.size() + (i - axis_index - 1)]; + } + value.data[static_cast(out_index)] = + data.data[static_cast( + (outer_index * axis_dim + position) * inner + inner_index)]; + Advance(&coordinate, value.shape); + } + } else if (node.op == "Reshape") { + if (node.inputs.size() != 1 || node.attrs.empty()) { + return false; + } + Value data; + if (!EvalRef(expr, node.inputs[0], source_shapes, memo, &data)) { + return false; + } + const int64_t allowzero = node.attrs[0]; + value.shape.assign(node.attrs.begin() + 1, node.attrs.end()); + for (size_t i = 0; i < value.shape.size(); ++i) { + if (value.shape[i] == 0 && allowzero == 0) { + if (i >= data.shape.size()) { + return false; + } + value.shape[i] = data.shape[i]; + } + } + int64_t inferred = -1; + for (size_t i = 0; i < value.shape.size(); ++i) { + if (value.shape[i] == -1) { + if (inferred >= 0) { + return false; + } + inferred = static_cast(i); + } else if (value.shape[i] < 0) { + return false; + } + } + const int64_t source_count = ElementCount(data.shape); + if (source_count < 0) { + return false; + } + if (inferred >= 0) { + int64_t known = 1; + for (int64_t dim : value.shape) { + if (dim > 0) { + known *= dim; + } + } + if (known == 0 || source_count % known != 0) { + return false; + } + value.shape[static_cast(inferred)] = source_count / known; + } + if (ElementCount(value.shape) != source_count) { + return false; + } + value.data = data.data; + } else if (node.op == "Slice") { + if (node.inputs.size() != 1 || node.attrs.empty()) { + return false; + } + Value data; + if (!EvalRef(expr, node.inputs[0], source_shapes, memo, &data)) { + return false; + } + const int64_t rank = static_cast(data.shape.size()); + const int64_t entries = node.attrs[0]; + if (entries < 0 || + static_cast(entries) * 3 + 1 != node.attrs.size()) { + return false; + } + std::vector starts(static_cast(rank), 0); + std::vector ends = data.shape; + for (int64_t entry = 0; entry < entries; ++entry) { + int64_t axis = node.attrs[static_cast(1 + entry * 3)]; + int64_t start = node.attrs[static_cast(2 + entry * 3)]; + int64_t end = node.attrs[static_cast(3 + entry * 3)]; + if (axis < 0) { + axis += rank; + } + if (axis < 0 || axis >= rank) { + return false; + } + const size_t axis_index = static_cast(axis); + const int64_t dim = data.shape[axis_index]; + if (start < 0) { + start += dim; + } + if (end < 0) { + end += dim; + } + starts[axis_index] = std::clamp(start, 0, dim); + ends[axis_index] = std::clamp(end, 0, dim); + } + value.shape = data.shape; + for (size_t i = 0; i < value.shape.size(); ++i) { + value.shape[i] = ends[i] - starts[i]; + if (value.shape[i] < 0) { + return false; + } + } + const int64_t count = ElementCount(value.shape); + if (count < 0) { + return false; + } + const std::vector strides = Strides(data.shape); + value.data.assign(static_cast(count), 0); + std::vector coordinate(value.shape.size(), 0); + for (int64_t i = 0; i < count; ++i) { + int64_t source_offset = 0; + for (size_t d = 0; d < coordinate.size(); ++d) { + source_offset += (coordinate[d] + starts[d]) * strides[d]; + } + value.data[static_cast(i)] = + data.data[static_cast(source_offset)]; + Advance(&coordinate, value.shape); + } + } else if (node.op == "Concat") { + if (node.inputs.empty() || node.attrs.empty()) { + return false; + } + std::vector operands(node.inputs.size()); + for (size_t i = 0; i < node.inputs.size(); ++i) { + if (!EvalRef(expr, node.inputs[i], source_shapes, memo, &operands[i])) { + return false; + } + if (operands[i].shape.size() != operands.front().shape.size()) { + return false; + } + } + const int64_t rank = static_cast(operands.front().shape.size()); + int64_t axis = node.attrs[0]; + if (axis < 0) { + axis += rank; + } + if (rank == 0 || axis < 0 || axis >= rank) { + return false; + } + const size_t axis_index = static_cast(axis); + value.shape = operands.front().shape; + value.shape[axis_index] = 0; + for (const Value& operand : operands) { + for (size_t d = 0; d < value.shape.size(); ++d) { + if (d == axis_index) { + continue; + } + if (operand.shape[d] != value.shape[d]) { + return false; + } + } + value.shape[axis_index] += operand.shape[axis_index]; + } + const int64_t count = ElementCount(value.shape); + if (count < 0) { + return false; + } + value.data.assign(static_cast(count), 0); + int64_t inner = 1; + for (size_t i = axis_index + 1; i < value.shape.size(); ++i) { + inner *= value.shape[i]; + } + int64_t outer = 1; + for (size_t i = 0; i < axis_index; ++i) { + outer *= value.shape[i]; + } + const int64_t row = value.shape[axis_index] * inner; + for (int64_t o = 0; o < outer; ++o) { + int64_t written = o * row; + for (const Value& operand : operands) { + const int64_t chunk = operand.shape[axis_index] * inner; + const int64_t begin = o * chunk; + std::copy(operand.data.begin() + begin, + operand.data.begin() + begin + chunk, + value.data.begin() + written); + written += chunk; + } + } + } else if (node.op == "Unsqueeze") { + if (node.inputs.size() != 1) { + return false; + } + Value data; + if (!EvalRef(expr, node.inputs[0], source_shapes, memo, &data)) { + return false; + } + const int64_t out_rank = static_cast(data.shape.size()) + + static_cast(node.attrs.size()); + std::vector axes; + if (!NormalizeAxes(node.attrs, out_rank, &axes)) { + return false; + } + value.shape.clear(); + size_t source = 0; + for (int64_t i = 0; i < out_rank; ++i) { + if (std::binary_search(axes.begin(), axes.end(), i)) { + value.shape.push_back(1); + } else { + if (source >= data.shape.size()) { + return false; + } + value.shape.push_back(data.shape[source++]); + } + } + if (source != data.shape.size()) { + return false; + } + value.data = data.data; + } else if (node.op == "Squeeze") { + if (node.inputs.size() != 1) { + return false; + } + Value data; + if (!EvalRef(expr, node.inputs[0], source_shapes, memo, &data)) { + return false; + } + std::vector axes; + if (node.attrs.empty()) { + for (size_t i = 0; i < data.shape.size(); ++i) { + if (data.shape[i] == 1) { + axes.push_back(static_cast(i)); + } + } + } else if (!NormalizeAxes(node.attrs, + static_cast(data.shape.size()), &axes)) { + return false; + } + value.shape.clear(); + for (size_t i = 0; i < data.shape.size(); ++i) { + if (std::binary_search(axes.begin(), axes.end(), + static_cast(i))) { + if (data.shape[i] != 1) { + return false; + } + continue; + } + value.shape.push_back(data.shape[i]); + } + value.data = data.data; + } else if (node.op == "ReduceProd") { + if (node.inputs.size() != 1 || node.attrs.size() < 2) { + return false; + } + Value data; + if (!EvalRef(expr, node.inputs[0], source_shapes, memo, &data)) { + return false; + } + const int64_t keepdims = node.attrs[0]; + const int64_t axis_count = node.attrs[1]; + if (axis_count < 0 || + static_cast(axis_count) + 2 != node.attrs.size()) { + return false; + } + std::vector axes; + const int64_t rank = static_cast(data.shape.size()); + if (axis_count == 0) { + for (int64_t i = 0; i < rank; ++i) { + axes.push_back(i); + } + } else if (!NormalizeAxes(std::vector(node.attrs.begin() + 2, + node.attrs.end()), + rank, &axes)) { + return false; + } + std::vector reduced(data.shape.size(), false); + for (int64_t axis : axes) { + reduced[static_cast(axis)] = true; + } + value.shape.clear(); + for (size_t i = 0; i < data.shape.size(); ++i) { + if (reduced[i]) { + if (keepdims != 0) { + value.shape.push_back(1); + } + } else { + value.shape.push_back(data.shape[i]); + } + } + const int64_t count = ElementCount(value.shape); + const int64_t source_count = ElementCount(data.shape); + if (count < 0 || source_count < 0) { + return false; + } + value.data.assign(static_cast(count), 1); + std::vector coordinate(data.shape.size(), 0); + for (int64_t i = 0; i < source_count; ++i) { + int64_t out_index = 0; + size_t out_dim = 0; + for (size_t d = 0; d < data.shape.size(); ++d) { + if (reduced[d]) { + continue; + } + out_index = out_index * value.shape[out_dim++] + coordinate[d]; + } + value.data[static_cast(out_index)] *= + data.data[static_cast(i)]; + Advance(&coordinate, data.shape); + } + } else if (node.op == "ConstantOfShape") { + if (node.inputs.size() != 1 || node.attrs.empty()) { + return false; + } + Value shape_value; + if (!EvalRef(expr, node.inputs[0], source_shapes, memo, &shape_value)) { + return false; + } + value.shape = shape_value.data; + const int64_t count = ElementCount(value.shape); + if (count < 0) { + return false; + } + value.data.assign(static_cast(count), node.attrs[0]); + } else if (node.op == "Cast" || node.op == "Floor" || node.op == "Ceil" || + node.op == "Identity") { + if (node.inputs.size() != 1 || + !EvalRef(expr, node.inputs[0], source_shapes, memo, &value)) { + return false; + } + } else if (node.op == "Add" || node.op == "Sub" || node.op == "Mul" || + node.op == "Div" || node.op == "Mod") { + if (node.inputs.size() != 2) { + return false; + } + Value lhs; + Value rhs; + if (!EvalRef(expr, node.inputs[0], source_shapes, memo, &lhs) || + !EvalRef(expr, node.inputs[1], source_shapes, memo, &rhs)) { + return false; + } + if (!BroadcastShapes(lhs.shape, rhs.shape, &value.shape)) { + return false; + } + const int64_t count = ElementCount(value.shape); + if (count < 0) { + return false; + } + value.data.assign(static_cast(count), 0); + std::vector coordinate(value.shape.size(), 0); + for (int64_t i = 0; i < count; ++i) { + const int64_t a = lhs.data[static_cast( + BroadcastOffset(lhs.shape, value.shape, coordinate))]; + const int64_t b = rhs.data[static_cast( + BroadcastOffset(rhs.shape, value.shape, coordinate))]; + int64_t result = 0; + if (node.op == "Add") { + result = a + b; + } else if (node.op == "Sub") { + result = a - b; + } else if (node.op == "Mul") { + result = a * b; + } else if (node.op == "Div") { + if (b == 0) { + return false; + } + result = a / b; + } else { + if (b == 0) { + return false; + } + result = a % b; + if (result != 0 && (result < 0) != (b < 0)) { + result += b; + } + } + value.data[static_cast(i)] = result; + Advance(&coordinate, value.shape); + } + } else { + return false; + } + + const int64_t count = ElementCount(value.shape); + if (count < 0 || static_cast(count) != value.data.size()) { + return false; + } + slot = value; + *out = value; + return true; +} + +bool EvalRef(const ShapeExpr& expr, const ShapeExprRef& ref, + const std::vector>& source_shapes, + std::vector>* memo, Value* out) { + if (!ref.is_node) { + if (ref.index < 0 || + static_cast(ref.index) >= source_shapes.size()) { + return false; + } + out->shape = {static_cast(source_shapes[ref.index].size())}; + out->data = source_shapes[ref.index]; + return true; + } + return EvalNode(expr, ref.index, source_shapes, memo, out); +} + +// --------------------------------------------------------------------------- +// Builder +// --------------------------------------------------------------------------- + +using Producers = std::unordered_map; +using InputIndices = std::unordered_map; +using InitializerValues = std::unordered_map; + +class ExprBuilder { + public: + ExprBuilder(const Producers& producers, const InputIndices& input_indices, + const std::unordered_set& initializer_names, + const InitializerValues& initializer_values) + : producers_(producers), + input_indices_(input_indices), + initializer_names_(initializer_names), + initializer_values_(initializer_values) {} + + ShapeExpr Build(const std::string& value_name) { + ShapeExpr expr; + const int root = BuildValue(value_name, &expr); + if (root < 0) { + return ShapeExpr{}; + } + expr.root = root; + return expr; + } + + private: + int BuildValue(const std::string& name, ShapeExpr* expr) { + if (name.empty()) { + return -1; + } + auto memo_it = memo_.find(name); + if (memo_it != memo_.end()) { + return memo_it->second; + } + if (!pending_.insert(name).second) { + return -1; + } + + int result = -1; + auto producer_it = producers_.find(name); + if (producer_it == producers_.end()) { + auto constant_it = initializer_values_.find(name); + if (initializer_names_.count(name) != 0 && + constant_it != initializer_values_.end()) { + ShapeExprNode node; + node.op = "Constant"; + node.attrs = constant_it->second.shape; + node.constant = constant_it->second.data; + result = PushNode(std::move(node), expr); + } + } else { + result = BuildNode(producer_it->second, expr); + } + + pending_.erase(name); + if (result < 0) { + return -1; + } + memo_.emplace(name, result); + return result; + } + + int BuildConstantNode(Ort::ConstNode node, ShapeExpr* expr) { + ShapeExprConstant constant; + Ort::ConstOpAttr attr; + if (node.GetAttributeByName("value", attr).IsOK() && attr) { + Ort::Value tensor{nullptr}; + if (!attr.GetTensorAttributeAsOrtValue(tensor).IsOK() || + !ReadConstantTensor(tensor, &constant)) { + return -1; + } + } else { + int64_t int_value = 0; + std::vector ints_value; + if (TryGetIntAttribute(node, "value_int", &int_value)) { + constant.shape = {}; + constant.data = {int_value}; + } else if (TryGetIntsAttribute(node, "value_ints", &ints_value)) { + constant.shape = {static_cast(ints_value.size())}; + constant.data = std::move(ints_value); + } else { + return -1; + } + } + ShapeExprNode out; + out.op = "Constant"; + out.attrs = std::move(constant.shape); + out.constant = std::move(constant.data); + return PushNode(std::move(out), expr); + } + + int BuildNode(Ort::ConstNode node, ShapeExpr* expr) { + if (!node || !IsOnnxDomain(node.GetDomain())) { + return -1; + } + const std::string op = node.GetOperatorType(); + const std::vector inputs = node.GetInputs(); + + if (op == "Constant") { + return BuildConstantNode(node, expr); + } + if (op == "Shape") { + if (inputs.size() != 1) { + return -1; + } + auto index_it = input_indices_.find(Name(inputs[0])); + if (index_it == input_indices_.end()) { + return -1; + } + ShapeExprNode out; + out.op = "Shape"; + out.inputs.push_back( + {false, static_cast(SourceSlot(index_it->second, expr))}); + return PushNode(std::move(out), expr); + } + if (op == "Gather") { + if (inputs.size() != 2) { + return -1; + } + int64_t axis = 0; + TryGetIntAttribute(node, "axis", &axis); + std::vector refs; + if (!AddInputRef(inputs[0], expr, &refs) || + !AddInputRef(inputs[1], expr, &refs)) { + return -1; + } + ShapeExprNode out; + out.op = "Gather"; + out.inputs = std::move(refs); + out.attrs = {axis}; + return PushNode(std::move(out), expr); + } + if (op == "Cast") { + if (inputs.size() != 1) { + return -1; + } + int64_t to = 0; + if (!TryGetIntAttribute(node, "to", &to) || + !IsIntElementType(static_cast(to))) { + return -1; + } + std::vector refs; + if (!AddInputRef(inputs[0], expr, &refs)) { + return -1; + } + ShapeExprNode out; + out.op = "Cast"; + out.inputs = std::move(refs); + out.attrs = {to}; + return PushNode(std::move(out), expr); + } + if (op == "Concat") { + if (inputs.empty()) { + return -1; + } + int64_t axis = 0; + TryGetIntAttribute(node, "axis", &axis); + std::vector refs; + for (Ort::ConstValueInfo input : inputs) { + if (!AddInputRef(input, expr, &refs)) { + return -1; + } + } + ShapeExprNode out; + out.op = "Concat"; + out.inputs = std::move(refs); + out.attrs = {axis}; + return PushNode(std::move(out), expr); + } + if (op == "Unsqueeze" || op == "Squeeze") { + if (inputs.empty() || inputs.size() > 2) { + return -1; + } + std::vector axes; + if (inputs.size() == 2) { + if (!ReadConstantOperand(inputs[1], expr, &axes)) { + return -1; + } + } else if (!TryGetIntsAttribute(node, "axes", &axes)) { + axes.clear(); // opset < 13 Squeeze without axes drops every 1-dim + } + if (op == "Unsqueeze" && axes.empty()) { + return -1; + } + std::vector refs; + if (!AddInputRef(inputs[0], expr, &refs)) { + return -1; + } + ShapeExprNode out; + out.op = op; + out.inputs = std::move(refs); + out.attrs = std::move(axes); + return PushNode(std::move(out), expr); + } + if (op == "Reshape") { + if (inputs.size() < 2) { + return -1; + } + std::vector shape; + if (!ReadConstantOperand(inputs[1], expr, &shape)) { + return -1; + } + int64_t allowzero = 0; + TryGetIntAttribute(node, "allowzero", &allowzero); + std::vector refs; + if (!AddInputRef(inputs[0], expr, &refs)) { + return -1; + } + ShapeExprNode out; + out.op = "Reshape"; + out.inputs = std::move(refs); + out.attrs.push_back(allowzero); + out.attrs.insert(out.attrs.end(), shape.begin(), shape.end()); + return PushNode(std::move(out), expr); + } + if (op == "Slice") { + std::vector starts; + std::vector ends; + std::vector axes; + std::vector steps; + if (TryGetIntsAttribute(node, "starts", &starts)) { + if (!TryGetIntsAttribute(node, "ends", &ends)) { + return -1; + } + TryGetIntsAttribute(node, "axes", &axes); + TryGetIntsAttribute(node, "steps", &steps); + } else { + if (inputs.size() < 3) { + return -1; + } + if (!ReadConstantOperand(inputs[1], expr, &starts) || + !ReadConstantOperand(inputs[2], expr, &ends)) { + return -1; + } + if (inputs.size() > 3 && inputs[3] != nullptr && + !ReadConstantOperand(inputs[3], expr, &axes)) { + return -1; + } + if (inputs.size() > 4 && inputs[4] != nullptr && + !ReadConstantOperand(inputs[4], expr, &steps)) { + return -1; + } + } + if (starts.size() != ends.size()) { + return -1; + } + if (axes.empty()) { + for (size_t i = 0; i < starts.size(); ++i) { + axes.push_back(static_cast(i)); + } + } + if (axes.size() != starts.size()) { + return -1; + } + if (!steps.empty()) { + if (steps.size() != starts.size()) { + return -1; + } + for (int64_t step : steps) { + if (step != 1) { + return -1; + } + } + } + std::vector refs; + if (!AddInputRef(inputs[0], expr, &refs)) { + return -1; + } + ShapeExprNode out; + out.op = "Slice"; + out.inputs = std::move(refs); + out.attrs.push_back(static_cast(starts.size())); + for (size_t i = 0; i < starts.size(); ++i) { + out.attrs.push_back(axes[i]); + out.attrs.push_back(starts[i]); + out.attrs.push_back(ends[i]); + } + return PushNode(std::move(out), expr); + } + if (op == "Add" || op == "Sub" || op == "Mul" || op == "Div" || + op == "Mod") { + if (inputs.size() != 2) { + return -1; + } + std::vector refs; + if (!AddInputRef(inputs[0], expr, &refs) || + !AddInputRef(inputs[1], expr, &refs)) { + return -1; + } + ShapeExprNode out; + out.op = op; + out.inputs = std::move(refs); + return PushNode(std::move(out), expr); + } + if (op == "Floor" || op == "Ceil" || op == "Identity") { + if (inputs.size() != 1) { + return -1; + } + std::vector refs; + if (!AddInputRef(inputs[0], expr, &refs)) { + return -1; + } + ShapeExprNode out; + out.op = op; + out.inputs = std::move(refs); + return PushNode(std::move(out), expr); + } + if (op == "ReduceProd") { + if (inputs.empty()) { + return -1; + } + std::vector axes; + if (inputs.size() > 1 && inputs[1] != nullptr) { + if (!ReadConstantOperand(inputs[1], expr, &axes)) { + return -1; + } + } else { + TryGetIntsAttribute(node, "axes", &axes); + } + int64_t keepdims = 1; + TryGetIntAttribute(node, "keepdims", &keepdims); + std::vector refs; + if (!AddInputRef(inputs[0], expr, &refs)) { + return -1; + } + ShapeExprNode out; + out.op = "ReduceProd"; + out.inputs = std::move(refs); + out.attrs.push_back(keepdims); + out.attrs.push_back(static_cast(axes.size())); + out.attrs.insert(out.attrs.end(), axes.begin(), axes.end()); + return PushNode(std::move(out), expr); + } + if (op == "ConstantOfShape") { + if (inputs.size() != 1) { + return -1; + } + std::vector refs; + if (!AddInputRef(inputs[0], expr, &refs)) { + return -1; + } + int64_t fill = 0; + Ort::ConstOpAttr attr; + if (node.GetAttributeByName("value", attr).IsOK() && attr) { + Ort::Value tensor{nullptr}; + ShapeExprConstant constant; + if (!attr.GetTensorAttributeAsOrtValue(tensor).IsOK() || + !ReadConstantTensor(tensor, &constant) || + constant.data.size() != 1) { + return -1; + } + fill = constant.data[0]; + } + ShapeExprNode out; + out.op = op; + out.inputs = std::move(refs); + out.attrs = {fill}; + return PushNode(std::move(out), expr); + } + return -1; + } + + bool AddInputRef(Ort::ConstValueInfo input, ShapeExpr* expr, + std::vector* refs) { + if (input == nullptr) { + return false; + } + const int index = BuildValue(Name(input), expr); + if (index < 0) { + return false; + } + refs->push_back({true, index}); + return true; + } + + bool ReadConstantOperand(Ort::ConstValueInfo input, ShapeExpr* expr, + std::vector* data) { + if (input == nullptr) { + return false; + } + const int index = BuildValue(Name(input), expr); + if (index < 0) { + return false; + } + const std::vector> no_sources; + std::vector> memo(expr->nodes.size()); + Value value; + if (!EvalNode(*expr, index, no_sources, &memo, &value)) { + return false; + } + *data = std::move(value.data); + return true; + } + + size_t SourceSlot(size_t fused_input_index, ShapeExpr* expr) { + for (size_t slot = 0; slot < expr->source_input_indices.size(); ++slot) { + if (expr->source_input_indices[slot] == fused_input_index) { + return slot; + } + } + expr->source_input_indices.push_back(fused_input_index); + return expr->source_input_indices.size() - 1; + } + + int PushNode(ShapeExprNode node, ShapeExpr* expr) { + if (expr->nodes.size() >= kMaxShapeExprNodes) { + return -1; + } + expr->nodes.push_back(std::move(node)); + return static_cast(expr->nodes.size() - 1); + } + + const Producers& producers_; + const InputIndices& input_indices_; + const std::unordered_set& initializer_names_; + const InitializerValues& initializer_values_; + std::unordered_map memo_; + std::unordered_set pending_; +}; + +} // namespace + +std::optional ReadShapeExprConstant( + Ort::ConstValueInfo value_info) { + if (value_info == nullptr || !value_info.IsConstantInitializer()) { + return std::nullopt; + } + Ort::ConstValue value{nullptr}; + Ort::Status status = value_info.GetInitializer(value); + if (!status.IsOK() || !value) { + return std::nullopt; + } + ShapeExprConstant constant; + if (!ReadConstantTensor(value, &constant)) { + return std::nullopt; + } + return constant; +} + +ShapeExpr BuildShapeExpr( + const std::string& value_name, + const std::unordered_map& producers, + const std::unordered_map& input_indices, + const std::unordered_set& initializer_names, + const std::unordered_map& + initializer_values) { + ExprBuilder builder(producers, input_indices, initializer_names, + initializer_values); + return builder.Build(value_name); +} + +bool EvaluateShapeExpr(const ShapeExpr& expr, + const std::vector>& source_shapes, + std::vector* result) { + if (expr.Empty() || result == nullptr || + expr.source_input_indices.size() != source_shapes.size()) { + return false; + } + std::vector> memo(expr.nodes.size()); + Value value; + if (!EvalNode(expr, expr.root, source_shapes, &memo, &value)) { + return false; + } + *result = std::move(value.data); + return true; +} + +} // namespace musa_ep diff --git a/musa/ep/src/fusion/shape_expr.h b/musa/ep/src/fusion/shape_expr.h new file mode 100644 index 0000000..f75f9df --- /dev/null +++ b/musa/ep/src/fusion/shape_expr.h @@ -0,0 +1,94 @@ +// Copyright (c) Moore Threads Technology Co., Ltd. All rights reserved. +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#pragma once + +#define ORT_API_MANUAL_INIT +#include "onnxruntime_cxx_api.h" +#undef ORT_API_MANUAL_INIT + +#include +#include +#include +#include +#include +#include + +namespace musa_ep { + +// A small interpreter for the shape-arithmetic chains that produce Tile repeats +// and MHTA Slice ends. Such a chain only reads tensor SHAPES and constants, so +// a fused node can recompute the value from the shapes of its inputs instead of +// consuming the tensor that ORT copied to the device (which would otherwise +// require a device-to-host read and a compute-stream synchronize). +// +// The expression is built at fusion-creation time and is read-only afterwards, +// so evaluation is thread-safe. +struct ShapeExprRef { + bool is_node = false; // node result vs source tensor shape + int index = 0; // node index or source index +}; + +struct ShapeExprNode { + std::string op; + std::vector inputs; + std::vector attrs; // op-specific constants (axes, etc.) + std::vector constant; // for kConstant +}; + +struct ShapeExpr { + std::vector nodes; + std::vector source_input_indices; // fused-input index per source + int root = -1; + + bool Empty() const { return root < 0; } +}; + +// A constant tensor that a shape chain reads: an int32/int64 initializer, or +// the tensor carried by a Constant node. +struct ShapeExprConstant { + std::vector shape; + std::vector data; +}; + +// Reads an int32/int64 initializer into ShapeExprConstant. Returns nullopt for +// anything that is not a constant integer tensor. +std::optional ReadShapeExprConstant( + Ort::ConstValueInfo value_info); + +// Builds an expression for `value_name`. `input_indices` maps fused input names +// to their index; every Shape source must already be a fused input, otherwise +// the expression cannot be recomputed inside the fusion and `expr` stays empty. +// Returns an empty expression for chains that are not provably shape-only. +// `initializer_values` carries the tensors of the initializer leaves. Names +// alone cannot be turned back into values through the C API, so a caller that +// wants initializer leaves resolved must supply them here; leaves that are +// listed in `initializer_names` but missing from `initializer_values` make the +// expression empty (the chain cannot be recomputed). +ShapeExpr BuildShapeExpr( + const std::string& value_name, + const std::unordered_map& producers, + const std::unordered_map& input_indices, + const std::unordered_set& initializer_names, + const std::unordered_map& + initializer_values = {}); + +// Evaluates the expression. `source_shapes[i]` is the runtime shape of fused +// input `source_input_indices[i]`. Returns false on any semantic error (rank, +// bounds, division by zero, element-count mismatch); callers must fall back to +// reading the tensor in that case. +bool EvaluateShapeExpr(const ShapeExpr& expr, + const std::vector>& source_shapes, + std::vector* result); + +} // namespace musa_ep diff --git a/musa/ep/src/fusion/tile_concat_fusion.cc b/musa/ep/src/fusion/tile_concat_fusion.cc index ec4fda1..6b2ca48 100644 --- a/musa/ep/src/fusion/tile_concat_fusion.cc +++ b/musa/ep/src/fusion/tile_concat_fusion.cc @@ -16,14 +16,17 @@ #include #include #include +#include #include #include #include #include #include +#include #include #include +#include "fusion/shape_expr.h" #include "kernels/shared_inc/op_kernel_common.h" #include "kernels/tensor/concat_impl.h" #include "kernels/tensor/tile_impl.h" @@ -35,11 +38,19 @@ constexpr size_t kConcatManySmallMaxWidthBytes = 4096; constexpr size_t kConcatRelaxedSmallInputCount = 2; constexpr size_t kConcatRelaxedSmallMaxWidthBytes = 32 * 1024; constexpr size_t kConcatManySmallMaxMapBytes = 4 * 1024 * 1024; +constexpr size_t kNoInputIndex = std::numeric_limits::max(); struct ConcatInputPlan { size_t data_input_index; + // Set unless the repeats producer was absorbed into the fused node; then the + // tensor is not an input of the fused node and the expression is the only way + // to obtain the value. size_t repeats_input_index; bool is_tile; + // Non-empty when repeats was proven to be a pure function of the shapes of + // the fused inputs: recompute it instead of reading the device tensor, which + // would cost a device-to-host copy plus a compute-stream synchronize. + musa_ep::ShapeExpr repeats_expr; }; struct TileConcatFusionCompute : FusionNodeCompute { @@ -94,8 +105,28 @@ struct TileConcatFusionCompute : FusionNodeCompute { std::vector effective_shape = input_shape; const void* effective_data = input.GetTensorRawData(); if (input_plan.is_tile) { - std::vector repeats = - ReadIntTensor(ctx, input_plan.repeats_input_index); + std::vector repeats; + bool recomputed = false; + if (!input_plan.repeats_expr.Empty()) { + std::vector> source_shapes; + source_shapes.reserve( + input_plan.repeats_expr.source_input_indices.size()); + for (size_t index : input_plan.repeats_expr.source_input_indices) { + source_shapes.push_back( + ctx.GetInput(index).GetTensorTypeAndShapeInfo().GetShape()); + } + recomputed = musa_ep::EvaluateShapeExpr(input_plan.repeats_expr, + source_shapes, &repeats); + } + if (!recomputed) { + if (input_plan.repeats_input_index == kNoInputIndex) { + cleanup_temps(); + return Ort::GetApi().CreateStatus( + ORT_EP_FAIL, + "TileConcat could not recompute the fused Tile repeats"); + } + repeats = ReadIntTensor(ctx, input_plan.repeats_input_index); + } if (repeats.size() < input_shape.size()) { repeats.insert(repeats.begin(), input_shape.size() - repeats.size(), 1); @@ -340,7 +371,6 @@ struct TileConcatFusionCompute : FusionNodeCompute { int64_t axis; size_t output_index; std::vector inputs; - }; bool IsOnnxDomain(const std::string& domain) { @@ -429,14 +459,31 @@ std::unique_ptr CreateTileConcatFusion( auto input_indices = FusedInputIndices(fused_node); auto output_indices = FusedOutputIndices(fused_node); + // The absorbed repeats chain may contain a Concat of its own (for example + // Concat([batch, 1])). The Concat this fusion produces is the one that takes + // a Tile output. Ort::ConstNode concat_node{nullptr}; for (Ort::ConstNode node : nodes) { - if (IsOnnxOp(node, "Concat")) { - if (concat_node) { - throw std::runtime_error("TileConcat expects one Concat node"); + if (!IsOnnxOp(node, "Concat")) { + continue; + } + bool consumes_tile = false; + for (Ort::ConstValueInfo input : node.GetInputs()) { + auto producer_it = producers.find(Name(input)); + if (producer_it != producers.end() && + IsOnnxOp(producer_it->second, "Tile")) { + consumes_tile = true; + break; } - concat_node = node; } + if (!consumes_tile) { + continue; + } + if (concat_node) { + throw std::runtime_error( + "TileConcat expects one Concat node consuming a Tile output"); + } + concat_node = node; } if (!concat_node) { throw std::runtime_error("TileConcat missing Concat node"); @@ -448,6 +495,33 @@ std::unique_ptr CreateTileConcatFusion( throw std::runtime_error("invalid TileConcat Concat node"); } + // Constant leaves of the repeats chain are resolved here: ORT keeps them in + // the fused graph's initializer list and also declares them as inputs of the + // fused node (drop_constant_initializers is false). Read both so the lookup + // does not depend on which of the two views the runtime populates. + std::unordered_set initializer_names; + std::unordered_map + initializer_values; + auto add_initializer = [&](Ort::ConstValueInfo value_info) { + if (value_info == nullptr || !value_info.IsConstantInitializer()) { + return; + } + const std::string initializer_name = Name(value_info); + if (initializer_name.empty() || initializer_names.count(initializer_name)) { + return; + } + if (auto constant = musa_ep::ReadShapeExprConstant(value_info)) { + initializer_names.insert(initializer_name); + initializer_values.emplace(initializer_name, std::move(*constant)); + } + }; + for (Ort::ConstValueInfo initializer : graph.GetInitializers()) { + add_initializer(initializer); + } + for (Ort::ConstValueInfo input : fused_node.GetInputs()) { + add_initializer(input); + } + std::vector plans; plans.reserve(concat_inputs.size()); for (Ort::ConstValueInfo concat_input : concat_inputs) { @@ -460,14 +534,36 @@ std::unique_ptr CreateTileConcatFusion( if (tile_inputs.size() != 2) { throw std::runtime_error("invalid TileConcat Tile node"); } + const std::string repeats_name = Name(tile_inputs[1]); + if (producers.find(repeats_name) != producers.end()) { + // The matcher absorbed the shape-only chain that produced repeats, so + // the value is not an input of this node any more. Recompute it from + // the shapes of the fused inputs. + musa_ep::ShapeExpr repeats_expr = + musa_ep::BuildShapeExpr(repeats_name, producers, input_indices, + initializer_names, initializer_values); + if (repeats_expr.Empty()) { + throw std::runtime_error( + "TileConcat repeats chain absorbed by the matcher is not " + "recomputable from the fused inputs"); + } + plans.push_back( + {GetIndex(input_indices, Name(tile_inputs[0]), "Tile data input"), + kNoInputIndex, true, std::move(repeats_expr)}); + continue; + } + const size_t repeats_input_index = + GetIndex(input_indices, repeats_name, "Tile repeats input"); + musa_ep::ShapeExpr repeats_expr = + musa_ep::BuildShapeExpr(repeats_name, producers, input_indices, + initializer_names, initializer_values); plans.push_back( {GetIndex(input_indices, Name(tile_inputs[0]), "Tile data input"), - GetIndex(input_indices, Name(tile_inputs[1]), "Tile repeats input"), - true}); + repeats_input_index, true, std::move(repeats_expr)}); } else { plans.push_back( - {GetIndex(input_indices, concat_input_name, "Concat input"), 0, - false}); + {GetIndex(input_indices, concat_input_name, "Concat input"), + kNoInputIndex, false}); } } diff --git a/musa/ep/src/fusion/tile_concat_fusion_matcher.cc b/musa/ep/src/fusion/tile_concat_fusion_matcher.cc index e254a58..f37805a 100644 --- a/musa/ep/src/fusion/tile_concat_fusion_matcher.cc +++ b/musa/ep/src/fusion/tile_concat_fusion_matcher.cc @@ -12,14 +12,7 @@ // limitations under the License. #include -#include -#include #include -#include -#include -#include -#include -#include #include #include #include @@ -28,10 +21,142 @@ #include "fusion/fusion_matcher.h" #include "fusion/fusion_matcher_utils.h" +#include "fusion/shape_expr.h" #include "graph/graph_utils.h" #include "plugin_ep_utils.h" namespace musa_ep { +namespace { + +// The ops a repeats value may be computed from. Every one of them reads only +// shapes and constants (see shape_expr.cc), so the fused node can recompute the +// value from the shapes of its inputs instead of consuming the tensor. +const std::unordered_set& ShapeOnlyOps() { + static const std::unordered_set kOps = {"Shape", + "Gather", + "Unsqueeze", + "Squeeze", + "Concat", + "Add", + "Sub", + "Mul", + "Div", + "Cast", + "Slice", + "Reshape", + "Identity", + "Mod", + "Floor", + "Ceil", + "ReduceProd", + "Constant", + "ConstantOfShape"}; + return kOps; +} + +struct ShapeChain { + std::vector nodes; // inputs precede their consumers + std::unordered_set node_ids; + // Values the chain reads that no node in the graph produces: initializers + // (which carry their tensor) or data leaves (which make the chain unusable). + std::unordered_map leaves; +}; + +bool CollectShapeChain( + const std::string& value_name, + const std::unordered_map& producers, + std::unordered_set& visiting, ShapeChain& chain) { + auto producer_it = producers.find(value_name); + if (producer_it == producers.end()) { + return true; + } + Ort::ConstNode node = producer_it->second; + if (!node || !IsOnnxDomain(node.GetDomain()) || + ShapeOnlyOps().count(node.GetOperatorType()) == 0) { + return false; + } + if (chain.node_ids.count(node.GetId()) != 0) { + return true; + } + if (!visiting.insert(value_name).second) { + return false; // cycle: not a shape chain we can unroll + } + for (Ort::ConstValueInfo input : node.GetInputs()) { + if (input == nullptr) { + continue; + } + const std::string name = Name(input); + if (name.empty()) { + continue; + } + if (producers.find(name) == producers.end()) { + chain.leaves.emplace(name, input); + continue; + } + if (!CollectShapeChain(name, producers, visiting, chain)) { + return false; + } + } + visiting.erase(value_name); + chain.node_ids.insert(node.GetId()); + chain.nodes.push_back(node); + return true; +} + +// The values a fused node would declare as inputs, before any chain is +// absorbed. A Shape source has to be one of them, otherwise absorbing the chain +// would add an input (and possibly a device copy) to the fused node. +std::unordered_map ExternalInputs( + const std::vector& nodes) { + std::unordered_set internal; + for (Ort::ConstNode node : nodes) { + for (Ort::ConstValueInfo output : node.GetOutputs()) { + if (output != nullptr) { + internal.insert(Name(output)); + } + } + } + + std::unordered_map inputs; + for (Ort::ConstNode node : nodes) { + for (Ort::ConstValueInfo input : node.GetInputs()) { + if (input == nullptr) { + continue; + } + const std::string name = Name(input); + if (name.empty() || internal.count(name) != 0) { + continue; + } + inputs.emplace(name, inputs.size()); + } + } + return inputs; +} + +bool ChainHasExternalConsumer( + const ShapeChain& chain, const std::unordered_set& fusion_ids, + const std::unordered_set& graph_output_names) { + std::unordered_set allowed = fusion_ids; + allowed.insert(chain.node_ids.begin(), chain.node_ids.end()); + for (Ort::ConstNode node : chain.nodes) { + for (Ort::ConstValueInfo output : node.GetOutputs()) { + if (output == nullptr) { + continue; + } + if (graph_output_names.count(Name(output)) != 0) { + return true; + } + for (const auto& consumer : output.GetConsumers()) { + if (allowed.count(consumer.node.GetId()) == 0) { + return true; + } + } + } + } + return false; +} + +} // namespace bool CanFuseTileConcat( Ort::ConstNode concat_node, @@ -51,9 +176,7 @@ bool CanFuseTileConcat( return false; } - std::unordered_set selected_node_ids; - fusion_nodes.clear(); - bool found_tile_input = false; + std::vector tile_nodes; for (size_t input_idx = 0; input_idx < concat_inputs.size(); ++input_idx) { auto producer_it = producers.find(Name(concat_inputs[input_idx])); if (producer_it == producers.end() || @@ -73,28 +196,96 @@ bool CanFuseTileConcat( static_cast(input_idx))) { return false; } + tile_nodes.push_back(tile_node); + } + + if (tile_nodes.empty()) { + return false; + } + + // A fused node cannot declare CPU inputs, so ORT copies a runtime repeats + // tensor onto the device. Recompute such a repeats value inside the fused + // node instead: absorb the shape-only chain that produces it, then read only + // the shapes of the fused inputs at run time (no device read, no sync). + std::vector base_nodes = tile_nodes; + base_nodes.push_back(concat_node); + const std::unordered_map base_inputs = + ExternalInputs(base_nodes); - // A fused node cannot declare CPU inputs, so a runtime repeats would be - // copied to the device by ORT's MemcpyTransformer. Leave such a Tile - // unfused: the standalone Tile kernel declares its repeats as an - // OrtMemTypeCPUInput and reads it without a device copy or stream sync. - if (producers.find(Name(tile_inputs[1])) != producers.end()) { + std::unordered_set fusion_ids; + for (Ort::ConstNode node : base_nodes) { + fusion_ids.insert(node.GetId()); + } + + ShapeChain chain; + std::vector chain_roots; + for (Ort::ConstNode tile_node : tile_nodes) { + const std::string repeats = Name(tile_node.GetInputs()[1]); + if (producers.find(repeats) == producers.end()) { + continue; // initializer or graph input: keep reading it as today + } + std::unordered_set visiting; + if (!CollectShapeChain(repeats, producers, visiting, chain)) { return false; } + chain_roots.push_back(repeats); + } - found_tile_input = true; - if (!AddFusionNode(tile_node, accepted_node_ids, selected_node_ids, - fusion_nodes)) { + if (!chain.nodes.empty() && + ChainHasExternalConsumer(chain, fusion_ids, graph_output_names)) { + return false; + } + + // Compile() picks the factory by scanning its detectors in a fixed order, and + // IsShapeReshapeFusionGraph is a loose presence check (Gather + Concat + + // Reshape). The fused body always carries this fusion's Concat, so a chain + // that also brings a Gather and a Reshape would be handed to the ShapeReshape + // factory. Leave such a Tile unfused instead of changing that detector. + if (!chain.nodes.empty()) { + bool has_gather = false; + bool has_reshape = false; + for (Ort::ConstNode node : chain.nodes) { + has_gather = has_gather || IsOnnxOp(node, "Gather"); + has_reshape = has_reshape || node.GetOperatorType() == "Reshape"; + } + if (has_gather && has_reshape) { return false; } } - if (!found_tile_input) { - return false; + if (!chain_roots.empty()) { + std::unordered_set initializer_names; + std::unordered_map initializer_values; + for (const auto& [name, value_info] : chain.leaves) { + auto constant = ReadShapeExprConstant(value_info); + if (!constant.has_value()) { + continue; + } + initializer_names.insert(name); + initializer_values.emplace(name, std::move(*constant)); + } + for (const std::string& repeats : chain_roots) { + if (BuildShapeExpr(repeats, producers, base_inputs, initializer_names, + initializer_values) + .Empty()) { + return false; + } + } } - if (!AddFusionNode(concat_node, accepted_node_ids, selected_node_ids, - fusion_nodes)) { - return false; + + fusion_nodes.clear(); + std::unordered_set selected_node_ids; + for (Ort::ConstNode node : base_nodes) { + if (!AddFusionNode(node, accepted_node_ids, selected_node_ids, + fusion_nodes)) { + return false; + } + } + for (Ort::ConstNode node : chain.nodes) { + if (!AddFusionNode(node, accepted_node_ids, selected_node_ids, + fusion_nodes)) { + return false; + } } std::sort(fusion_nodes.begin(), fusion_nodes.end(), [](Ort::ConstNode lhs, Ort::ConstNode rhs) { diff --git a/test/fusion/test_tile_concat_shape_chain.py b/test/fusion/test_tile_concat_shape_chain.py new file mode 100644 index 0000000..476c8e5 --- /dev/null +++ b/test/fusion/test_tile_concat_shape_chain.py @@ -0,0 +1,259 @@ +# Copyright (c) Moore Threads Technology Co., Ltd. All rights reserved. +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Tile repeats computed from a shape-only chain, absorbed into Tile+Concat. + +A fused node cannot declare a CPU input, so ORT copies a runtime repeats tensor +onto the device and the kernel would have to read it back with a stream +synchronize. Instead the matcher absorbs the Shape/Gather/... chain and the +kernel recomputes repeats from the shapes of its inputs. These tests cover the +cases that make that absorption observable: a repeats value that follows +another tensor's shape, a chain value with a consumer outside the fusion, and +inputs the evaluator has to reject. +""" + +import json +import os + +import numpy as np +import onnxruntime as ort +import pytest +from onnx import TensorProto, helper + +from op_test_utils import musa_devices + +INT64_MAX = 2**63 - 1 + + +def _build_model(name, nodes, inputs, outputs, initializers, opset=17): + graph = helper.make_graph(nodes, name, inputs, outputs, initializer=initializers) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", opset)]) + model.ir_version = min(model.ir_version, 10) + return model.SerializeToString() + + +def _int64_initializer(name, values, shape=None): + array = np.asarray(values, dtype=np.int64) + if shape is not None: + array = array.reshape(shape) + return helper.make_tensor( + name, TensorProto.INT64, list(array.shape), array.reshape(-1).tolist() + ) + + +def _float_initializer(name, values, shape): + array = np.asarray(values, dtype=np.float32).reshape(shape) + return helper.make_tensor( + name, TensorProto.FLOAT, list(array.shape), array.reshape(-1).tolist() + ) + + +def _session(model, use_musa, profile_prefix=None): + if not use_musa: + return ort.InferenceSession(model, providers=["CPUExecutionProvider"]) + devices = musa_devices() + assert devices, "no MUSA device: refusing a CPU-only session for the MUSA run" + options = ort.SessionOptions() + options.add_session_config_entry("session.disable_cpu_ep_fallback", "1") + if profile_prefix is not None: + options.enable_profiling = True + options.profile_file_prefix = profile_prefix + options.add_provider_for_devices(devices, {}) + return ort.InferenceSession(model, sess_options=options) + + +def _read_ops_by_provider(session): + profile_path = session.end_profiling() + try: + with open(profile_path, "r", encoding="utf-8") as handle: + events = json.load(handle) + finally: + if os.path.exists(profile_path): + os.remove(profile_path) + ops = {} + for event in events: + if event.get("cat") != "Node" or not event.get("name", "").endswith( + "_kernel_time" + ): + continue + args = event.get("args", {}) + ops.setdefault(args.get("provider"), set()).add(args.get("op_name")) + return ops + + +def _musa_ops(ops): + return ops.get("MUSAExecutionProvider", set()) + + +def _fused_ops(ops): + return { + op for op in _musa_ops(ops) if str(op).startswith("MUSAExecutionProvider_") + } + + +def _assert_tile_concat_absorbed(ops): + """The Tile and Concat ran inside the fused node, chain and all.""" + assert _fused_ops(ops) + assert "Tile" not in _musa_ops(ops) + assert "Concat" not in _musa_ops(ops) + + +def _batch_chain_nodes(repeats_source): + return [ + helper.make_node("Shape", ["query"], ["query_shape"]), + helper.make_node("Gather", ["query_shape", "zero"], ["batch"], axis=0), + *repeats_source, + helper.make_node("Tile", ["data", "repeats"], ["tiled"]), + helper.make_node("Concat", ["tiled", "query"], ["y"], axis=1), + ] + + +def _batch_chain_initializers(): + return [ + _int64_initializer("zero", [0]), + _int64_initializer("one", [1]), + _float_initializer("data", np.arange(4, dtype=np.float32), (1, 4)), + ] + + +def test_tile_concat_repeats_follows_another_tensors_shape(tmp_path): + """One session, two batch sizes: repeats moves while the Tile input does not. + + The Tile data stays [1, 4] and only `query`'s batch axis changes. The matcher + absorbs the chain, so repeats is no longer an input of the fused node and the + expression is the only way to obtain it. + """ + rng = np.random.default_rng(29) + nodes = _batch_chain_nodes( + [ + helper.make_node("Unsqueeze", ["batch", "zero"], ["batch_1d"]), + helper.make_node("Concat", ["batch_1d", "one"], ["repeats"], axis=0), + ] + ) + model = _build_model( + "tile_concat_repeats_shape_chain", + nodes, + [helper.make_tensor_value_info("query", TensorProto.FLOAT, ["batch", 4])], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, None)], + _batch_chain_initializers(), + ) + + cpu = _session(model, use_musa=False) + musa = _session(model, use_musa=True, profile_prefix=str(tmp_path / "shape_chain")) + for batch in (1, 3): + query = rng.standard_normal((batch, 4)).astype(np.float32) + expected = cpu.run(None, {"query": query})[0] + actual = musa.run(None, {"query": query})[0] + assert actual.shape == expected.shape == (batch, 8) + np.testing.assert_allclose(actual, expected, rtol=1e-6, atol=1e-6) + + _assert_tile_concat_absorbed(_read_ops_by_provider(musa)) + + +def test_tile_concat_repeats_chain_value_used_outside_fusion(): + """A chain value that is also a graph output must not be absorbed. + + A fused node owns its internal values, so absorbing a value another consumer + reads would either break that consumer or change the graph. The Tile falls + back to the standalone kernel (which takes repeats as a CPU input), the chain + stays in the graph, and both outputs stay correct. + """ + rng = np.random.default_rng(31) + nodes = _batch_chain_nodes( + [ + helper.make_node("Unsqueeze", ["batch", "zero"], ["batch_1d"]), + helper.make_node("Concat", ["batch_1d", "one"], ["repeats"], axis=0), + ] + ) + model = _build_model( + "tile_concat_repeats_chain_external_consumer", + nodes, + [helper.make_tensor_value_info("query", TensorProto.FLOAT, ["batch", 4])], + [ + helper.make_tensor_value_info("y", TensorProto.FLOAT, None), + helper.make_tensor_value_info("batch", TensorProto.INT64, None), + ], + _batch_chain_initializers(), + ) + + query = rng.standard_normal((2, 4)).astype(np.float32) + feeds = {"query": query} + expected = _session(model, use_musa=False).run(None, feeds) + actual = _session(model, use_musa=True).run(None, feeds) + assert len(actual) == 2 + np.testing.assert_allclose(actual[0], expected[0], rtol=1e-6, atol=1e-6) + np.testing.assert_array_equal(actual[1], expected[1]) + assert actual[1] == np.int64(2) + + +def test_tile_concat_repeats_chain_with_negative_slice(): + """Negative Slice bounds are resolved against the runtime shape.""" + rng = np.random.default_rng(37) + nodes = [ + helper.make_node("Shape", ["query"], ["query_shape"]), + helper.make_node( + "Slice", ["query_shape", "head_start", "head_end"], ["batch_dim"] + ), + helper.make_node("Concat", ["batch_dim", "one"], ["repeats"], axis=0), + helper.make_node("Tile", ["data", "repeats"], ["tiled"]), + helper.make_node("Concat", ["tiled", "query"], ["y"], axis=1), + ] + initializers = [ + _float_initializer("data", np.arange(4, dtype=np.float32), (1, 4)), + _int64_initializer("one", [1]), + _int64_initializer("head_start", [-2]), + _int64_initializer("head_end", [-1]), + ] + model = _build_model( + "tile_concat_repeats_negative_slice", + nodes, + [helper.make_tensor_value_info("query", TensorProto.FLOAT, ["batch", 4])], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, None)], + initializers, + ) + + for batch in (2, 5): + query = rng.standard_normal((batch, 4)).astype(np.float32) + feeds = {"query": query} + expected = _session(model, use_musa=False).run(None, feeds)[0] + actual = _session(model, use_musa=True).run(None, feeds)[0] + assert actual.shape == expected.shape == (batch, 8) + np.testing.assert_allclose(actual, expected, rtol=1e-6, atol=1e-6) + + +def test_tile_concat_repeats_chain_with_reshape_element_mismatch(): + """An unusable shape value fails loudly instead of running on stale data.""" + nodes = [ + helper.make_node("Shape", ["query"], ["query_shape"]), + helper.make_node("Reshape", ["query_shape", "three"], ["bad_repeats"]), + helper.make_node("Tile", ["data", "bad_repeats"], ["tiled"]), + helper.make_node("Concat", ["tiled", "query"], ["y"], axis=1), + ] + initializers = [ + _float_initializer("data", np.arange(2, dtype=np.float32), (1, 2)), + _int64_initializer("three", [3]), + ] + model = _build_model( + "tile_concat_repeats_reshape_mismatch", + nodes, + [helper.make_tensor_value_info("query", TensorProto.FLOAT, [1, 4])], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, None)], + initializers, + ) + + feeds = {"query": np.zeros((1, 4), dtype=np.float32)} + # Shape(query) holds two entries and the Reshape asks for three. The graph is + # valid ONNX, so the failure has to surface at run time, not as a crash. + with pytest.raises(Exception): + _session(model, use_musa=False).run(None, feeds) + with pytest.raises(Exception): + _session(model, use_musa=True).run(None, feeds) From 3c47ddc55a32034b0d46e2cbb58d305067d6c833 Mon Sep 17 00:00:00 2001 From: "yichao.hu" Date: Sun, 18 Oct 2026 13:31:40 -0400 Subject: [PATCH 27/52] fix(fusion): recompute Tile repeats without absorbing the chain 209e500 absorbed the shape-only chain that produces Tile repeats into the Tile+Concat fusion and vetoed the fusion when any chain value had a consumer outside it. In the four recommendation models in the model zoo every such chain value is also read by another branch (MathConcatLog and friends), so all 140/140/136/140 Tiles fell back and the fusion never ran. Absorb nothing: the fused node is Tile + Concat again, the chain stays in the graph, and repeats stays an input of the fused node. The matcher still proves that repeats is a pure function of the shapes of the fused inputs (shape-only whitelist, every Shape source already an input of Tile+Concat, every constant readable) and hands that expression to the fused node, which evaluates it from the input shapes instead of reading the tensor. That keeps the device-to-host copy and the compute-stream synchronize out of the steady state while the tensor remains the fallback. A Tile whose repeats cannot be proven keeps its standalone kernel and becomes an input of the fused Concat; per-Tile fallback replaces rejecting the whole Concat. Compile() only receives the fused subgraph, so it cannot rebuild the expression: GetCapability() stores the proofs in the EP keyed by the fusion boundary (its input and output value names), and Compile() maps the recorded source names onto the input indices of the fused node ORT built. Two vetoes disappear with the absorption: - a chain value consumed outside the fusion no longer vetoes anything, - a chain containing Gather + Reshape no longer vetoes anything. That check kept the fused body away from IsShapeReshapeFusionGraph, a loose presence check (Gather + Concat + Reshape) that Compile() runs before the TileConcat detector: while the chain was absorbed, the fused body carried the chain's Gather and Reshape next to this fusion's Concat and was handed to the ShapeReshape factory. The body now holds only Tile + Concat, so that collision cannot arise. Session creation now prints one MUSA_TILE_CONCAT_REPEATS line per runtime repeats value (shape_expr=1 when the expression arrived, 0 when the fused node will read the tensor) so the verification round can see which path ran. DoD: clean build on the test host (SDK 4.3.8) with no new compiler warnings; the .so sha256 is fc422405abfae83bf8251b812d4063cb6d85ecf40880c2fcf4cbbcd2698cd188. Unit tests are written but not executed: no GPU window this round. The 123 absorbed / 13-17 fallback counts are a static estimate and are not verified on the models. H1 (the recommendation-model residual comes from dropping this fusion) remains an unverified hypothesis. --- musa/ep/src/ep.cc | 8 +- musa/ep/src/ep.h | 17 ++ musa/ep/src/fusion/fusion_matcher.cc | 7 +- musa/ep/src/fusion/fusion_matcher.h | 16 +- musa/ep/src/fusion/fusion_node_compute.cc | 3 +- musa/ep/src/fusion/shape_expr.h | 12 + musa/ep/src/fusion/tile_concat_fusion.cc | 113 +++++-- musa/ep/src/fusion/tile_concat_fusion.h | 16 +- .../src/fusion/tile_concat_fusion_matcher.cc | 206 +++++++------ test/fusion/test_tile_concat_shape_chain.py | 282 +++++++++++++++--- 10 files changed, 513 insertions(+), 167 deletions(-) diff --git a/musa/ep/src/ep.cc b/musa/ep/src/ep.cc index 515cf60..66a01cc 100644 --- a/musa/ep/src/ep.cc +++ b/musa/ep/src/ep.cc @@ -377,8 +377,12 @@ MusaEp::GetCapabilityImpl(OrtEp* this_ptr, const OrtGraph* ort_graph, std::unordered_set graph_output_names = GetGraphOutputNames(graph); std::vector fusion_matches; + // Filled while the whole graph is still visible: Compile() only gets the + // fused subgraph, which no longer holds the chain a shape value comes from. + FusionShapeExprHints shape_expr_hints; if (!DisableAllFusions()) { - fusion_matches = FindFusionMatches(all_nodes, graph_output_names); + fusion_matches = + FindFusionMatches(all_nodes, graph_output_names, &shape_expr_hints); if (DisableShapeReshapeFusion()) { fusion_matches.erase( std::remove_if(fusion_matches.begin(), fusion_matches.end(), @@ -396,7 +400,9 @@ MusaEp::GetCapabilityImpl(OrtEp* this_ptr, const OrtGraph* ort_graph, "FindTileConcatFusions"; }), fusion_matches.end()); + shape_expr_hints.clear(); } + ep->AddShapeExprHints(std::move(shape_expr_hints)); RETURN_IF_ERROR(RegisterFusionMatches(ep->ep_api_, graph_support_info, fusion_matches)); } diff --git a/musa/ep/src/ep.h b/musa/ep/src/ep.h index 365d474..df59131 100644 --- a/musa/ep/src/ep.h +++ b/musa/ep/src/ep.h @@ -22,6 +22,7 @@ #include #include "ep_options.h" +#include "fusion/fusion_matcher.h" class FusionNodeCompute; class MusaEpFactory; @@ -47,6 +48,21 @@ class MusaEp : public OrtEp { return fusion_computes_; } + // Shape expressions a finder proved while it could still see the whole graph. + // Compile() only receives fused subgraphs, so GetCapability() stores them + // here. + const musa_ep::FusionShapeExprHints& ShapeExprHints() const { + return shape_expr_hints_; + } + void AddShapeExprHints(musa_ep::FusionShapeExprHints hints) { + for (auto& [key, values] : hints) { + auto& stored = shape_expr_hints_[key]; + for (auto& [name, hint] : values) { + stored.emplace(name, std::move(hint)); + } + } + } + private: static const char* ORT_API_CALL GetNameImpl(const OrtEp* this_ptr) noexcept; @@ -83,4 +99,5 @@ class MusaEp : public OrtEp { const OrtLogger& logger_; std::unordered_map> fusion_computes_; + musa_ep::FusionShapeExprHints shape_expr_hints_; }; diff --git a/musa/ep/src/fusion/fusion_matcher.cc b/musa/ep/src/fusion/fusion_matcher.cc index 51623eb..f1371f1 100644 --- a/musa/ep/src/fusion/fusion_matcher.cc +++ b/musa/ep/src/fusion/fusion_matcher.cc @@ -148,7 +148,8 @@ void AddFusionMatch(std::vector& matches, const char* finder, std::vector FindFusionMatches( const std::vector& all_nodes, - const std::unordered_set& graph_output_names) { + const std::unordered_set& graph_output_names, + FusionShapeExprHints* shape_expr_hints) { std::unordered_set accepted_node_ids; std::vector matches; @@ -326,8 +327,8 @@ std::vector FindFusionMatches( AddFusionMatch(matches, "FindShapeReshapeFusions", true, std::move(shape_reshape_fusions), accepted_node_ids); - auto tile_concat_fusions = - FindTileConcatFusions(all_nodes, graph_output_names, accepted_node_ids); + auto tile_concat_fusions = FindTileConcatFusions( + all_nodes, graph_output_names, accepted_node_ids, shape_expr_hints); AddFusionMatch(matches, "FindTileConcatFusions", false, std::move(tile_concat_fusions), accepted_node_ids); diff --git a/musa/ep/src/fusion/fusion_matcher.h b/musa/ep/src/fusion/fusion_matcher.h index 6129c5d..be7df39 100644 --- a/musa/ep/src/fusion/fusion_matcher.h +++ b/musa/ep/src/fusion/fusion_matcher.h @@ -18,13 +18,23 @@ #undef ORT_API_MANUAL_INIT #include +#include #include #include #include "fusion/fusion_dtype.h" +#include "fusion/shape_expr.h" namespace musa_ep { +// Shape expressions a finder proved on the whole graph, keyed first by a +// signature of the fusion boundary and then by the outer value name they +// compute. GetCapability() fills this in; Compile() reads it back because the +// subgraph it receives does not contain the nodes the value is computed from. +using FusionShapeExprHints = + std::unordered_map>; + struct FusionDTypeRejection { std::vector nodes; std::string reason; @@ -40,7 +50,8 @@ struct FusionMatch { std::vector FindFusionMatches( const std::vector& all_nodes, - const std::unordered_set& graph_output_names); + const std::unordered_set& graph_output_names, + FusionShapeExprHints* shape_expr_hints = nullptr); std::vector> FindSplitUnsqueezeConcatFusions( const std::vector& all_nodes, @@ -65,7 +76,8 @@ std::vector> FindSliceConcatFusions( std::vector> FindTileConcatFusions( const std::vector& all_nodes, const std::unordered_set& graph_output_names, - const std::unordered_set& accepted_node_ids); + const std::unordered_set& accepted_node_ids, + FusionShapeExprHints* shape_expr_hints = nullptr); std::vector> FindGemmActivationFusions( const std::vector& all_nodes, const std::unordered_set& graph_output_names, diff --git a/musa/ep/src/fusion/fusion_node_compute.cc b/musa/ep/src/fusion/fusion_node_compute.cc index 87c46ce..6f3eea4 100644 --- a/musa/ep/src/fusion/fusion_node_compute.cc +++ b/musa/ep/src/fusion/fusion_node_compute.cc @@ -420,7 +420,8 @@ OrtStatus* ORT_API_CALL MusaEp::CompileImpl( fusion_compute = CreateShapeReshapeFusion(graph, fused_node); finder_name = "FindShapeReshapeFusions"; } else if (IsTileConcatFusionGraph(graph)) { - fusion_compute = CreateTileConcatFusion(graph, fused_node); + fusion_compute = + CreateTileConcatFusion(graph, fused_node, &ep->ShapeExprHints()); finder_name = "FindTileConcatFusions"; } else if (IsSliceConcatFusionGraph(graph)) { fusion_compute = CreateSliceConcatFusion(graph, fused_node); diff --git a/musa/ep/src/fusion/shape_expr.h b/musa/ep/src/fusion/shape_expr.h index f75f9df..79cfdaf 100644 --- a/musa/ep/src/fusion/shape_expr.h +++ b/musa/ep/src/fusion/shape_expr.h @@ -54,6 +54,18 @@ struct ShapeExpr { bool Empty() const { return root < 0; } }; +// An expression computed before the fused node exists, carried to Compile(). +// +// GetCapability() sees the whole graph, so that is where a chain outside the +// fusion can be resolved. Compile() only receives the fused subgraph, so the +// proof has to travel with the fusion: `source_input_names[i]` is the name of +// the fused input whose shape fills source slot i, and Compile() translates +// those names into the indices of the fused node that ORT built. +struct ShapeExprHint { + ShapeExpr expr; + std::vector source_input_names; +}; + // A constant tensor that a shape chain reads: an int32/int64 initializer, or // the tensor carried by a Constant node. struct ShapeExprConstant { diff --git a/musa/ep/src/fusion/tile_concat_fusion.cc b/musa/ep/src/fusion/tile_concat_fusion.cc index 6b2ca48..b214c91 100644 --- a/musa/ep/src/fusion/tile_concat_fusion.cc +++ b/musa/ep/src/fusion/tile_concat_fusion.cc @@ -15,6 +15,7 @@ #include #include +#include #include #include #include @@ -42,14 +43,14 @@ constexpr size_t kNoInputIndex = std::numeric_limits::max(); struct ConcatInputPlan { size_t data_input_index; - // Set unless the repeats producer was absorbed into the fused node; then the - // tensor is not an input of the fused node and the expression is the only way - // to obtain the value. + // Index of the repeats input, or kNoInputIndex for a Concat input that is not + // a Tile. size_t repeats_input_index; bool is_tile; // Non-empty when repeats was proven to be a pure function of the shapes of // the fused inputs: recompute it instead of reading the device tensor, which - // would cost a device-to-host copy plus a compute-stream synchronize. + // would cost a device-to-host copy plus a compute-stream synchronize. The + // tensor stays an input, so a failed evaluation can still fall back to it. musa_ep::ShapeExpr repeats_expr; }; @@ -440,8 +441,45 @@ std::unordered_map BuildProducers( return producers; } +musa_ep::ShapeExpr RemapShapeExprHint( + const musa_ep::ShapeExprHint& hint, + const std::unordered_map& input_indices) { + if (hint.expr.Empty() || + hint.source_input_names.size() != hint.expr.source_input_indices.size()) { + return musa_ep::ShapeExpr{}; + } + musa_ep::ShapeExpr expr = hint.expr; + for (size_t slot = 0; slot < hint.source_input_names.size(); ++slot) { + auto index_it = input_indices.find(hint.source_input_names[slot]); + if (index_it == input_indices.end()) { + return musa_ep::ShapeExpr{}; // source is not a fused input after all + } + expr.source_input_indices[slot] = index_it->second; + } + return expr; +} + } // namespace +std::string TileConcatFusionKey(const std::vector& input_names, + const std::vector& output_names) { + auto normalize = [](std::vector names) { + names.erase( + std::remove_if(names.begin(), names.end(), + [](const std::string& name) { return name.empty(); }), + names.end()); + std::sort(names.begin(), names.end()); + names.erase(std::unique(names.begin(), names.end()), names.end()); + std::string key; + for (const std::string& name : names) { + key += name; + key += '\x1e'; + } + return key; + }; + return normalize(input_names) + "|" + normalize(output_names); +} + bool IsTileConcatFusionGraph(Ort::ConstGraph graph) { bool has_concat = false; bool has_tile = false; @@ -453,7 +491,8 @@ bool IsTileConcatFusionGraph(Ort::ConstGraph graph) { } std::unique_ptr CreateTileConcatFusion( - Ort::ConstGraph graph, Ort::ConstNode fused_node) { + Ort::ConstGraph graph, Ort::ConstNode fused_node, + const musa_ep::FusionShapeExprHints* shape_expr_hints) { std::vector nodes = graph.GetNodes(); auto producers = BuildProducers(nodes); auto input_indices = FusedInputIndices(fused_node); @@ -522,6 +561,30 @@ std::unique_ptr CreateTileConcatFusion( add_initializer(input); } + // The repeats expressions were proved on the whole graph in GetCapability(); + // this subgraph only holds Tile + Concat, so look them up by the fusion + // boundary. Both sides build the same key from names that ORT keeps. + const std::unordered_map* repeats_hints = + nullptr; + if (shape_expr_hints != nullptr) { + std::vector boundary_names; + for (Ort::ConstValueInfo input : fused_node.GetInputs()) { + boundary_names.push_back(Name(input)); + } + for (const std::string& initializer_name : initializer_names) { + boundary_names.push_back(initializer_name); + } + std::vector output_names; + for (Ort::ConstValueInfo output : fused_node.GetOutputs()) { + output_names.push_back(Name(output)); + } + auto hint_it = shape_expr_hints->find( + TileConcatFusionKey(boundary_names, output_names)); + if (hint_it != shape_expr_hints->end()) { + repeats_hints = &hint_it->second; + } + } + std::vector plans; plans.reserve(concat_inputs.size()); for (Ort::ConstValueInfo concat_input : concat_inputs) { @@ -535,28 +598,30 @@ std::unique_ptr CreateTileConcatFusion( throw std::runtime_error("invalid TileConcat Tile node"); } const std::string repeats_name = Name(tile_inputs[1]); - if (producers.find(repeats_name) != producers.end()) { - // The matcher absorbed the shape-only chain that produced repeats, so - // the value is not an input of this node any more. Recompute it from - // the shapes of the fused inputs. - musa_ep::ShapeExpr repeats_expr = + const size_t repeats_input_index = + GetIndex(input_indices, repeats_name, "Tile repeats input"); + musa_ep::ShapeExpr repeats_expr; + if (repeats_hints != nullptr) { + auto hint_it = repeats_hints->find(repeats_name); + if (hint_it != repeats_hints->end()) { + repeats_expr = RemapShapeExprHint(hint_it->second, input_indices); + } + } + if (repeats_expr.Empty()) { + // Secondary source of truth: a chain ORT did keep inside the fusion. + repeats_expr = musa_ep::BuildShapeExpr(repeats_name, producers, input_indices, initializer_names, initializer_values); - if (repeats_expr.Empty()) { - throw std::runtime_error( - "TileConcat repeats chain absorbed by the matcher is not " - "recomputable from the fused inputs"); - } - plans.push_back( - {GetIndex(input_indices, Name(tile_inputs[0]), "Tile data input"), - kNoInputIndex, true, std::move(repeats_expr)}); - continue; } - const size_t repeats_input_index = - GetIndex(input_indices, repeats_name, "Tile repeats input"); - musa_ep::ShapeExpr repeats_expr = - musa_ep::BuildShapeExpr(repeats_name, producers, input_indices, - initializer_names, initializer_values); + if (initializer_names.count(repeats_name) == 0) { + // One line per runtime repeats at session creation: it tells the + // verification round whether the matcher's expression arrived. A zero + // is correct but slower - the kernel then reads the tensor back with a + // device-to-host copy and a stream synchronize on every run. + std::fprintf(stderr, + "MUSA_TILE_CONCAT_REPEATS repeats=%s shape_expr=%d\n", + repeats_name.c_str(), repeats_expr.Empty() ? 0 : 1); + } plans.push_back( {GetIndex(input_indices, Name(tile_inputs[0]), "Tile data input"), repeats_input_index, true, std::move(repeats_expr)}); diff --git a/musa/ep/src/fusion/tile_concat_fusion.h b/musa/ep/src/fusion/tile_concat_fusion.h index dcb3a92..95e815f 100644 --- a/musa/ep/src/fusion/tile_concat_fusion.h +++ b/musa/ep/src/fusion/tile_concat_fusion.h @@ -18,9 +18,23 @@ #undef ORT_API_MANUAL_INIT #include +#include +#include +#include "fusion/fusion_matcher.h" #include "fusion/fusion_node_compute.h" bool IsTileConcatFusionGraph(Ort::ConstGraph graph); + +// Identifies one Tile+Concat fusion across GetCapability() and Compile(): both +// sides see the fusion boundary (its input and output value names) but not each +// other's nodes. Names are sorted and de-duplicated so the order ORT chooses +// for the fused node's inputs does not matter. +std::string TileConcatFusionKey(const std::vector& input_names, + const std::vector& output_names); + +// `shape_expr_hints` carries the repeats expressions the matcher proved on the +// whole graph; Compile() cannot see that chain any more. std::unique_ptr CreateTileConcatFusion( - Ort::ConstGraph graph, Ort::ConstNode fused_node); + Ort::ConstGraph graph, Ort::ConstNode fused_node, + const musa_ep::FusionShapeExprHints* shape_expr_hints = nullptr); diff --git a/musa/ep/src/fusion/tile_concat_fusion_matcher.cc b/musa/ep/src/fusion/tile_concat_fusion_matcher.cc index f37805a..7e1d41b 100644 --- a/musa/ep/src/fusion/tile_concat_fusion_matcher.cc +++ b/musa/ep/src/fusion/tile_concat_fusion_matcher.cc @@ -22,6 +22,7 @@ #include "fusion/fusion_matcher.h" #include "fusion/fusion_matcher_utils.h" #include "fusion/shape_expr.h" +#include "fusion/tile_concat_fusion.h" #include "graph/graph_utils.h" #include "plugin_ep_utils.h" @@ -55,7 +56,6 @@ const std::unordered_set& ShapeOnlyOps() { } struct ShapeChain { - std::vector nodes; // inputs precede their consumers std::unordered_set node_ids; // Values the chain reads that no node in the graph produces: initializers // (which carry their tensor) or data leaves (which make the chain unusable). @@ -99,13 +99,12 @@ bool CollectShapeChain( } visiting.erase(value_name); chain.node_ids.insert(node.GetId()); - chain.nodes.push_back(node); return true; } -// The values a fused node would declare as inputs, before any chain is -// absorbed. A Shape source has to be one of them, otherwise absorbing the chain -// would add an input (and possibly a device copy) to the fused node. +// The values a fused node declares as inputs. A Shape source has to be one of +// them: the fused node cannot gain an input that Tile+Concat do not already +// read, so a chain reaching outside would not be recomputable at run time. std::unordered_map ExternalInputs( const std::vector& nodes) { std::unordered_set internal; @@ -133,29 +132,6 @@ std::unordered_map ExternalInputs( return inputs; } -bool ChainHasExternalConsumer( - const ShapeChain& chain, const std::unordered_set& fusion_ids, - const std::unordered_set& graph_output_names) { - std::unordered_set allowed = fusion_ids; - allowed.insert(chain.node_ids.begin(), chain.node_ids.end()); - for (Ort::ConstNode node : chain.nodes) { - for (Ort::ConstValueInfo output : node.GetOutputs()) { - if (output == nullptr) { - continue; - } - if (graph_output_names.count(Name(output)) != 0) { - return true; - } - for (const auto& consumer : output.GetConsumers()) { - if (allowed.count(consumer.node.GetId()) == 0) { - return true; - } - } - } - } - return false; -} - } // namespace bool CanFuseTileConcat( @@ -163,7 +139,8 @@ bool CanFuseTileConcat( const std::unordered_map& producers, const std::unordered_set& graph_output_names, const std::unordered_set& accepted_node_ids, - std::vector& fusion_nodes) { + std::vector& fusion_nodes, + FusionShapeExprHints* shape_expr_hints) { if (!IsOnnxOp(concat_node, "Concat") || accepted_node_ids.count(concat_node.GetId()) != 0) { return false; @@ -204,59 +181,37 @@ bool CanFuseTileConcat( } // A fused node cannot declare CPU inputs, so ORT copies a runtime repeats - // tensor onto the device. Recompute such a repeats value inside the fused - // node instead: absorb the shape-only chain that produces it, then read only - // the shapes of the fused inputs at run time (no device read, no sync). - std::vector base_nodes = tile_nodes; - base_nodes.push_back(concat_node); - const std::unordered_map base_inputs = - ExternalInputs(base_nodes); - - std::unordered_set fusion_ids; - for (Ort::ConstNode node : base_nodes) { - fusion_ids.insert(node.GetId()); - } - - ShapeChain chain; - std::vector chain_roots; - for (Ort::ConstNode tile_node : tile_nodes) { - const std::string repeats = Name(tile_node.GetInputs()[1]); + // tensor onto the device. Do not absorb the shape-only chain that produces + // it: the dnn models also read its intermediate values from other branches + // (MathConcatLog and friends), a fused node owns every value it absorbs, and + // such a chain would have vetoed the fusion outright. Keep the chain in the + // graph, leave `repeats` an input of the fused node, and prove here that the + // value is a pure function of the shapes of the fused inputs. The fused node + // then recomputes it from those shapes (no device read, no stream sync) and + // still falls back to reading the tensor if evaluation ever fails. + // + // Compile() only receives the Tile+Concat subgraph, so the proof is handed + // over through `shape_expr_hints`: GetCapability() is the last point that can + // see the chain. A Tile whose repeats cannot be proven shape-only keeps its + // standalone kernel and becomes an input of the fused Concat. + std::unordered_set initializer_names; + std::unordered_map initializer_values; + std::vector recomputable(tile_nodes.size(), true); + for (size_t i = 0; i < tile_nodes.size(); ++i) { + const std::string repeats = Name(tile_nodes[i].GetInputs()[1]); if (producers.find(repeats) == producers.end()) { continue; // initializer or graph input: keep reading it as today } + ShapeChain chain; std::unordered_set visiting; if (!CollectShapeChain(repeats, producers, visiting, chain)) { - return false; - } - chain_roots.push_back(repeats); - } - - if (!chain.nodes.empty() && - ChainHasExternalConsumer(chain, fusion_ids, graph_output_names)) { - return false; - } - - // Compile() picks the factory by scanning its detectors in a fixed order, and - // IsShapeReshapeFusionGraph is a loose presence check (Gather + Concat + - // Reshape). The fused body always carries this fusion's Concat, so a chain - // that also brings a Gather and a Reshape would be handed to the ShapeReshape - // factory. Leave such a Tile unfused instead of changing that detector. - if (!chain.nodes.empty()) { - bool has_gather = false; - bool has_reshape = false; - for (Ort::ConstNode node : chain.nodes) { - has_gather = has_gather || IsOnnxOp(node, "Gather"); - has_reshape = has_reshape || node.GetOperatorType() == "Reshape"; - } - if (has_gather && has_reshape) { - return false; + recomputable[i] = false; // not a shape-only chain + continue; } - } - - if (!chain_roots.empty()) { - std::unordered_set initializer_names; - std::unordered_map initializer_values; for (const auto& [name, value_info] : chain.leaves) { + if (initializer_names.count(name) != 0) { + continue; + } auto constant = ReadShapeExprConstant(value_info); if (!constant.has_value()) { continue; @@ -264,47 +219,118 @@ bool CanFuseTileConcat( initializer_names.insert(name); initializer_values.emplace(name, std::move(*constant)); } - for (const std::string& repeats : chain_roots) { - if (BuildShapeExpr(repeats, producers, base_inputs, initializer_names, - initializer_values) - .Empty()) { - return false; + } + + // Dropping a Tile shrinks the fusion boundary, which can invalidate another + // Tile whose Shape source lived on the dropped Tile's input. Rebuild until + // the set stops shrinking. + std::vector fused_tiles; + std::unordered_map repeats_hints; + for (;;) { + fused_tiles.clear(); + for (size_t i = 0; i < tile_nodes.size(); ++i) { + if (recomputable[i]) { + fused_tiles.push_back(tile_nodes[i]); } } + if (fused_tiles.empty()) { + return false; + } + + std::vector base_nodes = fused_tiles; + base_nodes.push_back(concat_node); + const std::unordered_map base_inputs = + ExternalInputs(base_nodes); + std::vector input_names(base_inputs.size()); + for (const auto& [name, index] : base_inputs) { + input_names[index] = name; + } + + repeats_hints.clear(); + bool dropped = false; + for (size_t i = 0; i < tile_nodes.size(); ++i) { + if (!recomputable[i]) { + continue; + } + const std::string repeats = Name(tile_nodes[i].GetInputs()[1]); + if (producers.find(repeats) == producers.end()) { + continue; + } + ShapeExpr expr = BuildShapeExpr(repeats, producers, base_inputs, + initializer_names, initializer_values); + if (expr.Empty()) { + // A Shape source that is not an input of Tile+Concat cannot be + // recomputed inside the fusion. + recomputable[i] = false; + dropped = true; + continue; + } + ShapeExprHint hint; + hint.source_input_names.reserve(expr.source_input_indices.size()); + for (size_t source_index : expr.source_input_indices) { + hint.source_input_names.push_back(input_names[source_index]); + } + hint.expr = std::move(expr); + repeats_hints.emplace(repeats, std::move(hint)); + } + if (!dropped) { + break; + } } + // Compile() picks the factory by scanning its detectors in a fixed order, and + // IsShapeReshapeFusionGraph is a loose presence check (Gather + Concat + + // Reshape). That used to matter: while the repeats chain was absorbed, the + // fused body carried the chain's Gather and Reshape next to this fusion's + // Concat, so such a body was handed to the ShapeReshape factory and the Tile + // was left unfused. The chain now stays in the graph, the fused body holds + // only Tile + Concat, and the collision cannot arise. + const std::vector fusion_outputs = {Name(concat_outputs[0])}; + fusion_nodes.clear(); std::unordered_set selected_node_ids; - for (Ort::ConstNode node : base_nodes) { + for (Ort::ConstNode node : fused_tiles) { if (!AddFusionNode(node, accepted_node_ids, selected_node_ids, fusion_nodes)) { return false; } } - for (Ort::ConstNode node : chain.nodes) { - if (!AddFusionNode(node, accepted_node_ids, selected_node_ids, - fusion_nodes)) { - return false; - } + if (!AddFusionNode(concat_node, accepted_node_ids, selected_node_ids, + fusion_nodes)) { + return false; } std::sort(fusion_nodes.begin(), fusion_nodes.end(), [](Ort::ConstNode lhs, Ort::ConstNode rhs) { return lhs.GetId() < rhs.GetId(); }); - return FusionHasNoExternalPathBetweenSelectedNodes(fusion_nodes, - selected_node_ids); + if (!FusionHasNoExternalPathBetweenSelectedNodes(fusion_nodes, + selected_node_ids)) { + return false; + } + + if (shape_expr_hints != nullptr && !repeats_hints.empty()) { + std::vector boundary_names; + boundary_names.reserve(fusion_nodes.size() * 2); + for (const auto& [name, index] : ExternalInputs(fusion_nodes)) { + boundary_names.push_back(name); + } + (*shape_expr_hints)[TileConcatFusionKey(boundary_names, fusion_outputs)] = + std::move(repeats_hints); + } + return true; } std::vector> FindTileConcatFusions( const std::vector& all_nodes, const std::unordered_set& graph_output_names, - const std::unordered_set& accepted_node_ids) { + const std::unordered_set& accepted_node_ids, + FusionShapeExprHints* shape_expr_hints) { std::vector> fusions; auto producers = BuildProducerMap(all_nodes); for (Ort::ConstNode node : all_nodes) { std::vector fusion_nodes; if (!CanFuseTileConcat(node, producers, graph_output_names, - accepted_node_ids, fusion_nodes)) { + accepted_node_ids, fusion_nodes, shape_expr_hints)) { continue; } fusions.push_back(std::move(fusion_nodes)); diff --git a/test/fusion/test_tile_concat_shape_chain.py b/test/fusion/test_tile_concat_shape_chain.py index 476c8e5..0dfe2b1 100644 --- a/test/fusion/test_tile_concat_shape_chain.py +++ b/test/fusion/test_tile_concat_shape_chain.py @@ -10,19 +10,30 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -"""Tile repeats computed from a shape-only chain, absorbed into Tile+Concat. +"""Tile repeats computed from a shape-only chain that stays outside the fusion. A fused node cannot declare a CPU input, so ORT copies a runtime repeats tensor -onto the device and the kernel would have to read it back with a stream -synchronize. Instead the matcher absorbs the Shape/Gather/... chain and the -kernel recomputes repeats from the shapes of its inputs. These tests cover the -cases that make that absorption observable: a repeats value that follows -another tensor's shape, a chain value with a consumer outside the fusion, and -inputs the evaluator has to reject. +onto the device and the kernel would read it back with a device-to-host copy and +a compute-stream synchronize. The Tile+Concat matcher therefore proves that +repeats is a pure function of the shapes of the fused inputs and hands that +expression to the fused node, which recomputes the value instead of reading it. + +The chain itself is NOT absorbed: the dnn models read its intermediate values +from other branches, and a fused node owns every value it absorbs. These tests +pin that behaviour down: + +* a chain with Gather + Reshape whose intermediate value is also a graph output + still fuses (previously both conditions vetoed the fusion), +* repeats follows another tensor's shape while the Tile input keeps its shape, +* a chain that is not shape-only, and a chain whose Shape source is not an input + of Tile+Concat, leave the Tile standalone and still produce the right result. """ import json import os +import subprocess +import sys +from pathlib import Path import numpy as np import onnxruntime as ort @@ -31,7 +42,33 @@ from op_test_utils import musa_devices -INT64_MAX = 2**63 - 1 +DUMP_ENV = "ORT_MUSA_DUMP_RUNTIME_GRAPH_MERMAID" +DUMP_PATH_ENV = "ORT_MUSA_DUMP_RUNTIME_GRAPH_MERMAID_PATH" +REPEATS_DIAGNOSTIC = "MUSA_TILE_CONCAT_REPEATS" +FALLBACK_DIAGNOSTIC = "shape_expr=0" +RECOMPUTED_DIAGNOSTIC = "shape_expr=1" +OPS_DIR = Path(__file__).resolve().parents[1] / "ops" + +_SUBPROCESS_RUNNER = """ +import os +import sys + +import numpy as np +import onnxruntime as ort + +from op_test_utils import musa_devices + +model_path, feeds_path = sys.argv[1], sys.argv[2] +with open(model_path, "rb") as handle: + model = handle.read() +feeds = dict(np.load(feeds_path)) + +options = ort.SessionOptions() +options.add_session_config_entry("session.disable_cpu_ep_fallback", "0") +options.add_provider_for_devices(musa_devices(), {}) +session = ort.InferenceSession(model, sess_options=options) +session.run(None, feeds) +""" def _build_model(name, nodes, inputs, outputs, initializers, opset=17): @@ -50,6 +87,10 @@ def _int64_initializer(name, values, shape=None): ) +def _int64_scalar(name, value): + return helper.make_tensor(name, TensorProto.INT64, [], [value]) + + def _float_initializer(name, values, shape): array = np.asarray(values, dtype=np.float32).reshape(shape) return helper.make_tensor( @@ -63,7 +104,7 @@ def _session(model, use_musa, profile_prefix=None): devices = musa_devices() assert devices, "no MUSA device: refusing a CPU-only session for the MUSA run" options = ort.SessionOptions() - options.add_session_config_entry("session.disable_cpu_ep_fallback", "1") + options.add_session_config_entry("session.disable_cpu_ep_fallback", "0") if profile_prefix is not None: options.enable_profiling = True options.profile_file_prefix = profile_prefix @@ -100,26 +141,57 @@ def _fused_ops(ops): } -def _assert_tile_concat_absorbed(ops): - """The Tile and Concat ran inside the fused node, chain and all.""" - assert _fused_ops(ops) - assert "Tile" not in _musa_ops(ops) - assert "Concat" not in _musa_ops(ops) +def _run_in_subprocess(model, feeds, tmp_path): + """Run once in a fresh process and return (runtime dump, stderr). + + The runtime graph dump is written at process exit, and the diagnostics the + Tile+Concat factory prints are only emitted while the session is created, so + both have to be collected outside the pytest process. + """ + model_path = tmp_path / "model.onnx" + model_path.write_bytes(model) + feeds_path = tmp_path / "feeds.npz" + np.savez(feeds_path, **feeds) + dump_path = tmp_path / "runtime_graph.mmd" + env = dict(os.environ) + env[DUMP_ENV] = "1" + env[DUMP_PATH_ENV] = str(dump_path) + # The subprocess imports op_test_utils the way conftest.py makes it visible. + env["PYTHONPATH"] = os.pathsep.join( + [str(OPS_DIR)] + ([env["PYTHONPATH"]] if env.get("PYTHONPATH") else []) + ) + completed = subprocess.run( + [ + sys.executable, + "-c", + _SUBPROCESS_RUNNER, + str(model_path), + str(feeds_path), + ], + capture_output=True, + text=True, + env=env, + check=True, + ) + dump = dump_path.read_text(encoding="utf-8") if dump_path.exists() else "" + return dump, completed.stderr def _batch_chain_nodes(repeats_source): return [ helper.make_node("Shape", ["query"], ["query_shape"]), - helper.make_node("Gather", ["query_shape", "zero"], ["batch"], axis=0), + helper.make_node("Gather", ["query_shape", "zero_scalar"], ["batch"], axis=0), *repeats_source, helper.make_node("Tile", ["data", "repeats"], ["tiled"]), - helper.make_node("Concat", ["tiled", "query"], ["y"], axis=1), + helper.make_node("Concat", ["tiled", "query"], ["y_pre"], axis=1), + helper.make_node("Relu", ["y_pre"], ["y"]), ] def _batch_chain_initializers(): return [ _int64_initializer("zero", [0]), + _int64_scalar("zero_scalar", 0), _int64_initializer("one", [1]), _float_initializer("data", np.arange(4, dtype=np.float32), (1, 4)), ] @@ -128,9 +200,8 @@ def _batch_chain_initializers(): def test_tile_concat_repeats_follows_another_tensors_shape(tmp_path): """One session, two batch sizes: repeats moves while the Tile input does not. - The Tile data stays [1, 4] and only `query`'s batch axis changes. The matcher - absorbs the chain, so repeats is no longer an input of the fused node and the - expression is the only way to obtain it. + `data` stays [1, 4] and only `query`'s batch axis changes, so a constant + repeats would be wrong for one of the two runs. """ rng = np.random.default_rng(29) nodes = _batch_chain_nodes( @@ -140,7 +211,7 @@ def test_tile_concat_repeats_follows_another_tensors_shape(tmp_path): ] ) model = _build_model( - "tile_concat_repeats_shape_chain", + "tile_concat_repeats_follows_shape", nodes, [helper.make_tensor_value_info("query", TensorProto.FLOAT, ["batch", 4])], [helper.make_tensor_value_info("y", TensorProto.FLOAT, None)], @@ -156,48 +227,167 @@ def test_tile_concat_repeats_follows_another_tensors_shape(tmp_path): assert actual.shape == expected.shape == (batch, 8) np.testing.assert_allclose(actual, expected, rtol=1e-6, atol=1e-6) - _assert_tile_concat_absorbed(_read_ops_by_provider(musa)) + ops = _read_ops_by_provider(musa) + assert _fused_ops(ops) + assert "Tile" not in _musa_ops(ops) + + dump, stderr = _run_in_subprocess(model, {"query": query}, tmp_path) + assert "TileConcatFusionCompute" in dump + assert REPEATS_DIAGNOSTIC in stderr + assert RECOMPUTED_DIAGNOSTIC in stderr + assert FALLBACK_DIAGNOSTIC not in stderr -def test_tile_concat_repeats_chain_value_used_outside_fusion(): - """A chain value that is also a graph output must not be absorbed. +def test_tile_concat_repeats_chain_with_gather_reshape_and_external_consumer( + tmp_path, +): + """A Gather+Reshape chain whose intermediate value is a graph output fuses. - A fused node owns its internal values, so absorbing a value another consumer - reads would either break that consumer or change the graph. The Tile falls - back to the standalone kernel (which takes repeats as a CPU input), the chain - stays in the graph, and both outputs stay correct. + Both properties used to veto the fusion: a chain value with an outside + consumer, and a body that carries Gather + Concat + Reshape (which the loose + IsShapeReshapeFusionGraph detector could claim). The chain now stays in the + graph, the fused body is only Tile + Concat, and the fused node is still the + TileConcat compute. """ rng = np.random.default_rng(31) - nodes = _batch_chain_nodes( - [ - helper.make_node("Unsqueeze", ["batch", "zero"], ["batch_1d"]), - helper.make_node("Concat", ["batch_1d", "one"], ["repeats"], axis=0), - ] - ) + nodes = [ + helper.make_node("Shape", ["query"], ["query_shape"]), + helper.make_node("Gather", ["query_shape", "zero_scalar"], ["batch"], axis=0), + helper.make_node("Reshape", ["batch", "one_shape"], ["batch_1d"]), + helper.make_node("Concat", ["batch_1d", "one"], ["repeats"], axis=0), + helper.make_node("Tile", ["data", "repeats"], ["tiled"]), + helper.make_node("Concat", ["tiled", "query"], ["y_pre"], axis=1), + helper.make_node("Relu", ["y_pre"], ["y"]), + ] + initializers = _batch_chain_initializers() + [_int64_initializer("one_shape", [1])] model = _build_model( - "tile_concat_repeats_chain_external_consumer", + "tile_concat_repeats_gather_reshape_external_consumer", nodes, [helper.make_tensor_value_info("query", TensorProto.FLOAT, ["batch", 4])], [ helper.make_tensor_value_info("y", TensorProto.FLOAT, None), helper.make_tensor_value_info("batch", TensorProto.INT64, None), ], - _batch_chain_initializers(), + initializers, + ) + + cpu = _session(model, use_musa=False) + musa = _session( + model, use_musa=True, profile_prefix=str(tmp_path / "gather_reshape") + ) + for batch in (2, 5): + query = rng.standard_normal((batch, 4)).astype(np.float32) + feeds = {"query": query} + expected = cpu.run(None, feeds) + actual = musa.run(None, feeds) + assert len(actual) == 2 + np.testing.assert_allclose(actual[0], expected[0], rtol=1e-6, atol=1e-6) + np.testing.assert_array_equal(actual[1], expected[1]) + assert actual[1] == np.int64(batch) + + ops = _read_ops_by_provider(musa) + assert _fused_ops(ops) + assert "Tile" not in _musa_ops(ops) + + dump, stderr = _run_in_subprocess(model, {"query": query}, tmp_path) + # The factory ORT picked for this fused node, not its detector order. + assert "TileConcatFusionCompute" in dump + assert "ShapeReshapeFusionCompute" not in dump + assert RECOMPUTED_DIAGNOSTIC in stderr + assert FALLBACK_DIAGNOSTIC not in stderr + + +def test_tile_concat_repeats_chain_with_non_shape_op_stays_standalone(tmp_path): + """A chain with an op outside the shape-only set keeps the standalone Tile. + + `Max` is not shape arithmetic, so the matcher proves nothing. The Tile runs + on its own (reading repeats, as before), the Concat is not fused, and both + outputs still match the CPU. + """ + rng = np.random.default_rng(37) + nodes = [ + helper.make_node("Shape", ["query"], ["query_shape"]), + helper.make_node("Gather", ["query_shape", "zero_scalar"], ["batch"], axis=0), + helper.make_node("Max", ["batch", "one"], ["max_batch"]), + helper.make_node("Reshape", ["max_batch", "one_shape"], ["batch_1d"]), + helper.make_node("Concat", ["batch_1d", "one"], ["repeats"], axis=0), + helper.make_node("Tile", ["data", "repeats"], ["tiled"]), + helper.make_node("Concat", ["tiled", "query"], ["y_pre"], axis=1), + helper.make_node("Relu", ["y_pre"], ["y"]), + ] + initializers = _batch_chain_initializers() + [_int64_initializer("one_shape", [1])] + model = _build_model( + "tile_concat_repeats_non_shape_op", + nodes, + [helper.make_tensor_value_info("query", TensorProto.FLOAT, ["batch", 4])], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, None)], + initializers, ) - query = rng.standard_normal((2, 4)).astype(np.float32) + query = rng.standard_normal((3, 4)).astype(np.float32) feeds = {"query": query} - expected = _session(model, use_musa=False).run(None, feeds) - actual = _session(model, use_musa=True).run(None, feeds) - assert len(actual) == 2 - np.testing.assert_allclose(actual[0], expected[0], rtol=1e-6, atol=1e-6) - np.testing.assert_array_equal(actual[1], expected[1]) - assert actual[1] == np.int64(2) + expected = _session(model, use_musa=False).run(None, feeds)[0] + musa = _session( + model, use_musa=True, profile_prefix=str(tmp_path / "non_shape_op") + ) + actual = musa.run(None, feeds)[0] + assert actual.shape == expected.shape == (3, 8) + np.testing.assert_allclose(actual, expected, rtol=1e-6, atol=1e-6) + + musa_ops = _musa_ops(_read_ops_by_provider(musa)) + assert "Tile" in musa_ops + assert "Concat" in musa_ops + + +def test_tile_concat_repeats_source_outside_fusion_stays_standalone(tmp_path): + """A Shape source that Tile+Concat does not read cannot be recomputed. + + `other` only feeds the chain, so it is not an input of the fused node and + the expression would need a tensor the fusion never sees. The Tile keeps its + standalone kernel, which reads repeats, and the result stays correct. + """ + rng = np.random.default_rng(41) + nodes = [ + helper.make_node("Shape", ["other"], ["other_shape"]), + helper.make_node("Gather", ["other_shape", "zero_scalar"], ["other_dim"], axis=0), + helper.make_node("Unsqueeze", ["other_dim", "zero"], ["other_dim_1d"]), + helper.make_node("Concat", ["other_dim_1d", "one"], ["repeats"], axis=0), + helper.make_node("Tile", ["data", "repeats"], ["tiled"]), + helper.make_node("Concat", ["tiled", "query"], ["y_pre"], axis=1), + helper.make_node("Relu", ["y_pre"], ["y"]), + ] + initializers = _batch_chain_initializers() + model = _build_model( + "tile_concat_repeats_source_outside_fusion", + nodes, + [ + helper.make_tensor_value_info("query", TensorProto.FLOAT, [3, 4]), + helper.make_tensor_value_info("other", TensorProto.FLOAT, ["m"]), + ], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, None)], + initializers, + ) + + feeds = { + "query": rng.standard_normal((3, 4)).astype(np.float32), + "other": rng.standard_normal(3).astype(np.float32), + } + expected = _session(model, use_musa=False).run(None, feeds)[0] + musa = _session( + model, use_musa=True, profile_prefix=str(tmp_path / "source_outside") + ) + actual = musa.run(None, feeds)[0] + assert actual.shape == expected.shape == (3, 8) + np.testing.assert_allclose(actual, expected, rtol=1e-6, atol=1e-6) + + musa_ops = _musa_ops(_read_ops_by_provider(musa)) + assert "Tile" in musa_ops + assert "Concat" in musa_ops def test_tile_concat_repeats_chain_with_negative_slice(): """Negative Slice bounds are resolved against the runtime shape.""" - rng = np.random.default_rng(37) + rng = np.random.default_rng(43) nodes = [ helper.make_node("Shape", ["query"], ["query_shape"]), helper.make_node( @@ -205,7 +395,8 @@ def test_tile_concat_repeats_chain_with_negative_slice(): ), helper.make_node("Concat", ["batch_dim", "one"], ["repeats"], axis=0), helper.make_node("Tile", ["data", "repeats"], ["tiled"]), - helper.make_node("Concat", ["tiled", "query"], ["y"], axis=1), + helper.make_node("Concat", ["tiled", "query"], ["y_pre"], axis=1), + helper.make_node("Relu", ["y_pre"], ["y"]), ] initializers = [ _float_initializer("data", np.arange(4, dtype=np.float32), (1, 4)), @@ -236,7 +427,8 @@ def test_tile_concat_repeats_chain_with_reshape_element_mismatch(): helper.make_node("Shape", ["query"], ["query_shape"]), helper.make_node("Reshape", ["query_shape", "three"], ["bad_repeats"]), helper.make_node("Tile", ["data", "bad_repeats"], ["tiled"]), - helper.make_node("Concat", ["tiled", "query"], ["y"], axis=1), + helper.make_node("Concat", ["tiled", "query"], ["y_pre"], axis=1), + helper.make_node("Relu", ["y_pre"], ["y"]), ] initializers = [ _float_initializer("data", np.arange(2, dtype=np.float32), (1, 2)), From 1186493f4c1c818ad2b380f8f329c2f05b7dd3ce Mon Sep 17 00:00:00 2001 From: "yichao.hu" Date: Sun, 18 Oct 2026 14:57:51 -0400 Subject: [PATCH 28/52] fix(plugin): let Slice read device-resident shape metadata Drop the starts/ends/axes/steps CPUInput declarations added in fd8ca63. With them, ORT inserts a MemcpyToHost for every Slice execution whose metadata chain runs on the device; inside dien's two GRU Loop bodies this produced about 300 D2H copies per inference (303 vs 3 on master), six per iteration. The Slice kernel already reads those inputs through ReadIntTensors on the compute stream, so this restores the pre-fd8ca63 device-read path while keeping the small-copy ordering fix from that commit. The dynamic-shape metadata and boundary-copy tests, plus zipformer run twice, gate the buffer-reuse failure. --- musa/ep/src/kernels/tensor/slice.cc | 25 ++++++++----------------- test/ops/test_dynamic_shape_metadata.py | 9 ++++++--- 2 files changed, 14 insertions(+), 20 deletions(-) diff --git a/musa/ep/src/kernels/tensor/slice.cc b/musa/ep/src/kernels/tensor/slice.cc index fae9a75..53028f2 100644 --- a/musa/ep/src/kernels/tensor/slice.cc +++ b/musa/ep/src/kernels/tensor/slice.cc @@ -209,13 +209,12 @@ ONNX_OPERATOR_VERSIONED_KERNEL_EX( (Ort::KernelDefBuilder() .AddTypeConstraint("T", AllFixedSizeTensorTypes()) .AddTypeConstraint("Tind", IntTensorTypes()) - // starts/ends/axes/steps are shape metadata: keep them in host memory - // (the in-tree CUDA kernel declares the same) so ORT orders the - // CPU/device boundary itself instead of leaving it inside the kernel. - .SetInputMemType(1, OrtMemTypeCPUInput) - .SetInputMemType(2, OrtMemTypeCPUInput) - .SetInputMemType(3, OrtMemTypeCPUInput) - .SetInputMemType(4, OrtMemTypeCPUInput)), + // starts/ends/axes/steps are read by ReadIntTensors on the compute + // stream, so a device-produced metadata chain needs no per-execution + // host copy. Declaring them CPUInput made ORT insert a D2H copy for + // every execution, which showed up as six copies per GRU Loop + // iteration in dien (300+ per inference). + ), Slice) // opset 10 moved starts/ends/axes/steps from attributes to inputs; opset 11 @@ -225,22 +224,14 @@ ONNX_OPERATOR_VERSIONED_KERNEL_EX( Slice, kOnnxDomain, 10, 10, (Ort::KernelDefBuilder() .AddTypeConstraint("T", AllFixedSizeTensorTypesNoBFloat16()) - .AddTypeConstraint("Tind", IntTensorTypes()) - .SetInputMemType(1, OrtMemTypeCPUInput) - .SetInputMemType(2, OrtMemTypeCPUInput) - .SetInputMemType(3, OrtMemTypeCPUInput) - .SetInputMemType(4, OrtMemTypeCPUInput)), + .AddTypeConstraint("Tind", IntTensorTypes())), Slice) ONNX_OPERATOR_VERSIONED_KERNEL_EX( Slice, kOnnxDomain, 11, 12, (Ort::KernelDefBuilder() .AddTypeConstraint("T", AllFixedSizeTensorTypesNoBFloat16()) - .AddTypeConstraint("Tind", IntTensorTypes()) - .SetInputMemType(1, OrtMemTypeCPUInput) - .SetInputMemType(2, OrtMemTypeCPUInput) - .SetInputMemType(3, OrtMemTypeCPUInput) - .SetInputMemType(4, OrtMemTypeCPUInput)), + .AddTypeConstraint("Tind", IntTensorTypes())), Slice) // opset 1-9 keeps starts/ends/axes in node attributes; those schemas do not diff --git a/test/ops/test_dynamic_shape_metadata.py b/test/ops/test_dynamic_shape_metadata.py index 5c50e3b..ca64ec4 100644 --- a/test/ops/test_dynamic_shape_metadata.py +++ b/test/ops/test_dynamic_shape_metadata.py @@ -14,9 +14,12 @@ `Shape`/`Gather`/`Unsqueeze` arithmetic produces the runtime shapes of a graph. Where that arithmetic derives from a device tensor it must not be executed as a -device-side subgraph: the values belong on the CPU EP, and the operators that -consume them (`Slice` starts/ends/axes/steps, `Reshape` shape, ...) declare those -inputs as `OrtMemTypeCPUInput`, exactly like the in-tree CUDA kernels do. +device-side subgraph: the values belong on the CPU EP, and most operators that +consume them (`Reshape` shape, `Tile` repeats, ...) declare those inputs as +`OrtMemTypeCPUInput`, exactly like the in-tree CUDA kernels do. `Slice` no +longer declares starts/ends/axes/steps; its kernel reads them through +`ReadIntTensors` on the compute stream, so a device-resident metadata chain is +read correctly without ORT-inserting a host copy for every execution. Regression cover for sherpa_whisper_base / campplus_3dspeaker and zipformer_encoder_en: with ORT graph optimisation on, their `Slice` received From 3e38acbe3e86771bef2e432a02abe0a795abc258 Mon Sep 17 00:00:00 2001 From: "yichao.hu" Date: Sun, 18 Oct 2026 15:43:41 -0400 Subject: [PATCH 29/52] fix(fusion): key the MHTA Slice-end cache through shape-preserving ops ValueDependsOnlyOnInputShapes required every Shape source to be an input of the fusion. A recommendation model's raw INT32 mask end is Unsqueeze(Gather(Shape( Transpose(k_norm)))) where the Transpose output has no other consumer and its input is the fused K input, so the cache was disabled and the end was read from the device with a stream synchronize on every run (about 0.9 ms). Follow Transpose/Cast/Identity from a Shape source back to the fused inputs that determine its shape. The Shape value is then still a pure function of those input shapes, so the cache key stays sound; no fusion input or output is added and no copy is introduced. Adds a two-shape test that runs the same session with seqlen 5, 4 and 5 again, checks the mask follows K's shape, and a MUSA_MHTA_SLICE_END diagnostic line for verification rounds. --- musa/ep/src/fusion/fusion_matcher_utils.cc | 58 +++++++++++- musa/ep/src/fusion/fusion_matcher_utils.h | 10 +- ...hta_scaled_dot_product_attention_fusion.cc | 7 ++ ...hta_scaled_dot_product_attention_fusion.py | 94 +++++++++++++++++++ 4 files changed, 164 insertions(+), 5 deletions(-) diff --git a/musa/ep/src/fusion/fusion_matcher_utils.cc b/musa/ep/src/fusion/fusion_matcher_utils.cc index 87dbd2d..9096e0e 100644 --- a/musa/ep/src/fusion/fusion_matcher_utils.cc +++ b/musa/ep/src/fusion/fusion_matcher_utils.cc @@ -283,6 +283,56 @@ bool HasCastBoundary(const std::vector& fusion_nodes) { return false; } +// Ops whose output shape is a pure function of their input shapes. A Shape +// node reading their output can therefore be keyed on the shapes of the fused +// inputs that feed them, without the fused node taking a new input. +bool IsShapePreservingDataOp(const std::string& op) { + return op == "Transpose" || op == "Cast" || op == "Identity"; +} + +// Walk from `value_name` through shape-preserving producers to the fused +// inputs that determine its shape. Returns false when the walk leaves the +// fused boundary through anything whose output shape is data-dependent. +bool CollectShapePreservingSources( + const std::string& value_name, + const std::unordered_map& producers, + const std::unordered_map& input_indices, + const std::unordered_set& initializer_names, + std::unordered_set* sources) { + std::vector pending{value_name}; + std::unordered_set visited; + while (!pending.empty()) { + const std::string name = std::move(pending.back()); + pending.pop_back(); + if (!visited.insert(name).second) { + continue; + } + auto index_it = input_indices.find(name); + if (index_it != input_indices.end()) { + sources->insert(index_it->second); + continue; + } + if (initializer_names.count(name) != 0) { + continue; // a constant's shape is static and needs no runtime key + } + auto producer_it = producers.find(name); + if (producer_it == producers.end()) { + return false; + } + Ort::ConstNode node = producer_it->second; + if (!IsShapePreservingDataOp(node.GetOperatorType())) { + return false; + } + for (Ort::ConstValueInfo input : node.GetInputs()) { + const std::string input_name = input.GetName(); + if (!input_name.empty()) { + pending.push_back(input_name); + } + } + } + return true; +} + bool ValueDependsOnlyOnInputShapes( const std::string& value_name, const std::unordered_map& producers, @@ -325,8 +375,12 @@ bool ValueDependsOnlyOnInputShapes( const std::string source = node_inputs[0].GetName(); auto index_it = input_indices.find(source); if (index_it == input_indices.end()) { - // Shape of a constant is constant and needs no runtime key. - if (initializer_names.count(source) == 0) { + // Shape of a constant is constant and needs no runtime key. Otherwise + // follow shape-preserving producers (a Transpose kept only for its + // shape, for example) back to the fused inputs that determine it. + if (initializer_names.count(source) == 0 && + !CollectShapePreservingSources(source, producers, input_indices, + initializer_names, &sources)) { return false; } } else { diff --git a/musa/ep/src/fusion/fusion_matcher_utils.h b/musa/ep/src/fusion/fusion_matcher_utils.h index 09a404f..b485dc1 100644 --- a/musa/ep/src/fusion/fusion_matcher_utils.h +++ b/musa/ep/src/fusion/fusion_matcher_utils.h @@ -62,9 +62,13 @@ bool HasCastBoundary(const std::vector& fusion_nodes); // Returns true when `value_name` is produced only from shape arithmetic: // constants, plus Shape ops over inputs of the current fusion. The output is // then a pure function of the shapes of the fused inputs listed in -// `source_indices` (empty means the value is constant). Anything that could -// make the value depend on data, on a tensor outside the fusion, or on an op -// that is not shape-only returns false. +// `source_indices` (empty means the value is constant). A Shape whose source +// is not itself a fused input is accepted when the source is reached from one +// through shape-preserving ops only (Transpose with a perm attribute, Cast, +// Identity): its output shape, and therefore the Shape value, is still a pure +// function of those inputs' shapes. Anything that could make the value depend +// on data, on a tensor outside the fusion by another path, or on an op that is +// not shape-only returns false. bool ValueDependsOnlyOnInputShapes( const std::string& value_name, const std::unordered_map& producers, diff --git a/musa/ep/src/fusion/mhta_scaled_dot_product_attention_fusion.cc b/musa/ep/src/fusion/mhta_scaled_dot_product_attention_fusion.cc index 0de8806..b93cdda 100644 --- a/musa/ep/src/fusion/mhta_scaled_dot_product_attention_fusion.cc +++ b/musa/ep/src/fusion/mhta_scaled_dot_product_attention_fusion.cc @@ -17,6 +17,7 @@ #include #include +#include #include #include #include @@ -1494,6 +1495,12 @@ std::unique_ptr CreateMhtaScaledDotProductAttentionFusion( mask_end_cacheable = musa_ep::ValueDependsOnlyOnInputShapes( Name(slice_it->second.GetInputs()[2]), producers, fused_indices, initializer_names, &mask_end_source_indices); + // One line per fused raw-INT32-mask attention node at session + // creation: 0 means the end is read from the device every run. + std::fprintf(stderr, + "MUSA_MHTA_SLICE_END cacheable=%d sources=%zu\n", + mask_end_cacheable ? 1 : 0, + mask_end_source_indices.size()); } boolean_mask_int32 = true; } diff --git a/test/fusion/test_mhta_scaled_dot_product_attention_fusion.py b/test/fusion/test_mhta_scaled_dot_product_attention_fusion.py index 16284ed..af08f87 100644 --- a/test/fusion/test_mhta_scaled_dot_product_attention_fusion.py +++ b/test/fusion/test_mhta_scaled_dot_product_attention_fusion.py @@ -928,3 +928,97 @@ def test_mhta_scaled_dot_product_attention_sim_rank3_div_temperature_fusion( assert "Einsum" not in musa_ops assert "Div" not in musa_ops assert "Softmax" not in musa_ops + + +def _mhta_dynamic_end_model(batch, heads, head_dim): + """Raw INT32 mask whose Slice end comes from Shape(Transpose(K)). + + Q/K/V follow the plain bshd layout used by the other raw-mask tests. K is + [batch, heads, head_dim, seqlen] with a symbolic seqlen; a shape-only + `Transpose(K, perm=[0, 1, 3, 2])` (no other consumer, like + `Transpose_1` in ranking-gr) feeds Shape -> Gather(2) -> Unsqueeze, so the + mask Slice end equals the runtime seqlen. A cache keyed on K's shape must + therefore change when seqlen changes. + """ + def vi(name, elem_type, dims): + return helper.make_tensor_value_info(name, elem_type, dims) + + nodes = [ + helper.make_node("Transpose", ["K", "shape_only_perm"], ["KShapeOnly"]), + helper.make_node("Shape", ["KShapeOnly"], ["KShape"]), + helper.make_node("Gather", ["KShape", "seqlen_index"], ["Seqlen"], axis=0), + helper.make_node("Unsqueeze", ["Seqlen", "axis0"], ["Ends"]), + helper.make_node("Slice", ["RawMask", "starts", "Ends", "axes", "steps"], ["MaskSlice"]), + helper.make_node("Cast", ["MaskSlice"], ["MaskInt64"], to=TensorProto.INT64), + helper.make_node("Equal", ["MaskInt64", "one"], ["Mask"]), + helper.make_node("MatMul", ["Q", "K"], ["Score"]), + helper.make_node("Mul", ["Score", "scale"], ["Scaled"]), + helper.make_node("Where", ["Mask", "Scaled", "neg_inf"], ["MaskedScore"]), + helper.make_node("Softmax", ["MaskedScore"], ["Prob"], axis=-1), + helper.make_node("Where", ["Mask", "Prob", "zero"], ["MaskedProb"]), + helper.make_node("MatMul", ["MaskedProb", "V"], ["Y"]), + ] + initializers = [ + numpy_helper.from_array(np.asarray(v, dtype=dtype), name=name) + for name, v, dtype in [ + ("shape_only_perm", [0, 1, 3, 2], np.int64), + ("seqlen_index", np.asarray(2, dtype=np.int64), np.int64), + ("axis0", [0], np.int64), + ("starts", [0], np.int64), + ("axes", [3], np.int64), + ("steps", [1], np.int64), + ("one", np.asarray(1, dtype=np.int64), np.int64), + ("scale", np.asarray(1.0 / np.sqrt(head_dim), dtype=np.float32), np.float32), + ("neg_inf", np.asarray(-np.inf, dtype=np.float32), np.float32), + ("zero", np.asarray(0.0, dtype=np.float32), np.float32), + ] + ] + graph = helper.make_graph( + nodes, + "mhta_dynamic_slice_end_graph", + [ + vi("Q", TensorProto.FLOAT, [batch, heads, "seqlen", head_dim]), + vi("K", TensorProto.FLOAT, [batch, heads, head_dim, "seqlen"]), + vi("V", TensorProto.FLOAT, [batch, heads, "seqlen", head_dim]), + vi("RawMask", TensorProto.INT32, [batch, 1, "mask_rows", "mask_cols"]), + ], + [vi("Y", TensorProto.FLOAT, None)], + initializer=initializers, + ) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 17)]) + model.ir_version = min(model.ir_version, 10) + return model.SerializeToString() + + +def _mhta_dynamic_feeds(rng, batch, heads, head_dim, seqlen): + raw_mask = rng.integers(0, 3, (batch, 1, seqlen, seqlen + 2), dtype=np.int32) + raw_mask[:, :, -1, :seqlen] = 0 + return { + "Q": rng.standard_normal((batch, heads, seqlen, head_dim)).astype(np.float32), + "K": rng.standard_normal((batch, heads, head_dim, seqlen)).astype(np.float32), + "V": rng.standard_normal((batch, heads, seqlen, head_dim)).astype(np.float32), + "RawMask": raw_mask, + } + + +def test_mhta_slice_end_cache_follows_k_shape(capfd): + """The cached Slice end must be keyed on K's shape, not reused blindly.""" + rng = np.random.default_rng(61) + batch, heads, head_dim = 1, 3, 4 + model = _mhta_dynamic_end_model(batch, heads, head_dim) + + so = ort.SessionOptions() + so.add_session_config_entry("session.disable_cpu_ep_fallback", "1") + so.add_provider_for_devices(musa_devices(), {}) + capfd.readouterr() + musa_session = ort.InferenceSession(model, sess_options=so) + err = capfd.readouterr().err + assert "MUSA_MHTA_SLICE_END cacheable=1" in err, err + + for seqlen in (5, 4, 5): + feeds = _mhta_dynamic_feeds(rng, batch, heads, head_dim, seqlen) + cpu = ort.InferenceSession(model, providers=["CPUExecutionProvider"]).run(None, feeds) + out = musa_session.run(None, feeds) + np.testing.assert_allclose(out[0], cpu[0], rtol=1e-4, atol=1e-4, + err_msg=f"seqlen={seqlen}") + np.testing.assert_array_equal(out[0][:, :, -1, :], 0.0) From de62347b66dd833a722b672ea168d4f74e54ed0e Mon Sep 17 00:00:00 2001 From: "yichao.hu" Date: Sun, 18 Oct 2026 15:57:14 -0400 Subject: [PATCH 30/52] test(fusion): fix the Transpose perm input in the Slice-end shape test Transpose takes perm as an attribute, not as a second input. --- test/fusion/test_mhta_scaled_dot_product_attention_fusion.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/test/fusion/test_mhta_scaled_dot_product_attention_fusion.py b/test/fusion/test_mhta_scaled_dot_product_attention_fusion.py index af08f87..ee611a9 100644 --- a/test/fusion/test_mhta_scaled_dot_product_attention_fusion.py +++ b/test/fusion/test_mhta_scaled_dot_product_attention_fusion.py @@ -944,7 +944,7 @@ def vi(name, elem_type, dims): return helper.make_tensor_value_info(name, elem_type, dims) nodes = [ - helper.make_node("Transpose", ["K", "shape_only_perm"], ["KShapeOnly"]), + helper.make_node("Transpose", ["K"], ["KShapeOnly"], perm=[0, 1, 3, 2]), helper.make_node("Shape", ["KShapeOnly"], ["KShape"]), helper.make_node("Gather", ["KShape", "seqlen_index"], ["Seqlen"], axis=0), helper.make_node("Unsqueeze", ["Seqlen", "axis0"], ["Ends"]), @@ -961,7 +961,6 @@ def vi(name, elem_type, dims): initializers = [ numpy_helper.from_array(np.asarray(v, dtype=dtype), name=name) for name, v, dtype in [ - ("shape_only_perm", [0, 1, 3, 2], np.int64), ("seqlen_index", np.asarray(2, dtype=np.int64), np.int64), ("axis0", [0], np.int64), ("starts", [0], np.int64), From 07ef3d72b19b9facc52a58b0d333f43edd166035 Mon Sep 17 00:00:00 2001 From: "yichao.hu" Date: Sun, 18 Oct 2026 15:59:34 -0400 Subject: [PATCH 31/52] test(fusion): declare the output shape in the dynamic Slice-end test model onnx.checker requires a shape on every graph output; Y now carries the symbolic [batch, heads, seqlen, head_dim] shape. --- test/fusion/test_mhta_scaled_dot_product_attention_fusion.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/fusion/test_mhta_scaled_dot_product_attention_fusion.py b/test/fusion/test_mhta_scaled_dot_product_attention_fusion.py index ee611a9..f93a053 100644 --- a/test/fusion/test_mhta_scaled_dot_product_attention_fusion.py +++ b/test/fusion/test_mhta_scaled_dot_product_attention_fusion.py @@ -981,7 +981,7 @@ def vi(name, elem_type, dims): vi("V", TensorProto.FLOAT, [batch, heads, "seqlen", head_dim]), vi("RawMask", TensorProto.INT32, [batch, 1, "mask_rows", "mask_cols"]), ], - [vi("Y", TensorProto.FLOAT, None)], + [vi("Y", TensorProto.FLOAT, [batch, heads, "seqlen", head_dim])], initializer=initializers, ) model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 17)]) From df399e809bf3d8e16a8b39e58462b84875bebd05 Mon Sep 17 00:00:00 2001 From: "yichao.hu" Date: Sun, 18 Oct 2026 16:09:51 -0400 Subject: [PATCH 32/52] fix(fusion): prove the MHTA Slice-end cache in the matcher, not at Compile Compile() only receives the fused subgraph, so ValueDependsOnlyOnInputShapes() there could never see the end's producer chain and always returned false for a dynamic end. Every raw-INT32-mask attention node therefore read the end from the device with a stream synchronize on every run (about 0.9 ms on a recommendation model). Prove the dependency in the matcher, which still sees the whole graph, and hand the proven fused-input names to Compile through a boundary-keyed hint map, the same way TileConcat carries its shape expressions. Follow shape-preserving producers (Transpose/Cast/Identity), and accept a value shared with a fused input through another Transpose: both expose its shape up to a fixed permutation. Tests: the dynamic end test now asserts cacheable=1 and keeps running seqlen 5/4/5 in one session; two negative controls (an unrelated shape source, and a producer outside the whitelist) assert cacheable=0 and stay correct. --- musa/ep/src/ep.cc | 6 +- musa/ep/src/ep.h | 15 +++ musa/ep/src/fusion/fusion_matcher.cc | 7 +- musa/ep/src/fusion/fusion_matcher.h | 14 ++- musa/ep/src/fusion/fusion_matcher_utils.cc | 56 +++++++++ musa/ep/src/fusion/fusion_matcher_utils.h | 15 ++- musa/ep/src/fusion/fusion_node_compute.cc | 3 +- ...hta_scaled_dot_product_attention_fusion.cc | 69 +++++++++-- ...mhta_scaled_dot_product_attention_fusion.h | 6 +- ...ed_dot_product_attention_fusion_matcher.cc | 113 +++++++++++++++++- musa/ep/src/fusion/tile_concat_fusion.cc | 16 +-- ...hta_scaled_dot_product_attention_fusion.py | 95 +++++++++++---- 12 files changed, 344 insertions(+), 71 deletions(-) diff --git a/musa/ep/src/ep.cc b/musa/ep/src/ep.cc index 66a01cc..582de4c 100644 --- a/musa/ep/src/ep.cc +++ b/musa/ep/src/ep.cc @@ -380,9 +380,10 @@ MusaEp::GetCapabilityImpl(OrtEp* this_ptr, const OrtGraph* ort_graph, // Filled while the whole graph is still visible: Compile() only gets the // fused subgraph, which no longer holds the chain a shape value comes from. FusionShapeExprHints shape_expr_hints; + FusionMaskEndHints mask_end_hints; if (!DisableAllFusions()) { - fusion_matches = - FindFusionMatches(all_nodes, graph_output_names, &shape_expr_hints); + fusion_matches = FindFusionMatches(all_nodes, graph_output_names, + &shape_expr_hints, &mask_end_hints); if (DisableShapeReshapeFusion()) { fusion_matches.erase( std::remove_if(fusion_matches.begin(), fusion_matches.end(), @@ -403,6 +404,7 @@ MusaEp::GetCapabilityImpl(OrtEp* this_ptr, const OrtGraph* ort_graph, shape_expr_hints.clear(); } ep->AddShapeExprHints(std::move(shape_expr_hints)); + ep->AddMaskEndHints(std::move(mask_end_hints)); RETURN_IF_ERROR(RegisterFusionMatches(ep->ep_api_, graph_support_info, fusion_matches)); } diff --git a/musa/ep/src/ep.h b/musa/ep/src/ep.h index df59131..b679b54 100644 --- a/musa/ep/src/ep.h +++ b/musa/ep/src/ep.h @@ -63,6 +63,20 @@ class MusaEp : public OrtEp { } } + // Mask-end provenance a finder proved while it could still see the whole + // graph; Compile() remaps the source names to fused input indices. + const musa_ep::FusionMaskEndHints& MaskEndHints() const { + return mask_end_hints_; + } + void AddMaskEndHints(musa_ep::FusionMaskEndHints hints) { + for (auto& [key, values] : hints) { + auto& stored = mask_end_hints_[key]; + for (auto& [name, sources] : values) { + stored[name] = std::move(sources); + } + } + } + private: static const char* ORT_API_CALL GetNameImpl(const OrtEp* this_ptr) noexcept; @@ -100,4 +114,5 @@ class MusaEp : public OrtEp { std::unordered_map> fusion_computes_; musa_ep::FusionShapeExprHints shape_expr_hints_; + musa_ep::FusionMaskEndHints mask_end_hints_; }; diff --git a/musa/ep/src/fusion/fusion_matcher.cc b/musa/ep/src/fusion/fusion_matcher.cc index f1371f1..dccfcd3 100644 --- a/musa/ep/src/fusion/fusion_matcher.cc +++ b/musa/ep/src/fusion/fusion_matcher.cc @@ -149,7 +149,8 @@ void AddFusionMatch(std::vector& matches, const char* finder, std::vector FindFusionMatches( const std::vector& all_nodes, const std::unordered_set& graph_output_names, - FusionShapeExprHints* shape_expr_hints) { + FusionShapeExprHints* shape_expr_hints, + FusionMaskEndHints* mask_end_hints) { std::unordered_set accepted_node_ids; std::vector matches; @@ -171,8 +172,8 @@ std::vector FindFusionMatches( // local SDPA matcher so future ONNX-node-based expansions do not get // preempted by the narrower attention chain. auto mhta_scaled_dot_product_attention_fusions = - FindMhtaScaledDotProductAttentionFusions(all_nodes, graph_output_names, - accepted_node_ids); + FindMhtaScaledDotProductAttentionFusions( + all_nodes, graph_output_names, accepted_node_ids, mask_end_hints); AddFusionMatch(matches, "FindMhtaScaledDotProductAttentionFusions", false, std::move(mhta_scaled_dot_product_attention_fusions), accepted_node_ids); diff --git a/musa/ep/src/fusion/fusion_matcher.h b/musa/ep/src/fusion/fusion_matcher.h index be7df39..f31e2db 100644 --- a/musa/ep/src/fusion/fusion_matcher.h +++ b/musa/ep/src/fusion/fusion_matcher.h @@ -35,6 +35,14 @@ using FusionShapeExprHints = std::unordered_map>; +// Mask-end provenance a finder proved on the whole graph, keyed the same way: +// boundary signature -> end value name -> names of the fused inputs whose +// shapes determine that value. Compile() cannot see the end's producer chain +// and remaps these names to fused input indices. +using FusionMaskEndHints = + std::unordered_map>>; + struct FusionDTypeRejection { std::vector nodes; std::string reason; @@ -51,7 +59,8 @@ struct FusionMatch { std::vector FindFusionMatches( const std::vector& all_nodes, const std::unordered_set& graph_output_names, - FusionShapeExprHints* shape_expr_hints = nullptr); + FusionShapeExprHints* shape_expr_hints = nullptr, + FusionMaskEndHints* mask_end_hints = nullptr); std::vector> FindSplitUnsqueezeConcatFusions( const std::vector& all_nodes, @@ -178,7 +187,8 @@ std::vector> FindMhtaScaledDotProductAttentionFusions( const std::vector& all_nodes, const std::unordered_set& graph_output_names, - const std::unordered_set& accepted_node_ids); + const std::unordered_set& accepted_node_ids, + FusionMaskEndHints* mask_end_hints = nullptr); std::vector> FindQkvAttentionOutputProjectionFusions( const std::vector& all_nodes, diff --git a/musa/ep/src/fusion/fusion_matcher_utils.cc b/musa/ep/src/fusion/fusion_matcher_utils.cc index 9096e0e..2a8bea7 100644 --- a/musa/ep/src/fusion/fusion_matcher_utils.cc +++ b/musa/ep/src/fusion/fusion_matcher_utils.cc @@ -299,6 +299,20 @@ bool CollectShapePreservingSources( const std::unordered_map& input_indices, const std::unordered_set& initializer_names, std::unordered_set* sources) { + // A value can also be pinned down from the other side: when another + // Transpose reads the *same* value and produces one of the fusion's inputs, + // that input's shape is this value's shape under a fixed permutation (every + // valid ONNX Transpose perm is invertible), so it can key the cache too. + std::unordered_map> consumers; + for (const auto& [output, producer] : producers) { + for (Ort::ConstValueInfo input : producer.GetInputs()) { + const std::string input_name = input.GetName(); + if (!input_name.empty()) { + consumers[input_name].push_back(producer); + } + } + } + std::vector pending{value_name}; std::unordered_set visited; while (!pending.empty()) { @@ -315,6 +329,29 @@ bool CollectShapePreservingSources( if (initializer_names.count(name) != 0) { continue; // a constant's shape is static and needs no runtime key } + bool keyed_from_consumer = false; + auto consumer_it = consumers.find(name); + if (consumer_it != consumers.end()) { + for (Ort::ConstNode consumer : consumer_it->second) { + if (!IsOnnxOp(consumer, "Transpose") || + consumer.GetInputs().size() != 1) { + continue; + } + for (Ort::ConstValueInfo output : consumer.GetOutputs()) { + auto fused_it = input_indices.find(Name(output)); + if (fused_it != input_indices.end()) { + sources->insert(fused_it->second); + keyed_from_consumer = true; + } + } + if (keyed_from_consumer) { + break; + } + } + } + if (keyed_from_consumer) { + continue; + } auto producer_it = producers.find(name); if (producer_it == producers.end()) { return false; @@ -333,6 +370,25 @@ bool CollectShapePreservingSources( return true; } +std::string FusionBoundaryKey(const std::vector& input_names, + const std::vector& output_names) { + auto normalize = [](std::vector names) { + names.erase( + std::remove_if(names.begin(), names.end(), + [](const std::string& name) { return name.empty(); }), + names.end()); + std::sort(names.begin(), names.end()); + names.erase(std::unique(names.begin(), names.end()), names.end()); + std::string key; + for (const std::string& name : names) { + key += name; + key += '\x1e'; + } + return key; + }; + return normalize(input_names) + "|" + normalize(output_names); +} + bool ValueDependsOnlyOnInputShapes( const std::string& value_name, const std::unordered_map& producers, diff --git a/musa/ep/src/fusion/fusion_matcher_utils.h b/musa/ep/src/fusion/fusion_matcher_utils.h index b485dc1..b811542 100644 --- a/musa/ep/src/fusion/fusion_matcher_utils.h +++ b/musa/ep/src/fusion/fusion_matcher_utils.h @@ -59,16 +59,23 @@ bool ValueHasExternalConsumerOrGraphOutput( std::optional ReadScalarIntInitializer(Ort::ConstValueInfo value_info); bool HasCastBoundary(const std::vector& fusion_nodes); +// Sorted, de-duplicated signature of a fusion boundary (input and output value +// names). GetCapability() and Compile() build the same key from the boundary, +// which is all either side sees of the other. +std::string FusionBoundaryKey(const std::vector& input_names, + const std::vector& output_names); + // Returns true when `value_name` is produced only from shape arithmetic: // constants, plus Shape ops over inputs of the current fusion. The output is // then a pure function of the shapes of the fused inputs listed in // `source_indices` (empty means the value is constant). A Shape whose source // is not itself a fused input is accepted when the source is reached from one // through shape-preserving ops only (Transpose with a perm attribute, Cast, -// Identity): its output shape, and therefore the Shape value, is still a pure -// function of those inputs' shapes. Anything that could make the value depend -// on data, on a tensor outside the fusion by another path, or on an op that is -// not shape-only returns false. +// Identity), or when another Transpose reading the same value produces one of +// the fusion's inputs (both expose the value's shape up to a fixed +// permutation). Anything that could make the value depend on data, on a tensor +// outside the fusion by another path, or on an op that is not shape-only +// returns false. bool ValueDependsOnlyOnInputShapes( const std::string& value_name, const std::unordered_map& producers, diff --git a/musa/ep/src/fusion/fusion_node_compute.cc b/musa/ep/src/fusion/fusion_node_compute.cc index 6f3eea4..766387f 100644 --- a/musa/ep/src/fusion/fusion_node_compute.cc +++ b/musa/ep/src/fusion/fusion_node_compute.cc @@ -331,7 +331,8 @@ OrtStatus* ORT_API_CALL MusaEp::CompileImpl( finder_name = "FindMultiKqvMhaOutputProjectionFusions"; } else if (IsMhtaScaledDotProductAttentionFusionGraph(graph)) { fusion_compute = - CreateMhtaScaledDotProductAttentionFusion(graph, fused_node); + CreateMhtaScaledDotProductAttentionFusion(graph, fused_node, + &ep->MaskEndHints()); finder_name = "FindMhtaScaledDotProductAttentionFusions"; } else if (IsQkvAttentionOutputProjectionFusionGraph(graph)) { fusion_compute = diff --git a/musa/ep/src/fusion/mhta_scaled_dot_product_attention_fusion.cc b/musa/ep/src/fusion/mhta_scaled_dot_product_attention_fusion.cc index b93cdda..ecf0233 100644 --- a/musa/ep/src/fusion/mhta_scaled_dot_product_attention_fusion.cc +++ b/musa/ep/src/fusion/mhta_scaled_dot_product_attention_fusion.cc @@ -1363,8 +1363,60 @@ bool IsMhtaScaledDotProductAttentionFusionGraph(Ort::ConstGraph graph) { lseq_bf16_shd_boundary); } +namespace { + +// Compile() only holds the fused subgraph, so the matcher's proof that the +// end is a function of fused input shapes arrives through the hint map. Both +// sides key it by the same boundary signature. +void MaskEndCacheHint( + const musa_ep::FusionMaskEndHints* hints, Ort::ConstGraph graph, + Ort::ConstNode fused_node, + const std::unordered_map& fused_indices, + const std::string& end_name, bool* cacheable, + std::vector* source_indices) { + *cacheable = false; + source_indices->clear(); + if (hints == nullptr) { + return; + } + std::vector input_names; + for (Ort::ConstValueInfo input : fused_node.GetInputs()) { + input_names.push_back(Name(input)); + } + for (Ort::ConstValueInfo initializer : graph.GetInitializers()) { + input_names.push_back(Name(initializer)); + } + std::vector output_names; + for (Ort::ConstValueInfo output : fused_node.GetOutputs()) { + output_names.push_back(Name(output)); + } + auto hint_it = + hints->find(musa_ep::FusionBoundaryKey(input_names, output_names)); + if (hint_it == hints->end()) { + return; + } + auto end_it = hint_it->second.find(end_name); + if (end_it == hint_it->second.end() || end_it->second.empty()) { + return; + } + std::vector indices; + indices.reserve(end_it->second.size()); + for (const std::string& name : end_it->second) { + auto index_it = fused_indices.find(name); + if (index_it == fused_indices.end()) { + return; + } + indices.push_back(index_it->second); + } + *cacheable = true; + *source_indices = std::move(indices); +} + +} // namespace + std::unique_ptr CreateMhtaScaledDotProductAttentionFusion( - Ort::ConstGraph graph, Ort::ConstNode fused_node) { + Ort::ConstGraph graph, Ort::ConstNode fused_node, + const musa_ep::FusionMaskEndHints* mask_end_hints) { Ort::ConstNode value_matmul{nullptr}; Ort::ConstNode score_matmul{nullptr}; Ort::ConstNode mul_node{nullptr}; @@ -1405,11 +1457,6 @@ std::unique_ptr CreateMhtaScaledDotProductAttentionFusion( } auto fused_indices = FusedInputIndices(fused_node); - std::unordered_set initializer_names; - for (Ort::ConstValueInfo initializer : graph.GetInitializers()) { - initializer_names.insert(Name(initializer)); - } - if (where_nodes.size() == 2) { Ort::ConstNode post_where{nullptr}; for (Ort::ConstNode node : graph.GetNodes()) { @@ -1490,11 +1537,11 @@ std::unique_ptr CreateMhtaScaledDotProductAttentionFusion( end.has_value()) { static_mask_end = *end; } else { - mask_end_index = - InputIndex(fused_indices, slice_it->second.GetInputs()[2]); - mask_end_cacheable = musa_ep::ValueDependsOnlyOnInputShapes( - Name(slice_it->second.GetInputs()[2]), producers, fused_indices, - initializer_names, &mask_end_source_indices); + const std::string end_name = Name(slice_it->second.GetInputs()[2]); + mask_end_index = InputIndex(fused_indices, end_name); + MaskEndCacheHint(mask_end_hints, graph, fused_node, fused_indices, + end_name, &mask_end_cacheable, + &mask_end_source_indices); // One line per fused raw-INT32-mask attention node at session // creation: 0 means the end is read from the device every run. std::fprintf(stderr, diff --git a/musa/ep/src/fusion/mhta_scaled_dot_product_attention_fusion.h b/musa/ep/src/fusion/mhta_scaled_dot_product_attention_fusion.h index 94c2a1d..94efc47 100644 --- a/musa/ep/src/fusion/mhta_scaled_dot_product_attention_fusion.h +++ b/musa/ep/src/fusion/mhta_scaled_dot_product_attention_fusion.h @@ -19,8 +19,12 @@ #include "onnxruntime_cxx_api.h" #undef ORT_API_MANUAL_INIT +#include "fusion/fusion_matcher.h" #include "fusion/fusion_node_compute.h" bool IsMhtaScaledDotProductAttentionFusionGraph(Ort::ConstGraph graph); +// `mask_end_hints` carries the Slice-end provenance the matcher proved on the +// whole graph; Compile() cannot see the end's producer chain any more. std::unique_ptr CreateMhtaScaledDotProductAttentionFusion( - Ort::ConstGraph graph, Ort::ConstNode fused_node); + Ort::ConstGraph graph, Ort::ConstNode fused_node, + const musa_ep::FusionMaskEndHints* mask_end_hints = nullptr); diff --git a/musa/ep/src/fusion/mhta_scaled_dot_product_attention_fusion_matcher.cc b/musa/ep/src/fusion/mhta_scaled_dot_product_attention_fusion_matcher.cc index 342952f..3a3f94c 100644 --- a/musa/ep/src/fusion/mhta_scaled_dot_product_attention_fusion_matcher.cc +++ b/musa/ep/src/fusion/mhta_scaled_dot_product_attention_fusion_matcher.cc @@ -590,7 +590,8 @@ bool TryAppendRawInt32MaskPreprocessing( const std::unordered_set& graph_output_names, const std::unordered_set& accepted_node_ids, std::unordered_set& selected, - std::vector& fusion_nodes) { + std::vector& fusion_nodes, + std::string* mask_end_name = nullptr) { Ort::ConstNode equal{nullptr}; if (!GetProducer(boolean_mask, equal) || !IsOnnxOp(equal, "Equal")) { return true; @@ -649,6 +650,9 @@ bool TryAppendRawInt32MaskPreprocessing( if (!starts.has_value() || *starts != std::vector{0}) return false; if (!axes.has_value() || *axes != std::vector{3}) return false; if (!steps.has_value() || *steps != std::vector{1}) return false; + if (mask_end_name != nullptr) { + *mask_end_name = Name(slice_inputs[2]); + } for (Ort::ConstNode node : {slice, cast, equal}) { if (!AddFusionNode(node, accepted_node_ids, selected, fusion_nodes)) { @@ -714,7 +718,8 @@ bool CanFuseBooleanMhtaScaledDotProductAttention( Ort::ConstNode value_matmul, const std::unordered_set& graph_output_names, const std::unordered_set& accepted_node_ids, - std::vector& fusion_nodes) { + std::vector& fusion_nodes, + std::string* mask_end_name = nullptr) { if (!IsOnnxOp(value_matmul, "MatMul") || accepted_node_ids.count(value_matmul.GetId()) != 0) return false; @@ -808,7 +813,8 @@ bool CanFuseBooleanMhtaScaledDotProductAttention( } if (!TryAppendRawInt32MaskPreprocessing(pre_inputs[0], pre_where, post_where, graph_output_names, accepted_node_ids, - selected, fusion_nodes)) { + selected, fusion_nodes, + mask_end_name)) { return false; } TryAppendBshdBoundaryTransposes(score_matmul, value_matmul, @@ -1313,24 +1319,119 @@ bool CanFuseSimRank3MhtaScaledDotProductAttention( return true; } +// Prove, on the whole graph, that the raw-mask Slice end depends only on the +// shapes of the fusion's own inputs. Compile() only receives the fused +// subgraph, so the answer has to travel through the hint map, keyed by the +// fusion boundary both sides can see. +void AddMaskEndHint( + const std::string& end_name, + const std::vector& fusion_nodes, + const std::unordered_set& graph_output_names, + const std::unordered_map& producers, + const std::unordered_set& initializer_names, + musa_ep::FusionMaskEndHints* hints) { + if (hints == nullptr || end_name.empty() || fusion_nodes.empty()) { + return; + } + std::unordered_set internal; + std::unordered_set node_ids; + for (Ort::ConstNode node : fusion_nodes) { + node_ids.insert(node.GetId()); + for (Ort::ConstValueInfo output : node.GetOutputs()) { + if (output != nullptr) { + internal.insert(Name(output)); + } + } + } + std::vector input_names; + std::unordered_map input_indices; + for (Ort::ConstNode node : fusion_nodes) { + for (Ort::ConstValueInfo input : node.GetInputs()) { + if (input == nullptr) { + continue; + } + const std::string name = Name(input); + if (name.empty() || internal.count(name) != 0) { + continue; + } + if (input_indices.emplace(name, input_indices.size()).second) { + input_names.push_back(name); + } + } + } + std::vector output_names; + for (Ort::ConstNode node : fusion_nodes) { + for (Ort::ConstValueInfo output : node.GetOutputs()) { + if (output == nullptr) { + continue; + } + bool external = graph_output_names.count(Name(output)) != 0; + for (const auto& consumer : output.GetConsumers()) { + if (node_ids.count(consumer.node.GetId()) == 0) { + external = true; + } + } + if (external) { + output_names.push_back(Name(output)); + } + } + } + std::vector source_indices; + if (!musa_ep::ValueDependsOnlyOnInputShapes(end_name, producers, input_indices, + initializer_names, + &source_indices)) { + return; + } + if (source_indices.empty()) { + return; // constant end: nothing to key on + } + std::vector source_names; + source_names.reserve(source_indices.size()); + for (size_t index : source_indices) { + if (index >= input_names.size()) { + return; + } + source_names.push_back(input_names[index]); + } + (*hints)[musa_ep::FusionBoundaryKey(input_names, output_names)][end_name] = + std::move(source_names); +} + } // namespace std::vector> FindMhtaScaledDotProductAttentionFusions( const std::vector& all_nodes, const std::unordered_set& graph_output_names, - const std::unordered_set& accepted_node_ids) { + const std::unordered_set& accepted_node_ids, + musa_ep::FusionMaskEndHints* mask_end_hints) { std::vector> fusions; + const auto producers = BuildProducerMap(all_nodes); + std::unordered_set initializer_names; + for (const Ort::ConstNode& node : all_nodes) { + for (Ort::ConstValueInfo input : node.GetInputs()) { + if (IsConstantInitializerValueInfo(input)) { + initializer_names.insert(Name(input)); + } + } + } for (Ort::ConstNode node : all_nodes) { std::vector fusion_nodes; - if (CanFuseBooleanMhtaScaledDotProductAttention( - node, graph_output_names, accepted_node_ids, fusion_nodes) || + std::string mask_end_name; + const bool boolean_fused = CanFuseBooleanMhtaScaledDotProductAttention( + node, graph_output_names, accepted_node_ids, fusion_nodes, + &mask_end_name); + if (boolean_fused || CanFuseLseqPreMaskMhtaScaledDotProductAttention( node, graph_output_names, accepted_node_ids, fusion_nodes) || CanFuseMhtaScaledDotProductAttention(node, graph_output_names, accepted_node_ids, fusion_nodes) || CanFuseSimRank3MhtaScaledDotProductAttention( node, graph_output_names, accepted_node_ids, fusion_nodes)) { + if (boolean_fused) { + AddMaskEndHint(mask_end_name, fusion_nodes, graph_output_names, + producers, initializer_names, mask_end_hints); + } fusions.push_back(std::move(fusion_nodes)); } } diff --git a/musa/ep/src/fusion/tile_concat_fusion.cc b/musa/ep/src/fusion/tile_concat_fusion.cc index b214c91..a3b21b0 100644 --- a/musa/ep/src/fusion/tile_concat_fusion.cc +++ b/musa/ep/src/fusion/tile_concat_fusion.cc @@ -463,21 +463,7 @@ musa_ep::ShapeExpr RemapShapeExprHint( std::string TileConcatFusionKey(const std::vector& input_names, const std::vector& output_names) { - auto normalize = [](std::vector names) { - names.erase( - std::remove_if(names.begin(), names.end(), - [](const std::string& name) { return name.empty(); }), - names.end()); - std::sort(names.begin(), names.end()); - names.erase(std::unique(names.begin(), names.end()), names.end()); - std::string key; - for (const std::string& name : names) { - key += name; - key += '\x1e'; - } - return key; - }; - return normalize(input_names) + "|" + normalize(output_names); + return FusionBoundaryKey(input_names, output_names); } bool IsTileConcatFusionGraph(Ort::ConstGraph graph) { diff --git a/test/fusion/test_mhta_scaled_dot_product_attention_fusion.py b/test/fusion/test_mhta_scaled_dot_product_attention_fusion.py index f93a053..f9ffd9c 100644 --- a/test/fusion/test_mhta_scaled_dot_product_attention_fusion.py +++ b/test/fusion/test_mhta_scaled_dot_product_attention_fusion.py @@ -930,23 +930,42 @@ def test_mhta_scaled_dot_product_attention_sim_rank3_div_temperature_fusion( assert "Softmax" not in musa_ops -def _mhta_dynamic_end_model(batch, heads, head_dim): - """Raw INT32 mask whose Slice end comes from Shape(Transpose(K)). +def _mhta_dynamic_end_model(batch, heads, head_dim, end_source="transpose_k"): + """Raw INT32 mask whose Slice end comes from a runtime shape. Q/K/V follow the plain bshd layout used by the other raw-mask tests. K is - [batch, heads, head_dim, seqlen] with a symbolic seqlen; a shape-only - `Transpose(K, perm=[0, 1, 3, 2])` (no other consumer, like - `Transpose_1` in ranking-gr) feeds Shape -> Gather(2) -> Unsqueeze, so the - mask Slice end equals the runtime seqlen. A cache keyed on K's shape must - therefore change when seqlen changes. + [batch, heads, head_dim, seqlen] with a symbolic seqlen. `end_source` + selects what feeds the Slice end: + + - "transpose_k": Shape(Transpose(K, perm=[0, 1, 3, 2])) -> Gather -> Unsqueeze. + K is a fused input, so the matcher can prove the end from its shape (the + positive case, like a recommendation model's Transpose_1 chain). + - "transpose_other": the same chain reads a second input K2 that the + attention never touches, so the end cannot be keyed on any fused input. + - "abs_k": Shape(Abs(K)) - a producer that is not in the shape-preserving + whitelist, so the end is not provable by the conservative rules. """ def vi(name, elem_type, dims): return helper.make_tensor_value_info(name, elem_type, dims) - nodes = [ - helper.make_node("Transpose", ["K"], ["KShapeOnly"], perm=[0, 1, 3, 2]), - helper.make_node("Shape", ["KShapeOnly"], ["KShape"]), - helper.make_node("Gather", ["KShape", "seqlen_index"], ["Seqlen"], axis=0), + end_nodes = { + "transpose_k": [ + helper.make_node("Transpose", ["K"], ["KShapeOnly"], perm=[0, 1, 3, 2]), + helper.make_node("Shape", ["KShapeOnly"], ["KShape"]), + helper.make_node("Gather", ["KShape", "seqlen_index_2"], ["Seqlen"], axis=0), + ], + "transpose_other": [ + helper.make_node("Transpose", ["K2"], ["K2ShapeOnly"], perm=[0, 1, 3, 2]), + helper.make_node("Shape", ["K2ShapeOnly"], ["K2Shape"]), + helper.make_node("Gather", ["K2Shape", "seqlen_index_2"], ["Seqlen"], axis=0), + ], + "abs_k": [ + helper.make_node("Abs", ["K"], ["KAbs"]), + helper.make_node("Shape", ["KAbs"], ["KShape"]), + helper.make_node("Gather", ["KShape", "seqlen_index_3"], ["Seqlen"], axis=0), + ], + }[end_source] + nodes = end_nodes + [ helper.make_node("Unsqueeze", ["Seqlen", "axis0"], ["Ends"]), helper.make_node("Slice", ["RawMask", "starts", "Ends", "axes", "steps"], ["MaskSlice"]), helper.make_node("Cast", ["MaskSlice"], ["MaskInt64"], to=TensorProto.INT64), @@ -961,7 +980,8 @@ def vi(name, elem_type, dims): initializers = [ numpy_helper.from_array(np.asarray(v, dtype=dtype), name=name) for name, v, dtype in [ - ("seqlen_index", np.asarray(2, dtype=np.int64), np.int64), + ("seqlen_index_2", np.asarray(2, dtype=np.int64), np.int64), + ("seqlen_index_3", np.asarray(3, dtype=np.int64), np.int64), ("axis0", [0], np.int64), ("starts", [0], np.int64), ("axes", [3], np.int64), @@ -972,15 +992,20 @@ def vi(name, elem_type, dims): ("zero", np.asarray(0.0, dtype=np.float32), np.float32), ] ] + graph_inputs = [ + vi("Q", TensorProto.FLOAT, [batch, heads, "seqlen", head_dim]), + vi("K", TensorProto.FLOAT, [batch, heads, head_dim, "seqlen"]), + vi("V", TensorProto.FLOAT, [batch, heads, "seqlen", head_dim]), + vi("RawMask", TensorProto.INT32, [batch, 1, "mask_rows", "mask_cols"]), + ] + if end_source == "transpose_other": + graph_inputs.append( + vi("K2", TensorProto.FLOAT, [batch, heads, head_dim, "seqlen"]) + ) graph = helper.make_graph( nodes, "mhta_dynamic_slice_end_graph", - [ - vi("Q", TensorProto.FLOAT, [batch, heads, "seqlen", head_dim]), - vi("K", TensorProto.FLOAT, [batch, heads, head_dim, "seqlen"]), - vi("V", TensorProto.FLOAT, [batch, heads, "seqlen", head_dim]), - vi("RawMask", TensorProto.INT32, [batch, 1, "mask_rows", "mask_cols"]), - ], + graph_inputs, [vi("Y", TensorProto.FLOAT, [batch, heads, "seqlen", head_dim])], initializer=initializers, ) @@ -992,32 +1017,50 @@ def vi(name, elem_type, dims): def _mhta_dynamic_feeds(rng, batch, heads, head_dim, seqlen): raw_mask = rng.integers(0, 3, (batch, 1, seqlen, seqlen + 2), dtype=np.int32) raw_mask[:, :, -1, :seqlen] = 0 - return { + feeds = { "Q": rng.standard_normal((batch, heads, seqlen, head_dim)).astype(np.float32), "K": rng.standard_normal((batch, heads, head_dim, seqlen)).astype(np.float32), "V": rng.standard_normal((batch, heads, seqlen, head_dim)).astype(np.float32), "RawMask": raw_mask, } + return feeds -def test_mhta_slice_end_cache_follows_k_shape(capfd): - """The cached Slice end must be keyed on K's shape, not reused blindly.""" +def _run_dynamic_end_cache_case(capfd, end_source, expected_cacheable): rng = np.random.default_rng(61) batch, heads, head_dim = 1, 3, 4 - model = _mhta_dynamic_end_model(batch, heads, head_dim) - + model = _mhta_dynamic_end_model(batch, heads, head_dim, end_source) so = ort.SessionOptions() so.add_session_config_entry("session.disable_cpu_ep_fallback", "1") so.add_provider_for_devices(musa_devices(), {}) capfd.readouterr() musa_session = ort.InferenceSession(model, sess_options=so) err = capfd.readouterr().err - assert "MUSA_MHTA_SLICE_END cacheable=1" in err, err - + expected = f"MUSA_MHTA_SLICE_END cacheable={expected_cacheable}" + assert expected in err, err for seqlen in (5, 4, 5): feeds = _mhta_dynamic_feeds(rng, batch, heads, head_dim, seqlen) + if end_source == "transpose_other": + feeds["K2"] = rng.standard_normal( + (batch, heads, head_dim, seqlen) + ).astype(np.float32) cpu = ort.InferenceSession(model, providers=["CPUExecutionProvider"]).run(None, feeds) out = musa_session.run(None, feeds) np.testing.assert_allclose(out[0], cpu[0], rtol=1e-4, atol=1e-4, - err_msg=f"seqlen={seqlen}") + err_msg=f"{end_source} seqlen={seqlen}") np.testing.assert_array_equal(out[0][:, :, -1, :], 0.0) + + +def test_mhta_slice_end_cache_follows_k_shape(capfd): + """The cached Slice end must be keyed on K's shape, not reused blindly.""" + _run_dynamic_end_cache_case(capfd, "transpose_k", 1) + + +def test_mhta_slice_end_cache_rejects_unrelated_shape_source(capfd): + """An end derived from a tensor outside the fusion must not be cached.""" + _run_dynamic_end_cache_case(capfd, "transpose_other", 0) + + +def test_mhta_slice_end_cache_rejects_unproven_chain(capfd): + """A producer outside the shape-preserving whitelist must not be cached.""" + _run_dynamic_end_cache_case(capfd, "abs_k", 0) From a37c7c88efa731ebacecadfcf368edf77b1131fd Mon Sep 17 00:00:00 2001 From: "yichao.hu" Date: Sun, 18 Oct 2026 16:11:34 -0400 Subject: [PATCH 33/52] feat(fusion): fuse runtime Tile repeats behind a same-build switch A fused node cannot declare CPU inputs, so a runtime repeats tensor is copied onto the device and read back with a stream synchronize. 2d06d07 avoided that by leaving such Tiles standalone; on a multi-task recommendation model the 13 remaining Tiles cost about 0.26 ms end to end. With ORT_MUSA_TILE_CONCAT_RUNTIME_REPEATS=1 the matcher keeps those Tiles in the fusion and the kernel reads the shared repeats tensor once per execution (memoized by input index) instead of once per Tile. The switch is off by default so the tree keeps the safe fallback until a measurement round accepts the trade; a unit test covers the outside-source Tile under the switch. --- musa/ep/src/fusion/tile_concat_fusion.cc | 12 ++++- .../src/fusion/tile_concat_fusion_matcher.cc | 21 ++++++++- test/fusion/test_tile_concat_shape_chain.py | 46 +++++++++++++++++++ 3 files changed, 77 insertions(+), 2 deletions(-) diff --git a/musa/ep/src/fusion/tile_concat_fusion.cc b/musa/ep/src/fusion/tile_concat_fusion.cc index a3b21b0..dcefe3e 100644 --- a/musa/ep/src/fusion/tile_concat_fusion.cc +++ b/musa/ep/src/fusion/tile_concat_fusion.cc @@ -67,6 +67,10 @@ struct TileConcatFusionCompute : FusionNodeCompute { std::vector> shapes; std::vector input_data; std::vector> temp_buffers; + // Runtime repeats shared by several Tiles (a multi-task model's second group + // feeds 13 Tiles from one tensor): read the tensor back once per + // execution instead of once per Tile. + std::unordered_map> repeats_cache; shapes.reserve(inputs.size()); input_data.reserve(inputs.size()); @@ -126,7 +130,13 @@ struct TileConcatFusionCompute : FusionNodeCompute { ORT_EP_FAIL, "TileConcat could not recompute the fused Tile repeats"); } - repeats = ReadIntTensor(ctx, input_plan.repeats_input_index); + auto cached = repeats_cache.find(input_plan.repeats_input_index); + if (cached != repeats_cache.end()) { + repeats = cached->second; + } else { + repeats = ReadIntTensor(ctx, input_plan.repeats_input_index); + repeats_cache.emplace(input_plan.repeats_input_index, repeats); + } } if (repeats.size() < input_shape.size()) { repeats.insert(repeats.begin(), input_shape.size() - repeats.size(), diff --git a/musa/ep/src/fusion/tile_concat_fusion_matcher.cc b/musa/ep/src/fusion/tile_concat_fusion_matcher.cc index 7e1d41b..e93eb15 100644 --- a/musa/ep/src/fusion/tile_concat_fusion_matcher.cc +++ b/musa/ep/src/fusion/tile_concat_fusion_matcher.cc @@ -13,6 +13,8 @@ #include #include +#include +#include #include #include #include @@ -102,6 +104,16 @@ bool CollectShapeChain( return true; } +// Same-build A/B switch. With it enabled a Tile whose repeats is not provably +// shape-only is still fused: the fused kernel reads the repeats tensor once +// per execution (one device copy plus one read-back) instead of leaving the +// Tile standalone. Disabled by default so the tree keeps the safe fallback +// until a measurement round accepts the trade. +bool RuntimeRepeatsEnabled() { + const char* value = std::getenv("ORT_MUSA_TILE_CONCAT_RUNTIME_REPEATS"); + return value != nullptr && value[0] != '\0' && std::strcmp(value, "0") != 0; +} + // The values a fused node declares as inputs. A Shape source has to be one of // them: the fused node cannot gain an input that Tile+Concat do not already // read, so a chain reaching outside would not be recomputable at run time. @@ -205,7 +217,9 @@ bool CanFuseTileConcat( ShapeChain chain; std::unordered_set visiting; if (!CollectShapeChain(repeats, producers, visiting, chain)) { - recomputable[i] = false; // not a shape-only chain + // Not a shape-only chain. With the runtime-repeats switch on, keep the + // Tile fused and let the kernel read the repeats tensor itself. + recomputable[i] = RuntimeRepeatsEnabled(); continue; } for (const auto& [name, value_info] : chain.leaves) { @@ -259,6 +273,11 @@ bool CanFuseTileConcat( ShapeExpr expr = BuildShapeExpr(repeats, producers, base_inputs, initializer_names, initializer_values); if (expr.Empty()) { + if (RuntimeRepeatsEnabled()) { + // A Shape source that is not an input of Tile+Concat: keep the Tile + // fused and let the kernel read the repeats once per execution. + continue; + } // A Shape source that is not an input of Tile+Concat cannot be // recomputed inside the fusion. recomputable[i] = false; diff --git a/test/fusion/test_tile_concat_shape_chain.py b/test/fusion/test_tile_concat_shape_chain.py index 0dfe2b1..1fd207e 100644 --- a/test/fusion/test_tile_concat_shape_chain.py +++ b/test/fusion/test_tile_concat_shape_chain.py @@ -449,3 +449,49 @@ def test_tile_concat_repeats_chain_with_reshape_element_mismatch(): _session(model, use_musa=False).run(None, feeds) with pytest.raises(Exception): _session(model, use_musa=True).run(None, feeds) + + +def test_tile_concat_runtime_repeats_switch_fuses_outside_source( + tmp_path, monkeypatch +): + """With runtime repeats enabled the outside-source Tile joins the fusion. + + The fused kernel then reads the repeats tensor once per execution instead + of leaving the Tile standalone; the result must stay identical. + """ + monkeypatch.setenv("ORT_MUSA_TILE_CONCAT_RUNTIME_REPEATS", "1") + rng = np.random.default_rng(42) + nodes = [ + helper.make_node("Shape", ["other"], ["other_shape"]), + helper.make_node("Gather", ["other_shape", "zero_scalar"], ["other_dim"], axis=0), + helper.make_node("Unsqueeze", ["other_dim", "zero"], ["other_dim_1d"]), + helper.make_node("Concat", ["other_dim_1d", "one"], ["repeats"], axis=0), + helper.make_node("Tile", ["data", "repeats"], ["tiled"]), + helper.make_node("Concat", ["tiled", "query"], ["y_pre"], axis=1), + helper.make_node("Relu", ["y_pre"], ["y"]), + ] + initializers = _batch_chain_initializers() + model = _build_model( + "tile_concat_runtime_repeats_outside_source", + nodes, + [ + helper.make_tensor_value_info("query", TensorProto.FLOAT, [3, 4]), + helper.make_tensor_value_info("other", TensorProto.FLOAT, ["m"]), + ], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, None)], + initializers, + ) + feeds = { + "query": rng.standard_normal((3, 4)).astype(np.float32), + "other": rng.standard_normal(3).astype(np.float32), + } + expected = _session(model, use_musa=False).run(None, feeds)[0] + musa = _session( + model, use_musa=True, profile_prefix=str(tmp_path / "runtime_repeats") + ) + actual = musa.run(None, feeds)[0] + np.testing.assert_allclose(actual, expected, rtol=1e-6, atol=1e-6) + musa_ops = _musa_ops(_read_ops_by_provider(musa)) + assert any(str(op).startswith("MUSAExecutionProvider_") for op in musa_ops) + assert "Tile" not in musa_ops + assert "Concat" not in musa_ops From 2ff6d5b3789e5add461fb9bb7a54ddf2adb1273c Mon Sep 17 00:00:00 2001 From: "yichao.hu" Date: Sun, 18 Oct 2026 16:18:49 -0400 Subject: [PATCH 34/52] fix(fusion): compile fixes for the hint plumbing and boundary key InputIndex takes the end value info, not its name; tile_concat_fusion.cc needs the matcher-utils header for FusionBoundaryKey. --- .../src/fusion/mhta_scaled_dot_product_attention_fusion.cc | 5 +++-- musa/ep/src/fusion/tile_concat_fusion.cc | 1 + 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/musa/ep/src/fusion/mhta_scaled_dot_product_attention_fusion.cc b/musa/ep/src/fusion/mhta_scaled_dot_product_attention_fusion.cc index ecf0233..a4bedda 100644 --- a/musa/ep/src/fusion/mhta_scaled_dot_product_attention_fusion.cc +++ b/musa/ep/src/fusion/mhta_scaled_dot_product_attention_fusion.cc @@ -1537,8 +1537,9 @@ std::unique_ptr CreateMhtaScaledDotProductAttentionFusion( end.has_value()) { static_mask_end = *end; } else { - const std::string end_name = Name(slice_it->second.GetInputs()[2]); - mask_end_index = InputIndex(fused_indices, end_name); + Ort::ConstValueInfo end_value = slice_it->second.GetInputs()[2]; + const std::string end_name = Name(end_value); + mask_end_index = InputIndex(fused_indices, end_value); MaskEndCacheHint(mask_end_hints, graph, fused_node, fused_indices, end_name, &mask_end_cacheable, &mask_end_source_indices); diff --git a/musa/ep/src/fusion/tile_concat_fusion.cc b/musa/ep/src/fusion/tile_concat_fusion.cc index dcefe3e..7508b8b 100644 --- a/musa/ep/src/fusion/tile_concat_fusion.cc +++ b/musa/ep/src/fusion/tile_concat_fusion.cc @@ -27,6 +27,7 @@ #include #include +#include "fusion/fusion_matcher_utils.h" #include "fusion/shape_expr.h" #include "kernels/shared_inc/op_kernel_common.h" #include "kernels/tensor/concat_impl.h" From 85d26eac6e2e411d49954f96ca19e45d65fabfe4 Mon Sep 17 00:00:00 2001 From: "yichao.hu" Date: Sun, 18 Oct 2026 16:24:18 -0400 Subject: [PATCH 35/52] fix(fusion): qualify FusionBoundaryKey in tile_concat_fusion.cc TileConcatFusionKey lives at global scope; the key helper is in musa_ep. --- musa/ep/src/fusion/tile_concat_fusion.cc | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/musa/ep/src/fusion/tile_concat_fusion.cc b/musa/ep/src/fusion/tile_concat_fusion.cc index 7508b8b..48519e7 100644 --- a/musa/ep/src/fusion/tile_concat_fusion.cc +++ b/musa/ep/src/fusion/tile_concat_fusion.cc @@ -474,7 +474,7 @@ musa_ep::ShapeExpr RemapShapeExprHint( std::string TileConcatFusionKey(const std::vector& input_names, const std::vector& output_names) { - return FusionBoundaryKey(input_names, output_names); + return musa_ep::FusionBoundaryKey(input_names, output_names); } bool IsTileConcatFusionGraph(Ort::ConstGraph graph) { From 4d84cc5d3fd49fa0547eb46baaec76f9c0b4d7e7 Mon Sep 17 00:00:00 2001 From: "yichao.hu" Date: Sun, 18 Oct 2026 16:37:36 -0400 Subject: [PATCH 36/52] test(fusion): keep the CPU fallback for the dynamic mask-end tests The Slice-end chain (Shape/Gather/Transpose) legitimately stays on the CPU EP; disabling the fallback rejected the session after the matcher had already run. --- test/fusion/test_mhta_scaled_dot_product_attention_fusion.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/test/fusion/test_mhta_scaled_dot_product_attention_fusion.py b/test/fusion/test_mhta_scaled_dot_product_attention_fusion.py index f9ffd9c..b9508fc 100644 --- a/test/fusion/test_mhta_scaled_dot_product_attention_fusion.py +++ b/test/fusion/test_mhta_scaled_dot_product_attention_fusion.py @@ -1030,8 +1030,9 @@ def _run_dynamic_end_cache_case(capfd, end_source, expected_cacheable): rng = np.random.default_rng(61) batch, heads, head_dim = 1, 3, 4 model = _mhta_dynamic_end_model(batch, heads, head_dim, end_source) + # The end chain itself legitimately runs on the CPU EP, so this session + # must keep the CPU fallback enabled; only the attention body is fused. so = ort.SessionOptions() - so.add_session_config_entry("session.disable_cpu_ep_fallback", "1") so.add_provider_for_devices(musa_devices(), {}) capfd.readouterr() musa_session = ort.InferenceSession(model, sess_options=so) From d2b90c0f40f62a24c1778dce48c4592b82334e6f Mon Sep 17 00:00:00 2001 From: "yichao.hu" Date: Sun, 18 Oct 2026 16:38:15 -0400 Subject: [PATCH 37/52] diag(fusion): gate an MHTA Slice-end cache trace behind an env flag ORT_MUSA_MHTA_END_TRACE=1 prints hit/miss, key size, resolved value and Sk for the first 64 computes, so a measurement round can tell whether the cache is hit and whether the read is the remaining cost. --- ...hta_scaled_dot_product_attention_fusion.cc | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/musa/ep/src/fusion/mhta_scaled_dot_product_attention_fusion.cc b/musa/ep/src/fusion/mhta_scaled_dot_product_attention_fusion.cc index a4bedda..31cc41f 100644 --- a/musa/ep/src/fusion/mhta_scaled_dot_product_attention_fusion.cc +++ b/musa/ep/src/fusion/mhta_scaled_dot_product_attention_fusion.cc @@ -15,6 +15,7 @@ #include +#include #include #include #include @@ -814,9 +815,11 @@ OrtStatus* MhtaScaledDotProductAttentionFusionCompute::Compute( shape_key.push_back(-1); } } + const size_t key_size = shape_key.size(); std::lock_guard lock(mask_end_mutex_); const bool cache_hit = mask_end_cacheable_ && mask_end_cached_ && mask_end_cached_shape_key_ == shape_key; + bool reread = false; if (cache_hit) { mask_end_value = mask_end_cached_value_; } else { @@ -831,6 +834,22 @@ OrtStatus* MhtaScaledDotProductAttentionFusionCompute::Compute( // Not shape-derived after all: re-read once and refresh the cache. RETURN_IF_ERROR(ReadMaskEndValue(mask_end, stream, &mask_end_value)); mask_end_cached_value_ = mask_end_value; + reread = true; + } + if (std::getenv("ORT_MUSA_MHTA_END_TRACE") != nullptr) { + static std::atomic trace_left{64}; + if (trace_left.load() > 0) { + trace_left.fetch_sub(1); + std::fprintf( + stderr, + "MUSA_MHTA_END node=%p cacheable=%d cached=%d hit=%d reread=%d " + "key=%zu sel=%zu value=%lld seqlen_k=%lld\n", + static_cast(this), mask_end_cacheable_ ? 1 : 0, + mask_end_cached_ ? 1 : 0, cache_hit ? 1 : 0, reread ? 1 : 0, + key_size, mask_end_source_indices_.size(), + static_cast(mask_end_value), + static_cast(seqlen_k)); + } } } if (mask_end_value != seqlen_k) { From c0415db7215ff2642143c6994306c5e48efdde8b Mon Sep 17 00:00:00 2001 From: "yichao.hu" Date: Sun, 18 Oct 2026 16:49:09 -0400 Subject: [PATCH 38/52] fix(fusion): pass the proved mask-end cache hint on the kBhsd path too The boolean raw-INT32 branch has two construction sites. The bshd-strided-flash one forwarded the matcher's proof; the kBhsd fallback did not, so a recommendation model's fused nodes printed cacheable=1 at Compile and ran with cacheable=0, reading the Slice end from the device every run. The runtime trace added for the diagnosis showed the mismatch (compile cacheable=1, compute cacheable=0). --- musa/ep/src/fusion/mhta_scaled_dot_product_attention_fusion.cc | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/musa/ep/src/fusion/mhta_scaled_dot_product_attention_fusion.cc b/musa/ep/src/fusion/mhta_scaled_dot_product_attention_fusion.cc index 31cc41f..f9d2f99 100644 --- a/musa/ep/src/fusion/mhta_scaled_dot_product_attention_fusion.cc +++ b/musa/ep/src/fusion/mhta_scaled_dot_product_attention_fusion.cc @@ -1590,7 +1590,8 @@ std::unique_ptr CreateMhtaScaledDotProductAttentionFusion( InputIndex(fused_indices, value_inputs[1]), InputIndex(fused_indices, mask_input), scale, 0.0f, MhtaSdpaLayout::kBhsd, true, boolean_mask_int32, mask_end_index, - static_mask_end); + static_mask_end, false, false, mask_end_cacheable, + mask_end_source_indices); } if (where_nodes.size() == 1 && !add_node && !div_node && !einsum_node && From 18b8aac2c5190c37fee29e7d80f7446eff11177c Mon Sep 17 00:00:00 2001 From: "yichao.hu" Date: Sun, 18 Oct 2026 17:13:28 -0400 Subject: [PATCH 39/52] perf(kernel): make Reshape may-inplace instead of always copying c4 removed the hard input/output alias so a CPU-produced input could never be read as device memory; every Reshape then allocated a fresh output and copied (about 0.19 ms over 46 Reshapes on a graph model, 0.27 ms on a recommendation model). Advertise a mutable alias instead (KernelDefBuilder_AddInputOutputMutableAliases, the EP API's MayInplace): the allocator reuses the input buffer only when it is safe, Compute() skips the copy when src == dst, and every other case still copies. With MemcpyFromHost registered the kernel only sees device tensors, so the staged-input path keeps c4's safety. Adds a trace line and three tests (reuse, copy with a second consumer, staged CPU-produced input). --- musa/ep/src/kernels/tensor/reshape.cc | 20 ++++- test/ops/test_reshape_may_inplace.py | 108 ++++++++++++++++++++++++++ 2 files changed, 126 insertions(+), 2 deletions(-) create mode 100644 test/ops/test_reshape_may_inplace.py diff --git a/musa/ep/src/kernels/tensor/reshape.cc b/musa/ep/src/kernels/tensor/reshape.cc index 6dd3211..e202a76 100644 --- a/musa/ep/src/kernels/tensor/reshape.cc +++ b/musa/ep/src/kernels/tensor/reshape.cc @@ -11,6 +11,7 @@ // See the License for the specific language governing permissions and // limitations under the License. +#include #include #include "shared_inc/op_kernel_common.h" @@ -70,6 +71,16 @@ OrtStatus* Reshape::Compute(Ort::KernelContext& ctx) const { std::vector out_shape = ResolveReshapeOutputShape(shape0, ReadIntTensor(ctx, 1), allowzero_); Ort::UnownedValue y = ctx.GetOutput(0, out_shape); + const void* src = input0.GetTensorRawData(); + const void* dst = y.GetTensorMutableRawData(); + const bool inplace = src == dst; + if (std::getenv("ORT_MUSA_RESHAPE_TRACE") != nullptr) { + std::fprintf(stderr, "MUSA_RESHAPE inplace=%d\n", inplace ? 1 : 0); + } + if (inplace) { + // May-inplace: the allocator reused the input buffer for the output. + return nullptr; + } return CopyRawTensor(input0, y, input0.GetTensorSizeInBytes(), GetComputeStream(ctx)); } @@ -110,7 +121,11 @@ ONNX_OPERATOR_VERSIONED_KERNEL_EX( Reshape, kOnnxDomain, 5, 12, (Ort::KernelDefBuilder() .AddTypeConstraint("T", AllFixedSizeTensorTypesNoBFloat16()) - .SetInputMemType(1, OrtMemTypeCPUInput)), + .SetInputMemType(1, OrtMemTypeCPUInput) + // May-inplace: reuse the input buffer for the output whenever the + // allocator proves it safe. Compute() skips the copy when src == dst + // and copies otherwise, so a staged input keeps c4's safe path. + .AddInputOutputMutableAlias(0, 0)), Reshape) // The opset-21 schema only adds the int4/uint4 element types, which this @@ -120,7 +135,8 @@ ONNX_OPERATOR_VERSIONED_KERNEL_EX( Reshape, kOnnxDomain, 13, 21, (Ort::KernelDefBuilder() .AddTypeConstraint("T", AllFixedSizeTensorTypes()) - .SetInputMemType(1, OrtMemTypeCPUInput)), + .SetInputMemType(1, OrtMemTypeCPUInput) + .AddInputOutputMutableAlias(0, 0)), Reshape) ONNX_OPERATOR_VERSIONED_KERNEL_EX( diff --git a/test/ops/test_reshape_may_inplace.py b/test/ops/test_reshape_may_inplace.py new file mode 100644 index 0000000..ee88164 --- /dev/null +++ b/test/ops/test_reshape_may_inplace.py @@ -0,0 +1,108 @@ +# Copyright (c) Moore Threads Technology Co., Ltd. All rights reserved. +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Reshape may-inplace: reuse the input buffer only when the allocator allows. + +c4 removed the hard input/output alias so a CPU-produced input could never be +read as device memory. With MemcpyFromHost registered the Reshape kernel only +receives device tensors, so the kernel can advertise a *mutable* alias: ORT +reuses the buffer when it is safe, the kernel skips the copy when src == dst, +and every other case still copies (c4's safety is not weakened). +""" + +import numpy as np +import onnxruntime as ort +from onnx import TensorProto, helper, numpy_helper + +from op_test_utils import ( + build_graph_model, + musa_devices, + run_model_and_compare_with_cpu_fallback, +) + + +def _musa_session(model: bytes) -> ort.InferenceSession: + options = ort.SessionOptions() + options.add_session_config_entry("session.disable_cpu_ep_fallback", "1") + options.add_provider_for_devices(musa_devices(), {}) + return ort.InferenceSession(model, sess_options=options) + + +def test_reshape_may_inplace_reuses_single_consumer_device_output(capfd): + """Add's device output feeds only Reshape: the buffer can be reused.""" + rng = np.random.default_rng(7) + x = rng.standard_normal((2, 3)).astype(np.float32) + shape = np.array([3, 2], dtype=np.int64) + one = np.array(1.0, dtype=np.float32) + model = build_graph_model( + [ + helper.make_node("Add", ["X", "one"], ["added"]), + helper.make_node("Reshape", ["added", "shape"], ["Y"]), + ], + {"X": x, "shape": shape}, + [("Y", TensorProto.FLOAT)], + initializers=[numpy_helper.from_array(one, name="one")], + name="reshape_may_inplace_reuse_graph", + ) + + capfd.readouterr() + session = _musa_session(model) + out = session.run(None, {"X": x, "shape": shape})[0] + err = capfd.readouterr().err + np.testing.assert_allclose(out, x + 1.0, rtol=0, atol=0) + assert "MUSA_RESHAPE inplace=1" in err, err + + +def test_reshape_may_inplace_copies_with_another_consumer(capfd): + """Add's output has a second consumer, so Reshape must copy.""" + rng = np.random.default_rng(8) + x = rng.standard_normal((2, 3)).astype(np.float32) + shape = np.array([3, 2], dtype=np.int64) + one = np.array(1.0, dtype=np.float32) + model = build_graph_model( + [ + helper.make_node("Add", ["X", "one"], ["added"]), + helper.make_node("Reshape", ["added", "shape"], ["Y"]), + helper.make_node("Relu", ["added"], ["Z"]), + ], + {"X": x, "shape": shape}, + [("Y", TensorProto.FLOAT), ("Z", TensorProto.FLOAT)], + initializers=[numpy_helper.from_array(one, name="one")], + name="reshape_may_inplace_copy_graph", + ) + + capfd.readouterr() + session = _musa_session(model) + y, z = session.run(None, {"X": x, "shape": shape}) + err = capfd.readouterr().err + np.testing.assert_allclose(y, (x + 1.0).reshape(3, 2), rtol=0, atol=0) + np.testing.assert_allclose(z, np.maximum(x + 1.0, 0.0), rtol=0, atol=0) + assert "MUSA_RESHAPE inplace=0" in err, err + + +def test_reshape_after_cpu_produced_input_stays_correct(monkeypatch): + """A CPU-produced input is staged; the may-inplace alias stays safe.""" + x = np.random.default_rng(10).standard_normal((1, 4, 1029)).astype(np.float16) + shape = np.array([1, 4, 1029], dtype=np.int64) + model = build_graph_model( + [ + helper.make_node("Sin", ["X"], ["cpu_output"]), + helper.make_node("Reshape", ["cpu_output", "shape"], ["Y"]), + ], + {"X": x, "shape": shape}, + [("Y", TensorProto.FLOAT16)], + name="reshape_may_inplace_staged_graph", + ) + monkeypatch.setenv("ORT_MUSA_TEST_FAIL_RESHAPE_CPU_INPUT", "1") + run_model_and_compare_with_cpu_fallback( + model, {"X": x, "shape": shape}, disable_ep_error_fallback=True + ) From 1fd50ba5f204e30dc9d8757865acfeaae790d07c Mon Sep 17 00:00:00 2001 From: "yichao.hu" Date: Sun, 18 Oct 2026 17:26:51 -0400 Subject: [PATCH 40/52] test(kernel): fix the may-inplace assertions and enable the trace env The reuse case compared a (3,2) output with a (2,3) expected array, and both trace assertions ran without ORT_MUSA_RESHAPE_TRACE set. --- test/ops/test_reshape_may_inplace.py | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/test/ops/test_reshape_may_inplace.py b/test/ops/test_reshape_may_inplace.py index ee88164..65e201b 100644 --- a/test/ops/test_reshape_may_inplace.py +++ b/test/ops/test_reshape_may_inplace.py @@ -37,8 +37,11 @@ def _musa_session(model: bytes) -> ort.InferenceSession: return ort.InferenceSession(model, sess_options=options) -def test_reshape_may_inplace_reuses_single_consumer_device_output(capfd): +def test_reshape_may_inplace_reuses_single_consumer_device_output( + capfd, monkeypatch +): """Add's device output feeds only Reshape: the buffer can be reused.""" + monkeypatch.setenv("ORT_MUSA_RESHAPE_TRACE", "1") rng = np.random.default_rng(7) x = rng.standard_normal((2, 3)).astype(np.float32) shape = np.array([3, 2], dtype=np.int64) @@ -58,12 +61,15 @@ def test_reshape_may_inplace_reuses_single_consumer_device_output(capfd): session = _musa_session(model) out = session.run(None, {"X": x, "shape": shape})[0] err = capfd.readouterr().err - np.testing.assert_allclose(out, x + 1.0, rtol=0, atol=0) + np.testing.assert_allclose(out, (x + 1.0).reshape(3, 2), rtol=0, atol=0) assert "MUSA_RESHAPE inplace=1" in err, err -def test_reshape_may_inplace_copies_with_another_consumer(capfd): +def test_reshape_may_inplace_copies_with_another_consumer( + capfd, monkeypatch +): """Add's output has a second consumer, so Reshape must copy.""" + monkeypatch.setenv("ORT_MUSA_RESHAPE_TRACE", "1") rng = np.random.default_rng(8) x = rng.standard_normal((2, 3)).astype(np.float32) shape = np.array([3, 2], dtype=np.int64) From 89171ea2ca10a4ba06f6e31c53899a4f03066b86 Mon Sep 17 00:00:00 2001 From: "yichao.hu" Date: Sun, 18 Oct 2026 17:36:57 -0400 Subject: [PATCH 41/52] perf(kernel): restore Reshape's hard input/output alias c4 removed the alias before c5 existed. With MemcpyFromHost registered, ORT stages a CPU-produced input onto the device before the kernel runs, so Reshape only sees device tensors and the alias is safe again; Compute() skips the copy when src == dst and keeps the copy on the defensive different-pointer path. The may-inplace route showed no allocator reuse on the measured models. --- musa/ep/src/kernels/tensor/reshape.cc | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/musa/ep/src/kernels/tensor/reshape.cc b/musa/ep/src/kernels/tensor/reshape.cc index e202a76..5fc3108 100644 --- a/musa/ep/src/kernels/tensor/reshape.cc +++ b/musa/ep/src/kernels/tensor/reshape.cc @@ -122,10 +122,11 @@ ONNX_OPERATOR_VERSIONED_KERNEL_EX( (Ort::KernelDefBuilder() .AddTypeConstraint("T", AllFixedSizeTensorTypesNoBFloat16()) .SetInputMemType(1, OrtMemTypeCPUInput) - // May-inplace: reuse the input buffer for the output whenever the - // allocator proves it safe. Compute() skips the copy when src == dst - // and copies otherwise, so a staged input keeps c4's safe path. - .AddInputOutputMutableAlias(0, 0)), + // Hard alias: c5 registers MemcpyFromHost, so ORT stages any + // CPU-produced input before this kernel runs and Reshape only sees + // device tensors (c4 predates c5). Compute() skips the copy when + // src == dst and still copies on the defensive different-pointer path. + .AddInputOutputAlias(0, 0)), Reshape) // The opset-21 schema only adds the int4/uint4 element types, which this @@ -136,7 +137,7 @@ ONNX_OPERATOR_VERSIONED_KERNEL_EX( (Ort::KernelDefBuilder() .AddTypeConstraint("T", AllFixedSizeTensorTypes()) .SetInputMemType(1, OrtMemTypeCPUInput) - .AddInputOutputMutableAlias(0, 0)), + .AddInputOutputAlias(0, 0)), Reshape) ONNX_OPERATOR_VERSIONED_KERNEL_EX( From 1622d22071f527919752ffbe4c8ad82af123c5f6 Mon Sep 17 00:00:00 2001 From: "yichao.hu" Date: Sun, 18 Oct 2026 17:52:49 -0400 Subject: [PATCH 42/52] test(kernel): keep the may-inplace reuse output internal A graph output does not share a buffer with its input, so the Reshape output must feed another node for the alias to be observable. --- test/ops/test_reshape_may_inplace.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/test/ops/test_reshape_may_inplace.py b/test/ops/test_reshape_may_inplace.py index 65e201b..31f9f75 100644 --- a/test/ops/test_reshape_may_inplace.py +++ b/test/ops/test_reshape_may_inplace.py @@ -49,7 +49,10 @@ def test_reshape_may_inplace_reuses_single_consumer_device_output( model = build_graph_model( [ helper.make_node("Add", ["X", "one"], ["added"]), - helper.make_node("Reshape", ["added", "shape"], ["Y"]), + helper.make_node("Reshape", ["added", "shape"], ["reshaped"]), + # Keep the Reshape output internal: a graph output does not share + # a buffer with the input, so the alias cannot be observed there. + helper.make_node("Relu", ["reshaped"], ["Y"]), ], {"X": x, "shape": shape}, [("Y", TensorProto.FLOAT)], @@ -61,7 +64,9 @@ def test_reshape_may_inplace_reuses_single_consumer_device_output( session = _musa_session(model) out = session.run(None, {"X": x, "shape": shape})[0] err = capfd.readouterr().err - np.testing.assert_allclose(out, (x + 1.0).reshape(3, 2), rtol=0, atol=0) + np.testing.assert_allclose( + out, np.maximum((x + 1.0).reshape(3, 2), 0.0), rtol=0, atol=0 + ) assert "MUSA_RESHAPE inplace=1" in err, err From b9abcec9ab05e67c66accc18e9c9346c8590d202 Mon Sep 17 00:00:00 2001 From: "yichao.hu" Date: Mon, 19 Oct 2026 09:30:00 -0400 Subject: [PATCH 43/52] perf(fusion): cache shared Tile repeats expressions within each execution standalone cache, cherry-pickable onto ba650d0. Cache successful expression results by repeats input index before source-shape construction; preserve tensor fallback on evaluation failure and reset all cached values on each Compute. The isolated-index apply onto ba650d0 passed with no partner dependency. Stage A static checks pass; build, runtime and performance validation remain pending. --- musa/ep/src/fusion/tile_concat_fusion.cc | 32 ++++++++++++++++++------ 1 file changed, 24 insertions(+), 8 deletions(-) diff --git a/musa/ep/src/fusion/tile_concat_fusion.cc b/musa/ep/src/fusion/tile_concat_fusion.cc index 48519e7..a597987 100644 --- a/musa/ep/src/fusion/tile_concat_fusion.cc +++ b/musa/ep/src/fusion/tile_concat_fusion.cc @@ -72,6 +72,10 @@ struct TileConcatFusionCompute : FusionNodeCompute { // feeds 13 Tiles from one tensor): read the tensor back once per // execution instead of once per Tile. std::unordered_map> repeats_cache; + // Each repeats input has one expression and one set of source indices. + // Reuse successful evaluations only within this Compute: source shapes + // can change on the next run, and failures must keep the tensor fallback. + std::unordered_map> expr_repeats_cache; shapes.reserve(inputs.size()); input_data.reserve(inputs.size()); @@ -114,15 +118,27 @@ struct TileConcatFusionCompute : FusionNodeCompute { std::vector repeats; bool recomputed = false; if (!input_plan.repeats_expr.Empty()) { - std::vector> source_shapes; - source_shapes.reserve( - input_plan.repeats_expr.source_input_indices.size()); - for (size_t index : input_plan.repeats_expr.source_input_indices) { - source_shapes.push_back( - ctx.GetInput(index).GetTensorTypeAndShapeInfo().GetShape()); + auto cached = + expr_repeats_cache.find(input_plan.repeats_input_index); + if (cached != expr_repeats_cache.end()) { + repeats = cached->second; + recomputed = true; + } else { + std::vector> source_shapes; + source_shapes.reserve( + input_plan.repeats_expr.source_input_indices.size()); + for (size_t index : + input_plan.repeats_expr.source_input_indices) { + source_shapes.push_back( + ctx.GetInput(index).GetTensorTypeAndShapeInfo().GetShape()); + } + recomputed = musa_ep::EvaluateShapeExpr(input_plan.repeats_expr, + source_shapes, &repeats); + if (recomputed) { + expr_repeats_cache.emplace(input_plan.repeats_input_index, + repeats); + } } - recomputed = musa_ep::EvaluateShapeExpr(input_plan.repeats_expr, - source_shapes, &repeats); } if (!recomputed) { if (input_plan.repeats_input_index == kNoInputIndex) { From ba7ced62f7749f04edaaeb94a8dbe529e88896e4 Mon Sep 17 00:00:00 2001 From: "yichao.hu" Date: Mon, 19 Oct 2026 09:31:30 -0400 Subject: [PATCH 44/52] test(fusion): add opt-in per-execution Tile expression counters ORT_MUSA_TILE_CONCAT_EXPR_TRACE=1 prints actual evaluator calls and distinct repeats inputs once on every Compute exit. Disabled tracing does not count or allocate distinct-input entries. Keep instrumentation separate from the standalone cache commit for clean performance comparisons. Stage A formatting checks pass; runtime validation remains pending. --- musa/ep/src/fusion/tile_concat_fusion.cc | 27 ++++++++++++++++++++++++ 1 file changed, 27 insertions(+) diff --git a/musa/ep/src/fusion/tile_concat_fusion.cc b/musa/ep/src/fusion/tile_concat_fusion.cc index a597987..016d950 100644 --- a/musa/ep/src/fusion/tile_concat_fusion.cc +++ b/musa/ep/src/fusion/tile_concat_fusion.cc @@ -16,6 +16,7 @@ #include #include #include +#include #include #include #include @@ -61,6 +62,31 @@ struct TileConcatFusionCompute : FusionNodeCompute { : axis(axis), output_index(output_index), inputs(std::move(inputs)) {} OrtStatus* Compute(OrtKernelContext* kernel_context) const override { + // The guard reports every exit, including errors and empty outputs. With + // tracing off it neither counts evaluations nor allocates a distinct set. + struct ExprTrace { + bool enabled; + size_t evals = 0; + std::unordered_set distinct; + + void Record(size_t repeats_index) { + if (enabled) { + ++evals; + distinct.insert(repeats_index); + } + } + + ~ExprTrace() { + if (enabled) { + std::fprintf(stderr, + "MUSA_TILE_CONCAT_EXPR_EVALS evals=%zu distinct=%zu\n", + evals, distinct.size()); + } + } + }; + const char* trace_env = std::getenv("ORT_MUSA_TILE_CONCAT_EXPR_TRACE"); + ExprTrace expr_trace{trace_env != nullptr && trace_env[0] == '1' && + trace_env[1] == '\0'}; try { Ort::KernelContext ctx(kernel_context); musaStream_t stream = GetComputeStream(ctx); @@ -132,6 +158,7 @@ struct TileConcatFusionCompute : FusionNodeCompute { source_shapes.push_back( ctx.GetInput(index).GetTensorTypeAndShapeInfo().GetShape()); } + expr_trace.Record(input_plan.repeats_input_index); recomputed = musa_ep::EvaluateShapeExpr(input_plan.repeats_expr, source_shapes, &repeats); if (recomputed) { From 62090a294d717347ceeab883eb07e767cea004b8 Mon Sep 17 00:00:00 2001 From: "yichao.hu" Date: Mon, 19 Oct 2026 10:42:39 -0400 Subject: [PATCH 45/52] test(fusion): cover shared expression caching on the cache-only branch Add two Tiles sharing one shape-derived repeats value to the existing shape-chain suite. Reuse one session for batches 3 and 5, compare exactly with CPU, require fused execution and one evaluator call per run when tracing is enabled, and require silence when disabled. Disable runtime CPU retry for the new MUSA session and preserve every existing test and helper. Stage A diff, C++ formatting and Python syntax checks pass; build and full unit tests remain for the Supervisor. --- test/fusion/test_tile_concat_shape_chain.py | 63 +++++++++++++++++++++ 1 file changed, 63 insertions(+) diff --git a/test/fusion/test_tile_concat_shape_chain.py b/test/fusion/test_tile_concat_shape_chain.py index 1fd207e..200178b 100644 --- a/test/fusion/test_tile_concat_shape_chain.py +++ b/test/fusion/test_tile_concat_shape_chain.py @@ -495,3 +495,66 @@ def test_tile_concat_runtime_repeats_switch_fuses_outside_source( assert any(str(op).startswith("MUSAExecutionProvider_") for op in musa_ops) assert "Tile" not in musa_ops assert "Concat" not in musa_ops + + +@pytest.mark.parametrize("trace_enabled", [False, True]) +def test_tile_concat_repeats_shared_expression_cache( + tmp_path, capfd, monkeypatch, trace_enabled +): + """Two Tiles share one evaluation per run; a new batch recomputes it.""" + if trace_enabled: + monkeypatch.setenv("ORT_MUSA_TILE_CONCAT_EXPR_TRACE", "1") + else: + monkeypatch.delenv("ORT_MUSA_TILE_CONCAT_EXPR_TRACE", raising=False) + monkeypatch.delenv("ORT_MUSA_TILE_CONCAT_RUNTIME_REPEATS", raising=False) + nodes = [ + helper.make_node("Shape", ["query"], ["query_shape"]), + helper.make_node("Gather", ["query_shape", "zero_scalar"], ["batch"], axis=0), + helper.make_node("Reshape", ["batch", "one_shape"], ["batch_1d"]), + helper.make_node("Concat", ["batch_1d", "one"], ["repeats"], axis=0), + helper.make_node("Tile", ["data", "repeats"], ["tiled"]), + helper.make_node("Tile", ["data2", "repeats"], ["tiled2"]), + helper.make_node("Concat", ["tiled", "tiled2", "query"], ["y_pre"], axis=1), + helper.make_node("Relu", ["y_pre"], ["y"]), + ] + initializers = _batch_chain_initializers() + [ + _int64_initializer("one_shape", [1]), + _float_initializer("data2", np.arange(4, dtype=np.float32) + 10, (1, 4)), + ] + model = _build_model( + "tile_concat_shared_repeats_expression_cache", + nodes, + [helper.make_tensor_value_info("query", TensorProto.FLOAT, ["batch", 4])], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, None)], + initializers, + ) + cpu = _session(model, use_musa=False) + musa = _session( + model, use_musa=True, profile_prefix=str(tmp_path / "shared_expression") + ) + musa.disable_fallback() + factory_lines = [ + line + for line in capfd.readouterr().err.splitlines() + if REPEATS_DIAGNOSTIC in line + ] + assert len(factory_lines) == 2 + assert all(RECOMPUTED_DIAGNOSTIC in line for line in factory_lines) + rng = np.random.default_rng(47) + for batch in (3, 5): + feeds = {"query": rng.standard_normal((batch, 4)).astype(np.float32)} + expected = cpu.run(None, feeds)[0] + actual = musa.run(None, feeds)[0] + np.testing.assert_array_equal(actual, expected) + assert actual.shape == (batch, 12) + trace_lines = [ + line + for line in capfd.readouterr().err.splitlines() + if "MUSA_TILE_CONCAT_EXPR_EVALS" in line + ] + expected_trace = "MUSA_TILE_CONCAT_EXPR_EVALS evals=1 distinct=1" + assert trace_lines == ([expected_trace] if trace_enabled else []) + ops = _read_ops_by_provider(musa) + assert _fused_ops(ops) + assert "Tile" not in _musa_ops(ops) + assert "Concat" not in _musa_ops(ops) From e0e69c0ac50ca66c4a5c5b592447827f8bfb6a01 Mon Sep 17 00:00:00 2001 From: Albert Date: Mon, 26 Oct 2026 10:22:43 +0800 Subject: [PATCH 46/52] update doc --- musa/docs/fusion/split_concat.md | 11 ++++++----- musa/docs/fusion/split_unsqueeze_concat.md | 13 ++++++------- musa/docs/fusion/tile_concat.md | 2 +- 3 files changed, 13 insertions(+), 13 deletions(-) diff --git a/musa/docs/fusion/split_concat.md b/musa/docs/fusion/split_concat.md index 3c05028..07e78e7 100644 --- a/musa/docs/fusion/split_concat.md +++ b/musa/docs/fusion/split_concat.md @@ -13,7 +13,7 @@ This file is generated from the current C++ fusion source and matching fusion te - Runtime compute: `SplitConcatFusionCompute` - Runtime implementation: `musa/ep/src/fusion/split_concat_fusion.cc` - `drop_constant_initializers`: `false` -- Before-graph topology source: `musa/ep/src/fusion/split_concat_fusion_matcher.cc` finder source +- Before-graph topology source: `test/fusion/test_split_unsqueeze_concat_fusion.py::_build_model` ## Extracted ONNX Ops @@ -24,12 +24,13 @@ This file is generated from the current C++ fusion source and matching fusion te ```mermaid flowchart LR subgraph Before[Before fusion] - B0["Split"] - B1["Reshape"] - B0 --> B1 + B0["Shape"] + B1["Slice"] B2["Concat"] + B3["Reshape"] + B4["Transpose"] + B0 --> B1 B1 --> B2 - B3["Transpose"] B2 --> B3 end subgraph After[After fusion] diff --git a/musa/docs/fusion/split_unsqueeze_concat.md b/musa/docs/fusion/split_unsqueeze_concat.md index 657010e..a78e4cc 100644 --- a/musa/docs/fusion/split_unsqueeze_concat.md +++ b/musa/docs/fusion/split_unsqueeze_concat.md @@ -13,7 +13,7 @@ This file is generated from the current C++ fusion source and matching fusion te - Runtime compute: `SplitUnsqueezeConcatFusionCompute` - Runtime implementation: `musa/ep/src/fusion/split_unsqueeze_concat_fusion.cc` - `drop_constant_initializers`: `false` -- Before-graph topology source: `musa/ep/src/fusion/split_unsqueeze_concat_fusion_matcher.cc` finder source +- Before-graph topology source: `test/fusion/test_split_unsqueeze_concat_fusion.py::_build_model` ## Extracted ONNX Ops @@ -24,15 +24,14 @@ This file is generated from the current C++ fusion source and matching fusion te ```mermaid flowchart LR subgraph Before[Before fusion] - B0["Split"] - B1["Reshape"] + B0["Shape"] + B1["Slice"] + B2["Concat"] + B3["Reshape"] + B4["Transpose"] B0 --> B1 - B2["Unsqueeze"] B1 --> B2 - B3["Concat"] B2 --> B3 - B4["Transpose"] - B3 --> B4 end subgraph After[After fusion] A0["MUSA fused node"] diff --git a/musa/docs/fusion/tile_concat.md b/musa/docs/fusion/tile_concat.md index 9fd152b..f9eba26 100644 --- a/musa/docs/fusion/tile_concat.md +++ b/musa/docs/fusion/tile_concat.md @@ -17,7 +17,7 @@ This file is generated from the current C++ fusion source and matching fusion te ## Extracted ONNX Ops -`Concat`, `Tile`, `Gather` +`Concat`, `Tile` ## Mermaid From c94ec58f8ef0f3d8eb9dcdce90c1455d3209236b Mon Sep 17 00:00:00 2001 From: Albert Date: Mon, 26 Oct 2026 11:55:15 +0800 Subject: [PATCH 47/52] Feat(CI): add ci workflows --- .github/workflows/musa-ci.yml | 259 ++++++++++++++++++++++++++++++++++ 1 file changed, 259 insertions(+) create mode 100644 .github/workflows/musa-ci.yml diff --git a/.github/workflows/musa-ci.yml b/.github/workflows/musa-ci.yml new file mode 100644 index 0000000..69c58c5 --- /dev/null +++ b/.github/workflows/musa-ci.yml @@ -0,0 +1,259 @@ +name: MUSA CI + +on: + pull_request: + workflow_dispatch: + +permissions: + contents: read + actions: read + +concurrency: + group: musa-ci-${{ github.event.pull_request.number || github.ref }} + cancel-in-progress: true + +jobs: + test: + name: Unit Tests + runs-on: [self-hosted, musa, s5000] + timeout-minutes: 90 + env: + CI_ASSETS: /data/ci-assets + PYTHONUNBUFFERED: "1" + + steps: + - name: Checkout PR head + uses: actions/checkout@v4 + with: + ref: ${{ github.event.pull_request.head.sha || github.sha }} + fetch-depth: 1 + submodules: false + + - name: Configure MUSA toolkit environment + run: | + set -euo pipefail + echo "/usr/local/musa/bin" >> "${GITHUB_PATH}" + { + echo "MUSA_HOME=/usr/local/musa" + echo "LD_LIBRARY_PATH=/usr/local/musa/lib${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}" + echo "CPLUS_INCLUDE_PATH=/usr/include/c++/11:/usr/include/x86_64-linux-gnu/c++/11:/usr/include${CPLUS_INCLUDE_PATH:+:${CPLUS_INCLUDE_PATH}}" + } >> "${GITHUB_ENV}" + + echo "MUSA_HOME=/usr/local/musa" + echo "PATH includes /usr/local/musa/bin in subsequent steps" + echo "LD_LIBRARY_PATH=/usr/local/musa/lib${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}" + echo "CPLUS_INCLUDE_PATH=/usr/include/c++/11:/usr/include/x86_64-linux-gnu/c++/11:/usr/include${CPLUS_INCLUDE_PATH:+:${CPLUS_INCLUDE_PATH}}" + + - name: Prepare ONNX Runtime headers + run: | + set -euo pipefail + test -d "${CI_ASSETS}/onnxruntime" + rm -rf third_party/onnxruntime + cp -a "${CI_ASSETS}/onnxruntime" third_party/onnxruntime + + - name: Build plugin and wheel + run: | + set -euo pipefail + echo "Testing commit: $(git rev-parse HEAD)" + bash build.sh + + - name: Prepare ONNX Runtime SDK + run: | + set -euo pipefail + test -d "${CI_ASSETS}/ort-sdk" + rm -rf build/ort-sdk + cp -a "${CI_ASSETS}/ort-sdk" build/ort-sdk + + - name: Install wheel + run: | + set -euo pipefail + python -m pip install --no-deps "onnxruntime==1.26.0" + + shopt -s nullglob + wheels=(dist/onnxruntime_ep_musa-*.whl) + if (( ${#wheels[@]} != 1 )); then + echo "Expected exactly one MUSA wheel, found ${#wheels[@]}" >&2 + printf ' %s\n' "${wheels[@]}" >&2 + exit 1 + fi + python -m pip install --no-deps --force-reinstall "${wheels[0]}" + + - name: Run provider tests + run: bash test/run_all.sh + + - name: Package build artifacts + run: | + set -euo pipefail + mkdir -p "${RUNNER_TEMP}/musa-build" + tar -C build -czf "${RUNNER_TEMP}/musa-build/musa-build.tar.gz" \ + Release/libonnxruntime_providers_musa_plugin.so \ + ort-sdk + tar -tzf "${RUNNER_TEMP}/musa-build/musa-build.tar.gz" >/dev/null + + - name: Upload build artifacts + uses: actions/upload-artifact@v4 + with: + name: musa-build-${{ github.run_id }}-${{ github.run_attempt }} + path: ${{ runner.temp }}/musa-build/musa-build.tar.gz + if-no-files-found: error + retention-days: 7 + + model-zoo: + name: Model Zoo Inference Performance Test + needs: test + runs-on: [self-hosted, musa, s5000] + timeout-minutes: 720 + env: + CI_ASSETS: /data/ci-assets + MODEL_ZOO_WORK: /data/ci-work/onnx-model-zoo-${{ github.run_id }}-${{ github.run_attempt }} + PERF_RESULTS_DIR: /data/ci-work/onnx-model-zoo-${{ github.run_id }}-${{ github.run_attempt }}/performance-results + BASE_SOURCE_DIR: ${{ github.workspace }}/.ci/base-src + PYTHONUNBUFFERED: "1" + + steps: + - name: Checkout PR head + uses: actions/checkout@v4 + with: + ref: ${{ github.event.pull_request.head.sha || github.sha }} + fetch-depth: 1 + submodules: false + + - name: Configure MUSA toolkit environment + run: | + set -euo pipefail + echo "/usr/local/musa/bin" >> "${GITHUB_PATH}" + { + echo "MUSA_HOME=/usr/local/musa" + echo "LD_LIBRARY_PATH=/usr/local/musa/lib${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}" + echo "CPLUS_INCLUDE_PATH=/usr/include/c++/11:/usr/include/x86_64-linux-gnu/c++/11:/usr/include${CPLUS_INCLUDE_PATH:+:${CPLUS_INCLUDE_PATH}}" + } >> "${GITHUB_ENV}" + + - name: Download build artifacts + uses: actions/download-artifact@v4 + with: + name: musa-build-${{ github.run_id }}-${{ github.run_attempt }} + path: ${{ runner.temp }}/musa-build + + - name: Extract build artifacts + run: | + set -euo pipefail + mkdir -p build + tar -C build -xzf "${RUNNER_TEMP}/musa-build/musa-build.tar.gz" + + plugin="${GITHUB_WORKSPACE}/build/Release/libonnxruntime_providers_musa_plugin.so" + sdk_root="${GITHUB_WORKSPACE}/build/ort-sdk/onnxruntime-linux-x64-1.26.0" + test -f "${plugin}" + test -f "${sdk_root}/include/onnxruntime_c_api.h" + echo "Head commit: $(git rev-parse HEAD)" + sha256sum "${plugin}" + + - name: Checkout PR base + uses: actions/checkout@v4 + with: + ref: ${{ github.event.pull_request.base.sha || github.event.repository.default_branch }} + path: .ci/base-src + fetch-depth: 1 + submodules: false + + - name: Build base plugin + run: | + set -euo pipefail + cd "${BASE_SOURCE_DIR}" + rm -rf third_party/onnxruntime + cp -a "${CI_ASSETS}/onnxruntime" third_party/onnxruntime + + echo "Base commit: $(git rev-parse HEAD)" + bash build.sh --clean --no-wheel + + base_plugin="${BASE_SOURCE_DIR}/build/Release/libonnxruntime_providers_musa_plugin.so" + test -f "${base_plugin}" + sha256sum "${base_plugin}" + + - name: Prepare model zoo + run: | + set -euo pipefail + test -d "${CI_ASSETS}/onnx-model-zoo" + rm -rf "${MODEL_ZOO_WORK}" + mkdir -p "${MODEL_ZOO_WORK}" + cp -al "${CI_ASSETS}/onnx-model-zoo/." "${MODEL_ZOO_WORK}/" + rm -rf "${MODEL_ZOO_WORK}/cpp_runners" + cp -a "${CI_ASSETS}/onnx-model-zoo/cpp_runners" "${MODEL_ZOO_WORK}/cpp_runners" + rm -f "${MODEL_ZOO_WORK}/musa_ep_performance.md" + rm -f "${MODEL_ZOO_WORK}/musa_ep_performance_with_multi_stream.md" + mkdir -p "${PERF_RESULTS_DIR}" + test -f "${MODEL_ZOO_WORK}/JD/real_models/platform/unirank-v2-8layer-combo-v6-online_graph2.onnx" + + - name: Build model zoo runners + run: | + set -euo pipefail + ORT_ROOT="${GITHUB_WORKSPACE}/build/ort-sdk/onnxruntime-linux-x64-1.26.0" \ + bash "${MODEL_ZOO_WORK}/cpp_runners/build_onnx_runner.sh" + + - name: Benchmark base single-stream + run: | + set -euo pipefail + bash scripts/ci/run_model_zoo_benchmark.sh \ + base \ + "${BASE_SOURCE_DIR}/build/Release/libonnxruntime_providers_musa_plugin.so" \ + single \ + "${MODEL_ZOO_WORK}" \ + "${PERF_RESULTS_DIR}" + + - name: Benchmark base multi-stream + run: | + set -euo pipefail + bash scripts/ci/run_model_zoo_benchmark.sh \ + base \ + "${BASE_SOURCE_DIR}/build/Release/libonnxruntime_providers_musa_plugin.so" \ + multi \ + "${MODEL_ZOO_WORK}" \ + "${PERF_RESULTS_DIR}" + + - name: Benchmark head single-stream + run: | + set -euo pipefail + bash scripts/ci/run_model_zoo_benchmark.sh \ + head \ + "${GITHUB_WORKSPACE}/build/Release/libonnxruntime_providers_musa_plugin.so" \ + single \ + "${MODEL_ZOO_WORK}" \ + "${PERF_RESULTS_DIR}" + + - name: Benchmark head multi-stream + run: | + set -euo pipefail + bash scripts/ci/run_model_zoo_benchmark.sh \ + head \ + "${GITHUB_WORKSPACE}/build/Release/libonnxruntime_providers_musa_plugin.so" \ + multi \ + "${MODEL_ZOO_WORK}" \ + "${PERF_RESULTS_DIR}" + + - name: Compare base and head performance + run: | + set -euo pipefail + python3 scripts/ci/compare_model_zoo_performance.py \ + --baseline-single "${PERF_RESULTS_DIR}/base-single.md" \ + --candidate-single "${PERF_RESULTS_DIR}/head-single.md" \ + --baseline-multi "${PERF_RESULTS_DIR}/base-multi.md" \ + --candidate-multi "${PERF_RESULTS_DIR}/head-multi.md" \ + --threshold 10 \ + --output "${PERF_RESULTS_DIR}/model-zoo-performance-comparison.md" + + - name: Upload model zoo reports + if: always() + uses: actions/upload-artifact@v4 + with: + name: model-zoo-reports-${{ github.run_id }}-${{ github.run_attempt }} + path: ${{ env.PERF_RESULTS_DIR }}/ + if-no-files-found: ignore + retention-days: 7 + + - name: Cleanup model zoo work directory + if: always() + run: | + set -euo pipefail + case "${MODEL_ZOO_WORK}" in + /data/ci-work/onnx-model-zoo-*) rm -rf "${MODEL_ZOO_WORK}" ;; + *) echo "Refusing to remove unexpected path: ${MODEL_ZOO_WORK}" >&2; exit 1 ;; + esac From 0d3f66712a079256ef4aac2bc6e34e08e4a4b245 Mon Sep 17 00:00:00 2001 From: Albert Date: Mon, 26 Oct 2026 15:11:50 +0800 Subject: [PATCH 48/52] fix add scripts/ci --- scripts/ci/compare_model_zoo_performance.py | 379 ++++++++++++++++++++ scripts/ci/run_model_zoo_benchmark.sh | 116 ++++++ 2 files changed, 495 insertions(+) create mode 100755 scripts/ci/compare_model_zoo_performance.py create mode 100755 scripts/ci/run_model_zoo_benchmark.sh diff --git a/scripts/ci/compare_model_zoo_performance.py b/scripts/ci/compare_model_zoo_performance.py new file mode 100755 index 0000000..c6edf44 --- /dev/null +++ b/scripts/ci/compare_model_zoo_performance.py @@ -0,0 +1,379 @@ +#!/usr/bin/env python3 +"""Compare PR base and head model-zoo performance reports. + +Single-stream latency and multi-stream QPS are both compared with the same +percentage threshold. A candidate fails when it is slower than the base by +more than the threshold. +""" +from __future__ import annotations + +import argparse +import math +import os +import re +import sys +from dataclasses import dataclass +from pathlib import Path + + +@dataclass(frozen=True) +class ReportRow: + model: str + status: str + value: float | None + + +@dataclass(frozen=True) +class Comparison: + model: str + base_status: str + head_status: str + base_value: float | None + head_value: float | None + change_pct: float | None + result: str + + @property + def failed(self) -> bool: + return self.result in ( + "REGRESSION", + "CANDIDATE_SKIPPED", + "STATUS_FAILED", + "METRIC_MISSING", + ) + + +_NUMBER_RE = re.compile(r"[-+]?(?:\d+(?:\.\d*)?|\.\d+)(?:[eE][-+]?\d+)?") + + +def _split_markdown_row(line: str) -> list[str]: + return [cell.strip() for cell in line.strip().strip("|").split("|")] + + +def _parse_metric(value: str) -> float | None: + lowered = value.lower() + if any(word in lowered for word in ("unknown", "skipped", "failed", "n/a")): + return None + match = _NUMBER_RE.search(value) + if match is None: + return None + parsed = float(match.group()) + if not math.isfinite(parsed) or parsed <= 0: + return None + return parsed + + +def parse_report(path: Path, metric_column: str) -> dict[str, ReportRow]: + lines = path.read_text(encoding="utf-8").splitlines() + header_index = -1 + model_index = status_index = metric_index = -1 + + for index, line in enumerate(lines): + if not line.strip().startswith("|"): + continue + cells = _split_markdown_row(line) + if "Model" not in cells or "Status" not in cells or metric_column not in cells: + continue + header_index = index + model_index = cells.index("Model") + status_index = cells.index("Status") + metric_index = cells.index(metric_column) + break + + if header_index < 0: + raise ValueError(f"{path}: table with {metric_column!r} column not found") + + rows: dict[str, ReportRow] = {} + for line in lines[header_index + 1 :]: + if not line.strip().startswith("|"): + break + cells = _split_markdown_row(line) + if not cells or max(model_index, status_index, metric_index) >= len(cells): + continue + if set(cells[0]) <= set("-: "): + continue + model = cells[model_index] + if not model: + continue + if model in rows: + raise ValueError(f"{path}: duplicate model row: {model}") + rows[model] = ReportRow( + model=model, + status=cells[status_index], + value=_parse_metric(cells[metric_index]), + ) + + if not rows: + raise ValueError(f"{path}: no model rows found") + return rows + + +def compare_reports( + base: dict[str, ReportRow], + head: dict[str, ReportRow], + mode: str, + threshold_pct: float, +) -> list[Comparison]: + missing_from_head = sorted(set(base) - set(head)) + missing_from_base = sorted(set(head) - set(base)) + if missing_from_head or missing_from_base: + details: list[str] = [] + if missing_from_head: + details.append(f"missing from head: {', '.join(missing_from_head)}") + if missing_from_base: + details.append(f"missing from base: {', '.join(missing_from_base)}") + raise ValueError("model sets differ; " + "; ".join(details)) + + comparisons: list[Comparison] = [] + for model in base: + base_row = base[model] + head_row = head[model] + result = "PASS" + change_pct: float | None = None + + if head_row.status not in ("PASS", "SKIPPED"): + result = "STATUS_FAILED" + elif base_row.status == "SKIPPED" and head_row.status == "SKIPPED": + result = "NOT_MEASURED" + elif head_row.status == "SKIPPED": + result = "CANDIDATE_SKIPPED" + elif base_row.status == "SKIPPED": + result = "BASE_SKIPPED" + elif base_row.status != "PASS": + result = "STATUS_FAILED" + elif base_row.value is None or head_row.value is None: + result = "METRIC_MISSING" + else: + change_pct = (head_row.value / base_row.value - 1.0) * 100.0 + slower_pct = change_pct if mode == "single" else -change_pct + if slower_pct > threshold_pct + 1e-9: + result = "REGRESSION" + + comparisons.append( + Comparison( + model=model, + base_status=base_row.status, + head_status=head_row.status, + base_value=base_row.value, + head_value=head_row.value, + change_pct=change_pct, + result=result, + ) + ) + return comparisons + + +def _format_number(value: float | None) -> str: + if value is None: + return "n/a" + return f"{value:.6g}" + + +def _format_change(value: float | None) -> str: + if value is None: + return "n/a" + return f"{value:+.2f}%" + + +def _status_label(row: Comparison, mode: str) -> str: + if row.result == "PASS": + return "PASS" + if row.result == "REGRESSION": + metric = "latency" if mode == "single" else "QPS" + return f"REGRESSION ({metric} {_format_change(row.change_pct)})" + return row.result + + +def _table_for_mode( + title: str, + comparisons: list[Comparison], + mode: str, + base_metric: str, + head_metric: str, + change_metric: str, +) -> list[str]: + lines = [ + f"## {title}", + "", + f"| Model | Base {base_metric} | Head {head_metric} | {change_metric} | Result |", + "|---|---|---|---|---|", + ] + for row in comparisons: + lines.append( + "| " + + " | ".join( + [ + row.model.replace("|", "\\|"), + _format_number(row.base_value), + _format_number(row.head_value), + _format_change(row.change_pct), + _status_label(row, mode), + ] + ) + + " |" + ) + lines.append("") + return lines + + +def build_markdown( + single: list[Comparison], + multi: list[Comparison], + threshold_pct: float, +) -> tuple[str, str]: + failed = [row for row in single + multi if row.failed] + unmeasured = [row for row in single + multi if row.result == "NOT_MEASURED"] + model_count = len(single) + + summary_lines = [ + "## Model Zoo Performance Gate", + "", + f"- Models: {model_count}", + f"- Threshold: more than {threshold_pct:g}% slower than base", + "- Single stream metric: average latency (lower is better)", + "- Multi stream metric: QPS (higher is better)", + f"- Result: {'FAIL' if failed else 'PASS'}", + ] + if failed: + summary_lines.extend(["", "### Regressions", ""]) + for row in failed: + summary_lines.append( + f"- {row.model}: {row.result}, change {_format_change(row.change_pct)}" + ) + if unmeasured: + summary_lines.extend( + ["", f"- Skipped in both versions (not measured): {len(unmeasured)}"] + ) + summary = "\n".join(summary_lines) + "\n" + + lines = [ + "# Model Zoo Performance Comparison", + "", + "| Key | Value |", + "|---|---|", + f"| Threshold | {threshold_pct:g}% |", + "| Single stream metric | Average latency, lower is better |", + "| Multi stream metric | QPS, higher is better |", + f"| Models compared | {model_count} |", + f"| Regressions | {len(failed)} |", + f"| Result | {'FAIL' if failed else 'PASS'} |", + "", + ] + lines.extend( + _table_for_mode( + "Single Stream", + single, + "single", + "latency (ms)", + "latency (ms)", + "Latency change", + ) + ) + lines.extend( + _table_for_mode( + "Multi Stream (-s 4 -q 128)", + multi, + "multi", + "QPS", + "QPS", + "QPS change", + ) + ) + lines.extend(["## Result", ""]) + if failed: + lines.append( + f"FAIL: {len(failed)} model measurement(s) regressed by more than " + f"{threshold_pct:g}% or could not be measured." + ) + else: + lines.append( + f"PASS: no model regressed by more than {threshold_pct:g}% " + "in single-stream latency or multi-stream QPS." + ) + lines.append("") + return "\n".join(lines) + "\n", summary + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--baseline-single", required=True, type=Path) + parser.add_argument("--candidate-single", required=True, type=Path) + parser.add_argument("--baseline-multi", required=True, type=Path) + parser.add_argument("--candidate-multi", required=True, type=Path) + parser.add_argument( + "--threshold", + type=float, + default=10.0, + help="Maximum allowed slowdown in percent (default: 10).", + ) + parser.add_argument("--output", type=Path, default=None) + parser.add_argument( + "--summary", + type=Path, + default=Path(os.environ["GITHUB_STEP_SUMMARY"]) + if os.environ.get("GITHUB_STEP_SUMMARY") + else None, + help="Optional markdown summary file, appended if it exists.", + ) + args = parser.parse_args() + + if not math.isfinite(args.threshold) or args.threshold < 0: + parser.error("--threshold must be a non-negative number") + + try: + single = compare_reports( + parse_report(args.baseline_single, "Average Latency"), + parse_report(args.candidate_single, "Average Latency"), + "single", + args.threshold, + ) + multi = compare_reports( + parse_report(args.baseline_multi, "QPS"), + parse_report(args.candidate_multi, "QPS"), + "multi", + args.threshold, + ) + except (OSError, ValueError) as exc: + print(f"Failed to compare reports: {exc}", file=sys.stderr) + return 1 + + single_models = {row.model for row in single} + multi_models = {row.model for row in multi} + if single_models != multi_models: + missing_from_multi = sorted(single_models - multi_models) + missing_from_single = sorted(multi_models - single_models) + print( + "Single-stream and multi-stream model sets differ; " + f"single-only={missing_from_multi}, multi-only={missing_from_single}", + file=sys.stderr, + ) + return 1 + + markdown, summary = build_markdown(single, multi, args.threshold) + if args.output is not None: + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(markdown, encoding="utf-8") + print(f"Saved performance comparison to {args.output}") + else: + print(markdown, end="") + + if args.summary is not None: + with args.summary.open("a", encoding="utf-8") as stream: + stream.write(summary) + + failed = [row for row in single + multi if row.failed] + if failed: + print(f"Performance gate failed: {len(failed)} regression(s)", file=sys.stderr) + for row in failed: + print( + f" {row.model}: {row.result}, change {_format_change(row.change_pct)}", + file=sys.stderr, + ) + return 1 + + print("Performance gate passed") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/scripts/ci/run_model_zoo_benchmark.sh b/scripts/ci/run_model_zoo_benchmark.sh new file mode 100755 index 0000000..b210b23 --- /dev/null +++ b/scripts/ci/run_model_zoo_benchmark.sh @@ -0,0 +1,116 @@ +#!/usr/bin/env bash +set -euo pipefail + +usage() { + cat <<'EOF' +Usage: run_model_zoo_benchmark.sh LABEL PLUGIN MODE ZOO_DIR RESULTS_DIR + +Run one model-zoo performance pass and preserve its report and error logs. + +Arguments: + LABEL Version label used in result file names, for example base or head. + PLUGIN Absolute path to libonnxruntime_providers_musa_plugin.so. + MODE single or multi. + ZOO_DIR Prepared onnx-model-zoo working directory. + RESULTS_DIR Directory that receives the report and error logs. + +Modes: + single run_all_onnx.sh --ep musa + multi run_all_onnx.sh --ep musa -s 4 -q 128 +EOF +} + +if [[ $# -ne 5 ]]; then + usage >&2 + exit 2 +fi + +LABEL="$1" +PLUGIN="$2" +MODE="$3" +ZOO_DIR="$4" +RESULTS_DIR="$5" + +if [[ ! "$LABEL" =~ ^[A-Za-z0-9_.-]+$ ]]; then + echo "Invalid label: $LABEL" >&2 + exit 2 +fi +if [[ ! -f "$PLUGIN" ]]; then + echo "MUSA plugin does not exist: $PLUGIN" >&2 + exit 1 +fi +if [[ ! -d "$ZOO_DIR" ]]; then + echo "Model zoo directory does not exist: $ZOO_DIR" >&2 + exit 1 +fi + +case "$MODE" in + single) + REPORT_NAME="musa_ep_performance.md" + RUN_ARGS=(--ep musa) + ;; + multi) + REPORT_NAME="musa_ep_performance_with_multi_stream.md" + RUN_ARGS=(--ep musa -s 4 -q 128) + ;; + *) + echo "Mode must be single or multi: $MODE" >&2 + exit 2 + ;; +esac + +mkdir -p "$RESULTS_DIR" +export MUSA_EP_LIBRARY="$PLUGIN" +export MUSA_HOME="${MUSA_HOME:-/usr/local/musa}" +export MUSA_VISIBLE_DEVICES="${MUSA_VISIBLE_DEVICES:-0}" +export LD_LIBRARY_PATH="$MUSA_HOME/lib:$MUSA_HOME/lib64:${LD_LIBRARY_PATH:-}" + +echo "benchmark_version=${LABEL}" +echo "benchmark_mode=${MODE}" +echo "benchmark_plugin=${MUSA_EP_LIBRARY}" +echo "MUSA_VISIBLE_DEVICES=${MUSA_VISIBLE_DEVICES}" +sha256sum "$MUSA_EP_LIBRARY" + +cd "$ZOO_DIR" +set +e +bash run_all_onnx.sh "${RUN_ARGS[@]}" +run_status=$? +set -e + +report_path="$ZOO_DIR/$REPORT_NAME" +result_path="$RESULTS_DIR/$LABEL-$MODE.md" +if [[ -f "$report_path" ]]; then + cp "$report_path" "$result_path" + echo "Saved benchmark report to $result_path" +else + echo "Benchmark report was not generated: $report_path" >&2 + exit 1 +fi + +if [[ -d "$ZOO_DIR/error_logs" ]]; then + error_results="$RESULTS_DIR/$LABEL-$MODE-error-logs" + rm -rf "$error_results" + cp -a "$ZOO_DIR/error_logs" "$error_results" + echo "Saved benchmark error logs to $error_results" +fi + +if [[ "$run_status" -ne 0 ]]; then + echo "Benchmark failed with exit code $run_status" >&2 + exit "$run_status" +fi + +if ! grep -Fq '| Fail | 0 |' "$result_path"; then + echo "Benchmark reported one or more failed models: $result_path" >&2 + grep -F '| Fail |' "$result_path" >&2 || true + exit 1 +fi + +if [[ "$MODE" == "single" ]]; then + grep -Fq '| Mode | classic |' "$result_path" +else + grep -Fq '| Querys | 128 |' "$result_path" + grep -Fq '| Streams | 4 |' "$result_path" +fi + +grep -F '| Model count |' "$result_path" +echo "Benchmark completed: version=${LABEL} mode=${MODE}" From 72b8de3b15abfb9c324f91e9797bf70691801fcf Mon Sep 17 00:00:00 2001 From: Albert Date: Mon, 26 Oct 2026 17:42:39 +0800 Subject: [PATCH 49/52] fix --- scripts/ci/compare_model_zoo_performance.py | 94 ++++++++++++++++++--- 1 file changed, 82 insertions(+), 12 deletions(-) diff --git a/scripts/ci/compare_model_zoo_performance.py b/scripts/ci/compare_model_zoo_performance.py index c6edf44..de2c72d 100755 --- a/scripts/ci/compare_model_zoo_performance.py +++ b/scripts/ci/compare_model_zoo_performance.py @@ -175,6 +175,35 @@ def _format_change(value: float | None) -> str: return f"{value:+.2f}%" +def _mode_failure_label(mode: str) -> str: + return "Single-stream latency" if mode == "single" else "Multi-stream QPS" + + +def _failure_detail(row: Comparison, mode: str) -> str: + if row.result == "REGRESSION" and row.change_pct is not None: + if mode == "single": + return f"latency increased by {row.change_pct:.2f}%" + return f"QPS decreased by {abs(row.change_pct):.2f}%" + if row.change_pct is None: + return row.result + return f"{row.result}, change {_format_change(row.change_pct)}" + + +def _failure_lines(rows: list[Comparison], mode: str) -> list[str]: + return [f"- {row.model}: {_failure_detail(row, mode)}" for row in rows] + + +def _failure_counts(single: list[Comparison], multi: list[Comparison]) -> list[str]: + counts: list[str] = [] + single_failed = sum(row.failed for row in single) + multi_failed = sum(row.failed for row in multi) + if single_failed: + counts.append(f"{single_failed} single-stream latency") + if multi_failed: + counts.append(f"{multi_failed} multi-stream QPS") + return counts + + def _status_label(row: Comparison, mode: str) -> str: if row.result == "PASS": return "PASS" @@ -221,7 +250,9 @@ def build_markdown( multi: list[Comparison], threshold_pct: float, ) -> tuple[str, str]: - failed = [row for row in single + multi if row.failed] + single_failed = [row for row in single if row.failed] + multi_failed = [row for row in multi if row.failed] + failed = single_failed + multi_failed unmeasured = [row for row in single + multi if row.result == "NOT_MEASURED"] model_count = len(single) @@ -232,13 +263,28 @@ def build_markdown( f"- Threshold: more than {threshold_pct:g}% slower than base", "- Single stream metric: average latency (lower is better)", "- Multi stream metric: QPS (higher is better)", + f"- Single-stream latency failures: {len(single_failed)}", + f"- Multi-stream QPS failures: {len(multi_failed)}", f"- Result: {'FAIL' if failed else 'PASS'}", ] if failed: - summary_lines.extend(["", "### Regressions", ""]) - for row in failed: - summary_lines.append( - f"- {row.model}: {row.result}, change {_format_change(row.change_pct)}" + if single_failed: + summary_lines.extend( + [ + "", + f"### Single-stream latency failures ({len(single_failed)})", + "", + *_failure_lines(single_failed, "single"), + ] + ) + if multi_failed: + summary_lines.extend( + [ + "", + f"### Multi-stream QPS failures ({len(multi_failed)})", + "", + *_failure_lines(multi_failed, "multi"), + ] ) if unmeasured: summary_lines.extend( @@ -255,7 +301,8 @@ def build_markdown( "| Single stream metric | Average latency, lower is better |", "| Multi stream metric | QPS, higher is better |", f"| Models compared | {model_count} |", - f"| Regressions | {len(failed)} |", + f"| Single-stream latency failures | {len(single_failed)} |", + f"| Multi-stream QPS failures | {len(multi_failed)} |", f"| Result | {'FAIL' if failed else 'PASS'} |", "", ] @@ -282,8 +329,9 @@ def build_markdown( lines.extend(["## Result", ""]) if failed: lines.append( - f"FAIL: {len(failed)} model measurement(s) regressed by more than " - f"{threshold_pct:g}% or could not be measured." + f"FAIL: {len(failed)} failed measurement(s): " + + ", ".join(_failure_counts(single, multi)) + + "." ) else: lines.append( @@ -361,14 +409,36 @@ def main() -> int: with args.summary.open("a", encoding="utf-8") as stream: stream.write(summary) - failed = [row for row in single + multi if row.failed] + single_failed = [row for row in single if row.failed] + multi_failed = [row for row in multi if row.failed] + failed = single_failed + multi_failed if failed: - print(f"Performance gate failed: {len(failed)} regression(s)", file=sys.stderr) - for row in failed: + print( + f"Performance gate failed: {len(failed)} model measurement(s) failed", + file=sys.stderr, + ) + if single_failed: + print( + f" Single-stream latency: {len(single_failed)} model(s) failed", + file=sys.stderr, + ) + if multi_failed: + print( + f" Multi-stream QPS: {len(multi_failed)} model(s) failed", + file=sys.stderr, + ) + for mode, rows in (("single", single_failed), ("multi", multi_failed)): + if not rows: + continue print( - f" {row.model}: {row.result}, change {_format_change(row.change_pct)}", + f"\n{_mode_failure_label(mode)} failures ({len(rows)}):", file=sys.stderr, ) + for row in rows: + print( + f" - {row.model}: {_failure_detail(row, mode)}", + file=sys.stderr, + ) return 1 print("Performance gate passed") From bb043119b869f5ce59e9dad64a727800678545ee Mon Sep 17 00:00:00 2001 From: Albert Date: Mon, 26 Oct 2026 23:29:57 +0800 Subject: [PATCH 50/52] fix --- scripts/ci/compare_model_zoo_performance.py | 2 +- scripts/ci/run_model_zoo_benchmark.sh | 12 +++++++----- 2 files changed, 8 insertions(+), 6 deletions(-) diff --git a/scripts/ci/compare_model_zoo_performance.py b/scripts/ci/compare_model_zoo_performance.py index de2c72d..a0f7310 100755 --- a/scripts/ci/compare_model_zoo_performance.py +++ b/scripts/ci/compare_model_zoo_performance.py @@ -318,7 +318,7 @@ def build_markdown( ) lines.extend( _table_for_mode( - "Multi Stream (-s 4 -q 128)", + "Multi Stream (-s 4 -q 1024)", multi, "multi", "QPS", diff --git a/scripts/ci/run_model_zoo_benchmark.sh b/scripts/ci/run_model_zoo_benchmark.sh index b210b23..abcfd3b 100755 --- a/scripts/ci/run_model_zoo_benchmark.sh +++ b/scripts/ci/run_model_zoo_benchmark.sh @@ -15,8 +15,8 @@ Arguments: RESULTS_DIR Directory that receives the report and error logs. Modes: - single run_all_onnx.sh --ep musa - multi run_all_onnx.sh --ep musa -s 4 -q 128 + single run_all_onnx.sh --ep musa --warmup 20 --repeat 100 + multi run_all_onnx.sh --ep musa -s 4 -q 1024 EOF } @@ -47,11 +47,11 @@ fi case "$MODE" in single) REPORT_NAME="musa_ep_performance.md" - RUN_ARGS=(--ep musa) + RUN_ARGS=(--ep musa --warmup 20 --repeat 100) ;; multi) REPORT_NAME="musa_ep_performance_with_multi_stream.md" - RUN_ARGS=(--ep musa -s 4 -q 128) + RUN_ARGS=(--ep musa -s 4 -q 1024) ;; *) echo "Mode must be single or multi: $MODE" >&2 @@ -107,8 +107,10 @@ fi if [[ "$MODE" == "single" ]]; then grep -Fq '| Mode | classic |' "$result_path" + grep -Fq '| Warmup | 20 |' "$result_path" + grep -Fq '| Repeat | 100 |' "$result_path" else - grep -Fq '| Querys | 128 |' "$result_path" + grep -Fq '| Querys | 1024 |' "$result_path" grep -Fq '| Streams | 4 |' "$result_path" fi From f35d4f7e1124c00a8f27952505c726b84f82ad54 Mon Sep 17 00:00:00 2001 From: Albert Date: Tue, 27 Oct 2026 14:50:53 +0800 Subject: [PATCH 51/52] fix --- .github/workflows/musa-ci.yml | 160 +++---- scripts/ci/compare_model_zoo_performance.py | 449 -------------------- scripts/ci/run_model_zoo_benchmark.sh | 118 ----- 3 files changed, 67 insertions(+), 660 deletions(-) delete mode 100755 scripts/ci/compare_model_zoo_performance.py delete mode 100755 scripts/ci/run_model_zoo_benchmark.sh diff --git a/.github/workflows/musa-ci.yml b/.github/workflows/musa-ci.yml index 69c58c5..cda46e5 100644 --- a/.github/workflows/musa-ci.yml +++ b/.github/workflows/musa-ci.yml @@ -57,13 +57,6 @@ jobs: echo "Testing commit: $(git rev-parse HEAD)" bash build.sh - - name: Prepare ONNX Runtime SDK - run: | - set -euo pipefail - test -d "${CI_ASSETS}/ort-sdk" - rm -rf build/ort-sdk - cp -a "${CI_ASSETS}/ort-sdk" build/ort-sdk - - name: Install wheel run: | set -euo pipefail @@ -86,8 +79,7 @@ jobs: set -euo pipefail mkdir -p "${RUNNER_TEMP}/musa-build" tar -C build -czf "${RUNNER_TEMP}/musa-build/musa-build.tar.gz" \ - Release/libonnxruntime_providers_musa_plugin.so \ - ort-sdk + Release/libonnxruntime_providers_musa_plugin.so tar -tzf "${RUNNER_TEMP}/musa-build/musa-build.tar.gz" >/dev/null - name: Upload build artifacts @@ -98,16 +90,16 @@ jobs: if-no-files-found: error retention-days: 7 - model-zoo: - name: Model Zoo Inference Performance Test + model-zoo-cpu-ep-diff: + name: Model Zoo CPU EP Diff needs: test runs-on: [self-hosted, musa, s5000] timeout-minutes: 720 env: CI_ASSETS: /data/ci-assets - MODEL_ZOO_WORK: /data/ci-work/onnx-model-zoo-${{ github.run_id }}-${{ github.run_attempt }} - PERF_RESULTS_DIR: /data/ci-work/onnx-model-zoo-${{ github.run_id }}-${{ github.run_attempt }}/performance-results - BASE_SOURCE_DIR: ${{ github.workspace }}/.ci/base-src + MODEL_ZOO_WORK: /data/ci-work/onnx-model-zoo-cpu-ep-diff-${{ github.run_id }}-${{ github.run_attempt }} + RESULTS_DIR: /data/ci-work/onnx-model-zoo-cpu-ep-diff-${{ github.run_id }}-${{ github.run_attempt }}/results + MUSA_VISIBLE_DEVICES: "0" PYTHONUNBUFFERED: "1" steps: @@ -141,34 +133,10 @@ jobs: tar -C build -xzf "${RUNNER_TEMP}/musa-build/musa-build.tar.gz" plugin="${GITHUB_WORKSPACE}/build/Release/libonnxruntime_providers_musa_plugin.so" - sdk_root="${GITHUB_WORKSPACE}/build/ort-sdk/onnxruntime-linux-x64-1.26.0" test -f "${plugin}" - test -f "${sdk_root}/include/onnxruntime_c_api.h" echo "Head commit: $(git rev-parse HEAD)" sha256sum "${plugin}" - - name: Checkout PR base - uses: actions/checkout@v4 - with: - ref: ${{ github.event.pull_request.base.sha || github.event.repository.default_branch }} - path: .ci/base-src - fetch-depth: 1 - submodules: false - - - name: Build base plugin - run: | - set -euo pipefail - cd "${BASE_SOURCE_DIR}" - rm -rf third_party/onnxruntime - cp -a "${CI_ASSETS}/onnxruntime" third_party/onnxruntime - - echo "Base commit: $(git rev-parse HEAD)" - bash build.sh --clean --no-wheel - - base_plugin="${BASE_SOURCE_DIR}/build/Release/libonnxruntime_providers_musa_plugin.so" - test -f "${base_plugin}" - sha256sum "${base_plugin}" - - name: Prepare model zoo run: | set -euo pipefail @@ -178,74 +146,80 @@ jobs: cp -al "${CI_ASSETS}/onnx-model-zoo/." "${MODEL_ZOO_WORK}/" rm -rf "${MODEL_ZOO_WORK}/cpp_runners" cp -a "${CI_ASSETS}/onnx-model-zoo/cpp_runners" "${MODEL_ZOO_WORK}/cpp_runners" - rm -f "${MODEL_ZOO_WORK}/musa_ep_performance.md" - rm -f "${MODEL_ZOO_WORK}/musa_ep_performance_with_multi_stream.md" - mkdir -p "${PERF_RESULTS_DIR}" - test -f "${MODEL_ZOO_WORK}/JD/real_models/platform/unirank-v2-8layer-combo-v6-online_graph2.onnx" + mkdir -p "${RESULTS_DIR}" + test -f "${MODEL_ZOO_WORK}/run_all_onnx_cpu_ep_diff.sh" + test -f "${MODEL_ZOO_WORK}/cpp_runners/build_onnx_runner.sh" + test -f "/data/ci-assets/ort-sdk/onnxruntime-linux-x64-1.26.0/include/onnxruntime_c_api.h" - - name: Build model zoo runners + - name: Build model zoo CPU EP diff runner run: | set -euo pipefail - ORT_ROOT="${GITHUB_WORKSPACE}/build/ort-sdk/onnxruntime-linux-x64-1.26.0" \ - bash "${MODEL_ZOO_WORK}/cpp_runners/build_onnx_runner.sh" + cd "${MODEL_ZOO_WORK}" + ORT_ROOT=/data/ci-assets/ort-sdk/onnxruntime-linux-x64-1.26.0/ \ + bash cpp_runners/build_onnx_runner.sh + test -x cpp_runners/onnx_cpu_ep_diff - - name: Benchmark base single-stream + - name: Run MUSA-vs-CPU diff run: | set -euo pipefail - bash scripts/ci/run_model_zoo_benchmark.sh \ - base \ - "${BASE_SOURCE_DIR}/build/Release/libonnxruntime_providers_musa_plugin.so" \ - single \ - "${MODEL_ZOO_WORK}" \ - "${PERF_RESULTS_DIR}" + plugin="${GITHUB_WORKSPACE}/build/Release/libonnxruntime_providers_musa_plugin.so" + result_log="${RESULTS_DIR}/musa-cpu-ep-diff.log" + test -f "${plugin}" - - name: Benchmark base multi-stream - run: | - set -euo pipefail - bash scripts/ci/run_model_zoo_benchmark.sh \ - base \ - "${BASE_SOURCE_DIR}/build/Release/libonnxruntime_providers_musa_plugin.so" \ - multi \ - "${MODEL_ZOO_WORK}" \ - "${PERF_RESULTS_DIR}" + cd "${MODEL_ZOO_WORK}" + set +e + MUSA_EP_LIBRARY="${plugin}" \ + bash run_all_onnx_cpu_ep_diff.sh --ep musa -s 4 -q 16 2>&1 | tee "${result_log}" + pipeline_status=("${PIPESTATUS[@]}") + run_status=${pipeline_status[0]} + tee_status=${pipeline_status[1]} + set -e + + if [[ -d "${MODEL_ZOO_WORK}/musa_cpu_diff_logs" ]]; then + rm -rf "${RESULTS_DIR}/musa_cpu_diff_logs" + cp -a "${MODEL_ZOO_WORK}/musa_cpu_diff_logs" "${RESULTS_DIR}/musa_cpu_diff_logs" + fi - - name: Benchmark head single-stream - run: | - set -euo pipefail - bash scripts/ci/run_model_zoo_benchmark.sh \ - head \ - "${GITHUB_WORKSPACE}/build/Release/libonnxruntime_providers_musa_plugin.so" \ - single \ - "${MODEL_ZOO_WORK}" \ - "${PERF_RESULTS_DIR}" + summary="$(grep -E '^pass=[0-9]+ fail=[0-9]+ skip=[0-9]+ logs=' "${result_log}" | tail -n 1 || true)" + summary_pattern='^pass=([0-9]+) fail=([0-9]+) skip=([0-9]+) logs=' + if [[ ! "${summary}" =~ ${summary_pattern} ]]; then + echo "Could not find the final CPU EP diff summary in ${result_log}" >&2 + echo "Diff runner exit code: ${run_status}" >&2 + exit 1 + fi - - name: Benchmark head multi-stream - run: | - set -euo pipefail - bash scripts/ci/run_model_zoo_benchmark.sh \ - head \ - "${GITHUB_WORKSPACE}/build/Release/libonnxruntime_providers_musa_plugin.so" \ - multi \ - "${MODEL_ZOO_WORK}" \ - "${PERF_RESULTS_DIR}" + pass_count="${BASH_REMATCH[1]}" + fail_count="${BASH_REMATCH[2]}" + skip_count="${BASH_REMATCH[3]}" - - name: Compare base and head performance - run: | - set -euo pipefail - python3 scripts/ci/compare_model_zoo_performance.py \ - --baseline-single "${PERF_RESULTS_DIR}/base-single.md" \ - --candidate-single "${PERF_RESULTS_DIR}/head-single.md" \ - --baseline-multi "${PERF_RESULTS_DIR}/base-multi.md" \ - --candidate-multi "${PERF_RESULTS_DIR}/head-multi.md" \ - --threshold 10 \ - --output "${PERF_RESULTS_DIR}/model-zoo-performance-comparison.md" + { + echo "## Model Zoo CPU EP Diff" + echo '```text' + echo "${summary}" + echo '```' + } >> "${GITHUB_STEP_SUMMARY}" + + echo "CPU EP diff result: pass=${pass_count} fail=${fail_count} skip=${skip_count}" + if (( tee_status != 0 )); then + echo "Failed to write the CPU EP diff log (tee exit ${tee_status})" >&2 + exit 1 + fi + if (( fail_count != 0 )); then + echo "CPU EP diff gate failed: fail=${fail_count}" >&2 + exit 1 + fi + if (( run_status != 0 )); then + echo "CPU EP diff runner exited with code ${run_status} despite fail=0" >&2 + exit "${run_status}" + fi + echo "CPU EP diff gate passed: fail=0 (PASS and SKIPPED are both accepted)" - - name: Upload model zoo reports + - name: Upload CPU EP diff results if: always() uses: actions/upload-artifact@v4 with: - name: model-zoo-reports-${{ github.run_id }}-${{ github.run_attempt }} - path: ${{ env.PERF_RESULTS_DIR }}/ + name: model-zoo-cpu-ep-diff-${{ github.run_id }}-${{ github.run_attempt }} + path: ${{ env.RESULTS_DIR }}/ if-no-files-found: ignore retention-days: 7 @@ -254,6 +228,6 @@ jobs: run: | set -euo pipefail case "${MODEL_ZOO_WORK}" in - /data/ci-work/onnx-model-zoo-*) rm -rf "${MODEL_ZOO_WORK}" ;; + /data/ci-work/onnx-model-zoo-cpu-ep-diff-*) rm -rf "${MODEL_ZOO_WORK}" ;; *) echo "Refusing to remove unexpected path: ${MODEL_ZOO_WORK}" >&2; exit 1 ;; esac diff --git a/scripts/ci/compare_model_zoo_performance.py b/scripts/ci/compare_model_zoo_performance.py deleted file mode 100755 index a0f7310..0000000 --- a/scripts/ci/compare_model_zoo_performance.py +++ /dev/null @@ -1,449 +0,0 @@ -#!/usr/bin/env python3 -"""Compare PR base and head model-zoo performance reports. - -Single-stream latency and multi-stream QPS are both compared with the same -percentage threshold. A candidate fails when it is slower than the base by -more than the threshold. -""" -from __future__ import annotations - -import argparse -import math -import os -import re -import sys -from dataclasses import dataclass -from pathlib import Path - - -@dataclass(frozen=True) -class ReportRow: - model: str - status: str - value: float | None - - -@dataclass(frozen=True) -class Comparison: - model: str - base_status: str - head_status: str - base_value: float | None - head_value: float | None - change_pct: float | None - result: str - - @property - def failed(self) -> bool: - return self.result in ( - "REGRESSION", - "CANDIDATE_SKIPPED", - "STATUS_FAILED", - "METRIC_MISSING", - ) - - -_NUMBER_RE = re.compile(r"[-+]?(?:\d+(?:\.\d*)?|\.\d+)(?:[eE][-+]?\d+)?") - - -def _split_markdown_row(line: str) -> list[str]: - return [cell.strip() for cell in line.strip().strip("|").split("|")] - - -def _parse_metric(value: str) -> float | None: - lowered = value.lower() - if any(word in lowered for word in ("unknown", "skipped", "failed", "n/a")): - return None - match = _NUMBER_RE.search(value) - if match is None: - return None - parsed = float(match.group()) - if not math.isfinite(parsed) or parsed <= 0: - return None - return parsed - - -def parse_report(path: Path, metric_column: str) -> dict[str, ReportRow]: - lines = path.read_text(encoding="utf-8").splitlines() - header_index = -1 - model_index = status_index = metric_index = -1 - - for index, line in enumerate(lines): - if not line.strip().startswith("|"): - continue - cells = _split_markdown_row(line) - if "Model" not in cells or "Status" not in cells or metric_column not in cells: - continue - header_index = index - model_index = cells.index("Model") - status_index = cells.index("Status") - metric_index = cells.index(metric_column) - break - - if header_index < 0: - raise ValueError(f"{path}: table with {metric_column!r} column not found") - - rows: dict[str, ReportRow] = {} - for line in lines[header_index + 1 :]: - if not line.strip().startswith("|"): - break - cells = _split_markdown_row(line) - if not cells or max(model_index, status_index, metric_index) >= len(cells): - continue - if set(cells[0]) <= set("-: "): - continue - model = cells[model_index] - if not model: - continue - if model in rows: - raise ValueError(f"{path}: duplicate model row: {model}") - rows[model] = ReportRow( - model=model, - status=cells[status_index], - value=_parse_metric(cells[metric_index]), - ) - - if not rows: - raise ValueError(f"{path}: no model rows found") - return rows - - -def compare_reports( - base: dict[str, ReportRow], - head: dict[str, ReportRow], - mode: str, - threshold_pct: float, -) -> list[Comparison]: - missing_from_head = sorted(set(base) - set(head)) - missing_from_base = sorted(set(head) - set(base)) - if missing_from_head or missing_from_base: - details: list[str] = [] - if missing_from_head: - details.append(f"missing from head: {', '.join(missing_from_head)}") - if missing_from_base: - details.append(f"missing from base: {', '.join(missing_from_base)}") - raise ValueError("model sets differ; " + "; ".join(details)) - - comparisons: list[Comparison] = [] - for model in base: - base_row = base[model] - head_row = head[model] - result = "PASS" - change_pct: float | None = None - - if head_row.status not in ("PASS", "SKIPPED"): - result = "STATUS_FAILED" - elif base_row.status == "SKIPPED" and head_row.status == "SKIPPED": - result = "NOT_MEASURED" - elif head_row.status == "SKIPPED": - result = "CANDIDATE_SKIPPED" - elif base_row.status == "SKIPPED": - result = "BASE_SKIPPED" - elif base_row.status != "PASS": - result = "STATUS_FAILED" - elif base_row.value is None or head_row.value is None: - result = "METRIC_MISSING" - else: - change_pct = (head_row.value / base_row.value - 1.0) * 100.0 - slower_pct = change_pct if mode == "single" else -change_pct - if slower_pct > threshold_pct + 1e-9: - result = "REGRESSION" - - comparisons.append( - Comparison( - model=model, - base_status=base_row.status, - head_status=head_row.status, - base_value=base_row.value, - head_value=head_row.value, - change_pct=change_pct, - result=result, - ) - ) - return comparisons - - -def _format_number(value: float | None) -> str: - if value is None: - return "n/a" - return f"{value:.6g}" - - -def _format_change(value: float | None) -> str: - if value is None: - return "n/a" - return f"{value:+.2f}%" - - -def _mode_failure_label(mode: str) -> str: - return "Single-stream latency" if mode == "single" else "Multi-stream QPS" - - -def _failure_detail(row: Comparison, mode: str) -> str: - if row.result == "REGRESSION" and row.change_pct is not None: - if mode == "single": - return f"latency increased by {row.change_pct:.2f}%" - return f"QPS decreased by {abs(row.change_pct):.2f}%" - if row.change_pct is None: - return row.result - return f"{row.result}, change {_format_change(row.change_pct)}" - - -def _failure_lines(rows: list[Comparison], mode: str) -> list[str]: - return [f"- {row.model}: {_failure_detail(row, mode)}" for row in rows] - - -def _failure_counts(single: list[Comparison], multi: list[Comparison]) -> list[str]: - counts: list[str] = [] - single_failed = sum(row.failed for row in single) - multi_failed = sum(row.failed for row in multi) - if single_failed: - counts.append(f"{single_failed} single-stream latency") - if multi_failed: - counts.append(f"{multi_failed} multi-stream QPS") - return counts - - -def _status_label(row: Comparison, mode: str) -> str: - if row.result == "PASS": - return "PASS" - if row.result == "REGRESSION": - metric = "latency" if mode == "single" else "QPS" - return f"REGRESSION ({metric} {_format_change(row.change_pct)})" - return row.result - - -def _table_for_mode( - title: str, - comparisons: list[Comparison], - mode: str, - base_metric: str, - head_metric: str, - change_metric: str, -) -> list[str]: - lines = [ - f"## {title}", - "", - f"| Model | Base {base_metric} | Head {head_metric} | {change_metric} | Result |", - "|---|---|---|---|---|", - ] - for row in comparisons: - lines.append( - "| " - + " | ".join( - [ - row.model.replace("|", "\\|"), - _format_number(row.base_value), - _format_number(row.head_value), - _format_change(row.change_pct), - _status_label(row, mode), - ] - ) - + " |" - ) - lines.append("") - return lines - - -def build_markdown( - single: list[Comparison], - multi: list[Comparison], - threshold_pct: float, -) -> tuple[str, str]: - single_failed = [row for row in single if row.failed] - multi_failed = [row for row in multi if row.failed] - failed = single_failed + multi_failed - unmeasured = [row for row in single + multi if row.result == "NOT_MEASURED"] - model_count = len(single) - - summary_lines = [ - "## Model Zoo Performance Gate", - "", - f"- Models: {model_count}", - f"- Threshold: more than {threshold_pct:g}% slower than base", - "- Single stream metric: average latency (lower is better)", - "- Multi stream metric: QPS (higher is better)", - f"- Single-stream latency failures: {len(single_failed)}", - f"- Multi-stream QPS failures: {len(multi_failed)}", - f"- Result: {'FAIL' if failed else 'PASS'}", - ] - if failed: - if single_failed: - summary_lines.extend( - [ - "", - f"### Single-stream latency failures ({len(single_failed)})", - "", - *_failure_lines(single_failed, "single"), - ] - ) - if multi_failed: - summary_lines.extend( - [ - "", - f"### Multi-stream QPS failures ({len(multi_failed)})", - "", - *_failure_lines(multi_failed, "multi"), - ] - ) - if unmeasured: - summary_lines.extend( - ["", f"- Skipped in both versions (not measured): {len(unmeasured)}"] - ) - summary = "\n".join(summary_lines) + "\n" - - lines = [ - "# Model Zoo Performance Comparison", - "", - "| Key | Value |", - "|---|---|", - f"| Threshold | {threshold_pct:g}% |", - "| Single stream metric | Average latency, lower is better |", - "| Multi stream metric | QPS, higher is better |", - f"| Models compared | {model_count} |", - f"| Single-stream latency failures | {len(single_failed)} |", - f"| Multi-stream QPS failures | {len(multi_failed)} |", - f"| Result | {'FAIL' if failed else 'PASS'} |", - "", - ] - lines.extend( - _table_for_mode( - "Single Stream", - single, - "single", - "latency (ms)", - "latency (ms)", - "Latency change", - ) - ) - lines.extend( - _table_for_mode( - "Multi Stream (-s 4 -q 1024)", - multi, - "multi", - "QPS", - "QPS", - "QPS change", - ) - ) - lines.extend(["## Result", ""]) - if failed: - lines.append( - f"FAIL: {len(failed)} failed measurement(s): " - + ", ".join(_failure_counts(single, multi)) - + "." - ) - else: - lines.append( - f"PASS: no model regressed by more than {threshold_pct:g}% " - "in single-stream latency or multi-stream QPS." - ) - lines.append("") - return "\n".join(lines) + "\n", summary - - -def main() -> int: - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("--baseline-single", required=True, type=Path) - parser.add_argument("--candidate-single", required=True, type=Path) - parser.add_argument("--baseline-multi", required=True, type=Path) - parser.add_argument("--candidate-multi", required=True, type=Path) - parser.add_argument( - "--threshold", - type=float, - default=10.0, - help="Maximum allowed slowdown in percent (default: 10).", - ) - parser.add_argument("--output", type=Path, default=None) - parser.add_argument( - "--summary", - type=Path, - default=Path(os.environ["GITHUB_STEP_SUMMARY"]) - if os.environ.get("GITHUB_STEP_SUMMARY") - else None, - help="Optional markdown summary file, appended if it exists.", - ) - args = parser.parse_args() - - if not math.isfinite(args.threshold) or args.threshold < 0: - parser.error("--threshold must be a non-negative number") - - try: - single = compare_reports( - parse_report(args.baseline_single, "Average Latency"), - parse_report(args.candidate_single, "Average Latency"), - "single", - args.threshold, - ) - multi = compare_reports( - parse_report(args.baseline_multi, "QPS"), - parse_report(args.candidate_multi, "QPS"), - "multi", - args.threshold, - ) - except (OSError, ValueError) as exc: - print(f"Failed to compare reports: {exc}", file=sys.stderr) - return 1 - - single_models = {row.model for row in single} - multi_models = {row.model for row in multi} - if single_models != multi_models: - missing_from_multi = sorted(single_models - multi_models) - missing_from_single = sorted(multi_models - single_models) - print( - "Single-stream and multi-stream model sets differ; " - f"single-only={missing_from_multi}, multi-only={missing_from_single}", - file=sys.stderr, - ) - return 1 - - markdown, summary = build_markdown(single, multi, args.threshold) - if args.output is not None: - args.output.parent.mkdir(parents=True, exist_ok=True) - args.output.write_text(markdown, encoding="utf-8") - print(f"Saved performance comparison to {args.output}") - else: - print(markdown, end="") - - if args.summary is not None: - with args.summary.open("a", encoding="utf-8") as stream: - stream.write(summary) - - single_failed = [row for row in single if row.failed] - multi_failed = [row for row in multi if row.failed] - failed = single_failed + multi_failed - if failed: - print( - f"Performance gate failed: {len(failed)} model measurement(s) failed", - file=sys.stderr, - ) - if single_failed: - print( - f" Single-stream latency: {len(single_failed)} model(s) failed", - file=sys.stderr, - ) - if multi_failed: - print( - f" Multi-stream QPS: {len(multi_failed)} model(s) failed", - file=sys.stderr, - ) - for mode, rows in (("single", single_failed), ("multi", multi_failed)): - if not rows: - continue - print( - f"\n{_mode_failure_label(mode)} failures ({len(rows)}):", - file=sys.stderr, - ) - for row in rows: - print( - f" - {row.model}: {_failure_detail(row, mode)}", - file=sys.stderr, - ) - return 1 - - print("Performance gate passed") - return 0 - - -if __name__ == "__main__": - sys.exit(main()) diff --git a/scripts/ci/run_model_zoo_benchmark.sh b/scripts/ci/run_model_zoo_benchmark.sh deleted file mode 100755 index abcfd3b..0000000 --- a/scripts/ci/run_model_zoo_benchmark.sh +++ /dev/null @@ -1,118 +0,0 @@ -#!/usr/bin/env bash -set -euo pipefail - -usage() { - cat <<'EOF' -Usage: run_model_zoo_benchmark.sh LABEL PLUGIN MODE ZOO_DIR RESULTS_DIR - -Run one model-zoo performance pass and preserve its report and error logs. - -Arguments: - LABEL Version label used in result file names, for example base or head. - PLUGIN Absolute path to libonnxruntime_providers_musa_plugin.so. - MODE single or multi. - ZOO_DIR Prepared onnx-model-zoo working directory. - RESULTS_DIR Directory that receives the report and error logs. - -Modes: - single run_all_onnx.sh --ep musa --warmup 20 --repeat 100 - multi run_all_onnx.sh --ep musa -s 4 -q 1024 -EOF -} - -if [[ $# -ne 5 ]]; then - usage >&2 - exit 2 -fi - -LABEL="$1" -PLUGIN="$2" -MODE="$3" -ZOO_DIR="$4" -RESULTS_DIR="$5" - -if [[ ! "$LABEL" =~ ^[A-Za-z0-9_.-]+$ ]]; then - echo "Invalid label: $LABEL" >&2 - exit 2 -fi -if [[ ! -f "$PLUGIN" ]]; then - echo "MUSA plugin does not exist: $PLUGIN" >&2 - exit 1 -fi -if [[ ! -d "$ZOO_DIR" ]]; then - echo "Model zoo directory does not exist: $ZOO_DIR" >&2 - exit 1 -fi - -case "$MODE" in - single) - REPORT_NAME="musa_ep_performance.md" - RUN_ARGS=(--ep musa --warmup 20 --repeat 100) - ;; - multi) - REPORT_NAME="musa_ep_performance_with_multi_stream.md" - RUN_ARGS=(--ep musa -s 4 -q 1024) - ;; - *) - echo "Mode must be single or multi: $MODE" >&2 - exit 2 - ;; -esac - -mkdir -p "$RESULTS_DIR" -export MUSA_EP_LIBRARY="$PLUGIN" -export MUSA_HOME="${MUSA_HOME:-/usr/local/musa}" -export MUSA_VISIBLE_DEVICES="${MUSA_VISIBLE_DEVICES:-0}" -export LD_LIBRARY_PATH="$MUSA_HOME/lib:$MUSA_HOME/lib64:${LD_LIBRARY_PATH:-}" - -echo "benchmark_version=${LABEL}" -echo "benchmark_mode=${MODE}" -echo "benchmark_plugin=${MUSA_EP_LIBRARY}" -echo "MUSA_VISIBLE_DEVICES=${MUSA_VISIBLE_DEVICES}" -sha256sum "$MUSA_EP_LIBRARY" - -cd "$ZOO_DIR" -set +e -bash run_all_onnx.sh "${RUN_ARGS[@]}" -run_status=$? -set -e - -report_path="$ZOO_DIR/$REPORT_NAME" -result_path="$RESULTS_DIR/$LABEL-$MODE.md" -if [[ -f "$report_path" ]]; then - cp "$report_path" "$result_path" - echo "Saved benchmark report to $result_path" -else - echo "Benchmark report was not generated: $report_path" >&2 - exit 1 -fi - -if [[ -d "$ZOO_DIR/error_logs" ]]; then - error_results="$RESULTS_DIR/$LABEL-$MODE-error-logs" - rm -rf "$error_results" - cp -a "$ZOO_DIR/error_logs" "$error_results" - echo "Saved benchmark error logs to $error_results" -fi - -if [[ "$run_status" -ne 0 ]]; then - echo "Benchmark failed with exit code $run_status" >&2 - exit "$run_status" -fi - -if ! grep -Fq '| Fail | 0 |' "$result_path"; then - echo "Benchmark reported one or more failed models: $result_path" >&2 - grep -F '| Fail |' "$result_path" >&2 || true - exit 1 -fi - -if [[ "$MODE" == "single" ]]; then - grep -Fq '| Mode | classic |' "$result_path" - grep -Fq '| Warmup | 20 |' "$result_path" - grep -Fq '| Repeat | 100 |' "$result_path" -else - grep -Fq '| Querys | 1024 |' "$result_path" - grep -Fq '| Streams | 4 |' "$result_path" -fi - -grep -F '| Model count |' "$result_path" -echo "Benchmark completed: version=${LABEL} mode=${MODE}" From f954b5099e10ef6da26284196bf4e9a7c7627289 Mon Sep 17 00:00:00 2001 From: Albert Date: Tue, 27 Oct 2026 15:22:18 +0800 Subject: [PATCH 52/52] weekly --- .github/workflows/musa-weekly-cpu-ep-diff.yml | 177 ++++++++++++++++++ 1 file changed, 177 insertions(+) create mode 100644 .github/workflows/musa-weekly-cpu-ep-diff.yml diff --git a/.github/workflows/musa-weekly-cpu-ep-diff.yml b/.github/workflows/musa-weekly-cpu-ep-diff.yml new file mode 100644 index 0000000..a229422 --- /dev/null +++ b/.github/workflows/musa-weekly-cpu-ep-diff.yml @@ -0,0 +1,177 @@ +name: MUSA Weekly CPU EP Diff + +on: + schedule: + # Monday 06:00 Asia/Shanghai (UTC+8) = Sunday 22:00 UTC. + - cron: "0 22 * * 0" + workflow_dispatch: + +permissions: + contents: read + actions: read + +concurrency: + group: musa-weekly-cpu-ep-diff + cancel-in-progress: false + +jobs: + weekly-cpu-ep-diff: + name: Weekly Model Zoo CPU EP Diff + runs-on: [self-hosted, musa, s5000] + timeout-minutes: 720 + env: + CI_ASSETS: /data/ci-assets + MODEL_ZOO_WORK: /data/ci-work/onnx-model-zoo-weekly-cpu-ep-diff-${{ github.run_id }}-${{ github.run_attempt }} + RESULTS_DIR: /data/ci-work/onnx-model-zoo-weekly-cpu-ep-diff-${{ github.run_id }}-${{ github.run_attempt }}/results + MUSA_VISIBLE_DEVICES: "0" + PYTHONUNBUFFERED: "1" + + steps: + - name: Checkout master + uses: actions/checkout@v4 + with: + ref: master + fetch-depth: 1 + submodules: false + + - name: Configure MUSA toolkit environment + run: | + set -euo pipefail + echo "/usr/local/musa/bin" >> "${GITHUB_PATH}" + { + echo "MUSA_HOME=/usr/local/musa" + echo "LD_LIBRARY_PATH=/usr/local/musa/lib${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}" + echo "CPLUS_INCLUDE_PATH=/usr/include/c++/11:/usr/include/x86_64-linux-gnu/c++/11:/usr/include${CPLUS_INCLUDE_PATH:+:${CPLUS_INCLUDE_PATH}}" + } >> "${GITHUB_ENV}" + + - name: Prepare ONNX Runtime headers + run: | + set -euo pipefail + test -d "${CI_ASSETS}/onnxruntime" + rm -rf third_party/onnxruntime + cp -a "${CI_ASSETS}/onnxruntime" third_party/onnxruntime + + - name: Build master plugin + run: | + set -euo pipefail + echo "Master commit: $(git rev-parse HEAD)" + bash build.sh --clean --no-wheel + test -f build/Release/libonnxruntime_providers_musa_plugin.so + sha256sum build/Release/libonnxruntime_providers_musa_plugin.so + + - name: Prepare model zoo + run: | + set -euo pipefail + test -d "${CI_ASSETS}/onnx-model-zoo" + rm -rf "${MODEL_ZOO_WORK}" + mkdir -p "${MODEL_ZOO_WORK}" + cp -al "${CI_ASSETS}/onnx-model-zoo/." "${MODEL_ZOO_WORK}/" + rm -rf "${MODEL_ZOO_WORK}/cpp_runners" + cp -a "${CI_ASSETS}/onnx-model-zoo/cpp_runners" "${MODEL_ZOO_WORK}/cpp_runners" + mkdir -p "${RESULTS_DIR}" + test -f "${MODEL_ZOO_WORK}/run_all_onnx_cpu_ep_diff.sh" + test -f "${MODEL_ZOO_WORK}/cpp_runners/build_onnx_runner.sh" + test -f "/data/ci-assets/ort-sdk/onnxruntime-linux-x64-1.26.0/include/onnxruntime_c_api.h" + + - name: Build model zoo CPU EP diff runner + run: | + set -euo pipefail + cd "${MODEL_ZOO_WORK}" + ORT_ROOT=/data/ci-assets/ort-sdk/onnxruntime-linux-x64-1.26.0/ \ + bash cpp_runners/build_onnx_runner.sh + test -x cpp_runners/onnx_cpu_ep_diff + + - name: Run weekly MUSA-vs-CPU diff + run: | + set -euo pipefail + plugin="${GITHUB_WORKSPACE}/build/Release/libonnxruntime_providers_musa_plugin.so" + raw_log="${RESULTS_DIR}/musa-cpu-ep-diff.log" + report="${RESULTS_DIR}/musa-weekly-cpu-ep-diff-report.md" + test -f "${plugin}" + + commit="$(git rev-parse HEAD)" + plugin_sha="$(sha256sum "${plugin}" | awk '{print $1}')" + + cd "${MODEL_ZOO_WORK}" + set +e + MUSA_EP_LIBRARY="${plugin}" \ + bash run_all_onnx_cpu_ep_diff.sh --ep musa -s 4 -q 1024 2>&1 | tee "${raw_log}" + pipeline_status=("${PIPESTATUS[@]}") + run_status=${pipeline_status[0]} + tee_status=${pipeline_status[1]} + set -e + + summary="$(grep -E '^pass=[0-9]+ fail=[0-9]+ skip=[0-9]+ logs=' "${raw_log}" | tail -n 1 || true)" + summary_pattern='^pass=([0-9]+) fail=([0-9]+) skip=([0-9]+) logs=' + summary_found=0 + fail_count="" + if [[ "${summary}" =~ ${summary_pattern} ]]; then + summary_found=1 + fail_count="${BASH_REMATCH[2]}" + fi + + gate_status="FAIL" + if [[ "${summary_found}" -eq 1 && "${tee_status}" -eq 0 && "${fail_count}" -eq 0 ]]; then + gate_status="PASS" + fi + + { + echo "# Weekly Model Zoo CPU EP Diff" + echo + echo "| Field | Value |" + echo "| --- | --- |" + echo "| Result | **${gate_status}** |" + echo '| Branch | `master` |' + echo "| Commit | \`${commit}\` |" + echo '| Command | `run_all_onnx_cpu_ep_diff.sh --ep musa -s 4 -q 1024` |' + echo "| Plugin SHA256 | \`${plugin_sha}\` |" + echo "| Summary | \`${summary:-not found}\` |" + echo "| Runner exit | \`${run_status}\` |" + echo + echo "## Diff Output" + echo + echo '```text' + cat "${raw_log}" + echo '```' + } > "${report}" + + { + echo "## Weekly Model Zoo CPU EP Diff" + echo + echo "- Result: **${gate_status}**" + echo "- Commit: \`${commit}\`" + echo "- Summary: \`${summary:-not found}\`" + } >> "${GITHUB_STEP_SUMMARY}" + + if [[ "${summary_found}" -eq 0 ]]; then + echo "Could not find the final CPU EP diff summary in ${raw_log}" >&2 + exit 1 + fi + echo "Weekly CPU EP diff result: ${summary}" + if (( tee_status != 0 )); then + echo "Failed to write the weekly report log (tee exit ${tee_status})" >&2 + exit 1 + fi + if (( fail_count != 0 )); then + echo "Weekly CPU EP diff gate failed: fail=${fail_count}" >&2 + exit 1 + fi + echo "Weekly CPU EP diff passed: fail=0 (PASS and SKIPPED are accepted)" + + - name: Upload weekly report + if: always() + uses: actions/upload-artifact@v4 + with: + name: musa-weekly-cpu-ep-diff-report-${{ github.run_id }}-${{ github.run_attempt }} + path: ${{ env.RESULTS_DIR }}/musa-weekly-cpu-ep-diff-report.md + if-no-files-found: warn + retention-days: 30 + + - name: Cleanup model zoo work directory + if: always() + run: | + set -euo pipefail + case "${MODEL_ZOO_WORK}" in + /data/ci-work/onnx-model-zoo-weekly-cpu-ep-diff-*) rm -rf "${MODEL_ZOO_WORK}" ;; + *) echo "Refusing to remove unexpected path: ${MODEL_ZOO_WORK}" >&2; exit 1 ;; + esac