diff --git a/cmake/deps.txt b/cmake/deps.txt index f4ffa9fb2ec46..477383a418782 100644 --- a/cmake/deps.txt +++ b/cmake/deps.txt @@ -54,10 +54,10 @@ pytorch_cpuinfo;https://github.com/pytorch/cpuinfo/archive/4628dc060ce4e82345dc1 re2;https://github.com/google/re2/archive/refs/tags/2024-07-02.zip;646e1728269cde7fcef990bf4a8e87b047882e88 safeint;https://github.com/dcleblanc/SafeInt/archive/refs/tags/3.0.28.zip;23f252040ff6cb9f1fd18575b32fa8fb5928daac tensorboard;https://github.com/tensorflow/tensorboard/archive/373eb09e4c5d2b3cc2493f0949dc4be6b6a45e81.zip;67b833913605a4f3f499894ab11528a702c2b381 -cutlass;https://github.com/NVIDIA/cutlass/archive/refs/tags/v4.4.2.zip;4b0bae4428b84370407c0a71778b13dc2eee5be1 +cutlass;https://github.com/NVIDIA/cutlass/archive/refs/tags/v4.7.0.zip;51d4f1ba4b9977b94606e23bbc307463a6741383 extensions;https://github.com/microsoft/onnxruntime-extensions/archive/c24b7bab0c12f53da76d0c31b03b9f0f8ec8f3b4.zip;239063aee4946a9af147b473a4c3da78ba7413b4 directx_headers;https://github.com/microsoft/DirectX-Headers/archive/refs/tags/v1.613.1.zip;47653509a3371eabb156360f42faf582f314bf2e -cudnn_frontend;https://github.com/NVIDIA/cudnn-frontend/archive/refs/tags/v1.24.0.zip;a55a1980bf5c57692d66ae7bc3b39798f5535e1f +cudnn_frontend;https://github.com/NVIDIA/cudnn-frontend/archive/refs/tags/v1.27.0.zip;1e4c9a464d3437e388ab0163f3be068dba783c08 dawn;https://github.com/google/dawn/archive/refs/tags/v20260714.215939.zip;3056ed22d1606258ab43221b8c85b55b88614137 kleidiai;https://github.com/ARM-software/kleidiai/archive/refs/tags/v1.20.0.tar.gz;6895e72b3d5cf1173358164cb3d64c9d7d33cc84 # kleidiai-qmx is pinned to a specific commit as there are no tagged releases. When an appropriate tagged release becomes available, diff --git a/cmake/external/cutlass.cmake b/cmake/external/cutlass.cmake index cd9a9c5179615..b0a67deaeb963 100644 --- a/cmake/external/cutlass.cmake +++ b/cmake/external/cutlass.cmake @@ -4,7 +4,7 @@ onnxruntime_fetchcontent_declare( URL ${DEP_URL_cutlass} URL_HASH SHA1=${DEP_SHA1_cutlass} EXCLUDE_FROM_ALL - PATCH_COMMAND ${Patch_EXECUTABLE} --ignore-whitespace -p1 < ${PROJECT_SOURCE_DIR}/patches/cutlass/cutlass_4.4.2.patch + PATCH_COMMAND ${Patch_EXECUTABLE} --ignore-whitespace -p1 < ${PROJECT_SOURCE_DIR}/patches/cutlass/cutlass_4.7.0.patch ) # We only consume CUTLASS as a header-only dependency. Avoid FetchContent_MakeAvailable here diff --git a/cmake/onnxruntime_unittests.cmake b/cmake/onnxruntime_unittests.cmake index ea0d684b30b61..30fe6607d30d7 100644 --- a/cmake/onnxruntime_unittests.cmake +++ b/cmake/onnxruntime_unittests.cmake @@ -1066,6 +1066,7 @@ if (onnxruntime_ENABLE_CUDA_EP_INTERNAL_TESTS AND onnxruntime_BUILD_CUDA_EP_AS_P "${ONNXRUNTIME_ROOT}/core/providers/cuda/cudnn_common.cc" "${ONNXRUNTIME_ROOT}/core/providers/cuda/cudnn_loader.cc" "${ONNXRUNTIME_ROOT}/core/providers/cuda/cufft_loader.cc" + "${ONNXRUNTIME_ROOT}/core/providers/cuda/fpgeneric.cu" "${ONNXRUNTIME_ROOT}/core/providers/cuda/reduction/reduction_functions.cc" "${ONNXRUNTIME_ROOT}/core/providers/cuda/reduction/reduction_functions.cu" "${TEST_SRC_DIR}/providers/cuda/test_cases/cuda_plugin_test_shims.cc" diff --git a/cmake/patches/cudnn_frontend/cudnn_frontend_win_dynamic_loading.patch b/cmake/patches/cudnn_frontend/cudnn_frontend_win_dynamic_loading.patch index 96242982a7468..c815d6305ddf8 100644 --- a/cmake/patches/cudnn_frontend/cudnn_frontend_win_dynamic_loading.patch +++ b/cmake/patches/cudnn_frontend/cudnn_frontend_win_dynamic_loading.patch @@ -15,7 +15,16 @@ const char *error = reinterpret_cast(dlerror()); if (!handle || error) { // If opening the library fails, throw an exception with the error message -@@ -85,13 +92,22 @@ +@@ -82,5 +89,8 @@ + if (user_lib) { + if (user_lib[0] != '\0') { ++#ifdef _WIN32 ++ SetLastError(0); ++#endif + HMODULE handle = dlopen(user_lib, RTLD_NOW); + const char *error = reinterpret_cast(dlerror()); + if (!handle || error) { +@@ -85,13 +95,22 @@ dlerror(); // List of potential libcudart libraries (Adding major version to support python package) diff --git a/cmake/patches/cutlass/cutlass_4.4.2.patch b/cmake/patches/cutlass/cutlass_4.7.0.patch similarity index 72% rename from cmake/patches/cutlass/cutlass_4.4.2.patch rename to cmake/patches/cutlass/cutlass_4.7.0.patch index 6776eded10640..afb2d21472483 100644 --- a/cmake/patches/cutlass/cutlass_4.4.2.patch +++ b/cmake/patches/cutlass/cutlass_4.7.0.patch @@ -26,11 +26,17 @@ index a8af62be..22e7332d 100644 virtual Status memsetDeviceImpl( void* destination, ///< Device memory pointer to be filled void const* fill_value, ///< Value to be filled in the buffer +diff --git a/include/cutlass/gemm/kernel/sm100_static_tile_scheduler.hpp b/include/cutlass/gemm/kernel/sm100_static_tile_scheduler.hpp +index ae31620a..fd559cbb 100644 +--- a/include/cutlass/gemm/kernel/sm100_static_tile_scheduler.hpp ++++ b/include/cutlass/gemm/kernel/sm100_static_tile_scheduler.hpp +@@ -53 +52,0 @@ public: +- CUTLASS_HOST_DEVICE diff --git a/include/cutlass/exmy_base.h b/include/cutlass/exmy_base.h index be207a49..6028e01d 100644 --- a/include/cutlass/exmy_base.h +++ b/include/cutlass/exmy_base.h -@@ -1021,18 +1021,18 @@ struct float_exmy_base +@@ -1021,17 +1021,17 @@ struct float_exmy_base /// Floating point conversion CUTLASS_HOST_DEVICE @@ -51,4 +57,15 @@ index be207a49..6028e01d 100644 + explicit float_exmy_base(unsigned x) { storage = static_cast(this)->convert_from_float(float(x)).storage; } - +diff --git a/include/cutlass/subbyte_reference.h b/include/cutlass/subbyte_reference.h +--- a/include/cutlass/subbyte_reference.h ++++ b/include/cutlass/subbyte_reference.h +@@ -456,7 +456,7 @@ public: + // + Storage assumed; + #if (__CUDACC_VER_MAJOR__ > 12) || (__CUDACC_VER_MAJOR__ == 12 && __CUDACC_VER_MINOR__ >= 8) +- Storage original = __nv_atomic_load_n(ptr_, __NV_ATOMIC_RELAXED); ++ Storage original = __nv_atomic_load_n(ptr_, __NV_ATOMIC_RELAXED, __NV_THREAD_SCOPE_DEVICE); + #else + Storage original = *const_cast(ptr_); + #endif diff --git a/docs/ContribOperators.md b/docs/ContribOperators.md index acaa3b724fc51..8338a938903bd 100644 --- a/docs/ContribOperators.md +++ b/docs/ContribOperators.md @@ -7214,6 +7214,8 @@ No versioning maintained for experimental ops. The embedding layer takes input_ids (word IDs) and segment_ids (sentence IDs) to look up word_embedding, position_embedding, and segment_emedding; the embeddings are added then applied layer normalization using gamma and beta tensors. The input_ids and segment_ids remain int32. All embeddings, gamma, and beta tensors are converted to int8/uint8. The last input mask is optional. + segment_ids, segment_embedding, segment_embedding_scale, and segment_embedding_zero_point must either all be provided or all + be omitted. If mask is provided, mask index (that is position of first 0 in mask, or number of words will be calculated. #### Version @@ -7232,13 +7234,13 @@ No versioning maintained for experimental ops.
input_ids : T1
2D words IDs with shape (batch_size, sequence_length)
segment_ids (optional) : T1
-
2D segment IDs with shape (batch_size, sequence_length)
+
2D segment IDs with shape (batch_size, sequence_length). Part of the all-or-none segment input group.
word_embedding_quant : T2
2D with shape (,hidden_size)
position_embedding_quant : T2
2D with shape (, hidden_size)
segment_embedding (optional) : T2
-
2D with shape (, hidden_size)
+
2D with shape (, hidden_size). Part of the all-or-none segment input group.
gamma_quant : T2
1D gamma tensor for layer normalization with shape (hidden_size)
beta_quant : T2
@@ -7250,7 +7252,7 @@ No versioning maintained for experimental ops.
position_embedding_scale : T
Scale for position embeddings
segment_embedding_scale (optional) : T
-
Scale for segment embeddings
+
Scale for segment embeddings. Part of the all-or-none segment input group.
gamma_scale : T
Scale for 1D gamma tensor
beta_scale : T
@@ -7260,7 +7262,7 @@ No versioning maintained for experimental ops.
position_embedding_zero_point : T2
Zero point for position embeddings
segment_embedding_zero_point (optional) : T2
-
Zero Point for segment embeddings
+
Zero Point for segment embeddings. Part of the all-or-none segment input group.
gamma_zero_point : T2
Zero Point for 1D gamma tensor
beta_zero_point : T2
diff --git a/onnxruntime/contrib_ops/cpu/quantization/qembed_layer_norm.cc b/onnxruntime/contrib_ops/cpu/quantization/qembed_layer_norm.cc index 2907abbfe7c62..1bc16eacfe1e9 100644 --- a/onnxruntime/contrib_ops/cpu/quantization/qembed_layer_norm.cc +++ b/onnxruntime/contrib_ops/cpu/quantization/qembed_layer_norm.cc @@ -187,6 +187,8 @@ Status ComputeInternal(OpKernelContext* context, float epsilon) { } Status CheckQuantizedInputs(OpKernelContext* context, bool* is_signed_inputs) { + const Tensor* segment_ids_tensor = context->Input(1); + const Tensor* segment_embedding_tensor = context->Input(4); const Tensor* word_embedding_scale_tensor = context->Input(8); const Tensor* position_embedding_scale_tensor = context->Input(9); const Tensor* segment_embedding_scale_tensor = context->Input(10); @@ -198,8 +200,14 @@ Status CheckQuantizedInputs(OpKernelContext* context, bool* is_signed_inputs) { const Tensor* gamma_zero_point_tensor = context->Input(16); const Tensor* beta_zero_point_tensor = context->Input(17); + const bool has_segment_embedding = segment_ids_tensor != nullptr; + ORT_RETURN_IF(has_segment_embedding != (segment_embedding_tensor != nullptr) || + has_segment_embedding != (segment_embedding_scale_tensor != nullptr) || + has_segment_embedding != (segment_embedding_zero_point_tensor != nullptr), + "segment_ids, segment_embedding, segment_embedding_scale, and segment_embedding_zero_point " + "must either all be provided or all be omitted"); + bool word_embedding_is_signed_inputs = word_embedding_zero_point_tensor->IsDataType(); - bool has_segment_embedding = context->Input(1) != nullptr; if (!IsScalarOr1ElementVector(word_embedding_scale_tensor)) { return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, diff --git a/onnxruntime/contrib_ops/cuda/llm/cutlass_extensions/epilogue/collective/epilogue_moe_finalize.hpp b/onnxruntime/contrib_ops/cuda/llm/cutlass_extensions/epilogue/collective/epilogue_moe_finalize.hpp index 8ba877aa21a68..4aca09018377d 100644 --- a/onnxruntime/contrib_ops/cuda/llm/cutlass_extensions/epilogue/collective/epilogue_moe_finalize.hpp +++ b/onnxruntime/contrib_ops/cuda/llm/cutlass_extensions/epilogue/collective/epilogue_moe_finalize.hpp @@ -449,20 +449,20 @@ struct EpilogueMoeFusedFinalizeBuilder { // Dummy methods to perform different parts of TMA/Tensormap modifications - template + template CUTLASS_DEVICE void tensormaps_perform_update([[maybe_unused]] TensorMapStorage& shared_tensormaps, [[maybe_unused]] typename EpilogueOp::Params const& params, - [[maybe_unused]] cute::TmaDescriptor const* tensormap, [[maybe_unused]] ProblemShapeMNKL problem_shape, + [[maybe_unused]] TensorMaps const& tensormaps, [[maybe_unused]] ProblemShapeMNKL problem_shape, [[maybe_unused]] int32_t next_batch, [[maybe_unused]] int32_t warp_group_idx) { } - template + template CUTLASS_DEVICE void tensormaps_cp_fence_release([[maybe_unused]] TensorMapStorage& shared_tensormaps, - [[maybe_unused]] cute::TmaDescriptor const* tensormap, [[maybe_unused]] int32_t warp_group_idx) { + [[maybe_unused]] TensorMaps const& tensormaps, [[maybe_unused]] int32_t warp_group_idx) { } - template - CUTLASS_DEVICE void tensormaps_fence_acquire([[maybe_unused]] cute::TmaDescriptor const* tensormap) { + template + CUTLASS_DEVICE void tensormaps_fence_acquire([[maybe_unused]] TensorMaps const& tensormaps) { } }; diff --git a/onnxruntime/core/framework/tensor_external_data_info.cc b/onnxruntime/core/framework/tensor_external_data_info.cc index dfdb3ba962609..494671ba0c5a1 100644 --- a/onnxruntime/core/framework/tensor_external_data_info.cc +++ b/onnxruntime/core/framework/tensor_external_data_info.cc @@ -29,6 +29,7 @@ Status ExternalDataInfo::Create(const RepeatedPtrField& std::unique_ptr& external_data_info_result) { auto external_data_info = std::make_unique(); PrepackedInfos prepacked_infos; + bool has_location = false; const int input_size = input.size(); @@ -39,7 +40,11 @@ Status ExternalDataInfo::Create(const RepeatedPtrField& if (!stringmap.has_value()) return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "model format error! Need a value for the external data info"); - if (stringmap.key() == "location" && !stringmap.value().empty()) { + if (stringmap.key() == "location") { + ORT_RETURN_IF(has_location, + "model format error! TensorProto external data has duplicate 'location' entries"); + has_location = true; + ORT_RETURN_IF(stringmap.value().empty(), "model format error! External data location cannot be empty"); external_data_info->rel_path_ = ToWideString(stringmap.value()); } else if (stringmap.key() == "offset" && !stringmap.value().empty()) { ORT_RETURN_IF_ERROR(ParseStringWithClassicLocale(stringmap.value(), external_data_info->offset_)); diff --git a/onnxruntime/core/framework/tensorprotoutils.cc b/onnxruntime/core/framework/tensorprotoutils.cc index f19222d3f7b87..bc8e72e2ebf15 100644 --- a/onnxruntime/core/framework/tensorprotoutils.cc +++ b/onnxruntime/core/framework/tensorprotoutils.cc @@ -312,11 +312,13 @@ namespace utils { bool HasExternalDataInMemory(const ONNX_NAMESPACE::TensorProto& ten_proto) { if (HasExternalData(ten_proto)) { - // Retrieve the external data info for (const auto& entry : ten_proto.external_data()) { if (entry.key() == "location") { - PathString location = ToWideString(entry.value()); - return ((location == kTensorProtoLittleEndianMemoryAddressTag) || (location == kTensorProtoNativeEndianMemoryAddressTag)); + const PathString location = ToWideString(entry.value()); + if (location == kTensorProtoLittleEndianMemoryAddressTag || + location == kTensorProtoNativeEndianMemoryAddressTag) { + return true; + } } } } @@ -407,11 +409,12 @@ static bool HasPathComponentPrefix(const std::filesystem::path& prefix, const st /// /// Validation steps: /// 1. Reject empty paths -/// 2. Reject absolute paths (including Unix-style '/...' on Windows) -/// 3. Skip remaining checks on WASM if no filesystem is available -/// 4. Resolve `model_dir / external_data_path` to a canonical path (resolving symlinks for existing segments) -/// 5. Verify the canonical path is a prefix-child of the canonical model_dir (containment check) -/// 6. Verify the resolved file exists on disk +/// 2. Reject internal in-memory reference tags +/// 3. Reject absolute paths (including Unix-style '/...' on Windows) +/// 4. Skip remaining checks on WASM if no filesystem is available +/// 5. Resolve `model_dir / external_data_path` to a canonical path (resolving symlinks for existing segments) +/// 6. Verify the canonical path is a prefix-child of the canonical model_dir (containment check) +/// 7. Verify the resolved file exists on disk /// /// This function does NOT handle the symlinked-model fallback — that is the responsibility of /// ValidateExternalDataPath(), which calls this function as a first pass. @@ -420,13 +423,18 @@ Status ValidateExternalDataPathFromDir(const std::filesystem::path& model_dir, // Step 1: Reject empty external data paths. ORT_RETURN_IF(external_data_path.empty(), "Empty external data path not allowed"); - // Step 2: Reject absolute paths. + // Step 2: Reject internal in-memory reference tags. + ORT_RETURN_IF(external_data_path.native() == kTensorProtoLittleEndianMemoryAddressTag || + external_data_path.native() == kTensorProtoNativeEndianMemoryAddressTag, + "In-memory external data reference tag is not a valid file path"); + + // Step 3: Reject absolute paths. // Use !root_path().empty() to reject paths like '/some/path' even on Windows (where is_absolute() // requires a drive letter). ORT_RETURN_IF(!external_data_path.root_path().empty(), "Absolute path not allowed for external data location"); #if defined(__wasm__) - // Step 3 (WASM only): If we can't access the current working directory, assume the WASM environment + // Step 4 (WASM only): If we can't access the current working directory, assume the WASM environment // does not have a virtual filesystem and defer validation to an ExternalDataLoader for the WASM EP. std::error_code error_code; std::filesystem::current_path(error_code); @@ -435,7 +443,7 @@ Status ValidateExternalDataPathFromDir(const std::filesystem::path& model_dir, } #endif - // Step 4: Resolve both the model directory and the combined path to canonical forms. + // Step 5: Resolve both the model directory and the combined path to canonical forms. // WeaklyCanonicalPath resolves symlinks for existing path segments while lexically normalizing // non-existent trailing segments. std::filesystem::path resolved_dir = model_dir.empty() ? std::filesystem::path{"."} : model_dir; @@ -445,8 +453,8 @@ Status ValidateExternalDataPathFromDir(const std::filesystem::path& model_dir, ORT_RETURN_IF_ERROR(WeaklyCanonicalPath(resolved_dir, model_dir_canonical)); ORT_RETURN_IF_ERROR(WeaklyCanonicalPath(model_dir_canonical / external_data_path, external_data_path_canonical)); - // Step 5: Containment check — verify the resolved external data path starts with the model directory. - // Step 6: Existence check — verify the file actually exists on disk. + // Step 6: Containment check — verify the resolved external data path starts with the model directory. + // Step 7: Existence check — verify the file actually exists on disk. if (HasPathComponentPrefix(model_dir_canonical, external_data_path_canonical)) { bool path_exists = false; ORT_RETURN_IF_ERROR(PathExists(external_data_path_canonical, path_exists)); @@ -466,7 +474,7 @@ Status ValidateExternalDataPathFromDir(const std::filesystem::path& model_dir, /// Validation flow: /// 1. Try ValidateExternalDataPathFromDir against the model file's parent directory. /// If it passes, return success. -/// 2. If it fails due to empty/absolute external_data_path, return the error immediately +/// 2. If it fails due to an empty/absolute path or an in-memory reference tag, return the error immediately /// (these are input errors unrelated to the model location). /// 3. If model_path is empty (model loaded from bytes), wrap the error with context. /// 4. If model_path is a symlink, try the symlink fallback: @@ -494,9 +502,11 @@ Status ValidateExternalDataPath(const std::filesystem::path& model_path, } // --- Guard: Don't retry for input-validation errors --- - // Empty and absolute paths are always invalid regardless of model directory or symlinks. + // Empty paths, absolute paths, and in-memory reference tags are always invalid regardless of model directory. // Return the error directly without misleading "escapes directory" context. - if (external_data_path.empty() || !external_data_path.root_path().empty()) { + if (external_data_path.empty() || !external_data_path.root_path().empty() || + external_data_path.native() == kTensorProtoLittleEndianMemoryAddressTag || + external_data_path.native() == kTensorProtoNativeEndianMemoryAddressTag) { return status; } diff --git a/onnxruntime/core/graph/contrib_ops/quantization_defs.cc b/onnxruntime/core/graph/contrib_ops/quantization_defs.cc index 5a3cd86b04492..70cd688832b02 100644 --- a/onnxruntime/core/graph/contrib_ops/quantization_defs.cc +++ b/onnxruntime/core/graph/contrib_ops/quantization_defs.cc @@ -1011,6 +1011,8 @@ QEmbedLayerNormalization is the quantized fusion of embedding layer in BERT mode The embedding layer takes input_ids (word IDs) and segment_ids (sentence IDs) to look up word_embedding, position_embedding, and segment_emedding; the embeddings are added then applied layer normalization using gamma and beta tensors. The input_ids and segment_ids remain int32. All embeddings, gamma, and beta tensors are converted to int8/uint8. The last input mask is optional. +segment_ids, segment_embedding, segment_embedding_scale, and segment_embedding_zero_point must either all be provided or all +be omitted. If mask is provided, mask index (that is position of first 0 in mask, or number of words will be calculated.)DOC"; ONNX_MS_OPERATOR_SET_SCHEMA( @@ -1021,21 +1023,21 @@ ONNX_MS_OPERATOR_SET_SCHEMA( .Attr("epsilon", "The epsilon value to use to avoid division by zero.", AttributeProto::FLOAT, kDefaultEmbedLayerNormEpsilon) .Input(0, "input_ids", "2D words IDs with shape (batch_size, sequence_length)", "T1") - .Input(1, "segment_ids", "2D segment IDs with shape (batch_size, sequence_length)", "T1", OpSchema::Optional) + .Input(1, "segment_ids", "2D segment IDs with shape (batch_size, sequence_length). Part of the all-or-none segment input group.", "T1", OpSchema::Optional) .Input(2, "word_embedding_quant", "2D with shape (,hidden_size)", "T2") .Input(3, "position_embedding_quant", "2D with shape (, hidden_size)", "T2") - .Input(4, "segment_embedding", "2D with shape (, hidden_size)", "T2", OpSchema::Optional) + .Input(4, "segment_embedding", "2D with shape (, hidden_size). Part of the all-or-none segment input group.", "T2", OpSchema::Optional) .Input(5, "gamma_quant", "1D gamma tensor for layer normalization with shape (hidden_size)", "T2") .Input(6, "beta_quant", "1D beta tensor for layer normalization with shape (hidden_size)", "T2") .Input(7, "mask", "Mask", "T1", OpSchema::Optional) .Input(8, "word_embedding_scale", "Scale for word embeddings", "T") .Input(9, "position_embedding_scale", "Scale for position embeddings", "T") - .Input(10, "segment_embedding_scale", "Scale for segment embeddings", "T", OpSchema::Optional) + .Input(10, "segment_embedding_scale", "Scale for segment embeddings. Part of the all-or-none segment input group.", "T", OpSchema::Optional) .Input(11, "gamma_scale", "Scale for 1D gamma tensor", "T") .Input(12, "beta_scale", "Scale for 1D beta tensor", "T") .Input(13, "word_embedding_zero_point", "Zero point for word embeddings", "T2") .Input(14, "position_embedding_zero_point", "Zero point for position embeddings", "T2") - .Input(15, "segment_embedding_zero_point", "Zero Point for segment embeddings", "T2", OpSchema::Optional) + .Input(15, "segment_embedding_zero_point", "Zero Point for segment embeddings. Part of the all-or-none segment input group.", "T2", OpSchema::Optional) .Input(16, "gamma_zero_point", "Zero Point for 1D gamma tensor", "T2") .Input(17, "beta_zero_point", "Zero Point for 1D beta tensor", "T2") .Output(0, "layernorm_out", "LayerNorm Output", "T") @@ -1043,7 +1045,18 @@ ONNX_MS_OPERATOR_SET_SCHEMA( .TypeConstraint("T1", {"tensor(int32)"}, "Constrain mask index to integer types") .TypeConstraint("T2", {"tensor(int8)", "tensor(uint8)"}, "Constrain input and output types to int8 tensors.") .TypeConstraint("T", {"tensor(float)"}, "Constrain input and output types to float32 tensors.") - .TypeAndShapeInferenceFunction(EmbedLayerNormalizationShapeInference)); + .TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) { + const bool has_segment_ids = ctx.getInputType(1) != nullptr; + if (has_segment_ids != (ctx.getInputType(4) != nullptr) || + has_segment_ids != (ctx.getInputType(10) != nullptr) || + has_segment_ids != (ctx.getInputType(15) != nullptr)) { + fail_type_inference( + "segment_ids, segment_embedding, segment_embedding_scale, and segment_embedding_zero_point " + "must either all be provided or all be omitted"); + } + + EmbedLayerNormalizationShapeInference(ctx); + })); ONNX_MS_OPERATOR_SET_SCHEMA( QuantizeWithOrder, 1, diff --git a/onnxruntime/core/mlas/lib/aarch64/SymQgemmS8KernelNeon.S b/onnxruntime/core/mlas/lib/aarch64/SymQgemmS8KernelNeon.S index f236ef4ed1742..adcb4ba87a95c 100644 --- a/onnxruntime/core/mlas/lib/aarch64/SymQgemmS8KernelNeon.S +++ b/onnxruntime/core/mlas/lib/aarch64/SymQgemmS8KernelNeon.S @@ -140,64 +140,82 @@ M4_ProcessNextColumnLoop: add x11,x10,x7 // A3 M4_ComputeBlockLoop: + // Reduce each 8-lane int8 product group into int32 before the next multiply. + // Two -128 * -128 products can overflow an int16 accumulator; see #31573. smull v12.8h,v0.8b,v4.8b smull v13.8h,v0.8b,v5.8b smull v14.8h,v0.8b,v6.8b smull v15.8h,v0.8b,v7.8b - smlal v12.8h,v2.8b,v8.8b - smlal v13.8h,v2.8b,v9.8b - smlal v14.8h,v2.8b,v10.8b - smlal v15.8h,v2.8b,v11.8b - ldp d0,d2,[x10],#16 // Load A2 sadalp v16.4s,v12.8h sadalp v17.4s,v13.8h sadalp v18.4s,v14.8h sadalp v19.4s,v15.8h + smull v12.8h,v2.8b,v8.8b + smull v13.8h,v2.8b,v9.8b + smull v14.8h,v2.8b,v10.8b + smull v15.8h,v2.8b,v11.8b + sadalp v16.4s,v12.8h + sadalp v17.4s,v13.8h + sadalp v18.4s,v14.8h + sadalp v19.4s,v15.8h + ldp d0,d2,[x10],#16 // Load A2 sub x3,x3,#1 smull v12.8h,v1.8b,v4.8b smull v13.8h,v1.8b,v5.8b smull v14.8h,v1.8b,v6.8b smull v15.8h,v1.8b,v7.8b - smlal v12.8h,v3.8b,v8.8b - smlal v13.8h,v3.8b,v9.8b - smlal v14.8h,v3.8b,v10.8b - smlal v15.8h,v3.8b,v11.8b - ldp d1,d3,[x11],#16 // Load A3 sadalp v20.4s,v12.8h sadalp v21.4s,v13.8h sadalp v22.4s,v14.8h sadalp v23.4s,v15.8h + smull v12.8h,v3.8b,v8.8b + smull v13.8h,v3.8b,v9.8b + smull v14.8h,v3.8b,v10.8b + smull v15.8h,v3.8b,v11.8b + sadalp v20.4s,v12.8h + sadalp v21.4s,v13.8h + sadalp v22.4s,v14.8h + sadalp v23.4s,v15.8h + ldp d1,d3,[x11],#16 // Load A3 cbz x3,M4_ComputeBlockLoopFinish smull v12.8h,v0.8b,v4.8b smull v13.8h,v0.8b,v5.8b smull v14.8h,v0.8b,v6.8b smull v15.8h,v0.8b,v7.8b - smlal v12.8h,v2.8b,v8.8b - smlal v13.8h,v2.8b,v9.8b - smlal v14.8h,v2.8b,v10.8b - smlal v15.8h,v2.8b,v11.8b - ldp d0,d2,[x0],#16 // Load A0 next iter sadalp v24.4s,v12.8h sadalp v25.4s,v13.8h sadalp v26.4s,v14.8h sadalp v27.4s,v15.8h + smull v12.8h,v2.8b,v8.8b + smull v13.8h,v2.8b,v9.8b + smull v14.8h,v2.8b,v10.8b + smull v15.8h,v2.8b,v11.8b + sadalp v24.4s,v12.8h + sadalp v25.4s,v13.8h + sadalp v26.4s,v14.8h + sadalp v27.4s,v15.8h + ldp d0,d2,[x0],#16 // Load A0 next iter smull v12.8h,v1.8b,v4.8b smull v13.8h,v1.8b,v5.8b smull v14.8h,v1.8b,v6.8b smull v15.8h,v1.8b,v7.8b - smlal v12.8h,v3.8b,v8.8b + sadalp v28.4s,v12.8h + sadalp v29.4s,v13.8h + sadalp v30.4s,v14.8h + sadalp v31.4s,v15.8h + smull v12.8h,v3.8b,v8.8b + smull v13.8h,v3.8b,v9.8b + smull v14.8h,v3.8b,v10.8b + smull v15.8h,v3.8b,v11.8b + sadalp v28.4s,v12.8h + sadalp v29.4s,v13.8h + sadalp v30.4s,v14.8h + sadalp v31.4s,v15.8h ldp d4,d8,[x1],#64 // B - smlal v13.8h,v3.8b,v9.8b ldp d5,d9,[x1,#-48] - smlal v14.8h,v3.8b,v10.8b ldp d6,d10,[x1,#-32] - smlal v15.8h,v3.8b,v11.8b ldp d7,d11,[x1,#-16] - sadalp v28.4s,v12.8h ldp d1,d3,[x9],#16 // Load A1 next iter - sadalp v29.4s,v13.8h - sadalp v30.4s,v14.8h - sadalp v31.4s,v15.8h b M4_ComputeBlockLoop M4_ComputeBlockLoopFinish: @@ -205,23 +223,31 @@ M4_ComputeBlockLoopFinish: smull v13.8h,v0.8b,v5.8b smull v14.8h,v0.8b,v6.8b smull v15.8h,v0.8b,v7.8b - smlal v12.8h,v2.8b,v8.8b - smlal v13.8h,v2.8b,v9.8b - smlal v14.8h,v2.8b,v10.8b - smlal v15.8h,v2.8b,v11.8b - ld1 {v2.4s},[x13],#16 // load ColumnSumBuffer[0] sadalp v24.4s,v12.8h sadalp v25.4s,v13.8h sadalp v26.4s,v14.8h sadalp v27.4s,v15.8h + smull v12.8h,v2.8b,v8.8b + smull v13.8h,v2.8b,v9.8b + smull v14.8h,v2.8b,v10.8b + smull v15.8h,v2.8b,v11.8b + sadalp v24.4s,v12.8h + sadalp v25.4s,v13.8h + sadalp v26.4s,v14.8h + sadalp v27.4s,v15.8h + ld1 {v2.4s},[x13],#16 // load ColumnSumBuffer[0] smull v12.8h,v1.8b,v4.8b smull v13.8h,v1.8b,v5.8b smull v14.8h,v1.8b,v6.8b smull v15.8h,v1.8b,v7.8b - smlal v12.8h,v3.8b,v8.8b - smlal v13.8h,v3.8b,v9.8b - smlal v14.8h,v3.8b,v10.8b - smlal v15.8h,v3.8b,v11.8b + sadalp v28.4s,v12.8h + sadalp v29.4s,v13.8h + sadalp v30.4s,v14.8h + sadalp v31.4s,v15.8h + smull v12.8h,v3.8b,v8.8b + smull v13.8h,v3.8b,v9.8b + smull v14.8h,v3.8b,v10.8b + smull v15.8h,v3.8b,v11.8b sadalp v28.4s,v12.8h sadalp v29.4s,v13.8h sadalp v30.4s,v14.8h @@ -336,53 +362,57 @@ M2_ComputeBlockLoop: smull v29.8h,v0.8b,v5.8b smull v30.8h,v0.8b,v6.8b smull v31.8h,v0.8b,v7.8b - cbz x3,M2_ComputeBlockLoopFinish - smlal v28.8h,v2.8b,v24.8b - smlal v29.8h,v2.8b,v25.8b - smlal v30.8h,v2.8b,v26.8b - smlal v31.8h,v2.8b,v27.8b - ldp d0,d2,[x0],#16 // Load A0 sadalp v16.4s,v28.8h sadalp v17.4s,v29.8h sadalp v18.4s,v30.8h sadalp v19.4s,v31.8h + smull v28.8h,v2.8b,v24.8b + smull v29.8h,v2.8b,v25.8b + smull v30.8h,v2.8b,v26.8b + smull v31.8h,v2.8b,v27.8b + sadalp v16.4s,v28.8h + sadalp v17.4s,v29.8h + sadalp v18.4s,v30.8h + sadalp v19.4s,v31.8h + cbz x3,M2_ComputeBlockLoopFinish + ldp d0,d2,[x0],#16 // Load A0 smull v28.8h,v1.8b,v4.8b smull v29.8h,v1.8b,v5.8b smull v30.8h,v1.8b,v6.8b smull v31.8h,v1.8b,v7.8b - smlal v28.8h,v3.8b,v24.8b + sadalp v20.4s,v28.8h + sadalp v21.4s,v29.8h + sadalp v22.4s,v30.8h + sadalp v23.4s,v31.8h + smull v28.8h,v3.8b,v24.8b + smull v29.8h,v3.8b,v25.8b + smull v30.8h,v3.8b,v26.8b + smull v31.8h,v3.8b,v27.8b + sadalp v20.4s,v28.8h + sadalp v21.4s,v29.8h + sadalp v22.4s,v30.8h + sadalp v23.4s,v31.8h ldp d4,d24,[x1],#16 // B - smlal v29.8h,v3.8b,v25.8b ldp d5,d25,[x1],#16 - smlal v30.8h,v3.8b,v26.8b ldp d6,d26,[x1],#16 - smlal v31.8h,v3.8b,v27.8b ldp d7,d27,[x1],#16 - sadalp v20.4s,v28.8h ldp d1,d3,[x9],#16 // Load A1 - sadalp v21.4s,v29.8h - sadalp v22.4s,v30.8h - sadalp v23.4s,v31.8h b M2_ComputeBlockLoop M2_ComputeBlockLoopFinish: ld1 {v0.4s},[x13],#16 // load ColumnSumBuffer[0] - smlal v28.8h,v2.8b,v24.8b - smlal v29.8h,v2.8b,v25.8b - smlal v30.8h,v2.8b,v26.8b - smlal v31.8h,v2.8b,v27.8b - sadalp v16.4s,v28.8h - sadalp v17.4s,v29.8h - sadalp v18.4s,v30.8h - sadalp v19.4s,v31.8h smull v28.8h,v1.8b,v4.8b smull v29.8h,v1.8b,v5.8b smull v30.8h,v1.8b,v6.8b smull v31.8h,v1.8b,v7.8b - smlal v28.8h,v3.8b,v24.8b - smlal v29.8h,v3.8b,v25.8b - smlal v30.8h,v3.8b,v26.8b - smlal v31.8h,v3.8b,v27.8b + sadalp v20.4s,v28.8h + sadalp v21.4s,v29.8h + sadalp v22.4s,v30.8h + sadalp v23.4s,v31.8h + smull v28.8h,v3.8b,v24.8b + smull v29.8h,v3.8b,v25.8b + smull v30.8h,v3.8b,v26.8b + smull v31.8h,v3.8b,v27.8b sadalp v20.4s,v28.8h sadalp v21.4s,v29.8h sadalp v22.4s,v30.8h @@ -470,32 +500,34 @@ M1_ComputeBlockLoop: sub x3,x3,#1 smull v20.8h,v0.8b,v4.8b smull v21.8h,v0.8b,v5.8b - cbz x3,M1_ComputeBlockLoopFinish smull v22.8h,v0.8b,v6.8b smull v23.8h,v0.8b,v7.8b - smlal v20.8h,v2.8b,v24.8b - ldp d4,d24,[x1],#16 // B - smlal v21.8h,v2.8b,v25.8b - ldp d5,d25,[x1],#16 - smlal v22.8h,v2.8b,v26.8b - ldp d6,d26,[x1],#16 - smlal v23.8h,v2.8b,v27.8b - ldp d0,d2,[x0],#16 // A0 sadalp v16.4s,v20.8h sadalp v17.4s,v21.8h - ldp d7,d27,[x1],#16 sadalp v18.4s,v22.8h sadalp v19.4s,v23.8h + cbz x3,M1_ComputeBlockLoopFinish + smull v20.8h,v2.8b,v24.8b + smull v21.8h,v2.8b,v25.8b + smull v22.8h,v2.8b,v26.8b + smull v23.8h,v2.8b,v27.8b + sadalp v16.4s,v20.8h + sadalp v17.4s,v21.8h + sadalp v18.4s,v22.8h + sadalp v19.4s,v23.8h + ldp d4,d24,[x1],#16 // B + ldp d5,d25,[x1],#16 + ldp d6,d26,[x1],#16 + ldp d7,d27,[x1],#16 + ldp d0,d2,[x0],#16 // A0 b M1_ComputeBlockLoop M1_ComputeBlockLoopFinish: ld1 {v4.4s},[x13],#16 // load ColumnSumBuffer[0] - smull v22.8h,v0.8b,v6.8b - smull v23.8h,v0.8b,v7.8b - smlal v20.8h,v2.8b,v24.8b - smlal v21.8h,v2.8b,v25.8b - smlal v22.8h,v2.8b,v26.8b - smlal v23.8h,v2.8b,v27.8b + smull v20.8h,v2.8b,v24.8b + smull v21.8h,v2.8b,v25.8b + smull v22.8h,v2.8b,v26.8b + smull v23.8h,v2.8b,v27.8b sadalp v16.4s,v20.8h sadalp v17.4s,v21.8h sadalp v18.4s,v22.8h diff --git a/onnxruntime/core/mlas/lib/arm64/SymQgemmS8KernelNeon.asm b/onnxruntime/core/mlas/lib/arm64/SymQgemmS8KernelNeon.asm index 4770b071dd84d..ce63415029c11 100644 --- a/onnxruntime/core/mlas/lib/arm64/SymQgemmS8KernelNeon.asm +++ b/onnxruntime/core/mlas/lib/arm64/SymQgemmS8KernelNeon.asm @@ -140,64 +140,82 @@ M4_ProcessNextColumnLoop add x11,x10,x7 // A3 M4_ComputeBlockLoop + // Reduce each 8-lane int8 product group into int32 before the next multiply. + // Two -128 * -128 products can overflow an int16 accumulator; see #31573. smull v12.8h,v0.8b,v4.8b smull v13.8h,v0.8b,v5.8b smull v14.8h,v0.8b,v6.8b smull v15.8h,v0.8b,v7.8b - smlal v12.8h,v2.8b,v8.8b - smlal v13.8h,v2.8b,v9.8b - smlal v14.8h,v2.8b,v10.8b - smlal v15.8h,v2.8b,v11.8b - ldp d0,d2,[x10],#16 // Load A2 sadalp v16.4s,v12.8h sadalp v17.4s,v13.8h sadalp v18.4s,v14.8h sadalp v19.4s,v15.8h + smull v12.8h,v2.8b,v8.8b + smull v13.8h,v2.8b,v9.8b + smull v14.8h,v2.8b,v10.8b + smull v15.8h,v2.8b,v11.8b + sadalp v16.4s,v12.8h + sadalp v17.4s,v13.8h + sadalp v18.4s,v14.8h + sadalp v19.4s,v15.8h + ldp d0,d2,[x10],#16 // Load A2 sub x3,x3,#1 smull v12.8h,v1.8b,v4.8b smull v13.8h,v1.8b,v5.8b smull v14.8h,v1.8b,v6.8b smull v15.8h,v1.8b,v7.8b - smlal v12.8h,v3.8b,v8.8b - smlal v13.8h,v3.8b,v9.8b - smlal v14.8h,v3.8b,v10.8b - smlal v15.8h,v3.8b,v11.8b - ldp d1,d3,[x11],#16 // Load A3 sadalp v20.4s,v12.8h sadalp v21.4s,v13.8h sadalp v22.4s,v14.8h sadalp v23.4s,v15.8h + smull v12.8h,v3.8b,v8.8b + smull v13.8h,v3.8b,v9.8b + smull v14.8h,v3.8b,v10.8b + smull v15.8h,v3.8b,v11.8b + sadalp v20.4s,v12.8h + sadalp v21.4s,v13.8h + sadalp v22.4s,v14.8h + sadalp v23.4s,v15.8h + ldp d1,d3,[x11],#16 // Load A3 cbz x3,M4_ComputeBlockLoopFinish smull v12.8h,v0.8b,v4.8b smull v13.8h,v0.8b,v5.8b smull v14.8h,v0.8b,v6.8b smull v15.8h,v0.8b,v7.8b - smlal v12.8h,v2.8b,v8.8b - smlal v13.8h,v2.8b,v9.8b - smlal v14.8h,v2.8b,v10.8b - smlal v15.8h,v2.8b,v11.8b - ldp d0,d2,[x0],#16 // Load A0 next iter sadalp v24.4s,v12.8h sadalp v25.4s,v13.8h sadalp v26.4s,v14.8h sadalp v27.4s,v15.8h + smull v12.8h,v2.8b,v8.8b + smull v13.8h,v2.8b,v9.8b + smull v14.8h,v2.8b,v10.8b + smull v15.8h,v2.8b,v11.8b + sadalp v24.4s,v12.8h + sadalp v25.4s,v13.8h + sadalp v26.4s,v14.8h + sadalp v27.4s,v15.8h + ldp d0,d2,[x0],#16 // Load A0 next iter smull v12.8h,v1.8b,v4.8b smull v13.8h,v1.8b,v5.8b smull v14.8h,v1.8b,v6.8b smull v15.8h,v1.8b,v7.8b - smlal v12.8h,v3.8b,v8.8b + sadalp v28.4s,v12.8h + sadalp v29.4s,v13.8h + sadalp v30.4s,v14.8h + sadalp v31.4s,v15.8h + smull v12.8h,v3.8b,v8.8b + smull v13.8h,v3.8b,v9.8b + smull v14.8h,v3.8b,v10.8b + smull v15.8h,v3.8b,v11.8b + sadalp v28.4s,v12.8h + sadalp v29.4s,v13.8h + sadalp v30.4s,v14.8h + sadalp v31.4s,v15.8h ldp d4,d8,[x1],#64 // B - smlal v13.8h,v3.8b,v9.8b ldp d5,d9,[x1,#-48] - smlal v14.8h,v3.8b,v10.8b ldp d6,d10,[x1,#-32] - smlal v15.8h,v3.8b,v11.8b ldp d7,d11,[x1,#-16] - sadalp v28.4s,v12.8h ldp d1,d3,[x9],#16 // Load A1 next iter - sadalp v29.4s,v13.8h - sadalp v30.4s,v14.8h - sadalp v31.4s,v15.8h b M4_ComputeBlockLoop M4_ComputeBlockLoopFinish @@ -205,23 +223,31 @@ M4_ComputeBlockLoopFinish smull v13.8h,v0.8b,v5.8b smull v14.8h,v0.8b,v6.8b smull v15.8h,v0.8b,v7.8b - smlal v12.8h,v2.8b,v8.8b - smlal v13.8h,v2.8b,v9.8b - smlal v14.8h,v2.8b,v10.8b - smlal v15.8h,v2.8b,v11.8b - ld1 {v2.4s},[x13],#16 // load ColumnSumBuffer[0] sadalp v24.4s,v12.8h sadalp v25.4s,v13.8h sadalp v26.4s,v14.8h sadalp v27.4s,v15.8h + smull v12.8h,v2.8b,v8.8b + smull v13.8h,v2.8b,v9.8b + smull v14.8h,v2.8b,v10.8b + smull v15.8h,v2.8b,v11.8b + sadalp v24.4s,v12.8h + sadalp v25.4s,v13.8h + sadalp v26.4s,v14.8h + sadalp v27.4s,v15.8h + ld1 {v2.4s},[x13],#16 // load ColumnSumBuffer[0] smull v12.8h,v1.8b,v4.8b smull v13.8h,v1.8b,v5.8b smull v14.8h,v1.8b,v6.8b smull v15.8h,v1.8b,v7.8b - smlal v12.8h,v3.8b,v8.8b - smlal v13.8h,v3.8b,v9.8b - smlal v14.8h,v3.8b,v10.8b - smlal v15.8h,v3.8b,v11.8b + sadalp v28.4s,v12.8h + sadalp v29.4s,v13.8h + sadalp v30.4s,v14.8h + sadalp v31.4s,v15.8h + smull v12.8h,v3.8b,v8.8b + smull v13.8h,v3.8b,v9.8b + smull v14.8h,v3.8b,v10.8b + smull v15.8h,v3.8b,v11.8b sadalp v28.4s,v12.8h sadalp v29.4s,v13.8h sadalp v30.4s,v14.8h @@ -336,53 +362,57 @@ M2_ComputeBlockLoop smull v29.8h,v0.8b,v5.8b smull v30.8h,v0.8b,v6.8b smull v31.8h,v0.8b,v7.8b - cbz x3,M2_ComputeBlockLoopFinish - smlal v28.8h,v2.8b,v24.8b - smlal v29.8h,v2.8b,v25.8b - smlal v30.8h,v2.8b,v26.8b - smlal v31.8h,v2.8b,v27.8b - ldp d0,d2,[x0],#16 // Load A0 sadalp v16.4s,v28.8h sadalp v17.4s,v29.8h sadalp v18.4s,v30.8h sadalp v19.4s,v31.8h + smull v28.8h,v2.8b,v24.8b + smull v29.8h,v2.8b,v25.8b + smull v30.8h,v2.8b,v26.8b + smull v31.8h,v2.8b,v27.8b + sadalp v16.4s,v28.8h + sadalp v17.4s,v29.8h + sadalp v18.4s,v30.8h + sadalp v19.4s,v31.8h + cbz x3,M2_ComputeBlockLoopFinish + ldp d0,d2,[x0],#16 // Load A0 smull v28.8h,v1.8b,v4.8b smull v29.8h,v1.8b,v5.8b smull v30.8h,v1.8b,v6.8b smull v31.8h,v1.8b,v7.8b - smlal v28.8h,v3.8b,v24.8b + sadalp v20.4s,v28.8h + sadalp v21.4s,v29.8h + sadalp v22.4s,v30.8h + sadalp v23.4s,v31.8h + smull v28.8h,v3.8b,v24.8b + smull v29.8h,v3.8b,v25.8b + smull v30.8h,v3.8b,v26.8b + smull v31.8h,v3.8b,v27.8b + sadalp v20.4s,v28.8h + sadalp v21.4s,v29.8h + sadalp v22.4s,v30.8h + sadalp v23.4s,v31.8h ldp d4,d24,[x1],#16 // B - smlal v29.8h,v3.8b,v25.8b ldp d5,d25,[x1],#16 - smlal v30.8h,v3.8b,v26.8b ldp d6,d26,[x1],#16 - smlal v31.8h,v3.8b,v27.8b ldp d7,d27,[x1],#16 - sadalp v20.4s,v28.8h ldp d1,d3,[x9],#16 // Load A1 - sadalp v21.4s,v29.8h - sadalp v22.4s,v30.8h - sadalp v23.4s,v31.8h b M2_ComputeBlockLoop M2_ComputeBlockLoopFinish ld1 {v0.4s},[x13],#16 // load ColumnSumBuffer[0] - smlal v28.8h,v2.8b,v24.8b - smlal v29.8h,v2.8b,v25.8b - smlal v30.8h,v2.8b,v26.8b - smlal v31.8h,v2.8b,v27.8b - sadalp v16.4s,v28.8h - sadalp v17.4s,v29.8h - sadalp v18.4s,v30.8h - sadalp v19.4s,v31.8h smull v28.8h,v1.8b,v4.8b smull v29.8h,v1.8b,v5.8b smull v30.8h,v1.8b,v6.8b smull v31.8h,v1.8b,v7.8b - smlal v28.8h,v3.8b,v24.8b - smlal v29.8h,v3.8b,v25.8b - smlal v30.8h,v3.8b,v26.8b - smlal v31.8h,v3.8b,v27.8b + sadalp v20.4s,v28.8h + sadalp v21.4s,v29.8h + sadalp v22.4s,v30.8h + sadalp v23.4s,v31.8h + smull v28.8h,v3.8b,v24.8b + smull v29.8h,v3.8b,v25.8b + smull v30.8h,v3.8b,v26.8b + smull v31.8h,v3.8b,v27.8b sadalp v20.4s,v28.8h sadalp v21.4s,v29.8h sadalp v22.4s,v30.8h @@ -470,32 +500,34 @@ M1_ComputeBlockLoop sub x3,x3,#1 smull v20.8h,v0.8b,v4.8b smull v21.8h,v0.8b,v5.8b - cbz x3,M1_ComputeBlockLoopFinish smull v22.8h,v0.8b,v6.8b smull v23.8h,v0.8b,v7.8b - smlal v20.8h,v2.8b,v24.8b - ldp d4,d24,[x1],#16 // B - smlal v21.8h,v2.8b,v25.8b - ldp d5,d25,[x1],#16 - smlal v22.8h,v2.8b,v26.8b - ldp d6,d26,[x1],#16 - smlal v23.8h,v2.8b,v27.8b - ldp d0,d2,[x0],#16 // A0 sadalp v16.4s,v20.8h sadalp v17.4s,v21.8h - ldp d7,d27,[x1],#16 sadalp v18.4s,v22.8h sadalp v19.4s,v23.8h + cbz x3,M1_ComputeBlockLoopFinish + smull v20.8h,v2.8b,v24.8b + smull v21.8h,v2.8b,v25.8b + smull v22.8h,v2.8b,v26.8b + smull v23.8h,v2.8b,v27.8b + sadalp v16.4s,v20.8h + sadalp v17.4s,v21.8h + sadalp v18.4s,v22.8h + sadalp v19.4s,v23.8h + ldp d4,d24,[x1],#16 // B + ldp d5,d25,[x1],#16 + ldp d6,d26,[x1],#16 + ldp d7,d27,[x1],#16 + ldp d0,d2,[x0],#16 // A0 b M1_ComputeBlockLoop M1_ComputeBlockLoopFinish ld1 {v4.4s},[x13],#16 // load ColumnSumBuffer[0] - smull v22.8h,v0.8b,v6.8b - smull v23.8h,v0.8b,v7.8b - smlal v20.8h,v2.8b,v24.8b - smlal v21.8h,v2.8b,v25.8b - smlal v22.8h,v2.8b,v26.8b - smlal v23.8h,v2.8b,v27.8b + smull v20.8h,v2.8b,v24.8b + smull v21.8h,v2.8b,v25.8b + smull v22.8h,v2.8b,v26.8b + smull v23.8h,v2.8b,v27.8b sadalp v16.4s,v20.8h sadalp v17.4s,v21.8h sadalp v18.4s,v22.8h diff --git a/onnxruntime/core/optimizer/fuse_initializers_transformer.cc b/onnxruntime/core/optimizer/fuse_initializers_transformer.cc index 9a24687d9a4d9..6af10a23eafbf 100644 --- a/onnxruntime/core/optimizer/fuse_initializers_transformer.cc +++ b/onnxruntime/core/optimizer/fuse_initializers_transformer.cc @@ -57,8 +57,8 @@ static bool IsNodeValidForFusion(const Graph& graph, const Node& node, const onnxruntime::MLDataType tensor_type, const onnxruntime::MLDataType output_type) { - // Node must have initialized tensor - if (!(graph.IsInitializedTensor(node.InputDefs()[0]->Name()))) return false; + // Node input must be a constant initializer. Initializers that are also graph inputs can be overridden. + if (graph_utils::GetConstantInitializer(graph, node.InputDefs()[0]->Name()) == nullptr) return false; // Initialzed tensor must be of tensor_type if (!(DataTypeImpl::TypeFromProto(*(node.InputDefs()[0]->TypeAsProto())) == tensor_type)) return false; diff --git a/onnxruntime/core/optimizer/graph_transformer_utils.cc b/onnxruntime/core/optimizer/graph_transformer_utils.cc index 4788d132a8473..038b0928acba6 100644 --- a/onnxruntime/core/optimizer/graph_transformer_utils.cc +++ b/onnxruntime/core/optimizer/graph_transformer_utils.cc @@ -349,6 +349,11 @@ InlinedVector> GenerateTransformers( onnxruntime::kAclExecutionProvider, onnxruntime::kCudaExecutionProvider, onnxruntime::kDmlExecutionProvider}; + const InlinedHashSet cpu_acl_cuda_dml_webgpu_eps = {onnxruntime::kCpuExecutionProvider, + onnxruntime::kAclExecutionProvider, + onnxruntime::kCudaExecutionProvider, + onnxruntime::kDmlExecutionProvider, + onnxruntime::kWebGpuExecutionProvider}; const InlinedHashSet cpu_acl_cuda_dml_js_webgpu_eps = {onnxruntime::kCpuExecutionProvider, onnxruntime::kAclExecutionProvider, onnxruntime::kCudaExecutionProvider, @@ -401,7 +406,7 @@ InlinedVector> GenerateTransformers( transformers.emplace_back(std::make_unique(cpu_acl_js_webgpu_eps)); - transformers.emplace_back(std::make_unique(cpu_acl_cuda_dml_eps, level)); + transformers.emplace_back(std::make_unique(cpu_acl_cuda_dml_webgpu_eps, level)); transformers.emplace_back(std::make_unique(cpu_acl_cuda_dml_eps, level)); transformers.emplace_back(std::make_unique(cpu_cuda_eps)); transformers.emplace_back(std::make_unique(cpu_acl_cuda_dml_eps)); @@ -409,7 +414,7 @@ InlinedVector> GenerateTransformers( transformers.emplace_back(std::make_unique(cpu_cuda_eps)); transformers.emplace_back(std::make_unique(cpu_cuda_eps)); transformers.emplace_back(std::make_unique(cpu_cuda_dml_eps)); - transformers.emplace_back(std::make_unique(cpu_acl_cuda_dml_eps)); + transformers.emplace_back(std::make_unique(cpu_acl_cuda_dml_webgpu_eps)); transformers.emplace_back(std::make_unique(cuda_eps)); // Run MatMulAddFusion again after *AttentionFusion transforms with `preserve_attention_pattern = false`, // to cleanup the remaining MatMul-Add that were part of the attention pattern but not detected or fused. diff --git a/onnxruntime/core/optimizer/initializer.cc b/onnxruntime/core/optimizer/initializer.cc index 05523e82a3c4e..369d8bbe8723c 100644 --- a/onnxruntime/core/optimizer/initializer.cc +++ b/onnxruntime/core/optimizer/initializer.cc @@ -494,8 +494,9 @@ struct ScaleByAxis { void Initializer::scale_by_axis(const Initializer& scalers, int axis, bool column_major) { ORT_ENFORCE(axis >= 0, "Axis must be non-negative"); - const size_t block_size = narrow(data_->Shape().SizeFromDimension(gsl::narrow_cast(axis))); - const size_t num_blocks = size() / block_size; + const auto axis_value = gsl::narrow_cast(axis); + const size_t block_size = narrow(data_->Shape().SizeFromDimension(axis_value)); + const size_t num_blocks = narrow(data_->Shape().SizeToDimension(axis_value)); ORT_ENFORCE(scalers.size() == 1 || (column_major ? scalers.size() == block_size : scalers.size() == num_blocks), "Invalid other(scalers) size"); diff --git a/onnxruntime/core/providers/cpu/llm/tensorscatter.cc b/onnxruntime/core/providers/cpu/llm/tensorscatter.cc index a90c596ecf7ec..4f0262db06798 100644 --- a/onnxruntime/core/providers/cpu/llm/tensorscatter.cc +++ b/onnxruntime/core/providers/cpu/llm/tensorscatter.cc @@ -75,6 +75,18 @@ Status TensorScatter::Compute(OpKernelContext* context) const { write_indices = write_indices_tensor->Data(); } + if (write_indices != nullptr) { + for (int64_t batch_idx = 0; batch_idx < batch_size; ++batch_idx) { + const int64_t wi = write_indices[batch_idx]; + ORT_ENFORCE(wi >= 0, "TensorScatter: write_indices[", batch_idx, "] = ", wi, " is negative"); + if (!circular_) { + ORT_ENFORCE(wi <= max_sequence_length - sequence_length, + "TensorScatter linear mode: write_indices[", batch_idx, "] + sequence_length (", + wi, " + ", sequence_length, ") exceeds max_sequence_length (", max_sequence_length, ")"); + } + } + } + // Allocate output with the same shape as past_cache. Tensor* present_cache = context->Output(0, cache_shape); @@ -83,12 +95,16 @@ Status TensorScatter::Compute(OpKernelContext* context) const { const size_t total_bytes = SafeInt(cache_shape.Size()) * element_size; const auto* src_raw = past_cache->DataRaw(); auto* dst_raw = present_cache->MutableDataRaw(); - if (dst_raw != src_raw) { + if (dst_raw != src_raw && total_bytes > 0) { LOGS(context->Logger(), WARNING) << "TensorScatter: in-place optimization not activated, copying past_cache to present_cache (" << total_bytes << " bytes)"; memcpy(dst_raw, src_raw, total_bytes); } + if (sequence_length == 0) { + return Status::OK(); + } + // Step 2: Scatter the update into present_cache. // // Layout: (batch_size, D1, ..., D_{axis-1}, max_seq_len, D_{axis+1}, ..., D_{n-1}) @@ -126,7 +142,6 @@ Status TensorScatter::Compute(OpKernelContext* context) const { for (int64_t p = 0; p < prefix_count; ++p) { int64_t batch_idx = p / prefix_stride_for_batch; int64_t wi = (write_indices != nullptr) ? write_indices[batch_idx] : 0; - ORT_ENFORCE(wi >= 0, "TensorScatter: write_indices[", batch_idx, "] = ", wi, " is negative"); ptrdiff_t update_offset = static_cast(SafeInt(p) * update_axis_stride); ptrdiff_t cache_offset = static_cast(SafeInt(p) * cache_axis_stride); @@ -134,17 +149,16 @@ Status TensorScatter::Compute(OpKernelContext* context) const { uint8_t* cache_base = dst_bytes + cache_offset; if (!circular_) { - ORT_ENFORCE(wi + sequence_length <= max_sequence_length, - "TensorScatter linear mode: write_indices[", batch_idx, "] + sequence_length (", - wi, " + ", sequence_length, ") exceeds max_sequence_length (", max_sequence_length, ")"); // Single contiguous memcpy for the whole slice. ptrdiff_t wi_offset = static_cast(SafeInt(wi) * suffix_bytes); size_t copy_len = SafeInt(sequence_length) * suffix_bytes; memcpy(cache_base + wi_offset, update_base, copy_len); } else { // Circular: each sequence position wraps independently. + const int64_t wi_mod = wi % max_sequence_length; + const int64_t distance_to_end = max_sequence_length - wi_mod; for (int64_t s = 0; s < sequence_length; ++s) { - int64_t cache_pos = (wi + s) % max_sequence_length; + const int64_t cache_pos = s >= distance_to_end ? s - distance_to_end : wi_mod + s; ptrdiff_t dst_off = static_cast(SafeInt(cache_pos) * suffix_bytes); ptrdiff_t src_off = static_cast(SafeInt(s) * suffix_bytes); memcpy(cache_base + dst_off, update_base + src_off, suffix_bytes); diff --git a/onnxruntime/core/providers/cpu/ml/normalizer.cc b/onnxruntime/core/providers/cpu/ml/normalizer.cc index 68ca49dee1958..76178341e60d9 100644 --- a/onnxruntime/core/providers/cpu/ml/normalizer.cc +++ b/onnxruntime/core/providers/cpu/ml/normalizer.cc @@ -128,6 +128,11 @@ Status Normalizer::Normalize(OpKernelContext* context) const { } const auto& x_dims = x_shape.GetDims(); + if (x_dims.empty()) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "Input to Normalizer must have rank 1 or 2. Got rank 0."); + } + int64_t num_batches = x_dims.size() == 1 ? 1 : x_dims[0]; int64_t batch_size = x_dims.size() == 1 ? x_dims[0] : x_dims[1]; diff --git a/onnxruntime/core/providers/cpu/nn/conv.cc b/onnxruntime/core/providers/cpu/nn/conv.cc index 07d3e2c4d55a2..bccf67fd68ab2 100644 --- a/onnxruntime/core/providers/cpu/nn/conv.cc +++ b/onnxruntime/core/providers/cpu/nn/conv.cc @@ -72,6 +72,9 @@ Status Conv::Compute(OpKernelContext* context) const { const int64_t M = W->Shape()[0]; ORT_RETURN_IF_ERROR(conv_attrs_.ValidateInputShape(X, W)); + ORT_RETURN_IF_NOT(B == nullptr || (B->Shape().NumDimensions() == 1 && B->Shape().Size() == M), + "Conv : bias must be a 1D tensor of size output_channels (", M, ")"); + TensorShapeVector kernel_shape; ORT_RETURN_IF_ERROR(conv_attrs_.ComputeKernelShape(W->Shape(), kernel_shape)); @@ -255,6 +258,9 @@ Status Conv::Compute(OpKernelContext* context) const { const int64_t C = channels_last_ ? X->Shape().GetDims().back() : X->Shape()[1]; const int64_t M = W->Shape()[0]; + ORT_RETURN_IF_NOT(B == nullptr || (B->Shape().NumDimensions() == 1 && B->Shape().Size() == M), + "Conv : bias must be a 1D tensor of size output_channels (", M, ")"); + TensorShapeVector kernel_shape; ORT_RETURN_IF_ERROR(conv_attrs_.ComputeKernelShape(W->Shape(), kernel_shape)); const size_t kernel_rank = kernel_shape.size(); diff --git a/onnxruntime/core/providers/cpu/reduction/reduction_ops.cc b/onnxruntime/core/providers/cpu/reduction/reduction_ops.cc index bb7a50bec6f99..f59ad5b0ba696 100644 --- a/onnxruntime/core/providers/cpu/reduction/reduction_ops.cc +++ b/onnxruntime/core/providers/cpu/reduction/reduction_ops.cc @@ -885,9 +885,12 @@ bool check_and_reduce_empty_set_input(OpKernelContext* ctx, const gsl::spanInput(1); - auto nDims = static_cast(axes_tensor->Shape()[0]); - const auto* data = axes_tensor->Data(); - input_axes.insert(input_axes.begin(), data, data + nDims); + if (axes_tensor != nullptr) { + ORT_ENFORCE(axes_tensor->Shape().NumDimensions() == 1, + "An axes tensor must be a vector tensor."); + const auto axes_data = axes_tensor->DataAsSpan(); + input_axes.assign(axes_data.begin(), axes_data.end()); + } } else { input_axes.resize(axes.size()); std::copy(axes.begin(), axes.end(), input_axes.begin()); @@ -925,7 +928,9 @@ inline void ApplyNoopEmptyAxesElementwise(OpKernelContext* ctx) { Tensor* Y = ctx->Output(0, shape); if constexpr (!ReduceAggTraits::kHasPreOp && !ReduceAggTraits::kHasPostOp) { - std::memcpy(Y->MutableDataRaw(), X->DataRaw(), X->SizeInBytes()); + if (X->SizeInBytes() > 0) { + std::memcpy(Y->MutableDataRaw(), X->DataRaw(), X->SizeInBytes()); + } } else { using Tin = typename AGG::input_type; @@ -968,10 +973,6 @@ template void CommonReduce1Loop(OpKernelContext* ctx, const gsl::span& axes_, int64_t keepdims_, bool noop_with_empty_axes) { - if (check_and_reduce_empty_set_input(ctx, axes_, keepdims_ != 0)) { - return; - } - TensorShapeVector tmp_axes; auto effective_axes = GetEffectiveAxes(ctx, axes_, tmp_axes); if (effective_axes.empty() && noop_with_empty_axes) { @@ -979,6 +980,10 @@ void CommonReduce1Loop(OpKernelContext* ctx, return; } + if (check_and_reduce_empty_set_input(ctx, axes_, keepdims_ != 0)) { + return; + } + FastReduceKind fast_kind; TensorShapeVector fast_shape; TensorShapeVector output_shape; @@ -1013,10 +1018,6 @@ template void CommonReduce2Loops(OpKernelContext* ctx, const gsl::span& axes_, int64_t keepdims_, bool noop_with_empty_axes) { - if (check_and_reduce_empty_set_input(ctx, axes_, keepdims_ != 0)) { - return; - } - TensorShapeVector tmp_axes; auto effective_axes = GetEffectiveAxes(ctx, axes_, tmp_axes); if (effective_axes.empty() && noop_with_empty_axes) { @@ -1024,6 +1025,10 @@ void CommonReduce2Loops(OpKernelContext* ctx, return; } + if (check_and_reduce_empty_set_input(ctx, axes_, keepdims_ != 0)) { + return; + } + FastReduceKind fast_kind; TensorShapeVector fast_shape, output_shape, fast_axes; if (CommonFastReduce(ctx, axes_, keepdims_, noop_with_empty_axes, diff --git a/onnxruntime/core/providers/cpu/tensor/scatter_nd.cc b/onnxruntime/core/providers/cpu/tensor/scatter_nd.cc index a7e2f19bfe147..b80f37f47871d 100644 --- a/onnxruntime/core/providers/cpu/tensor/scatter_nd.cc +++ b/onnxruntime/core/providers/cpu/tensor/scatter_nd.cc @@ -61,11 +61,13 @@ struct Prepare { TData* output_base; uint64_t element_to_copy; std::vector element_offsets; + bool serialize_updates; Prepare() : input_base(nullptr), output_base(nullptr), element_to_copy(0), - element_offsets(0) {} + element_offsets(0), + serialize_updates(false) {} }; // struct Prepare template @@ -89,7 +91,7 @@ Status PrepareForCompute(OpKernelContext* context, Prepare& p) { auto last_indice_dimension = indice_shape[indice_shape.NumDimensions() - 1]; // Re-use input for output. If input/output Tensor* are the same, do not copy. - if (src_base != dst_base) { + if (src_base != dst_base && input_tensor->Shape().Size() > 0) { if (is_string_type) { const auto* str_begin = input_tensor->Data(); const std::string* str_end = str_begin + input_shape.Size(); @@ -108,8 +110,9 @@ Status PrepareForCompute(OpKernelContext* context, Prepare& p) { } p.element_to_copy = input_shape.SizeFromDimension(onnxruntime::narrow(last_indice_dimension)); + p.serialize_updates = last_indice_dimension == 0; const int64_t* indice_offset = indice_tensor->Data(); - auto offset_count = indice_shape.Size() / last_indice_dimension; // Times to copy + auto offset_count = indice_shape.SizeToDimension(indice_shape.NumDimensions() - 1); // Times to copy p.element_offsets.assign(onnxruntime::narrow(offset_count), 0LL); p.input_base = update_tensor->Data(); @@ -302,6 +305,12 @@ struct ScatterNDDispatchTarget { Status operator()(OpKernelContext* context, concurrency::ThreadPool* tp, ScatterND::Reduction reduction) const { Prepare prepare; ORT_RETURN_IF_ERROR(PrepareForCompute(context, prepare)); + if (prepare.element_to_copy == 0 || prepare.element_offsets.empty()) { + return Status::OK(); + } + if (prepare.serialize_updates) { + tp = nullptr; + } auto lambda = [&](ptrdiff_t i) { switch (reduction) { diff --git a/onnxruntime/core/providers/cpu/tensor/split.h b/onnxruntime/core/providers/cpu/tensor/split.h index cf22ff8106f0d..02e425475228f 100644 --- a/onnxruntime/core/providers/cpu/tensor/split.h +++ b/onnxruntime/core/providers/cpu/tensor/split.h @@ -61,11 +61,23 @@ class SplitBase { } split_sizes = std::vector(static_cast(num_outputs), split_dim_size / num_outputs); } else { - int64_t split_size_sum = split_size_sum_; - if (split_size_sum == -1) { - split_size_sum = std::accumulate(split_sizes.cbegin(), split_sizes.cend(), 0LL); + int64_t remaining_split_size = split_dim_size; + for (int64_t s : split_sizes) { + if (s < 0) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "Invalid negative value in 'split'. All split sizes must be >= 0."); + } + if (s > remaining_split_size) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "Invalid value in 'split'. Split size ", s, + " exceeds the remaining size of the selected axis, ", remaining_split_size, "."); + } + remaining_split_size -= s; } - if (split_sizes.size() != static_cast(num_outputs) || split_size_sum != split_dim_size) + + const int64_t split_size_sum = + split_size_sum_ == -1 ? split_dim_size - remaining_split_size : split_size_sum_; + if (split_sizes.size() != static_cast(num_outputs) || remaining_split_size != 0) return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Cannot split using values in 'split' attribute. Axis=", axis_, " Input shape=", input_shape, @@ -86,9 +98,9 @@ class SplitBase { if (num_inputs == 1) { // optional if (info.GetAttrs("split", split_sizes_).IsOK()) { - split_size_sum_ = std::accumulate(split_sizes_.cbegin(), split_sizes_.cend(), 0LL); ORT_ENFORCE(std::all_of(split_sizes_.cbegin(), split_sizes_.cend(), [](int64_t value) { return value >= 0; }), "Invalid value in 'split' attribute. All values must be > 0"); + split_size_sum_ = std::accumulate(split_sizes_.cbegin(), split_sizes_.cend(), SafeInt{0}); } } diff --git a/onnxruntime/core/providers/cuda/tensor/scatter_nd.cc b/onnxruntime/core/providers/cuda/tensor/scatter_nd.cc index e6359cc048048..530160b94f43c 100644 --- a/onnxruntime/core/providers/cuda/tensor/scatter_nd.cc +++ b/onnxruntime/core/providers/cuda/tensor/scatter_nd.cc @@ -90,18 +90,21 @@ Status ScatterNDDisjointAndNoReduction::ComputeInternal(OpKernelContext* context const void* input_data = input_tensor->DataRaw(); void* output_data = output_tensor->MutableDataRaw(); - if (input_data != output_data) { + if (input_data != output_data && input_tensor->SizeInBytes() > 0) { // TODO: Run benchmarks to determine if a dedicated kernel doing data copy will be faster than invoking cudaMemcpy ? CUDA_RETURN_IF_ERROR( cudaMemcpyAsync(output_data, input_data, input_tensor->SizeInBytes(), cudaMemcpyDeviceToDevice, Stream(context))); } + const auto num_indices = indices_shape.SizeToDimension(indices_shape.NumDimensions() - 1); + auto last_index_dimension = indices_shape[indices_shape.NumDimensions() - 1]; + const auto num_update_elements = input_shape.SizeFromDimension(last_index_dimension); + // Bail out early - if (indices_shape.Size() == 0) { + if (num_indices == 0 || num_update_elements == 0) { return Status::OK(); } - auto last_index_dimension = indices_shape[indices_shape.NumDimensions() - 1]; size_t element_size = input_tensor->DataType()->Size(); // We need element counts for each dimension and the input dim value for each dimension @@ -118,12 +121,12 @@ Status ScatterNDDisjointAndNoReduction::ComputeInternal(OpKernelContext* context Stream(context), output_data, element_size, - indices_shape.Size() / static_cast(last_index_dimension), + onnxruntime::narrow(num_indices), indices_tensor->Data(), // only int64_t is supported for indices as per the onnx spec last_index_dimension, element_counts_and_input_dims, updates_tensor->DataRaw(), - input_shape.SizeFromDimension(last_index_dimension))); + onnxruntime::narrow(num_update_elements))); return Status::OK(); } @@ -145,7 +148,7 @@ Status ScatterNDWithAtomicReduction::ComputeInternal(OpKernelContext* context) c const void* input_data = input_tensor->DataRaw(); void* output_data = output_tensor->MutableDataRaw(); - if (input_data != output_data) { + if (input_data != output_data && input_tensor->SizeInBytes() > 0) { // TODO: Run benchmarks to determine if a dedicated kernel doing data copy will // be faster than invoking cudaMemcpy ? CUDA_RETURN_IF_ERROR( @@ -153,12 +156,15 @@ Status ScatterNDWithAtomicReduction::ComputeInternal(OpKernelContext* context) c cudaMemcpyDeviceToDevice, Stream(context))); } + const auto num_indices = indices_shape.SizeToDimension(indices_shape.NumDimensions() - 1); + auto last_index_dimension = indices_shape[indices_shape.NumDimensions() - 1]; + const auto num_update_elements = input_shape.SizeFromDimension(last_index_dimension); + // Bail out early - if (indices_shape.Size() == 0) { + if (num_indices == 0 || num_update_elements == 0) { return Status::OK(); } - auto last_index_dimension = indices_shape[indices_shape.NumDimensions() - 1]; ElementCountsAndInputDimsSpanOrGpu element_counts_and_input_dims; CudaAsyncBuffer element_counts_and_input_dims_gpu(this); ORT_RETURN_IF_ERROR(InitializeElementCountsAndInputDimsSpanOrGpu(last_index_dimension, input_shape, @@ -173,12 +179,12 @@ Status ScatterNDWithAtomicReduction::ComputeInternal(OpKernelContext* context) c Stream(context), output_data, element_size, - indices_shape.Size() / static_cast(last_index_dimension), + onnxruntime::narrow(num_indices), indices_tensor->Data(), // only int64_t is supported for indices as per the onnx spec last_index_dimension, element_counts_and_input_dims, updates_tensor->DataRaw(), - input_shape.SizeFromDimension(last_index_dimension))); + onnxruntime::narrow(num_update_elements))); } break; case ScatterNDReduction::Add: case ScatterNDReduction::Min: @@ -189,12 +195,12 @@ Status ScatterNDWithAtomicReduction::ComputeInternal(OpKernelContext* context) c Stream(context), output_data, element_type, - indices_shape.Size() / static_cast(last_index_dimension), + onnxruntime::narrow(num_indices), indices_tensor->Data(), // only int64_t is supported for indices as per the onnx spec last_index_dimension, element_counts_and_input_dims, updates_tensor->DataRaw(), - input_shape.SizeFromDimension(last_index_dimension), + onnxruntime::narrow(num_update_elements), reduction_)); } break; default: diff --git a/onnxruntime/core/providers/cuda/tensor/split.cc b/onnxruntime/core/providers/cuda/tensor/split.cc index 06b0c7e50f919..54596b7d02a4d 100644 --- a/onnxruntime/core/providers/cuda/tensor/split.cc +++ b/onnxruntime/core/providers/cuda/tensor/split.cc @@ -85,11 +85,22 @@ Status SplitKernel::PrepareForComputeLocal(const TensorShape& input_shape, } split_sizes = std::vector(static_cast(num_outputs), split_dim_size / num_outputs); } else { - int64_t split_size_sum = split_size_sum_; - if (split_size_sum == -1) { - split_size_sum = std::accumulate(split_sizes.cbegin(), split_sizes.cend(), 0LL); + int64_t remaining_split_size = split_dim_size; + for (int64_t s : split_sizes) { + if (s < 0) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "Invalid negative value in 'split'. All split sizes must be >= 0."); + } + if (s > remaining_split_size) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "Invalid value in 'split'. Split size ", s, + " exceeds the remaining size of the selected axis, ", remaining_split_size, "."); + } + remaining_split_size -= s; } - if (split_sizes.size() != static_cast(num_outputs) || split_size_sum != split_dim_size) { + + const int64_t split_size_sum = split_dim_size - remaining_split_size; + if (split_sizes.size() != static_cast(num_outputs) || remaining_split_size != 0) { return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Cannot split using values in 'split' attribute. Axis=", axis_, " Input shape=", input_shape, diff --git a/onnxruntime/python/onnxruntime_pybind_state.cc b/onnxruntime/python/onnxruntime_pybind_state.cc index 683975f64b2d8..404604bb1fd69 100644 --- a/onnxruntime/python/onnxruntime_pybind_state.cc +++ b/onnxruntime/python/onnxruntime_pybind_state.cc @@ -4,6 +4,7 @@ #include #include +#include #include "python/onnxruntime_pybind_exceptions.h" #include "python/onnxruntime_pybind_mlvalue.h" #include "python/onnxruntime_pybind_model_compiler.h" @@ -139,7 +140,122 @@ AdaptedProviderOptions AdaptProviderOptionsForRegisteredPluginEp(const std::stri using PyCallback = std::function, py::object user_data, std::string)>; +struct PendingPythonReleases { + std::mutex mutex; + std::vector objects; + std::vector cleanup_thread_objects; + bool callback_scheduled = false; + bool cleanup_thread_running = false; +}; + +PendingPythonReleases& GetPendingPythonReleases() { + static auto* pending_releases = new PendingPythonReleases(); + return *pending_releases; +} + +int DrainPendingPythonReleases(void*) { + std::vector objects; + auto& pending_releases = GetPendingPythonReleases(); + { + std::lock_guard lock{pending_releases.mutex}; + objects.swap(pending_releases.objects); + pending_releases.callback_scheduled = false; + } + + for (PyObject* object : objects) { + Py_DECREF(object); + } + + return 0; +} + +void DrainCleanupThreadPythonReleases() { + for (;;) { + std::vector objects; + auto& pending_releases = GetPendingPythonReleases(); + { + std::lock_guard lock{pending_releases.mutex}; + if (pending_releases.cleanup_thread_objects.empty()) { + pending_releases.cleanup_thread_running = false; + return; + } + objects.swap(pending_releases.cleanup_thread_objects); + } + + { + py::gil_scoped_acquire acquire; + for (PyObject* object : objects) { + Py_DECREF(object); + } + } + } +} + +void ReleasePythonObjectOnCleanupThread(PyObject* object) noexcept { + auto& pending_releases = GetPendingPythonReleases(); + std::lock_guard lock{pending_releases.mutex}; + ORT_TRY { + pending_releases.cleanup_thread_objects.push_back(object); + } + ORT_CATCH(...) { + // Retaining the reference is safer than releasing it on the ORT worker. + return; + } + + if (!pending_releases.cleanup_thread_running) { + pending_releases.cleanup_thread_running = true; + std::thread* cleanup_thread = nullptr; + ORT_TRY { + cleanup_thread = new std::thread(DrainCleanupThreadPythonReleases); + } + ORT_CATCH(...) { + pending_releases.cleanup_thread_running = false; + return; + } + + ORT_TRY { + cleanup_thread->detach(); + delete cleanup_thread; + } + ORT_CATCH(...) { + // The valid thread remains joinable and will drain the queue. Leak its + // handle rather than destroying it and terminating the process. + } + } +} + +void DeferPythonRelease(PyObject* object) noexcept { + if (object == nullptr) { + return; + } + + bool release_on_cleanup_thread = false; + ORT_TRY { + auto& pending_releases = GetPendingPythonReleases(); + std::lock_guard lock{pending_releases.mutex}; + pending_releases.objects.push_back(object); + if (!pending_releases.callback_scheduled && + Py_AddPendingCall(DrainPendingPythonReleases, nullptr) != 0) { + pending_releases.objects.pop_back(); + release_on_cleanup_thread = true; + } else { + pending_releases.callback_scheduled = true; + } + } + ORT_CATCH(...) { + release_on_cleanup_thread = true; + } + + if (release_on_cleanup_thread) { + ReleasePythonObjectOnCleanupThread(object); + } +} + struct AsyncResource { + std::vector feed_objects; + py::object session; + py::object run_options; + std::vector feeds; std::vector feeds_raw; @@ -156,6 +272,7 @@ struct AsyncResource { py::object user_data; void ReserveFeeds(size_t sz) { + feed_objects.reserve(sz); feeds.reserve(sz); feeds_raw.reserve(sz); feed_names.reserve(sz); @@ -183,33 +300,44 @@ void AsyncCallback(void* user_data, OrtValue** outputs, size_t num_outputs, OrtS auto invoke_callback = [&]() { std::unique_ptr async_resource{reinterpret_cast(user_data)}; - Ort::Status status(ort_status); + PyObject* session = async_resource->session.release().ptr(); - // return on error - if (!status.IsOK()) { - async_resource->callback({}, async_resource->user_data, status.GetErrorMessage()); - return; - } - - std::vector rfetch; - rfetch.reserve(num_outputs); - size_t pos = 0; - for (size_t ith = 0; ith < num_outputs; ++ith) { - const auto& fet = *outputs[ith]; - if (fet.IsAllocated()) { - if (fet.IsTensor()) { - rfetch.push_back(AddTensorAsPyObj(fet, nullptr, nullptr)); - } else if (fet.IsSparseTensor()) { - rfetch.push_back(GetPyObjectFromSparseTensor(pos, fet, nullptr)); - } else { - rfetch.push_back(AddNonTensorAsPyObj(fet, nullptr, nullptr)); - } + try { + Ort::Status status(ort_status); + if (!status.IsOK()) { + async_resource->callback({}, async_resource->user_data, status.GetErrorMessage()); } else { - rfetch.push_back(py::none()); + std::vector rfetch; + rfetch.reserve(num_outputs); + size_t pos = 0; + for (size_t ith = 0; ith < num_outputs; ++ith) { + const auto& fet = *outputs[ith]; + if (fet.IsAllocated()) { + if (fet.IsTensor()) { + rfetch.push_back(AddTensorAsPyObj(fet, nullptr, nullptr)); + } else if (fet.IsSparseTensor()) { + rfetch.push_back(GetPyObjectFromSparseTensor(pos, fet, nullptr)); + } else { + rfetch.push_back(AddNonTensorAsPyObj(fet, nullptr, nullptr)); + } + } else { + rfetch.push_back(py::none()); + } + ++pos; + } + async_resource->callback(rfetch, async_resource->user_data, ""); } - ++pos; + } catch (py::error_already_set& ex) { + ex.discard_as_unraisable("onnxruntime.InferenceSession.run_async callback"); + } catch (const std::exception& ex) { + PyErr_SetString(PyExc_RuntimeError, ex.what()); + PyErr_WriteUnraisable(Py_None); + } catch (...) { + PyErr_SetString(PyExc_RuntimeError, "Unknown exception in run_async callback"); + PyErr_WriteUnraisable(Py_None); } - async_resource->callback(rfetch, async_resource->user_data, ""); + + DeferPythonRelease(session); }; if (PyGILState_Check()) { @@ -2884,30 +3012,38 @@ including arg name, arg type (contains both type and shape).)pbdoc") return result; }) .def("run_async", - [](PyInferenceSession* sess, + [](py::object session, const std::vector& output_names, const std::map& pyfeeds, PyCallback callback, py::object user_data = {}, - RunOptions* run_options = nullptr) + py::object run_options = py::none()) -> void { - if (run_options != nullptr && !run_options->active_adapters.empty()) { + auto* sess = session.cast(); + auto* run_options_ptr = run_options.is_none() + ? nullptr + : run_options.cast(); + if (run_options_ptr != nullptr && !run_options_ptr->active_adapters.empty()) { LOGS(*sess->GetSessionHandle()->GetLogger(), WARNING) << "run_async has active adapters specified, but won't have an effect"; } std::unique_ptr async_resource = std::make_unique(); + async_resource->session = std::move(session); + async_resource->run_options = std::move(run_options); async_resource->callback = callback; async_resource->user_data = user_data; // prepare feeds async_resource->ReserveFeeds(pyfeeds.size()); for (const auto& feed : pyfeeds) { if (!feed.second.is(py::none())) { + async_resource->feed_objects.push_back(feed.second); OrtValue ml_value; auto px = sess->GetSessionHandle()->GetModelInputs(); if (!px.first.IsOK() || !px.second) { throw std::runtime_error("Either failed to get model inputs from the session object or the input def list was null"); } - CreateGenericMLValue(px.second, GetAllocator(), feed.first, feed.second, &ml_value); + CreateGenericMLValue(px.second, GetAllocator(), feed.first, + async_resource->feed_objects.back(), &ml_value); ThrowIfPyErrOccured(); async_resource->feeds.push_back(ml_value); async_resource->feeds_raw.push_back(&async_resource->feeds.back()); @@ -2922,7 +3058,7 @@ including arg name, arg type (contains both type and shape).)pbdoc") async_resource->fetch_names_raw.push_back(async_resource->fetch_names.back().c_str()); async_resource->fetches_raw.push_back({}); } - const RunOptions* run_async_option = run_options ? run_options : &async_resource->default_run_option; + const RunOptions* run_async_option = run_options_ptr ? run_options_ptr : &async_resource->default_run_option; common::Status status = sess->GetSessionHandle()->RunAsync(run_async_option, gsl::span(async_resource->feed_names_raw.data(), async_resource->feed_names_raw.size()), gsl::span(async_resource->feeds_raw.data(), async_resource->feeds_raw.size()), diff --git a/onnxruntime/test/contrib_ops/activation_op_test.cc b/onnxruntime/test/contrib_ops/activation_op_test.cc index 061fffa572be2..f4223206cdc1d 100644 --- a/onnxruntime/test/contrib_ops/activation_op_test.cc +++ b/onnxruntime/test/contrib_ops/activation_op_test.cc @@ -60,6 +60,22 @@ TEST_F(ActivationOpTest, Gelu) { } #endif +TEST_F(ActivationOpTest, Gelu_half) { + const std::vector& X = input_values_fp16[0]; + std::vector Y; + Y.reserve(X.size()); + for (const MLFloat16& x_half : X) { + const float x = x_half.ToFloat(); + Y.push_back(MLFloat16(x * 0.5f * (1.0f + std::erf(x * static_cast(M_SQRT1_2))))); + } + + OpTester tester("Gelu", 1, onnxruntime::kMSDomain); + const std::vector dims{1, 1, static_cast(X.size())}; + tester.AddInput("X", dims, X); + tester.AddOutput("Y", dims, Y); + tester.Run(OpTester::ExpectResult::kExpectSuccess, "", {kTensorrtExecutionProvider}); +} + #if defined(USE_DNNL) std::vector expected_output_bfloat16(const std::vector& input_data) { std::vector output; diff --git a/onnxruntime/test/contrib_ops/qembed_layer_norm_op_test.cc b/onnxruntime/test/contrib_ops/qembed_layer_norm_op_test.cc index 7ffb99e488dd1..9093d98b83c94 100644 --- a/onnxruntime/test/contrib_ops/qembed_layer_norm_op_test.cc +++ b/onnxruntime/test/contrib_ops/qembed_layer_norm_op_test.cc @@ -14,7 +14,12 @@ namespace { template static void RunTest(const embedlayernorm::OpData& data, - float accuracy_threshold = 0.25f) { + float accuracy_threshold = 0.25f, + uint8_t segment_inputs_mask = 0x0F) { + constexpr uint8_t kSegmentIds = 1 << 0; + constexpr uint8_t kSegmentEmbedding = 1 << 1; + constexpr uint8_t kSegmentEmbeddingScale = 1 << 2; + constexpr uint8_t kSegmentEmbeddingZeroPoint = 1 << 3; ASSERT_TRUE(data.word_embedding_data.size() % data.hidden_size == 0); ASSERT_TRUE(data.position_embedding_data.size() % data.hidden_size == 0); ASSERT_TRUE(data.segment_embedding_data.size() % data.hidden_size == 0); @@ -64,7 +69,7 @@ static void RunTest(const embedlayernorm::OpData& data, // Operator inputs passed in at int32_t: tester.AddInput("input_ids", input_ids_dims, data.input_ids_data); - if (data.has_segment) { + if (data.has_segment && (segment_inputs_mask & kSegmentIds) != 0) { tester.AddInput("segment_ids", segment_ids_dims, data.segment_ids_data); } else { tester.AddOptionalInputEdge(); @@ -79,7 +84,7 @@ static void RunTest(const embedlayernorm::OpData& data, position_embedding_dims, position_embedding_data_quant, /*is_initializer=*/true); - if (data.has_segment) { + if (data.has_segment && (segment_inputs_mask & kSegmentEmbedding) != 0) { tester.AddInput("segment_embedding_data", segment_embedding_dims, segment_embedding_data_quant, @@ -111,7 +116,7 @@ static void RunTest(const embedlayernorm::OpData& data, /*dims=*/{}, {position_embedding_params.scale}, /*is_initializer=*/true); - if (data.has_segment) { + if (data.has_segment && (segment_inputs_mask & kSegmentEmbeddingScale) != 0) { tester.AddInput("segment_embedding_scale", /*dims=*/{}, {segment_embedding_params.scale}, @@ -137,7 +142,7 @@ static void RunTest(const embedlayernorm::OpData& data, /*dims=*/{}, {position_embedding_params.zero_point}, /*is_initializer=*/true); - if (data.has_segment) { + if (data.has_segment && (segment_inputs_mask & kSegmentEmbeddingZeroPoint) != 0) { tester.AddInput("segment_embedding_zero_point", /*dims=*/{}, {segment_embedding_params.zero_point}, @@ -165,7 +170,12 @@ static void RunTest(const embedlayernorm::OpData& data, // Attributes: tester.AddAttribute("epsilon", embedlayernorm::kEpsilon); - tester.Run(); + if (segment_inputs_mask == 0 || segment_inputs_mask == 0x0F) { + tester.Run(); + } else { + tester.Run(OpTester::ExpectResult::kExpectFailure, + "must either all be provided or all be omitted"); + } } } // namespace @@ -175,6 +185,15 @@ TEST(QEmbedLayerNormTest, EmbedLayerNormBatch1) { RunTest(embedlayernorm::EmbedLayerNormBatch1()); } +TEST(QEmbedLayerNormTest, PartialSegmentInputsRejected) { + const auto data = embedlayernorm::EmbedLayerNormBatch1(); + ASSERT_TRUE(data.has_segment); + + for (uint8_t segment_inputs_mask = 1; segment_inputs_mask < 0x0F; ++segment_inputs_mask) { + RunTest(data, 0.25f, segment_inputs_mask); + } +} + TEST(QEmbedLayerNormTest, EmbedLayerNormBatch1_Float16) { RunTest(embedlayernorm::EmbedLayerNormBatch1(), /*use_float16=*/true); RunTest(embedlayernorm::EmbedLayerNormBatch1(), /*use_float16=*/true); diff --git a/onnxruntime/test/framework/tensorutils_test.cc b/onnxruntime/test/framework/tensorutils_test.cc index 34d1a9d906563..d238627f8e6df 100644 --- a/onnxruntime/test/framework/tensorutils_test.cc +++ b/onnxruntime/test/framework/tensorutils_test.cc @@ -143,6 +143,50 @@ TEST(TensorProtoUtilsTest, ParseExternalDataInfoOffsetAndLength) { // TODO should ExternalDataInfo::Create() also reject negative offset values? } +TEST(TensorProtoUtilsTest, ExternalDataInfoRejectsDuplicateLocations) { + const std::string memory_tag = ToUTF8String(utils::kTensorProtoNativeEndianMemoryAddressTag); + const std::array, 3> locations{{ + {"data.bin", "data.bin"}, + {"data.bin", memory_tag}, + {memory_tag, "data.bin"}, + }}; + + for (const auto& [first_location, second_location] : locations) { + ONNX_NAMESPACE::TensorProto tensor_proto; + tensor_proto.set_data_location(ONNX_NAMESPACE::TensorProto_DataLocation_EXTERNAL); + auto* first = tensor_proto.add_external_data(); + first->set_key("location"); + first->set_value(first_location); + auto* second = tensor_proto.add_external_data(); + second->set_key("location"); + second->set_value(second_location); + + std::unique_ptr external_data_info; + const Status status = ExternalDataInfo::Create(tensor_proto.external_data(), external_data_info); + ASSERT_STATUS_NOT_OK_AND_HAS_SUBSTR(status, "duplicate 'location'"); + } +} + +TEST(TensorProtoUtilsTest, HasExternalDataInMemoryChecksAllLocations) { + auto make_tensor = [](const std::string& first_location, const std::string& second_location) { + ONNX_NAMESPACE::TensorProto tensor_proto; + tensor_proto.set_data_type(ONNX_NAMESPACE::TensorProto_DataType_FLOAT); + tensor_proto.set_data_location(ONNX_NAMESPACE::TensorProto_DataLocation_EXTERNAL); + auto* first = tensor_proto.add_external_data(); + first->set_key("location"); + first->set_value(first_location); + auto* second = tensor_proto.add_external_data(); + second->set_key("location"); + second->set_value(second_location); + return tensor_proto; + }; + + const std::string memory_tag = ToUTF8String(utils::kTensorProtoNativeEndianMemoryAddressTag); + EXPECT_TRUE(utils::HasExternalDataInMemory(make_tensor("data.bin", memory_tag))); + EXPECT_TRUE(utils::HasExternalDataInMemory(make_tensor(memory_tag, "data.bin"))); + EXPECT_TRUE(utils::HasExternalDataInMemory(make_tensor(memory_tag, ""))); +} + // Test ExternalData functionality TEST(TensorProtoUtilsTest, SetExternalDataInformation) { ONNX_NAMESPACE::TensorProto tensor_proto; @@ -725,6 +769,18 @@ class PathValidationTest : public ::testing::Test { std::vector other_files_; }; +TEST_F(PathValidationTest, ValidateExternalDataPathRejectsMemoryTags) { + for (const auto* memory_tag : {utils::kTensorProtoLittleEndianMemoryAddressTag, + utils::kTensorProtoNativeEndianMemoryAddressTag}) { + const Status status = utils::ValidateExternalDataPathFromDir(base_dir_, std::filesystem::path{memory_tag}); + ASSERT_STATUS_NOT_OK_AND_HAS_SUBSTR(status, "In-memory external data reference tag"); + + const Status wrapper_status = + utils::ValidateExternalDataPath(base_dir_ / "model.onnx", std::filesystem::path{memory_tag}); + ASSERT_STATUS_NOT_OK_AND_HAS_SUBSTR(wrapper_status, "In-memory external data reference tag"); + } +} + // Test cases for ValidateExternalDataPath. TEST_F(PathValidationTest, ValidateExternalDataPath) { std::filesystem::path model_path = base_dir_ / "model.onnx"; diff --git a/onnxruntime/test/mlas/unittest/test_symm_qgemm_fixture.h b/onnxruntime/test/mlas/unittest/test_symm_qgemm_fixture.h index 36d32b41bf8c6..8adaa3925c72d 100644 --- a/onnxruntime/test/mlas/unittest/test_symm_qgemm_fixture.h +++ b/onnxruntime/test/mlas/unittest/test_symm_qgemm_fixture.h @@ -4,7 +4,6 @@ #pragma once -#include "core/mlas/lib/mlasi.h" #include "test_symm_qgemm.h" // @@ -93,16 +92,6 @@ class SymmQgemmS8SignedInputTest : public MlasTestFixture // needed for std::transform #include "gtest/gtest.h" #include "test/unittest_util/framework_test_utils.h" +#include "test/unittest_util/graph_transform_test_builder.h" #include "test/test_environment.h" #include "test/util/include/default_providers.h" #include "test/util/include/asserts.h" @@ -511,5 +512,36 @@ TEST(TransformerTest, FuseFp16InitializersWithGraphOutputs) { ASSERT_STATUS_OK(session.Run(inputs, output_names, &outputs)); } // FuseFp16InitializersWithGraphOutputs +TEST(TransformerTest, FuseInitializersSkipsOverridableInitializer) { + auto& logger = DefaultLoggingManager().DefaultLogger(); + Model model("FuseInitializersSkipsOverridableInitializer", false, ModelMetaData(), PathString(), + IOnnxRuntimeOpSchemaRegistryList(), {{kOnnxDomain, 23}}, {}, logger); + Graph& graph = model.MainGraph(); + ModelTestBuilder builder(graph); + + NodeArg* initializer = builder.MakeInitializer({1}, {MLFloat16(1.0f)}); + NodeArg* cast_output = builder.MakeIntermediate(std::vector{1}); + builder.AddNode("Cast", {initializer}, {cast_output}) + .AddAttribute("to", static_cast(ONNX_NAMESPACE::TensorProto_DataType_FLOAT)); + NodeArg* output = builder.MakeOutput(std::vector{1}); + builder.AddNode("Neg", {cast_output}, {output}); + + graph.SetInputs({initializer}); + graph.SetOutputs({output}); + ASSERT_STATUS_OK(graph.Resolve()); + ASSERT_TRUE(graph.IsInitializedTensor(initializer->Name())); + ASSERT_EQ(nullptr, graph_utils::GetConstantInitializer(graph, initializer->Name())); + + FuseInitializersTransformer transformer("TransformerTest.FusedInitializers", + DataTypeImpl::GetTensorType(), + DataTypeImpl::GetTensorType()); + bool modified = false; + ASSERT_STATUS_OK(transformer.Apply(graph, modified, logger)); + + EXPECT_FALSE(modified); + EXPECT_EQ(1, CountOpsInGraph(graph)["Cast"]); + EXPECT_EQ(1, CountOpsInGraph(graph)["Neg"]); +} + } // namespace test } // namespace onnxruntime diff --git a/onnxruntime/test/optimizer/graph_transform_test.cc b/onnxruntime/test/optimizer/graph_transform_test.cc index 525bf62de5812..127d1c2347159 100644 --- a/onnxruntime/test/optimizer/graph_transform_test.cc +++ b/onnxruntime/test/optimizer/graph_transform_test.cc @@ -7382,6 +7382,106 @@ TEST_F(GraphTransformationTests, BiasGeluFusionCurrentOpsetTest) { ModelOptions{kAllowReleasedOpsetsOnly, /*strict_shape_type_inference*/ false})); } +#if !defined(DISABLE_CONTRIB_OPS) +// Regression test for the WebGPU entry added to the Level-2 GeluFusion allowlist +// (cpu_acl_cuda_dml_webgpu_eps in graph_transformer_utils.cc). +TEST_F(GraphTransformationTests, GeluFusionWebGpu) { + constexpr const ORTCHAR_T* model_uri = MODEL_FOLDER "fusion/gelu.onnx"; + std::shared_ptr p_model; + ASSERT_STATUS_OK(Model::Load(model_uri, p_model, nullptr, *logger_)); + Graph& graph = p_model->MainGraph(); +#if defined(USE_WEBGPU) + const std::string expected_ep = kWebGpuExecutionProvider; +#else + const std::string expected_ep = kCpuExecutionProvider; +#endif + for (auto& node : graph.Nodes()) { + node.SetExecutionProviderType(expected_ep); + } + + SessionOptions session_options; + auto cpu_ep = std::make_unique(CPUExecutionProviderInfo()); + const InlinedHashSet gelu_transformer_names = {"GeluFusionL1", "GeluFusionL2"}; + onnxruntime::GraphTransformerManager graph_transformation_mgr{5}; + for (auto level : {TransformerLevel::Level1, TransformerLevel::Level2}) { + for (auto& transformer : optimizer_utils::GenerateTransformers(level, session_options, *cpu_ep, *logger_, {})) { + if (gelu_transformer_names.count(transformer->Name()) != 0) { + ASSERT_STATUS_OK(graph_transformation_mgr.Register(std::move(transformer), level)); + } + } + } + ASSERT_STATUS_OK(graph_transformation_mgr.ApplyTransformers(graph, TransformerLevel::Level1, *logger_)); + ASSERT_STATUS_OK(graph_transformation_mgr.ApplyTransformers(graph, TransformerLevel::Level2, *logger_)); + + std::map op_to_count = CountOpsInGraph(graph); + ASSERT_EQ(op_to_count["com.microsoft.Gelu"], 1); + ASSERT_EQ(op_to_count["Div"], 0); + ASSERT_EQ(op_to_count["Erf"], 0); + ASSERT_EQ(op_to_count["Add"], 0); + ASSERT_EQ(op_to_count["Mul"], 0); + + const Node* gelu_node = nullptr; + for (auto& node : graph.Nodes()) { + if (node.OpType() == "Gelu" && node.Domain() == kMSDomain) { + gelu_node = &node; + break; + } + } + ASSERT_NE(gelu_node, nullptr); + EXPECT_EQ(gelu_node->GetExecutionProviderType(), expected_ep); +} + +// Regression test for the WebGPU entry added to the BiasGeluFusion allowlist. +TEST_F(GraphTransformationTests, BiasGeluFusionWebGpu) { + constexpr const ORTCHAR_T* model_uri = MODEL_FOLDER "fusion/bias_gelu_fusion.onnx"; + std::shared_ptr p_model; + ASSERT_STATUS_OK(Model::Load(model_uri, p_model, nullptr, *logger_)); + Graph& graph = p_model->MainGraph(); +#if defined(USE_WEBGPU) + const std::string expected_ep = kWebGpuExecutionProvider; +#else + const std::string expected_ep = kCpuExecutionProvider; +#endif + for (auto& node : graph.Nodes()) { + node.SetExecutionProviderType(expected_ep); + } + + SessionOptions session_options; + auto cpu_ep = std::make_unique(CPUExecutionProviderInfo()); + const InlinedHashSet gelu_transformer_names = { + "GeluFusionL1", "GeluFusionL2", "BiasGeluFusion"}; + onnxruntime::GraphTransformerManager graph_transformation_mgr{5}; + for (auto level : {TransformerLevel::Level1, TransformerLevel::Level2}) { + for (auto& transformer : optimizer_utils::GenerateTransformers(level, session_options, *cpu_ep, *logger_, {})) { + if (gelu_transformer_names.count(transformer->Name()) != 0) { + ASSERT_STATUS_OK(graph_transformation_mgr.Register(std::move(transformer), level)); + } + } + } + ASSERT_STATUS_OK(graph_transformation_mgr.ApplyTransformers(graph, TransformerLevel::Level1, *logger_)); + ASSERT_STATUS_OK(graph_transformation_mgr.ApplyTransformers(graph, TransformerLevel::Level2, *logger_)); + + std::map op_to_count = CountOpsInGraph(graph); + ASSERT_EQ(op_to_count["com.microsoft.BiasGelu"], 1); + ASSERT_EQ(op_to_count["com.microsoft.Gelu"], 0); + ASSERT_EQ(op_to_count["Gelu"], 0); + ASSERT_EQ(op_to_count["Add"], 0); + ASSERT_EQ(op_to_count["Div"], 0); + ASSERT_EQ(op_to_count["Erf"], 0); + ASSERT_EQ(op_to_count["Mul"], 0); + + const Node* bias_gelu_node = nullptr; + for (auto& node : graph.Nodes()) { + if (node.OpType() == "BiasGelu" && node.Domain() == kMSDomain) { + bias_gelu_node = &node; + break; + } + } + ASSERT_NE(bias_gelu_node, nullptr); + EXPECT_EQ(bias_gelu_node->GetExecutionProviderType(), expected_ep); +} +#endif // !defined(DISABLE_CONTRIB_OPS) + TEST_F(GraphTransformationTests, MatMulAddFusionCurrentOpsetTest) { // MatMul + Add -> Gemm fusion int current_opset = GetCurrentOnnxOpset(); diff --git a/onnxruntime/test/optimizer/initializer_test.cc b/onnxruntime/test/optimizer/initializer_test.cc index 9cd87c6f5b40f..acf0bbff7304e 100644 --- a/onnxruntime/test/optimizer/initializer_test.cc +++ b/onnxruntime/test/optimizer/initializer_test.cc @@ -13,6 +13,7 @@ #include "gtest/gtest.h" #include "core/common/common.h" +#include "core/common/span_utils.h" #include "core/framework/endian_utils.h" #include "test/util/include/asserts.h" #include "test/util/include/file_util.h" @@ -275,6 +276,41 @@ TEST(OptimizerInitializerTest, DataField) { TestInitializerDataField(); } +TEST(OptimizerInitializerTest, ScaleByAxisEmptyTensor) { + { + Initializer target(ONNX_NAMESPACE::TensorProto_DataType_FLOAT, "target", AsSpan({2, 0})); + Initializer scalers(ONNX_NAMESPACE::TensorProto_DataType_FLOAT, "scalers", AsSpan({2})); + EXPECT_NO_THROW(target.scale_by_axis(scalers, 1)); + EXPECT_EQ(target.size(), 0u); + } + + { + Initializer target(ONNX_NAMESPACE::TensorProto_DataType_FLOAT, "target", AsSpan({2, 0})); + Initializer scalers(ONNX_NAMESPACE::TensorProto_DataType_FLOAT, "scalers", AsSpan({0})); + EXPECT_NO_THROW(target.scale_by_axis(scalers, 1, true)); + EXPECT_EQ(target.size(), 0u); + } + + { + Initializer target(ONNX_NAMESPACE::TensorProto_DataType_FLOAT, "target", AsSpan({0})); + Initializer scalar(ONNX_NAMESPACE::TensorProto_DataType_FLOAT, "scalar", AsSpan({})); + EXPECT_NO_THROW(target.scale_by_axis(scalar, 0)); + EXPECT_EQ(target.size(), 0u); + } + + { + Initializer target(ONNX_NAMESPACE::TensorProto_DataType_FLOAT, "target", AsSpan({2, 0})); + Initializer invalid_scalers(ONNX_NAMESPACE::TensorProto_DataType_FLOAT, "invalid_scalers", AsSpan({3})); + EXPECT_THROW(target.scale_by_axis(invalid_scalers, 1), OnnxRuntimeException); + } + + { + Initializer target(ONNX_NAMESPACE::TensorProto_DataType_FLOAT, "target", AsSpan({2, 0})); + Initializer scalers(ONNX_NAMESPACE::TensorProto_DataType_DOUBLE, "scalers", AsSpan({2})); + EXPECT_THROW(target.scale_by_axis(scalers, 1), OnnxRuntimeException); + } +} + // An in-memory external-data initializer with no registered OrtValue must load without a model_path. TEST(OptimizerInitializerTest, InMemoryExternalDataWithoutOrtValueOrModelPath) { std::vector backing(64); // 256 bytes > kSmallTensorExternalDataThreshold diff --git a/onnxruntime/test/providers/cpu/llm/tensorscatter_op_test.cc b/onnxruntime/test/providers/cpu/llm/tensorscatter_op_test.cc index 5b2af54d309da..7d3d010c93084 100644 --- a/onnxruntime/test/providers/cpu/llm/tensorscatter_op_test.cc +++ b/onnxruntime/test/providers/cpu/llm/tensorscatter_op_test.cc @@ -1,6 +1,8 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. +#include + #include "gtest/gtest.h" #include "test/providers/provider_test_utils.h" @@ -335,6 +337,21 @@ TEST(TensorScatterTest, Linear_OutOfBoundsWriteIndex) { {}, nullptr, &execution_providers); } +TEST(TensorScatterTest, Linear_WriteIndexAdditionOverflow) { + OpTester test("TensorScatter", 24); + test.AddAttribute("mode", "linear"); + + test.AddInput("past_cache", {1, 4, 1}, {0, 0, 0, 0}); + test.AddInput("update", {1, 2, 1}, {1, 2}); + test.AddInput("write_indices", {1}, {std::numeric_limits::max()}); + test.AddOutput("present_cache", {1, 4, 1}, {0, 0, 0, 0}); + + std::vector> execution_providers; + execution_providers.push_back(DefaultCpuExecutionProvider()); + test.Run(OpTester::ExpectResult::kExpectFailure, "exceeds max_sequence_length", + {}, nullptr, &execution_providers); +} + // Circular mode: negative write_indices should still fail. // Run CPU-only: CUDA validates asynchronously via CUDA_KERNEL_ASSERT. TEST(TensorScatterTest, Circular_NegativeWriteIndex) { @@ -354,6 +371,131 @@ TEST(TensorScatterTest, Circular_NegativeWriteIndex) { {}, nullptr, &execution_providers); } +TEST(TensorScatterTest, Circular_LargeWriteIndexWrapsWithoutOverflow) { + OpTester test("TensorScatter", 24); + test.AddAttribute("mode", "circular"); + + test.AddInput("past_cache", {1, 4, 1}, {0, 0, 0, 0}); + test.AddInput("update", {1, 2, 1}, {1, 2}); + test.AddInput("write_indices", {1}, {std::numeric_limits::max()}); + test.AddOutput("present_cache", {1, 4, 1}, {2, 0, 0, 1}); + + std::vector> execution_providers; + execution_providers.push_back(DefaultCpuExecutionProvider()); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers); +} + +TEST(TensorScatterTest, Circular_ZeroSequenceLengthIsNoOp) { + OpTester test("TensorScatter", 24); + test.AddAttribute("mode", "circular"); + + test.AddInput("past_cache", {1, 0, 1}, {}); + test.AddInput("update", {1, 0, 1}, {}); + test.AddInput("write_indices", {1}, {std::numeric_limits::max()}); + test.AddOutput("present_cache", {1, 0, 1}, {}); + + std::vector> execution_providers; + execution_providers.push_back(DefaultCpuExecutionProvider()); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers); +} + +TEST(TensorScatterTest, ZeroSequenceLengthWithoutWriteIndicesSkipsLargeBatch) { + OpTester test("TensorScatter", 24); + test.AddAttribute("mode", "circular"); + + constexpr int64_t large_batch = std::numeric_limits::max(); + test.AddInput("past_cache", {large_batch, 0, 1}, {}); + test.AddInput("update", {large_batch, 0, 1}, {}); + test.AddOptionalInputEdge(); + test.AddOutput("present_cache", {large_batch, 0, 1}, {}); + + std::vector> execution_providers; + execution_providers.push_back(DefaultCpuExecutionProvider()); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers); +} + +TEST(TensorScatterTest, ExplicitWriteIndicesValidatedWithZeroPrefixDimension) { + OpTester test("TensorScatter", 24); + test.AddAttribute("axis", 2); + test.AddAttribute("mode", "circular"); + + test.AddInput("past_cache", {2, 0, 4, 1}, {}); + test.AddInput("update", {2, 0, 1, 1}, {}); + test.AddInput("write_indices", {2}, {-1, 0}); + test.AddOutput("present_cache", {2, 0, 4, 1}, {}); + + std::vector> execution_providers; + execution_providers.push_back(DefaultCpuExecutionProvider()); + test.Run(OpTester::ExpectResult::kExpectFailure, "is negative", + {}, nullptr, &execution_providers); +} + +TEST(TensorScatterTest, Circular_ZeroSequenceLengthPreservesCache) { + OpTester test("TensorScatter", 24); + test.AddAttribute("mode", "circular"); + + test.AddInput("past_cache", {1, 4, 1}, {1, 2, 3, 4}); + test.AddInput("update", {1, 0, 1}, {}); + test.AddInput("write_indices", {1}, {std::numeric_limits::max()}); + test.AddOutput("present_cache", {1, 4, 1}, {1, 2, 3, 4}); + + std::vector> execution_providers; + execution_providers.push_back(DefaultCpuExecutionProvider()); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers); +} + +TEST(TensorScatterTest, Linear_ZeroSequenceLengthPreservesCache) { + OpTester test("TensorScatter", 24); + test.AddAttribute("mode", "linear"); + + test.AddInput("past_cache", {1, 4, 1}, {1, 2, 3, 4}); + test.AddInput("update", {1, 0, 1}, {}); + test.AddInput("write_indices", {1}, {4}); + test.AddOutput("present_cache", {1, 4, 1}, {1, 2, 3, 4}); + + std::vector> execution_providers; + execution_providers.push_back(DefaultCpuExecutionProvider()); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers); +} + +TEST(TensorScatterTest, Linear_ZeroSequenceLengthRejectsOutOfBoundsIndex) { + OpTester test("TensorScatter", 24); + test.AddAttribute("mode", "linear"); + + test.AddInput("past_cache", {1, 4, 1}, {1, 2, 3, 4}); + test.AddInput("update", {1, 0, 1}, {}); + test.AddInput("write_indices", {1}, {std::numeric_limits::max()}); + test.AddOutput("present_cache", {1, 4, 1}, {1, 2, 3, 4}); + + std::vector> execution_providers; + execution_providers.push_back(DefaultCpuExecutionProvider()); + test.Run(OpTester::ExpectResult::kExpectFailure, "exceeds max_sequence_length", + {}, nullptr, &execution_providers); +} + +static void RunZeroSequenceLengthNegativeWriteIndexTest(const std::string& mode) { + OpTester test("TensorScatter", 24); + test.AddAttribute("mode", mode); + + test.AddInput("past_cache", {1, 4, 1}, {1, 2, 3, 4}); + test.AddInput("update", {1, 0, 1}, {}); + test.AddInput("write_indices", {1}, {-1}); + test.AddOutput("present_cache", {1, 4, 1}, {1, 2, 3, 4}); + + std::vector> execution_providers; + execution_providers.push_back(DefaultCpuExecutionProvider()); + test.Run(OpTester::ExpectResult::kExpectFailure, "is negative", + {}, nullptr, &execution_providers); +} + +TEST(TensorScatterTest, Linear_ZeroSequenceLengthRejectsNegativeIndex) { + RunZeroSequenceLengthNegativeWriteIndexTest("linear"); +} + +TEST(TensorScatterTest, Circular_ZeroSequenceLengthRejectsNegativeIndex) { + RunZeroSequenceLengthNegativeWriteIndexTest("circular"); +} + // The CPU kernel only supports fixed-size element types (matching the CUDA // kernel's type constraint). Non-fixed-size element types such as string are // intentionally excluded because the kernel operates on raw memory buffers. diff --git a/onnxruntime/test/providers/cpu/ml/normalizer_test.cc b/onnxruntime/test/providers/cpu/ml/normalizer_test.cc index d23696ddc5d85..e12a53576d8b4 100644 --- a/onnxruntime/test/providers/cpu/ml/normalizer_test.cc +++ b/onnxruntime/test/providers/cpu/ml/normalizer_test.cc @@ -187,5 +187,13 @@ TEST(Normalizer, InvalidNorm) { RunTest(input, dims, output, "InvalidNormValue", OpTester::ExpectResult::kExpectFailure); } +TEST(Normalizer, ScalarInputRejected) { + for (const char* const norm : {"MAX", "L1", "L2"}) { + RunTest({1.0f}, {}, {0.0f}, norm, + OpTester::ExpectResult::kExpectFailure, + "must have rank 1 or 2"); + } +} + } // namespace test } // namespace onnxruntime diff --git a/onnxruntime/test/providers/cpu/nn/conv_op_test.cc b/onnxruntime/test/providers/cpu/nn/conv_op_test.cc index 25d37846a2028..2165c28027f62 100644 --- a/onnxruntime/test/providers/cpu/nn/conv_op_test.cc +++ b/onnxruntime/test/providers/cpu/nn/conv_op_test.cc @@ -1671,5 +1671,30 @@ TEST(ConvTest, Conv2D_ZeroDilation_Dml) { .RunWithConfig(); } +TEST(ConvTest, Conv2D_InvalidBiasSize) { + OpTester test("Conv", 22); + test.AddInput("X", {1, 1, 1, 1}, {1.0f}); + test.AddInput("W", {4, 1, 1, 1}, std::vector(4, 0.0f)); + test.AddInput("B", {1}, {0.0f}); + test.AddOutput("Y", {1, 4, 1, 1}, std::vector(4, 0.0f)); + test.ConfigEp(DefaultCpuExecutionProvider()) + .Config(OpTester::ExpectResult::kExpectFailure, + "bias must be a 1D tensor of size output_channels") + .RunWithConfig(); +} + +TEST(ConvTest, Conv2D_InvalidBiasRank) { + OpTester test("Conv", 22); + test.AddShapeToTensorData(false); + test.AddInput("X", {1, 1, 1, 1}, {1.0f}); + test.AddInput("W", {4, 1, 1, 1}, std::vector(4, 0.0f)); + test.AddInput("B", {1, 4}, std::vector(4, 0.0f)); + test.AddOutput("Y", {1, 4, 1, 1}, std::vector(4, 0.0f)); + test.ConfigEp(DefaultCpuExecutionProvider()) + .Config(OpTester::ExpectResult::kExpectFailure, + "bias must be a 1D tensor of size output_channels") + .RunWithConfig(); +} + } // namespace test } // namespace onnxruntime diff --git a/onnxruntime/test/providers/cpu/reduction/reduction_ops_test.cc b/onnxruntime/test/providers/cpu/reduction/reduction_ops_test.cc index 163c350cbbb36..55cad7cd529fd 100644 --- a/onnxruntime/test/providers/cpu/reduction/reduction_ops_test.cc +++ b/onnxruntime/test/providers/cpu/reduction/reduction_ops_test.cc @@ -6353,6 +6353,55 @@ void test_empty_set(const std::string& op, int opset, bool axes_as_input, float }); } +TEST(ReductionOpTest, EmptySetMissingOptionalAxesReducesAllDimensions) { + OpTester test("ReduceSum", 20); + test.AddInput("data", {2, 0, 4}, {}); + test.AddOptionalInputEdge(); + test.AddOutput("reduced", {1, 1, 1}, {0.0f}); + test.ConfigEp(DefaultCpuExecutionProvider()).RunWithConfig(); +} + +TEST(ReductionOpTest, EmptySetAxesMustBeVector) { + const std::vector, std::vector>> axes_cases{ + {{}, {1}}, + {{2, 0}, {}}, + }; + for (const char* const op : {"ReduceSum", "ReduceLogSumExp"}) { + for (const auto& [axes_shape, axes_data] : axes_cases) { + OpTester test(op, 20); + test.AddInput("data", {2, 0, 4}, {}); + test.AddInput("axes", axes_shape, axes_data); + test.AddOutput("reduced", {1, 1, 1}, {0.0f}); + test.Config(OpTester::ExpectResult::kExpectFailure, "An axes tensor must be a vector tensor.") + .ConfigEp(DefaultCpuExecutionProvider()) + .RunWithConfig(); + } + } +} + +void TestEmptySetNoopWithEmptyAxes(const std::string& op, bool omit_axes) { + OpTester test(op, 20); + test.AddInput("data", {2, 0, 4}, {}); + if (omit_axes) { + test.AddOptionalInputEdge(); + } else { + test.AddInput("axes", {0}, {}, true); + } + test.AddAttribute("noop_with_empty_axes", 1); + test.AddOutput("reduced", {2, 0, 4}, {}); + test.ConfigEp(DefaultCpuExecutionProvider()).RunWithConfig(); +} + +TEST(ReductionOpTest, EmptySetNoopWithMissingAxes) { + TestEmptySetNoopWithEmptyAxes("ReduceSum", true); + TestEmptySetNoopWithEmptyAxes("ReduceLogSumExp", true); +} + +TEST(ReductionOpTest, EmptySetNoopWithEmptyAxes) { + TestEmptySetNoopWithEmptyAxes("ReduceSum", false); + TestEmptySetNoopWithEmptyAxes("ReduceLogSumExp", false); +} + TEST(ReductionOpTest, empty_set_ReduceL1) { test_empty_set("ReduceL1", 20, true, 0); } diff --git a/onnxruntime/test/providers/cpu/tensor/scatter_nd_op_test.cc b/onnxruntime/test/providers/cpu/tensor/scatter_nd_op_test.cc index 69788633af660..332b4c56aaba0 100644 --- a/onnxruntime/test/providers/cpu/tensor/scatter_nd_op_test.cc +++ b/onnxruntime/test/providers/cpu/tensor/scatter_nd_op_test.cc @@ -269,5 +269,36 @@ TEST(ScatterNDOpTest, ScatterND_empty_indices) { test1.Run(OpTester::ExpectResult::kExpectSuccess, "", {kDmlExecutionProvider}); } +TEST(ScatterNDOpTest, ScatterND_zero_index_depth_updates_entire_tensor) { + OpTester test("ScatterND", 18); + test.AddInput("data", {2, 3}, {1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f}); + test.AddInput("indices", {1, 0}, {}); + test.AddInput("updates", {1, 2, 3}, {10.0f, 20.0f, 30.0f, 40.0f, 50.0f, 60.0f}); + test.AddOutput("output", {2, 3}, {10.0f, 20.0f, 30.0f, 40.0f, 50.0f, 60.0f}); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", + {kTensorrtExecutionProvider, kWebGpuExecutionProvider}); +} + +TEST(ScatterNDOpTest, ScatterND_zero_index_depth_adds_multiple_updates) { + OpTester test("ScatterND", 18); + test.AddAttribute("reduction", "add"); + test.AddInput("data", {2}, {1.0f, 2.0f}); + test.AddInput("indices", {2, 0}, {}); + test.AddInput("updates", {2, 2}, {10.0f, 20.0f, 100.0f, 200.0f}); + test.AddOutput("output", {2}, {111.0f, 222.0f}); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", + {kTensorrtExecutionProvider, kWebGpuExecutionProvider}); +} + +TEST(ScatterNDOpTest, ScatterND_zero_index_depth_empty_data) { + OpTester test("ScatterND", 18); + test.AddInput("data", {0, 3}, {}); + test.AddInput("indices", {1, 0}, {}); + test.AddInput("updates", {1, 0, 3}, {}); + test.AddOutput("output", {0, 3}, {}); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", + {kTensorrtExecutionProvider, kWebGpuExecutionProvider}); +} + } // namespace test } // namespace onnxruntime diff --git a/onnxruntime/test/providers/cpu/tensor/split_op_test.cc b/onnxruntime/test/providers/cpu/tensor/split_op_test.cc index bb0a945612188..0ee94a2903d34 100644 --- a/onnxruntime/test/providers/cpu/tensor/split_op_test.cc +++ b/onnxruntime/test/providers/cpu/tensor/split_op_test.cc @@ -1,6 +1,8 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. +#include + #include "gtest/gtest.h" #include "core/framework/to_tensor_proto_element_type.h" #include "test/providers/provider_test_utils.h" @@ -949,5 +951,22 @@ TEST(SplitOperatorTest, InvalidValueInSplitInput_NegativeEntry_NegativeAxis) { {}, nullptr, &execution_providers); } +TEST(SplitOperatorTest, InvalidValueInSplitInput_Overflow) { + OpTester test("Split", 13, onnxruntime::kOnnxDomain); + test.AddAttribute("axis", 0); + test.AddInput("input", {4, 2}, {1.f, 2.f, 3.f, 4.f, 5.f, 6.f, 7.f, 8.f}); + test.AddInput("split", {3}, {6, std::numeric_limits::max(), std::numeric_limits::max()}, + /*is_initializer=*/false); + test.AddOutput("output0", {1, 2}, {0.f, 0.f}); + test.AddOutput("output1", {1, 2}, {0.f, 0.f}); + test.AddOutput("output2", {1, 2}, {0.f, 0.f}); + + std::vector> execution_providers; + execution_providers.push_back(DefaultCpuExecutionProvider()); + test.Run(OpTester::ExpectResult::kExpectFailure, + "exceeds the remaining size of the selected axis", + {}, nullptr, &execution_providers); +} + } // namespace test } // namespace onnxruntime diff --git a/onnxruntime/test/python/onnxruntime_test_python.py b/onnxruntime/test/python/onnxruntime_test_python.py index d60acfc8bc34f..5fbf908e7f3d6 100644 --- a/onnxruntime/test/python/onnxruntime_test_python.py +++ b/onnxruntime/test/python/onnxruntime_test_python.py @@ -12,7 +12,9 @@ import queue import sys import threading +import time import unittest +import weakref import numpy as np from helper import get_name @@ -685,7 +687,9 @@ def test_run_model(self): def test_run_async(self): event = threading.Event() + allow_callback = threading.Event() output_expected = np.array([[1.0, 4.0], [9.0, 16.0], [25.0, 36.0]], dtype=np.float32) + input_ref = None class MyData: def __init__(self, id): @@ -697,6 +701,8 @@ def get_id(self): my_data = MyData(123456) def callback(res: np.ndarray, data: MyData, err: str) -> None: + self.assertTrue(allow_callback.wait(10)) + self.assertIsNotNone(input_ref()) self.assertEqual(len(err), 0) self.assertEqual(len(res), 1) self.assertEqual(data.get_id(), 123456) @@ -709,10 +715,22 @@ def callback(res: np.ndarray, data: MyData, err: str) -> None: sess = onnxrt.InferenceSession(get_name("mul_1.onnx"), so, providers=available_providers) x = np.array([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]], dtype=np.float32) - sess.run_async(["Y"], {"X": x}, callback, my_data) + input_ref = weakref.ref(x) + run_options = onnxrt.RunOptions() + sess.run_async(["Y"], {"X": x}, callback, my_data, run_options) + del x + del run_options + del sess + gc.collect() + allow_callback.set() event.wait(10) # timeout in 10 sec self.assertTrue(event.is_set()) + deadline = time.monotonic() + 10 + while input_ref() is not None and time.monotonic() < deadline: + gc.collect() + time.sleep(0.01) + self.assertIsNone(input_ref()) def test_run_model_from_bytes(self): with open(get_name("mul_1.onnx"), "rb") as f: diff --git a/tools/ci_build/github/azure-pipelines/stages/plugin-linux-cuda-stage.yml b/tools/ci_build/github/azure-pipelines/stages/plugin-linux-cuda-stage.yml index 4eeb7c490af26..70114ccfcc52a 100644 --- a/tools/ci_build/github/azure-pipelines/stages/plugin-linux-cuda-stage.yml +++ b/tools/ci_build/github/azure-pipelines/stages/plugin-linux-cuda-stage.yml @@ -130,6 +130,7 @@ stages: -p ${{ parameters.docker_python_exe_path }} -v ${{ parameters.cuda_version }} -a "${{ parameters.cmake_cuda_archs }}" + -h "${{ parameters.arch }}" workingDirectory: $(Build.SourcesDirectory) displayName: 'Build CUDA Plugin (Python ${{ parameters.python_version }}, ${{ parameters.arch }}, CUDA ${{ parameters.cuda_version }})' env: diff --git a/tools/ci_build/github/azure-pipelines/templates/set-plugin-ep-build-variables-step.yml b/tools/ci_build/github/azure-pipelines/templates/set-plugin-ep-build-variables-step.yml index 97202651e1841..9f16f62fcccda 100644 --- a/tools/ci_build/github/azure-pipelines/templates/set-plugin-ep-build-variables-step.yml +++ b/tools/ci_build/github/azure-pipelines/templates/set-plugin-ep-build-variables-step.yml @@ -28,5 +28,9 @@ steps: # Use 'script' (not 'bash') so this works on both Linux and Windows agents. # On Linux aarch64 agents UsePythonVersion@0 is unavailable, so we call the configured # Python executable directly instead of using PythonScript@0. -- script: ${{ parameters.python_command }} "$(Build.SourcesDirectory)/tools/ci_build/set_plugin_ep_build_variables.py" "${{ parameters.package_version }}" "${{ parameters.version_file }}" +- script: >- + ${{ parameters.python_command }} + "$(Build.SourcesDirectory)/tools/ci_build/set_plugin_ep_build_variables.py" + --package-version "${{ parameters.package_version }}" + --version-file "${{ parameters.version_file }}" displayName: 'Set plugin EP package version string' diff --git a/tools/ci_build/github/linux/build_cuda_plugin_package.sh b/tools/ci_build/github/linux/build_cuda_plugin_package.sh index f12811ba087c4..a8bc55df45aca 100755 --- a/tools/ci_build/github/linux/build_cuda_plugin_package.sh +++ b/tools/ci_build/github/linux/build_cuda_plugin_package.sh @@ -9,8 +9,9 @@ DOCKER_IMAGE="onnxruntimecuda128pluginbuild" PYTHON_EXE="/opt/python/cp312-cp312/bin/python3.12" CUDA_VERSION="12.8" CMAKE_CUDA_ARCHS="86" +arch="x64" -while getopts "i:c:p:v:a:" parameter_Option +while getopts "i:c:p:v:a:h:" parameter_Option do case "${parameter_Option}" in i) DOCKER_IMAGE=${OPTARG};; @@ -18,11 +19,19 @@ c) BUILD_CONFIG=${OPTARG};; p) PYTHON_EXE=${OPTARG};; v) CUDA_VERSION=${OPTARG};; a) CMAKE_CUDA_ARCHS=${OPTARG};; -*) echo "Usage: $0 -i [-c ] [-p ] [-v ] [-a ]" +h) arch=${OPTARG};; +*) echo "Usage: $0 -i [-c ] [-p ] [-v ] [-a ] [-h ]" exit 1;; esac done +# Update parallel to the minimum of the number of processors and 8 for aarch64 to avoid OOM errors during build. +if [ "$arch" = "aarch64" ]; then + PARALLEL=$(( $(nproc) < 8 ? $(nproc) : 8 )) +else + PARALLEL="" +fi + PYTHON_BIN_DIR=$(dirname "${PYTHON_EXE}") # Derive SHORT_CUDA_VERSION (e.g., 12.8 from 12.8, 13.0 from 13.0) @@ -54,7 +63,7 @@ docker run --rm \ --build_dir /build \ --config ${BUILD_CONFIG} \ --skip_submodule_sync \ - --parallel \ + --parallel ${PARALLEL} \ --nvcc_threads 1 \ --flash_nvcc_threads 1 \ --use_binskim_compliant_compile_flags \ diff --git a/tools/ci_build/set_plugin_ep_build_variables.py b/tools/ci_build/set_plugin_ep_build_variables.py index 017f827779553..9b97aff06bab5 100644 --- a/tools/ci_build/set_plugin_ep_build_variables.py +++ b/tools/ci_build/set_plugin_ep_build_variables.py @@ -2,36 +2,47 @@ # Copyright (c) Microsoft Corporation. All rights reserved. # Licensed under the MIT License. -"""Set plugin EP package version variables for Azure Pipelines. - -Usage: - python set_plugin_ep_build_variables.py - -Where: - package_version: 'release', 'rc', or 'dev' - version_file_rel: path relative to BUILD_SOURCESDIRECTORY of the VERSION_NUMBER file -""" +"""Set plugin EP package version variables for Azure Pipelines.""" +import argparse +import datetime import os import re import subprocess import sys -def main(): - if len(sys.argv) != 3: - print(f"Usage: {sys.argv[0]} ") - sys.exit(1) +class AzurePipelinesArgumentParser(argparse.ArgumentParser): + def error(self, message): + print(f"##vso[task.logissue type=error]{message}") + self.exit(2) + - package_version = sys.argv[1] - version_file_rel = sys.argv[2] +def parse_arguments(): + parser = AzurePipelinesArgumentParser(description=__doc__) + parser.add_argument( + "--package-version", + choices=("release", "rc", "dev"), + required=True, + help="Package version type.", + ) + parser.add_argument( + "--version-file", + required=True, + help="Path to the VERSION_NUMBER file, relative to BUILD_SOURCESDIRECTORY.", + ) + return parser.parse_args() - if not version_file_rel: - print("##vso[task.logissue type=error]version_file parameter is empty.") + +def main(): + args = parse_arguments() + + src_root = os.environ.get("BUILD_SOURCESDIRECTORY") + if not src_root: + print("##vso[task.logissue type=error]BUILD_SOURCESDIRECTORY is not set.") sys.exit(1) - src_root = os.environ.get("BUILD_SOURCESDIRECTORY", "") - version_file = os.path.join(src_root, version_file_rel) + version_file = os.path.join(src_root, args.version_file) if not os.path.isfile(version_file): print(f"##vso[task.logissue type=error]Cannot find version number file at: {version_file}") sys.exit(1) @@ -44,19 +55,19 @@ def main(): sys.exit(1) print(f"Original version: {original_ver}") - print(f"Package version type: {package_version}") + print(f"Package version type: {args.package_version}") - if package_version == "release": + if args.package_version == "release": version_string = original_ver python_version = original_ver - elif package_version == "rc": + elif args.package_version == "rc": # RC versioning is not yet implemented. Fail the build to prevent publishing # an ambiguous version without an RC number. print("##vso[task.logissue type=error]RC versioning is not yet implemented. Use 'dev' or 'release' instead.") sys.exit(1) - elif package_version == "dev": + elif args.package_version == "dev": try: commit_sha = ( subprocess.check_output( @@ -66,25 +77,27 @@ def main(): .decode("utf-8") .strip() ) - date_str = ( + commit_timestamp = int( subprocess.check_output( - ["git", "show", "-s", "--format=%cd", "--date=format:%Y%m%d", "HEAD"], + ["git", "show", "-s", "--format=%ct", "HEAD"], cwd=src_root, ) .decode("utf-8") .strip() ) + date_time_str = datetime.datetime.fromtimestamp( + commit_timestamp, + tz=datetime.timezone.utc, + ).strftime("%Y%m%d%H%M%S") except Exception as e: print(f"##vso[task.logissue type=error]Failed to get git info: {e}") sys.exit(1) - version_string = f"{original_ver}-dev.{date_str}+{commit_sha}" - python_version = f"{original_ver}.dev{date_str}" - else: - print( - f"##vso[task.logissue type=error]Unknown package_version '{package_version}'. Must be 'release', 'rc', or 'dev'." - ) - sys.exit(1) + # The UTC commit timestamp determines dev-version precedence. Distinct commits made in the same second have + # equal SemVer precedence and identical Python versions. If that is a problem, we can add a unique identifier + # such as Azure Pipelines Build.BuildId to the precedence-bearing portion of each version. + version_string = f"{original_ver}-dev.{date_time_str}+{commit_sha}" + python_version = f"{original_ver}.dev{date_time_str}" print(f"Plugin package version string: {version_string}") print(f"Plugin Python package version string: {python_version}")