[FORK] [FEATURE] cpu: add inner product with sparse packed weights

3.5 squash list:

[FORK][FIX] separate _blk_off() and _blk_off_sparse() to fix perf issues
[FORK][FEATURE] cpu: add inner product with sparse packed weights

[FORK][FIX] changed comp_tile_len data type from int16_t to int
[FORK][FEATURE] cpu: add inner product with sparse packed weights

Author: jianan-gu

include/oneapi/dnnl/dnnl.h

@@ -976,6 +976,19 @@ dnnl_status_t DNNL_API dnnl_memory_desc_create_with_packed_encoding(
         dnnl_memory_desc_t *memory_desc, int ndims, const dnnl_dims_t dims,
         dnnl_data_type_t data_type, dnnl_dim_t nnz);
 #endif
+/// Initializes a sparse descriptor.
+///
+/// @param memory_desc Output memory descriptor.
+/// @param encoding Encoding.
+/// @param ndims Number of dimensions.
+/// @param dims Array of dimensions.
+/// @param data_type Elements data type.
+/// @returns #dnnl_success on success and a status describing the error
+///     otherwise.
+dnnl_status_t DNNL_API dnnl_memory_desc_create_sparse(
+        dnnl_memory_desc_t *memory_desc,
+        dnnl_sparse_encoding_t encoding, int ndims,
+        const dnnl_dims_t dims, dnnl_data_type_t data_type);
 
 /// Creates a memory descriptor for a region inside an area
 /// described by an existing memory descriptor.
@@ -1230,7 +1243,6 @@ size_t DNNL_API dnnl_memory_desc_get_size(const_dnnl_memory_desc_t memory_desc);
 size_t DNNL_API dnnl_memory_desc_get_size_v2(
         const_dnnl_memory_desc_t memory_desc, int index);
 #endif
-
 /// Returns the size of data type.
 ///
 /// @param data_type Data type.
@@ -1283,7 +1295,6 @@ dnnl_status_t DNNL_API dnnl_memory_create_v2(dnnl_memory_t *memory,
         const_dnnl_memory_desc_t memory_desc, dnnl_engine_t engine,
         int nhandles, void **handles);
 #endif
-
 /// Returns the memory descriptor for a memory object.
 ///
 /// @param memory Memory object.
