La red bitónica: estructura y coste

Ken Batcher publicó en 1968 una red de ordenación construida sobre una idea con nombre extraño y consecuencias muy concretas: una secuencia que sube y luego baja se puede partir, con un único paso de comparadores independientes, en dos mitades tales que todo elemento de la primera es menor o igual que todo elemento de la segunda, y que además siguen siendo secuencias del mismo tipo. Aplicar eso recursivamente ordena. Y como la recursión es sobre la estructura y no sobre los datos, se desenrolla en un patrón fijo de índices que es exactamente lo que una GPU sabe ejecutar.

Qué es una secuencia bitónica

Una secuencia es bitónica si crece hasta un punto y a partir de ahí decrece, o si es una rotación cíclica de una secuencia así. Las dos formas degeneradas cuentan: una secuencia enteramente creciente es bitónica —el tramo decreciente está vacío— y una decreciente también.

bitonica:      1  4  7  9  8  5  3  2        sube y baja
bitonica:      8  5  3  2  1  4  7  9        rotacion de la anterior
bitonica:      1  2  3  4  5  6  7  8        caso degenerado
no bitonica:   1  5  2  6  3  7  4  8        sube y baja cuatro veces

El interés está en un hecho que no es evidente: dos secuencias ordenadas, una en sentido creciente y otra en sentido decreciente, concatenadas, forman una secuencia bitónica. Eso da la vía de construcción.

La partición bitónica

Sea una secuencia bitónica de longitud n par. Se comparan las posiciones i e i + n/2 para todo i menor que n/2, y se intercambian si están desordenadas.

entrada bitonica:   3  5  8  9  7  4  2  1
comparar i con i+4: (3,7) (5,4) (8,2) (9,1)
tras el paso:       3  4  2  1  7  5  8  9
                    |__ mitad baja __|  |__ mitad alta __|

Tras ese único paso ocurren tres cosas a la vez. Todo elemento de la mitad baja es menor o igual que todo elemento de la mitad alta. La mitad baja sigue siendo bitónica. La mitad alta también. Las tres propiedades juntas son el teorema de Batcher, y su demostración más limpia usa el principio del cero y el uno: con ceros y unos, una secuencia bitónica es de la forma 0*1*0* o 1*0*1*, y comprobar los pocos casos posibles es mecánico.

La consecuencia es que ordenar una secuencia bitónica de n elementos requiere log2(n) pasos: uno con distancia n/2, otro con distancia n/4 sobre cada mitad, y así hasta distancia 1. Todos los comparadores de un mismo paso son independientes, así que cada paso es un dispatch o una iteración con barrera.

Esa operación —convertir una bitónica en una ordenada— se llama fusión bitónica.

La red completa

Falta llegar a una secuencia bitónica desde una entrada arbitraria, y ahí entra la recursión estructural. Se ordenan los pares adyacentes de forma alterna: el primero creciente, el segundo decreciente, el tercero creciente. Cada dos pares se obtiene entonces una secuencia bitónica de cuatro, que se fusiona alternando el sentido; cada dos grupos de cuatro se obtiene una bitónica de ocho, y así hasta el array completo.

La red queda descrita por dos bucles anidados. El exterior recorre las etapas, con k tomando los valores 2, 4, 8, hasta n: k es la longitud de la secuencia bitónica que se está fusionando. El interior recorre los pasos de esa fusión, con j tomando los valores k/2, k/4, hasta 1: j es la distancia entre los elementos que se comparan.

for (let k = 2; k <= n; k <<= 1) {          // etapa
  for (let j = k >> 1; j > 0; j >>= 1) {    // paso
    unPaso(j, k);                            // n/2 comparadores independientes
  }
}

El número de pasos totales, con L = log2(n), es 1 + 2 + ... + L, es decir L·(L+1)/2. Cada paso ejecuta n/2 comparadores. De ahí sale el coste:

El tamaño de la red es (n/4)·L·(L+1), o sea O(n·log²n), y la profundidad es L·(L+1)/2, o sea O(log²n). Para un millón de elementos, 110 millones de comparaciones frente a los 20 millones de un mergesort secuencial: un factor cinco y medio de trabajo extra. A cambio, la profundidad es 210 en vez de un millón.

La columna que de verdad duele en WebGPU es la de pasos, porque en la implementación ingenua cada paso es un dispatch, y 210 dispatches encadenados con su sobrecoste son varios milisegundos antes de comparar nada. Reducir ese número es el objetivo entero de la implementación.

La aritmética de índices

La traducción a código de la estructura anterior es sorprendentemente compacta, y se apoya en dos operaciones bit a bit.

El compañero de la posición i en el paso de distancia j es i XOR j. Con j potencia de dos, el o exclusivo alterna el bit correspondiente: si ese bit era cero, el compañero está j posiciones más adelante; si era uno, j posiciones más atrás. Cada comparador se procesa una sola vez si se exige que el compañero esté por delante.

El sentido del comparador en la etapa k lo da el bit correspondiente a k en el índice: si i AND k es cero, el bloque se ordena en sentido creciente; si no, decreciente. Eso es lo que produce la alternancia de sentidos que construye las bitónicas del siguiente nivel, y en la última etapa, con k = n, i AND k es cero para todos los índices y toda la secuencia sale creciente.

// La red completa en JavaScript, para verificar los indices antes
// de escribir el shader. Verificable con el principio del 0 y el 1.
function bitonico(a) {
  const n = a.length;                    // n debe ser potencia de dos
  for (let k = 2; k <= n; k <<= 1) {
    for (let j = k >> 1; j > 0; j >>= 1) {
      for (let i = 0; i < n; i++) {
        const l = i ^ j;
        if (l <= i) continue;            // procesar cada par una sola vez
        const creciente = (i & k) === 0;
        if ((a[i] > a[l]) === creciente) {
          const t = a[i]; a[i] = a[l]; a[l] = t;
        }
      }
    }
  }
  return a;
}

El bucle sobre i es el que desaparece en la GPU: se convierte en una invocación por índice, y la condición l <= i hace que la mitad de las invocaciones no haga nada. Es un desperdicio del 50 por ciento que se puede evitar lanzando n/2 invocaciones y calculando el índice con desplazamientos de bits, a costa de una aritmética menos legible.

La restricción que arrastra toda la construcción es que n tiene que ser potencia de dos. La red no está definida para otros tamaños. La solución universal es rellenar hasta la siguiente potencia de dos con un valor centinela que quede al final: el máximo del tipo si se ordena de forma creciente. Con un millón de elementos reales hay que rellenar hasta 1.048.576, que casualmente son 1.048.576, pero con un millón y uno habría que ir hasta 2.097.152 y se ordenaría el doble de datos. Ese salto es una de las razones para preferir un radix sort cuando el tamaño es arbitrario.