For launch patch, not merge

aten/src/ATen/native/CPUBlas.cpp

@@ -985,8 +985,10 @@ static inline dnnl::memory::data_type get_dnnl_dtype(ScalarType dtype) {
     return dnnl::memory::data_type::u8;
   } else if (dtype == ScalarType::Char) {
     return dnnl::memory::data_type::s8;
+  } else if (dtype == ScalarType::Float8_e4m3fn) {
+    return dnnl::memory::data_type::f8_e4m3;
   } else {
-    TORCH_CHECK(false, "get_dnnl_dtype expects float/bfloat16/half/int8 tensor input");
+    TORCH_CHECK(false, "get_dnnl_dtype expects float/bfloat16/half/int8/f8_e4m3 tensor input");
   }
 }
 
@@ -1172,6 +1174,9 @@ struct Brgemm : public KernelCache <BrgemmKey, GemmHelper> {
     } else if (dtype == ScalarType::Char) {
       static bool s8_support = dnnl::get_effective_cpu_isa() >= dnnl::cpu_isa::avx512_core_vnni;
       return s8_support;
+    } else if (dtype == ScalarType::Float8_e4m3fn) {
+      static bool f8_support = dnnl::get_effective_cpu_isa() >= dnnl::cpu_isa::avx512_core_amx;
+      return f8_support;
     }
     return false;
   }
@@ -1225,6 +1230,9 @@ struct Pack : public KernelCache <PackKey, pack_t> {
     } else if (dtype == ScalarType::Byte || dtype == ScalarType::Char) {
       static bool bit8_pack = dnnl::get_effective_cpu_isa() >= dnnl::cpu_isa::avx512_core_amx;
       return bit8_pack;
+    } else if (dtype == ScalarType::Float8_e4m3fn) {
+      static bool fp8_pack = dnnl::get_effective_cpu_isa() >= dnnl::cpu_isa::avx512_core_amx;
+      return fp8_pack;
     }
     return false;
   }
@@ -1398,6 +1406,31 @@ void brgemm(
     "I8 Brgemm is only supported on X64 when oneDNN ukernel is enabled and `amx` is supported");
 }
 
