Address review comments

triton-lang · Sep 20, 2024 · c1d1a3b · c1d1a3b
1 parent 29e8324
commit c1d1a3b
Show file tree

Hide file tree

Showing 27 changed files with 73 additions and 204 deletions.
diff --git a/third_party/amd/CMakeLists.txt b/third_party/amd/CMakeLists.txt
@@ -3,7 +3,7 @@ include_directories(${CMAKE_CURRENT_BINARY_DIR}/include)
 add_subdirectory(include)
 add_subdirectory(lib)
 if(TRITON_BUILD_PYTHON_MODULE)
-  add_triton_plugin(TritonAMD ${CMAKE_CURRENT_SOURCE_DIR}/python/triton_amd.cc LINK_LIBS TritonAMDGPUToLLVM TritonAMDGPUTransforms AMDGPUToLLVM)
+  add_triton_plugin(TritonAMD ${CMAKE_CURRENT_SOURCE_DIR}/python/triton_amd.cc LINK_LIBS TritonAMDGPUToLLVM TritonAMDGPUTransforms TritonAMDGPUDialectToLLVM)
 endif()
 if(TRITON_BUILD_UT)
   add_subdirectory(unittest)

diff --git a/third_party/amd/backend/compiler.py b/third_party/amd/backend/compiler.py
@@ -212,7 +212,6 @@ def make_llir(src, metadata, options):
         ##    For now it is used as a controller for developers only.
         __HIP_FTZ = True
         amd.passes.ttgpuir.add_to_llvmir(pm, options.arch, __HIP_FTZ)
-        amd.passes.ttgpuir.add_amdgpu_to_llvm(pm)
         passes.common.add_canonicalizer(pm)
         passes.common.add_cse(pm)
 

diff --git a/third_party/amd/include/AMDGPUToLLVM/AMDGPUToLLVMPass.h b/third_party/amd/include/AMDGPUToLLVM/AMDGPUToLLVMPass.h
diff --git a/third_party/amd/include/AMDGPUToLLVM/CMakeLists.txt b/third_party/amd/include/AMDGPUToLLVM/CMakeLists.txt
diff --git a/third_party/amd/include/AMDGPUToLLVM/Passes.h b/third_party/amd/include/AMDGPUToLLVM/Passes.h
diff --git a/third_party/amd/include/AMDGPUToLLVM/Passes.td b/third_party/amd/include/AMDGPUToLLVM/Passes.td
diff --git a/third_party/amd/include/CMakeLists.txt b/third_party/amd/include/CMakeLists.txt
@@ -1,4 +1,3 @@
-add_subdirectory(AMDGPUToLLVM)
+add_subdirectory(Dialect)
 add_subdirectory(TritonAMDGPUToLLVM)
 add_subdirectory(TritonAMDGPUTransforms)
