From ade58cf78929f43d3f8566c234fdedeacd3b063c Mon Sep 17 00:00:00 2001 From: Changming Sun Date: Fri, 19 May 2023 21:48:21 -0700 Subject: [PATCH] Cherry-pick 4 commits to rel-1.15.0 branch (#16021) ### Description Cherry-pick 4 commits to rel-1.15.0 branch #14834 #15727 #16010 #16011 --- .pipelines/nuget_config/x64/packages.config | 2 +- .pipelines/nuget_config/x86/packages.config | 2 +- cmake/CMakeLists.txt | 5 +- cmake/external/dml.cmake | 2 +- docs/OperatorKernels.md | 2 + .../platform/EigenNonBlockingThreadPool.h | 21 + .../onnxruntime/core/platform/ort_spin_lock.h | 35 + .../contrib_ops/cuda/bert/attention.cc | 1 + .../core/optimizer/graph_transformer_utils.cc | 2 + .../src/External/DirectMLHelpers/ApiHelpers.h | 6 +- .../src/External/DirectMLHelpers/ApiTraits.h | 18 +- .../External/DirectMLHelpers/DirectMLSchema.h | 29 + .../src/External/DirectMLHelpers/DirectMLX.h | 120 +-- .../DirectMLHelpers/GeneratedSchemaHelpers.h | 28 + .../src/Operators/DmlOperatorAttention.cpp | 893 ++++++++++-------- .../DmlOperatorMultiHeadAttention.cpp | 281 ++++++ .../src/Operators/OperatorRegistration.cpp | 3 + .../src/Operators/OperatorUtility.cpp | 17 +- .../dml/OperatorAuthorHelper/Attributes.h | 2 + .../OperatorAuthorHelper/OperatorHelper.cpp | 108 +++ .../dml/OperatorAuthorHelper/OperatorHelper.h | 35 +- .../OperatorAuthorHelper/OperatorVersions.h | 1 + .../test/contrib_ops/attention_op_test.cc | 63 +- .../contrib_ops/attention_op_test_helper.cc | 2 +- .../contrib_ops/attention_op_test_helper.h | 2 +- .../multihead_attention_op_test.cc | 28 +- .../test/providers/provider_test_utils.cc | 2 +- packages.config | 2 +- tools/ci_build/build.py | 5 + .../azure-pipelines/win-ci-pipeline.yml | 2 +- .../nuget/generate_nuspec_for_native_nuget.py | 2 +- 31 files changed, 1210 insertions(+), 511 deletions(-) create mode 100644 include/onnxruntime/core/platform/ort_spin_lock.h create mode 100644 onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/DmlOperatorMultiHeadAttention.cpp diff --git a/.pipelines/nuget_config/x64/packages.config b/.pipelines/nuget_config/x64/packages.config index 5dfeb0836f..8eef0b5bac 100644 --- a/.pipelines/nuget_config/x64/packages.config +++ b/.pipelines/nuget_config/x64/packages.config @@ -1,6 +1,6 @@  - + diff --git a/.pipelines/nuget_config/x86/packages.config b/.pipelines/nuget_config/x86/packages.config index f8e63133f4..81f97948f1 100644 --- a/.pipelines/nuget_config/x86/packages.config +++ b/.pipelines/nuget_config/x86/packages.config @@ -1,6 +1,6 @@  - + diff --git a/cmake/CMakeLists.txt b/cmake/CMakeLists.txt index 17fdc9a86e..2e54063598 100644 --- a/cmake/CMakeLists.txt +++ b/cmake/CMakeLists.txt @@ -226,6 +226,7 @@ option(onnxruntime_BUILD_CACHE "onnxruntime build with cache" OFF) cmake_dependent_option(MSVC_Z7_OVERRIDE "replacing /Zi and /ZI with /Z7 when using MSVC with CCache" ON "onnxruntime_BUILD_CACHE; MSVC" OFF) option(onnxruntime_USE_AZURE "Build with azure inferencing support" OFF) +option(onnxruntime_USE_LOCK_FREE_QUEUE "Build with lock-free task queue for threadpool." OFF) # ENABLE_TRAINING includes all training functionality # The following 2 entry points @@ -743,7 +744,9 @@ if (onnxruntime_USE_AZURE) list(APPEND ORT_PROVIDER_CMAKE_FLAGS -Donnxruntime_USE_AZURE=1) list(APPEND ONNXRUNTIME_PROVIDER_NAMES azure) endif() - +if (onnxruntime_USE_LOCK_FREE_QUEUE) + add_compile_definitions(USE_LOCK_FREE_QUEUE) +endif() if (onnxruntime_ENABLE_LAZY_TENSOR) # To support LazyTensor, ORT needs to call Python function from C/C++. diff --git a/cmake/external/dml.cmake b/cmake/external/dml.cmake index 4dec75deca..1ee3a06b41 100644 --- a/cmake/external/dml.cmake +++ b/cmake/external/dml.cmake @@ -41,7 +41,7 @@ if (NOT onnxruntime_USE_CUSTOM_DIRECTML) set(NUGET_CONFIG ${PROJECT_SOURCE_DIR}/../NuGet.config) set(PACKAGES_CONFIG ${PROJECT_SOURCE_DIR}/../packages.config) get_filename_component(PACKAGES_DIR ${CMAKE_CURRENT_BINARY_DIR}/../packages ABSOLUTE) - set(DML_PACKAGE_DIR ${PACKAGES_DIR}/Microsoft.AI.DirectML.1.11.0) + set(DML_PACKAGE_DIR ${PACKAGES_DIR}/Microsoft.AI.DirectML.1.12.0) # Restore nuget packages, which will pull down the DirectML redist package. add_custom_command( diff --git a/docs/OperatorKernels.md b/docs/OperatorKernels.md index 97051e99ef..d4b2310fa1 100644 --- a/docs/OperatorKernels.md +++ b/docs/OperatorKernels.md @@ -1193,8 +1193,10 @@ Do not modify directly.* |DequantizeLinear|*in* x:**T1**
*in* x_scale:**T2**
*in* x_zero_point:**T1**
*out* y:**T2**|1+|**T1** = tensor(float)
**T2** = tensor(uint8)| |EmbedLayerNormalization|*in* input_ids:**T1**
*in* segment_ids:**T1**
*in* word_embedding:**T**
*in* position_embedding:**T**
*in* segment_embedding:**T**
*in* gamma:**T**
*in* beta:**T**
*in* mask:**T1**
*in* position_ids:**T1**
*out* output:**T**
*out* mask_index:**T1**
*out* embedding_sum:**T**|1+|**T** = tensor(float), tensor(float16)| |FusedMatMul|*in* A:**T**
*in* B:**T**
*out* Y:**T**|1+|**T** = tensor(float), tensor(float16)| +|FusedMatMulActivation|*in* A:**T**
*in* B:**T**
*out* Y:**T**|1+|**T** = tensor(float), tensor(float16)| |Gelu|*in* X:**T**
*out* Y:**T**|1+|**T** = tensor(float), tensor(float16)| |GroupNorm|*in* X:**T**
*in* gamma:**M**
*in* beta:**M**
*out* Y:**T**|1+|**M** = tensor(float), tensor(float16)
**T** = tensor(float), tensor(float16)| +|MultiHeadAttention|*in* query:**T**
*in* key:**T**
*in* value:**T**
*in* bias:**T**
*in* key_padding_mask:**M**
*in* relative_position_bias:**T**
*in* past_key:**T**
*in* past_value:**T**
*out* output:**T**
*out* present_key:**T**
*out* present_value:**T**|1+|**M** = tensor(int32)
**T** = tensor(float), tensor(float16)| |NhwcConv|*in* X:**T**
*in* W:**T**
*in* B:**T**
*out* Y:**T**|1+|**T** = tensor(float), tensor(float16)| |QLinearAdd|*in* A:**T**
*in* A_scale:**tensor(float)**
*in* A_zero_point:**T**
*in* B:**T**
*in* B_scale:**tensor(float)**
*in* B_zero_point:**T**
*in* C_scale:**tensor(float)**
*in* C_zero_point:**T**
*out* C:**T**|1+|**T** = tensor(int8), tensor(uint8)| |QLinearSigmoid|*in* X:**T**
*in* X_scale:**tensor(float)**
*in* X_zero_point:**T**
*in* Y_scale:**tensor(float)**
*in* Y_zero_point:**T**
*out* Y:**T**|1+|**T** = tensor(int8), tensor(uint8)| diff --git a/include/onnxruntime/core/platform/EigenNonBlockingThreadPool.h b/include/onnxruntime/core/platform/EigenNonBlockingThreadPool.h index d3cc805965..7f0046d137 100644 --- a/include/onnxruntime/core/platform/EigenNonBlockingThreadPool.h +++ b/include/onnxruntime/core/platform/EigenNonBlockingThreadPool.h @@ -44,6 +44,7 @@ #include "core/common/inlined_containers_fwd.h" #include "core/common/spin_pause.h" #include "core/platform/ort_mutex.h" +#include "core/platform/ort_spin_lock.h" #include "core/platform/Barrier.h" // ORT thread pool overview @@ -449,7 +450,11 @@ class RunQueue { // PushBack adds w at the end of the queue. // If queue is full returns w, otherwise returns default-constructed Work. Work PushBack(Work w) { +#ifdef USE_LOCK_FREE_QUEUE + std::lock_guard mtx(spin_lock_); +#else std::lock_guard lock(mutex_); +#endif unsigned back = back_.load(std::memory_order_relaxed); Elem& e = array_[(back - 1) & kMask]; ElemState s = e.state.load(std::memory_order_relaxed); @@ -469,7 +474,11 @@ class RunQueue { // with w_idx. Typically the tag will be a per-thread ID to distinguish work // submitted from different threads. PushResult PushBackWithTag(Work w, Tag tag, unsigned& w_idx) { +#ifdef USE_LOCK_FREE_QUEUE + std::lock_guard mtx(spin_lock_); +#else std::lock_guard lock(mutex_); +#endif unsigned back = back_.load(std::memory_order_relaxed); w_idx = (back - 1) & kMask; Elem& e = array_[w_idx]; @@ -490,7 +499,11 @@ class RunQueue { Work PopBack() { if (Empty()) return Work(); +#ifdef USE_LOCK_FREE_QUEUE + std::lock_guard mtx(spin_lock_); +#else std::lock_guard lock(mutex_); +#endif unsigned back; Elem* e; ElemState s; @@ -532,7 +545,11 @@ class RunQueue { bool RevokeWithTag(Tag tag, unsigned w_idx) { bool revoked = false; +#ifdef USE_LOCK_FREE_QUEUE + std::lock_guard mtx(spin_lock_); +#else std::lock_guard lock(mutex_); +#endif Elem& e = array_[w_idx]; ElemState s = e.state.load(std::memory_order_relaxed); @@ -604,7 +621,11 @@ class RunQueue { Work w; }; +#ifdef USE_LOCK_FREE_QUEUE + OrtSpinLock spin_lock_; +#else OrtMutex mutex_; +#endif // Low log(kSize) + 1 bits in front_ and back_ contain rolling index of // front/back, respectively. The remaining bits contain modification counters diff --git a/include/onnxruntime/core/platform/ort_spin_lock.h b/include/onnxruntime/core/platform/ort_spin_lock.h new file mode 100644 index 0000000000..db80abe1ee --- /dev/null +++ b/include/onnxruntime/core/platform/ort_spin_lock.h @@ -0,0 +1,35 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. +#pragma once +#include "core/common/spin_pause.h" +#include + +namespace onnxruntime { +/* +OrtSpinLock implemented mutex semantic "lock-freely", +calling thread will not be put to sleep on blocked, +which reduces cpu usage on context switching. +*/ +struct OrtSpinLock { + using LockState = enum { Locked = 0, + Unlocked }; + + void lock() noexcept { + LockState state = Unlocked; + while (!state_.compare_exchange_weak(state, Locked, std::memory_order_acq_rel, std::memory_order_relaxed)) { + state = Unlocked; + concurrency::SpinPause(); // pause and retry + } + } + bool try_lock() noexcept { + LockState state = Unlocked; + return state_.compare_exchange_weak(state, Locked, std::memory_order_acq_rel, std::memory_order_relaxed); + } + void unlock() noexcept { + state_.store(Unlocked, std::memory_order_release); + } + + private: + std::atomic state_{Unlocked}; +}; +} // namespace onnxruntime \ No newline at end of file diff --git a/onnxruntime/contrib_ops/cuda/bert/attention.cc b/onnxruntime/contrib_ops/cuda/bert/attention.cc index def1508ca2..bd1498c0f9 100644 --- a/onnxruntime/contrib_ops/cuda/bert/attention.cc +++ b/onnxruntime/contrib_ops/cuda/bert/attention.cc @@ -164,6 +164,7 @@ Status Attention::ComputeInternal(OpKernelContext* context) const { has_memory_efficient_attention(sm, sizeof(T) == 2); #else constexpr bool use_memory_efficient_attention = false; + ORT_UNUSED_VARIABLE(is_mask_1d_key_seq_len_start); #endif cublasHandle_t cublas = GetCublasHandle(context); diff --git a/onnxruntime/core/optimizer/graph_transformer_utils.cc b/onnxruntime/core/optimizer/graph_transformer_utils.cc index 73403f2c62..7196465e5d 100644 --- a/onnxruntime/core/optimizer/graph_transformer_utils.cc +++ b/onnxruntime/core/optimizer/graph_transformer_utils.cc @@ -182,6 +182,7 @@ InlinedVector> GenerateTransformers( #ifndef DISABLE_CONTRIB_OPS const InlinedHashSet cpu_ep = {onnxruntime::kCpuExecutionProvider}; #endif + const InlinedHashSet dml_ep = {onnxruntime::kDmlExecutionProvider}; switch (level) { case TransformerLevel::Level1: { // RewriteRule optimizations are the simplest (they generally remove unnecessary nodes and are cheap to run) @@ -308,6 +309,7 @@ InlinedVector> GenerateTransformers( transformers.emplace_back(std::make_unique(cpu_cuda_dml_rocm_eps)); transformers.emplace_back(std::make_unique(cpu_cuda_dml_rocm_eps)); + transformers.emplace_back(std::make_unique(dml_ep)); // GeluApproximation has side effects which may change results. It needs to be manually enabled, // or alternatively the model can be updated offline using a model conversion script diff --git a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/External/DirectMLHelpers/ApiHelpers.h b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/External/DirectMLHelpers/ApiHelpers.h index 76ca37bd05..9a1c23093f 100644 --- a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/External/DirectMLHelpers/ApiHelpers.h +++ b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/External/DirectMLHelpers/ApiHelpers.h @@ -183,13 +183,13 @@ private: static T RoundUpToMultiple(T value, T multiple) { static_assert(std::is_integral_v); - + T remainder = value % multiple; if (remainder != 0) { value += multiple - remainder; } - + return value; } @@ -231,4 +231,4 @@ private: // allocated memory if the fixed stack array is exhausted. FixedBucket m_fixed; std::deque m_dynamic; -}; \ No newline at end of file +}; diff --git a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/External/DirectMLHelpers/ApiTraits.h b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/External/DirectMLHelpers/ApiTraits.h index 8b5cf36937..c75b662af7 100644 --- a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/External/DirectMLHelpers/ApiTraits.h +++ b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/External/DirectMLHelpers/ApiTraits.h @@ -1155,6 +1155,11 @@ struct OperatorDescTraits static constexpr DML_OPERATOR_TYPE Type = DML_OPERATOR_ACTIVATION_GELU; }; +template <> +struct OperatorDescTraits +{ + static constexpr DML_OPERATOR_TYPE Type = DML_OPERATOR_MULTIHEAD_ATTENTION; +}; template struct OperatorTypeTraits @@ -2139,14 +2144,20 @@ struct OperatorTypeTraits<(DML_OPERATOR_TYPE)DML_OPERATOR_ACTIVATION_GELU> using DescType = DML_ACTIVATION_GELU_OPERATOR_DESC; }; +template <> +struct OperatorTypeTraits<(DML_OPERATOR_TYPE)DML_OPERATOR_MULTIHEAD_ATTENTION> +{ + using DescType = DML_MULTIHEAD_ATTENTION_OPERATOR_DESC; +}; + // Calls a visitor functor, supplying an empty operator desc corresponding to the given DML_OPERATOR_TYPE as // the first argument. -// +// // For example: // Visit(DML_OPERATOR_ELEMENT_WISE_IDENTITY, [](auto tag) { // using T = decltype(tag); // T is one of the DML_*_OPERATOR_DESC structs // }); -// +// #pragma warning(push) #pragma warning(disable:4702) template @@ -2432,6 +2443,8 @@ auto OperatorTypeVisitor(DML_OPERATOR_TYPE type, Visitor&& visitor, Ts&&... args return std::invoke(std::forward(visitor), DML_RESAMPLE_GRAD1_OPERATOR_DESC{}, std::forward(args)...); case DML_OPERATOR_DIAGONAL_MATRIX1: return std::invoke(std::forward(visitor), DML_DIAGONAL_MATRIX1_OPERATOR_DESC{}, std::forward(args)...); + case DML_OPERATOR_MULTIHEAD_ATTENTION: + return std::invoke(std::forward(visitor), DML_MULTIHEAD_ATTENTION_OPERATOR_DESC{}, std::forward(args)...); case DML_OPERATOR_ACTIVATION_ELU: return std::invoke(std::forward(visitor), DML_ACTIVATION_ELU_OPERATOR_DESC{}, std::forward(args)...); case DML_OPERATOR_ACTIVATION_CELU: @@ -2633,6 +2646,7 @@ inline gsl::czstring ToString(DML_OPERATOR_TYPE value) case DML_OPERATOR_RESAMPLE2: return "DML_OPERATOR_RESAMPLE2"; case DML_OPERATOR_RESAMPLE_GRAD1: return "DML_OPERATOR_RESAMPLE_GRAD1"; case DML_OPERATOR_DIAGONAL_MATRIX1: return "DML_OPERATOR_DIAGONAL_MATRIX1"; + case DML_OPERATOR_MULTIHEAD_ATTENTION: return "DML_OPERATOR_MULTIHEAD_ATTENTION"; default: assert(false); return ""; diff --git a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/External/DirectMLHelpers/DirectMLSchema.h b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/External/DirectMLHelpers/DirectMLSchema.h index 42de619a87..1ebd52d4ed 100644 --- a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/External/DirectMLHelpers/DirectMLSchema.h +++ b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/External/DirectMLHelpers/DirectMLSchema.h @@ -2302,6 +2302,35 @@ constexpr DML_OPERATOR_SCHEMA DML_DIAGONAL_MATRIX1_OPERATOR_SCHEMA{ DML_DIAGONAL_MATRIX1_OPERATOR_SCHEMA_FIELDS, }; +constexpr DML_SCHEMA_FIELD DML_MULTIHEAD_ATTENTION_OPERATOR_SCHEMA_FIELDS[18] { + DML_SCHEMA_FIELD { DML_SCHEMA_FIELD_KIND_INPUT_TENSOR, DML_SCHEMA_FIELD_TYPE_TENSOR_DESC, "QueryTensor", true }, + DML_SCHEMA_FIELD { DML_SCHEMA_FIELD_KIND_INPUT_TENSOR, DML_SCHEMA_FIELD_TYPE_TENSOR_DESC, "KeyTensor", true }, + DML_SCHEMA_FIELD { DML_SCHEMA_FIELD_KIND_INPUT_TENSOR, DML_SCHEMA_FIELD_TYPE_TENSOR_DESC, "ValueTensor", true }, + DML_SCHEMA_FIELD { DML_SCHEMA_FIELD_KIND_INPUT_TENSOR, DML_SCHEMA_FIELD_TYPE_TENSOR_DESC, "StackedQueryKeyTensor", true }, + DML_SCHEMA_FIELD { DML_SCHEMA_FIELD_KIND_INPUT_TENSOR, DML_SCHEMA_FIELD_TYPE_TENSOR_DESC, "StackedKeyValueTensor", true }, + DML_SCHEMA_FIELD { DML_SCHEMA_FIELD_KIND_INPUT_TENSOR, DML_SCHEMA_FIELD_TYPE_TENSOR_DESC, "StackedQueryKeyValueTensor", true }, + DML_SCHEMA_FIELD { DML_SCHEMA_FIELD_KIND_INPUT_TENSOR, DML_SCHEMA_FIELD_TYPE_TENSOR_DESC, "BiasTensor", true }, + DML_SCHEMA_FIELD { DML_SCHEMA_FIELD_KIND_INPUT_TENSOR, DML_SCHEMA_FIELD_TYPE_TENSOR_DESC, "MaskTensor", true }, + DML_SCHEMA_FIELD { DML_SCHEMA_FIELD_KIND_INPUT_TENSOR, DML_SCHEMA_FIELD_TYPE_TENSOR_DESC, "RelativePositionBiasTensor", true }, + DML_SCHEMA_FIELD { DML_SCHEMA_FIELD_KIND_INPUT_TENSOR, DML_SCHEMA_FIELD_TYPE_TENSOR_DESC, "PastKeyTensor", true }, + DML_SCHEMA_FIELD { DML_SCHEMA_FIELD_KIND_INPUT_TENSOR, DML_SCHEMA_FIELD_TYPE_TENSOR_DESC, "PastValueTensor", true }, + DML_SCHEMA_FIELD { DML_SCHEMA_FIELD_KIND_OUTPUT_TENSOR, DML_SCHEMA_FIELD_TYPE_TENSOR_DESC, "OutputTensor", false }, + DML_SCHEMA_FIELD { DML_SCHEMA_FIELD_KIND_OUTPUT_TENSOR, DML_SCHEMA_FIELD_TYPE_TENSOR_DESC, "OutputPresentKeyTensor", true }, + DML_SCHEMA_FIELD { DML_SCHEMA_FIELD_KIND_OUTPUT_TENSOR, DML_SCHEMA_FIELD_TYPE_TENSOR_DESC, "OutputPresentValueTensor", true }, + DML_SCHEMA_FIELD { DML_SCHEMA_FIELD_KIND_ATTRIBUTE, DML_SCHEMA_FIELD_TYPE_FLOAT, "Scale", false }, + DML_SCHEMA_FIELD { DML_SCHEMA_FIELD_KIND_ATTRIBUTE, DML_SCHEMA_FIELD_TYPE_FLOAT, "MaskFilterValue", false }, + DML_SCHEMA_FIELD { DML_SCHEMA_FIELD_KIND_ATTRIBUTE, DML_SCHEMA_FIELD_TYPE_UINT, "HeadCount", false }, + DML_SCHEMA_FIELD { DML_SCHEMA_FIELD_KIND_ATTRIBUTE, DML_SCHEMA_FIELD_TYPE_UINT, "MaskType", false }, +}; + +constexpr DML_OPERATOR_SCHEMA DML_MULTIHEAD_ATTENTION_OPERATOR_SCHEMA { + "DML_OPERATOR_MULTIHEAD_ATTENTION", + DML_OPERATOR_MULTIHEAD_ATTENTION, + DML_SCHEMA_OPERATOR_SUPPORT_FLAG_NONE, + 18, + DML_MULTIHEAD_ATTENTION_OPERATOR_SCHEMA_FIELDS, +}; + constexpr DML_SCHEMA_FIELD DML_ACTIVATION_ELU_OPERATOR_SCHEMA_FIELDS[3] { DML_SCHEMA_FIELD { DML_SCHEMA_FIELD_KIND_INPUT_TENSOR, DML_SCHEMA_FIELD_TYPE_TENSOR_DESC, "InputTensor", false }, DML_SCHEMA_FIELD { DML_SCHEMA_FIELD_KIND_OUTPUT_TENSOR, DML_SCHEMA_FIELD_TYPE_TENSOR_DESC, "OutputTensor", false }, diff --git a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/External/DirectMLHelpers/DirectMLX.h b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/External/DirectMLHelpers/DirectMLX.h index f7feb5ff21..7e050eef50 100644 --- a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/External/DirectMLHelpers/DirectMLX.h +++ b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/External/DirectMLHelpers/DirectMLX.h @@ -202,7 +202,7 @@ namespace dml }; } -#if DMLX_USE_ABSEIL +#if DMLX_USE_ABSEIL template using Optional = absl::optional; @@ -231,7 +231,7 @@ namespace dml #elif DMLX_USE_GSL template using Span = gsl::span; - #else + #else template using Span = dml::detail::span; #endif @@ -245,11 +245,11 @@ namespace dml #define DMLX_THROW(_hr) THROW_HR(_hr) #else #define DMLX_THROW_IF_FAILED(_hr) if (FAILED(_hr)) { throw std::runtime_error(#_hr); } - #define DMLX_THROW(_hr) throw std::runtime_error(#_hr); + #define DMLX_THROW(_hr) throw std::runtime_error(#_hr); #endif #else #define DMLX_THROW_IF_FAILED(_hr) if (FAILED(_hr)) { std::abort(); } - #define DMLX_THROW(_hr) { std::abort(); } + #define DMLX_THROW(_hr) { std::abort(); } #endif class Graph; @@ -307,7 +307,7 @@ namespace dml // (0, 2, ..., n, 1). This is often referred to as "NHWC" or "interleaved channel" layout. This is useful, // for example, when applied to 2D Convolution to produce outputs in an NHWC layout (as opposed to NCHW, which // is the DirectML default for 2D Convolution). - // + // // Examples of the transposes produced by this policy: // NCW -> NWC // NCHW -> NHWC @@ -713,7 +713,7 @@ namespace dml // Represents an activation to be fused with an existing operator. The meaning of param1 and param2 depend on the // activation to be fused. - // + // // For HARD_SIGMOID, LINEAR, PARAMETRIC_SOFTPLUS, and SCALED_TANH: param1 = Alpha and param2 = Beta // For ELU, LEAKY_RELU, THRESHOLDED_RELU, and CELU: param1 = Alpha. param2 is unused. // For SCALED_ELU, param1 = Alpha and param2 = Gamma. @@ -1858,13 +1858,13 @@ namespace dml } // Helper for setting parameters for the Convolution operator. Sample usage: - // + // // auto conv = dml::ConvolutionBuilder(...) // .StartPadding(...) // .EndPadding(...) // .Strides(...) // .Build(); - // + // // Parameters left unspecified will be defaulted with the same values as dml::Convolution(). class ConvolutionBuilder { @@ -2114,9 +2114,9 @@ namespace dml return output; } - // + // // TODO: LpPooling - // + // // --------------------------------------------------------------------------------------------------------------- @@ -2203,13 +2203,13 @@ namespace dml } // Helper for setting parameters for the MaxPooling operator. Sample usage: - // + // // auto [out, outIndices] = dml::MaxPoolingBuilder(...) // .StartPadding(...) // .EndPadding(...) // .OutputIndices(...) // .Build(); - // + // // Parameters left unspecified will be defaulted with the same values as dml::MaxPooling(). class MaxPoolingBuilder { @@ -2251,13 +2251,13 @@ namespace dml // --------------------------------------------------------------------------------------------------------------- - // + // // TODO: MaxUnpooling - // + // - // + // // TODO: ROIPooling - // + // inline Expression Slice( Expression input, @@ -2683,7 +2683,7 @@ namespace dml { detail::GraphBuilder* builder = input.Impl()->GetGraphBuilder(); TensorDesc inputTensor = input.Impl()->GetOutputDesc(); - + assert(inputTensor.sizes.size() == 4); dml::TensorDesc::Dimensions outputSizes = { @@ -2691,7 +2691,7 @@ namespace dml inputTensor.sizes[1] * blockSize * blockSize, inputTensor.sizes[2] / blockSize, inputTensor.sizes[3] / blockSize - }; + }; TensorDesc outputTensor(inputTensor.dataType, outputSizes, builder->GetTensorPolicy()); @@ -2715,7 +2715,7 @@ namespace dml { detail::GraphBuilder* builder = input.Impl()->GetGraphBuilder(); TensorDesc inputTensor = input.Impl()->GetOutputDesc(); - + assert(inputTensor.sizes.size() == 4); dml::TensorDesc::Dimensions outputSizes = { @@ -2771,7 +2771,7 @@ namespace dml struct TopKOutputs { Expression value; - Expression index; + Expression index; }; inline TopKOutputs TopK(Expression input, uint32_t axis, uint32_t k, DML_AXIS_DIRECTION axisDirection) @@ -2909,14 +2909,14 @@ namespace dml desc.VarianceTensor = varianceTensor.AsPtr(); desc.ScaleTensor = scaleTensor.AsPtr(); desc.Epsilon = epsilon; - + desc.OutputGradientTensor = outputGradientTensor.AsPtr(); desc.OutputScaleGradientTensor = outputScaleGradientTensor.AsPtr(); desc.OutputBiasGradientTensor = outputBiasGradientTensor.AsPtr(); - + dml::detail::NodeOutput* const inputs[] = { input.Impl(), inputGradient.Impl(), mean.Impl(), variance.Impl(), scale.Impl() }; dml::detail::NodeID node = builder->CreateOperatorNode(DML_OPERATOR_BATCH_NORMALIZATION_GRAD, &desc, inputs); - + BatchNormalizationGradOutputs outputValues; outputValues.gradient = builder->CreateNodeOutput(node, 0, *desc.OutputGradientTensor); outputValues.scaleGradient = builder->CreateNodeOutput(node, 1, *desc.OutputScaleGradientTensor); @@ -2932,7 +2932,7 @@ namespace dml { Expression output; Expression mean; - Expression variance; + Expression variance; }; inline BatchNormalizationTrainingOutputs BatchNormalizationTraining( @@ -3005,14 +3005,14 @@ namespace dml desc.VarianceTensor = varianceTensor.AsPtr(); desc.ScaleTensor = scaleTensor.AsPtr(); desc.Epsilon = epsilon; - + desc.OutputGradientTensor = outputGradientTensor.AsPtr(); desc.OutputScaleGradientTensor = outputScaleGradientTensor.AsPtr(); desc.OutputBiasGradientTensor = outputBiasGradientTensor.AsPtr(); - + dml::detail::NodeOutput* const inputs[] = { input.Impl(), inputGradient.Impl(), mean.Impl(), variance.Impl(), scale.Impl() }; dml::detail::NodeID node = builder->CreateOperatorNode(DML_OPERATOR_BATCH_NORMALIZATION_TRAINING_GRAD, &desc, inputs); - + BatchNormalizationGradOutputs outputValues; outputValues.gradient = builder->CreateNodeOutput(node, 0, *desc.OutputGradientTensor); outputValues.scaleGradient = builder->CreateNodeOutput(node, 1, *desc.OutputScaleGradientTensor); @@ -3099,17 +3099,17 @@ namespace dml return output; } - // + // // TODO: LpNormalization - // + // - // + // // TODO: RNN - // + // - // + // // TODO: LSTM - // + // enum class GRUOutputOptions { @@ -3121,7 +3121,7 @@ namespace dml struct GRUOutputs { Expression sequence; - Expression single; + Expression single; }; inline GRUOutputs GRU( @@ -3230,7 +3230,7 @@ namespace dml return { outputSequenceExpr, outputSingleExpr }; } - // + // // TODO: DiagonalMatrix // @@ -3442,33 +3442,33 @@ namespace dml return output; } - // + // // TODO: MatrixMultiplyInteger - // + // - // + // // TODO: QuantizedLinearMatrixMultiply - // + // - // + // // TODO: ConvolutionInteger - // + // - // + // // TODO: QuantizedLinearConvolution - // + // - // + // // TODO: ReluGrad - // + // - // + // // TODO: AveragePoolingGrad - // + // - // + // // TODO: MaxPoolingGrad - // + // struct RandomGeneratorOutputs { @@ -3496,7 +3496,7 @@ namespace dml // Input and output state have the same TensorDesc. desc.OutputStateTensor = inputStateTensor.AsPtr(); } - + RandomGeneratorOutputs out; detail::NodeOutput* const inputs[] = { inputState.Impl() }; @@ -3537,7 +3537,7 @@ namespace dml desc.InputTensor = inputTensor.AsPtr(); desc.OutputCountTensor = outputCountTensor.AsPtr(); desc.OutputCoordinatesTensor = outputCoordinatesTensor.AsPtr(); - + NonZeroCoordinatesOutputs output; detail::NodeOutput* const inputs[] = { input.Impl() }; @@ -3640,17 +3640,17 @@ namespace dml return output; } - // + // // TODO: AdamOptimizer - // + // - // + // // TODO: Argmin - // + // - // + // // TODO: Argmax - // + // #if DML_TARGET_VERSION >= 0x4000 @@ -3694,7 +3694,7 @@ namespace dml desc.ROITensor = roiTensor.AsPtr(); desc.BatchIndicesTensor = batchIndicesTensor.AsPtr(); desc.OutputTensor = outputTensor.AsPtr(); - desc.ReductionFunction = reductionFunction; + desc.ReductionFunction = reductionFunction; desc.InterpolationMode = interpolationMode; desc.SpatialScaleX = spatialScaleX; desc.SpatialScaleY = spatialScaleY; @@ -3763,7 +3763,7 @@ namespace dml outputGradientTensor = TensorDesc(inputGradientTensor.dataType, outputGradientSizes, builder->GetTensorPolicy()); } - + TensorDesc outputROIGradientTensor = computeOutputROIGradient ? TensorDesc(roiTensor.dataType, roiTensor.sizes, builder->GetTensorPolicy()) : TensorDesc(); assert(!computeOutputROIGradient || outputROIGradientTensor.sizes == roiTensor.sizes); @@ -3774,7 +3774,7 @@ namespace dml desc.BatchIndicesTensor = batchIndicesTensor.AsPtr(); desc.OutputGradientTensor = computeOutputGradient ? outputGradientTensor.AsPtr() : nullptr; desc.OutputROIGradientTensor = computeOutputROIGradient ? outputROIGradientTensor.AsPtr() : nullptr; - desc.ReductionFunction = reductionFunction; + desc.ReductionFunction = reductionFunction; desc.InterpolationMode = interpolationMode; desc.SpatialScaleX = spatialScaleX; desc.SpatialScaleY = spatialScaleY; diff --git a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/External/DirectMLHelpers/GeneratedSchemaHelpers.h b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/External/DirectMLHelpers/GeneratedSchemaHelpers.h index aaf02ca146..833871de0b 100644 --- a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/External/DirectMLHelpers/GeneratedSchemaHelpers.h +++ b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/External/DirectMLHelpers/GeneratedSchemaHelpers.h @@ -1418,6 +1418,29 @@ inline std::vector GetFields(const DML_DIAGONAL_MATRIX1_OPERATOR_ OperatorField(&DML_DIAGONAL_MATRIX1_OPERATOR_SCHEMA.Fields[5], ToOperatorFieldType(static_cast(desc.DiagonalFillEnd))), }; } +inline std::vector GetFields(const DML_MULTIHEAD_ATTENTION_OPERATOR_DESC& desc) +{ + return { + OperatorField(&DML_MULTIHEAD_ATTENTION_OPERATOR_SCHEMA.Fields[0], ToOperatorFieldType(static_cast(desc.QueryTensor))), + OperatorField(&DML_MULTIHEAD_ATTENTION_OPERATOR_SCHEMA.Fields[1], ToOperatorFieldType(static_cast(desc.KeyTensor))), + OperatorField(&DML_MULTIHEAD_ATTENTION_OPERATOR_SCHEMA.Fields[2], ToOperatorFieldType(static_cast(desc.ValueTensor))), + OperatorField(&DML_MULTIHEAD_ATTENTION_OPERATOR_SCHEMA.Fields[3], ToOperatorFieldType(static_cast(desc.StackedQueryKeyTensor))), + OperatorField(&DML_MULTIHEAD_ATTENTION_OPERATOR_SCHEMA.Fields[4], ToOperatorFieldType(static_cast(desc.StackedKeyValueTensor))), + OperatorField(&DML_MULTIHEAD_ATTENTION_OPERATOR_SCHEMA.Fields[5], ToOperatorFieldType(static_cast(desc.StackedQueryKeyValueTensor))), + OperatorField(&DML_MULTIHEAD_ATTENTION_OPERATOR_SCHEMA.Fields[6], ToOperatorFieldType(static_cast(desc.BiasTensor))), + OperatorField(&DML_MULTIHEAD_ATTENTION_OPERATOR_SCHEMA.Fields[7], ToOperatorFieldType(static_cast(desc.MaskTensor))), + OperatorField(&DML_MULTIHEAD_ATTENTION_OPERATOR_SCHEMA.Fields[8], ToOperatorFieldType(static_cast(desc.RelativePositionBiasTensor))), + OperatorField(&DML_MULTIHEAD_ATTENTION_OPERATOR_SCHEMA.Fields[9], ToOperatorFieldType(static_cast(desc.PastKeyTensor))), + OperatorField(&DML_MULTIHEAD_ATTENTION_OPERATOR_SCHEMA.Fields[10], ToOperatorFieldType(static_cast(desc.PastValueTensor))), + OperatorField(&DML_MULTIHEAD_ATTENTION_OPERATOR_SCHEMA.Fields[11], ToOperatorFieldType(static_cast(desc.OutputTensor))), + OperatorField(&DML_MULTIHEAD_ATTENTION_OPERATOR_SCHEMA.Fields[12], ToOperatorFieldType(static_cast(desc.OutputPresentKeyTensor))), + OperatorField(&DML_MULTIHEAD_ATTENTION_OPERATOR_SCHEMA.Fields[13], ToOperatorFieldType(static_cast(desc.OutputPresentValueTensor))), + OperatorField(&DML_MULTIHEAD_ATTENTION_OPERATOR_SCHEMA.Fields[14], ToOperatorFieldType(static_cast(desc.Scale))), + OperatorField(&DML_MULTIHEAD_ATTENTION_OPERATOR_SCHEMA.Fields[15], ToOperatorFieldType(static_cast(desc.MaskFilterValue))), + OperatorField(&DML_MULTIHEAD_ATTENTION_OPERATOR_SCHEMA.Fields[16], ToOperatorFieldType(static_cast(desc.HeadCount))), + OperatorField(&DML_MULTIHEAD_ATTENTION_OPERATOR_SCHEMA.Fields[17], ToOperatorFieldType(static_cast(desc.MaskType))), + }; +} inline std::vector GetFields(const DML_ACTIVATION_ELU_OPERATOR_DESC& desc) { return { @@ -1753,6 +1776,7 @@ inline const DML_OPERATOR_SCHEMA& GetSchema(DML_OPERATOR_TYPE operatorType) case DML_OPERATOR_RESAMPLE2: return DML_RESAMPLE2_OPERATOR_SCHEMA; case DML_OPERATOR_RESAMPLE_GRAD1: return DML_RESAMPLE_GRAD1_OPERATOR_SCHEMA; case DML_OPERATOR_DIAGONAL_MATRIX1: return DML_DIAGONAL_MATRIX1_OPERATOR_SCHEMA; + case DML_OPERATOR_MULTIHEAD_ATTENTION: return DML_MULTIHEAD_ATTENTION_OPERATOR_SCHEMA; case DML_OPERATOR_ACTIVATION_ELU: return DML_ACTIVATION_ELU_OPERATOR_SCHEMA; case DML_OPERATOR_ACTIVATION_CELU: return DML_ACTIVATION_CELU_OPERATOR_SCHEMA; case DML_OPERATOR_ACTIVATION_HARDMAX: return DML_ACTIVATION_HARDMAX_OPERATOR_SCHEMA; @@ -2346,6 +2370,10 @@ inline AbstractOperatorDesc ConvertOperatorDesc(const DML_OPERATOR_DESC& opDesc) return AbstractOperatorDesc( &DML_DIAGONAL_MATRIX1_OPERATOR_SCHEMA, GetFields(*static_cast(opDesc.Desc))); + case DML_OPERATOR_MULTIHEAD_ATTENTION: + return AbstractOperatorDesc( + &DML_MULTIHEAD_ATTENTION_OPERATOR_SCHEMA, + GetFields(*static_cast(opDesc.Desc))); case DML_OPERATOR_ACTIVATION_ELU: return AbstractOperatorDesc( &DML_ACTIVATION_ELU_OPERATOR_SCHEMA, diff --git a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/DmlOperatorAttention.cpp b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/DmlOperatorAttention.cpp index 132b099aab..bbebb4a333 100644 --- a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/DmlOperatorAttention.cpp +++ b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/DmlOperatorAttention.cpp @@ -10,32 +10,20 @@ Abbreviations: B is batch_size, S is sequence_length, W is hidden_size M A B C // M, A, B, and C are Inputs | \ | / - Cast Gemm + | Gemm | / | \ | / | \ | / | \ | Slice Slice Slice - Identity | | | + | | | | | | | | | Identity Identity Identity // The identities are used to transpose NCHW -> NHCW while | | | | // keeping the GEMM strides as NCHW to better target metacommands | | | | - | ----- | - ----------- | | - \ | | - Gemm | - | | - | | - Softmax | - | / - | / - \ / - \ / - Gemm - | - ActivationLinear - | - Output // Final output + ----------------- MHA ----- + | + | + Output // Final output This kernel creates a DML_GRAPH, as mentioned above. For reference, refer to this Doc: @@ -49,51 +37,165 @@ public: DmlOperatorAttention(const MLOperatorKernelCreationContext& kernelCreationContext) : DmlOperator(kernelCreationContext) { - ML_CHECK_VALID_ARGUMENT(kernelCreationContext.GetInputCount() >= 3); + enum DmlInputIndex : uint32_t + { + mhaQueryIndex, + mhaKeyIndex, + mhaValueIndex, + mhaStackedQueryKeyIndex, + mhaStackedKeyValueIndex, + mhaStackedQueryKeyValueIndex, + mhaBiasIndex, + mhaMaskIndex, + mhaRelativePositionBiasIndex, + mhaPastKeyIndex, + mhaPastValueIndex, + mhaInputCount, + }; + + enum InputIndex : uint32_t + { + inputIndex, + weightsIndex, + biasIndex, + maskIndex, + pastIndex, + relativePositionBiasIndex, + pastSequenceLengthIndex, + inputCount, + }; + + enum OutputIndex : uint32_t + { + outputIndex, + outputCount, + }; + + ML_CHECK_VALID_ARGUMENT(kernelCreationContext.GetInputCount() >= 2); + ML_CHECK_VALID_ARGUMENT(kernelCreationContext.GetOutputCount() >= 1); + + const uint32_t dmlInputIndex = inputIndex; + const uint32_t dmlWeightsIndex = weightsIndex; + const uint32_t dmlBiasIndex = biasIndex; + const uint32_t dmlMaskIndex = maskIndex; + const uint32_t dmlRelativePositionBiasIndex = relativePositionBiasIndex; + + const bool hasBias = kernelCreationContext.IsInputValid(biasIndex); + const bool hasMask = kernelCreationContext.IsInputValid(maskIndex); + const bool hasUnpaddedBounds = hasMask && kernelCreationContext.GetInputTensorDimensionCount(maskIndex) == 1; + const bool hasRelativePositionBias = kernelCreationContext.IsInputValid(relativePositionBiasIndex); + DmlOperator::Initialize(kernelCreationContext, std::nullopt, std::nullopt, std::nullopt, std::nullopt, 1); - std::vector inputTensorShape = kernelCreationContext.GetTensorShapeDescription().GetInputTensorShape(0); - std::vector weightTensorShape = kernelCreationContext.GetTensorShapeDescription().GetInputTensorShape(1); - std::vector biasTensorShape = kernelCreationContext.GetTensorShapeDescription().GetInputTensorShape(2); - std::vector maskIndexTensorShape = kernelCreationContext.GetTensorShapeDescription().GetInputTensorShape(3); - ML_CHECK_VALID_ARGUMENT(inputTensorShape.size() == 3); - ML_CHECK_VALID_ARGUMENT(weightTensorShape.size() == 2); - ML_CHECK_VALID_ARGUMENT(biasTensorShape.size() == 1); - ML_CHECK_VALID_ARGUMENT(weightTensorShape[1] == biasTensorShape[0]); - ML_CHECK_VALID_ARGUMENT(biasTensorShape[0] % 3 == 0); - ML_CHECK_VALID_ARGUMENT(inputTensorShape[2] == weightTensorShape[0]); - // TODO: fix Attention kernel when maskIndexTensorShape is 1 - // https://microsoft.visualstudio.com/OS/_workitems/edit/41893987 - ML_CHECK_VALID_ARGUMENT(maskIndexTensorShape.size() > 1 && maskIndexTensorShape.size() <= 4); - const uint32_t batchSize = inputTensorShape[0]; - const uint32_t sequenceLength = inputTensorShape[1]; - const uint32_t hiddenSize = biasTensorShape[0] / 3; const uint32_t numHeads = gsl::narrow_cast(kernelCreationContext.GetAttribute(AttrName::NumHeads)); ML_CHECK_VALID_ARGUMENT(numHeads > 0); // to avoid process crash because of division by zero. - ML_CHECK_VALID_ARGUMENT(hiddenSize % numHeads == 0); + + auto inputTensorShape = m_inputTensorDescs[dmlInputIndex].GetSizes(); + ML_CHECK_VALID_ARGUMENT(inputTensorShape.size() == 3); + + auto weightTensorShape = m_inputTensorDescs[dmlWeightsIndex].GetSizes(); + ML_CHECK_VALID_ARGUMENT(weightTensorShape.size() == 2); + ML_CHECK_VALID_ARGUMENT(weightTensorShape[0] == inputTensorShape[2]); + + const auto qkvHiddenSizes = kernelCreationContext.GetOptionalAttributeVectorInt32(AttrName::QkvHiddenSizes); + if (hasBias) + { + auto biasTensorShape = m_inputTensorDescs[dmlBiasIndex].GetSizes(); + ML_CHECK_VALID_ARGUMENT(biasTensorShape.size() == 1); + ML_CHECK_VALID_ARGUMENT(weightTensorShape[1] == biasTensorShape[0]); + + if (qkvHiddenSizes.empty()) + { + ML_CHECK_VALID_ARGUMENT(biasTensorShape[0] % 3 == 0); + } + } + + if (!qkvHiddenSizes.empty()) + { + ML_CHECK_VALID_ARGUMENT(qkvHiddenSizes.size() == 3); + ML_CHECK_VALID_ARGUMENT(qkvHiddenSizes[0] == qkvHiddenSizes[1]); + } + else + { + ML_CHECK_VALID_ARGUMENT(weightTensorShape[1] % 3 == 0); + } + + const uint32_t hiddenSize = qkvHiddenSizes.empty() ? weightTensorShape[1] / 3 : qkvHiddenSizes[0]; + const uint32_t vHiddenSize = qkvHiddenSizes.empty() ? weightTensorShape[1] / 3 : qkvHiddenSizes[2]; const uint32_t headSize = hiddenSize / numHeads; + const uint32_t vHeadSize = vHiddenSize / numHeads; + const uint32_t batchSize = inputTensorShape[0]; + const uint32_t sequenceLength = inputTensorShape[1]; - uint32_t desiredWeightTensorShape[3] = {batchSize, weightTensorShape[0], 3 * hiddenSize}; - uint32_t desiredBiasTensorShape[3] = {batchSize, sequenceLength, 3 * hiddenSize}; - MLOperatorTensorDataType dataType = kernelCreationContext.GetInputEdgeDescription(0).tensorDataType; + uint32_t desiredWeightTensorShape[3] = {batchSize, weightTensorShape[0], hiddenSize + hiddenSize + vHiddenSize}; + MLOperatorTensorDataType dataType = kernelCreationContext.GetInputEdgeDescription(inputIndex).tensorDataType; - // overwrite weightTensorDesc - m_inputTensorDescs[1] = TensorDesc::ConstructBroadcastedTensorDesc(dataType, desiredWeightTensorShape, weightTensorShape); + m_inputTensorDescs[dmlWeightsIndex] = TensorDesc::ConstructBroadcastedTensorDesc(dataType, desiredWeightTensorShape, weightTensorShape); - // overwrite biasTensorDesc - m_inputTensorDescs[2] = TensorDesc::ConstructBroadcastedTensorDesc(dataType, desiredBiasTensorShape, biasTensorShape); + uint32_t desiredBiasTensorShape[3] = {batchSize, sequenceLength, hiddenSize + hiddenSize + vHiddenSize}; + if (hasBias) + { + auto biasTensorShape = m_inputTensorDescs[dmlBiasIndex].GetSizes(); + m_inputTensorDescs[dmlBiasIndex] = TensorDesc::ConstructBroadcastedTensorDesc(dataType, desiredBiasTensorShape, biasTensorShape); + } - // overwrite maskIndexTensorDesc - uint32_t maskIndexDimensionCount = gsl::narrow_cast(maskIndexTensorShape.size()); - maskIndexTensorShape.insert(maskIndexTensorShape.begin() + 1, 4 - maskIndexDimensionCount, 1); - uint32_t desiredMaskIndexShape[4] {batchSize, numHeads, sequenceLength, sequenceLength}; - MLOperatorTensorDataType maskTensorDataType = kernelCreationContext.GetInputEdgeDescription(3).tensorDataType; - m_inputTensorDescs[3] = TensorDesc::ConstructBroadcastedTensorDesc(maskTensorDataType, desiredMaskIndexShape, maskIndexTensorShape); + MLOperatorTensorDataType maskTensorDataType = MLOperatorTensorDataType::Undefined; + bool hasMaxSequenceMask = false; + DML_MULTIHEAD_ATTENTION_MASK_TYPE maskType = DML_MULTIHEAD_ATTENTION_MASK_TYPE_NONE; + if (hasMask) + { + if (hasUnpaddedBounds) + { + auto unpaddedKeyBoundsShape = m_inputTensorDescs[dmlMaskIndex].GetSizes(); + ML_CHECK_VALID_ARGUMENT(unpaddedKeyBoundsShape.size() == 1); - // overwrite output tensor desc - uint32_t outputTensorShape[4] = {batchSize, sequenceLength, numHeads, headSize}; - uint32_t outputTensorStrides[4] = {sequenceLength * numHeads * headSize, headSize, headSize * sequenceLength, 1}; - m_outputTensorDescs[0] = TensorDesc(GetDmlDataTypeFromMlDataType(dataType), outputTensorShape, outputTensorStrides, 0); + const uint32_t batchGroupCount = unpaddedKeyBoundsShape[0] / batchSize; + ML_CHECK_VALID_ARGUMENT(batchGroupCount == 1 || batchGroupCount == 2); + + uint32_t desiredShape[2] = {batchGroupCount, batchSize}; + m_inputTensorDescs[dmlMaskIndex] = TensorDesc( + m_inputTensorDescs[dmlMaskIndex].GetDmlDataType(), + desiredShape); + + maskType = batchGroupCount == 1 + ? DML_MULTIHEAD_ATTENTION_MASK_TYPE_KEY_SEQUENCE_LENGTH + : DML_MULTIHEAD_ATTENTION_MASK_TYPE_KEY_SEQUENCE_END_START; + } + else + { + auto maskIndexTensorShape = m_inputTensorDescs[dmlMaskIndex].GetSizes(); + ML_CHECK_VALID_ARGUMENT(maskIndexTensorShape.size() > 1 && maskIndexTensorShape.size() <= 4); + + maskType = DML_MULTIHEAD_ATTENTION_MASK_TYPE_BOOLEAN; + std::vector reshapedMaskIndexTensorShape(maskIndexTensorShape.begin(), maskIndexTensorShape.end()); + if (maskIndexTensorShape.size() == 4 && maskIndexTensorShape[2] != sequenceLength) + { + hasMaxSequenceMask = true; + ML_CHECK_VALID_ARGUMENT(maskIndexTensorShape[2] == maskIndexTensorShape[3]); + const uint32_t maxSequenceLength = maskIndexTensorShape[2]; + uint32_t desiredMaskIndexShape[4] {batchSize, numHeads, maxSequenceLength, maxSequenceLength}; + maskTensorDataType = kernelCreationContext.GetInputEdgeDescription(maskIndex).tensorDataType; + m_inputTensorDescs[dmlMaskIndex] = TensorDesc::ConstructBroadcastedTensorDesc(maskTensorDataType, desiredMaskIndexShape, reshapedMaskIndexTensorShape); + } + else + { + uint32_t maskIndexDimensionCount = gsl::narrow_cast(maskIndexTensorShape.size()); + reshapedMaskIndexTensorShape.insert(reshapedMaskIndexTensorShape.begin() + 1, 4 - maskIndexDimensionCount, 1); + uint32_t desiredMaskIndexShape[4] {batchSize, numHeads, sequenceLength, sequenceLength}; + maskTensorDataType = kernelCreationContext.GetInputEdgeDescription(maskIndex).tensorDataType; + m_inputTensorDescs[dmlMaskIndex] = TensorDesc::ConstructBroadcastedTensorDesc(maskTensorDataType, desiredMaskIndexShape, reshapedMaskIndexTensorShape); + } + } + } + + if (hasRelativePositionBias) + { + auto relativePositionBiasTensorShape = m_inputTensorDescs[dmlRelativePositionBiasIndex].GetSizes(); + ML_CHECK_VALID_ARGUMENT(relativePositionBiasTensorShape.size() == 4); + ML_CHECK_VALID_ARGUMENT(relativePositionBiasTensorShape[0] == inputTensorShape[0]); + ML_CHECK_VALID_ARGUMENT(relativePositionBiasTensorShape[1] == numHeads); + ML_CHECK_VALID_ARGUMENT(relativePositionBiasTensorShape[2] == inputTensorShape[1]); + } TensorDesc firstGemmOutputTensorDesc = TensorDesc::ConstructDefaultTensorDesc(dataType, desiredBiasTensorShape); DML_TENSOR_DESC namedFirstGemmOutputTensorDesc = firstGemmOutputTensorDesc.GetDmlDesc(); @@ -101,345 +203,335 @@ public: std::vector inputDescs = GetDmlInputDescs(); std::vector outputDescs = GetDmlOutputDescs(); - DML_GEMM_OPERATOR_DESC xWeightOperatorDesc = {}; - xWeightOperatorDesc.ATensor = &inputDescs[0]; - xWeightOperatorDesc.BTensor = &inputDescs[1]; - xWeightOperatorDesc.CTensor = &inputDescs[2]; - xWeightOperatorDesc.OutputTensor = &namedFirstGemmOutputTensorDesc; - xWeightOperatorDesc.TransA = DML_MATRIX_TRANSFORM_NONE; - xWeightOperatorDesc.TransB = DML_MATRIX_TRANSFORM_NONE; - xWeightOperatorDesc.Alpha = 1.0f; - xWeightOperatorDesc.Beta = 1.0f; - xWeightOperatorDesc.FusedActivation = nullptr; - const DML_OPERATOR_DESC xWeightDesc {DML_OPERATOR_GEMM, &xWeightOperatorDesc}; + DML_GEMM_OPERATOR_DESC gemmOperatorDesc = {}; + gemmOperatorDesc.ATensor = &inputDescs[0]; + gemmOperatorDesc.BTensor = &inputDescs[1]; - - std::array querySlicedTensorShape {batchSize, sequenceLength, hiddenSize}; - TensorDesc querySlicedInputTensorDesc = TensorDesc::ConstructDefaultTensorDesc(dataType, querySlicedTensorShape); - DML_TENSOR_DESC namedQuerySlicedInputTensorDesc = querySlicedInputTensorDesc.GetDmlDesc(); - - std::array querySliceOffset = {0, 0, 0}; - std::array keySliceOffset = {0, 0, hiddenSize}; - std::array valueSliceOffset = {0, 0, 2 * hiddenSize}; - std::array sliceSize = {batchSize, sequenceLength, hiddenSize}; - std::array strides = {1, 1, 1}; - DML_SLICE1_OPERATOR_DESC querySlicedOperatorDesc = {}; - querySlicedOperatorDesc.InputTensor = &namedFirstGemmOutputTensorDesc; - querySlicedOperatorDesc.OutputTensor = &namedQuerySlicedInputTensorDesc; - querySlicedOperatorDesc.DimensionCount = gsl::narrow_cast(querySlicedTensorShape.size()); - querySlicedOperatorDesc.InputWindowOffsets = querySliceOffset.data(); - querySlicedOperatorDesc.InputWindowSizes = sliceSize.data(); - querySlicedOperatorDesc.InputWindowStrides = strides.data(); - const DML_OPERATOR_DESC querySlicedDesc = { DML_OPERATOR_SLICE1, &querySlicedOperatorDesc }; - - DML_SLICE1_OPERATOR_DESC keySlicedOperatorDesc = {}; - keySlicedOperatorDesc.InputTensor = &namedFirstGemmOutputTensorDesc; - keySlicedOperatorDesc.OutputTensor = &namedQuerySlicedInputTensorDesc; - keySlicedOperatorDesc.DimensionCount = gsl::narrow_cast(querySlicedTensorShape.size()); - keySlicedOperatorDesc.InputWindowOffsets = keySliceOffset.data(); - keySlicedOperatorDesc.InputWindowSizes = sliceSize.data(); - keySlicedOperatorDesc.InputWindowStrides = strides.data(); - const DML_OPERATOR_DESC keySlicedDesc = { DML_OPERATOR_SLICE1, &keySlicedOperatorDesc }; - - DML_SLICE1_OPERATOR_DESC valueSlicedOperatorDesc = {}; - valueSlicedOperatorDesc.InputTensor = &namedFirstGemmOutputTensorDesc; - valueSlicedOperatorDesc.OutputTensor = &namedQuerySlicedInputTensorDesc; - valueSlicedOperatorDesc.DimensionCount = gsl::narrow_cast(querySlicedTensorShape.size()); - valueSlicedOperatorDesc.InputWindowOffsets = valueSliceOffset.data(); - valueSlicedOperatorDesc.InputWindowSizes = sliceSize.data(); - valueSlicedOperatorDesc.InputWindowStrides = strides.data(); - const DML_OPERATOR_DESC valueSlicedDesc = { DML_OPERATOR_SLICE1, &valueSlicedOperatorDesc}; - - TensorDesc castedMaskIndexTensorDesc = TensorDesc::ConstructDefaultTensorDesc(dataType, desiredMaskIndexShape); - DML_TENSOR_DESC namedCastedMaskIndexTensorDesc = castedMaskIndexTensorDesc.GetDmlDesc(); - - DML_CAST_OPERATOR_DESC castMaskIndexOperatorDesc = {}; - castMaskIndexOperatorDesc.InputTensor = &inputDescs[3]; - castMaskIndexOperatorDesc.OutputTensor = &namedCastedMaskIndexTensorDesc; - const DML_OPERATOR_DESC castMaskIndexDesc = {DML_OPERATOR_CAST, &castMaskIndexOperatorDesc}; - - // The attention fusion in ORT expects this to be number to -10000. - // https://github.com/microsoft/onnxruntime/blob/main/onnxruntime/core/optimizer/attention_fusion_helper.h#L604 - // The decomposed Attention performs: (M - 1.0) * -10000.0, where M is the 4th input of the Attention node. - // Above equation can be written as (M * -1000) + 10000.0 - DML_SCALE_BIAS scaleBias = {}; - scaleBias.Scale = -10000.0f; - scaleBias.Bias = 10000.0f; - DML_ELEMENT_WISE_IDENTITY_OPERATOR_DESC maskOperatorDesc = {}; - maskOperatorDesc.InputTensor = &namedCastedMaskIndexTensorDesc; - maskOperatorDesc.OutputTensor = &namedCastedMaskIndexTensorDesc; - maskOperatorDesc.ScaleBias = &scaleBias; - const DML_OPERATOR_DESC maskDesc = {DML_OPERATOR_ELEMENT_WISE_IDENTITY, &maskOperatorDesc}; - - // original reshaped shape: [batchSize, seqenceLength, numHeads, headSize] - // transposed shape to [0, 2, 1, 3] -> [batchSize, numHeads, sequenceLength, headSize] - uint32_t reshapedTransposedQueryTensorShape[4] = {batchSize, numHeads, sequenceLength, headSize}; - uint32_t reshapedTransposedQueryTensorStrides[4] = {sequenceLength * numHeads * headSize, headSize, numHeads * headSize, 1}; - TensorDesc reshapedTransposedQueryTensorDesc = TensorDesc( - GetDmlDataTypeFromMlDataType(dataType), - reshapedTransposedQueryTensorShape, - reshapedTransposedQueryTensorStrides); - DML_TENSOR_DESC namedReshapedTransposedQueryTensorDesc = reshapedTransposedQueryTensorDesc.GetDmlDesc(); - - TensorDesc reshapedTransposedQueryOutputTensorDesc = TensorDesc( - GetDmlDataTypeFromMlDataType(dataType), - reshapedTransposedQueryTensorShape); - DML_TENSOR_DESC namedReshapedTransposedQueryOutputTensorDesc = reshapedTransposedQueryOutputTensorDesc.GetDmlDesc(); - - DML_ELEMENT_WISE_IDENTITY_OPERATOR_DESC transposedQueryOperatorDesc{}; - transposedQueryOperatorDesc.InputTensor = &namedReshapedTransposedQueryTensorDesc; - transposedQueryOperatorDesc.OutputTensor = &namedReshapedTransposedQueryOutputTensorDesc; - const DML_OPERATOR_DESC transposedQueryDesc {DML_OPERATOR_ELEMENT_WISE_IDENTITY, &transposedQueryOperatorDesc}; - - uint32_t reshapedTransposedKeyTensorShape[4] = {batchSize, numHeads, headSize, sequenceLength}; - uint32_t reshapedTransposedKeyTensorStrides[4] = {sequenceLength * numHeads * headSize, headSize, 1, numHeads * headSize}; - TensorDesc reshapedTransposedKeyTensorDesc = TensorDesc( - GetDmlDataTypeFromMlDataType(dataType), - reshapedTransposedKeyTensorShape, - reshapedTransposedKeyTensorStrides); - DML_TENSOR_DESC namedReshapedTransposedKeyTensorDesc = reshapedTransposedKeyTensorDesc.GetDmlDesc(); - - TensorDesc reshapedTransposedKeyOutputTensorDesc = TensorDesc( - GetDmlDataTypeFromMlDataType(dataType), - reshapedTransposedKeyTensorShape); - DML_TENSOR_DESC namedReshapedTransposedKeyOutputTensorDesc = reshapedTransposedKeyOutputTensorDesc.GetDmlDesc(); - - DML_ELEMENT_WISE_IDENTITY_OPERATOR_DESC transposedKeyOperatorDesc{}; - transposedKeyOperatorDesc.InputTensor = &namedReshapedTransposedKeyTensorDesc; - transposedKeyOperatorDesc.OutputTensor = &namedReshapedTransposedKeyOutputTensorDesc; - const DML_OPERATOR_DESC transposedKeyDesc {DML_OPERATOR_ELEMENT_WISE_IDENTITY, &transposedKeyOperatorDesc}; - - uint32_t queryKeyTensorShape[4] = {batchSize, numHeads, sequenceLength, sequenceLength}; - TensorDesc queryKeyTensorDesc = TensorDesc::ConstructDefaultTensorDesc(dataType, queryKeyTensorShape); - DML_TENSOR_DESC namedQueryKeyTensorDesc = queryKeyTensorDesc.GetDmlDesc(); - - float alpha = static_cast(1 / sqrt(headSize)); - DML_GEMM_OPERATOR_DESC attentionScoreOperatorDesc = {}; - attentionScoreOperatorDesc.ATensor = &namedReshapedTransposedQueryOutputTensorDesc; - attentionScoreOperatorDesc.BTensor = &namedReshapedTransposedKeyOutputTensorDesc; - attentionScoreOperatorDesc.CTensor = &namedCastedMaskIndexTensorDesc; - attentionScoreOperatorDesc.OutputTensor = &namedQueryKeyTensorDesc; - attentionScoreOperatorDesc.TransA = DML_MATRIX_TRANSFORM_NONE; - attentionScoreOperatorDesc.TransB = DML_MATRIX_TRANSFORM_NONE; - attentionScoreOperatorDesc.Alpha = alpha; - attentionScoreOperatorDesc.Beta = 0.0f; - attentionScoreOperatorDesc.FusedActivation = nullptr; - const DML_OPERATOR_DESC attentionScoreDesc {DML_OPERATOR_GEMM, &attentionScoreOperatorDesc}; - - std::array axes = {3}; - DML_ACTIVATION_SOFTMAX1_OPERATOR_DESC softmaxOperatorDesc = {}; - softmaxOperatorDesc.InputTensor = &namedQueryKeyTensorDesc; - softmaxOperatorDesc.OutputTensor = &namedQueryKeyTensorDesc; - softmaxOperatorDesc.AxisCount = gsl::narrow_cast(axes.size()); - softmaxOperatorDesc.Axes = axes.data(); - const DML_OPERATOR_DESC softmaxDesc = {DML_OPERATOR_ACTIVATION_SOFTMAX1, &softmaxOperatorDesc}; - - uint32_t reshapedTransposedOutputTensorShape[4] {batchSize, numHeads, sequenceLength, headSize}; - uint32_t reshapedTransposedOutputTensorStrides[4] {sequenceLength * numHeads * headSize, headSize * sequenceLength, headSize, 1}; - TensorDesc reshapedTransposedOutputTensorDesc = TensorDesc( - GetDmlDataTypeFromMlDataType(dataType), - reshapedTransposedOutputTensorShape, - reshapedTransposedOutputTensorStrides, - 0 // guaranteedBaseOffsetAlignment - ); - DML_TENSOR_DESC namedReshapedTransposedOutputTensorDesc = reshapedTransposedOutputTensorDesc.GetDmlDesc(); - - DML_GEMM_OPERATOR_DESC attentionWeightOperatorDesc = {}; - attentionWeightOperatorDesc.ATensor = &namedQueryKeyTensorDesc; - attentionWeightOperatorDesc.BTensor = &namedReshapedTransposedQueryOutputTensorDesc; - attentionWeightOperatorDesc.CTensor = nullptr; - attentionWeightOperatorDesc.OutputTensor = &namedReshapedTransposedOutputTensorDesc; - attentionWeightOperatorDesc.TransA = DML_MATRIX_TRANSFORM_NONE; - attentionWeightOperatorDesc.TransB = DML_MATRIX_TRANSFORM_NONE; - attentionWeightOperatorDesc.Alpha = 1.0f; - attentionWeightOperatorDesc.Beta = 0.0f; - attentionWeightOperatorDesc.FusedActivation = nullptr; - const DML_OPERATOR_DESC attentionWeightDesc {DML_OPERATOR_GEMM, &attentionWeightOperatorDesc}; - - TensorDesc transposedOutputTensorDesc = TensorDesc( - m_outputTensorDescs[0].GetDmlDataType(), - m_outputTensorDescs[0].GetSizes(), - std::nullopt, - 0 // guaranteedBaseOffsetAlignment - ); - DML_TENSOR_DESC namedTransposedOutputTensorDesc = transposedOutputTensorDesc.GetDmlDesc(); - - DML_ACTIVATION_LINEAR_OPERATOR_DESC outputOperatorDesc = {}; - outputOperatorDesc.Alpha = 1.0f; - outputOperatorDesc.Beta = 0.0f; - outputOperatorDesc.InputTensor = &outputDescs[0]; - outputOperatorDesc.OutputTensor = &namedTransposedOutputTensorDesc; - const DML_OPERATOR_DESC outputDesc {DML_OPERATOR_ACTIVATION_LINEAR, &outputOperatorDesc}; - - enum NodeIndex : uint32_t + if (hasBias) { - xWeight, - querySlice, - keySlice, - valueSlice, - queryTranspose, - keyTranspose, - attentionScore, - softmax, - valueTranspose, - attentionWeight, - castMaskIndex, - mask, - output, - count, + gemmOperatorDesc.CTensor = &inputDescs[2]; + } + + gemmOperatorDesc.OutputTensor = &namedFirstGemmOutputTensorDesc; + gemmOperatorDesc.TransA = DML_MATRIX_TRANSFORM_NONE; + gemmOperatorDesc.TransB = DML_MATRIX_TRANSFORM_NONE; + gemmOperatorDesc.Alpha = 1.0f; + gemmOperatorDesc.Beta = 1.0f; + gemmOperatorDesc.FusedActivation = nullptr; + const DML_OPERATOR_DESC gemmDesc {DML_OPERATOR_GEMM, &gemmOperatorDesc}; + + std::array queryKeySlicedTensorShape {batchSize, sequenceLength, hiddenSize + hiddenSize}; + TensorDesc queryKeySlicedInputTensorDesc = TensorDesc::ConstructDefaultTensorDesc(dataType, queryKeySlicedTensorShape); + DML_TENSOR_DESC namedQueryKeySlicedInputTensorDesc = queryKeySlicedInputTensorDesc.GetDmlDesc(); + + std::array valueSlicedTensorShape {batchSize, sequenceLength, vHiddenSize}; + TensorDesc valueSlicedInputTensorDesc = TensorDesc::ConstructDefaultTensorDesc(dataType, valueSlicedTensorShape); + DML_TENSOR_DESC namedValueSlicedInputTensorDesc = valueSlicedInputTensorDesc.GetDmlDesc(); + + // Transpose slice QK from [batchSize, sequenceLength, 2, numHeads, headSize] to [batchSize, sequenceLength, numHeads, 2, headSize] + std::array queryKeyTransposedTensorShape {batchSize, sequenceLength, numHeads, 2, headSize}; + std::array queryKeyTransposedStrides { + sequenceLength * numHeads * 2 * headSize, + numHeads * 2 * headSize, + headSize, + numHeads * headSize, + 1, }; + TensorDesc queryKeyTransposedInputTensorDesc = TensorDesc( + m_inputTensorDescs[dmlInputIndex].GetDmlDataType(), + queryKeyTransposedTensorShape, + queryKeyTransposedStrides); + DML_TENSOR_DESC namedQueryKeyTransposedInputTensorDesc = queryKeyTransposedInputTensorDesc.GetDmlDesc(); + + TensorDesc queryKeyTransposedOutputTensorDesc = TensorDesc( + m_inputTensorDescs[dmlInputIndex].GetDmlDataType(), + queryKeyTransposedTensorShape); + DML_TENSOR_DESC namedQueryKeyTransposedOutputTensorDesc = queryKeyTransposedOutputTensorDesc.GetDmlDesc(); + + // Transpose QKV from [batchSize, sequenceLength, 3, numHeads, headSize] to [batchSize, sequenceLength, numHeads, 3, headSize] + std::array queryKeyValueTransposedTensorShape {batchSize, sequenceLength, numHeads, 3, headSize}; + std::array queryKeyValueTransposedStrides { + sequenceLength * numHeads * 3 * headSize, + numHeads * 3 * headSize, + headSize, + numHeads * headSize, + 1, + }; + + TensorDesc queryKeyValueTransposedInputTensorDesc = TensorDesc( + m_inputTensorDescs[dmlInputIndex].GetDmlDataType(), + queryKeyValueTransposedTensorShape, + queryKeyValueTransposedStrides); + DML_TENSOR_DESC namedQueryKeyValueTransposedInputTensorDesc = queryKeyValueTransposedInputTensorDesc.GetDmlDesc(); + + TensorDesc queryKeyValueTransposedOutputTensorDesc = TensorDesc( + m_inputTensorDescs[dmlInputIndex].GetDmlDataType(), + queryKeyValueTransposedTensorShape); + DML_TENSOR_DESC namedQueryKeyValueTransposedOutputTensorDesc = queryKeyValueTransposedOutputTensorDesc.GetDmlDesc(); + + std::array queryKeySliceOffset = {0, 0, 0}; + std::array queryKeySliceSize = {batchSize, sequenceLength, hiddenSize + hiddenSize}; + std::array queryKeySliceStrides = {1, 1, 1}; + + std::array valueSliceOffset = {0, 0, 2 * hiddenSize}; + std::array valueSliceSize = {batchSize, sequenceLength, vHiddenSize}; + std::array valueSliceStrides = {1, 1, 1}; + const bool hasSlicedValue = hiddenSize != vHiddenSize; + + // We need to slice the value tensor when its hidden size is different from the query and key + DML_SLICE1_OPERATOR_DESC queryKeySlicedOperatorDesc = {}; + DML_SLICE1_OPERATOR_DESC valueSlicedOperatorDesc = {}; + DML_ELEMENT_WISE_IDENTITY_OPERATOR_DESC transposeOperatorDesc = {}; + if (hasSlicedValue) + { + queryKeySlicedOperatorDesc.InputTensor = &namedFirstGemmOutputTensorDesc; + queryKeySlicedOperatorDesc.OutputTensor = &namedQueryKeySlicedInputTensorDesc; + queryKeySlicedOperatorDesc.DimensionCount = gsl::narrow_cast(queryKeySlicedTensorShape.size()); + queryKeySlicedOperatorDesc.InputWindowOffsets = queryKeySliceOffset.data(); + queryKeySlicedOperatorDesc.InputWindowSizes = queryKeySliceSize.data(); + queryKeySlicedOperatorDesc.InputWindowStrides = queryKeySliceStrides.data(); + + valueSlicedOperatorDesc.InputTensor = &namedFirstGemmOutputTensorDesc; + valueSlicedOperatorDesc.OutputTensor = &namedValueSlicedInputTensorDesc; + valueSlicedOperatorDesc.DimensionCount = gsl::narrow_cast(valueSlicedTensorShape.size()); + valueSlicedOperatorDesc.InputWindowOffsets = valueSliceOffset.data(); + valueSlicedOperatorDesc.InputWindowSizes = valueSliceSize.data(); + valueSlicedOperatorDesc.InputWindowStrides = valueSliceStrides.data(); + + transposeOperatorDesc.InputTensor = &namedQueryKeyTransposedInputTensorDesc; + transposeOperatorDesc.OutputTensor = &namedQueryKeyTransposedOutputTensorDesc; + } + else + { + // When Q/K/V all have the same hidden size, we just have to transpose it before sending it to MHA + transposeOperatorDesc.InputTensor = &namedQueryKeyValueTransposedInputTensorDesc; + transposeOperatorDesc.OutputTensor = &namedQueryKeyValueTransposedOutputTensorDesc; + } + const DML_OPERATOR_DESC queryKeySlicedDesc = { DML_OPERATOR_SLICE1, &queryKeySlicedOperatorDesc}; + const DML_OPERATOR_DESC valueSlicedDesc = { DML_OPERATOR_SLICE1, &valueSlicedOperatorDesc}; + const DML_OPERATOR_DESC transposedDesc = { DML_OPERATOR_ELEMENT_WISE_IDENTITY, &transposeOperatorDesc}; + + std::array maskSliceOutputShape {batchSize, numHeads, sequenceLength, sequenceLength}; + std::array maskSliceStrides = {1, 1, 1, 1}; + std::array maskSliceOffsets = {0, 0, 0, 0}; + TensorDesc maskSliceOutputTensorDesc; + DML_TENSOR_DESC namedMaskSliceOutputTensorDesc; + + DML_SLICE1_OPERATOR_DESC maskSlicedOperatorDesc = {}; + if (hasMaxSequenceMask) + { + maskSliceOutputTensorDesc = TensorDesc::ConstructDefaultTensorDesc(maskTensorDataType, maskSliceOutputShape); + namedMaskSliceOutputTensorDesc = maskSliceOutputTensorDesc.GetDmlDesc(); + maskSlicedOperatorDesc.InputTensor = &inputDescs[dmlMaskIndex]; + maskSlicedOperatorDesc.OutputTensor = &namedMaskSliceOutputTensorDesc; + maskSlicedOperatorDesc.DimensionCount = gsl::narrow_cast(maskSliceOutputShape.size()); + maskSlicedOperatorDesc.InputWindowOffsets = maskSliceOffsets.data(); + maskSlicedOperatorDesc.InputWindowSizes = maskSliceOutputShape.data(); + maskSlicedOperatorDesc.InputWindowStrides = maskSliceStrides.data(); + } + const DML_OPERATOR_DESC maskSlicedDesc = { DML_OPERATOR_SLICE1, &maskSlicedOperatorDesc}; + + DML_MULTIHEAD_ATTENTION_OPERATOR_DESC mhaOperatorDesc = {}; + mhaOperatorDesc.ValueTensor = hasSlicedValue ? &namedValueSlicedInputTensorDesc : nullptr; + mhaOperatorDesc.StackedQueryKeyTensor = hasSlicedValue ? &namedQueryKeyTransposedOutputTensorDesc : nullptr; + mhaOperatorDesc.StackedQueryKeyValueTensor = hasSlicedValue ? nullptr : &namedQueryKeyValueTransposedOutputTensorDesc; + + if (hasMaxSequenceMask) + { + mhaOperatorDesc.MaskTensor = &namedMaskSliceOutputTensorDesc; + } + else + { + mhaOperatorDesc.MaskTensor = hasMask ? &inputDescs[dmlMaskIndex] : nullptr; + } + + mhaOperatorDesc.RelativePositionBiasTensor = hasRelativePositionBias ? &inputDescs[dmlRelativePositionBiasIndex] : nullptr; + mhaOperatorDesc.OutputTensor = &outputDescs[outputIndex]; + mhaOperatorDesc.Scale = kernelCreationContext.GetOptionalAttribute(AttrName::Scale, gsl::narrow_cast(1.0f / std::sqrt(headSize))); + mhaOperatorDesc.MaskFilterValue = kernelCreationContext.GetOptionalAttribute(AttrName::MaskFilterValue, -10'000.0f); + mhaOperatorDesc.HeadCount = numHeads; + mhaOperatorDesc.MaskType = maskType; + const DML_OPERATOR_DESC mhaDesc = { DML_OPERATOR_MULTIHEAD_ATTENTION, &mhaOperatorDesc }; + + // Construct the graph + std::vector inputEdges; + std::vector intermediateEdges; + std::vector outputEdges; + + std::vector opDescs = { + &gemmDesc, + &mhaDesc, + }; + + uint32_t currentNodeIndex = 0; + const uint32_t gemmNodeIndex = currentNodeIndex++; + const uint32_t mhaNodeIndex = currentNodeIndex++; + + uint32_t valueSliceNodeIndex = 0; + uint32_t queryKeySliceNodeIndex = 0; + uint32_t queryKeyTransposedNodeIndex = 0; + uint32_t queryKeyValueTransposedNodeIndex = 0; + if (hasSlicedValue) + { + opDescs.push_back(&queryKeySlicedDesc); + queryKeySliceNodeIndex = currentNodeIndex++; + + opDescs.push_back(&valueSlicedDesc); + valueSliceNodeIndex = currentNodeIndex++; + + opDescs.push_back(&transposedDesc); + queryKeyTransposedNodeIndex = currentNodeIndex++; + } + else + { + opDescs.push_back(&transposedDesc); + queryKeyValueTransposedNodeIndex = currentNodeIndex++; + } + + uint32_t maskSliceNodeIndex = 0; + if (hasMaxSequenceMask) + { + opDescs.push_back(&maskSlicedDesc); + maskSliceNodeIndex = currentNodeIndex++; + } + + DML_INPUT_GRAPH_EDGE_DESC inputToGemmEdge = {}; + inputToGemmEdge.GraphInputIndex = dmlInputIndex; + inputToGemmEdge.ToNodeIndex = gemmNodeIndex; + inputToGemmEdge.ToNodeInputIndex = 0; + inputEdges.push_back(inputToGemmEdge); + + DML_INPUT_GRAPH_EDGE_DESC weightToGemmEdge = {}; + weightToGemmEdge.GraphInputIndex = dmlWeightsIndex; + weightToGemmEdge.ToNodeIndex = gemmNodeIndex; + weightToGemmEdge.ToNodeInputIndex = 1; + inputEdges.push_back(weightToGemmEdge); + + if (hasBias) + { + DML_INPUT_GRAPH_EDGE_DESC biasToGemmEdge = {}; + biasToGemmEdge.GraphInputIndex = dmlBiasIndex; + biasToGemmEdge.ToNodeIndex = gemmNodeIndex; + biasToGemmEdge.ToNodeInputIndex = 2; + inputEdges.push_back(biasToGemmEdge); + } + + if (hasMask) + { + if (hasUnpaddedBounds) + { + DML_INPUT_GRAPH_EDGE_DESC maskToMhaEdge = {}; + maskToMhaEdge.GraphInputIndex = dmlMaskIndex; + maskToMhaEdge.ToNodeIndex = mhaNodeIndex; + maskToMhaEdge.ToNodeInputIndex = mhaMaskIndex; + inputEdges.push_back(maskToMhaEdge); + } + else if (hasMaxSequenceMask) + { + DML_INPUT_GRAPH_EDGE_DESC maskToMaskSliceEdge = {}; + maskToMaskSliceEdge.GraphInputIndex = dmlMaskIndex; + maskToMaskSliceEdge.ToNodeIndex = maskSliceNodeIndex; + maskToMaskSliceEdge.ToNodeInputIndex = 0; + inputEdges.push_back(maskToMaskSliceEdge); + + DML_INTERMEDIATE_GRAPH_EDGE_DESC maskSliceToMhaEdge = {}; + maskSliceToMhaEdge.FromNodeIndex = maskSliceNodeIndex; + maskSliceToMhaEdge.FromNodeOutputIndex = 0; + maskSliceToMhaEdge.ToNodeIndex = mhaNodeIndex; + maskSliceToMhaEdge.ToNodeInputIndex = mhaMaskIndex; + intermediateEdges.push_back(maskSliceToMhaEdge); + } + else + { + DML_INPUT_GRAPH_EDGE_DESC maskToMhaEdge = {}; + maskToMhaEdge.GraphInputIndex = dmlMaskIndex; + maskToMhaEdge.ToNodeIndex = mhaNodeIndex; + maskToMhaEdge.ToNodeInputIndex = mhaMaskIndex; + inputEdges.push_back(maskToMhaEdge); + } + } + + if (hasRelativePositionBias) + { + DML_INPUT_GRAPH_EDGE_DESC relativePositionBiasToMhaEdge = {}; + relativePositionBiasToMhaEdge.GraphInputIndex = dmlRelativePositionBiasIndex; + relativePositionBiasToMhaEdge.ToNodeIndex = mhaNodeIndex; + relativePositionBiasToMhaEdge.ToNodeInputIndex = mhaRelativePositionBiasIndex; + inputEdges.push_back(relativePositionBiasToMhaEdge); + } + + if (hasSlicedValue) + { + // We need to slice QK and V, and transpose QK + DML_INTERMEDIATE_GRAPH_EDGE_DESC gemmToQueryKeySliceEdge = {}; + gemmToQueryKeySliceEdge.FromNodeIndex = gemmNodeIndex; + gemmToQueryKeySliceEdge.FromNodeOutputIndex = 0; + gemmToQueryKeySliceEdge.ToNodeIndex = queryKeySliceNodeIndex; + gemmToQueryKeySliceEdge.ToNodeInputIndex = 0; + intermediateEdges.push_back(gemmToQueryKeySliceEdge); + + DML_INTERMEDIATE_GRAPH_EDGE_DESC queryKeySliceToTransposeEdge = {}; + queryKeySliceToTransposeEdge.FromNodeIndex = queryKeySliceNodeIndex; + queryKeySliceToTransposeEdge.FromNodeOutputIndex = 0; + queryKeySliceToTransposeEdge.ToNodeIndex = queryKeyTransposedNodeIndex; + queryKeySliceToTransposeEdge.ToNodeInputIndex = 0; + intermediateEdges.push_back(queryKeySliceToTransposeEdge); + + DML_INTERMEDIATE_GRAPH_EDGE_DESC queryKeyTransposedToMhaEdge = {}; + queryKeyTransposedToMhaEdge.FromNodeIndex = queryKeyTransposedNodeIndex; + queryKeyTransposedToMhaEdge.FromNodeOutputIndex = 0; + queryKeyTransposedToMhaEdge.ToNodeIndex = mhaNodeIndex; + queryKeyTransposedToMhaEdge.ToNodeInputIndex = mhaStackedQueryKeyIndex; + intermediateEdges.push_back(queryKeyTransposedToMhaEdge); + + DML_INTERMEDIATE_GRAPH_EDGE_DESC gemmToValueSliceEdge = {}; + gemmToValueSliceEdge.FromNodeIndex = gemmNodeIndex; + gemmToValueSliceEdge.FromNodeOutputIndex = 0; + gemmToValueSliceEdge.ToNodeIndex = valueSliceNodeIndex; + gemmToValueSliceEdge.ToNodeInputIndex = 0; + intermediateEdges.push_back(gemmToValueSliceEdge); + + DML_INTERMEDIATE_GRAPH_EDGE_DESC valueSliceToMhaEdge = {}; + valueSliceToMhaEdge.FromNodeIndex = valueSliceNodeIndex; + valueSliceToMhaEdge.FromNodeOutputIndex = 0; + valueSliceToMhaEdge.ToNodeIndex = mhaNodeIndex; + valueSliceToMhaEdge.ToNodeInputIndex = mhaValueIndex; + intermediateEdges.push_back(valueSliceToMhaEdge); + } + else + { + DML_INTERMEDIATE_GRAPH_EDGE_DESC gemmToQueryKeyValueTransposeEdge = {}; + gemmToQueryKeyValueTransposeEdge.FromNodeIndex = gemmNodeIndex; + gemmToQueryKeyValueTransposeEdge.FromNodeOutputIndex = 0; + gemmToQueryKeyValueTransposeEdge.ToNodeIndex = queryKeyValueTransposedNodeIndex; + gemmToQueryKeyValueTransposeEdge.ToNodeInputIndex = 0; + intermediateEdges.push_back(gemmToQueryKeyValueTransposeEdge); + + // All we need to do here is transpose the stacked QKV tensor into something DML supports + DML_INTERMEDIATE_GRAPH_EDGE_DESC queryKeyValueTransposedToMhaEdge = {}; + queryKeyValueTransposedToMhaEdge.FromNodeIndex = queryKeyValueTransposedNodeIndex; + queryKeyValueTransposedToMhaEdge.FromNodeOutputIndex = 0; + queryKeyValueTransposedToMhaEdge.ToNodeIndex = mhaNodeIndex; + queryKeyValueTransposedToMhaEdge.ToNodeInputIndex = mhaStackedQueryKeyValueIndex; + intermediateEdges.push_back(queryKeyValueTransposedToMhaEdge); + } + + DML_OUTPUT_GRAPH_EDGE_DESC mhaToOutputEdge = {}; + mhaToOutputEdge.FromNodeIndex = mhaNodeIndex; + mhaToOutputEdge.FromNodeOutputIndex = 0; + mhaToOutputEdge.GraphOutputIndex = 0; + outputEdges.push_back(mhaToOutputEdge); + MLOperatorGraphDesc operatorGraphDesc = {}; - std::array opDescs = { - &xWeightDesc, - &querySlicedDesc, - &keySlicedDesc, - &valueSlicedDesc, - &transposedQueryDesc, - &transposedKeyDesc, - &attentionScoreDesc, - &softmaxDesc, - &transposedQueryDesc, - &attentionWeightDesc, - &castMaskIndexDesc, - &maskDesc, - &outputDesc - }; - operatorGraphDesc.nodeCount = NodeIndex::count; - operatorGraphDesc.nodesAsOpDesc = opDescs.data(); - - // set input edges - std::pair nodeToNodeInputIndex[4] { - {NodeIndex::xWeight, 0}, - {NodeIndex::xWeight, 1}, - {NodeIndex::xWeight, 2}, - {NodeIndex::castMaskIndex, 0} - }; - std::array inputEdges; - for (uint32_t inputIndex = 0; inputIndex < inputEdges.size(); inputIndex++) - { - DML_INPUT_GRAPH_EDGE_DESC inputEdge = {}; - inputEdge.GraphInputIndex = inputIndex; - inputEdge.ToNodeIndex = nodeToNodeInputIndex[inputIndex].first; - inputEdge.ToNodeInputIndex = nodeToNodeInputIndex[inputIndex].second; - inputEdges[inputIndex] = inputEdge; - } operatorGraphDesc.inputEdgeCount = gsl::narrow_cast(inputEdges.size()); operatorGraphDesc.inputEdges = inputEdges.data(); - - // set intermediate edges - std::vector intermediateEdges; - - DML_INTERMEDIATE_GRAPH_EDGE_DESC gemmToQuerySliceEdge = {}; - gemmToQuerySliceEdge.FromNodeIndex = NodeIndex::xWeight; - gemmToQuerySliceEdge.FromNodeOutputIndex = 0; - gemmToQuerySliceEdge.ToNodeIndex = NodeIndex::querySlice; - gemmToQuerySliceEdge.ToNodeInputIndex = 0; - intermediateEdges.push_back(gemmToQuerySliceEdge); - - DML_INTERMEDIATE_GRAPH_EDGE_DESC gemmToKeySliceEdge = {}; - gemmToKeySliceEdge.FromNodeIndex = NodeIndex::xWeight; - gemmToKeySliceEdge.FromNodeOutputIndex = 0; - gemmToKeySliceEdge.ToNodeIndex = NodeIndex::keySlice; - gemmToKeySliceEdge.ToNodeInputIndex = 0; - intermediateEdges.push_back(gemmToKeySliceEdge); - - DML_INTERMEDIATE_GRAPH_EDGE_DESC gemmToValueSliceEdge = {}; - gemmToValueSliceEdge.FromNodeIndex = NodeIndex::xWeight; - gemmToValueSliceEdge.FromNodeOutputIndex = 0; - gemmToValueSliceEdge.ToNodeIndex = NodeIndex::valueSlice; - gemmToValueSliceEdge.ToNodeInputIndex = 0; - intermediateEdges.push_back(gemmToValueSliceEdge); - - DML_INTERMEDIATE_GRAPH_EDGE_DESC querySliceToQueryTranspose = {}; - querySliceToQueryTranspose.FromNodeIndex = NodeIndex::querySlice; - querySliceToQueryTranspose.FromNodeOutputIndex = 0; - querySliceToQueryTranspose.ToNodeIndex = NodeIndex::queryTranspose; - querySliceToQueryTranspose.ToNodeInputIndex = 0; - intermediateEdges.push_back(querySliceToQueryTranspose); - - DML_INTERMEDIATE_GRAPH_EDGE_DESC keySliceToKeyTranspose = {}; - keySliceToKeyTranspose.FromNodeIndex = NodeIndex::keySlice; - keySliceToKeyTranspose.FromNodeOutputIndex = 0; - keySliceToKeyTranspose.ToNodeIndex = NodeIndex::keyTranspose; - keySliceToKeyTranspose.ToNodeInputIndex = 0; - intermediateEdges.push_back(keySliceToKeyTranspose); - - DML_INTERMEDIATE_GRAPH_EDGE_DESC queryTransposeToGemm = {}; - queryTransposeToGemm.FromNodeIndex = NodeIndex::queryTranspose; - queryTransposeToGemm.FromNodeOutputIndex = 0; - queryTransposeToGemm.ToNodeIndex = NodeIndex::attentionScore; - queryTransposeToGemm.ToNodeInputIndex = 0; - intermediateEdges.push_back(queryTransposeToGemm); - - DML_INTERMEDIATE_GRAPH_EDGE_DESC keyTransposeToGemm = {}; - keyTransposeToGemm.FromNodeIndex = NodeIndex::keyTranspose; - keyTransposeToGemm.FromNodeOutputIndex = 0; - keyTransposeToGemm.ToNodeIndex = NodeIndex::attentionScore; - keyTransposeToGemm.ToNodeInputIndex = 1; - intermediateEdges.push_back(keyTransposeToGemm); - - DML_INTERMEDIATE_GRAPH_EDGE_DESC castedMaskIndexToIdentity = {}; - castedMaskIndexToIdentity.FromNodeIndex = NodeIndex::castMaskIndex; - castedMaskIndexToIdentity.FromNodeOutputIndex = 0; - castedMaskIndexToIdentity.ToNodeIndex = NodeIndex::mask; - castedMaskIndexToIdentity.ToNodeInputIndex = 0; - intermediateEdges.push_back(castedMaskIndexToIdentity); - - DML_INTERMEDIATE_GRAPH_EDGE_DESC maskToGemm = {}; - maskToGemm.FromNodeIndex = NodeIndex::mask; - maskToGemm.FromNodeOutputIndex = 0; - maskToGemm.ToNodeIndex = NodeIndex::attentionScore; - maskToGemm.ToNodeInputIndex = 2; - intermediateEdges.push_back(maskToGemm); - - DML_INTERMEDIATE_GRAPH_EDGE_DESC attentionScoreToSoftmax = {}; - attentionScoreToSoftmax.FromNodeIndex = NodeIndex::attentionScore; - attentionScoreToSoftmax.FromNodeOutputIndex = 0; - attentionScoreToSoftmax.ToNodeIndex = NodeIndex::softmax; - attentionScoreToSoftmax.ToNodeInputIndex = 0; - intermediateEdges.push_back(attentionScoreToSoftmax); - - DML_INTERMEDIATE_GRAPH_EDGE_DESC softmaxToGemm = {}; - softmaxToGemm.FromNodeIndex = NodeIndex::softmax; - softmaxToGemm.FromNodeOutputIndex = 0; - softmaxToGemm.ToNodeIndex = NodeIndex::attentionWeight; - softmaxToGemm.ToNodeInputIndex = 0; - intermediateEdges.push_back(softmaxToGemm); - - DML_INTERMEDIATE_GRAPH_EDGE_DESC valueSliceToValueTranspose = {}; - valueSliceToValueTranspose.FromNodeIndex = NodeIndex::valueSlice; - valueSliceToValueTranspose.FromNodeOutputIndex = 0; - valueSliceToValueTranspose.ToNodeIndex = NodeIndex::valueTranspose; - valueSliceToValueTranspose.ToNodeInputIndex = 0; - intermediateEdges.push_back(valueSliceToValueTranspose); - - DML_INTERMEDIATE_GRAPH_EDGE_DESC valueTransposeToGemm = {}; - valueTransposeToGemm.FromNodeIndex = NodeIndex::valueTranspose; - valueTransposeToGemm.FromNodeOutputIndex = 0; - valueTransposeToGemm.ToNodeIndex = NodeIndex::attentionWeight; - valueTransposeToGemm.ToNodeInputIndex = 1; - intermediateEdges.push_back(valueTransposeToGemm); - - DML_INTERMEDIATE_GRAPH_EDGE_DESC gemmToIdentity = {}; - gemmToIdentity.FromNodeIndex = NodeIndex::attentionWeight; - gemmToIdentity.FromNodeOutputIndex = 0; - gemmToIdentity.ToNodeIndex = NodeIndex::output; - gemmToIdentity.ToNodeInputIndex = 0; - intermediateEdges.push_back(gemmToIdentity); - operatorGraphDesc.intermediateEdgeCount = gsl::narrow_cast(intermediateEdges.size()); operatorGraphDesc.intermediateEdges = intermediateEdges.data(); - - // set the output edges - std::array outputEdges; - DML_OUTPUT_GRAPH_EDGE_DESC outputEdge = {}; - outputEdge.FromNodeIndex = NodeIndex::output; - outputEdge.FromNodeOutputIndex = 0; - outputEdge.GraphOutputIndex = 0; - outputEdges[0] = outputEdge; operatorGraphDesc.outputEdgeCount = gsl::narrow_cast(outputEdges.size()); operatorGraphDesc.outputEdges = outputEdges.data(); + operatorGraphDesc.nodeCount = gsl::narrow_cast(opDescs.size()); + operatorGraphDesc.nodesAsOpDesc = opDescs.data(); SetDmlOperatorGraphDesc(std::move(operatorGraphDesc), kernelCreationContext); } @@ -448,32 +540,37 @@ public: void CALLBACK QueryAttention(IMLOperatorSupportQueryContextPrivate* context, /*out*/ bool* isSupported) { *isSupported = false; - // Fall back to CPU if input 'past' and 'relative_position_bias' is present because there is no current use case for this. - // and it will make the implementation more complex. - // Also fall back to CPU if output 'present' is present for same reason as above. - if (context->GetInputCount() > 4 || context->GetOutputCount() > 1) - { - return; - } - // Checking input count alone is not sufficient to fallback to CPU if input 'past' and 'relative_position_bias' is present - // because input 'mask_index', 'past', and 'relative_position_bias' all are optional. - if (context->IsInputValid(4) || context->IsInputValid(5)) - { - return; - } - // Fall back to CPU if attibute 'qkv_hidden_sizes' is present or - // if value of attribute 'unidirectional' is 1, because of same reason as above. - MLOperatorAttributes attributes(context); - if (attributes.HasAttribute(AttrName::QkvHiddenSizes, MLOperatorAttributeType::IntArray)) + // `past` input tensor is not supported yet + if (context->IsInputValid(4)) { return; } + // `past_sequence_length` input tensor is not supported yet + if (context->IsInputValid(6)) + { + return; + } + + // `present` output tensor is not supported yet + if (context->IsOutputValid(1)) + { + return; + } + + // `unidirectional == 1` is not supported yet + MLOperatorAttributes attributes(context); if (attributes.GetOptionalAttribute(AttrName::Unidirectional, 0) != 0) { return; } + // `do_rotary == 1` is not supported yet + if (attributes.GetOptionalAttribute(AttrName::DoRotary, 0) != 0) + { + return; + } + *isSupported = true; } diff --git a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/DmlOperatorMultiHeadAttention.cpp b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/DmlOperatorMultiHeadAttention.cpp new file mode 100644 index 0000000000..9c1a7baeaa --- /dev/null +++ b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/DmlOperatorMultiHeadAttention.cpp @@ -0,0 +1,281 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "precomp.h" + +namespace Dml +{ +class DmlOperatorMultiHeadAttention : public DmlOperator +{ +public: + DmlOperatorMultiHeadAttention(const MLOperatorKernelCreationContext& kernelCreationContext) + : DmlOperator(kernelCreationContext) + { + enum InputIndex : uint32_t + { + queryIndex, + keyIndex, + valueIndex, + biasIndex, + maskIndex, + relativePositionBiasIndex, + pastKeyIndex, + pastValueIndex, + inputCount, + }; + + enum DmlInputIndex : uint32_t + { + dmlQueryIndex, + dmlKeyIndex, + dmlValueIndex, + dmlStackedQueryKeyIndex, + dmlStackedKeyValueIndex, + dmlStackedQueryKeyValueIndex, + dmlBiasIndex, + dmlMaskIndex, + dmlRelativePositionBiasIndex, + dmlPastKeyIndex, + dmlPastValueIndex, + dmlInputCount, + }; + + enum OutputIndex : uint32_t + { + outputIndex, + outputPresentKeyIndex, + outputPresentValueIndex, + outputCount, + }; + + ML_CHECK_VALID_ARGUMENT(kernelCreationContext.GetInputCount() >= 1); + ML_CHECK_VALID_ARGUMENT(kernelCreationContext.GetOutputCount() >= 1); + + const bool keyValueIsPast = kernelCreationContext.IsInputValid(keyIndex) && kernelCreationContext.GetInputTensorDimensionCount(keyIndex) == 4; + const bool hasValue = kernelCreationContext.IsInputValid(valueIndex) && !keyValueIsPast; + const bool hasBias = kernelCreationContext.IsInputValid(biasIndex); + const bool hasMask = kernelCreationContext.IsInputValid(maskIndex); + const bool hasRelativePositionBias = kernelCreationContext.IsInputValid(relativePositionBiasIndex); + const bool hasPastKey = keyValueIsPast || kernelCreationContext.IsInputValid(pastKeyIndex); + const bool hasPastValue = keyValueIsPast || kernelCreationContext.IsInputValid(pastValueIndex); + const bool hasPresentKeyOutput = kernelCreationContext.IsOutputValid(outputPresentKeyIndex); + const bool hasPresentValueOutput = kernelCreationContext.IsOutputValid(outputPresentValueIndex); + const bool stackedQkv = kernelCreationContext.GetInputTensorDimensionCount(queryIndex) == 5; + const bool stackedKv = kernelCreationContext.IsInputValid(keyIndex) && kernelCreationContext.GetInputTensorDimensionCount(keyIndex) == 5; + const bool hasKey = !stackedKv && !keyValueIsPast && kernelCreationContext.IsInputValid(keyIndex); + + std::vector> inputIndices = { + stackedQkv ? std::nullopt : std::optional(queryIndex), + hasKey ? std::optional(keyIndex) : std::nullopt, + hasValue ? std::optional(valueIndex) : std::nullopt, + std::nullopt, + stackedKv ? std::optional(keyIndex) : std::nullopt, + stackedQkv ? std::optional(queryIndex) : std::nullopt, + biasIndex, + hasMask ? std::optional(maskIndex) : std::nullopt, + relativePositionBiasIndex, + keyValueIsPast ? keyIndex : pastKeyIndex, + keyValueIsPast ? valueIndex : pastValueIndex, + }; + + std::vector> outputIndices = { + outputIndex, + outputPresentKeyIndex, + outputPresentValueIndex, + }; + DmlOperator::Initialize(kernelCreationContext, inputIndices, outputIndices, std::nullopt, std::nullopt, 1); + + ML_CHECK_VALID_ARGUMENT(!stackedQkv || m_inputTensorDescs[dmlStackedQueryKeyValueIndex].GetDimensionCount() == 5); + ML_CHECK_VALID_ARGUMENT(stackedQkv || m_inputTensorDescs[dmlQueryIndex].GetDimensionCount() == 3); + ML_CHECK_VALID_ARGUMENT(!hasKey || m_inputTensorDescs[dmlKeyIndex].GetDimensionCount() == 3); + ML_CHECK_VALID_ARGUMENT(!hasValue || m_inputTensorDescs[dmlValueIndex].GetDimensionCount() == 3); + ML_CHECK_VALID_ARGUMENT(!hasPastKey || m_inputTensorDescs[dmlPastKeyIndex].GetDimensionCount() == 4); + ML_CHECK_VALID_ARGUMENT(!hasPastValue || m_inputTensorDescs[dmlPastValueIndex].GetDimensionCount() == 4); + + const uint32_t batchSize = stackedQkv + ? m_inputTensorDescs[dmlStackedQueryKeyValueIndex].GetSizes()[0] + : m_inputTensorDescs[dmlQueryIndex].GetSizes()[0]; + + const uint32_t numHeads = gsl::narrow_cast(kernelCreationContext.GetAttribute(AttrName::NumHeads)); + const uint32_t headSize = stackedQkv + ? m_inputTensorDescs[dmlStackedQueryKeyValueIndex].GetSizes()[4] + : m_inputTensorDescs[dmlQueryIndex].GetSizes()[2] / numHeads; + + const uint32_t sequenceLength = stackedQkv + ? m_inputTensorDescs[dmlStackedQueryKeyValueIndex].GetSizes()[1] + : m_inputTensorDescs[dmlQueryIndex].GetSizes()[1]; + + uint32_t kvSequenceLength; + if (hasKey) + { + kvSequenceLength = m_inputTensorDescs[dmlKeyIndex].GetSizes()[1]; + } + else if (stackedKv) + { + kvSequenceLength = m_inputTensorDescs[dmlStackedKeyValueIndex].GetSizes()[1]; + } + else if (hasPastKey) + { + kvSequenceLength = m_inputTensorDescs[dmlPastKeyIndex].GetSizes()[2]; + } + else + { + kvSequenceLength = sequenceLength; + } + + const uint32_t hiddenSize = numHeads * headSize; + const uint32_t vHiddenSize = hasValue ? m_inputTensorDescs[dmlValueIndex].GetSizes()[2] : hiddenSize; + const uint32_t pastSequenceLength = hasPastKey ? m_inputTensorDescs[dmlPastKeyIndex].GetSizes()[2] : 0; + const uint32_t totalSequenceLength = kvSequenceLength + pastSequenceLength; + + if (stackedQkv) + { + auto stackedQkvSizes = m_inputTensorDescs[dmlStackedQueryKeyValueIndex].GetSizes(); + ML_CHECK_VALID_ARGUMENT(stackedQkvSizes[0] == batchSize); + ML_CHECK_VALID_ARGUMENT(stackedQkvSizes[1] == sequenceLength); + ML_CHECK_VALID_ARGUMENT(stackedQkvSizes[2] == numHeads); + ML_CHECK_VALID_ARGUMENT(stackedQkvSizes[3] == 3); + ML_CHECK_VALID_ARGUMENT(stackedQkvSizes[4] == headSize); + } + else + { + auto querySizes = m_inputTensorDescs[dmlQueryIndex].GetSizes(); + ML_CHECK_VALID_ARGUMENT(querySizes[0] == batchSize); + ML_CHECK_VALID_ARGUMENT(querySizes[1] == sequenceLength); + ML_CHECK_VALID_ARGUMENT(querySizes[2] == hiddenSize); + } + + if (hasKey) + { + ML_CHECK_VALID_ARGUMENT(m_inputTensorDescs[dmlKeyIndex].GetDimensionCount() == 3); + + auto keySizes = m_inputTensorDescs[dmlKeyIndex].GetSizes(); + ML_CHECK_VALID_ARGUMENT(keySizes[0] == batchSize); + ML_CHECK_VALID_ARGUMENT(keySizes[1] == kvSequenceLength); + ML_CHECK_VALID_ARGUMENT(keySizes[2] == hiddenSize); + } + + if (hasValue) + { + auto valueSizes = m_inputTensorDescs[dmlValueIndex].GetSizes(); + ML_CHECK_VALID_ARGUMENT(valueSizes[0] == batchSize); + ML_CHECK_VALID_ARGUMENT(valueSizes[1] == kvSequenceLength); + ML_CHECK_VALID_ARGUMENT(valueSizes[2] == vHiddenSize); + } + + if (stackedKv) + { + ML_CHECK_VALID_ARGUMENT(m_inputTensorDescs[dmlStackedKeyValueIndex].GetDimensionCount() == 5); + + auto stackedKvSizes = m_inputTensorDescs[dmlStackedKeyValueIndex].GetSizes(); + ML_CHECK_VALID_ARGUMENT(stackedKvSizes[0] == batchSize); + ML_CHECK_VALID_ARGUMENT(stackedKvSizes[1] == kvSequenceLength); + ML_CHECK_VALID_ARGUMENT(stackedKvSizes[2] == numHeads); + ML_CHECK_VALID_ARGUMENT(stackedKvSizes[3] == 2); + ML_CHECK_VALID_ARGUMENT(stackedKvSizes[4] == headSize); + } + + if (hasBias) + { + ML_CHECK_VALID_ARGUMENT(m_inputTensorDescs[dmlBiasIndex].GetDimensionCount() == 1); + ML_CHECK_VALID_ARGUMENT(m_inputTensorDescs[dmlBiasIndex].GetSizes()[0] == hiddenSize + hiddenSize + vHiddenSize); + } + + DML_MULTIHEAD_ATTENTION_MASK_TYPE maskType = DML_MULTIHEAD_ATTENTION_MASK_TYPE_NONE; + if (hasMask) + { + if (kernelCreationContext.GetInputTensorDimensionCount(maskIndex) == 1) + { + const auto unpaddedKeyBoundsShape = m_inputTensorDescs[dmlMaskIndex].GetSizes(); + ML_CHECK_VALID_ARGUMENT(unpaddedKeyBoundsShape.size() == 1); + ML_CHECK_VALID_ARGUMENT(unpaddedKeyBoundsShape[0] == batchSize || unpaddedKeyBoundsShape[0] == batchSize * 3 + 2); + + maskType = unpaddedKeyBoundsShape[0] == batchSize + ? DML_MULTIHEAD_ATTENTION_MASK_TYPE_KEY_SEQUENCE_LENGTH + : DML_MULTIHEAD_ATTENTION_MASK_TYPE_KEY_QUERY_SEQUENCE_LENGTH_START_END; + + if (maskType == DML_MULTIHEAD_ATTENTION_MASK_TYPE_KEY_SEQUENCE_LENGTH) + { + uint32_t desiredShape[2] = {1, batchSize}; + m_inputTensorDescs[dmlMaskIndex] = TensorDesc( + m_inputTensorDescs[dmlMaskIndex].GetDmlDataType(), + desiredShape); + } + } + else + { + const auto keyPaddingMaskTensorShape = m_inputTensorDescs[dmlMaskIndex].GetSizes(); + ML_CHECK_VALID_ARGUMENT(keyPaddingMaskTensorShape.size() == 2); + ML_CHECK_VALID_ARGUMENT(keyPaddingMaskTensorShape[0] == batchSize); + ML_CHECK_VALID_ARGUMENT(keyPaddingMaskTensorShape[1] == kvSequenceLength); + + const uint32_t actualShape[4] = {batchSize, 1, 1, kvSequenceLength}; + const uint32_t desiredShape[4] = {batchSize, numHeads, sequenceLength, kvSequenceLength}; + + m_inputTensorDescs[dmlMaskIndex] = TensorDesc::ConstructBroadcastedTensorDesc( + m_inputTensorDescs[dmlMaskIndex].GetMlOperatorDataType(), + desiredShape, + actualShape); + + maskType = DML_MULTIHEAD_ATTENTION_MASK_TYPE_BOOLEAN; + } + } + + if (hasRelativePositionBias) + { + ML_CHECK_VALID_ARGUMENT(m_inputTensorDescs[dmlRelativePositionBiasIndex].GetDimensionCount() == 4); + + auto relativePositionBiasSizes = m_inputTensorDescs[dmlRelativePositionBiasIndex].GetSizes(); + ML_CHECK_VALID_ARGUMENT(relativePositionBiasSizes[0] == batchSize); + ML_CHECK_VALID_ARGUMENT(relativePositionBiasSizes[1] == numHeads); + ML_CHECK_VALID_ARGUMENT(relativePositionBiasSizes[2] == sequenceLength); + ML_CHECK_VALID_ARGUMENT(relativePositionBiasSizes[3] == totalSequenceLength); + } + + if (hasPastKey) + { + auto pastKeySizes = m_inputTensorDescs[dmlPastKeyIndex].GetSizes(); + ML_CHECK_VALID_ARGUMENT(pastKeySizes[0] == batchSize); + ML_CHECK_VALID_ARGUMENT(pastKeySizes[1] == numHeads); + ML_CHECK_VALID_ARGUMENT(pastKeySizes[2] == pastSequenceLength); + ML_CHECK_VALID_ARGUMENT(pastKeySizes[3] == headSize); + } + + if (hasPastValue) + { + auto pastValueSizes = m_inputTensorDescs[dmlPastValueIndex].GetSizes(); + ML_CHECK_VALID_ARGUMENT(pastValueSizes[0] == batchSize); + ML_CHECK_VALID_ARGUMENT(pastValueSizes[1] == numHeads); + ML_CHECK_VALID_ARGUMENT(pastValueSizes[2] == pastSequenceLength); + ML_CHECK_VALID_ARGUMENT(pastValueSizes[3] == headSize); + } + + std::vector inputDescs = GetDmlInputDescs(); + std::vector outputDescs = GetDmlOutputDescs(); + + DML_MULTIHEAD_ATTENTION_OPERATOR_DESC mhaDesc = {}; + mhaDesc.QueryTensor = stackedQkv ? nullptr : &inputDescs[dmlQueryIndex]; + mhaDesc.KeyTensor = hasKey ? &inputDescs[dmlKeyIndex] : nullptr; + mhaDesc.ValueTensor = hasValue ? &inputDescs[dmlValueIndex] : nullptr; + mhaDesc.StackedKeyValueTensor = stackedKv ? &inputDescs[dmlStackedKeyValueIndex] : nullptr; + mhaDesc.StackedQueryKeyValueTensor = stackedQkv ? &inputDescs[dmlStackedQueryKeyValueIndex] : nullptr; + mhaDesc.BiasTensor = hasBias ? &inputDescs[dmlBiasIndex] : nullptr; + mhaDesc.MaskTensor = hasMask ? &inputDescs[dmlMaskIndex] : nullptr; + mhaDesc.RelativePositionBiasTensor = hasRelativePositionBias ? &inputDescs[dmlRelativePositionBiasIndex] : nullptr; + mhaDesc.PastKeyTensor = hasPastKey ? &inputDescs[dmlPastKeyIndex] : nullptr; + mhaDesc.PastValueTensor = hasPastValue ? &inputDescs[dmlPastValueIndex] : nullptr; + mhaDesc.OutputTensor = &outputDescs[outputIndex]; + mhaDesc.OutputPresentKeyTensor = hasPresentKeyOutput ? &outputDescs[outputPresentKeyIndex] : nullptr; + mhaDesc.OutputPresentValueTensor = hasPresentValueOutput ? &outputDescs[outputPresentValueIndex] : nullptr; + mhaDesc.Scale = kernelCreationContext.GetOptionalAttribute(AttrName::Scale, gsl::narrow_cast(1.0f / std::sqrt(headSize))); + mhaDesc.MaskFilterValue = kernelCreationContext.GetOptionalAttribute(AttrName::MaskFilterValue, -10'000.0f); + mhaDesc.HeadCount = numHeads; + mhaDesc.MaskType = maskType; + + DML_OPERATOR_DESC opDesc = { DML_OPERATOR_MULTIHEAD_ATTENTION, &mhaDesc }; + SetDmlOperatorDesc(opDesc, kernelCreationContext); + } +}; + +DML_OP_DEFINE_CREATION_FUNCTION(MultiHeadAttention, DmlOperatorMultiHeadAttention); +} // namespace Dml diff --git a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/OperatorRegistration.cpp b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/OperatorRegistration.cpp index fd0ad8385f..51f99d5488 100644 --- a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/OperatorRegistration.cpp +++ b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/OperatorRegistration.cpp @@ -386,6 +386,7 @@ DML_OP_EXTERN_CREATION_FUNCTION(Trilu); DML_OP_EXTERN_CREATION_FUNCTION(Shape); DML_OP_EXTERN_CREATION_FUNCTION(Size); DML_OP_EXTERN_CREATION_FUNCTION(Attention); +DML_OP_EXTERN_CREATION_FUNCTION(MultiHeadAttention); DML_OP_EXTERN_CREATION_FUNCTION(NonZero); DML_OP_EXTERN_CREATION_FUNCTION(QuickGelu); @@ -871,8 +872,10 @@ constexpr static OperatorRegistrationInformation operatorRegistrationInformation {REG_INFO_MS( 1, Gelu, typeNameListDefault, supportedTypeListFloat16to32, DmlGraphSupport::Supported)}, {REG_INFO_MS( 1, BiasGelu, typeNameListDefault, supportedTypeListFloat16to32, DmlGraphSupport::Supported)}, {REG_INFO_MS( 1, FusedMatMul, typeNameListDefault, supportedTypeListFloat16to32, DmlGraphSupport::Supported)}, + {REG_INFO_MS( 1, FusedMatMulActivation, typeNameListDefault, supportedTypeListFloat16to32, DmlGraphSupport::Supported)}, {REG_INFO_MS( 1, QLinearSigmoid, typeNameListDefault, supportedTypeListQLinearSigmoid, DmlGraphSupport::Supported, requiredConstantCpuInputs(), std::nullopt, QueryQLinearSigmoid)}, {REG_INFO_MS( 1, Attention, typeNameListAttention, supportedTypeListAttention, DmlGraphSupport::Supported, requiredConstantCpuInputs(), std::nullopt, QueryAttention)}, + {REG_INFO_MS( 1, MultiHeadAttention, typeNameListAttention, supportedTypeListAttention, DmlGraphSupport::Supported)}, {REG_INFO( 10, IsInf, typeNameListTwo, supportedTypeListIsInf, DmlGraphSupport::Supported)}, {REG_INFO( 10, Mod, typeNameListDefault, supportedTypeListNumericDefault, DmlGraphSupport::Supported)}, diff --git a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/OperatorUtility.cpp b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/OperatorUtility.cpp index f83ad5a3a4..d8290bbdae 100644 --- a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/OperatorUtility.cpp +++ b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/OperatorUtility.cpp @@ -364,19 +364,10 @@ namespace Dml int onnxAxis = HandleNegativeAxis(kernelInfo.GetOptionalAttribute(AttrName::GraphFusedAxis, -1), onnxDimCount); auto dmlAdjustedAxis = GetDmlAdjustedAxis(onnxAxis, onnxDimCount, kernelInfo.GetTensorShapeDescription().GetInputTensorDimensionCount(0)); - - // If the axis is supported by Softmax, use this version instead since it's more likely to be supported by metacommands - if (dmlAdjustedAxis == onnxDimCount - 1) - { - activation.desc.activationType = DML_OPERATOR_ACTIVATION_SOFTMAX; - } - else - { - activation.desc.activationType = DML_OPERATOR_ACTIVATION_SOFTMAX1; - activation.dmlAxes.push_back(dmlAdjustedAxis); - activation.desc.params.softmax1.Axes = activation.dmlAxes.data(); - activation.desc.params.softmax1.AxisCount = gsl::narrow_cast(activation.dmlAxes.size()); - } + activation.desc.activationType = DML_OPERATOR_ACTIVATION_SOFTMAX1; + activation.dmlAxes.push_back(dmlAdjustedAxis); + activation.desc.params.softmax1.Axes = activation.dmlAxes.data(); + activation.desc.params.softmax1.AxisCount = gsl::narrow_cast(activation.dmlAxes.size()); } else { diff --git a/onnxruntime/core/providers/dml/OperatorAuthorHelper/Attributes.h b/onnxruntime/core/providers/dml/OperatorAuthorHelper/Attributes.h index b47e0c5ed1..3c5ae4f764 100644 --- a/onnxruntime/core/providers/dml/OperatorAuthorHelper/Attributes.h +++ b/onnxruntime/core/providers/dml/OperatorAuthorHelper/Attributes.h @@ -113,6 +113,8 @@ namespace AttrName static constexpr const char* FusedBeta = "fused_beta"; static constexpr const char* FusedGamma = "fused_gamma"; static constexpr const char* FusedRatio = "fused_ratio"; + static constexpr const char* MaskFilterValue = "mask_filter_value"; + static constexpr const char* DoRotary = "do_rotary"; static constexpr const char* Activation = "activation"; static constexpr const char* Groups = "groups"; diff --git a/onnxruntime/core/providers/dml/OperatorAuthorHelper/OperatorHelper.cpp b/onnxruntime/core/providers/dml/OperatorAuthorHelper/OperatorHelper.cpp index 1ea466b771..1b9ac2a605 100644 --- a/onnxruntime/core/providers/dml/OperatorAuthorHelper/OperatorHelper.cpp +++ b/onnxruntime/core/providers/dml/OperatorAuthorHelper/OperatorHelper.cpp @@ -2536,6 +2536,114 @@ namespace OperatorHelper return outputShapes; } + std::vector MultiHeadAttentionHelper::GetOutputShapes(const MLShapeInferenceContext& shapeInfo) const + { + ML_CHECK_VALID_ARGUMENT(shapeInfo.GetInputCount() >= 1); + + auto queryShape = shapeInfo.GetInputTensorShape(0); + ML_CHECK_VALID_ARGUMENT(queryShape.size() == 3 || queryShape.size() == 5); + + const uint32_t batchSize = queryShape[0]; + const uint32_t sequenceLength = queryShape[1]; + uint32_t kvSequenceLength = 0; + uint32_t vHiddenSize = 0; + uint32_t headSize = 0; + + if (shapeInfo.IsInputValid(2)) + { + auto valueShape = shapeInfo.GetInputTensorShape(2); + ML_CHECK_VALID_ARGUMENT(queryShape.size() == 3); + headSize = queryShape[2] / m_numHeads; + + if (valueShape.size() == 3) + { + kvSequenceLength = valueShape[1]; + vHiddenSize = valueShape[2]; + } + else + { + ML_CHECK_VALID_ARGUMENT(valueShape.size() == 4); + const uint32_t vHeadSize = valueShape[3]; + kvSequenceLength = valueShape[2]; + vHiddenSize = vHeadSize * m_numHeads; + } + } + else if (shapeInfo.IsInputValid(1)) + { + auto keyShape = shapeInfo.GetInputTensorShape(1); + ML_CHECK_VALID_ARGUMENT(keyShape.size() == 5); + kvSequenceLength = keyShape[1]; + vHiddenSize = queryShape[2]; + headSize = keyShape[4]; + } + else + { + ML_CHECK_VALID_ARGUMENT(queryShape.size() == 5); + kvSequenceLength = queryShape[1]; + headSize = queryShape[4]; + vHiddenSize = headSize * m_numHeads; + } + + std::vector outputShapes(3); + outputShapes[0] = EdgeShapes({batchSize, sequenceLength, vHiddenSize}); + + uint32_t totalSequenceLength = kvSequenceLength; + if (shapeInfo.IsInputValid(6)) + { + ML_CHECK_VALID_ARGUMENT(shapeInfo.GetInputTensorDimensionCount(6) == 4); + const uint32_t pastSequenceLength = shapeInfo.GetInputTensorShape(6)[2]; + totalSequenceLength += pastSequenceLength; + } + + if (shapeInfo.IsOutputValid(1)) + { + outputShapes[1] = EdgeShapes({batchSize, m_numHeads, totalSequenceLength, headSize}); + } + + if (shapeInfo.IsOutputValid(2)) + { + outputShapes[2] = EdgeShapes({batchSize, m_numHeads, totalSequenceLength, headSize}); + } + + return outputShapes; + } + + void MultiHeadAttentionHelper::Initialize(const IKernelInformationAdapter& kernelInformation) + { + m_numHeads = gsl::narrow_cast(kernelInformation.GetAttributes().GetAttribute(AttrName::NumHeads)); + } + + std::vector AttentionHelper::GetOutputShapes(const MLShapeInferenceContext& shapeInfo) const + { + ML_CHECK_VALID_ARGUMENT(shapeInfo.GetInputCount() >= 2); + + auto queryShape = shapeInfo.GetInputTensorShape(0); + ML_CHECK_VALID_ARGUMENT(queryShape.size() == 3); + + auto weightShape = shapeInfo.GetInputTensorShape(1); + ML_CHECK_VALID_ARGUMENT(weightShape.size() == 2); + + if (m_qkvHiddenSizes.empty()) + { + ML_CHECK_VALID_ARGUMENT(weightShape[1] % 3 == 0); + } + else + { + ML_CHECK_VALID_ARGUMENT(m_qkvHiddenSizes.size() == 3); + } + + const uint32_t batchSize = queryShape[0]; + const uint32_t sequenceLength = queryShape[1]; + const uint32_t vHiddenSize = m_qkvHiddenSizes.empty() ? weightShape[1] / 3 : m_qkvHiddenSizes[2]; + + return { EdgeShapes({batchSize, sequenceLength, vHiddenSize}) }; + } + + void AttentionHelper::Initialize(const IKernelInformationAdapter& kernelInformation) + { + m_qkvHiddenSizes = kernelInformation.GetAttributes().GetOptionalAttributeVectorInt32(AttrName::QkvHiddenSizes); + } + std::vector SkipLayerNormHelper::GetOutputShapes(const MLShapeInferenceContext& shapeInfo) const { ML_CHECK_VALID_ARGUMENT(shapeInfo.GetInputCount() >= 3); diff --git a/onnxruntime/core/providers/dml/OperatorAuthorHelper/OperatorHelper.h b/onnxruntime/core/providers/dml/OperatorAuthorHelper/OperatorHelper.h index 0206e483d9..5cc611904e 100644 --- a/onnxruntime/core/providers/dml/OperatorAuthorHelper/OperatorHelper.h +++ b/onnxruntime/core/providers/dml/OperatorAuthorHelper/OperatorHelper.h @@ -1419,6 +1419,38 @@ public: std::vector GetOutputShapes(const MLShapeInferenceContext& shapeInfo) const; }; +class MultiHeadAttentionHelper +{ +public: + template + MultiHeadAttentionHelper(const Info_t& info, const Shape_t& shapeInfo) + { + Initialize(KernelInformationAdapter(info)); + } + + std::vector GetOutputShapes(const MLShapeInferenceContext& shapeInfo) const; + +private: + void Initialize(const IKernelInformationAdapter& kernelInformation); + uint32_t m_numHeads; +}; + +class AttentionHelper +{ +public: + template + AttentionHelper(const Info_t& info, const Shape_t& shapeInfo) + { + Initialize(KernelInformationAdapter(info)); + } + + std::vector GetOutputShapes(const MLShapeInferenceContext& shapeInfo) const; + +private: + void Initialize(const IKernelInformationAdapter& kernelInformation); + std::vector m_qkvHiddenSizes; +}; + class SkipLayerNormHelper { public: @@ -1554,7 +1586,8 @@ using ShapeInferenceHelper_Affine = GetOutputShapeAsInputShapeHelper; using ShapeInferenceHelper_QuantizeLinear = GetOutputShapeAsInputShapeHelper; using ShapeInferenceHelper_DequantizeLinear = GetOutputShapeAsInputShapeHelper; using ShapeInferenceHelper_QLinearSigmoid = GetOutputShapeAsInputShapeHelper; -using ShapeInferenceHelper_Attention = GetOutputShapeAsInputShapeHelper; +using ShapeInferenceHelper_Attention = AttentionHelper; +using ShapeInferenceHelper_MultiHeadAttention = MultiHeadAttentionHelper; using ShapeInferenceHelper_Sign = GetBroadcastedOutputShapeHelper; using ShapeInferenceHelper_IsNaN = GetBroadcastedOutputShapeHelper; using ShapeInferenceHelper_Erf = GetBroadcastedOutputShapeHelper; diff --git a/onnxruntime/core/providers/dml/OperatorAuthorHelper/OperatorVersions.h b/onnxruntime/core/providers/dml/OperatorAuthorHelper/OperatorVersions.h index 054faad5ba..490c6e5d1e 100644 --- a/onnxruntime/core/providers/dml/OperatorAuthorHelper/OperatorVersions.h +++ b/onnxruntime/core/providers/dml/OperatorAuthorHelper/OperatorVersions.h @@ -423,6 +423,7 @@ namespace OperatorHelper static const int sc_sinceVer_FusedMatMulActivation = 1; static const int sc_sinceVer_QLinearSigmoid = 1; static const int sc_sinceVer_Attention = 1; + static const int sc_sinceVer_MultiHeadAttention = 1; static const int sc_sinceVer_SkipLayerNormalization = 1; static const int sc_sinceVer_EmbedLayerNormalization = 1; static const int sc_sinceVer_BiasSplitGelu = 1; diff --git a/onnxruntime/test/contrib_ops/attention_op_test.cc b/onnxruntime/test/contrib_ops/attention_op_test.cc index 709fbcb378..3b5d83f0d4 100644 --- a/onnxruntime/test/contrib_ops/attention_op_test.cc +++ b/onnxruntime/test/contrib_ops/attention_op_test.cc @@ -58,6 +58,7 @@ static void RunAttentionTest( const bool disable_cpu = false, const bool disable_cuda = false, const bool disable_rocm = false, + const bool disable_dml = false, std::vector qkv_sizes = {}, const std::vector& relative_position_bias_data = {}, int kv_sequence_length = 0, @@ -72,9 +73,10 @@ static void RunAttentionTest( bool enable_cuda = HasCudaEnvironment(min_cuda_architecture) && !is_weights_constant && !disable_cuda; bool enable_rocm = (nullptr != DefaultRocmExecutionProvider().get()) && !is_weights_constant && !disable_rocm; bool enable_cpu = (nullptr != DefaultCpuExecutionProvider().get()) && !use_float16 && !disable_cpu; + bool enable_dml = (nullptr != DefaultDmlExecutionProvider().get()) && !disable_dml; int head_size = hidden_size / number_of_heads; - if (enable_cpu || enable_cuda || enable_rocm) { + if (enable_cpu || enable_cuda || enable_rocm || enable_dml) { OpTester tester("Attention", 1, onnxruntime::kMSDomain); tester.AddAttribute("num_heads", static_cast(number_of_heads)); tester.AddAttribute("unidirectional", static_cast(is_unidirectional ? 1 : 0)); @@ -242,6 +244,12 @@ static void RunAttentionTest( execution_providers.push_back(DefaultCpuExecutionProvider()); tester.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers); } + + if (enable_dml) { + std::vector> execution_providers; + execution_providers.push_back(DefaultDmlExecutionProvider()); + tester.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers); + } } } @@ -267,6 +275,7 @@ static void RunAttentionTest( const bool disable_cpu = false, const bool disable_cuda = false, const bool disable_rocm = false, + const bool disable_dml = false, const std::vector qkv_sizes = {}, const std::vector& relative_position_bias_data = {}, int kv_sequence_length = 0, @@ -277,13 +286,13 @@ static void RunAttentionTest( batch_size, sequence_length, hidden_size, number_of_heads, use_float16, is_unidirectional, use_past_state, past_sequence_length, past_data, present_data, mask_type, input_hidden_size, max_sequence_length, - disable_cpu, disable_cuda, disable_rocm, qkv_sizes, relative_position_bias_data, + disable_cpu, disable_cuda, disable_rocm, disable_dml, qkv_sizes, relative_position_bias_data, kv_sequence_length, past_present_share_buffer, use_scale, do_neox_rotary); RunAttentionTest(input_data, weights_data, true, bias_data, mask_index_data, output_data, batch_size, sequence_length, hidden_size, number_of_heads, use_float16, is_unidirectional, use_past_state, past_sequence_length, past_data, present_data, mask_type, input_hidden_size, max_sequence_length, - disable_cpu, disable_cuda, disable_rocm, qkv_sizes, relative_position_bias_data, + disable_cpu, disable_cuda, disable_rocm, disable_dml, qkv_sizes, relative_position_bias_data, kv_sequence_length, past_present_share_buffer, use_scale, do_neox_rotary); } @@ -354,7 +363,7 @@ TEST(AttentionTest, AttentionBatch1WithQKVAttr1) { RunAttentionTest(input_data, weight_data, bias_data, mask_index_data, output_data, batch_size, sequence_length, hidden_size, number_of_heads, false, false, false, 0, nullptr, nullptr, AttentionMaskType::MASK_1D_KEY_SEQ_LEN, 0, - 0, false, false, disable_rocm, qkv_sizes); + 0, false, false, disable_rocm, false, qkv_sizes); } TEST(AttentionTest, AttentionBatch1WithQKVAttr2) { @@ -392,7 +401,7 @@ TEST(AttentionTest, AttentionBatch1WithQKVAttr2) { RunAttentionTest(input_data, weight_data, bias_data, mask_index_data, output_data, batch_size, sequence_length, hidden_size, number_of_heads, false, false, false, 0, nullptr, nullptr, AttentionMaskType::MASK_1D_KEY_SEQ_LEN, 0, - 0, false, false, disable_rocm, qkv_sizes); + 0, false, false, disable_rocm, false, qkv_sizes); } TEST(AttentionTest, AttentionBatch1RelativePositionBias) { @@ -429,10 +438,11 @@ TEST(AttentionTest, AttentionBatch1RelativePositionBias) { constexpr bool disable_cpu = false; constexpr bool disable_cuda = false; constexpr bool disable_rocm = false; + constexpr bool disable_dml = false; RunAttentionTest(input_data, weight_data, bias_data, mask_index_data, output_data, batch_size, sequence_length, hidden_size, number_of_heads, false, false, false, 0, nullptr, nullptr, AttentionMaskType::MASK_1D_KEY_SEQ_LEN, 0, - 0, disable_cpu, disable_cuda, disable_rocm, qkv_sizes, relative_position_bias); + 0, disable_cpu, disable_cuda, disable_rocm, disable_dml, qkv_sizes, relative_position_bias); } TEST(AttentionTest, AttentionBatch2RelativePositionBias) { @@ -474,10 +484,11 @@ TEST(AttentionTest, AttentionBatch2RelativePositionBias) { constexpr bool disable_cpu = false; constexpr bool disable_cuda = false; constexpr bool disable_rocm = false; + constexpr bool disable_dml = false; RunAttentionTest(input_data, weight_data, bias_data, mask_index_data, output_data, batch_size, sequence_length, hidden_size, number_of_heads, false, false, false, 0, nullptr, nullptr, AttentionMaskType::MASK_1D_KEY_SEQ_LEN, 0, - 0, disable_cpu, disable_cuda, disable_rocm, qkv_sizes, relative_position_bias); + 0, disable_cpu, disable_cuda, disable_rocm, disable_dml, qkv_sizes, relative_position_bias); } TEST(AttentionTest, AttentionBatch1_Float16) { @@ -817,10 +828,14 @@ void RawAttentionEmptyPastState(bool past_present_share_buffer) { batch_size, sequence_length, hidden_size, number_of_heads, false, is_unidirectional, use_past_state, past_sequence_length, &past_data, &present_data); } else { + // TODO: Unskip when fixed #41968513 + // DML doesn't support past_present_share_buffer for Attention yet + constexpr bool disable_dml = true; + RunAttentionTest(input_data, weight_data, bias_data, mask_index_data, output_data, batch_size, sequence_length, hidden_size, number_of_heads, false, is_unidirectional, use_past_state, past_sequence_length, &past_data, &present_data, - AttentionMaskType::MASK_1D_KEY_SEQ_LEN, 0, sequence_length, true, false, true, {}, {}, 0, + AttentionMaskType::MASK_1D_KEY_SEQ_LEN, 0, sequence_length, true, false, true, disable_dml, {}, {}, 0, true); } } @@ -1075,11 +1090,15 @@ void RawAttentionPastStateBatch1(bool past_present_share_buffer) { batch_size, sequence_length, hidden_size, number_of_heads, false, is_unidirectional, use_past_state, past_sequence_length, &past_data, &present_data); } else { + // TODO: Unskip when fixed #41968513 + // DML doesn't support past_present_share_buffer for Attention yet + constexpr bool disable_dml = true; + RunAttentionTest(input_data, weight_data, bias_data, mask_index_data, output_data, batch_size, sequence_length, hidden_size, number_of_heads, false, is_unidirectional, use_past_state, past_sequence_length, &past_data, &present_data, AttentionMaskType::MASK_1D_KEY_SEQ_LEN, 0, past_sequence_length + sequence_length + 4, - true, false, true, {}, {}, 0, true); + true, false, true, disable_dml, {}, {}, 0, true); } } @@ -1204,11 +1223,15 @@ void RawAttentionPastStateBatch2(bool past_present_share_buffer) { batch_size, sequence_length, hidden_size, number_of_heads, false, is_unidirectional, use_past_state, past_sequence_length, &past_data, &present_data); } else { + // TODO: Unskip when fixed #41968513 + // DML doesn't support past_present_share_buffer for Attention yet + constexpr bool disable_dml = true; + RunAttentionTest(input_data, weight_data, bias_data, mask_index_data, output_data, batch_size, sequence_length, hidden_size, number_of_heads, false, is_unidirectional, use_past_state, past_sequence_length, &past_data, &present_data, AttentionMaskType::MASK_1D_KEY_SEQ_LEN, 0, past_sequence_length + sequence_length, - true, false, true, {}, {}, 0, true); + true, false, true, disable_dml, {}, {}, 0, true); } } @@ -1324,12 +1347,16 @@ void RawAttentionPastStateBatch2WithPadding(bool past_present_share_buffer) { use_past_state, past_sequence_length, &past_data, &present_data, AttentionMaskType::MASK_1D_END_START); } else { + // TODO: Unskip when fixed #41968513 + // DML doesn't support past_present_share_buffer for Attention yet + constexpr bool disable_dml = true; + RunAttentionTest(input_data, weight_data, bias_data, mask_index_data, output_data, batch_size, sequence_length, hidden_size, number_of_heads, false, is_unidirectional, use_past_state, past_sequence_length, &past_data, &present_data, AttentionMaskType::MASK_1D_END_START, 0, past_sequence_length + sequence_length + 4, - true, false, true, {}, {}, 0, true); + true, false, true, disable_dml, {}, {}, 0, true); } } @@ -1716,7 +1743,7 @@ TEST(AttentionTest, AttentionWithNormFactor) { batch_size, sequence_length, hidden_size, number_of_heads, use_float16, is_unidirectional, use_past_state, past_sequence_length, past_data, present_data, AttentionMaskType::MASK_2D_KEY_PADDING, 0 /*input_hidden_size*/, 0 /*max_sequence_length*/, - false /*disable_cpu*/, false /*disable_cuda*/, true /*disable_rocm*/, {} /*qkv_sizes*/, + false /*disable_cpu*/, false /*disable_cuda*/, true /*disable_rocm*/, false /*disable_dml*/, {} /*qkv_sizes*/, {} /*relative_position_bias_data*/, 0 /*kv_sequence_length*/, false /*past_present_share_buffer*/, true /*use_scale*/); } @@ -1757,11 +1784,16 @@ TEST(AttentionTest, AttentionWithNeoXRotaryEmbedding) { int past_sequence_length = 0; const std::vector* past_data = nullptr; const std::vector* present_data = nullptr; + + // TODO: Unskip when fixed #41968513 + // DML doesn't support do_rotary for Attention yet + constexpr bool disable_dml = true; + RunAttentionTest(input_data, weight_data, bias_data, mask_index_data, output_data, batch_size, sequence_length, hidden_size, number_of_heads, use_float16, is_unidirectional, use_past_state, past_sequence_length, past_data, present_data, AttentionMaskType::MASK_2D_KEY_PADDING, 0 /*input_hidden_size*/, 0 /*max_sequence_length*/, - true /*disable_cpu*/, false /*disable_cuda*/, true /*disable_rocm*/, {} /*qkv_sizes*/, + true /*disable_cpu*/, false /*disable_cuda*/, true /*disable_rocm*/, disable_dml, {} /*qkv_sizes*/, {} /*relative_position_bias_data*/, 0 /*kv_sequence_length*/, false /*past_present_share_buffer*/, true /*use_scale*/, true /*use_neox_rotary_embedding*/); } @@ -2191,7 +2223,8 @@ static void RunModelWithRandomInput( bool enable_cuda = HasCudaEnvironment(is_float16 ? 530 : 0); bool enable_rocm = (nullptr != DefaultRocmExecutionProvider().get()); bool enable_cpu = (nullptr != DefaultCpuExecutionProvider().get() && !is_float16); - if (enable_cuda || enable_rocm) { + bool enable_dml = (nullptr != DefaultDmlExecutionProvider().get()); + if (enable_cuda || enable_rocm || enable_dml) { OpTester test("Attention", 1, onnxruntime::kMSDomain); test.AddAttribute("num_heads", num_heads); if (is_float16) { @@ -2208,6 +2241,8 @@ static void RunModelWithRandomInput( std::vector> execution_providers; if (enable_cuda) { execution_providers.push_back(DefaultCudaExecutionProvider()); + } else if (enable_dml) { + execution_providers.push_back(DefaultDmlExecutionProvider()); } else { execution_providers.push_back(DefaultRocmExecutionProvider()); } diff --git a/onnxruntime/test/contrib_ops/attention_op_test_helper.cc b/onnxruntime/test/contrib_ops/attention_op_test_helper.cc index a4757cbaeb..bac91aa544 100644 --- a/onnxruntime/test/contrib_ops/attention_op_test_helper.cc +++ b/onnxruntime/test/contrib_ops/attention_op_test_helper.cc @@ -7,7 +7,7 @@ namespace onnxruntime { namespace test { -#ifndef _MSC_VER +#if !defined(_MSC_VER) || defined(USE_DML) void GetWeight_64_3_64(std::vector& weight_data) { weight_data = { -0.004707f, -0.006775f, 0.0009236f, 0.003067f, -0.00806f, 0.00779f, 0.0004425f, 0.00846f, 0.00048f, diff --git a/onnxruntime/test/contrib_ops/attention_op_test_helper.h b/onnxruntime/test/contrib_ops/attention_op_test_helper.h index 21cee275fd..0e2241e23e 100644 --- a/onnxruntime/test/contrib_ops/attention_op_test_helper.h +++ b/onnxruntime/test/contrib_ops/attention_op_test_helper.h @@ -43,7 +43,7 @@ struct AttentionTestData { }; // Disable some tests in Windows since prefast build might crash with large test data. -#ifndef _MSC_VER +#if !defined(_MSC_VER) || defined(USE_DML) // Return packed weights and bias for input projection. void GetAttentionWeight(std::vector& weight_data, int elements = 64 * 3 * 64, int offset = 0, int step = 1); void GetAttentionBias(std::vector& bias_data, int elements = 3 * 64, int offset = 0, int step = 1); diff --git a/onnxruntime/test/contrib_ops/multihead_attention_op_test.cc b/onnxruntime/test/contrib_ops/multihead_attention_op_test.cc index 9ca8e72b58..5cd42e815f 100644 --- a/onnxruntime/test/contrib_ops/multihead_attention_op_test.cc +++ b/onnxruntime/test/contrib_ops/multihead_attention_op_test.cc @@ -55,8 +55,8 @@ static void RunMultiHeadAttentionTest( bool use_float16 = false, bool disable_cpu = false, // some cases not supported in cpu right now. bool disable_cuda = false, - bool disable_rocm = DISABLE_ROCM) // not supported in rocm right now. -{ + bool disable_rocm = DISABLE_ROCM, // not supported in rocm right now. + bool disable_dml = false) { kv_sequence_length = (kv_sequence_length == 0 ? sequence_length : kv_sequence_length); int min_cuda_architecture = use_float16 ? 750 : 0; @@ -64,8 +64,9 @@ static void RunMultiHeadAttentionTest( // rocm mha is required to work with TunableOp Enabled bool enable_rocm = (nullptr != DefaultRocmExecutionProvider(/*test_tunable_op=*/true).get()) && !disable_rocm; bool enable_cpu = (nullptr != DefaultCpuExecutionProvider().get()) && !use_float16 && !disable_cpu; + bool enable_dml = (nullptr != DefaultDmlExecutionProvider().get()) && !disable_dml; - if (enable_cpu || enable_cuda || enable_rocm) { + if (enable_cpu || enable_cuda || enable_rocm || enable_dml) { OpTester tester("MultiHeadAttention", 1, onnxruntime::kMSDomain); tester.AddAttribute("num_heads", static_cast(num_heads)); tester.AddAttribute("mask_filter_value", static_cast(-10000.0f)); @@ -255,6 +256,12 @@ static void RunMultiHeadAttentionTest( execution_providers.push_back(DefaultCpuExecutionProvider()); tester.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers); } + + if (enable_dml) { + std::vector> execution_providers; + execution_providers.push_back(DefaultDmlExecutionProvider()); + tester.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers); + } } } @@ -284,7 +291,8 @@ static void RunMultiHeadAttentionKernel( bool is_static_kv = true, bool disable_cpu = false, // some cases not supported in cpu right now. bool disable_cuda = false, - bool disable_rocm = DISABLE_ROCM) { + bool disable_rocm = DISABLE_ROCM, + bool disable_dml = false) { if (kernel_type == AttentionKernelType::AttentionKernel_Default) { ScopedEnvironmentVariables scoped_env_vars{ EnvVarMap{ @@ -296,7 +304,7 @@ static void RunMultiHeadAttentionKernel( query_data, key_data, value_data, kv_data, qkv_data, bias_data, rel_pos_bias_data, past_key_data, past_value_data, present_key_data, present_value_data, key_padding_mask_data, mask_type, output_data, num_heads, batch_size, sequence_length, kv_sequence_length, - hidden_size, v_hidden_size, is_static_kv, use_float16, disable_cpu, disable_cuda, disable_rocm); + hidden_size, v_hidden_size, is_static_kv, use_float16, disable_cpu, disable_cuda, disable_rocm, disable_dml); return; } @@ -311,7 +319,7 @@ static void RunMultiHeadAttentionKernel( query_data, key_data, value_data, kv_data, qkv_data, bias_data, rel_pos_bias_data, past_key_data, past_value_data, present_key_data, present_value_data, key_padding_mask_data, mask_type, output_data, num_heads, batch_size, sequence_length, kv_sequence_length, - hidden_size, v_hidden_size, is_static_kv, use_float16, disable_cpu, disable_cuda, disable_rocm); + hidden_size, v_hidden_size, is_static_kv, use_float16, disable_cpu, disable_cuda, disable_rocm, disable_dml); return; } @@ -326,7 +334,7 @@ static void RunMultiHeadAttentionKernel( query_data, key_data, value_data, kv_data, qkv_data, bias_data, rel_pos_bias_data, past_key_data, past_value_data, present_key_data, present_value_data, key_padding_mask_data, mask_type, output_data, num_heads, batch_size, sequence_length, kv_sequence_length, - hidden_size, v_hidden_size, is_static_kv, use_float16, disable_cpu, disable_cuda, disable_rocm); + hidden_size, v_hidden_size, is_static_kv, use_float16, disable_cpu, disable_cuda, disable_rocm, disable_dml); return; } @@ -342,7 +350,7 @@ static void RunMultiHeadAttentionKernel( query_data, key_data, value_data, kv_data, qkv_data, bias_data, rel_pos_bias_data, past_key_data, past_value_data, present_key_data, present_value_data, key_padding_mask_data, mask_type, output_data, num_heads, batch_size, sequence_length, kv_sequence_length, - hidden_size, v_hidden_size, is_static_kv, use_float16, disable_cpu, disable_cuda, disable_rocm); + hidden_size, v_hidden_size, is_static_kv, use_float16, disable_cpu, disable_cuda, disable_rocm, disable_dml); return; } #endif @@ -358,7 +366,7 @@ static void RunMultiHeadAttentionKernel( query_data, key_data, value_data, kv_data, qkv_data, bias_data, rel_pos_bias_data, past_key_data, past_value_data, present_key_data, present_value_data, key_padding_mask_data, mask_type, output_data, num_heads, batch_size, sequence_length, kv_sequence_length, - hidden_size, v_hidden_size, is_static_kv, use_float16, disable_cpu, disable_cuda, disable_rocm); + hidden_size, v_hidden_size, is_static_kv, use_float16, disable_cpu, disable_cuda, disable_rocm, disable_dml); } } @@ -444,7 +452,7 @@ static void RunMultiHeadAttentionTests(AttentionTestData& data, bool disable_cpu } } -#ifndef _MSC_VER +#if !defined(_MSC_VER) || defined(USE_DML) // Test fused cross attention kernel // It requires head_size > 32 and head_size <= 64 for T4 GPU; hidden_size == v_hidden_size. TEST(MultiHeadAttentionTest, CrossAttention_Batch2_HeadSize40) { diff --git a/onnxruntime/test/providers/provider_test_utils.cc b/onnxruntime/test/providers/provider_test_utils.cc index 6396355642..a572b9eaa0 100644 --- a/onnxruntime/test/providers/provider_test_utils.cc +++ b/onnxruntime/test/providers/provider_test_utils.cc @@ -342,7 +342,7 @@ struct TensorCheck { #if defined(USE_TENSORRT) || defined(ENABLE_TRAINING_CORE) || defined(USE_CUDA) || defined(USE_ROCM) threshold = 0.005f; #elif defined(USE_DML) - threshold = 0.008f; + threshold = 0.02f; #endif for (int i = 0; i < size; ++i) { if (std::isnan(f_expected[i])) { diff --git a/packages.config b/packages.config index 97d03b55b6..b2c918c414 100644 --- a/packages.config +++ b/packages.config @@ -1,6 +1,6 @@  - + diff --git a/tools/ci_build/build.py b/tools/ci_build/build.py index 32d74ebe43..f88be53ee7 100644 --- a/tools/ci_build/build.py +++ b/tools/ci_build/build.py @@ -677,6 +677,8 @@ def parse_arguments(): parser.add_argument("--use_cache", action="store_true", help="Use compiler cache in CI") + parser.add_argument("--use_lock_free_queue", action="store_true", help="Use lock-free task queue for threadpool.") + if not is_windows(): parser.add_argument( "--allow_running_as_root", @@ -1307,6 +1309,9 @@ def generate_build_tree( if args.use_azure: add_default_definition(cmake_extra_defines, "onnxruntime_USE_AZURE", "ON") + if args.use_lock_free_queue: + add_default_definition(cmake_extra_defines, "onnxruntime_USE_LOCK_FREE_QUEUE", "ON") + cmake_args += [f"-D{define}" for define in cmake_extra_defines] cmake_args += cmake_extra_args diff --git a/tools/ci_build/github/azure-pipelines/win-ci-pipeline.yml b/tools/ci_build/github/azure-pipelines/win-ci-pipeline.yml index 708ea540d4..94e33a583e 100644 --- a/tools/ci_build/github/azure-pipelines/win-ci-pipeline.yml +++ b/tools/ci_build/github/azure-pipelines/win-ci-pipeline.yml @@ -200,7 +200,7 @@ stages: BuildConfig: 'RelWithDebInfo' EnvSetupScript: setup_env_azure.bat buildArch: x64 - additionalBuildFlags: --use_azure + additionalBuildFlags: --use_azure --use_lock_free_queue msbuildPlatform: x64 isX86: false job_name_suffix: x64_release_azure diff --git a/tools/nuget/generate_nuspec_for_native_nuget.py b/tools/nuget/generate_nuspec_for_native_nuget.py index f6cfcc3ce2..7445150b2e 100644 --- a/tools/nuget/generate_nuspec_for_native_nuget.py +++ b/tools/nuget/generate_nuspec_for_native_nuget.py @@ -190,7 +190,7 @@ def generate_repo_url(line_list, repo_url, commit_id): def generate_dependencies(xml_text, package_name, version): - dml_dependency = '' + dml_dependency = '' if package_name == "Microsoft.AI.MachineLearning": xml_text.append("")