@@ -1395,7 +1406,6 @@ dnnl_status_t DNNL_API dnnl_memory_unmap_data(
 dnnl_status_t DNNL_API dnnl_memory_unmap_data_v2(
         const_dnnl_memory_t memory, void *mapped_ptr, int index);
 #endif
-
 /// Returns memory object's data handle.
 ///
 /// @param memory Memory object.

include/oneapi/dnnl/dnnl.hpp

@@ -923,6 +923,7 @@ struct memory : public handle<dnnl_memory_t> {
         /// Format kind for sparse tensors.
         sparse = dnnl_format_kind_sparse,
 #endif
+        sparsed = dnnl_format_sparse,
         /// A special format kind that indicates that tensor format is opaque.
         opaque = dnnl_format_kind_opaque,
     };
@@ -2753,7 +2754,6 @@ struct memory : public handle<dnnl_memory_t> {
     /// A memory descriptor.
     struct desc : public handle<dnnl_memory_desc_t> {
         using handle<dnnl_memory_desc_t>::handle;
-
         friend struct memory;
 
         /// Constructs a zero (empty) memory descriptor. Such a memory
@@ -2944,6 +2944,31 @@ struct memory : public handle<dnnl_memory_t> {
             reset(md);
         }
 
+        /// @fork
+        /// Copy constructor for memory::desc
+        /// Ensures deep copy (underlying C structure is copied as well)
+        /// To preserve behavior of 2.x oneDNN versions
+        ///
+        /// @param desc memory descriptor to copy.
+        desc(const memory::desc& adesc) {
+            auto cdesc = adesc.get();
+            dnnl_memory_desc_t cloned_md = nullptr;
+            dnnl_memory_desc_clone(&cloned_md, cdesc);
+
+            reset(cloned_md);
+        }
+
+        desc sparse_desc(const dims &adims, data_type adata_type,
+                bool allow_empty = false) {
+            dnnl_memory_desc_t md = nullptr;
+            dnnl_status_t status = dnnl_memory_desc_create_sparse(&md, dnnl_sparse_encoding_packed,
+                    (int)adims.size(), adims.data(), convert_to_c(adata_type));
+
+            if (!allow_empty)
+                error::wrap_c_api(status,
+                        "could not construct a memory descriptor with sparse format");
+            return desc(md);
+        }
         /// Constructs a memory descriptor for a region inside an area
         /// described by this memory descriptor.
         //
@@ -3192,9 +3217,9 @@ struct memory : public handle<dnnl_memory_t> {
         /// Returns the data type of the memory descriptor.
         ///
         /// @returns The data type.
-        memory::data_type get_data_type() const {
-            return query_data_type(query::data_type);
-        }
+        // memory::data_type get_data_type() const {
+        //     return query_data_type(query::data_type);
+        // }
 #endif
 
         /// Returns the format kind of the memory descriptor.
@@ -3209,6 +3234,30 @@ struct memory : public handle<dnnl_memory_t> {
                     : dnnl::memory::format_kind::undef;
         }
 
+        /// Returns the format kind of the memory descriptor.
+        ///
+        /// @returns the format kind.
+        dnnl_sparse_encoding_t get_sparse_encoding() const {
+            dnnl_sparse_encoding_t sparse_encoding;
+            dnnl_status_t status = dnnl_memory_desc_query(
+                get(), dnnl_query_sparse_encoding, &sparse_encoding);
+            return status == dnnl_success
+                    ? sparse_encoding
+                    : dnnl_sparse_encoding_undef;
+        }
+
+        /// Returns the data type of the memory descriptor.
+        ///
+        /// @returns The data type.
+        memory::data_type get_data_type() const {
+            dnnl_data_type_t data_type;
+            dnnl_status_t status = dnnl_memory_desc_query(
+                    get(), dnnl_query_data_type, &data_type);
+            return status == dnnl_success
+                    ? static_cast<dnnl::memory::data_type>(data_type)
+                    : dnnl::memory::data_type::undef;
+        }
+
         /// Returns dimensions of the memory descriptor.
         ///
         /// Potentially expensive due to the data copy involved.
@@ -3386,6 +3435,44 @@ struct memory : public handle<dnnl_memory_t> {
         reset(result);
     }
 #else
+    /// Constructs a memory object.
+    ///
+    /// Unless @p handle is equal to #DNNL_MEMORY_NONE, the constructed memory
+    /// object will have the underlying buffer set. In this case, the buffer
+    /// will be initialized as if #dnnl::memory::set_data_handle() had been
+    /// called.
+    ///
+    /// @sa memory::set_data_handle()
+    ///
+    /// @param md Memory descriptor.
+    /// @param aengine Engine to store the data on.
+    /// @param handle Handle of the memory buffer to use.
+    ///     - A pointer to the user-allocated buffer. In this case the library
+    ///       doesn't own the buffer.
+    ///     - The #DNNL_MEMORY_ALLOCATE special value. Instructs the library to
+    ///       allocate the buffer for the memory object. In this case the
+    ///       library owns the buffer.
+    ///     - #DNNL_MEMORY_NONE to create dnnl::memory without an underlying
+    ///       buffer.
+    // memory(const desc &md, const engine &aengine, void *handle) {
+    //     dnnl_memory_t result;
+    //     error::wrap_c_api(
+    //             dnnl_memory_create(&result, md.get(), aengine.get(), handle),
+    //             "could not create a memory object");
+    //     reset(result);
+    // }
+
+    /// Constructs a memory object.
+    ///
+    /// The underlying buffer(s) for the memory will be allocated by the
+    /// library.
+    ///
+    /// @param md Memory descriptor.
+    /// @param aengine Engine to store the data on.
+    // memory(const desc &md, const engine &aengine)
+    //     : memory(md, aengine, DNNL_MEMORY_ALLOCATE) {}
+#endif
+
     /// Constructs a memory object.
     ///
     /// Unless @p handle is equal to #DNNL_MEMORY_NONE, the constructed memory
@@ -3413,15 +3500,11 @@ struct memory : public handle<dnnl_memory_t> {
         reset(result);
     }
 
-    /// Constructs a memory object.
-    ///
     /// The underlying buffer for the memory will be allocated by the library.
-    ///
     /// @param md Memory descriptor.
     /// @param aengine Engine to store the data on.
     memory(const desc &md, const engine &aengine)
         : memory(md, aengine, DNNL_MEMORY_ALLOCATE) {}
-#endif
 
     /// Returns the associated memory descriptor.
     desc get_desc() const {

include/oneapi/dnnl/dnnl_debug.h

@@ -44,6 +44,7 @@ const char DNNL_API *dnnl_fmt_tag2str(dnnl_format_tag_t v);
 const char DNNL_API *dnnl_prop_kind2str(dnnl_prop_kind_t v);
 const char DNNL_API *dnnl_prim_kind2str(dnnl_primitive_kind_t v);
 const char DNNL_API *dnnl_alg_kind2str(dnnl_alg_kind_t v);
+const char DNNL_API *dnnl_sparse_encoding2str(dnnl_sparse_encoding_t v);
 const char DNNL_API *dnnl_rnn_flags2str(dnnl_rnn_flags_t v);
 const char DNNL_API *dnnl_rnn_direction2str(dnnl_rnn_direction_t v);
 const char DNNL_API *dnnl_scratchpad_mode2str(dnnl_scratchpad_mode_t v);

include/oneapi/dnnl/dnnl_types.h

@@ -55,6 +55,8 @@ typedef enum {
     /// Format kind for sparse tensors.
     dnnl_format_kind_sparse,
 #endif
+    /// Format for sparse data.
+    dnnl_format_sparse,
     /// Parameter to allow internal only format kinds without undefined
     /// behavior. This parameter is chosen to be valid for so long as
     /// sizeof(int) >= 2.
@@ -2322,6 +2324,52 @@ typedef struct dnnl_memory_desc *dnnl_memory_desc_t;
 /// A memory descriptor handle.
 typedef const struct dnnl_memory_desc *const_dnnl_memory_desc_t;
 
+/// Sparse encodings.
+typedef enum {
+    dnnl_sparse_encoding_undef = 0,
+    dnnl_sparse_encoding_any,
+    dnnl_sparse_encoding_packed,
+    dnnl_sparse_encoding_csr,
+    dnnl_sparse_encoding_coo,
+} dnnl_sparse_encoding_t;
+
+/* typedef struct dnnl_sparse_desc *dnnl_sparse_desc_t; */
+/* typedef const struct dnnl_sparse_desc *const_dnnl_sparse_desc_t; */
+
+/// Flags for memory special features
+typedef enum {
+    dnnl_memory_extra_flag_none = 0x0U,
+    /// Indicates the weights have an additional buffer, that depends on the
+    /// @p compensation_mask.
+    ///
+    /// For instance, in 4D case with the compensation mask equals (1 << 0)
+    /// the additional buffer would consist of OC values:
+    /// O[oc : 0,OC] =
+    ///  -128 * SUM(ic : 0,IC; kh : 0,KH; kw : 0,KW){ weights(oc, ic, kh, kw) }
+    dnnl_memory_extra_flag_compensation_conv_s8s8 = 0x1U,
+    dnnl_memory_extra_flag_scale_adjust = 0x2U,
+    dnnl_memory_extra_flag_rnn_u8s8_compensation = 0x4U,
+    dnnl_memory_extra_flag_gpu_rnn_u8s8_compensation
+    = dnnl_memory_extra_flag_rnn_u8s8_compensation,
+    dnnl_memory_extra_flag_compensation_conv_asymmetric_src = 0x8U,
+    dnnl_memory_extra_flag_rnn_s8s8_compensation = 0x16U,
+} dnnl_memory_extra_flags_t;
+
+/// Description of extra information stored in memory
+typedef struct {
+    /// The flags contain arbitrary extra information, such as compensation.
+    /// @sa dnnl_memory_extra_flags_t
+    uint64_t flags;
+    /// Compensation mask
+    int compensation_mask;
+    /// Scale applied to the data
+    float scale_adjust;
+    /// Compensation mask for asymmetric quantization
+    int asymm_compensation_mask;
+    /// For future backwards compatibility
+    char reserved[60];
+} dnnl_memory_extra_desc_t;
+
 /// @struct dnnl_memory
 /// An opaque structure to describe a memory.
 struct dnnl_memory;
@@ -2836,6 +2884,8 @@ typedef enum {
     dnnl_query_num_handles_s32, ///< Number of buffers required for a memory
 ///  descriptor
 #endif
+    dnnl_query_sparse_encoding,
+
     // Max value to prevent UB for internal use only dnnl_query_t
     dnnl_query_max = 0x7fff,
 } dnnl_query_t;

scripts/generate_dnnl_debug.py

@@ -129,6 +129,7 @@ def header_benchdnn(body):
 #ifdef DNNL_EXPERIMENTAL_SPARSE
 const char *sparse_encoding2str(dnnl_sparse_encoding_t encoding);
 #endif
+const char *sparse_encoding2str(dnnl_sparse_encoding_t encoding);
 
 /* engine kind */
 const char *engine_kind2str(dnnl_engine_kind_t kind);
@@ -183,6 +184,9 @@ def source_benchdnn(body):
     return dnnl_sparse_encoding2str(encoding);
 }
 #endif
+const char *sparse_encoding2str(dnnl_sparse_encoding_t encoding) {
+    return dnnl_sparse_encoding2str(encoding);
+}
 
 const char *engine_kind2str(dnnl_engine_kind_t kind) {
     return dnnl_engine_kind2str(kind);

src/common/c_types_map.hpp

@@ -217,13 +217,13 @@ const sparse_encoding_t packed = dnnl_packed;
 } // namespace sparse_encoding
 #else
 // Declare dummy values to avoid guarding internal implementation.
-using sparse_encoding_t = int;
-namespace sparse_encoding {
-const sparse_encoding_t undef = 0;
-const sparse_encoding_t csr = 1;
-const sparse_encoding_t packed = 2;
-const sparse_encoding_t coo = 3;
-} // namespace sparse_encoding
+// using sparse_encoding_t = int;
+// namespace sparse_encoding {
+// const sparse_encoding_t undef = 0;
+// const sparse_encoding_t csr = 1;
+// const sparse_encoding_t packed = 2;
+// const sparse_encoding_t coo = 3;
+// } // namespace sparse_encoding
 #endif
 
 using format_kind_t = dnnl_format_kind_t;
@@ -235,14 +235,15 @@ const format_kind_t opaque = dnnl_format_kind_opaque;
 #ifdef DNNL_EXPERIMENTAL_SPARSE
 const format_kind_t sparse = dnnl_format_kind_sparse;
 #else
-const format_kind_t sparse = static_cast<format_kind_t>(4);
+// const format_kind_t sparse = static_cast<format_kind_t>(4);
 #endif
 
 // Internal only format kinds.
 const format_kind_t internal_only_start = (format_kind_t)(1 << 8);
 const format_kind_t wino = internal_only_start;
 const format_kind_t rnn_packed = (format_kind_t)(internal_only_start + 1);
 const format_kind_t cublaslt_blocked = (format_kind_t)(internal_only_start + 2);
+const format_kind_t sparse = dnnl_format_sparse;
 } // namespace format_kind
 
 #ifdef DNNL_EXPERIMENTAL_PROFILING
@@ -1915,6 +1916,15 @@ const rnn_flags_t diff_weights_overwrite
         = dnnl_rnn_flags_diff_weights_overwrite;
 } // namespace rnn_flags
 
+using sparse_encoding_t = dnnl_sparse_encoding_t;
+namespace sparse_encoding {
+const sparse_encoding_t undef = dnnl_sparse_encoding_undef;
+const sparse_encoding_t any = dnnl_sparse_encoding_any;
+const sparse_encoding_t packed = dnnl_sparse_encoding_packed;
+const sparse_encoding_t csr = dnnl_sparse_encoding_csr;
+const sparse_encoding_t coo = dnnl_sparse_encoding_coo;
+} // namespace sparse_encoding
+
 using engine_kind_t = dnnl_engine_kind_t;
 namespace engine_kind {
 const engine_kind_t any_engine = dnnl_any_engine;
@@ -2051,15 +2061,16 @@ const query_t sparse_encoding = dnnl_query_sparse_encoding;
 const query_t nnz_s64 = dnnl_query_nnz_s64;
 const query_t num_handles_s32 = dnnl_query_num_handles_s32;
 #else
-const query_t sparse_encoding = static_cast<query_t>(266);
-const query_t nnz_s64 = static_cast<query_t>(267);
-const query_t num_handles_s32 = static_cast<query_t>(268);
+// const query_t sparse_encoding = static_cast<query_t>(266);
+// const query_t nnz_s64 = static_cast<query_t>(267);
+// const query_t num_handles_s32 = static_cast<query_t>(268);
 #endif
 
 // Internal only query kinds.
 const query_t internal_only_start = (query_t)(1 << 12);
 const query_t zero_pad_d = internal_only_start;
 const query_t preferred_gpu_threads_per_eu = (query_t)(internal_only_start + 1);
+const query_t sparse_encoding = dnnl_query_sparse_encoding;
 } // namespace query
 
 using rnn_direction_t = dnnl_rnn_direction_t;

src/common/dnnl_debug.cpp

@@ -54,6 +54,7 @@ const char *dnnl_fmt_kind2str(dnnl_format_kind_t v) {
             || v == format_kind::cublaslt_blocked)
         return "opaque";
     if (v == dnnl_format_kind_max) return "max";
+    if (v == dnnl_format_sparse) return "format_sparse";
     assert(!"unknown fmt_kind");
     return "unknown fmt_kind";
 }

src/common/dnnl_debug_autogenerated.cpp

@@ -1854,6 +1854,14 @@ const char *dnnl_alg_kind2str(dnnl_alg_kind_t v) {
     return "unknown alg_kind";
 }
 
+const char *dnnl_sparse_encoding2str(dnnl_sparse_encoding_t v) {
+    if (v == dnnl_sparse_encoding_undef) return "undef";
+    if (v == dnnl_sparse_encoding_any) return "any";
+    if (v == dnnl_sparse_encoding_packed) return "sparse_encoding_packed";
+    assert(!"unknown sparse_encoding");
+    return "unknown sparse_encoding";
+}
+
 const char *dnnl_rnn_flags2str(dnnl_rnn_flags_t v) {
     if (v == dnnl_rnn_flags_undef) return "undef";
     if (v == dnnl_rnn_flags_diff_weights_overwrite) return "rnn_flags_diff_weights_overwrite";