JavaScriptでバイトニックソート

GPGPUでバイトニックソートを実装しようと思いましたが、アルゴリズムの理解のためにまずはJavaScriptで実装してみました。
といっても英語版WikipediaのバイトニックソートのJava実装をそのままJavaScriptで実装しただけですが...

function kernel(x, p, q) {
  const d = 1 << (p - q);

  for (let i = 0; i < x.length; i++) {
    const up = ((i >> p) & 2) === 0;
    if ((i & d) == 0 && (x[i] > x[i | d]) === up) {
      const tmp = x[i];
      x[i] = x[i | d];
      x[i | d] = tmp;
    }
  }
}

function bitonicsort(x, n) {
  for (let i = 0; i < n; i++) {
    for(let j = 0; j <=i; j++) {
      kernel(x, i, j);
    }
  }
}

以下の使用例では0から255までの数字がシャッフルして入った配列をbitonicsort(x, n)でソートしています。

const n = 8;
const size = 2**n;
const x = Array.from({length: size}, (v, i) => );

// shuffle
x.forEach((v, i) => {
  const j = Math.floor(Math.random() * size);
  const tmp = x[i];
  x[i] = x[j];
  x[j] = tmp;
});

console.log(x);
// [75, 16, 59, 28, 110, 199, 188, 144, 10, ...]

bitonicsort(x, n);

console.log(x);
// [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, ...]

バイトニックソートの解説をいろいろ読んでなんとなく感覚は掴めたものの具体的な実装がいまいちピンとこなかったですが、英語版Wikipediaにある次の画像とJava実装が一番わかりやすい印象です。

(Bitonic sorter - Wikipediaより引用)

画像は左から右にソートが順に進んでいく様子を表しており、右端がソートが完全に終わった状態になります。
青い部分は昇順、緑色の部分は降順になるようにソートが行われ、矢印が値を比較しあうインデックスを表しています。

画像を横方向に見ると大きく4つのブロックに分かれています。これがbitonicsort関数の外側のforループに該当します。画像では配列の長さが2^4=16なので4回forループを回す感じですね。

さらに、各ブロックの中を見てみると左のブロックから順に1, 2, 3, 4と4つのサブブロックを含んでいます。これがbitonicsort関数の内側のforループに対応しています。つまりkernel関数は1つのサブブロックの中の処理を行っていることになります。

kernel関数では最初にd = 1 << (p - q)を求めています。この値がある配列のインデックスから比較を行うインデックスまでの距離になります。例えば画像の4つ目のブロックの1つ目のサブブロックではp = 3, q = 0 になるのでd = 1 << (3 - 0) = 8となり、各インデックスでは8つ離れたインデックスと値と比較します。

kernel関数のforループ内で実際に比較処理を行い必要があれば値をスワップしています。

forループの最初のup = ((i >> p) & 2) === 0でそのインデックスでは昇順、降順のどちらで値を並べるのかを決定しています。画像を見ると1つ目のブロックではインデックスが2進むごとに降順・昇順が入れ替わり、2つ目のブロックでは4ごと、3つ目では8、4つ目では16と2^(p+1)ごとに昇順・降順が入れ替わっています。そのため(i >> p) & 2でp + 1ビット目の値を調べることで昇順に並べるか、降順に並べるかを判断できます。

あるインデックスに対してそれと比較するインデックスはdだけ離れています。そのため2つのインデックスをビットで考えると値はp - qビット目の値が0か1かの違いだけになります(一番右端のビットを0ビット目として数えています)。配列の先頭に近いほうにあるインデックスをiとすると比較対象のインデックスはp - qビット目が1になるのでそのインデックスはi | dになります。つまり、if文内の(i & d) == 0は配列の先頭に近いほうにあるインデックスかどうかを確認しており、x[i] > x[i | d]で2つの値の大小を確認していることになります。

---

先に紹介した実装だとkernel関数内で値をスワップしていますが、GPGPUで2つのバッファに交互に書き込むことを意識した場合、こんな感じになるのかなというのも実装したので載せておきます。

function kernel(x, p, q) {
  const y = new Array(x.length);
  const d = 1 << (p - q);

  for (let i = 0; i < x.length; i++) {
    const up = ((i >> p) & 2) === 0;

    if ((i & d) == 0 && (x[i] > x[i + d]) === up) {
      y[i] =  x[i + d];
    } else if ((i & d) == d && (x[i - d] > x[i]) === up) {
      y[i] = x[i - d];
    } else {
      y[i] = x[i];
    }
  }
  return y;
}

function bitonicsort(x, n) {
  for (let i = 0; i < n; i++) {
    for(let j = 0; j <=i; j++) {
      x = kernel(x, i, j);
    }
  }
  return x;
}

const n = 8;
const size = 2**n;
const x = Array.from({length: size}, (v, i) => i);

// shuffle
x.forEach((v, i) => {
  const j = Math.floor(Math.random() * size);
  const tmp = x[i];
  x[i] = x[j];
  x[j] = tmp;
});

console.log(x);

const y = bitonicsort(x, n);

console.log(y);