|
| 1 | +/// \file libxsmm_gemm.cpp |
| 2 | +/// The ONLY translation unit that includes <libxsmm.h>. Isolating the libxsmm |
| 3 | +/// include here keeps its macros (libxsmm_macros.h) from leaking into any |
| 4 | +/// TiledArray header. When TILEDARRAY_HAS_LIBXSMM is undefined, libxsmm_gemm_le64 |
| 5 | +/// is still defined here as a `return false` stub so callers link unconditionally. |
| 6 | + |
| 7 | +#include "TiledArray/math/libxsmm_gemm.h" |
| 8 | + |
| 9 | +#ifdef TILEDARRAY_HAS_LIBXSMM |
| 10 | +#include <cstdlib> |
| 11 | +#include <cstring> |
| 12 | +#include <mutex> |
| 13 | +#include <libxsmm.h> |
| 14 | +#endif |
| 15 | + |
| 16 | +namespace TiledArray::detail { |
| 17 | + |
| 18 | +/// max(M,N,K) cutoff above which we keep using the vendor BLAS. |
| 19 | +static constexpr std::int64_t libxsmm_gemm_max_dim = 64; |
| 20 | + |
| 21 | +#ifdef TILEDARRAY_HAS_LIBXSMM |
| 22 | +/// Runtime master switch for the libxsmm fast path. Even in a libxsmm-enabled |
| 23 | +/// build, exporting `TA_LIBXSMM=0` (also accepts `off`/`OFF`/`false`/`no`) |
| 24 | +/// routes EVERY strided micro-GEMM back through the vendor BLAS path, i.e. |
| 25 | +/// libxsmm_gemm_le64() returns false for all shapes. Unset or any other value |
| 26 | +/// => libxsmm ON. Read from the environment once, on first use, and cached. |
| 27 | +static bool libxsmm_runtime_enabled() { |
| 28 | + static const bool enabled = [] { |
| 29 | + const char* v = std::getenv("TA_LIBXSMM"); |
| 30 | + if (v == nullptr || *v == '\0') return true; // default ON when compiled in |
| 31 | + return !(std::strcmp(v, "0") == 0 || std::strcmp(v, "off") == 0 || |
| 32 | + std::strcmp(v, "OFF") == 0 || std::strcmp(v, "false") == 0 || |
| 33 | + std::strcmp(v, "no") == 0); |
| 34 | + }(); |
| 35 | + return enabled; |
| 36 | +} |
| 37 | +#endif |
| 38 | + |
| 39 | +bool libxsmm_gemm_le64(bool trans_a, bool trans_b, std::int64_t m, |
| 40 | + std::int64_t n, std::int64_t k, double alpha, |
| 41 | + const double* a, std::int64_t lda, const double* b, |
| 42 | + std::int64_t ldb, double beta, double* c, |
| 43 | + std::int64_t ldc) { |
| 44 | +#ifdef TILEDARRAY_HAS_LIBXSMM |
| 45 | + // Runtime master switch: TA_LIBXSMM=0 sends everything back to vendor BLAS. |
| 46 | + if (!libxsmm_runtime_enabled()) return false; |
| 47 | + // libxsmm only for small shapes; max(M,N,K) <= 64. |
| 48 | + if (m > libxsmm_gemm_max_dim || n > libxsmm_gemm_max_dim || |
| 49 | + k > libxsmm_gemm_max_dim) |
| 50 | + return false; |
| 51 | + // libxsmm SMM has no alpha and only beta in {0,1} (LIBXSMM_GEMM_NO_BYPASS). |
| 52 | + if (alpha != 1.0) return false; |
| 53 | + if (beta != 0.0 && beta != 1.0) return false; |
| 54 | + |
| 55 | + static std::once_flag init_flag; |
| 56 | + std::call_once(init_flag, [] { |
| 57 | + // libxsmm's own verbose dispatch/JIT stats are part of the profiling |
| 58 | + // result, so fold them into TA_PROFILE: when profiling is on and the user |
| 59 | + // has NOT pinned LIBXSMM_VERBOSE explicitly, enable libxsmm verbosity here, |
| 60 | + // BEFORE libxsmm_init() parses the environment. TA_PROFILE>=1 -> a concise |
| 61 | + // exit summary (version + registry "gemm=<n>" kernel count); TA_PROFILE>=2 |
| 62 | + // -> verbose per-kernel JIT events. Must run before libxsmm_init(). |
| 63 | + if (std::getenv("LIBXSMM_VERBOSE") == nullptr) { |
| 64 | + const char* p = std::getenv("TA_PROFILE"); |
| 65 | + const int lvl = (p != nullptr) ? std::atoi(p) : 0; |
| 66 | + if (lvl >= 2) |
| 67 | + setenv("LIBXSMM_VERBOSE", "3", /*overwrite=*/0); |
| 68 | + else if (lvl >= 1) |
| 69 | + setenv("LIBXSMM_VERBOSE", "2", /*overwrite=*/0); |
| 70 | + } |
| 71 | + libxsmm_init(); |
| 72 | + }); |
| 73 | + |
| 74 | + // Mirror blas::gemm's row-major -> col-major mapping: it realizes the result |
| 75 | + // as a column-major GEMM (op_b, op_a, n, m, k) with operands (b, a) swapped. |
| 76 | + // libxsmm is column-major, so: A'=b (ld=ldb), B'=a (ld=lda), dims (n, m, k), |
| 77 | + // TRANS_A from op_b, TRANS_B from op_a. |
| 78 | + const libxsmm_bitfield flags = |
| 79 | + (trans_b ? LIBXSMM_GEMM_FLAG_TRANS_A : 0) | |
| 80 | + (trans_a ? LIBXSMM_GEMM_FLAG_TRANS_B : 0) | |
| 81 | + (beta == 0.0 ? LIBXSMM_GEMM_FLAG_BETA_0 : 0); |
| 82 | + |
| 83 | + const libxsmm_gemm_shape shape = libxsmm_create_gemm_shape( |
| 84 | + static_cast<libxsmm_blasint>(n), static_cast<libxsmm_blasint>(m), |
| 85 | + static_cast<libxsmm_blasint>(k), static_cast<libxsmm_blasint>(ldb), |
| 86 | + static_cast<libxsmm_blasint>(lda), static_cast<libxsmm_blasint>(ldc), |
| 87 | + LIBXSMM_DATATYPE_F64, LIBXSMM_DATATYPE_F64, LIBXSMM_DATATYPE_F64, |
| 88 | + LIBXSMM_DATATYPE_F64); |
| 89 | + |
| 90 | + const libxsmm_gemmfunction kernel = libxsmm_dispatch_gemm( |
| 91 | + shape, flags, static_cast<libxsmm_bitfield>(LIBXSMM_GEMM_PREFETCH_NONE)); |
| 92 | + if (kernel == nullptr) return false; // shape not JIT-able -> fall back |
| 93 | + |
| 94 | + libxsmm_gemm_param param; |
| 95 | + std::memset(¶m, 0, sizeof param); |
| 96 | + param.a.primary = const_cast<double*>(b); // A' = b |
| 97 | + param.b.primary = const_cast<double*>(a); // B' = a |
| 98 | + param.c.primary = c; |
| 99 | + kernel(¶m); |
| 100 | + return true; |
| 101 | +#else |
| 102 | + (void)trans_a; (void)trans_b; (void)m; (void)n; (void)k; (void)alpha; |
| 103 | + (void)a; (void)lda; (void)b; (void)ldb; (void)beta; (void)c; (void)ldc; |
| 104 | + return false; |
| 105 | +#endif |
| 106 | +} |
| 107 | + |
| 108 | +} // namespace TiledArray::detail |
0 commit comments