Improve AVX512 feature detection

examples/blas.cc

@@ -0,0 +1,221 @@
+// Copyright 2024 Justine Alexandra Roberts Tunney
+//
+// Permission to use, copy, modify, and/or distribute this software for
+// any purpose with or without fee is hereby granted, provided that the
+// above copyright notice and this permission notice appear in all copies.
+//
+// THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL
+// WARRANTIES WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED
+// WARRANTIES OF MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE
+// AUTHOR BE LIABLE FOR ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL
+// DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR
+// PROFITS, WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR OTHER
+// TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR
+// PERFORMANCE OF THIS SOFTWARE.
+
+#include <unistd.h>
+#include <cinttypes>
+#include <cmath>
+#include <cstdio>
+#include <cstdlib>
+#include <ctime>
+#include "libc/assert.h"
+
+// high performance high accuracy matrix multiplication in ansi c
+
+#define MATH __target_clones("avx512f,fma")
+
+namespace {
+namespace ansiBLAS {
+
+static constexpr int KN = 8;
+
+struct Vector {
+  double v[KN];
+};
+
+Vector load(const float *p) {
+  Vector x;
+  for (int i = 0; i < KN; ++i)
+    x.v[i] = p[i];
+  return x;
+}
+
+Vector madd(Vector x, Vector y, Vector s) {
+  for (int i = 0; i < KN; ++i)
+    s.v[i] = fma(x.v[i], y.v[i], s.v[i]);
+  return s;
+}
+
+float hsum(Vector x) {
+  double s = 0;
+  for (int i = 0; i < KN; ++i)
+    s += x.v[i];
+  return s;
+}
+
+struct ansiBLAS {
+ public:
+  ansiBLAS(int k, const float *A, int lda, const float *B, int ldb, float *C,
+           int ldc, int ith, int nth)
+      : k(k),
+        A(A),
+        lda(lda),
+        B(B),
+        ldb(ldb),
+        C(C),
+        ldc(ldc),
+        ith(ith),
+        nth(nth) {
+  }
+
+  void matmul(int m, int n) {
+    mnpack(0, m, 0, n);
+  }
+
+ private:
+  void mnpack(int m0, int m, int n0, int n) {
+    int mc, nc, mp, np;
+    if (m - m0 <= 0 || n - n0 <= 0)
+      return;
+    if (m - m0 >= 4 && n - n0 >= 3) {
+      mc = 4;
+      nc = 3;
+      gemm<4, 3>(m0, m, n0, n);
+    } else {
+      mc = 1;
+      nc = 1;
+      gemm<1, 1>(m0, m, n0, n);
+    }
+    mp = m0 + (m - m0) / mc * mc;
+    np = n0 + (n - n0) / nc * nc;
+    mnpack(mp, m, n0, np);
+    mnpack(m0, m, np, n);
+  }
+
+  template <int RM, int RN>
+  MATH void gemm(int m0, int m, int n0, int n) {
+    int ytiles = (m - m0) / RM;
+    int xtiles = (n - n0) / RN;
+    int tiles = xtiles * ytiles;
+    int duty = (tiles + nth - 1) / nth;
+    int start = duty * ith;
+    int end = start + duty;
+    if (end > tiles)
+      end = tiles;
+    for (int job = start; job < end; ++job) {
+      int ii = m0 + job / xtiles * RM;
+      int jj = n0 + job % xtiles * RN;
+      Vector Cv[RN][RM] = {};
+      for (int l = 0; l < k; l += KN)
+        for (int j = 0; j < RN; ++j)
+          for (int i = 0; i < RM; ++i)
+            Cv[j][i] = madd(load(A + lda * (ii + i) + l),  //
+                            load(B + ldb * (jj + j) + l),  //
+                            Cv[j][i]);
+      for (int j = 0; j < RN; ++j)
+        for (int i = 0; i < RM; ++i)
+          C[ldc * (jj + j) + (ii + i)] = hsum(Cv[j][i]);
+    }
+  }
+
+  const int k;
+  const float *const A;
+  const int lda;
+  const float *const B;
+  const int ldb;
+  float *const C;
+  const int ldc;
+  const int ith;
+  const int nth;
+};
+
+void sgemm(int m, int n, int k,      //
+           const float *A, int lda,  //
+           const float *B, int ldb,  //
+           float *C, int ldc) {
+  int nth = sysconf(_SC_NPROCESSORS_ONLN);
+#pragma omp parallel for
+  for (int ith = 0; ith < nth; ++ith) {
+    ansiBLAS tb{k, A, lda, B, ldb, C, ldc, ith, nth};
+    tb.matmul(m, n);
+  }
+}
+
+}  // namespace ansiBLAS
+
+long micros(void) {
+  struct timespec ts;
+  clock_gettime(CLOCK_REALTIME, &ts);
+  return ts.tv_sec * 1000000 + (ts.tv_nsec + 999) / 1000;
+}
+
+unsigned rand32(void) {
+  /* Knuth, D.E., "The Art of Computer Programming," Vol 2,
+     Seminumerical Algorithms, Third Edition, Addison-Wesley, 1998,
+     p. 106 (line 26) & p. 108 */
+  static unsigned long long lcg = 1;
+  lcg *= 6364136223846793005;
+  lcg += 1442695040888963407;
+  return lcg >> 32;
+}
+
+float float01(unsigned x) {  // (0,1)
+  return 1.f / 8388608 * ((x >> 9) + .5f);
+}
+
+float numba(void) {  // (-1,1)
+  return float01(rand32()) * 2 - 1;
+}
+
+void fill(int m, int n, float *A, int lda) {
+  for (int j = 0; j < n; ++j)
+    for (int i = 0; i < m; ++i)
+      A[lda * j + i] = numba();
+}
+
+float *new_matrix(int m, int n, int *lda) {
+  void *ptr = 0;
+  int b = 64 / sizeof(float);
+  *lda = (n + b - 1) & -b;
+  posix_memalign(&ptr, 4096, sizeof(float) * m * *lda);
+  return (float *)ptr;
+}
+
+}  // namespace
+
+void barrier(void) {
+}
+void (*pBarrier)(void) = barrier;
+
+#define BENCH(x)                                                           \
+  do {                                                                     \
+    x;                                                                     \
+    int N = 10;                                                            \
+    long t1 = micros();                                                    \
+    for (long i = 0; i < N; ++i) {                                         \
+      pBarrier();                                                          \
+      x;                                                                   \
+    }                                                                      \
+    long t2 = micros();                                                    \
+    printf("%8" PRId64 " µs %s %g gigaflops\n", (t2 - t1 + N - 1) / N, #x, \
+           1e6 / ((t2 - t1 + N - 1) / N) * m * n * k * 2 * 1e-9);          \
+  } while (0)
+
+int main() {
+  int m = 1024;
+  int n = 1024;
+  int k = 1024;
+  int lda, ldb, ldc;
+  float *A = new_matrix(m, k, &lda);
+  float *B = new_matrix(n, k, &ldb);
+  float *C = new_matrix(n, m, &ldc);
+  fill(k, n, A, lda);
+  fill(k, m, B, ldb);
+  BENCH(ansiBLAS::sgemm(m, n, k, A, lda, B, ldb, C, ldc));
+  assert(C[0] == -0x1.20902ap+4);
+  assert(C[1] == -0x1.bf7726p+4);
+  free(C);
+  free(B);
+  free(A);
+}