GPUでソートを行う必要があり、GPUソートにおける代表格であるバイトニックソートについて調べた。

ググってみると色んなサイトにソースやアルゴリズムの図が載ってるけど、いまいちよく解らない。
図を見ると概念レベルの事はすぐ分かるけど、図というものの宿命としてそのスペースが限られてるためか、どこの図も要素数8個とかで少なすぎて、要素数を2^nに抽象化した際のアルゴリズムをイメージしづらい。
あと、処理の流れの図が、前半と後半で分けてあって、アルゴリズムが大きく分けて2フェーズからなると勘違いしやすい図も多い。
いやまぁこれはあくまでも私の頭の問題ですけどね。

個人的に一番分かりやすかったのは米Wikipediaのバイトニックソートのページ。
ここにも図が載ってるけど、分かりやすかったのは図ではなくて、その下に載ってるPythonのソース。

しかしそのソース、純粋なアルゴリズムの理解にはかなり役立つとはいえ、バリバリ再帰つかってるし、リストを連結したりしてるから、GPUだとどんな流れで処理すればいいのかということが分かりにくい。

ということで、まず再帰とリストの連結を取り除いたバージョンを書いてみて、さらにそこからGPUだとこんな感じになるよねというバージョンを書いてみた。
（ソースは本文の最後に載せてます。）

以下ソースの解説。

WikipediaBitonicSorter がWikipediaに載ってたソースをクラス化したもので、 MyFirstBitonicSorter がその名のとおり私が最初に書いたバイトニックソート（再帰とリストの連結なし）で、 GPULikeBitonicSorter がGPUだとこんな処理の流れになるよねバージョンのバイトニックソートである。

もちろん「GPU Like」という事で実際はCPUで処理を行うけど、同じような処理をCUDAやらOpenCLやらOpenGL向けに書けばちゃんと動くと思う。
ただし、共有メモリを使ったり、メモリアクセスをコアレッシングしたり、スレッドの分岐方向をなるべく揃えたりするような工夫はしてないので、このままGPUに持って行っても性能は出し切れないと思う。
そういう最適化は実際にGPUに持っていった後しかできない。

今回のソースはCPUで動作するので速度を測っても仕方ないんだけど、一応速度を測るようにした。
うちの環境では、要素数 1048576 (1024 * 1024) でどれも3分前後かかる中、組み込み関数のsorted（中身はCで書かれたクイックソートかな？）は2秒弱となかなかのタイムだった。
実際にGPUに持って行ってどのくらいのタイムが出るかそのうち試したい。

import random
import timeit


class WikipediaBitonicSorter(object):
    def __init__(self):
        pass
    
    def bitonic_sort(self, data):
        self.comparison_count = 0
        self.swap_count = 0
        return self.__bitonic_sort(True, data)
    
    def __bitonic_sort(self, up, x):
        if len(x)<=1:
            return x
        else: 
            first = self.__bitonic_sort(True,x[:len(x)/2])
            second = self.__bitonic_sort(False,x[len(x)/2:])
            return self.bitonic_merge(up,first+second)
 
    def bitonic_merge(self, up, x): 
        # assume input x is bitonic, and sorted list is returned 
        if len(x) == 1:
            return x
        else:
            self.bitonic_compare(up,x)
            first = self.bitonic_merge(up,x[:len(x)/2])
            second = self.bitonic_merge(up,x[len(x)/2:])
            return first + second
 
    def bitonic_compare(self, up, x):
        dist = len(x)/2
        for i in xrange(dist):
            self.comparison_count += 1
            if (x[i] > x[i+dist]) == up:
                self.swap_count += 1
                x[i], x[i+dist] = x[i+dist], x[i] #swap


