PR tensorflow#252: Support BatchMatMul and fusions with quantized bat…

…ched matrices
Intel-tensorflow · Jul 14, 2023 · d3b1dbb · d3b1dbb
1 parent cd46629
commit d3b1dbb
Show file tree

Hide file tree

Showing 9 changed files with 1,189 additions and 152 deletions.
diff --git a/tensorflow/core/kernels/dequantize_op.cc b/tensorflow/core/kernels/dequantize_op.cc
@@ -70,12 +70,6 @@ class DequantizeOp : public OpKernel {
                   errors::InvalidArgument("Mode string must be 'MIN_COMBINED',"
                                           " 'MIN_FIRST', or 'SCALED', is '" +
                                           mode_string + "'"));
-    } else {
-      OP_REQUIRES(
-          ctx, (mode_string == "MIN_COMBINED"),
-          errors::InvalidArgument("When output type is bfloat16, Mode"
-                                  " string must be 'MIN_COMBINED', is '" +
-                                  mode_string + "'"));
     }
 
     if (mode_string == "MIN_COMBINED") {

diff --git a/tensorflow/core/kernels/mkl/BUILD b/tensorflow/core/kernels/mkl/BUILD
@@ -634,3 +634,17 @@ tf_cc_test_mkl(
         "@com_google_absl//absl/strings",
     ] + MKL_TEST_DEPS,
 )
+
+tf_cc_test_mkl(
+    name = "onednn_fused_batchmatmul_op_test",
+    size = "medium",
+    srcs = ["onednn_fused_batchmatmul_op_test.cc"],
+    linkstatic = 1,
+    deps = [
+        ":mkl_batch_matmul_op",
+        ":mkl_kernel_util",
+        "//tensorflow/core:direct_session",
+        "//tensorflow/core/kernels:quantization_utils",
+        "@com_google_absl//absl/strings",
+    ] + MKL_TEST_DEPS,
+)