PaddlePaddle
diff --git a/‎paddle/fluid/framework/new_executor/collect_shape_manager.cc
Copy file name to clipboard
+44-25Lines changed: 44 additions & 25 deletions b/‎paddle/fluid/framework/new_executor/collect_shape_manager.cc
Copy file name to clipboard
+44-25Lines changed: 44 additions & 25 deletions
diff --git a/‎paddle/fluid/pir/transforms/tensorrt/trt_op_marker_pass.cc
Copy file name to clipboard
+30Lines changed: 30 additions & 0 deletions b/‎paddle/fluid/pir/transforms/tensorrt/trt_op_marker_pass.cc
Copy file name to clipboard
+30Lines changed: 30 additions & 0 deletions
diff --git a/‎python/paddle/tensorrt/converter.py
Copy file name to clipboard
-1Lines changed: 0 additions & 1 deletion b/‎python/paddle/tensorrt/converter.py
Copy file name to clipboard
-1Lines changed: 0 additions & 1 deletion
diff --git a/‎python/paddle/tensorrt/export.py
Copy file name to clipboard
+4-1Lines changed: 4 additions & 1 deletion b/‎python/paddle/tensorrt/export.py
Copy file name to clipboard
+4-1Lines changed: 4 additions & 1 deletion
diff --git a/‎python/paddle/tensorrt/impls/manipulation.py
Copy file name to clipboard
+10Lines changed: 10 additions & 0 deletions b/‎python/paddle/tensorrt/impls/manipulation.py
Copy file name to clipboard
+10Lines changed: 10 additions & 0 deletions
diff --git a/‎python/paddle/tensorrt/impls/others.py
Copy file name to clipboard
+51Lines changed: 51 additions & 0 deletions b/‎python/paddle/tensorrt/impls/others.py
Copy file name to clipboard
+51Lines changed: 51 additions & 0 deletions
diff --git a/‎python/paddle/tensorrt/util.py
Copy file name to clipboard
+8-1Lines changed: 8 additions & 1 deletion b/‎python/paddle/tensorrt/util.py
Copy file name to clipboard
+8-1Lines changed: 8 additions & 1 deletion
diff --git a/‎test/tensorrt/tensorrt_test_base.py
Copy file name to clipboard
+25-2Lines changed: 25 additions & 2 deletions b/‎test/tensorrt/tensorrt_test_base.py
Copy file name to clipboard
+25-2Lines changed: 25 additions & 2 deletions
@@ -34,7 +34,9 @@ void CollectShapeManager::CollectShapeInfo(
     auto *var = scope->FindVar(var_name);
     if (!var || !var->IsType<phi::DenseTensor>()) continue;
     auto tensor = var->Get<phi::DenseTensor>();
-    if (!tensor.initialized()) continue;
+    if (!tensor.initialized() && !instr->NoNeedBuffer().count(input.first)) {
+      continue;
+    }
     paddle::platform::DeviceContextPool &pool =
         paddle::platform::DeviceContextPool::Instance();
 #if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
@@ -124,36 +126,53 @@ void CollectShapeManager::StatisticShapeRangeInfo() {
         for (auto const &it : shape_data) {
           auto val = it.first;
           auto shapes = it.second;
+
           std::vector<int32_t> min_shape(shapes[0].begin(), shapes[0].end());
           std::vector<int32_t> max_shape(shapes[0].begin(), shapes[0].end());
           std::vector<int32_t> opt_shape(shapes[0].begin(), shapes[0].end());
+          // Applicable to scenarios where min/opt/max are specified;
+          if (shapes.size() == 3) {
+            for (size_t d = 0; d < shapes[0].size(); ++d) {
+              std::vector<int32_t> dim_values;
+              for (const auto &shape : shapes) {
+                dim_values.push_back(shape[d]);
+              }
+              std::sort(dim_values.begin(), dim_values.end());
+              min_shape[d] = dim_values[0];
+              opt_shape[d] = dim_values[1];
+              max_shape[d] = dim_values[2];
+            }
+            min_data[val] = min_shape;
+            max_data[val] = max_shape;
+            opt_data[val] = opt_shape;
+          } else {
+            // suitable for scenarios where shape is automatically collected.
+            auto ShapeMaxFreq =
+                [](const std::map<int32_t, int32_t> &m) -> int32_t {
+              std::vector<std::pair<int32_t, int32_t>> counter;
+              for (auto &it : m) counter.emplace_back(it);
+              std::sort(counter.begin(),
+                        counter.end(),
+                        [](std::pair<int32_t, int32_t> &a,
+                           std::pair<int32_t, int32_t> &b) {
+                          return a.second > b.second;
+                        });
+              return counter[0].first;
+            };
 
-          auto ShapeMaxFreq =
-              [](const std::map<int32_t, int32_t> &m) -> int32_t {
-            std::vector<std::pair<int32_t, int32_t>> counter;
-            for (auto &it : m) counter.emplace_back(it);
-            std::sort(counter.begin(),
-                      counter.end(),
-                      [](std::pair<int32_t, int32_t> &a,
-                         std::pair<int32_t, int32_t> &b) {
-                        return a.second > b.second;
-                      });
-            return counter[0].first;
-          };
-
-          for (size_t d = 0; d < shapes[0].size(); ++d) {
-            std::map<int32_t, int32_t> counter;
-            for (auto &shape : shapes) {
-              counter[shape[d]] += 1;
-              if (shape[d] < min_shape[d]) min_shape[d] = shape[d];
-              if (shape[d] > max_shape[d]) max_shape[d] = shape[d];
+            for (size_t d = 0; d < shapes[0].size(); ++d) {
+              std::map<int32_t, int32_t> counter;
+              for (auto &shape : shapes) {
+                counter[shape[d]] += 1;
+                if (shape[d] < min_shape[d]) min_shape[d] = shape[d];
+                if (shape[d] > max_shape[d]) max_shape[d] = shape[d];
+              }
+              opt_shape[d] = ShapeMaxFreq(counter);
             }
-            opt_shape[d] = ShapeMaxFreq(counter);
+            min_data[val] = min_shape;
+            max_data[val] = max_shape;
+            opt_data[val] = opt_shape;
           }
-
-          min_data[val] = min_shape;
-          max_data[val] = max_shape;
-          opt_data[val] = opt_shape;
         }
       };
   extract_min_max_opt(min_shapes_, max_shapes_, opt_shapes_, shape_info_);
 
@@ -94,6 +94,7 @@ DEFINE_GENERAL_PATTERN(Flip, paddle::dialect::FlipOp)
 DEFINE_GENERAL_PATTERN(Mish, paddle::dialect::MishOp)
 DEFINE_GENERAL_PATTERN(AssignValue, paddle::dialect::AssignValueOp)
 DEFINE_GENERAL_PATTERN(AssignValue_, paddle::dialect::AssignValue_Op)
+DEFINE_GENERAL_PATTERN(Numel, paddle::dialect::NumelOp)
 #undef DEFINE_GENERAL_PATTERN
 
 // Add ReduceCommonOpPattern base class to simplify code
@@ -2191,6 +2192,33 @@ class InstanceNormOpPattern
   }
 };
 
+class AffineChannelOpPattern
+    : public pir::OpRewritePattern<paddle::dialect::AffineChannelOp> {
+ public:
+  using pir::OpRewritePattern<
+      paddle::dialect::AffineChannelOp>::OpRewritePattern;
+  bool MatchAndRewrite(paddle::dialect::AffineChannelOp op,
+                       pir::PatternRewriter &rewriter) const override {
+    if (op->HasAttribute(kCanRunTrtAttr) &&
+        op->attribute<pir::BoolAttribute>(kCanRunTrtAttr).data()) {
+      return false;
+    }
+    if (!op->HasAttribute("data_layout")) {
+      VLOG(3) << "pd_op.affine_channel must has data_layout";
+      return false;
+    }
+    pir::Value x = op.operand_source(0);
+    auto x_shape = pir::GetShapeFromValue(x);
+    if (x_shape.size() == 2) {
+      VLOG(3) << "pd_op.affine_channel x.shape can not be 2";
+      return false;
+    }
+
+    op->set_attribute(kCanRunTrtAttr, rewriter.bool_attr(true));
+    return true;
+  }
+};
+
 class TrtOpMarkerPass : public pir::PatternRewritePass {
  public:
   TrtOpMarkerPass() : pir::PatternRewritePass("trt_op_marker_pass", 2) {}
@@ -2245,6 +2273,7 @@ class TrtOpMarkerPass : public pir::PatternRewritePass {
     ADD_PATTERN(Mish)
     ADD_PATTERN(AssignValue)
     ADD_PATTERN(AssignValue_)
+    ADD_PATTERN(Numel)
 #if IS_TRT_VERSION_GE(8600)
     ADD_PATTERN(Layer_norm)
 #endif
@@ -2321,6 +2350,7 @@ class TrtOpMarkerPass : public pir::PatternRewritePass {
     ps.Add(std::make_unique<CeluOpPattern>(context));
     ps.Add(std::make_unique<OneHotOpPattern>(context));
     ps.Add(std::make_unique<InstanceNormOpPattern>(context));
+    ps.Add(std::make_unique<AffineChannelOpPattern>(context));
     return ps;
   }
 };
 
@@ -294,7 +294,6 @@ def convert_subgraph_to_trt(self, program, group_op):
                     max_shape = get_value_shape_range_info(
                         value, False, paddle.base.core.ShapeMode.kMAX
                     )
-
                     if trt_input.is_shape_tensor:
                         min_value = get_value_shape_range_info(
                             value, True, paddle.base.core.ShapeMode.kMIN
 
@@ -248,17 +248,20 @@ def convert_to_trt(program, trt_config, scope):
     with paddle.pir_utils.IrGuard():
         min_shape_feed = {}
         max_shape_feed = {}
+        opt_shape_feed = {}
         for i, input_instance in enumerate(trt_config.inputs):
             # get fake inputs
-            min_data, _, max_data = input_instance.generate_input_data()
+            min_data, opt_data, max_data = input_instance.generate_input_data()
             program_with_output = program.list_vars()[-1]
             min_shape_feed[feed_name[i]] = min_data
+            opt_shape_feed[feed_name[i]] = opt_data
             max_shape_feed[feed_name[i]] = max_data
 
             # run warmup for collecting shape
         program = warmup_shape_infer(
             program,
             min_shape_feed=min_shape_feed,
+            opt_shape_feed=opt_shape_feed,
             max_shape_feed=max_shape_feed,
             scope=scope,
         )
 
@@ -947,3 +947,13 @@ def roll_converter(network, paddle_op, inputs):
             )
 
     return layer.get_output(0)
+
+
+@converter_registry.register("pd_op.numel", trt_version="8.x")
+def numel_converter(network, paddle_op, inputs):
+    input_tensor = inputs[0]
+    shape_tensor = network.add_shape(input_tensor).get_output(0)
+    layer = network.add_reduce(
+        shape_tensor, trt.ReduceOperation.PROD, axes=1, keep_dims=False
+    )
+    return layer.get_output(0)
@@ -301,3 +301,54 @@ def share_data_converter(network, paddle_op, inputs):
     identity_layer = network.add_identity(x)
 
     return identity_layer.get_output(0)
+
+
+@converter_registry.register("pd_op.affine_channel", trt_version="8.x")
+def affine_channel_converter(network, paddle_op, inputs):
+    x, scale_weights, bias_weights = inputs
+    data_layout = paddle_op.attrs().get("data_layout")
+
+    if data_layout == "NCHW":
+        channel_axis = 1
+        x_input = x
+    elif data_layout == "NHWC":
+        # Permute NHWC to NCHW
+        shuffle_layer1 = network.add_shuffle(x)
+        shuffle_layer1.first_transpose = (0, 3, 1, 2)
+        x_input = shuffle_layer1.get_output(0)
+        channel_axis = 1
+    else:
+        raise ValueError(f"affine_channel: Unsupported layout: {data_layout}")
+
+    if not isinstance(scale_weights, trt.Weights):
+        raise TypeError("affine_channel requires scale as trt.Weights")
+    if not isinstance(bias_weights, trt.Weights):
+        raise TypeError("affine_channel requires bias as trt.Weights")
+
+    if scale_weights.size != bias_weights.size:
+        raise ValueError(
+            f"affine_channel: scale.size({scale_weights.size}) != bias.size({bias_weights.size})"
+        )
+
+    power_array = np.ones((scale_weights.size,), dtype=np.float32)
+    power_weights = trt.Weights(power_array)
+
+    layer = network.add_scale_nd(
+        input=x_input,
+        mode=trt.ScaleMode.CHANNEL,
+        shift=bias_weights,
+        scale=scale_weights,
+        power=power_weights,
+        channel_axis=channel_axis,
+    )
+    if not layer:
+        raise RuntimeError("affine_channel: add_scale_nd failed.")
+
+    out_tensor = layer.get_output(0)
+
+    if data_layout == "NHWC":
+        shuffle_layer2 = network.add_shuffle(out_tensor)
+        shuffle_layer2.first_transpose = (0, 2, 3, 1)
+        out_tensor = shuffle_layer2.get_output(0)
+
+    return out_tensor
@@ -94,7 +94,9 @@ def predict_program(program, feed_data, fetch_var_list, scope=None):
             return output
 
 
-def warmup_shape_infer(program, min_shape_feed, max_shape_feed, scope=None):
+def warmup_shape_infer(
+    program, min_shape_feed, opt_shape_feed, max_shape_feed, scope=None
+):
     paddle.framework.set_flags({"FLAGS_enable_collect_shape": True})
     with paddle.pir_utils.IrGuard():
         with paddle.static.program_guard(program):
@@ -103,6 +105,9 @@ def warmup_shape_infer(program, min_shape_feed, max_shape_feed, scope=None):
             for _ in range(1):
                 executor.run(program, feed=min_shape_feed, scope=scope)
 
+            for _ in range(1):
+                executor.run(program, feed=opt_shape_feed, scope=scope)
+
             # Run the program with input_data_max_shape (fake max_shape input)
             for _ in range(1):
                 executor.run(program, feed=max_shape_feed, scope=scope)
@@ -120,6 +125,7 @@ def warmup_shape_infer(program, min_shape_feed, max_shape_feed, scope=None):
                 )
             )
     paddle.framework.set_flags({"FLAGS_enable_collect_shape": False})
