Merge remote-tracking branch 'origin/master' into ort_training_for_merge_to_master

This commit is contained in:
Edward Chen 2020-04-24 02:27:27 +00:00
commit 5a790a4b42
15 changed files with 566 additions and 68 deletions

View file

@ -166,11 +166,26 @@ class ThreadPool {
static void TryParallelFor(concurrency::ThreadPool* tp, std::ptrdiff_t total, const TensorOpCost& cost_per_unit,
const std::function<void(std::ptrdiff_t first, std::ptrdiff_t last)>& 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<void(std::ptrdiff_t, std::ptrdiff_t)>& fn);
static void TryParallelFor(concurrency::ThreadPool* tp, std::ptrdiff_t total, const SchedulingParams& scheduling_params,
const std::function<void(std::ptrdiff_t, std::ptrdiff_t)>& fn) {
const std::function<void(std::ptrdiff_t first, std::ptrdiff_t last)>& 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<void(std::ptrdiff_t)> fn);
#ifdef _OPENMP
template <typename F>
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 <typename F>
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<ThreadPoolTempl<Env>> eigen_threadpool_;
std::unique_ptr<ThreadPoolTempl<Env> > eigen_threadpool_;
#ifndef _OPENMP
std::unique_ptr<Eigen::ThreadPoolDevice> threadpool_device_;
#endif
@ -309,7 +337,7 @@ class ThreadPool {
*WorkRemaining = WorkPerThread;
}
}
};
}; // namespace concurrency
} // namespace concurrency
} // namespace onnxruntime

View file

