[Java] Add Dataset based on MemorySegment

[ldematte]

This PR adds the ability to define a Dataset directly over a MemorySegment, "wrapping" it instead of allocating a new one.

• Depends on [Java] Add Java API benchmarks #1033 and [Java] Encapsulate on-heap float arrays into Dataset #1024
• The new API has a Object memorySegment parameter, as we target Java 21 for the API (but 22 for the implementation); it works but it's definitely a hack and we need to sort this out
  • As discussed, we want to keep targeting Java 21 for the API. This means the API will return a MethodHandle, and the Java 22 implementation will use it to return a factory method to build a Dataset from a MemorySegment.
  • This factory method can then be used as shown in the tests (see the DatasetHelper convenience class/method).
• Benchmarks show a sizeable speedup -- it is still tiny related to the "big picture" (index build time), but there is an improvement and above all we avoid a whole new copy of the input data (halving the memory requirements).

Fixes #698

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[ldematte]

I've refactored the code a bit to better handle MemorySegment allocation/deallocation; now I am able to benchmark Dataset creation separately:

Benchmark                                          (dims)  (size)  Mode  Cnt      Score       Error  Units
CagraIndexBenchmarks.testDatasetFromHeap             1024     100  avgt    5  45798.418 ± 12296.887  ns/op
CagraIndexBenchmarks.testDatasetFromMemorySegment    1024     100  avgt    5      2.736 ±     0.016  ns/op

Of course this is a tiny part of building an index; even when I use a small index (size 100, dims 1024) the time to create the index still dominates

Benchmark                                           (dims)  (size)   Mode  Cnt    Score   Error  Units
CagraIndexBenchmarks.testIndexingFromHeap             1024     100  thrpt    5  123.619 ± 1.854  ops/s
CagraIndexBenchmarks.testIndexingFromMemorySegment    1024     100  thrpt    5  184.753 ± 2.229  ops/s

But there still is a benefit, and I have the poorest of the GPUs on my dev machine (it's a 2070 mobile).

[ldematte]

When I bump up data size (-p size=N), of course the margin becomes smaller and smaller as the index build time dominates; still, if the data comes from off-heap memory (e.g. memory mapped file), we potentially use half the memory for the input data (or 1/3 in some cases), which I think is the key point.

[ChrisHegarty]

Thanks for doing this @ldematte, I think that this is case for revisiting the Java 21/22 decision that we previously made. I'll take a closer look and give it some thought.

[mythrocks]

Hmm. There seem to be java test failures here:

Error:  Failures: 
Error:    CagraBuildAndSearchIT.testIndexingAndSearchingFlow:145->queryAndCompare:464->CuVSTestCase.checkResults:148 expected:<[{0=0.83774555, 2=0.3590463, 3=0.038782578}, {0=0.12472608, 1=0.31918612, 2=0.21700792}, {0=0.48305473, 2=0.20332818, 3=0.047766715}, {0=0.59063464, 1=0.15224178, 3=0.5986642}]> but was:<[{0=0.23431994, 2=0.23431946, 3=0.23431946}, {0=0.68692017, 2=0.6869197, 3=0.6869197}, {0=0.33261782, 2=0.3326173, 3=0.3326173}, {0=0.3379389, 2=0.33793885, 3=0.33793885}]>
Error:  Errors: 
Error:    CagraBuildAndSearchIT.testIndexing:233->runConcurrently:85 ? Execution java.lang.AssertionError: Exception while executing runnable: java.lang.RuntimeException: cuvsCagraBuild returned 0[RAFT failure at file=/tmp/conda-bld-output/bld/rattler-build_libcuvs/work/cpp/src/neighbors/detail/cagra/graph_core.cuh line=1406: Could not generate an intermediate CAGRA graph because the initial kNN graph contains too many invalid or duplicated neighbor nodes. This error can occur, for example, if too many overflows occur during the norm computation between the dataset vectors.
Obtained 6 stack frames
#1 in /opt/conda/envs/java/lib/jvm/lib/server/../../../libcuvs.so(+0x49117d) [0x7c03a02ec17d]
#2 in /opt/conda/envs/java/lib/jvm/lib/server/../../../libcuvs.so: void cuvs::neighbors::cagra::detail::graph::optimize<unsigned int, raft::host_device_accessor<std::experimental::default_accessor<unsigned int>, (raft::memory_type)0> >(raft::resources const&, std::experimental::mdspan<unsigned int, std::experimental::extents<long, 18446744073709551615ul, 18446744073709551615ul>, std::experimental::layout_right, raft::host_device_accessor<std::experimental::default_accessor<unsigned int>, (raft::memory_type)0> >, std::experimental::mdspan<unsigned int, std::experimental::extents<long, 18446744073709551615ul, 18446744073709551615ul>, std::experimental::layout_right, raft::host_device_accessor<std::experimental::default_accessor<unsigned int>, (raft::memory_type)0> >, bool, bool) +0x1aa3 [0x7c03a0d3b4d3]
#3 in /opt/conda/envs/java/lib/jvm/lib/server/../../../libcuvs.so: cuvs::neighbors::cagra::index<float, unsigned int> cuvs::neighbors::cagra::detail::build<float, unsigned int, raft::host_device_accessor<std::experimental::default_accessor<float const>, (raft::memory_type)2> >(raft::resources const&, cuvs::neighbors::cagra::index_params const&, std::experimental::mdspan<float const, std::experimental::extents<long, 18446744073709551615ul, 18446744073709551615ul>, std::experimental::layout_right, raft::host_device_accessor<std::experimental::default_accessor<float const>, (raft::memory_type)2> >) +0xd81 [0x7c03a0d629b1]
#4 in /opt/conda/envs/java/lib/jvm/lib/server/../../../libcuvs.so: cuvs::neighbors::cagra::build(raft::resources const&, cuvs::neighbors::cagra::index_params const&, std::experimental::mdspan<float const, std::experimental::extents<long, 18446744073709551615ul, 18446744073709551615ul>, std::experimental::layout_right, raft::host_device_accessor<std::experimental::default_accessor<float const>, (raft::memory_type)2> >) +0x21 [0x7c03a0d1d761]
#5 in /opt/conda/envs/java/lib/jvm/lib/server/../../../libcuvs_c.so: cuvsCagraBuild +0x806 [0x7c04352dbb96]
#6 in [0x7c045c0a6ec7]

[ldematte]

Hmm. There seem to be java test failures here:

Yeah, it's a chicken and egg problem :) I fixed those in another PR (#1089) but that PR depends on this to be merged first. These tests randomly fail because they are using resources and threading wrong.
I will cherry pick the test fixes from #1089 so we can unblock the situation.

[ldematte]

(BTW, the issue is already there, it's just random that it happened to fail in this PR -- better fix it sooner than later)
@mythrocks, I cherry picked the test fix, it is good to go again.

[mythrocks]

/ok to test 8582abb

[ldematte]

@mythrocks is there anything else I can do to unblock this PR?

[mythrocks]

/ok to test 63c4932

[mythrocks]

The CI builds should be fixed once #1114 is resolved.

[benfred]

/ok to test 2e65a55

[mythrocks]

/merge