-add_subdirectory(Dialect)
diff --git a/third_party/amd/include/Dialect/AMDGPU/IR/CMakeLists.txt b/third_party/amd/include/Dialect/AMDGPU/IR/CMakeLists.txt
diff --git a/third_party/amd/include/Dialect/CMakeLists.txt b/third_party/amd/include/Dialect/CMakeLists.txt
@@ -1 +1 @@
-add_subdirectory(AMDGPU)
+add_subdirectory(TritonAMDGPU)
diff --git a/...amd/include/Dialect/AMDGPU/CMakeLists.txt → ...clude/Dialect/TritonAMDGPU/CMakeLists.txt b/...amd/include/Dialect/AMDGPU/CMakeLists.txt → ...clude/Dialect/TritonAMDGPU/CMakeLists.txt
diff --git a/third_party/amd/include/Dialect/TritonAMDGPU/IR/CMakeLists.txt b/third_party/amd/include/Dialect/TritonAMDGPU/IR/CMakeLists.txt
@@ -0,0 +1,16 @@
+set(MLIR_BINARY_DIR ${CMAKE_BINARY_DIR})
+
+set(LLVM_TARGET_DEFINITIONS TritonAMDGPUOps.td)
+mlir_tablegen(Dialect.h.inc -gen-dialect-decls -dialect=amdgpu)
+mlir_tablegen(Dialect.cpp.inc -gen-dialect-defs -dialect=amdgpu)
+mlir_tablegen(OpsConversions.inc -gen-llvmir-conversions)
+mlir_tablegen(Ops.h.inc -gen-op-decls)
+mlir_tablegen(Ops.cpp.inc -gen-op-defs)
+add_mlir_doc(TritonAMDGPUDialect TritonAMDGPUDialect dialects/ -gen-dialect-doc)
+add_mlir_doc(TritonAMDGPUOps TritonAMDGPUOps dialects/ -gen-op-doc)
+add_public_tablegen_target(TritonAMDGPUTableGen)
+
+set(LLVM_TARGET_DEFINITIONS TritonAMDGPUAttrDefs.td)
+mlir_tablegen(TritonAMDGPUAttrDefs.h.inc -gen-attrdef-decls)
+mlir_tablegen(TritonAMDGPUAttrDefs.cpp.inc -gen-attrdef-defs)
+add_public_tablegen_target(TritonAMDGPUAttrDefsIncGen)
diff --git a/...y/amd/include/Dialect/AMDGPU/IR/Dialect.h → ...include/Dialect/TritonAMDGPU/IR/Dialect.h b/...y/amd/include/Dialect/AMDGPU/IR/Dialect.h → ...include/Dialect/TritonAMDGPU/IR/Dialect.h
@@ -24,17 +24,20 @@
 #ifndef TRITON_DIALECT_AMDGPU_IR_DIALECT_H_
 #define TRITON_DIALECT_AMDGPU_IR_DIALECT_H_
 
-#include "amd/include/Dialect/AMDGPU/IR/Dialect.h.inc"
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/Dialect.h"
+#include "mlir/IR/PatternMatch.h"
+// clang-format off
+#include "amd/include/Dialect/TritonAMDGPU/IR/Dialect.h.inc"
+// clang-format on
 
 #define GET_ATTRDEF_CLASSES
-#include "amd/include/Dialect/AMDGPU/IR/AMDGPUAttrDefs.h.inc"
+#include "amd/include/Dialect/TritonAMDGPU/IR/TritonAMDGPUAttrDefs.h.inc"
 
 #define GET_OP_CLASSES
