mbutrovich opened a new pull request, #12058: URL: https://github.com/apache/arrow-rs/pull/12058
# Which issue does this PR close? N/A. # Rationale for this change Checked `add`, `sub` and `mul` on integers [call `try_binary` or `try_unary`](https://github.com/apache/arrow-rs/blob/b231cf15efa5dbbbe4ae8503849aa6ab3cfc4e0c/arrow-arith/src/numeric.rs#L337-L343), which compute one row at a time and return at the first overflow. That early return keeps LLVM from vectorizing the loop, while the wrapping kernels vectorize. DataFusion calls the checked kernels when `fail_on_overflow` is set ([`binary.rs`](https://github.com/apache/datafusion/blob/32bd81fb226ccf1ae5df67ff340d77baa3e73668/datafusion/physical-expr/src/expressions/binary.rs#L641-L651)), and Comet calls them for Spark's ANSI mode ([`checked_arithmetic.rs`](https://github.com/apache/datafusion-comet/blob/79b5ddd50ba2e3f605399d641e4c0b3be84c20ef/native/spark-expr/src/math_funcs/checked_arithmetic.rs#L146-L148)). # What changes are included in this PR? - Integer `add`, `sub` and `mul` compute the wrapping result and an overflow flag for every slot in one pass, without branching. If no slot overflowed, that is the result. The pass is an `inline(never)` function, because inlined into `integer_op` under thin LTO it keeps the flag in memory and doesn't vectorize. - If any slot overflowed, they run the existing checked kernel, so errors and their messages don't change. Null slots can hold any value, so a null slot that overflows also takes this path and still returns `Ok`. - Signed `add` and `sub` test sign bits, because the signed `overflowing_add` and `overflowing_sub` lower to `llvm.sadd.with.overflow` and `llvm.ssub.with.overflow`, which LLVM doesn't vectorize. Signed `mul` of 8 to 32 bits multiplies in the next wider type. `mul` of `i64` and of unsigned types uses `overflowing_mul`, which doesn't branch but stays scalar. - Null slots of the result hold the wrapping result instead of zero, as they already do for `add_wrapping`. - af931c5b0 adds the benchmark cases before the change, so the benchmark bot can compare against it. # Are these changes tested? `test_checked_integer_edges` in `arrow-arith/src/numeric.rs` compares `add`, `sub` and `mul` with the scalar checked operations for every pair of edge values of each integer type, with arrays on both sides and with a scalar on either side, and with overflowing values in null slots. It also checks each pair as a one-row slice, because an overflow in one row sends the whole array to the checked kernel and would hide a missed overflow in another row. `test_checked_integer_overflow_position` puts the overflow at the end of arrays whose lengths fall on and around the vector loop boundaries. `test_checked_integer_slice_null_scalar_and_lengths` covers sliced input, a null scalar, empty input and mismatched lengths. These tests pass before and after the change. TODO: chart from the bot runs. TODO: one sentence with the bot's speedups for checked `add`, `sub` and `mul`, and the cost of the `_add_checked_overflow` and `_add_checked_null_slot_overflow` cases. The `_add_checked_overflow` and `_add_checked_null_slot_overflow` cases pay for the new pass and then the checked kernel. DataFusion with `fail_on_overflow` and Comet in ANSI mode turn an overflow into an error, so the first case runs at most once per partition before the query fails. The second case repeats on every batch whose null slots hold values that overflow, and doesn't error. <details><summary>Numbers</summary> TODO: speedups over af931c5b0 for every `int32_` and `int64_` case on Neoverse-V2 and Xeon, default and aligned, with links to the bot comments. The wrapping cases don't call the changed code and show each run's noise. Speedups over af931c5b0 on Apple M5 Max with thin LTO, at 8192 rows. The aligned run adds `-C llvm-args=-align-all-functions=7 -C llvm-args=-align-all-nofallthru-blocks=7`. | Case | `Int32` | aligned | `Int64` | aligned | |---|---|---|---|---| | `add_checked(0)` | 3.76x | 3.79x | 2.15x | 2.14x | | `add_checked(0.1)` | 5.54x | 5.58x | 2.90x | 2.94x | | `add_checked_scalar(0)` | 4.18x | 4.20x | 2.60x | 2.61x | | `add_checked_scalar(0.1)` | 5.90x | 5.96x | 3.37x | 3.46x | | `subtract_checked(0)` | 3.80x | 3.76x | 2.18x | 2.15x | | `subtract_checked(0.1)` | 5.43x | 5.57x | 2.89x | 2.91x | | `multiply_checked(0)` | 2.53x | 2.53x | 0.96x | 0.97x | | `multiply_checked(0.1)` | 4.22x | 4.15x | 1.53x | 1.50x | | `add_checked_overflow` | 0.79x | 0.78x | 0.74x | 0.74x | | `add_checked_null_slot_overflow` | 0.86x | 0.87x | 0.75x | 0.77x | | wrapping cases (unchanged) | 0.96x to 1.03x | 0.98x to 1.06x | 0.97x to 1.01x | 0.98x to 1.00x | `Int64` multiply stays scalar, and without nulls it costs about the same as before. The two overflow cases pay for the new pass and then the checked kernel. </details> <details><summary>Codegen</summary> From the thin-LTO `arithmetic_kernels` binary on aarch64, the checked `Int32` add loop handles 32 rows per iteration. One of its eight lanes of 4 rows: ```asm add.4s v29, v18, v2 ; wrapping sum of 4 rows eor.16b v2, v29, v2 ; sum ^ l eor.16b v18, v29, v18 ; sum ^ r and.16b v2, v2, v18 cmlt.4s v2, v2, #0 ; sign bit set: overflow ... orr.16b v0, v0, v2 ; accumulate the flag ``` Before this change, the same loop did one `adds` and `b.vs` per row. The thin-LTO binary has vector loops for every integer type and operation except `mul` of `Int64` and `UInt64`. Before linking, the same functions have vector code for `neoverse-n1` and for `x86-64-v3` on `x86_64-unknown-linux-gnu`. On baseline `x86-64`, `mul` of `Int32` also stays scalar, because SSE2 has no signed 32-bit to 64-bit multiply. </details> # Are there any significant user-facing changes? No. ## AI Disclosure I wrote this PR with help from Claude but understand and support the core approach, implementation, and test coverage. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
