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

Author: megemini

cmake/cinn.cmake

@@ -177,12 +177,6 @@ target_link_libraries(cinnapi op_dialect pir phi)
 add_dependencies(cinnapi op_dialect pir phi)
 
 add_dependencies(cinnapi python)
-if(LINUX)
-  target_link_libraries(cinnapi "-Wl,--unresolved-symbols=ignore-all")
-elseif(APPLE)
-  target_link_libraries(cinnapi "-Wl,-undefined,dynamic_lookup")
-endif()
-
 if(WITH_MKL)
   target_link_libraries(cinnapi cinn_mklml)
   add_dependencies(cinnapi cinn_mklml)

cmake/cinn/core.cmake

@@ -19,7 +19,10 @@ function(cinn_cc_library TARGET_NAME)
     endif()
 
     if(cinn_cc_library_DEPS)
-      # Don't need link libwarpctc.so
+      if("${cinn_cc_library_DEPS};" MATCHES "python;")
+        list(REMOVE_ITEM cinn_cc_library_DEPS python)
+        add_dependencies(${TARGET_NAME} python)
+      endif()
       target_link_libraries(${TARGET_NAME} ${cinn_cc_library_DEPS})
       add_dependencies(${TARGET_NAME} ${cinn_cc_library_DEPS})
     endif()

cmake/external/flashattn.cmake

@@ -103,6 +103,10 @@ else()
         CACHE FILEPATH "flash-attn Library" FORCE)
   endif()
 
+  if(NOT DEFINED FA_JOB_POOLS_COMPILE)
+    set(FA_JOB_POOLS_COMPILE 4)
+  endif()
+
   if(CMAKE_CXX_COMPILER_ID STREQUAL "Clang"
      OR CMAKE_CXX_COMPILER_ID STREQUAL "AppleClang"
      OR WIN32)
@@ -172,7 +176,7 @@ else()
                -DCMAKE_POSITION_INDEPENDENT_CODE=ON
                -DCMAKE_BUILD_TYPE=${THIRD_PARTY_BUILD_TYPE}
                -DCMAKE_JOB_POOL_COMPILE:STRING=compile
-               -DCMAKE_JOB_POOLS:STRING=compile=4
+               -DCMAKE_JOB_POOLS:STRING=compile=${FA_JOB_POOLS_COMPILE}
                -DNVCC_ARCH_BIN=${FA_NVCC_ARCH_BIN}
                ${EXTERNAL_OPTIONAL_ARGS}
     CMAKE_CACHE_ARGS

cmake/external/xpu.cmake

@@ -30,9 +30,9 @@ if(NOT DEFINED XPU_XRE_BASE_VERSION)
   set(XPU_XRE_BASE_VERSION "4.32.0.1")
 endif()
 if(NOT DEFINED XPU_XHPC_BASE_DATE)
-  set(XPU_XHPC_BASE_DATE "20240626")
+  set(XPU_XHPC_BASE_DATE "20240704")
 endif()
-set(XPU_XCCL_BASE_VERSION "1.2.1.2")
+set(XPU_XCCL_BASE_VERSION "1.2.3")
 if(NOT DEFINED XPU_XFT_BASE_VERSION)
   set(XPU_XFT_BASE_VERSION "20230602")
 endif()
@@ -43,7 +43,7 @@ set(XPU_XRE_BASE_URL
 )
 
 set(XPU_XCCL_BASE_URL
-    "https://klx-sdk-release-public.su.bcebos.com/xccl/release/${XPU_XCCL_BASE_VERSION}"
+    "https://klx-sdk-release-public.su.bcebos.com/xccl/release/${XPU_XCCL_BASE_VERSION}.1"
 )
 
 if(NOT XPU_XFT_BASE_URL)
@@ -60,21 +60,28 @@ if(WITH_XPTI)
 endif()
 
 if(WITH_XPU_XRE5)
