Fix gather/scatter fast division bug by huuanhhuyn · Pull Request #3059 · NVIDIA/raft

huuanhhuyn · 2026-06-24T08:44:01Z

raft::matrix::gather/scatter crashes at int32 boundary even when it advertises int64 type -> issue 3055. This PR fixes this issue.

int64 is still not supported for performance reason with int128.

closes #3055

copy-pr-bot · 2026-06-24T08:44:04Z

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.

tarang-jain

LGTM! Thanks for the fix.

tarang-jain · 2026-06-24T16:40:42Z

Actually, can we add a test for this diff? We could either add a test for FastIntDiv directly, or we can test the raft::matrix::gather when N * D exceeds INT_MAX. I am more inclined to have tests for both -- testing FastIntDiv should be quick where we just compare with the C++ operators ( / and %). But the more important one is to test the raft::matrix::gather for those larger matrices.

huuanhhuyn · 2026-06-26T08:24:47Z

            inplace_inputs_i64);
+GATHER_TEST((GatherTest<false, false, true, float, int64_t, int64_t>),
+            GatherInplaceTestCI64I64,
+            inplace_inputs_i64_i32_max);  // slow test, 8GB allocation, reproduces https://github.com/rapidsai/raft/issues/3055


@tfeher @tarang-jain Only float type is possible here (8GB allocation, 100ms). I would suggest to not have it. I already added a regression test for the fix in fast_int_div.cu

I think it is fine to spend 100ms on this.

tfeher

Thanks Huy for the fix, LGTM!

tfeher · 2026-06-26T09:53:23Z

            inplace_inputs_i64);
+GATHER_TEST((GatherTest<false, false, true, float, int64_t, int64_t>),
+            GatherInplaceTestCI64I64,
+            inplace_inputs_i64_i32_max);  // slow test, 8GB allocation, reproduces https://github.com/rapidsai/raft/issues/3055


I think it is fine to spend 100ms on this.

tfeher · 2026-06-26T15:22:37Z

/ok to test 9acb86f

aamijar · 2026-06-26T21:38:03Z

/ok to test 33ec5d7

aamijar · 2026-06-26T22:20:19Z

/ok to test b18e428

aamijar · 2026-06-27T03:54:45Z

/ok to test dc2ede3

- gather/scatter crashes at int32 boundary even when it advertises int64 type. This change fixes this issue. - int64 is still not supported for performance reason with int128.

huuanhhuyn · 2026-06-29T09:08:59Z

/ok to test 9f3d2cb

huuanhhuyn · 2026-06-29T09:14:28Z

/ok to test d87dcdd

huuanhhuyn · 2026-06-29T11:00:24Z

/merge

tfeher · 2026-06-29T11:12:10Z

/merge

tfeher

LGTM

tfeher · 2026-06-29T11:14:52Z

/merge

gforsyth · 2026-06-29T12:53:03Z

/merge

aamijar · 2026-06-29T15:28:32Z

Hello @huuanhhuyn, going forward let's remember to build cuvs and build cuml at the least to ensure that raft doesn't cause any breakages. This PR is causing some breakages currently.

Add explicit template parameters for FastIntDiv, follow up to NVIDIA/raft#3059. Authors: - Anupam (https://github.com/aamijar) Approvers: - Jim Crist-Harif (https://github.com/jcrist) URL: #8299

…around) The shuffle_dataset path used an out-of-place gather into a temporary buffer to work around an illegal memory access in raft's in-place gather overload when n_rows * row_len exceeded 2^31 (32-bit index overflow). That bug is now fixed upstream in raft (NVIDIA/raft#3059, closes #3055), which the cuvs raft pin now includes. Revert to the in-place gather to drop the extra full-size temporary allocation and copy.

@mfoerste4

#3093) This PR allows FastIntDiv to support the full 64-bit range values. Within the 32-bit range values, the fast implementation is selected. When exceeds 32-bit, it fallbacks during runtime to the normal division. The overhead of runtime check is measured to be negligible. FastIntDiv is used by in-place gather() where number of columns (divisor) is set once and the / and % divisions are called a lot more frequent to obtain row and col numbers from the given flat index. Therefore, the benchmark compares the fast int division against the normal division at 100 divisors and 1 million numerators for 32-bit and 64-bit range. ``` ---------------------------------------------------------------------------------------- Benchmark Time CPU Iterations ---------------------------------------------------------------------------------------- fast_int_div_i32/FastIntDiv/int32/manual_time 0.138 ms 0.893 ms 5086 native_div_i32/NativeIntDiv/int32/manual_time 0.186 ms 0.941 ms 3755 fast_int_div_i64/FastIntDiv/int64/manual_time 0.762 ms 1.52 ms 921 native_div_i64/NativeIntDiv/int64/manual_time 0.742 ms 1.50 ms 948 ``` @mfoerste4 Original version without the current changes and [this PR](#3059) -> Only negligible runtime overhead (~4%) is added. ``` ---------------------------------------------------------------------------------------- Benchmark Time CPU Iterations ---------------------------------------------------------------------------------------- fast_int_div_i32/FastIntDiv/int32/manual_time 0.131 ms 0.223 ms 5371 native_div_i32/NativeIntDiv/int32/manual_time 0.180 ms 0.273 ms 3884 ``` The [gather() bench](https://github.com/huuanhhuyn/raft/blob/fc9681c1821fc4a1559b8970a8e814f17569f84a/cpp/bench/prims/matrix/gather.cu) gives a similar comparison. The added runtime check is only a negligible overhead Authors: - https://github.com/huuanhhuyn Approvers: - Malte Förster (https://github.com/mfoerste4) - Divye Gala (https://github.com/divyegala) URL: #3093

