Merge branch 'PaddlePaddle:develop' into pir7

Author: longranger2

.gitignore

@@ -108,6 +108,8 @@ paddle/fluid/pir/dialect/operator/ir/pd_api.*
 paddle/fluid/pir/dialect/operator/ir/op_decomp.cc
 paddle/fluid/pir/dialect/operator/ir/pd_op_vjp.cc
 paddle/fluid/pir/dialect/operator/ir/pd_op.*
+paddle/fluid/pir/dialect/operator/ir/pd_onednn_op.*
+paddle/fluid/pir/dialect/operator/ir/pd_onednn_op_info.*
 paddle/fluid/pir/dialect/operator/ir/pd_op_bwd.*
 paddle/fluid/pir/dialect/operator/ir/pd_op_fused.*
 paddle/fluid/pir/dialect/operator/ir/pd_op_fused_bwd.*

CMakeLists.txt

@@ -54,7 +54,8 @@ option(WITH_XPU "Compile PaddlePaddle with BAIDU KUNLUN XPU" OFF)
 option(WITH_XPU_KP "Compile PaddlePaddle with BAIDU XPU compiler " OFF)
 option(WITH_XPU_XFT "Compile PaddlePaddle with BAIDU XPU-XFT" OFF)
 option(WITH_XPU_PLUGIN "Compile PaddlePaddle with BAIDU XPU plugin" OFF)
-option(WITH_XPU_XHPC "Compile PaddlePaddle with BAIDU XPU-HPC library" OFF)
+option(WITH_XPU_XHPC "Compile PaddlePaddle with BAIDU XPU-HPC library"
+       ${WITH_XPU})
 option(WITH_WIN_DUMP_DBG "Compile with windows core dump debug mode" OFF)
 option(WITH_ROCM "Compile PaddlePaddle with ROCM platform" OFF)
 option(WITH_IPU "Compile PaddlePaddle with Graphcore IPU" OFF)

cmake/external/xpu.cmake

@@ -26,12 +26,12 @@ set(XPU_XBLAS_LIB_NAME "libxpu_blas.so")
 set(XPU_XFA_LIB_NAME "libxpu_flash_attention.so")
 
 if(NOT DEFINED XPU_BASE_DATE)
-  set(XPU_BASE_DATE "20231203")
+  set(XPU_BASE_DATE "20231218")
 endif()
 if(NOT DEFINED XPU_XHPC_BASE_DATE)
-  set(XPU_XHPC_BASE_DATE "20231215")
+  set(XPU_XHPC_BASE_DATE "20231229")
 endif()
-set(XPU_XCCL_BASE_VERSION "1.1.7.1")
+set(XPU_XCCL_BASE_VERSION "1.1.8.1")
 if(NOT DEFINED XPU_XFT_BASE_VERSION)
   set(XPU_XFT_BASE_VERSION "20230602")
 endif()

cmake/inference_lib.cmake

@@ -237,6 +237,16 @@ copy_part_of_thrid_party(inference_lib_dist ${PADDLE_INFERENCE_INSTALL_DIR})
 
 set(src_dir "${PADDLE_SOURCE_DIR}/paddle/fluid")
 
+if(WIN32)
+  set(paddle_common_lib ${PADDLE_BINARY_DIR}/paddle/common/common.*)
+else()
+  set(paddle_common_lib ${PADDLE_BINARY_DIR}/paddle/common/libcommon.*)
+endif()
+copy(
+  inference_lib_dist
+  SRCS ${paddle_common_lib}
+  DSTS ${PADDLE_INFERENCE_INSTALL_DIR}/paddle/lib)
+
 if(WIN32)
   if(WITH_STATIC_LIB)
     set(paddle_inference_lib
@@ -268,11 +278,6 @@ else()
       SRCS ${paddle_phi_lib}
       DSTS ${PADDLE_INFERENCE_INSTALL_DIR}/paddle/lib)
   endif()
-  set(paddle_common_lib ${PADDLE_BINARY_DIR}/paddle/common/libcommon.*)
-  copy(
-    inference_lib_dist
-    SRCS ${paddle_common_lib}
-    DSTS ${PADDLE_INFERENCE_INSTALL_DIR}/paddle/lib)
 endif()
 
 copy(
@@ -323,10 +328,13 @@ copy(
   inference_lib_dist
   SRCS ${PADDLE_SOURCE_DIR}/paddle/phi/core/visit_type.h
   DSTS ${PADDLE_INFERENCE_INSTALL_DIR}/paddle/include/paddle/phi/core/)
+
 copy(
   inference_lib_dist
-  SRCS ${PADDLE_SOURCE_DIR}/paddle/phi/core/hostdevice.h
-  DSTS ${PADDLE_INFERENCE_INSTALL_DIR}/paddle/include/paddle/phi/core/)
+  SRCS ${PADDLE_SOURCE_DIR}/paddle/phi/core/distributed/type_defs.h
+  DSTS ${PADDLE_INFERENCE_INSTALL_DIR}/paddle/include/paddle/phi/core/distributed/
+)
+
 copy(
   inference_lib_dist
   SRCS ${PADDLE_SOURCE_DIR}/paddle/fluid/platform/init_phi.h

paddle/cinn/backends/codegen_cuda_host.cc

@@ -198,13 +198,23 @@ llvm::Value* CodeGenCUDA_Host::LowerHostFunc(const ir::_LoweredFunc_* func) {
                  [](auto& arg) { return std::addressof(arg); });
   // @}
 
+  // Set local scope table
+  CHECK_EQ(ll_function_args.size(), func->args.size());
+  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_;

paddle/cinn/backends/codegen_cuda_host.h

@@ -53,7 +53,7 @@ class CodeGenCUDA_Host : public CodeGenLLVM {
     } else if (op->name == runtime::intrinsic::call_cuda_kernel) {
       return LowerCUDAKernelCall(op);
     } else {
-      CINN_NOT_IMPLEMENTED;
+      return CodeGenLLVM::Visit(op);
     }
   }
 