-  set(XPU_XRE_BASE_VERSION "5.0.3.1")
+  set(XPU_XRE_BASE_VERSION "5.0.11.1")
   set(XPU_XRE_BASE_URL
       "https://klx-sdk-release-public.su.bcebos.com/xre/kl3-release/${XPU_XRE_BASE_VERSION}"
   )
+  set(XPU_XCCL_BASE_URL
+      "https://klx-sdk-release-public.su.bcebos.com/xccl/release/${XPU_XCCL_BASE_VERSION}.2"
+  )
 endif()
 
 if(WITH_XCCL_RDMA)
-  set(XPU_XCCL_PREFIX "xccl_rdma")
+  set(XPU_XCCL_PREFIX "xccl_rdma-")
 else()
-  set(XPU_XCCL_PREFIX "xccl_socket")
+  set(XPU_XCCL_PREFIX "xccl_socket-")
+  # NOTE(lijin23): socket has not been supported for XPU3, so the xccl output name was changed.
+  if(WITH_XPU_XRE5)
+    set(XPU_XCCL_PREFIX "xccl_")
+  endif()
 endif()
 
 if(WITH_AARCH64)
   set(XPU_XRE_DIR_NAME "xre-kylin_aarch64")
-  set(XPU_XCCL_DIR_NAME "${XPU_XCCL_PREFIX}-kylin_aarch64")
+  set(XPU_XCCL_DIR_NAME "${XPU_XCCL_PREFIX}kylin_aarch64")
   set(XPU_XFT_DIR_NAME "") # TODO: xft has no kylin output at now.
 elseif(WITH_SUNWAY)
   set(XPU_XRE_DIR_NAME "xre-deepin_sw6_64")
@@ -87,11 +94,11 @@ elseif(WITH_BDCENTOS)
   else()
     set(XPU_XRE_DIR_NAME "xre-bdcentos_x86_64")
   endif()
-  set(XPU_XCCL_DIR_NAME "${XPU_XCCL_PREFIX}-bdcentos_x86_64")
+  set(XPU_XCCL_DIR_NAME "${XPU_XCCL_PREFIX}bdcentos_x86_64")
   set(XPU_XFT_DIR_NAME "xft_bdcentos6u3_x86_64_gcc82")
 elseif(WITH_CENTOS)
   set(XPU_XRE_DIR_NAME "xre-centos7_x86_64")
-  set(XPU_XCCL_DIR_NAME "${XPU_XCCL_PREFIX}-bdcentos_x86_64")
+  set(XPU_XCCL_DIR_NAME "${XPU_XCCL_PREFIX}bdcentos_x86_64")
   set(XPU_XFT_DIR_NAME "xft_bdcentos6u3_x86_64_gcc82")
 else()
   # Ubuntu as default
@@ -102,7 +109,7 @@ else()
     set(XPU_XRE_DIR_NAME "xre-ubuntu_1604_x86_64")
     set(XPU_XHPC_DIR_NAME "xhpc-ubuntu1604_x86_64")
   endif()
-  set(XPU_XCCL_DIR_NAME "${XPU_XCCL_PREFIX}-ubuntu_x86_64")
+  set(XPU_XCCL_DIR_NAME "${XPU_XCCL_PREFIX}ubuntu_x86_64")
   set(XPU_XFT_DIR_NAME "xft_ubuntu1604_x86_64")
 endif()
 

cmake/generic.cmake

@@ -709,6 +709,11 @@ function(nv_test TARGET_NAME)
     get_property(os_dependency_modules GLOBAL PROPERTY OS_DEPENDENCY_MODULES)
     target_link_libraries(${TARGET_NAME} ${nv_test_DEPS}
                           ${os_dependency_modules} paddle_gtest_main phi)
+    if(WIN32)
+      target_link_libraries(${TARGET_NAME} ${PYTHON_LIBRARIES})
+    else()
+      target_link_libraries(${TARGET_NAME} python)
+    endif()
     add_dependencies(${TARGET_NAME} ${nv_test_DEPS} paddle_gtest_main)
     common_link(${TARGET_NAME})
     add_test(${TARGET_NAME} ${TARGET_NAME})

