diff --git a/onnxruntime/core/framework/sequential_executor.cc b/onnxruntime/core/framework/sequential_executor.cc index ea7f1397c9..0cc7294a46 100644 --- a/onnxruntime/core/framework/sequential_executor.cc +++ b/onnxruntime/core/framework/sequential_executor.cc @@ -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 diff --git a/onnxruntime/core/optimizer/layout_transformation/layout_transformation.cc b/onnxruntime/core/optimizer/layout_transformation/layout_transformation.cc index 085a02c7c4..7953cde668 100644 --- a/onnxruntime/core/optimizer/layout_transformation/layout_transformation.cc +++ b/onnxruntime/core/optimizer/layout_transformation/layout_transformation.cc @@ -44,7 +44,8 @@ const std::unordered_set& GetCUDALayoutSensitiveOps() { "AveragePool", "GridSample", "DepthToSpace", - "SpaceToDepth"}; + "SpaceToDepth", + "LRN"}; }(); return cuda_nhwc_ops; } diff --git a/onnxruntime/core/providers/cuda/cuda_nhwc_kernels.cc b/onnxruntime/core/providers/cuda/cuda_nhwc_kernels.cc index da7802fe8d..8fdcaacdb0 100644 --- a/onnxruntime/core/providers/cuda/cuda_nhwc_kernels.cc +++ b/onnxruntime/core/providers/cuda/cuda_nhwc_kernels.cc @@ -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, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, }; for (auto& function_table_entry : nhwc_function_table) { diff --git a/onnxruntime/core/providers/cuda/nn/lrn.cc b/onnxruntime/core/providers/cuda/nn/lrn.cc index 6fcdec74d8..788299b5eb 100644 --- a/onnxruntime/core/providers/cuda/nn/lrn.cc +++ b/onnxruntime/core/providers/cuda/nn/lrn.cc @@ -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()), \ - LRN); + LRN); -#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()), \ - LRN); + LRN); -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 -LRN::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 +LRN::LRN(const OpKernelInfo& info) : CudaKernel(info) { int64_t size; ORT_ENFORCE(info.GetAttr("size", &size).IsOK()); ORT_ENFORCE(size > 0); @@ -58,8 +68,8 @@ LRN::LRN(const OpKernelInfo& info) : CudaKernel(info) { .IsOK()); } -template -Status LRN::ComputeInternal(OpKernelContext* context) const { +template +Status LRN::ComputeInternal(OpKernelContext* context) const { typedef typename ToCudaType::MappedType CudaT; const Tensor* X = context->Input(0); @@ -71,7 +81,7 @@ Status LRN::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())); + ORT_RETURN_IF_ERROR(x_tensor.Set(X->Shape().GetDims(), CudnnTensor::GetDataType(), Layout == NHWC)); const auto one = Consts::One; const auto zero = Consts::Zero; diff --git a/onnxruntime/core/providers/cuda/nn/lrn.h b/onnxruntime/core/providers/cuda/nn/lrn.h index 319e323c72..31b2819ccc 100644 --- a/onnxruntime/core/providers/cuda/nn/lrn.h +++ b/onnxruntime/core/providers/cuda/nn/lrn.h @@ -20,7 +20,7 @@ class CudnnLRNDescriptor final { cudnnLRNDescriptor_t desc_; }; -template +template class LRN : public CudaKernel { public: LRN(const OpKernelInfo& info);