BLI: refactor IndexMask for better performance and memory usage #104629

Merged
Jacques Lucke merged 254 commits from JacquesLucke/blender:index-mask-refactor into main 2023-05-24 18:11:47 +02:00
6 changed files with 167 additions and 73 deletions
Showing only changes of commit aed0ec5d65 - Show all commits

View File

@ -165,15 +165,13 @@ class IndexMask {
IndexMask slice(IndexRange range) const;
void foreach_span(FunctionRef<void(OffsetSpan<int64_t, int16_t>)> fn) const;
template<typename Fn> void foreach_span(Fn &&fn) const;
template<typename Fn> void foreach_range(Fn &&fn) const;
template<typename Fn> void foreach_span_or_range(Fn &&fn) const;
template<typename Fn> void foreach_index(Fn &&fn) const;
void foreach_span_parallel(int64_t grain_size,
FunctionRef<void(OffsetSpan<int64_t, int16_t>)> fn) const;
void foreach_span_parallel(int64_t grain_size,
FunctionRef<void(OffsetSpan<int64_t, int16_t>, IndexRange)> fn) const;
template<typename Fn> void foreach_span_or_range_index(Fn &&fn) const;
template<typename Fn> void foreach_span_parallel(int64_t grain_size, Fn &&fn) const;
template<typename Fn> void foreach_span_or_range_parallel(int64_t grain_size, Fn &&fn) const;
template<typename T> static IndexMask from_indices(Span<T> indices, IndexMaskMemory &memory);
static IndexMask from_bits(BitSpan bits, IndexMaskMemory &memory, int64_t offset = 0);
@ -205,6 +203,9 @@ class IndexMask {
const IndexMaskData &data() const;
IndexMaskData &data_for_inplace_construction();
private:
void foreach_span_impl(FunctionRef<void(OffsetSpan<int64_t, int16_t>)> fn) const;
};
std::ostream &operator<<(std::ostream &stream, const IndexMask &mask);
@ -633,11 +634,39 @@ template<typename Fn> inline void ChunkSlice::foreach_span(Fn &&fn) const
}
}
template<typename Fn>
constexpr bool has_mask_segment_and_start_parameter =
std::is_invocable_r_v<void, Fn, OffsetSpan<int64_t, int16_t>, int64_t> ||
std::is_invocable_r_v<void, Fn, IndexRange, int64_t>;
template<typename Fn> inline void IndexMask::foreach_index(Fn &&fn) const
{
this->foreach_span([&](const OffsetSpan<int64_t, int16_t> indices) {
for (const int64_t index : indices) {
fn(index);
this->foreach_span([&](const OffsetSpan<int64_t, int16_t> indices, const int64_t start) {
if constexpr (std::is_invocable_r_v<void, Fn, int64_t, int64_t>) {
for (const int64_t i : indices.index_range()) {
fn(indices[i], start + i);
}
}
else {
for (const int64_t index : indices) {
fn(index);
}
}
});
}
template<typename Fn> inline void IndexMask::foreach_span_or_range_index(Fn &&fn) const
{
this->foreach_span_or_range([&](const auto mask_segment, const int64_t start) {
if constexpr (std::is_invocable_r_v<void, Fn, int64_t, int64_t>) {
for (const int64_t i : mask_segment.index_range()) {
fn(mask_segment[i], start + i);
}
}
else {
for (const int64_t index : mask_segment) {
fn(index);
}
}
});
}
@ -645,46 +674,90 @@ template<typename Fn> inline void IndexMask::foreach_index(Fn &&fn) const
template<typename Fn> inline void IndexMask::foreach_span_or_range(Fn &&fn) const
{
IndexRangeChecker is_index_mask;
this->foreach_span([&, is_index_mask](const OffsetSpan<int64_t, int16_t> indices) {
if (is_index_mask.check(indices.base_span())) {
fn(IndexRange(indices[0], indices.size()));
}
else {
fn(indices);
}
});
this->foreach_span(
[&, is_index_mask](const OffsetSpan<int64_t, int16_t> mask_segment, const int64_t start) {
if (is_index_mask.check(mask_segment.base_span())) {
const IndexRange range(mask_segment[0], mask_segment.size());
if constexpr (has_mask_segment_and_start_parameter<Fn>) {
fn(range, start);
}
else {
fn(range);
}
}
else {
if constexpr (has_mask_segment_and_start_parameter<Fn>) {
fn(mask_segment, start);
}
else {
fn(mask_segment);
}
}
});
}
template<typename Fn> inline void IndexMask::foreach_span(Fn &&fn) const
{
if constexpr (has_mask_segment_and_start_parameter<Fn>) {
this->foreach_span_impl(
[&, counter = int64_t(0)](const OffsetSpan<int64_t, int16_t> mask_segment) mutable {
fn(mask_segment, counter);
counter += mask_segment.size();
});
}
else {
this->foreach_span_impl(fn);
}
}
template<typename Fn> inline void IndexMask::foreach_range(Fn &&fn) const
{
this->foreach_span([&](const OffsetSpan<int64_t, int16_t> indices) {
this->foreach_span([&](const OffsetSpan<int64_t, int16_t> indices, int64_t start) {
Span<int16_t> base_indices = indices.base_span();
while (!base_indices.is_empty()) {
const int64_t next_range_size = unique_sorted_indices::find_size_of_next_range(base_indices);
fn(IndexRange(int64_t(base_indices[0]) + indices.offset(), next_range_size));
const IndexRange range(int64_t(base_indices[0]) + indices.offset(), next_range_size);
if constexpr (has_mask_segment_and_start_parameter<Fn>) {
fn(range, start);
}
else {
fn(range);
}
start += next_range_size;
base_indices = base_indices.drop_front(next_range_size);
}
});
}
inline void IndexMask::foreach_span_parallel(
const int64_t grain_size, const FunctionRef<void(OffsetSpan<int64_t, int16_t>)> fn) const
template<typename Fn>
inline void IndexMask::foreach_span_parallel(const int64_t grain_size, Fn &&fn) const
{
threading::parallel_for(this->index_range(), grain_size, [&](const IndexRange range) {
const IndexMask sub_mask = this->slice(range);
sub_mask.foreach_span(fn);
sub_mask.foreach_span(
[&](const OffsetSpan<int64_t, int16_t> mask_segment, const int64_t start) {
if constexpr (has_mask_segment_and_start_parameter<Fn>) {
fn(mask_segment, start + range.start());
}
else {
fn(mask_segment);
}
});
});
}
inline void IndexMask::foreach_span_parallel(
int64_t grain_size, FunctionRef<void(OffsetSpan<int64_t, int16_t>, IndexRange)> fn) const
template<typename Fn>
void IndexMask::foreach_span_or_range_parallel(int64_t grain_size, Fn &&fn) const
{
threading::parallel_for(this->index_range(), grain_size, [&](const IndexRange range) {
const IndexMask sub_mask = this->slice(range);
int64_t counter = 0;
sub_mask.foreach_span([&](const OffsetSpan<int64_t, int16_t> mask_segment) {
fn(mask_segment, range.slice(counter, mask_segment.size()));
counter += mask_segment.size();
sub_mask.foreach_span_or_range([&](const auto mask_segment, const int64_t start) {
if constexpr (has_mask_segment_and_start_parameter<Fn>) {
fn(mask_segment, start + range.start());
}
else {
fn(mask_segment);
}
});
});
}
@ -713,6 +786,7 @@ inline IndexMask IndexMask::from_predicate(const IndexMask &universe,
IndexMaskMemory &memory,
Fn &&predicate)
{
UNUSED_VARS(grain_size);
Vector<int64_t> indices;
universe.foreach_index([&](const int64_t index) {
if (predicate(index)) {

View File

@ -49,10 +49,10 @@ void GVectorArray::extend(const int64_t index, const GSpan values)
void GVectorArray::extend(IndexMask mask, const GVVectorArray &values)
{
for (const int i : mask) {
mask.foreach_index([&](const int64_t i) {
GVArray_For_GVVectorArrayIndex array{values, i};
this->extend(i, GVArray(&array));
}
});
}
void GVectorArray::extend(IndexMask mask, const GVectorArray &values)
@ -63,11 +63,11 @@ void GVectorArray::extend(IndexMask mask, const GVectorArray &values)
void GVectorArray::clear(IndexMask mask)
{
for (const int64_t i : mask) {
mask.foreach_index([&](const int64_t i) {
Item &item = items_[i];
type_.destruct_n(item.start, item.length);
item.length = 0;
}
});
}
GMutableSpan GVectorArray::operator[](const int64_t index)

View File

@ -10,34 +10,34 @@ namespace blender {
void GVArrayImpl::materialize(const IndexMask mask, void *dst) const
{
for (const int64_t i : mask) {
mask.foreach_span_or_range_index([&](const int64_t i) {
void *elem_dst = POINTER_OFFSET(dst, type_->size() * i);
this->get(i, elem_dst);
}
});
}
void GVArrayImpl::materialize_to_uninitialized(const IndexMask mask, void *dst) const
{
for (const int64_t i : mask) {
mask.foreach_span_or_range_index([&](const int64_t i) {
void *elem_dst = POINTER_OFFSET(dst, type_->size() * i);
this->get_to_uninitialized(i, elem_dst);
}
});
}
void GVArrayImpl::materialize_compressed(IndexMask mask, void *dst) const
{
for (const int64_t i : mask.index_range()) {
void *elem_dst = POINTER_OFFSET(dst, type_->size() * i);
this->get(mask[i], elem_dst);
}
mask.foreach_span_or_range_index([&](const int64_t i, const int64_t i_in_mask) {
void *elem_dst = POINTER_OFFSET(dst, type_->size() * i_in_mask);
this->get(i, elem_dst);
});
}
void GVArrayImpl::materialize_compressed_to_uninitialized(IndexMask mask, void *dst) const
{
for (const int64_t i : mask.index_range()) {
void *elem_dst = POINTER_OFFSET(dst, type_->size() * i);
this->get_to_uninitialized(mask[i], elem_dst);
}
mask.foreach_span_or_range_index([&](const int64_t i, const int64_t i_in_mask) {
void *elem_dst = POINTER_OFFSET(dst, type_->size() * i_in_mask);
this->get_to_uninitialized(i, elem_dst);
});
}
void GVArrayImpl::get(const int64_t index, void *r_value) const
@ -497,18 +497,11 @@ class GVArrayImpl_For_SlicedGVArray : public GVArrayImpl {
void materialize_compressed_to_uninitialized(const IndexMask mask, void *dst) const override
{
if (mask.is_range()) {
const IndexRange mask_range = mask.as_range();
const IndexRange offset_mask_range{mask_range.start() + offset_, mask_range.size()};
varray_.materialize_compressed_to_uninitialized(offset_mask_range, dst);
}
else {
Vector<int64_t, 32> offset_mask_indices(mask.size());
for (const int64_t i : mask.index_range()) {
offset_mask_indices[i] = mask[i] + offset_;
}
varray_.materialize_compressed_to_uninitialized(offset_mask_indices.as_span(), dst);
}
mask.foreach_range([&](const IndexRange mask_segment, const int64_t start) {
const IndexRange offset_mask_range = mask_segment.shift(offset_);
varray_.materialize_compressed_to_uninitialized(offset_mask_range,
POINTER_OFFSET(dst, type_->size() * start));
});
}
};

View File

@ -337,7 +337,7 @@ template<typename T> void from_index_mask(const IndexMask &mask, MutableSpan<T>
} // namespace unique_sorted_indices
void IndexMask::foreach_span(const FunctionRef<void(OffsetSpan<int64_t, int16_t>)> fn) const
void IndexMask::foreach_span_impl(const FunctionRef<void(OffsetSpan<int64_t, int16_t>)> fn) const
{
if (data_.indices_num == 0) {
return;

View File

@ -109,7 +109,9 @@ TEST(cpp_type, DefaultConstruction)
EXPECT_EQ(buffer[1], default_constructed_value);
EXPECT_EQ(buffer[2], default_constructed_value);
EXPECT_EQ(buffer[3], 0);
CPPType_TestType.default_construct_indices((void *)buffer, {2, 5, 7});
IndexMaskMemory memory;
CPPType_TestType.default_construct_indices((void *)buffer,
IndexMask::from_indices<int>({2, 5, 7}, memory));
EXPECT_EQ(buffer[2], default_constructed_value);
EXPECT_EQ(buffer[4], 0);
EXPECT_EQ(buffer[5], default_constructed_value);
@ -136,7 +138,9 @@ TEST(cpp_type, ValueInitialize)
EXPECT_EQ(buffer[1], default_constructed_value);
EXPECT_EQ(buffer[2], default_constructed_value);
EXPECT_EQ(buffer[3], 0);
CPPType_TestType.value_initialize_indices((void *)buffer, {2, 5, 7});
IndexMaskMemory memory;
CPPType_TestType.value_initialize_indices((void *)buffer,
IndexMask::from_indices<int>({2, 5, 7}, memory));
EXPECT_EQ(buffer[2], default_constructed_value);
EXPECT_EQ(buffer[4], 0);
EXPECT_EQ(buffer[5], default_constructed_value);
@ -163,7 +167,9 @@ TEST(cpp_type, Destruct)
EXPECT_EQ(buffer[1], destructed_value);
EXPECT_EQ(buffer[2], destructed_value);
EXPECT_EQ(buffer[3], 0);
CPPType_TestType.destruct_indices((void *)buffer, {2, 5, 7});
IndexMaskMemory memory;
CPPType_TestType.destruct_indices((void *)buffer,
IndexMask::from_indices<int>({2, 5, 7}, memory));
EXPECT_EQ(buffer[2], destructed_value);
EXPECT_EQ(buffer[4], 0);
EXPECT_EQ(buffer[5], destructed_value);
@ -188,7 +194,9 @@ TEST(cpp_type, CopyToUninitialized)
EXPECT_EQ(buffer2[2], copy_constructed_value);
EXPECT_EQ(buffer1[3], 0);
EXPECT_EQ(buffer2[3], 0);
CPPType_TestType.copy_construct_indices((void *)buffer1, (void *)buffer2, {2, 5, 7});
IndexMaskMemory memory;
CPPType_TestType.copy_construct_indices(
(void *)buffer1, (void *)buffer2, IndexMask::from_indices<int>({2, 5, 7}, memory));
EXPECT_EQ(buffer1[2], copy_constructed_from_value);
EXPECT_EQ(buffer2[2], copy_constructed_value);
EXPECT_EQ(buffer1[4], 0);
@ -219,7 +227,9 @@ TEST(cpp_type, CopyToInitialized)
EXPECT_EQ(buffer2[2], copy_assigned_value);
EXPECT_EQ(buffer1[3], 0);
EXPECT_EQ(buffer2[3], 0);
CPPType_TestType.copy_assign_indices((void *)buffer1, (void *)buffer2, {2, 5, 7});
IndexMaskMemory memory;
CPPType_TestType.copy_assign_indices(
(void *)buffer1, (void *)buffer2, IndexMask::from_indices<int>({2, 5, 7}, memory));
EXPECT_EQ(buffer1[2], copy_assigned_from_value);
EXPECT_EQ(buffer2[2], copy_assigned_value);
EXPECT_EQ(buffer1[4], 0);
@ -250,7 +260,9 @@ TEST(cpp_type, RelocateToUninitialized)
EXPECT_EQ(buffer2[2], move_constructed_value);
EXPECT_EQ(buffer1[3], 0);
EXPECT_EQ(buffer2[3], 0);
CPPType_TestType.relocate_construct_indices((void *)buffer1, (void *)buffer2, {2, 5, 7});
IndexMaskMemory memory;
CPPType_TestType.relocate_construct_indices(
(void *)buffer1, (void *)buffer2, IndexMask::from_indices<int>({2, 5, 7}, memory));
EXPECT_EQ(buffer1[2], destructed_value);
EXPECT_EQ(buffer2[2], move_constructed_value);
EXPECT_EQ(buffer1[4], 0);
@ -281,7 +293,9 @@ TEST(cpp_type, RelocateToInitialized)
EXPECT_EQ(buffer2[2], move_assigned_value);
EXPECT_EQ(buffer1[3], 0);
EXPECT_EQ(buffer2[3], 0);
CPPType_TestType.relocate_assign_indices((void *)buffer1, (void *)buffer2, {2, 5, 7});
IndexMaskMemory memory;
CPPType_TestType.relocate_assign_indices(
(void *)buffer1, (void *)buffer2, IndexMask::from_indices<int>({2, 5, 7}, memory));
EXPECT_EQ(buffer1[2], destructed_value);
EXPECT_EQ(buffer2[2], move_assigned_value);
EXPECT_EQ(buffer1[4], 0);
@ -308,7 +322,9 @@ TEST(cpp_type, FillInitialized)
EXPECT_EQ(buffer2[3], 0);
buffer1 = 0;
CPPType_TestType.fill_assign_indices((void *)&buffer1, (void *)buffer2, {1, 6, 8});
IndexMaskMemory memory;
CPPType_TestType.fill_assign_indices(
(void *)&buffer1, (void *)buffer2, IndexMask::from_indices<int>({1, 6, 8}, memory));
EXPECT_EQ(buffer1, copy_assigned_from_value);
EXPECT_EQ(buffer2[0], copy_assigned_value);
EXPECT_EQ(buffer2[1], copy_assigned_value);
@ -334,7 +350,9 @@ TEST(cpp_type, FillUninitialized)
EXPECT_EQ(buffer2[3], 0);
buffer1 = 0;
CPPType_TestType.fill_construct_indices((void *)&buffer1, (void *)buffer2, {1, 6, 8});
IndexMaskMemory memory;
CPPType_TestType.fill_construct_indices(
(void *)&buffer1, (void *)buffer2, IndexMask::from_indices<int>({1, 6, 8}, memory));
EXPECT_EQ(buffer1, copy_constructed_from_value);
EXPECT_EQ(buffer2[0], copy_constructed_value);
EXPECT_EQ(buffer2[1], copy_constructed_value);
@ -385,7 +403,9 @@ TEST(cpp_type, CopyAssignCompressed)
{
std::array<std::string, 5> array = {"a", "b", "c", "d", "e"};
std::array<std::string, 3> array_compressed;
CPPType::get<std::string>().copy_assign_compressed(&array, &array_compressed, {0, 2, 3});
IndexMaskMemory memory;
CPPType::get<std::string>().copy_assign_compressed(
&array, &array_compressed, IndexMask::from_indices<int>({0, 2, 3}, memory));
EXPECT_EQ(array_compressed[0], "a");
EXPECT_EQ(array_compressed[1], "c");
EXPECT_EQ(array_compressed[2], "d");

View File

@ -183,15 +183,18 @@ TEST(virtual_array, MutableToImmutable)
TEST(virtual_array, MaterializeCompressed)
{
IndexMaskMemory memory;
{
std::array<int, 10> array = {0, 10, 20, 30, 40, 50, 60, 70, 80, 90};
VArray<int> varray = VArray<int>::ForSpan(array);
std::array<int, 3> compressed_array;
varray.materialize_compressed({3, 6, 7}, compressed_array);
varray.materialize_compressed(IndexMask::from_indices<int>({3, 6, 7}, memory),
compressed_array);
EXPECT_EQ(compressed_array[0], 30);
EXPECT_EQ(compressed_array[1], 60);
EXPECT_EQ(compressed_array[2], 70);
varray.materialize_compressed_to_uninitialized({2, 8, 9}, compressed_array);
varray.materialize_compressed_to_uninitialized(IndexMask::from_indices<int>({2, 8, 9}, memory),
compressed_array);
EXPECT_EQ(compressed_array[0], 20);
EXPECT_EQ(compressed_array[1], 80);
EXPECT_EQ(compressed_array[2], 90);
@ -199,12 +202,14 @@ TEST(virtual_array, MaterializeCompressed)
{
VArray<int> varray = VArray<int>::ForSingle(4, 10);
std::array<int, 3> compressed_array;
varray.materialize_compressed({2, 6, 7}, compressed_array);
varray.materialize_compressed(IndexMask::from_indices<int>({2, 6, 7}, memory),
compressed_array);
EXPECT_EQ(compressed_array[0], 4);
EXPECT_EQ(compressed_array[1], 4);
EXPECT_EQ(compressed_array[2], 4);
compressed_array.fill(0);
varray.materialize_compressed_to_uninitialized({0, 1, 2}, compressed_array);
varray.materialize_compressed_to_uninitialized(IndexMask::from_indices<int>({0, 1, 2}, memory),
compressed_array);
EXPECT_EQ(compressed_array[0], 4);
EXPECT_EQ(compressed_array[1], 4);
EXPECT_EQ(compressed_array[2], 4);
@ -212,11 +217,13 @@ TEST(virtual_array, MaterializeCompressed)
{
VArray<int> varray = VArray<int>::ForFunc(10, [](const int64_t i) { return int(i * i); });
std::array<int, 3> compressed_array;
varray.materialize_compressed({5, 7, 8}, compressed_array);
varray.materialize_compressed(IndexMask::from_indices<int>({5, 7, 8}, memory),
compressed_array);
EXPECT_EQ(compressed_array[0], 25);
EXPECT_EQ(compressed_array[1], 49);
EXPECT_EQ(compressed_array[2], 64);
varray.materialize_compressed_to_uninitialized({1, 2, 3}, compressed_array);
varray.materialize_compressed_to_uninitialized(IndexMask::from_indices<int>({1, 2, 3}, memory),
compressed_array);
EXPECT_EQ(compressed_array[0], 1);
EXPECT_EQ(compressed_array[1], 4);
EXPECT_EQ(compressed_array[2], 9);