【CINN】Add device query components for cinn

paddle/cinn/common/CMakeLists.txt

@@ -19,7 +19,9 @@ gather_srcs(
   arithmatic.cc
   cas.cc
   union_find.cc
-  python_interpreter_guard.cc)
+  python_interpreter_guard.cc
+  dev_info_manager.cc
+  nvgpu_dev_info.cc)
 
 message(STATUS "srcs: ${cinnapi_src}")
 

paddle/cinn/common/dev_info_base.h

@@ -0,0 +1,26 @@
+// Copyright (c) 2023 CINN 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
+
+namespace cinn {
+namespace common {
+
+class DevInfoBase {
+ public:
+  virtual ~DevInfoBase() = default;
+};
+
+}  // namespace common
+}  // namespace cinn

paddle/cinn/common/dev_info_manager.cc

@@ -0,0 +1,48 @@
+// Copyright (c) 2023 CINN 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.
+
+#include <glog/logging.h>
+
+#include "paddle/cinn/common/dev_info_base.h"
+#include "paddle/cinn/common/dev_info_manager.h"
+
+namespace cinn {
+namespace common {
+
+DevInfoMgr::DevInfoMgr(Target::Arch arch, int device_num)
+    : arch_(arch), device_num_(device_num) {
+  switch (arch) {
+    case Target::Arch::ARM:
+    case Target::Arch::X86:
+    case Target::Arch::Unk:
+      impl_ = std::make_unique<DevInfoBase>();
+      break;
+    case Target::Arch::NVGPU:
+#ifdef CINN_WITH_CUDA
+      impl_ = std::make_unique<NVGPUDevInfo>(device_num);
+#endif
+    default:
+      CHECK(false) << "Current device can't be recognized!\n";
+      break;
+  }
+}
+
+std::unique_ptr<DevInfoMgr> DevInfoMgr::GetDevInfo(Target::Arch arch,
+                                                   int device_num) {
+  std::unique_ptr<DevInfoMgr> ret(new DevInfoMgr(arch, device_num));
+  return ret;
+}
+
+}  // namespace common
+}  // namespace cinn

paddle/cinn/common/dev_info_manager.h

@@ -0,0 +1,53 @@
+// Copyright (c) 2023 CINN 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 <memory>
+
+#include "paddle/cinn/common/macros.h"
+#include "paddle/cinn/common/nvgpu_dev_info.h"
+#include "paddle/cinn/common/target.h"
+
+namespace cinn {
+namespace common {
+
+class DevInfoMgr final {
+ private:
+  explicit DevInfoMgr(Target::Arch arch = Target::Arch::Unk,
+                      int device_num = 0);
+  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);
+
+// 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());
+  }
+  RET_TYPE* operator->() {
+    CHECK(!std::is_void<RET_TYPE>()) << "Current device can't be recognized!\n";
+    return dynamic_cast<RET_TYPE*>(impl_.get());
+  }
+#endif
+};
+
+}  // namespace common
+}  // namespace cinn

paddle/cinn/common/nvgpu_dev_info.cc

@@ -0,0 +1,50 @@
+// Copyright (c) 2023 CINN 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.
+#ifdef CINN_WITH_CUDA
+#include "paddle/cinn/common/nvgpu_dev_info.h"
+
+namespace cinn {
+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];
+  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];
+  return ret;
+}
+
+int NVGPUDevInfo::GetMultiProcessorCount() const {
+  return prop_->multiProcessorCount;
+}
+
+int NVGPUDevInfo::GetMaxThreadsPerMultiProcessor() const {
+  return prop_->maxThreadsPerMultiProcessor;
+}
+
+int NVGPUDevInfo::GetMaxThreadsPerBlock() const {
+  return prop_->maxThreadsPerBlock;
+}
+
+}  // namespace common
+}  // namespace cinn
+#endif

paddle/cinn/common/nvgpu_dev_info.h

@@ -0,0 +1,49 @@
+// Copyright (c) 2023 CINN 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
+
+#ifdef CINN_WITH_CUDA
+
+#include <ostream>
+#include <string>
+#include <vector>
+
+#include "paddle/cinn/backends/cuda_util.h"
+#include "paddle/cinn/common/dev_info_base.h"
+#include "paddle/cinn/common/macros.h"
+#include "paddle/cinn/common/target.h"
+
+namespace cinn {
+namespace common {
+
+class NVGPUDevInfo : public DevInfoBase {
+ public:
+  explicit NVGPUDevInfo(int device) : device_num_(device) {
+    CUDA_CALL(cudaGetDeviceProperties(prop_, device));
+  }
+
+  std::array<int, 3> GetMaxGridDims() const;
+  std::array<int, 3> GetMaxBlockDims() const;
+  int GetMultiProcessorCount() const;
+  int GetMaxThreadsPerMultiProcessor() const;
+  int GetMaxThreadsPerBlock() const;
+
+ private:
+  int device_num_;
+  cudaDeviceProp* prop_;
+};
+}  // namespace common
+}  // namespace cinn
+#endif

paddle/cinn/common/target.cc

@@ -26,6 +26,13 @@
 namespace cinn {
 namespace common {
 
+Target::Target(OS o,
+               Arch a,
+               Bit b,
+               const std::vector<Feature> &features,
+               const std::vector<Lib> &libs)
+    : os(o), arch(a), bits(b), features(features), libs(libs) {}
+
 bool Target::operator==(const Target &other) const {
   return os == other.os &&      //
          arch == other.arch &&  //

paddle/cinn/common/target.h

@@ -71,8 +71,7 @@ struct Target {
                   Arch a = Arch::Unk,
                   Bit b = Bit::Unk,
                   const std::vector<Feature>& features = {},
-                  const std::vector<Lib>& libs = {})
-      : os(o), arch(a), bits(b), features(features), libs(libs) {}
+                  const std::vector<Lib>& libs = {});
 
   bool defined() const {
     return os != OS::Unk && arch != Arch::Unk && bits != Bit::Unk;