Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[lang] Build sparse matrix using ndarray #6563

Merged
merged 10 commits into from
Nov 14, 2022
Merged

[lang] Build sparse matrix using ndarray #6563

Show file tree
Hide file tree
Changes from 8 commits
Commits
Show all changes
10 commits
Select commit Hold shift + click to select a range
5a16aa7
replace cpu sparse matrix builder's array with ndarray
FantasyVR Nov 9, 2022
52029da
sparse matrix builder supports cuda arch
FantasyVR Nov 9, 2022
dd7f38c
simplify
FantasyVR Nov 9, 2022
2f1d60b
fix float insertion bug and clean code
FantasyVR Nov 10, 2022
80a670b
clean
FantasyVR Nov 10, 2022
a352891
rm
FantasyVR Nov 10, 2022
9862717
clean
FantasyVR Nov 10, 2022
177d486
Update taichi/program/sparse_matrix.cpp
FantasyVR Nov 10, 2022
26698e0
rm
FantasyVR Nov 11, 2022
b6f938f
enable offlinecach
FantasyVR Nov 11, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion misc/sparse_matrix.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
import taichi as ti

ti.init(arch=ti.x64)
ti.init(arch=ti.x64, debug=True, offline_cache=False)
FantasyVR marked this conversation as resolved.
Show resolved Hide resolved

n = 8

Expand Down
2 changes: 1 addition & 1 deletion python/taichi/lang/kernel_impl.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -673,7 +673,7 @@ def func__(*args):
elif isinstance(needed, sparse_matrix_builder):
# Pass only the base pointer of the ti.types.sparse_matrix_builder() argument
launch_ctx.set_arg_uint(actual_argument_slot,
v._get_addr())
v._get_ndarray_addr())
elif isinstance(needed,
ndarray_type.NdarrayType) and isinstance(
v, taichi.lang._ndarray.Ndarray):
Expand Down
4 changes: 4 additions & 0 deletions python/taichi/linalg/sparse_matrix.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -277,6 +277,10 @@ def _get_addr(self):
"""Get the address of the sparse matrix"""
return self.ptr.get_addr()

def _get_ndarray_addr(self):
"""Get the address of the ndarray"""
return self.ptr.get_ndarray_data_ptr()

