From c0e817ff161a558591180ab8d81412f39fa940d2 Mon Sep 17 00:00:00 2001 From: Yufeng Li Date: Thu, 23 Apr 2020 14:18:59 -0700 Subject: [PATCH 1/5] Fix a bug in skiplayernorm fusion pattern 2 (#3660) For skiplayernorm fusion pattern 2, its input[0] should be equal to the input[0] of Add_1, but is overridden by the input[0] of Add_2. --- .../core/optimizer/skip_layer_norm_fusion.cc | 42 ++++++++++-------- .../test/optimizer/graph_transform_test.cc | 33 ++++++++++++++ .../skip_layer_norm_input_output_check.onnx | Bin 0 -> 1149 bytes 3 files changed, 56 insertions(+), 19 deletions(-) create mode 100644 onnxruntime/test/testdata/transform/fusion/skip_layer_norm_input_output_check.onnx diff --git a/onnxruntime/core/optimizer/skip_layer_norm_fusion.cc b/onnxruntime/core/optimizer/skip_layer_norm_fusion.cc index c76781d28f..ccf9ea8aeb 100644 --- a/onnxruntime/core/optimizer/skip_layer_norm_fusion.cc +++ b/onnxruntime/core/optimizer/skip_layer_norm_fusion.cc @@ -57,13 +57,18 @@ static bool CheckFirstAdd(Node& add, ProviderType providertype) { // Add2 is the 2nd add of the to be fused sub-graph // The 1st input should be a 3D tensor // The 2nd input should be a 1D constant value -static bool CheckSecondAdd(Node& add, ProviderType providertype) { +static bool CheckSecondAdd(Graph& graph, Node& add, ProviderType providertype) { if (providertype != add.GetExecutionProviderType() || !IsSupportedDataType(add) || add.GetOutputEdgesCount() != 1) { return false; } + // The 2nd input should be a constant value + if (!graph_utils::NodeArgIsConstant(graph, *(add.MutableInputDefs()[1]))) { + return false; + } + // Check the input dimensions of the "Add" node. const TensorShapeProto* add_input1_shape = add.MutableInputDefs()[0]->Shape(); const TensorShapeProto* add_input2_shape = add.MutableInputDefs()[1]->Shape(); @@ -84,8 +89,8 @@ Skip Layer Normalization will fuse Add + LayerNormalization into one node, and a Before fusion: Format 1: - [Sub1] [Sub2] - \ / + [Sub1] C [Sub2] + \ / / Add2 / \ / Add1 @@ -93,10 +98,10 @@ Format 1: LayerNormalization Format 2: - [Sub1] [Sub2] - \ / - \ Add2 - \ / + [Sub1] [Sub2] C + \ \ / + \ Add2 + \ / Add1 | LayerNormalization @@ -131,10 +136,9 @@ Status SkipLayerNormFusion::ApplyImpl(Graph& graph, bool& modified, int graph_le const auto& node_topology_list = graph_viewer.GetNodesInTopologicalOrder(); std::vector> nodes_to_remove; for (auto node_index : node_topology_list) { - nodes_to_remove.clear(); Node* p_layernorm = graph.GetNode(node_index); if (p_layernorm == nullptr) - continue; // we removed the node as part of an earlier fusion. + continue; // node was removed in an earlier fusion. Node& ln_node = *p_layernorm; ORT_RETURN_IF_ERROR(Recurse(ln_node, modified, graph_level, logger)); @@ -167,7 +171,7 @@ Status SkipLayerNormFusion::ApplyImpl(Graph& graph, bool& modified, int graph_le p_add2 = const_cast(&edges[1]->GetNode()); if (CheckFirstAdd(*p_add1, ln_node.GetExecutionProviderType()) && - CheckSecondAdd(*p_add2, ln_node.GetExecutionProviderType())) { + CheckSecondAdd(graph, *p_add2, ln_node.GetExecutionProviderType())) { matched_format = Format::Format1; } } @@ -183,7 +187,7 @@ Status SkipLayerNormFusion::ApplyImpl(Graph& graph, bool& modified, int graph_le p_add2 = const_cast(&edges[1]->GetNode()); if (CheckFirstAdd(*p_add1, ln_node.GetExecutionProviderType()) && - CheckSecondAdd(*p_add2, ln_node.GetExecutionProviderType())) { + CheckSecondAdd(graph, *p_add2, ln_node.GetExecutionProviderType())) { matched_format = Format::Format2; } } @@ -230,18 +234,18 @@ Status SkipLayerNormFusion::ApplyImpl(Graph& graph, bool& modified, int graph_le "SkipLayerNormalization", "fused SkipLayerNorm subgraphs ", skip_layer_norm_input_defs, - {}, {}, kMSDomain); + ln_node.MutableOutputDefs(), {}, kMSDomain); // Assign provider to this new node. Provider should be same as the provider for old node. skip_layer_norm_node.SetExecutionProviderType(ln_node.GetExecutionProviderType()); - - // move input edges to add (first in list) across to the layer_norm_node. - // move output definitions and output edges from mul_node (last in list) to layer_norm_node. - // remove all the other nodes. - graph_utils::FinalizeNodeFusion(graph, nodes_to_remove, skip_layer_norm_node); - - modified = true; } + for (const auto& node : nodes_to_remove) { + graph_utils::RemoveNodeOutputEdges(graph, node); + graph.RemoveNode(node.get().Index()); + } + + modified = true; + return Status::OK(); } } // namespace onnxruntime diff --git a/onnxruntime/test/optimizer/graph_transform_test.cc b/onnxruntime/test/optimizer/graph_transform_test.cc index f5e0ec7ec5..caa5beb6bf 100644 --- a/onnxruntime/test/optimizer/graph_transform_test.cc +++ b/onnxruntime/test/optimizer/graph_transform_test.cc @@ -1609,6 +1609,39 @@ TEST_F(GraphTransformationTests, SkipLayerNormFusionTest) { TestSkipLayerNormFusion(MODEL_FOLDER "fusion/skip_layer_norm_format3_no_fusion.onnx", 1, 1, 0, logger_.get()); } +TEST_F(GraphTransformationTests, SkipLayerNormFusion_Input_Output_Check) { + auto model_uri = MODEL_FOLDER "fusion/skip_layer_norm_input_output_check.onnx"; + std::shared_ptr p_model; + ASSERT_STATUS_OK(Model::Load(model_uri, p_model, nullptr, *logger_)); + Graph& graph = p_model->MainGraph(); + + onnxruntime::GraphTransformerManager graph_transformation_mgr{5}; + graph_transformation_mgr.Register(onnxruntime::make_unique(), TransformerLevel::Level2); + graph_transformation_mgr.Register(onnxruntime::make_unique(), TransformerLevel::Level2); + auto ret = graph_transformation_mgr.ApplyTransformers(graph, TransformerLevel::Level2, *logger_); + ASSERT_TRUE(ret.IsOK()); + + for (Node& node : graph.Nodes()) { + if (node.OpType() == "SkipLayerNormalization") { + // check inputs + std::vector& input_defs = node.MutableInputDefs(); + EXPECT_EQ(input_defs.size(), 5u) << "SkipLayerNormalization number of inputs does not equal to 5. Got:" << node.InputDefs().size(); + EXPECT_EQ(input_defs[0]->Name(), "input.1"); + EXPECT_EQ(input_defs[1]->Name(), "6"); + EXPECT_EQ(input_defs[2]->Name(), "1"); + EXPECT_EQ(input_defs[3]->Name(), "2"); + EXPECT_EQ(input_defs[4]->Name(), "4"); + + // check outputs + std::vector& output_defs = node.MutableOutputDefs(); + EXPECT_EQ(node.OutputDefs().size(), 1u) << "SkipLayerNormalization number of outputs does not equal to 1. Got:" << node.OutputDefs().size(); + EXPECT_EQ(output_defs[0]->Name(), "19"); + } else { + EXPECT_EQ(node.OpType(), "MatMul") << "Unexpected node: " << node.OpType() << "," << node.Name(); + } + } +} + TEST_F(GraphTransformationTests, EmbedLayerNormFusionFormat1) { auto model_uri = MODEL_FOLDER "fusion/embed_layer_norm_format1.onnx"; std::shared_ptr p_model; diff --git a/onnxruntime/test/testdata/transform/fusion/skip_layer_norm_input_output_check.onnx b/onnxruntime/test/testdata/transform/fusion/skip_layer_norm_input_output_check.onnx new file mode 100644 index 0000000000000000000000000000000000000000..555c701239f9c9ac38a9ae9cd42e37d5b4b070b5 GIT binary patch literal 1149 zcma)5Z%k8H6z^?IdAmdCUUbUa$g(1pG{E}W(tq^cc5IDWWmJrAu$Z?3&8%X<^4iJ9 z7P26c!JIJ*YS9Sjvc;I}i+(}gF-&C(A=$F|Wkf$Xh=v#@A>o6%S#R5QZpoH#laq7a z`Q>-d@0{~4r{nWGy5#m)%U(evR#m(2a@c_LBhikyTqR<)#mpnSz}1E1x_FypHt4bu zUWbt#BO8w#Lf-4Gt&(WaFo$1yMnw)DIkTW83o?R}M=qhDA>0~o3D<=~QORn+&%R7+ zu{Lxl-04f@4jMuE_(jIU$i=Hgvyitl-YnS+8s@N>WkdnL(H`xTLs40(HgLI^-ISLI zwZ+3GPLBdwl}hCsL0Y^4vpblHjA;?_g6$tjc7ujF>||LB%Ub`X&an!0X%)7AsBr$T ziioLK4}? zna}j(z$o34Q_!WiB(ixkpF!B*(rvi}U1m$xOdVdIS$O5Xh+G*y)X^T3O(KW&s1}KV za+RvbyBf<;ENRN>j_Hhs!$rnsI?_VRS#vo5gC55dg7 zMRGNGlHNb9g|=899i1DXoYn}@l@a>LrKab1FDg^BU(o5vQhHWwqK#!fdU(rWW$XM7 zm@Ivl9OzY1=_fDz+V_N9awO;u^3N`e7v)qe1;Tc&Fk7Y<*lsE`EKDIE;O? z?9)+bE?*`OQe%pH;W$lB1(n8{W*C1Ifmlf+^<5X~o%um(>WGmt?+DzkNkZqs9y(R{ z95J3d0jeb<@E0ifyGxb$t)nzyZBk4{4Iui)uiWeJhwy|tYt_Z%eKfA)KN z=6c9^<1XcwZ(pU|vz;*Z+NkoCWe7G_4$$x|LN6441dT;QA?;ih}hh#e04L2u_PE#JghAWF+ v*!Iep^7cdtv>p3Rsl3z+O(lQ7c%H|44)Fz8&7r>xH>Y2m>%Wj%wf_01rN?X_ literal 0 HcmV?d00001 From 2d2375aa230b28a8bb56c217f1f1f5221e9fb2be Mon Sep 17 00:00:00 2001 From: Sheil Kumar Date: Thu, 23 Apr 2020 14:27:18 -0700 Subject: [PATCH 2/5] swap float16/float (#3663) Co-authored-by: Sheil Kumar --- onnxruntime/core/framework/onnxruntime_map_type_info.cc | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/onnxruntime/core/framework/onnxruntime_map_type_info.cc b/onnxruntime/core/framework/onnxruntime_map_type_info.cc index ec08108afb..e0b5ffc538 100644 --- a/onnxruntime/core/framework/onnxruntime_map_type_info.cc +++ b/onnxruntime/core/framework/onnxruntime_map_type_info.cc @@ -15,8 +15,8 @@ ToONNXTensorElementDataType(ONNX_NAMESPACE::TensorProto_DataType data_type) { switch (data_type) { case TensorType::TensorProto_DataType_BOOL: { return ONNXTensorElementDataType::ONNX_TENSOR_ELEMENT_DATA_TYPE_BOOL; } case TensorType::TensorProto_DataType_STRING: { return ONNXTensorElementDataType::ONNX_TENSOR_ELEMENT_DATA_TYPE_STRING; } // maps to c++ type std::string - case TensorType::TensorProto_DataType_FLOAT16: { return ONNXTensorElementDataType::ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT; } // maps to c type float - case TensorType::TensorProto_DataType_FLOAT: { return ONNXTensorElementDataType::ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT16; } + case TensorType::TensorProto_DataType_FLOAT16: { return ONNXTensorElementDataType::ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT16; } // maps to c type float16 + case TensorType::TensorProto_DataType_FLOAT: { return ONNXTensorElementDataType::ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT; } // maps to c type float case TensorType::TensorProto_DataType_DOUBLE: { return ONNXTensorElementDataType::ONNX_TENSOR_ELEMENT_DATA_TYPE_DOUBLE; } // maps to c type double case TensorType::TensorProto_DataType_INT8: { return ONNXTensorElementDataType::ONNX_TENSOR_ELEMENT_DATA_TYPE_INT8; } // maps to c type int8_t case TensorType::TensorProto_DataType_INT16: { return ONNXTensorElementDataType::ONNX_TENSOR_ELEMENT_DATA_TYPE_INT16; } // maps to c type int16_t From bae1dd7f0409968caf2c8f33038185878c2b3687 Mon Sep 17 00:00:00 2001 From: Ori Levari Date: Thu, 23 Apr 2020 15:37:32 -0700 Subject: [PATCH 3/5] add test for LearningModel creation from missing model path (#3661) --- winml/test/api/LearningModelAPITest.cpp | 14 +++++++++++++- winml/test/api/LearningModelAPITest.h | 2 ++ 2 files changed, 15 insertions(+), 1 deletion(-) diff --git a/winml/test/api/LearningModelAPITest.cpp b/winml/test/api/LearningModelAPITest.cpp index a6cd9ff2d6..24abbc361e 100644 --- a/winml/test/api/LearningModelAPITest.cpp +++ b/winml/test/api/LearningModelAPITest.cpp @@ -26,6 +26,17 @@ static void CreateModelFromFilePath() { WINML_EXPECT_NO_THROW(APITest::LoadModel(L"squeezenet_modifiedforruntimestests.onnx", learningModel)); } +static void CreateModelFileNotFound() { + LearningModel learningModel = nullptr; + + WINML_EXPECT_THROW_SPECIFIC( + APITest::LoadModel(L"missing_model.onnx", learningModel), + winrt::hresult_error, + [](const winrt::hresult_error& e) -> bool { + return e.code() == __HRESULT_FROM_WIN32(ERROR_FILE_NOT_FOUND); + }); +} + static void CreateModelFromIStorage() { std::wstring path = FileHelpers::GetModulePath() + L"squeezenet_modifiedforruntimestests.onnx"; auto storageFile = ws::StorageFile::GetFileFromPathAsync(path).get(); @@ -242,7 +253,7 @@ static void CloseModelNoNewSessions() { WINML_EXPECT_NO_THROW(learningModel.Close()); LearningModelSession session = nullptr; WINML_EXPECT_THROW_SPECIFIC( - session = LearningModelSession(learningModel);, + session = LearningModelSession(learningModel), winrt::hresult_error, [](const winrt::hresult_error& e) -> bool { return e.code() == E_INVALIDARG; @@ -255,6 +266,7 @@ const LearningModelApiTestsApi& getapi() { LearningModelAPITestsClassSetup, LearningModelAPITestsGpuMethodSetup, CreateModelFromFilePath, + CreateModelFileNotFound, CreateModelFromIStorage, CreateModelFromIStorageOutsideCwd, CreateModelFromIStream, diff --git a/winml/test/api/LearningModelAPITest.h b/winml/test/api/LearningModelAPITest.h index 376fb661f7..6881c99b31 100644 --- a/winml/test/api/LearningModelAPITest.h +++ b/winml/test/api/LearningModelAPITest.h @@ -7,6 +7,7 @@ struct LearningModelApiTestsApi SetupClass LearningModelAPITestsClassSetup; SetupTest LearningModelAPITestsGpuMethodSetup; VoidTest CreateModelFromFilePath; + VoidTest CreateModelFileNotFound; VoidTest CreateModelFromIStorage; VoidTest CreateModelFromIStorageOutsideCwd; VoidTest CreateModelFromIStream; @@ -27,6 +28,7 @@ WINML_TEST_CLASS_BEGIN(LearningModelAPITests) WINML_TEST_CLASS_SETUP_CLASS(LearningModelAPITestsClassSetup) WINML_TEST_CLASS_BEGIN_TESTS WINML_TEST(LearningModelAPITests, CreateModelFromFilePath) +WINML_TEST(LearningModelAPITests, CreateModelFileNotFound) WINML_TEST(LearningModelAPITests, CreateModelFromIStorage) WINML_TEST(LearningModelAPITests, CreateModelFromIStorageOutsideCwd) WINML_TEST(LearningModelAPITests, CreateModelFromIStream) From 2659f205cc597aeb997c87fb0b843b5b649ba73d Mon Sep 17 00:00:00 2001 From: Du Li Date: Thu, 23 Apr 2020 17:21:48 -0700 Subject: [PATCH 4/5] Complex multiplication and conjugate contrib ops (#3384) * adding ComplexMulConj * Adding fp16 support. * adding a util func --- .../contrib_ops/cuda/math/complex_mul.cc | 63 +++++++ .../contrib_ops/cuda/math/complex_mul.h | 22 +++ .../contrib_ops/cuda/math/complex_mul_impl.cu | 174 ++++++++++++++++++ .../contrib_ops/cuda/math/complex_mul_impl.h | 32 ++++ .../contrib_ops/cuda_contrib_kernels.cc | 8 + .../core/graph/contrib_ops/contrib_defs.cc | 18 ++ .../test/contrib_ops/element_wise_ops_test.cc | 104 +++++++++++ 7 files changed, 421 insertions(+) create mode 100644 onnxruntime/contrib_ops/cuda/math/complex_mul.cc create mode 100644 onnxruntime/contrib_ops/cuda/math/complex_mul.h create mode 100644 onnxruntime/contrib_ops/cuda/math/complex_mul_impl.cu create mode 100644 onnxruntime/contrib_ops/cuda/math/complex_mul_impl.h diff --git a/onnxruntime/contrib_ops/cuda/math/complex_mul.cc b/onnxruntime/contrib_ops/cuda/math/complex_mul.cc new file mode 100644 index 0000000000..1365009619 --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/math/complex_mul.cc @@ -0,0 +1,63 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "complex_mul.h" +#include "complex_mul_impl.h" + +namespace onnxruntime { +namespace contrib { +namespace cuda { +#define REGISTER_KERNEL_TYPED(T) \ + ONNX_OPERATOR_TYPED_KERNEL_EX( \ + ComplexMul, \ + kMSDomain, \ + 1, \ + T, \ + kCudaExecutionProvider, \ + KernelDefBuilder() \ + .TypeConstraint("T", DataTypeImpl::GetTensorType()), \ + ComplexMul); \ + ONNX_OPERATOR_TYPED_KERNEL_EX( \ + ComplexMulConj, \ + kMSDomain, \ + 1, \ + T, \ + kCudaExecutionProvider, \ + KernelDefBuilder() \ + .TypeConstraint("T", DataTypeImpl::GetTensorType()), \ + ComplexMul); + +REGISTER_KERNEL_TYPED(float) +REGISTER_KERNEL_TYPED(MLFloat16) + +template +Status ComplexMul::ComputeInternal(OpKernelContext* context) const { + for (int index = 0; index < context->InputCount(); ++index) { + const Tensor* input = context->Input(index); + TensorShape shape = input->Shape(); + int64_t last_dimension = shape[shape.NumDimensions() - 1]; + ORT_ENFORCE(last_dimension == 2, "The input_", index, " last dimension is supposed to be 2, but get ", last_dimension); + } + + BinaryElementwisePreparation prepare; + Prepare(context, &prepare); + ComplexMul_Impl::MappedType>( + prepare.output_rank_or_simple_broadcast, + &prepare.lhs_padded_strides, + reinterpret_cast::MappedType*>(prepare.lhs_tensor->template Data()), + &prepare.rhs_padded_strides, + reinterpret_cast::MappedType*>(prepare.rhs_tensor->template Data()), + &prepare.fdm_output_strides, + prepare.fdm_H, + prepare.fdm_C, + reinterpret_cast::MappedType*>(prepare.output_tensor->template MutableData()), + prepare.output_tensor->Shape().Size(), + prepare.lhs_tensor->Shape().Size(), + prepare.rhs_tensor->Shape().Size(), + is_conj); + return Status::OK(); +} + +} // namespace cuda +} // namespace contrib +} // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cuda/math/complex_mul.h b/onnxruntime/contrib_ops/cuda/math/complex_mul.h new file mode 100644 index 0000000000..b8cefdd4db --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/math/complex_mul.h @@ -0,0 +1,22 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once +#include "core/providers/cuda/math/binary_elementwise_ops.h" + +namespace onnxruntime { +namespace contrib { +namespace cuda { + +using namespace ::onnxruntime::cuda; + +template +class ComplexMul : public BinaryElementwise { + public: + ComplexMul(const OpKernelInfo info) : BinaryElementwise{info} {} + Status ComputeInternal(OpKernelContext* context) const override; +}; + +} // namespace cuda +} // namespace contrib +} // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cuda/math/complex_mul_impl.cu b/onnxruntime/contrib_ops/cuda/math/complex_mul_impl.cu new file mode 100644 index 0000000000..5b75df928e --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/math/complex_mul_impl.cu @@ -0,0 +1,174 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "complex_mul.h" +#include "complex_mul_impl.h" +#include "core/providers/cuda/cu_inc/common.cuh" +#include "core/providers/cuda/math/binary_elementwise_ops.h" + +namespace onnxruntime { +namespace contrib { +namespace cuda { +template +__device__ __inline__ void _ComplexMul(T a0, T a1, T b0, T b1, T* output_data, bool is_conj) { + if (is_conj) { + T out_real = a0 * b0 + a1 * b1; + T out_imag = a1 * b0 - a0 * b1; + output_data[0] = out_real; + output_data[1] = out_imag; + } else { + T out_real = a0 * b0 - a1 * b1; + T out_imag = a0 * b1 + a1 * b0; + output_data[0] = out_real; + output_data[1] = out_imag; + } +}; + +// broadcast by computing output coordinate from offset, using fast_divmod +template +__global__ void _ElementWiseWithStrideTwo( + int32_t output_rank, + const TArray lhs_padded_strides, + const T* lhs_data, + const TArray rhs_padded_strides, + const T* rhs_data, + const TArray fdm_output_strides, + T* output_data, + CUDA_LONG N, + int64_t lhs_size, + int64_t rhs_size, + bool is_conj) { + CUDA_LONG start = NumElementsPerThread * NumThreadsPerBlock * blockIdx.x + threadIdx.x; + T a[NumElementsPerThread]; + T b[NumElementsPerThread]; + T c[NumElementsPerThread]; + T d[NumElementsPerThread]; + + CUDA_LONG id = start; +#pragma unroll + for (int i = 0; i < NumElementsPerThread; i++) { + if (id < N / 2) { + CUDA_LONG lhs_index = (lhs_need_compute ? 0 : id); + CUDA_LONG rhs_index = (rhs_need_compute ? 0 : id); + // compute indexes with broadcasting rules: https://github.com/onnx/onnx/blob/master/docs/Broadcasting.md + CUDA_LONG offset = id; +#pragma unroll + for (auto dim = 0; dim < fdm_output_strides.GetCapacity(); dim++) { + if (dim >= output_rank) { + break; + } + int q, r; + fdm_output_strides.data_[dim].divmod(offset, q, r); + if (lhs_need_compute) { + lhs_index += static_cast(lhs_padded_strides.data_[dim]) * q; + } + + if (rhs_need_compute) { + rhs_index += static_cast(rhs_padded_strides.data_[dim]) * q; + } + offset = r; + } + + a[i] = lhs_data[(2 * lhs_index) % lhs_size]; + b[i] = lhs_data[(2 * lhs_index + 1) % lhs_size]; + c[i] = rhs_data[(2 * rhs_index) % rhs_size]; + d[i] = rhs_data[(2 * rhs_index + 1) % rhs_size]; + + id += NumThreadsPerBlock; + } + } + + id = start; +#pragma unroll + for (int i = 0; i < NumElementsPerThread; i++) { + if (id < N / 2) { + _ComplexMul(a[i], b[i], c[i], d[i], &output_data[2 * id], is_conj); + id += NumThreadsPerBlock; + } + } +}; + +template +void ComplexMul_Impl( + int32_t output_rank_or_simple_broadcast, + const TArray* lhs_padded_strides, + const T* lhs_data, + const TArray* rhs_padded_strides, + const T* rhs_data, + const TArray* fdm_output_strides, + const onnxruntime::cuda::fast_divmod& fdm_H, + const onnxruntime::cuda::fast_divmod& fdm_C, + T* output_data, + int64_t count, + int64_t lhs_size, + int64_t rhs_size, + bool is_conj) { + if (count == 0) // special case where there's a dim value of 0 in the output shape + return; + + int blocksPerGrid = static_cast(CeilDiv(count, GridDim::maxThreadsPerBlock * GridDim::maxElementsPerThread)); + CUDA_LONG N = static_cast(count); + + if (lhs_padded_strides && rhs_padded_strides && lhs_padded_strides->size_ && rhs_padded_strides->size_) + _ElementWiseWithStrideTwo<<>>( + output_rank_or_simple_broadcast, + *lhs_padded_strides, + lhs_data, + *rhs_padded_strides, + rhs_data, + *fdm_output_strides, + output_data, + N, + lhs_size, + rhs_size, + is_conj); + else if (lhs_padded_strides && lhs_padded_strides->size_) + _ElementWiseWithStrideTwo<<>>( + output_rank_or_simple_broadcast, + *lhs_padded_strides, + lhs_data, + *rhs_padded_strides, + rhs_data, + *fdm_output_strides, + output_data, + N, + lhs_size, + rhs_size, + is_conj); + else + _ElementWiseWithStrideTwo<<>>( + output_rank_or_simple_broadcast, + *lhs_padded_strides, + lhs_data, + *rhs_padded_strides, + rhs_data, + *fdm_output_strides, + output_data, + N, + lhs_size, + rhs_size, + is_conj); +}; + +#define SPECIALIZE_STACKEDCOMPLEXMUL_IMPL(T) \ + template void ComplexMul_Impl( \ + int32_t output_rank_or_simple_broadcast, \ + const TArray* lhs_padded_strides, \ + const T* lhs_data, \ + const TArray* rhs_padded_strides, \ + const T* rhs_data, \ + const TArray* fdm_output_strides, \ + const onnxruntime::cuda::fast_divmod& fdm_H, \ + const onnxruntime::cuda::fast_divmod& fdm_C, \ + T* output_data, \ + int64_t count, \ + int64_t lhs_size, \ + int64_t rhs_size, \ + bool is_conj); + +SPECIALIZE_STACKEDCOMPLEXMUL_IMPL(float) +SPECIALIZE_STACKEDCOMPLEXMUL_IMPL(half) + +} // namespace cuda +} // namespace contrib +} // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cuda/math/complex_mul_impl.h b/onnxruntime/contrib_ops/cuda/math/complex_mul_impl.h new file mode 100644 index 0000000000..d48eea9a9f --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/math/complex_mul_impl.h @@ -0,0 +1,32 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once +#include "complex_mul.h" +#include "core/providers/cuda/shared_inc/cuda_utils.h" + +namespace onnxruntime { +namespace contrib { +namespace cuda { + +using namespace ::onnxruntime::cuda; + +template +void ComplexMul_Impl( + int32_t output_rank_or_simple_broadcast, + const TArray* lhs_padded_strides, + const T* lhs_data, + const TArray* rhs_padded_strides, + const T* rhs_data, + const TArray* fdm_output_strides, + const onnxruntime::cuda::fast_divmod& fdm_H, + const onnxruntime::cuda::fast_divmod& fdm_C, + T* output_data, + int64_t count, + int64_t lhs_size, + int64_t rhs_size, + bool is_conj); + +} // namespace cuda +} // namespace contrib +} // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cuda_contrib_kernels.cc b/onnxruntime/contrib_ops/cuda_contrib_kernels.cc index eef175e7b8..b8bed20eae 100644 --- a/onnxruntime/contrib_ops/cuda_contrib_kernels.cc +++ b/onnxruntime/contrib_ops/cuda_contrib_kernels.cc @@ -23,6 +23,10 @@ class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1 class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, float, Irfft); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, double, Irfft); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, MLFloat16, Irfft); +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, float, ComplexMul); +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, MLFloat16, ComplexMul); +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, float, ComplexMulConj); +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, MLFloat16, ComplexMulConj); // These ops were experimental ops in onnx domain which have been removed now. We add them here as // contrib ops to maintain backward compatibility @@ -78,6 +82,10 @@ Status RegisterCudaContribKernels(KernelRegistry& kernel_registry) { BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, // These ops were experimental ops in onnx domain which have been removed now. We add them here as // contrib ops to maintain backward compatibility diff --git a/onnxruntime/core/graph/contrib_ops/contrib_defs.cc b/onnxruntime/core/graph/contrib_ops/contrib_defs.cc index 598154ac1e..4097f232bb 100644 --- a/onnxruntime/core/graph/contrib_ops/contrib_defs.cc +++ b/onnxruntime/core/graph/contrib_ops/contrib_defs.cc @@ -994,6 +994,24 @@ Sample echo operator.)DOC"); .Output(0, "Y", "output tensor", "T") .TypeConstraint("T", {"tensor(float)", "tensor(double)", "tensor(float16)"}, "Constrain input and output types to float or half tensors."); + ONNX_CONTRIB_OPERATOR_SCHEMA(ComplexMul) + .SetDomain(kMSDomain) + .SinceVersion(1) + .SetDoc(R"DOC()DOC") + .Input(0, "A", "input_0", "T") + .Input(1, "B", "input_1", "T") + .Output(0, "C", "output tensor", "T") + .TypeConstraint("T", {"tensor(float)", "tensor(double)", "tensor(float16)"}, "Constrain input and output types to float or half tensors."); + + ONNX_CONTRIB_OPERATOR_SCHEMA(ComplexMulConj) + .SetDomain(kMSDomain) + .SinceVersion(1) + .SetDoc(R"DOC()DOC") + .Input(0, "A", "input_0", "T") + .Input(1, "B", "input_1", "T") + .Output(0, "C", "output tensor", "T") + .TypeConstraint("T", {"tensor(float)", "tensor(double)", "tensor(float16)"}, "Constrain input and output types to float or half tensors."); + ONNX_CONTRIB_OPERATOR_SCHEMA(ConvTransposeWithDynamicPads) .SetDomain(kMSDomain) .SinceVersion(1) diff --git a/onnxruntime/test/contrib_ops/element_wise_ops_test.cc b/onnxruntime/test/contrib_ops/element_wise_ops_test.cc index 10da620c53..a7debe845e 100644 --- a/onnxruntime/test/contrib_ops/element_wise_ops_test.cc +++ b/onnxruntime/test/contrib_ops/element_wise_ops_test.cc @@ -114,5 +114,109 @@ TEST(BiasGeluTest, Two_One_Dim) { RunBiasGeluTest(input_a_data, input_b_data, {2, 4}, {4}); } +TEST(MathOpTest, ComplexMul) { + if (DefaultCudaExecutionProvider() == nullptr) return; + + std::vector input_a_data = { + -0.5f, 0.6f}; + + std::vector input_b_data = { + 0.8f, -0.5f, 0.0f, 1.f, + 0.5f, 0.2f, 0.3f, -0.6f}; + + std::vector output_data = { + -0.10f, 0.73f, + -0.60f, -0.50f, + -0.37f, 0.20f, + 0.21f, 0.48f}; + + OpTester tester("ComplexMul", 1, onnxruntime::kMSDomain); + tester.AddInput("A", {1, 2}, input_a_data); + tester.AddInput("B", {4, 2}, input_b_data); + tester.AddOutput("C", {4, 2}, output_data); + + std::vector> execution_providers; + execution_providers.push_back(DefaultCudaExecutionProvider()); + tester.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers); +} + +TEST(MathOpTest, ComplexMulConj) { + if (DefaultCudaExecutionProvider() == nullptr) return; + + std::vector input_a_data = { + -0.5f, 0.6f}; + + std::vector input_b_data = { + 0.8f, -0.5f, 0.0f, 1.f, + 0.5f, 0.2f, 0.3f, -0.6f}; + + std::vector output_data = { + -0.70f, 0.23f, + 0.60f, 0.50f, + -0.13f, 0.40f, + -0.51f, -0.12f}; + + OpTester tester("ComplexMulConj", 1, onnxruntime::kMSDomain); + tester.AddInput("A", {1, 2}, input_a_data); + tester.AddInput("B", {4, 2}, input_b_data); + tester.AddOutput("C", {4, 2}, output_data); + + std::vector> execution_providers; + execution_providers.push_back(DefaultCudaExecutionProvider()); + tester.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers); +} + +TEST(MathOpTest, ComplexMul_fp16) { + if (DefaultCudaExecutionProvider() == nullptr) return; + + std::vector input_a_data = { + MLFloat16(math::floatToHalf(-0.5f)), MLFloat16(math::floatToHalf(0.6f))}; + + std::vector input_b_data = { + MLFloat16(math::floatToHalf(0.8f)), MLFloat16(math::floatToHalf(-0.5f)), MLFloat16(math::floatToHalf(0.0f)), MLFloat16(math::floatToHalf(1.f)), + MLFloat16(math::floatToHalf(0.5f)), MLFloat16(math::floatToHalf(0.2f)), MLFloat16(math::floatToHalf(0.3f)), MLFloat16(math::floatToHalf(-0.6f))}; + + std::vector output_data = { + MLFloat16(math::floatToHalf(-0.10f)), MLFloat16(math::floatToHalf(0.73f)), + MLFloat16(math::floatToHalf(-0.60f)), MLFloat16(math::floatToHalf(-0.50f)), + MLFloat16(math::floatToHalf(-0.37f)), MLFloat16(math::floatToHalf(0.20f)), + MLFloat16(math::floatToHalf(0.21f)), MLFloat16(math::floatToHalf(0.48f))}; + + OpTester tester("ComplexMul", 1, onnxruntime::kMSDomain); + tester.AddInput("A", {1, 2}, input_a_data); + tester.AddInput("B", {4, 2}, input_b_data); + tester.AddOutput("C", {4, 2}, output_data); + + std::vector> execution_providers; + execution_providers.push_back(DefaultCudaExecutionProvider()); + tester.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers); +} + +TEST(MathOpTest, ComplexMulConj_fp16) { + if (DefaultCudaExecutionProvider() == nullptr) return; + + std::vector input_a_data = { + MLFloat16(math::floatToHalf(-0.5f)), MLFloat16(math::floatToHalf(0.6f))}; + + std::vector input_b_data = { + MLFloat16(math::floatToHalf(0.8f)), MLFloat16(math::floatToHalf(-0.5f)), MLFloat16(math::floatToHalf(0.0f)), MLFloat16(math::floatToHalf(1.f)), + MLFloat16(math::floatToHalf(0.5f)), MLFloat16(math::floatToHalf(0.2f)), MLFloat16(math::floatToHalf(0.3f)), MLFloat16(math::floatToHalf(-0.6f))}; + + std::vector output_data = { + MLFloat16(math::floatToHalf(-0.70f)), MLFloat16(math::floatToHalf(0.23f)), + MLFloat16(math::floatToHalf(0.60f)), MLFloat16(math::floatToHalf(0.50f)), + MLFloat16(math::floatToHalf(-0.13f)), MLFloat16(math::floatToHalf(0.40f)), + MLFloat16(math::floatToHalf(-0.51f)), MLFloat16(math::floatToHalf(-0.12f))}; + + OpTester tester("ComplexMulConj", 1, onnxruntime::kMSDomain); + tester.AddInput("A", {1, 2}, input_a_data); + tester.AddInput("B", {4, 2}, input_b_data); + tester.AddOutput("C", {4, 2}, output_data); + + std::vector> execution_providers; + execution_providers.push_back(DefaultCudaExecutionProvider()); + tester.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers); +} + } // namespace test } // namespace onnxruntime From 939d03666045f6e597b3f68e38a41d79400c5740 Mon Sep 17 00:00:00 2001 From: Pranav Sharma Date: Thu, 23 Apr 2020 18:24:46 -0700 Subject: [PATCH 5/5] Add omp impl for tryparallelfor and modify gelu to use fastgelu impl. (#3667) * Add omp impl for tryparallelfor and modify gelu to use fastgelu impl. * Address PR comments. --- .../onnxruntime/core/platform/threadpool.h | 58 ++++++++++++++----- onnxruntime/contrib_ops/cpu/activations.h | 45 +++++++++----- 2 files changed, 72 insertions(+), 31 deletions(-) diff --git a/include/onnxruntime/core/platform/threadpool.h b/include/onnxruntime/core/platform/threadpool.h index 6cd9d81df3..7b776a9e53 100644 --- a/include/onnxruntime/core/platform/threadpool.h +++ b/include/onnxruntime/core/platform/threadpool.h @@ -166,11 +166,26 @@ class ThreadPool { static void TryParallelFor(concurrency::ThreadPool* tp, std::ptrdiff_t total, const TensorOpCost& cost_per_unit, const std::function& fn) { +#ifdef _OPENMP + ORT_UNUSED_PARAMETER(cost_per_unit); + std::ptrdiff_t num_threads = concurrency::ThreadPool::NumThreads(tp); + if (total < num_threads) { + num_threads = total; + } +#pragma omp parallel for + for (std::ptrdiff_t i = 0; i < num_threads; i++) { + std::ptrdiff_t start, work_remaining; + PartitionWork(i, num_threads, total, &start, &work_remaining); + std::ptrdiff_t end = start + work_remaining; + fn(start, end); + } +#else if (tp == nullptr) { fn(0, total); return; } tp->ParallelFor(total, cost_per_unit, fn); +#endif } // Similar to ParallelFor above, but takes the specified scheduling strategy @@ -180,13 +195,28 @@ class ThreadPool { const std::function& fn); static void TryParallelFor(concurrency::ThreadPool* tp, std::ptrdiff_t total, const SchedulingParams& scheduling_params, - const std::function& fn) { + const std::function& fn) { +#ifdef _OPENMP + ORT_UNUSED_PARAMETER(scheduling_params); + std::ptrdiff_t num_threads = concurrency::ThreadPool::NumThreads(tp); + if (total < num_threads) { + num_threads = total; + } +#pragma omp parallel for + for (std::ptrdiff_t i = 0; i < num_threads; i++) { + std::ptrdiff_t start, work_remaining; + PartitionWork(i, num_threads, total, &start, &work_remaining); + std::ptrdiff_t end = start + work_remaining; + fn(start, end); + } +#else if (tp == nullptr) { fn(0, total); return; } tp->ParallelFor(total, scheduling_params, fn); - } +#endif + } // namespace concurrency // Prefer using this API to get the number of threads unless you know what you're doing. // This API takes into account if openmp is enabled/disabled and if the thread pool ptr is nullptr. @@ -208,16 +238,6 @@ class ThreadPool { // cutting them by halves void SimpleParallelFor(std::ptrdiff_t total, std::function fn); -#ifdef _OPENMP - template - inline static void TryBatchParallelFor(ThreadPool*, std::ptrdiff_t total, F&& fn, std::ptrdiff_t /*num_batches*/) { -#pragma omp parallel for - for (std::ptrdiff_t i = 0; i < total; ++i) { - fn(i); - } - } -#else - /** * Tries to call the given function in parallel, with calls split into (num_batches) batches. *\param num_batches If it is zero, it will be replaced to the value of NumThreads(). @@ -230,6 +250,14 @@ class ThreadPool { **/ template inline static void TryBatchParallelFor(ThreadPool* tp, std::ptrdiff_t total, F&& fn, std::ptrdiff_t num_batches) { +#ifdef _OPENMP + ORT_UNUSED_PARAMETER(tp); + ORT_UNUSED_PARAMETER(num_batches); +#pragma omp parallel for + for (std::ptrdiff_t i = 0; i < total; ++i) { + fn(i); + } +#else if (tp == nullptr) { for (std::ptrdiff_t i = 0; i < total; ++i) { // In many cases, fn can be inlined here. @@ -264,8 +292,8 @@ class ThreadPool { fn(i); } }); - } #endif + } #ifndef _OPENMP //Deprecated. Please avoid using Eigen Tensor because it will blow up binary size quickly. @@ -291,7 +319,7 @@ class ThreadPool { Eigen::ThreadPoolInterface* underlying_threadpool_; // eigen_threadpool_ is instantiated and owned by thread::ThreadPool if // user_threadpool is not in the constructor. - std::unique_ptr> eigen_threadpool_; + std::unique_ptr > eigen_threadpool_; #ifndef _OPENMP std::unique_ptr threadpool_device_; #endif @@ -309,7 +337,7 @@ class ThreadPool { *WorkRemaining = WorkPerThread; } } -}; +}; // namespace concurrency } // namespace concurrency } // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cpu/activations.h b/onnxruntime/contrib_ops/cpu/activations.h index 4d2baee3bd..828e023bb2 100644 --- a/onnxruntime/contrib_ops/cpu/activations.h +++ b/onnxruntime/contrib_ops/cpu/activations.h @@ -37,23 +37,36 @@ class Gelu : public OpKernel { Gelu(const OpKernelInfo& info) : OpKernel(info) {} Status Compute(OpKernelContext* context) const override { - const auto* X = context->Input(0); - Tensor* Y = context->Output(0, X->Shape()); + const Tensor* input = context->Input(0); + const T* input_data = input->template Data(); + + Tensor* output = context->Output(0, input->Shape()); + T* output_data = output->template MutableData(); + concurrency::ThreadPool* tp = context->GetOperatorThreadPool(); - const int64_t input_size = X->Shape().Size(); - std::ptrdiff_t batch_size = static_cast(input_size); - //The cost comes from microbenchmark(manual tunning). - const double cost = 10.0; - const T* data = X->template Data(); - T* output = Y->template MutableData(); - concurrency::ThreadPool::TryParallelFor(tp, batch_size, cost, [data, output](ptrdiff_t first, ptrdiff_t last) { - ptrdiff_t len = last - first; - onnxruntime::ConstEigenVectorArrayMap xm(data + first, len); - onnxruntime::EigenVectorArrayMap ym(output + first, len); - ym = xm * static_cast(M_SQRT1_2); - MlasComputeErf(output, output, len); - ym = xm * 0.5f * (ym + 1.0f); - }); + int64_t elem_count = input->Shape().Size(); + static const int64_t length_per_task = 4096; // this number comes from FastGelu. + int64_t task_count = (elem_count + length_per_task - 1) / length_per_task; + concurrency::ThreadPool::TryBatchParallelFor( + tp, static_cast(task_count), + [&](ptrdiff_t task_idx) { + const auto start = task_idx * length_per_task; + const T* p_input = input_data + start; + T* p_output = output_data + start; + int64_t count = std::min(length_per_task, elem_count - start); + + for (int64_t i = 0; i < count; i++) { + T value = p_input[i]; + p_output[i] = value * static_cast(M_SQRT1_2); + } + + MlasComputeErf(p_output, p_output, count); + + for (int64_t i = 0; i < count; i++) { + p_output[i] = 0.5f * p_input[i] * (p_output[i] + 1.0f); + } + }, + 0); return Status::OK(); } };