mirror of
https://github.com/NVIDIA/cutlass.git
synced 2026-06-29 19:07:07 +00:00
137 lines
5.7 KiB
Python
137 lines
5.7 KiB
Python
#################################################################################################
|
|
#
|
|
# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
|
# SPDX-License-Identifier: BSD-3-Clause
|
|
#
|
|
# Redistribution and use in source and binary forms, with or without
|
|
# modification, are permitted provided that the following conditions are met:
|
|
#
|
|
# 1. Redistributions of source code must retain the above copyright notice, this
|
|
# list of conditions and the following disclaimer.
|
|
#
|
|
# 2. Redistributions in binary form must reproduce the above copyright notice,
|
|
# this list of conditions and the following disclaimer in the documentation
|
|
# and/or other materials provided with the distribution.
|
|
#
|
|
# 3. Neither the name of the copyright holder nor the names of its
|
|
# contributors may be used to endorse or promote products derived from
|
|
# this software without specific prior written permission.
|
|
#
|
|
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
|
|
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
|
|
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
|
|
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
|
|
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
|
|
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
|
|
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
|
|
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
|
|
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
|
|
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
|
#
|
|
#################################################################################################
|
|
|
|
import sparse_utils as su
|
|
import cutlass
|
|
import torch
|
|
from cutlass.cute.runtime import from_dlpack
|
|
import numpy as np
|
|
import pytest
|
|
|
|
|
|
@pytest.mark.L0
|
|
def test_sparse_cpu():
|
|
M = 128
|
|
N = 32
|
|
K = 32
|
|
L = 1
|
|
debug = False
|
|
# generate sparse tensor
|
|
a = torch.empty(M, K).random_(-5, 5).to(torch.float16)
|
|
sparse_utils = su.SparseUtils(M, K, L, cutlass.Float16)
|
|
if debug:
|
|
sparse_utils.use_specific_meta_data()
|
|
a_gen_from_cpu = sparse_utils.generate_sparse_4_2_tensor_with_tensor(a, True)
|
|
# print(a_gen_from_cpu)
|
|
# generate compressed tensor and meta data
|
|
a_compressed_cpu = torch.empty(M, K // 2).to(torch.float16)
|
|
meta_data_cpu = torch.empty(M, K // 4 // 8).to(torch.uint32)
|
|
compressor = su.Compressor(M, K, L)
|
|
compressor.compress(a_gen_from_cpu, a_compressed_cpu, meta_data_cpu, True)
|
|
# # test with gemm
|
|
b = torch.empty(N, K).random_(-5, 5).to(torch.float16).cuda()
|
|
d = torch.empty(M, N).zero_().to(torch.float16).cuda()
|
|
b_tensor = from_dlpack(b)
|
|
d_tensor = from_dlpack(d)
|
|
a_compressed_cpu_tensor = from_dlpack(a_compressed_cpu.cuda())
|
|
meta_data_cpu_tensor = from_dlpack(meta_data_cpu.cuda())
|
|
sparse_emulation = su.SparseEmulation(M, N, K, 1)
|
|
sparse_emulation(a_compressed_cpu_tensor, b_tensor, d_tensor, meta_data_cpu_tensor)
|
|
|
|
ref = torch.einsum("mk,nk->mn", a_gen_from_cpu.cpu(), b.cpu())
|
|
if debug:
|
|
a_ori = a_gen_from_cpu.cpu().numpy()
|
|
np.savetxt("a.txt", a_ori, fmt="%f")
|
|
a_compressed_cpu_ori = a_compressed_cpu.cpu().numpy()
|
|
np.savetxt("a_compressed_cpu.txt", a_compressed_cpu_ori, fmt="%f")
|
|
meta_data_cpu_ori = meta_data_cpu.cpu().numpy()
|
|
np.savetxt("meta_data_cpu.txt", meta_data_cpu_ori, fmt="%f")
|
|
d_ori = d.cpu().numpy()
|
|
np.savetxt("d.txt", d_ori, fmt="%f")
|
|
ref_ori = ref.cpu().numpy()
|
|
np.savetxt("ref.txt", ref_ori, fmt="%f")
|
|
torch.testing.assert_close(d.cpu(), ref)
|
|
print("cpu d == ref")
|
|
|
|
|
|
@pytest.mark.L0
|
|
def test_sparse_cuda():
|
|
M = 128
|
|
N = 32
|
|
K = 32
|
|
L = 1
|
|
debug = False
|
|
sparse_utils = su.SparseUtils(M, K, L, cutlass.Float16)
|
|
if debug:
|
|
sparse_utils.use_specific_meta_data()
|
|
# generate sparse tensor
|
|
a = torch.empty(M, K).random_(-5, 5).to(torch.float16).cuda()
|
|
a_gen_from_cuda = sparse_utils.generate_4_2_sparse_tensor(False)
|
|
# print(a_gen_from_cuda)
|
|
# generate compressed tensor and meta data
|
|
a_compressed_cuda = torch.empty(M, K // 2).to(torch.float16).cuda()
|
|
meta_data_cuda = torch.empty(M, K // 4 // 8).to(torch.uint32).cuda()
|
|
compressor = su.Compressor(M, K, L)
|
|
compressor.compress(a_gen_from_cuda, a_compressed_cuda, meta_data_cuda, False)
|
|
# test with gemm
|
|
b = torch.empty(N, K).random_(-5, 5).to(torch.float16).cuda()
|
|
d = torch.empty(M, N).zero_().to(torch.float16).cuda()
|
|
b_tensor = from_dlpack(b)
|
|
d_tensor = from_dlpack(d)
|
|
a_compressed_cuda_tensor = from_dlpack(a_compressed_cuda)
|
|
meta_data_cuda_tensor = from_dlpack(meta_data_cuda)
|
|
sparse_emulation = su.SparseEmulation(M, N, K, 1)
|
|
sparse_emulation(
|
|
a_compressed_cuda_tensor, b_tensor, d_tensor, meta_data_cuda_tensor
|
|
)
|
|
|
|
ref = torch.einsum("mk,nk->mn", a_gen_from_cuda.cpu(), b.cpu())
|
|
if debug:
|
|
a_ori = a_gen_from_cuda.cpu().numpy()
|
|
np.savetxt("a.txt", a_ori, fmt="%f")
|
|
a_compressed_cuda_ori = a_compressed_cuda.cpu().numpy()
|
|
np.savetxt("a_compressed_cuda.txt", a_compressed_cuda_ori, fmt="%f")
|
|
meta_data_cuda_ori = meta_data_cuda.cpu().numpy()
|
|
np.savetxt("meta_data_cuda.txt", meta_data_cuda_ori, fmt="%f")
|
|
d_ori = d.cpu().numpy()
|
|
np.savetxt("d.txt", d_ori, fmt="%f")
|
|
ref_ori = ref.cpu().numpy()
|
|
np.savetxt("ref.txt", ref_ori, fmt="%f")
|
|
torch.testing.assert_close(d.cpu(), ref)
|
|
print("cuda d == ref")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
cutlass.cuda.initialize_cuda_context()
|
|
test_sparse_cpu()
|
|
test_sparse_cuda()
|