In pytorch, how can I parallelize of a set of boolean functions that are executed (on a GPU) repeatedly?

I doubt that this is the best/fastest solution, but using torch.compile does provide acceleration. I haven't yet tested when scaling up to thousands of Boolean functions.

import torch
from time import time

A = torch.tensor([True, False, True]).to('cuda')  # Initial values.
B = torch.tensor([False, True, True]).to('cuda')  # Values don't matter. Will write over them in the first iteration

@torch.compile
def process(A, B, n_steps):
    for step in range(n_steps):

        # Use values in A to compute new values in B.
        # How to run the three lines below IN PARALLEL?
        B[0] = torch.logical_and(torch.logical_or( A[0], A[1]), A[2])  # func1: Y0 = X0 | X1 & X2
        B[1] = torch.logical_or( torch.logical_or( A[0], A[1]), A[2])  # func2: Y1 = X0 | X1 | X2
        B[2] = torch.logical_and(torch.logical_and(A[0], A[1]), A[2])  # func3: Y2 = X0 & X1 & X2
        # Only after the above three lines above finish their computation (and B has new values), should the lines below be run.

        # Use values in B to compute new values in A.
        # Note that the functions below are identical to the ones above (which may allow for some additional acceleration?)
        # How to run the three lines below IN PARALLEL?
        A[0] = torch.logical_and(torch.logical_or( B[0], B[1]), B[2])  # func1: Y0 = X0 | X1 & X2
        A[1] = torch.logical_or( torch.logical_or( B[0], B[1]), B[2])  # func2: Y1 = X0 | X1 | X2
        A[2] = torch.logical_and(torch.logical_and(B[0], B[1]), B[2])  # func3: Y2 = X0 & X1 & X2
        # Only after the above three lines above finish their computation (and A has new values), should the next loop be run.

    return A

# First run is slow due to compilation
t_start = time()
A = process(A, B, 100)
print(f'{time()-t_start} seconds')

# Runs are faster subsequently, and can be looped over to effectively increase n_steps
t_start = time()
A = process(A, B, 100)
print(f'{time()-t_start} seconds')

Output without @torch.compile decorator:

First run time: 0.12589144706726074 seconds
Second run time: 0.059000492095947266 seconds

Output with @torch.compile decorator:

First run time: 18.201257467269897 seconds
Second run time: 0.007639169692993164 seconds