Implementing Putmask #667
Implementing Putmask #667
Conversation
ipdemes
added
the
category:new-feature
|
This PR implements logic for putmask funtion. |
| }; | ||
|
|
||
| template <> | ||
| struct BoolMaskPolicy<VariantKind::GPU, false> { |
There was a problem hiding this comment.
does this need to be specific to boolean masks? this policy cannot be used without materializing a boolean mask in a case where, for example, the kernel needs to run when the input is not NaN. it'd be more general if this policy took a predicate that evaluates to a boolean value on each point.
There was a problem hiding this comment.
Makes sense. I will make it more general and call it parallel_loop, probably.
cunumeric/deferred.py
Outdated
| is_scalar_value = False | ||
| if values.shape != self.shape and values.size == 1: | ||
| if values.shape == (): | ||
| values = values._convert_future_to_regionfield(True) |
There was a problem hiding this comment.
is there a reason that you do this? can't you just promote this scalar store to match the shape of self?
There was a problem hiding this comment.
replaced this with broadcasting
cunumeric/deferred.py
Outdated
| task.add_input(mask.base) | ||
| task.add_input(values.base) | ||
| task.add_output(self.base) | ||
| task.add_scalar_arg(is_scalar_value, bool) # value is scalar |
There was a problem hiding this comment.
again, is there a reason that you don't do numpy broadcasting on the scalar value? that would make the code a lot simpler.
There was a problem hiding this comment.
makes sense. replaced the logic with broadcasting
cunumeric/deferred.py
Outdated
| task.add_broadcast(values.base) | ||
| else: | ||
| task.add_alignment(self.base, values.base) | ||
| task.add_broadcast(self.base, axes=range(1, self.ndim)) |
There was a problem hiding this comment.
I don't understand this broadcast constraint... why is this necessary?
There was a problem hiding this comment.
removed. I don't need it anymore
| using namespace Legion; | ||
| using namespace legate; | ||
|
|
||
| template <VariantKind KIND, LegateTypeCode CODE, int DIM, int VDIM, bool SCALAR_VALUE = false> |
There was a problem hiding this comment.
like I said in the other comment, if you do numpy broadcasting on the values, you can remove the SCALAR_VALUE == true case.
cunumeric/module.py
Outdated
|
|
||
| if a.dtype != values.dtype: | ||
| values = values._warn_and_convert(a.dtype) | ||
| if values.shape != a.shape and values.size != 1: |
There was a problem hiding this comment.
the first condition should be values.size != size. Here's an example you can handle without wrapping or tiling but simply using numpy broadcasting:
import numpy as np
a = np.full((3, 3), 3)
np.putmask(a, np.full((3, 3), True), np.full(3, 10))
There was a problem hiding this comment.
unless stated otherwise, it's safe to assume that all related array arguments follow the numpy broadcasting semantics.
There was a problem hiding this comment.
I adding the check if arrays can be broadcasted. We will cal wrap only for the case when they can't
tests/integration/test_putmask.py
Outdated
| num.putmask(num_arr, num_mask, num_val) | ||
| assert np.array_equal(np_arr, num_arr) | ||
|
|
||
| # val is different shape |
There was a problem hiding this comment.
you need to subdivide this test into two: 1) when the shapes are different but the sizes are the same, 2) when both the shapes and sizes are different.
| has_set_value = set_value is not None and set_value.size == 1 | ||
| if has_set_value: | ||
| mask = DeferredArray( | ||
| self.runtime, | ||
| base=key_store, | ||
| dtype=self.dtype, | ||
| ) | ||
| rhs.putmask(mask, set_value) | ||
| return False, rhs, rhs, self |
There was a problem hiding this comment.
I'm not a big fan of this code. the name of the function is _create_indexing_array, but this code makes it do more than just creating indexing arrays. and in fact, the function was already complicated enough to warrant a refactoring, which I was thinking of doing at some point. I guess I'm fine with accepting this edit for this PR, but keep in mind that _create_indexing_array can use some refactoring for better maintainability.
cunumeric/module.py
Outdated
| @@ -3506,6 +3506,58 @@ def put( | |||
| a.put(indices=indices, values=values, mode=mode) | |||
|
|
|||
|
|
|||
| def _can_broadcast_to_shape(self: ndarray, shape: NdShape) -> bool: | |||
There was a problem hiding this comment.
can't we just use np.broadcast_shapes instead of this custom code?
There was a problem hiding this comment.
sure, will use it instead with `try\ except"
No description provided.