+
     return exe_program
 
 
@@ -192,6 +198,7 @@ def weight_to_tensor(network, paddle_value, trt_tensor, use_op_name):
         "pd_op.batch_norm_",
         "pd_op.layer_norm",
         "pd_op.depthwise_conv2d_transpose",
+        "pd_op.affine_channel",
     ]
     if use_op_name in forbid_cast_op:
         return trt_tensor
 
@@ -39,6 +39,7 @@ def __init__(self, methodName='runTest'):
         self.api_args = None
         self.program_config = None
         self.min_shape = None
+        self.opt_shape = None
         self.max_shape = None
         self.target_marker_op = ""
         self.dynamic_shape_data = {}
@@ -62,6 +63,7 @@ def create_fake_program(self):
                     ].items():
                         if (
                             feed_name in self.min_shape.keys()
+                            and feed_name in self.opt_shape.keys()
                             and feed_name in self.max_shape.keys()
                         ):
                             input_shape_without_dynamic_dim = (
@@ -89,11 +91,15 @@ def create_fake_program(self):
                     api_args[feed_name] = new_list_args
                 else:
                     empty_min_max_shape = (
-                        self.min_shape is None or self.max_shape is None
+                        self.min_shape is None
+                        or self.max_shape is None
+                        or self.opt_shape is None
                     )
+
                     if (
                         not empty_min_max_shape
                         and feed_name in self.min_shape.keys()
+                        and feed_name in self.opt_shape.keys()
                         and feed_name in self.max_shape.keys()
                     ):
                         # dynamic shape condition
@@ -181,6 +187,7 @@ def check_trt_result(self, rtol=1e-4, atol=1e-4, precision_mode="fp32"):
             output_expected = self.run_program(main_program, fetch_list)
 
             min_shape_data = dict()  # noqa: C408
+            opt_shape_data = dict()  # noqa: C408
             max_shape_data = dict()  # noqa: C408
             for feed_name in self.program_config["feed_list"]:
                 if self.api_args[feed_name] is None:
@@ -190,11 +197,13 @@ def check_trt_result(self, rtol=1e-4, atol=1e-4, precision_mode="fp32"):
                     if (
                         feed_name not in self.min_shape.keys()
                         and feed_name not in self.max_shape.keys()
+                        and feed_name not in self.opt_shape.keys()
                     ):
                         for sub_feed_name, sub_feed_value in self.api_args[
                             feed_name
                         ].items():
                             min_shape_data[sub_feed_name] = sub_feed_value
+                            opt_shape_data[sub_feed_name] = sub_feed_value
                             max_shape_data[sub_feed_name] = sub_feed_value
                             continue
                     else:
@@ -206,6 +215,11 @@ def check_trt_result(self, rtol=1e-4, atol=1e-4, precision_mode="fp32"):
                             ).astype(
                                 self.api_args[feed_name][sub_feed_name].dtype
                             )