-#include "amd/include/Dialect/AMDGPU/IR/Ops.h.inc"
+#include "amd/include/Dialect/TritonAMDGPU/IR/Ops.h.inc"
 
 namespace mlir {
 namespace triton {

diff --git a/...clude/Dialect/AMDGPU/IR/AMDGPUAttrDefs.td → ...t/TritonAMDGPU/IR/TritonAMDGPUAttrDefs.td b/...clude/Dialect/AMDGPU/IR/AMDGPUAttrDefs.td → ...t/TritonAMDGPU/IR/TritonAMDGPUAttrDefs.td
@@ -21,15 +21,15 @@
  * SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
  */
 
-#ifndef AMDGPU_ATTRDEFS
-#define AMDGPU_ATTRDEFS
+#ifndef TRITON_AMDGPU_ATTRDEFS
+#define TRITON_AMDGPU_ATTRDEFS
 
 include "mlir/IR/AttrTypeBase.td"
-include "AMDGPUDialect.td"
+include "TritonAMDGPUDialect.td"
 
-class AMDGPU_Attr<string name, list<Trait> traits = [],
+class TritonAMDGPU_Attr<string name, list<Trait> traits = [],
                      string baseCppClass = "::mlir::Attribute">
-  : AttrDef<AMDGPU_Dialect, name, traits, baseCppClass> {
+  : AttrDef<TritonAMDGPU_Dialect, name, traits, baseCppClass> {
 }
 
 #endif
diff --git a/...nclude/Dialect/AMDGPU/IR/AMDGPUDialect.td → ...ct/TritonAMDGPU/IR/TritonAMDGPUDialect.td b/...nclude/Dialect/AMDGPU/IR/AMDGPUDialect.td → ...ct/TritonAMDGPU/IR/TritonAMDGPUDialect.td
@@ -21,17 +21,17 @@
  * SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
  */
 
-#ifndef AMDGPU_DIALECT
-#define AMDGPU_DIALECT
+#ifndef TRITON_AMDGPU_DIALECT
+#define TRITON_AMDGPU_DIALECT
 
 include "mlir/IR/OpBase.td"
 
-def AMDGPU_Dialect : Dialect {
+def TritonAMDGPU_Dialect : Dialect {
   let name = "amdgpu";
   let cppNamespace = "::mlir::triton::amdgpu";
 
   let description = [{
-    AMDGPU Dialect.
+    TritonAMDGPU Dialect hosts AMD specific ops at TritonGPU abstraction level.
   }];
 
   let dependentDialects = [];

diff --git a/...md/include/Dialect/AMDGPU/IR/AMDGPUOps.td → ...ialect/TritonAMDGPU/IR/TritonAMDGPUOps.td b/...md/include/Dialect/AMDGPU/IR/AMDGPUOps.td → ...ialect/TritonAMDGPU/IR/TritonAMDGPUOps.td
@@ -22,14 +22,14 @@
  */
 
 
-#ifndef AMDGPU_OPS
-#define AMDGPU_OPS
+#ifndef TRITON_AMDGPU_OPS
+#define TRITON_AMDGPU_OPS
 
 include "mlir/IR/OpBase.td"
 include "mlir/IR/EnumAttr.td"
 include "mlir/Dialect/LLVMIR/LLVMOpBase.td"
 include "mlir/Interfaces/InferTypeOpInterface.td"
-include "AMDGPUDialect.td"
-include "AMDGPUAttrDefs.td"
+include "TritonAMDGPUDialect.td"
+include "TritonAMDGPUAttrDefs.td"
 
 #endif
diff --git a/third_party/amd/lib/AMDGPUToLLVM/AMDGPUToLLVMPass.cpp b/third_party/amd/lib/AMDGPUToLLVM/AMDGPUToLLVMPass.cpp
diff --git a/third_party/amd/lib/AMDGPUToLLVM/CMakeLists.txt b/third_party/amd/lib/AMDGPUToLLVM/CMakeLists.txt
diff --git a/third_party/amd/lib/CMakeLists.txt b/third_party/amd/lib/CMakeLists.txt
@@ -1,4 +1,4 @@
-add_subdirectory(AMDGPUToLLVM)
 add_subdirectory(Dialect)
 add_subdirectory(TritonAMDGPUToLLVM)
+add_subdirectory(TritonAMDGPUDialectToLLVM)
 add_subdirectory(TritonAMDGPUTransforms)
diff --git a/third_party/amd/lib/Dialect/CMakeLists.txt b/third_party/amd/lib/Dialect/CMakeLists.txt
@@ -1 +1 @@
-add_subdirectory(AMDGPU)
+add_subdirectory(TritonAMDGPU)
diff --git a/...rty/amd/lib/Dialect/AMDGPU/CMakeLists.txt → ...d/lib/Dialect/TritonAMDGPU/CMakeLists.txt b/...rty/amd/lib/Dialect/AMDGPU/CMakeLists.txt → ...d/lib/Dialect/TritonAMDGPU/CMakeLists.txt
diff --git a/.../amd/lib/Dialect/AMDGPU/IR/CMakeLists.txt → ...ib/Dialect/TritonAMDGPU/IR/CMakeLists.txt b/.../amd/lib/Dialect/AMDGPU/IR/CMakeLists.txt → ...ib/Dialect/TritonAMDGPU/IR/CMakeLists.txt
@@ -1,9 +1,9 @@
-add_triton_library(AMDGPUIR
+add_triton_library(TritonAMDGPUIR
   Dialect.cpp
 
   DEPENDS
-  AMDGPUTableGen
-  AMDGPUAttrDefsIncGen
+  TritonAMDGPUTableGen
+  TritonAMDGPUAttrDefsIncGen
 
   LINK_LIBS PUBLIC
   MLIRLLVMDialect

diff --git a/...rty/amd/lib/Dialect/AMDGPU/IR/Dialect.cpp → ...d/lib/Dialect/TritonAMDGPU/IR/Dialect.cpp b/...rty/amd/lib/Dialect/AMDGPU/IR/Dialect.cpp → ...d/lib/Dialect/TritonAMDGPU/IR/Dialect.cpp
@@ -25,24 +25,24 @@
 #include "mlir/IR/OpImplementation.h"
 
 // clang-format off
-#include "Dialect/AMDGPU/IR/Dialect.h"
-#include "Dialect/AMDGPU/IR/Dialect.cpp.inc"
+#include "Dialect/TritonAMDGPU/IR/Dialect.h"
+#include "Dialect/TritonAMDGPU/IR/Dialect.cpp.inc"
 // clang-format on
 
 using namespace mlir;
 using namespace mlir::triton::amdgpu;
 
-void mlir::triton::amdgpu::AMDGPUDialect::initialize() {
+void mlir::triton::amdgpu::TritonAMDGPUDialect::initialize() {
   addAttributes<
 #define GET_ATTRDEF_LIST
-#include "Dialect/AMDGPU/IR/AMDGPUAttrDefs.cpp.inc"
+#include "Dialect/TritonAMDGPU/IR/TritonAMDGPUAttrDefs.cpp.inc"
       >();
 
   addOperations<
 #define GET_OP_LIST
-#include "Dialect/AMDGPU/IR/Ops.cpp.inc"
+#include "Dialect/TritonAMDGPU/IR/Ops.cpp.inc"
       >();
 }
 
 #define GET_OP_CLASSES
-#include "Dialect/AMDGPU/IR/Ops.cpp.inc"
+#include "Dialect/TritonAMDGPU/IR/Ops.cpp.inc"
diff --git a/third_party/amd/lib/TritonAMDGPUDialectToLLVM/CMakeLists.txt b/third_party/amd/lib/TritonAMDGPUDialectToLLVM/CMakeLists.txt
@@ -0,0 +1,6 @@
+add_triton_library(TritonAMDGPUDialectToLLVM
+    TritonAMDGPUToLLVMPatterns.cpp
+
+    DEPENDS
+    TritonAMDGPUIR
+)
diff --git a/third_party/amd/lib/TritonAMDGPUDialectToLLVM/TritonAMDGPUToLLVMPatterns.cpp b/third_party/amd/lib/TritonAMDGPUDialectToLLVM/TritonAMDGPUToLLVMPatterns.cpp
@@ -0,0 +1,9 @@
+#include "triton/Conversion/TritonGPUToLLVM/PatternTritonGPUOpToLLVM.h"
+
+namespace mlir::triton::AMD {
+void populateTritonAMDGPUToLLVMPatterns(LLVMTypeConverter &typeConverter,
+                                        RewritePatternSet &patterns,
+                                        PatternBenefit benefit) {
+  // TODO: Insert TrtionAMDGPU dialect patterns.
+}
+} // namespace mlir::triton::AMD
diff --git a/third_party/amd/lib/TritonAMDGPUToLLVM/PatternTritonGPUOpToLLVM.h b/third_party/amd/lib/TritonAMDGPUToLLVM/PatternTritonGPUOpToLLVM.h
@@ -30,6 +30,9 @@ void populateLoadStoreOpToLLVMPatterns(LLVMTypeConverter &typeConverter,
 void populateSPMDOpToLLVMPattern(LLVMTypeConverter &typeConverter,
                                  RewritePatternSet &patterns,
                                  PatternBenefit benefit);
+void populateTritonAMDGPUToLLVMPatterns(LLVMTypeConverter &typeConverter,
+                                        RewritePatternSet &patterns,
+                                        PatternBenefit benefit);
 
 } // namespace mlir::triton::AMD
 

diff --git a/third_party/amd/lib/TritonAMDGPUToLLVM/TritonGPUToLLVM.cpp b/third_party/amd/lib/TritonAMDGPUToLLVM/TritonGPUToLLVM.cpp
@@ -200,6 +200,10 @@ struct ConvertTritonAMDGPUToLLVM
     mlir::triton::populateSPMDOpToLLVMPattern(typeConverter, patterns,
                                               targetInfo, commonBenefit);
     AMD::populateSPMDOpToLLVMPattern(typeConverter, patterns, AMDBenefit);
+
+    mlir::triton::AMD::populateTritonAMDGPUToLLVMPatterns(typeConverter,
+                                                          patterns, AMDBenefit);
+
     // TODO(thomas): this should probably be done in a separate step to not
     // interfere with our own lowering of arith ops. Add arith/math's patterns
     // to help convert scalar expression to LLVM.