x86-64 assembly for computing the length of a number's Collatz sequence

collatz-faster.s

# We implement the space-time tradeoff described by Scollo [0, Section IV.D], and use the notation
# given in [1].  The method allows us to compute k iterations of a function f using precomputed
# lookup tables of size O(2^k).  Instead of using HOTPO directly, we define f such that
# - f(n) = (3n+1)/2, if n is odd, and
# - f(n) = n/2,      if n is even.
# Hence, in the case of an odd n, f performs two HOTPO operations.  Let f^k(n) denote the result of
# applying f k-times to n, i.e. f^k(n) = f(f(...f(n)...)).

# As precomputation, we prepare two arrays c and d of size 2^k each.  For i = 0, ..., 2^k - 1,
# - d[i] = result of applying f^k to i,
# - c[i] = number of odd integers occurring during the computation of d[i].

# Using these arrays, we can compute f^k(n) as follows: First we split n into two parts as n = 2^k *
# a + b such that b consists of the lower k bits of n.  Then we have f^k(n) = 3^c[b] * a + d[b].
# Since f performs two HOTPO operations for odd integers, the number of HOTPO operations is k +
# c[b].

# In the implementation, we use the method above with k = 17 to repeatedly apply f^k to the input
# until the result is < 2^k.  Then, we use another lookup table to get the remaining stopping time.
# Note that, since every number >= 2^k needs at least k divisions by two to reach 1, we cannot
# accidentally run "past 1" by iterating f^k.

# [0]: https://www.dmi.unict.it/~scollo/seminars/gridpa2007/CR3x+1paper.pdf
# [1]: https://en.wikipedia.org/wiki/Collatz_conjecture#Time%E2%80%93space_tradeoff


.intel_syntax noprefix                  # Because it is more beautiful.
.section .note.GNU-stack,"",@progbits   # We do not need nor want an executable stack.


.section .rodata                        # The lookup tables are read-only data,
                                        # so they belong in .rodata.

# Lookup table (1) containing the stopping times for all number < 2^k.
# Each entry is a 16 bit value.
.type lookup_tab, @object
lookup_tab:
.incbin "collatz-fast-lookup.bin"

# Lookup table (2) of size 2^k for computing k iterations of f.
# Each entry is a triple packed into a 64 bit word:
# - 27 bits  [0..26]: 3^c[b] (so that we do not need to compute the power at runtime)
# - 27 bits [32..58]: d[b]
# -  5 bits [59..63]: c[b]
.type acc_tab0, @object
acc_tab:
.incbin "collatz-fast-acc-tab-combined.bin"


.text                                   # Section for code


# Function computing the stopping time of a number
.global collatz
.type collatz, @function
collatz:

xor rcx, rcx                            # initialize count with 0
lea r11, [rip + acc_tab]                # load address of lookup table (2)
                                        # - using rip-relative addressing for position-independent
                                        #    code; we like PIE :)
                                        # - compute the address once, and not in the loop
jmp test                                # go to loop condition

# loop: compute f^k and keep track of the number of HOTPO operations
loop_start:

# - the current number n is in rdi
# - split it as n = 2^k * a + b,
# - and use table (2) to compute k iterations of f applied to n
mov r8, rdi
and r8, 0x1ffff                         # compute b := n mod 2^k
mov rsi, QWORD PTR [r8 * 8 + r11]       # load entry of table (2) at index b

mov eax, esi                            # extract 3^c[b] from table lookup
                                        # - upper 32 bits of rax get cleared by this
shr rdi, 17                             # compute a := (n - b) / 2^k
mul rdi                                 # compute 3^c[b] * a
                                        # - mul rdi computes rdx:rax <- rdi * rax
                                        # - we assume the result fits into 64 bits
                                        # - mul has the highest latency of the instructions used
                                        #   here (3 on many modern CPUs)
                                        # -> compute it as early as possible in the loop

mov rdi, rsi                            # extract d[b] from table lookup
shr rdi, 32
and edi, 0x07ffffff                     #   -> d[b]

add rcx, 17                             # increment count by k

mov r9, rsi                             # extract c[b] from table lookup
shr rsi, 59                             #   -> c[b]
add rcx, rsi                            # increment count by c[b]

add rdi, rax                            # compute new n := 3^c[b] * a + d[b]

test:                                   # loop while the value is >= 2^k
test rdi, 0xfffffffffffe0000
jnz loop_start


lookup:                                 # the value in rdi is now < 2^k, so we can lookup the
                                        #   (remaining) stopping time from table (2)
lea rax, [rip + lookup_tab]             # compute address of lookup table (1)
movzx rax, WORD PTR [rdi * 2 + rax]     # load stopping time from table
add rax, rcx                            # add the computed stopping time from the first part
ret                                     # done

collatz-small.s

.intel_syntax noprefix                  # because it is more beautiful
.section .note.GNU-stack,"",@progbits   # we do not need nor want an executable stack

.text

.global collatz
collatz:
movabs rax, -1                # initialize counter with -1 (we increment at the beginning of the loop)
                              # - store in rax, since it's also our return value
loop:
add rax, 1                    # increment counter
lea rsi, [2 * rdi + rdi + 1]  # odd case:  n |-> 3n + 1   (1)
shr rdi, 1                    # even case: n |-> n/2      (2)
                              # - if ZF set, then n was 1, so we are done and should return
                              # - if CF set, then n was odd, so we should use (1)
cmovc rdi, rsi                # if CF set, copy the result of (1)
jnz loop                      # if ZF not set, go to the beginning of the loop
ret                           # return

generate-lookup-tables.py

#!/usr/bin/env python3

import struct

def hotpo(n):
    if n & 1:
        return 3 * n + 1
    else:
        return n // 2


def delay(n):
    count = 0
    while n != 1:
        n = hotpo(n)
        count += 1
    return count


def f(n):
    if n & 1:
        return (3 * n + 1) // 2
    else:
        return n // 2


def f_k(n, k):
    odds = 0
    for _ in range(k):
        if n & 1:
            odds += 1
        n = f(n)
    return n, odds


def main():
    k = 17

    lookup_tab = [0] + [delay(n) for n in range(1, 1 << k)]
    elem_size_lookup_tab = max(lookup_tab).bit_length()
    assert elem_size_lookup_tab <= 16

    acc_tab0, acc_tab1 = zip(*(f_k(n, k) for n in range(1 << k)))
    elem_size_acc_tab0 = max(acc_tab0).bit_length()
    elem_size_acc_tab1 = max(acc_tab1).bit_length()
    assert elem_size_acc_tab0 == 27, elem_size_acc_tab0
    assert elem_size_acc_tab1 == 5, elem_size_acc_tab1

    acc_tab1a = [3 ** n for n in acc_tab1]
    elem_size_acc_tab1a = max(acc_tab1a).bit_length()
    assert elem_size_acc_tab1a == 27, elem_size_acc_tab1a

    acc_tab_combined = [
            (y << 59) | (x << 32) | z
            for x, y, z in zip(acc_tab0, acc_tab1, acc_tab1a)
        ]
    elem_size_acc_tab_combined = max(acc_tab_combined).bit_length()
    assert elem_size_acc_tab_combined <= 64

    with open('collatz-fast-lookup.bin', 'wb') as f:
        f.write(struct.pack(f'<{1 << k}H', *lookup_tab))
    with open('collatz-fast-acc-tab-combined.bin', 'wb') as f:
        f.write(struct.pack(f'<{1 << k}Q', *acc_tab_combined))


if __name__ == '__main__':
    main()

LICENSE

MIT License

Copyright (c) 2021 Lennart Braun

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.