From 63695ad9513a31ca34e6bf96cbb85f6f76ed1e8a Mon Sep 17 00:00:00 2001
From: guopeian <guopeian1@hisilicon.com>
Date: Thu, 16 Jan 2025 15:17:05 +0800
Subject: [PATCH] temp

---
 tf_adapter/kernels/geop_npu.cc | 52 +++++++++++++++++++++++++++++-----
 tf_adapter/kernels/geop_npu.h  |  3 +-
 2 files changed, 47 insertions(+), 8 deletions(-)

diff --git a/tf_adapter/kernels/geop_npu.cc b/tf_adapter/kernels/geop_npu.cc
index 80fa15a3b..778bec37b 100644
--- a/tf_adapter/kernels/geop_npu.cc
+++ b/tf_adapter/kernels/geop_npu.cc
@@ -215,6 +215,18 @@ Status BuildStringOutput(geDataUniquePtr data_ptr, size_t output_size, Tensor &c
   return Status::OK();
 }
 
+Status BuildFakeOutputTensorInfo(OpKernelContext *ctx) {
+  int num_outputs = ctx->num_outputs();
+  ADP_LOG(INFO) << "BuildOutputTensorInfo, num_outputs:" << num_outputs;
+  for (int32_t i = 0; i < num_outputs; i++) {
+    const DataType out_type = ctx->op_kernel().output_type(i);
+    std::vector<int64> dims;
+    TensorShape out_shape(dims);
+    ctx->set_output(i, Tensor(out_type, out_shape));
+  }
+  return Status::OK();
+}
+
 Status BuildOutputTensorInfo(OpKernelContext *ctx, std::vector<ge::Tensor> &outputs) {
   // ctx is not nullptr
   int num_outputs = ctx->num_outputs();
@@ -933,11 +945,22 @@ PartialTensorShape GeOp::MakeCompatShape(const PartialTensorShape &a, const Part
   if (a.dims() != b.dims()) {
     return kUnknownRankShape;
   }
-  return MakeUnknownShape(b.dims());
+  return MakeUnknownShape(a, b);
 }
 
-void GeOp::InitGraphShape(OpKernelContext *const ctx) {
+PartialTensorShape GeOp::MakeUnknownShape(const int32_t &size) const {
+  const static auto kUnknownRankShape = PartialTensorShape();
+  static constexpr int64 kUnknownDim = -1;
+  std::vector<int64> dims(size, kUnknownDim);
+  PartialTensorShape out_shape;
+  auto status = PartialTensorShape::MakePartialShape(dims.data(),
+      static_cast<int32_t>(dims.size()), &out_shape);
+  return status.ok() ? out_shape : kUnknownRankShape;
+}
+
+bool GeOp::InitGraphShape(OpKernelContext *const ctx) {
   mutex_lock lock{graph_handler_.graph_mu};
+  bool is_init = false;
   for (size_t i = 0UL; i < static_cast<size_t>(ctx->num_inputs()); i++) {
     auto &shape = input_shapes_vec_[i];
     auto &value_shape = ctx->input(static_cast<int32_t>(i)).shape();
@@ -948,9 +971,11 @@ void GeOp::InitGraphShape(OpKernelContext *const ctx) {
       } else {
         shape = value_shape;
       }
+      is_init = true;
       ADP_LOG(INFO) << "Init input " << i << " shape to " << shape.value().DebugString();
     }
   }
+  return is_init;
 }
 
 bool GeOp::MaybeUpdateShape(OpKernelContext *const ctx) {
@@ -1034,13 +1059,20 @@ Status GeOp::DoGraphParser(ge::ComputeGraphPtr &compute_graph, FunctionLibraryDe
   return Status::OK();
 }
 
-PartialTensorShape GeOp::MakeUnknownShape(const int32_t &size) const {
+PartialTensorShape GeOp::MakeUnknownShape(const PartialTensorShape &a, const PartialTensorShape &b) const {
   const static auto kUnknownRankShape = PartialTensorShape();
   static constexpr int64 kUnknownDim = -1;
-  std::vector<int64> dims(size, kUnknownDim);
+  std::vector<int64> result_dims;
+  for (int32_t i = 0; i < a.dims(); i++) {
+    if (a.dim_size(i) != b.dim_size(i)) {
+      result_dims.emplace_back(kUnknownDim);
+      continue;
+    }
+    result_dims.emplace_back(a.dim_size(i));
+  }
   PartialTensorShape out_shape;
-  auto status = PartialTensorShape::MakePartialShape(dims.data(),
-      static_cast<int32_t>(dims.size()), &out_shape);
+  auto status = PartialTensorShape::MakePartialShape(result_dims.data(),
+      static_cast<int32_t>(result_dims.size()), &out_shape);
   return status.ok() ? out_shape : kUnknownRankShape;
 }
 
@@ -1438,7 +1470,7 @@ void GeOp::ComputeAsync(OpKernelContext *ctx, DoneCallback done) {
   OP_REQUIRES_ASYNC(ctx, (!((is_set_dynamic_config) && (compile_dynamic_mode_ == "1"))),
       errors::Internal("Option compile_dynamic_mode cannot set when set input_shape, dynamic_dims and dynamic_node_type."),
       done);
-  InitGraphShape(ctx);
+  bool is_init = InitGraphShape(ctx);
   if (is_aoe_) {
     OP_REQUIRES_ASYNC(ctx, !is_set_dynamic_config,
         errors::Internal("Dynamic input config can not use with mstuning."), done);
@@ -1447,6 +1479,12 @@ void GeOp::ComputeAsync(OpKernelContext *ctx, DoneCallback done) {
         errors::Internal("RunTuning fail.\n", ge::GEGetErrorMsgV2().GetString()), done);
     ADP_LOG(INFO) << ctx->op_kernel().name() << " RunTuning finish.";
   }
+  if (is_init) {
+    ADP_LOG(INFO) << "[GEOP] First step need not execute";
+    BuildFakeOutputTensorInfo(ctx);
+    done();
+    return;
+  }
   OP_REQUIRES_OK_ASYNC(ctx,
       CompileAndRunGraph(ctx, input_vec, inputs, input_shapes, done), done);
   return;
diff --git a/tf_adapter/kernels/geop_npu.h b/tf_adapter/kernels/geop_npu.h
index fc44c4022..5fe3ee608 100644
--- a/tf_adapter/kernels/geop_npu.h
+++ b/tf_adapter/kernels/geop_npu.h
@@ -185,13 +185,14 @@ public:
   bool IsDynamicGetNext(const Node *node);
 
   void ChangeChannelNameAttr(NodeDef &node_def) const;
-  void InitGraphShape(OpKernelContext *const ctx);
+  bool InitGraphShape(OpKernelContext *const ctx);
   bool IsDynamicConfig();
 
   PartialTensorShape MakeCompatShape(const PartialTensorShape &a, const PartialTensorShape &b) const;
 
   bool MaybeUpdateShape(OpKernelContext *const ctx);
   PartialTensorShape MakeUnknownShape(const int32_t &size) const;
+  PartialTensorShape MakeUnknownShape(const PartialTensorShape &a, const PartialTensorShape &b) const;
   Status ProcessForDiffNodeTypes(Graph &graph, bool &is_initialize, bool &is_allreduce);
 
   void ProcessGetNextNode(const Node *node);
-- 
Gitee