【CINN】refactor cuda query component (#59096)

Courtesy-Xs · web-flow · commit 619c16a57265 · 2023-11-22T10:36:54.000+08:00
* refactor cuda query component

* fix bugs

* fix review comments

* fix bugs

* fix cuda error

* bind should be used on NVGPU device
diff --git a/paddle/cinn/common/CMakeLists.txt b/paddle/cinn/common/CMakeLists.txt
@@ -20,7 +20,6 @@ gather_srcs(
   cas.cc
   union_find.cc
   python_interpreter_guard.cc
-  dev_info_manager.cc
   nvgpu_dev_info.cc)
 
 message(STATUS "srcs: ${cinnapi_src}")
diff --git a/paddle/cinn/common/dev_info_base.h b/paddle/cinn/common/dev_info_base.h
@@ -19,7 +19,11 @@ namespace common {
 
 class DevInfoBase {
  public:
+  explicit DevInfoBase(int device_num = 0) : device_num_(device_num) {}
   virtual ~DevInfoBase() = default;
+
+ protected:
+  int device_num_;
 };
 
 }  // namespace common
diff --git a/paddle/cinn/common/dev_info_manager.cc b/paddle/cinn/common/dev_info_manager.cc
diff --git a/paddle/cinn/common/dev_info_manager.h b/paddle/cinn/common/dev_info_manager.h
@@ -16,37 +16,51 @@
 
 #include <memory>
 
+#include "paddle/cinn/common/dev_info_base.h"
 #include "paddle/cinn/common/macros.h"
 #include "paddle/cinn/common/nvgpu_dev_info.h"
 #include "paddle/cinn/common/target.h"
 
 namespace cinn {
 namespace common {
 
+template <Target::Arch arch>
+struct GetDevType {
+  using DevType = DevInfoBase;
+};
+
+// Extra device should be added here
+class NVGPUDevInfo;
+template <>
+struct GetDevType<Target::Arch::NVGPU> {
+  using DevType = NVGPUDevInfo;
+};
+
+template <Target::Arch arch>
 class DevInfoMgr final {
  private:
-  explicit DevInfoMgr(Target::Arch arch = Target::Arch::Unk,
-                      int device_num = 0);
+  explicit DevInfoMgr(int device_num = 0) : device_num_(device_num) {
+    impl_ = std::make_unique<typename GetDevType<arch>::DevType>(device_num);
+  }
+
   std::unique_ptr<DevInfoBase> impl_;
-  Target::Arch arch_;
   int device_num_;
 
  public:
-  static std::unique_ptr<DevInfoMgr> GetDevInfo(
-      Target::Arch arch = Target::Arch::NVGPU, int device_num = 0);
+  static DevInfoMgr<arch> GetDevInfo(int device_num = 0) {
+    return DevInfoMgr(device_num);
+  }
 
-// Extra device should be added here
-#ifdef CINN_WITH_CUDA
-  using RET_TYPE = NVGPUDevInfo;
-  const RET_TYPE* operator->() const {
-    CHECK(!std::is_void<RET_TYPE>()) << "Current device can't be recognized!\n";
-    return dynamic_cast<const RET_TYPE*>(impl_.get());
+  using RetType = typename GetDevType<arch>::DevType;
+
+  const RetType* operator->() const {
+    CHECK(!std::is_void<RetType>()) << "Current device can't be recognized!\n";
+    return dynamic_cast<const RetType*>(impl_.get());
   }
-  RET_TYPE* operator->() {
-    CHECK(!std::is_void<RET_TYPE>()) << "Current device can't be recognized!\n";
-    return dynamic_cast<RET_TYPE*>(impl_.get());
+  RetType* operator->() {
+    CHECK(!std::is_void<RetType>()) << "Current device can't be recognized!\n";
+    return dynamic_cast<RetType*>(impl_.get());
   }
-#endif
 };
 
 }  // namespace common
diff --git a/paddle/cinn/common/nvgpu_dev_info.cc b/paddle/cinn/common/nvgpu_dev_info.cc
@@ -19,30 +19,30 @@ namespace common {
 
 std::array<int, 3> NVGPUDevInfo::GetMaxGridDims() const {
   std::array<int, 3> ret;
-  ret[0] = prop_->maxGridSize[0];
-  ret[1] = prop_->maxGridSize[1];
-  ret[2] = prop_->maxGridSize[2];
+  ret[0] = prop_.maxGridSize[0];
+  ret[1] = prop_.maxGridSize[1];
+  ret[2] = prop_.maxGridSize[2];
   return ret;
 }
 
 std::array<int, 3> NVGPUDevInfo::GetMaxBlockDims() const {
   std::array<int, 3> ret;
-  ret[0] = prop_->maxThreadsDim[0];
-  ret[1] = prop_->maxThreadsDim[1];
-  ret[2] = prop_->maxThreadsDim[2];
+  ret[0] = prop_.maxThreadsDim[0];
+  ret[1] = prop_.maxThreadsDim[1];
+  ret[2] = prop_.maxThreadsDim[2];
   return ret;
 }
 
 int NVGPUDevInfo::GetMultiProcessorCount() const {
-  return prop_->multiProcessorCount;
+  return prop_.multiProcessorCount;
 }
 
 int NVGPUDevInfo::GetMaxThreadsPerMultiProcessor() const {
-  return prop_->maxThreadsPerMultiProcessor;
+  return prop_.maxThreadsPerMultiProcessor;
 }
 
 int NVGPUDevInfo::GetMaxThreadsPerBlock() const {
-  return prop_->maxThreadsPerBlock;
+  return prop_.maxThreadsPerBlock;
 }
 
 }  // namespace common
diff --git a/paddle/cinn/common/nvgpu_dev_info.h b/paddle/cinn/common/nvgpu_dev_info.h
@@ -30,8 +30,8 @@ namespace common {
 
 class NVGPUDevInfo : public DevInfoBase {
  public:
-  explicit NVGPUDevInfo(int device) : device_num_(device) {
-    CUDA_CALL(cudaGetDeviceProperties(prop_, device));
+  explicit NVGPUDevInfo(int device_num = 0) : DevInfoBase(device_num) {
+    CUDA_CALL(cudaGetDeviceProperties(&prop_, device_num));
   }
 
   std::array<int, 3> GetMaxGridDims() const;
@@ -41,8 +41,7 @@ class NVGPUDevInfo : public DevInfoBase {
   int GetMaxThreadsPerBlock() const;
 
  private:
-  int device_num_;
-  cudaDeviceProp* prop_;
+  cudaDeviceProp prop_;
 };
 }  // namespace common
 }  // namespace cinn
diff --git a/paddle/cinn/common/target.cc b/paddle/cinn/common/target.cc
@@ -215,18 +215,6 @@ int GetMaxThreads() {
   return max_threads;
 }
 
-std::array<int, 3> GetCUDAMaxBlockDims() {
-  std::array<int, 3> ret;
-#ifdef CINN_WITH_CUDA
-  cudaDeviceProp prop;
-  CUDA_CALL(cudaGetDeviceProperties(&prop, 0));
-  ret[0] = prop.maxThreadsDim[0];
-  ret[1] = prop.maxThreadsDim[1];
-  ret[2] = prop.maxThreadsDim[2];
-#endif
-  return ret;
-}
-
 int GetMaxBlocks() {
   // cudaDeviceGetAttribute ( int* value, cudaDeviceAttr attr, int  device )
   int max_blocks = 1;
@@ -243,18 +231,6 @@ int GetMaxBlocks() {
   return max_blocks;
 }
 
-std::array<int, 3> GetCUDAMaxGridDims() {
-  std::array<int, 3> ret;
-#ifdef CINN_WITH_CUDA
-  cudaDeviceProp prop;
-  CUDA_CALL(cudaGetDeviceProperties(&prop, 0));
-  ret[0] = prop.maxGridSize[0];
-  ret[1] = prop.maxGridSize[1];
-  ret[2] = prop.maxGridSize[2];
-#endif
-  return ret;
-}
-
 const Target &DefaultTarget() {
 #ifdef CINN_WITH_CUDA
   return DefaultNVGPUTarget();
diff --git a/paddle/cinn/common/target.h b/paddle/cinn/common/target.h
@@ -113,10 +113,6 @@ int GetMaxThreads();
 
 int GetMaxBlocks();
 
-std::array<int, 3> GetCUDAMaxBlockDims();
-
-std::array<int, 3> GetCUDAMaxGridDims();
-
 std::ostream& operator<<(std::ostream& os, Target::Arch arch);
 
 }  // namespace common
diff --git a/paddle/cinn/ir/schedule/ir_schedule.cc b/paddle/cinn/ir/schedule/ir_schedule.cc
@@ -27,6 +27,7 @@
 
 #include "paddle/cinn/common/cas.h"
 #include "paddle/cinn/common/common.h"
+#include "paddle/cinn/common/dev_info_manager.h"
 #include "paddle/cinn/common/ir_util.h"
 #include "paddle/cinn/common/target.h"
 #include "paddle/cinn/ir/dy_schedule/ir_schedule.h"
@@ -356,6 +357,7 @@ void StScheduleImpl::Unroll(const Expr& loop) {
 }
 
 void StScheduleImpl::Bind(const Expr& loop, const std::string& thread_axis) {
+#ifdef CINN_WITH_CUDA
   static std::set<std::string> thread_axes = {"blockIdx.x",
                                               "blockIdx.y",
                                               "blockIdx.z",
@@ -365,8 +367,10 @@ void StScheduleImpl::Bind(const Expr& loop, const std::string& thread_axis) {
   CHECK(thread_axes.count(thread_axis))
       << "thread_axis " << thread_axis << " is not supported";
   int offset = thread_axis.back() - 'x';
-  const std::array<int, 3> kMaxBlockDims = common::GetCUDAMaxBlockDims();
-  const std::array<int, 3> kMaxGridDims = common::GetCUDAMaxGridDims();
+  auto cur_dev_info =
+      common::DevInfoMgr<common::Target::Arch::NVGPU>::GetDevInfo(0);
+  const std::array<int, 3> kMaxBlockDims = cur_dev_info->GetMaxBlockDims();
+  const std::array<int, 3> kMaxGridDims = cur_dev_info->GetMaxGridDims();
   auto check_offset = [&](const char& c) -> bool {
     auto extent = loop.As<ir::For>()->extent.as_int32();
     return extent <= (c == 'b' ? kMaxGridDims[offset] : kMaxBlockDims[offset]);
@@ -380,6 +384,7 @@ void StScheduleImpl::Bind(const Expr& loop, const std::string& thread_axis) {
         << "Invalid Bind! The extent of loop is out of range on block size!\n";
     MutateForType(loop, ForType::GPUThread, offset);
   }
+#endif
 }
 
 // The struct used to mutate new rfactor forloop and its' schedule block.