huuanhhuyn requested a review from a team as a code owner June 24, 2026 08:44

github-project-automation Bot added this to Unstructured Data Processing Jun 24, 2026

huuanhhuyn mentioned this pull request Jun 24, 2026

[BUG] raft::matrix::gather (in-place overload): illegal memory access when N*D ≥ 2^31 (32-bit index overflow) #3055

Closed

tfeher added bug Something isn't working breaking Breaking change labels Jun 24, 2026

tarang-jain approved these changes Jun 24, 2026

View reviewed changes

huuanhhuyn requested a review from a team as a code owner June 26, 2026 08:23

huuanhhuyn commented Jun 26, 2026

View reviewed changes

tfeher approved these changes Jun 26, 2026

View reviewed changes

aamijar requested review from a team as code owners June 26, 2026 21:37

huuanhhuyn added 4 commits June 29, 2026 08:39

Fix gather/scatter fast div bug

618a22a

- gather/scatter crashes at int32 boundary even when it advertises int64 type. This change fixes this issue. - int64 is still not supported for performance reason with int128.

Add unit tests for FastIntDiv

c6f82a2

Remote heavy UT due to OOM

180ccef

precommit run

9f3d2cb

huuanhhuyn force-pushed the bug_3055 branch from dc2ede3 to 9f3d2cb Compare June 29, 2026 08:47

Remove gather test leftover

d87dcdd

tfeher approved these changes Jun 29, 2026

View reviewed changes

rapids-bot Bot merged commit 30cb3e0 into NVIDIA:main Jun 29, 2026
79 checks passed

github-project-automation Bot moved this to Done in Unstructured Data Processing Jun 29, 2026

This was referenced Jun 29, 2026

Fix FastIntDiv tparam NVIDIA/cuml#8299

Merged

Fix FastIntDiv tparam NVIDIA/cuvs#2278

Merged

huuanhhuyn deleted the bug_3055 branch June 30, 2026 07:24

huuanhhuyn mentioned this pull request Jun 30, 2026

Fix gather index type in ann_ivf_flat test NVIDIA/cuvs#2279

Closed

josephine-wolf-oberholtzer added this to Unstructured Data Processing Jul 1, 2026

josephine-wolf-oberholtzer moved this to Done in Unstructured Data Processing Jul 1, 2026

mfoerste4 mentioned this pull request Aug 4, 2026

FastIntDiv: fallback to normal division when values exceed 32-bit ran… #3093

Merged

Uh oh!

Conversation

huuanhhuyn commented Jun 24, 2026 • edited by tfeher Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

copy-pr-bot Bot commented Jun 24, 2026

Uh oh!

tarang-jain left a comment

Choose a reason for hiding this comment

Uh oh!

tarang-jain commented Jun 24, 2026

Uh oh!

huuanhhuyn Jun 26, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

tfeher Jun 26, 2026

Choose a reason for hiding this comment

Uh oh!

tfeher left a comment

Choose a reason for hiding this comment

Uh oh!

tfeher Jun 26, 2026

Choose a reason for hiding this comment

Uh oh!

tfeher commented Jun 26, 2026

Uh oh!

aamijar commented Jun 26, 2026

Uh oh!

aamijar commented Jun 26, 2026

Uh oh!

aamijar commented Jun 27, 2026

Uh oh!

huuanhhuyn commented Jun 29, 2026

Uh oh!

huuanhhuyn commented Jun 29, 2026

Uh oh!

huuanhhuyn commented Jun 29, 2026

Uh oh!

tfeher commented Jun 29, 2026

Uh oh!

tfeher left a comment

Choose a reason for hiding this comment

Uh oh!

tfeher commented Jun 29, 2026

Uh oh!

gforsyth commented Jun 29, 2026

Uh oh!

Uh oh!

aamijar commented Jun 29, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

huuanhhuyn commented Jun 24, 2026 •

edited by tfeher

Loading

huuanhhuyn Jun 26, 2026 •

edited

Loading