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_TRANSFORM_ACL_IR_ACL_ALLOCATOR_H_ 18 #define MINDSPORE_CCSRC_TRANSFORM_ACL_IR_ACL_ALLOCATOR_H_ 19 20 #include <memory> 21 #include <string> 22 #include "utils/hash_map.h" 23 #include "runtime/hardware/device_context_manager.h" 24 #include "acl/acl_rt_allocator.h" 25 26 namespace mindspore { 27 namespace transform { 28 class AclAllocator { 29 public: AclAllocator(void * stream)30 explicit AclAllocator(void *stream) : stream_(stream) { 31 auto ms_context = MsContext::GetInstance(); 32 MS_EXCEPTION_IF_NULL(ms_context); 33 auto device_id = ms_context->get_param<uint32_t>(MS_CTX_DEVICE_ID); 34 auto device_target = ms_context->get_param<std::string>(MS_CTX_DEVICE_TARGET); 35 device_context_ = device::DeviceContextManager::GetInstance().GetOrCreateDeviceContext({device_target, device_id}); 36 MS_EXCEPTION_IF_NULL(device_context_); 37 } 38 ~AclAllocator() = default; 39 40 // Acl register func. 41 static void *AllocFunc(void *obj, size_t size); 42 static void *AllocAdviseFunc(void *obj, size_t size, void *addr); 43 static void FreeFunc(void *obj, void *block); 44 static void *GetAddrFromBlock(void *block); 45 set_allocator_desc(const aclrtAllocatorDesc & allocator_desc)46 void set_allocator_desc(const aclrtAllocatorDesc &allocator_desc) { allocator_desc_ = allocator_desc; } allocator_desc()47 aclrtAllocatorDesc allocator_desc() { return allocator_desc_; } stream()48 void *stream() { return stream_; } 49 50 private: 51 void *stream_{nullptr}; 52 device::DeviceContext *device_context_{nullptr}; 53 aclrtAllocatorDesc allocator_desc_{nullptr}; 54 }; 55 using AclAllocatorPtr = std::shared_ptr<AclAllocator>; 56 57 class AclAllocatorRegister { 58 public: 59 // return singlgeton instance 60 static AclAllocatorRegister &Instance(); 61 void RegisterAllocator(void *stream); 62 ~AclAllocatorRegister(); 63 64 private: 65 AclAllocatorRegister() = default; 66 AclAllocatorPtr NewAclAllocator(void *stream); 67 void FreeAclAllocatorRes(const AclAllocatorPtr &allocator_obj); 68 69 mindspore::HashMap<void *, AclAllocatorPtr> allocator_map_; 70 }; 71 } // namespace transform 72 } // namespace mindspore 73 74 #endif // MINDSPORE_CCSRC_TRANSFORM_ACL_IR_ACL_ALLOCATOR_H_ 75