[IR]Lower pd op to kernel dialect

paddle/fluid/ir/CMakeLists.txt

@@ -1,2 +1,3 @@
 add_subdirectory(interface)
 add_subdirectory(dialect)
+add_subdirectory(pass)

paddle/fluid/ir/dialect/kernel_attribute.h

@@ -0,0 +1,38 @@
+// Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#pragma once
+
+#include "paddle/fluid/ir/dialect/kernel_attribute_storage.h"
+#include "paddle/ir/core/attribute.h"
+#include "paddle/phi/core/enforce.h"
+
+namespace paddle {
+namespace dialect {
+
+class KernelAttribute : public ir::Attribute {
+ public:
+  using Attribute::Attribute;
+
+  DECLARE_ATTRIBUTE_UTILITY_FUNCTOR(KernelAttribute, KernelAttributeStorage);
+
+  bool operator<(const KernelAttribute &right) const {
+    return storage() < right.storage();
+  }
+
+  phi::KernelKey data() const { return storage()->GetAsKey(); }
+};
+
+}  // namespace dialect
+}  // namespace paddle

paddle/fluid/ir/dialect/kernel_attribute_storage.h

@@ -0,0 +1,48 @@
+// Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#pragma once
+
+#include "paddle/ir/core/attribute.h"
+#include "paddle/ir/core/utils.h"
+#include "paddle/phi/common/data_type.h"
+#include "paddle/phi/core/kernel_factory.h"
+
+namespace paddle {
+namespace dialect {
+
+struct KernelAttributeStorage : public ir::AttributeStorage {
+  using ParamKey = phi::KernelKey;
+
+  explicit KernelAttributeStorage(const ParamKey &key) { kernel_key_ = key; }
+
+  static KernelAttributeStorage *Construct(ParamKey key) {
+    return new KernelAttributeStorage(key);
+  }
+
+  static std::size_t HashValue(const ParamKey &key) {
+    auto t = phi::KernelKey::Hash()(key);
+    return t;
+  }
+
+  bool operator==(const ParamKey &key) const { return kernel_key_ == key; }
+
+  ParamKey GetAsKey() const { return kernel_key_; }
+
+ private:
+  phi::KernelKey kernel_key_;
+};
+
+}  // namespace dialect
+}  // namespace paddle

paddle/fluid/ir/dialect/pd_kernel_dialect.cc → paddle/fluid/ir/dialect/kernel_dialect.cc

@@ -12,15 +12,16 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-#include "paddle/fluid/ir/dialect/pd_kernel_dialect.h"
+#include "paddle/fluid/ir/dialect/kernel_dialect.h"
+#include "paddle/fluid/ir/dialect/kernel_op.h"
 #include "paddle/fluid/ir/dialect/pd_attribute.h"
-#include "paddle/fluid/ir/dialect/pd_kernel_op.h"
 // NOTE(zhangbo9674): File pd_op.h is generated by op_gen.py, see details in
 // paddle/fluid/ir/dialect/CMakeLists.txt.
 #include "paddle/fluid/framework/convert_utils.h"
 #include "paddle/fluid/framework/data_type.h"
-#include "paddle/fluid/ir/dialect/pd_kernel_type.h"
-#include "paddle/fluid/ir/dialect/pd_kernel_type_storage.h"
+#include "paddle/fluid/ir/dialect/kernel_attribute.h"
+#include "paddle/fluid/ir/dialect/kernel_type.h"
+#include "paddle/fluid/ir/dialect/kernel_type_storage.h"
 #include "paddle/fluid/ir/dialect/pd_op.h"
 #include "paddle/fluid/ir/dialect/utils.h"
 #include "paddle/ir/core/dialect_interface.h"
@@ -38,13 +39,10 @@ void PaddleKernelDialect::initialize() {
   RegisterTypes<paddle::dialect::AllocatedDenseTensorType>();
   RegisterOps<dialect::PhiKernelOp>();
 
-  // RegisterAttributes<paddle::dialect::IntArrayAttribute,
-  //                    paddle::dialect::DataTypeAttribute,
-  //                    paddle::dialect::PlaceAttribute,
-  //                    paddle::dialect::DataLayoutAttribute>();
+  RegisterAttributes<paddle::dialect::KernelAttribute>();
 }
 
-void PaddleKernelDialect::PrintType(ir::Type type, std::ostream &os) {
+void PaddleKernelDialect::PrintType(ir::Type type, std::ostream &os) const {
   AllocatedDenseTensorType tensor_type =
       type.dyn_cast<AllocatedDenseTensorType>();
 
@@ -58,5 +56,13 @@ void PaddleKernelDialect::PrintType(ir::Type type, std::ostream &os) {
   os << ">";
 }
 
+void PaddleKernelDialect::PrintAttribute(ir::Attribute attr,
+                                         std::ostream &os) const {
+  phi::KernelKey kernel = attr.dyn_cast<KernelAttribute>().data();
+
+  os << "<backend:" << kernel.backend() << "|layout:" << kernel.layout()
+     << "|dtype:" << kernel.dtype() << ">";
+}
+
 }  // namespace dialect
 }  // namespace paddle

