ValeevGroup / tiledarray

A massively-parallel, block-sparse tensor framework written in C++
GNU General Public License v3.0
254 stars 52 forks source link

function for element-wise operations with coordinate index #452

Open Marclie opened 4 months ago

Marclie commented 4 months ago

Adds a function forall that behaves like foreach_inplace, but also passes in the coordinate index with each tile. I've used this code in my projects with tiledarray to cleanly extract, fill, and modify the upper diagonal elements of a tiledarray object.

The function uses foreach_inplace and a recursive loop for iteration. The recursion helps the user focus on element-wise operations without explicitly creating loops for each dimension that use the lobound and upbound of each tile.

I added a test in tests/foreach.cpp to show its usage and ensure it works as intended. If you find this feature useful and everything looks good, feel free to merge. If this feature already exists elsewhere, please let me know. Thank you for your consideration!

evaleev commented 4 months ago

@Marclie thanks for the contribution!

We need to address several issues first:

I'm not actually sure if the proposed

vector<double> vec(n*n*n);
std::generate(v.begin(), v.end(), std::rand);
forall(array, [&vec] (auto& tile, auto& index) {
  size_t i = index[0], j = index[1], k = index[2];
  if (i <= j && j <= k) {
    tile[index] = std::sqrt(vec[i*n*n+j*n+k]);
  } else {
    tile[index] = 0.0;
  }
});

is any shorter than this

vector<double> vec(n*n*n);
std::generate(v.begin(), v.end(), std::rand);
foreach_inplace(array, [&vec] (auto& tile) {
  for(auto& idx: tile.range()) {
    size_t i = idx[0], j = idx[1], k = idx[2];
    if (i <= j && j <= k) {
      tile[idx] = std::sqrt(vec[i*n*n+j*n+k]);
    } else {
      tile[idx] = 0.0;
    }
  }
});

If you think there is a value to element wise operations we need to extend the foreach code to support such Op's by dispatching to tile-wise and element-wise paths based on the traits of Op ...

P.S. See also the implementation of DistArray::init_elements.

Marclie commented 4 months ago

@evaleev Thank you for your review!

I did not realize we could directly iterate over the tile ranges like that. Useful! We have been using upbound and lobound for everything and I came up with this hacky workaround. I modified the code to enable the foreach_inplace functions depending on the traits of Op; I also now directly use the iterators over the tile ranges.

However, I do agree that my wrapper is no more cleaner than the syntax for direct iteration over the range, and it may not add anything more substantial. Perhaps explicitly restricting the type traits of Op in this PR could still be useful for preventing users from passing an incorrect callback function to foreach_inplace, although I think that could be better done with static_assert that would give descriptive error messages.

Once more, the tile iteration is extremely useful to know. I am happy with closing my pull request knowing this functionality already exists in a clean syntax, and I'll be refactoring my projects with it.

evaleev commented 4 months ago

@Marclie I think the template parameter constraints are useful, I'm certainly happy to merge them in ...

re element wise ops: I'm currently on the fence about them, for the following reasons:

I agree that newbies want to use element-wise operations, and for constructing a DistArray we support this (DistArray ::init_elements()) but for foreach we need to think of the design more carefully.