paddle/cinn/backends/codegen_cuda_util.h

@@ -31,6 +31,7 @@ namespace backends {
 #define KERNEL_ARGS "kernel_args"
 #define KERNEL_ARGS_NUM "kernel_args_num"
 #define KERNEL_STREAM "kernel_stream"
+#define TENSOR_SHAPE_ARGS "tensor_shape_args"
 
 /**
  * Split a CINN Module into two separate modules, one cantains the host
@@ -150,7 +151,8 @@ struct CollectBucketStrategyHostFunctionVisitor
       : CollectHostFunctionVisitor(module_name),
         kernel_args_(KERNEL_ARGS, type_of<void*>()),
         kernel_args_num_(KERNEL_ARGS_NUM, type_of<int>()),
-        kernel_stream_(KERNEL_STREAM, type_of<void*>()) {}
+        kernel_stream_(KERNEL_STREAM, type_of<void*>()),
+        tensor_shape_args_(TENSOR_SHAPE_ARGS, type_of<int32_t**>()) {}
 
   std::tuple<ir::Module, ir::Module> operator()(Expr* expr) {
     ir::IRMutator<>::Visit(expr, expr);
@@ -181,6 +183,25 @@ struct CollectBucketStrategyHostFunctionVisitor
                                 {});
     host_module_builder.AddFunctionWithoutOptim(
         host_func.as_lowered_func_ref());
+
+    // Parse LoweredFunc to infer output tensor's shape
+    std::vector<ir::Expr> infer_shape_func_body_stmts(arg_defs_);
+    infer_shape_func_body_stmts.insert(
+        infer_shape_func_body_stmts.end(),
+        op->infer_shape_func.as_lowered_func()->body);
+
+    std::vector<ir::Argument> infer_shape_arguments = {
+        ir::Argument(kernel_args_, ir::Argument::IO::kOutput),
+        ir::Argument(kernel_args_num_, ir::Argument::IO::kInput),
+        ir::Argument(tensor_shape_args_, ir::Argument::IO::kOutput)};
+
+    ir::Expr host_infer_shape_func =
+        ir::_LoweredFunc_::Make(op->infer_shape_func.as_lowered_func()->name,
+                                infer_shape_arguments,
+                                ir::Block::Make(infer_shape_func_body_stmts),
+                                {});
+    host_module_builder.AddFunctionWithoutOptim(
+        host_infer_shape_func.as_lowered_func_ref());
   }
 
   void ProcessLoweredFunc(ir::Expr func, ir::Expr predicate);
@@ -199,6 +220,7 @@ struct CollectBucketStrategyHostFunctionVisitor
   ir::Var kernel_args_;
   ir::Var kernel_args_num_;
   ir::Var kernel_stream_;
+  ir::Var tensor_shape_args_;
 };
 
 }  // namespace detail

paddle/cinn/backends/llvm/codegen_llvm.cc

@@ -818,7 +818,8 @@ llvm::Value *CodeGenLLVM::Visit(const ir::_Var_ *op) {
   // TODO(fc500110) hard coding
   if (LLVM_WillVarLowerAsPointer(op->name)) {
     result = value;
-  } else if (value->getType()->isPointerTy()) {
+  } else if (value->getType()->isPointerTy() &&
+             !value->getType()->getPointerElementType()->isPointerTy()) {
     result = Load(value, op->name + "_load");
   } else {
     result = value;

paddle/cinn/common/CMakeLists.txt

@@ -22,7 +22,10 @@ gather_srcs(
   python_interpreter_guard.cc
   nvgpu_dev_info.cc
   integer_set.cc
-  dim_expr_simplify.cc)
+  dim_expr_simplify.cc
+  dim_expr_converter.cc
+  broadcast_tree.cc
+  dim_expr_util.cc)
 
 cinn_cc_test(test_equation_graph_topo_walker SRCS
              equation_graph_topo_walker_test.cc DEPS gtest glog)
@@ -47,6 +50,10 @@ if(WITH_CUDA)
                gtest glog)
 endif()
 if(NOT CINN_ONLY)
+  cinn_cc_test(dim_expr_util_test SRCS dim_expr_util_test.cc DEPS cinncore)
   cinn_cc_test(dim_expr_simplify_test SRCS dim_expr_simplify_test.cc DEPS
                cinncore)
+  cinn_cc_test(dim_expr_converter_test SRCS dim_expr_converter_test.cc DEPS
+               cinncore)
+  cinn_cc_test(broadcast_tree_test SRCS broadcast_tree_test.cc DEPS cinncore)
 endif()