[Paddle Inference] Add take_along_axis trt converter (#48358)

paddle/fluid/inference/api/analysis_predictor.cc

@@ -2341,6 +2341,7 @@ USE_TRT_CONVERTER(tanh_shrink)
 USE_TRT_CONVERTER(logsigmoid)
 USE_TRT_CONVERTER(lookup_table)
 USE_TRT_CONVERTER(expand_v2)
+USE_TRT_CONVERTER(take_along_axis)
 #if PADDLE_WITH_CUSPARSELT && IS_TRT_VERSION_GE(8000)
 USE_TRT_CONVERTER(sparse_fc)
 USE_TRT_CONVERTER(sparse_multihead_matmul)

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

@@ -82,6 +82,7 @@ list(
   celu_op.cc
   layernorm_shift_partition_op.cc
   tanhshrink_op.cc
+  take_along_axis_op.cc
   logsigmoid_op.cc
   preln_layernorm_shift_partition_op.cc
   merge_layernorm_op.cc

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

+/* Copyright (c) 2022 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 inference {
+namespace tensorrt {
+
+/*
+ * TakeAlongAxis Op
+ */
+class TakeAlongAxisOpConverter : public OpConverter {
+ public:
+  void operator()(const framework::proto::OpDesc& op,
+                  const framework::Scope& scope,
+                  bool test_mode) override {
+    // AddGatherV2 is supported by the trt version of 8.2.
+#if IS_TRT_VERSION_GE(8200)
+    VLOG(3) << "convert take_along_axis op to tensorrt take_along_axis layer";
+    framework::OpDesc op_desc(op, nullptr);
+    const auto input_tensor = engine_->GetITensor(op_desc.Input("Input")[0]);
+    const auto index_tensor = engine_->GetITensor(op_desc.Input("Index")[0]);
+    auto output_name = op_desc.Output("Result")[0];
+
+    int axis = 0;
+    if (op_desc.HasAttr("Axis")) {
+      axis = PADDLE_GET_CONST(int, op_desc.GetAttr("Axis"));
+    }
+    auto input_dims = input_tensor->getDimensions();
+    int NbDims = input_dims.nbDims;
+    if (axis < 0) axis = axis + NbDims;
+
+    auto layer = TRT_ENGINE_ADD_LAYER(engine_,
+                                      GatherV2,
+                                      *input_tensor,
+                                      *index_tensor,
+                                      nvinfer1::GatherMode::kELEMENT);
+    layer->setGatherAxis(axis);
+
+    RreplenishLayerAndOutput(
+        layer, "take_along_axis", {output_name}, test_mode);
+#endif
+  }
+};
+
+}  // namespace tensorrt
+}  // namespace inference
+}  // namespace paddle
+
+REGISTER_TRT_OP_CONVERTER(take_along_axis, TakeAlongAxisOpConverter);

paddle/fluid/inference/tensorrt/op_teller.cc

@@ -596,6 +596,36 @@ struct SimpleOpTypeSetTeller : public Teller {
 #endif
     }
 
+    if (op_type == "take_along_axis") {
+#if IS_TRT_VERSION_GE(8200)
+      if (!with_dynamic_shape) return false;
+      auto* block = desc.Block();
+      auto input_var_name = desc.Input("Input")[0];
+      auto index_var_name = desc.Input("Index")[0];
+      auto* input_var_desc = block->FindVar(input_var_name);
+      auto* index_var_desc = block->FindVar(index_var_name);
+
+      // The index input must be int32 datatype.
+      if (index_var_desc->GetDataType() !=
+          paddle::framework::proto::VarType_Type::VarType_Type_INT32) {
+        VLOG(3) << "take_along_axis op Index input data type must be int32";
+        return false;
+      }
+
+      const auto input_shape = input_var_desc->GetShape();
+      const auto index_shape = index_var_desc->GetShape();
+      if (input_shape.size() != index_shape.size()) {
+        VLOG(3) << "take_along_axis op Index input dims size ["
+                << index_shape.size() << " ] not equal to input dims size ["
+                << input_shape.size() << "]";
+        return false;
+      }
+#else
+      VLOG(3) << "take_along_axis op is only supported by trt8.2 above ";
+      return false;
+#endif
+    }
+
     if (op_type == "anchor_generator") {
       if (!with_dynamic_shape) return false;
     }
