From 48dc94ea52efa34a60d797dce123b317996cc806 Mon Sep 17 00:00:00 2001 From: LeiWang1999 Date: Wed, 2 Oct 2024 16:20:24 +0000 Subject: [PATCH] Refactor test_general_matmul_tilelang_impl.py and test_general_matmul_tilelang_kernel.py to use centered random values for input tensors --- .../operators/test_general_matmul_tilelang_impl.py | 12 ++++++------ .../operators/test_general_matmul_tilelang_kernel.py | 4 ++-- 2 files changed, 8 insertions(+), 8 deletions(-) diff --git a/testing/python/operators/test_general_matmul_tilelang_impl.py b/testing/python/operators/test_general_matmul_tilelang_impl.py index 03150f74..5192325d 100644 --- a/testing/python/operators/test_general_matmul_tilelang_impl.py +++ b/testing/python/operators/test_general_matmul_tilelang_impl.py @@ -53,8 +53,8 @@ def assert_matmul_blocked_correctness(M, # src_code is the generated cuda source assert src_code is not None - A = torch.rand(M, K, device="cuda", dtype=getattr(torch, in_dtype)) - B = torch.rand(N, K, device="cuda", dtype=getattr(torch, in_dtype)) + A = torch.rand(M, K, device="cuda", dtype=getattr(torch, in_dtype)) - 0.5 + B = torch.rand(N, K, device="cuda", dtype=getattr(torch, in_dtype)) - 0.5 C = torch.zeros(M, N, device="cuda", dtype=getattr(torch, accum_dtype)) mod = tl.Profiler(mod, params, [], tl.TensorSupplyType.Integer) @@ -111,8 +111,8 @@ def assert_matmul_macro_tensorcore_correctness( # src_code represents generated cuda source assert src_code is not None - A = torch.rand(M, K, device="cuda", dtype=getattr(torch, in_dtype)) - B = torch.rand(N, K, device="cuda", dtype=getattr(torch, in_dtype)) + A = torch.rand(M, K, device="cuda", dtype=getattr(torch, in_dtype)) - 0.5 + B = torch.rand(N, K, device="cuda", dtype=getattr(torch, in_dtype)) - 0.5 C = torch.zeros(M, N, device="cuda", dtype=getattr(torch, accum_dtype)) mod = tl.Profiler(mod, params, [], tl.TensorSupplyType.Integer) @@ -170,8 +170,8 @@ def assert_tl_matmul_with_ladder_weight_only_transform_correctness( # src_code is the generated cuda source assert src_code is not None - A = torch.rand(M, K, device="cuda", dtype=getattr(torch, in_dtype)) - B = torch.rand(N, K, device="cuda", dtype=getattr(torch, in_dtype)) + A = torch.rand(M, K, device="cuda", dtype=getattr(torch, in_dtype)) - 0.5 + B = torch.rand(N, K, device="cuda", dtype=getattr(torch, in_dtype)) - 0.5 C = torch.zeros(M, N, device="cuda", dtype=getattr(torch, accum_dtype)) ladder_permutate_config = bitblas.ops.LadderPermutateConfig( diff --git a/testing/python/operators/test_general_matmul_tilelang_kernel.py b/testing/python/operators/test_general_matmul_tilelang_kernel.py index 5e59ef04..9308a942 100644 --- a/testing/python/operators/test_general_matmul_tilelang_kernel.py +++ b/testing/python/operators/test_general_matmul_tilelang_kernel.py @@ -267,8 +267,8 @@ def assert_matmul_fine_grained_apply_config_correctness( # src_code is the generated cuda source assert src_code is not None - A = torch.rand(M, K, device="cuda", dtype=getattr(torch, in_dtype)) - B = torch.rand(N, K, device="cuda", dtype=getattr(torch, in_dtype)) + A = torch.rand(M, K, device="cuda", dtype=getattr(torch, in_dtype)) - 0.5 + B = torch.rand(N, K, device="cuda", dtype=getattr(torch, in_dtype)) - 0.5 C = torch.zeros(M, N, device="cuda", dtype=getattr(torch, accum_dtype)) mod = tl.Profiler(mod, params, [], tl.TensorSupplyType.Integer)