def print_triplets(self):
"""Print the triplets stored in the builder"""
self.ptr.print_triplets()
Expand Down
49 changes: 22 additions & 27 deletions taichi/program/sparse_matrix.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -65,53 +65,48 @@ SparseMatrixBuilder::SparseMatrixBuilder(int rows,
int cols,
int max_num_triplets,
DataType dtype,
const std::string &storage_format)
const std::string &storage_format,
Program *prog)
: rows_(rows),
cols_(cols),
max_num_triplets_(max_num_triplets),
dtype_(dtype),
storage_format_(storage_format) {
storage_format_(storage_format),
prog_(prog) {
auto element_size = data_type_size(dtype);
TI_ASSERT((element_size == 4 || element_size == 8));
data_base_ptr_ =
std::make_unique<uchar[]>(max_num_triplets_ * 3 * element_size);
ndarray_data_base_ptr_ = std::make_unique<Ndarray>(
prog_, dtype_, std::vector<int>{3 * (int)max_num_triplets_ + 1});
}

template <typename T, typename G>
void SparseMatrixBuilder::print_template() {
void SparseMatrixBuilder::print_triplets() {
num_triplets_ = ndarray_data_base_ptr_->read_int(std::vector<int>{0});
fmt::print("n={}, m={}, num_triplets={} (max={})\n", rows_, cols_,
num_triplets_, max_num_triplets_);
T *data = reinterpret_cast<T *>(data_base_ptr_.get());
for (int64 i = 0; i < num_triplets_; i++) {
fmt::print("({}, {}) val={}\n", ((G *)data)[i * 3], ((G *)data)[i * 3 + 1],
taichi_union_cast<T>(data[i * 3 + 2]));
for (int i = 0; i < num_triplets_; i++) {
auto idx = 3 * i + 1;
auto row = ndarray_data_base_ptr_->read_int(std::vector<int>{idx});
auto col = ndarray_data_base_ptr_->read_int(std::vector<int>{idx + 1});
auto val = ndarray_data_base_ptr_->read_float(std::vector<int>{idx + 2});
fmt::print("[{}, {}] = {}\n", row, col, val);
}
fmt::print("\n");
}

void SparseMatrixBuilder::print_triplets() {
auto element_size = data_type_size(dtype_);
switch (element_size) {
case 4:
print_template<float32, int32>();
break;
case 8:
print_template<float64, int64>();
break;
default:
TI_ERROR("Unsupported sparse matrix data type!");
break;
}
intptr_t SparseMatrixBuilder::get_ndarray_data_ptr() const {
return prog_->get_ndarray_data_ptr_as_int(ndarray_data_base_ptr_.get());
}

template <typename T, typename G>
void SparseMatrixBuilder::build_template(std::unique_ptr<SparseMatrix> &m) {
using V = Eigen::Triplet<T>;
std::vector<V> triplets;
T *data = reinterpret_cast<T *>(data_base_ptr_.get());
auto ptr = get_ndarray_data_ptr();
G *data = reinterpret_cast<G *>(ptr);
num_triplets_ = data[0];
data += 1;
for (int i = 0; i < num_triplets_; i++) {
triplets.push_back(V(((G *)data)[i * 3], ((G *)data)[i * 3 + 1],
taichi_union_cast<T>(data[i * 3 + 2])));
triplets.push_back(
V(data[i * 3], data[i * 3 + 1], taichi_union_cast<T>(data[i * 3 + 2])));
}
m->build_triplets(static_cast<void *>(&triplets));
clear();
Expand Down
11 changes: 6 additions & 5 deletions taichi/program/sparse_matrix.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,30 +19,31 @@ class SparseMatrixBuilder {
int cols,
int max_num_triplets,
DataType dtype,
const std::string &storage_format);
const std::string &storage_format,
Program *prog);

void print_triplets();

intptr_t get_ndarray_data_ptr() const;

std::unique_ptr<SparseMatrix> build();

void clear();

private:
template <typename T, typename G>
void print_template();

template <typename T, typename G>
void build_template(std::unique_ptr<SparseMatrix> &);

private:
uint64 num_triplets_{0};
std::unique_ptr<uchar[]> data_base_ptr_{nullptr};
std::unique_ptr<Ndarray> ndarray_data_base_ptr_{nullptr};
int rows_{0};
int cols_{0};
uint64 max_num_triplets_{0};
bool built_{false};
DataType dtype_{PrimitiveType::f32};
std::string storage_format_{"col_major"};
Program *prog_{nullptr};
};

class SparseMatrix {
Expand Down
8 changes: 5 additions & 3 deletions taichi/python/export_lang.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -398,10 +398,11 @@ void export_lang(py::module &m) {
.def("create_sparse_matrix_builder",
[](Program *program, int n, int m, uint64 max_num_entries,
DataType dtype, const std::string &storage_format) {
TI_ERROR_IF(!arch_is_cpu(program->this_thread_config().arch),
"SparseMatrix Builder only supports CPU for now.");
TI_ERROR_IF(!arch_is_cpu(program->this_thread_config().arch) &&
!arch_is_cuda(program->this_thread_config().arch),
"SparseMatrix only supports CPU and CUDA for now.");
return SparseMatrixBuilder(n, m, max_num_entries, dtype,
storage_format);
storage_format, program);
})
.def("create_sparse_matrix",
[](Program *program, int n, int m, DataType dtype,
Expand Down Expand Up @@ -1204,6 +1205,7 @@ void export_lang(py::module &m) {
// Sparse Matrix
py::class_<SparseMatrixBuilder>(m, "SparseMatrixBuilder")
.def("print_triplets", &SparseMatrixBuilder::print_triplets)
.def("get_ndarray_data_ptr", &SparseMatrixBuilder::get_ndarray_data_ptr)
.def("build", &SparseMatrixBuilder::build)
.def("get_addr", [](SparseMatrixBuilder *mat) { return uint64(mat); });

Expand Down
22 changes: 11 additions & 11 deletions taichi/runtime/llvm/runtime_module/internal_functions.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,15 +9,15 @@
} \
} while (0)

#define ATOMIC_INSERT(T) \
do { \
auto base_ptr = (int64 *)base_ptr_; \
int64 *num_triplets = base_ptr; \
auto data_base_ptr = *(T **)(base_ptr + 1); \
auto triplet_id = atomic_add_i64(num_triplets, 1); \
data_base_ptr[triplet_id * 3] = i; \
data_base_ptr[triplet_id * 3 + 1] = j; \
data_base_ptr[triplet_id * 3 + 2] = taichi_union_cast<T>(value); \
#define ATOMIC_INSERT(T) \
do { \
auto base_ptr = reinterpret_cast<int##T *>(base_ptr_); \
int##T *num_triplets = base_ptr; \
auto data_base_ptr = base_ptr + 1; \
auto triplet_id = atomic_add_i##T(num_triplets, 1); \
data_base_ptr[triplet_id * 3] = i; \
data_base_ptr[triplet_id * 3 + 1] = j; \
data_base_ptr[triplet_id * 3 + 2] = taichi_union_cast<int##T>(value); \
} while (0);

i32 do_nothing(RuntimeContext *context) {
Expand All @@ -36,7 +36,7 @@ i32 insert_triplet_f32(RuntimeContext *context,
int i,
int j,
float value) {
ATOMIC_INSERT(int32);
ATOMIC_INSERT(32);
return 0;
}

Expand All @@ -45,7 +45,7 @@ i32 insert_triplet_f64(RuntimeContext *context,
int i,
int j,
float64 value) {
ATOMIC_INSERT(int64);
ATOMIC_INSERT(64);
return 0;
}

Expand Down
42 changes: 12 additions & 30 deletions tests/python/test_sparse_linear_solver.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,40 +4,19 @@
import taichi as ti
from tests import test_utils

"""
The symmetric positive definite matrix is created in matlab using the following script:
A = diag([1,2,3,4]);
OrthM = [1 0 1 0; -1 -2 0 1; 0 1 -1 0; 0, 1, 0 1];
U = orth(OrthM);
Aarray = U * A * U';
b = [1,2,3,4]';
res = inv(A) * b;
""" # yapf: disable

Aarray = np.array([[
2.73999501130921, 0.518002544441220, 0.745119303009342, 0.0508907745638859
], [0.518002544441220, 1.45111665837647, 0.757997555750432, 0.290885785873098],
[
0.745119303009342, 0.757997555750432, 2.96711176987733,
-0.518002544441220
],
[
0.0508907745638859, 0.290885785873098,
-0.518002544441220, 2.84177656043698
]])

res = np.array([
-0.0754984396447588, 0.469972700892492, 1.18527357933586, 1.57686870529319
])


@pytest.mark.parametrize("dtype", [ti.f32])
@pytest.mark.parametrize("solver_type", ["LLT", "LDLT", "LU"])
@pytest.mark.parametrize("ordering", ["AMD", "COLAMD"])
@test_utils.test(arch=ti.cpu)
def test_sparse_LLT_solver(dtype, solver_type, ordering):
n = 4
Abuilder = ti.linalg.SparseMatrixBuilder(n, n, max_num_triplets=100)
n = 10
A = np.random.rand(n, n)
A_psd = np.dot(A, A.transpose())
dtype = ti.float32
solver_type = "LLT"
ordering = "AMD"
FantasyVR marked this conversation as resolved.
Show resolved Hide resolved
Abuilder = ti.linalg.SparseMatrixBuilder(n, n, max_num_triplets=300)
b = ti.field(ti.f32, shape=n)

@ti.kernel
Expand All @@ -48,16 +27,19 @@ def fill(Abuilder: ti.types.sparse_matrix_builder(),
for i in range(n):
b[i] = i + 1

fill(Abuilder, Aarray, b)
fill(Abuilder, A_psd, b)
A = Abuilder.build()
print(A)
solver = ti.linalg.SparseSolver(dtype=dtype,
solver_type=solver_type,
ordering=ordering)
solver.analyze_pattern(A)
solver.factorize(A)
x = solver.solve(b)

res = np.linalg.solve(A_psd, b.to_numpy())
for i in range(n):
assert x[i] == test_utils.approx(res[i])
assert x[i] == test_utils.approx(res[i], rel=1.0)


@test_utils.test(arch=ti.cuda)
Expand Down
26 changes: 13 additions & 13 deletions tests/python/test_sparse_matrix.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,7 @@
(ti.f32, 'row_major'),
(ti.f64, 'col_major'),
(ti.f64, 'row_major')])
@test_utils.test(arch=ti.cpu)
@test_utils.test(arch=ti.cpu, offline_cache=False)
FantasyVR marked this conversation as resolved.
Show resolved Hide resolved
def test_sparse_matrix_builder_deprecated_anno(dtype, storage_format):
n = 8
Abuilder = ti.linalg.SparseMatrixBuilder(n,
Expand All @@ -34,7 +34,7 @@ def fill(Abuilder: ti.types.sparse_matrix_builder()):
(ti.f32, 'row_major'),
(ti.f64, 'col_major'),
(ti.f64, 'row_major')])
@test_utils.test(arch=ti.cpu)
@test_utils.test(arch=ti.cpu, offline_cache=False)
def test_sparse_matrix_builder(dtype, storage_format):
n = 8
Abuilder = ti.linalg.SparseMatrixBuilder(n,
Expand All @@ -59,7 +59,7 @@ def fill(Abuilder: ti.types.sparse_matrix_builder()):
(ti.f32, 'row_major'),
(ti.f64, 'col_major'),
(ti.f64, 'row_major')])
@test_utils.test(arch=ti.cpu)
@test_utils.test(arch=ti.cpu, offline_cache=False)
def test_build_sparse_matrix_frome_ndarray(dtype, storage_format):
n = 8
triplets = ti.Vector.ndarray(n=3, dtype=ti.f32, shape=n)
Expand All @@ -85,7 +85,7 @@ def fill(triplets: ti.types.ndarray()):
(ti.f32, 'row_major'),
(ti.f64, 'col_major'),
(ti.f64, 'row_major')])
@test_utils.test(arch=ti.cpu)
@test_utils.test(arch=ti.cpu, offline_cache=False)
def test_sparse_matrix_shape(dtype, storage_format):
n, m = 8, 9
Abuilder = ti.linalg.SparseMatrixBuilder(n,
Expand All @@ -108,7 +108,7 @@ def fill(Abuilder: ti.types.sparse_matrix_builder()):
(ti.f32, 'row_major'),
(ti.f64, 'col_major'),
(ti.f64, 'row_major')])
@test_utils.test(arch=ti.cpu)
@test_utils.test(arch=ti.cpu, offline_cache=False)
def test_sparse_matrix_element_access(dtype, storage_format):
n = 8
Abuilder = ti.linalg.SparseMatrixBuilder(n,
Expand All @@ -132,7 +132,7 @@ def fill(Abuilder: ti.types.sparse_matrix_builder()):
(ti.f32, 'row_major'),
(ti.f64, 'col_major'),
(ti.f64, 'row_major')])
@test_utils.test(arch=ti.cpu)
@test_utils.test(arch=ti.cpu, offline_cache=False)
def test_sparse_matrix_element_modify(dtype, storage_format):
n = 8
Abuilder = ti.linalg.SparseMatrixBuilder(n,
Expand All @@ -156,7 +156,7 @@ def fill(Abuilder: ti.types.sparse_matrix_builder()):
(ti.f32, 'row_major'),
(ti.f64, 'col_major'),
(ti.f64, 'row_major')])
@test_utils.test(arch=ti.cpu)
@test_utils.test(arch=ti.cpu, offline_cache=False)
def test_sparse_matrix_addition(dtype, storage_format):
n = 8
Abuilder = ti.linalg.SparseMatrixBuilder(n,
Expand Down Expand Up @@ -190,7 +190,7 @@ def fill(Abuilder: ti.types.sparse_matrix_builder(),
(ti.f32, 'row_major'),
(ti.f64, 'col_major'),
(ti.f64, 'row_major')])
@test_utils.test(arch=ti.cpu)
@test_utils.test(arch=ti.cpu, offline_cache=False)
def test_sparse_matrix_subtraction(dtype, storage_format):
n = 8
Abuilder = ti.linalg.SparseMatrixBuilder(n,
Expand Down Expand Up @@ -224,7 +224,7 @@ def fill(Abuilder: ti.types.sparse_matrix_builder(),
(ti.f32, 'row_major'),
(ti.f64, 'col_major'),
(ti.f64, 'row_major')])
@test_utils.test(arch=ti.cpu)
@test_utils.test(arch=ti.cpu, offline_cache=False)
def test_sparse_matrix_scalar_multiplication(dtype, storage_format):
n = 8
Abuilder = ti.linalg.SparseMatrixBuilder(n,
Expand All @@ -250,7 +250,7 @@ def fill(Abuilder: ti.types.sparse_matrix_builder()):
(ti.f32, 'row_major'),
(ti.f64, 'col_major'),
(ti.f64, 'row_major')])
@test_utils.test(arch=ti.cpu)
@test_utils.test(arch=ti.cpu, offline_cache=False)
def test_sparse_matrix_transpose(dtype, storage_format):
n = 8
Abuilder = ti.linalg.SparseMatrixBuilder(n,
Expand All @@ -276,7 +276,7 @@ def fill(Abuilder: ti.types.sparse_matrix_builder()):
(ti.f32, 'row_major'),
(ti.f64, 'col_major'),
(ti.f64, 'row_major')])
@test_utils.test(arch=ti.cpu)
@test_utils.test(arch=ti.cpu, offline_cache=False)
def test_sparse_matrix_elementwise_multiplication(dtype, storage_format):
n = 8
Abuilder = ti.linalg.SparseMatrixBuilder(n,
Expand Down Expand Up @@ -310,7 +310,7 @@ def fill(Abuilder: ti.types.sparse_matrix_builder(),
(ti.f32, 'row_major'),
(ti.f64, 'col_major'),
(ti.f64, 'row_major')])
@test_utils.test(arch=ti.cpu)
@test_utils.test(arch=ti.cpu, offline_cache=False)
def test_sparse_matrix_multiplication(dtype, storage_format):
n = 2
Abuilder = ti.linalg.SparseMatrixBuilder(n,
Expand Down Expand Up @@ -345,7 +345,7 @@ def fill(Abuilder: ti.types.sparse_matrix_builder(),
(ti.f32, 'row_major'),
(ti.f64, 'col_major'),
(ti.f64, 'row_major')])
@test_utils.test(arch=ti.cpu)
@test_utils.test(arch=ti.cpu, offline_cache=False)
def test_sparse_matrix_nonsymmetric_multiplication(dtype, storage_format):
n, k, m = 2, 3, 4
Abuilder = ti.linalg.SparseMatrixBuilder(n,
Expand Down
Loading