Merge commit '9a38395c82bc2b86b5dbbc67ba5a1f5f2dde53e1'

.gitignore

@@ -6,7 +6,7 @@ python/build/
 python/triton.egg-info/
 python/triton/_C/libtriton.pyd
 python/triton/_C/libtriton.so
-python/triton/backends/cuda
+python/triton/backends/nvidia
 python/triton/backends/xpu
 
 # Python caches

.gitmodules

@@ -0,0 +1,3 @@
+[submodule "third_party/amd"]
+	path = third_party/amd
+	url = https://github.com/ptillet/triton.git

CMakeLists.txt

@@ -19,6 +19,8 @@ if(NOT WIN32)
   list(APPEND CMAKE_MODULE_PATH "${CMAKE_CURRENT_SOURCE_DIR}/cmake")
 endif()
 
+
+
 # Options
 option(TRITON_BUILD_TUTORIALS "Build C++ Triton tutorials" ON)
 option(TRITON_BUILD_PYTHON_MODULE "Build Python Triton bindings" OFF)
@@ -82,8 +84,41 @@ include(TableGen) # required by AddMLIR
 include(AddLLVM)
 include(AddMLIR)
 
+# Utilities
+function(add_triton_object name)
+  cmake_parse_arguments(ARG "" "" "DEPENDS;LINK_LIBS" ${ARGN})
+  add_library(${name} OBJECT)
+  target_sources(${name}
+    PRIVATE ${ARG_UNPARSED_ARGUMENTS}
+    INTERFACE $<TARGET_OBJECTS:${name}>
+  )
+
+
+  # add_library(${name} OBJECT ${ARG_UNPARSED_ARGUMENTS})
+  if(ARG_DEPENDS)
+    add_dependencies(${name} ${ARG_DEPENDS})
+  endif()
+  if(ARG_LINK_LIBS)
+    target_link_libraries(${name} PUBLIC ${ARG_LINK_LIBS})
+  endif()
+endfunction(add_triton_object)
+
+set_property(GLOBAL PROPERTY TRITON_LIBS "")
+function(add_triton_library name)
+  set_property(GLOBAL APPEND PROPERTY TRITON_LIBS ${name})
+  add_triton_object(${name} ${ARGN})
+  llvm_update_compile_flags(${name})
+endfunction()
+
+set_property(GLOBAL PROPERTY TRITON_PLUGINS "")
+function(add_triton_plugin name)
+  set_property(GLOBAL APPEND PROPERTY TRITON_PLUGINS ${name})
+  add_triton_object(${name} ${ARGN})
+endfunction()
+
+
 # Disable warnings that show up in external code (gtest;pybind11)
-set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Werror -Wno-covered-switch-default")
+set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Werror -Wno-covered-switch-default -fvisibility=hidden")
 
 include_directories(".")
 include_directories(${MLIR_INCLUDE_DIRS})
@@ -99,9 +134,6 @@ add_subdirectory(lib)
 set(TRITON_SOURCE_DIR "${CMAKE_CURRENT_SOURCE_DIR}")
 set(TRITON_BINARY_DIR "${CMAKE_CURRENT_BINARY_DIR}")
 
-get_property(dialect_libs GLOBAL PROPERTY MLIR_DIALECT_LIBS)
-get_property(conversion_libs GLOBAL PROPERTY MLIR_CONVERSION_LIBS)
-
 # TODO: Figure out which target is sufficient to fix errors; triton is
 # apparently not enough. Currently set linking libstdc++fs for all targets
 # to support some old version GCC compilers like 8.3.0.
@@ -128,33 +160,26 @@ if(TRITON_BUILD_PYTHON_MODULE)
     add_link_options(${Python3_LINK_OPTIONS})
   endif()
 
-  set(TRITON_CODEGEN_BACKENDS "xpu")
   foreach(CODEGEN_BACKEND ${TRITON_CODEGEN_BACKENDS})
     add_subdirectory(third_party/${CODEGEN_BACKEND})
   endforeach()
 
+
+  get_property(triton_libs GLOBAL PROPERTY TRITON_LIBS)
+  get_property(triton_plugins GLOBAL PROPERTY TRITON_PLUGINS)
   set(TRITON_LIBRARIES
-    TritonIR
-    TritonAnalysis
-    TritonTransforms
-    TritonToTritonGPU
-    TritonGPUIR
-    TritonGPUTransforms
-    TritonLLVMIR
-    TritonNvidiaGPUIR
-    MLIRAMDGPUDialect
-    TritonAnalysis
-    NVGPUToLLVM
-    TritonNvidiaGPUTransforms
-    TritonGPUToLLVM
+    ${triton_libs}
+    ${triton_plugins}
     TritonSPIRV
+
+    # mlir
+    MLIRAMDGPUDialect
     MLIRNVVMDialect
     MLIRNVVMToLLVMIRTranslation
     MLIRGPUToNVVMTransforms
     MLIRGPUToGPURuntimeTransforms
     MLIRGPUTransforms
-
-    # optimizations
+    MLIRIR
     MLIRControlFlowToLLVM
     MLIRBytecodeWriter
     MLIRPass
@@ -166,7 +191,9 @@ if(TRITON_BUILD_PYTHON_MODULE)
     MLIRROCDLToLLVMIRTranslation
     MLIRGENXToLLVMIRTranslation
     MLIRGPUDialect
-    MLIRIR
+    MLIRSCFToControlFlow
+    MLIRIndexToLLVM
+    MLIRGPUToROCDLTransforms
 
     # LLVM
     LLVMPasses
@@ -180,12 +207,14 @@ if(TRITON_BUILD_PYTHON_MODULE)
   )
 
   # Define triton library
