Merge branch 'develop' of https://github.com/PaddlePaddle/Paddle into…

… dist_api

paddle/infrt/CMakeLists.txt

@@ -82,7 +82,6 @@ add_subdirectory(tensor)
 add_subdirectory(support)
 add_subdirectory(external_kernels)
 add_subdirectory(paddle)
-add_subdirectory(naive)
 add_subdirectory(tests)
 
 
@@ -99,14 +98,15 @@ set(infrt_mlir_incs
         trt_ops_inc
         )
 if (INFRT_WITH_PTEN)
+    set(pten_libs pten)
     set(infrt_mlir_incs ${infrt_mlir_incs}
         MLIRinfrt_pten_tensorIncGen
         MLIRinfrt_pten_baseIncGen
         )
 endif()
 
-cc_library(infrt SHARED SRCS ${infrt_src} DEPS glog boost ${mlir_libs} paddle_framework_proto infrt_naive)
-cc_library(infrt_static SRCS ${infrt_src} DEPS glog boost ${mlir_libs} paddle_framework_proto)
+cc_library(infrt SHARED SRCS ${infrt_src} DEPS glog boost ${mlir_libs} ${pten_libs} paddle_framework_proto infrt_naive)
+cc_library(infrt_static SRCS ${infrt_src} DEPS glog boost ${mlir_libs} ${pten_libs} paddle_framework_proto)
 add_dependencies(infrt ${infrt_mlir_incs} mlir-headers)
 
 add_custom_target(test_infrt_exec DEPENDS ${INFRT_TEST_TARGETS})

paddle/infrt/backends/host/pten_allocator.h