+                            opt_shape_data[sub_feed_name] = np.random.randn(
+                                *self.opt_shape[feed_name][i]
+                            ).astype(
+                                self.api_args[feed_name][sub_feed_name].dtype
+                            )
                             max_shape_data[sub_feed_name] = np.random.randn(
                                 *self.max_shape[feed_name][i]
                             ).astype(
@@ -216,29 +230,38 @@ def check_trt_result(self, rtol=1e-4, atol=1e-4, precision_mode="fp32"):
                     if (
                         feed_name not in self.min_shape.keys()
                         and feed_name not in self.max_shape.keys()
+                        and feed_name not in self.opt_shape.keys()
                     ):
                         min_shape_data[feed_name] = self.api_args[feed_name]
+                        opt_shape_data[feed_name] = self.api_args[feed_name]
                         max_shape_data[feed_name] = self.api_args[feed_name]
                         continue
                     else:
                         if self.dynamic_shape_data:
                             min_shape_data[feed_name] = self.dynamic_shape_data[
                                 feed_name
                             ](self.min_shape[feed_name])
+                            opt_shape_data[feed_name] = self.dynamic_shape_data[
+                                feed_name
+                            ](self.opt_shape[feed_name])
                             max_shape_data[feed_name] = self.dynamic_shape_data[
                                 feed_name
                             ](self.max_shape[feed_name])
                         else:
                             min_shape_data[feed_name] = np.random.randn(
                                 *self.min_shape[feed_name]
                             ).astype(self.api_args[feed_name].dtype)
+                            opt_shape_data[feed_name] = np.random.randn(
+                                *self.opt_shape[feed_name]
+                            ).astype(self.api_args[feed_name].dtype)
                             max_shape_data[feed_name] = np.random.randn(
                                 *self.max_shape[feed_name]
                             ).astype(self.api_args[feed_name].dtype)
             scope = paddle.static.global_scope()
             main_program = warmup_shape_infer(
                 main_program,
                 min_shape_feed=min_shape_data,
+                opt_shape_feed=opt_shape_data,
                 max_shape_feed=max_shape_data,
                 scope=scope,
             )
@@ -262,7 +285,7 @@ def check_trt_result(self, rtol=1e-4, atol=1e-4, precision_mode="fp32"):
 
             input = Input(
                 min_input_shape=self.min_shape,
-                optim_input_shape=self.min_shape,
+                optim_input_shape=self.opt_shape,
                 max_input_shape=self.max_shape,
             )
             trt_config = TensorRTConfig(inputs=[input])
Original file line number	Diff line number	Diff line change
`@@ -294,7 +294,6 @@ def convert_subgraph_to_trt(self, program, group_op):`
`294`	`294`	`max_shape = get_value_shape_range_info(`
`295`	`295`	`value, False, paddle.base.core.ShapeMode.kMAX`
`296`	`296`	`)`
`297`		`-`
`298`	`297`	`if trt_input.is_shape_tensor:`
`299`	`298`	`min_value = get_value_shape_range_info(`
`300`	`299`	`value, True, paddle.base.core.ShapeMode.kMIN`