Method SetOpMemDesc

tensorflow/core/util/mkl_util.h:1429–1432

Set memory descriptor of an operation in terms of dimensions and memory format. E.g., For Conv2D, the dimensions would be same as user dimensions but memory::format_tag would be dnnl::any because we want OneDNN to choose the best layout/format for given input dimensions.

  /// but memory::format_tag would be dnnl::any because we want OneDNN to
  /// choose the best layout/format for given input dimensions.
  inline void SetOpMemDesc(const memory::dims& dim, memory::format_tag fm) {
    // TODO(nhasabni): can we remove dynamic memory allocation?
    op_md_ = new memory::desc(dim, MklDnnType<T>(), fm);
  }

  /// Get function for memory descriptor for an operation
  inline const memory::desc& GetOpMemDesc() const { return *op_md_; }