@@ -0,0 +1,33 @@
+/* Copyright (c) 2022 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/pten/core/allocator.h"
+
+namespace infrt {
+namespace backends {
+
+class CpuPtenAllocator : public pten::Allocator {
+ public:
+  static void deleter(pten::Allocation* ptr) { ::operator delete(ptr); }
+
+  AllocationPtr Allocate(size_t bytes_size) {
+    return AllocationPtr(
+        new pten::Allocation(::operator new(bytes_size),
+                             bytes_size,
+                             pten::Place(pten::AllocationType::CPU)),
+        deleter);
+  }
+};
+
+}  // namespace backends
+}  // namespace infrt

paddle/infrt/backends/host/pten_context.h

@@ -0,0 +1,26 @@
+/* Copyright (c) 2022 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/pten/backends/cpu/cpu_context.h"
+
+namespace infrt {
+namespace backends {
+
+class CpuPtenContext : public pten::CPUContext {
+ public:
+  using Base = pten::CPUContext;
+  using pten::CPUContext::SetEigenDevice;
+};
+
+}  // namespace backends
+}  // namespace infrt

paddle/infrt/dialect/dense_tensor.td

@@ -112,6 +112,7 @@ def LoadParamsOp : DT_Op<"load_params", [NoSideEffect]> {
   let verifier = ?;
 }
 
+
 def TensorMapGetTensorOp : DT_Op<"tensor_map_get_tensor", [NoSideEffect]> {
   let summary = "dt.tensor_map_get_tensor operation";
 
@@ -122,7 +123,7 @@ def TensorMapGetTensorOp : DT_Op<"tensor_map_get_tensor", [NoSideEffect]> {
   // input path of model params.
   let arguments = (ins
           TensorMapType:$map,
-          StringType:$name
+          StrAttr:$name
           );
   let results = (outs TensorType:$output);
   let assemblyFormat = "`(` operands `)` attr-dict `->` type($output)";

paddle/infrt/dialect/pten/CMakeLists.txt

@@ -5,6 +5,7 @@ endif()
 #mlir_tablegen_on(infrt_pten_base DIALECT pten)
 add_mlir_dialect(infrt_pten_base pten)
 add_mlir_dialect(infrt_pten_tensor pten_dt)
+add_mlir_dialect(infrt_pten_kernel pten_kernel)
 #mlir_tablegen_on(infrt_pten_tensor)
 
 gather_srcs(infrt_src SRCS

paddle/infrt/dialect/pten/infrt_pten_kernel.td

@@ -0,0 +1,26 @@
+#ifndef PTEN_KERNEL
+#define PTEN_KERNEL
+
+include "paddle/infrt/dialect/pten/infrt_pten_tensor.td"
+
+def PTEN_KernelDialect : Dialect {
+  let name = "pten_kernel";
+
+  let description = [{
+    The PTEN Kernel dialect.
+  }];
+
+  let cppNamespace = "::infrt::pten";
+}
+
+// PTEN Kernel related ops.
+class PDT_Kernel<string mnemonic, list<OpTrait> traits = []> : Op<PTEN_KernelDialect, mnemonic, !listconcat(traits, [IsolatedFromAbove])> {
+}
+
+def FakeKernelOp : PDT_Kernel<"pten.matmul.host.fp32"> {
+  let arguments = (ins CPU_Context:$dev_ctx, TensorType:$x, TensorType:$y, BoolAttr:$transpose_x, BoolAttr:$transpose_y);
+  let results = (outs TensorType:$output);
+}
+
+#endif
+

paddle/infrt/dialect/pten/infrt_pten_tensor.h

@@ -33,6 +33,7 @@
 #include "paddle/infrt/dialect/pten/infrt_pten_tensorTypes.h.inc"
 
 #include "paddle/infrt/dialect/dense_tensor.h"
+#include "paddle/infrt/dialect/pten/pten_base.h"
 // NOLINT
 #define GET_OP_CLASSES
 #include "paddle/infrt/dialect/pten/infrt_pten_tensor.h.inc"

paddle/infrt/dialect/pten/infrt_pten_tensor.td

@@ -21,84 +21,36 @@ def PTEN_DenseTensorDialect : Dialect {
 class PDT_Op<string mnemonic, list<OpTrait> traits = []> : Op<PTEN_DenseTensorDialect, mnemonic, !listconcat(traits, [IsolatedFromAbove])> {
 }
 
-class CreateUninitTensorOp<string dtype>
-      : PDT_Op<"create_uninit_tensor." # dtype, [NoSideEffect]> {
-  let summary = "pdt.create_uninit_tensor operation";
-
-  let description = [{
-      An operation that creates an uninitialized tensor.
-  }];
-
-  let arguments = (ins I64ArrayAttr:$shape);
-  let results = (outs TensorType:$output);
-}
-
-class CreateInitedTensorOp<string dtype, Attr array_attr>
-      : PDT_Op<"create_inited_tensor." #dtype, [NoSideEffect]> {
-  let summary = "pdt.create_inited_tensor operation";
-
-  let description = [{
-      An operation that creates an tensor with shape and values assigned.
-  }];
-
-  let arguments = (ins I64ArrayAttr:$shape, array_attr:$values);
+class CreateDenseTensorOp<string place, string dtype, string layout> 
+      : PDT_Op<"create_dense_tensor." # place # "." # dtype # "." # layout, [NoSideEffect]> {
+  let arguments = (ins CPU_Allocator:$allocator, I64ArrayAttr:$dims, I64ArrayAttr:$lod);
   let results = (outs TensorType:$output);
 }
 
-def PrintTensorOp : PDT_Op<"print_tensor"> {
-  let summary = "pdt.print_tensor operation";
-
-  let description = [{
-    An operation that prints a tensor.
-  }];
-
-  let arguments = (ins TensorType:$input);
-  let results = (outs);
-  let assemblyFormat = "`(` $input `:` type($input) `)` attr-dict";
-}
-
-class FillTensor<string dtype, Attr attr_type> :
-      PDT_Op<"fill_tensor." # dtype> {
-  let summary = "dt.fill_tensor operation";
-
-  let description = [{
-      An operation that fills an input tensor with a values.
-  }];
-
+class FillDenseTensorOp<Attr attr_type, string dtype> :
+      PDT_Op<"fill_dense_tensor." # dtype> {
   let arguments = (ins
       TensorType:$input,
       attr_type:$value
   );
   let results = (outs);
-
-  let assemblyFormat = "`(` $input `:` type($input) `)`  attr-dict";
 }
 
-class FillTensorWithConstantOp<string dtype> :
-      PDT_Op<"fill_tensor_with_constant." # dtype> {
-  let summary = "dt.fill_tensor_with_constant operation";
-
-  let description = [{
-      An operation that fills an input tensor with a single value.
-  }];
-
-  let arguments = (ins
-      TensorType:$input,
-      AnyAttr:$value
-  );
-  let results = (outs);
-
-  let assemblyFormat = "`(` $input `:` type($input) `)`  attr-dict";
+class CreateCPUAllocatorOp
+      : PDT_Op<"create_allocator." # "cpu", [NoSideEffect]> {
+  let arguments = (ins);
+  let results = (outs CPU_Allocator:$output);
 }
 
-foreach dtype = ["ui8", "ui16", "ui32", "ui64", "i32", "f32", "f64", "i64"] in {
-    def PDT_CreateUninitTensorOp_#dtype : CreateUninitTensorOp<dtype>;
-    def PDT_FillTensorWithConstantOp_#dtype : FillTensorWithConstantOp<dtype>;
+class CreateCPUContextOp
+      : PDT_Op<"create_context." # "cpu", [NoSideEffect]> {
+  let arguments = (ins);
+  let results = (outs CPU_Context:$output);
 }
 
-def PDT_FillTensor_f32: FillTensor<"f32", F32ArrayAttr>;
-def PDT_FillTensor_i32: FillTensor<"i32", I32ArrayAttr>;
-def PDT_CreateInitedTensorOp_f32 : CreateInitedTensorOp<"f32", F32ArrayAttr>;
-def PDT_CreateInitedTensorOp_i32 : CreateInitedTensorOp<"i32", I32ArrayAttr>;
+def PDT_CreateDenseTensorOp_cpu_f32_nchw : CreateDenseTensorOp<"cpu", "f32", "nchw">;
+def PDT_FillDenseTensorOp_f32 : FillDenseTensorOp<F32ArrayAttr, "f32">;
+def PDT_CreateAllocatorOp_cpu : CreateCPUAllocatorOp;
+def PDT_CreateContextOp_cpu : CreateCPUContextOp;
 
 #endif

paddle/infrt/dialect/pten/pten_base.cc

@@ -29,7 +29,23 @@ namespace pten {
 
 void PTENDialect::printType(::mlir::Type type,
                             mlir::DialectAsmPrinter& os) const {
-  Dialect::printType(type, os);
+  if (type.isa<CPUAllocatorType>()) {
+    os << "CPU_Allocator";
+    return;
+  }
+  if (type.isa<GPUAllocatorType>()) {
+    os << "GPU_Allocator";
+    return;
+  }
+  if (type.isa<CPUContextType>()) {
+    os << "CPU_Context";
+    return;
+  }
+  if (type.isa<GPUContextType>()) {
+    os << "GPU_Context";
+    return;
+  }
+  llvm_unreachable("unexpected 'allocator/context' type kind");
 }
 
 void PTENDialect::initialize() {
@@ -46,14 +62,16 @@ void PTENDialect::initialize() {
 mlir::Type PTENDialect::parseType(mlir::DialectAsmParser& parser) const {
   llvm::StringRef keyword;
   if (parser.parseKeyword(&keyword)) return mlir::Type();
-  if (keyword == "allocator_CPU") {
+  if (keyword == "CPU_allocator") {
     return CPUAllocatorType::get(parser.getContext());
-  } else if (keyword == "allocator_GPU") {
+  } else if (keyword == "GPU_allocator") {
     return GPUAllocatorType::get(parser.getContext());
-  } else if (keyword == "context_CPU") {
+  } else if (keyword == "CPU_context") {
     return CPUContextType::get(parser.getContext());
-  } else if (keyword == "context_GPU") {
+  } else if (keyword == "GPU_context") {
     return GPUContextType::get(parser.getContext());
+  } else {
+    llvm_unreachable("unexpected 'allocator/context' type kind");
   }
 
   return mlir::Type();

paddle/infrt/host_context/kernel_registry.h

@@ -14,6 +14,7 @@
 
 #pragma once
 
+#include <functional>
 #include <memory>
 #include <string>
 #include <vector>
@@ -23,7 +24,7 @@ namespace host_context {
 
 class KernelFrame;
 
-using KernelImplementation = void (*)(KernelFrame *frame);
+using KernelImplementation = std::function<void(KernelFrame *frame)>;
 
 /**
  * Hold the kernels registered in the system.

paddle/infrt/host_context/kernel_registry_test.cc

@@ -28,7 +28,7 @@ TEST(KernelRegistry, basic) {
   std::string key = "infrt.test.add.i32";
   registry.AddKernel(key, INFRT_KERNEL(add_i32));
 
-  auto* kernel_impl = registry.GetKernel(key);
+  const auto& kernel_impl = registry.GetKernel(key);
   ASSERT_TRUE(kernel_impl);
 
   ValueRef a(1);

paddle/infrt/host_context/mlir_exec.cc

@@ -28,6 +28,9 @@
 #include "paddle/infrt/kernel/tensor_kernels.h"
 #include "paddle/infrt/kernel/tensor_shape_kernels.h"
 #include "paddle/infrt/kernel/test_kernels.h"
+#ifdef INFRT_WITH_PTEN
+#include "paddle/infrt/kernel/pten/registry.h"
+#endif
 
 static llvm::cl::list<std::string> cl_shared_libs(  // NOLINT
     "shared_libs",
@@ -53,6 +56,9 @@ int main(int argc, char** argv) {
   kernel::RegisterTensorShapeKernels(&registry);
   kernel::RegisterTensorKernels(&registry);
   kernel::RegisterControlFlowKernels(&registry);
+#ifdef INFRT_WITH_PTEN
+  kernel::RegisterPtenKernels(&registry);
+#endif
 
   // load extra shared library
   for (const auto& lib_path : cl_shared_libs) {

paddle/infrt/host_context/paddle_mlir.cc

@@ -171,7 +171,7 @@ void MLIRModelGenImpl::UpdateModelParams(
                       builder_,
                       &precision_);
       mlir::Type type_ = mlir::RankedTensorType::get(dims, precision_);
-      auto op = builder_.create<infrt::dt::GetParamOp>(
+      auto op = builder_.create<infrt::dt::TensorMapGetTensorOp>(
           mlir::UnknownLoc::get(context_), type_, map, name);
       params_map_.insert(std::pair<std::string, mlir::Value>(
           var_desc.name(), op.getOperation()->getResult(0)));
@@ -224,15 +224,14 @@ llvm::SmallVector<mlir::Value, 4> MLIRModelGenImpl::GetOpInputValue(
     const infrt::paddle::framework_proto::OpDesc &op_) {
   llvm::SmallVector<mlir::Value, 4> operands;
 
-  std::vector<std::string> inputs_info = {};
+  std::unordered_map<std::string, uint8_t> inputs_info = {};
   if (pd_dialect_inputs_info_map_.count(op_.type()))
     inputs_info = pd_dialect_inputs_info_map_.at(op_.type());
 
   for (int var_idx = 0; var_idx < op_.inputs_size(); ++var_idx) {
     auto &var = op_.inputs(var_idx);
     if (!var.arguments().empty()) {
-      if (!std::count(inputs_info.begin(), inputs_info.end(), var.parameter()))
-        continue;
+      if (!inputs_info.count(var.parameter())) continue;
       operands.push_back((params_map_[var.arguments()[0]]));
     }
   }
@@ -243,17 +242,15 @@ llvm::SmallVector<mlir::Type, 4> MLIRModelGenImpl::GetOpOutputType(
     const infrt::paddle::framework_proto::OpDesc &op_) {
   llvm::SmallVector<mlir::Type, 4> resultTypes;
 
-  std::vector<std::string> pd_dialect_outputs_info = {};
+  std::unordered_map<std::string, uint8_t> pd_dialect_outputs_info = {};
   if (pd_dialect_outputs_info_map_.count(op_.type()))
     pd_dialect_outputs_info = pd_dialect_outputs_info_map_.at(op_.type());
 
   // update op outputs info
   for (int var_idx = 0; var_idx < op_.outputs_size(); ++var_idx) {
     auto &var_name = op_.outputs(var_idx).arguments()[0];
 
-    if (!std::count(pd_dialect_outputs_info.begin(),
-                    pd_dialect_outputs_info.end(),
-                    op_.outputs(var_idx).parameter()))
+    if (!pd_dialect_outputs_info.count(op_.outputs(var_idx).parameter()))
       continue;
 
     // update persistable tensors