paddle/cinn/ast_gen_ius/tensor_group.cc

@@ -215,27 +215,5 @@ absl::flat_hash_map<std::string, ir::Tensor> TensorGroup::AllocateBuffers() {
   return name_to_tensor_;
 }
 
-void StageMapShareMemory(const poly::StageMap& stages) {
-  absl::flat_hash_map<std::string, ir::_Tensor_*> tensor_map;
-  for (auto& stage : stages) {
-    tensor_map[stage.second->tensor()->name] = stage.second->tensor();
-  }
-  for (auto& stage : stages) {
-    if (!stage.second->tensor()->buffer.defined() &&
-        !stage.second->meta.tensors_to_share_buffer_with.empty()) {
-      for (auto& str : stage.second->meta.tensors_to_share_buffer_with) {
-        if (tensor_map[str]->buffer.defined()) {
-          auto edited_shape = tensor_map[str]->buffer->shape;
-          stage.second->tensor()->Bind(tensor_map[str]->buffer);
-          tensor_map[str]->buffer->shape = edited_shape;
-          VLOG(3) << "Stage Tensor " << stage.second->tensor()->name
-                  << " bind buffer to " << tensor_map[str]->name << " , "
-                  << tensor_map[str]->buffer->name;
-        }
-      }
-    }
-  }
-}
-
 }  // namespace ast_gen_ius
 }  // namespace cinn

paddle/cinn/backends/CMakeLists.txt

