diff --git a/bae/sparse/__init__.py b/bae/sparse/__init__.py index ac7d4e7..bd08f26 100644 --- a/bae/sparse/__init__.py +++ b/bae/sparse/__init__.py @@ -1,5 +1,3 @@ -from .bsr import * -from .bsr_cuda import * from .py_ops import * try: from .solve import * diff --git a/bae/sparse/sparse_op_cpp.cpp b/bae/sparse/sparse_op_cpp.cpp deleted file mode 100644 index 81c6e6a..0000000 --- a/bae/sparse/sparse_op_cpp.cpp +++ /dev/null @@ -1,245 +0,0 @@ -#include -#include -#include -#include -#include -#include - -#include -#if AT_MKL_ENABLED() && (!defined(_WIN32)) -#define AT_USE_MKL_SPARSE() 1 -#else -#define AT_USE_MKL_SPARSE() 0 -#endif - -template -std::tuple scan_symbol(const index_t *crow_indices_ptr, - const index_t *ccol_indices_ptr, - const index_t *row_indices_ptr, - const index_t *col_indices_ptr, - int64_t sm, - int64_t sp) -{ - std::vector> kkijs[sm]; - at::parallel_for(0, sm, at::internal::GRAIN_SIZE, [&](int64_t start, int64_t end) { - // auto start = 0; - // auto end = sm; - for (const auto i : c10::irange(start, end)) - { - for (const auto k1 : c10::irange(crow_indices_ptr[i], crow_indices_ptr[i + 1])) - { - for (const auto j : c10::irange(sp)) - { - index_t k2 = ccol_indices_ptr[j]; - if (k2 == ccol_indices_ptr[j + 1]) - continue; - while (row_indices_ptr[k2] < col_indices_ptr[k1] && k2 < ccol_indices_ptr[j + 1] - 1) - { - k2 += 1; - } - if (row_indices_ptr[k2] == col_indices_ptr[k1]) - { - kkijs[i].push_back(std::make_tuple(k1, k2, i, j)); - } - } - } - } - }); - // get nnz - auto nnz = 0; - auto nuijs = 0; - for (const auto i : c10::irange(sm)) - { - nnz += kkijs[i].size(); - } - // std::cout << "nnz: " << nnz << std::endl; - // flatten kkijs - index_t *data_ptr = (index_t *)malloc(nnz * 5 * sizeof(index_t)); - index_t *unique_ijs = (index_t *)malloc(nnz * 2 * sizeof(index_t)); - std::unordered_map unique_ij_map; - - int cur = 0; - for (const auto i : c10::irange(sm)) - { - for (const auto &kkij : kkijs[i]) - { - auto k1 = std::get<0>(kkij); - auto k2 = std::get<1>(kkij); - auto i = std::get<2>(kkij); - auto j = std::get<3>(kkij); - auto ij = i * sm + j; - if (auto search = unique_ij_map.find(ij); search != unique_ij_map.end()) - { - data_ptr[cur * 5 + 4] = search->second; - } - else - { - unique_ijs[2*nuijs] = i; - unique_ijs[2*nuijs+1] = j; - data_ptr[cur * 5 + 4] = nuijs; - unique_ij_map[ij] = nuijs; - nuijs += 1; - } - - data_ptr[cur * 5] = k1; - data_ptr[cur * 5 + 1] = k2; - data_ptr[cur * 5 + 2] = i; - data_ptr[cur * 5 + 3] = j; - cur += 1; - } - } - // torch::Tensor result = torch::from_blob(data_ptr, {nnz, 5}); - // torch::Tensor coo_indices = torch::from_blob(unique_ijs.data(), {unique_ijs.size(), 2}); - return std::make_tuple(data_ptr, nnz, unique_ijs, nuijs); -} - -int called_count = 0; -torch::Tensor sources; -torch::Tensor coo_indices; - -torch::Tensor sparse_bsr_mm(const torch::Tensor &bsr, const torch::Tensor &bsc) -{ - auto crow_indices = bsr.crow_indices().cpu().contiguous(); - auto col_indices = bsr.col_indices().cpu().contiguous(); - auto csr_values = bsr.values(); - - auto ccol_indices = bsc.ccol_indices().cpu().contiguous(); - auto row_indices = bsc.row_indices().cpu().contiguous(); - auto csc_values = bsc.values(); - - TORCH_CHECK_EQ(bsr.ndimension(), 2); - TORCH_CHECK_EQ(bsc.ndimension(), 2); - - auto m = bsr.size(-2); - auto n = bsr.size(-1); - auto p = bsc.size(-1); - int dm; - int dn; - int dp; - if (bsr.layout() == at::kSparseCsr && bsc.layout() == at::kSparseCsc) - { - TORCH_CHECK_EQ(csr_values.ndimension(), 1); - TORCH_CHECK_EQ(csc_values.ndimension(), 1); - dm = 1; - dn = 1; - dp = 1; - } - else { - dm = csr_values.size(-2); - dn = csr_values.size(-1); - dp = csc_values.size(-1); - } - auto sm = m / dm; - auto sn = n / dn; - auto sp = p / dp; - TORCH_CHECK_EQ(dm * sm, m); - TORCH_CHECK_EQ(dn * sn, n); - TORCH_CHECK_EQ(dp * sp, p); - - - int nnz; - int nuijs; - if (called_count == 0) {AT_DISPATCH_INDEX_TYPES( - crow_indices.scalar_type(), - "bsr_mm_crow_indices", - [&]() - { - auto symbols = scan_symbol(crow_indices.data_ptr(), - ccol_indices.data_ptr(), - row_indices.data_ptr(), - col_indices.data_ptr(), - sm, sp); - auto data_ptr = std::get<0>(symbols); - auto unique_ijs = std::get<2>(symbols); - // for (const auto i : c10::irange(std::get<1>(symbols))) - // { - // std::cout << data_ptr[i * 5] << " " << data_ptr[i * 5 + 1] << " " << data_ptr[i * 5 + 2] << " " << data_ptr[i * 5 + 3] << " " << data_ptr[i * 5 + 4] << std::endl; - // } - nnz = std::get<1>(symbols); - nuijs = std::get<3>(symbols); - sources = torch::from_blob(data_ptr, {nnz, 5}, crow_indices.options()).to(csr_values.device()); - // std::cout << "index: " << index << std::endl; - //print unique_ijs - // for (const auto i : c10::irange(nuijs)) - // { - // std::cout << unique_ijs[i * 2] << " " << unique_ijs[i * 2 + 1] << std::endl; - // } - coo_indices = torch::from_blob(unique_ijs, {nuijs, 2}, crow_indices.options()).to(csr_values.device()); - });} - // std::cout << "nuijs: " << nuijs << std::endl; - // std::cout << "coo_indices[..., 0]: " << coo_indices.index({"...", 0}) << std::endl; - // print coo_indices using tensor accessor - // for (const auto i : c10::irange(nuijs)) { - // std::cout << coo_indices[i][0].item() << " " << coo_indices[i][1].item() << std::endl; - // } - auto index = sources.index({"...", 4}); - auto prod = torch::bmm(csr_values.index({sources.index({"...", 0})}), csc_values.index({sources.index({"...", 1})})); - auto reduced = torch::zeros({nuijs, dm, dp}, prod.options()); - reduced.scatter_add_(0, index.index({"...", torch::indexing::None, torch::indexing::None}).expand_as(prod), prod); - auto row_res = coo_indices.index({"...", 0}); - auto col_res = coo_indices.index({"...", 1}); - // std::cout << "sm, sp: " << sm << " " << sp << std::endl; - // std::cout << "row_res: " << row_res << std::endl; - // std::cout << "col_res: " << col_res << std::endl; - auto crow_res = at::_convert_indices_from_coo_to_csr(row_res.contiguous(), sm, row_res.dtype() == at::kInt); - // auto ccol_res = at::_convert_indices_from_coo_to_csr(col_res, sm, col_res.dtype() == at::kInt); - // std::cout << "crow_res: " << crow_res << std::endl; - // std::cout << "ccol_res: " << ccol_res << std::endl; - // std::cout << "reduced: " << reduced << std::endl; - // at::IntArrayRef(); - // int64_t *size = (int64_t *)malloc(2 * sizeof(int64_t)); - // size[0] = sm; - // size[1] = sp; - // auto dummy_coo = torch::sparse_coo_tensor(coo_indices.mT(), torch::zeros({nuijs}, row_res.options()), {sm, sp}, at::TensorOptions().dtype(row_res.dtype()).device(row_res.device()).layout(at::kSparse)); - // auto crow_res = dummy_coo.to_sparse_csr().crow_indices().to(reduced.device()); - called_count += 1; - return at::_sparse_compressed_tensor_unsafe(crow_res.to(reduced.device()), col_res.to(reduced.device()), reduced.to(reduced.device()), {m, p}, at::TensorOptions().dtype(reduced.dtype()).device(reduced.device()).layout(at::kSparseBsr)); -} - -torch::Tensor sparse_bsr_csr_mm(const torch::Tensor &a, const torch::Tensor &b) -{ - if (a.layout() == at::kSparseBsr && b.layout() == at::kSparseBsc) - { - // auto output = (*_sparse_bsr_bsc_matmul)({a, b}); - // return output.toTensor(); - return sparse_bsr_mm(a, b); - } -# if AT_USE_MKL_SPARSE() - return at::native::_sparse_csr_mm(a, b); -# else - // TORCH_CHECK(false, "MKL Sparse is not enabled."); - if (a.layout() == at::kSparseBsr && b.layout() == at::kStrided) - { - int64_t b_shape[2] = {a.values().size(-1), b.size(-1)}; - torch::Tensor bsc = b.to_sparse_bsc(b_shape); - // std::cout << "b_shape: " << b_shape[0] << " " << b_shape[1] << std::endl; - return sparse_bsr_mm(a, bsc).to_dense(); - } -# endif -} - -PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { - m.def("sparse_bsr_mm", &sparse_bsr_mm, "Sparse BSR CSR matrix multiplication"); -} -TORCH_LIBRARY_IMPL(aten, SparseCsrCPU, m) -{ - // hide stderr - // freopen("/dev/null", "a", stderr); - // auto module = pybind11::module::import("pypose.sparse.ops"); - // auto object = module.attr("bsr_bsc_matmul"); - // auto script_function = object.cast(); - // _sparse_bsr_bsc_matmul = script_function.function_; - // m.impl("mm", sparse_bsr_csr_mm); - // resume stderr - // freopen("/dev/tty", "a", stderr); -} - - -/* -the file that calls the cuBLAS API for matrix multiplication is located at aten/src/ATen/native/cuda/Blas.cpp. In this file, there is a function called matmul_kernel that dispatches different cuBLAS functions based on the input tensor shapes and data types. For example, if the input tensors are 2-dimensional and have floating-point values, the function will call cublasSgemm or cublasDgemm for single-precision or double-precision arithmetic, respectively. - -the file that calls the MKL API for matrix multiplication is located at aten/src/ATen/native/LinearAlgebra.cpp. In this file, there is a function called bmm_out_or_baddbmm_ that dispatches different MKL functions based on the input tensor shapes and data types. For example, if the input tensors are 2-dimensional and have floating-point values, the function will call cblas_sgemm or cblas_dgemm for single-precision or double-precision arithmetic, respectively. - -Blas: same file but in different function -at::blas::gemm -*/ \ No newline at end of file diff --git a/bae/sparse/sparse_op_cuda.cpp b/bae/sparse/sparse_op_cuda.cpp deleted file mode 100755 index ddf2e28..0000000 --- a/bae/sparse/sparse_op_cuda.cpp +++ /dev/null @@ -1,266 +0,0 @@ -#include -#include -#include -#include -#include -#include - -#include -#if AT_MKL_ENABLED() && (!defined(_WIN32)) -#define AT_USE_MKL_SPARSE() 1 -#else -#define AT_USE_MKL_SPARSE() 0 -#endif - -#if defined(USE_ROCM) -#include -#include - -#define HIP_CHECK(stat) \ - { \ - if (stat != hipSuccess) \ - { \ - std::cerr << "Error: hip error in line " << __LINE__ << std::endl; \ - } \ - } - -#define ROCSPARSE_CHECK(stat) \ - { \ - if (stat != rocsparse_status_success) \ - { \ - std::cerr << "Error: rocsparse error in line " << __LINE__ << std::endl; \ - } \ - } - - -torch::Tensor sparse_bsr_mm_rocm(const torch::Tensor a, const torch::Tensor b) -{ - rocsparse_int m = a.size(0); - rocsparse_int n = b.size(1); - TORCH_CHECK(a.size(1) == b.size(0), "matrices are not aligned"); - rocsparse_int k = a.size(1); - TORCH_CHECK(a.values().size(1) == a.values().size(2), "a must have square blocks"); - TORCH_CHECK(b.values().size(1) == b.values().size(2), "b must have square blocks"); - TORCH_CHECK(a.values().size(2) == b.values().size(1), "a and b must have compatible block sizes"); - rocsparse_int block_dim = a.values().size(2); - rocsparse_int nnzb_A = a.values().size(0); - rocsparse_int nnzb_B = b.values().size(0); - rocsparse_int nnzb_D = 0; - rocsparse_int mb = m / block_dim; - rocsparse_int nb = n / block_dim; - rocsparse_int kb = k / block_dim; - - torch::Tensor crow_A = a.crow_indices(); - torch::Tensor col_A = a.col_indices(); - torch::Tensor crow_B = b.crow_indices(); - torch::Tensor col_B = b.col_indices(); - if (crow_A.dtype() != torch::kInt32) - { - crow_A = crow_A.to(torch::kInt32); - col_A = col_A.to(torch::kInt32); - } - if (crow_B.dtype() != torch::kInt32) - { - crow_B = crow_B.to(torch::kInt32); - col_B = col_B.to(torch::kInt32); - } - rocsparse_int *bsr_row_ptr_A = crow_A.data(); - rocsparse_int *bsr_col_ind_A = col_A.data(); - rocsparse_int *bsr_row_ptr_B = crow_B.data(); - rocsparse_int *bsr_col_ind_B = col_B.data(); - rocsparse_int *bsr_row_ptr_D = nullptr; - rocsparse_int *bsr_col_ind_D = nullptr; - - // rocSPARSE handle - rocsparse_handle handle; - ROCSPARSE_CHECK(rocsparse_create_handle(&handle)); - - // Initialize scalar multipliers - float alpha = 1.0f; - float beta = 0.0f; - - // Create matrix descriptors - rocsparse_mat_descr descr_A; - rocsparse_mat_descr descr_B; - rocsparse_mat_descr descr_C; - rocsparse_mat_descr descr_D; - - rocsparse_create_mat_descr(&descr_A); - rocsparse_create_mat_descr(&descr_B); - rocsparse_create_mat_descr(&descr_C); - rocsparse_create_mat_descr(&descr_D); - - // Create matrix info structure - rocsparse_mat_info info_C; - rocsparse_create_mat_info(&info_C); - - // Set pointer mode - rocsparse_set_pointer_mode(handle, rocsparse_pointer_mode_host); - - // Query rocsparse for the required buffer size - size_t buffer_size; - - rocsparse_sbsrgemm_buffer_size(handle, - rocsparse_direction_row, - rocsparse_operation_none, - rocsparse_operation_none, - mb, - nb, - kb, - block_dim, - &alpha, - descr_A, - nnzb_A, - bsr_row_ptr_A, - bsr_col_ind_A, - descr_B, - nnzb_B, - bsr_row_ptr_B, - bsr_col_ind_B, - &beta, - descr_D, - nnzb_D, - bsr_row_ptr_D, - bsr_col_ind_D, - info_C, - &buffer_size); - - // Allocate buffer - void *buffer; - hipMalloc(&buffer, buffer_size); - - // Obtain number of total non-zero block entries in C and block row pointers of C - torch::TensorOptions int_tensor_option = torch::TensorOptions().dtype(torch::kInt32).device(crow_A.device()); - rocsparse_int nnzb_C; - torch::Tensor crow_C = torch::empty({mb + 1}, int_tensor_option); - rocsparse_int *bsr_row_ptr_C = crow_C.data(); - - rocsparse_bsrgemm_nnzb(handle, - rocsparse_direction_row, - rocsparse_operation_none, - rocsparse_operation_none, - mb, - nb, - kb, - block_dim, - descr_A, - nnzb_A, - bsr_row_ptr_A, - bsr_col_ind_A, - descr_B, - nnzb_B, - bsr_row_ptr_B, - bsr_col_ind_B, - descr_D, - nnzb_D, - bsr_row_ptr_D, - bsr_col_ind_D, - descr_C, - bsr_row_ptr_C, - &nnzb_C, - info_C, - buffer); - - // Compute block column indices and values of C - torch::Tensor col_C = torch::empty({nnzb_C}, int_tensor_option); - torch::Tensor val_C = torch::empty({nnzb_C, block_dim, block_dim}, torch::TensorOptions().dtype(torch::kFloat32).device(crow_A.device())); - rocsparse_int *bsr_col_ind_C = col_C.data(); - - float *bsr_val_A = a.values().data(); - float *bsr_val_B = b.values().data(); - float *bsr_val_C = val_C.data(); - float *bsr_val_D = nullptr; - - rocsparse_sbsrgemm(handle, - rocsparse_direction_row, - rocsparse_operation_none, - rocsparse_operation_none, - mb, - nb, - kb, - block_dim, - &alpha, - descr_A, - nnzb_A, - bsr_val_A, - bsr_row_ptr_A, - bsr_col_ind_A, - descr_B, - nnzb_B, - bsr_val_B, - bsr_row_ptr_B, - bsr_col_ind_B, - &beta, - descr_D, - nnzb_D, - bsr_val_D, - bsr_row_ptr_D, - bsr_col_ind_D, - descr_C, - bsr_val_C, - bsr_row_ptr_C, - bsr_col_ind_C, - info_C, - buffer); - ROCSPARSE_CHECK(rocsparse_destroy_mat_descr(descr_A)); - ROCSPARSE_CHECK(rocsparse_destroy_mat_descr(descr_B)); - ROCSPARSE_CHECK(rocsparse_destroy_mat_descr(descr_C)); - ROCSPARSE_CHECK(rocsparse_destroy_mat_descr(descr_D)); - ROCSPARSE_CHECK(rocsparse_destroy_handle(handle)); - hipFree(buffer); - return at::_sparse_compressed_tensor_unsafe(crow_C, col_C, val_C, {m, n}, at::TensorOptions().dtype(val_C.dtype()).device(val_C.device()).layout(at::kSparseBsr)); -} - - -#endif - -#define CHECK_CUSPARSE(call) \ -{ \ - cusparseStatus_t err; \ - if ((err = (call)) != CUSPARSE_STATUS_SUCCESS) \ - { \ - fprintf(stderr, "Got error %d at %s:%d\n", err, __FILE__, __LINE__); \ - cudaError_t cuda_err = cudaGetLastError(); \ - if (cuda_err != cudaSuccess) \ - { \ - fprintf(stderr, " CUDA error \"%s\" also detected\n", \ - cudaGetErrorString(cuda_err)); \ - } \ - exit(1); \ - } \ -} - -#define CHECK_CUDA(x) TORCH_CHECK(x.device().is_cuda(), #x " must be a CUDA tensor") - -torch::Tensor sparse_bsr_mm_cuda(const torch::Tensor, const torch::Tensor); - -torch::Tensor sparse_bsr_csr_mm(const torch::Tensor &a, const torch::Tensor &b) -{ - - if (a.layout() == at::kSparseBsr && b.layout() == at::kSparseBsc) - { - CHECK_CUDA(a); - CHECK_CUDA(b); -#if defined(USE_ROCM) - return sparse_bsr_mm_rocm(a, b); -#else - return sparse_bsr_mm_cuda(a, b); -#endif - } -} - -PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {} - -TORCH_LIBRARY_IMPL(aten, SparseCsrCUDA, m) -{ - // m.impl("mm", sparse_bsr_csr_mm); -} - -/* -the file that calls the cuBLAS API for matrix multiplication is located at aten/src/ATen/native/cuda/Blas.cpp. In this file, there is a function called matmul_kernel that dispatches different cuBLAS functions based on the input tensor shapes and data types. For example, if the input tensors are 2-dimensional and have floating-point values, the function will call cublasSgemm or cublasDgemm for single-precision or double-precision arithmetic, respectively. - -the file that calls the MKL API for matrix multiplication is located at aten/src/ATen/native/LinearAlgebra.cpp. In this file, there is a function called bmm_out_or_baddbmm_ that dispatches different MKL functions based on the input tensor shapes and data types. For example, if the input tensors are 2-dimensional and have floating-point values, the function will call cblas_sgemm or cblas_dgemm for single-precision or double-precision arithmetic, respectively. - -Blas: same file but in different function -at::blas::gemm -*/ \ No newline at end of file diff --git a/bae/sparse/sparse_op_cuda_kernel.cu b/bae/sparse/sparse_op_cuda_kernel.cu deleted file mode 100755 index c56993d..0000000 --- a/bae/sparse/sparse_op_cuda_kernel.cu +++ /dev/null @@ -1,248 +0,0 @@ -#include - -#include -#include -#include -#include - -#include -#include -#include - -#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous") - -template -__global__ void scan_symbol_kernel( - const index_t *coo_indices_ptr, - const index_t *crow_indices_ptr, - const index_t *ccol_indices_ptr, - const index_t *row_indices_ptr, - const index_t *col_indices_ptr, - int64_t sm, - int64_t sp, - index_t *out_data_ptr) -{ - int64_t count; - int64_t max_partial = sp; - int nnza = crow_indices_ptr[sm]; - index_t col_indices_ptr_k1; - for (int k1 = blockIdx.x * blockDim.x + threadIdx.x; k1 < nnza; k1 += blockDim.x * gridDim.x) - { - int64_t i = coo_indices_ptr[k1]; - // int64_t i = blockIdx.x; // each block is responsible for a row - // TODO: Move crow indices to shared memory - // for (auto k1 = crow_indices_ptr[i] + threadIdx.x; k1 < crow_indices_ptr[i + 1]; k1 = k1 + blockDim.x) // TODO: allocate cache in case of more than one iter - count = 0; - col_indices_ptr_k1 = col_indices_ptr[k1]; - for (int64_t j = 0; j < sp; j++) - { - index_t k2 = ccol_indices_ptr[j]; - index_t col_end = ccol_indices_ptr[j + 1]; - if (k2 == col_end) - continue; - while (row_indices_ptr[k2] < col_indices_ptr_k1 && k2 < col_end - 1) - { - k2 += 1; - } - if (row_indices_ptr[k2] == col_indices_ptr_k1) - { - int global_pos = k1 * max_partial + count; - out_data_ptr[nnza + global_pos * 4 + 0] = k1; - out_data_ptr[nnza + global_pos * 4 + 1] = k2; - out_data_ptr[nnza + global_pos * 4 + 2] = i; - out_data_ptr[nnza + global_pos * 4 + 3] = j; - count += 1; - } - } - out_data_ptr[k1] = count; - } - - __syncthreads(); -} - -template -__global__ void flatten_var_len(index_t *in_data_ptr, count_t stride, - index_t *offsets, index_t *elements_per_row, index_t *out_data_ptr, size_t size_per_cell) -{ - // use memcpy to copy data from in_data_ptr to out_data_ptr - // in_data_ptr is a 2D array of size num_rows * (stride * size_per_cell) - // out_data_ptr is a 2D array of size (total_elements * size_per_cell) - // offsets is a 1D array of size num_rows - // elements_per_row is a 1D array of size num_rows - // stride is the number of max elements per row - // size_per_cell is the size of each element in bytes - int i = blockIdx.x * blockDim.x + threadIdx.x; - - int offset = offsets[i]; - int num_elements = elements_per_row[i]; - memcpy(out_data_ptr + offset * size_per_cell, in_data_ptr + i * stride * size_per_cell, num_elements * size_per_cell * sizeof(index_t)); -} - -torch::Tensor sparse_bsr_mm_cuda(const torch::Tensor bsr, const torch::Tensor bsc) -{ - auto crow_indices = bsr.crow_indices(); - auto col_indices = bsr.col_indices(); - auto csr_values = bsr.values(); - - auto ccol_indices = bsc.ccol_indices(); - auto row_indices = bsc.row_indices(); - auto csc_values = bsc.values(); - CHECK_CONTIGUOUS(crow_indices); - CHECK_CONTIGUOUS(col_indices); - CHECK_CONTIGUOUS(ccol_indices); - CHECK_CONTIGUOUS(row_indices); - - TORCH_CHECK_EQ(bsr.ndimension(), 2); - TORCH_CHECK_EQ(bsc.ndimension(), 2); - - auto m = bsr.size(-2); - auto n = bsr.size(-1); - auto p = bsc.size(-1); - int dm; - int dn; - int dp; - if (bsr.layout() == at::kSparseCsr && bsc.layout() == at::kSparseCsc) - { - TORCH_CHECK_EQ(csr_values.ndimension(), 1); - TORCH_CHECK_EQ(csc_values.ndimension(), 1); - dm = 1; - dn = 1; - dp = 1; - } - else { - dm = csr_values.size(-2); - dn = csr_values.size(-1); - dp = csc_values.size(-1); - } - auto sm = m / dm; - auto sn = n / dn; - auto sp = p / dp; - TORCH_CHECK_EQ(dm * sm, m); - TORCH_CHECK_EQ(dn * sn, n); - TORCH_CHECK_EQ(dp * sp, p); - - torch::Tensor cooa = at::_convert_indices_from_csr_to_coo(crow_indices, col_indices, crow_indices.dtype() == at::kInt, false); - // std::cout << "cooa: " << cooa << std::endl; - torch::Tensor sources; - torch::Tensor coo_indices; - torch::Tensor index; - int nuijs; - AT_DISPATCH_INDEX_TYPES( - crow_indices.scalar_type(), - "bsr_mm_crow_indices", - [&]() { // maybe useful: CppTypeToScalarType - auto nnza = col_indices.size(0); - if (nnza == 0) - { - int total_partials = 0; - nuijs = 0; - sources = torch::empty({total_partials, 4}, torch::TensorOptions(crow_indices.dtype()).device(crow_indices.device())); - index = torch::empty({total_partials}, torch::TensorOptions(crow_indices.dtype()).device(crow_indices.device())); - coo_indices = torch::empty({0, 2}, torch::TensorOptions(crow_indices.dtype()).device(crow_indices.device())); - return; - } - index_t *out_data_ptr; - cudaMalloc(&out_data_ptr, sizeof(index_t) * 5 * (nnza * sp)); - int numSMs; - cudaDeviceGetAttribute(&numSMs, cudaDevAttrMultiProcessorCount, 0); // Perform SAXPY on 1M elements saxpy<<<32*numSMs, 256>>>(1 << 20, 2.0, x, y); - scan_symbol_kernel<<<32 * numSMs, 256>>>( - cooa.data(), - crow_indices.data(), - ccol_indices.data(), - row_indices.data(), - col_indices.data(), - sm, sp, out_data_ptr); - cudaDeviceSynchronize(); - // Determine temporary device storage requirements - void *d_temp_storage = NULL; - index_t *offsets; - cudaMalloc(&offsets, nnza * sizeof(index_t)); - thrust::exclusive_scan(thrust::device, out_data_ptr, out_data_ptr + nnza, offsets); - // auto total_partials = out_data_ptr[num_items - 1] + offsets[num_items - 1]; - index_t total_partials = 0; - index_t tmp; - cudaMemcpy(&total_partials, out_data_ptr + nnza - 1, sizeof(index_t), cudaMemcpyDeviceToHost); - cudaMemcpy(&tmp, offsets + nnza - 1, sizeof(index_t), cudaMemcpyDeviceToHost); - total_partials += tmp; - // std::cout << "total_partials: " << total_partials << std::endl; - - // Allocate kkijs - sources = torch::empty({total_partials, 4}, torch::TensorOptions(crow_indices.dtype()).device(crow_indices.device())); - index = torch::empty({total_partials}, torch::TensorOptions(crow_indices.dtype()).device(crow_indices.device())); - // torch::zeros({total_partials, 4}, torch::TensorOptions(index_t).device(crow_indices.device())); // TODO - if (total_partials == 0) - { - coo_indices = torch::empty({0, 2}, torch::TensorOptions(crow_indices.dtype()).device(crow_indices.device())); - return; - } - // Copy data to sources - // std::cout << "!!!.....!!!!" << std::endl; - flatten_var_len<<>>(out_data_ptr + nnza, - sp, - offsets, - out_data_ptr, - sources.data(), - 4); - cudaDeviceSynchronize(); - // std::cout << "111111" << std::endl; - // assign_rank<<<1, total_partials>>>(sources.data(), sm, sp, total_partials, coo_indices.data(), nuijs_ptr, index.data()); - - torch::Tensor ijs = sources.index({torch::indexing::Slice(0, total_partials), 2}) * sp + sources.index({torch::indexing::Slice(0, total_partials), 3}); - torch::Tensor partial_indices = torch::arange(0, total_partials, torch::TensorOptions().dtype(crow_indices.dtype()).device(crow_indices.device())); - // sort by key - // https://nvidia.github.io/cccl/thrust/api/groups/group__sorting.html#function-sort-by-key - thrust::sort_by_key(thrust::device, ijs.data(), ijs.data() + total_partials, partial_indices.data()); - - // struct equals : public thrust::binary_function - // { - // index_t operator()(index_t x, index_t y) { return !(x ^ y); } - // }; - // adj diff - thrust::adjacent_difference(thrust::device, ijs.data(), ijs.data() + total_partials, index.data()); - index.clamp_(0, 1); - index[0] = 1; - - torch::Tensor uijs = ijs.masked_select(index.to(torch::kBool)); // TODO: check if casting is necessary - coo_indices = torch::empty({uijs.size(0), 2}, torch::TensorOptions(crow_indices.dtype()).device(crow_indices.device())); - coo_indices.index({torch::indexing::Slice(0, uijs.size(0)), 0}) = uijs / sp; - coo_indices.index({torch::indexing::Slice(0, uijs.size(0)), 1}) = uijs % sp; - - // - // https://nvidia.github.io/cccl/thrust/api/groups/group__prefixsums.html#function-inclusive-scan - thrust::inclusive_scan(thrust::device, index.data(), index.data() + total_partials, index.data()); - nuijs = index[total_partials - 1].item(); - // std::cout << "nuijs: " << nuijs << std::endl; - index -= 1; // TODO: should be more efficient - // TORCH_CHECK_EQ(nuijs, uijs.size(0)); - // index[partial_indices] = index; - index = torch::empty_like(index).scatter_(0, partial_indices, index); - - cudaFree(out_data_ptr); - cudaFree(offsets); - }); - // std::cout << "coo_indices[..., 0]: " << coo_indices.index({"...", 0}) << std::endl; - // print coo_indices using tensor accessor - // for (int j = 0; j < sp; j++js)) { - // std::cout << coo_indices[i][0].item() << " " << coo_indices[i][1].item() << std::endl; - // } - auto prod = torch::bmm(csr_values.index({sources.index({"...", 0})}), csc_values.index({sources.index({"...", 1})})); - auto reduced = torch::zeros({nuijs, dm, dp}, prod.options()); - reduced.scatter_add_(0, index.index({"...", torch::indexing::None, torch::indexing::None}).expand_as(prod), prod); - auto row_res = coo_indices.index({"...", 0}).contiguous(); - auto col_res = coo_indices.index({"...", 1}).contiguous(); - // std::cout << "sm, sp: " << sm << " " << sp << std::endl; - // std::cout << "row_res: " << row_res << std::endl; - // std::cout << "col_res: " << col_res << std::endl; - auto crow_res = at::_convert_indices_from_coo_to_csr(row_res, sm, row_res.dtype() == at::kInt); - // auto ccol_res = at::_convert_indices_from_coo_to_csr(col_res, sm, col_res.dtype() == at::kInt); - // std::cout << "crow_res: " << crow_res << std::endl; - // std::cout << "ccol_res: " << ccol_res << std::endl; - // std::cout << "reduced: " << reduced << std::endl; - // at::IntArrayRef(); - // int64_t *size = (int64_t *)malloc(2 * sizeof(int64_t)); - // size[0] = sm; - // size[1] = sp; - // auto dummy_coo = torch::sparse_coo_tensor(coo_indices.mT(), torch::zeros({nuijs}, row_res.options()), {sm, sp}, at::TensorOptions().dtype(row_res.dtype()).device(row_res.device()).layout(at::kSparse)); - // auto crow_res = dummy_coo.to_sparse_csr().crow_indices().to(reduced.device()); - return at::_sparse_compressed_tensor_unsafe(crow_res.to(reduced.device()), col_res.to(reduced.device()), reduced.to(reduced.device()), {m, p}, at::TensorOptions().dtype(reduced.dtype()).device(reduced.device()).layout(at::kSparseBsr)); -} diff --git a/setup.py b/setup.py index b176ba5..0495f1c 100644 --- a/setup.py +++ b/setup.py @@ -5,7 +5,7 @@ from packaging.specifiers import SpecifierSet from packaging.version import Version from setuptools import setup, find_packages -from torch.utils.cpp_extension import CppExtension, CUDAExtension, BuildExtension +from torch.utils.cpp_extension import CUDAExtension, BuildExtension VERSION = "0.2.5" SUPPORTED_CUDSS_SPECIFIER = SpecifierSet("<=0.7.1.6") @@ -107,17 +107,6 @@ def resolve_cudss_link_inputs(cudss_root): # Common extensions ext_modules = [ - CppExtension( - 'bae.sparse.bsr', - [os.path.join('bae', 'sparse', 'sparse_op_cpp.cpp')] - ), - CUDAExtension( - 'bae.sparse.bsr_cuda', - [ - os.path.join('bae', 'sparse', 'sparse_op_cuda.cpp'), - os.path.join('bae', 'sparse', 'sparse_op_cuda_kernel.cu') - ] - ), CUDAExtension( 'bae.sparse.spgemm', [os.path.join('bae', 'sparse', 'cusparse_wrapper.cpp')] diff --git a/tests/sparse/test_bsr.py b/tests/sparse/test_bsr.py index e9d11bb..be59b00 100644 --- a/tests/sparse/test_bsr.py +++ b/tests/sparse/test_bsr.py @@ -1,7 +1,6 @@ from functools import partial import torch import pytest -from bae.sparse import bsr, bsr_cuda from bae.sparse import diagonal_op_ from torchvision.transforms import Compose