fix bug of fuse shape ops to generate_shape (#62587)

PaddlePaddle · Mar 11, 2024 · cf96b67 · cf96b67
1 parent 04f5f59
commit cf96b67
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/paddle/cinn/hlir/dialect/operator/transforms/fuse_shape_ops_into_generate_shape_op_pass.cc b/paddle/cinn/hlir/dialect/operator/transforms/fuse_shape_ops_into_generate_shape_op_pass.cc
@@ -26,6 +26,7 @@
 #include "paddle/fluid/pir/dialect/operator/ir/op_attribute.h"
 #include "paddle/fluid/pir/dialect/operator/ir/op_type.h"
 #include "paddle/fluid/pir/dialect/operator/ir/pd_op.h"
+#include "paddle/fluid/pir/transforms/transform_general_functions.h"
 #include "paddle/pir/include/core/builtin_dialect.h"
 #include "paddle/pir/include/dialect/shape/utils/dim_expr.h"
 #include "paddle/pir/include/dialect/shape/utils/shape_analysis.h"
@@ -57,8 +58,8 @@ std::vector<pir::Value> FindSourceDenseTensorOfDimTensor(
         // find input dimension tensor;
         pir::Operation* owner = value.defining_op();
         if (owner == nullptr) return;
-        for (int i = 0; i < owner->num_operands(); ++i) {
-          Visit(owner->operand_source(i));
+        for (auto input_value : pir::GetUsedExternalValue(*owner)) {
+          Visit(input_value);
         }
       };
   const auto& IsDimTensorOrListDimExpr = symbol::Overloaded{