@@ -7,6 +7,7 @@ gather_srcs(
   codegen_c.cc
   codegen_c_x86.cc
   codegen_cuda_host.cc
+  codegen_invoke_module.cc
   extern_func_emitter.cc
   extern_func_emitter_builtin.cc
   function_prototype.cc

paddle/cinn/backends/codegen_cuda_host.cc

@@ -194,72 +194,6 @@ llvm::Value* CodeGenCUDA_Host::LowerGPUKernelLauncher(
   return function;
 }
 
-llvm::Value* CodeGenCUDA_Host::LowerHostFunc(const ir::_LoweredFunc_* func) {
-  // Create the function
-  // @{
-  auto* function_type = GenFunctionTypeFromCinnFunction(func, true);
-  f_ = llvm::Function::Create(
-      function_type, llvm::Function::ExternalLinkage, func->name, m_);
-  f_->setCallingConv(llvm::CallingConv::C);
-  f_->setHasUWTable();
-
-  std::vector<llvm::Value*> ll_function_args;
-  std::transform(f_->arg_begin(),
-                 f_->arg_end(),
-                 std::back_inserter(ll_function_args),
-                 [](auto& arg) { return std::addressof(arg); });
-  // @}
-
-  // Set local scope table
-  PADDLE_ENFORCE_EQ(ll_function_args.size(),
-                    func->args.size(),
-                    phi::errors::InvalidArgument(
-                        "The number of arguments is not equal to the number of "
-                        "function arguments"));
-  for (int i = 0; i < ll_function_args.size(); ++i) {
-    SetVar(func->args[i].name(), ll_function_args[i]);
-  }
-  llvm::BasicBlock* entry = llvm::BasicBlock::Create(
-      /*Context=*/b_->getContext(),
-      /*Name=*/"entry",
-      /*Parent=*/f_,
-      /*InsertBefore=*/nullptr);
-  b_->SetInsertPoint(entry);
-  CodeGenLLVM::Visit(&func->body);
-
-  // Reset local scope table
-  for (const ir::Argument& func_arg : func->args) {
-    symbol_table_->Erase(func_arg.name());
-  }
-  RetVoid();
-
-  return f_;
-}
-
-llvm::Value* CodeGenCUDA_Host::LowerParseArgsValueCall(
-    const ir::Call* call_ir) {
-  auto ret_type = CinnTypeToLLVMType(Int(64), m_);
-  std::vector<llvm::Type*> args_type;
-  PADDLE_ENFORCE_EQ(
-      call_ir->read_args.size(),
-      2,
-      phi::errors::InvalidArgument(
-          "The number of arguments of ParseArgsValue should be 2"));
-  CHECK(call_ir->read_args[0].is_var() &&
-        call_ir->read_args[0].as_var()->type().is_cpp_handle());
-  CHECK(call_ir->read_args[1].type().is_int(32));
-  args_type.push_back(CinnTypeToLLVMType(type_of<void*>(), m_));
-  args_type.push_back(CinnTypeToLLVMType(type_of<int32_t>(), m_));
-
-  auto func_type = llvm::FunctionType::get(ret_type, args_type, false);
-  auto call_func = m_->getOrInsertFunction(call_ir->name, func_type);
-
-  std::vector<llvm::Value*> call_args;
-  call_args.push_back(std::addressof(*f_->arg_begin()));
-  call_args.push_back(b_->getInt32(call_ir->read_args[1].as_int32()));
-  return b_->CreateCall(call_func, call_args);
-}
-
 llvm::Value* CodeGenCUDA_Host::LowerCUDAKernelCall(const ir::Call* call_ir) {
   std::vector<llvm::Value*> ll_function_args;
   std::transform(f_->arg_begin(),

paddle/cinn/backends/codegen_cuda_host.h

@@ -13,15 +13,9 @@
 // limitations under the License.
 
 #pragma once
-
-#include <absl/container/flat_hash_map.h>
-
 #include <memory>
-#include <string>
-#include <tuple>
-#include <vector>
 
-#include "paddle/cinn/backends/llvm/codegen_llvm.h"
+#include "paddle/cinn/backends/codegen_invoke_module.h"
 #include "paddle/cinn/runtime/intrinsic.h"
 
 PD_DECLARE_bool(cinn_bucket_compile);
@@ -30,30 +24,29 @@ namespace cinn {
 namespace backends {
 
 /**
- * CodeGenCUDA takes a CINN Module with host functions and output a LLVM module.
+ * CodeGenCUDA_Host takes a CINN Module with CUDA host functions and output a
+ * LLVM module.
  */
-class CodeGenCUDA_Host : public CodeGenLLVM {
+class CodeGenCUDA_Host : public CodeGenHost {
  public:
   explicit CodeGenCUDA_Host(llvm::Module *m,
                             llvm::IRBuilder<> *b,
                             const std::shared_ptr<SymbolTable> &vars = nullptr)
-      : CodeGenLLVM(m, b, vars) {}
+      : CodeGenHost(m, b, vars) {}
 
-  using CodeGenLLVM::Visit;
+  // TODO(Hongqing-work): remove this after we clear some old codes.
   llvm::Value *Visit(const ir::_LoweredFunc_ *func) override {
     if (FLAGS_cinn_bucket_compile) {
-      return LowerHostFunc(func);
+      return CodeGenHost::Visit(func);
     }
     return LowerGPUKernelLauncher(func);
   }
 
   llvm::Value *Visit(const ir::Call *op) override {
-    if (op->name == runtime::intrinsic::get_value_in_cuda_kernel_args) {
-      return LowerParseArgsValueCall(op);
-    } else if (op->name == runtime::intrinsic::call_cuda_kernel) {
+    if (op->name == runtime::intrinsic::call_cuda_kernel) {
       return LowerCUDAKernelCall(op);
     } else {
-      return CodeGenLLVM::Visit(op);
+      return CodeGenHost::Visit(op);
     }
   }
 
@@ -73,10 +66,6 @@ class CodeGenCUDA_Host : public CodeGenLLVM {
    */
   llvm::Value *LowerGPUKernelLauncher(const ir::_LoweredFunc_ *func);
 
-  llvm::Value *LowerHostFunc(const ir::_LoweredFunc_ *func);
-
-  llvm::Value *LowerParseArgsValueCall(const ir::Call *call_ir);
-
   llvm::Value *LowerCUDAKernelCall(const ir::Call *op);
 };