1 /** 2 * Copyright 2023 Huawei Technologies Co., Ltd 3 * 4 * Licensed under the Apache License, Version 2.0 (the "License"); 5 * you may not use this file except in compliance with the License. 6 * You may obtain a copy of the License at 7 * 8 * http://www.apache.org/licenses/LICENSE-2.0 9 * 10 * Unless required by applicable law or agreed to in writing, software 11 * distributed under the License is distributed on an "AS IS" BASIS, 12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 * See the License for the specific language governing permissions and 14 * limitations under the License. 15 */ 16 17 #ifndef MINDSPORE_CCSRC_BACKEND_KERNEL_COMPILER_AKG_GPU_DYNAMIC_AKG_GPU_KERNEL_MOD_H_ 18 #define MINDSPORE_CCSRC_BACKEND_KERNEL_COMPILER_AKG_GPU_DYNAMIC_AKG_GPU_KERNEL_MOD_H_ 19 #include <cuda.h> 20 #include <string> 21 #include <vector> 22 #include <unordered_set> 23 #include <unordered_map> 24 #include <map> 25 #include <memory> 26 #include <utility> 27 #include "plugin/device/gpu/kernel/dynamic_akg/dynamic_utils.h" 28 #include "kernel/kernel.h" 29 #include "plugin/device/gpu/kernel/gpu_kernel_mod.h" 30 #include "plugin/device/gpu/kernel/akg/akg_gpu_kernel_mod.h" 31 #include "kernel/common_utils.h" 32 33 namespace mindspore { 34 namespace kernel { 35 constexpr auto kMappingUpdated = "updated"; 36 constexpr auto kBlockIdxX = "blockIdx.x"; 37 constexpr auto kBlockIdxY = "blockIdx.y"; 38 constexpr auto kBlockIdxZ = "blockIdx.z"; 39 constexpr auto kThreadIdxX = "threadIdx.x"; 40 constexpr auto kThreadIdxY = "threadIdx.y"; 41 constexpr auto kThreadIdxZ = "threadIdx.z"; 42 43 class DynamicAkgGpuKernelManager { 44 public: 45 DynamicAkgGpuKernelManager(); ~DynamicAkgGpuKernelManager()46 virtual ~DynamicAkgGpuKernelManager() { 47 for (auto iter = infotable_.begin(); iter != infotable_.end(); ++iter) { 48 CUresult ret = cuModuleUnload(iter->second->module_); 49 if (ret != CUDA_SUCCESS && ret != CUDA_ERROR_DEINITIALIZED) { 50 const char *msg = nullptr; 51 cuGetErrorName(ret, &msg); 52 MS_LOG(ERROR) << "Unload GPU Module failed. cuModuleUnload error message: " << msg; 53 } 54 } 55 } 56 CUresult GetCUResult(const char *kernel_content, bool force_reload, std::vector<uint32_t> *thread_info, 57 CUfunction *func, const std::string kernel_name); 58 CUresult GetFunction(const KernelPackPtr &kernel_pack, bool force_reload, std::vector<uint32_t> *thread_info, 59 CUfunction *func, const std::string kernel_name); 60 61 private: 62 std::unordered_map<std::string, GpuKernelMetaPtr> infotable_; 63 }; 64 using DynamicAkgGpuKernelManagerPtr = std::shared_ptr<DynamicAkgGpuKernelManager>; 65 66 class DynamicAkgGpuKernelMod : public GpuKernelMod { 67 public: 68 explicit DynamicAkgGpuKernelMod(const KernelPackPtr &kernel_pack); ~DynamicAkgGpuKernelMod()69 virtual ~DynamicAkgGpuKernelMod() {} 70 Init(const std::vector<KernelTensor * > &,const std::vector<KernelTensor * > &)71 bool Init(const std::vector<KernelTensor *> &, const std::vector<KernelTensor *> &) override { return true; }; 72 73 int Resize(const std::vector<KernelTensor *> &inputs, const std::vector<KernelTensor *> &outputs) override; 74 75 bool Launch(const std::vector<KernelTensor *> &inputs, const std::vector<KernelTensor *> &workspace, 76 const std::vector<KernelTensor *> &outputs, void *stream_ptr) override; 77 78 void Initialize(); 79 void CheckJsonParsed(); 80 void InitAkgKernelImpls(); 81 void UpdateStaticShapeMappingInfo(); 82 void UpdateShapeList(const std::vector<KernelTensor *> &inputs, const std::vector<KernelTensor *> &outputs); SetKernelDynamicStatus(bool is_dynamic)83 void SetKernelDynamicStatus(bool is_dynamic) { is_dynamic_ = is_dynamic; } 84 GetKernelModType()85 enum KernelModType GetKernelModType() const override { return KernelModType::DynamicAkgCpuKernelMod; } 86 87 static DynamicAkgGpuKernelManagerPtr kernel_manager_; GetOpSupport()88 std::vector<KernelAttr> GetOpSupport() override { return {}; } 89 std::string kernel_name_; 90 91 private: 92 KernelPackPtr kernel_pack_; 93 std::vector<uint32_t> thread_info_; 94 CUfunction kernel_addr_{nullptr}; 95 bool is_dynamic_{false}; 96 std::vector<std::vector<int64_t>> shape_list_; 97 nlohmann::json parsed_js_; 98 std::vector<int64_t> arg_size_vec_; 99 100 AkgKernelImplInfoPtr kernel_impl_; 101 std::unordered_map<std::string, AkgKernelImplInfoPtr> kernel_map_; 102 AkgKernelImplInfoPtr SelectKernelImpl(); 103 }; 104 } // namespace kernel 105 } // namespace mindspore 106 107 #endif // MINDSPORE_CCSRC_BACKEND_KERNEL_COMPILER_AKG_GPU_DYNAMIC_AKG_GPU_KERNEL_MOD_H_ 108