[X86][AMX] Add unit tests for AMX feature

[yubing]

Summary

Currently, we only test compilation results of AMX cases.

[yubing]

Take notes for myself:
Once spr is supported in testing infra, llvm-test-suite/configure should be modified as well to check if the cpu support amx.

SingleSource/UnitTests/Vector/AMXINT8/CMakeLists.txt

3 ↗

(On Diff #493105)

should place a FIXME here.

[LuoYuanke]

SingleSource/UnitTests/Vector/AMXINT8/CMakeLists.txt

5 ↗

(On Diff #493105)

Should we create Matrix folder under UnitTests and put the AMX test case to Matrix folder?

SingleSource/UnitTests/Vector/AMXINT8/t_gemm_bf16.cpp

9 ↗

(On Diff #493105)

Maybe format the function name and parameter following the LLVM coding style.

39 ↗

(On Diff #493105)

__ is reserved by compiler. We may rename it to bfloat16_t?

149 ↗

(On Diff #493105)

Check the err_num for inner_product.

SingleSource/UnitTests/Vector/AMXINT8/t_gemm_int8.cpp

129 ↗

(On Diff #493105)

Check the err_num for inner_product.

[yubing]

SingleSource/UnitTests/Vector/AMXINT8/t_gemm_int8.cpp

129 ↗

(On Diff #493105)

we don't need to check errnum here.
see line131: err_num = memcmp(gC, gD, sizeof(gC))?(1+err_num):err_num;

[yubing]

SingleSource/UnitTests/Vector/AMXINT8/t_gemm_bf16.cpp

9 ↗

(On Diff #493105)

it is not necessary since other testcases don't follow LLVM coding style.

[yubing]

move into Matrix dir

[yubing]

@craig.topper hi, can we use sapphirerapids server to run pre-CI during community review?

[craig.topper]

I think this is the patch that added the checks for AVX512 support. https://reviews.llvm.org/D38182

SingleSource/UnitTests/Matrix/AMXINT8/CMakeLists.txt

5

Matirx -> Matrix

SingleSource/UnitTests/Matrix/CMakeLists.txt

10

Matirx -> Matrix

SingleSource/UnitTests/Matrix/Makefile

1

Matirx -> Matrix

[yubing]

This comment was removed by yubing.

[yubing]

Matirx -> Matrix

[LuoYuanke]

SingleSource/UnitTests/Matrix/Makefile

12

Maybe just check HAVE_X86_AVX512F_INSTRUCTIONS?

[LuoYuanke]

LGTM.

[yubing]

https://github.com/llvm/llvm-test-suite/commit/afc7e6588b3b43351e49f697fe48283ec5056c11 has landed.
i use "arc land" to land it successfully, of course the commit carry "Differential Revision: https://reviews.llvm.org/D142845"

[MatzeB]

SingleSource/UnitTests/Matrix/AMXINT8/t_gemm_bf16.reference_output

2

I just noticed this in the testsuite. This exit code indicates an "Illegal instruction". Is this test really supposed to run for people and do we really expect it to hit an illegal instruction? This feels silly...

[MatzeB]

SingleSource/UnitTests/Matrix/AMXINT8/t_gemm_bf16.reference_output

2

(I am trying to cleanup the test-suite tests to only ever exit with code 0 when working properly...)

[LuoYuanke]

SingleSource/UnitTests/Matrix/AMXINT8/t_gemm_bf16.reference_output

2

This test case would generate AMX instructions and is supposed to run on SPR machine. Since we have have SPR machine in buildbot, we just check it can be built by compiler.

SingleSource/UnitTests/CMakeLists.txt

 include(CheckCCompilerFlag)
 
 add_subdirectory(C++11)
 add_subdirectory(Float)
+add_subdirectory(Matrix)
 add_subdirectory(SignlessTypes)
 add_subdirectory(Threads)
 add_subdirectory(Vector)
 add_subdirectory(Vectorizer)
 
 list(APPEND CFLAGS -Wno-implicit-function-declaration -Wno-implicit-int)
 
 # FIXME: Disable SJLJ tests for now, until EH edges are represented.

SingleSource/UnitTests/Matrix/AMXINT8/CMakeLists.txt

+list(APPEND CPPFLAGS -I ${CMAKE_SOURCE_DIR}/${MATRIX_MAIN_DIR})
+list(APPEND LDFLAGS -lm)
+list(APPEND CFLAGS "-march=sapphirerapids")
+list(APPEND CFLAGS -fms-extensions)
+llvm_singlesource(PREFIX "Matirx-AMXINT8-")

SingleSource/UnitTests/Matrix/AMXINT8/Makefile

+# SingleSource/UnitTests/Matrix/AMXINT8/Makefile
+
+DIRS = 
+LEVEL = ../../../..
+CFLAGS += -fms-extensions -march=native -mamx-int8 -mamx-bf16 -I${SourceDir}/..
+LDFLAGS += -lm
+
+include $(LEVEL)/SingleSource/Makefile.singlesrc
+
+TARGET_FLAGS += -march=native -mamx-int8 -mamx-bf16
+LCCFLAGS += -march=native -mamx-int8 -mamx-bf16 

SingleSource/UnitTests/Matrix/AMXINT8/t_gemm_bf16.cpp

+#include <immintrin.h>
+#include <string.h>
+#include <cassert>
+#include <iostream>
+
+
+#define TILE_SZ 16
+__attribute__((noinline))
+void inner_product2(int *A_mem, int *B_mem, int *C_mem, int M, int N, int K) {
+  // tiling
+  const int m = M / TILE_SZ;
+  const int n = N / TILE_SZ;
+  const int k = K / TILE_SZ;
+
+  assert(!(M % TILE_SZ) && !(N % TILE_SZ) && !(K % TILE_SZ));
+  for (int i = 0; i < m; i++)
+    for (int j = 0; j < n; j++) {
+      __tile1024i c = {TILE_SZ, TILE_SZ*sizeof(int)};
+      __tile_zero(&c);
+      for (int l = 0; l < k; l++) {
+        __tile1024i a = {TILE_SZ, TILE_SZ*sizeof(int)};
+        __tile1024i b = {TILE_SZ, TILE_SZ*sizeof(int)};
+        __tile_loadd(&a, A_mem+(i*TILE_SZ)*K+l*TILE_SZ, K*sizeof(int));
+        __tile_loadd(&b, B_mem+(l*TILE_SZ)*N+j*TILE_SZ, N*sizeof(int));
+        __tile_dpbf16ps(&c, a, b);
+      }
+      __tile_stored(C_mem+(i*TILE_SZ)*M+j*TILE_SZ, N*sizeof(int), c);
+    }
+}
+
+#define TILE_M 16
+#define TILE_N 16
+#define TILE_K 16
+#define M_ACC 2
+#define N_ACC 2
+#define KPACK 1
+typedef int type_t;
+typedef int res_type_t;
+
+template <short Row, short Col>
+class Tile {
+public:
+  __tile1024i& getTile() {
+    return tile;
+  }
+private:
+  __tile1024i tile {Row, Col*sizeof(int)};
+};
+
+template <short Row, short Col>
+static void tilezero(Tile<Row, Col> &tile) {
+  __tile_zero(&tile.getTile());
+}
+
+template <short Row, short Col>
+static void tileload(Tile<Row, Col> &dst, const void *base, size_t stride) {
+  __tile_loadd(&dst.getTile(), base, stride);
+}
+
+template <short Row, short Col>
+static void tdp(Tile<Row, Col> &dst, Tile<Row, Col> src1, Tile<Row, Col> src2) {
+  __tile_dpbf16ps(&dst.getTile(), src1.getTile(), src2.getTile());
+}
+
+template <short Row, short Col>
+void tilestore(Tile<Row, Col> &src, void *base, size_t stride) {
+  __tile_stored(base, stride, src.getTile());
+}
+
+__attribute__((noinline))
+void inner_product(int *A_mem, int *B_mem, int *C_mem, int M, int N, int K) {
+  // tiling
+  assert(M%(TILE_M*M_ACC) == 0 && N%(TILE_N*N_ACC) == 0 && K%TILE_K == 0);
+  for (int n = 0; n < N; n += N_ACC*TILE_N) {
+    for (int m = 0; m < M; m += M_ACC*TILE_M) {
+      Tile<TILE_M, TILE_N> tC[M_ACC][N_ACC];
+      Tile<TILE_M, TILE_K> tA[M_ACC];
+      Tile<TILE_K, TILE_N> tB;
+
+      for (int n_acc = 0; n_acc < N_ACC; ++n_acc)
+        for (int m_acc = 0; m_acc < M_ACC; ++m_acc)
+          tilezero(tC[m_acc][n_acc]);
+
+      for (int k = 0; k < K; k += TILE_K) {
+        for (int n_acc = 0; n_acc < N_ACC; ++n_acc) {
+          tileload(tB, B_mem+k*N+n+n_acc*TILE_N, N*sizeof(type_t)*KPACK);
+          for (int m_acc = 0; m_acc < M_ACC; ++m_acc) {
+            if (n_acc == 0)
+              tileload(tA[m_acc], A_mem+(m + m_acc*TILE_M)*K+k, K*sizeof(type_t));
+            tdp(tC[m_acc][n_acc], tA[m_acc], tB);
+            if (k == K - TILE_K) {
+              int mc = m + m_acc*TILE_M, nc = n + n_acc*TILE_N;
+              tilestore(tC[m_acc][n_acc], C_mem+mc*N+nc, N*sizeof(res_type_t));
+            }
+          }
+        }
+      }
+    }
+  } 
+}
+
+float make_fp32(short x)
+{
+  unsigned int y = x;
+  y = y << 16;
+  float *res = reinterpret_cast<float*>(&y);
+  return *res;
+}
+
+unsigned short make_bf16(float x)
+{
+  int *res = reinterpret_cast<int*>(&x);
+  *res = *res >> 16;
+  return (unsigned short)*res;
+}
+
+void inner_product_ref(int *A_mem, int *B_mem, int *C_mem, int M, int N, int K) {
+  // tiling
+  for (int m = 0; m < M; m++)
+    for (int n = 0; n < N; n++) {
+      for (int k = 0; k < K; k++) {
+        short *va = (short *)(A_mem + m*K + k);
+        short *vb = (short *)(B_mem + k*N + n);
+        float acc = *((float*)(C_mem + m*N + n));
+        for (int i = 0; i < 2; i++) {
+          acc += (make_fp32(va[i]) * make_fp32(vb[i]));
+        }
+        *((float*)(C_mem + m*N + n))= acc;
+      }
+    }
+}
+
+#define SIZE 128 
+int gA[SIZE][SIZE], gB[SIZE][SIZE], gC[SIZE][SIZE], gD[SIZE][SIZE];
+int main() {
+  for (int i = 0; i < SIZE; i++)
+    for (int j = 0; j < SIZE; j++) {
+      unsigned short *pAbh = (unsigned short *)&gA[i][j];
+      *pAbh = make_bf16(0.1f * (i+j));
+      *(pAbh+1) = make_bf16(0.1f * (i+j));
+      unsigned short *pBbh = (unsigned short *)&gB[i][j];
+      *pBbh = make_bf16(0.1f * (i+j));
+      *(pBbh+1) = make_bf16(0.1f * (i+j));
+    }
+
+  inner_product_ref((int *)gA, (int *)gB, (int *)gC, SIZE, SIZE, SIZE);
+  inner_product((int *)gA, (int *)gB, (int *)gD, SIZE, SIZE, SIZE);
+  int err_num = memcmp(gC, gD, sizeof(gC))?1:0;
+  inner_product2((int *)gA, (int *)gB, (int *)gD, SIZE, SIZE, SIZE);
+  err_num = memcmp(gC, gD, sizeof(gC))?(1+err_num):err_num;
+  if (err_num == 0)
+    std::cout << "PASSED\n";
+  else
+    std::cout << "FAILED\n";
+}
+

SingleSource/UnitTests/Matrix/AMXINT8/t_gemm_bf16.reference_output

+exit 132

SingleSource/UnitTests/Matrix/AMXINT8/t_gemm_int8.cpp

+#include <immintrin.h>
+#include <string.h>
+#include <cassert>
+#include <iostream>
+
+
+#define TILE_SZ 16
+__attribute__((noinline))
+void inner_product2(int *A_mem, int *B_mem, int *C_mem, int M, int N, int K) {
+  // tiling
+  const int m = M / TILE_SZ;
+  const int n = N / TILE_SZ;
+  const int k = K / TILE_SZ;
+
+  assert(!(M % TILE_SZ) && !(N % TILE_SZ) && !(K % TILE_SZ));
+  for (int i = 0; i < m; i++)
+    for (int j = 0; j < n; j++) {
+      __tile1024i c = {TILE_SZ, TILE_SZ*sizeof(int)};
+      __tile_zero(&c);
+      for (int l = 0; l < k; l++) {
+        __tile1024i a = {TILE_SZ, TILE_SZ*sizeof(int)};
+        __tile1024i b = {TILE_SZ, TILE_SZ*sizeof(int)};
+        __tile_loadd(&a, A_mem+(i*TILE_SZ)*K+l*TILE_SZ, K*sizeof(int));
+        __tile_loadd(&b, B_mem+(l*TILE_SZ)*N+j*TILE_SZ, N*sizeof(int));
+        __tile_dpbssd(&c, a, b);
+      }
+      __tile_stored(C_mem+(i*TILE_SZ)*M+j*TILE_SZ, N*sizeof(int), c);
+    }
+}
+
+#define TILE_M 16
+#define TILE_N 16
+#define TILE_K 16
+#define M_ACC 2
+#define N_ACC 2
+#define KPACK 1
+typedef int type_t;
+typedef int res_type_t;
+
+template <short Row, short Col>
+class Tile {
+public:
+  __tile1024i& getTile() {
+    return tile;
+  }
+private:
+  __tile1024i tile {Row, Col*sizeof(int)};
+};
+
+template <short Row, short Col>
+static void tilezero(Tile<Row, Col> &tile) {
+  __tile_zero(&tile.getTile());
+}
+
+template <short Row, short Col>
+static void tileload(Tile<Row, Col> &dst, const void *base, size_t stride) {
+  __tile_loadd(&dst.getTile(), base, stride);
+}
+
+template <short Row, short Col>
+static void tdp(Tile<Row, Col> &dst, Tile<Row, Col> src1, Tile<Row, Col> src2) {
+  __tile_dpbssd(&dst.getTile(), src1.getTile(), src2.getTile());
+}
+
+template <short Row, short Col>
+void tilestore(Tile<Row, Col> &src, void *base, size_t stride) {
+  __tile_stored(base, stride, src.getTile());
+}
+
+__attribute__((noinline))
+void inner_product(int *A_mem, int *B_mem, int *C_mem, int M, int N, int K) {
+  // tiling
+  assert(M%(TILE_M*M_ACC) == 0 && N%(TILE_N*N_ACC) == 0 && K%TILE_K == 0);
+  for (int n = 0; n < N; n += N_ACC*TILE_N) {
+    for (int m = 0; m < M; m += M_ACC*TILE_M) {
+      Tile<TILE_M, TILE_N> tC[M_ACC][N_ACC];
+      Tile<TILE_M, TILE_K> tA[M_ACC];
+      Tile<TILE_K, TILE_N> tB;
+
+      for (int n_acc = 0; n_acc < N_ACC; ++n_acc)
+        for (int m_acc = 0; m_acc < M_ACC; ++m_acc)
+          tilezero(tC[m_acc][n_acc]);
+
+      for (int k = 0; k < K; k += TILE_K) {
+        for (int n_acc = 0; n_acc < N_ACC; ++n_acc) {
+          tileload(tB, B_mem+k*N+n+n_acc*TILE_N, N*sizeof(type_t)*KPACK);
+          for (int m_acc = 0; m_acc < M_ACC; ++m_acc) {
+            if (n_acc == 0)
+              tileload(tA[m_acc], A_mem+(m + m_acc*TILE_M)*K+k, K*sizeof(type_t));
+            tdp(tC[m_acc][n_acc], tA[m_acc], tB);
+            if (k == K - TILE_K) {
+              int mc = m + m_acc*TILE_M, nc = n + n_acc*TILE_N;
+              tilestore(tC[m_acc][n_acc], C_mem+mc*N+nc, N*sizeof(res_type_t));
+            }
+          }
+        }
+      }
+    }
+  } 
+}
+
+void inner_product_ref(int *A_mem, int *B_mem, int *C_mem, int M, int N, int K) {
+  // tiling
+  for (int m = 0; m < M; m++)
+    for (int n = 0; n < N; n++) {
+      for (int k = 0; k < K; k++) {
+        char *va = (char *)(A_mem + m*K + k);
+        char *vb = (char *)(B_mem + k*N + n);
+        int acc = *(C_mem + m*N + n);
+        for (int i = 0; i < 4; i++) {
+          acc += (va[i] * vb[i]);
+        }
+        *(C_mem + m*N + n) = acc;
+      }
+    }
+}
+
+#define SIZE 128 
+int gA[SIZE][SIZE], gB[SIZE][SIZE], gC[SIZE][SIZE], gD[SIZE][SIZE];
+int main() {
+  for (int i = 0; i < SIZE; i++)
+    for (int j = 0; j < SIZE; j++) {
+      gA[i][j] = 0x01010101 * i * j;
+      gB[i][j] = 0x01010101 * i * j;
+    }
+
+  inner_product_ref((int *)gA, (int *)gB, (int *)gC, SIZE, SIZE, SIZE);
+  inner_product((int *)gA, (int *)gB, (int *)gD, SIZE, SIZE, SIZE);
+  int err_num = memcmp(gC, gD, sizeof(gC))?1:0;
+  inner_product2((int *)gA, (int *)gB, (int *)gD, SIZE, SIZE, SIZE);
+  err_num = memcmp(gC, gD, sizeof(gC))?(1+err_num):err_num;
+  if (err_num == 0)
+    std::cout << "PASSED\n";
+  else
+    std::cout << "FAILED\n";
+}
+

SingleSource/UnitTests/Matrix/AMXINT8/t_gemm_int8.reference_output

+exit 132

SingleSource/UnitTests/Matrix/CMakeLists.txt

+set(MATRIX_MAIN_DIR SingleSource/UnitTests/Matrix)
+
+if(CMAKE_C_COMPILER_ID STREQUAL "Clang")
+  if(ARCH STREQUAL "x86")
+    if(X86CPU_ARCH STREQUAL "skylake-avx512")
+      add_subdirectory(AMXINT8)
+    endif()
+  endif()
+endif()
+llvm_singlesource(PREFIX "Matirx-")

SingleSource/UnitTests/Matrix/Makefile

+# SingleSource/UnitTests/Matirx/Makefile
+LEVEL = ../../..
+
+include $(LEVEL)/Makefile.config
+
+DIRS = 
+
+ifeq ($(CC_UNDER_TEST_IS_CLANG), 1)
+# FIXME: Once testing infrastructure runs on sapphirerapids, we should change following line.
+# currenly we only test compilation results of AMX cases
+ifeq ($(HAVE_X86_AVX512DQ_INSTRUCTIONS), 1)
+DIRS += AMXINT8
+endif
+endif
+
+include $(LEVEL)/SingleSource/Makefile.singlesrc
+
+