paddle/fluid/ir/dialect/pd_kernel_dialect.h → paddle/fluid/ir/dialect/kernel_dialect.h

@@ -27,7 +27,9 @@ class PaddleKernelDialect : public ir::Dialect {
 
   static const char* name() { return "pd_kernel"; }
 
-  void PrintType(ir::Type type, std::ostream& os);
+  void PrintType(ir::Type type, std::ostream& os) const override;
+
+  void PrintAttribute(ir::Attribute attr, std::ostream& os) const override;
 
  private:
   void initialize();

paddle/fluid/ir/dialect/pd_kernel_op.cc → paddle/fluid/ir/dialect/kernel_op.cc

@@ -12,7 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-#include "paddle/fluid/ir/dialect/pd_kernel_op.h"
+#include "paddle/fluid/ir/dialect/kernel_op.h"
 
 namespace paddle {
 namespace dialect {
@@ -23,7 +23,8 @@ const char *PhiKernelOp::attributes_name[attributes_num] = {
 void PhiKernelOp::Verify(const std::vector<ir::OpResult> &inputs,
                          const std::vector<ir::Type> &outputs,
                          const ir::AttributeMap &attributes) {
-  VLOG(4) << "Verifying inputs, outputs and attributes for: SetParameterOp.";
+  VLOG(4) << "Verifying inputs, outputs and attributes for: PhiKernelOp.";
+
   // Verify inputs type:
 
   // Verify if attributes contain attribute name in attributes_name:

paddle/fluid/ir/dialect/pd_kernel_type.cc → paddle/fluid/ir/dialect/kernel_type.cc

@@ -12,7 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-#include "paddle/fluid/ir/dialect/pd_kernel_type.h"
+#include "paddle/fluid/ir/dialect/kernel_type.h"
 
 namespace paddle {
 namespace dialect {

paddle/fluid/ir/dialect/pd_kernel_type.h → paddle/fluid/ir/dialect/kernel_type.h

@@ -14,7 +14,7 @@
 
 #pragma once
 
-#include "paddle/fluid/ir/dialect/pd_kernel_type_storage.h"
+#include "paddle/fluid/ir/dialect/kernel_type_storage.h"
 #include "paddle/fluid/ir/dialect/pd_type.h"
 #include "paddle/ir/core/type.h"
 
@@ -31,18 +31,18 @@ class AllocatedDenseTensorType : public ir::Type {
                                AllocatedDenseTensorTypeStorage);
 
   static AllocatedDenseTensorType get(ir::IrContext *ctx,
-                                      phi::Place place,
+                                      const phi::Place &place,
                                       dialect::DenseTensorType type) {
     return ir::TypeManager::template get<AllocatedDenseTensorType>(
         ctx, place, type);
   }
 
   static AllocatedDenseTensorType get(ir::IrContext *ctx,
-                                      phi::Place place,
-                                      ir::Type dtype,
-                                      phi::DDim dims,
-                                      phi::DataLayout layout,
-                                      phi::LoD lod,
+                                      const phi::Place &place,
+                                      const ir::Type &dtype,
+                                      const phi::DDim &dims,
+                                      const phi::DataLayout &layout,
+                                      const phi::LoD &lod,
                                       size_t offset) {
     dialect::DenseTensorType dense_tensor_type =
         dialect::DenseTensorType::get(ctx, dtype, dims, layout, lod, offset);

paddle/fluid/ir/dialect/op_gen.py

@@ -110,7 +110,8 @@ class {op_name} : public ir::Op<{op_name}{interfaces}{traits}> {{
   std::vector<paddle::dialect::OpInputInfo> inputs = {{ {inputs} }};
   std::vector<paddle::dialect::OpAttributeInfo> attributes = {{ {attributes} }};
   std::vector<paddle::dialect::OpOutputInfo> outputs = {{ {outputs} }};
-  paddle::dialect::OpRunTimeInfo run_time_info = OpRunTimeInfo("{infer_meta_func}", {{"{infer_meta_param}"}}, {{"{kernel_func}"}}, {{"{kernel_param}"}}, {{{inplace}}}, {{{view}}});
+  paddle::dialect::OpRunTimeInfo run_time_info = OpRunTimeInfo("{infer_meta_func}", {{"{infer_meta_param}"}}, {{"{kernel_func}"}}, {{"{kernel_param}"}}, {{"{kernel_key_dtype}"}}, {{{inplace}}}, {{{view}}});
+
   return std::make_tuple(inputs, attributes, outputs, run_time_info);
 }}
 """
@@ -1488,9 +1489,14 @@ def OpGenerator(
 
             kernel_func_str = ""
             kernel_param_str = ""
+            kernel_key_dtype = ""
             if op_kernel_map is not None:
                 kernel_func_str = '", "'.join(op_kernel_map['func'])
                 kernel_param_str = '", "'.join(op_kernel_map['param'])
+                if 'data_type' in op_kernel_map and op_kernel_map['data_type']:
+                    kernel_key_dtype = '", "'.join(
+                        op_kernel_map['data_type']['candidates']
+                    )
 
             inplace_str = ""
             view_str = ""
@@ -1513,6 +1519,7 @@ def OpGenerator(
                 infer_meta_param=infer_meta_param_str,
                 kernel_func=kernel_func_str,
                 kernel_param=kernel_param_str,
+                kernel_key_dtype=kernel_key_dtype,
                 inplace=inplace_str,
                 view=view_str,
             )

paddle/fluid/ir/dialect/pd_type_storage.h

@@ -112,15 +112,6 @@ struct DenseTensorTypeStorage : public ir::TypeStorage {
     return ParamKey(dtype_, dims_, layout_, lod_, offset_) == key;
   }
 
-  bool operator==(const DenseTensorTypeStorage &storage) const {
-    return ParamKey(dtype_, dims_, layout_, lod_, offset_) ==
-           ParamKey(storage.dtype_,
-                    storage.dims_,
-                    storage.layout_,
-                    storage.lod_,
-                    storage.offset_);
-  }
-
   ParamKey GetAsKey() const {
     return ParamKey(dtype_, dims_, layout_, lod_, offset_);
   }

paddle/fluid/ir/dialect/utils.h

@@ -144,18 +144,21 @@ struct OpRunTimeInfo {
   std::vector<std::string> infer_meta_param;
   std::vector<std::string> kernel_func;
   std::vector<std::string> kernel_param;
+  std::vector<std::string> kernel_key_dtype;
   std::vector<std::pair<std::string, std::string>> inplace;
   std::vector<std::pair<std::string, std::string>> view;
   OpRunTimeInfo(std::string infer_meta_func,
                 std::vector<std::string> infer_meta_param,
                 std::vector<std::string> kernel_func,
                 std::vector<std::string> kernel_param,
+                std::vector<std::string> dtype,
                 std::vector<std::pair<std::string, std::string>> inplace,
                 std::vector<std::pair<std::string, std::string>> view)
       : infer_meta_func(infer_meta_func),
         infer_meta_param(infer_meta_param),
         kernel_func(kernel_func),
         kernel_param(kernel_param),
+        kernel_key_dtype(dtype),
         inplace(inplace),
         view(view) {}
 };

paddle/fluid/ir/interface/infershape.h

@@ -19,19 +19,15 @@
 class InferShapeInterface : public ir::OpInterfaceBase<InferShapeInterface> {
  public:
   struct Concept {
-    explicit Concept(void (*infer_shape)(ir::Operation *,
-                                         phi::InferMetaContext *))
+    explicit Concept(void (*infer_shape)(phi::InferMetaContext *))
         : infer_shape_(infer_shape) {}
-    void (*infer_shape_)(ir::Operation *, phi::InferMetaContext *);
+    void (*infer_shape_)(phi::InferMetaContext *);
   };
 
   template <class ConcreteOp>
   struct Model : public Concept {
-    static void InferShape(ir::Operation *op,
-                           phi::InferMetaContext *infer_meta) {
-      ConcreteOp concret_op = op->dyn_cast<ConcreteOp>();
-      if (concret_op == nullptr) throw("concret_op is nullptr");
-      concret_op.InferShape(infer_meta);
+    static void InferShape(phi::InferMetaContext *infer_meta) {
+      return ConcreteOp::InferShape(infer_meta);
     }
 
     Model() : Concept(InferShape) {}
@@ -41,7 +37,7 @@ class InferShapeInterface : public ir::OpInterfaceBase<InferShapeInterface> {
       : ir::OpInterfaceBase<InferShapeInterface>(op), impl_(impl) {}
 
   void InferShape(phi::InferMetaContext *infer_meta) {
-    impl_->infer_shape_(operation(), infer_meta);
+    impl_->infer_shape_(infer_meta);
   }
 
  private:

paddle/fluid/ir/pass/CMakeLists.txt

@@ -0,0 +1,7 @@
+# All source files of pd_dialect, except for the source file of op, which is generated in the compilation directory.
+file(GLOB PD_PASS_SRCS "*.cc")
+
+cc_library(
+  pd_op_to_kernel_pass
+  SRCS ${PD_PASS_SRCS}
+  DEPS new_ir phi_utils)