[AMD-AIE] Build microkernels for AMD-AIE (#1109)

Signed-off-by: Abhishek Varma <[email protected]>
Co-authored-by: erwei-xilinx <[email protected]>

CMakeLists.txt

@@ -243,6 +243,7 @@ endif()
 add_subdirectory(include)
 add_subdirectory(lib)
 add_subdirectory(aie_runtime_lib)
+add_subdirectory(aie_kernels)
 add_subdirectory(tools)
 add_subdirectory(runtime_lib)
 

aie_kernels/CMakeLists.txt

@@ -0,0 +1,30 @@
+#
+# This file is licensed under the Apache License v2.0 with LLVM Exceptions.
+# See https://llvm.org/LICENSE.txt for license information.
+# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+#
+# (c) Copyright 2024 Xilinx Inc.
+
+# Stuff into the build area:
+add_custom_target(aie-kernels ALL)
+
+function(add_aie_kernels arch)
+  if(DEFINED VITIS_ROOT)
+    if(EXISTS "${PROJECT_SOURCE_DIR}/tools/chess-clang/xchesscc_wrapper")
+        add_custom_target(${arch}_kernels ALL DEPENDS ${CMAKE_CURRENT_BINARY_DIR}/mm.o)
+        add_custom_command(OUTPUT ${CMAKE_CURRENT_BINARY_DIR}/mm.o
+                      COMMAND ${PROJECT_SOURCE_DIR}/tools/chess-clang/xchesscc_wrapper ${arch} -I ${VITIS_ROOT}/aietools/include -c ${CMAKE_CURRENT_SOURCE_DIR}/mm.cc -o ${CMAKE_CURRENT_BINARY_DIR}/mm.o
+                      DEPENDS ${CMAKE_CURRENT_SOURCE_DIR}/mm.cc)
+        add_dependencies(aie-kernels ${arch}_kernels)
+
+        install(FILES ${CMAKE_CURRENT_BINARY_DIR}/mm.o DESTINATION ${CMAKE_INSTALL_PREFIX}/aie_kernels/)
+    else()
+      message(STATUS "Could not find xchesscc_wrapper inside mlir-aie's build directory")
+    endif()
+  else()
+      message(STATUS "Tool 'xchesscc' not found. Skipping compilation of microkernel .o files.")
+  endif()
+
+endfunction()
+
+add_aie_kernels(AIE2)

aie_kernels/mm.cc

@@ -0,0 +1,62 @@
+//===- mm.cc ----------------------------------------------000---*- C++ -*-===//
+//
+//  THIS IS WIP AND HAS BEEN TAKEN AND MODIFIED FROM MLIR-AIE.
+//
+// This file is licensed under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+// Copyright (C) 2024, Advanced Micro Devices, Inc.
+//
+//===----------------------------------------------------------------------===//
+
+#define __AIENGINE__ 2
+#define NOCPP
+#define __AIEARCH__ 20
+
+#include <stdint.h>
+#include <stdio.h>
+#include <stdlib.h>
+#include <type_traits>
+
+#define REL_WRITE 0
+#define REL_READ 1
+
+#include <aie_api/aie.hpp>
+
+template <typename T_in, typename T_out, int M, int K, int N>
+void matmul_scalar(T_in *a, unsigned offsetA, T_in *b, unsigned offsetB,
+                   T_out *c, unsigned offsetC) {
+  event0();
+  for (int row = 0; row < M; row++) {
+    for (int col = 0; col < N; col++) {
+      T_out running_sum = 0;
+      for (int i = 0; i < K; i++) {
+        running_sum += a[offsetA + row * K + i] * b[offsetB + i * N + col];
+      }
+      c[offsetC + row * N + col] += running_sum;
+    }
+  }
+  event1();
+}
+
+extern "C" {
+
+#define combos(X)                                                              \
+  X(int16, i16, int16, i16, 4, 4, 4)                                           \
+  X(int32, i32, int32, i32, 4, 4, 4)                                           \
+  X(bfloat16, bf16, bfloat16, bf16, 4, 8, 4)                                   \
+  X(bfloat16, bf16, float, f32, 4, 8, 4)
+
+#define matmul_scalar_c_func(ctype_in, mlir_type_in, ctype_out, mlir_type_out, \
+                             r, s, t)                                          \
+  void matmul_scalar_##mlir_type_in##_##mlir_type_out(                         \
+      ctype_in *a_in, unsigned offsetA, ctype_in *b_in, unsigned offsetB,      \
+      ctype_out *c_out, unsigned offsetC) {                                    \
+    matmul_scalar<ctype_in, ctype_out, 4, 4, 4>(a_in, offsetA, b_in, offsetB,  \
+                                                c_out, offsetC);               \
+  }
+
+combos(matmul_scalar_c_func)
+
+} // extern "C"