saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-06-26 03:00:54 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Add some safety check for conv op (#16839)

Browse source 
### 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).
This commit is contained in:
Changming Sun 2023-07-27 16:37:55 -07:00 committed by GitHub
parent e67547b978
commit 161a9d1d6d
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
5 changed files with 38 additions and 38 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
cmake/CMakeLists.txt
View file

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

10
onnxruntime/contrib_ops/cpu/nchwc_ops.cc
View file

@ -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<float>();
auto y_data = Y->MutableDataAsSpan<float>();
// 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<float>();
if (y_data != sum_data) {
memcpy(y_data, sum_data, SafeInt<size_t>(sum_shape.Size()) * sizeof(float));
auto sum_data = Sum->DataAsSpan<float>();
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<float>(),
W->Data<float>(),
B != nullptr ? B->Data<float>() : nullptr,
y_data,
y_data.data(),
&activation_,
Sum == nullptr,
context->GetOperatorThreadPool());

10
onnxruntime/core/providers/common.h
View file

@ -5,6 +5,7 @@
#include <cstdint>
#include <functional>
#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<int64_t> legacy_target_size = (SafeInt<int64_t>(in_dim) + stride - 1) / stride;
SafeInt<int64_t> 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<int64_t>(static_cast<double>(in_dim + pad_head + pad_tail - dkernel) / stride + 1);
const SafeInt<int64_t> dkernel = SafeInt<int64_t>(dilation) * (kernel - 1) + 1;
int64_t dkernel_value = SafeInt<int64_t>(in_dim) + pad_head + pad_tail - dkernel;
return static_cast<int64_t>(static_cast<double>(dkernel_value) / stride + 1);
}
inline Status ComputePadAndOutputShape(const int64_t in_dim,

1
onnxruntime/core/providers/coreml/model/model.mm
View file

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

54
onnxruntime/core/providers/cpu/nn/conv.cc
View file

@ -195,18 +195,18 @@ Status Conv<float>::Compute(OpKernelContext* context) const {
AllocatorPtr alloc;
ORT_RETURN_IF_ERROR(context->GetTempSpaceAllocator(&alloc));
const auto* Xdata = X->Data<float>();
auto Xdata = X->DataAsSpan<float>();
const auto* Bdata = B != nullptr ? B->Data<float>() : nullptr;
auto* Ydata = Y->MutableData<float>();
auto Ydata = Y->MutableDataAsSpan<float>();
// 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<float>();
if (Ydata != sum_data) {
memcpy(Ydata, sum_data, SafeInt<size_t>(sum_shape.Size()) * sizeof(float));
auto sum_data = Sum->DataAsSpan<float>();
if (Ydata.data() != sum_data.data()) {
gsl::copy(sum_data, Ydata);
}
Beta = 1.0f;
}
@ -218,16 +218,16 @@ Status Conv<float>::Compute(OpKernelContext* context) const {
size_t WorkingBufferSize;
MlasConvPrepare(&Parameters,
kernel_rank,
static_cast<size_t>(N),
static_cast<size_t>(conv_attrs_.group),
static_cast<size_t>(C / conv_attrs_.group),
narrow<size_t>(N),
narrow<size_t>(conv_attrs_.group),
narrow<size_t>(C / conv_attrs_.group),
input_shape.GetDims().data(),
kernel_shape.data(),
dilations.data(),
pads.data(),
strides.data(),
output_shape.GetDims().data(),
static_cast<size_t>(M / conv_attrs_.group),
narrow<size_t>(M / conv_attrs_.group),
&activation_,
&WorkingBufferSize,
Beta,
@ -238,30 +238,28 @@ Status Conv<float>::Compute(OpKernelContext* context) const {
BufferUniquePtr working_buffer(working_data, BufferDeleter(std::move(alloc)));
MlasConv(&Parameters,
Xdata,
Xdata.data(),
W->Data<float>(),
Bdata,
static_cast<float*>(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<int64_t> X_offset = SafeInt<int64_t>(C) / conv_attrs_.group * input_image_size;
const SafeInt<int64_t> Y_offset = SafeInt<int64_t>(Y->Shape().Size()) / Y->Shape()[0] / conv_attrs_.group;
const SafeInt<int64_t> W_offset = SafeInt<int64_t>(W->Shape().Size()) / conv_attrs_.group;
const SafeInt<int64_t> kernel_dim = SafeInt<int64_t>(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<size_t>(col_buffer_size));
BufferUniquePtr col_buffer(col_data, BufferDeleter(std::move(alloc)));
auto* col_buffer_data = static_cast<float*>(col_buffer.get());
auto col_data = IAllocator::MakeUniquePtr<float>(alloc, narrow<size_t>(col_buffer_size));
auto w_data = W->DataAsSpan<float>();
for (int image_id = 0; image_id < N; ++image_id) {
for (int group_id = 0; group_id < conv_attrs_.group; ++group_id) {
math::Im2col<float, StorageOrder::NCHW>()(
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<float>::Compute(OpKernelContext* context) const {
strides.data(),
dilations.data(),
pads.data(),
static_cast<int>(kernel_shape.size()),
col_buffer_data);
narrow<int>(kernel_shape.size()),
col_data.get());
math::Gemm<float>(
CblasNoTrans,
@ -279,17 +277,17 @@ Status Conv<float>::Compute(OpKernelContext* context) const {
narrow<ptrdiff_t>(output_image_size),
narrow<ptrdiff_t>(kernel_dim),
1,
W->Data<float>() + 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<size_t>(M), narrow<size_t>(output_image_size), narrow<size_t>(output_image_size));
MlasActivation(&activation_, Ydata.data(), Bdata, narrow<size_t>(M), narrow<size_t>(output_image_size), narrow<size_t>(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);
}
}