benchdnn: fix and limit scratchpad check to CPU

oneapi-src · Jan 23, 2025 · 64cc1c9 · 64cc1c9
1 parent d702ab2
commit 64cc1c9
Showing 1 changed file with 5 additions and 3 deletions.
diff --git a/tests/benchdnn/dnnl_common.cpp b/tests/benchdnn/dnnl_common.cpp
@@ -1141,13 +1141,15 @@ static int check_total_size(
         //
         // 0.75 value supposed to be experimental and might be adjusted.
         static constexpr float scratch_trh = 0.75f;
-        if (check_mem_size_args.scratchpad_size
-                > scratch_trh * total_size_cpu) {
+        if (is_cpu()
+                && check_mem_size_args.scratchpad_size
+                        > scratch_trh * check_mem_size_args.total_size_device) {
             BENCHDNN_PRINT(2,
                     "[CHECK_MEM][%s]: CPU scratchpad size `%zu` exceeded a "
                     "given threshold `%zu`.\n",
                     dir_c_str(), check_mem_size_args.scratchpad_size,
-                    (size_t)(scratch_trh * total_size_cpu));
+                    (size_t)(scratch_trh
+                            * check_mem_size_args.total_size_device));
             res->state = FAILED;
         } else {
             res->state = SKIPPED;