experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Merge pull request #42849 from Intel-tensorflow:dnn0x_clean_quantize

Browse Source 
PiperOrigin-RevId: 331178405
Change-Id: Id163829724709e5e9aa6eead0625287b40079860
This commit is contained in:
TensorFlower Gardener 2020-09-11 10:49:31 -07:00
commit 7d35f9e0ee
2 changed files with 30 additions and 71 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

47
tensorflow/core/kernels/mkl/mkl_dequantize_op.cc
View File

@ -26,7 +26,6 @@ limitations under the License.
#include "tensorflow/core/kernels/meta_support.h" #include "tensorflow/core/kernels/meta_support.h"
#include "tensorflow/core/kernels/quantization_utils.h" #include "tensorflow/core/kernels/quantization_utils.h"
#include "tensorflow/core/lib/core/errors.h" #include "tensorflow/core/lib/core/errors.h"
#include "tensorflow/core/util/mkl_types.h"
#include "tensorflow/core/util/mkl_util.h" #include "tensorflow/core/util/mkl_util.h"
using mkldnn::primitive_attr; using mkldnn::primitive_attr;
@ -51,7 +50,7 @@ class MklDequantizeOp : public OpKernel {
void Compute(OpKernelContext* ctx) override { void Compute(OpKernelContext* ctx) override {
try { try {
// Using CPU device // Using CPU device
auto cpu_engine = engine(ENGINE_CPU, 0); auto cpu_engine = engine(engine::kind::cpu, 0);
// Get the inputs // Get the inputs
const Tensor& src_tensor = MklGetInput(ctx, kSrcIndex); const Tensor& src_tensor = MklGetInput(ctx, kSrcIndex);
@ -82,10 +81,10 @@ class MklDequantizeOp : public OpKernel {
// construct input TF layout. For TF layout, although input shape // construct input TF layout. For TF layout, although input shape
// (src_dims) required is in MKL-DNN order, the layout is Tensorflow's // (src_dims) required is in MKL-DNN order, the layout is Tensorflow's
// layout // layout
auto src_md = auto src_md = src_mkl_shape.IsMklTensor()
src_mkl_shape.IsMklTensor() ? src_mkl_shape.GetMklLayout()
? src_mkl_shape.GetMklLayout() : memory::desc(src_dims, MklDnnType<T>(),
: memory::desc(src_dims, MklDnnType<T>(), MEMORY_FORMAT::nhwc); memory::format_tag::nhwc);
src.SetUsrMem(src_md, &src_tensor); src.SetUsrMem(src_md, &src_tensor);
src.SetUsrMemDataHandle(&src_tensor, reorder_stream); src.SetUsrMemDataHandle(&src_tensor, reorder_stream);
@ -93,14 +92,6 @@ class MklDequantizeOp : public OpKernel {
Tensor* output_tensor = nullptr; Tensor* output_tensor = nullptr;
MklDnnShape output_mkl_shape; MklDnnShape output_mkl_shape;
TensorShape output_tf_shape; TensorShape output_tf_shape;
#ifndef ENABLE_MKLDNN_V1
memory::desc dst_md =
src_mkl_shape.IsMklTensor()
? memory::desc(src_dims, MklDnnType<float>(),
static_cast<MEMORY_FORMAT>(src_md.data.format))
: memory::desc(src_dims, MklDnnType<float>(),
MEMORY_FORMAT::nhwc);
#else
memory::desc dst_md = memory::desc(); memory::desc dst_md = memory::desc();
if (src_mkl_shape.IsMklTensor()) { if (src_mkl_shape.IsMklTensor()) {
dst_md = memory::desc(src_mkl_shape.GetMklLayout().data); dst_md = memory::desc(src_mkl_shape.GetMklLayout().data);
@ -108,10 +99,9 @@ class MklDequantizeOp : public OpKernel {
// same .data field but different type. // same .data field but different type.
dst_md.data.data_type = memory::convert_to_c(MklDnnType<float>()); dst_md.data.data_type = memory::convert_to_c(MklDnnType<float>());
} else { } else {
dst_md = dst_md = memory::desc(src_dims, MklDnnType<float>(),
memory::desc(src_dims, MklDnnType<float>(), MEMORY_FORMAT::nhwc); memory::format_tag::nhwc);
} }
#endif // !ENABLE_MKLDNN_V1
// If input is MKL shape, output is also MKL shape. // If input is MKL shape, output is also MKL shape.
// If input is TF shape, output is also TF shape. // If input is TF shape, output is also TF shape.
@ -122,8 +112,7 @@ class MklDequantizeOp : public OpKernel {
output_mkl_shape.SetTfLayout(src_mkl_shape.GetDimension(), output_mkl_shape.SetTfLayout(src_mkl_shape.GetDimension(),
src_mkl_shape.GetSizesAsMklDnnDims(), src_mkl_shape.GetSizesAsMklDnnDims(),
src_mkl_shape.GetTfDataFormat()); src_mkl_shape.GetTfDataFormat());
output_tf_shape.AddDim(GET_MEMORY_SIZE_FROM_MD(dst_md, cpu_engine) / output_tf_shape.AddDim(dst_md.get_size() / sizeof(float));
sizeof(float));
} else { } else {
output_mkl_shape.SetMklTensor(false); output_mkl_shape.SetMklTensor(false);
output_tf_shape = MklDnnDimsToTFShape(output_dims); output_tf_shape = MklDnnDimsToTFShape(output_dims);
@ -155,29 +144,17 @@ class MklDequantizeOp : public OpKernel {
scales.push_back(scale_factor); scales.push_back(scale_factor);
primitive_attr attr; primitive_attr attr;
attr.set_output_scales(0, scales); attr.set_output_scales(0, scales);
#ifndef ENABLE_MKLDNN_V1
// MKL-DNN 1.0 does not provide set_int_output_round_mode() API.
// Also it does not define round_nearest (enum).
attr.set_int_output_round_mode(mkldnn::round_mode::round_nearest);
#endif // !ENABLE_MKLDNN_V1
std::vector<primitive> net; std::vector<primitive> net;
// Create reorder primitive and then execute. // Create reorder primitive and then execute.
auto reorder_pd = REORDER_PD_CONSTRUCTOR_WITH_ATTR( auto reorder_pd =
GET_MEMORY_PRIMITIVE_DESC_FROM_MEM_PTR(src.GetUsrMem()), ReorderPd(cpu_engine, src.GetUsrMem()->get_desc(), cpu_engine,
GET_MEMORY_PRIMITIVE_DESC_FROM_MEM_PTR(dst.GetUsrMem()), cpu_engine, dst.GetUsrMem()->get_desc(), attr);
attr);
#ifdef ENABLE_MKLDNN_V1
net.push_back(reorder(reorder_pd)); net.push_back(reorder(reorder_pd));
std::vector<std::unordered_map<int, memory>> reorder_net_args; std::vector<std::unordered_map<int, memory>> reorder_net_args;
reorder_net_args.push_back({{MKLDNN_ARG_FROM, *src.GetUsrMem()}, reorder_net_args.push_back({{MKLDNN_ARG_FROM, *src.GetUsrMem()},
{ MKLDNN_ARG_TO, {MKLDNN_ARG_TO, *dst.GetUsrMem()}});
*dst.GetUsrMem() }});
execute_primitives(net, reorder_stream, reorder_net_args); execute_primitives(net, reorder_stream, reorder_net_args);
#else
net.push_back(reorder(reorder_pd, *src.GetUsrMem(), *dst.GetUsrMem()));
reorder_stream->submit(net);
#endif // ENABLE_MKLDNN_V1
} catch (mkldnn::error& e) { } catch (mkldnn::error& e) {
string error_msg = "Status: " + std::to_string(e.status) + string error_msg = "Status: " + std::to_string(e.status) +
", message: " + string(e.message) + ", in file " + ", message: " + string(e.message) + ", in file " +

54
tensorflow/core/kernels/mkl/mkl_quantize_op.cc
View File

@ -25,7 +25,6 @@ limitations under the License.
#include "tensorflow/core/graph/mkl_graph_util.h" #include "tensorflow/core/graph/mkl_graph_util.h"
#include "tensorflow/core/lib/core/errors.h" #include "tensorflow/core/lib/core/errors.h"
#include "tensorflow/core/platform/logging.h" #include "tensorflow/core/platform/logging.h"
#include "tensorflow/core/util/mkl_types.h"
#include "tensorflow/core/util/mkl_util.h" #include "tensorflow/core/util/mkl_util.h"
using mkldnn::primitive_attr; using mkldnn::primitive_attr;
@ -77,7 +76,7 @@ class MklReorderWithScalePrimitive : public MklPrimitive {
public: public:
explicit MklReorderWithScalePrimitive( explicit MklReorderWithScalePrimitive(
const MklReorderWithScaleFwdParams& fwdParams) const MklReorderWithScaleFwdParams& fwdParams)
: MklPrimitive(engine(ENGINE_CPU, 0)) { : MklPrimitive(engine(engine::kind::cpu, 0)) {
// Create reorder primitive // Create reorder primitive
Setup(fwdParams); Setup(fwdParams);
} }
@ -95,11 +94,7 @@ class MklReorderWithScalePrimitive : public MklPrimitive {
context_.src_mem->set_data_handle(src_data); context_.src_mem->set_data_handle(src_data);
context_.dst_mem->set_data_handle(dst_data); context_.dst_mem->set_data_handle(dst_data);
#endif // ENABLE_MKLDNN_THREADPOOL #endif // ENABLE_MKLDNN_THREADPOOL
#ifndef ENABLE_MKLDNN_V1
reorder_stream->submit(context_.net);
#else
context_.reorder_prim->execute(*reorder_stream, context_.prim_args); context_.reorder_prim->execute(*reorder_stream, context_.prim_args);
#endif // !ENABLE_MKLDNN_V1
// After execution, set data handle back. // After execution, set data handle back.
context_.src_mem->set_data_handle(DummyData); context_.src_mem->set_data_handle(DummyData);
context_.dst_mem->set_data_handle(DummyData); context_.dst_mem->set_data_handle(DummyData);
@ -119,11 +114,7 @@ class MklReorderWithScalePrimitive : public MklPrimitive {
// Stream and primitive vector // Stream and primitive vector
std::shared_ptr<mkldnn::stream> reorder_stream; std::shared_ptr<mkldnn::stream> reorder_stream;
#ifndef ENABLE_MKLDNN_V1
std::vector<mkldnn::primitive> net;
#else
std::unordered_map<int, mkldnn::memory> prim_args; std::unordered_map<int, mkldnn::memory> prim_args;
#endif // !ENABLE_MKLDNN_V1
ReorderContext() ReorderContext()
: src_mem(nullptr), : src_mem(nullptr),
@ -135,10 +126,10 @@ class MklReorderWithScalePrimitive : public MklPrimitive {
// Reorder primitive setup // Reorder primitive setup
void Setup(const MklReorderWithScaleFwdParams& fwdParams) { void Setup(const MklReorderWithScaleFwdParams& fwdParams) {
// Create memory descriptors for reorder data with specified format // Create memory descriptors for reorder data with specified format
context_.src_mem.reset(new MEMORY_CONSTRUCTOR_USING_MD( context_.src_mem.reset(
fwdParams.src_md, cpu_engine_, DummyData)); new memory(fwdParams.src_md, cpu_engine_, DummyData));
context_.dst_mem.reset(new MEMORY_CONSTRUCTOR_USING_MD( context_.dst_mem.reset(
fwdParams.dst_md, cpu_engine_, DummyData)); new memory(fwdParams.dst_md, cpu_engine_, DummyData));
// Check if there is any fusion as post-ops // Check if there is any fusion as post-ops
auto const& post_op_params = fwdParams.post_op_params; auto const& post_op_params = fwdParams.post_op_params;
@ -150,21 +141,14 @@ class MklReorderWithScalePrimitive : public MklPrimitive {
scales.push_back(post_op_params.param[0]); scales.push_back(post_op_params.param[0]);
post_ops_attr.set_output_scales(0, scales); post_ops_attr.set_output_scales(0, scales);
context_.reorder_pd.reset(new REORDER_PD_CONSTRUCTOR_WITH_ATTR( context_.reorder_pd.reset(
GET_MEMORY_PRIMITIVE_DESC_FROM_MEM_PTR(context_.src_mem), new ReorderPd(cpu_engine_, context_.src_mem->get_desc(), cpu_engine_,
GET_MEMORY_PRIMITIVE_DESC_FROM_MEM_PTR(context_.dst_mem), cpu_engine_, context_.dst_mem->get_desc(), post_ops_attr));
post_ops_attr));
// Create reorder primitive // Create reorder primitive
#ifndef ENABLE_MKLDNN_V1
context_.reorder_prim.reset(new reorder(
*context_.reorder_pd, *context_.src_mem, *context_.dst_mem));
context_.net.push_back(*context_.reorder_prim);
#else
context_.reorder_prim.reset(new reorder(*context_.reorder_pd)); context_.reorder_prim.reset(new reorder(*context_.reorder_pd));
context_.prim_args.insert({MKLDNN_ARG_FROM, *context_.src_mem}); context_.prim_args.insert({MKLDNN_ARG_FROM, *context_.src_mem});
context_.prim_args.insert({MKLDNN_ARG_TO, *context_.dst_mem}); context_.prim_args.insert({MKLDNN_ARG_TO, *context_.dst_mem});
#endif // !ENABLE_MKLDNN_V1
} }
}; };
@ -278,7 +262,7 @@ class MklQuantizeV2Op : public OpKernel {
"Scalar calculation in MKL is supported only for" "Scalar calculation in MKL is supported only for"
"MIN_FIRST mode for now.")); "MIN_FIRST mode for now."));
auto cpu_engine = engine(ENGINE_CPU, 0); auto cpu_engine = engine(engine::kind::cpu, 0);
const Tensor& input = ctx->input(0); const Tensor& input = ctx->input(0);
const unsigned int src_idx = 0; const unsigned int src_idx = 0;
const Tensor& src_tensor = MklGetInput(ctx, src_idx); const Tensor& src_tensor = MklGetInput(ctx, src_idx);
@ -344,7 +328,7 @@ class MklQuantizeV2Op : public OpKernel {
max_range = std::max(input_max_range, min_range + epsilon); max_range = std::max(input_max_range, min_range + epsilon);
// Clamping the max_range to zero since max_range can also be negative. // Clamping the max_range to zero since max_range can also be negative.
max_range = std::max(0.0f, max_range); max_range = std::max(0.0f, max_range);
auto cpu_engine = engine(ENGINE_CPU, 0); auto cpu_engine = engine(engine::kind::cpu, 0);
const Tensor& src_tensor = MklGetInput(ctx, src_idx); const Tensor& src_tensor = MklGetInput(ctx, src_idx);
MklDnnShape src_mkl_shape; MklDnnShape src_mkl_shape;
GetMklShape(ctx, src_idx, &src_mkl_shape); GetMklShape(ctx, src_idx, &src_mkl_shape);
@ -355,25 +339,25 @@ class MklQuantizeV2Op : public OpKernel {
: TFShapeToMklDnnDims(src_tensor.shape()); : TFShapeToMklDnnDims(src_tensor.shape());
auto output_dims = src_dims; auto output_dims = src_dims;
// Set the dst layout to be the best mkl layout based on dims and type. // Set the dst layout to be the best mkl layout based on dims and type.
MEMORY_FORMAT dst_layout_type; memory::format_tag dst_layout_type;
switch (src_tf_shape.dims()) { switch (src_tf_shape.dims()) {
case 0: case 0:
ComputeScalar(ctx, min_range, max_range); ComputeScalar(ctx, min_range, max_range);
return; return;
case 1: case 1:
dst_layout_type = MEMORY_FORMAT::x; dst_layout_type = memory::format_tag::x;
break; break;
case 2: case 2:
dst_layout_type = MEMORY_FORMAT::nc; dst_layout_type = memory::format_tag::nc;
break; break;
case 3: case 3:
dst_layout_type = MEMORY_FORMAT::tnc; dst_layout_type = memory::format_tag::tnc;
break; break;
case 4: case 4:
dst_layout_type = MEMORY_FORMAT::nhwc; dst_layout_type = memory::format_tag::nhwc;
break; break;
case 5: case 5:
dst_layout_type = MEMORY_FORMAT::ndhwc; dst_layout_type = memory::format_tag::ndhwc;
break; break;
default: default:
OP_REQUIRES_OK(ctx, OP_REQUIRES_OK(ctx,
@ -417,9 +401,7 @@ class MklQuantizeV2Op : public OpKernel {
memory::desc dst_md = memory::desc dst_md =
memory::desc(src_dims, MklDnnType<T>(), dst_layout_type); memory::desc(src_dims, MklDnnType<T>(), dst_layout_type);
#ifndef ENABLE_MKLDNN_V1
auto dst_pd = memory::primitive_desc(dst_md, cpu_engine);
#endif // !ENABLE_MKLDNN_V1
// Standard shape assignments for layout pass // Standard shape assignments for layout pass
MklDnnShape output_mkl_shape; MklDnnShape output_mkl_shape;
TensorShape output_tf_shape; TensorShape output_tf_shape;