@ -37,38 +37,36 @@ class Gelu : public OpKernel {
Gelu(const OpKernelInfo& info) : OpKernel(info) {}
Status Compute(OpKernelContext* context) const override {
const auto* X = context->Input<Tensor>(0);
Tensor* Y = context->Output(0, X->Shape());
concurrency::ThreadPool* tp = context->GetOperatorThreadPool();
if (nullptr != tp) {
const T* input = X->template Data<T>();
T* output = Y->template MutableData<T>();
int task_count = tp->NumThreads();
int64_t elem_count = X->Shape().Size();
if (elem_count > task_count) {
tp->SimpleParallelFor(task_count, [input,
output,
elem_count,
task_count](std::ptrdiff_t i) {
int64_t elem_inx_start = i * elem_count / task_count;
int64_t elem_inx_end = (i + 1) * elem_count / task_count;
for (int64_t elem_inx = elem_inx_start; elem_inx < elem_inx_end; elem_inx++) {
output[elem_inx] = input[elem_inx] * static_cast<float>(M_SQRT1_2);
}
MlasComputeErf(output + elem_inx_start, output + elem_inx_start, elem_inx_end - elem_inx_start);
for (int64_t elem_inx = elem_inx_start; elem_inx < elem_inx_end; elem_inx++) {
output[elem_inx] = 0.5f * input[elem_inx] * (output[elem_inx] + 1.0f);
}
});
return Status::OK();
}
}
const Tensor* input = context->Input<Tensor>(0);
const T* input_data = input->template Data<T>();
EIGEN_X_VAR(xm);
EIGEN_Y_VAR(ym);
ym = xm * static_cast<float>(M_SQRT1_2);
MlasComputeErf(Y->template MutableData<T>(), Y->template MutableData<T>(), X->Shape().Size());
ym = xm * 0.5f * (ym + 1.0f);
Tensor* output = context->Output(0, input->Shape());
T* output_data = output->template MutableData<T>();
concurrency::ThreadPool* tp = context->GetOperatorThreadPool();
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<int32_t>(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<T>(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();
}
};

View file

@ -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<T>()), \
ComplexMul<T, false>); \
ONNX_OPERATOR_TYPED_KERNEL_EX( \
ComplexMulConj, \
kMSDomain, \
1, \
T, \
kCudaExecutionProvider, \
KernelDefBuilder() \
.TypeConstraint("T", DataTypeImpl::GetTensorType<T>()), \
ComplexMul<T, true>);
REGISTER_KERNEL_TYPED(float)
REGISTER_KERNEL_TYPED(MLFloat16)
template <typename T, bool is_conj>
Status ComplexMul<T, is_conj>::ComputeInternal(OpKernelContext* context) const {
for (int index = 0; index < context->InputCount(); ++index) {
const Tensor* input = context->Input<Tensor>(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<typename ToCudaType<T>::MappedType>(
prepare.output_rank_or_simple_broadcast,
&prepare.lhs_padded_strides,
reinterpret_cast<const typename ToCudaType<T>::MappedType*>(prepare.lhs_tensor->template Data<T>()),
&prepare.rhs_padded_strides,
reinterpret_cast<const typename ToCudaType<T>::MappedType*>(prepare.rhs_tensor->template Data<T>()),
&prepare.fdm_output_strides,
prepare.fdm_H,
prepare.fdm_C,
reinterpret_cast<typename ToCudaType<T>::MappedType*>(prepare.output_tensor->template MutableData<T>()),
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

View file

@ -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 <typename T, bool is_conj>
class ComplexMul : public BinaryElementwise<ShouldBroadcast> {
public:
ComplexMul(const OpKernelInfo info) : BinaryElementwise{info} {}
Status ComputeInternal(OpKernelContext* context) const override;
};
} // namespace cuda
} // namespace contrib
} // namespace onnxruntime

View file

@ -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 <typename T>
__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 <typename T, bool lhs_need_compute, bool rhs_need_compute, int NumThreadsPerBlock, int NumElementsPerThread>
__global__ void _ElementWiseWithStrideTwo(
int32_t output_rank,
const TArray<int64_t> lhs_padded_strides,
const T* lhs_data,
const TArray<int64_t> rhs_padded_strides,
const T* rhs_data,
const TArray<fast_divmod> 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<int>(lhs_padded_strides.data_[dim]) * q;
}
if (rhs_need_compute) {
rhs_index += static_cast<int>(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 <typename T>
void ComplexMul_Impl(
int32_t output_rank_or_simple_broadcast,
const TArray<int64_t>* lhs_padded_strides,
const T* lhs_data,
const TArray<int64_t>* rhs_padded_strides,
const T* rhs_data,
const TArray<onnxruntime::cuda::fast_divmod>* 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<int>(CeilDiv(count, GridDim::maxThreadsPerBlock * GridDim::maxElementsPerThread));
CUDA_LONG N = static_cast<CUDA_LONG>(count);
if (lhs_padded_strides && rhs_padded_strides && lhs_padded_strides->size_ && rhs_padded_strides->size_)
_ElementWiseWithStrideTwo<T, true, true, GridDim::maxThreadsPerBlock, GridDim::maxElementsPerThread><<<blocksPerGrid, GridDim::maxThreadsPerBlock, 0>>>(
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<T, true, false, GridDim::maxThreadsPerBlock, GridDim::maxElementsPerThread><<<blocksPerGrid, GridDim::maxThreadsPerBlock, 0>>>(
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<T, false, true, GridDim::maxThreadsPerBlock, GridDim::maxElementsPerThread><<<blocksPerGrid, GridDim::maxThreadsPerBlock, 0>>>(
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<T>( \
int32_t output_rank_or_simple_broadcast, \
const TArray<int64_t>* lhs_padded_strides, \
const T* lhs_data, \
const TArray<int64_t>* rhs_padded_strides, \
const T* rhs_data, \
const TArray<onnxruntime::cuda::fast_divmod>* 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

View file

@ -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 <typename T>
void ComplexMul_Impl(
int32_t output_rank_or_simple_broadcast,
const TArray<int64_t>* lhs_padded_strides,
const T* lhs_data,
const TArray<int64_t>* rhs_padded_strides,
const T* rhs_data,
const TArray<onnxruntime::cuda::fast_divmod>* 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

View file

@ -26,6 +26,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
@ -84,6 +88,10 @@ Status RegisterCudaContribKernels(KernelRegistry& kernel_registry) {
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, float, Irfft)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, double, Irfft)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, MLFloat16, Irfft)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, float, ComplexMul)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, MLFloat16, ComplexMul)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, float, ComplexMulConj)>,
BuildKernelCreateInfo<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

View file

@ -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

View file

@ -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)

View file

@ -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<std::reference_wrapper<Node>> 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<Node*>(&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<Node*>(&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

View file

@ -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<float> input_a_data = {
-0.5f, 0.6f};
std::vector<float> input_b_data = {
0.8f, -0.5f, 0.0f, 1.f,
0.5f, 0.2f, 0.3f, -0.6f};
std::vector<float> 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<float>("A", {1, 2}, input_a_data);
tester.AddInput<float>("B", {4, 2}, input_b_data);
tester.AddOutput<float>("C", {4, 2}, output_data);
std::vector<std::unique_ptr<IExecutionProvider>> execution_providers;
execution_providers.push_back(DefaultCudaExecutionProvider());
tester.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers);
}
TEST(MathOpTest, ComplexMulConj) {
if (DefaultCudaExecutionProvider() == nullptr) return;
std::vector<float> input_a_data = {
-0.5f, 0.6f};
std::vector<float> input_b_data = {
0.8f, -0.5f, 0.0f, 1.f,
0.5f, 0.2f, 0.3f, -0.6f};
std::vector<float> 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<float>("A", {1, 2}, input_a_data);
tester.AddInput<float>("B", {4, 2}, input_b_data);
tester.AddOutput<float>("C", {4, 2}, output_data);
std::vector<std::unique_ptr<IExecutionProvider>> 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<MLFloat16> input_a_data = {
MLFloat16(math::floatToHalf(-0.5f)), MLFloat16(math::floatToHalf(0.6f))};
std::vector<MLFloat16> 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<MLFloat16> 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<MLFloat16>("A", {1, 2}, input_a_data);
tester.AddInput<MLFloat16>("B", {4, 2}, input_b_data);
tester.AddOutput<MLFloat16>("C", {4, 2}, output_data);
std::vector<std::unique_ptr<IExecutionProvider>> 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<MLFloat16> input_a_data = {
MLFloat16(math::floatToHalf(-0.5f)), MLFloat16(math::floatToHalf(0.6f))};
std::vector<MLFloat16> 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<MLFloat16> 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<MLFloat16>("A", {1, 2}, input_a_data);
tester.AddInput<MLFloat16>("B", {4, 2}, input_b_data);
tester.AddOutput<MLFloat16>("C", {4, 2}, output_data);
std::vector<std::unique_ptr<IExecutionProvider>> execution_providers;
execution_providers.push_back(DefaultCudaExecutionProvider());
tester.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers);
}
} // namespace test
} // namespace onnxruntime

View file

@ -1723,6 +1723,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<Model> 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<LayerNormFusion>(), TransformerLevel::Level2);
graph_transformation_mgr.Register(onnxruntime::make_unique<SkipLayerNormFusion>(), 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<NodeArg*>& 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<NodeArg*>& 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<Model> p_model;

View file

@ -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,

View file

@ -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)