[NewIR]Add cinn RuntimeDialect and JitKernelOp

paddle/cinn/hlir/CMakeLists.txt

@@ -3,3 +3,4 @@ add_subdirectory(pe)
 add_subdirectory(op)
 add_subdirectory(pass)
 add_subdirectory(kernels)
+add_subdirectory(dialect)

paddle/cinn/hlir/dialect/CMakeLists.txt

@@ -0,0 +1,6 @@
+# TODO(Aurelius84): new_ir_compiler depends on pd_dialect and could
+# not found under CINN_ONLY mode
+if(NOT CINN_ONLY)
+  cinn_cc_library(cinn_dialect SRCS runtime_dialect.cc jit_kernel_op.cc DEPS
+                  pd_dialect)
+endif()

paddle/cinn/hlir/dialect/jit_kernel_op.cc

@@ -0,0 +1,44 @@
+// 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.
+
+#include "paddle/cinn/hlir/dialect/jit_kernel_op.h"
+
+#include "paddle/ir/core/builtin_attribute.h"
+#include "paddle/ir/core/enforce.h"
+
+namespace cinn {
+namespace dialect {
+
+const char* JitKernelOp::attributes_name[attributes_num] = {kAttrName};
+
+void JitKernelOp::Verify() {
+  VLOG(4) << "Verifying inputs, outputs and attributes for: JitKernelOp.";
+
+  auto& attributes = this->attributes();
+
+  IR_ENFORCE(attributes.count(kAttrName) > 0 &&
+                 attributes.at(kAttrName).isa<::ir::PointerAttribute>(),
+             "Type of attribute: instruction is not right.");
+}
+
+hlir::framework::Instruction* JitKernelOp::instruction() {
+  void* ptr =
+      attributes().at(kAttrName).dyn_cast<ir::PointerAttribute>().data();
+  return reinterpret_cast<hlir::framework::Instruction*>(ptr);
+}
+
+}  // namespace dialect
+}  // namespace cinn
+
+IR_DEFINE_EXPLICIT_TYPE_ID(cinn::dialect::JitKernelOp)

paddle/cinn/hlir/dialect/jit_kernel_op.h

@@ -0,0 +1,60 @@
+// 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/op_base.h"
+
+namespace cinn {
+
+namespace hlir {
+namespace framework {
+class Instruction;
+}  // namespace framework
+}  // namespace hlir
+
+namespace dialect {
+
+/*
+ * TODO(Aurelius84): THIS IS NOT FINAL STATE!
+ *   JitKernel is unified runtime operation to represent
+ *   jit compiled function ptr from backend, such as
+ *   nvrct.
+
+ *   Ideally, JitKernel should only contains ArrayAttribute
+ *   with each element is PointerAttribute, which is jit
+ *   function ptr indeed.
+
+ *   Currently, we regard hlir::framework::Instruction
+ *   temporarily, and will spilt executor information like
+ *   scope, inputs, outputs into InterpretorCore module.
+*/
+class JitKernelOp : public ::ir::Op<JitKernelOp> {
+ public:
+  using Op::Op;
+  static const char* name() { return "cinn.jit_kernel"; }
+  // TODO(Aurelius84): Think deeply what should contains
+  static constexpr uint32_t attributes_num = 1;
+  static constexpr char* kAttrName = "instruction";
+  static const char* attributes_name[attributes_num];
+
+  hlir::framework::Instruction* instruction();
+
+  void Verify();
+};
+
+}  // namespace dialect
+}  // namespace cinn
+
+IR_DECLARE_EXPLICIT_TYPE_ID(cinn::dialect::JitKernelOp)

paddle/cinn/hlir/dialect/runtime_dialect.cc

@@ -0,0 +1,32 @@
+// 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.
+
+#include "paddle/cinn/hlir/dialect/runtime_dialect.h"
+#include "paddle/cinn/hlir/dialect/jit_kernel_op.h"
+
+namespace cinn {
+namespace dialect {
+
+RuntimeDialect::RuntimeDialect(::ir::IrContext* context)
+    : ::ir::Dialect(
+          name(), context, ::ir::TypeId::get<cinn::dialect::RuntimeDialect>()) {
+  this->initialize();
+}
+
+void RuntimeDialect::initialize() { RegisterOps<cinn::dialect::JitKernelOp>(); }
+
+}  // namespace dialect
+}  // namespace cinn
+
+IR_DEFINE_EXPLICIT_TYPE_ID(cinn::dialect::RuntimeDialect)

