Merge pull request #13 from Intel-tensorflow/guizili/matmul_fusion

Add matmul_grad_filter and biasaddgrad fusion

Author: chuanqi129

tensorflow/core/framework/common_shape_fns.cc

@@ -130,6 +130,36 @@ Status MatMulShape(shape_inference::InferenceContext* c) {
   return Status::OK();
 }
 
+Status MatMulGradFilterShape(shape_inference::InferenceContext* c) {
+  ShapeHandle a;
+  TF_RETURN_IF_ERROR(c->WithRank(c->input(0), 2, &a));
+
+  ShapeHandle b;
+  TF_RETURN_IF_ERROR(c->WithRank(c->input(1), 2, &b));
+
+  bool transpose_a, transpose_b;
+  TF_RETURN_IF_ERROR(c->GetAttr("transpose_a", &transpose_a));
+  TF_RETURN_IF_ERROR(c->GetAttr("transpose_b", &transpose_b));
+  DimensionHandle output_rows = transpose_a ? c->Dim(a, 0) : c->Dim(a, 1);
+  DimensionHandle output_cols = c->Dim(b, 1);
+
+  if (transpose_b) {
+    auto tmp = output_rows;
+    output_rows = output_cols;
+    output_cols = tmp;
+  }
+
+  // Validate that the inner shapes are compatible.
+  DimensionHandle inner_a = transpose_a ? c->Dim(a, 1) : c->Dim(a, 0);
+  DimensionHandle inner_b = c->Dim(b, 0);
+  DimensionHandle merged;
+  TF_RETURN_IF_ERROR(c->Merge(inner_a, inner_b, &merged));
+
+  c->set_output(0, c->Matrix(output_rows, output_cols));
+  c->set_output(1, c->Vector(output_cols));
+  return Status::OK();
+}
+
 namespace {
 
 // Validate that an Einsum subscript contains exactly one or zero ellipsis; and

tensorflow/core/framework/common_shape_fns.h

@@ -102,6 +102,7 @@ Status MakeShapeFromFormat(TensorFormat format, DimensionOrConstant N,
 
 // Shape function for MatMul-like operations.
 Status MatMulShape(shape_inference::InferenceContext* c);
+Status MatMulGradFilterShape(shape_inference::InferenceContext* c);
 
 // Shape function for Batched MatMul-like operations with broadcasting across
 // batch dimensions.

tensorflow/core/graph/mkl_layout_pass.cc

@@ -275,6 +275,7 @@ class MklLayoutRewritePass : public GraphOptimizationPass {
     csinfo_.fused_conv2d = "_FusedConv2D";
     csinfo_.fused_depthwise_conv2d = "_FusedDepthwiseConv2dNative";
     csinfo_.fused_matmul = "_FusedMatMul";
+    csinfo_.fused_matmul_grad = "_FusedMatMulGrad";
     csinfo_.identity = "Identity";
     csinfo_.leakyrelu = "LeakyRelu";
     csinfo_.leakyrelu_grad = "LeakyReluGrad";
@@ -298,6 +299,7 @@ class MklLayoutRewritePass : public GraphOptimizationPass {
     csinfo_.mkl_fused_conv2d = "_MklFusedConv2D";
     csinfo_.mkl_fused_depthwise_conv2d = "_MklFusedDepthwiseConv2dNative";
     csinfo_.mkl_fused_matmul = "_MklFusedMatMul";
+    csinfo_.mkl_fused_matmul_grad = "_MklFusedMatMulGrad";
     csinfo_.mkl_pad_with_conv2d = "_MklPadWithConv2D";
     csinfo_.mkl_pad_with_fused_conv2d = "_MklPadWithFusedConv2D";
     csinfo_.pad = "Pad";
@@ -487,6 +489,9 @@ class MklLayoutRewritePass : public GraphOptimizationPass {
                       kRewriteForLayoutPropagation});
     rinfo_.push_back({csinfo_.fused_matmul, csinfo_.mkl_fused_matmul,
                       CopyAttrsAllCheckConstFilter, FusedMatMulRewrite});
+    rinfo_.push_back({csinfo_.fused_matmul_grad, csinfo_.mkl_fused_matmul_grad,
+                      CopyAttrsAll, AlwaysRewrite,
+                      kRewriteForLayoutPropagation});
 
     rinfo_.push_back({csinfo_.identity,
                       mkl_op_registry::GetMklOpName(csinfo_.identity),
@@ -933,6 +938,7 @@ class MklLayoutRewritePass : public GraphOptimizationPass {
     string fused_conv2d;
     string fused_depthwise_conv2d;
     string fused_matmul;
+    string fused_matmul_grad;
     string identity;
     string leakyrelu;
     string leakyrelu_grad;
@@ -954,6 +960,7 @@ class MklLayoutRewritePass : public GraphOptimizationPass {
     string mkl_fused_conv2d;
     string mkl_fused_depthwise_conv2d;
     string mkl_fused_matmul;
+    string mkl_fused_matmul_grad;
     string mkl_pad_with_conv2d;
     string mkl_pad_with_fused_conv2d;
     string mul;
@@ -3742,6 +3749,7 @@ MklLayoutRewritePass::CheckForNodeRewrite(const Node* n) const {
       n->type_string() != csinfo_.fused_conv2d &&
       n->type_string() != csinfo_.fused_depthwise_conv2d &&
       n->type_string() != csinfo_.fused_matmul &&
+      n->type_string() != csinfo_.fused_matmul_grad &&
       !mkl_op_registry::IsMklOp(mkl_op_registry::GetMklOpName(n->type_string()),
                                 T)) {
     return nullptr;

tensorflow/core/graph/mkl_layout_pass_test.cc

@@ -2009,6 +2009,30 @@ REGISTER_TEST_ALL_TYPES(NodeRewrite_FusedMatMul_Positive)
 REGISTER_TEST_ALL_TYPES(NodeRewrite_FusedMatMul_Negative);
 #undef REGISTER_TEST
 
+// Test set: _FusedMatMulGrad -> MklFusedMatMulGrad rewrite tests
+#define REGISTER_TEST(NAME, T, INPUT)                                          \
+  TEST_F(MklLayoutPassTest, NAME##_##T) {                                      \
+  InitGraph(                                                                   \
+      "node { name: 'A' op: '" #INPUT "'}"                                     \
+      "node { name: 'B' op: '" #INPUT "'}"                                     \
+      "node { name: 'D' op: '_FusedMatMulGrad'"                                \
+      " attr { key: 'T'                value { type:" #T  "} }"                \
+      " attr { key: 'transpose_a'      value { b: false } }"                   \
+      " attr { key: 'transpose_b'      value { b: false } }"                   \
+      " attr { key: 'fused_ops'        value { list: {s: 'BiasAddGrad'} } }"   \
+      " input: ['A', 'B']}"                                                    \
+      "node { name: 'Z' op: 'Zeta'"                                            \
+      " attr {key: 'T'                 value { type: " #T " } }"               \
+      " input: ['D']}");                                                       \
+  EXPECT_EQ(DoMklLayoutOptimizationPass(),                                     \
+            "A(" #INPUT ");B(" #INPUT ");D(_MklFusedMatMulGrad);"              \
+            "DMT/_0(Const);DMT/_1(Const);Z(Zeta)"                              \
+            "|A->D;A:control->DMT/_0:control;A:control->DMT/_1:control;"       \
+            "B->D:1;D->Z;DMT/_0->D:2;DMT/_1->D:3");                            \
+}
+REGISTER_TEST_ALL_TYPES(NodeRewrite_FusedMatMulGrad_Positive);
+#undef REGISTER_TEST
+
 // Merge test for PadWithFusedConv2D Op with BiasAdd fusion
 // padding is VALID type
 // A = input(image), B = input(paddings), C = Pad(A, B) = input of conv2D,