Adds libxsmm support for micro-GEMMs

CMakeLists.txt

@@ -154,6 +154,10 @@ if (TA_STRIDED_DGEMM_COUNT)
   add_compile_definitions(TA_STRIDED_DGEMM_COUNT)
 endif()
 
+option(TA_LIBXSMM
+       "Fetch+build libxsmm from source and route small strided ToT micro-GEMMs (ce+e, ce+ce, scale) through its JIT fast path; falls back to vendor BLAS for shapes max(M,N,K)>64. Runtime toggle: TA_LIBXSMM=0"
+       OFF)
+
 option(TA_EXPERT "TiledArray Expert mode: disables automatically downloading or building dependencies" OFF)
 
 redefaultable_option(TA_WERROR "Treat compiler warnings as errors when compiling TiledArray's own translation units (does not propagate to consumers of installed TiledArray targets)" OFF)
@@ -366,6 +370,9 @@ include(${PROJECT_SOURCE_DIR}/cmake/modules/FindOrFetchBTAS.cmake)
 if(TA_SCALAPACK)
   include(external/scalapackpp.cmake)
 endif()
+if(TA_LIBXSMM)
+  include(external/libxsmm.cmake)
+endif()
 
 # other optional deps:
 # 2. TTG

external/libxsmm.cmake