@@ -2399,6 +2429,7 @@ struct SimpleOpTypeSetTeller : public Teller {
       "squeeze2",
       "unsqueeze2",
       "layernorm_shift_partition",
+      "take_along_axis",
       "tanh_shrink",
       "logsigmoid",
       "preln_layernorm_shift_partition",
@@ -2530,6 +2561,7 @@ struct SimpleOpTypeSetTeller : public Teller {
       "fused_token_prune",
       "layernorm_shift_partition",
       "tanh_shrink",
+      "take_along_axis",
       "logsigmoid",
       "preln_layernorm_shift_partition",
       "merge_layernorm",

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

+# Copyright (c) 2022 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 TrtConvertTakeAlongAxisTest(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]['Axis']:
+            return False
+        if len(inputs['input_data'].shape) != len(inputs['index_data'].shape):
+            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.zeros(index).astype(np.int32)
+
+        def generate_input3(axis):
+            return np.array([axis]).astype(np.int32)
+
+        for shape in [[32], [3, 64], [1, 64, 16], [1, 64, 16, 32]]:
+            for index in [[1], [1, 1], [1, 1, 2], [1, 1, 1, 1]]:
+                for axis in [0, 1, 2, 3]:
+                    self.shape = shape
+                    self.axis = axis
+                    dics = [{"Axis": axis}]
+                    ops_config = [
+                        {
+                            "op_type": "take_along_axis",
+                            "op_inputs": {
+                                "Input": ["input_data"],
+                                "Index": ["index_data"],
+                            },
+                            "op_outputs": {"Result": ["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, 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": [3, 64],
+                    "index_data": [1, 1],
+                }
+                self.dynamic_shape.max_input_shape = {
+                    "input_data": [3, 64],
+                    "index_data": [1, 1],
+                }
+                self.dynamic_shape.opt_input_shape = {
+                    "input_data": [3, 64],
+                    "index_data": [1, 1],
+                }
+            elif len(self.shape) == 3:
+                self.dynamic_shape.min_input_shape = {
+                    "input_data": [1, 64, 16],
+                    "index_data": [1, 1, 2],
+                }
+                self.dynamic_shape.max_input_shape = {
+                    "input_data": [1, 64, 16],
+                    "index_data": [1, 1, 2],
+                }
+                self.dynamic_shape.opt_input_shape = {
+                    "input_data": [1, 64, 16],
+                    "index_data": [1, 1, 2],
+                }
+            elif len(self.shape) == 4:
+                self.dynamic_shape.min_input_shape = {
+                    "input_data": [1, 64, 16, 32],
+                    "index_data": [1, 1, 1, 1],
+                }
+                self.dynamic_shape.max_input_shape = {
+                    "input_data": [1, 64, 16, 32],
+                    "index_data": [1, 1, 1, 1],
+                }
+                self.dynamic_shape.opt_input_shape = {
+                    "input_data": [1, 64, 16, 32],
+                    "index_data": [1, 1, 1, 1],
+                }
+
+        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):
+            ver = paddle_infer.get_trt_compile_version()
+            if (
+                ver[0] * 1000 + ver[1] * 100 + ver[2] * 10 > 8200
+                and dynamic_shape
+            ):
+                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-5
+
+        # 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 add_skip_trt_case(self):
+        pass
+
+    def test(self):
+        self.add_skip_trt_case()
+        self.run_test()
+
+
+if __name__ == "__main__":
+    unittest.main()