Use Catanzaro's algorithm for non-power-of-two interleaves

Author: abadams

src/CodeGen_LLVM.cpp

@@ -2211,6 +2211,8 @@ Value *CodeGen_LLVM::interleave_vectors(const std::vector<Value *> &vecs) {
     }
     int vec_elements = get_vector_num_elements(vecs[0]->getType());
 
+    int factor = gcd(vec_elements, (int)vecs.size());
+
     if (vecs.size() == 1) {
         return vecs[0];
     } else if (vecs.size() == 2) {
@@ -2221,57 +2223,97 @@ Value *CodeGen_LLVM::interleave_vectors(const std::vector<Value *> &vecs) {
             indices[i] = i % 2 == 0 ? i / 2 : i / 2 + vec_elements;
         }
         return optimization_fence(shuffle_vectors(a, b, indices));
-    } else {
-        // Grab the even and odd elements of vecs.
-        vector<Value *> even_vecs;
-        vector<Value *> odd_vecs;
-        for (size_t i = 0; i < vecs.size(); i++) {
-            if (i % 2 == 0) {
-                even_vecs.push_back(vecs[i]);
-            } else {
-                odd_vecs.push_back(vecs[i]);
+    } else if (factor == 1) {
+        // The number of vectors and the vector length is
+        // coprime. (E.g. interleaving an odd number of vectors of some
+        // power-of-two length). Use the algorithm from "A Decomposition for
+        // In-place Matrix Transposition" by Catanzaro et al.
+        std::vector<Value *> v = vecs;
+
+        // Using unary shuffles, get each element into the right ultimate
+        // lane. This works out without collisions because the number of vectors
+        // and the length of each vector is coprime.
+        const int num_vecs = (int)v.size();
+        std::vector<int> shuffle(vec_elements);
+        for (int i = 0; i < num_vecs; i++) {
+            for (int j = 0; j < vec_elements; j++) {
+                int k = j * num_vecs + i;
+                shuffle[k % vec_elements] = j;
             }
+            v[i] = shuffle_vectors(v[i], v[i], shuffle);
         }
 
-        // If the number of vecs is odd, save the last one for later.
-        Value *last = nullptr;
-        if (even_vecs.size() > odd_vecs.size()) {
-            last = even_vecs.back();
-            even_vecs.pop_back();
+        // We intentionally don't put an optimization fence after the unary
+        // shuffles, because some architectures have a two-way shuffle, so it
+        // helps to fuse the unary shuffle into the first layer of two-way
+        // blends below.
+
+        // Now we need to transfer the elements across the vectors. If we
+        // reorder the vectors, this becomes a rotation across the vectors of a
+        // different amount per lane.
+        std::vector<Value *> new_v(v.size());
+        for (int i = 0; i < num_vecs; i++) {
+            int j = (i * vec_elements) % num_vecs;
+            new_v[i] = v[j];
         }
-        internal_assert(even_vecs.size() == odd_vecs.size());
+        v.swap(new_v);
 
-        // Interleave the even and odd parts.
-        Value *even = interleave_vectors(even_vecs);
-        Value *odd = interleave_vectors(odd_vecs);
+        std::vector<int> rotation(vec_elements, 0);
+        for (int i = 0; i < vec_elements; i++) {
+            int k = (i * num_vecs) % vec_elements;
+            rotation[k] = (i * num_vecs) / vec_elements;
+        }
+        internal_assert(rotation[0] == 0);
 
-        if (last) {
-            int result_elements = vec_elements * vecs.size();
+        // We'll handle each bit of the rotation one at a time with a two-way
+        // shuffle.
+        int d = 1;
+        while (d < num_vecs) {
 
-            // Interleave even and odd, leaving a space for the last element.
-            vector<int> indices(result_elements, -1);
-            for (int i = 0, idx = 0; i < result_elements; i++) {
-                if (i % vecs.size() < vecs.size() - 1) {
-                    indices[i] = idx % 2 == 0 ? idx / 2 : idx / 2 + vec_elements * even_vecs.size();
-                    idx++;
-                }
+            for (int i = 0; i < vec_elements; i++) {
+                shuffle[i] = ((rotation[i] & d) == 0) ? i : (i + vec_elements);
             }
-            Value *even_odd = shuffle_vectors(even, odd, indices);
 
-            // Interleave the last vector into the result.
-            last = slice_vector(last, 0, result_elements);
-            for (int i = 0; i < result_elements; i++) {
-                if (i % vecs.size() < vecs.size() - 1) {
-                    indices[i] = i;
-                } else {
-                    indices[i] = i / vecs.size() + result_elements;
-                }
+            for (int i = 0; i < num_vecs; i++) {
+                int j = (i + num_vecs - d) % num_vecs;
+                new_v[i] = shuffle_vectors(v[i], v[j], shuffle);
             }
 
-            return shuffle_vectors(even_odd, last, indices);
-        } else {
-            return interleave_vectors({even, odd});
+            v.swap(new_v);
+
+            d *= 2;
         }
+
+        return concat_vectors(v);
+
+    } else {
+        // The number of vectors shares a factor with the length of the
+        // vectors. Pick some large factor of the number of vectors, interleave
+        // in separate groups, and then interleave the results.
+        const int n = (int)vecs.size();
+        int f = 1;
+        for (int i = 2; i < n; i++) {
+            if (n % i == 0) {
+                f = i;
+                break;
+            }
+        }
+
+        internal_assert(f > 1 && f < n);
+
+        vector<vector<Value *>> groups(f);
+        for (size_t i = 0; i < vecs.size(); i++) {
+            groups[i % f].push_back(vecs[i]);
+        }
+
+        // Interleave each group
+        vector<Value *> interleaved(f);
+        for (int i = 0; i < f; i++) {
+            interleaved[i] = optimization_fence(interleave_vectors(groups[i]));
+        }
+
+        // Interleave the result
+        return interleave_vectors(interleaved);
     }
 }
 

test/performance/CMakeLists.txt

@@ -16,6 +16,7 @@ tests(GROUPS performance
       fast_pow.cpp
       fast_sine_cosine.cpp
       gpu_half_throughput.cpp
+      interleave.cpp
       jit_stress.cpp
       lots_of_inputs.cpp
       memcpy.cpp

test/performance/interleave.cpp

@@ -0,0 +1,159 @@
+#include "Halide.h"
+#include "halide_benchmark.h"
+#include "halide_test_dirs.h"
+
+#include <cstdio>
+
+using namespace Halide;
+using namespace Halide::Tools;
+
+struct Result {
+    int type_size, factor;
+    double bandwidth;
+};
+
+template<typename T>
+Result test_interleave(int factor, const Target &t) {
+    const int N = 8192;
+    Buffer<T> in(N, factor), out(N * factor);
+
+    for (int y = 0; y < factor; y++) {
+        for (int x = 0; x < N; x++) {
+            in(x, y) = (T)(x * factor + y);
+        }
+    }
+
+    Func output;
+    Var x, y;
+
+    output(x) = in(x / factor, x % factor);
+
+    Var xi, yi;
+    output.unroll(x, factor, TailStrategy::RoundUp).vectorize(x, t.natural_vector_size<T>(), TailStrategy::RoundUp);
+    output.output_buffer().dim(0).set_min(0);
+
+    output.compile_jit();
+
+    output.realize(out);
+
+    double time = benchmark(20, 20, [&]() {
+        output.realize(out);
+    });
+
+    for (int y = 0; y < factor; y++) {
+        for (int x = 0; x < N; x++) {
+            uint64_t actual = out(x * factor + y), correct = in(x, y);
+            if (actual != correct) {
+                std::cerr << "For factor " << factor
+                          << "out(" << x << " * " << factor << " + " << y << ") = "
+                          << actual << " instead of " << correct << "\n";
+                exit(1);
+            }
+        }
+    }
+
+    // Uncomment to dump asm for inspection
+    // output.compile_to_assembly("/dev/stdout",
+    // std::vector<Argument>{in}, "interleave", t);
+
+    return Result{(int)sizeof(T), factor, out.size_in_bytes() / (1.0e9 * time)};
+}
+
+template<typename T>
+Result test_deinterleave(int factor, const Target &t) {
+    const int N = 8192;
+    Buffer<T> in(N * factor), out(N, factor);
+
+    for (int x = 0; x < N; x++) {
+        for (int y = 0; y < factor; y++) {
+            in(x * factor + y) = (T)(x + y * N);
+        }
+    }
+
+    Func output;
+    Var x, y;
+
+    output(x, y) = in(x * factor + y);
+
+    Var xi, yi;
+    output.reorder(y, x).bound(y, 0, factor).unroll(y).vectorize(x, t.natural_vector_size<T>(), TailStrategy::RoundUp);
+    // output.output_buffer().dim(0).set_min(0);
+
+    output.compile_jit();
+
+    output.realize(out);
+
+    double time = benchmark(20, 20, [&]() {
+        output.realize(out);
+    });
+
+    for (int y = 0; y < factor; y++) {
+        for (int x = 0; x < N; x++) {
+            uint64_t actual = out(x, y), correct = in(x * factor + y);
+            if (actual != correct) {
+                std::cerr << "For factor " << factor
+                          << "out(" << x << ", " << y << ") = "
+                          << actual << " instead of " << correct << "\n";
+                exit(1);
+            }
+        }
+    }
+
+    // Uncomment to dump asm for inspection
+    output.compile_to_assembly("/dev/stdout",
+    std::vector<Argument>{in}, "interleave", t);
+
+    return Result{(int)sizeof(T), factor, out.size_in_bytes() / (1.0e9 * time)};
+}
+
+int main(int argc, char **argv) {
+    Target target = get_jit_target_from_environment();
+    if (target.arch == Target::WebAssembly) {
+        printf("[SKIP] Performance tests are meaningless and/or misleading under WebAssembly interpreter.\n");
+        return 0;
+    }
+
+    // Set the target features to use for dumping to assembly
+    target.set_features({Target::NoRuntime, Target::NoAsserts, Target::NoBoundsQuery});
+
+    std::cout << "\nbytes, interleave factor, interleave bandwidth (GB/s), deinterleave bandwidth (GB/s):\n";
+#if 0
+    for (int t : {1, 2, 4, 8}) {
+        for (int f = 2; f < 16; f++) {
+#else
+     {
+         {
+            int t = 1, f = 4;
+#endif
+            Result r1, r2;
+            switch (t) {
+            case 1:
+                r1 = test_interleave<uint8_t>(f, target);
+                r2 = test_deinterleave<uint8_t>(f, target);
+                break;
+            case 2:
+                r1 = test_interleave<uint16_t>(f, target);
+                r2 = test_deinterleave<uint16_t>(f, target);
+                break;
+            case 4:
+                r1 = test_interleave<uint32_t>(f, target);
+                r2 = test_deinterleave<uint32_t>(f, target);
+                break;
+            case 8:
+                r1 = test_interleave<uint64_t>(f, target);
+                r2 = test_deinterleave<uint64_t>(f, target);
+                break;
+            default:
+                break;
+            }
+            std::cout << r1.type_size << " "
+                      << r1.factor << " "
+                      << r1.bandwidth << " "
+                      << r2.bandwidth << "\n";
+
+        }
+    }
+
+    printf("Success!\n");
+    return 0;
+}