class MyFirstBitonicSorter(object):
    def __init__(self):
        pass
    
    def bitonic_sort(self, data):
        self.comparison_count = 0
        self.swap_count = 0
        
        self.data = data
        chunk_size = 2
        while chunk_size <= len(self.data):
            for chunk_index in xrange(len(self.data) / chunk_size):
                _from = chunk_size * chunk_index
                _to = _from + chunk_size - 1
                self.sort(chunk_index % 2 == 0, _from, _to)
            chunk_size *= 2
        return self.data
    
    def sort(self, up, _from, _to):
        chunk_size = _to - _from + 1
        if chunk_size == 2:
            self.compare(up, _from, _to)
        else:
            split = 1
            while split < chunk_size:
                sub_chunk_size = chunk_size / split
                for sub_chunk_index in xrange(split):
                    for i in xrange(sub_chunk_size / 2):
                        a = sub_chunk_size * sub_chunk_index + i + _from
                        b = a + sub_chunk_size / 2
                        self.compare(up, a, b)
                split *= 2
            
    def compare(self, up, a, b):
        self.comparison_count += 1
        if (self.data[a] > self.data[b]) == up:
            self.swap_count += 1
            self.data[a], self.data[b] = self.data[b], self.data[a]
        

class GPULikeBitonicSorter(object):
    def __init__(self):
        pass
    
    def bitonic_sort(self, data):
        self.comparison_count = 0
        self.swap_count = 0
        self.launch_kernel_count = 0
        
        chunk_size = 2
        while chunk_size <= len(data):
            sub_chunk_size = chunk_size
            while sub_chunk_size > 1:
                self.launch_kernel(data, chunk_size, sub_chunk_size)
                sub_chunk_size /= 2
            chunk_size *= 2
        return data
    
    def launch_kernel(self, data, chunk_size, sub_chunk_size):
        self.launch_kernel_count += 1
        for i in xrange(len(data) / 2):
            self.kernel(data, len(data), chunk_size, sub_chunk_size, i)
        
    def kernel(self, data, data_size, chunk_size, sub_chunk_size, thread_index):
        half_chunk_size = chunk_size / 2
        chunk_index = thread_index / half_chunk_size
        
        half_sub_chunk_size = sub_chunk_size / 2
        sub_chunk_index = thread_index / half_sub_chunk_size
        
        up = chunk_index % 2 == 0
        a = sub_chunk_size * sub_chunk_index + thread_index % half_sub_chunk_size
        b = a + half_sub_chunk_size
        
        self.comparison_count += 1
        if (data[a] > data[b]) == up:
            self.swap_count += 1
            data[a], data[b] = data[b], data[a]
            
        
if __name__ == "__main__":
    #random.seed(12345)
    
    data_size = 128 * 128
    trial_count = 3
    
    data = [random.randint(0, data_size) for i in xrange(data_size)]
    results = []
    times = []
    
    wbs = WikipediaBitonicSorter()
    times.append(
        timeit.timeit(
            setup = 'from __main__ import results, data, wbs',
            stmt = 'results.append(wbs.bitonic_sort(data[:]))',
            number = trial_count
        )
    )
    print 'wbs.comparison_count = %d' % wbs.comparison_count
    print 'wbs.swap_count = %d' % wbs.swap_count
    print 'wbs time: %f sec' % times[-1]
    print

    mfbs = MyFirstBitonicSorter()
    times.append(
        timeit.timeit(
            setup = 'from __main__ import results, data, mfbs',
            stmt = 'results.append(mfbs.bitonic_sort(data[:]))',
            number = trial_count
        )
    )
    print 'mfbs.comparison_count = %d' % mfbs.comparison_count
    print 'mfbs.swap_count = %d' % mfbs.swap_count
    print 'mfbs time: %f sec' % times[-1]
    print
    
    glbs = GPULikeBitonicSorter()
    times.append(
        timeit.timeit(
            setup = 'from __main__ import results, data, glbs',
            stmt = 'results.append(glbs.bitonic_sort(data[:]))',
            number = trial_count
        )
    )
    print 'glbs.comparison_count = %d' % glbs.comparison_count
    print 'glbs.swap_count = %d' % glbs.swap_count
    print 'glbs.launch_kernel_count = %d' % glbs.launch_kernel_count
    print 'glbs time: %f sec' % times[-1]
    print
    
    times.append(
        timeit.timeit(
            setup = 'from __main__ import results, data',
            stmt = 'results.append(sorted(data))',
            number = trial_count
        )
    )
    print 'builtin function sorted time: %f sec' % times[-1]
    print
    
    for i in xrange(len(results) - 1):
        if results[i] != results[i+1]:
            print 'Result has a difference.'