未验证 提交 a63bcf9a 编写于 作者: M Michał Gallus 提交者: GitHub

[DNNL][INT8][FP32] MatMul (#23395)

* Initial FP32 DNNL MatMul Implementation

* Implement int8 DNNL MatMul

* Unify in-kernel-naming, clean UTs

* MatmuL: Introduce op caching

* Final adjustments

test=develop

* Remove dy_graph disablement

test=develop

* Change dnnl header name to new one

test=develop

* Contrain multi head check to prevent fails

test=develop

* Resolve dnnl header problems on MAC CI

* Variable namings to kernel and skip_grad_ci added

test=develop

* Prevent MAC CI from failing

* Prevent windows build from failing

test=develop

* Modify UTs to conform to the rules

* Modify MatMul aux functions namings

test=develop
上级 a2315d37
...@@ -17,6 +17,9 @@ limitations under the License. */ ...@@ -17,6 +17,9 @@ limitations under the License. */
#include <vector> #include <vector>
#include "paddle/fluid/framework/op_registry.h" #include "paddle/fluid/framework/op_registry.h"
#include "paddle/fluid/operators/math/blas.h" #include "paddle/fluid/operators/math/blas.h"
#ifdef PADDLE_WITH_MKLDNN
#include "paddle/fluid/platform/mkldnn_helper.h"
#endif
namespace paddle { namespace paddle {
namespace operators { namespace operators {
...@@ -409,6 +412,21 @@ class MatMulOp : public framework::OperatorWithKernel { ...@@ -409,6 +412,21 @@ class MatMulOp : public framework::OperatorWithKernel {
context->SetOutputDim("Out", framework::make_ddim(dim_out)); context->SetOutputDim("Out", framework::make_ddim(dim_out));
context->ShareLoD("X", /*->*/ "Out"); context->ShareLoD("X", /*->*/ "Out");
} }
framework::OpKernelType GetExpectedKernelType(
const framework::ExecutionContext &ctx) const override {
auto input_data_type = OperatorWithKernel::IndicateVarDataType(ctx, "X");
#ifdef PADDLE_WITH_MKLDNN
using mkldnn::memory;
if (platform::CanMKLDNNBeUsed(ctx)) {
return framework::OpKernelType(input_data_type, ctx.GetPlace(),
framework::DataLayout::kMKLDNN,
framework::LibraryType::kMKLDNN);
}
#endif
return framework::OpKernelType(input_data_type, ctx.GetPlace());
}
}; };
class MatMulOpMaker : public framework::OpProtoAndCheckerMaker { class MatMulOpMaker : public framework::OpProtoAndCheckerMaker {
...@@ -426,6 +444,30 @@ class MatMulOpMaker : public framework::OpProtoAndCheckerMaker { ...@@ -426,6 +444,30 @@ class MatMulOpMaker : public framework::OpProtoAndCheckerMaker {
)DOC") )DOC")
.SetDefault(false); .SetDefault(false);
AddAttr<float>("alpha", "The scale of Out").SetDefault(1.0f); AddAttr<float>("alpha", "The scale of Out").SetDefault(1.0f);
AddAttr<bool>(
"use_mkldnn",
"(bool, default false) Indicates if MKL-DNN kernel will be used")
.SetDefault(false);
/* int8 parameters */
AddAttr<bool>("use_quantizer",
"(bool, default false) "
"Set to true for operators that should be quantized and use "
"int8 kernel. "
"Only used on CPU.")
.SetDefault(false);
AddAttr<float>("Scale_x",
"(float, default 1.0f), The quantize scale of X tensor")
.SetDefault(1.0f);
AddAttr<float>("Scale_y",
"(float, default 1.0f), The quantize scale of Y tensor")
.SetDefault(1.0f);
AddAttr<float>("Scale_out",
"(float, default 1.0f), The quantize scale of output data")
.SetDefault(1.0f);
AddAttr<bool>("force_fp32_output",
"(bool, default false) Force INT8 kernel output FP32, only "
"used in MKL-DNN INT8")
.SetDefault(false);
#if defined(PADDLE_WITH_MKLML) && !defined(PADDLE_WITH_CUDA) #if defined(PADDLE_WITH_MKLML) && !defined(PADDLE_WITH_CUDA)
AddAttr<int>("head_number", "The number of heads of the matrix") AddAttr<int>("head_number", "The number of heads of the matrix")
.SetDefault(1); .SetDefault(1);
......
/* Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License. */
#include "mkldnn.hpp"
#include "paddle/fluid/framework/op_registry.h"
#include "paddle/fluid/framework/tensor.h"
#include "paddle/fluid/operators/math/blas.h"
#include "paddle/fluid/platform/mkldnn_helper.h"
namespace paddle {
namespace operators {
using dnnl::memory;
using dnnl::primitive;
using platform::to_void_cast;
using framework::DataLayout;
using platform::GetMKLDNNFormat;
using platform::MKLDNNGetDataType;
using platform::MKLDNNDeviceContext;
using framework::ExecutionContext;
using Tensor = framework::Tensor;
// Get row matrix shape from a vector shape. If the rank of x_dim > 1, the
// original x_dim is returned.
static framework::DDim RowMatrixDimsFromVector(const framework::DDim& x_dim) {
return x_dim.size() > 1 ? x_dim : framework::make_ddim({1, x_dim[0]});
}
// Get column matrix shape from a vector shape. If the ran of y_dim > 1, the
// original y_dim is returned.
static framework::DDim ColumnMatrixDimsFromVector(
const framework::DDim& y_dim) {
return y_dim.size() > 1 ? y_dim : framework::make_ddim({y_dim[0], 1});
}
template <typename XT, typename YT, typename OT>
class MatMulFactory {
public:
void CreateAndExecute(const ExecutionContext& ctx) {
SetDNNLEngine(ctx);
if (IsInitialized()) {
UpdateDataPointers(ctx);
Execute();
SetOutputFormat(ctx);
return;
}
CreateMemories(ctx);
CreatePrimitive(ctx);
Execute();
SetOutputFormat(ctx);
SetInitialized();
}
private:
struct MatMulDims {
const memory::dim BS, M, N, K;
};
void SetDNNLEngine(const ExecutionContext& ctx) {
auto& dev_ctx =
ctx.template device_context<platform::MKLDNNDeviceContext>();
engine_ = dev_ctx.GetEngine();
}
template <typename T>
dnnl::memory CreateMemory(const memory::dims& dims,
const memory::dims& strides, const T* data) {
auto md = memory::desc(dims, MKLDNNGetDataType<T>(), strides);
return dnnl::memory(md, engine_, to_void_cast(data));
}
MatMulDims GetMatmulDims(const ExecutionContext& ctx) {
auto mat_dim_x = math::CreateMatrixDescriptor(
RowMatrixDimsFromVector(ctx.Input<Tensor>("X")->dims()), 0,
ctx.Attr<bool>("transpose_X"));
auto mat_dim_y = math::CreateMatrixDescriptor(
ColumnMatrixDimsFromVector(ctx.Input<Tensor>("Y")->dims()), 0,
ctx.Attr<bool>("transpose_Y"));
const auto x_bs = mat_dim_x.batch_size_;
const auto y_bs = mat_dim_y.batch_size_;
PADDLE_ENFORCE_EQ(x_bs > 0 && y_bs > 0 && x_bs != y_bs, false,
platform::errors::InvalidArgument(
"If batch sizes of X and Y are positive,"
"they have to be equal."));
// Store 1 if both batches are zero, otherwise save the nonzero batch
const memory::dim BS = x_bs || y_bs ? std::max(x_bs, y_bs) : 1;
const memory::dim M = mat_dim_x.height_;
const memory::dim N = mat_dim_y.width_;
const memory::dim K = mat_dim_x.width_;
return {BS, M, N, K};
}
void CreateMemories(const ExecutionContext& ctx) {
auto matmul_dims = GetMatmulDims(ctx);
auto BS = matmul_dims.BS;
auto M = matmul_dims.M;
auto N = matmul_dims.N;
auto K = matmul_dims.K;
bool x_trans = ctx.Attr<bool>("transpose_X");
bool y_trans = ctx.Attr<bool>("transpose_Y");
typedef memory::dims dims;
dims x_dims = {BS, M, K};
dims y_dims = {BS, K, N};
dims out_dims = {BS, M, N};
// Translate transA and transB
dims x_strides = !x_trans ? dims{M * K, K, 1} : dims{M * K, 1, M};
dims y_strides = !y_trans ? dims{N * K, N, 1} : dims{N * K, 1, K};
dims out_strides = {M * N, N, 1};
x_mem_ =
CreateMemory<XT>(x_dims, x_strides, ctx.Input<Tensor>("X")->data<XT>());
y_mem_ =
CreateMemory<YT>(y_dims, y_strides, ctx.Input<Tensor>("Y")->data<YT>());
out_mem_ = CreateMemory<OT>(
out_dims, out_strides,
ctx.Output<Tensor>("Out")->mutable_data<OT>(ctx.GetPlace()));
}
float ComputeOutputScale(const ExecutionContext& ctx) {
float scale_x = ctx.Attr<float>("Scale_x");
float scale_y = ctx.Attr<float>("Scale_y");
bool force_fp32_out = ctx.Attr<bool>("force_fp32_output");
float scale_out = force_fp32_out ? 1.f : ctx.Attr<float>("Scale_out");
float alpha = ctx.Attr<float>("alpha");
return alpha * scale_out / (scale_x * scale_y);
}
void CreatePrimitive(const ExecutionContext& ctx) {
dnnl::primitive_attr attr;
float scale_out = ComputeOutputScale(ctx);
if (scale_out != 1.0f) {
constexpr unsigned tensor_wide_scale = 0;
attr.set_output_scales(tensor_wide_scale, {scale_out});
}
auto matmul_d = dnnl::matmul::desc(x_mem_.get_desc(), y_mem_.get_desc(),
out_mem_.get_desc());
auto matmul_pd = dnnl::matmul::primitive_desc(matmul_d, attr, engine_);
matmul_prim_ = dnnl::matmul(matmul_pd);
}
void Execute() {
dnnl::stream stream(engine_);
matmul_prim_.execute(stream, {
{MKLDNN_ARG_SRC, x_mem_},
{MKLDNN_ARG_WEIGHTS, y_mem_},
{MKLDNN_ARG_DST, out_mem_},
});
stream.wait();
}
void SetOutputFormat(const ExecutionContext& ctx) {
using platform::MKLDNNFormatForSize;
auto* out = ctx.Output<Tensor>("Out");
auto format =
MKLDNNFormatForSize(out->dims().size(), MKLDNNMemoryFormat::nchw);
out->set_format(format);
out->set_layout(DataLayout::kMKLDNN);
}
void UpdateDataPointers(const ExecutionContext& ctx) {
auto* x = ctx.Input<Tensor>("X");
auto* y = ctx.Input<Tensor>("Y");
auto* out = ctx.Output<Tensor>("Out");
x_mem_.set_data_handle(to_void_cast(x->data<XT>()));
y_mem_.set_data_handle(to_void_cast(y->data<YT>()));
out_mem_.set_data_handle(out->mutable_data<OT>(ctx.GetPlace()));
}
// If initialized, x memory should've been already initialized
bool IsInitialized() { return initialized_; }
void SetInitialized() { initialized_ = true; }
private:
dnnl::engine engine_;
dnnl::memory x_mem_;
dnnl::memory y_mem_;
dnnl::memory out_mem_;
dnnl::matmul matmul_prim_;
bool initialized_ = false;
};
template <typename XT, typename YT, typename OT>
static std::shared_ptr<MatMulFactory<XT, YT, OT>> GetPrimitiveFactory(
const ExecutionContext& ctx) {
const auto x_dims = framework::vectorize<int>(ctx.Input<Tensor>("X")->dims());
const auto y_dims = framework::vectorize<int>(ctx.Input<Tensor>("Y")->dims());
const auto& out_name = ctx.OutputName("Out");
const auto& dev_ctx = ctx.template device_context<MKLDNNDeviceContext>();
const std::string key =
platform::CreateKey(platform::ThreadIDasStr(), x_dims, y_dims, out_name);
auto factory =
std::static_pointer_cast<MatMulFactory<XT, YT, OT>>(dev_ctx.GetBlob(key));
if (factory == nullptr) {
factory = std::make_shared<MatMulFactory<XT, YT, OT>>();
dev_ctx.SetBlob(key, factory);
}
return factory;
}
template <typename T>
constexpr bool IsInt8() {
return std::is_same<T, int8_t>::value || std::is_same<T, uint8_t>::value;
}
// Choose appropriate primitive factory implementation based on inferred
// output type (uint8, int8 or float).
template <typename XT, typename YT>
static void ExecuteMatMul(const ExecutionContext& ctx) {
constexpr bool is_int8 = IsInt8<XT>();
const bool force_fp32_output = ctx.Attr<bool>("force_fp32_output");
constexpr bool fuse_relu = false; // TODO(intel): Enable eltwise fuses
if (!is_int8 || force_fp32_output) {
GetPrimitiveFactory<XT, YT, float>(ctx)->CreateAndExecute(ctx);
} else if (fuse_relu) {
GetPrimitiveFactory<XT, YT, uint8_t>(ctx)->CreateAndExecute(ctx);
} else {
GetPrimitiveFactory<XT, YT, int8_t>(ctx)->CreateAndExecute(ctx);
}
}
template <typename T>
class DNNLMatMulKernel : public framework::OpKernel<T> {
public:
void Compute(const framework::ExecutionContext& ctx) const override {
if (ctx.HasAttr("head_number")) {
PADDLE_ENFORCE_EQ(ctx.Attr<int>("head_number"), 1,
platform::errors::Unimplemented(
"DNNL matmul doesn't support multiple heads."));
}
ExecuteMatMul<T, T>(ctx);
}
};
} // namespace operators
} // namespace paddle
namespace ops = paddle::operators;
REGISTER_OP_KERNEL(matmul, MKLDNN, ::paddle::platform::CPUPlace,
ops::DNNLMatMulKernel<float>, ops::DNNLMatMulKernel<int8_t>,
ops::DNNLMatMulKernel<uint8_t>);
...@@ -13,12 +13,12 @@ See the License for the specific language governing permissions and ...@@ -13,12 +13,12 @@ See the License for the specific language governing permissions and
limitations under the License. */ limitations under the License. */
#pragma once #pragma once
#include <mkldnn.h>
#include <algorithm> #include <algorithm>
#include <memory> #include <memory>
#include <string> #include <string>
#include <utility> #include <utility>
#include <vector> #include <vector>
#include "mkldnn.hpp"
#include "paddle/fluid/framework/operator.h" #include "paddle/fluid/framework/operator.h"
#include "paddle/fluid/platform/place.h" #include "paddle/fluid/platform/place.h"
namespace paddle { namespace paddle {
......
# Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from __future__ import print_function
import unittest, os
import numpy as np
from paddle.fluid.tests.unittests.op_test import OpTest, skip_check_grad_ci
@skip_check_grad_ci(reason="DNNL's MatMul doesn't implemend grad kernel.")
class TestDnnlMatMulOp(OpTest):
def generate_data(self):
self.x = np.random.random((25, 2, 2)).astype("float32")
self.y = np.random.random((25, 2, 2)).astype("float32")
self.alpha = 1.0
self.out = self.alpha * np.matmul(self.x, self.y)
def set_attributes(self):
self.alpha = self.alpha if hasattr(self, 'alpha') else 1.0
self.attrs = {'alpha': self.alpha}
def setUp(self):
# Set max isa, otherwise fails on SKX and earlier
os.environ["DNNL_MAX_CPU_ISA"] = "AVX"
self.op_type = "matmul"
self._cpu_only = True
self.use_mkldnn = True
self.generate_data()
self.set_attributes()
self.attrs['use_mkldnn'] = True
self.inputs = {'X': self.x, 'Y': self.y}
self.outputs = {'Out': self.out}
def test_check_output(self):
self.check_output()
class TestDnnlMatMulOpAlpha(TestDnnlMatMulOp):
def generate_data(self):
self.x = np.random.random((17, 2, 3)).astype("float32")
self.y = np.random.random((17, 3, 2)).astype("float32")
self.alpha = 2.0
self.out = self.alpha * np.matmul(self.x, self.y)
class TestDnnlMatMulOp2D(TestDnnlMatMulOp):
def print_tensor(self, name, tensor):
print(name)
print(tensor)
def generate_data(self):
self.x = np.random.random((12, 9)).astype("float32")
self.y = np.random.random((9, 12)).astype("float32")
self.out = np.matmul(self.x, self.y)
class TestDnnlMatMulOpTransposeX(TestDnnlMatMulOp):
def generate_data(self):
self.x = np.random.random((12, 9)).astype("float32")
self.y = np.random.random((12, 9)).astype("float32")
self.out = np.matmul(np.transpose(self.x), self.y)
def set_attributes(self):
self.attrs = {'transpose_X': True}
class TestDnnlMatMulOpTransposeY(TestDnnlMatMulOp):
def generate_data(self):
self.x = np.random.random((12, 9)).astype("float32")
self.y = np.random.random((12, 9)).astype("float32")
self.out = np.matmul(self.x, np.transpose(self.y))
def set_attributes(self):
self.attrs = {'transpose_Y': True}
class TestDnnlMatMulOpTransposeY3D(TestDnnlMatMulOp):
def generate_data(self):
self.x = np.random.random((17, 3, 2)).astype("float32")
self.y = np.random.random((17, 3, 2)).astype("float32")
self.out = np.matmul(self.x, np.transpose(self.y, (0, 2, 1)))
def set_attributes(self):
self.attrs = {'transpose_Y': True}
class TestDnnlMatMulOpInt8NoScales(TestDnnlMatMulOp):
def generate_data(self):
self.x = np.random.random((12, 9)).astype("int8")
self.y = np.random.random((9, 12)).astype("int8")
self.out = np.matmul(self.x, self.y)
class TestDnnlMatMulOpInt8(TestDnnlMatMulOp):
def quantize(self, tensor):
scale = 127. / np.abs(np.amax(tensor))
quantized = np.round(scale * tensor).astype("int8")
return scale, quantized
def generate_data(self):
x_float = np.random.random((12, 9)).astype("float32")
self.x_scale, self.x = self.quantize(x_float)
y_float = np.random.random((9, 12)).astype("float32")
self.y_scale, self.y = self.quantize(y_float)
out_float = np.matmul(x_float, y_float)
self.out_scale, self.out = self.quantize(out_float)
def set_attributes(self):
self.attrs = {
'Scale_x': self.x_scale,
'Scale_y': self.y_scale,
'Scale_out': self.out_scale,
}
def test_check_output(self):
int_atol = 1
self.check_output(atol=int_atol)
class TestDnnlMatMulOpInt8ForceFP32(TestDnnlMatMulOpInt8):
def generate_data(self):
x_float = np.random.random((12, 9)).astype("float32")
self.x_scale, self.x = self.quantize(x_float)
y_float = np.random.random((9, 12)).astype("float32")
self.y_scale, self.y = self.quantize(y_float)
out_float = np.matmul(x_float, y_float)
self.out = out_float
def set_attributes(self):
self.attrs = {
'Scale_x': self.x_scale,
'Scale_y': self.y_scale,
'force_fp32_output': True
}
class TestDnnlMatMulOpInt8ForceFP32BasicScales(TestDnnlMatMulOp):
def generate_data(self):
self.x = np.random.randint(0, 3, (12, 9)).astype("int8")
self.y = np.random.randint(0, 3, (9, 12)).astype("int8")
self.out = np.matmul(self.x, self.y).astype("float32")
def set_attributes(self):
self.attrs = {'force_fp32_output': True}
if __name__ == "__main__":
unittest.main()
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册