From 161a9d1d6d46dd894be86383bca7fce8d36957cc Mon Sep 17 00:00:00 2001 From: Changming Sun Date: Thu, 27 Jul 2023 16:37:55 -0700 Subject: [PATCH] Add some safety check for conv op (#16839) ### Description Add some safety check for conv op. It is to validate if the attributes coming from a conv op are in a valid range. (shouldn't be too large or too small). --- cmake/CMakeLists.txt | 1 + onnxruntime/contrib_ops/cpu/nchwc_ops.cc | 10 ++-- onnxruntime/core/providers/common.h | 10 ++-- .../core/providers/coreml/model/model.mm | 1 - onnxruntime/core/providers/cpu/nn/conv.cc | 54 +++++++++---------- 5 files changed, 38 insertions(+), 38 deletions(-) diff --git a/cmake/CMakeLists.txt b/cmake/CMakeLists.txt index dce2b756c1..104b8ac213 100644 --- a/cmake/CMakeLists.txt +++ b/cmake/CMakeLists.txt @@ -926,6 +926,7 @@ function(onnxruntime_set_compile_flags target_name) # float16.h:90:12: error: ‘tmp’ is used uninitialized list(APPEND ORT_HIP_WARNING_FLAGS -Wno-uninitialized) + list(APPEND ORT_HIP_WARNING_FLAGS -Wno-deprecated-copy) # some #pragma unroll will fail, do not treat them as error # #warning must not be treated as error diff --git a/onnxruntime/contrib_ops/cpu/nchwc_ops.cc b/onnxruntime/contrib_ops/cpu/nchwc_ops.cc index ad4a7a591a..13748b43b1 100644 --- a/onnxruntime/contrib_ops/cpu/nchwc_ops.cc +++ b/onnxruntime/contrib_ops/cpu/nchwc_ops.cc @@ -189,16 +189,16 @@ Status NchwcConv::Compute(OpKernelContext* context) const { TensorShape input_shape = X->Shape().Slice(2); ORT_RETURN_IF_ERROR(conv_attrs_.InferPadsAndOutputShape(input_shape, kernel_shape, strides, dilations, pads, Y_dims)); auto* Y = context->Output(0, Y_dims); - auto* y_data = Y->MutableData(); + auto y_data = Y->MutableDataAsSpan(); // Check for the optional Conv/Sum fusion. if (Sum != nullptr) { const auto& sum_shape = Sum->Shape(); ORT_RETURN_IF_NOT(Y->Shape() == sum_shape, "output and sum shape must match"); // If the output was not allocated inplace with the sum tensor, then copy here. - const auto* sum_data = Sum->Data(); - if (y_data != sum_data) { - memcpy(y_data, sum_data, SafeInt(sum_shape.Size()) * sizeof(float)); + auto sum_data = Sum->DataAsSpan(); + if (y_data.data() != sum_data.data()) { + gsl::copy(sum_data, y_data); } } @@ -213,7 +213,7 @@ Status NchwcConv::Compute(OpKernelContext* context) const { X->Data(), W->Data(), B != nullptr ? B->Data() : nullptr, - y_data, + y_data.data(), &activation_, Sum == nullptr, context->GetOperatorThreadPool()); diff --git a/onnxruntime/core/providers/common.h b/onnxruntime/core/providers/common.h index 6990bb9f63..16c8e2045f 100644 --- a/onnxruntime/core/providers/common.h +++ b/onnxruntime/core/providers/common.h @@ -5,6 +5,7 @@ #include #include +#include "core/common/safeint.h" #ifndef SHARED_PROVIDER #include "core/common/common.h" @@ -103,8 +104,8 @@ inline Status ComputePad(const int64_t in_dim, // The ONNX spec says if `auto_pad` attribute is set, pad until the `legacy_target_size` // is `ceil (in_dim / stride)`. The following line of code is essentially just that and // is retained as is - int64_t legacy_target_size = (in_dim + stride - 1) / stride; - int64_t pad_needed = (legacy_target_size - 1) * stride + kernel - in_dim; + SafeInt legacy_target_size = (SafeInt(in_dim) + stride - 1) / stride; + SafeInt pad_needed = (legacy_target_size - 1) * stride + kernel - in_dim; // make sure padding is symmetric if (force_symmetric_auto_padding) { // Inlining math::roundUpPow2() from util/math.h to avoid bringing in the transitive dependencies. @@ -128,8 +129,9 @@ inline Status ComputePad(const int64_t in_dim, constexpr inline int64_t ComputeOutputShape(const int64_t in_dim, const int64_t stride, const int64_t kernel, const int64_t dilation, const int64_t pad_head, const int64_t pad_tail) { - const int64_t dkernel = dilation * (kernel - 1) + 1; - return static_cast(static_cast(in_dim + pad_head + pad_tail - dkernel) / stride + 1); + const SafeInt dkernel = SafeInt(dilation) * (kernel - 1) + 1; + int64_t dkernel_value = SafeInt(in_dim) + pad_head + pad_tail - dkernel; + return static_cast(static_cast(dkernel_value) / stride + 1); } inline Status ComputePadAndOutputShape(const int64_t in_dim, diff --git a/onnxruntime/core/providers/coreml/model/model.mm b/onnxruntime/core/providers/coreml/model/model.mm index cb2f70294c..1ce5acd2c1 100644 --- a/onnxruntime/core/providers/coreml/model/model.mm +++ b/onnxruntime/core/providers/coreml/model/model.mm @@ -8,7 +8,6 @@ #include "core/common/common.h" #include "core/common/logging/logging.h" #include "core/graph/onnx_protobuf.h" -#include "core/providers/common.h" #include "core/providers/coreml/builders/helper.h" #include "core/providers/coreml/coreml_provider_factory.h" #include "host_utils.h" diff --git a/onnxruntime/core/providers/cpu/nn/conv.cc b/onnxruntime/core/providers/cpu/nn/conv.cc index 1a40b05cc3..51dfb143fb 100644 --- a/onnxruntime/core/providers/cpu/nn/conv.cc +++ b/onnxruntime/core/providers/cpu/nn/conv.cc @@ -195,18 +195,18 @@ Status Conv::Compute(OpKernelContext* context) const { AllocatorPtr alloc; ORT_RETURN_IF_ERROR(context->GetTempSpaceAllocator(&alloc)); - const auto* Xdata = X->Data(); + auto Xdata = X->DataAsSpan(); const auto* Bdata = B != nullptr ? B->Data() : nullptr; - auto* Ydata = Y->MutableData(); + auto Ydata = Y->MutableDataAsSpan(); // Check for the optional Conv/Sum fusion. float Beta = 0.0f; if (Sum != nullptr) { const auto& sum_shape = Sum->Shape(); ORT_RETURN_IF_NOT(Y->Shape() == sum_shape, "output and sum shape must match"); // If the output was not allocated inplace with the sum tensor, then copy here. - const auto* sum_data = Sum->Data(); - if (Ydata != sum_data) { - memcpy(Ydata, sum_data, SafeInt(sum_shape.Size()) * sizeof(float)); + auto sum_data = Sum->DataAsSpan(); + if (Ydata.data() != sum_data.data()) { + gsl::copy(sum_data, Ydata); } Beta = 1.0f; } @@ -218,16 +218,16 @@ Status Conv::Compute(OpKernelContext* context) const { size_t WorkingBufferSize; MlasConvPrepare(&Parameters, kernel_rank, - static_cast(N), - static_cast(conv_attrs_.group), - static_cast(C / conv_attrs_.group), + narrow(N), + narrow(conv_attrs_.group), + narrow(C / conv_attrs_.group), input_shape.GetDims().data(), kernel_shape.data(), dilations.data(), pads.data(), strides.data(), output_shape.GetDims().data(), - static_cast(M / conv_attrs_.group), + narrow(M / conv_attrs_.group), &activation_, &WorkingBufferSize, Beta, @@ -238,30 +238,28 @@ Status Conv::Compute(OpKernelContext* context) const { BufferUniquePtr working_buffer(working_data, BufferDeleter(std::move(alloc))); MlasConv(&Parameters, - Xdata, + Xdata.data(), W->Data(), Bdata, static_cast(working_buffer.get()), - Ydata, + Ydata.data(), thread_pool); } else { const int64_t input_image_size = input_shape.Size(); const int64_t output_image_size = output_shape.Size(); const int64_t kernel_size = TensorShape(kernel_shape).Size(); - const int64_t X_offset = C / conv_attrs_.group * input_image_size; - const int64_t Y_offset = Y->Shape().Size() / Y->Shape()[0] / conv_attrs_.group; - const int64_t W_offset = W->Shape().Size() / conv_attrs_.group; - const int64_t kernel_dim = C / conv_attrs_.group * kernel_size; + const SafeInt X_offset = SafeInt(C) / conv_attrs_.group * input_image_size; + const SafeInt Y_offset = SafeInt(Y->Shape().Size()) / Y->Shape()[0] / conv_attrs_.group; + const SafeInt W_offset = SafeInt(W->Shape().Size()) / conv_attrs_.group; + const SafeInt kernel_dim = SafeInt(C) / conv_attrs_.group * kernel_size; const int64_t col_buffer_size = kernel_dim * output_image_size; - auto* col_data = alloc->Alloc(sizeof(float) * SafeInt(col_buffer_size)); - BufferUniquePtr col_buffer(col_data, BufferDeleter(std::move(alloc))); - auto* col_buffer_data = static_cast(col_buffer.get()); - + auto col_data = IAllocator::MakeUniquePtr(alloc, narrow(col_buffer_size)); + auto w_data = W->DataAsSpan(); for (int image_id = 0; image_id < N; ++image_id) { for (int group_id = 0; group_id < conv_attrs_.group; ++group_id) { math::Im2col()( - Xdata + group_id * X_offset, + &Xdata[group_id * X_offset], input_shape.GetDims().data(), output_shape.GetDims().data(), kernel_dim, @@ -269,8 +267,8 @@ Status Conv::Compute(OpKernelContext* context) const { strides.data(), dilations.data(), pads.data(), - static_cast(kernel_shape.size()), - col_buffer_data); + narrow(kernel_shape.size()), + col_data.get()); math::Gemm( CblasNoTrans, @@ -279,17 +277,17 @@ Status Conv::Compute(OpKernelContext* context) const { narrow(output_image_size), narrow(kernel_dim), 1, - W->Data() + group_id * W_offset, - col_buffer_data, + &w_data[group_id * W_offset], + col_data.get(), Beta, - Ydata + group_id * Y_offset, + &Ydata[group_id * Y_offset], thread_pool); } - MlasActivation(&activation_, Ydata, Bdata, narrow(M), narrow(output_image_size), narrow(output_image_size)); + MlasActivation(&activation_, Ydata.data(), Bdata, narrow(M), narrow(output_image_size), narrow(output_image_size)); - Xdata += X_offset * conv_attrs_.group; - Ydata += Y_offset * conv_attrs_.group; + Xdata = Xdata.subspan(X_offset * conv_attrs_.group); + Ydata = Ydata.subspan(Y_offset * conv_attrs_.group); } }