+  string(JOIN "," TRITON_BACKENDS_TUPLE ${TRITON_CODEGEN_BACKENDS})
+  set(TRITON_BACKENDS_TUPLE "(${TRITON_BACKENDS_TUPLE})")
+  add_compile_definitions(TRITON_BACKENDS_TUPLE=${TRITON_BACKENDS_TUPLE})
   add_library(triton SHARED ${PYTHON_SRC_PATH}/main.cc
                   ${PYTHON_SRC_PATH}/ir.cc
                   ${PYTHON_SRC_PATH}/passes.cc
                   ${PYTHON_SRC_PATH}/interpreter.cc
-                  ${PYTHON_SRC_PATH}/llvm.cc
-                  ${CMAKE_CURRENT_SOURCE_DIR}/third_party/xpu/triton_xpu.cc)
+                  ${PYTHON_SRC_PATH}/llvm.cc)
 
   # Link triton with its dependencies
   target_link_libraries(triton PUBLIC ${TRITON_LIBRARIES})
@@ -195,7 +224,6 @@ if(TRITON_BUILD_PYTHON_MODULE)
     target_link_libraries(triton PRIVATE z)
   endif()
   target_link_options(triton PRIVATE ${LLVM_LDFLAGS} ${GenISAIntrinsics_LDFLAGS})
-  set_target_properties(triton PROPERTIES INTERFACE_LINK_LIBRARIES "")
 endif()
 
 if(UNIX AND NOT APPLE)

bin/CMakeLists.txt