paddle/cinn/hlir/dialect/runtime_dialect.h

@@ -0,0 +1,35 @@
+// 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/dialect.h"
+
+namespace cinn {
+namespace dialect {
+
+class RuntimeDialect : public ::ir::Dialect {
+ public:
+  explicit RuntimeDialect(::ir::IrContext* context);
+
+  static const char* name() { return "cinn"; }
+
+ private:
+  void initialize();
+};
+
+}  // namespace dialect
+}  // namespace cinn
+
+IR_DECLARE_EXPLICIT_TYPE_ID(cinn::dialect::RuntimeDialect)

paddle/cinn/hlir/framework/CMakeLists.txt

@@ -26,6 +26,8 @@ gather_srcs(
 if(NOT CINN_ONLY)
   cinn_cc_library(new_ir_compiler SRCS new_ir_compiler.cc DEPS cinncore
                   pd_dialect)
+  cinn_cc_library(convert_to_dialect SRCS convert_to_dialect.cc DEPS cinncore
+                  cinn_dialect)
 endif()
 
 if(WITH_CUDA)

paddle/cinn/hlir/framework/convert_to_dialect.cc

@@ -0,0 +1,55 @@
+// 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.
+
+#include "paddle/cinn/hlir/framework/convert_to_dialect.h"
+
+#include <string>
+#include <unordered_map>
+
+#include "paddle/cinn/hlir/dialect/jit_kernel_op.h"
+#include "paddle/cinn/hlir/dialect/runtime_dialect.h"
+#include "paddle/cinn/hlir/framework/program.h"
+#include "paddle/ir/core/builtin_attribute.h"
+#include "paddle/ir/core/program.h"
+
+namespace cinn {
+namespace hlir {
+namespace framework {
+
+std::unique_ptr<::ir::Program> ConvertToRuntimeDialect(
+    const hlir::framework::Program& program) {
+  ::ir::IrContext* ctx = ::ir::IrContext::Instance();
+  ctx->GetOrRegisterDialect<cinn::dialect::RuntimeDialect>();
+  auto ir_program = std::make_unique<::ir::Program>(ctx);
+
+  std::string jit_op_name = dialect::JitKernelOp::name();
+  ::ir::OpInfo op_info = ctx->GetRegisteredOpInfo(jit_op_name);
+
+  auto& instrs = program.GetRunInstructions();
+  for (auto& instr : instrs) {
+    std::unordered_map<std::string, ::ir::Attribute> op_attrs{
+        {dialect::JitKernelOp::kAttrName,
+         ::ir::PointerAttribute::get(ctx, instr.get())},
+    };
+
+    ::ir::Operation* cinn_op =
+        ::ir::Operation::Create({}, op_attrs, {}, op_info);
+    ir_program->block()->push_back(cinn_op);
+  }
+  return std::move(ir_program);
+}
+
+}  // namespace framework
+}  // namespace hlir
+}  // namespace cinn

paddle/cinn/hlir/framework/convert_to_dialect.h

@@ -0,0 +1,33 @@
+// 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 <memory>
+
+namespace ir {
+class Program;
+}  // namespace ir
+
+namespace cinn {
+namespace hlir {
+namespace framework {
+class Program;
+
+std::unique_ptr<::ir::Program> ConvertToRuntimeDialect(
+    const hlir::framework::Program& program);
+
+}  // namespace framework
+}  // namespace hlir
+}  // namespace cinn

paddle/cinn/hlir/framework/program.h

@@ -56,10 +56,11 @@ class Program {
    */
   size_t size() const { return instrs_.size(); }
 
-  const std::vector<std::unique_ptr<Instruction>>& GetPreRunInstructions() {
+  const std::vector<std::unique_ptr<Instruction>>& GetPreRunInstructions()
+      const {
     return prerun_instrs_;
   }
-  const std::vector<std::unique_ptr<Instruction>>& GetRunInstructions() {
+  const std::vector<std::unique_ptr<Instruction>>& GetRunInstructions() const {
     return instrs_;
   }
 

test/cpp/ir/cinn/CMakeLists.txt

@@ -5,6 +5,7 @@ if(WITH_TESTING AND WITH_CINN)
     new_ir_compiler_test.cc
     DEPS
     new_ir_compiler
+    convert_to_dialect
     ir
     phi
     gtest