diff --git a/paddle/fluid/inference/tensorrt/op_teller.cc b/paddle/fluid/inference/tensorrt/op_teller.cc index a6eeffbb877d537d76ef15a4469337be51273aa6..5d056bb0da86b55ddf66d24ea843929eea53c9a0 100644 --- a/paddle/fluid/inference/tensorrt/op_teller.cc +++ b/paddle/fluid/inference/tensorrt/op_teller.cc @@ -534,6 +534,16 @@ struct SimpleOpTypeSetTeller : public Teller { "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 datatype. + if (index_var_desc->GetDataType() != + paddle::framework::proto::VarType_Type::VarType_Type_INT32) { + VLOG(3) << "gather op Index input data type must be int32"; + return false; + } #if !IS_TRT_VERSION_GE(7000) auto* x_var_desc = block->FindVar(desc.Input("X")[0]); const auto x_shape = x_var_desc->GetShape(); diff --git a/python/paddle/fluid/tests/unittests/ir/inference/test_trt_convert_gather.py b/python/paddle/fluid/tests/unittests/ir/inference/test_trt_convert_gather.py index 5405f114651061e88f76ef731f6412388f77fe39..25d0d48c8c3bf3dcc94efda97b297112475d2585 100644 --- a/python/paddle/fluid/tests/unittests/ir/inference/test_trt_convert_gather.py +++ b/python/paddle/fluid/tests/unittests/ir/inference/test_trt_convert_gather.py @@ -42,6 +42,9 @@ class TrtConvertGatherTest(TrtLayerAutoScanTest): 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) @@ -57,44 +60,48 @@ class TrtConvertGatherTest(TrtLayerAutoScanTest): "Index": ["index_data"], "Axis": ["axis_data"] }]: - self.shape = shape - self.axis = axis - self.input_num = len(input) - dics = [{"overwrite": overwrite, "axis": axis}] - ops_config = [{ - "op_type": "gather", - "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, index)), - } if len(input) == 2 else { - "input_data": - TensorConfig(data_gen=partial( - generate_input1, shape)), - "index_data": - TensorConfig(data_gen=partial( - generate_input2, index)), - "axis_data": - TensorConfig(data_gen=partial( - generate_input3, axis)), - }, - outputs=["output_data"]) - - yield program_config + 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 = [{"overwrite": overwrite, "axis": axis}] + ops_config = [{ + "op_type": "gather", + "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 == + True else generate_input4, index)), + } if len(input) == 2 else { + "input_data": + TensorConfig(data_gen=partial( + generate_input1, shape)), + "index_data": + TensorConfig(data_gen=partial( + generate_input2, index)), + "axis_data": + TensorConfig(data_gen=partial( + generate_input3, axis)), + }, + outputs=["output_data"]) + + yield program_config def sample_predictor_configs( self, program_config) -> (paddle_infer.Config, List[int], float): @@ -162,7 +169,7 @@ class TrtConvertGatherTest(TrtLayerAutoScanTest): if self.input_num == 3: return 0, 5 else: - if dynamic_shape: + if dynamic_shape and self.index_type_int32 == True: return 1, 3 else: return 0, 4