[BUILD] When set, use MLIR_DIR for finding both MLIR and LLVM (#755)

CMakeLists.txt

@@ -49,7 +49,8 @@ endif()
 ##########
 # LLVM
 ##########
-if("${LLVM_LIBRARY_DIR}" STREQUAL "")
+if (NOT MLIR_DIR)
+  if(NOT LLVM_LIBRARY_DIR)
     if(WIN32)
       find_package(LLVM 13 REQUIRED COMPONENTS nvptx amdgpu)
 
@@ -127,7 +128,8 @@ libLLVMCoroutines.a
       libLLVMAnalysis.a
     )
   endif()
-include_directories(${LLVM_INCLUDE_DIRS})
+  set (MLIR_DIR ${LLVM_LIBRARY_DIR}/cmake/mlir)
+endif()
 
 # Python module
 if(TRITON_BUILD_PYTHON_MODULE)
@@ -161,10 +163,10 @@ endif()
 
 
 # MLIR
-find_package(MLIR REQUIRED CONFIG PATHS ${LLVM_LIBRARY_DIR}/cmake/mlir)
+find_package(MLIR REQUIRED CONFIG PATHS ${MLIR_DIR})
 
-list(APPEND CMAKE_MODULE_PATH ${LLVM_LIBRARY_DIR}/cmake/llvm)
-list(APPEND CMAKE_MODULE_PATH ${LLVM_LIBRARY_DIR}/cmake/mlir)
+list(APPEND CMAKE_MODULE_PATH "${MLIR_CMAKE_DIR}")
+list(APPEND CMAKE_MODULE_PATH "${LLVM_CMAKE_DIR}")
 
 include(TableGen) # required by AddMLIR
 include(AddLLVM)
@@ -172,6 +174,7 @@ include(AddMLIR)
 # include(HandleLLVMOptions) # human-friendly error message
 
 include_directories(${MLIR_INCLUDE_DIRS})
+include_directories(${LLVM_INCLUDE_DIRS})
 include_directories(${PROJECT_SOURCE_DIR}/include)
 include_directories(${PROJECT_BINARY_DIR}/include) # Tablegen'd files
 # link_directories(${LLVM_LIBRARY_DIR})