Home
last modified time | relevance | path

Searched refs:fwdParams (Results 1 – 8 of 8) sorted by relevance

/external/tensorflow/tensorflow/core/kernels/mkl/
Dmkl_softmax_op.cc49 explicit MklSoftmaxPrimitive(const MklSoftmaxParams& fwdParams) in MklSoftmaxPrimitive() argument
51 Setup(fwdParams); in MklSoftmaxPrimitive()
114 void Setup(const MklSoftmaxParams& fwdParams) { in Setup() argument
116 auto src_format = MklTensorFormatToMklDnnDataFormat(fwdParams.src_fmt); in Setup()
118 new memory::desc({fwdParams.src_dims}, MklDnnType<T>(), src_format)); in Setup()
122 prop_kind::forward_scoring, *context_.src_md, fwdParams.axis)); in Setup()
146 static MklSoftmaxPrimitive<T>* Get(const MklSoftmaxParams& fwdParams) { in Get() argument
151 fwdParams)); in Get()
153 softmax_forward = new MklSoftmaxPrimitive<T>(fwdParams); in Get()
155 fwdParams, softmax_forward); in Get()
[all …]
Dmkl_quantize_op.cc78 const MklReorderWithScaleFwdParams& fwdParams) in MklReorderWithScalePrimitive() argument
81 Setup(fwdParams); in MklReorderWithScalePrimitive()
127 void Setup(const MklReorderWithScaleFwdParams& fwdParams) { in Setup() argument
130 new memory(fwdParams.src_md, cpu_engine_, DummyData)); in Setup()
132 new memory(fwdParams.dst_md, cpu_engine_, DummyData)); in Setup()
135 auto const& post_op_params = fwdParams.post_op_params; in Setup()
160 const MklReorderWithScaleFwdParams& fwdParams) { in Get() argument
164 from, to, fwdParams)); in Get()
166 reorderPrim = new MklReorderWithScalePrimitive(fwdParams); in Get()
168 from, to, reorderPrim, fwdParams); in Get()
[all …]
Dmkl_pooling_ops_common.cc31 void MklPoolingFwdPrimitive<T>::Setup(const MklPoolingParams& fwdParams) { in Setup() argument
32 DCHECK(fwdParams.alg_kind == mkldnn::algorithm::pooling_max || in Setup()
33 fwdParams.alg_kind == mkldnn::algorithm::pooling_avg || in Setup()
34 fwdParams.alg_kind == mkldnn::algorithm::pooling_avg_include_padding || in Setup()
35 fwdParams.alg_kind == mkldnn::algorithm::pooling_avg_exclude_padding) in Setup()
38 context_.alg_kind = fwdParams.alg_kind; in Setup()
39 context_.prop_kind = fwdParams.prop_kind; in Setup()
46 context_.src_md.reset(new memory::desc(fwdParams.src_md.data)); in Setup()
47 context_.dst_md.reset(new memory::desc({fwdParams.dst_dims}, MklDnnType<T>(), in Setup()
48 fwdParams.native_format in Setup()
[all …]
Dmkl_pooling_ops_common.h75 explicit MklPoolingFwdPrimitive(const MklPoolingParams& fwdParams) in MklPoolingFwdPrimitive() argument
77 if (context_.fwd == nullptr) Setup(fwdParams); in MklPoolingFwdPrimitive()
97 void Setup(const MklPoolingParams& fwdParams);
156 static MklPoolingFwdPrimitive<T>* Get(const MklPoolingParams& fwdParams) { in Get() argument
162 fwdParams)); in Get()
165 pooling_forward = new MklPoolingFwdPrimitive<T>(fwdParams); in Get()
167 fwdParams, pooling_forward); in Get()
185 static string CreateKey(const MklPoolingParams& fwdParams) { in CreateKey() argument
189 key_creator.AddAsKey(fwdParams.src_dims); in CreateKey()
190 key_creator.AddAsKey(fwdParams.dst_dims); in CreateKey()
[all …]
Dmkl_relu_op.cc62 explicit MklEltwiseFwdPrimitive(const MklEltwiseFwdParams<T>& fwdParams) in MklEltwiseFwdPrimitive() argument
66 Setup(fwdParams); in MklEltwiseFwdPrimitive()
136 void Setup(const MklEltwiseFwdParams<T>& fwdParams) { in Setup() argument
138 context_.src_md.reset(new memory::desc(fwdParams.src_md.data)); in Setup()
143 prop_kind::forward, fwdParams.alg_kind, *context_.src_md, in Setup()
144 fwdParams.alpha, fwdParams.beta)); in Setup()
168 const MklEltwiseFwdParams<T>& fwdParams) { in Get() argument
174 fwdParams)); in Get()
176 eltwise_forward = new MklEltwiseFwdPrimitive<T>(fwdParams); in Get()
178 fwdParams, eltwise_forward); in Get()
[all …]
Dmkl_fused_batch_norm_op.cc65 explicit MklFusedBatchNormFwdPrimitive(const MklBatchNormFwdParams& fwdParams) in MklFusedBatchNormFwdPrimitive() argument
67 if (context_.bn_fwd == nullptr) Setup(fwdParams); in MklFusedBatchNormFwdPrimitive()
185 void Setup(const MklBatchNormFwdParams& fwdParams) { in Setup() argument
187 fwdParams.training in Setup()
190 context_.pkind = fwdParams.training ? prop_kind::forward_training in Setup()
193 if (fwdParams.activation_mode == FusedBNActivationMode::kRelu) { in Setup()
197 auto src_md = fwdParams.src_md; in Setup()
200 context_.pkind, src_md, fwdParams.eps, in Setup()
211 memory::dims s_dims = {2, fwdParams.depth}; in Setup()
212 memory::dims m_dims = {1, fwdParams.depth}; in Setup()
[all …]
Dmkl_avgpooling_op.cc113 MklPoolingParams fwdParams( in Compute() local
119 pooling_fwd = MklPoolingFwdPrimitiveFactory<T>::Get(fwdParams); in Compute()
Dmkl_maxpooling_op.cc139 MklPoolingParams fwdParams( in Compute() local
144 pooling_fwd = MklPoolingFwdPrimitiveFactory<T>::Get(fwdParams); in Compute()