| 1 | /******************************************************************************* |
| 2 | * Copyright 2016-2018 Intel Corporation |
| 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 CPU_PRIMITIVE_HPP |
| 18 | #define CPU_PRIMITIVE_HPP |
| 19 | |
| 20 | #include "mkldnn.h" |
| 21 | |
| 22 | #include "c_types_map.hpp" |
| 23 | #include "memory_tracking.hpp" |
| 24 | #include "primitive.hpp" |
| 25 | #include "scratchpad.hpp" |
| 26 | |
| 27 | #define CTX_IN_MEM(type, arg) static_cast<type>(ctx.input(arg)) |
| 28 | #define CTX_OUT_MEM(type, arg) static_cast<type>(ctx.output(arg)) |
| 29 | |
| 30 | namespace mkldnn { |
| 31 | namespace impl { |
| 32 | namespace cpu { |
| 33 | |
| 34 | struct cpu_memory_t; |
| 35 | |
| 36 | struct cpu_primitive_t: public primitive_t { |
| 37 | cpu_primitive_t(const primitive_desc_t *pd, |
| 38 | bool use_global_scratchpad = false) |
| 39 | : primitive_t(pd) |
| 40 | , scratchpad_buffer_(nullptr) |
| 41 | , global_scratchpad_(nullptr) |
| 42 | { |
| 43 | const size_t scratchpad_size = |
| 44 | this->pd()->scratchpad_size(scratchpad_mode::library); |
| 45 | |
| 46 | if (scratchpad_size) { |
| 47 | if (use_global_scratchpad) |
| 48 | global_scratchpad_ = create_scratchpad(scratchpad_size); |
| 49 | else |
| 50 | scratchpad_buffer_ = malloc(scratchpad_size, 64); |
| 51 | } |
| 52 | } |
| 53 | |
| 54 | virtual ~cpu_primitive_t() { |
| 55 | delete global_scratchpad_; |
| 56 | free(scratchpad_buffer_); |
| 57 | } |
| 58 | |
| 59 | protected: |
| 60 | memory_tracking::grantor_t scratchpad(const exec_ctx_t &ctx) const { |
| 61 | void *ptr = nullptr; |
| 62 | if (pd()->attr()->scratchpad_mode_ == scratchpad_mode::user) { |
| 63 | ptr = CTX_OUT_MEM(void *, MKLDNN_ARG_SCRATCHPAD); |
| 64 | } else { |
| 65 | ptr = global_scratchpad_ |
| 66 | ? global_scratchpad_->get() : scratchpad_buffer_; |
| 67 | } |
| 68 | |
| 69 | return pd()->scratchpad_registry().grantor(ptr); |
| 70 | } |
| 71 | |
| 72 | private: |
| 73 | void *scratchpad_buffer_; |
| 74 | scratchpad_t *global_scratchpad_; |
| 75 | }; |
| 76 | |
| 77 | } |
| 78 | } |
| 79 | } |
| 80 | |
| 81 | #endif |
| 82 | |
| 83 | // vim: et ts=4 sw=4 cindent cino^=l0,\:0,N-s |
| 84 | |