【custom】Add CustomEngine instruction ... (#70549)

xiaoguoguo626807 · web-flow · commit 8b7f151fd785 · 2025-01-02T10:33:22.000+08:00
* modify assign inplace

* delete floor_divide decomp

* delete floor_divide decomp

* modify prim test

* modify prim test

* add chunk_id guard

* add engine decomp

* modify flag

* add all

* add print_flag

* tmp customengine

* add cutomengine dialect

* modify ci bug

* modify ci

* modify ci

* modify windows ci

* modify windows dll bug

* test

* add paddle_with_export of test dir cmake_txt

* change to IR_DEFINE_EXPLICIT_TEST_TYPE_ID

* fix windows ci bug

* add IR_API

* recover comment

* move loadcustomlib and add base func for loadcustomenginelib

* add loadcustomEnginelib test

* add test

* rollback noneed test

* rollback

* modify declare

* modify declare

* add cmake cond

* windows bug

* reply comment

* add customengine execute

* optimize

* fix conflict
diff --git a/paddle/fluid/custom_engine/custom_engine_ext.h b/paddle/fluid/custom_engine/custom_engine_ext.h
@@ -16,6 +16,7 @@
 #if !defined(_WIN32)
 #include <cstddef>
 #include <cstring>
+#include <unordered_map>
 
 #include "paddle/phi/backends/device_ext.h"
 
@@ -24,12 +25,32 @@ extern "C" {
 #endif
 
 typedef struct C_Operation_st* C_Operation;
+typedef struct C_IrContext_st* C_IrContext;
+typedef struct C_Value_st* C_Value;
+typedef struct C_Place_st* C_Place;
+typedef struct C_KernelKey_st* C_KernelKey;
+typedef struct C_Block_st* C_Block;
+typedef struct std::unordered_map<C_Operation_st, C_Operation_st>*
+    C_Operation_Map;
+typedef struct std::unordered_map<C_Value_st, C_Value_st>* C_Value_Map;
+typedef struct C_CustomEngineInstruction_st* C_CustomEngineInstruction;
+
+struct C_CustomEngineLowerParams {
+  C_IrContext ir_context;
+  C_Operation operation;
+  C_KernelKey kernel_key;
+  C_Place place;
+  C_Operation_Map map_op_pair;
+  C_Value_Map map_value_pair;
+  C_Block block;
+};
 
 struct C_CustomEngineInterface {
   size_t size;
-  C_Status (*graph_engine_build)();
-  C_Status (*graph_engine_execute)();
-  C_Status (*custom_engine_op_lower)();
+  C_Status (*register_custom_engine_op)();
+  C_Status (*graph_engine_build)(C_CustomEngineInstruction);
+  C_Status (*graph_engine_execute)(C_CustomEngineInstruction);
+  C_Status (*custom_engine_op_lower)(C_CustomEngineLowerParams*);
 };
 
 struct CustomEngineParams {
diff --git a/paddle/fluid/custom_engine/custom_engine_load.cc b/paddle/fluid/custom_engine/custom_engine_load.cc
@@ -40,6 +40,19 @@ void LoadCustomEngineLib(const std::string& dso_lib_path,
     paddle::custom_engine::CustomEngineManager::Instance()
         ->SetCustomEngineInterface(engine_params->interface);
 
+    auto* interface = paddle::custom_engine::CustomEngineManager::Instance()
+                          ->GetCustomEngineInterface();
+
+    // register custom engine op
+    if (interface->register_custom_engine_op) {
+      interface->register_custom_engine_op();
+    } else {
+      LOG(WARNING) << "Skipped lib [" << dso_lib_path
+                   << "]. register_custom_engine_op is not set!!! please check "
+                      "the version "
+                      "compatibility between PaddlePaddle and Custom Engine.";
+    }
+
   } else {
     LOG(WARNING) << "Skipped lib [" << dso_lib_path
                  << "]. Wrong engine parameters!!! please check the version "
diff --git a/paddle/fluid/custom_engine/fake_cpu_engine.h b/paddle/fluid/custom_engine/fake_cpu_engine.h
diff --git a/paddle/fluid/framework/new_executor/CMakeLists.txt b/paddle/fluid/framework/new_executor/CMakeLists.txt
@@ -14,6 +14,11 @@ if(NOT WITH_ONEDNN)
     ${CMAKE_CURRENT_SOURCE_DIR}/instruction/onednn/onednn_mixed_instruction.cc)
 endif()
 
+if(NOT WITH_CUSTOM_DEVICE)
+  list(REMOVE_ITEM standalone_executor_srcs
+       ${CMAKE_CURRENT_SOURCE_DIR}/instruction/custom_engine_instruction.cc)
+endif()
+
 if(NOT TENSORRT_FOUND OR NOT WITH_TENSORRT)
   list(REMOVE_ITEM standalone_executor_srcs
        ${CMAKE_CURRENT_SOURCE_DIR}/instruction/tensorrt_engine_instruction.cc)
diff --git a/paddle/fluid/framework/new_executor/instruction/custom_engine_instruction.cc b/paddle/fluid/framework/new_executor/instruction/custom_engine_instruction.cc
@@ -0,0 +1,83 @@
+// Copyright (c) 2024 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
+//
+//     https://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/fluid/framework/new_executor/instruction/custom_engine_instruction.h"
+#include "paddle/fluid/custom_engine/custom_engine_ext.h"
+#include "paddle/fluid/custom_engine/custom_engine_manager.h"
+#include "paddle/fluid/framework/new_executor/instruction/instruction_util.h"
+#include "paddle/fluid/pir/dialect/operator/ir/op_dialect.h"
+
+namespace paddle {
+namespace framework {
+
+CustomEngineInstruction::CustomEngineInstruction(
+    size_t id,
+    const phi::Place &place,
+    ::pir::Operation *op,
+    const ValueExecutionInfo *value_exec_info)
+    : InstructionBase(id, place), value_exec_info_(value_exec_info) {
+  PADDLE_ENFORCE_EQ(paddle::dialect::IsCustomEngineOp(op),
+                    true,
+                    ::common::errors::InvalidArgument(
+                        "The Op to construct CustomEngineInstruction must be a "
+                        "custom engine op.  "
+                        "but got op is %d",
+                        op->name()));
+
+  auto op_attributes = op->attributes();
+  op_ = op;
+  VLOG(6) << "Start Build custom engine";
+  interface_ = paddle::custom_engine::CustomEngineManager::Instance()
+                   ->GetCustomEngineInterface();
+  if (interface_ && interface_->graph_engine_build) {
+    interface_->graph_engine_build(
+        reinterpret_cast<C_CustomEngineInstruction>(this));
+  } else {
+    PADDLE_THROW(common::errors::Unimplemented(
+        "CustomEngineInstruction's C_CustomEngineInterface->graph_engine_build "
+        "not implemented"));
+  }
+
+  PADDLE_ENFORCE_NOT_NULL(
+      custom_engine_,
+      ::common::errors::InvalidArgument(
+          "custom_engine_ should not be nullptr after graph_engine_build"));
+
+  VLOG(6) << "Finish build engine for: " << op_name_;
+
+  SetDeviceContext(
+      ParseDeviceContext(op,
+                         phi::DeviceContextPool::Instance().Get(place),
+                         place,
+                         GetExecutionStream(),
+                         GetStreamPriority()));
+  VLOG(6) << "finish process device context";
+
+  InitInputsOutputsIds(op, *value_exec_info_);
+  VLOG(6) << "finish process inputs outputs index";
+}
+
+void CustomEngineInstruction::Run() {
+  if (interface_ && interface_->graph_engine_execute) {
+    interface_->graph_engine_execute(
+        reinterpret_cast<C_CustomEngineInstruction>(this));
+  } else {
+    PADDLE_THROW(common::errors::Unimplemented(
+        "CustomEngineInstruction's C_CustomEngineInterface->graph_engine_run "
+        "not implemented"));
+  }
+}
+
+}  // namespace framework
+}  // namespace paddle
diff --git a/paddle/fluid/framework/new_executor/instruction/custom_engine_instruction.h b/paddle/fluid/framework/new_executor/instruction/custom_engine_instruction.h
@@ -0,0 +1,72 @@
+// Copyright (c) 2024 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
+//
+//     https://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/custom_engine/custom_engine_ext.h"
+#include "paddle/fluid/framework/new_executor/instruction/instruction_base.h"
+#include "paddle/fluid/framework/new_executor/pir_adaptor/pir_adaptor_util.h"
+#include "paddle/phi/core/platform/device_context.h"
+namespace pir {
+class Operation;
+}  // namespace pir
+
+namespace paddle {
+namespace framework {
+class Scope;
+class ValueExecutionInfo;
+
+class CustomEngineInstruction : public InstructionBase {
+ public:
+  typedef void (*DeletePtr)(void*);
+  CustomEngineInstruction(size_t id,
+                          const phi::Place& place,
+                          ::pir::Operation* op,
+                          const ValueExecutionInfo* value_exec_info);
+
+  ::pir::Operation* Operation() const override { return op_; }
+  const ValueExecutionInfo* GetValueExecutionInfo() const {
+    return value_exec_info_;
+  }
+
+  void Run() override;
+
+  void SetName(const std::string& name) { op_name_ = name; }
+  const std::string& Name() const override { return op_name_; }
+
+  void SetCustomEngine(void* custom_engine) { custom_engine_ = custom_engine; }
+
+  void* CustomEngine() { return custom_engine_; }
+
+  void SetCustomEngineDeleter(DeletePtr custom_engine_deleter) {
+    custom_engine_deleter_ = custom_engine_deleter;
+  }
+
+  ~CustomEngineInstruction() override {
+    if (custom_engine_deleter_) {
+      custom_engine_deleter_(custom_engine_);
+    }
+  }
+
+ private:
+  phi::Place place_;
+  std::string op_name_ = "custom_engine.base_op";
+  ::pir::Operation* op_{nullptr};  // not owned
+
+  C_CustomEngineInterface* interface_;
+  const ValueExecutionInfo* value_exec_info_;  // not owned
+  void* custom_engine_;                        // not owned
+  DeletePtr custom_engine_deleter_{nullptr};   // not owned
+};
+}  // namespace framework
+}  // namespace paddle
diff --git a/paddle/fluid/framework/new_executor/pir_interpreter.cc b/paddle/fluid/framework/new_executor/pir_interpreter.cc
@@ -45,7 +45,9 @@
 #ifdef PADDLE_WITH_CINN
 #include "paddle/fluid/framework/new_executor/instruction/cinn_jit_instruction.h"
 #endif
-
+#ifdef PADDLE_WITH_CUSTOM_DEVICE
+#include "paddle/fluid/framework/new_executor/instruction/custom_engine_instruction.h"
+#endif
 #include "paddle/fluid/framework/new_executor/instruction/builtin_combine_instruction.h"
 #include "paddle/fluid/framework/new_executor/instruction/control_flow/assert_instruction.h"
 #include "paddle/fluid/framework/new_executor/instruction/control_flow/has_elements_instruction.h"
@@ -959,9 +961,18 @@ void PirInterpreter::BuildInstruction() {
       vec_instruction_base_.emplace_back(
           std::make_unique<CustomKernelInstruction>(
               op_idx++, place_, &op, *(value_exe_info_.get())));
+    } else if (paddle::dialect::IsCustomEngineOp(&op)) {
+#ifdef PADDLE_WITH_CUSTOM_DEVICE
+      CREATE_INSTR(CustomEngineInstruction);
+#else
+      PADDLE_THROW(common::errors::PreconditionNotMet(
+          "Program has CustomEngineOp and must compile Paddle use "
+          "-DWITH_CUSTOM_DEVICE=ON"));
+#endif
     } else {
       PADDLE_THROW(common::errors::Unimplemented(
-          "Now only support pd_kernel, onednn_kernel, custom_kernel, trt_op "
+          "Now only support pd_kernel, onednn_kernel, custom_kernel, trt_op, "
+          "custom_engine_op "
           "and cinn dialect."));
     }
   }
diff --git a/paddle/fluid/pir/transforms/pd_op_to_kernel_pass.cc b/paddle/fluid/pir/transforms/pd_op_to_kernel_pass.cc
@@ -58,6 +58,10 @@
 #include "paddle/cinn/hlir/framework/pir/utils.h"
 #endif
 
+#ifdef PADDLE_WITH_CUSTOM_DEVICE
+#include "paddle/fluid/custom_engine/custom_engine_manager.h"
+#endif
+
 #ifdef PADDLE_WITH_DNNL
 #include "paddle/fluid/pir/dialect/operator/ir/onednn_op.h"
 #include "paddle/fluid/pir/dialect/operator/ir/op_onednn_dialect.h"
@@ -2444,6 +2448,35 @@ void HandleForTensorRTOp(
   block->push_back(op);
 }
 
+#ifdef PADDLE_WITH_CUSTOM_DEVICE
+void HandleForCustomEngineOp(
+    pir::IrContext* ctx,
+    pir::Operation* op_item,
+    phi::KernelKey* kernel_key,
+    phi::Place place,
+    std::unordered_map<pir::Operation*, pir::Operation*>* map_op_pair,
+    std::unordered_map<pir::Value, pir::Value>* map_value_pair,
+    pir::Block* block,
+    C_CustomEngineInterface* interface) {
+  if (interface->custom_engine_op_lower) {
+    struct C_CustomEngineLowerParams lower_params {
+      reinterpret_cast<C_IrContext>(ctx),
+          reinterpret_cast<C_Operation>(op_item),
+          reinterpret_cast<C_KernelKey>(kernel_key),
+          reinterpret_cast<C_Place>(&place),
+          reinterpret_cast<C_Operation_Map>(map_op_pair),
+          reinterpret_cast<C_Value_Map>(map_value_pair),
+          reinterpret_cast<C_Block>(block)
+    };
+    VLOG(6) << "Handle CustomEngineOp while lowering to kernel pass";
+    interface->custom_engine_op_lower(&lower_params);
+  } else {
+    PADDLE_THROW(common::errors::Unimplemented(
+        "CustomEngineInstruction's "
+        "C_CustomEngineInterface->custom_engine_op_lower not implemented"));
+  }
+}
+#endif
 std::vector<pir::Type> BuildOutputs(
     pir::Operation* op_item,
     const std::string& kernel_fn_str,
@@ -3460,7 +3493,21 @@ void ProcessBlock(
                           new_block);
       continue;
     }
-
+#ifdef PADDLE_WITH_CUSTOM_DEVICE
+    if (paddle::dialect::IsCustomEngineOp(op_item)) {
+      auto* interface = paddle::custom_engine::CustomEngineManager::Instance()
+                            ->GetCustomEngineInterface();
+      HandleForCustomEngineOp(ctx,
+                              op_item,
+                              &kernel_key,
+                              place,
+                              map_op_pair,
+                              map_value_pair,
+                              new_block,
+                              interface);
+      continue;
+    }
+#endif
 #ifdef PADDLE_WITH_DNNL
     if (op_item->HasTrait<OneDNNTrait>() &&
         kernel_key.backend() != phi::Backend::ONEDNN) {
diff --git a/paddle/phi/backends/device_ext.h b/paddle/phi/backends/device_ext.h
@@ -14,6 +14,7 @@
 
 #pragma once
 #if !defined(_WIN32)
+#include <stdint.h>
 #include <cstddef>
 #include <cstring>
 
diff --git a/python/setup.py.in b/python/setup.py.in
@@ -1164,8 +1164,10 @@ headers = (
     list(find_files('infer_symbolic_shape.h','@PADDLE_SOURCE_DIR@/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/'))+
     #pir headers
     list(find_files('lexer.h','@PADDLE_SOURCE_DIR@/paddle/pir/src/core/parser/'))+
-    list(find_files('token.h','@PADDLE_SOURCE_DIR@/paddle/pir/src/core/parser/')))
-
+    list(find_files('token.h','@PADDLE_SOURCE_DIR@/paddle/pir/src/core/parser/'))+
+    #new_executor headers
+    list(find_files('pir_adaptor_util.h','@PADDLE_SOURCE_DIR@/paddle/fluid/framework/new_executor/pir_adaptor/'))+
+    list(find_files('custom_engine_instruction.h','@PADDLE_SOURCE_DIR@/paddle/fluid/framework/new_executor/instruction/')))
 jit_layer_headers = ['layer.h', 'serializer.h', 'serializer_utils.h', 'all.h', 'function.h']
 for f in jit_layer_headers:
     headers += list(find_files(f, '@PADDLE_SOURCE_DIR@/paddle/fluid/jit', recursive=True))
diff --git a/test/cpp/pir/custom_engine/CMakeLists.txt b/test/cpp/pir/custom_engine/CMakeLists.txt
@@ -1,5 +1,7 @@
 if(WITH_CUSTOM_DEVICE)
-  paddle_test(test_custom_engine_loadlib SRCS test_custom_engine_loadlib.cc)
+  paddle_test(test_custom_engine_loadlib SRCS test_custom_engine_loadlib.cc
+              custom_engine_op.cc)
+
 endif()
 paddle_test(test_custom_engine SRCS test_custom_engine_operation.cc
             custom_engine_op.cc)
diff --git a/test/cpp/pir/custom_engine/custom_engine_op.h b/test/cpp/pir/custom_engine/custom_engine_op.h
diff --git a/test/cpp/pir/custom_engine/fake_cpu_engine.h b/test/cpp/pir/custom_engine/fake_cpu_engine.h
diff --git a/test/cpp/pir/custom_engine/fake_cpu_engine_base.h b/test/cpp/pir/custom_engine/fake_cpu_engine_base.h
diff --git a/test/cpp/pir/custom_engine/test_custom_engine_loadlib.cc b/test/cpp/pir/custom_engine/test_custom_engine_loadlib.cc
