From 00cf573549ccfba5f46669be052f7099f7f24966 Mon Sep 17 00:00:00 2001
From: q00475944 <qianzehong@huawei.com>
Date: Tue, 22 Jun 2021 16:34:26 +0800
Subject: [PATCH] x

---
 tf_adapter/kernels/lru_cache_v2_ops.cc        | 43 +++++++++++++++++++
 tf_adapter/ops/npu_aicore_ops.cc              | 22 ++++++++++
 .../npu_bridge/estimator/npu_aicore_ops.py    | 16 +++++++
 3 files changed, 81 insertions(+)
 create mode 100644 tf_adapter/kernels/lru_cache_v2_ops.cc

diff --git a/tf_adapter/kernels/lru_cache_v2_ops.cc b/tf_adapter/kernels/lru_cache_v2_ops.cc
new file mode 100644
index 000000000..15764c741
--- /dev/null
+++ b/tf_adapter/kernels/lru_cache_v2_ops.cc
@@ -0,0 +1,43 @@
+/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+
+Copyright (C) 2021-2021. Huawei Technologies Co., Ltd. All rights reserved.
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/core/framework/op_kernel.h"
+#include "tf_adapter/common/adp_logger.h"
+
+namespace tensorflow {
+class LRUCacheV2Op : public OpKernel {
+    public:
+        explicit LRUCacheV2Op(OpKernelConstruction *context) : OpKernel(context) {}
+        ~LRUCacheV2Op() override = default;
+        void Compute(OpKernelContext *context) override {
+            ADP_LOG(INFO) << "LRUCacheV2Op Compute, num_inputs: " << context->num_inputs();
+        }
+        bool IsExpensive() override { return false; }
+};
+
+REGISTER_KERNEL_BUILDER(Name("LRUCacheV2").Device(DEVICE_CPU), LRUCacheV2Op);
+}  // namespace tensorflow
diff --git a/tf_adapter/ops/npu_aicore_ops.cc b/tf_adapter/ops/npu_aicore_ops.cc
index 3dfed09bc..4ab661bea 100644
--- a/tf_adapter/ops/npu_aicore_ops.cc
+++ b/tf_adapter/ops/npu_aicore_ops.cc
@@ -328,5 +328,27 @@ REGISTER_OP("PReluGrad")
         c->set_output(1, c->input(2));
         return Status::OK();
     });
+
+REGISTER_OP("LRUCacheV2")
+    .Input("index_list: T")
+    .Input("data: T")
+    .Input("cache: T")
+    .Input("tag: T")
+    .Input("is_last_call: T")
+    .Output("data: T")
+    .Output("index_offset_list: T")
+    .Output("not_in_cache_index_list: T")
+    .Output("not_in_cache_number: T")
+    .Attr("pre_route_count: int")
+    .SetIsStateful()
+    .SetShapeFn([](shape_inference::InferenceContext *c) {
+        c->set_output(0, c->input(1));
+        c->set_output(1, c->input(0));
+        c->set_output(2, c->input(0));
+        auto output_scalar_shape = 
+            c->MakeShape({1,});
+        c->set_output(3, output_scalar_shape);
+        return Status::OK();
+    });
 }  // namespace
 } // namespace tensorflow
\ No newline at end of file
diff --git a/tf_adapter/python/npu_bridge/estimator/npu_aicore_ops.py b/tf_adapter/python/npu_bridge/estimator/npu_aicore_ops.py
index cc4264b86..cb2daf34a 100644
--- a/tf_adapter/python/npu_bridge/estimator/npu_aicore_ops.py
+++ b/tf_adapter/python/npu_bridge/estimator/npu_aicore_ops.py
@@ -149,4 +149,20 @@ def _DropOutDoMaskV3Grad(op, grad):
     result = npu_aicore_ops.drop_out_do_mask_v3(grad, op.inputs[1],  op.inputs[2])
     return [result, None, None]
 
+def lru_cache_v2(index_list, data, cache, tag, is_last_call, pre_route_count, name=None):
+    """
+    LRUCacheV2 op
+    """
+    index_list = ops.convert_to_tensor(index_list, name="index_list")
+    data = ops.convert_to_tensor(data, name="data")
+    cache = ops.convert_to_tensor(cache, name="cache")
+    tag = ops.convert_to_tensor(tag, name="tag")
+    is_last_call = ops.convert_to_tensor(is_last_call, name="is_last_call")
+    data,
+    index_offset_list,
+    not_in_cache_index_list,
+    not_in_cache_number = npu_aicore_ops.lru_cache_v2(index_list, data, cache, tag, is_last_call,
+                                         pre_route_count, name=name)
+    return [data, index_offset_list, not_in_cache_index_list, not_in_cache_number]
+
 # go/tf-wildcard-import
-- 
Gitee