saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-10 17:37:14 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Add support for LRN NHWC OPs (#19866)

Browse source 
Support LRN NHWC in the CUDA EP.

### Motivation and Context
Add support for all NHWC OPs to avoid NHWC/NCHW Layout transformation
This commit is contained in:
Markus Tavenrath 2024-03-14 01:52:07 +01:00 committed by GitHub
parent 9f08f8d5b2
commit f42e6ad61e
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 65 additions and 31 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

26
onnxruntime/core/framework/sequential_executor.cc
View file

@ -306,18 +306,20 @@ class KernelScope {
#endif
#ifdef ENABLE_NVTX_PROFILE
auto& node = kernel_.Node();
profile::NvtxRangeCreator& forward_range = session_scope_.forward_range_;
profile::NvtxRangeCreator& backward_range = session_scope_.backward_range_;
if (node.Description() != "Backward pass" && !forward_range.IsBeginCalled()) {
// Start timing forward pass when encountering the first forward node.
forward_range.Begin();
} else if (node.Description() == "Backward pass" && !backward_range.IsBeginCalled() &&
forward_range.IsBeginCalled()) {
// Start timing backward pass when encountering the first backward node.
// In the meanwhile, forward range ends.
forward_range.End();
backward_range.Begin();
{
auto& node = kernel_.Node();
profile::NvtxRangeCreator& forward_range = session_scope_.forward_range_;
profile::NvtxRangeCreator& backward_range = session_scope_.backward_range_;
if (node.Description() != "Backward pass" && !forward_range.IsBeginCalled()) {
// Start timing forward pass when encountering the first forward node.
forward_range.Begin();
} else if (node.Description() == "Backward pass" && !backward_range.IsBeginCalled() &&
forward_range.IsBeginCalled()) {
// Start timing backward pass when encountering the first backward node.
// In the meanwhile, forward range ends.
forward_range.End();
backward_range.Begin();
}
}
#endif

3
onnxruntime/core/optimizer/layout_transformation/layout_transformation.cc
View file

@ -44,7 +44,8 @@ const std::unordered_set<std::string_view>& GetCUDALayoutSensitiveOps() {
"AveragePool",
"GridSample",
"DepthToSpace",
"SpaceToDepth"};
"SpaceToDepth",
"LRN"};
}();
return cuda_nhwc_ops;
}

21
onnxruntime/core/providers/cuda/cuda_nhwc_kernels.cc
View file

