| 263 | } |
| 264 | |
| 265 | void AllocateWorkspaceTensor( |
| 266 | OpKernelContext* context, |
| 267 | const lrn_forward::primitive_desc& lrn_fwd_prim_desc, |
| 268 | MklDnnData<uint8>* dnn_data_wksp) { |
| 269 | DCHECK(dnn_data_wksp != nullptr); |
| 270 | Tensor* workspace_tensor = nullptr; |
| 271 | MEMORY_PRIMITIVE_DESC workspace_pd = |
| 272 | lrn_fwd_prim_desc.PRIMITIVE_DESC_WORKSPACE; |
| 273 | size_t workspace_bytes = workspace_pd.get_size(); |
| 274 | MklDnnShape workspace_mkl_shape; |
| 275 | // the workspace tensor is a uint8 tensor that has |
| 276 | // exactly the number of bytes necessary |
| 277 | workspace_mkl_shape.SetMklTensor(false); |
| 278 | TensorShape workspace_tf_shape; |
| 279 | workspace_tf_shape.AddDim(workspace_bytes); |
| 280 | AllocateOutputSetMklShape(context, kIdxWorkspace, &workspace_tensor, |
| 281 | workspace_tf_shape, workspace_mkl_shape); |
| 282 | DCHECK(workspace_tensor != nullptr); |
| 283 | dnn_data_wksp->SetUsrMem(workspace_pd, workspace_tensor); |
| 284 | } |
| 285 | |
| 286 | void SanityCheckInputs(OpKernelContext* context) { |
| 287 | const Tensor& src_tensor = MklGetInput(context, kIdxInput); |
nothing calls this directly
no test coverage detected