Skip to content

Commit 24ba646

Browse files
committed
[Code/Python] Vector fixes && Cuda matrix-to-vec reduce fixes
1 parent cb0e8d3 commit 24ba646

File tree

6 files changed

+34
-26
lines changed

6 files changed

+34
-26
lines changed

cubool/sources/cuda/kernels/spreduce.cuh

Lines changed: 9 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -143,7 +143,6 @@ namespace cubool {
143143
using VectorType = details::SpVector<IndexType, AllocType>;
144144

145145
VectorType operator()(const MatrixType& m) {
146-
auto nrows = m.m_rows;
147146
auto ncols = m.m_cols;
148147
auto nvals = m.m_vals;
149148

@@ -156,33 +155,34 @@ namespace cubool {
156155
__device__ (IndexType i) {
157156
auto j = colIndex[i];
158157

159-
atomicOr((mask + j).get(), (IndexType) 1);
158+
mask[j] = 0x1u;
160159
}
161160
);
162161

163162
// Count nnz of the vector
164163
auto resultSize = thrust::reduce(mask.begin(), mask.end(), (IndexType) 0);
165164

166165
ContainerType<index> result(resultSize);
167-
ContainerType<index> offsets(mask.size());
168-
169-
// Evaluate write offsets of the result
170-
thrust::exclusive_scan(mask.begin(), mask.end(), offsets.begin(), (IndexType) 0, thrust::plus<IndexType>());
166+
ContainerType<index> order(1);
167+
thrust::fill(order.begin(), order.end(), (IndexType) 0);
171168

172169
// For each value of the mask write in the result buffer if it not zero
173170
thrust::for_each(thrust::counting_iterator<IndexType>(0), thrust::counting_iterator<IndexType>(ncols),
174-
[mask = mask.data(), result = result.data(), offsets = offsets.data()]
171+
[mask = mask.data(), result = result.data(), order = order.data()]
175172
__device__ (IndexType i) {
176173
if (mask[i] > 0) {
177-
auto offset = offsets[i];
174+
auto offset = atomicAdd(order.get(), (IndexType) 1);
178175
auto colId = i;
179176

180177
result[offset] = colId;
181178
}
182179
}
183180
);
184181

185-
return VectorType(std::move(result), nrows, resultSize);
182+
// Sort values within vector
183+
thrust::sort(result.begin(), result.end());
184+
185+
return VectorType(std::move(result), ncols, resultSize);
186186
}
187187

188188
};

python/features.py

Lines changed: 17 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -111,5 +111,20 @@ def gen_matrix(size, seed):
111111
print(vals)
112112
print(a.equals(a))
113113

114-
matrix = pycubool.Matrix.generate(shape=(4, 4), density=0.5)
115-
print(matrix)
114+
matrix = pycubool.Matrix.generate(shape=(5, 4), density=0.5)
115+
print(matrix, matrix.extract_row(1), matrix.extract_col(2), sep="\n")
116+
117+
matrix = pycubool.Matrix.generate(shape=(5, 4), density=0.5)
118+
vector = pycubool.Vector.generate(nrows=4, density=0.6)
119+
print(matrix, vector, matrix.mxv(vector), sep="\n")
120+
121+
matrix = pycubool.Matrix.generate(shape=(5, 4), density=0.5)
122+
vector = pycubool.Vector.generate(nrows=5, density=0.6)
123+
print(matrix, vector, vector.vxm(matrix), sep="\n")
124+
125+
matrix = pycubool.Matrix.generate(shape=(10, 6), density=0.2)
126+
print(matrix, matrix.reduce_vector(transpose=False), matrix.reduce_vector(transpose=True), sep="\n")
127+
128+
129+
130+

python/pycubool/__init__.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121
from .wrapper import *
2222
from .utils import *
2323
from .matrix import *
24+
from .vector import *
2425
from .io import *
2526
from .gviz import *
2627

python/pycubool/__main__.py

Lines changed: 0 additions & 11 deletions
This file was deleted.

python/pycubool/matrix.py

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -493,7 +493,7 @@ def extract_matrix(self, i, j, shape, out=None, time_check=False):
493493

494494
def extract_row(self, i, out=None):
495495
if out is None:
496-
out = Vector.empty(self.ncols)
496+
out = vector.Vector.empty(self.ncols)
497497

498498
status = wrapper.loaded_dll.cuBool_Matrix_ExtractRow(
499499
out.hnd,
@@ -507,7 +507,7 @@ def extract_row(self, i, out=None):
507507

508508
def extract_col(self, j, out=None):
509509
if out is None:
510-
out = Vector.empty(self.nrows)
510+
out = vector.Vector.empty(self.nrows)
511511

512512
status = wrapper.loaded_dll.cuBool_Matrix_ExtractCol(
513513
out.hnd,
@@ -563,7 +563,7 @@ def mxm(self, other, out=None, accumulate=False, time_check=False):
563563

564564
def mxv(self, other, out=None, time_check=False):
565565
if out is None:
566-
out = Vector.empty(self.nrows)
566+
out = vector.Vector.empty(self.nrows)
567567

568568
status = wrapper.loaded_dll.cuBool_MxV(
569569
out.hnd,
@@ -689,7 +689,7 @@ def reduce(self, time_check=False):
689689
def reduce_vector(self, out=None, transpose=False, time_check=False):
690690
if out is None:
691691
nrows = self.ncols if transpose else self.nrows
692-
out = Vector.empty(nrows)
692+
out = vector.Vector.empty(nrows)
693693

694694
status = wrapper.loaded_dll.cuBool_Matrix_Reduce(
695695
out.hnd,

python/pycubool/vector.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -229,6 +229,9 @@ def ewiseadd(self, other, out=None, time_check=False):
229229
ctypes.c_uint(bridge.get_ewiseadd_hints(time_check=time_check))
230230
)
231231

232+
bridge.check(status)
233+
return out
234+
232235
def reduce(self, time_check=False):
233236
value = ctypes.c_uint(0)
234237

0 commit comments

Comments
 (0)