@ -91,6 +91,15 @@ class ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSInter
class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSInternalNHWCDomain, 13, DepthToSpace);
class ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSInternalNHWCDomain, 1, 12, SpaceToDepth);
class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSInternalNHWCDomain, 13, SpaceToDepth);
class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(
kCudaExecutionProvider, kMSInternalNHWCDomain, 1, 12, float, LRN);
class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(
kCudaExecutionProvider, kMSInternalNHWCDomain, 1, 12, double, LRN);
class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(
kCudaExecutionProvider, kMSInternalNHWCDomain, 1, 12, MLFloat16, LRN);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSInternalNHWCDomain, 13, float, LRN);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSInternalNHWCDomain, 13, double, LRN);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSInternalNHWCDomain, 13, MLFloat16, LRN);
Status RegisterCudaNhwcKernels(KernelRegistry& kernel_registry) {
static const BuildKernelCreateInfoFn nhwc_function_table[] = {
@ -187,6 +196,18 @@ Status RegisterCudaNhwcKernels(KernelRegistry& kernel_registry) {
1, 12, SpaceToDepth)>,
BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSInternalNHWCDomain,
13, SpaceToDepth)>,
BuildKernelCreateInfo<ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(
kCudaExecutionProvider, kMSInternalNHWCDomain, 1, 12, float, LRN)>,
BuildKernelCreateInfo<ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(
kCudaExecutionProvider, kMSInternalNHWCDomain, 1, 12, double, LRN)>,
BuildKernelCreateInfo<ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(
kCudaExecutionProvider, kMSInternalNHWCDomain, 1, 12, MLFloat16, LRN)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(
kCudaExecutionProvider, kMSInternalNHWCDomain, 13, float, LRN)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(
kCudaExecutionProvider, kMSInternalNHWCDomain, 13, double, LRN)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(
kCudaExecutionProvider, kMSInternalNHWCDomain, 13, MLFloat16, LRN)>,
};
for (auto& function_table_entry : nhwc_function_table) {

44
onnxruntime/core/providers/cuda/nn/lrn.cc
View file

@ -6,37 +6,47 @@
namespace onnxruntime {
namespace cuda {
#define REGISTER_KERNEL_VERSIONED_TYPED(START_VER, END_VER, T) \
#define REGISTER_KERNEL_VERSIONED_TYPED(START_VER, END_VER, T, DOMAIN, LAYOUT) \
ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_EX( \
LRN, \
kOnnxDomain, \
DOMAIN, \
START_VER, \
END_VER, \
T, \
kCudaExecutionProvider, \
(*KernelDefBuilder::Create()).TypeConstraint("T", DataTypeImpl::GetTensorType<T>()), \
LRN<T>);
LRN<T, LAYOUT>);
#define REGISTER_KERNEL_TYPED(VER, T) \
#define REGISTER_KERNEL_TYPED(VER, T, DOMAIN, LAYOUT) \
ONNX_OPERATOR_TYPED_KERNEL_EX( \
LRN, \
kOnnxDomain, \
DOMAIN, \
VER, \
T, \
kCudaExecutionProvider, \
(*KernelDefBuilder::Create()).TypeConstraint("T", DataTypeImpl::GetTensorType<T>()), \
LRN<T>);
LRN<T, LAYOUT>);
REGISTER_KERNEL_VERSIONED_TYPED(1, 12, float)
REGISTER_KERNEL_VERSIONED_TYPED(1, 12, double)
REGISTER_KERNEL_VERSIONED_TYPED(1, 12, MLFloat16)
REGISTER_KERNEL_VERSIONED_TYPED(1, 12, float, kOnnxDomain, false)
REGISTER_KERNEL_VERSIONED_TYPED(1, 12, double, kOnnxDomain, false)
REGISTER_KERNEL_VERSIONED_TYPED(1, 12, MLFloat16, kOnnxDomain, false)
REGISTER_KERNEL_TYPED(13, float)
REGISTER_KERNEL_TYPED(13, double)
REGISTER_KERNEL_TYPED(13, MLFloat16)
REGISTER_KERNEL_TYPED(13, float, kOnnxDomain, false)
REGISTER_KERNEL_TYPED(13, double, kOnnxDomain, false)
REGISTER_KERNEL_TYPED(13, MLFloat16, kOnnxDomain, false)
template <typename T>
LRN<T>::LRN(const OpKernelInfo& info) : CudaKernel(info) {
#ifdef ENABLE_CUDA_NHWC_OPS
REGISTER_KERNEL_VERSIONED_TYPED(1, 12, float, kMSInternalNHWCDomain, true)
REGISTER_KERNEL_VERSIONED_TYPED(1, 12, double, kMSInternalNHWCDomain, true)
REGISTER_KERNEL_VERSIONED_TYPED(1, 12, MLFloat16, kMSInternalNHWCDomain, true)
REGISTER_KERNEL_TYPED(13, float, kMSInternalNHWCDomain, true)
REGISTER_KERNEL_TYPED(13, double, kMSInternalNHWCDomain, true)
REGISTER_KERNEL_TYPED(13, MLFloat16, kMSInternalNHWCDomain, true)
#endif
template <typename T, bool Layout>
LRN<T, Layout>::LRN(const OpKernelInfo& info) : CudaKernel(info) {
int64_t size;
ORT_ENFORCE(info.GetAttr<int64_t>("size", &size).IsOK());
ORT_ENFORCE(size > 0);
@ -58,8 +68,8 @@ LRN<T>::LRN(const OpKernelInfo& info) : CudaKernel(info) {
.IsOK());
}
template <typename T>
Status LRN<T>::ComputeInternal(OpKernelContext* context) const {
template <typename T, bool Layout>
Status LRN<T, Layout>::ComputeInternal(OpKernelContext* context) const {
typedef typename ToCudaType<T>::MappedType CudaT;
const Tensor* X = context->Input<Tensor>(0);
@ -71,7 +81,7 @@ Status LRN<T>::ComputeInternal(OpKernelContext* context) const {
Tensor* Y = context->Output(0, X->Shape());
CudnnTensor x_tensor;
ORT_RETURN_IF_ERROR(x_tensor.Set(X->Shape().GetDims(), CudnnTensor::GetDataType<CudaT>()));
ORT_RETURN_IF_ERROR(x_tensor.Set(X->Shape().GetDims(), CudnnTensor::GetDataType<CudaT>(), Layout == NHWC));
const auto one = Consts<CudaT>::One;
const auto zero = Consts<CudaT>::Zero;

2
onnxruntime/core/providers/cuda/nn/lrn.h
View file

@ -20,7 +20,7 @@ class CudnnLRNDescriptor final {
cudnnLRNDescriptor_t desc_;
};
template <typename T>
template <typename T, bool Layout>
class LRN : public CudaKernel {
public:
LRN(const OpKernelInfo& info);