• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1 /**
2  * Copyright 2021 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_MINDSPORE_CCSRC_RUNTIME_DEVICE_ASCEND_ASCEND_LAUNCH_TRANSDATA_H_
18 #define MINDSPORE_MINDSPORE_CCSRC_RUNTIME_DEVICE_ASCEND_ASCEND_LAUNCH_TRANSDATA_H_
19 
20 #include <vector>
21 #include <memory>
22 #include <string>
23 #include "runtime/device/ascend/ascend_launch_kernel.h"
24 
25 namespace mindspore::device::ascend {
26 class AscendLaunchTransData : public AscendLaunchKernel {
27  public:
AscendLaunchTransData(void * stream,TypeId dtype,size_t total_size,std::string src_format,std::string dst_format,std::vector<size_t> host_shape)28   AscendLaunchTransData(void *stream, TypeId dtype, size_t total_size, std::string src_format, std::string dst_format,
29                         std::vector<size_t> host_shape)
30       : AscendLaunchKernel(stream),
31         dtype_(dtype),
32         total_size_(total_size),
33         transdata_graph_(nullptr),
34         input_addr_(nullptr),
35         src_format_(src_format),
36         dst_format_(dst_format),
37         shape_(host_shape) {}
38 
39   ~AscendLaunchTransData() override = default;
40 
SetInputAddr(uint8_t * input_addr)41   void SetInputAddr(uint8_t *input_addr) override { input_addr_ = input_addr; }
42   void FreeDeviceMem(void *addr) override;
43   size_t AlignSizeForLaunchKernel(size_t size) override;
44   uint8_t *AllocDeviceMem(size_t size) override;
45   void KernelSelect(const std::shared_ptr<session::KernelGraph> &kernel_graph) override;
46   void KernelBuild(const std::shared_ptr<session::KernelGraph> &kernel_graph) override;
47 
48   void LaunchOpKernel() override;
49   void FreeLaunchDeviceMem() override;
50 
51  protected:
52   TypeId dtype_;
53   size_t total_size_;
54   std::shared_ptr<session::KernelGraph> transdata_graph_;
55   uint8_t *input_addr_;
56   std::string src_format_;
57   std::string dst_format_;
58   std::vector<size_t> shape_;
59 
60  private:
61   std::shared_ptr<session::KernelGraph> ObtainTransDataKernelGraph();
62   void ConstructKernelGraphAndSetAttr();
63 };
64 }  // namespace mindspore::device::ascend
65 
66 #endif  // MINDSPORE_MINDSPORE_CCSRC_RUNTIME_DEVICE_ASCEND_ASCEND_LAUNCH_TRANSDATA_H_
67