[DimExpr] DimExpr support print

paddle/pir/dialect/shape/utils/dim_expr.cc

@@ -124,4 +124,64 @@ bool DimExpr::operator!=(const DimExpr& other) const {
   return !(*this == other);
 }
 
+namespace {
+
+std::string ToStringImpl(std::int64_t dim_expr) {
+  return std::to_string(dim_expr);
+}
+
+std::string ToStringImpl(const std::string& dim_expr) { return dim_expr; }
+
+std::string ToStringImpl(const Negative<DimExpr>& dim_expr) {
+  return "-" + ToString(dim_expr->data);
+}
+
+std::string ToStringImpl(const Reciprocal<DimExpr>& dim_expr) {
+  return "1 / (" + ToString(dim_expr->data) + ")";
+}
+
+std::string ListDimExprToString(const List<DimExpr>& dim_exprs,
+                                const std::string& delim = ", ") {
+  std::string ret;
+  for (std::size_t i = 0; i < dim_exprs->size(); ++i) {
+    if (i > 0) {
+      ret += delim;
+    }
+    ret += ToString(dim_exprs->at(i));
+  }
+  return ret;
+}
+
+std::string ToStringImpl(const Add<DimExpr>& dim_expr) {
+  return "Add(" + ListDimExprToString(dim_expr.operands, ", ") + ")";
+}
+
+std::string ToStringImpl(const Mul<DimExpr>& dim_expr) {
+  return "Mul(" + ListDimExprToString(dim_expr.operands, ", ") + ")";
+}
+
+std::string ToStringImpl(const Max<DimExpr>& dim_expr) {
+  return "Max(" + ListDimExprToString(dim_expr.operands, ", ") + ")";
+}
+
+std::string ToStringImpl(const Min<DimExpr>& dim_expr) {
+  return "Min(" + ListDimExprToString(dim_expr.operands, ", ") + ")";
+}
+
+std::string ToStringImpl(const Broadcast<DimExpr>& dim_expr) {
+  return "Broadcast(" + ListDimExprToString(dim_expr.operands, ", ") + ")";
+}
+
+}  // namespace
+
+std::string ToString(const DimExpr& dim_expr) {
+  return std::visit([](const auto& impl) { return ToStringImpl(impl); },
+                    dim_expr.variant());
+}
+
+std::ostream& operator<<(std::ostream& stream, const DimExpr& dim_expr) {
+  stream << ToString(dim_expr);
+  return stream;
+}
+
 }  // namespace symbol

paddle/pir/dialect/shape/utils/dim_expr.h

@@ -17,6 +17,7 @@
 #include <functional>
 #include <memory>
 #include <optional>
+#include <ostream>
 #include <string>
 #include <variant>
 #include <vector>
@@ -248,4 +249,8 @@ class ShapeOrData {
 
 using ShapeOrDataDimExprs = ShapeOrData<DimExpr>;
 
+IR_API std::string ToString(const DimExpr& dim_expr);
+
+IR_API std::ostream& operator<<(std::ostream&, const DimExpr& dim_expr);
+
 }  // namespace symbol

test/cpp/pir/shape_dialect/symbol_dim_expr_test.cc

@@ -111,4 +111,17 @@ TEST(DimExpr, equal) {
             builder.Broadcast(DimExpr("S0"), constant1));
 }
 
+TEST(DimExpr, print) {
+  DimExprBuilder builder{nullptr};
+  DimExpr sym0 = DimExpr("S0");
+  DimExpr sym1 = DimExpr("S1");
+  ASSERT_EQ((ToString(sym0 + sym1)), "Add(S0, S1)");
+  ASSERT_EQ((ToString(sym0 - sym1)), "Add(S0, -S1)");
+  ASSERT_EQ((ToString(sym0 * sym1)), "Mul(S0, S1)");
+  ASSERT_EQ((ToString(sym0 / sym1)), "Mul(S0, 1 / (S1))");
+  ASSERT_EQ((ToString(builder.Max(sym0, sym1))), "Max(S0, S1)");
+  ASSERT_EQ((ToString(builder.Min(sym0, sym1))), "Min(S0, S1)");
+  ASSERT_EQ((ToString(builder.Broadcast(sym0, sym1))), "Broadcast(S0, S1)");
+}
+
 }  // namespace symbol::test