+void brgemm(
+    int64_t M,
+    int64_t N,
+    int64_t K,
+    int64_t ld_a,
+    int64_t ld_b,
+    int64_t ld_c,
+    const bool add_C,
+    const at::Float8_e4m3fn* A,
+    const at::Float8_e4m3fn* B,
+    float* C,
+    bool is_vnni) {
+#if defined(ONEDNN_UKERNEL_ENABLED)
+  if (is_vnni && Brgemm::device_check(ScalarType::Float8_e4m3fn)) {
+    Brgemm::call<at::Float8_e4m3fn, at::Float8_e4m3fn, float>(
+      M, N, K, ld_a, ld_b, ld_c, add_C, A, B, C);
+    return;
+  }
+#endif
+  // raise an error if the path is not supported
+  TORCH_CHECK(false,
+    "F8 Brgemm is only supported on X64 when oneDNN ukernel is enabled and `amx` is supported");
+}
+
+
 void brgemm_release(bool is_vnni) {
 #if defined(ONEDNN_UKERNEL_ENABLED)
   if (is_vnni) {

aten/src/ATen/native/CPUBlas.h

@@ -284,6 +284,21 @@ TORCH_API void brgemm(
     int32_t* C,
     bool is_vnni = true);
 
+TORCH_API void brgemm(
+    int64_t M,
+    int64_t N,
+    int64_t K,
+    int64_t ld_a,
+    int64_t ld_b,
+    int64_t ld_c,
+    const bool add_C,
+    const at::Float8_e4m3fn* A,
+    const at::Float8_e4m3fn* B,
+    float* C,
+    bool is_vnni = true);
+
+#define CPUBLAS_BRGEMM_F8F8F32
+
 // Release brgemm hardware context
 TORCH_API void brgemm_release(bool is_vnni = true);
 

setup.py

@@ -1184,6 +1184,7 @@ def make_relative_rpath_args(path: str) -> list[str]:
     entry_points = {
         "console_scripts": [
             "torchrun = torch.distributed.run:main",
+            "torch-xeon-launcher = torch.backends.xeon.run_cpu:main",
         ],
         "torchrun.logs_specs": [
             "default = torch.distributed.elastic.multiprocessing:DefaultLogsSpecs",

test/backends/xeon/test_launch.py

@@ -17,38 +17,57 @@ def tearDown(self):
         shutil.rmtree(self._test_dir)
 
     def test_cpu_info(self):
-        lscpu_info = """# The following is the parsable format, which can be fed to other
-# programs. Each different item in every column has an unique ID
-# starting from zero.
-# CPU,Core,Socket,Node
-0,0,0,0
-1,1,0,0
-2,2,0,0
-3,3,0,0
-4,4,1,1
-5,5,1,1
-6,6,1,1
-7,7,1,1
-8,0,0,0
-9,1,0,0
-10,2,0,0
-11,3,0,0
-12,4,1,1
-13,5,1,1
-14,6,1,1
-15,7,1,1
+        lscpu_info = """
+CPU NODE SOCKET CORE L1d:L1i:L2:L3 ONLINE    MAXMHZ   MINMHZ      MHZ
+  0    0      0    0 0:0:0:0          yes 5000.0000 800.0000 2400.000
+  1    0      0    1 0:0:0:0          yes 5000.0000 800.0000 2400.000
+  2    0      0    2 0:0:0:0          yes 5000.0000 800.0000 2400.000
+  3    0      0    3 0:0:0:0          yes 5000.0000 800.0000 2400.000
+  4    1      1    4 0:0:0:0          yes 5000.0000 800.0000 2400.000
+  5    1      1    5 0:0:0:0          yes 5000.0000 800.0000 2400.000
+  6    1      1    6 0:0:0:0          yes 5000.0000 800.0000 2400.000
+  7    1      1    7 0:0:0:0          yes 5000.0000 800.0000 2400.000
+  8    0      0    0 0:0:0:0          yes 5000.0000 800.0000 2400.000
+  9    0      0    1 0:0:0:0          yes 5000.0000 800.0000 2400.000
+ 10    0      0    2 0:0:0:0          yes 5000.0000 800.0000 2400.000
+ 11    0      0    3 0:0:0:0          yes 5000.0000 800.0000 2400.000
+ 12    1      1    4 0:0:0:0          yes 5000.0000 800.0000 2400.000
+ 13    1      1    5 0:0:0:0          yes 5000.0000 800.0000 2400.000
+ 14    1      1    6 0:0:0:0          yes 5000.0000 800.0000 2400.000
+ 15    1      1    7 0:0:0:0          yes 5000.0000 800.0000 2400.000
 """
-        from torch.backends.xeon.run_cpu import _CPUinfo
+        from torch.backends.xeon._cpu_info import CPUPoolList
 
-        cpuinfo = _CPUinfo(lscpu_info)
-        assert cpuinfo._physical_core_nums() == 8
-        assert cpuinfo._logical_core_nums() == 16
-        assert cpuinfo.get_node_physical_cores(0) == [0, 1, 2, 3]
-        assert cpuinfo.get_node_physical_cores(1) == [4, 5, 6, 7]
-        assert cpuinfo.get_node_logical_cores(0) == [0, 1, 2, 3, 8, 9, 10, 11]
-        assert cpuinfo.get_node_logical_cores(1) == [4, 5, 6, 7, 12, 13, 14, 15]
-        assert cpuinfo.get_all_physical_cores() == [0, 1, 2, 3, 4, 5, 6, 7]
-        assert cpuinfo.get_all_logical_cores() == [
+        cpupool = CPUPoolList(lscpu_txt=lscpu_info)
+        assert [c.cpu for c in cpupool.pool_all] == [
+            0,
+            1,
+            2,
+            3,
+            4,
+            5,
+            6,
+            7,
+            8,  # noqa: Q003
+            9,
+            10,
+            11,
+            12,
+            13,
+            14,
+            15,
+        ]
+        assert [c.cpu for c in cpupool.pool_all if c.is_physical_core] == [
+            0,
+            1,
+            2,
+            3,
+            4,
+            5,  # noqa: Q003
+            6,
+            7,
+        ]
+        assert [c.cpu for c in cpupool.pool_all if c.node == 0] == [
             0,
             1,
             2,
@@ -57,6 +76,8 @@ def test_cpu_info(self):
             9,
             10,
             11,
+        ]
+        assert [c.cpu for c in cpupool.pool_all if c.node == 1] == [
             4,
             5,
             6,
@@ -66,22 +87,40 @@ def test_cpu_info(self):
             14,
             15,
         ]
-        assert cpuinfo.numa_aware_check([0, 1, 2, 3]) == [0]
-        assert cpuinfo.numa_aware_check([4, 5, 6, 7]) == [1]
-        assert cpuinfo.numa_aware_check([2, 3, 4, 5]) == [0, 1]
+        assert [
+            c.cpu for c in cpupool.pool_all if c.node == 0 and c.is_physical_core
+        ] == [0, 1, 2, 3]
+        assert [
+            c.cpu for c in cpupool.pool_all if c.node == 1 and c.is_physical_core
+        ] == [4, 5, 6, 7]
+
+    def test_multi_threads_module(self):
+        num = 0
+        with subprocess.Popen(
+            f'python -m torch.backends.xeon.run_cpu --ninstances 4 --memory-allocator default \
+            --omp-runtime default --multi-task-manager none --log-dir {self._test_dir} --no-python echo "test"',
+            shell=True,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.STDOUT,
+        ) as p:
+            for line in p.stdout.readlines():
+                segs = str(line, "utf-8").strip().split(":")
+                if segs[-1].strip() == "test":
+                    num += 1
+        assert num == 4, "Failed to launch multiple instances for inference"
 
-    def test_multi_threads(self):
+    def test_multi_threads_command(self):
         num = 0
         with subprocess.Popen(
-            f"python -m torch.backends.xeon.run_cpu --ninstances 4 --use-default-allocator \
-            --disable-iomp --disable-numactl --disable-taskset --log-path {self._test_dir} --no-python pwd",
+            f'torch-xeon-launcher --ninstances 4 --memory-allocator default \
+            --omp-runtime default --multi-task-manager none --log-dir {self._test_dir} --no-python echo "test"',
             shell=True,
             stdout=subprocess.PIPE,
             stderr=subprocess.STDOUT,
         ) as p:
             for line in p.stdout.readlines():
-                segs = str(line, "utf-8").strip().split("-")
-                if segs[-1].strip() == "pwd":
+                segs = str(line, "utf-8").strip().split(":")
+                if segs[-1].strip() == "test":
                     num += 1
         assert num == 4, "Failed to launch multiple instances for inference"
 

third_party/mkl-dnn.BUILD

@@ -15,9 +15,10 @@ _DNNL_RUNTIME_OMP = {
     "#cmakedefine DNNL_ENABLE_STACK_CHECKER": "#undef DNNL_ENABLE_STACK_CHECKER",
     "#cmakedefine DNNL_EXPERIMENTAL_UKERNEL": "/* undef DNNL_EXPERIMENTAL_UKERNEL */",
     "#cmakedefine DNNL_EXPERIMENTAL": "#undef DNNL_EXPERIMENTAL",
-    "#cmakedefine DNNL_EXPERIMENTAL_SPARSE": "#undef DNNL_EXPERIMENTAL_SPARSE",
     "#cmakedefine ONEDNN_BUILD_GRAPH": "#undef ONEDNN_BUILD_GRAPH",
     "#cmakedefine DNNL_EXPERIMENTAL_PROFILING": "#undef DNNL_EXPERIMENTAL_PROFILING",
+    "#cmakedefine DNNL_EXPERIMENTAL_LOGGING": "#undef DNNL_EXPERIMENTAL_LOGGING",
+    "#cmakedefine DNNL_EXPERIMENTAL_SYCL_KERNEL_COMPILER": "#undef DNNL_EXPERIMENTAL_SYCL_KERNEL_COMPILER",
     "#cmakedefine DNNL_DISABLE_GPU_REF_KERNELS": "#undef DNNL_DISABLE_GPU_REF_KERNELS",
     "#cmakedefine01 BUILD_TRAINING": "#define BUILD_TRAINING 1",
     "#cmakedefine01 BUILD_INFERENCE": "#define BUILD_INFERENCE 0",
@@ -49,8 +50,6 @@ _DNNL_RUNTIME_OMP = {
     "#cmakedefine01 BUILD_AVX512": "#define BUILD_AVX512 0",
     "#cmakedefine01 BUILD_AMX": "#define BUILD_AMX 0",
     "#cmakedefine01 BUILD_PRIMITIVE_GPU_ISA_ALL": "#define BUILD_PRIMITIVE_GPU_ISA_ALL 1",
-    "#cmakedefine01 BUILD_GEN9": "#define BUILD_GEN9 0",
-    "#cmakedefine01 BUILD_GEN11": "#define BUILD_GEN11 0",
     "#cmakedefine01 BUILD_XELP": "#define BUILD_XELP 0",
     "#cmakedefine01 BUILD_XEHPG": "#define BUILD_XEHPG 0",
     "#cmakedefine01 BUILD_XEHPC": "#define BUILD_XEHPC 0",
@@ -70,8 +69,8 @@ template_rule(
     out = "include/oneapi/dnnl/dnnl_version.h",
     substitutions = {
         "@DNNL_VERSION_MAJOR@": "3",
-        "@DNNL_VERSION_MINOR@": "7",
-        "@DNNL_VERSION_PATCH@": "1",
+        "@DNNL_VERSION_MINOR@": "9",
+        "@DNNL_VERSION_PATCH@": "0",
     },
 )
 
@@ -86,7 +85,7 @@ template_rule(
     name = "include_dnnl_version_hash",
     src = "include/oneapi/dnnl/dnnl_version_hash.h.in",
     out = "include/oneapi/dnnl/dnnl_version_hash.h",
-    substitutions = {"@DNNL_VERSION_HASH@": "8d263e693366ef8db40acc569cc7d8edf644556d",}
+    substitutions = {"@DNNL_VERSION_HASH@": "5545f4170185277820039ccc7a540bc802b5f5d9",}
 )
 
 cc_library(
@@ -111,6 +110,7 @@ cc_library(
         "src/common/*.hpp",
         "src/common/**/**/*.h",
         "src/common/ittnotify/jitprofiling.h",
+        "third_party/**/*.h",
     ], exclude=[
         "src/cpu/aarch64/**/*.hpp",
         "src/cpu/aarch64/**/*.h",
@@ -141,7 +141,7 @@ cc_library(
         "src/",
         "src/common/",
         "src/cpu/",
-        "src/cpu/x64/xbyak/",
+        "third_party/",
     ],
     visibility = ["//visibility:public"],
     linkopts = [

torch/backends/__init__.py

@@ -138,4 +138,5 @@ def __init__(self, m, name):
     openmp as openmp,
     opt_einsum as opt_einsum,
     quantized as quantized,
+    xeon as xeon,
 )

torch/backends/xeon/__init__.py

@@ -0,0 +1 @@
+from . import run_cpu