From eea87f4b020f4edbe48b5270057339dfb8a1c93b Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E6=9D=8E=E5=BE=B7=E9=92=8A?= <lidezhao1@huawei.com>
Date: Sat, 22 Feb 2025 17:35:31 +0800
Subject: [PATCH] add tfa npu_aicore_ops

---
 tf_adapter/kernels/aicore/npu_aicore_ops.cc   | 10 ++++++++++
 tf_adapter/ops/aicore/npu_aicore_ops.cc       | 14 ++++++++++++++
 .../python/npu_bridge/npu_cpu/npu_cpu_ops.py  | 19 +++++++++++++++++++
 3 files changed, 43 insertions(+)

diff --git a/tf_adapter/kernels/aicore/npu_aicore_ops.cc b/tf_adapter/kernels/aicore/npu_aicore_ops.cc
index 35c12ac69..4d9cef70f 100644
--- a/tf_adapter/kernels/aicore/npu_aicore_ops.cc
+++ b/tf_adapter/kernels/aicore/npu_aicore_ops.cc
@@ -73,6 +73,16 @@ public:
 REGISTER_KERNEL_BUILDER(Name("EmbeddingHashTableLookupOrInsert")
 .Device(tensorflow::DEVICE_CPU), EmbeddingHashTableLookupOrInsertOp);
 
+class StatelessDropoutOp : public tensorflow::OpKernel {
+public:
+  explicit StatelessDropoutOp(tensorflow::OpKernelConstruction *context) : OpKernel(context) {}
+  ~StatelessDropoutOp() override {}
+  void Compute(tensorflow::OpKernelContext *context) override {}
+};
+
+REGISTER_KERNEL_BUILDER(Name("StatelessDropout")
+.Device(tensorflow::DEVICE_CPU), StatelessDropoutOp);
+
 REGISTER_KERNEL_BUILDER(
     Name("FastGelu")
 .
diff --git a/tf_adapter/ops/aicore/npu_aicore_ops.cc b/tf_adapter/ops/aicore/npu_aicore_ops.cc
index 077d83788..29f8003c8 100644
--- a/tf_adapter/ops/aicore/npu_aicore_ops.cc
+++ b/tf_adapter/ops/aicore/npu_aicore_ops.cc
@@ -820,6 +820,20 @@ REGISTER_OP("TabulateFusionSeAGrad")
       return Status::OK();
     });
 
+REGISTER_OP("StatelessDropout")
+    .Input("x: T")
+    .Input("noise_shape: int64")
+    .Input("p: T")
+    .Input("seed: int64")
+    .Input("offset: int64")
+    .Output("y: T")
+    .Attr("T: {float16, float32, bfloat16}")
+    .SetIsStateful()
+    .SetShapeFn([](shape_inference::InferenceContext *c) {
+      c->set_output(0, c->input(0));
+      return Status::OK();
+    });
+
 REGISTER_OP("TabulateFusion")
     .Input("table:T")
     .Input("table_info:T")
diff --git a/tf_adapter/python/npu_bridge/npu_cpu/npu_cpu_ops.py b/tf_adapter/python/npu_bridge/npu_cpu/npu_cpu_ops.py
index 273c67241..d1aae2c37 100644
--- a/tf_adapter/python/npu_bridge/npu_cpu/npu_cpu_ops.py
+++ b/tf_adapter/python/npu_bridge/npu_cpu/npu_cpu_ops.py
@@ -379,6 +379,25 @@ def device_feature_mapping(feature_id):
     return result
 
 
+## 提供host侧StatelessDropout功能
+#  @param x float32,float16,bfloat16 类型
+#  @param noise_shape int64 类型
+#  @param p float32,float16,bfloat16 类型
+#  @param seed int64 类型
+#  @param offset int64 类型
+#  @return values float32,float16,bfloat16 类型
+def stateless_dropout(x, noise_shape, p, seed, offset):
+    """ host stateless_dropout. """
+    result = gen_npu_cpu_ops.StatelessDropout(
+        x=x,
+        noise_shape=noise_shape,
+        p=p,
+        seed=seed,
+        offset=offset   
+    )
+    return result
+
+
 ## 提供device侧EmbeddingFeatureMappingV2功能
 #  @param talbe_name string 类型
 #  @param feature_id int64 类型
-- 
Gitee