Home
last modified time | relevance | path

Searched refs:ws_mem (Results 1 – 3 of 3) sorted by relevance

/external/tensorflow/tensorflow/core/kernels/mkl/
Dmkl_pooling_ops_common.cc70 context_.ws_mem.reset(new memory(context_.fwd_pd.get()->workspace_desc(), in Setup()
74 {MKLDNN_ARG_WORKSPACE, *context_.ws_mem}}); in Setup()
97 context_.ws_mem->set_data_handle(ws_data, *fwd_stream); in Execute()
107 context_.ws_mem->set_data_handle(ws_data); in Execute()
119 context_.ws_mem->set_data_handle(DummyData); in Execute()
171 context_.ws_mem.reset( in Setup()
175 {MKLDNN_ARG_WORKSPACE, *context_.ws_mem}, in Setup()
198 context_.ws_mem->set_data_handle(const_cast<void*>(ws_data), *bwd_stream); in Execute()
206 context_.ws_mem->set_data_handle(const_cast<void*>(ws_data)); in Execute()
217 context_.ws_mem->set_data_handle(DummyData); in Execute()
Dmkl_fused_batch_norm_op.cc100 context_.ws_mem->set_data_handle(workspace_data, *fwd_stream); in Execute()
117 context_.ws_mem->set_data_handle(workspace_data); in Execute()
137 context_.ws_mem->set_data_handle(DummyData); in Execute()
162 std::shared_ptr<mkldnn::memory> ws_mem; member
181 ws_mem(nullptr), in BatchNormFwdContext()
230 context_.ws_mem.reset(new memory(context_.fwd_pd->workspace_desc(), in Setup()
256 {MKLDNN_ARG_WORKSPACE, *context_.ws_mem}}); in Setup()
272 {MKLDNN_ARG_WORKSPACE, *context_.ws_mem}}); in Setup()
291 {MKLDNN_ARG_WORKSPACE, *context_.ws_mem}}); in Setup()
307 {MKLDNN_ARG_WORKSPACE, *context_.ws_mem}}); in Setup()
Dmkl_pooling_ops_common.h117 std::shared_ptr<mkldnn::memory> ws_mem; member
140 ws_mem(nullptr), in PoolingFwdContext()
258 std::shared_ptr<mkldnn::memory> ws_mem; member
283 ws_mem(nullptr), in PoolingBwdContext()