add index select op (#51498)

* add index select op

* add to op teller

* add trt version control

* delete useless code

paddle/fluid/inference/api/analysis_predictor.cc

@@ -2546,6 +2546,7 @@ USE_TRT_CONVERTER(grid_sampler)
 #endif
 #if IS_TRT_VERSION_GE(8200)
 USE_TRT_CONVERTER(set_value)
+USE_TRT_CONVERTER(index_select);
 USE_TRT_CONVERTER(temporal_shift)
 #endif
 #if PADDLE_WITH_CUSPARSELT && IS_TRT_VERSION_GE(8000)

paddle/fluid/inference/tensorrt/convert/CMakeLists.txt

@@ -46,6 +46,7 @@ list(
   hard_swish_op.cc
   clip_op.cc
   gather_op.cc
+  index_select_op.cc
   anchor_generator_op.cc
   yolo_box_op.cc
   yolo_box_head_op.cc

paddle/fluid/inference/tensorrt/convert/index_select_op.cc

+/* Copyright (c) 2018 PaddlePaddle 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. */
+
+#include "paddle/fluid/inference/tensorrt/convert/op_converter.h"
+
+namespace paddle {
+namespace framework {
+class Scope;
+
+namespace proto {
+class OpDesc;
+}  // namespace proto
+}  // namespace framework
+}  // namespace paddle
+
+namespace paddle {
+namespace inference {
+namespace tensorrt {
+
+/*
+ * Gather Op
+ */
+class IndexSelectConverter : public OpConverter {
+ public:
+  void operator()(const framework::proto::OpDesc& op,
+                  const framework::Scope& scope,
+                  bool test_mode) override {
+    VLOG(3) << "convert a fluid index select op to tensorrt index select layer";
+
+    framework::OpDesc op_desc(op, nullptr);
+    std::string input_name = op_desc.Input("X").front();
+    std::string index_name = op_desc.Input("Index").front();
+    std::string output_name = op_desc.Output("Out").front();
+    const auto input_tensor = engine_->GetITensor(input_name);
+    const auto index_tensor = engine_->GetITensor(index_name);
+
+    int axis = 0;
+    if (op_desc.HasAttr("dim")) {
+      axis = PADDLE_GET_CONST(int, op_desc.GetAttr("dim"));
+    }
+
+    auto reshape_layer = TRT_ENGINE_ADD_LAYER(engine_, Shuffle, *index_tensor);
+
+    nvinfer1::Dims index_shape{};
+    index_shape.nbDims = 1;
+    index_shape.d[0] = -1;
+
+    reshape_layer->setReshapeDimensions(index_shape);
+    reshape_layer->setName(
+        ("Index select: Shuffle: (Output: " + output_name + ")").c_str());
+
+    auto layer = TRT_ENGINE_ADD_LAYER(
+        engine_, Gather, *input_tensor, *reshape_layer->getOutput(0), axis);
+    layer->setNbElementWiseDims(0);
+
+    RreplenishLayerAndOutput(layer, "index_select", {output_name}, test_mode);
+  }
+};
+
+}  // namespace tensorrt
+}  // namespace inference
+}  // namespace paddle
+
+REGISTER_TRT_OP_CONVERTER(index_select, IndexSelectConverter);

paddle/fluid/inference/tensorrt/op_teller.cc

@@ -76,6 +76,8 @@ struct SimpleOpTypeSetTeller : public Teller {
     teller_set.insert("round");
     int8_teller_set.insert("round");
     teller_set.insert("set_value");
+    teller_set.insert("index_select");
+    int8_teller_set.insert("index_select");
 #endif
   }
 
@@ -650,7 +652,36 @@ struct SimpleOpTypeSetTeller : public Teller {
       }
 #endif
     }
+    if (op_type == "index_select") {
+#if !IS_TRT_VERSION_GE(8200)
+      return false;
+#endif
+      auto gather_inputs = desc.Inputs();
+      if (!with_dynamic_shape) {
+        return false;
+      } else {
+        auto* block = desc.Block();
+        if (block == nullptr) {
+          VLOG(3) << "The block desc is nullptr, we can't continue to analyze. "
+                     "Developers need to check whether block_desc is passed in "
+                     "the pass.";
+          return false;
+        }
+
+        auto index_var_name = desc.Input("Index")[0];
+        auto* index_var_desc = block->FindVar(index_var_name);
 
+        // The index input must be int32 or int64 datatype.
+        if (index_var_desc->GetDataType() !=
+                paddle::framework::proto::VarType_Type::VarType_Type_INT32 &&
+            index_var_desc->GetDataType() !=
+                paddle::framework::proto::VarType_Type::VarType_Type_INT64) {
+          VLOG(3)
+              << "Index select op Index input data type must be int32 or int64";
+          return false;
+        }
+      }
+    }
     if (op_type == "take_along_axis") {
 #if IS_TRT_VERSION_GE(8200)
       if (!with_dynamic_shape) return false;

python/paddle/fluid/tests/unittests/ir/inference/test_trt_convert_index_select.py

+# Copyright (c) 2021 PaddlePaddle 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.
+
+import unittest
+from functools import partial
+from typing import List
+
+import numpy as np
+from program_config import ProgramConfig, TensorConfig
+from trt_layer_auto_scan_test import TrtLayerAutoScanTest
+
+import paddle.inference as paddle_infer
+
+
+class TrtConvertIndexSelectTest(TrtLayerAutoScanTest):
+    def is_program_valid(self, program_config: ProgramConfig) -> bool:
+        inputs = program_config.inputs
+        attrs = [
+            program_config.ops[i].attrs for i in range(len(program_config.ops))
+        ]
+        if len(inputs['input_data'].shape) <= attrs[0]['dim']:
+            return False
+
+        return True
+
+    def sample_program_configs(self):
+        def generate_input1(shape):
+            return np.random.random(shape).astype(np.float32)
+
+        def generate_input2(index):
+            return np.array(index).astype(np.int32)
+
+        def generate_input4(index):
+            return np.array(index).astype(np.int64)
+
+        def generate_input3(axis):
+            return np.array([axis]).astype(np.int32)
+
+        for shape in [[32, 64, 16, 32]]:
+            for index in [[1, 4], [4, 8]]:
+                for axis in [0, 1, 2, 3]:
+                    for overwrite in [True, False]:
+                        for input in [
+                            {"X": ["input_data"], "Index": ["index_data"]}
+                        ]:
+                            for index_type_int32 in [True, False]:
+                                self.shape = shape
+                                self.axis = axis
+                                self.input_num = len(input)
+                                self.index_type_int32 = index_type_int32
+                                dics = [{"dim": axis}]
+                                ops_config = [
+                                    {
+                                        "op_type": "index_select",
+                                        "op_inputs": input,
+                                        "op_outputs": {"Out": ["output_data"]},
+                                        "op_attrs": dics[0],
+                                    }
+                                ]
+                                ops = self.generate_op_config(ops_config)
+
+                                program_config = ProgramConfig(
+                                    ops=ops,
+                                    weights={},
+                                    inputs={
+                                        "input_data": TensorConfig(
+                                            data_gen=partial(
+                                                generate_input1, shape
+                                            )
+                                        ),
+                                        "index_data": TensorConfig(
+                                            data_gen=partial(
+                                                generate_input2
+                                                if index_type_int32
+                                                else generate_input4,
+                                                index,
+                                            )
+                                        ),
+                                    },
+                                    outputs=["output_data"],
+                                )
+
+                                yield program_config
+
+    def sample_predictor_configs(
+        self, program_config
+    ) -> (paddle_infer.Config, List[int], float):
+        def generate_dynamic_shape(attrs):
+            if len(self.shape) == 1:
+                self.dynamic_shape.min_input_shape = {
+                    "input_data": [4],
+                    "index_data": [1],
+                }
+                self.dynamic_shape.max_input_shape = {
+                    "input_data": [128],
+                    "index_data": [4],
+                }
+                self.dynamic_shape.opt_input_shape = {
+                    "input_data": [16],
+                    "index_data": [2],
+                }
+            elif len(self.shape) == 2:
+                self.dynamic_shape.min_input_shape = {
+                    "input_data": [2, 4],
+                    "index_data": [1],
+                }
+                self.dynamic_shape.max_input_shape = {
+                    "input_data": [256, 256],
+                    "index_data": [4],
+                }
+                self.dynamic_shape.opt_input_shape = {
+                    "input_data": [64, 32],
+                    "index_data": [2],
+                }
+            elif len(self.shape) == 3:
+                self.dynamic_shape.min_input_shape = {
+                    "input_data": [2, 4, 4],
+                    "index_data": [1],
+                }
+                self.dynamic_shape.max_input_shape = {
+                    "input_data": [128, 256, 256],
+                    "index_data": [4],
+                }
+                self.dynamic_shape.opt_input_shape = {
+                    "input_data": [16, 64, 32],
+                    "index_data": [2],
+                }
+            elif len(self.shape) == 4:
+                self.dynamic_shape.min_input_shape = {
+                    "input_data": [2, 4, 4, 2],
+                    "index_data": [1],
+                }
+                self.dynamic_shape.max_input_shape = {
+                    "input_data": [128, 256, 64, 128],
+                    "index_data": [4],
+                }
+                self.dynamic_shape.opt_input_shape = {
+                    "input_data": [16, 64, 16, 32],
+                    "index_data": [2],
+                }
+
+        def clear_dynamic_shape():
+            self.dynamic_shape.max_input_shape = {}
+            self.dynamic_shape.min_input_shape = {}
+            self.dynamic_shape.opt_input_shape = {}
+
+        def generate_trt_nodes_num(dynamic_shape):
+            if dynamic_shape:
+                ver = paddle_infer.get_trt_compile_version()
+                if ver[0] * 1000 + ver[1] * 100 + ver[2] * 10 < 8200:
+                    return 0, 4
+                return 1, 3
+            else:
+                return 0, 4
+
+        attrs = [
+            program_config.ops[i].attrs for i in range(len(program_config.ops))
+        ]
+
+        # for static_shape
+        clear_dynamic_shape()
+        self.trt_param.precision = paddle_infer.PrecisionType.Float32
+        yield self.create_inference_config(), generate_trt_nodes_num(
+            False
+        ), 1e-5
+        self.trt_param.precision = paddle_infer.PrecisionType.Half
+        yield self.create_inference_config(), generate_trt_nodes_num(
+            False
+        ), 1e-3
+
+        # for dynamic_shape
+        generate_dynamic_shape(attrs)
+        self.trt_param.precision = paddle_infer.PrecisionType.Float32
+        yield self.create_inference_config(), generate_trt_nodes_num(True), 1e-5
+        self.trt_param.precision = paddle_infer.PrecisionType.Half
+        yield self.create_inference_config(), generate_trt_nodes_num(True), 1e-3
+
+    def test(self):
+        self.trt_param.workspace_size = 1 << 60
+        self.run_test()
+
+
+if __name__ == "__main__":
+    unittest.main()