Implement AdvancedIndexingTask using thrust::copy_if

[manopapad]

See #1215 (comment). Quoting relevant part here:

I also have a quick question about memory usage. Since it worked fine on the other cluster I wanted to see the maximum size of array that will work. I noticed that the operation a_train = a[train_mask] requires peak memory of about 2.5 times the size of the original matrix a, when the mask train_mask selects 50% of the rows (train_mask is 1D array).

My expectation was a memory footprint of about 1.5x (1x for a + 0.5x for a_train). Can you confirm if this 2.5x peak memory usage is expected for boolean mask indexing? If so, could you briefly explain why the temporary memory requirement is so high?

The current implementation of the "masked copy" operation first calculates the offsets of the non-zeroes (which uses an array of size equal to that of the original mask array), then creates the output array and uses the offsets to fill it in. Therefore, the memory overhead you quote is expected, under the existing code. See https://github.com/nv-legate/cupynumeric/blob/main/src/cupynumeric/index/advanced_indexing.cu#L120.

But possibly for the case of a dense array we could use thrust::copy_if, which has lower memory requirements.

[aschaffer]

Offsets are necessary for the cases when (1) the mask is "compressed"; i.e., the size of mask is less than input size (|mask| < |input|). This is because in such cases the mask (index array) doesn't contain enough information to cover all the slices that do satisfy the mask.

Hence, an "expanded" index array is necessary to cover those. This is what the offsets are used for (after the exclusive_scan); otherwise, the slicing would include much fewer entries than what's intended.

However, in the cases (2), where the mask is explicitly specifying the entries to include ("un-compressed") the offsets would not be necessary.

The problem, though, is how to differentiate between cases (1) and (2). Unfortunately, it cannot be done by comparing |mask| vs. |input| sizes, because when the runtime reaches the C++ layer, they have the same size (see assert(index_rect == input_rect); here) regardless of the mask having been compressed or un-compressed in the Python layer.

If there were a way to differentiate between (1) and (2) in the C++ layer then the copy_if() would provide a less memory consuming solution (the majority of tests in tests/integration/test_advanced_indexing.py fall under the un-compressed case (2)).

[dinodeep]

@aschaffer Sorry to bring this up after a long time, but regarding this statement

The problem, though, is how to differentiate between cases (1) and (2).

Wouldn't it be possible to do this by checking if args.key_dim < DIM here? My understanding is that (1) only occurs when the boolean mask is applied to a subset of the dimensions of the array that it is masking. The Python side does project it here; however, we still pass in the original dimension size of the keys in args.key_dim as a scalar argument here. Would this be sufficient? Let me know if my understanding is incorrect on this

[dinodeep]

@manopapad Looking back on the original issue at the top here, the motivation for the thrust::copy_if implementation was to reduce memory requirements, not necessarily improve the performance of the task implementation.

The new scan + masked getitem implementation https://github.com/nv-legate/cupynumeric.internal/pull/1708 in this PR won't resolve the memory usage issue because it still requires allocating a large buffer to store indices that we then perform the masked getitem on.

To resolve the problem mentioned in this issue (which we don't really hit in any of the performance benchmarking right now), it may be worth considering implementing AdvancedIndexing via thrust::copy_if, and then, we can provide the user the means to switch between implementations depending on if they are trying to optimize for performance or memory usage.

It would be ideal if we had a means to inspect the memory usage at mapping time (so we could appropriately set pool bounds) and automatically choose an implementation based on how much memory is available (if there is enough for scan + masked getitem, then use that implementation, otherwise, use the thrust::copy_if implementation); however, I'm not sure how feasible that would be. I recall a while ago discussions on if Legate could have a means to query how much remaining available memory there is on the GPU, but I'm not sure where they went. (note: this would be also helpful for choosing between other cases of fast, memory-hungry implementations vs. slower implementations that don't use as much memory such as "legacy" matmul vs. batched matmul)

If the above is not possible or difficult to implement, we could just create an option to allow the user to decide. Because we've been adding more flags that modify the performance vs. correctness/completion, it could be useful to have an overall mode configuration in cupynumeric CUPYNUMERIC_PERFORMANCE_MODE="(fast|conservative|debug|...)" that automatically sets the related cupynumeric settings based on how we define the modes to make it easier for the user to switch between various execution patterns/algorithm selection.

[ipdemes]

@dinodeep should we close this Issue? I think your implementation shows better performance

[ipdemes]

@dinodeep can we close this?

[dinodeep]

I'll let @manopapad decide if we should close this or not. It is true that the scan + masked getitem implementation is faster; however, there is still a significant memory benefit to having a thrust::copy_if implementation (despite it being slower), so I'm not sure where we should hold onto that information. It would be ideal to note at least somewhere that there is a less memory-hungry algorithm for AdvancedIndexing, but not sure if leaving this issue open is necessary for that

[manopapad]

I think we can keep this open, as an optional alternative to the current implementation, but we don't need to prioritize it.

[ipdemes]

I created a draft PR https://github.com/nv-legate/cupynumeric.internal/pull/2171. You can see performance comparison numbers there as well