From 03803f20fdaafd9b82bf4881e8e832c9e6c234d1 Mon Sep 17 00:00:00 2001
From: Leo Chen <chenqiuliang@baidu.com>
Date: Tue, 23 Mar 2021 14:40:38 +0800
Subject: [PATCH] [NPU] support list of tensor input (#31801)

* support list of tensor as npu input

* add comment

* fix typo

* fix typo
---
 paddle/fluid/operators/npu_op_runner.cc | 15 +++++++++++++++
 paddle/fluid/operators/npu_op_runner.h  |  2 ++
 2 files changed, 17 insertions(+)

diff --git a/paddle/fluid/operators/npu_op_runner.cc b/paddle/fluid/operators/npu_op_runner.cc
index 7ccc6b60c1..f7c3597da4 100644
--- a/paddle/fluid/operators/npu_op_runner.cc
+++ b/paddle/fluid/operators/npu_op_runner.cc
@@ -185,6 +185,21 @@ NpuOpRunner &NpuOpRunner::AddInputs(const std::vector<Tensor> &tensors) {
   return *this;
 }
 
+// NOTE(zhiqiu): For operators whose input is a list (such as concat, stack),
+// It is needed to set the name of each input tensor.
+NpuOpRunner &NpuOpRunner::AddInputNames(const std::vector<std::string> &names) {
+  PADDLE_ENFORCE_EQ(names.size(), input_descs_.size(),
+                    platform::errors::InvalidArgument(
+                        "The size of input names should be "
+                        "equal to the size of input descs, but got the size "
+                        "of input names is %d, the size of input descs is %d.",
+                        names.size(), input_descs_.size()));
+  for (size_t i = 0; i < names.size(); ++i) {
+    aclSetTensorDescName(input_descs_[i], names[i].c_str());
+  }
+  return *this;
+}
+
 NpuOpRunner &NpuOpRunner::AddOutputs(const std::vector<Tensor> &tensors) {
   for (auto tensor : tensors) {
     // create aclTensorDesc
diff --git a/paddle/fluid/operators/npu_op_runner.h b/paddle/fluid/operators/npu_op_runner.h
index 97468364ee..2ead0d5b54 100644
--- a/paddle/fluid/operators/npu_op_runner.h
+++ b/paddle/fluid/operators/npu_op_runner.h
@@ -53,6 +53,8 @@ class NpuOpRunner {
 
   NpuOpRunner &AddInputs(const std::vector<Tensor> &tensors);
 
+  NpuOpRunner &AddInputNames(const std::vector<std::string> &names);
+
   NpuOpRunner &AddOutputs(const std::vector<Tensor> &tensors);
 
   aclTensorDesc *GetInputDesc(size_t index);
-- 
GitLab