@@ -12,6 +12,7 @@ target_link_libraries(triton-opt PRIVATE
   TritonNvidiaGPUTransforms
   ${dialect_libs}
   ${conversion_libs}
+  ${triton_libs}
   # tests
   TritonTestAnalysis
   # MLIR core
@@ -33,6 +34,7 @@ target_link_libraries(triton-reduce PRIVATE
   TritonNvidiaGPUTransforms
   ${dialect_libs}
   ${conversion_libs}
+  ${triton_libs}
   # tests
   TritonTestAnalysis
   # MLIR core

include/triton/Target/PTX/TmaMetadata.h

@@ -24,7 +24,7 @@
 #ifndef TRITON_TARGET_PTX_TMAMETADATA_H
 #define TRITON_TARGET_PTX_TMAMETADATA_H
 
-#include "third_party/cuda/backend/include/cuda.h"
+#include "third_party/nvidia/backend/include/cuda.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/Format.h"

lib/Analysis/CMakeLists.txt

@@ -1,4 +1,4 @@
-add_mlir_library(TritonAnalysis
+add_triton_library(TritonAnalysis
   AxisInfo.cpp
   Allocation.cpp
   Membar.cpp
@@ -10,7 +10,6 @@ add_mlir_library(TritonAnalysis
   TritonGPUAttrDefsIncGen
 
   LINK_LIBS PUBLIC
-  ASMBuilder
   MLIRAnalysis
   MLIRLLVMDialect
   TritonIR

lib/Conversion/NVGPUToLLVM/CMakeLists.txt

@@ -1,16 +1,9 @@
-add_mlir_conversion_library(NVGPUToLLVM
+add_triton_library(NVGPUToLLVM
     NVGPUToLLVMPass.cpp
 
-    ADDITIONAL_HEADER_DIRS
-    ${PROJECT_SOURCE_DIR}/include/triton/Conversion/NVGPUToLLVM
-    ${PROJECT_BINARY_DIR}/include/triton/Conversion/NVGPUToLLVM
-
     DEPENDS
     NVGPUConversionPassIncGen
 
-    LINK_COMPONENTS
-    Core
-
     LINK_LIBS PUBLIC
     MLIRIR
     MLIRPass

lib/Conversion/TritonGPUToLLVM/CMakeLists.txt

@@ -1,10 +1,4 @@
-# Separate out PTX/GCN builders to avoid cyclic dependencies as TritonAnalysis
-# depends on it.
-set(LLVM_OPTIONAL_SOURCES
-    PTXAsmFormat.cpp
-  )
-
-add_mlir_conversion_library(TritonGPUToLLVM
+add_triton_library(TritonGPUToLLVM
     ConvertLayoutOpToLLVM/SharedToDotOperandDPAS.cpp
     ConvertLayoutOpToLLVM/SharedToDotOperandFMA.cpp
     ConvertLayoutOpToLLVM/SharedToDotOperandMMAv1.cpp
@@ -16,8 +10,8 @@ add_mlir_conversion_library(TritonGPUToLLVM
     DotOpToLLVM/MMAv2.cpp
     DotOpToLLVM/WGMMA.cpp
     DotOpToLLVM.cpp
-    ElementwiseOpToLLVM.cpp
     HistogramOpToLLVM.cpp
+    ElementwiseOpToLLVM.cpp
     LoadStoreOpToLLVM.cpp
     BarrierOpToLLVM.cpp
     TritonGPUToLLVM.cpp
@@ -30,19 +24,12 @@ add_mlir_conversion_library(TritonGPUToLLVM
     TensorPtrOpsToLLVM.cpp
     ClusterOpsToLLVM.cpp
     RegReallocOpToLLVM.cpp
-
-    ADDITIONAL_HEADER_DIRS
-    ${PROJECT_SOURCE_DIR}/include/triton/Conversion/TritonGPUToLLVM
-    ${PROJECT_BINARY_DIR}/include/triton/Conversion/TritonGPUToLLVM
+    PTXAsmFormat.cpp
 
     DEPENDS
     TritonGPUConversionPassIncGen
 
-    LINK_COMPONENTS
-    Core
-
     LINK_LIBS PUBLIC
-    ASMBuilder
     MLIRIR
     MLIRPass
     MLIRGENXDialect
@@ -58,14 +45,3 @@ add_mlir_conversion_library(TritonGPUToLLVM
     TritonNvidiaGPUTransforms
     NVGPUIR
 )
-
-add_mlir_library(ASMBuilder
-  PTXAsmFormat.cpp
-
-  DEPENDS
-  TritonTableGen
-
-  LINK_LIBS PUBLIC
-  MLIRAnalysis
-  MLIRLLVMDialect
-)

lib/Conversion/TritonToTritonGPU/CMakeLists.txt

@@ -1,16 +1,9 @@
-add_mlir_conversion_library(TritonToTritonGPU
+add_triton_library(TritonToTritonGPU
     TritonToTritonGPUPass.cpp
 
-    ADDITIONAL_HEADER_DIRS
-    ${PROJECT_SOURCE_DIR}/include/triton/Conversion/TritonToTritonGPU
-    ${PROJECT_BINARY_DIR}/include/triton/Conversion/TritonToTritonGPU
-
     DEPENDS
     TritonConversionPassIncGen
 
-    LINK_COMPONENTS
-    Core
-
     LINK_LIBS PUBLIC
     MLIRIR
     MLIRPass

lib/Dialect/NVGPU/IR/CMakeLists.txt

@@ -1,4 +1,4 @@
-add_mlir_dialect_library(NVGPUIR
+add_triton_library(NVGPUIR
   Dialect.cpp
 
   DEPENDS

lib/Dialect/Triton/IR/CMakeLists.txt

@@ -1,4 +1,4 @@
-add_mlir_dialect_library(TritonIR
+add_triton_library(TritonIR
   Dialect.cpp
   Ops.cpp
   Types.cpp

lib/Dialect/Triton/Transforms/CMakeLists.txt

@@ -2,7 +2,7 @@ set(LLVM_TARGET_DEFINITIONS Combine.td)
 mlir_tablegen(TritonCombine.inc -gen-rewriters)
 add_public_tablegen_target(TritonCombineIncGen)
 
-add_mlir_dialect_library(TritonTransforms
+add_triton_library(TritonTransforms
   Combine.cpp
   ReorderBroadcast.cpp
   RewriteTensorPointer.cpp

lib/Dialect/TritonGPU/IR/CMakeLists.txt

@@ -1,4 +1,4 @@
-add_mlir_dialect_library(TritonGPUIR
+add_triton_library(TritonGPUIR
   Dialect.cpp
   Traits.cpp
   Types.cpp

lib/Dialect/TritonGPU/Transforms/CMakeLists.txt

@@ -1,4 +1,4 @@
-add_mlir_dialect_library(TritonGPUTransforms
+add_triton_library(TritonGPUTransforms
   AccelerateMatmul.cpp
   Coalesce.cpp
   DecomposeConversions.cpp

lib/Dialect/TritonNvidiaGPU/IR/CMakeLists.txt

@@ -1,4 +1,4 @@
-add_mlir_dialect_library(TritonNvidiaGPUIR
+add_triton_library(TritonNvidiaGPUIR
   Dialect.cpp
   Ops.cpp
   Traits.cpp

lib/Dialect/TritonNvidiaGPU/Transforms/CMakeLists.txt

@@ -1,4 +1,4 @@
-add_mlir_dialect_library(TritonNvidiaGPUTransforms
+add_triton_library(TritonNvidiaGPUTransforms
   MaterializeLoadStore.cpp
   PlanCTA.cpp
   WSDecomposing.cpp

lib/Target/LLVMIR/CMakeLists.txt

@@ -1,10 +1,7 @@
-add_mlir_translation_library(TritonLLVMIR
+add_triton_library(TritonLLVMIR
         LLVMDIScope.cpp
         LLVMIRBreakPhiStruct.cpp
 
-        LINK_COMPONENTS
-        Core
-
         DEPENDS
         LLVMIRIncGen