benchdnn: mem check: touch up

tests/benchdnn/binary/binary.cpp

@@ -294,6 +294,8 @@ int checkit(std::vector<benchdnn_dnnl_wrapper_t<dnnl_primitive_t>> &v_prim,
 
 int doit(const std::vector<benchdnn_dnnl_wrapper_t<dnnl_primitive_t>> &v_prim,
         const prb_t *prb, res_t *res) {
+    set_zmalloc_max_expected_size(res->mem_size_args.zmalloc_expected_size);
+
     const auto &prim = v_prim[0];
 
     dnn_mem_map_t mem_map, ref_mem_map;

tests/benchdnn/bnorm/bnorm.cpp

@@ -733,6 +733,8 @@ int checkit(std::vector<benchdnn_dnnl_wrapper_t<dnnl_primitive_t>> &v_prim,
 
 int doit(const std::vector<benchdnn_dnnl_wrapper_t<dnnl_primitive_t>> &v_prim,
         const prb_t *prb, res_t *res) {
+    set_zmalloc_max_expected_size(res->mem_size_args.zmalloc_expected_size);
+
     const auto &prim = prb->dir & FLAG_FWD ? v_prim[0] : v_prim[1];
 
     dnn_mem_map_t mem_map, ref_mem_map;

tests/benchdnn/common.cpp

@@ -267,23 +267,33 @@ struct memory_registry_t {
                 "allocation size: %s\n",
                 smart_bytes(size).c_str(), smart_bytes(total_size_).c_str());
         allocations_.erase(ptr);
+        if (has_warned_) exit(2);
     }
 
     void set_expected_max(size_t size) {
-        expected_max_ = size;
+        constexpr float expected_trh = 1.1f; // Smooth out small allocations.
+        expected_max_ = static_cast<size_t>(expected_trh * size);
         has_warned_ = false;
         warn_size_check();
     }
 
 private:
     size_t size() const { return total_size_; }
     void warn_size_check() {
-        if (expected_max_ != unset_ && !has_warned_
-                && total_size_ > expected_max_) {
-            // Switch to WARNING once existing failures are resolved
+        const bool is_max_set = expected_max_ != unset_;
+        // Verify the total amount of allocated memory when it starts exceeding
+        // 1 GB threshold. Small amount of memory is highly unlikely cause OOM.
+        // There's an idea to add a portion of RAM into account as well, keep
+        // only 1 GB so far to check if it proves working well.
+        const bool is_total_size_big = total_size_ >= 1024 * 1024 * 1024;
+        const bool is_total_size_unexpected = total_size_ > expected_max_;
+        if (!has_warned_ && is_max_set && is_total_size_big
+                && is_total_size_unexpected) {
+            // TODO: switch to WARNING once existing failures are resolved
             BENCHDNN_PRINT(0,
-                    "[CHECK_MEM][INFO]: memory use underestimated, "
-                    "zmalloc allocations exceed %s\n",
+                    "[CHECK_MEM][ERROR]: Memory use is underestimated. Current "
+                    "allocation size: %s; expected size: %s.\n",
+                    smart_bytes(total_size_).c_str(),
                     smart_bytes(expected_max_).c_str());
             // Prevent spamming logs with subsequent overflowing allocations;
             has_warned_ = true;

tests/benchdnn/concat/concat.cpp

@@ -224,6 +224,8 @@ int checkit(std::vector<benchdnn_dnnl_wrapper_t<dnnl_primitive_t>> &v_prim,
 
 int doit(const std::vector<benchdnn_dnnl_wrapper_t<dnnl_primitive_t>> &v_prim,
         const prb_t *prb, res_t *res) {
+    set_zmalloc_max_expected_size(res->mem_size_args.zmalloc_expected_size);
+
     const auto &prim = v_prim[0];
 
     dnn_mem_map_t mem_map, ref_mem_map;

tests/benchdnn/conv/conv.cpp

@@ -578,6 +578,8 @@ int checkit(std::vector<benchdnn_dnnl_wrapper_t<dnnl_primitive_t>> &v_prim,
 
 int doit(const std::vector<benchdnn_dnnl_wrapper_t<dnnl_primitive_t>> &v_prim,
         const prb_t *prb, res_t *res) {
+    set_zmalloc_max_expected_size(res->mem_size_args.zmalloc_expected_size);
+
     const auto &prim = v_prim[0];
     const auto &prim_ref = v_prim[1];
 

tests/benchdnn/deconv/deconv.cpp

@@ -558,6 +558,8 @@ int checkit(std::vector<benchdnn_dnnl_wrapper_t<dnnl_primitive_t>> &v_prim,
 
 int doit(const std::vector<benchdnn_dnnl_wrapper_t<dnnl_primitive_t>> &v_prim,
         const prb_t *prb, res_t *res) {
+    set_zmalloc_max_expected_size(res->mem_size_args.zmalloc_expected_size);
+
     const auto &prim = v_prim[0];
     const auto &prim_ref = v_prim[1];
 

tests/benchdnn/dnnl_common.cpp

@@ -1095,6 +1095,12 @@ int get_gpu_cache_size(size_t &cache_size) {
 // If no, indicate that the system won't make it and drop `prim_ref` falling
 // back to stock reference.
 int check_total_size(res_t *res, dnnl_primitive_t prim_ref) {
+    // Skip the check if it is disabled.
+    if (!mem_check) return OK;
+
+    // Skip the check if the test object won't be executed.
+    if (!has_bench_mode_bit(mode_bit_t::exec)) return OK;
+
     static size_t cpu_device_capacity = get_cpu_ram_size();
     static size_t gpu_device_capacity = 0;
     static size_t gpu_max_alloc_capacity = 0;
@@ -1210,10 +1216,10 @@ int check_total_size(res_t *res, dnnl_primitive_t prim_ref) {
     bool fits_cpu_ram = cpu_and_device_size
             <= (is_cpu() ? benchdnn_cpu_limit : benchdnn_combined_limit);
 
-    set_zmalloc_max_expected_size(
-            is_cpu() ? cpu_and_device_size : total_size_cpu);
-    // Check combined size against CPU capacity as the simpler method to account
-    // for integrated devices and mapping/unmapping memory.
+    // Save the expected value to set at `doit`, otherwise, the check doesn't
+    // work correctly. See `zmalloc_expected_size` comment.
+    res->mem_size_args.zmalloc_expected_size
+            = is_cpu() ? cpu_and_device_size : total_size_cpu;
 
     if (!fits_cpu_ram) {
         std::string prim_ref_msg

tests/benchdnn/eltwise/eltwise.cpp

@@ -513,6 +513,8 @@ int checkit(std::vector<benchdnn_dnnl_wrapper_t<dnnl_primitive_t>> &v_prim,
 
 int doit(const std::vector<benchdnn_dnnl_wrapper_t<dnnl_primitive_t>> &v_prim,
         const prb_t *prb, res_t *res) {
+    set_zmalloc_max_expected_size(res->mem_size_args.zmalloc_expected_size);
+
     const auto &prim = prb->dir & FLAG_FWD ? v_prim[0] : v_prim[1];
 
     dnn_mem_map_t mem_map, ref_mem_map;

tests/benchdnn/gnorm/gnorm.cpp

@@ -716,6 +716,8 @@ int checkit(std::vector<benchdnn_dnnl_wrapper_t<dnnl_primitive_t>> &v_prim,
 
 int doit(const std::vector<benchdnn_dnnl_wrapper_t<dnnl_primitive_t>> &v_prim,
         const prb_t *prb, res_t *res) {
+    set_zmalloc_max_expected_size(res->mem_size_args.zmalloc_expected_size);
+
     const auto &prim = v_prim[0];
 
     dnn_mem_map_t mem_map, ref_mem_map;

tests/benchdnn/ip/ip.cpp

@@ -419,6 +419,8 @@ int checkit(std::vector<benchdnn_dnnl_wrapper_t<dnnl_primitive_t>> &v_prim,
 
 int doit(const std::vector<benchdnn_dnnl_wrapper_t<dnnl_primitive_t>> &v_prim,
         const prb_t *prb, res_t *res) {
+    set_zmalloc_max_expected_size(res->mem_size_args.zmalloc_expected_size);
+
     const auto &prim = v_prim[0];
     const auto &prim_ref = v_prim[1];
 

tests/benchdnn/lnorm/lnorm.cpp

@@ -676,6 +676,8 @@ int checkit(std::vector<benchdnn_dnnl_wrapper_t<dnnl_primitive_t>> &v_prim,
 
 int doit(const std::vector<benchdnn_dnnl_wrapper_t<dnnl_primitive_t>> &v_prim,
         const prb_t *prb, res_t *res) {
+    set_zmalloc_max_expected_size(res->mem_size_args.zmalloc_expected_size);
+
     const auto &prim = v_prim[0];
 
     dnn_mem_map_t mem_map, ref_mem_map;

tests/benchdnn/lrn/lrn.cpp

@@ -222,6 +222,8 @@ int checkit(std::vector<benchdnn_dnnl_wrapper_t<dnnl_primitive_t>> &v_prim,
 
 int doit(const std::vector<benchdnn_dnnl_wrapper_t<dnnl_primitive_t>> &v_prim,
         const prb_t *prb, res_t *res) {
+    set_zmalloc_max_expected_size(res->mem_size_args.zmalloc_expected_size);
+
     const auto &prim = prb->dir & FLAG_FWD ? v_prim[0] : v_prim[1];
 
     dnn_mem_map_t mem_map, ref_mem_map;

tests/benchdnn/matmul/matmul.cpp

@@ -975,6 +975,8 @@ std::vector<data_kind_t> get_kinds_to_check(const prb_t *prb) {
 
 int doit(const std::vector<benchdnn_dnnl_wrapper_t<dnnl_primitive_t>> &v_prim,
         const prb_t *prb, res_t *res) {
+    set_zmalloc_max_expected_size(res->mem_size_args.zmalloc_expected_size);
+
     const auto &prim = v_prim[0];
     const auto &prim_ref = v_prim[1];
 

tests/benchdnn/pool/pool.cpp

@@ -361,6 +361,8 @@ int checkit(std::vector<benchdnn_dnnl_wrapper_t<dnnl_primitive_t>> &v_prim,
 
 int doit(const std::vector<benchdnn_dnnl_wrapper_t<dnnl_primitive_t>> &v_prim,
         const prb_t *prb, res_t *res) {
+    set_zmalloc_max_expected_size(res->mem_size_args.zmalloc_expected_size);
+
     const auto &prim = prb->dir & FLAG_FWD ? v_prim[0] : v_prim[1];
 
     dnn_mem_map_t mem_map, ref_mem_map;

tests/benchdnn/prelu/prelu.cpp

@@ -249,6 +249,8 @@ int checkit(std::vector<benchdnn_dnnl_wrapper_t<dnnl_primitive_t>> &v_prim,
 
 int doit(const std::vector<benchdnn_dnnl_wrapper_t<dnnl_primitive_t>> &v_prim,
         const prb_t *prb, res_t *res) {
+    set_zmalloc_max_expected_size(res->mem_size_args.zmalloc_expected_size);
+
     const auto &prim = v_prim[0];
 
     dnn_mem_map_t mem_map, ref_mem_map;

tests/benchdnn/reduction/reduction.cpp

@@ -335,6 +335,8 @@ int checkit(std::vector<benchdnn_dnnl_wrapper_t<dnnl_primitive_t>> &v_prim,
 
 int doit(const std::vector<benchdnn_dnnl_wrapper_t<dnnl_primitive_t>> &v_prim,
         const prb_t *prb, res_t *res) {
+    set_zmalloc_max_expected_size(res->mem_size_args.zmalloc_expected_size);
+
     const auto &prim = v_prim[0];
 
     dnn_mem_map_t mem_map, ref_mem_map;

tests/benchdnn/reorder/reorder.cpp

@@ -589,6 +589,8 @@ int checkit(std::vector<benchdnn_dnnl_wrapper_t<dnnl_primitive_t>> &v_prim,
 
 int doit(const std::vector<benchdnn_dnnl_wrapper_t<dnnl_primitive_t>> &v_prim,
         const prb_t *prb, res_t *res) {
+    set_zmalloc_max_expected_size(res->mem_size_args.zmalloc_expected_size);
+
     const auto &prim = v_prim[0];
 
     dnn_mem_map_t mem_map, ref_mem_map;

tests/benchdnn/resampling/resampling.cpp

@@ -262,6 +262,8 @@ int checkit(std::vector<benchdnn_dnnl_wrapper_t<dnnl_primitive_t>> &v_prim,
 
 int doit(const std::vector<benchdnn_dnnl_wrapper_t<dnnl_primitive_t>> &v_prim,
         const prb_t *prb, res_t *res) {
+    set_zmalloc_max_expected_size(res->mem_size_args.zmalloc_expected_size);
+
     const auto &prim = v_prim[0];
 
     dnn_mem_map_t mem_map, ref_mem_map;

tests/benchdnn/rnn/rnn.cpp

@@ -1266,6 +1266,8 @@ int checkit(std::vector<benchdnn_dnnl_wrapper_t<dnnl_primitive_t>> &v_prim,
 
 int doit(const std::vector<benchdnn_dnnl_wrapper_t<dnnl_primitive_t>> &v_prim,
         const prb_t &prb, res_t *res) {
+    set_zmalloc_max_expected_size(res->mem_size_args.zmalloc_expected_size);
+
     const auto &prim = prb.prop != dnnl_backward ? v_prim[0] : v_prim[1];
 
     dnn_mem_map_t mem_map, ref_mem_map;

tests/benchdnn/shuffle/shuffle.cpp

@@ -194,6 +194,8 @@ int checkit(std::vector<benchdnn_dnnl_wrapper_t<dnnl_primitive_t>> &v_prim,
 
 int doit(const std::vector<benchdnn_dnnl_wrapper_t<dnnl_primitive_t>> &v_prim,
         const prb_t *prb, res_t *res) {
+    set_zmalloc_max_expected_size(res->mem_size_args.zmalloc_expected_size);
+
     const auto &prim = v_prim[0];
 
     dnn_mem_map_t mem_map, ref_mem_map;

tests/benchdnn/softmax/softmax.cpp

@@ -436,6 +436,8 @@ int checkit(std::vector<benchdnn_dnnl_wrapper_t<dnnl_primitive_t>> &v_prim,
 
 int doit(const std::vector<benchdnn_dnnl_wrapper_t<dnnl_primitive_t>> &v_prim,
         const prb_t *prb, res_t *res) {
+    set_zmalloc_max_expected_size(res->mem_size_args.zmalloc_expected_size);
+
     const auto &prim = v_prim[0];
 
     dnn_mem_map_t mem_map, ref_mem_map;

tests/benchdnn/sum/sum.cpp

@@ -187,6 +187,8 @@ int checkit(std::vector<benchdnn_dnnl_wrapper_t<dnnl_primitive_t>> &v_prim,
 
 int doit(const std::vector<benchdnn_dnnl_wrapper_t<dnnl_primitive_t>> &v_prim,
         const prb_t *prb, res_t *res) {
+    set_zmalloc_max_expected_size(res->mem_size_args.zmalloc_expected_size);
+
     const auto &prim = v_prim[0];
 
     dnn_mem_map_t mem_map, ref_mem_map;

tests/benchdnn/utils/res.hpp

@@ -92,6 +92,11 @@ struct check_mem_size_args_t {
     size_t total_ref_md_size[2] = {0, 0};
     // `scratchpad_size` specifies a scratchpad size for specific checks.
     size_t scratchpad_size = 0;
+    // A setting for zmalloc_registry. It's stashed inside `check_total_size`
+    // call and used later in `doit` due to parallel mode as, otherwise, all
+    // test objects will be validated against the numbers from the last created
+    // test object.
+    size_t zmalloc_expected_size = 0;
 };
 
 struct res_t {