@@ -0,0 +1,154 @@
+##
+## Fetch + build libxsmm from source, and expose it as the TiledArray_LIBXSMM
+## INTERFACE target. Enabled by -DTA_LIBXSMM=ON.
+##
+## libxsmm provides a JIT small-GEMM fast path for the strided tensor-of-tensors
+## micro-GEMMs (ce+e, ce+ce, scale). Unlike most TA deps, libxsmm's canonical
+## build is a GNU Makefile (not CMake), so this uses ExternalProject_Add with a
+## custom `make ... install` build command rather than CMAKE_ARGS.
+##
+## No system install is assumed: if libxsmm is not found via LIBXSMM_INSTALL_DIR
+## (an optional hint), it is cloned and built from source under
+## ${FETCHCONTENT_BASE_DIR}. There is intentionally NO TA_LIBXSMM_ROOT knob.
+##
+
+# Optional: reuse a pre-built libxsmm ONLY if the user explicitly points at one
+# via -DLIBXSMM_INSTALL_DIR=... We deliberately do NOT search default system
+# paths (no NO_DEFAULT_PATH omission): a stray /usr/local install must never be
+# picked up silently. Absent an explicit hint, libxsmm is always fetched+built.
+# clear any stale value cached by a prior configure (e.g. before this guard)
+unset(_LIBXSMM_INSTALL_DIR CACHE)
+set(_LIBXSMM_PREBUILT _LIBXSMM_PREBUILT-NOTFOUND)
+if (DEFINED LIBXSMM_INSTALL_DIR)
+  find_path(_LIBXSMM_PREBUILT NAMES include/libxsmm.h lib/libxsmm.a
+            HINTS ${LIBXSMM_INSTALL_DIR} NO_DEFAULT_PATH)
+endif ()
+
+if (_LIBXSMM_PREBUILT)
+
+    set(_LIBXSMM_INSTALL_DIR ${_LIBXSMM_PREBUILT})
+    message(STATUS "libxsmm found at ${_LIBXSMM_INSTALL_DIR}")
+
+elseif (TA_EXPERT)
+
+    message("** libxsmm was not found")
+    message(STATUS "** Downloading and building libxsmm is explicitly disabled in EXPERT mode")
+    message(FATAL_ERROR "** Either provide a pre-built libxsmm via -DLIBXSMM_INSTALL_DIR=... or disable -DTA_LIBXSMM=OFF")
+
+else ()
+
+    include(ExternalProject)
+
+    # libxsmm is a C library; make sure CMAKE_C_COMPILER is configured
+    enable_language(C)
+
+    set(EXTERNAL_SOURCE_DIR  ${FETCHCONTENT_BASE_DIR}/libxsmm-src)
+    set(_LIBXSMM_INSTALL_DIR ${FETCHCONTENT_BASE_DIR}/libxsmm-install)
+
+    if (NOT LIBXSMM_URL)
+        set(LIBXSMM_URL https://github.com/libxsmm/libxsmm.git)
+    endif (NOT LIBXSMM_URL)
+    if (NOT LIBXSMM_TAG)
+        set(LIBXSMM_TAG ${TA_TRACKED_LIBXSMM_TAG})
+    endif (NOT LIBXSMM_TAG)
+
+    message("** Will clone libxsmm from ${LIBXSMM_URL}")
+
+    # Compiler for libxsmm's sub-make. libxsmm builds with -target
+    # <arch>-apple-macos, which makes a bare CommandLineTools clang stop
+    # auto-injecting the macOS SDK sysroot, so it cannot find system headers
+    # (pthread.h) or libSystem at link time (and CMAKE_OSX_SYSROOT is often
+    # empty). On Apple, build libxsmm with the /usr/bin/{cc,c++} xcrun shims,
+    # which always resolve the active SDK for both compile and link; the
+    # resulting libxsmm.a is C-ABI-compatible with the rest of TiledArray.
+    # Elsewhere, honor the project's configured compilers.
+    if (APPLE)
+        set(_libxsmm_cc /usr/bin/cc)
+        set(_libxsmm_cxx /usr/bin/c++)
+    else ()
+        set(_libxsmm_cc ${CMAKE_C_COMPILER})
+        set(_libxsmm_cxx ${CMAKE_CXX_COMPILER})
+    endif ()
+
+    # parallelism for the libxsmm sub-make (overridable; default mirrors the
+    # project's typical build budget rather than hardcoding into the command)
+    set(LIBXSMM_BUILD_NJOBS 6 CACHE STRING "Parallel jobs for the libxsmm sub-make")
+
+    # libxsmm Make knobs:
+    #   STATIC=1   build libxsmm.a (we link the archive into tiledarray)
+    #   FORTRAN=0  skip the Fortran interface (no gfortran needed)
+    #   BLAS=0     do not wrap an external BLAS (we only use the JIT SMM path,
+    #              and TA already links its own BLAS); avoids a second BLAS dep
+    #   PREFIX=... install headers+lib into our private prefix
+    set(LIBXSMM_BUILD_BYPRODUCTS "${_LIBXSMM_INSTALL_DIR}/lib/libxsmm.a")
+    message(STATUS "custom target libxsmm is expected to build these byproducts: ${LIBXSMM_BUILD_BYPRODUCTS}")
+
+    ExternalProject_Add(libxsmm
+            PREFIX ${FETCHCONTENT_BASE_DIR}
+            STAMP_DIR ${FETCHCONTENT_BASE_DIR}/libxsmm-ep-artifacts
+            TMP_DIR ${FETCHCONTENT_BASE_DIR}/libxsmm-ep-artifacts  # in case CMAKE_INSTALL_PREFIX is not writable
+            #--Download step--------------
+            DOWNLOAD_DIR ${EXTERNAL_SOURCE_DIR}
+            GIT_REPOSITORY ${LIBXSMM_URL}
+            GIT_TAG ${LIBXSMM_TAG}
+            #--Configure step------------- (none: libxsmm uses a plain Makefile)
+            SOURCE_DIR ${EXTERNAL_SOURCE_DIR}
+            UPDATE_DISCONNECTED 1
+            BUILD_IN_SOURCE 1
+            CONFIGURE_COMMAND ""
+            #--Build step----------------- build + install in one make invocation
+            BUILD_COMMAND make -j${LIBXSMM_BUILD_NJOBS} STATIC=1 FORTRAN=0 BLAS=0
+                          CC=${_libxsmm_cc} CXX=${_libxsmm_cxx} AR=${CMAKE_AR}
+                          PREFIX=${_LIBXSMM_INSTALL_DIR} install
+            BUILD_BYPRODUCTS ${LIBXSMM_BUILD_BYPRODUCTS}
+            #--Install step--------------- (done by BUILD_COMMAND above)
+            INSTALL_COMMAND ${CMAKE_COMMAND} -E echo "libxsmm installed to ${_LIBXSMM_INSTALL_DIR}"
+            #--Custom targets-------------
+            STEP_TARGETS build
+            )
+
+    # the include dir must exist at configure time so the INTERFACE target's
+    # BUILD_INTERFACE include path validates (it is populated at build time)
+    execute_process(COMMAND ${CMAKE_COMMAND} -E make_directory "${_LIBXSMM_INSTALL_DIR}/include")
+
+    # build libxsmm before any TiledArray translation unit links
+    add_dependencies(External-tiledarray libxsmm)
+
+endif (_LIBXSMM_INSTALL_DIR)
+
+# Fold libxsmm's static lib + headers into TiledArray's OWN install prefix, so
+# the exported TiledArray config is self-contained and does not reference TA's
+# build tree (which a downstream find_package(TiledArray) consumer like MPQC
+# would otherwise link against -- and which breaks once the build tree is
+# wiped/relocated). Done for both the fetched and the prebuilt cases so the TA
+# install is identical either way.
+install(FILES "${_LIBXSMM_INSTALL_DIR}/lib/libxsmm.a"
+        DESTINATION "${TILEDARRAY_INSTALL_LIBDIR}" COMPONENT tiledarray)
+install(DIRECTORY "${_LIBXSMM_INSTALL_DIR}/include/"
+        DESTINATION "${TILEDARRAY_INSTALL_INCLUDEDIR}" COMPONENT tiledarray)
+
+# Synthetic target carrying the include dir, the static archive, and the gating
+# define. PUBLIC propagation (via _TILEDARRAY_DEPENDENCIES) makes the libxsmm.a
+# link requirement reach consumers (libtiledarray is a static archive, so its
+# undefined libxsmm symbols are resolved at the consumer's final link), plus
+# TILEDARRAY_HAS_LIBXSMM + the include path. The link/include paths are split
+# into BUILD_INTERFACE (TA's build tree) and INSTALL_INTERFACE (the installed
+# copy above), so the exported config never references the build tree. The
+# include INSTALL_INTERFACE is relative (CMake prepends the import prefix); the
+# link library INSTALL_INTERFACE must be an absolute path to the installed
+# archive -- CMake does NOT prepend the import prefix to INTERFACE_LINK_LIBRARIES
+# entries, so a relative path there would be resolved against the consumer's cwd
+# and fail to link. Absolute CMAKE_INSTALL_PREFIX is leak-free (the install tree
+# is the stable final location, unlike the build tree).
+add_library(TiledArray_LIBXSMM INTERFACE)
+set_target_properties(TiledArray_LIBXSMM
+        PROPERTIES
+        INTERFACE_INCLUDE_DIRECTORIES
+        "$<BUILD_INTERFACE:${_LIBXSMM_INSTALL_DIR}/include>;$<INSTALL_INTERFACE:${TILEDARRAY_INSTALL_INCLUDEDIR}>"
+        INTERFACE_LINK_LIBRARIES
+        "$<BUILD_INTERFACE:${_LIBXSMM_INSTALL_DIR}/lib/libxsmm.a>;$<INSTALL_INTERFACE:${CMAKE_INSTALL_PREFIX}/${TILEDARRAY_INSTALL_LIBDIR}/libxsmm.a>;${CMAKE_DL_LIBS}"
+        INTERFACE_COMPILE_DEFINITIONS
+        "TILEDARRAY_HAS_LIBXSMM"
+        )
+
+install(TARGETS TiledArray_LIBXSMM EXPORT tiledarray COMPONENT tiledarray)

external/versions.cmake

@@ -23,6 +23,14 @@ set(TA_TRACKED_BTAS_PREVIOUS_TAG 7e64fbad97c76f316f313f4c8ed3fca5445da15f)
 set(TA_TRACKED_LIBRETT_TAG 6eed30d4dd2a5aa58840fe895dcffd80be7fbece)
 set(TA_TRACKED_LIBRETT_PREVIOUS_TAG 354e0ccee54aeb2f191c3ce2c617ebf437e49d83)
 
+# libxsmm: pin a recent `main` commit, NOT the 1.17 release tag (2021). The
+# modern dispatch API used by TiledArray/math/libxsmm_gemm.h
+# (libxsmm_dispatch_gemm/libxsmm_create_gemm_shape/libxsmm_gemm_param) does not
+# exist in 1.17, and 1.17 predates most of libxsmm's AArch64/Apple-Silicon JIT
+# work. This SHA is the version validated on Apple M2 (reports as "1.17-3808").
+set(TA_TRACKED_LIBXSMM_TAG c14cbc6f8bc7964f8c5190a3a16b8cace03e5889)
+set(TA_TRACKED_LIBXSMM_PREVIOUS_TAG c14cbc6f8bc7964f8c5190a3a16b8cace03e5889)
+
 set(TA_TRACKED_UMPIRE-CXX-ALLOCATOR_TAG 1ba7f5f0aa99438826dd1c6bc1cd396080b9d608)
 set(TA_TRACKED_UMPIRE-CXX-ALLOCATOR_PREVIOUS_TAG 0f8144f19897766d0f117f7353221d4e3b8b1178)
 

src/CMakeLists.txt

@@ -135,6 +135,7 @@ TiledArray/external/madness.h
 TiledArray/host/env.h
 TiledArray/math/blas.h
 TiledArray/math/gemm_helper.h
+TiledArray/math/libxsmm_gemm.h
 TiledArray/math/outer.h
 TiledArray/math/parallel_gemm.h
 TiledArray/math/partial_reduce.h
@@ -215,6 +216,7 @@ set(TILEDARRAY_SOURCE_FILES
         TiledArray/einsum/index.cpp
         TiledArray/expressions/permopt.cpp
         TiledArray/host/env.cpp
+        TiledArray/math/libxsmm_gemm.cpp
         TiledArray/math/linalg/basic.cpp
         TiledArray/math/linalg/rank-local.cpp
         TiledArray/tensor/print.cpp
@@ -303,6 +305,9 @@ endif(TILEDARRAY_HAS_CUDA OR TILEDARRAY_HAS_HIP)
 if( TARGET TiledArray_SCALAPACK )
   list(APPEND _TILEDARRAY_DEPENDENCIES TiledArray_SCALAPACK)
 endif()
+if( TARGET TiledArray_LIBXSMM )
+  list(APPEND _TILEDARRAY_DEPENDENCIES TiledArray_LIBXSMM)
+endif()
 list(APPEND _TILEDARRAY_DEPENDENCIES "${LAPACK_LIBRARIES}")
 
 if( TARGET ttg-parsec )

src/TiledArray/math/libxsmm_gemm.cpp

@@ -0,0 +1,114 @@
+/// \file libxsmm_gemm.cpp
+/// The ONLY translation unit that includes <libxsmm.h>. Isolating the libxsmm
+/// include here keeps its macros (libxsmm_macros.h) from leaking into any
+/// TiledArray header. When TILEDARRAY_HAS_LIBXSMM is undefined, libxsmm_gemm_le64
+/// is still defined here as a `return false` stub so callers link unconditionally.
+
+#include "TiledArray/math/libxsmm_gemm.h"
+
+#ifdef TILEDARRAY_HAS_LIBXSMM
+#include <cstdlib>
+#include <cstring>
+#include <limits>
+#include <mutex>
+#include <libxsmm.h>
+#endif
+
+namespace TiledArray::detail {
+
+/// max(M,N,K) cutoff above which we keep using the vendor BLAS.
+static constexpr std::int64_t libxsmm_gemm_max_dim = 64;
+
+#ifdef TILEDARRAY_HAS_LIBXSMM
+/// Runtime master switch for the libxsmm fast path. Even in a libxsmm-enabled
+/// build, exporting `TA_LIBXSMM=0` (also accepts `off`/`OFF`/`false`/`no`)
+/// routes EVERY strided micro-GEMM back through the vendor BLAS path, i.e.
+/// libxsmm_gemm_le64() returns false for all shapes. Unset or any other value
+/// => libxsmm ON. Read from the environment once, on first use, and cached.
+static bool libxsmm_runtime_enabled() {
+  static const bool enabled = [] {
+    const char* v = std::getenv("TA_LIBXSMM");
+    if (v == nullptr || *v == '\0') return true;  // default ON when compiled in
+    return !(std::strcmp(v, "0") == 0 || std::strcmp(v, "off") == 0 ||
+             std::strcmp(v, "OFF") == 0 || std::strcmp(v, "false") == 0 ||
+             std::strcmp(v, "no") == 0);
+  }();
+  return enabled;
+}
+#endif
+
+bool libxsmm_gemm_le64(bool trans_a, bool trans_b, std::int64_t m,
+                       std::int64_t n, std::int64_t k, double alpha,
+                       const double* a, std::int64_t lda, const double* b,
+                       std::int64_t ldb, double beta, double* c,
+                       std::int64_t ldc) {
+#ifdef TILEDARRAY_HAS_LIBXSMM
+  // Runtime master switch: TA_LIBXSMM=0 sends everything back to vendor BLAS.
+  if (!libxsmm_runtime_enabled()) return false;
+  // libxsmm only for small shapes; max(M,N,K) <= 64.
+  if (m > libxsmm_gemm_max_dim || n > libxsmm_gemm_max_dim ||
+      k > libxsmm_gemm_max_dim)
+    return false;
+  // libxsmm SMM has no alpha and only beta in {0,1} (LIBXSMM_GEMM_NO_BYPASS).
+  if (alpha != 1.0) return false;
+  if (beta != 0.0 && beta != 1.0) return false;
+  // libxsmm_blasint is 32-bit; refuse leading dims that would narrow silently.
+  // (M,N,K are already <=64; lda/ldb/ldc are strides and unbounded in general.)
+  constexpr std::int64_t bi_max =
+      static_cast<std::int64_t>(std::numeric_limits<libxsmm_blasint>::max());
+  if (lda > bi_max || ldb > bi_max || ldc > bi_max) return false;
+
+  static std::once_flag init_flag;
+  std::call_once(init_flag, [] {
+    // libxsmm's own verbose dispatch/JIT stats are part of the profiling
+    // result, so fold them into TA_PROFILE: when profiling is on and the user
+    // has NOT pinned LIBXSMM_VERBOSE explicitly, enable libxsmm verbosity here,
+    // BEFORE libxsmm_init() parses the environment. TA_PROFILE>=1 -> a concise
+    // exit summary (version + registry "gemm=<n>" kernel count); TA_PROFILE>=2
+    // -> verbose per-kernel JIT events. Must run before libxsmm_init().
+    if (std::getenv("LIBXSMM_VERBOSE") == nullptr) {
+      const char* p = std::getenv("TA_PROFILE");
+      const int lvl = (p != nullptr) ? std::atoi(p) : 0;
+      if (lvl >= 2)
+        setenv("LIBXSMM_VERBOSE", "3", /*overwrite=*/0);
+      else if (lvl >= 1)
+        setenv("LIBXSMM_VERBOSE", "2", /*overwrite=*/0);
+    }
+    libxsmm_init();
+  });
+
+  // Mirror blas::gemm's row-major -> col-major mapping: it realizes the result
+  // as a column-major GEMM (op_b, op_a, n, m, k) with operands (b, a) swapped.
+  // libxsmm is column-major, so: A'=b (ld=ldb), B'=a (ld=lda), dims (n, m, k),
+  // TRANS_A from op_b, TRANS_B from op_a.
+  const libxsmm_bitfield flags =
+      (trans_b ? LIBXSMM_GEMM_FLAG_TRANS_A : 0) |
+      (trans_a ? LIBXSMM_GEMM_FLAG_TRANS_B : 0) |
+      (beta == 0.0 ? LIBXSMM_GEMM_FLAG_BETA_0 : 0);
+
+  const libxsmm_gemm_shape shape = libxsmm_create_gemm_shape(
+      static_cast<libxsmm_blasint>(n), static_cast<libxsmm_blasint>(m),
+      static_cast<libxsmm_blasint>(k), static_cast<libxsmm_blasint>(ldb),
+      static_cast<libxsmm_blasint>(lda), static_cast<libxsmm_blasint>(ldc),
+      LIBXSMM_DATATYPE_F64, LIBXSMM_DATATYPE_F64, LIBXSMM_DATATYPE_F64,
+      LIBXSMM_DATATYPE_F64);
+
+  const libxsmm_gemmfunction kernel = libxsmm_dispatch_gemm(
+      shape, flags, static_cast<libxsmm_bitfield>(LIBXSMM_GEMM_PREFETCH_NONE));
+  if (kernel == nullptr) return false;  // shape not JIT-able -> fall back
+
+  libxsmm_gemm_param param;
+  std::memset(&param, 0, sizeof param);
+  param.a.primary = const_cast<double*>(b);  // A' = b
+  param.b.primary = const_cast<double*>(a);  // B' = a
+  param.c.primary = c;
+  kernel(&param);
+  return true;
+#else
+  (void)trans_a; (void)trans_b; (void)m; (void)n; (void)k; (void)alpha;
+  (void)a; (void)lda; (void)b; (void)ldb; (void)beta; (void)c; (void)ldc;
+  return false;
+#endif
+}
+
+}  // namespace TiledArray::detail

src/TiledArray/math/libxsmm_gemm.h

@@ -0,0 +1,37 @@
+#ifndef TILEDARRAY_MATH_LIBXSMM_GEMM_H__INCLUDED
+#define TILEDARRAY_MATH_LIBXSMM_GEMM_H__INCLUDED
+
+/// \file libxsmm_gemm.h
+/// Optional libxsmm fast path for small strided tensor-of-tensors micro-GEMMs.
+/// This header is DECLARATION-ONLY: the implementation lives in libxsmm_gemm.cpp,
+/// which is the single translation unit that includes <libxsmm.h>. Keeping the
+/// libxsmm include out of this header is deliberate -- <libxsmm.h> pulls in
+/// libxsmm_macros.h, whose macros otherwise leak into TiledArray headers that
+/// transitively include this one (e.g. arena_einsum.h -> ... -> math/vector_op.h,
+/// breaking detail::is_scalar_v). Callers just see a plain function.
+
+#include <cstdint>
+
+// N.B. namespace TiledArray::detail (NOT TiledArray::math::detail): introducing
+// a TiledArray::math::detail namespace would hijack unqualified `detail::` name
+// lookup inside TiledArray::math headers (e.g. vector_op.h's detail::is_scalar_v,
+// which lives in TiledArray::detail), breaking their compilation.
+namespace TiledArray::detail {
+
+/// Computes C(m x n) [+]= alpha * op_a(A) . op_b(B) in **row-major** layout with
+/// leading dims lda/ldb/ldc, i.e. exactly TiledArray::math::blas::gemm (double)
+/// semantics. \p trans_a / \p trans_b are the transpose flags (true == Trans).
+///
+/// \return true iff libxsmm performed the GEMM. Returns false (caller must fall
+/// back to blas::gemm) when: built without libxsmm (TILEDARRAY_HAS_LIBXSMM
+/// undefined), the runtime switch `TA_LIBXSMM=0` is set, max(m,n,k) > 64,
+/// alpha != 1, beta not in {0,1}, or libxsmm could not JIT this shape.
+bool libxsmm_gemm_le64(bool trans_a, bool trans_b, std::int64_t m,
+                       std::int64_t n, std::int64_t k, double alpha,
+                       const double* a, std::int64_t lda, const double* b,
+                       std::int64_t ldb, double beta, double* c,
+                       std::int64_t ldc);
+
+}  // namespace TiledArray::detail
+
+#endif  // TILEDARRAY_MATH_LIBXSMM_GEMM_H__INCLUDED