saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-29 20:14:01 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

[oneDNN ep] matmulinteger postop fusion (#12354)

Browse source 
* MatMulInteger + post op fusion

This fuses MatMulInteger with upto 32 binary/elementwise
operators if running on the oneDNN execution provider.

Signed-off-by: George Nash <george.nash@intel.com>

* Remove the un-needed transformer

The MatMulIntegerToFloat transformer is not needed since
the transform done is handled by the MatMulIntegerBinaryEltwise
transformer code.

Signed-off-by: George Nash <george.nash@intel.com>

* Refactor of the post op trasformer code

This separates the code that finds the post op
nodes for MatMul and MatMulInteger to reduce code
repetition.

Signed-off-by: George Nash <george.nash@intel.com>

* Minor cleanup based on cpplint

resolved unused-variable build failure

Signed-off-by: George Nash <george.nash@intel.com>
This commit is contained in:
George Nash 2022-08-02 20:42:34 -07:00 committed by GitHub
parent 5d610bc8eb
commit 26dc09417b
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
6 changed files with 1656 additions and 53 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

111
onnxruntime/core/providers/dnnl/subgraph/dnnl_matmul_integer.cc
View file

@ -4,6 +4,11 @@
#include "dnnl_matmul_integer.h"
#include "dnnl_subgraph.h"
#include "dnnl_subgraph_primitive.h"
#include "dnnl_util.h"
#include <unordered_set>
#include <vector>
#include <string>
namespace onnxruntime {
namespace ort_dnnl {
@ -11,8 +16,28 @@ namespace ort_dnnl {
DnnlMatMulInteger::DnnlMatMulInteger() {}
void DnnlMatMulInteger::CreatePrimitive(DnnlSubgraphPrimitive& sp, DnnlNode& node) {
std::unordered_set<std::string> binary_ops = {"Add", "Div", "Mul", "Sub"};
std::unordered_set<std::string> elementwise_ops = {"Abs", "Elu", "Exp", "LeakyRelu", "Log", "Relu",
"Round", "Sigmoid", "Softplus", "Sqrt", "Tanh"};
auto eng = sp.GetEngine();
bool has_postop_fusion = false;
std::vector<std::string> post_ops;
if (node.OpType() == "MatMulIntegerPostOps") {
has_postop_fusion = true;
post_ops = node.GetPostOps();
int binary_count = 0;
// Check we have enough inputs for MatMul and the binary post ops
for (size_t i = 0; i < post_ops.size(); ++i) {
if (binary_ops.count(post_ops[i]) != 0) {
assert(node.Input(IN_BINARY_0 + binary_count).Exists());
binary_count++;
}
}
}
auto src_dims = sp.GetMemory(node.Input(IN_A)).get_desc().dims();
auto weights_dims = sp.GetMemory(node.Input(IN_B)).get_desc().dims();
@ -52,6 +77,67 @@ void DnnlMatMulInteger::CreatePrimitive(DnnlSubgraphPrimitive& sp, DnnlNode& nod
matmul_attr.set_zero_points(DNNL_ARG_WEIGHTS, /* mask */ 0, {DNNL_RUNTIME_S32_VAL});
}
/*
create a post op binary with possible unsqueezing in order to make sure onednn properly broadcast
current limitation
1. is no unsqueeze for matmul output as it is not exposed due to post op fusion
2. the third input has to be reordered to plain format
(eg, no memory format propagation if the third input is internal to subgraph)
3. adding 1s to front (unsqueeze/expand) in logical dims would possibly fail if
physical layout is not plain format
*/
if (has_postop_fusion) {
int binary_count = 0;
dnnl::post_ops ops;
for (size_t i = 0; i < post_ops.size(); ++i) {
dnnl::algorithm algo = dnnl_util::OrtOperatorToDnnlAlgorithm(post_ops[i]);
// Handle Binary post ops including the input memory
if (binary_ops.count(post_ops[i]) != 0) {
auto ori_binary_md = sp.GetMemory(node.Input(IN_BINARY_0 + binary_count).Name()).get_desc();
auto ori_binary_dims = ori_binary_md.dims();
auto binary_mem_dims = ori_binary_dims;
if (ori_binary_dims.size() != output_shape.size()) {
if (ori_binary_dims.size() > output_shape.size()) {
ORT_THROW("add fusion with matmul output broadcasting by unsqueezing is not supported");
}
// expand the input (from the binary op) if needed to support broadcasting
while (binary_mem_dims.size() < output_shape.size()) {
binary_mem_dims.insert(binary_mem_dims.begin(), 1);
}
}
// expand the dims by 1s (should always be possible)
// will throw exception if not possible
auto binary_md = ori_binary_md.reshape(binary_mem_dims);
// Possible improvment: use format any to choose the best layout
ops.append_binary(algo, binary_md);
binary_count++;
// Handle Elementwise post ops. Some of these require obtaining an 'alpha' attribute
} else if (elementwise_ops.count(post_ops[i]) != 0) {
float post_op_alpha = 0.0;
switch (algo) {
case dnnl::algorithm::eltwise_relu: {
// Need to check operator since both Relu and LeakyRelu are covered by algorithm::eltwise_relu
if (post_ops[i] == "LeakyRelu") {
post_op_alpha = GetFloatAttr(node, "alpha", /*default_alpha*/ 0.01f);
} else {
post_op_alpha = 0.0;
}
break;
}
case dnnl::algorithm::eltwise_elu: {
post_op_alpha = GetFloatAttr(node, "alpha", /*default_alpha*/ 1.0f);
break;
}
default:
post_op_alpha = 0.0;
}
ops.append_eltwise(1.0f, algo, post_op_alpha, 0.0f);
}
}
matmul_attr.set_post_ops(ops);
}
auto matmul_d = dnnl::matmul::desc(src_md, weights_md, dst_md);
auto matmul_pd = dnnl::matmul::primitive_desc(matmul_d, matmul_attr, eng);
@ -79,10 +165,33 @@ void DnnlMatMulInteger::CreatePrimitive(DnnlSubgraphPrimitive& sp, DnnlNode& nod
mem_map[DNNL_ARG_ATTR_ZERO_POINTS | DNNL_ARG_WEIGHTS] = zp_B_mem_s32;
}
sp.AddPrimitive(matmul_prim, mem_map);
if (has_postop_fusion) {
// add to memory map for extra binary inputs
int binary_count = 0;
for (size_t i = 0; i < post_ops.size(); ++i) {
if (binary_ops.count(post_ops[i]) != 0) {
dnnl::algorithm algo;
dnnl::memory::desc binary_mem_desc;
matmul_pd.get_primitive_attr().get_post_ops().get_params_binary(static_cast<int>(i), algo, binary_mem_desc);
auto binary_post_op_mem = sp.GetMemoryAndReshape(node.Input(IN_BINARY_0 + binary_count), binary_mem_desc, eng);
mem_map[DNNL_ARG_ATTR_MULTIPLE_POST_OP(static_cast<int>(i)) | DNNL_ARG_SRC_1] = binary_post_op_mem;
binary_count++;
}
}
}
sp.AddPrimitive(matmul_prim, mem_map /*, {DNNL_ARG_SRC, DNNL_ARG_WEIGHTS, DNNL_ARG_DST}*/);
sp.SetMemory(node.Output(OUT_Y), matmul_dst_mem);
}
float DnnlMatMulInteger::GetFloatAttr(DnnlNode& node, std::string attr_name, float default_value) {
auto attr = node.Attributes().find(attr_name);
if (attr != node.Attributes().end()) {
return attr->second().f();
}
return default_value;
}
} // namespace ort_dnnl
} // namespace onnxruntime

8
onnxruntime/core/providers/dnnl/subgraph/dnnl_matmul_integer.h
View file

@ -5,6 +5,8 @@
#include "dnnl_subgraph.h"
#include "dnnl_subgraph_primitive.h"
#include <string>
namespace onnxruntime {
namespace ort_dnnl {
@ -14,13 +16,17 @@ class DnnlMatMulInteger {
IN_A = 0,
IN_B = 1,
IN_A_ZERO_POINT = 2,
IN_B_ZERO_POINT = 3
IN_B_ZERO_POINT = 3,
IN_BINARY_0 = 4 // the first binary input due to matmul + binary fusion
};
enum OutputTensors : int { OUT_Y = 0 };
DnnlMatMulInteger();
void CreatePrimitive(DnnlSubgraphPrimitive& sp, DnnlNode& node);
private:
float GetFloatAttr(DnnlNode& node, std::string attr_name, float default_value);
};
} // namespace ort_dnnl

2
onnxruntime/core/providers/dnnl/subgraph/dnnl_subgraph_primitive.cc
View file

@ -200,7 +200,7 @@ void DnnlSubgraphPrimitive::AddKernels() {
// https://github.com/microsoft/onnxruntime/blob/master/docs/ContribOperators.md#com.microsoft.FusedMatMul
} else if (node.OpType() == "MatMul" || node.OpType() == "MatMulPostOps" || node.OpType() == "FusedMatMul") {
DnnlMatMul().CreatePrimitive(*this, node);
} else if (node.OpType() == "MatMulInteger") {
} else if (node.OpType() == "MatMulInteger" || node.OpType() == "MatMulIntegerPostOps") {
DnnlMatMulInteger().CreatePrimitive(*this, node);
} else if (pool_ops.count(node.OpType())) {
DnnlPool().CreatePrimitive(*this, node);

188
onnxruntime/core/providers/dnnl/subgraph/dnnl_subgraph_transformer.cc
View file

@ -12,6 +12,7 @@
#include <sstream>
#include <unordered_set>
#include <utility>
#include <memory>
namespace onnxruntime {
namespace ort_dnnl {
@ -23,6 +24,7 @@ void DnnlGraphTransformer::Apply(DnnlSubgraph& subgraph, const onnxruntime::Grap
Gelu(subgraph, onnx_subgraph_viewer);
FastGelu(subgraph, onnx_subgraph_viewer);
RemoveMatMulIntegerZP(subgraph, onnx_subgraph_viewer);
MatMulIntegerBinaryEltwise(subgraph);
}
//resolve a fusion by replacing old_indices nodes with a new_node
@ -687,17 +689,12 @@ void DnnlGraphTransformer::ConvRelu(DnnlSubgraph& subgraph) {
}
void DnnlGraphTransformer::MatMulBinaryEltwise(DnnlSubgraph& subgraph) {
std::unordered_set<std::string> binary_ops = {"Add", "Div", "Mul" , "Sub"};
std::unordered_set<std::string> elementwise_ops = {"Abs", "Elu", "Exp", "LeakyRelu", "Log", "Relu",
"Round", "Sigmoid", "Softplus", "Sqrt", "Tanh"};
static int fused_index = 0;
size_t max_index = subgraph.GetMaxNodeIndex();
for (size_t index = 0; index < max_index; index++) {
std::vector<size_t> matmul_binary_eltwize_indices = {};
auto dnnl_node = subgraph.GetDnnlNode(index);
auto attr_node = dnnl_node;
bool attribute_flag = false;
if (dnnl_node == nullptr || dnnl_node->OpType() != "MatMul") {
continue;
@ -709,50 +706,11 @@ void DnnlGraphTransformer::MatMulBinaryEltwise(DnnlSubgraph& subgraph) {
auto fused_node_inputs = dnnl_node->Inputs();
matmul_binary_eltwize_indices.push_back(dnnl_node->Index());
// Upto 32 post-ops are supported. Since the initial MatMul is not a
// post op we check for MAX_POST_OP_COUNT + 1
const size_t MAX_POST_OP_COUNT = 32;
while (matmul_binary_eltwize_indices.size() < MAX_POST_OP_COUNT + 1) {
if (!IsNodeFusable(subgraph, dnnl_node)) {
break;
}
auto next_dnnl_node = dnnl_node->Output(0).GetConsumers()[0].GetNode();
if (next_dnnl_node == nullptr) {
break;
}
auto next_type = next_dnnl_node->OpType();
bool is_binary_op = !(binary_ops.count(next_type) == 0);
bool is_eltwise_op = !(elementwise_ops.count(next_type) == 0);
if (!is_binary_op && !is_eltwise_op) {
break;
}
if (is_binary_op) {
if (dnnl_node->Output(0).Name() == next_dnnl_node->Inputs()[0]->Name()) {
fused_node_inputs.push_back(next_dnnl_node->Inputs()[1]);
} else {
if (next_dnnl_node->OpType() == "Div" || next_dnnl_node->OpType() == "Sub") {
break;
}
fused_node_inputs.push_back(next_dnnl_node->Inputs()[0]);
}
} else if (is_eltwise_op) {
// We only support a single node with an "alpha" attribute. If we see Elu or
// LeakyRelu set attr_node and attribute_flag. If the attribute_flag is set break
// out of the while loop looking for additional post ops. Since the next op cannot
// be part of the postop fusion.
if (next_dnnl_node->OpType() == "Elu" || next_dnnl_node->OpType() == "LeakyRelu") {
if (attribute_flag) break;
attr_node = next_dnnl_node;
attribute_flag = true;
}
}
dnnl_node = next_dnnl_node;
matmul_binary_eltwize_indices.push_back(dnnl_node->Index());
}
dnnl_node = FuseBinaryEltwisePostOps(subgraph,
dnnl_node,
matmul_binary_eltwize_indices,
fused_node_inputs,
attr_node);
if (!(matmul_binary_eltwize_indices.size() > 1)) {
matmul_binary_eltwize_indices.clear();
@ -845,5 +803,137 @@ void DnnlGraphTransformer::RemoveMatMulIntegerZP(DnnlSubgraph& subgraph, const o
}
}
void DnnlGraphTransformer::MatMulIntegerBinaryEltwise(DnnlSubgraph& subgraph) {
static int fused_index = 0;
size_t max_index = subgraph.GetMaxNodeIndex();
for (size_t index = 0; index < max_index; index++) {
std::vector<size_t> matmul_binary_eltwize_indices = {};
auto dnnl_node = subgraph.GetDnnlNode(index);
auto attr_node = dnnl_node;
if (dnnl_node == nullptr || dnnl_node->OpType() != "MatMulInteger") {
continue;
}
if (!IsNodeFusable(subgraph, dnnl_node)) {
continue;
}
auto fused_node_inputs = dnnl_node->Inputs();
matmul_binary_eltwize_indices.push_back(dnnl_node->Index());
auto next_dnnl_node = dnnl_node->Output(0).GetConsumers()[0].GetNode();
if (next_dnnl_node == nullptr) {
continue;
}
if (next_dnnl_node->OpType() != "Cast") {
continue;
}
if (!IsNodeFusable(subgraph, next_dnnl_node)) {
continue;
}
dnnl_node = next_dnnl_node;
matmul_binary_eltwize_indices.push_back(dnnl_node->Index());
dnnl_node = FuseBinaryEltwisePostOps(subgraph,
dnnl_node,
matmul_binary_eltwize_indices,
fused_node_inputs,
attr_node);
if (!(matmul_binary_eltwize_indices.size() > 1)) {
matmul_binary_eltwize_indices.clear();
continue;
}
//construct new node
auto fused_node = std::make_unique<DnnlNode>();
fused_node->Name() = "MatMulIntegerPostOps_fusion" + std::to_string(fused_index++);
std::string fused_node_name = "MatMulIntegerPostOps";
for (size_t i : matmul_binary_eltwize_indices) {
if (subgraph.GetDnnlNode(i)->OpType() != "MatMulInteger" && subgraph.GetDnnlNode(i)->OpType() != "Cast") {
fused_node->AppendPostOp(subgraph.GetDnnlNode(i)->OpType());
}
}
fused_node->OpType() = fused_node_name;
fused_node->Inputs() = fused_node_inputs;
fused_node->Outputs() = {dnnl_node->Outputs()[0]};
fused_node->Attributes().insert(attr_node->Attributes());
if (debug_log_) {
std::stringstream ss;
for (size_t i : matmul_binary_eltwize_indices) {
ss << subgraph.GetDnnlNode(i)->OpType() << "[" << subgraph.GetDnnlNode(i)->Name() << "] ";
}
LOGS_DEFAULT(ERROR) << fused_node->OpType() << "[" << fused_node->Name() << "] fusion of " << ss.str();
}
// insert new node, remove original nodes, connect new edges
ResolveFusion(subgraph, matmul_binary_eltwize_indices, std::move(fused_node));
}
}
DnnlNode* DnnlGraphTransformer::FuseBinaryEltwisePostOps(DnnlSubgraph& subgraph,
DnnlNode* node,
std::vector<size_t>& indices,
std::vector<DnnlTensor*>& fused_node_inputs,
DnnlNode*& attr_node) {
std::unordered_set<std::string> binary_ops = {"Add", "Div", "Mul", "Sub"};
std::unordered_set<std::string> elementwise_ops = {"Abs", "Elu", "Exp", "LeakyRelu", "Log", "Relu",
"Round", "Sigmoid", "Softplus", "Sqrt", "Tanh"};
// Upto 32 post-ops are supported by oneDNN framework.
const size_t MAX_POST_OP_COUNT = 32;
bool attribute_flag = false;
auto dnnl_node = node;
size_t post_op_count = 0;
while (post_op_count < MAX_POST_OP_COUNT) {
if (!IsNodeFusable(subgraph, dnnl_node)) {
break;
}
auto next_dnnl_node = dnnl_node->Output(0).GetConsumers()[0].GetNode();
if (next_dnnl_node == nullptr) {
break;
}
auto next_type = next_dnnl_node->OpType();
bool is_binary_op = !(binary_ops.count(next_type) == 0);
bool is_eltwise_op = !(elementwise_ops.count(next_type) == 0);
if (!is_binary_op && !is_eltwise_op) {
break;
}
if (is_binary_op) {
if (dnnl_node->Output(0).Name() == next_dnnl_node->Inputs()[0]->Name()) {
fused_node_inputs.push_back(next_dnnl_node->Inputs()[1]);
} else {
// oneDNN can only fuse binary post op for the second input. Due to the fact
// that division and subraction are not associative we can not fuse them if the
// input for that node is the first input.
if (next_dnnl_node->OpType() == "Div" || next_dnnl_node->OpType() == "Sub") {
break;
}
fused_node_inputs.push_back(next_dnnl_node->Inputs()[0]);
}
} else if (is_eltwise_op) {
// We only support a single node with an "alpha" attribute. If we see Elu or
// LeakyRelu set attr_node and attribute_flag. If the attribute_flag is set break
// out of the while loop looking for additional post ops. Since the next op cannot
// be part of the postop fusion.
if (next_dnnl_node->OpType() == "Elu" || next_dnnl_node->OpType() == "LeakyRelu") {
if (attribute_flag) break;
attr_node = next_dnnl_node;
attribute_flag = true;
}
}
dnnl_node = next_dnnl_node;
indices.push_back(dnnl_node->Index());
post_op_count++;
}
return dnnl_node;
}
} // namespace ort_dnnl
} // namespace onnxruntime

12
onnxruntime/core/providers/dnnl/subgraph/dnnl_subgraph_transformer.h
View file

@ -29,8 +29,18 @@ class DnnlGraphTransformer {
bool IsInitilizedWithExpectedValue(const onnxruntime::GraphViewer& onnx_subgraph_viewer, DnnlTensor& input_arg, float expected_value);
void ConvRelu(DnnlSubgraph& subgraph);
void MatMulBinaryEltwise(DnnlSubgraph& subgraph);
void MatMulAdd(DnnlSubgraph& subgraph);
void RemoveMatMulIntegerZP(DnnlSubgraph& subgraph, const onnxruntime::GraphViewer& onnx_subgraph_viewer);
void MatMulIntegerBinaryEltwise(DnnlSubgraph& subgraph);
// Function used to identify and fuse post ops
//
// @param[in] subgraph the DnnlSubgrapy that we are searching for possible fusions
// @param[in] node is the first node to check if it contains a binary or an elementwise op
// @param[in/out] indicies list of all the indicies for the nodes that will be fused
// @param[in/out] fused_node_inputs list of all the inputs that will be part of the fused node
// @param[in/out] attr_node this node contains the attributes that will be passed onto the final fused node
//
// @return a pointer to the node after the last identified binary/elementwise fusion
DnnlNode* FuseBinaryEltwisePostOps(DnnlSubgraph& subgraph, DnnlNode* node, std::vector<size_t>& indices, std::vector<DnnlTensor*>& fused_node_inputs, DnnlNode*& attr_node);
// This function checks a few things
// - the node in question has a single output
// - The output of the node is only consumed by a one other node

1388
onnxruntime/test/providers/dnnl/transformer/matmulinteger_post_op_transform_test.cc Normal file
View file

File diff suppressed because it is too large Load diff