Refine print log for symbolic shape (#61607)

zyfncg · web-flow · commit 19d13df72a8a · 2024-02-06T20:17:51.000+08:00
* refine print log for symbolic shape

* fix complie

* fix test complie

* fix test
diff --git a/paddle/fluid/pir/dialect/operator/ir/control_flow_op.cc b/paddle/fluid/pir/dialect/operator/ir/control_flow_op.cc
@@ -173,12 +173,14 @@ void IfOp::Print(pir::IrPrinter &printer) {
   printer.AddIndentation();
   for (auto &item : true_block()) {
     printer.PrintOperation(&item);
+    os << "\n";
   }
   printer.DecreaseIndentation();
   os << printer.indentation() << "} else {\n";
   printer.AddIndentation();
   for (auto &item : false_block()) {
     printer.PrintOperation(&item);
+    os << "\n";
   }
   printer.DecreaseIndentation();
   os << printer.indentation() << "}";
@@ -371,6 +373,7 @@ void WhileOp::Print(pir::IrPrinter &printer) {
   printer.AddIndentation();
   for (auto &item : body()) {
     printer.PrintOperation(&item);
+    os << "\n";
   }
   printer.DecreaseIndentation();
   os << printer.indentation() << "}";
diff --git a/paddle/fluid/pir/transforms/shape_optimization_pass.cc b/paddle/fluid/pir/transforms/shape_optimization_pass.cc
@@ -14,6 +14,7 @@
 
 #include "paddle/fluid/pir/transforms/shape_optimization_pass.h"
 #include "paddle/fluid/pir/dialect/operator/ir/pd_op.h"
+#include "paddle/pir/core/dialect.h"
 #include "paddle/pir/dialect/shape/ir/shape_attribute.h"
 #include "paddle/pir/pass/pass_manager.h"
 #include "paddle/pir/pass/pass_registry.h"
@@ -25,12 +26,10 @@ using PassPipelineRunner =
     std::function<bool(pir::PassManager&, pir::ModuleOp)>;
 
 void PrintProgram(pir::ModuleOp m, std::string mgs) {
-  std::ostringstream print_stream;
-  print_stream << "\n\n";
-  m.program()->Print(print_stream);
-  print_stream << "\n\n";
+  ShapeConstraintIRAnalysis& shape_analysis =
+      ShapeAnalysisManager::Instance().Get(m.program());
   VLOG(3) << "===================== " << mgs << " =====================\n"
-          << print_stream.str();
+          << pir::CustomPrintHelper(*m.program(), shape_analysis.PrintHook());
 }
 
 void DebugPrintOpInfo(
diff --git a/paddle/pir/core/ir_printer.cc b/paddle/pir/core/ir_printer.cc
@@ -168,14 +168,11 @@ void IrPrinter::PrintOperation(Operation* op) {
   if (auto* dialect = op->dialect()) {
     if (auto print_fn = dialect->PrintOperation(op)) {
       print_fn(op, *this);
-      os << newline;
       return;
     }
   }
 
   PrintGeneralOperation(op);
-
-  os << newline;
 }
 
 void IrPrinter::PrintOperationWithNoRegion(Operation* op) {
@@ -221,6 +218,7 @@ void IrPrinter::PrintBlock(const Block& block) {
   AddIndentation();
   for (auto& item : block) {
     PrintOperation(&item);
+    os << "\n";
   }
   DecreaseIndentation();
   os << indentation() << "}\n";
diff --git a/paddle/pir/dialect/shape/ir/shape_dialect.cc b/paddle/pir/dialect/shape/ir/shape_dialect.cc
@@ -29,33 +29,7 @@ void ShapeDialect::initialize() {
 }
 
 void ShapeDialect::PrintAttribute(pir::Attribute attr, std::ostream &os) const {
-  if (attr.isa<SymbolAttribute>()) {
-    SymbolAttribute symbol_attr = attr.dyn_cast<SymbolAttribute>();
-    if (symbol_attr.data().isa<symbol::TensorListShapeOrDataDimExprs>()) return;
-    os << "(shape_data)";
-    os << "[";
-    for (size_t i = 0; i < symbol_attr.data().shape().size(); ++i) {
-      if (i != symbol_attr.data().shape().size() - 1) {
-        os << symbol::ToString(symbol_attr.data().shape()[i]) << ",";
-      } else {
-        os << symbol::ToString(symbol_attr.data().shape()[i]);
-      }
-    }
-    os << "]_[";
-    if (symbol_attr.data().data().has_value()) {
-      for (size_t i = 0; i < symbol_attr.data().data().value().size(); ++i) {
-        if (i != symbol_attr.data().data().value().size() - 1) {
-          os << symbol::ToString(symbol_attr.data().data().value()[i]) << ",";
-        } else {
-          os << symbol::ToString(symbol_attr.data().data().value()[i]);
-        }
-      }
-    } else {
-      os << "nullopt";
-    }
-
-    os << "]";
-  }
+  return;
 }
 
 }  // namespace pir::shape
diff --git a/paddle/pir/dialect/shape/utils/shape_analysis.cc b/paddle/pir/dialect/shape/utils/shape_analysis.cc
@@ -28,14 +28,6 @@ static std::string GetValueId(Value val) {
 
 ShapeConstraintIRAnalysis::ShapeConstraintIRAnalysis(ModuleOp m) : m_(m) {}
 
-ShapeConstraintIRAnalysis::ShapeConstraintIRAnalysis(
-    std::shared_ptr<pir::Program>&& program)
-    : ShapeConstraintIRAnalysis(program->module_op()) {
-  program_ = std::move(program);
-}
-ShapeConstraintIRAnalysis::ShapeConstraintIRAnalysis(pir::IrContext* ctx)
-    : ShapeConstraintIRAnalysis(std::make_shared<pir::Program>(ctx)) {}
-
 void ShapeConstraintIRAnalysis::Init() {
   value_to_shape_or_data_.clear();
   next_sym_idx_ = 0;
@@ -216,6 +208,26 @@ bool ShapeConstraintIRAnalysis::IsSameNumel(Value lhs, Value rhs) const {
                         static_cast<int>(rhs_type.GetRank()));
 }
 
+pir::PrintHooks ShapeConstraintIRAnalysis::PrintHook() const {
+  pir::PrintHooks print_hook;
+  print_hook.op_print_hook = [&](Operation* op, IrPrinter& printer) {
+    printer.IrPrinter::PrintOperation(op);
+    printer.os << " { ";
+    for (uint32_t i = 0; i < op->num_results(); ++i) {
+      if (this->HasShapeOrDataForValue(op->result(i))) {
+        printer.os << "(" << this->GetShapeOrDataForValue(op->result(i)) << ")";
+      } else {
+        printer.os << "()";
+      }
+      if (i < op->num_results() - 1) {
+        printer.os << ", ";
+      }
+    }
+    printer.os << " }";
+  };
+  return print_hook;
+}
+
 ShapeAnalysisManager& ShapeAnalysisManager::Instance() {
   static ShapeAnalysisManager instance;
   return instance;
diff --git a/paddle/pir/dialect/shape/utils/shape_analysis.h b/paddle/pir/dialect/shape/utils/shape_analysis.h
@@ -30,10 +30,6 @@ class IR_API ShapeConstraintIRAnalysis {
  public:
   explicit ShapeConstraintIRAnalysis(ModuleOp m);
 
-  explicit ShapeConstraintIRAnalysis(std::shared_ptr<pir::Program>&& program);
-
-  explicit ShapeConstraintIRAnalysis(pir::IrContext* ctx);
-
   void Init();
 
   const std::string GetNextSymName();
@@ -77,9 +73,10 @@ class IR_API ShapeConstraintIRAnalysis {
   // Returns true if the two value have the same number elements.
   bool IsSameNumel(Value lhs, Value rhs) const;
 
+  pir::PrintHooks PrintHook() const;
+
  private:
   ModuleOp m_;
-  std::shared_ptr<pir::Program> program_;
 
   int64_t next_sym_idx_ = 0;
 
diff --git a/test/cpp/pir/cinn/adt/map_expr_test.cc b/test/cpp/pir/cinn/adt/map_expr_test.cc
@@ -73,7 +73,6 @@ TEST(MapExpr, ElementWise_Fusion_0) {
       value1, builder.Build<paddle::dialect::ExpOp>(value2).result(0));
 
   ::pir::PassManager pass_manager(ctx);
-  auto shape_analysis = std::make_shared<pir::ShapeConstraintIRAnalysis>(ctx);
 
   // TODO(@jiahy0825): use CreateShapeOptimizationPass() instead of
   // CreateInferSymbolicShapePass() which is a fake pass
diff --git a/test/cpp/pir/core/add_dialect_parser_test.cc b/test/cpp/pir/core/add_dialect_parser_test.cc
@@ -102,7 +102,7 @@ TEST(IrParserTest, AddAttribute) {
   std::string op_str =
       "(%0) = \"builtin.parameter\" () "
       "{parameter_name:\"conv2d_0.w_0\",test:(tp.char)a} : () -> "
-      "pd_op.tensor<64x3x7x7xf32>\n";
+      "pd_op.tensor<64x3x7x7xf32>";
   std::stringstream ss;
   ss << op_str;
   pir::IrParser* parser = new pir::IrParser(ctx, ss);
diff --git a/test/cpp/pir/core/ir_op_test.cc b/test/cpp/pir/core/ir_op_test.cc
@@ -64,7 +64,7 @@ TEST(op_test, region_test) {
   // (3) Test custom operation printer
   std::stringstream ss;
   op1->Print(ss);
-  EXPECT_EQ(ss.str(), "(%0) = \"test.operation1\" ()\n");
+  EXPECT_EQ(ss.str(), "(%0) = \"test.operation1\" ()");
 
   region.push_back(new pir::Block());
   region.push_front(new pir::Block());
diff --git a/test/cpp/pir/core/ir_printer_test.cc b/test/cpp/pir/core/ir_printer_test.cc
@@ -61,7 +61,6 @@ TEST(printer_test, custom_hooks) {
     printer.PrintAttributeMap(op);
     printer.os << " :";
     printer.PrintOpReturnType(op);
-    printer.os << "\n";
   };
 
   hooks.attribute_print_hook = [](pir::Attribute attr,

Original file line number	Diff line number	Diff line change
`@@ -168,14 +168,11 @@ void IrPrinter::PrintOperation(Operation* op) {`
`168`	`168`	`if (auto* dialect = op->dialect()) {`
`169`	`169`	`if (auto print_fn = dialect->PrintOperation(op)) {`
`170`	`170`	`print_fn(op, *this);`
`171`		`- os << newline;`
`172`	`171`	`return;`
`173`	`172`	`}`
`174`	`173`	`}`
`175`	`174`
`176`	`175`	`PrintGeneralOperation(op);`
`177`		`-`
`178`		`- os << newline;`
`179`	`176`	`}`
`180`	`177`
`181`	`178`	`void IrPrinter::PrintOperationWithNoRegion(Operation* op) {`
`@@ -221,6 +218,7 @@ void IrPrinter::PrintBlock(const Block& block) {`
`221`	`218`	`AddIndentation();`
`222`	`219`	`for (auto& item : block) {`
`223`	`220`	`PrintOperation(&item);`
	`221`	`+ os << "\n";`
`224`	`222`	`}`
`225`	`223`	`DecreaseIndentation();`
`226`	`224`	`os << indentation() << "}\n";`