[PIR] move pir::DenseTensorType registration from OperatorDialect to BuiltinDialect (#62491)

huangjiyi · web-flow · commit a726f8253ac0 · 2024-03-07T14:58:36.000+08:00
diff --git a/paddle/fluid/pir/dialect/operator/ir/op_dialect.cc b/paddle/fluid/pir/dialect/operator/ir/op_dialect.cc
@@ -205,15 +205,7 @@ OperatorDialect::OperatorDialect(pir::IrContext* ctx)
 void PrintTypeImpl(pir::Type type, std::ostream& os) {
   os << type.dialect().name();
   os << '.';
-  if (auto tensor_type = type.dyn_cast<DenseTensorType>()) {
-    os << "tensor<";
-    for (auto d : common::vectorize(tensor_type.dims())) {
-      os << d;
-      os << "x";
-    }
-    tensor_type.dtype().Print(os);
-    os << ">";
-  } else if (auto selected_rows_type = type.dyn_cast<SelectedRowsType>()) {
+  if (auto selected_rows_type = type.dyn_cast<SelectedRowsType>()) {
     os << "selectedrows<";
     for (auto d : common::vectorize(selected_rows_type.dims())) {
       os << d;
@@ -266,8 +258,7 @@ void PrintOperationImpl(pir::Operation* op,
 }
 
 void OperatorDialect::initialize() {
-  RegisterTypes<paddle::dialect::DenseTensorType,
-                paddle::dialect::SelectedRowsType,
+  RegisterTypes<paddle::dialect::SelectedRowsType,
                 paddle::dialect::DenseTensorArrayType>();
 
   RegisterAttributes<paddle::dialect::IntArrayAttribute,
@@ -328,35 +319,6 @@ void OperatorDialect::PrintAttribute(pir::Attribute attr,
   PrintAttributeImpl(attr, os);
 }
 
-pir::Type OperatorDialect::ParseType(pir::IrParser& parser) {  // NOLINT
-  parser.ConsumeAToken("pd_op.tensor");
-  parser.ConsumeAToken("<");
-  std::vector<int> dim{};
-  Token dim_token = parser.PeekToken();
-  while (dim_token.token_type_ == DIGIT) {
-    dim_token = parser.ConsumeToken();
-    dim.push_back(atoi(dim_token.val_.c_str()));
-    std::string peek_token_val = parser.PeekToken().val_;
-    if (peek_token_val[0] != 'x') {
-      break;
-    }
-    parser.ConsumeToken();
-    parser.lexer->Unget(static_cast<int>(peek_token_val.size() - 1));
-    if (parser.PeekToken().token_type_ != DIGIT) {
-      break;
-    }
-  }
-  phi::DDim ddim = common::make_ddim(dim);
-  pir::Type dtype = parser.ParseType();
-  std::vector<std::vector<size_t>> lod;
-  std::vector<size_t> lodv;
-  lodv.push_back(0);
-  lod.push_back(lodv);
-  parser.ConsumeAToken(">");
-  return DenseTensorType::get(
-      parser.ctx, dtype, ddim, phi::DataLayout::UNDEFINED, lod, 0);
-}
-
 pir::Attribute OperatorDialect::ParseAttribute(
     pir::IrParser& parser) {  // NOLINT
   std::string type_name = parser.ConsumeToken().val_;
diff --git a/paddle/fluid/pir/dialect/operator/ir/op_dialect.h b/paddle/fluid/pir/dialect/operator/ir/op_dialect.h
@@ -29,7 +29,6 @@ class TEST_API OperatorDialect : public pir::Dialect {
 
   static const char* name() { return "pd_op"; }
 
-  pir::Type ParseType(pir::IrParser& parser) override;            // NOLINT
   pir::Attribute ParseAttribute(pir::IrParser& parser) override;  // NOLINT
 
   void PrintType(pir::Type type, std::ostream& os) const override;
diff --git a/paddle/fluid/pir/dialect/operator/ir/op_onednn_dialect.cc b/paddle/fluid/pir/dialect/operator/ir/op_onednn_dialect.cc
@@ -68,15 +68,7 @@ void OneDNNOperatorDialect::initialize() {
 void OneDNNOperatorDialect::PrintType(pir::Type type, std::ostream &os) const {
   os << type.dialect().name();
   os << '.';
-  if (auto tensor_type = type.dyn_cast<DenseTensorType>()) {
-    os << "tensor<";
-    for (auto d : common::vectorize(tensor_type.dims())) {
-      os << d;
-      os << "x";
-    }
-    tensor_type.dtype().Print(os);
-    os << ">";
-  } else if (auto selected_rows_type = type.dyn_cast<SelectedRowsType>()) {
+  if (auto selected_rows_type = type.dyn_cast<SelectedRowsType>()) {
     os << "selectedrows<";
     for (auto d : common::vectorize(selected_rows_type.dims())) {
       os << d;
@@ -117,35 +109,6 @@ void OneDNNOperatorDialect::PrintAttribute(pir::Attribute attr,
   }
 }
 
-pir::Type OneDNNOperatorDialect::ParseType(pir::IrParser &parser) {  // NOLINT
-  parser.ConsumeAToken("pd_op.tensor");
-  parser.ConsumeAToken("<");
-  std::vector<int> dim{};
-  Token dim_token = parser.PeekToken();
-  while (dim_token.token_type_ == DIGIT) {
-    dim_token = parser.ConsumeToken();
-    dim.push_back(atoi(dim_token.val_.c_str()));
-    std::string peek_token_val = parser.PeekToken().val_;
-    if (peek_token_val[0] != 'x') {
-      break;
-    }
-    parser.ConsumeToken();
-    parser.lexer->Unget(static_cast<int>(peek_token_val.size() - 1));
-    if (parser.PeekToken().token_type_ != DIGIT) {
-      break;
-    }
-  }
-  phi::DDim ddim = common::make_ddim(dim);
-  pir::Type dtype = parser.ParseType();
-  std::vector<std::vector<size_t>> lod;
-  std::vector<size_t> lodv;
-  lodv.push_back(0);
-  lod.push_back(lodv);
-  parser.ConsumeAToken(">");
-  return DenseTensorType::get(
-      parser.ctx, dtype, ddim, phi::DataLayout::UNDEFINED, lod, 0);
-}
-
 pir::Attribute OneDNNOperatorDialect::ParseAttribute(
     pir::IrParser &parser) {  // NOLINT
   std::string type_name = parser.ConsumeToken().val_;
diff --git a/paddle/fluid/pir/dialect/operator/ir/op_onednn_dialect.h b/paddle/fluid/pir/dialect/operator/ir/op_onednn_dialect.h
@@ -25,7 +25,6 @@ class OneDNNOperatorDialect : public pir::Dialect {
 
   static const char* name() { return "onednn_op"; }
 
-  pir::Type ParseType(pir::IrParser& parser) override;            // NOLINT
   pir::Attribute ParseAttribute(pir::IrParser& parser) override;  // NOLINT
 
   void PrintType(pir::Type type, std::ostream& os) const override;
diff --git a/paddle/fluid/pybind/pir.cc b/paddle/fluid/pybind/pir.cc
@@ -1543,10 +1543,10 @@ void BindUtils(pybind11::module *m) {
 
                 >>> print(pir_program)
                 {
-                 (%0) = "pd_op.data" () {dtype:(pd_op.DataType)float32,is_persistable:[false],name:"x",place:(pd_op.Place)Place(undefined:0),shape:(pd_op.IntArray)[4,4],stop_gradient:[false]} : () -> pd_op.tensor<4x4xf32>
-                 (%1) = "pd_op.matmul" (%0, %0) {is_persistable:[false],stop_gradient:[false],transpose_x:false,transpose_y:false} : (pd_op.tensor<4x4xf32>, pd_op.tensor<4x4xf32>) -> pd_op.tensor<4x4xf32>
-                 (%2) = "pd_op.add" (%1, %1) {is_persistable:[false],stop_gradient:[false]} : (pd_op.tensor<4x4xf32>, pd_op.tensor<4x4xf32>) -> pd_op.tensor<4x4xf32>
-                 (%3) = "pd_op.tanh" (%2) {is_persistable:[false],stop_gradient:[false]} : (pd_op.tensor<4x4xf32>) -> pd_op.tensor<4x4xf32>
+                 (%0) = "pd_op.data" () {dtype:(pd_op.DataType)float32,is_persistable:[false],name:"x",place:(pd_op.Place)Place(undefined:0),shape:(pd_op.IntArray)[4,4],stop_gradient:[false]} : () -> builtin.tensor<4x4xf32>
+                 (%1) = "pd_op.matmul" (%0, %0) {is_persistable:[false],stop_gradient:[false],transpose_x:false,transpose_y:false} : (builtin.tensor<4x4xf32>, builtin.tensor<4x4xf32>) -> builtin.tensor<4x4xf32>
+                 (%2) = "pd_op.add" (%1, %1) {is_persistable:[false],stop_gradient:[false]} : (builtin.tensor<4x4xf32>, builtin.tensor<4x4xf32>) -> builtin.tensor<4x4xf32>
+                 (%3) = "pd_op.tanh" (%2) {is_persistable:[false],stop_gradient:[false]} : (builtin.tensor<4x4xf32>) -> builtin.tensor<4x4xf32>
                 }
 
 
@@ -1618,14 +1618,14 @@ void BindUtils(pybind11::module *m) {
 
                 >>> print(pir_program)
                 {
-                 (%0) = "pd_op.data" () {dtype:(pd_op.DataType)float32,is_persistable:[false],name:"x",place:(pd_op.Place)Place(undefined:0),shape:(pd_op.IntArray)[4,4],stop_gradient:[false]} : () -> pd_op.tensor<4x4xf32>
-                 (%1) = "pd_op.matmul" (%0, %0) {is_persistable:[false],stop_gradient:[false],transpose_x:false,transpose_y:false} : (pd_op.tensor<4x4xf32>, pd_op.tensor<4x4xf32>) -> pd_op.tensor<4x4xf32>
-                 (%2) = "pd_op.add" (%1, %1) {is_persistable:[false],stop_gradient:[false]} : (pd_op.tensor<4x4xf32>, pd_op.tensor<4x4xf32>) -> pd_op.tensor<4x4xf32>
-                 (%3) = "pd_op.tanh" (%2) {is_persistable:[false],stop_gradient:[false]} : (pd_op.tensor<4x4xf32>) -> pd_op.tensor<4x4xf32>
+                 (%0) = "pd_op.data" () {dtype:(pd_op.DataType)float32,is_persistable:[false],name:"x",place:(pd_op.Place)Place(undefined:0),shape:(pd_op.IntArray)[4,4],stop_gradient:[false]} : () -> builtin.tensor<4x4xf32>
+                 (%1) = "pd_op.matmul" (%0, %0) {is_persistable:[false],stop_gradient:[false],transpose_x:false,transpose_y:false} : (builtin.tensor<4x4xf32>, builtin.tensor<4x4xf32>) -> builtin.tensor<4x4xf32>
+                 (%2) = "pd_op.add" (%1, %1) {is_persistable:[false],stop_gradient:[false]} : (builtin.tensor<4x4xf32>, builtin.tensor<4x4xf32>) -> builtin.tensor<4x4xf32>
+                 (%3) = "pd_op.tanh" (%2) {is_persistable:[false],stop_gradient:[false]} : (builtin.tensor<4x4xf32>) -> builtin.tensor<4x4xf32>
                 }
 
                 >>> print(mappings)
-                {'matmul_v2_0.tmp_0': [Value(define_op_name=pd_op.matmul, index=0, dtype=pd_op.tensor<4x4xf32>)], 'x': [Value(define_op_name=pd_op.data, index=0, dtype=pd_op.tensor<4x4xf32>)], 'tanh_0.tmp_0': [Value(define_op_name=pd_op.tanh, index=0, dtype=pd_op.tensor<4x4xf32>)], 'elementwise_add_0': [Value(define_op_name=pd_op.add, index=0, dtype=pd_op.tensor<4x4xf32>)]}
+                {'matmul_v2_0.tmp_0': [Value(define_op_name=pd_op.matmul, index=0, dtype=builtin.tensor<4x4xf32>)], 'x': [Value(define_op_name=pd_op.data, index=0, dtype=builtin.tensor<4x4xf32>)], 'tanh_0.tmp_0': [Value(define_op_name=pd_op.tanh, index=0, dtype=builtin.tensor<4x4xf32>)], 'elementwise_add_0': [Value(define_op_name=pd_op.add, index=0, dtype=builtin.tensor<4x4xf32>)]}
     )DOC");
 
   m->def("clear_pir_compiler_manager", []() {
diff --git a/paddle/pir/include/core/builtin_dialect.h b/paddle/pir/include/core/builtin_dialect.h
@@ -24,14 +24,17 @@ namespace pir {
 ///
 class IR_API BuiltinDialect : public pir::Dialect {
  public:
-  explicit BuiltinDialect(pir::IrContext *context);
+  explicit BuiltinDialect(pir::IrContext* context);
   ///
   /// \brief Each Dialect needs to provide a name function to return the name of
   /// the Dialect.
   ///
   /// \return The name of this Dialect.
   ///
-  static const char *name() { return "builtin"; }
+  static const char* name() { return "builtin"; }
+
+  pir::Type ParseType(pir::IrParser& parser) override;  // NOLINT
+  void PrintType(pir::Type type, std::ostream& os) const override;
 
  private:
   void initialize();
diff --git a/paddle/pir/src/core/builtin_dialect.cc b/paddle/pir/src/core/builtin_dialect.cc
@@ -13,12 +13,16 @@
 // limitations under the License.
 
 #include "paddle/pir/include/core/builtin_dialect.h"
+
+#include "paddle/common/ddim.h"
+#include "paddle/common/layout.h"
 #include "paddle/pir/include/core/builtin_attribute.h"
 #include "paddle/pir/include/core/builtin_op.h"
 #include "paddle/pir/include/core/builtin_type.h"
+#include "paddle/pir/include/core/parser/ir_parser.h"
 
 namespace pir {
-BuiltinDialect::BuiltinDialect(IrContext *context)
+BuiltinDialect::BuiltinDialect(IrContext* context)
     : Dialect(name(), context, TypeId::get<BuiltinDialect>()) {
   initialize();
 }
@@ -38,7 +42,8 @@ void BuiltinDialect::initialize() {
                 BoolType,
                 Complex64Type,
                 Complex128Type,
-                VectorType>();
+                VectorType,
+                DenseTensorType>();
 
   RegisterAttributes<StrAttribute,
                      BoolAttribute,
@@ -64,6 +69,49 @@ void BuiltinDialect::initialize() {
               ConstantOp>();
 }
 
+pir::Type BuiltinDialect::ParseType(pir::IrParser& parser) {  // NOLINT
+  parser.ConsumeAToken("builtin.tensor");
+  parser.ConsumeAToken("<");
+  std::vector<int> dim{};
+  Token dim_token = parser.PeekToken();
+  while (dim_token.token_type_ == DIGIT) {
+    dim_token = parser.ConsumeToken();
+    dim.push_back(atoi(dim_token.val_.c_str()));
+    std::string peek_token_val = parser.PeekToken().val_;
+    if (peek_token_val[0] != 'x') {
+      break;
+    }
+    parser.ConsumeToken();
+    parser.lexer->Unget(static_cast<int>(peek_token_val.size() - 1));
+    if (parser.PeekToken().token_type_ != DIGIT) {
+      break;
+    }
+  }
+  pir::DDim ddim = common::make_ddim(dim);
+  pir::Type dtype = parser.ParseType();
+  std::vector<std::vector<size_t>> lod;
+  std::vector<size_t> lodv;
+  lodv.push_back(0);
+  lod.push_back(lodv);
+  parser.ConsumeAToken(">");
+  return DenseTensorType::get(
+      parser.ctx, dtype, ddim, pir::DataLayout::UNDEFINED, lod, 0);
+}
+
+void BuiltinDialect::PrintType(pir::Type type, std::ostream& os) const {
+  os << type.dialect().name();
+  os << '.';
+  if (auto tensor_type = type.dyn_cast<DenseTensorType>()) {
+    os << "tensor<";
+    for (auto d : common::vectorize(tensor_type.dims())) {
+      os << d;
+      os << "x";
+    }
+    tensor_type.dtype().Print(os);
+    os << ">";
+  }
+}
+
 }  // namespace pir
 
 IR_DEFINE_EXPLICIT_TYPE_ID(pir::BuiltinDialect)
diff --git a/test/cpp/pir/core/TestParserText.txt b/test/cpp/pir/core/TestParserText.txt
@@ -27,14 +27,14 @@ f32
 //END
 
 //CHECK type
-pd_op.tensor<256xf32>
+builtin.tensor<256xf32>
 //END
 
 //CHECK program
 {
- (%0) = "builtin.parameter" () {parameter_name:"conv2d_0.w_0"} : () -> pd_op.tensor<64x3x7x7xf32>
- (%1) = "pd_op.feed" () {col:(Int32)0,is_persistable:[false],name:"data",stop_gradient:[true]} : () -> pd_op.tensor<-1x3x224x224xf32>
- (%2) = "pd_op.conv2d" (%1, %0) {data_format:"NCHW",dilations:[(Int32)1,(Int32)1],groups:(Int32)1,is_persistable:[false],padding_algorithm:"EXPLICIT",paddings:[(Int32)3,(Int32)3],stop_gradient:[false],strides:[(Int32)2,(Int32)2]} : (pd_op.tensor<-1x3x224x224xf32>, pd_op.tensor<64x3x7x7xf32>) -> pd_op.tensor<-1x64x112x112xf32>
+ (%0) = "builtin.parameter" () {parameter_name:"conv2d_0.w_0"} : () -> builtin.tensor<64x3x7x7xf32>
+ (%1) = "pd_op.feed" () {col:(Int32)0,is_persistable:[false],name:"data",stop_gradient:[true]} : () -> builtin.tensor<-1x3x224x224xf32>
+ (%2) = "pd_op.conv2d" (%1, %0) {data_format:"NCHW",dilations:[(Int32)1,(Int32)1],groups:(Int32)1,is_persistable:[false],padding_algorithm:"EXPLICIT",paddings:[(Int32)3,(Int32)3],stop_gradient:[false],strides:[(Int32)2,(Int32)2]} : (builtin.tensor<-1x3x224x224xf32>, builtin.tensor<64x3x7x7xf32>) -> builtin.tensor<-1x64x112x112xf32>
 }
 //END
 
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>";
+      "builtin.tensor<64x3x7x7xf32>";
   std::stringstream ss;
   ss << op_str;
   pir::IrParser* parser = new pir::IrParser(ctx, ss);
diff --git a/test/ir/pir/cinn/symbolic/simple_llama.config b/test/ir/pir/cinn/symbolic/simple_llama.config
diff --git a/test/ir/pir/cinn/symbolic/test_llama_group_log_softmax.py b/test/ir/pir/cinn/symbolic/test_llama_group_log_softmax.py
diff --git a/test/ir/pir/test_ir_pybind.py b/test/ir/pir/test_ir_pybind.py
diff --git a/test/legacy_test/test_fused_rotary_position_embedding.py b/test/legacy_test/test_fused_rotary_position_embedding.py
