max/lean4-htt
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 2

feat: optimize lean_nat_shiftr for scalars (#8268)

Browse source 
This PR optimizes lean_nat_shiftr for scalar operands. The new compiler
converts Nat divisions into right shifts, so this now shows up as hot in
some profiles.
This commit is contained in:
Cameron Zwarich 2025-05-10 18:39:59 -07:00 • committed by GitHub
parent ddf5512c9a
commit 575b4786f9
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 200 additions and 2 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

13
src/include/lean/lean.h
View file

@ -1424,11 +1424,22 @@ static inline lean_obj_res lean_nat_lxor(b_lean_obj_arg a1, b_lean_obj_arg a2) {
}
LEAN_EXPORT lean_obj_res lean_nat_shiftl(b_lean_obj_arg a1, b_lean_obj_arg a2);
LEAN_EXPORT lean_obj_res lean_nat_shiftr(b_lean_obj_arg a1, b_lean_obj_arg a2);
LEAN_EXPORT lean_obj_res lean_nat_big_shiftr(b_lean_obj_arg a1, b_lean_obj_arg a2);
LEAN_EXPORT lean_obj_res lean_nat_pow(b_lean_obj_arg a1, b_lean_obj_arg a2);
LEAN_EXPORT lean_obj_res lean_nat_gcd(b_lean_obj_arg a1, b_lean_obj_arg a2);
LEAN_EXPORT lean_obj_res lean_nat_log2(b_lean_obj_arg a);
static inline lean_obj_res lean_nat_shiftr(b_lean_obj_arg a1, b_lean_obj_arg a2) {
if (LEAN_LIKELY(lean_is_scalar(a1) && lean_is_scalar(a2))) {
size_t s1 = lean_unbox(a1);
size_t s2 = lean_unbox(a2);
size_t r = (s2 < sizeof(size_t)*8) ? s1 >> s2 : 0;
return lean_box(r);
} else {
return lean_nat_big_shiftr(a1, a2);
}
}
/* Integers */
#define LEAN_MAX_SMALL_INT (sizeof(void*) == 8 ? INT_MAX : (INT_MAX >> 1))

2
src/runtime/object.cpp
View file

@ -1462,7 +1462,7 @@ extern "C" LEAN_EXPORT lean_obj_res lean_nat_shiftl(b_lean_obj_arg a1, b_lean_ob
return mpz_to_nat(r);
}
extern "C" LEAN_EXPORT lean_obj_res lean_nat_shiftr(b_lean_obj_arg a1, b_lean_obj_arg a2) {
extern "C" LEAN_EXPORT lean_obj_res lean_nat_big_shiftr(b_lean_obj_arg a1, b_lean_obj_arg a2) {
if (!lean_is_scalar(a2)) {
return lean_box(0); // This large of an exponent must be 0.
}

19
tests/compiler/nat_shiftr.lean Normal file
View file

@ -0,0 +1,19 @@
def test (a : Nat) : IO Unit :=
for b in #[0, 1, 14, 15, 16, 17, 31, 32, 33, 63, 64, 65] do
IO.println f!"{a >>> b}"
def main : IO Unit := do
test 0
test 1
test 0xff
test 0x100
test 0x101
test 0xffff
test 0x1000_0
test 0x1000_1
test 0xffff_ffff
test 0x1_0000_0000
test 0x1_0000_0001
test 0xffff_ffff_ffff_ffff
test 0x1_0000_0000_0000_0000
test 0x1_0000_0000_0000_0001

168
tests/compiler/nat_shiftr.lean.expected.out Normal file
View file

@ -0,0 +1,168 @@
0
0
0
0
0
0
0
0
0
0
0
0
1
0
0
0
0
0
0
0
0
0
0
0
255
127
0
0
0
0
0
0
0
0
0
0
256
128
0
0
0
0
0
0
0
0
0
0
257
128
0
0
0
0
0
0
0
0
0
0
65535
32767
3
1
0
0
0
0
0
0
0
0
65536
32768
4
2
1
0
0
0
0
0
0
0
65537
32768
4
2
1
0
0
0
0
0
0
0
4294967295
2147483647
262143
131071
65535
32767
1
0
0
0
0
0
4294967296
2147483648
262144
131072
65536
32768
2
1
0
0
0
0
4294967297
2147483648
262144
131072
65536
32768
2
1
0
0
0
0
18446744073709551615
9223372036854775807
1125899906842623
562949953421311
281474976710655
140737488355327
8589934591
4294967295
2147483647
1
0
0
18446744073709551616
9223372036854775808
1125899906842624
562949953421312
281474976710656
140737488355328
8589934592
4294967296
2147483648
2
1
0
18446744073709551617
9223372036854775808
1125899906842624
562949953421312
281474976710656
140737488355328
8589934592
4294967296
2147483648
2
1
0