前の項目で解説した、ステージ（stage）とパス（pass）はシリアルな処理であり、データ並列化を行う設計とはなっていません。従って忠実にバイトニック整列を実装するのであれば以下のループを避けることはできません。

passes = 0
distance = 0
numberOfStages = np.uint32(np.log2(DATA_SIZE)) #(1)

for i in range(numberOfStages-1): #(2)
    for j in range(passes+1):
        distance = 1 << (i - j)
        stageDistance = 1 << i
        #OpenCLキューにコマンドを挿入

    passes+=1 //(3)

バイトニックマージについては、パスによってスワップする要素の距離を調整えます。スワップの距離は以下のように、ステージからパスを引いた値の累乗となります。

for j in range(passes):
    distance = 1 << (numberOfStages - 1 - j)

    #OpenCLキューにコマンドを挿入します。
    #distanceをカーネルのパラメータとする。

データの要素数が64とする場合はステージとパスについては以下の順序で推移していきます。

Stage = 0 Pass = 0
Stage = 1 Pass = 0
Stage = 1 Pass = 1
Stage = 2 Pass = 0
Stage = 2 Pass = 1
Stage = 2 Pass = 2
Stage = 3 Pass = 0
Stage = 3 Pass = 1
Stage = 3 Pass = 2
Stage = 3 Pass = 3
Stage = 4 Pass = 0
Stage = 4 Pass = 1
Stage = 4 Pass = 2
Stage = 4 Pass = 3
Stage = 4 Pass = 4

データ要素数が16の場合は、3つのステージ（log2 N - 1）を処理します。

Stage = 0 Pass = 0
Stage = 1 Pass = 0
Stage = 1 Pass = 1
Stage = 2 Pass = 0
Stage = 2 Pass = 1
Stage = 2 Pass = 2

ホストコードについてはステージとパスがどのように算出されるかを追ってみてネットワーク図に書いてみるとわかるものなので、コードに目を通す前に図を書いて過程を理解するようにしてください。

ホストコードと、カーネルのソースコードは以下のようになります。カーネルのソースは次の項目で解説します。

BitonicSortTest.py. 

import pyopencl as cl
import numpy as np
from numpy.random import *
import time

# set the seed to 100
np.random.seed(4)

KERNEL_BITONIC = "bitonic"
KERNEL_MERGE = "merge"

DATA_SIZE = 2097152
global_work_size = (DATA_SIZE, )

data = randint(0, DATA_SIZE, DATA_SIZE).astype(np.uint32)
out = np.zeros(DATA_SIZE).astype(np.uint32)

print(data)

devices = [cl.get_platforms()[0].get_devices(cl.device_type.GPU)[0]]

ctx = cl.Context(devices)
queue = cl.CommandQueue(ctx)
mf = cl.mem_flags
data_mem = cl.Buffer(ctx, mf.USE_HOST_PTR, hostbuf=data)

opt_string = "-Dsize="+str(DATA_SIZE)
options = [opt_string]

program = cl.Program(ctx, """
    __kernel void merge(
            __global uint* data,
            const uint distance) {
        uint gid = get_global_id(0);
        uint lid = get_local_id(0);
        uint cmp_mask;

        int in_range = isless(gid % (distance << 1),distance);

        if(in_range) {
            uint left = data[gid];
            uint right = data[gid+distance];
            cmp_mask = left < right ? 1 : 0;
            data[gid] = select(left,right,cmp_mask);
            data[gid+distance] = select(right,left,cmp_mask);
            //printf("merge. d=%d, gdiv=%d, left=%d, right=%d, data[gid]=%d, data[gid+d]=%d\\n", distance,gid % (distance << 1),left,right,data[gid],data[gid+distance]);
        }
    }

    __kernel void bitonic(
            __global uint* data,
            const uint distance,
            const uint stageDistance) {

        uint gid = get_global_id(0);
        uint lid = get_local_id(0);

        int in_range = isless(gid % (distance << 1),distance);

        if(in_range) {

            uint middle = stageDistance << 1;
            int dir_mask = isgreaterequal(gid%(middle*2),middle);

            uint left = data[gid];
            uint right = data[gid+distance];
            uint cmp_mask;

            if(dir_mask) {
                cmp_mask = left < right ? 1 : 0;
            } else {
                cmp_mask = left > right ? 1 : 0;
            }
            data[gid] = select(left,right,cmp_mask);
            data[gid+distance] = select(right,left,cmp_mask);

            //printf("d=%d, gdiv=%d, left=%d, right=%d, data[gid]=%d, data[gid+d]=%d\\n", distance,gid % (distance << 1),left,right,data[gid],data[gid+distance]);
        }
    }

    """).build(options=options)


kernel_bitonic = cl.Kernel(program, name=KERNEL_BITONIC)
kernel_merge = cl.Kernel(program, name=KERNEL_MERGE)

passes = 0
distance = 0
numberOfStages = np.uint32(np.log2(DATA_SIZE))

st = time.time()

for i in range(numberOfStages-1):
    for j in range(passes+1):
        distance = 1 << (i - j)
        stageDistance = 1 << i

        local_work_size = (np.clip(distance << 1, 0, 256), )
        kernel_bitonic.set_arg(0, data_mem)
        kernel_bitonic.set_arg(1, np.int32(distance))
        kernel_bitonic.set_arg(2, np.int32(stageDistance))

        cl.enqueue_nd_range_kernel(
            queue=queue,
            kernel=kernel_bitonic,
            global_work_size=global_work_size,
            local_work_size=local_work_size)

        #print("stage = " + str(i) + " pass = " + str(j) + " distance = " + str(distance))

    passes += 1


passes = numberOfStages
j = 0

for j in range(passes):
    distance = 1 << (numberOfStages - 1 - j)
    local_work_size = (np.clip(distance << 1, 0, 256), )
    kernel_merge.set_arg(0, data_mem)
    kernel_merge.set_arg(1, np.uint32(distance))
    #print("distance = " + str(distance) + " pass = " + str(j))
    cl.enqueue_nd_range_kernel(
        queue=queue,
        kernel=kernel_merge,
        global_work_size=global_work_size,
        local_work_size=local_work_size)


et = time.time()


out = cl.enqueue_map_buffer(
    queue=queue,
    buf=data_mem,
    flags=cl.mem_flags.WRITE_ONLY,
    offset=0,
    shape=(DATA_SIZE,),
    dtype=np.uint32)


#print(out)
print("%s" % (et-st) )