Fix Transpose and MatMul fusion code to check the input datatypes as … (#7147)

* Fix Transpose and MatMul fusion code to check the input datatypes as FusedMatMul only supports floating point datatypes.

* Added testcases to make sure that the int32/int64 datatypes prevent Transport-MatMul fusion.
This commit is contained in:
satyajandhyala 2021-03-28 09:24:12 -07:00 committed by GitHub
parent 65ce5f07b3
commit 90294b9c43
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
5 changed files with 51 additions and 1 deletions

View file

@ -168,6 +168,15 @@ static Node* ReorderCastAndTranspose(Graph& graph, Node* cast,
return &new_transpose;
}
// Check whether the element_type is an allowed FusedMatMul data type or not.
static bool IsAllowedFusedMatMulDataType(ONNX_NAMESPACE::TensorProto_DataType element_type)
{
return element_type == ONNX_NAMESPACE::TensorProto_DataType_FLOAT ||
element_type == ONNX_NAMESPACE::TensorProto_DataType_FLOAT16 ||
element_type == ONNX_NAMESPACE::TensorProto_DataType_DOUBLE ||
element_type == ONNX_NAMESPACE::TensorProto_DataType_BFLOAT16;
}
/*********************************************************************************************
Case I: The followin is a scenario where Transpose output feeds MatMul. The Transpose input can be either on the left or right.
@ -277,9 +286,17 @@ Status MatmulTransposeFusion::ApplyImpl(Graph& graph, bool& modified, int graph_
}
NodeArg* left_input = node.MutableInputDefs()[0];
auto left_type = left_input->TypeAsProto()->tensor_type().elem_type();
if (!IsAllowedFusedMatMulDataType(static_cast<ONNX_NAMESPACE::TensorProto_DataType>(left_type))) {
continue;
}
auto left = GetTransposeNodeFromOutput(graph, *left_input);
NodeArg* right_input = node.MutableInputDefs()[1];
auto right_type = right_input->TypeAsProto()->tensor_type().elem_type();
if (!IsAllowedFusedMatMulDataType(static_cast<ONNX_NAMESPACE::TensorProto_DataType>(right_type))) {
continue;
}
auto right = GetTransposeNodeFromOutput(graph, *right_input);
if (!left) {

View file

@ -918,10 +918,12 @@ TEST_F(GraphTransformationTests, TransposeMatmulFusionOnThreeTranspose) {
ASSERT_TRUE(static_cast<bool>(node.GetAttributes().at("transB").i()));
}
TEST_F(GraphTransformationTests, TransposeMatmulNoFusionOnInvalidPerm) {
TEST_F(GraphTransformationTests, TransposeMatmulNoFusionOnInvalidInput) {
const std::vector<PathString> model_uris = {
MODEL_FOLDER "fusion/transpose_matmul_4d_fusion_invalid_perm.onnx",
MODEL_FOLDER "fusion/transpose_matmul_4d_fusion_invalid_default_perm.onnx",
MODEL_FOLDER "fusion/transpose_matmul_4d_fusion_invalid_datatype_int32.onnx",
MODEL_FOLDER "fusion/transpose_matmul_4d_fusion_invalid_datatype_int64.onnx",
};
for (const auto& model_uri : model_uris) {
std::shared_ptr<Model> p_model;

View file

@ -221,3 +221,34 @@ def gen_transpose_fusion_with_cast(model_path):
gen_transpose_fusion_with_cast(
"transpose_cast_matmul_4d_fusion")
def gen_transpose_fusion_invalid_datatype(model_path, datatype):
nodes = [
helper.make_node(
"Transpose",
["input_0"],
["transposed_input_0"],
perm = [0, 1, 3, 2]),
helper.make_node(
"MatMul",
["transposed_input_0", "input_1"],
["output"])
]
inputs = [
helper.make_tensor_value_info(
"input_0", datatype, [2, 3, 'K', 'M']),
helper.make_tensor_value_info(
"input_1", datatype, [2, 3, 'K', 'N'])
]
outputs = [
helper.make_tensor_value_info(
"output", datatype, [2, 3, 'M', 'N'])
]
save(model_path, nodes, inputs, outputs, [])
gen_transpose_fusion_invalid_datatype("transpose_matmul_4d_fusion_invalid_datatype_int32.onnx", TensorProto.INT32)
gen_transpose_fusion_invalid_datatype("transpose_matmul_4d_fusion_invalid_datatype_int64.onnx", TensorProto.INT64)