Skip to content

Commit

Permalink
Merge pull request #59881 from tensorflow/venkat-patch-3
Browse files Browse the repository at this point in the history 
r2.12 cherry-pick: [INTEL oneDNN][Bug Fix] Add index-validity check for min-max tensors for quantized oneDNN ops and related tests #59581
  • Loading branch information
@nitins17
nitins17 committed Mar 2, 2023
2 parents c6dfc24 + ab9d5a3 commit 5966cce
Show file tree
Hide file tree
Showing 14 changed files with 209 additions and 148 deletions.
23 changes: 12 additions & 11 deletions tensorflow/core/kernels/mkl/mkl_avgpooling_op.cc
View file
Expand Up @@ -161,13 +161,13 @@ class MklAvgPoolingOp : public MklPoolingForwardOpBase<T> {
output_min_mkl_shape, this->native_format_);
AllocateOutputSetMklShape(context, 2, &output_max, {},
output_max_mkl_shape, this->native_format_);
output_min->flat<float>()(0) = min_input;
output_max->flat<float>()(0) = max_input;
output_min->scalar<float>()() = min_input;
output_max->scalar<float>()() = max_input;
}
} catch (dnnl::error& e) {
string error_msg = "Status: " + std::to_string(e.status) +
", message: " + string(e.message) + ", in file " +
string(__FILE__) + ":" + std::to_string(__LINE__);
string error_msg = "Status: " + std::to_string(e.status) + ", message: " +
string(e.message) + ", in file " + string(__FILE__) +
":" + std::to_string(__LINE__);
OP_REQUIRES_OK(
context,
errors::Aborted("Operation received an exception:", error_msg));
Expand Down Expand Up @@ -230,9 +230,10 @@ class MklAvgPoolingGradOp : public MklPoolingBackwardOpBase<T> {
memory::dims orig_input_dims_mkl_order =
orig_input_mkl_shape.IsMklTensor()
? orig_input_mkl_shape.GetSizesAsMklDnnDims()
: is_pool2d
? TFShapeToMklDnnDimsInNCHW(output_shape, this->data_format_tf_)
: TFShapeToMklDnnDimsInNCDHW(output_shape, this->data_format_tf_);
: is_pool2d ? TFShapeToMklDnnDimsInNCHW(output_shape,
this->data_format_tf_)
: TFShapeToMklDnnDimsInNCDHW(output_shape,
this->data_format_tf_);

memory::dims diff_dst_dims =
grad_mkl_shape.IsMklTensor()
Expand Down Expand Up @@ -297,9 +298,9 @@ class MklAvgPoolingGradOp : public MklPoolingBackwardOpBase<T> {
pooling_bwd->Execute(diff_dst_data, diff_src_data, nullptr,
bwd_cpu_stream);
} catch (dnnl::error& e) {
string error_msg = "Status: " + std::to_string(e.status) +
", message: " + string(e.message) + ", in file " +
string(__FILE__) + ":" + std::to_string(__LINE__);
string error_msg = "Status: " + std::to_string(e.status) + ", message: " +
string(e.message) + ", in file " + string(__FILE__) +
":" + std::to_string(__LINE__);
OP_REQUIRES_OK(context, errors::Aborted("Compute received an exception:",
error_msg));
}
Expand Down
44 changes: 28 additions & 16 deletions tensorflow/core/kernels/mkl/mkl_concat_op.cc
View file
Expand Up @@ -17,7 +17,6 @@ limitations under the License.
#include <unordered_map>
#include <vector>

#include "third_party/eigen3/unsupported/Eigen/CXX11/Tensor"
#include "dnnl.hpp"
#include "tensorflow/core/framework/bounds_check.h"
#include "tensorflow/core/framework/op_kernel.h"
Expand All @@ -32,6 +31,7 @@ limitations under the License.
#include "tensorflow/core/lib/core/status.h"
#include "tensorflow/core/platform/types.h"
#include "tensorflow/core/util/mkl_util.h"
#include "third_party/eigen3/unsupported/Eigen/CXX11/Tensor"
#ifdef DNNL_AARCH64_USE_ACL
#include "tensorflow/core/platform/mutex.h"
#endif
Expand Down Expand Up @@ -108,8 +108,8 @@ class EigenConcatBaseOp : public OpKernel {
float overall_min = std::numeric_limits<float>::max();
float overall_max = std::numeric_limits<float>::lowest();
for (int i = 0; i < N; ++i) {
const float input_min = input_mins[i].flat<float>()(0);
const float input_max = input_maxes[i].flat<float>()(0);
const float input_min = input_mins[i].scalar<float>()();
const float input_max = input_maxes[i].scalar<float>()();
input_mins_and_maxes->emplace_back(input_min, input_max);
overall_min = std::min(overall_min, input_min);
overall_max = std::max(overall_max, input_max);
Expand Down Expand Up @@ -187,13 +187,12 @@ class EigenConcatBaseOp : public OpKernel {
const auto in = values[i];
const bool in_is_scalar = TensorShapeUtils::IsScalar(input_shapes[i]);
OP_REQUIRES(
c,
(input_shapes[i].dims() == input_dims) ||
(input_is_scalar && in_is_scalar),
c, (input_shapes[i].dims() == input_dims) ||
(input_is_scalar && in_is_scalar),
errors::InvalidArgument(
"ConcatOp : Ranks of all input tensors should match: shape[0] = ",
input_shape.DebugString(), " vs. shape[", i,
"] = ", input_shapes[i].DebugString()));
input_shape.DebugString(), " vs. shape[", i, "] = ",
input_shapes[i].DebugString()));
if (in.NumElements() > 0) {
int64 inputs_flat_dim1 = in.NumElements() / inputs_flat_dim0;
inputs_flat.emplace_back(new typename TTypes<T, 2>::ConstMatrix(
Expand Down Expand Up @@ -227,11 +226,11 @@ class EigenConcatBaseOp : public OpKernel {
if (quantized_input) {
Tensor* output_min_tensor = nullptr;
OP_REQUIRES_OK(c, c->allocate_output(1, {}, &output_min_tensor));
output_min_tensor->flat<float>()(0) = output_min;
output_min_tensor->scalar<float>()() = output_min;

Tensor* output_max_tensor = nullptr;
OP_REQUIRES_OK(c, c->allocate_output(2, {}, &output_max_tensor));
output_max_tensor->flat<float>()(0) = output_max;
output_max_tensor->scalar<float>()() = output_max;
}
}
};
Expand Down Expand Up @@ -568,12 +567,25 @@ class MklConcatOp : public OpKernel {
errors::InvalidArgument(
"QuantizedConcatOp : Expected maxes input list length ",
input_maxes.size(), " to equal values length ", N));
float input_min = input_mins[0].flat<float>()(0);
float input_max = input_maxes[0].flat<float>()(0);

for (int i = 0; i < N; i++) {
OP_REQUIRES(context,
TensorShapeUtils::IsScalar(input_mins[i].shape()),
errors::InvalidArgument("`input_mins[", i,
"]` must be rank 0 but is rank ",
input_mins[i].dims()));
OP_REQUIRES(context,
TensorShapeUtils::IsScalar(input_maxes[i].shape()),
errors::InvalidArgument("`input_maxes[", i,
"]` must be rank 0 but is rank ",
input_maxes[i].dims()));
}
float input_min = input_mins[0].scalar<float>()();
float input_max = input_maxes[0].scalar<float>()();
const float eps = 1.0e-6;
for (int i = 1; i < N; ++i) {
float min = input_mins[i].flat<float>()(0);
float max = input_maxes[i].flat<float>()(0);
float min = input_mins[i].scalar<float>()();
float max = input_maxes[i].scalar<float>()();

if (fabs(input_min - min) > eps || fabs(input_max - max) > eps) {
invoke_eigen = true;
Expand Down Expand Up @@ -805,8 +817,8 @@ class MklConcatOp : public OpKernel {
output_max_mkl_shape, native_format);
// All input tensors should have the same range, just use the
// first one
output_min->flat<float>()(0) = input_mins[0].flat<float>()(0);
output_max->flat<float>()(0) = input_maxes[0].flat<float>()(0);
output_min->scalar<float>()() = input_mins[0].scalar<float>()();
output_max->scalar<float>()() = input_maxes[0].scalar<float>()();
}
} else {
MklDnnShape dnn_shape_dst;
Expand Down
77 changes: 58 additions & 19 deletions tensorflow/core/kernels/mkl/mkl_conv_ops.cc
View file
Expand Up @@ -1845,11 +1845,10 @@ class MklQuantizedConvOp
/*native_format*/ true>::Compute(context);

// Compute additional outputs: min/max scalars.

const float min_input =
context->input(min_input_idx_).template flat<float>()(0);
context->input(min_input_idx_).template scalar<float>()();
const float max_input =
context->input(max_input_idx_).template flat<float>()(0);
context->input(max_input_idx_).template scalar<float>()();

Tensor* output_min = nullptr;
Tensor* output_max = nullptr;
Expand All @@ -1859,18 +1858,18 @@ class MklQuantizedConvOp
OP_REQUIRES_OK(context, context->allocate_output(2, {}, &output_max));
// This is the case the convolution and requantization are fused.
output_min->flat<float>()(0) =
context->input(min_freezed_output_idx_).template flat<float>()(0);
context->input(min_freezed_output_idx_).template scalar<float>()();
output_max->flat<float>()(0) =
context->input(max_freezed_output_idx_).template flat<float>()(0);
context->input(max_freezed_output_idx_).template scalar<float>()();
} else {
const Tensor& min_filter = context->input(min_filter_idx_);
const Tensor& max_filter = context->input(max_filter_idx_);
if (min_filter.dims() == 0) {
float min_output_value;
float max_output_value;
MklQuantizationRangeForMultiplication<Tinput, qint8, qint32>(
min_input, max_input, min_filter.flat<float>()(0),
max_filter.flat<float>()(0), &min_output_value, &max_output_value);
min_input, max_input, min_filter.scalar<float>()(),
max_filter.scalar<float>()(), &min_output_value, &max_output_value);
OP_REQUIRES_OK(context, context->allocate_output(1, {}, &output_min));
OP_REQUIRES_OK(context, context->allocate_output(2, {}, &output_max));
output_min->flat<float>()(0) = min_output_value;
Expand Down Expand Up @@ -1903,18 +1902,25 @@ class MklQuantizedConvOp
if (std::is_same<Toutput, quint8>::value ||
std::is_same<Toutput, qint8>::value) {
const float min_input =
context->input(min_input_idx_).template flat<float>()(0);
context->input(min_input_idx_).template scalar<float>()();
const float max_input =
context->input(max_input_idx_).template flat<float>()(0);
context->input(max_input_idx_).template scalar<float>()();
const Tensor& min_filter_vector = context->input(min_filter_idx_);
const Tensor& max_filter_vector = context->input(max_filter_idx_);
OP_REQUIRES(
context,
((min_filter_vector.NumElements() > 0) &&
(max_filter_vector.NumElements() > 0) &&
(min_filter_vector.shape() == max_filter_vector.shape())),
errors::InvalidArgument("`min_ and max_filter` must have same"
"shape and contain at least one element."));

// min_freezed_output and max_freezed_output are the actual range
// for the output.
const float min_freezed_output =
context->input(min_freezed_output_idx_).template flat<float>()(0);
context->input(min_freezed_output_idx_).template scalar<float>()();
const float max_freezed_output =
context->input(max_freezed_output_idx_).template flat<float>()(0);
context->input(max_freezed_output_idx_).template scalar<float>()();

float int_output_limit =
std::is_same<Toutput, quint8>::value ? 255.0f : 127.0f;
Expand Down Expand Up @@ -1959,14 +1965,47 @@ class MklQuantizedConvOp
bool summand_condition =
(summand_dt == DT_QINT8) || (summand_dt == DT_QUINT8);
DCHECK((summand_condition));

const Tensor& min_freezed_output_tensor =
context->input(min_freezed_output_idx_);
const Tensor& max_freezed_output_tensor =
context->input(max_freezed_output_idx_);
OP_REQUIRES(
context,
TensorShapeUtils::IsScalar(min_freezed_output_tensor.shape()),
errors::InvalidArgument(
"`min_freezed_output` must be rank 0 but is rank ",
min_freezed_output_tensor.dims()));
OP_REQUIRES(
context,
TensorShapeUtils::IsScalar(max_freezed_output_tensor.shape()),
errors::InvalidArgument(
"`max_freezed_output` must be rank 0 but is rank ",
max_freezed_output_tensor.dims()));
const Tensor& min_freezed_summand_tensor =
context->input(min_summand_idx_);
const Tensor& max_freezed_summand_tensor =
context->input(max_summand_idx_);
OP_REQUIRES(
context,
TensorShapeUtils::IsScalar(min_freezed_summand_tensor.shape()),
errors::InvalidArgument(
"`min_freezed_summand` must be rank 0 but is rank ",
min_freezed_summand_tensor.dims()));
OP_REQUIRES(
context,
TensorShapeUtils::IsScalar(max_freezed_summand_tensor.shape()),
errors::InvalidArgument(
"`max_freezed_summand` must be rank 0 but is rank ",
max_freezed_summand_tensor.dims()));
const float min_freezed_output =
context->input(min_freezed_output_idx_).template flat<float>()(0);
min_freezed_output_tensor.template scalar<float>()();
const float max_freezed_output =
context->input(max_freezed_output_idx_).template flat<float>()(0);
max_freezed_output_tensor.template scalar<float>()();
const float min_freezed_summand =
context->input(min_summand_idx_).template flat<float>()(0);
min_freezed_summand_tensor.template scalar<float>()();
const float max_freezed_summand =
context->input(max_summand_idx_).template flat<float>()(0);
max_freezed_summand_tensor.template scalar<float>()();

float output_range = std::max(std::abs(min_freezed_output),
std::abs(max_freezed_output));
Expand Down Expand Up @@ -2052,9 +2091,9 @@ class MklQuantizedConvOp
"Current fusion requires summand to be float"));
// We need to compute scale for the summand
const float min_input =
context->input(min_input_idx_).template flat<float>()(0);
context->input(min_input_idx_).template scalar<float>()();
const float max_input =
context->input(max_input_idx_).template flat<float>()(0);
context->input(max_input_idx_).template scalar<float>()();
const Tensor& min_filter_vector = context->input(min_filter_idx_);
const Tensor& max_filter_vector = context->input(max_filter_idx_);
const float* min_filter = min_filter_vector.flat<float>().data();
Expand Down Expand Up @@ -2105,9 +2144,9 @@ class MklQuantizedConvOp
}

const float min_input =
context->input(min_input_idx_).template flat<float>()(0);
context->input(min_input_idx_).template scalar<float>()();
const float max_input =
context->input(max_input_idx_).template flat<float>()(0);
context->input(max_input_idx_).template scalar<float>()();
const Tensor& min_filter_vector = context->input(min_filter_idx_);
const Tensor& max_filter_vector = context->input(max_filter_idx_);
const float* min_filter = min_filter_vector.flat<float>().data();
Expand Down
4 changes: 2 additions & 2 deletions tensorflow/core/kernels/mkl/mkl_dequantize_op.cc
View file
Expand Up @@ -56,8 +56,8 @@ class MklDequantizeOp : public OpKernel {

// Get the inputs
const Tensor& src_tensor = ctx->input(kSrcIndex);
const float min_range = ctx->input(kMinIndex).template flat<float>()(0);
const float max_range = ctx->input(kMaxIndex).template flat<float>()(0);
const float min_range = ctx->input(kMinIndex).template scalar<float>()();
const float max_range = ctx->input(kMaxIndex).template scalar<float>()();

// Get MklShape
auto src_tf_shape = src_tensor.shape();
Expand Down
8 changes: 4 additions & 4 deletions tensorflow/core/kernels/mkl/mkl_dequantize_op_test.cc
View file
Expand Up @@ -42,9 +42,9 @@ TEST_F(MklDequantizeOpTest, small) {
AddInputFromArray<quint8>(TensorShape({1, 2, 2, 2}),
{0, 10, 50, 40, 25, 115, 190, 255});
// min_range = 0
AddInputFromArray<float>(TensorShape({1}), {0});
AddInputFromArray<float>(TensorShape({}), {0});
// max_range = 200
AddInputFromArray<float>(TensorShape({1}), {200.0f});
AddInputFromArray<float>(TensorShape({}), {200.0f});
TF_ASSERT_OK(RunOpKernel());
Tensor expected(allocator(), DT_FLOAT, TensorShape({1, 2, 2, 2}));
test::FillValues<float>(&expected,
Expand Down Expand Up @@ -84,9 +84,9 @@ TEST_F(MklDequantizeOpTest, MKLInput) {
AddInputFromArray<quint8>(TensorShape({1, 2, 2, 2}),
{0, 10, 50, 40, 25, 115, 190, 255});
// min_range = 0
AddInputFromArray<float>(TensorShape({1}), {0});
AddInputFromArray<float>(TensorShape({}), {0});
// max_range = 200
AddInputFromArray<float>(TensorShape({1}), {200.0f});
AddInputFromArray<float>(TensorShape({}), {200.0f});
TF_ASSERT_OK(RunOpKernel());
Tensor expected(allocator(), DT_FLOAT, TensorShape({1, 2, 2, 2}));
test::FillValues<float>(&expected,
Expand Down
16 changes: 8 additions & 8 deletions tensorflow/core/kernels/mkl/mkl_maxpooling_op.cc
View file
Expand Up @@ -180,8 +180,8 @@ class MklMaxPoolingOp : public MklPoolingForwardOpBase<T> {
output_min_mkl_shape, this->native_format_);
AllocateOutputSetMklShape(context, 2, &output_max, {},
output_max_mkl_shape, this->native_format_);
output_min->flat<float>()(0) = min_input;
output_max->flat<float>()(0) = max_input;
output_min->scalar<float>()() = min_input;
output_max->scalar<float>()() = max_input;
} else {
MklDnnData<uint8> dnn_data_wksp(&cpu_engine_);
AllocateWorkspaceTensor(context, *(pooling_fwd->GetPoolingFwdPd()),
Expand All @@ -193,9 +193,9 @@ class MklMaxPoolingOp : public MklPoolingForwardOpBase<T> {
pooling_fwd->Execute(src_data, dst_data, ws_data, fwd_cpu_stream);
}
} catch (dnnl::error& e) {
string error_msg = "Status: " + std::to_string(e.status) +
", message: " + string(e.message) + ", in file " +
string(__FILE__) + ":" + std::to_string(__LINE__);
string error_msg = "Status: " + std::to_string(e.status) + ", message: " +
string(e.message) + ", in file " + string(__FILE__) +
":" + std::to_string(__LINE__);
OP_REQUIRES_OK(context, errors::Aborted("Compute received an exception:",
error_msg));
}
Expand Down Expand Up @@ -356,9 +356,9 @@ class MklMaxPoolingGradOp : public MklPoolingBackwardOpBase<T> {
pooling_bwd->Execute(diff_dst_data, diff_src_data, ws_data,
bwd_cpu_stream);
} catch (dnnl::error& e) {
string error_msg = "Status:" + std::to_string(e.status) +
", message: " + string(e.message) + ". in file " +
string(__FILE__) + ":" + std::to_string(__LINE__);
string error_msg = "Status:" + std::to_string(e.status) + ", message: " +
string(e.message) + ". in file " + string(__FILE__) +
":" + std::to_string(__LINE__);
OP_REQUIRES_OK(context, errors::Aborted("Compute received an exception:",
error_msg));
}
Expand Down

0 comments on commit 5966cce

Please sign in to comment.