Task- und Datenparallelität mit Rust

Vektorisierung

Die präsentierte Lösung skaliert, die absolute Performance ist allerdings nicht besonders gut. Um schnelleren Code zu realisieren, gilt es, auch die Vektoreinheiten heutiger Prozessoren (z.B. SSE und AVX) auszunutzen. Hierfür ist das bisher gewählte Speicherlayout nicht optimal. Vielmehr müssten die einzelnen Komponenten des Vektors linear hintereinander angeordnet sein. Damit sich pro Iteration der x-Abstand zwischen mehreren Köpern bestimmen lässt, müssten alle x-Werte linear hintereinander liegen. Daher ist der Wechsel von einem Array of Structs (siehe oben) zu einem Struct of Arrays (siehe unten) notwendig:

typedef struct {
f32 x[nParticles];
f32 y[nParticles];
f32 z[nParticles];
f32 vx[nParticles];
f32 vy[nParticles];
f32 vz[nParticles];
} NBody;

Um nun die Berechnung zu vektorisieren, kann man sich auf den Compiler verlassen und hoffen, dass dieser das Potenzial der Vektorisierung erkennt. In OpenMP besteht zudem die Möglichkeit, dem Compiler Tipps zu geben, um ihm die Arbeit zu erleichtern. Reicht das nicht, können sogenannte Intrinsics helfen, die Assembler-Instruktionen in eine Hochsprache einblenden. Beispielsweise stellt _mm256_add_ps (m256 a, m256 b) eine C-Funktion dar, die eine Vektoraddition durchführt, wobei die Elemente aus einfachgenauen Fließkommazahlen bestehen und der Vektor insgesamt 256 Bit groß ist.

Rust bietet analog zu C entsprechende Intrinsics an, wodurch erfahrende Nutzer ihren Code schnell von C nach Rust portieren können. Allerdings ist der Schritt zur direkten Assembler-Programmierung nicht mehr weit, da dort die entsprechende AVX-Instruktion addps lautet. Zudem sind Intrinsics sehr hardwarespezifisch. Der Wechsel zu einer Hardware mit einem anderen Instruktionssatz erfordert zumindest eine teilweise Neuentwicklung des Codes.

In folgenden Abschnitt wird eine Spracherweiterung (RFC 2366) verwendet, die eine portable Vektorisierung ermöglicht. Diese Erweiterung ist zurzeit nur in der nightly-Version des Rust-Compilers enthalten. Wann und in welcher Form sie in einer stable-Version des Compilers zur Verfügung steht, ist noch nicht klar. Die Grundidee ist, portable Vektor-Datentypen zu definieren. So stellt der Datentyp f32x8 einen Vektor dar, der aus acht einfachgenauen Fließkommazahlen besteht. Die Zeichen vor dem x beschreiben quasi den Basisdatentyp, während die Zahl danach die Häufigkeit der Einträge widerspiegelt. Der Rust-Compiler stellt bereits Implementierungen für Basis-Operationen zur Verfügung, sodass beispielsweise die Addition zweier Vektoren recht einfach ist und sich folgendermaßen implementieren lässt:

pub fn add(a: f32x8, b: f32x8) -> f32x8 {
a+b
}

Wie dies auf den verschiedenen Plattformen (x86, aarch64 etc.) umgesetzt wird, ist für Entwickler nicht relevant. Solange gewährleistet ist, dass alle Zielplattformen eine Vektoreinheit besitzen und acht einfachgenaue Fließkommazahlen addieren können, können Entwickler diese Variante als plattformunabhängig betrachten.

Um das N-Body-Problem zu vektorisieren, wurden die Datenstrukturen umgeschrieben und ein Struct-of-Arrays-Ansatz gewählt. Zudem sind die Felder so definiert, dass die Genauigkeit der Einträge einfach zu ändern ist. Hier nun ein verbessertes Layout der Daten zur Lösung des N-Body-Problems:

pub struct Array<T>([T; N_PARTICLES_SOA]);

pub struct StructOfArrays<T> {
pub x: Array<T>,
pub y: Array<T>,
pub z: Array<T>
}

pub struct NBodySoA {
position: StructOfArrays<PrecisionSoA>,
velocity: StructOfArrays<PrecisionSoA>
}

Im Ausschnitt fehlt (aus Platzgründen) die Definition eines Iterators über die Datenstruktur StructOfArrays. Bei jeder Iteration liefert er einen Vektor bestehend aus (x, y, z) zurück. Die einzelnen Einträge des Vektors bestehen wiederum aus Einträgen (hier f32x8), die von der Vektoreinheit des Prozessors bearbeitbar sind.

Für die neue sequenzielle Lösung (siehe unten) sind daher zwei Iteratoren von Nöten, die im Gleichschritt über die Geschwindigkeit und Position laufen. In der ersten Zeile der Beispiellösung werden diese Iteratoren durch die Methode zip zu einem neuen Iterator verknüpft werden, der gleichmäßig über Geschwindigkeit und Position läuft. Die Lösung besitzt keine maschinenspezifische Instruktion. Das Initialisieren der Vektoren mit dem Wert 0 und das Auslesen einzelner Einträge werden durch die Methoden splat und extract hardwareunabhängig abstrahiert.

position.iter().zip(velocity.iter_mut()).for_each(|((pix, piy, piz), (vix, viy, viz))| {
let mut fx: PrecisionSoA = PrecisionSoA::splat(0.0);
let mut fy: PrecisionSoA = PrecisionSoA::splat(0.0);
let mut fz: PrecisionSoA = PrecisionSoA::splat(0.0);

position.iter().for_each(|(pjx, pjy, pjz)| {
// Newton’s law of universal gravity calculation.
let mut dx: PrecisionSoA = PrecisionSoA::splat(0.0);
let mut dy: PrecisionSoA = PrecisionSoA::splat(0.0);
let mut dz: PrecisionSoA = PrecisionSoA::splat(0.0);

for lane in 0..PrecisionSoA::lanes() {
dx += *pjx - PrecisionSoA::splat(pix.extract(lane));
dy += *pjy - PrecisionSoA::splat(piy.extract(lane));
dz += *pjz - PrecisionSoA::splat(piz.extract(lane));
}

let n2 = dx*dx + dy*dy + dz*dz;
let power = 1.0 / (n2.sqrt() * n2);

fx += dx*power;
fy += dy*power;
fz += dz*power;
});

*vix += fx * dt;
*viy += fy * dt;
*viz += fz * dt;
});

Mit Rayon lässt sich diese Lösung nun parallelisieren. Allerdings kommen selbst definierte Iteratoren zur Anwendung, für die es keine direkte Umsetzung für die parallele Bearbeitung mit Rayon gibt. Es besteht natürlich die Möglichkeit, diese zu implementieren. Darüber hinaus bietet Rayon aber auch die Option, die Iteratoren für die sequenzielle Bearbeitung entsprechend umzuwandeln. Dies erhöht den Overhead – und stellt daher keine optimale Lösung dar –, reduziert aber den Entwicklungsaufwand. Im Beispiel wird nur die äußere Schleife durch die Methode zur Umwandlung der Iteratoren (par_bridge()) ergänzt:

position.iter().zip(velocity.iter_mut()).par_bridge().for_each(|((pix, piy, piz), (vix, viy, viz))| {
...
}
Diagramm zu den Leistungsmessungen der Mehrkörpersimulation.

Die Ergebnisse aus den Leistungsmessungen zeigen, dass die Umwandlung in diesem Fall effizient ist und sich eine Eigenentwicklung daher kaum lohnt.