Improve error reporting in GEMM benchmark

NVIDIA · Feb 18, 2025 · 4088b47 · 4088b47
1 parent 1b49bdf
commit 4088b47
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/warp/examples/benchmarks/benchmark_gemm.py b/warp/examples/benchmarks/benchmark_gemm.py
@@ -95,7 +95,11 @@ def benchmark_warp(A: wp.array, B: wp.array, config: List[int], warm_up: int, it
 
     # check output
     if warm_up > 0:
-        assert np.allclose(output.numpy(), A.numpy() @ B.numpy(), atol=1e-3, rtol=1e-3)
+        try:
+            np.testing.assert_allclose(output.numpy(), A.numpy() @ B.numpy(), atol=1e-3, rtol=1e-3)
+        except AssertionError as e:
+            print(f"Failed with {TILE_M=}, {TILE_N=}, {TILE_K=}, {BLOCK_DIM=}")
+            raise e
 
     # benchmark
     with wp.ScopedTimer("warp", print=False, synchronize=True, cuda_filter=wp.TIMING_KERNEL) as timer: