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.

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

function for element-wise operations with coordinate index #452

Open
wants to merge 4 commits into
base: master
Choose a base branch
from

Conversation

Marclie
Copy link

@Marclie Marclie commented May 28, 2024

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!

Marclie and others added 2 commits May 28, 2024 13:23
…istArray objects that keeps track of the coordinate index for conditional operations.
@evaleev
Copy link
Member

evaleev commented May 31, 2024

@Marclie thanks for the contribution!

We need to address several issues first:

  • forall is new/alien terminology where foreach is already established (and been baked into std since early 90s). I suggest we stick with foreach. Since your function is actually in-place then it should be foreach_inplace.
  • Manual row-major iteration over the tile range, using non-inlinable recursion via std::function should be replaced by direct iteration over the range: for(auto& idx: tile.range()) { ...

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
Copy link
Author

Marclie commented May 31, 2024

@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
Copy link
Member

evaleev commented May 31, 2024

@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:

  • no single "natural" signature E.g., for foreach_inplace() should the signature be
    • (Tile&, ElementIndex&) -> void, like you propose, or
    • (Tile::value_type&, ElementIndex&) -> void, or
    • smth else?
  • interaction with tile-level Op variants ... currently tile Op can return void or a real number, the latter interpreted as tile norm. It's not obvious which of the tile-wise op signatures should we use for sparse tensors when we receive an element-wise op from the user.

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.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants