Incrementing an index based of a condition

Zama Concrete
1

Hi,

I am pretty new to FHE and so my question may be naive. I am trying to convert the quick sort algorithm towards FHE equivalent using the concrete compiler. I’m struggling on this bit of code i want to adapt :

for j in range(low, high):
                if arr[j] <= pivot:
                    i += 1
                    arr[i], arr[j] = arr[j], arr[i]

I’ve adapted the if so it has no branching to be able to swap, but can’t find a way to do it for the incrementation of i. Below is a minimal code that reproduces the error I can’t figure out completely

@fhe.compiler({"array": "encrypted"})
def function(array):
    i = 0
    for j in range(0, len(array)):
        cmp = array[j] <= array[len(array) - 1]
        #i+=cmp 
        # or
        #i = (i+1) * cmp + i * (1 - cmp)
        tmp = array[i]
        array[i] = array[j] * cmp + array[i] * (1 - cmp)
        array[j] = array[j] * (1 - cmp) + tmp * cmp
    return array

This does compile, but if i do uncomment on of the i incrementation methods i get (end of traceback):

File “/home/nthsmn/Documents/tests/quicksort_circuit_chunked.py”, line 18, in function
tmp = array[i]
File “/home/nthsmn/.local/lib/python3.10/site-packages/concrete/fhe/tracing/tracer.py”, line 800, in getitem
raise ValueError(message)
ValueError: Tracer<output=EncryptedTensor<uint4, shape=(20,)>> cannot be indexed with Tracer<output=EncryptedScalar>

Does anyone have a clever way of fixing this ?

Thanks a lot !

2

The i value is a clear constant → not an encrypted value. During compilation the loops will be unrolled and all the values that i, j take will become clear constants for the compiler. You’ll get something like:

i = 0
j = 0
swap(array[j] <= array[len(array) - 1], array[j], array[j])
i = 0
j = 1
swap(array[j] <= array[len(array) - 1], array[j], array[j])
i = 0
j = 2
swap(array[j] <= array[len(array) - 1], array[j], array[j])
..etc..

You can not add an encrypted value to i since it is a constant. Doing the addition would turn i into an encrypted variable. That would be doable in FHE, but even if the compiler supported it (it does not for now), the indexing of an encrypted array with an encrypted value is hugely expensive in FHE (you are currently indexing with constants which is fast).

For sorting algorithms I suggest you look at Sorting networks. Sorting networks perform sorting without branching. A top-k selection network - a modified sorting network - is implemented in the KNN model in Concrete ML: https://github.com/zama-ai/concrete-ml/blob/0229cf7541394d4f2b38eb915961f808a902cd80/src/concrete/ml/sklearn/base.py#L1958

You could simply remove the comparisons variable and the test on it and it will turn this function into a sorting network.