[Re-landing 68111] Add JIT graph fuser for oneDNN Graph API (Preview4.1)

aten/src/ATen/core/interned_strings.h

@@ -43,6 +43,8 @@ namespace c10 {
   _(prim, FusionGroup)               \
   _(prim, CudaFusionGroup)           \
   _(prim, CudaFusionGuard)           \
+  _(prim, oneDNNFusionGroup)         \
+  _(prim, oneDNNFusionGuard)         \
   _(prim, FunctionalGraph)           \
   _(prim, add_optional)              \
   _(prim, view_copy)                 \
@@ -314,6 +316,7 @@ namespace c10 {
   _(attr, cache_id)                  \
   _(attr, new_axis)                  \
   _(attr, warn_id)                   \
+  _(attr, output_layouts)            \
   _(attr, allowzero)                 \
   _(attr, seen_none)
 

caffe2/CMakeLists.txt

@@ -654,6 +654,22 @@ if(NOT INTERN_BUILD_MOBILE OR NOT BUILD_CAFFE2_MOBILE)
     include(../mlc/mlc_build.cmake)
   endif()
 
+  if(MKLDNN_FOUND AND NOT WIN32 AND NOT APPLE AND NOT BUILD_LITE_INTERPRETER AND NOT INTERN_BUILD_MOBILE)
+    list(APPEND Caffe2_CPU_SRCS
+      ${TORCH_SRC_DIR}/csrc/jit/codegen/onednn/LlgaTensorImpl.cpp
+      ${TORCH_SRC_DIR}/csrc/jit/codegen/onednn/graph_fuser.cpp
+      ${TORCH_SRC_DIR}/csrc/jit/codegen/onednn/graph_rewriter.cpp
+      ${TORCH_SRC_DIR}/csrc/jit/codegen/onednn/graph_helper.cpp
+      ${TORCH_SRC_DIR}/csrc/jit/codegen/onednn/register_interface.cpp
+      ${TORCH_SRC_DIR}/csrc/jit/codegen/onednn/interface.cpp
+      ${TORCH_SRC_DIR}/csrc/jit/codegen/onednn/kernel.cpp
+      ${TORCH_SRC_DIR}/csrc/jit/codegen/onednn/defer_size_check.cpp
+      ${TORCH_SRC_DIR}/csrc/jit/codegen/onednn/layout_propagation.cpp
+      ${TORCH_SRC_DIR}/csrc/jit/codegen/onednn/prepare_binary.cpp
+      ${TORCH_SRC_DIR}/csrc/jit/codegen/onednn/guard_shape.cpp
+    )
+  endif()
+
   if(USE_ROCM)
     list(APPEND Caffe2_HIP_SRCS ${Caffe2_GPU_HIP_JIT_FUSERS_SRCS})
     if(USE_NCCL)

cmake/Dependencies.cmake

@@ -1772,6 +1772,9 @@ if(NOT INTERN_BUILD_MOBILE)
       if(BUILD_CAFFE2_OPS)
         set(CAFFE2_USE_MKLDNN ON)
         list(APPEND Caffe2_PUBLIC_DEPENDENCY_LIBS caffe2::mkldnn)
+        if(NOT WIN32 AND NOT APPLE)
+          list(APPEND Caffe2_PUBLIC_DEPENDENCY_LIBS caffe2::dnnl_graph)
+        endif()
       endif(BUILD_CAFFE2_OPS)
     else()
       message(WARNING "MKLDNN could not be found.")

cmake/Modules/FindMKLDNN.cmake

@@ -18,21 +18,30 @@ SET(MKLDNN_LIBRARIES)
 SET(MKLDNN_INCLUDE_DIR)
 
 SET(IDEEP_ROOT "${PROJECT_SOURCE_DIR}/third_party/ideep")
-SET(MKLDNN_ROOT "${IDEEP_ROOT}/mkl-dnn/third_party/oneDNN")
+SET(MKLDNN_ROOT "${PROJECT_SOURCE_DIR}/third_party/ideep/mkl-dnn/third_party/oneDNN")
+IF (NOT APPLE AND NOT WIN32)
+  SET(LLGA_ROOT "${PROJECT_SOURCE_DIR}/third_party/ideep/mkl-dnn")
+ENDIF(NOT APPLE AND NOT WIN32)
 
 FIND_PACKAGE(BLAS)
 FIND_PATH(IDEEP_INCLUDE_DIR ideep.hpp PATHS ${IDEEP_ROOT} PATH_SUFFIXES include)
+IF (NOT APPLE AND NOT WIN32)
+  FIND_PATH(LLGA_INCLUDE_DIR oneapi/dnnl/dnnl_graph.hpp PATHS ${LLGA_ROOT} PATH_SUFFIXES include)
+ENDIF(NOT APPLE AND NOT WIN32)
 FIND_PATH(MKLDNN_INCLUDE_DIR dnnl.hpp dnnl.h PATHS ${MKLDNN_ROOT} PATH_SUFFIXES include)
 IF (NOT MKLDNN_INCLUDE_DIR)
   EXECUTE_PROCESS(COMMAND git${CMAKE_EXECUTABLE_SUFFIX} submodule update --init --jobs 0 mkl-dnn WORKING_DIRECTORY ${IDEEP_ROOT})
-  FIND_PATH(MKLDNN_INCLUDE_DIR mkldnn.hpp mkldnn.h PATHS ${MKLDNN_ROOT} PATH_SUFFIXES include)
+  FIND_PATH(MKLDNN_INCLUDE_DIR dnnl.hpp dnnl.h PATHS ${MKLDNN_ROOT} PATH_SUFFIXES include)
 ENDIF(NOT MKLDNN_INCLUDE_DIR)
 
 IF (NOT IDEEP_INCLUDE_DIR OR NOT MKLDNN_INCLUDE_DIR)
   MESSAGE(STATUS "MKLDNN source files not found!")
   RETURN()
 ENDIF(NOT IDEEP_INCLUDE_DIR OR NOT MKLDNN_INCLUDE_DIR)
 LIST(APPEND MKLDNN_INCLUDE_DIR ${IDEEP_INCLUDE_DIR})
+IF (NOT APPLE AND NOT WIN32)
+  LIST(APPEND MKLDNN_INCLUDE_DIR ${LLGA_INCLUDE_DIR})
+ENDIF(NOT APPLE AND NOT WIN32)
 IF(MKL_FOUND)
   ADD_DEFINITIONS(-DIDEEP_USE_MKL)
   # Append to mkldnn dependencies
@@ -66,6 +75,7 @@ SET(DNNL_BUILD_TESTS FALSE CACHE BOOL "" FORCE)
 SET(DNNL_BUILD_EXAMPLES FALSE CACHE BOOL "" FORCE)
 SET(DNNL_LIBRARY_TYPE STATIC CACHE STRING "" FORCE)
 SET(DNNL_ENABLE_PRIMITIVE_CACHE TRUE CACHE BOOL "" FORCE)
+set(DNNL_GRAPH_LIBRARY_TYPE STATIC CACHE STRING "" FORCE)
 IF(MKLDNN_USE_NATIVE_ARCH)  # Disable HostOpts in MKLDNN unless MKLDNN_USE_NATIVE_ARCH is set.
   SET(DNNL_ARCH_OPT_FLAGS "HostOpts" CACHE STRING "" FORCE)
 ELSE()
@@ -78,7 +88,20 @@ ELSE()
   ENDIF()
 ENDIF()
 
-ADD_SUBDIRECTORY(${MKLDNN_ROOT})
+IF (NOT APPLE AND NOT WIN32)
+  ADD_SUBDIRECTORY(${LLGA_ROOT})
+  IF(NOT TARGET dnnl_graph)
+    MESSAGE("Failed to include LLGA target")
+    RETURN()
+  ENDIF(NOT TARGET dnnl_graph)
+
+  IF(CMAKE_COMPILER_IS_GNUCC)
+    TARGET_COMPILE_OPTIONS(dnnl_graph PRIVATE -Wno-maybe-uninitialized)
+    TARGET_COMPILE_OPTIONS(dnnl_graph PRIVATE -Wno-strict-overflow)
+    TARGET_COMPILE_OPTIONS(dnnl_graph PRIVATE -Wno-error=strict-overflow)
+  ENDIF(CMAKE_COMPILER_IS_GNUCC)
+ENDIF(NOT APPLE AND NOT WIN32)
+
 IF(NOT TARGET dnnl)
   MESSAGE("Failed to include MKL-DNN target")
   RETURN()
@@ -89,6 +112,10 @@ IF(NOT APPLE AND CMAKE_COMPILER_IS_GNUCC)
   TARGET_COMPILE_OPTIONS(dnnl PRIVATE -Wno-strict-overflow)
   TARGET_COMPILE_OPTIONS(dnnl PRIVATE -Wno-error=strict-overflow)
 ENDIF(NOT APPLE AND CMAKE_COMPILER_IS_GNUCC)
+LIST(APPEND MKLDNN_LIBRARIES ${MKL_OPENMP_LIBRARY})
+IF (NOT APPLE AND NOT WIN32)
+  LIST(APPEND MKLDNN_LIBRARIES "$<TARGET_FILE:dnnl_graph>")
+ENDIF(NOT APPLE AND NOT WIN32)
 LIST(APPEND MKLDNN_LIBRARIES dnnl)
 
 SET(MKLDNN_FOUND TRUE)

cmake/public/mkldnn.cmake

@@ -16,3 +16,15 @@ set_property(
 set_property(
   TARGET caffe2::mkldnn PROPERTY INTERFACE_LINK_LIBRARIES
   ${MKLDNN_LIBRARIES})
+if(NOT APPLE AND NOT WIN32)
+  if(NOT TARGET caffe2::dnnl_graph)
+    add_library(caffe2::dnnl_graph INTERFACE IMPORTED)
+  endif()
+
+  set_property(
+    TARGET caffe2::dnnl_graph PROPERTY INTERFACE_INCLUDE_DIRECTORIES
+    ${MKLDNN_INCLUDE_DIR})
+  set_property(
+    TARGET caffe2::dnnl_graph PROPERTY INTERFACE_LINK_LIBRARIES
+    ${MKLDNN_LIBRARIES})
+endif()

docs/source/jit.rst

@@ -61,6 +61,8 @@ Creating TorchScript Code
     ScriptFunction
     freeze
     optimize_for_inference
+    enable_onednn_fusion
+    onednn_fusion_enabled
     set_fusion_strategy
     save
     load