benchdnn: matmul: ref: switch weights to ba

Author: dzarukin

tests/benchdnn/dnnl_memory.cpp

@@ -68,6 +68,18 @@ dnn_mem_t::dnn_mem_t(const_dnnl_memory_desc_t md, dnnl_data_type_t dt,
     }
 }
 
+dnn_mem_t::dnn_mem_t(const_dnnl_memory_desc_t md, dnnl_data_type_t dt,
+        const dnnl_dims_t strides, dnnl_engine_t engine) {
+    const int ndims = query_md_ndims(md);
+    if (ndims > 0) {
+        auto status = dnnl_memory_desc_create_with_strides(
+                &md_, ndims, query_md_dims(md), dt, strides);
+        (void)status;
+        assert(status == dnnl_success);
+        active_ = (initialize(engine) == OK);
+    }
+}
+
 dnn_mem_t::dnn_mem_t(int ndims, const dnnl_dims_t dims, dnnl_data_type_t dt,
         const std::string &tag, dnnl_engine_t engine) {
     if (ndims > 0) {

tests/benchdnn/dnnl_memory.hpp

@@ -47,8 +47,11 @@ struct dnn_mem_t {
     dnn_mem_t() { map(); }
     dnn_mem_t(const_dnnl_memory_desc_t md, dnnl_engine_t engine,
             const handle_info_t &handle_info = handle_info_t::allocate());
+
     dnn_mem_t(const_dnnl_memory_desc_t md, dnnl_data_type_t dt,
             const std::string &tag, dnnl_engine_t engine);
+    dnn_mem_t(const_dnnl_memory_desc_t md, dnnl_data_type_t dt,
+            const dnnl_dims_t strides, dnnl_engine_t engine);
 
     dnn_mem_t(int ndims, const dnnl_dims_t dims, dnnl_data_type_t dt,
             const std::string &tag, dnnl_engine_t engine);

tests/benchdnn/matmul/matmul.cpp

@@ -841,9 +841,24 @@ int init_ref_memory_args(dnn_mem_map_t &ref_mem_map, dnn_mem_map_t &mem_map,
         } else
 #endif
         {
-            // Scratchpad memory relates to a primitive. If reference needs it,
-            // use switch below to define a memory desc for it.
-            if (exec_arg != DNNL_ARG_SCRATCHPAD) {
+            if (exec_arg == DNNL_ARG_WEIGHTS) {
+                // Switch the format tag from "ab" to "ba" but to handle batched
+                // cases, use strides instead.
+                const auto ndims = mem.ndims();
+                const auto &dims = mem.dims();
+                dnnl_dims_t strides {};
+                dnnl_dim_t stride = 1;
+                for (int d = ndims - 2; d >= 0; d--) {
+                    strides[d] = stride * dims[d + 1];
+                    stride = strides[d];
+                }
+                strides[ndims - 2] = 1;
+                strides[ndims - 1] = dims[ndims - 2];
+                ref_mem_map.emplace(exec_arg,
+                        dnn_mem_t(mem.md_, dnnl_f32, strides, ref_engine));
+            } else if (exec_arg != DNNL_ARG_SCRATCHPAD) {
+                // Scratchpad memory relates to a primitive. If reference needs
+                // it, use switch below to define a memory desc for it.
                 ref_mem_map.emplace(exec_arg,
                         dnn_mem_t(mem.md_, dnnl_f32, tag::abx, ref_engine));
             }

tests/benchdnn/matmul/matmul.hpp

@@ -280,10 +280,6 @@ inline int64_t src_off_f(const prb_t *prb, int64_t mb, int64_t m, int64_t k) {
     return (mb * prb->m + m) * prb->k + k;
 }
 
-inline int64_t wei_off_f(const prb_t *prb, int64_t mb, int64_t k, int64_t n) {
-    return (mb * prb->k + k) * prb->n + n;
-}
-
 inline int64_t dst_off_f(const prb_t *prb, int64_t mb, int64_t m, int64_t n) {
     return (mb * prb->m + m) * prb->n + n;
 }

tests/benchdnn/matmul/ref_matmul.cpp

@@ -22,6 +22,13 @@
 
 namespace matmul {
 
+int64_t wei_ab_off_f(const prb_t *prb, int64_t mb, int64_t k, int64_t n) {
+    return (mb * prb->k + k) * prb->n + n;
+}
+int64_t wei_ba_off_f(const prb_t *prb, int64_t mb, int64_t k, int64_t n) {
+    return (mb * prb->n + n) * prb->k + k;
+}
+
 void compute_ref_matmul(const prb_t *prb, const args_t &args) {
     const dnn_mem_t &src_m = args.find(DNNL_ARG_SRC);
     const dnn_mem_t &wei_m = args.find(DNNL_ARG_WEIGHTS);
@@ -130,8 +137,9 @@ void compute_ref_matmul(const prb_t *prb, const args_t &args) {
         for (int64_t gK = 0; gK < n_k_groups; gK++) {
             const auto src_gK_off
                     = src_off_f(prb, src_mb, m, gK * smallest_k_group);
+            // Note: scales/zero-points are still always in `tag::abx` format.
             const auto wei_gK_off
-                    = wei_off_f(prb, wei_mb, gK * smallest_k_group, n);
+                    = wei_ab_off_f(prb, wei_mb, gK * smallest_k_group, n);
 
             if (has_src_zp && !has_src_single_zp) {
                 const auto src_zp_idx = src_m.get_idx(
@@ -158,8 +166,8 @@ void compute_ref_matmul(const prb_t *prb, const args_t &args) {
             for (int64_t k = 0; k < smallest_k_group; ++k) {
                 const auto src_off
                         = src_off_f(prb, src_mb, m, gK * smallest_k_group + k);
-                const auto wei_off
-                        = wei_off_f(prb, wei_mb, gK * smallest_k_group + k, n);
+                const auto wei_off = wei_ba_off_f(
+                        prb, wei_mb, gK * smallest_k_group + k, n);
 
                 auto s = src_scale * (src_m.get_elem(src_off) - src_zp);
                 auto w = wei_scale * (wei_m.get_elem(wei_off) - wei_zp);
@@ -292,7 +300,7 @@ void compute_ref_sparse_matmul(const prb_t *prb, const args_t &args) {
 
                 for (int64_t k = row_start; k < row_end; k++) {
                     const int64_t wei_idx
-                            = wei_off_f(prb, mb, src_indices[k], n);
+                            = wei_ba_off_f(prb, mb, src_indices[k], n);
                     const float src_val = src_m.get_elem(k, 0);
                     const float wei_val = wei_m.get_elem(wei_idx);
                     dst_val += src_val * wei_val;