Merge pull request #52261 from cfRod:mkldnn-acl-matmul

PiperOrigin-RevId: 401770330 Change-Id: Ia2a129693626db3ad8ae0d036f9c192575421742
tensorflow · Oct 8, 2021 · 1a811e3 · 1a811e3
2 parents 545fd4f + 6b8ff16
commit 1a811e3
Show file tree

Hide file tree

Showing 3 changed files with 558 additions and 1 deletion.
diff --git a/tensorflow/core/kernels/mkl/mkl_batch_matmul_op.cc b/tensorflow/core/kernels/mkl/mkl_batch_matmul_op.cc
@@ -140,6 +140,15 @@ class BatchMatMulMkl : public OpKernel {
     MklBatchMatMulHelper bmm;
     auto params = bmm.CreateMatMulParams(lhs.shape(), rhs.shape(), out_shape,
                                          adj_x_, adj_y_);
+
+#ifdef DNNL_AARCH64_USE_ACL
+    // ACL does not support reuse of primitives with different data.
+    // For matmul, the previous approach (PR #47775) of using Tensor addresses
+    // does not work, as the addresses are re-used in matmul with different data
+    // The counter  ensure we still benefit from caching via SetMklMatmul().
+    static int counter = 1;
+    params->aarch64_counter = counter++;
+#endif
     // Create or retrieve matmul primitive from cache.
     MklMatMulPrimitive<Scalar>* matmul_prim =
         MklMatMulPrimitiveFactory<Scalar>::Get(

diff --git a/tensorflow/core/kernels/mkl/mkl_matmul_ops_common.h b/tensorflow/core/kernels/mkl/mkl_matmul_ops_common.h
@@ -536,6 +536,9 @@ struct MklMatMulParams {
   memory::dims a_strides;
   memory::dims b_strides;
   memory::dims c_strides;
+#ifdef DNNL_AARCH64_USE_ACL
+  int aarch64_counter;
+#endif
 
   MklMatMulParams(memory::dims a_dims, memory::dims b_dims, memory::dims c_dims,
                   memory::dims a_strides, memory::dims b_strides,
@@ -697,7 +700,9 @@ class MklMatMulPrimitiveFactory : public MklPrimitiveFactory<T> {
     key_creator.AddAsKey(params.b_strides);
     key_creator.AddAsKey(params.c_strides);
     key_creator.AddAsKey(typeid(T).name());
-
+#ifdef DNNL_AARCH64_USE_ACL
+    key_creator.AddAsKey(params.aarch64_counter);
+#endif
     return key_creator.GetKey();
   }