()
| 8 | |
| 9 | @pytest.mark.L0 |
| 10 | def test_sparse_cpu(): |
| 11 | M = 128 |
| 12 | N = 32 |
| 13 | K = 32 |
| 14 | L = 1 |
| 15 | debug = False |
| 16 | # generate sparse tensor |
| 17 | a = torch.empty(M, K).random_(-5, 5).to(torch.float16) |
| 18 | sparse_utils = su.SparseUtils(M, K, L, cutlass.Float16) |
| 19 | if debug: |
| 20 | sparse_utils.use_specific_meta_data() |
| 21 | a_gen_from_cpu = sparse_utils.generate_sparse_4_2_tensor_with_tensor(a, True) |
| 22 | # print(a_gen_from_cpu) |
| 23 | # generate compressed tensor and meta data |
| 24 | a_compressed_cpu = torch.empty(M, K // 2).to(torch.float16) |
| 25 | meta_data_cpu = torch.empty(M, K // 4 // 8).to(torch.uint32) |
| 26 | compressor = su.Compressor(M, K, L) |
| 27 | compressor.compress(a_gen_from_cpu, a_compressed_cpu, meta_data_cpu, True) |
| 28 | # # test with gemm |
| 29 | b = torch.empty(N, K).random_(-5, 5).to(torch.float16).cuda() |
| 30 | d = torch.empty(M, N).zero_().to(torch.float16).cuda() |
| 31 | b_tensor = from_dlpack(b) |
| 32 | d_tensor = from_dlpack(d) |
| 33 | a_compressed_cpu_tensor = from_dlpack(a_compressed_cpu.cuda()) |
| 34 | meta_data_cpu_tensor = from_dlpack(meta_data_cpu.cuda()) |
| 35 | sparse_emulation = su.SparseEmulation(M, N, K, 1) |
| 36 | sparse_emulation(a_compressed_cpu_tensor, b_tensor, d_tensor, meta_data_cpu_tensor) |
| 37 | |
| 38 | ref = torch.einsum("mk,nk->mn", a_gen_from_cpu.cpu(), b.cpu()) |
| 39 | if debug: |
| 40 | a_ori = a_gen_from_cpu.cpu().numpy() |
| 41 | np.savetxt("a.txt", a_ori, fmt="%f") |
| 42 | a_compressed_cpu_ori = a_compressed_cpu.cpu().numpy() |
| 43 | np.savetxt("a_compressed_cpu.txt", a_compressed_cpu_ori, fmt="%f") |
| 44 | meta_data_cpu_ori = meta_data_cpu.cpu().numpy() |
| 45 | np.savetxt("meta_data_cpu.txt", meta_data_cpu_ori, fmt="%f") |
| 46 | d_ori = d.cpu().numpy() |
| 47 | np.savetxt("d.txt", d_ori, fmt="%f") |
| 48 | ref_ori = ref.cpu().numpy() |
| 49 | np.savetxt("ref.txt", ref_ori, fmt="%f") |
| 50 | torch.testing.assert_close(d.cpu(), ref) |
| 51 | print("cpu d == ref") |
| 52 | |
| 53 | |
| 54 | @pytest.mark.L0 |
no test coverage detected