Pi Pico - Calculate first 1,000,000 primes in Assembly

pico_sieve.py

# pico_sieve.py
# Raspberry Pi Pico  - Sieve of Eratosthenes demo
# Calculate all the prime numbers within a range of low integers.
# Based on:
#https://courses.ideate.cmu.edu/16-223/f2022/text/code/pico-numerical.html

import time,array
import math,gc
from usys import exit

machine.freq(270_000_000)
print('CPU Hz',machine.freq())

size = 1_000_000
stop = int(math.sqrt(size))
flags32 = array.array('I',0 for _ in range(32000)) # 1_000_000 / 32
flags = bytearray(1) #size

# Walk up through the array, identifying all prime numbers and removing their
# multiples from the set of integers by setting the corresponding flag to False.
# Note that the scan can stop once values are reached for which no unique
# multiples are included in the array.

@micropython.viper
def primes(stop:int,size:int):
    f = ptr8(flags)
    for i in range(2, stop): #int(math.sqrt(size))):
        if f[i] == 0:
            # This value is a prime, now remove all the multiples from the set.
            # Note that this marking may begin at i*i since all lower multiples will
            # have already been removed.
            multiple = i * i
            while multiple < size:
                f[multiple] = 1
                multiple += i
              

@micropython.viper
def primes32(stop:int,size:int):
    f = ptr32(flags32)
    for i in range(2, stop):
        word = i>>5
        bit  = i-((i>>5)<<5)
        if (f[word] >> bit) & 1 == 0:
            # This value is a prime, now remove all the multiples from the set.
            # Note that this marking may begin at i*i since all lower multiples will
            # have already been removed.
            multiple = i * i
            while multiple < size:
                word = multiple >> 5
                bit  = multiple - ((multiple>>5)<<5)
                f[word] = f[word] | (1<<bit)
                multiple += i
                
@micropython.asm_thumb
def primes32_asm(r0,r1,r2):  # stop, size, flags[]
    mov(r3,r2)
    mov(r4,r0)
    mov(r5,0)
    label(CLEAR_LOOP)
    str(r5,[r3,0])        # zero out array
    add(r3,4)
    sub(r4,1)
    bne(CLEAR_LOOP)
    mov(r3,2)             # r3  = i = 2
    label(FOR_LOOP)
    asr(r4,r3,5)          # r4 = word
    lsl(r5,r4,5)
    sub(r5,r3,r5)         # r5 = bit
    lsl(r4,r4,2)
    add(r6,r2,r4)         # f[word]
    ldr(r6,[r6,0])
    asr(r6,r5)            # f[word] >> bit
    mov(r7,1)
    and_(r6,r7)
    cmp(r6,0)
    bne(SKIP)
    mov(r7,r3)
    mul(r7,r7)            # multiple = i * i
    label(WHILE_LOOP)
    asr(r4,r7,5)          # word = multiple >> 5
    lsl(r5,r4,5)
    sub(r5,r7,r5)         # r5 = bit = multiple - ((multiple>>5)<<5)
    mov(r6,1)   
    lsl(r6,r5)            # 1<<bit
    lsl(r4,r4,2)          #
    add(r5,r2,r4)         # f[word]
    ldr(r4,[r5,0])
    orr(r4,r6)
    str(r4,[r5,0])
    add(r7,r7,r3)         # multiple += i   
    cmp(r7,r1)            # multiple < size
    blt(WHILE_LOOP)    
    label(SKIP)
    add(r3,1)
    cmp(r3,r0)            # i < stop
    blt(FOR_LOOP)
    label(EXIT)
    

@micropython.asm_thumb
def primes_asm(r0,r1,r2): # stop, size, flags[]
    mov(r7,1)
    mov(r3,2)             # r3  = i = 2
    label(LOOP)
    add(r6,r2,r3)         #       
    ldrb(r4,[r6,0])       # r4  = f[i]
    cmp(r4,0)
    bne(SKIP)
    mov(r5,r3)            # r5 = multiple
    mul(r5,r5)            #      multiple = i * i
    label(LOOP2)
    cmp(r5,r1)            # while multiple < size
    bge(SKIP)
    add(r6,r2,r5)
    strb(r7,[r6,0])       # f[multiple] = 1
    add(r5,r5,r3)         # multiple += i
    b(LOOP2)    
    label(SKIP)
    add(r3,1)             # i += 1
    cmp(r3,r0)            # i < stop
    blt(LOOP)
    mov(r0,r3)
    
def print32(size):
    for i in range(2,size):
        word = i >> 5
        bit = i-((i>>5)<<5)
        if (flags32[word] >> bit) & 1 == 0:
            print(i,word,bit)
    
def print_bytes(size):
    for i in range(2, size):
        if flags[i] == 0:
            print(i)

def five_sec(): # number of loops to 1,000,000 in 5 seconds
    i=0
    start = time.ticks_ms()
    while time.ticks_diff(time.ticks_ms(),start) < 5000: 
        primes32_asm(stop,size,flags32)
        i+=1
    return i

start = time.ticks_us()
primes32_asm(stop,size,flags32)
end = time.ticks_us()
print("First 1,000,000 primes: ", (end - start)/1_000_000, "seconds.")
print("In five seconds,", five_sec(), "loops")


#print32(30)
#print_bytes(30)

gc.collect()
print('Memory Free:',gc.mem_free())

output:
CPU Hz 270000000
First 1,000,000 primes: 0.12587 seconds.
In five seconds, 40 loops
Memory Free: 57312