From 3a97a8a765bb8755d37d4e31a96ad8b259231999 Mon Sep 17 00:00:00 2001
From: Renato Golin <rengolin@systemcall.eu>
Date: Wed, 24 May 2023 12:47:19 +0100
Subject: [PATCH] Fix TODO to use libxsmm for VNNI check (NFC)

Does not change the semantics, just makes sure to use a libxsmm call
instead of macros.

Working towards #561
---
 lib/TPP/CMakeLists.txt |  1 +
 lib/TPP/VNNIUtils.cpp  | 16 +++++-----------
 2 files changed, 6 insertions(+), 11 deletions(-)

diff --git a/lib/TPP/CMakeLists.txt b/lib/TPP/CMakeLists.txt
index dc047ccde..0aa738d72 100644
--- a/lib/TPP/CMakeLists.txt
+++ b/lib/TPP/CMakeLists.txt
@@ -58,6 +58,7 @@ add_mlir_library(MLIRTPP
 
 target_include_directories(MLIRTPP
   PUBLIC
+    ${XSMM_INCLUDE_DIRS}
     $<BUILD_INTERFACE:${TPP_GEN_INCLUDE_DIR}>
     $<BUILD_INTERFACE:${TPP_MAIN_INCLUDE_DIR}>
 )
diff --git a/lib/TPP/VNNIUtils.cpp b/lib/TPP/VNNIUtils.cpp
index ad7c166fc..ed0e6ea91 100644
--- a/lib/TPP/VNNIUtils.cpp
+++ b/lib/TPP/VNNIUtils.cpp
@@ -11,23 +11,17 @@
 #include "mlir/IR/TypeUtilities.h"
 #include "mlir/IR/Types.h"
 
+#include "libxsmm.h"
+
 namespace mlir {
 namespace vnni {
 namespace utils {
 
 std::optional<int64_t> getVnniBlockingFactor(Type type) {
   auto elementType = getElementTypeOrSelf(type);
-  if (!elementType.isBF16())
-    return std::nullopt;
-  // @ TODO we should call LIBXSMM here to query the 
-  // VNNI packing factor we need to match to
-#if defined(__x86_64__)
-  return 2;
-#elif defined(__aarch64__)
-  return 4;
-#else
-#error Unsupported architecture
-#endif
+  if (elementType.isBF16())
+    return libxsmm_cpuid_dot_pack_factor(LIBXSMM_DATATYPE_BF16);
+  return std::nullopt;
 }
 
 bool isBF16Type(Type type) {