Init amx for amx is enabled and override onednn isa flag (#1184) (#1227)

* init amx for amx is enabled and override onednn isa flag

* optimize code.

* avoid double amx initial.

* update code.

* add ut for ipex control onednn isa

* add more ut

* use ideep.hpp

* fix headers

* rename function onednn_isa_to_string

* skip isa override ut if not sync ipex isa to onednn

* optimze dnnl dependency & fix build issue.

* disable sync_isa_level_to_dnnl

* add env switch to enable dnnl isa level sync

* fix dyndisp UT issue.

* sync onednn isa level as default.

* fix no return issue.

* rename debug only env var.

* generalize manual setup onednn isa code.

* throw error when enable amx on non-linux platform and throw warninig for un-expected isa level

Co-authored-by: Han, Xu <xu.han@intel.com>

Co-authored-by: Han, Xu <xu.han@intel.com>

Author: zhuhaozhe

csrc/cpu/aten/utils/isa_help.cpp

@@ -1,8 +1,12 @@
 #include "isa_help.h"
 
+#include <dnnl.hpp>
+
 namespace torch_ipex {
 namespace cpu {
 
+using namespace dnnl;
+
 DEFINE_DISPATCH(get_current_isa_level_kernel_stub);
 
 // get_current_isa_level_kernel_impl
@@ -23,5 +27,31 @@ std::string get_highest_binary_support_isa_level() {
   return CPUCapabilityToString(level);
 }
 
+const char* OneDNNIsaLevelToString(cpu_isa isa) {
+  // convert dnnl::cpu_isa to string
+  switch (isa) {
+    case cpu_isa::avx2:
+      return "AVX2";
+    case cpu_isa::avx2_vnni:
+      return "AVX2_VNNI";
+    case cpu_isa::avx512_core:
+      return "AVX512";
+    case cpu_isa::avx512_core_vnni:
+      return "AVX512_VNNI";
+    case cpu_isa::avx512_core_bf16:
+      return "AVX512_BF16";
+    case cpu_isa::avx512_core_amx:
+      return "AMX";
+
+    default:
+      return "WrongLevel";
+  }
+}
+
+std::string get_current_onednn_isa_level() {
+  cpu_isa onednn_isa_level = get_effective_cpu_isa();
+  return OneDNNIsaLevelToString(onednn_isa_level);
+}
+
 } // namespace cpu
 } // namespace torch_ipex

csrc/cpu/aten/utils/isa_help.h

@@ -8,6 +8,8 @@
 namespace torch_ipex {
 namespace cpu {
 
+std::string get_current_onednn_isa_level();
+
 std::string get_current_isa_level();
 std::string get_highest_cpu_support_isa_level();
 std::string get_highest_binary_support_isa_level();

csrc/cpu/dyndisp/DispatchStub.cpp

@@ -4,13 +4,17 @@
 
 #include "../cpu/isa/cpu_feature.hpp"
 
+#include <dnnl.hpp>
+
 #include <algorithm>
 #include <cstdlib>
 #include <cstring>
 
 namespace torch_ipex {
 namespace cpu {
 
+using namespace dnnl;
+
 const char* CPUCapabilityToString(CPUCapability isa) {
   switch (isa) {
     case CPUCapability::DEFAULT:
@@ -79,6 +83,62 @@ CPUCapability _get_highest_binary_support_isa_level() {
   return CPUCapability::DEFAULT;
 }
 
+static bool _load_sync_onednn_isa_setting() {
+  // _IPEX_NOT_SYNC_ONEDNN_ISA is debug only env.
+  auto envar = std::getenv("_IPEX_NOT_SYNC_ONEDNN_ISA");
+  if (envar) {
+    if (strcmp(envar, "1") == 0) {
+      return true;
+    }
+  }
+  return false;
+}
+
+bool check_not_sync_onednn_isa_level() {
+  static bool b_not_sync = _load_sync_onednn_isa_setting();
+
+  return b_not_sync;
+}
+
+status set_current_cpu_isa_level_to_onednn(cpu_isa isa) {
+  if (check_not_sync_onednn_isa_level()) {
+    return status::unimplemented;
+  }
+
+  return set_max_cpu_isa(isa);
+}
+
+cpu_isa ipex_isa_to_onednn_isa(CPUCapability ipex_isa) {
+  switch (ipex_isa) {
+    case CPUCapability::DEFAULT:
+      return cpu_isa::all;
+    case CPUCapability::AVX2:
+      return cpu_isa::avx2;
+    case CPUCapability::AVX2_VNNI:
+      return cpu_isa::avx2_vnni;
+    case CPUCapability::AVX512:
+      return cpu_isa::avx512_core;
+    case CPUCapability::AVX512_VNNI:
+      return cpu_isa::avx512_core_vnni;
+    case CPUCapability::AVX512_BF16:
+      return cpu_isa::avx512_core_bf16;
+    case CPUCapability::AMX:
+      return cpu_isa::avx512_core_amx;
+    case CPUCapability::NUM_OPTIONS:
+      TORCH_WARN("DispatchStub: OutOfBoundaryISALevel for IPEX");
+      return cpu_isa::all;
+
+    default:
+      auto ipex_isa_str = CPUCapabilityToString(ipex_isa);
+      auto msg = c10::str(
+          "DispatchStub: No corresponding onednn isa for ",
+          ipex_isa_str,
+          "Please consider check whether this ISA is supported by oneDNN");
+      TORCH_WARN(msg);
+      return cpu_isa::all;
+  }
+}
+
 static CPUCapability compute_cpu_capability() {
   CPUCapability highest_cpu_supported_isa_level =
       _get_highest_cpu_support_isa_level();
@@ -119,6 +179,9 @@ static CPUCapability compute_cpu_capability() {
   CPUCapability max_support_isa_level = std::min(
       highest_cpu_supported_isa_level, highest_binary_supported_isa_level);
   if (b_manual_setup) {
+    cpu_isa manual_onednn_isa = ipex_isa_to_onednn_isa(manual_setup_isa_level);
+    set_current_cpu_isa_level_to_onednn(manual_onednn_isa);
+
     if (manual_setup_isa_level <= max_support_isa_level) {
       return manual_setup_isa_level;
     }

csrc/cpu/dyndisp/DispatchStub.h

@@ -63,6 +63,8 @@ const char* CPUCapabilityToString(CPUCapability isa);
 CPUCapability _get_highest_cpu_support_isa_level();
 CPUCapability _get_highest_binary_support_isa_level();
 
+bool check_not_sync_onednn_isa_level();
+
 CPUCapability get_cpu_capability();
 
 template <typename FnPtr, typename T>

csrc/cpu/isa/cpu_feature.cpp

@@ -2,6 +2,13 @@
 #include <stdio.h>
 #include "embedded_function.h"
 
+#ifdef __linux__
+#include <sys/syscall.h>
+#include <unistd.h>
+#endif
+
+#include <ATen/ATen.h>
+
 namespace torch_ipex {
 namespace cpu {
 CPUFeature::CPUFeature() {
@@ -277,6 +284,44 @@ bool CPUFeature::os_amx() {
   return false;
 }
 
+#ifdef __linux__
+
+#define XFEATURE_XTILECFG 17
+#define XFEATURE_XTILEDATA 18
+#define XFEATURE_MASK_XTILECFG (1 << XFEATURE_XTILECFG)
+#define XFEATURE_MASK_XTILEDATA (1 << XFEATURE_XTILEDATA)
+#define XFEATURE_MASK_XTILE (XFEATURE_MASK_XTILECFG | XFEATURE_MASK_XTILEDATA)
+#define ARCH_GET_XCOMP_PERM 0x1022
+#define ARCH_REQ_XCOMP_PERM 0x1023
+
+bool CPUFeature::init_amx() {
+  unsigned long bitmask = 0;
+  long status = syscall(SYS_arch_prctl, ARCH_GET_XCOMP_PERM, &bitmask);
+  if (0 != status)
+    return false;
+  if (bitmask & XFEATURE_MASK_XTILEDATA)
+    return true;
+
+  status = syscall(SYS_arch_prctl, ARCH_REQ_XCOMP_PERM, XFEATURE_XTILEDATA);
+  if (0 != status)
+    return false; // XFEATURE_XTILEDATA setup is failed, TMUL usage is not
+                  // allowed
+  status = syscall(SYS_arch_prctl, ARCH_GET_XCOMP_PERM, &bitmask);
+
+  // XFEATURE_XTILEDATA setup is failed, can't use TMUL
+  if (0 != status || !(bitmask & XFEATURE_MASK_XTILEDATA))
+    return false;
+
+  // XFEATURE_XTILEDATA set successfully, TMUL usage is allowed
+  return true;
+}
+#else
+bool CPUFeature::init_amx() {
+  AT_ERROR("DispatchStub: only support init amx on Linux now");
+  return false;
+}
+#endif
+
 bool CPUFeature::isa_level_avx2() {
   static bool b_is_support = os_avx2() && cpuid_avx2() && cpuid_fma();
   return b_is_support;
@@ -304,9 +349,19 @@ bool CPUFeature::isa_level_avx512_bf16() {
   return b_is_support;
 }
 
+bool CPUFeature::_do_check_and_init_amx() {
+  bool b_is_support = isa_level_avx512_bf16() && os_amx() && cpuid_amx_bf16() &&
+      cpuid_amx_int8() && cpuid_amx_tile();
+  if (b_is_support) {
+    b_is_support = init_amx();
+  }
+  return b_is_support;
+}
+
 bool CPUFeature::isa_level_amx() {
-  static bool b_is_support = isa_level_avx512_bf16() && os_amx() &&
-      cpuid_amx_bf16() && cpuid_amx_int8() && cpuid_amx_tile();
+  // check and init in a funtion, avoid to double init.
+  static bool b_is_support = _do_check_and_init_amx();
+
   return b_is_support;
 }
 

csrc/cpu/isa/cpu_feature.hpp

@@ -113,6 +113,8 @@ class CPUFeature {
   MICRO_CLASS_MEMBER_DECL(amx_bf16);
   MICRO_CLASS_MEMBER_DECL(amx_tile);
   MICRO_CLASS_MEMBER_DECL(amx_int8);
+  bool init_amx();
+  bool _do_check_and_init_amx();
 
  public:
   MICRO_CLASS_CHECK_FUNC(amx_bf16);

intel_extension_for_pytorch/csrc/python/cpu/init_python_bindings.cpp

@@ -61,6 +61,16 @@ void InitIpexModuleBindings(py::module m) {
     return get_current_isa_level();
   });
 
+  m.def("_get_current_onednn_isa_level", []() {
+    using namespace torch_ipex::cpu;
+    return get_current_onednn_isa_level();
+  });
+
+  m.def("_check_not_sync_onednn_isa_level", []() {
+    using namespace torch_ipex::cpu;
+    return check_not_sync_onednn_isa_level();
+  });
+
   m.def("_get_highest_cpu_support_isa_level", []() {
     using namespace torch_ipex::cpu;
     return get_highest_cpu_support_isa_level();

tests/cpu/test_dyndisp.py

@@ -1,5 +1,6 @@
 import unittest
 import os
+import subprocess
 
 import intel_extension_for_pytorch._C as core
 
@@ -36,6 +37,9 @@ def get_highest_binary_support_isa_level():
 def get_highest_cpu_support_isa_level():
     return core._get_highest_cpu_support_isa_level().lower()
 
+def check_not_sync_onednn_isa_level():
+    return core._check_not_sync_onednn_isa_level()
+
 class TestDynDisp(unittest.TestCase):
 
     def test_manual_select_kernel(self):
@@ -69,5 +73,22 @@ def test_dyndisp_in_supported_set(self):
         self.assertTrue(expected_isa)
         return        
 
+    @unittest.skipIf(check_not_sync_onednn_isa_level(), 'skip this if not sync onednn isa level')
+    def test_ipex_set_onednn_isa_level(self):
+        command = 'ATEN_CPU_CAPABILITY=avx2 python -c "import torch; import intel_extension_for_pytorch._C as core; print(core._get_current_onednn_isa_level())" '
+        with subprocess.Popen(command, shell=True, stdout=subprocess.PIPE, stderr=subprocess.STDOUT) as p:
+          out = p.stdout.readlines()
+          onednn_isa_level = str(out[-1], 'utf-8').strip()
+          self.assertTrue(onednn_isa_level == 'AVX2')
+
+    @unittest.skipIf(check_not_sync_onednn_isa_level(), 'skip this if not sync onednn isa level')
+    def test_onednn_do_not_set_isa_level(self):
+        command = 'ONEDNN_MAX_CPU_ISA=avx2 python -c "import torch; import intel_extension_for_pytorch._C as core; print(core._get_current_isa_level().lower())" '
+        cur_ipex_isa = get_currnet_isa_level()
+        with subprocess.Popen(command, shell=True, stdout=subprocess.PIPE, stderr=subprocess.STDOUT) as p:
+          out = p.stdout.readlines()
+          cur_ipex_isa_1 = str(out[-1], 'utf-8').strip()
+          self.assertTrue(cur_ipex_isa == cur_ipex_isa_1)
+
 if __name__ == '__main__':
     unittest.main()